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

#include "nccl.h"
#include "gin/gin_host.h"
#include "gin/gin_host_proxy.h"
#include "alloc.h"
#include "checks.h"
#include "gdrwrap.h"
#include "nccl_device/gin/proxy/gin_proxy_device_host_common.h"
#include "nccl_device/gin/gin_device_common.h"
#include "compiler.h"
#include "comm.h"
#include "proxy_gpucontext/gpucontext.h"

NCCL_PARAM(GinProxyQueueSize, "GIN_PROXY_QUEUE_SIZE", -1);
NCCL_PARAM(GinProxyPollBatch, "GIN_PROXY_POLL_BATCH", 32);
extern int64_t ncclParamIbDataDirect();
extern int64_t ncclParamDmaBufEnable();

struct ginProxyGfdState {
  ncclGinProxyOp_t op;
  uint16_t counterId;
  int done;
  void* request;
};

// a member might be on the GPU, if it has a *GdrHandle counterpart
struct ginProxyHostGpuCtx {
  int contextId;
  size_t queueSize;

  // size = nRanks * queueSize
  ncclGinProxyGfd_t* queues;
  void* cisGdrHandle;
  // Produced Indices, one per rank. Only accessed by the GPU side, here only for freeing
  uint32_t* pis;
  // Consumed Indices, one per rank
  uint32_t* cis;
  // to decrease the number of reads/writes to cis which might be on the GPU
  uint32_t* cisShadow;
  // Seen Indices one per rank
  uint32_t* sis;

  // same size as queues
  struct ginProxyGfdState* states;
  // same size as queues
  uint64_t* inlines;
  // inlines is registered as a memory region with the GIN plugin
  void* inlinesMhandle;
  void* inlinesGinHandle;

  uint32_t* lastIssuedGet; // per-rank index of most recent get
  uint32_t* lastVisibleGet; // per-rank index of last get for which the payload is guaranteed visible (via flush GFD)
};

struct ginProxyCtx {
  void* collComm;
  int nRanks;
  ncclNetDeviceHandle_t* devHandle;

  // GPU queues, if GDR on the GPU, else on the CPU
  // Queue size, must be a power of 2
  struct ginProxyHostGpuCtx* hostGpuCtx;

  void* countersGdrHandle;
  uint64_t* counters;
  uint64_t* countersDev;
  CUmemGenericAllocationHandle signalsCumemhandle;
  void* signalsMhandle;
  void* signalsGinHandle;
  uint64_t* signalsDev;
  uint64_t* signalOffsetsDev;
  bool hasError;
  int nContexts;
  int nCountersPerContext;
  int nSignalsPerContext;
  int pollBatch;
  void* rmaCtx; // from plugin
};

static ncclRma_t* rmaBackend;

static ncclResult_t getDmaBufFd(void* addr, size_t length, int* fd, bool forceNonDataDirect = false) {
  if (ncclParamDmaBufEnable() == 0) return ncclInvalidUsage;

#if CUDA_VERSION >= 11070
  static size_t hostPageSize = ncclOsGetPageSize();
  size_t alignedSize = length;
  ALIGN_SIZE(alignedSize, hostPageSize);

#if CUDA_VERSION >= 12080
  if (ncclParamIbDataDirect() && !forceNonDataDirect) {
    CUresult status =
      pfn_cuMemGetHandleForAddressRange((void*)fd, (CUdeviceptr)addr, alignedSize, CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD,
                                        CU_MEM_RANGE_FLAG_DMA_BUF_MAPPING_TYPE_PCIE);
    if (status == CUDA_SUCCESS) return ncclSuccess;
  }
#endif
  CUresult status = pfn_cuMemGetHandleForAddressRange((void*)fd, (CUdeviceptr)addr, alignedSize,
                                                      CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, 0);
  if (status == CUDA_SUCCESS) return ncclSuccess;
#endif

  return ncclInvalidUsage;
}

