Support nccl-test with mscclpp-nccl on H100 GPUs

apps/nccl/include/nccl.h

@@ -69,6 +69,15 @@ typedef struct ncclConfig_v21700 {
         NCCL_CONFIG_UNDEF_INT                             /* splitShare */     \
   }
 
+/* NCCL malloc and free function for all types of NCCL optimizations
+ * (e.g. user buffer registration). The actual allocated size might
+ * be larger than requested due to granularity requirement. */
+ncclResult_t  ncclMemAlloc(void** ptr, size_t size);
+ncclResult_t pncclMemAlloc(void** ptr, size_t size);
+
+ncclResult_t  ncclMemFree(void *ptr);
+ncclResult_t pncclMemFree(void *ptr);
+
 /* Return the NCCL_VERSION_CODE of the NCCL library in the supplied integer.
  * This integer is coded with the MAJOR, MINOR and PATCH level of the
  * NCCL library

apps/nccl/src/nccl.cu

@@ -33,6 +33,9 @@
 //                             mscclpp::Transport::IB3, mscclpp::Transport::IB4, mscclpp::Transport::IB5,
 //                             mscclpp::Transport::IB6, mscclpp::Transport::IB7};
 
+// Declare the global map to store associations between raw pointer and shared pointer
+std::unordered_map<void*, std::shared_ptr<char>> ptrMap;
+
 struct channelKey {
   const void* buff;
   size_t bytes;
@@ -644,3 +647,38 @@ NCCL_API ncclResult_t ncclGroupEnd() {
   // Do nothing
   return ncclSuccess;
 }
+
+NCCL_API ncclResult_t ncclCommRegister(const ncclComm_t, void*, size_t, void**) {
+  //TODO: Implementation
+  return ncclSuccess;
+}
+
+NCCL_API ncclResult_t ncclCommDeregister(const ncclComm_t, void*) {
+  //TODO: Implementation
+  return ncclSuccess;
+}
+
+ncclResult_t ncclMemAlloc(void** ptr, size_t size) {
+  // Allocate memory using mscclpp::allocSharedPhysicalCuda
+  auto rawPtr = mscclpp::allocSharedPhysicalCuda<char>(size);
+  if (rawPtr == nullptr) {
+    return ncclInternalError;
+  }
+
+  ptrMap[rawPtr.get()] = rawPtr;
+
+  // Return the pointer
+  *ptr = rawPtr.get();
+  return ncclSuccess;
+}
+
+ncclResult_t ncclMemFree(void* ptr) {
+  auto ptrIt = ptrMap.find(ptr);
+  if (ptrIt != ptrMap.end()) {
+    ptrMap.erase(ptrIt);
+    return ncclSuccess;
+  }
+
+  // Pointer not found
+  return ncclInternalError;
+}