/*************************************************************************
 * 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
 *************************************************************************/

#include "common.h"
#include "p2p_resiliency.h"

char ncclIbIfName[MAX_IF_NAME_SIZE + 1];
union ncclSocketAddress ncclIbIfAddr;

int ncclNMergedIbDevs = -1;
int ncclNIbDevs = -1;
struct ncclIbMergedDev ncclIbMergedDevs[MAX_IB_VDEVS];
struct ncclIbDev ncclIbDevs[MAX_IB_DEVS];
int ncclIbRelaxedOrderingEnabled = 0;
uint64_t ncclIbSpeedChangeCounter = 0;

ncclProfilerCallback_t ncclProfilerFunction;

NCCL_PARAM(IbSplitDataOnQps, "IB_SPLIT_DATA_ON_QPS", 0);
NCCL_PARAM(IbPrepostReceiveWorkRequests, "IB_PREPOST_RECEIVE_WORK_REQUESTS", -2);
NCCL_PARAM(IbAsyncEvents, "IB_RETURN_ASYNC_EVENTS", 1);
NCCL_PARAM(IbEventBasedLb, "IB_EVENT_BASED_LB", 0);
NCCL_PARAM(IbEventBasedLbRemote, "IB_EVENT_BASED_LB_REMOTE", 1);

extern int ncclParamIbReceiverSideMatchingScheme();
extern int ncclParamIbOooRq();
extern int ncclParamIbResiliencyPortFailover();

ncclResult_t ncclIbStatsCheckFatalCount(struct ncclIbStats* stat, const char* funcName) {
  if (ncclParamIbAsyncEvents() && COMPILER_ATOMIC_LOAD(&stat->fatalErrorCount, std::memory_order_relaxed)) {
    WARN("communicator encountered a fatal error (detected in %s)", funcName);
    return ncclSystemError;
  }
  return ncclSuccess;
}

struct ncclIbNetCommDevBase* ncclIbGetNetCommDevBase(ncclIbNetCommBase* base, int devIndex) {
  if (base->isSend) {
    struct ncclIbSendComm* sComm = (struct ncclIbSendComm*)base;
    return &sComm->devs[devIndex].base;
  } else {
    struct ncclIbRecvComm* rComm = (struct ncclIbRecvComm*)base;
    return &rComm->devs[devIndex].base;
  }
}

ncclResult_t ncclIbBaseCommInit(struct ncclIbNetCommBase* baseComm, bool isSend) {
  for (int i = 0; i < NCCL_IB_MAX_QPS; i++) {
    baseComm->qps[i].devIndex = -1;
    baseComm->qps[i].remDevIdx = -1;
    baseComm->activeQps[i] = &baseComm->qps[i];
    baseComm->qps[i].eceSupported = 0;
    baseComm->qps[i].ece = {0};
    memset(&baseComm->qps[i].initAttr, 0, sizeof(baseComm->qps[i].initAttr));
    memset(&baseComm->qps[i].rtrAttr, 0, sizeof(baseComm->qps[i].rtrAttr));
    memset(&baseComm->qps[i].rtsAttr, 0, sizeof(baseComm->qps[i].rtsAttr));
  }
  baseComm->nqps = -1;
  baseComm->splitDataOnQps = ncclParamIbSplitDataOnQps();
  baseComm->nDataQps = -1;
  baseComm->isSend = isSend;
  baseComm->ready = 0;
  baseComm->speedChangeCounter = 0;
  baseComm->totalSpeed = 0;
  memset(baseComm->weights, 0, sizeof(baseComm->weights));
  memset(baseComm->devSpeeds, 0, sizeof(baseComm->devSpeeds));

  NCCLCHECK(ncclIbResiliencyInit(baseComm, &baseComm->resiliency));
  baseComm->recvMatchingScheme =
    ncclParamIbReceiverSideMatchingScheme() == -2 ? BY_INDEX : ncclParamIbReceiverSideMatchingScheme();

  if (ncclParamIbOooRq() || (ncclParamIbResiliencyPortFailover() == 1)) {
    baseComm->recvMatchingScheme = BY_ID;
    if (ncclParamIbReceiverSideMatchingScheme() == BY_INDEX) {
      INFO(NCCL_NET, "NET/IB: %s: Overriding matching scheme to ID-based (%d)", __func__, BY_ID);
    }
  }

  return ncclSuccess;
}