static ncclResult_t proxyGinPollCompletions(void* collComm, struct ginProxyCtx* ctx,
                                            struct ginProxyHostGpuCtx* hostGpuCtx) {
  for (int targetRank = 0; targetRank < ctx->nRanks; targetRank++) {
    // loop on all seen but unconsumed GFDs
    for (uint32_t i = hostGpuCtx->cisShadow[targetRank]; i < hostGpuCtx->sis[targetRank]; i++) {
      uint32_t idx = i & (hostGpuCtx->queueSize - 1);
      struct ginProxyGfdState* state = &hostGpuCtx->states[targetRank * hostGpuCtx->queueSize + idx];
      // no need to poll if already done
      if (!state->done) {
        ncclResult_t res = rmaBackend->test(collComm, state->request, &state->done);
        if (res != ncclSuccess) {
          ctx->hasError = true;
          WARN("Error on GFD test %d - stateIdx: %lu, request: %p", res, state - hostGpuCtx->states, state->request);
          return res;
        }
        if (state->done) {
          TRACE(NCCL_NET, "GFD completed - contextId: %d, stateIdx: %lu, request: %p", hostGpuCtx->contextId,
                state - hostGpuCtx->states, state->request);
          // update the counter specified in the GFD
          if (state->op & ncclGinProxyOpWithCounter) {
            // Atomic load and atomic store are used here to ensure that they are volatile.
            // The GPU kernel is not allowed to reset the counter whilst there are outstanding operations.
            // Thus the add does not need to be atomic.
            int contextId = hostGpuCtx->contextId;
            uint64_t* counterPtr = &ctx->counters[contextId * ctx->nCountersPerContext + state->counterId];
            uint64_t oldValue = COMPILER_ATOMIC_LOAD(counterPtr, std::memory_order_relaxed);
            COMPILER_ATOMIC_STORE(counterPtr, oldValue + 1, std::memory_order_relaxed);
            TRACE(NCCL_NET, "Updated counter %d to %ld for context %d", state->counterId, *counterPtr, contextId);
          }
        }
      }
      // allow holes in the CI space to get resolved
      if (state->done && i == hostGpuCtx->cisShadow[targetRank]) {
        // tell the GPU that we have consumed the GFD
        COMPILER_ATOMIC_STORE(&hostGpuCtx->cis[targetRank], ++hostGpuCtx->cisShadow[targetRank],
                              std::memory_order_relaxed);
        TRACE(NCCL_NET, "Updated cis[%u] to %u for context %d", targetRank, hostGpuCtx->cisShadow[targetRank],
              hostGpuCtx->contextId);
      }
    }
  }

  return ncclSuccess;
}

static inline uint64_t extractSignalVal(ncclGinProxyGfd_t* gfd) {
  uint64_t signalVal = gfd->qword[ncclGinProxyGfdCompletion].completion.signalValLow;
  signalVal |= (uint64_t)gfd->qword[ncclGinProxyGfdSignalVal].signalVal.signalValLow2 << 16;
  signalVal |= (uint64_t)gfd->qword[ncclGinProxyGfdSignalVal].signalVal.signalValHigh << 32;
  return signalVal;
}

static ncclGinProxyOp_t extractOp(ncclGinProxyGfd_t* gfd) {
  return (ncclGinProxyOp_t)gfd->qword[ncclGinProxyGfdHeaderExt].headerExt.op;
}

static bool extractIsStrongSignal(ncclGinProxyGfd_t* gfd) {
  // Prior versions only had strong signals
  if (gfd->qword[ncclGinProxyGfdHeader].header.version < 2) return true;
  return gfd->qword[ncclGinProxyGfdSignalVal].signalVal.isStrongSignal != 0;
}

// True if the queue slot at sis[targetRank] holds a GFD (its header flag is set by the producer).
static bool isGfdAvailable(struct ginProxyHostGpuCtx* hostGpuCtx, int targetRank) {
  ncclGinProxyGfd_t* q = hostGpuCtx->queues + targetRank * hostGpuCtx->queueSize;
  uint32_t idx = hostGpuCtx->sis[targetRank] & (hostGpuCtx->queueSize - 1);
  ncclGinProxyQword_t header;
  COMPILER_ATOMIC_LOAD_DEST(&q[idx].qword[ncclGinProxyGfdHeader].raw, &header.raw, std::memory_order_relaxed);
  return header.flag.v != 0;
}

