2.21.5-1

Add support for IB SHARP 1PPN operation with user buffers.
Improve support for MNNVL, add NVLS support and multi-clique support.
 * Detect the NVLS clique through NVML
 * Exchange XML between peers in the same NVLS clique and fuse XMLs
   before creating the topology graph.
 * Rework bootstrap allgather algorithms to allow for large allgather
   operations intra-node (XML exchange).
Net/IB: add support for dynamic GID detection.
 * Automatically select RoCEv2/IPv4 interface by default. Allow to
   select IPv6 or even the network/mask.
Reduce NVLS memory usage.
 * Add stepSize as property of a connection to allow for different
   sizes on different peers; set it to 128K for NVLink SHARP.
Improve tuner loading
 * Look for more paths, be more consistent with the network device
   plugin.
 * Also search for tuner support inside the net plugin.
Improve tuner API
 * Add context to support multi-device per process.
Add magic number around comm object to detect comm corruption.
 * Add some basic check around communicators so that we can report a
   problem when a communicator gets corrupted or a wrong comm pointer
   is passed to NCCL.
Fix net/IB error path. Github PR #1164
Fix collnet rail mapping with split comm.
Fix packet reordering issue causing bootstrap mismatch
 * Use a different tag in ncclTransportP2pSetup for the connectInfo
   exchange and the following barrier.
Fix hang when crossNic is inconsistent between ranks.
Fix minCompCap/maxCompCap computation. Github issue #1184

src/include/argcheck.h

@@ -11,6 +11,7 @@
 #include "info.h"
 
 ncclResult_t PtrCheck(void* ptr, const char* opname, const char* ptrname);
+ncclResult_t CommCheck(struct ncclComm* ptr, const char* opname, const char* ptrname);
 ncclResult_t ArgsCheck(struct ncclInfo* info);
 ncclResult_t CudaPtrCheck(const void* pointer, struct ncclComm* comm, const char* ptrname, const char* opname);
 

src/include/bootstrap.h