ncclResult_t ncclIbRecvCommInit(struct ncclIbRecvComm* recvComm) {
  NCCLCHECK(ncclIbBaseCommInit(&recvComm->base, false));
  recvComm->ibRecvWorkRequest = {.wr_id = NCCL_IB_RECV_WR_ID_DUMMY, .next = NULL, .sg_list = NULL, .num_sge = 0};

  recvComm->prepostReceiveWorkRequests =
    (ncclParamIbPrepostReceiveWorkRequests() == -2) ? false : ncclParamIbPrepostReceiveWorkRequests();

  if (recvComm->base.resiliency) {
    if (ncclParamIbPrepostReceiveWorkRequests() == 0) {
      INFO(NCCL_NET, "NET/IB: %s: Overriding pre-posting to true (1).", __func__);
    }
    recvComm->prepostReceiveWorkRequests = true;
  }
  if (ncclParamIbOooRq()) {
    if (ncclParamIbPrepostReceiveWorkRequests() == 0) {
      INFO(NCCL_NET, "NET/IB: %s: OOO RQ is enabled, Overriding pre-posting to true (1).", __func__);
    }
    recvComm->prepostReceiveWorkRequests = true;
  }

  INFO(NCCL_NET, "NET/IB: %s: Receive work requests will be %s", __func__,
       recvComm->prepostReceiveWorkRequests ? "pre-posted" : "posted on-demand");
  return ncclSuccess;
}

ncclResult_t ncclIbSendCommInit(struct ncclIbSendComm* sendComm) {
  NCCLCHECK(ncclIbBaseCommInit(&sendComm->base, true));
  return ncclSuccess;
}

static ncclResult_t ncclIbEventGidChange(struct ncclIbDev* dev) {
  INFO(NCCL_NET, "NET/IB: %s: GID table changed on %s:%d", __func__, dev->devName, dev->portNum);
  if (dev->gidInfo.link_layer != IBV_LINK_LAYER_ETHERNET) {
    INFO(NCCL_NET, "NET/IB : %s:%d link is not Ethernet; ignoring GID change", dev->devName, dev->portNum);
    return ncclSuccess;
  }
  struct ibv_port_attr portAttr;
  ncclResult_t res = wrap_ibv_query_port(dev->context, dev->portNum, &portAttr);
  if (res != ncclSuccess) {
    WARN("NET/IB : %s:%d query_port failed during GID change (res=%d)", dev->devName, dev->portNum, (int)res);
    return res;
  }

  std::lock_guard<std::mutex> lock(dev->mutex);
  int oldIdx = dev->gidInfo.localGidIndex;
  union ibv_gid oldGid = dev->gidInfo.localGid;
  res = ncclIbGidInfoQuery(dev->context, dev->portNum, &portAttr, &dev->gidInfo);
  if (res != ncclSuccess) {
    WARN("NET/IB : %s:%d GID info query failed (%d) during GID change", dev->devName, dev->portNum, (int)res);
    return res;
  }
  dev->portAttr = portAttr; // publish only once everything succeeded
  char oldGidStr[INET6_ADDRSTRLEN] = "";
  char newGidStr[INET6_ADDRSTRLEN] = "";
  ibvGetGidStr(&oldGid, oldGidStr, sizeof(oldGidStr));
  ibvGetGidStr(&dev->gidInfo.localGid, newGidStr, sizeof(newGidStr));
  INFO(NCCL_NET, "NET/IB : %s:%d GID refreshed: idx %d -> %d, gid %s -> %s", dev->devName, dev->portNum, oldIdx,
       dev->gidInfo.localGidIndex, oldGidStr, newGidStr);
  return ncclSuccess;
}

static void ncclIbUpdateDeviceSpeed(struct ncclIbDev* dev) {
  uint64_t oldSpeed = COMPILER_ATOMIC_LOAD(&dev->currSpeed, std::memory_order_relaxed);
  uint64_t newSpeed = 0;
  if (wrap_ibv_query_port_speed(dev->context, dev->portNum, &newSpeed) != ncclSuccess) return;

  // In cases of port failover, the speed change event is ignored and will be considered as a port fail event
  if (newSpeed == 0) return;
  // ibv_query_port_speed returns speed in granularity of 100 Mbps
  newSpeed *= 100;

  if (newSpeed != oldSpeed) {
    INFO(NCCL_NET, "NET/IB : %s:%d speed change detected: %lu -> %lu Mbps", dev->devName, dev->portNum,
         (unsigned long)oldSpeed, (unsigned long)newSpeed);
    COMPILER_ATOMIC_STORE(&dev->currSpeed, newSpeed, std::memory_order_relaxed);
    COMPILER_ATOMIC_FETCH_ADD(&ncclIbSpeedChangeCounter, (uint64_t)1, std::memory_order_release);
  }
}