static int proxyGinPollGfd(struct ginProxyCtx* ctx, ginProxyHostGpuCtx* hostGpuCtx, int targetRank,
                           ncclGinProxyGfd_t* gfd, struct ginProxyGfdState** state) {
  if (!isGfdAvailable(hostGpuCtx, targetRank)) {
    return 0;
  }

  ncclGinProxyGfd_t* q = hostGpuCtx->queues + targetRank * hostGpuCtx->queueSize;
  uint32_t idx = hostGpuCtx->sis[targetRank] & (hostGpuCtx->queueSize - 1);
  ncclGinProxyQword_t qword;

  // We know for sure that the first qword is there, copy it.
  gfd->qword[ncclGinProxyGfdHeader] = q[idx].qword[ncclGinProxyGfdHeader];
  // Wait for and copy the other qwords.
  for (int k = 1; k < ncclGinProxyGfdQwords; k++) {
    do {
      COMPILER_ATOMIC_LOAD_DEST(&q[idx].qword[k].raw, &qword.raw, std::memory_order_relaxed);
    } while (qword.flag.v == 0);
    gfd->qword[k] = qword;
  }
  // Now we have the full GFD in the local struct.

  // Reset the GFD in the queue. This ensures that the proxy doesn't try to process the GFD again.
  for (int k = 0; k < ncclGinProxyGfdQwords; k++) {
    COMPILER_ATOMIC_STORE(&q[idx].qword[k].raw, 0ULL, std::memory_order_relaxed);
  }

  // set the counter_id into the state
  uint32_t stateIdx = targetRank * hostGpuCtx->queueSize + idx;
  *state = &hostGpuCtx->states[stateIdx];
  (*state)->op = extractOp(gfd);
  (*state)->counterId = gfd->qword[ncclGinProxyGfdCompletion].completion.counterId;
  (*state)->done = 0;
  (*state)->request = NULL;

  TRACE(NCCL_NET,
        "GFD on context %d to target PE %d raw idx: %u, idx: %u - op: %#x, size: %lu, srcOff: %lu, dstOff: %lu, "
        "srcHandle: %lu, dstHandle: %lu, counterId: %u, signalId: %u, stateIdx: %u",
        hostGpuCtx->contextId, targetRank, hostGpuCtx->sis[targetRank], idx, extractOp(gfd),
        gfd->qword[ncclGinProxyGfdHeader].header.size, gfd->qword[ncclGinProxyGfdSrcOff].srcOff.srcOff,
        gfd->qword[ncclGinProxyGfdDstOff].dstOff.dstOff, gfd->qword[ncclGinProxyGfdSrcHandle].srcHandle.srcHandle,
        gfd->qword[ncclGinProxyGfdDstHandle].dstHandle.dstHandle,
        gfd->qword[ncclGinProxyGfdCompletion].completion.counterId,
        gfd->qword[ncclGinProxyGfdCompletion].completion.signalId, stateIdx);

  hostGpuCtx->sis[targetRank]++;

  return 1;
}

static int mapGfdOpToSignalOp(ncclGinProxyGfd_t* gfd) {
  ncclGinProxyOp_t op = extractOp(gfd);
  uint8_t signalOp = op & (ncclGinProxyOpWithSignalInc | ncclGinProxyOpWithSignalAdd);
  switch (signalOp) {
  case ncclGinProxyOpWithSignalInc:
    return NCCL_NET_SIGNAL_OP_INC;
  case ncclGinProxyOpWithSignalAdd:
    return NCCL_NET_SIGNAL_OP_ADD;
  default:
    return -1;
  }
}

