/*************************************************************************
 * SPDX-FileCopyrightText: Copyright (c) 2016-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * See LICENSE.txt for more license information
 *************************************************************************/

#ifndef NET_IB_CONNECT_H_
#define NET_IB_CONNECT_H_

#include "common.h"
#include "ibvwrap.h"

struct ncclIbQpCreateAttr {
  void* qpContext;
  enum ibv_qp_type type;
  bool oooRq;
  struct ibv_cq* cq;
  struct ibv_pd* pd;
  uint32_t maxRecvWorkRequest;
  uint32_t maxSendWorkRequest;
};

// Per-QP connection metatdata
struct ncclIbQpInfo {
  uint32_t qpn;

  // Fields needed for ece (enhanced connection establishment)
  struct ibv_ece ece;
  int ece_supported;

  // The index of the device on which the QP was created. Allows the sender and
  // receiver side to have asymmetric device configuration, meaning the sender
  // and receiver can use different number of devices.
  int devIndex;
};

struct ncclIbResiliencyInfo {
  // QPs used for probing of data transfers in case of QP/device failures.
  struct ncclIbQpInfo probingQpsInfo[NCCL_IB_MAX_DEVS_PER_NIC];
  // QPs used for recovery protocol after QP/device failures.
  struct ncclIbQpInfo portRecoveryQpsInfo[NCCL_IB_MAX_DEVS_PER_NIC];
};

// Structure used to hold information needed to establish the communication
// between the sender and receiver.
// The structure is populated during the connection establishment phase and
// populated by each side of the connection before being sent to the remote
// peer. The remote peer uses the information passed to it from its peer to
// create and initialize its local resources.
struct ncclIbConnectionMetadata {
  struct ncclIbQpInfo qpInfo[NCCL_IB_MAX_QPS];
  struct ncclIbResiliencyInfo resiliencyInfo;
  struct ncclIbDevInfo devs[NCCL_IB_MAX_DEVS_PER_NIC];
  char devName[MAX_MERGED_DEV_NAME];
  // An address for a registered memory to be accessed by the peer. The address
  // can be accessed using RDMA using the key specified in ncclIbDevInfo::rkey.
  // The sender side gets in this member, from the receiver, the address of the
  // memory to which the sender writes the sizes of the data transfers that
  // the sender sends.
  // The receiver side gets in this member, from the sender, the address of the
  // memory to which the receiver writes the CTS messages.
  uint64_t addr;
  uint64_t remSpeedBufAddr;
  uint32_t remSpeedBufRkey;
  int ndevs;
  int tc;
  int sl;
};

ncclResult_t ncclIbQpCreate(struct ncclIbQp* qp, struct ncclIbQpCreateAttr* createQpAttrs);
ncclResult_t ncclIbQpInit(struct ncclIbQp* qp);
ncclResult_t ncclIbQpRtr(struct ncclIbQp* qp);
ncclResult_t ncclIbQpRts(struct ncclIbQp* qp);
ncclResult_t ncclIbQpReset(struct ncclIbQp* qp);
ncclResult_t ncclIbQpError(struct ncclIbQp* qp);

ncclResult_t ncclIbPostReceiveWorkRequestsOnQp(struct ncclIbRecvComm* recvComm, ncclIbQp* dataQp);

#endif // NET_IB_CONNECT_H_