std::thread ncclIbAsyncThread;
void* ncclIbAsyncThreadMain(void* args) {
  struct ncclIbDev* dev = (struct ncclIbDev*)args;
  while (1) {
    struct ibv_async_event event;
    if (ncclSuccess != wrap_ibv_get_async_event(dev->context, &event)) break;
    char* str;
    struct ibv_cq* cq = event.element.cq;    // only valid if CQ error
    struct ibv_qp* qp = event.element.qp;    // only valid if QP error
    struct ibv_srq* srq = event.element.srq; // only valid if SRQ error
    if (ncclSuccess != wrap_ibv_event_type_str(&str, event.event_type)) break;
    switch (event.event_type) {
    case IBV_EVENT_DEVICE_FATAL:
      // the above is device fatal error
      WARN("NET/IB : %s:%d async fatal event: %s", dev->devName, dev->portNum, str);
      ncclIbDevFatalError(dev);
      break;
    case IBV_EVENT_CQ_ERR:
      WARN("NET/IB : %s:%d async fatal event on CQ (%p) handle=%u cqe=%d: %s", dev->devName, dev->portNum, cq,
           (unsigned)cq->handle, cq->cqe, str);
      ncclIbCqFatalError(cq);
      break;
    case IBV_EVENT_QP_FATAL:
    case IBV_EVENT_QP_REQ_ERR:
    case IBV_EVENT_QP_ACCESS_ERR:
      {
        struct ibv_qp_attr qpAttr;
        memset(&qpAttr, 0, sizeof(qpAttr));
        struct ibv_qp_init_attr qpInitAttr;
        memset(&qpInitAttr, 0, sizeof(qpInitAttr));
        (void)wrap_ibv_query_qp(qp, &qpAttr, IBV_QP_STATE | IBV_QP_SQ_PSN | IBV_QP_RQ_PSN, &qpInitAttr);
        WARN("NET/IB : %s:%d async fatal event on QP (%p) qpn=%u handle=%u send_cq=%u recv_cq=%u state=%d sq_psn=%u "
             "rq_psn=%u: %s",
             dev->devName, dev->portNum, qp, qp->qp_num, (unsigned)qp->handle,
             qp->send_cq ? (unsigned)qp->send_cq->handle : 0u, qp->recv_cq ? (unsigned)qp->recv_cq->handle : 0u,
             (int)qpAttr.qp_state, (unsigned)qpAttr.sq_psn, (unsigned)qpAttr.rq_psn, str);
        ncclIbQpFatalError(qp);
        break;
      }
    case IBV_EVENT_SRQ_ERR:
      // SRQ are not used in NCCL
      WARN("NET/IB : %s:%d async fatal event on SRQ, unused for now (%p): %s", dev->devName, dev->portNum, srq, str);
      break;
    case IBV_EVENT_GID_CHANGE:
      if (ncclIbEventGidChange(dev) != ncclSuccess) {
        WARN("NET/IB : %s:%d marking device with fatal error after GID-change event handler failed", dev->devName,
             dev->portNum);
        ncclIbDevFatalError(dev);
      }
      break;
    case IBV_EVENT_DEVICE_SPEED_CHANGE:
      ncclIbUpdateDeviceSpeed(dev);
      break;
    case IBV_EVENT_PATH_MIG_ERR:
    case IBV_EVENT_PORT_ERR:
    case IBV_EVENT_PATH_MIG:
    case IBV_EVENT_PORT_ACTIVE:
    case IBV_EVENT_SQ_DRAINED:
    case IBV_EVENT_LID_CHANGE:
    case IBV_EVENT_PKEY_CHANGE:
    case IBV_EVENT_SM_CHANGE:
    case IBV_EVENT_QP_LAST_WQE_REACHED:
    case IBV_EVENT_CLIENT_REREGISTER:
    case IBV_EVENT_SRQ_LIMIT_REACHED:
      // the above are non-fatal
      WARN("NET/IB : %s:%d Got non-fatal async event: %s(%d)", dev->devName, dev->portNum, str, event.event_type);
      break;
    case IBV_EVENT_COMM_EST:
      break;
    default:
      WARN("NET/IB : %s:%d unknown event type (%d)", dev->devName, dev->portNum, event.event_type);
      break;
    }
    // acknowledgment needs to happen last to avoid user-after-free
    if (ncclSuccess != wrap_ibv_ack_async_event(&event)) break;
  }
  return NULL;
}

ncclNet_t ncclNetIb = {
  "IB",
  ncclIbInit,
  ncclIbDevices,
  ncclIbGetProperties,
  ncclIbListen,
  ncclIbConnect,
  ncclIbAccept,
  ncclIbRegMr,
  ncclIbRegMrDmaBuf,
  ncclIbDeregMr,
  ncclIbIsend,
  ncclIbIrecv,
  ncclIbIflush,
  ncclIbTest,
  ncclIbCloseSend,
  ncclIbCloseRecv,
  ncclIbCloseListen,
  NULL /* getDeviceMr */,
  NULL /* irecvConsumed */,
  ncclIbMakeVDevice,
  ncclIbFinalize,
  ncclIbSetNetAttr,
};