static ncclResult_t proxyGinProcessGfd(struct ginProxyCtx* ctx, struct ginProxyHostGpuCtx* hostGpuCtx, int targetRank,
                                       ncclGinProxyGfd_t* gfd, struct ginProxyGfdState* state, bool isLastInBatch) {
  int signalOp;
  uint64_t signalVal;
  // Aggregate only when this peer's next GFD is already queued and this isn't the last op in the batch.
  // Otherwise, these ops might remain unsubmitted until the next batch.
  bool aggregate = isGfdAvailable(hostGpuCtx, targetRank) && !isLastInBatch;
  uint32_t optFlags = aggregate ? ncclRmaOptFlagsAggregateRequests : ncclRmaOptFlagsDefault;

  // Handle VA Signal operations (signal-only, no PUT)
  if (extractOp(gfd) & ncclGinProxyOpVASignal) {
    uint64_t signalOff = gfd->qword[ncclGinProxyGfdVASignalOff].vaSignalOff.vaSignalOff;
    void* signalHandle = (void*)(uint64_t)gfd->qword[ncclGinProxyGfdVASignalHandle].vaSignalHandle.vaSignalHandle;
    signalVal = extractSignalVal(gfd);
    signalOp = mapGfdOpToSignalOp(gfd);
    NCCLCHECK(rmaBackend->iputSignal(ctx->rmaCtx, hostGpuCtx->contextId, 0, nullptr, 0, 0, nullptr, targetRank,
                                     signalOff, signalHandle, signalVal, signalOp, extractIsStrongSignal(gfd), optFlags,
                                     &state->request));
    return ncclSuccess;
  }

  if (extractOp(gfd) & ncclGinProxyOpGet) {
    uint64_t srcOff = gfd->qword[ncclGinProxyGfdSrcOff].srcOff.srcOff;
    void* srcHandle = (void*)(uint64_t)gfd->qword[ncclGinProxyGfdSrcHandle].srcHandle.srcHandle;
    uint64_t dstOff = gfd->qword[ncclGinProxyGfdDstOff].dstOff.dstOff;
    void* dstHandle = (void*)(uint64_t)gfd->qword[ncclGinProxyGfdDstHandle].dstHandle.dstHandle;
    uint64_t size = gfd->qword[ncclGinProxyGfdHeader].header.size;
    if (!rmaBackend->iget) {
      WARN("GIN plugin does not support GET");
      return ncclInvalidUsage;
    }
    NCCLCHECK(rmaBackend->iget(ctx->rmaCtx, hostGpuCtx->contextId, srcOff, srcHandle, size, dstOff, dstHandle,
                               targetRank, optFlags, &state->request));
    return ncclSuccess;
  }

  if (extractOp(gfd) & ncclGinProxyOpFlush) {
    if (!rmaBackend->iflush) {
      WARN("GIN plugin does not support FLUSH");
      return ncclInvalidUsage;
    }
    NCCLCHECK(rmaBackend->iflush(ctx->rmaCtx, hostGpuCtx->contextId, ctx->signalsGinHandle, targetRank,
                                 &state->request));
    if (state->request == NULL) {
      state->done = 1;
    }
    return ncclSuccess;
  }

  uint64_t size = gfd->qword[ncclGinProxyGfdHeader].header.size;
  uint64_t srcOff;
  void* srcHandle;
  if (extractOp(gfd) & ncclGinProxyOpWithInline) {
    uint64_t* inlineVal = &hostGpuCtx->inlines[state - hostGpuCtx->states];
    srcOff = (uint64_t)&inlineVal[0] - (uint64_t)hostGpuCtx->inlines;
    // reconstruct the inline value from the two qwords
    *inlineVal = gfd->qword[ncclGinProxyGfdInlineLow].inlineLow.inlineValLow;
    if (size > 4) *inlineVal |= (uint64_t)gfd->qword[ncclGinProxyGfdInlineLow].inlineLow.inlineValLow2 << 32;
    if (size > 6) *inlineVal |= (uint64_t)gfd->qword[ncclGinProxyGfdInlineHigh].inlineHigh.inlineValHigh << 48;
    srcHandle = hostGpuCtx->inlinesMhandle;
  } else {
    srcOff = gfd->qword[ncclGinProxyGfdSrcOff].srcOff.srcOff;
    srcHandle = (void*)(uint64_t)gfd->qword[ncclGinProxyGfdSrcHandle].srcHandle.srcHandle;
  }
  uint64_t dstOff = gfd->qword[ncclGinProxyGfdDstOff].dstOff.dstOff;
  void* dstHandle = (void*)(uint64_t)gfd->qword[ncclGinProxyGfdDstHandle].dstHandle.dstHandle;

  ncclGinProxyOp_t op = extractOp(gfd);
  switch (op & ncclGinProxyOpBaseMask) {
  case ncclGinProxyOpPut:
    signalOp = mapGfdOpToSignalOp(gfd);
    if (signalOp == -1) {
      // First cast from 63 bits to 64 bits and then to void * to avoid warnings
      NCCLCHECK(rmaBackend->iput(ctx->rmaCtx, hostGpuCtx->contextId, srcOff, srcHandle, size, dstOff, dstHandle,
                                 targetRank, optFlags, &state->request));
    } else {
      // Reconstruct the signal value
      signalVal = extractSignalVal(gfd);
      uint64_t signalOff =
        (gfd->qword[ncclGinProxyGfdCompletion].completion.signalId + hostGpuCtx->contextId * ctx->nSignalsPerContext) *
        sizeof(uint64_t);
      NCCLCHECK(rmaBackend->iputSignal(ctx->rmaCtx, hostGpuCtx->contextId, srcOff, srcHandle, size, dstOff, dstHandle,
                                       targetRank, signalOff, ctx->signalsGinHandle, signalVal, signalOp,
                                       extractIsStrongSignal(gfd), optFlags, &state->request));
    }
    break;
  default:
    // this error should already have been checked in pollGfd
    WARN("Unexpected GIN proxy operation %d", op);
    return ncclInternalError;
  }
  TRACE(NCCL_NET, "GFD submitted into GIN plugin - contextId: %d, stateIdx: %lu, request: %p", hostGpuCtx->contextId,
        state - hostGpuCtx->states, state->request);
  return ncclSuccess;
}

struct ncclGinProxyListenComm {
  int dev;
  void* listenComm;
};

static ncclResult_t ncclGinProxyGetProperties(ncclGinProperties_t* ginProps) {
  ginProps->supportsStrongSignals = true;
  ginProps->supportsVASignals = true;
  return ncclSuccess;
}

static ncclResult_t ncclGinProxyListen(void* ctx, int dev, void* handle, void** listenComm) {
  ncclResult_t ret = ncclSuccess;
  struct ncclGinProxyListenComm* lComm;
  NCCLCHECK(ncclCalloc(&lComm, 1));
  lComm->dev = dev;
  NCCLCHECKGOTO(rmaBackend->listen(ctx, dev, handle, &lComm->listenComm), ret, end);

end:
  if (ret != ncclSuccess) free(lComm);
  else *listenComm = lComm;
  return ret;
}