@@ -24,7 +24,9 @@ ncclResult_t bootstrapSplit(struct ncclBootstrapHandle* handle, struct ncclComm*
 ncclResult_t bootstrapAllGather(void* commState, void* allData, int size);
 ncclResult_t bootstrapSend(void* commState, int peer, int tag, void* data, int size);
 ncclResult_t bootstrapRecv(void* commState, int peer, int tag, void* data, int size);
-ncclResult_t bootstrapBarrier(void* commState, int *ranks, int rank, int nranks, int tag);
+ncclResult_t bootstrapBarrier(void* commState, int rank, int nranks, int tag);
+ncclResult_t bootstrapBroadcast(void* commState, int rank, int nranks, int root, void* bcastData, int size);
+ncclResult_t bootstrapIntraNodeBarrier(void* commState, int *ranks, int rank, int nranks, int tag);
 ncclResult_t bootstrapIntraNodeAllGather(void* commState, int *ranks, int rank, int nranks, void* allData, int size);
 ncclResult_t bootstrapIntraNodeBroadcast(void* commState, int *ranks, int rank, int nranks, int root, void* bcastData, int size);
 ncclResult_t bootstrapClose(void* commState);

src/include/comm.h

@@ -87,6 +87,12 @@ struct ncclNodeRanks {
   int* localRankToRank;
 };
 
+struct cliqueInfo {
+  int id;
+  int size;
+  int *ranks;
+};
+
 struct ncclDestructor {
   struct ncclDestructor* next;
   void* obj;
@@ -165,6 +171,14 @@ struct ncclNvlsMcHandleList {
   size_t size;
 };
 
+struct ncclCollnetHandleList {
+  struct ncclCollnetHandleList *next;
+  void* collnetHandle;
+  size_t size;
+  const void* buffer;
+  struct ncclProxyConnector* proxyconn;
+};
+
 struct ncclKernelPlan {
   // A kernel plan is also a callback that reclaims itself. Hence this must
   // be the first member.
@@ -188,6 +202,7 @@ struct ncclKernelPlan {
 
   struct ncclIntruQueue<struct ncclPointerList, &ncclPointerList::next> ipcMemQueue;
   struct ncclIntruQueue<struct ncclNvlsMcHandleList, &ncclNvlsMcHandleList::next> nvlsMcHandleQueue;
+  struct ncclIntruQueue<struct ncclCollnetHandleList, &ncclCollnetHandleList::next> collnetHandleQueue;
 
   struct Channel {
     int nWork;
@@ -202,7 +217,10 @@ struct ncclKernelPlan {
   size_t maxBytesPerChannel;
 };
 
+#define NCCL_MAGIC 0x0280028002800280 // Nickel atomic number is 28.
+
 struct ncclComm {
+  uint64_t startMagic;
   struct ncclMemoryStack memPermanent, memScoped;
   // List of destructors to run when comm is destructed
   struct ncclDestructor* destructorHead;
@@ -245,7 +263,10 @@ struct ncclComm {
   int* localRankToRank;
   // localRanks and localRanktoRank for all nodes
   struct ncclNodeRanks* nodeRanks;
-  int MNNVL; // MNNVL: Multi-Node NVLink
+  // MNNVL: Multi-Node NVLink
+  int MNNVL; // true when MNNVL is available
+  struct cliqueInfo clique; // Our MNNVL clique information
+  int cliqueRank; // Our rank within the MNNVL clique
 
   bool checkPointers;
   bool dmaBufSupport;
@@ -257,7 +278,6 @@ struct ncclComm {
   int nChannels; // connection nChannels
   int collChannels; // enqueue nChannels
   int nvlsChannels; // enqueue nChannels
-  int collNetChannels;
   // Channels (per peer) for p2p
   int p2pnChannels;
   int p2pnChannelsPerPeer;
@@ -269,6 +289,7 @@ struct ncclComm {
   // Buffer sizes
   int buffSizes[NCCL_NUM_PROTOCOLS];
   int p2pChunkSize;
+  int nvlsChunkSize;
 
   // Algorithm/Protocols thresholds
   ssize_t threadThresholds[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
@@ -315,11 +336,11 @@ struct ncclComm {
   int proxyRefCountOld; /* store proxy post-atomic-sub refcount */
   // Whether this communicator uses collNet
   int collNetSupport;
+  bool collNetRegSupport;
   uint8_t collNetSupportMatrix[4/*sum,prod,min,max*/][ncclNumTypes];
   int intraHighestTransportType;
   int* collNetHeads;
   int collNetHeadsNum;
-  int collNetHeadsUniqueNum;
   int* collNetDenseToUserRank;
   int* collNetUserToDenseRank;
   /* sharable collNet proxy progress resource. */
@@ -336,6 +357,7 @@ struct ncclComm {
   struct ncclMemoryPool memPool_ncclKernelPlan;
   struct ncclMemoryPool memPool_ncclPointerList;
   struct ncclMemoryPool memPool_ncclNvlsHandleList;
+  struct ncclMemoryPool memPool_ncclCollnetHandleList;
   // Next comm in this thread's active ncclGroup[Start|End](). Holds "0x1" when
   // this comm is not yet in a group.
   struct ncclComm* groupNext;
@@ -368,8 +390,10 @@ struct ncclComm {
 
   // Tuning plugin
   ncclTuner_t* tuner;
+  void *tunerContext;
   // buffer registration cache
   struct ncclRegCache regCache;
+  uint64_t endMagic;
 };
 
 enum ncclLaunchMode {

src/include/cudawrap.h

@@ -20,10 +20,6 @@ extern int ncclCuMemEnable();
 // Handle type used for cuMemCreate()
 extern CUmemAllocationHandleType ncclCuMemHandleType;
 
-#else
-typedef CUresult (CUDAAPI *PFN_cuInit_v2000)(unsigned int Flags);
-typedef CUresult (CUDAAPI *PFN_cuDriverGetVersion_v2020)(int *driverVersion);
-typedef CUresult (CUDAAPI *PFN_cuGetProcAddress_v11030)(const char *symbol, void **pfn, int driverVersion, cuuint64_t flags);
 #endif
 
 #define CUPFN(symbol) pfn_##symbol
@@ -69,53 +65,47 @@ typedef CUresult (CUDAAPI *PFN_cuGetProcAddress_v11030)(const char *symbol, void
     }									\
 } while(0)
 
-#define DECLARE_CUDA_PFN_EXTERN(symbol,version) extern PFN_##symbol##_v##version pfn_##symbol
+#define DECLARE_CUDA_PFN_EXTERN(symbol) extern PFN_##symbol pfn_##symbol
 
 #if CUDART_VERSION >= 11030
 /* CUDA Driver functions loaded with cuGetProcAddress for versioning */
-DECLARE_CUDA_PFN_EXTERN(cuDeviceGet, 2000);
-DECLARE_CUDA_PFN_EXTERN(cuDeviceGetAttribute, 2000);
-DECLARE_CUDA_PFN_EXTERN(cuGetErrorString, 6000);
-DECLARE_CUDA_PFN_EXTERN(cuGetErrorName, 6000);
-DECLARE_CUDA_PFN_EXTERN(cuMemGetAddressRange, 3020);
-DECLARE_CUDA_PFN_EXTERN(cuCtxCreate, 3020);
-DECLARE_CUDA_PFN_EXTERN(cuCtxDestroy, 4000);
-DECLARE_CUDA_PFN_EXTERN(cuCtxGetCurrent, 4000);
-DECLARE_CUDA_PFN_EXTERN(cuCtxSetCurrent, 4000);
-DECLARE_CUDA_PFN_EXTERN(cuCtxGetDevice, 2000);
-DECLARE_CUDA_PFN_EXTERN(cuPointerGetAttribute, 4000);
+DECLARE_CUDA_PFN_EXTERN(cuDeviceGet);
+DECLARE_CUDA_PFN_EXTERN(cuDeviceGetAttribute);
+DECLARE_CUDA_PFN_EXTERN(cuGetErrorString);
+DECLARE_CUDA_PFN_EXTERN(cuGetErrorName);
+DECLARE_CUDA_PFN_EXTERN(cuMemGetAddressRange);
+DECLARE_CUDA_PFN_EXTERN(cuCtxCreate);
+DECLARE_CUDA_PFN_EXTERN(cuCtxDestroy);
+DECLARE_CUDA_PFN_EXTERN(cuCtxGetCurrent);
+DECLARE_CUDA_PFN_EXTERN(cuCtxSetCurrent);
+DECLARE_CUDA_PFN_EXTERN(cuCtxGetDevice);
+DECLARE_CUDA_PFN_EXTERN(cuPointerGetAttribute);
 // cuMem API support
-DECLARE_CUDA_PFN_EXTERN(cuMemAddressReserve, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemAddressFree, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemCreate, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemGetAllocationGranularity, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemExportToShareableHandle, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemImportFromShareableHandle, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemMap, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemRelease, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemRetainAllocationHandle, 11000);
-DECLARE_CUDA_PFN_EXTERN(cuMemSetAccess, 10020);
-DECLARE_CUDA_PFN_EXTERN(cuMemUnmap, 10020);
+DECLARE_CUDA_PFN_EXTERN(cuMemAddressReserve);
+DECLARE_CUDA_PFN_EXTERN(cuMemAddressFree);
+DECLARE_CUDA_PFN_EXTERN(cuMemCreate);
+DECLARE_CUDA_PFN_EXTERN(cuMemGetAllocationGranularity);
+DECLARE_CUDA_PFN_EXTERN(cuMemExportToShareableHandle);
+DECLARE_CUDA_PFN_EXTERN(cuMemImportFromShareableHandle);
+DECLARE_CUDA_PFN_EXTERN(cuMemMap);
+DECLARE_CUDA_PFN_EXTERN(cuMemRelease);
+DECLARE_CUDA_PFN_EXTERN(cuMemRetainAllocationHandle);
+DECLARE_CUDA_PFN_EXTERN(cuMemSetAccess);
+DECLARE_CUDA_PFN_EXTERN(cuMemUnmap);
 #if CUDA_VERSION >= 11070
-DECLARE_CUDA_PFN_EXTERN(cuMemGetHandleForAddressRange, 11070); // DMA-BUF support
+DECLARE_CUDA_PFN_EXTERN(cuMemGetHandleForAddressRange); // DMA-BUF support
 #endif
 #if CUDA_VERSION >= 12010
 /* NVSwitch Multicast support */
-DECLARE_CUDA_PFN_EXTERN(cuMulticastAddDevice, 12010);
-DECLARE_CUDA_PFN_EXTERN(cuMulticastBindMem, 12010);
-DECLARE_CUDA_PFN_EXTERN(cuMulticastBindAddr, 12010);
-DECLARE_CUDA_PFN_EXTERN(cuMulticastCreate, 12010);
-DECLARE_CUDA_PFN_EXTERN(cuMulticastGetGranularity, 12010);
-DECLARE_CUDA_PFN_EXTERN(cuMulticastUnbind, 12010);
+DECLARE_CUDA_PFN_EXTERN(cuMulticastAddDevice);
+DECLARE_CUDA_PFN_EXTERN(cuMulticastBindMem);
+DECLARE_CUDA_PFN_EXTERN(cuMulticastBindAddr);
+DECLARE_CUDA_PFN_EXTERN(cuMulticastCreate);
+DECLARE_CUDA_PFN_EXTERN(cuMulticastGetGranularity);
+DECLARE_CUDA_PFN_EXTERN(cuMulticastUnbind);
 #endif
 #endif
 
-/* CUDA Driver functions loaded with dlsym() */
-DECLARE_CUDA_PFN_EXTERN(cuInit, 2000);
-DECLARE_CUDA_PFN_EXTERN(cuDriverGetVersion, 2020);
-DECLARE_CUDA_PFN_EXTERN(cuGetProcAddress, 11030);
-
-
 ncclResult_t ncclCudaLibraryInit(void);
 
 extern int ncclCudaDriverVersionCache;

src/include/device.h

@@ -84,6 +84,15 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK
 #define NCCL_IPC_READ     0x10
 #define NCCL_NVLS_MIN_POLL 0x20
 
+#define NCCL_MAX_COLLNET_SIZE (1L << 29)
+
+enum ncclRegBufferType {
+  NCCL_REGULAR_BUFFER = 0,
+  NCCL_IPC_REG_BUFFER = 1,
+  NCCL_NVLS_REG_BUFFER = 2,
+  NCCL_COLLNET_REG_BUFFER = 3
+};
+
 struct ncclConnInfo {
   // Regular comm mechanism
   char *buffs[NCCL_NUM_PROTOCOLS]; // Local for recv, remote for send
@@ -93,6 +102,7 @@ struct ncclConnInfo {
 
   int flags;          // Direct communication / other flags
   int shared;         // Buffers are shared
+  int stepSize;       // Step size for the SIMPLE buffer
   void **ptrExchange; // Pointer exchange for direct communication
   uint64_t* redOpArgExchange; // PreOp scaler exchange for direct pull case
 
@@ -157,7 +167,7 @@ struct ncclDirect {
   int down[NCCL_MAX_DIRECT_ARITY];
 };
 
-#define NCCL_MAX_NVLS_ARITY 8
+#define NCCL_MAX_NVLS_ARITY 32
 #define NCCL_MAX_NVLS_TREE_ARITY 3
 struct ncclNvls {
   int out;
@@ -171,6 +181,12 @@ struct ncclNvls {
   int nNodes;
 };
 
+#if __CUDA_ARCH__ >= 900
+#define NCCL_MAX_ARITY NCCL_MAX_NVLS_ARITY
+#else
+#define NCCL_MAX_ARITY NCCL_MAX_DIRECT_ARITY
+#endif
+
 #define NCCL_MAX_CONNS 2
 struct ncclChannelPeer {
   struct ncclConnector send[NCCL_MAX_CONNS];
@@ -212,9 +228,10 @@ struct ncclWorkElem {
   union {
     uint8_t flagBits;
     struct {
-      uint8_t isUsed:1, redOpArgIsPtr:1, regUsed:1, oneNode:1;
+      uint8_t isUsed:1, redOpArgIsPtr:1, oneNode:1;
     };
   };
+  uint8_t regUsed;
   uint8_t nWarps;
   uint8_t direct;
   uint32_t root;

src/include/graph.h

@@ -31,10 +31,10 @@ ncclResult_t ncclTopoGetNvbGpus(struct ncclTopoSystem* system, int rank, int* nr
 int ncclTopoPathAllNVLink(struct ncclTopoSystem* system);
 
 // Query topology
-ncclResult_t ncclTopoGetNetDev(struct ncclComm* comm, int rank, struct ncclTopoGraph* graph, int channelId, int peerRank, int* net, int* proxyRank);
+ncclResult_t ncclTopoGetNetDev(struct ncclComm* comm, int rank, struct ncclTopoGraph* graph, int channelId, int peerRank, int64_t* id, int* dev, int* proxyRank);
 ncclResult_t ncclTopoCheckP2p(struct ncclTopoSystem* system, int64_t id1, int64_t id2, int* p2p, int *read, int* intermediateRank);
 ncclResult_t ncclTopoCheckMNNVL(struct ncclTopoSystem* system, struct ncclPeerInfo* info1, struct ncclPeerInfo* info2, int* ret);
-ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int netDev, int read, int* useGdr);
+ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int64_t netId, int read, int* useGdr);
 ncclResult_t ncclTopoNeedFlush(struct ncclTopoSystem* system, int64_t busId, int* flush);
 ncclResult_t ncclTopoCheckNet(struct ncclTopoSystem* system, int64_t id1, int64_t id2, int* net);
 int ncclPxnDisable(struct ncclComm* comm);
@@ -56,8 +56,8 @@ ncclResult_t ncclTopoCpuType(struct ncclTopoSystem* system, int* arch, int* vend
 ncclResult_t ncclTopoGetGpuCount(struct ncclTopoSystem* system, int* count);
 ncclResult_t ncclTopoGetNetCount(struct ncclTopoSystem* system, int* count);
 ncclResult_t ncclTopoGetNvsCount(struct ncclTopoSystem* system, int* count);
-ncclResult_t ncclTopoGetLocalNet(struct ncclTopoSystem* system, int rank, int channelId, int* id);
-ncclResult_t ncclTopoGetLocalGpu(struct ncclTopoSystem* system, int net, int* gpuIndex);
+ncclResult_t ncclTopoGetLocalNet(struct ncclTopoSystem* system, int rank, int channelId, int64_t* id, int* dev);
+ncclResult_t ncclTopoGetLocalGpu(struct ncclTopoSystem* system, int64_t netId, int* gpuIndex);
 ncclResult_t getLocalNetCountByBw(struct ncclTopoSystem* system, int gpu, int *count);
 
 #define NCCL_TOPO_MAX_NODES 256
@@ -88,7 +88,7 @@ struct ncclTopoGraph {
   int sameChannels;
   int nHops;
   int intra[MAXCHANNELS*NCCL_TOPO_MAX_NODES];
-  int inter[MAXCHANNELS*2];
+  int64_t inter[MAXCHANNELS*2];
 };
 ncclResult_t ncclTopoCompute(struct ncclTopoSystem* system, struct ncclTopoGraph* graph);
 
@@ -110,7 +110,7 @@ struct ncclTopoRanks {
 ncclResult_t ncclTopoPreset(struct ncclComm* comm, struct ncclTopoGraph** graphs, struct ncclTopoRanks* topoRanks);
 
 ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePatterns,
-    struct ncclTopoRanks** allTopoRanks, int* rings, struct ncclTopoGraph** graphs);
+    struct ncclTopoRanks** allTopoRanks, int* rings, struct ncclTopoGraph** graphs, struct ncclComm* parent);
 
 ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCompCap, struct ncclTopoGraph** graphs);
 #include "info.h"

src/include/info.h

@@ -31,13 +31,6 @@ typedef enum : uint8_t {
   ncclPatternRecv
 } ncclPattern_t;
 
-enum ncclRegBufferType {
-  NCCL_REGULAR_BUFFER = 0,
-  NCCL_IPC_REG_BUFFER = 1,
-  NCCL_NVLS_REG_BUFFER = 2,
-  NCCL_REG_BUFFER_NUM = 3
-};
-
 // Used to pass NCCL call information between functions
 struct ncclInfo {
   ncclFunc_t coll;
@@ -70,6 +63,9 @@ struct ncclInfo {
   ncclRegBufferType regBufType;
   void* regBufSend[NCCL_MAX_LOCAL_RANKS];
   void* regBufRecv[NCCL_MAX_LOCAL_RANKS];
+  // collnet buffer reg handles
+  void* sendMhandle;
+  void* recvMhandle;
   // Need to initialize
   int nThreads;
   int nChannels;

src/include/nccl_common.h

@@ -8,7 +8,7 @@
 #define NCCL_DEBUG_H_
 
 typedef enum {NCCL_LOG_NONE=0, NCCL_LOG_VERSION=1, NCCL_LOG_WARN=2, NCCL_LOG_INFO=3, NCCL_LOG_ABORT=4, NCCL_LOG_TRACE=5} ncclDebugLogLevel;
-typedef enum {NCCL_INIT=1, NCCL_COLL=2, NCCL_P2P=4, NCCL_SHM=8, NCCL_NET=16, NCCL_GRAPH=32, NCCL_TUNING=64, NCCL_ENV=128, NCCL_ALLOC=256, NCCL_CALL=512, NCCL_PROXY=1024, NCCL_NVLS=2048, NCCL_ALL=~0} ncclDebugLogSubSys;
+typedef enum {NCCL_INIT=1, NCCL_COLL=2, NCCL_P2P=4, NCCL_SHM=8, NCCL_NET=16, NCCL_GRAPH=32, NCCL_TUNING=64, NCCL_ENV=128, NCCL_ALLOC=256, NCCL_CALL=512, NCCL_PROXY=1024, NCCL_NVLS=2048, NCCL_BOOTSTRAP=4096, NCCL_REG=8192, NCCL_ALL=~0} ncclDebugLogSubSys;
 
 typedef void (*ncclDebugLogger_t)(ncclDebugLogLevel level, unsigned long flags, const char *file, int line, const char *fmt, ...);
 

src/include/nccl_tuner.h

@@ -17,13 +17,17 @@ typedef struct {
   const char* name;
 
   // Initializes tuner states.
-  // nRanks: number of ranks in current communicator. Each communicator initialize its own tuner.
-  // nNodes: number of nodes in current communicator.
-  // logFunction: a logFunction can be useful to integrate logging together with NCCL core.
-  ncclResult_t (*init)(size_t nRanks, size_t nNodes, ncclDebugLogger_t logFunction);
+  // Inputs:
+  //   - nRanks: number of ranks in current communicator. Each communicator initialize its own tuner.
+  //   - nNodes: number of nodes in current communicator.
+  //   - logFunction: a logFunction can be useful to integrate logging together with NCCL core.
+  // Outputs:
+  //   - context: tuner context object
+  ncclResult_t (*init)(size_t nRanks, size_t nNodes, ncclDebugLogger_t logFunction, void **context);
 
   // Gets info (algo, protocol, number of ctas and threads) for a given collective.
   // Inputs:
+  //   - context: tuner context object
   //   - collType: collective type , e.g., allreduce, allgather…
   //   - nBytes: collective size in bytes
   //   - collNetTypeSupport: whether collnet supports this type
@@ -40,16 +44,17 @@ typedef struct {
   // Also, the plugin is allowed to not set any output, or set only the
   // algorithm and protocol, but not only the algorithm or only the protocol.
   // Unset fields will be set automatically by NCCL.
-  ncclResult_t (*getCollInfo)(ncclFunc_t collType, size_t nBytes,
+  ncclResult_t (*getCollInfo)(void* context, ncclFunc_t collType, size_t nBytes,
                               int collNetSupport, int nvlsSupport, int numPipeOps,
                               int *algorithm, int *protocol, int* nChannels);
 
   // Terminates the plugin and cleans up any resources that the plugin allocated.
-  ncclResult_t (*destroy)();
-} ncclTuner_v1_t;
+  // context: tuner context object
+  ncclResult_t (*destroy)(void* context);
+} ncclTuner_v2_t;
 
-typedef ncclTuner_v1_t ncclTuner_t;
+typedef ncclTuner_v2_t ncclTuner_t;
 
-#define NCCL_TUNER_PLUGIN_SYMBOL "ncclTunerPlugin_v1"
+#define NCCL_TUNER_PLUGIN_SYMBOL "ncclTunerPlugin_v2"
 
 #endif

src/include/proxy.h

@@ -33,7 +33,6 @@ union ncclProxyOpSpecifics {
 
 struct ncclProxyOp {
   struct ncclProxyConnection* connection;
-  void* buffer;
   ssize_t nbytes;
   uint64_t opCount;
   int root;
@@ -49,6 +48,11 @@ struct ncclProxyOp {
   uint8_t /*ncclPattern_t*/ pattern;
   uint8_t protocol;
   uint8_t reg;
+  // collnet buffer reg handles
+  void* sendMhandle;
+  void* recvMhandle;
+  uint8_t* sendbuff;
+  uint8_t* recvbuff;
 
   union ncclProxyOpSpecifics specifics;
 
@@ -58,8 +62,14 @@ struct ncclProxyOp {
 struct ncclProxySubArgs {
   struct ncclProxyConnection* connection;
   int reg;
-  void* buffer;
+  // p2p mhandle
   void* mhandle;
+  // collnet handles
+  void* sendMhandle;
+  void* recvMhandle;
+  uint8_t* sendbuff;
+  uint8_t* recvbuff;
+  size_t offset;
   int channelId;
   int nsteps;
   ssize_t nbytes;
@@ -88,6 +98,10 @@ struct ncclProxyArgs {
   int sliceSteps;
   int chunkSteps;
   int chunkSize;
+  size_t totalSendSize;
+  size_t totalRecvSize;
+  size_t sendSizePerRound;
+  size_t recvSizePerRound;
   uint8_t /*ncclDataType_t*/ dtype;
   uint8_t /*ncclDevRedOp_t*/ redOp;
   uint8_t /*ncclPattern_t*/ pattern;
@@ -302,6 +316,8 @@ enum ncclProxyMsgType {
   ncclProxyMsgAbort = 7,
   ncclProxyMsgStop = 8,
   ncclProxyMsgGetFd = 9, // cuMem API support (UDS)
+  ncclProxyMsgRegister = 10,
+  ncclProxyMsgDeregister = 11
 };
 
 // This function is called by a client of the proxy that needs to invoke any of the non-progress proxyOp types

src/include/register.h

@@ -5,7 +5,8 @@ enum {
   NET_REG_COMPLETE = 0x01,
   NVLS_REG_COMPLETE = 0x02,
   NVLS_REG_POSSIBLE = 0x04,
-  NVLS_REG_NO_SUPPORT = 0x08
+  NVLS_REG_NO_SUPPORT = 0x08,
+  COLLNET_REG_COMPLETE = 0x10
 };
 
 struct ncclReg {
@@ -26,6 +27,9 @@ struct ncclReg {
   int dev;
   CUmemGenericAllocationHandle mcHandle;
   uintptr_t caddrs[NCCL_MAX_LOCAL_RANKS]; /* use to check if NVLS buffers match among intra-node ranks */
+  // collnet reg
+  void* collnetHandle;
+  struct ncclProxyConnector* proxyconn;
 };
 
 struct ncclRegCache {

src/include/socket.h

@@ -92,6 +92,7 @@ ncclResult_t ncclSocketProgress(int op, struct ncclSocket* sock, void* ptr, int
 ncclResult_t ncclSocketWait(int op, struct ncclSocket* sock, void* ptr, int size, int* offset);
 ncclResult_t ncclSocketSend(struct ncclSocket* sock, void* ptr, int size);
 ncclResult_t ncclSocketRecv(struct ncclSocket* sock, void* ptr, int size);
+ncclResult_t ncclSocketSendRecv(struct ncclSocket* sendSock, void* sendPtr, int sendSize, struct ncclSocket* recvSock, void* recvPtr, int recvSize);
 ncclResult_t ncclSocketTryRecv(struct ncclSocket* sock, void* ptr, int size, int* closed, bool blocking);
 ncclResult_t ncclSocketClose(struct ncclSocket* sock);
 #endif

src/include/transport.h

@@ -95,6 +95,8 @@ struct ncclTransportComm {
   ncclResult_t (*proxyConnect)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState, void* reqBuff, int reqSize, void* respBuff, int respSize, int* done);
   ncclResult_t (*proxyFree)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState);
   ncclResult_t (*proxyProgress)(struct ncclProxyState* proxyState, struct ncclProxyArgs*);
+  ncclResult_t (*proxyRegister)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState, void* reqBuff, int reqSize, void* respBuff, int respSize, int* done);
+  ncclResult_t (*proxyDeregister)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState, void* reqBuff, int reqSize, int* done);
 };
 
 struct ncclTransport {
@@ -107,15 +109,6 @@ struct ncclTransport {
 ncclResult_t ncclTransportP2pConnect(struct ncclComm* comm, int channelId, int nrecv, int* peerRecv, int nsend, int* peerSend, int connIndex);
 ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType=NULL);
 
-// Currently we only support POSIX_FILE_DESCRIPTOR handle exchange
-#define USE_POSIX_FD 1
-
-#if USE_POSIX_FD
-#define NVLS_CU_MEM_HANDLE_TYPE CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
-#else
-#define NVLS_CU_MEM_HANDLE_TYPE CU_MEM_HANDLE_TYPE_NONE
-#endif
-
 ncclResult_t ncclNvlsInit(struct ncclComm* comm);
 ncclResult_t ncclNvlsSetup(struct ncclComm* comm, struct ncclComm* parent);
 ncclResult_t ncclNvlsGraphRegisterBuffer(struct ncclComm *comm, struct ncclKernelPlan *plan, const void *sendbuff, void *recvbuff, size_t sendbuffSize, size_t recvbuffSize, bool *outRegBufUsed, void **outRegBufSend, void **outRegBufRecv);
@@ -124,7 +117,10 @@ ncclResult_t ncclNvlsDeregBuffer(CUmemGenericAllocationHandle *mcHandler, CUdevi
 ncclResult_t ncclNvlsFree(struct ncclComm* comm);
 
 enum { collNetRecv=0, collNetSend=1 };
-int ncclTransportCollNetSetup(struct ncclComm* comm, struct ncclTopoGraph* collNetGraph, struct ncclChannel* channel, int masterRank, int masterPeer, int collNetGraphChannelId, int type);
+int ncclTransportCollNetSetup(struct ncclComm* comm, struct ncclTopoGraph* collNetGraph, struct ncclChannel* channel, int masterRank, int masterPeer, int collNetGraphChannelId, int type, ncclConnect* connect);
 ncclResult_t ncclTransportCollNetCheck(struct ncclComm* comm, int collNetSetupFail);
 ncclResult_t ncclTransportCollNetFree(struct ncclComm* comm);
+ncclResult_t ncclCollnetLocalRegisterBuffer(struct ncclComm* comm, const void* userbuff, size_t buffSize, int type, int* outRegBufUsed, void** outHandle);
+ncclResult_t ncclCollnetGraphRegisterBuffer(struct ncclComm* comm, struct ncclKernelPlan *plan, const void* userbuff, size_t buffSize, int type, int* outRegBufFlag, void** outHandle);
+ncclResult_t ncclCollnetDeregBuffer(struct ncclComm* comm, struct ncclProxyConnector* proxyconn, void* handle);
 #endif

src/include/tuner.h

@@ -15,8 +15,8 @@
 // Attempts to load NCCL tuner from environmental variable.
 // Returns ncclSuccess if the correct tuner symbol has been found and
 // successully loaded.  Otherwise returns an error and also logs the error.
-ncclResult_t ncclLoadTunerPlugin(ncclTuner_t** tuner);
+ncclResult_t ncclTunerPluginLoad(ncclTuner_t** tuner);
 
 // Cleans up NCCL tuner plugin.
-ncclResult_t ncclCloseTunerPlugin(ncclTuner_t** tuner);
+ncclResult_t ncclTunerPluginUnload(ncclTuner_t** tuner);
 #endif