static ncclResult_t ncclGinProxyCloseListen(void* listenComm) {
  struct ncclGinProxyListenComm* lComm = (struct ncclGinProxyListenComm*)listenComm;
  NCCLCHECK(rmaBackend->closeListen(lComm->listenComm));
  free(lComm);
  return ncclSuccess;
}

struct ncclGinProxyCollComm {
  ncclNetProperties_t props;
  int nRanks;
  void* collComm;
};

static ncclResult_t ncclGinProxyConnect(void* ctx, void* handles[], int nranks, int rank, void* listenComm,
                                        void** collComm) {
  ncclResult_t ret = ncclSuccess;
  struct ncclGinProxyCollComm* cComm = NULL;
  struct ncclGinProxyListenComm* lComm = (struct ncclGinProxyListenComm*)listenComm;
  NCCLCHECK(ncclCalloc(&cComm, 1));
  cComm->nRanks = nranks;
  NCCLCHECKGOTO(rmaBackend->getProperties(lComm->dev, &cComm->props), ret, end);
  NCCLCHECKGOTO(rmaBackend->connect(ctx, handles, nranks, rank, lComm->listenComm, &cComm->collComm), ret, end);

end:
  if (ret != ncclSuccess) free(cComm);
  else *collComm = cComm;
  return ret;
}

static ncclResult_t ncclGinProxyCloseColl(void* collComm) {
  struct ncclGinProxyCollComm* cComm = (struct ncclGinProxyCollComm*)collComm;
  NCCLCHECK(rmaBackend->closeColl(cComm->collComm));
  free(cComm);
  return ncclSuccess;
}

// Check if the GIN plugin supports DMA-BUF, if so we can try to get the DMA-BUF handle from CUDA,
// if that fails we fallback to non-DMA-BUF
static ncclResult_t ncclGinProxyRegMrSym(void* ginCtx, void* addr, size_t size, int type, uint64_t mrFlags,
                                         void** mhandle, void** ginHandle) {
  struct ncclGinProxyCollComm* cComm = (struct ncclGinProxyCollComm*)ginCtx;
  if (type == NCCL_PTR_HOST) {
    NCCLCHECK(rmaBackend->regMrSym(cComm->collComm, addr, size, type, mrFlags, mhandle));
  } else if (type == NCCL_PTR_CUDA) {
    ncclResult_t dmabufResult = ncclInvalidUsage;
    if (ncclParamDmaBufEnable() && (cComm->props.ptrSupport & NCCL_PTR_DMABUF)) {
      ncclResult_t registrationResult = ncclSuccess;
      int dmabufFd = -1;
      dmabufResult = getDmaBufFd(addr, size, &dmabufFd);
      if (dmabufResult == ncclSuccess) {
        registrationResult =
          rmaBackend->regMrSymDmaBuf(cComm->collComm, addr, size, type, 0, dmabufFd, mrFlags, mhandle);
        close(dmabufFd);
      }
      if (registrationResult != ncclSuccess) {
        // This code path assumes if one MR enters this path, all others will too.
        // Mixed usage of DataDirect and non-DataDirect breaks GIN ordering guarantees.
        dmabufFd = -1;
        dmabufResult = getDmaBufFd(addr, size, &dmabufFd, true);
        if (dmabufResult == ncclSuccess) {
          NCCLCHECK(rmaBackend->regMrSymDmaBuf(cComm->collComm, addr, size, type, 0, dmabufFd, mrFlags, mhandle));
          close(dmabufFd);
        }
      }
    }
    // Fallback to non-DMA-BUF if the DMA-BUF handle is not supported
    if (dmabufResult != ncclSuccess) {
      NCCLCHECK(rmaBackend->regMrSym(cComm->collComm, addr, size, type, mrFlags, mhandle));
    }
  } else {
    return ncclInvalidUsage;
  }
  *ginHandle = *mhandle;
  return ncclSuccess;
}

static ncclResult_t ncclGinProxyDeregMrSym(void* collComm, void* mhandle) {
  struct ncclGinProxyCollComm* cComm = (struct ncclGinProxyCollComm*)collComm;
  // Deregister the memory region with the GIN plugin
  NCCLCHECK(rmaBackend->deregMrSym(cComm->collComm, mhandle));
  return ncclSuccess;
}

static uint64_t isPowerOfTwo(uint64_t n) {
  return (n > 0) && ((n & (n - 1)) == 0);
}

static ncclResult_t ncclGinProxyCreateContext(void* collComm, ncclGinConfig_t* config, void** outGinCtx,
                                              ncclNetDeviceHandle_t** outDevHandle) {
  if (config->backendVersion < 0 || config->backendVersion > NCCL_GIN_PROXY_GPU_CONTEXT_VERSION) {
    WARN("Invalid GIN proxy backend version %d", config->backendVersion);
    return ncclInternalError;
  }

  struct ncclGinProxyCollComm* cComm = (struct ncclGinProxyCollComm*)collComm;
  void* devGpuCtxArray_h = nullptr;

  if (!ncclGdrCopy) INFO(NCCL_NET, "GIN Proxy will not be using GDRCopy");

  struct ginProxyCtx* proxyCtx = NULL;
  NCCLCHECK(ncclCalloc(&proxyCtx, 1));

  proxyCtx->collComm = cComm->collComm;
  proxyCtx->nRanks = cComm->nRanks;
  int nContexts = proxyCtx->nContexts = config->nContexts;

  ncclRmaConfig_t rmaConfig = {config->nContexts, config->trafficClass, config->rankStride};
  NCCLCHECK(rmaBackend->createContext(cComm->collComm, &rmaConfig, &proxyCtx->rmaCtx));

  // Parse poll batch size
  int64_t pollBatchParam = ncclParamGinProxyPollBatch();
  proxyCtx->pollBatch = (pollBatchParam < 1) ? 1 : (int)pollBatchParam;

  // Sanitize the queue size
  uint64_t queueSize = ncclParamGinProxyQueueSize();
  uint32_t maxRequests = NCCL_NET_MAX_REQUESTS * cComm->props.maxRecvs;
  if (queueSize == -1) {
    queueSize = maxRequests;
  }
  if (queueSize > maxRequests) {
    INFO(NCCL_NET,
         "NCCL_GIN_PROXY_QUEUE_SIZE is greater than the maximum outstanding requests in the GIN "
         "plugin (%d), using the default/maximum value instead",
         maxRequests);
    queueSize = maxRequests;
  }
  if (queueSize < 1) {
    INFO(NCCL_NET, "NCCL_GIN_PROXY_QUEUE_SIZE is less than 1, using the default/maximum value instead");
    queueSize = maxRequests;
  }
  if (!isPowerOfTwo(queueSize)) {
    INFO(NCCL_NET, "NCCL_GIN_PROXY_QUEUE_SIZE is not a power of two, using the default/maximum value instead");
    queueSize = maxRequests;
  }

  if (config->nCounters) {
    // Allocate the counters on the GPU or CPU depending on GDR
    NCCLCHECK(allocMemCPUAccessible(&proxyCtx->counters, &proxyCtx->countersDev, config->nCounters * nContexts,
                                    CU_MEMHOSTALLOC_WRITECOMBINED, &proxyCtx->countersGdrHandle, NULL));
  }
  proxyCtx->nCountersPerContext = config->nCounters;

  // Allocate the signals on the GPU and then register the memory region with the GIN plugin.
  // Enforcing strong ordering on the signals mr is vital to ensure ordering between puts and
  // signals.
  if (config->nSignals) {
    size_t signalsBufSize = config->nSignals * nContexts * sizeof(uint64_t);
    NCCLCHECK(ncclCuMemAlloc((void**)&proxyCtx->signalsDev, &proxyCtx->signalsCumemhandle, CU_MEM_HANDLE_TYPE_NONE,
                             signalsBufSize, NULL));
    CUDACHECK(cudaMemset(proxyCtx->signalsDev, 0, signalsBufSize));
    uint64_t mrFlags = NCCL_NET_MR_FLAG_FORCE_SO;
    if (config->backendVersion >= 2) {
      // Signals were always reset in v1.
      mrFlags |= NCCL_NET_MR_FLAG_SIGNAL_NEVER_RESET;
    }
    NCCLCHECK(ncclGinProxyRegMrSym(collComm, proxyCtx->signalsDev, signalsBufSize, NCCL_PTR_CUDA, mrFlags,
                                   &proxyCtx->signalsMhandle, &proxyCtx->signalsGinHandle));
  }
  proxyCtx->nSignalsPerContext = config->nSignals;

  // Allocate offset arrays for reset-without-zeroing (GPU memory, zeroed)
  if (config->nSignals) {
    NCCLCHECK(ncclCudaCalloc(&proxyCtx->signalOffsetsDev, config->nSignals * nContexts, NULL));
  }

  NCCLCHECK(ncclCalloc(&proxyCtx->hostGpuCtx, nContexts));
  const size_t gpuCtxArraySize = nContexts * NCCL_GIN_PROXY_GPU_CONTEXT_MAX_SIZE;
  NCCLCHECK(ncclCalloc(((char**)&devGpuCtxArray_h), gpuCtxArraySize));
  for (int contextId = 0; contextId < nContexts; contextId++) {
    struct ginProxyHostGpuCtx* hostGpuCtx = proxyCtx->hostGpuCtx + contextId;
    hostGpuCtx->contextId = contextId;
    hostGpuCtx->queueSize = queueSize;
    size_t queuesLength = hostGpuCtx->queueSize * cComm->nRanks;
    NCCLCHECK(ncclCalloc(&hostGpuCtx->states, queuesLength));
    NCCLCHECK(ncclCalloc(&hostGpuCtx->cisShadow, cComm->nRanks));
    NCCLCHECK(ncclCalloc(&hostGpuCtx->sis, cComm->nRanks));
    NCCLCHECK(ncclCalloc(&hostGpuCtx->inlines, queuesLength));
    NCCLCHECK(ncclGinProxyRegMrSym(collComm, hostGpuCtx->inlines, queuesLength * sizeof(uint64_t), NCCL_PTR_HOST, 0,
                                   &hostGpuCtx->inlinesMhandle, &hostGpuCtx->inlinesGinHandle));
    NCCLCHECK(ncclCudaCalloc(&hostGpuCtx->pis, cComm->nRanks, NULL));
    NCCLCHECK(ncclCudaCalloc(&hostGpuCtx->lastIssuedGet, cComm->nRanks, NULL));
    NCCLCHECK(ncclCudaCalloc(&hostGpuCtx->lastVisibleGet, cComm->nRanks, NULL));

    // Allocate the GFD queues, CIs, counters, signals and test/wait variables on the either the CPU
    // or GPU.
    uint32_t* cisDev;
    ncclGinProxyGfd_t* queuesDev;
    NCCLCHECK(allocMemCPUAccessible(&hostGpuCtx->queues, &queuesDev, queuesLength, 0, NULL, NULL, true /*forceHost*/));
    NCCLCHECK(allocMemCPUAccessible(&hostGpuCtx->cis, &cisDev, cComm->nRanks, CU_MEMHOSTALLOC_WRITECOMBINED,
                                    &hostGpuCtx->cisGdrHandle, NULL));

    NCCLCHECK(ncclGinProxyGpuCtx_init(
      config->backendVersion, devGpuCtxArray_h, contextId, cComm->nRanks, hostGpuCtx->queueSize, queuesDev,
      hostGpuCtx->pis, cisDev, proxyCtx->countersDev + contextId * config->nCounters,
      proxyCtx->signalsDev + contextId * config->nSignals, proxyCtx->signalOffsetsDev + contextId * config->nSignals,
      hostGpuCtx->lastIssuedGet, hostGpuCtx->lastVisibleGet));
  }

  void* devGpuCtxArray_d = nullptr;
  NCCLCHECK(ncclCudaCalloc(((char**)&devGpuCtxArray_d), gpuCtxArraySize, NULL));
  // Copy the proxy's devGpuCtx to the GPU
  NCCLCHECK(ncclCudaMemcpy(devGpuCtxArray_d, devGpuCtxArray_h, gpuCtxArraySize));

  ncclNetDeviceHandle_t* devHandle = NULL;
  NCCLCHECK(ncclCalloc(&devHandle, 1));
  devHandle->netDeviceType = NCCL_NET_DEVICE_GIN_PROXY;
  devHandle->netDeviceVersion = NCCL_GIN_PROXY_VERSION;
  devHandle->handle = devGpuCtxArray_d;
  devHandle->size = 0;
  devHandle->needsProxyProgress = 1;

  proxyCtx->devHandle = devHandle;

  *outDevHandle = devHandle;
  *outGinCtx = proxyCtx;

  free(devGpuCtxArray_h);

  return ncclSuccess;
}

static ncclResult_t ncclGinProxyDestroyContext(void* ginCtx) {
  if (!ginCtx) return ncclSuccess;
  struct ginProxyCtx* ctx = (struct ginProxyCtx*)ginCtx;

  NCCLCHECK(rmaBackend->destroyContext(ctx->rmaCtx));

  // Free counters
  if (ctx) {
    if (ctx->counters || ctx->countersGdrHandle) {
      NCCLCHECK(freeMemCPUAccessible(ctx->counters, ctx->countersGdrHandle, NULL));
    }

    // Free signals
    if (ctx->collComm && ctx->signalsMhandle) rmaBackend->deregMrSym(ctx->collComm, ctx->signalsMhandle);
    if (ctx->signalsDev) NCCLCHECK(ncclCudaFree(ctx->signalsDev, NULL));

    // Free hostGpuCtx and its allocations
    if (ctx->hostGpuCtx) {
      for (int contextId = 0; contextId < ctx->nContexts; contextId++) {
        struct ginProxyHostGpuCtx* hostGpuCtx = ctx->hostGpuCtx + contextId;
        if (hostGpuCtx->cisShadow) free(hostGpuCtx->cisShadow);
        if (hostGpuCtx->sis) free(hostGpuCtx->sis);
        if (hostGpuCtx->pis) NCCLCHECK(ncclCudaFree(hostGpuCtx->pis, NULL));
        if (hostGpuCtx->lastIssuedGet) NCCLCHECK(ncclCudaFree(hostGpuCtx->lastIssuedGet, NULL));
        if (hostGpuCtx->lastVisibleGet) NCCLCHECK(ncclCudaFree(hostGpuCtx->lastVisibleGet, NULL));
        if (hostGpuCtx->states) free(hostGpuCtx->states);
        if (hostGpuCtx->inlines) free(hostGpuCtx->inlines);
        if (ctx->collComm && hostGpuCtx->inlinesMhandle) {
          rmaBackend->deregMrSym(ctx->collComm, hostGpuCtx->inlinesMhandle);
        }
        if (hostGpuCtx->queues) NCCLCHECK(freeMemCPUAccessible(hostGpuCtx->queues, NULL, NULL));
        if (hostGpuCtx->cis || hostGpuCtx->cisGdrHandle) {
          NCCLCHECK(freeMemCPUAccessible(hostGpuCtx->cis, hostGpuCtx->cisGdrHandle, NULL));
        }
      }
      free(ctx->hostGpuCtx);
    }

    ncclNetDeviceHandle_t* devHandle = (ncclNetDeviceHandle_t*)ctx->devHandle;
    if (devHandle) {
      if (devHandle->handle) NCCLCHECK(ncclCudaFree((void*)devHandle->handle, NULL));
      free(devHandle);
    }

    free(ctx);
  }

  return ncclSuccess;
}

static ncclResult_t ncclGinProxyProgress(void* ginCtx) {
  struct ginProxyCtx* ctx = (struct ginProxyCtx*)ginCtx;

  for (int contextId = 0; contextId < ctx->nContexts; contextId++) {
    struct ginProxyHostGpuCtx* hostGpuCtx = ctx->hostGpuCtx + contextId;
    NCCLCHECK(proxyGinPollCompletions(ctx->collComm, ctx, hostGpuCtx));
    for (int targetRank = 0; targetRank < ctx->nRanks; targetRank++) {
      for (int p = 0; p < ctx->pollBatch; p++) {
        ncclGinProxyGfd_t gfd;
        struct ginProxyGfdState* state = NULL;
        if (!proxyGinPollGfd(ctx, hostGpuCtx, targetRank, &gfd, &state)) break;
        bool isLastInBatch = (p == ctx->pollBatch - 1);
        ncclResult_t ret = proxyGinProcessGfd(ctx, hostGpuCtx, targetRank, &gfd, state, isLastInBatch);
        if (ret) ctx->hasError = ret;
        NCCLCHECK(ret);
      }
    }
  }
  if (rmaBackend->rmaProgress) rmaBackend->rmaProgress(ctx->rmaCtx);

  return ncclSuccess;
}

static ncclResult_t ncclGinProxyQueryLastError(void* ginCtx, bool* hasError) {
  struct ginProxyCtx* ctx = (struct ginProxyCtx*)ginCtx;
  *hasError = ctx->hasError;
  if (ctx->hasError == ncclSuccess && rmaBackend->queryLastError) {
    NCCLCHECK(rmaBackend->queryLastError(ctx->rmaCtx, hasError));
  }
  return ncclSuccess;
}

ncclGin_t ncclGinProxy{
  NULL, // Will map directly to the plugin: name
  NULL, // Will map directly to the plugin: init()
  NULL, // Will map directly to the plugin: devices()
  ncclGinProxyGetProperties,
  NULL, // Will map directly to the plugin: getProperties()
  ncclGinProxyListen,
  ncclGinProxyConnect,
  ncclGinProxyCreateContext,
  ncclGinProxyRegMrSym,
  NULL, // regMrSymDmaBuf() is not used by upper layer at the moment, hidden in RegMrSym.
  ncclGinProxyDeregMrSym,
  ncclGinProxyDestroyContext,
  ncclGinProxyCloseColl,
  ncclGinProxyCloseListen,
  ncclGinProxyProgress,
  ncclGinProxyQueryLastError,
  NULL  // Will map directly to the plugin: finalize()
};

int ncclGinProxyVersion = -1;

ncclResult_t ncclGinProxyInit(struct ncclComm* comm) {
  // Use GIN proxy implementation, based on the RMA plugin and
  // GPU-host communication queues.
  rmaBackend = comm->rmaState.rmaProxyState.ncclRma;
  if (!rmaBackend) {
    WARN("GIN Proxy: RMA backend not found");
    return ncclInvalidArgument;
  }
  ncclGinProxy.name = rmaBackend->name;
  ncclGinProxy.init = rmaBackend->init;
  ncclGinProxy.devices = rmaBackend->devices;
  ncclGinProxy.getProperties = rmaBackend->getProperties;
  ncclGinProxy.finalize = rmaBackend->finalize;
  ncclGinProxyVersion = comm->rmaState.rmaProxyState.rmaVersion;
  return ncclSuccess;
}
