Add support for AMD AINIC within RCCL default internal network plugin. (#2078)

* Added support for AMD ROCm net-ib alongside vanilla net-ib, with auto-generation to detect conflicts early during NCCL sync and enable future customizations.
* Integrated AMD AINIC support in RCCL for out-of-the-box usage, leveraging performance improvements by default, channel pinning for optimal pipeline performance, and extended support for 32B in-line CTS messages.
* Implemented internal derivation of AINIC-specific flags when RCCL AINIC environment parameter is set, and checks before initializing AINIC net-ib methods.
* Included snapshot of auto-generated ROCm net-ib file (src/transport/net_ib_rocm.cc) for reference.
* Fixed typos in RCCL param API (RCCL_AINIC_ROCE) and dlclose.
* Updated plugin loading logic:
* Load internal ROCmIB plugin only when NCCL_NET_PLUGIN is not set.
* Load default internal net-ib only when not AINIC and no external plugin env is set.

src/include/ionic/ionicdvcore.h

@@ -0,0 +1,20 @@
+#ifndef NCCL_IONICDV_CORE_H_
+#define NCCL_IONICDV_CORE_H_
+
+/* Basic ionic direct verbs structs.
+ * Needed to dynamically load ionic direct verbs functions without
+ * explicit including of ionic direct verbs header.
+ */
+
+#include <stddef.h>
+#include <stdint.h>
+#include <sys/types.h>
+#include <unistd.h>
+#include "ibvwrap.h"
+
+enum ionicdv_reg_udma_mask {
+    IONIC_UDMA_MASK_LOW    = 1,
+    IONIC_UDMA_MASK_HIGH   = 2
+};
+
+#endif  // NCCL_IONICDV_CORE_H_

src/include/ionic/ionicdvsymbols.h

@@ -0,0 +1,16 @@
+#ifndef NCCL_IONICDV_SYMBOLS_H_
+#define NCCL_IONICDV_SYMBOLS_H_
+
+#include "ionic/ionicdvcore.h"
+#include "nccl.h"
+
+/* Ionic Direct Verbs Function Pointers*/
+struct ncclIonicdvSymbols  {
+  int (*ionicdv_internal_qp_set_gda)(struct ibv_qp *qp, bool enable_send, bool enable_recv);
+  int (*ionicdv_internal_pd_set_udma_mask)(struct ibv_pd *ibpd, uint8_t udma_mask);
+};
+
+/* Constructs ionic direct verbs symbols per rdma-core linking or dynamic loading mode */
+ncclResult_t buildIonicdvSymbols(struct ncclIonicdvSymbols* ionicdvSymbols);
+
+#endif  // NCCL_IONICDV_SYMBOLS_H_

src/include/ionic/ionicdvwrap.h

@@ -0,0 +1,17 @@
+#ifndef NCCL_IONICDVWRAP_H_
+#define NCCL_IONICDVWRAP_H_
+
+#include <arpa/inet.h>
+#include <netinet/in.h>
+#include "ionic/ionicdvcore.h"
+#include "core.h"
+#include "ibvwrap.h"
+#include <sys/types.h>
+#include <unistd.h>
+
+ncclResult_t wrap_ionicdv_symbols(void);
+/* NCCL wrappers of ionic direct verbs functions */
+ncclResult_t wrap_ionicdv_qp_set_gda(struct ibv_qp *ibqp, bool enable_send, bool enable_recv);
+ncclResult_t wrap_ionicdv_pd_set_udma_mask(struct ibv_pd *ibpd, uint8_t udma_mask);
+
+#endif // NCCL_IONICDVWRAP_H_

src/include/net.h

@@ -25,4 +25,9 @@ extern ncclNet_t ncclNetSocket;
 
 extern ncclResult_t rcclNetP2pPolicy(void* handle, int isP2p);
 
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+extern ncclNet_t rocmNetIb;
+extern ncclResult_t rcclRocmNetP2pPolicy(void* handle, int isP2p);
+#endif
+
 #endif

src/misc/ionicdvsymbols.cc

@@ -0,0 +1,60 @@
+#include <sys/types.h>
+#include <unistd.h>
+
+#include "ionic/ionicdvsymbols.h"
+
+/* ionicdv dynamic loading mode. Symbols are loaded from shared objects. */
+#include <dlfcn.h>
+#include "core.h"
+
+// IONICDV Library versioning
+#define IONIC_VERSION "IONIC_1.0"
+
+ncclResult_t buildIonicdvSymbols(struct ncclIonicdvSymbols* ionicdvSymbols) {
+  static void* ionicdvhandle = NULL;
+  void* tmp;
+  void** cast;
+
+  ionicdvhandle = dlopen("libionic.so", RTLD_NOW);
+  if (!ionicdvhandle) {
+    ionicdvhandle = dlopen("libionic.so.1", RTLD_NOW);
+    if (!ionicdvhandle) {
+      INFO(NCCL_INIT, "Failed to open libionic.so[.1]");
+      goto teardown;
+    }
+  }
+
+#define LOAD_SYM(handle, symbol, funcptr) do {           \
+    cast = (void**)&funcptr;                             \
+    tmp = dlvsym(handle, symbol, IONIC_VERSION);       \
+    if (tmp == NULL) {                                   \
+      WARN("dlvsym failed on %s - %s version %s", symbol, dlerror(), IONIC_VERSION);  \
+      goto teardown;                                     \
+    } else {                                             \
+      WARN("dlvsym loaded successfully for %s - version %s", symbol, IONIC_VERSION);  \
+    }                                                    \
+    *cast = tmp;                                         \
+  } while (0)
+
+// Attempt to load a specific symbol version - fail silently
+#define LOAD_SYM_VERSION(handle, symbol, funcptr, version) do {  \
+    cast = (void**)&funcptr;                                     \
+    *cast = dlvsym(handle, symbol, version);                     \
+    if (*cast == NULL) {                                         \
+      INFO(NCCL_NET, "dlvsym failed on %s - %s version %s", symbol, dlerror(), version);  \
+    }                                                            \
+  } while (0)
+
+  LOAD_SYM(ionicdvhandle, "ionic_dv_qp_set_gda", ionicdvSymbols->ionicdv_internal_qp_set_gda);
+  LOAD_SYM(ionicdvhandle, "ionic_dv_pd_set_udma_mask", ionicdvSymbols->ionicdv_internal_pd_set_udma_mask);
+  INFO(NCCL_INIT, "Loaded dlvsym from libionic.so[.1]");
+
+  return ncclSuccess;
+
+teardown:
+  ionicdvSymbols->ionicdv_internal_qp_set_gda = NULL;
+  ionicdvSymbols->ionicdv_internal_pd_set_udma_mask = NULL;
+
+  if (ionicdvhandle != NULL) dlclose(ionicdvhandle);
+  return ncclSystemError;
+}

src/misc/ionicdvwrap.cc

@@ -0,0 +1,59 @@
+#include "ionic/ionicdvwrap.h"
+#include <sys/types.h>
+#include <unistd.h>
+#include "param.h"
+
+#include "ionic/ionicdvcore.h"
+#include "ionic/ionicdvsymbols.h"
+
+static pthread_once_t initOnceControl = PTHREAD_ONCE_INIT;
+static ncclResult_t initResult;
+struct ncclIonicdvSymbols ionicdvSymbols;
+
+extern int64_t rcclParamAinicRoce();
+
+ncclResult_t wrap_ionicdv_symbols(void) {
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+  if (rcclParamAinicRoce() == 1) {
+    pthread_once(&initOnceControl,
+                 [](){ initResult = buildIonicdvSymbols(&ionicdvSymbols); });
+    return initResult;
+  }
+#endif
+  // simply return for unsupported platform/NIC.
+  return ncclSuccess;
+}
+
+/* CHECK_NOT_NULL: helper macro to check for NULL symbol */
+#define CHECK_NOT_NULL(container, internal_name) \
+  if (container.internal_name == NULL) { \
+     WARN("lib wrapper not initialized."); \
+     return ncclInternalError; \
+  }
+
+#define IONICDV_INT_CHECK_RET_ERRNO(container, internal_name, call, success_retval, name) \
+  CHECK_NOT_NULL(container, internal_name); \
+  int ret = container.call; \
+  if (ret != success_retval) { \
+    INFO(NCCL_NET, "Call to " name " failed with error %s errno %d", strerror(ret), ret); \
+    return ncclSystemError; \
+  } else { \
+    INFO(NCCL_NET, "Call to " name " success with ret %d", ret); \
+  } \
+  return ncclSuccess;
+
+ncclResult_t wrap_ionicdv_qp_set_gda(struct ibv_qp *qp, bool enable_send, bool enable_recv) {
+  if (ionicdvSymbols.ionicdv_internal_qp_set_gda == NULL) {
+    errno = EOPNOTSUPP;
+    return ncclSystemError;
+  }
+  IONICDV_INT_CHECK_RET_ERRNO(ionicdvSymbols, ionicdv_internal_qp_set_gda, ionicdv_internal_qp_set_gda(qp, enable_send, enable_recv), 0, "ionic_dv_qp_set_gda");
+}
+
+ncclResult_t wrap_ionicdv_pd_set_udma_mask(struct ibv_pd *ibpd, uint8_t udma_mask) {
+  if (ionicdvSymbols.ionicdv_internal_pd_set_udma_mask == NULL) {
+    errno = EOPNOTSUPP;
+    return ncclSystemError;
+  }
+  IONICDV_INT_CHECK_RET_ERRNO(ionicdvSymbols, ionicdv_internal_pd_set_udma_mask, ionicdv_internal_pd_set_udma_mask(ibpd, udma_mask), 0, "ionic_dv_pd_set_udma_mask");
+}

src/plugin/net.cc

@@ -30,6 +30,8 @@ extern getNcclCollNet_t getNcclCollNet_v8;
 extern getNcclCollNet_t getNcclCollNet_v9;
 extern getNcclCollNet_t getNcclCollNet_v10;
 
+extern int64_t rcclParamAinicRoce();
+
 NCCL_PARAM(NetPluginRefCount, "NET_PLUGIN_REF_COUNT", 1);
 #define NCCL_NET_VERSION_COUNT 5
 int ncclNetVersion[NCCL_NET_VERSION_COUNT] = {10, 9, 8, 7, 6};
@@ -244,8 +246,18 @@ static void initPluginLibsOnceFunc() {
   }
 
   // Add 2 internal ib and socket plugins
-  netPluginLibs[pluginCounter].ncclNet = &ncclNetIb;
-  netPluginLibs[pluginCounter++].ncclNetPluginState = ncclNetPluginStateInitReady;
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+  if ((rcclParamAinicRoce() == 1) && !(envNetPlugin)) {
+    // For AINIC add rocm internal ib instead of default internal ib
+    netPluginLibs[pluginCounter].ncclNet = &rocmNetIb;
+    netPluginLibs[pluginCounter++].ncclNetPluginState = ncclNetPluginStateInitReady;
+  } else {
+#endif
+    netPluginLibs[pluginCounter].ncclNet = &ncclNetIb;
+    netPluginLibs[pluginCounter++].ncclNetPluginState = ncclNetPluginStateInitReady;
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+  }
+#endif
   netPluginLibs[pluginCounter].ncclNet = &ncclNetSocket;
   netPluginLibs[pluginCounter++].ncclNetPluginState = ncclNetPluginStateInitReady;
   pluginCount = pluginCounter;

src/transport/net.cc

@@ -29,8 +29,6 @@
 
 static_assert(sizeof(ncclNetHandle_t) <= CONNECT_SIZE, "NET Connect info is too large");
 
-#define RCCL_ANP_PLUGIN_STR  "RCCL-ANP"
-
 #define NCCL_NET_MAP_HOSTMEM 0
 #define NCCL_NET_MAP_DEVMEM 1
 #define NCCL_NET_MAP_SHARED_HOSTMEM 2
@@ -199,6 +197,7 @@ struct setupReq {
 };
 
 NCCL_PARAM(NetOptionalRecvCompletion, "NET_OPTIONAL_RECV_COMPLETION", 1);
+RCCL_PARAM(AinicRoce, "AINIC_ROCE", 0);
 
 static_assert(sizeof(ncclNetHandle_t) + sizeof(int) <= CONNECT_SIZE, "Not large enough ncclConnect to hold ncclNetHandle_t and useGdr flag");
 // Forward declaration
@@ -769,12 +768,12 @@ static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, str
   ncclNet_ctxt_t ncclNetCtxt = {};
   struct sendNetResources* resources = (struct sendNetResources*)(connection->transportResources);
   ncclNetCommConfig_t commConfig = {0};
+  bool rcclAinicRoce = ((rcclParamAinicRoce() == 1) ? true : false);
   if (reqSize != sizeof(netSendConnectArgs)) return ncclInternalError;
   ncclResult_t ret = ncclSuccess;
   netSendConnectArgs* req = (netSendConnectArgs*) reqBuff;
   commConfig.trafficClass = req->trafficClass == NCCL_CONFIG_UNDEF_INT ? NCCL_NET_TRAFFIC_CLASS_UNDEF : req->trafficClass;
   NCCLCHECK(ncclNetGetDeviceHandle(resources->netDeviceType, resources->netDeviceVersion, false /*isRecv*/, &resources->netDeviceHandle));
-  bool rccl_anp = !(strcmp(proxyState->ncclNet->name, RCCL_ANP_PLUGIN_STR));
   
   // Only call rcclNetP2pPolicy for ncclNetIb
   if (proxyState->ncclNet == &ncclNetIb) {
@@ -804,7 +803,7 @@ static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, str
         comms->activeConnect[resources->channelId] = (resources->tpLocalRank + 1);
       if (comms->sendComm[resources->channelId] == NULL
           && comms->activeConnect[resources->channelId] == (resources->tpLocalRank + 1)) {
-        if (rccl_anp) {
+        if (rcclAinicRoce) {
           ncclNetCtxt.chId = resources->channelId;
           ret = proxyState->ncclNet->connect(resources->netDev, &commConfig, req->handle,
             comms->sendComm + resources->channelId, (ncclNetDeviceHandle_t **)&ncclNetCtxt);
@@ -816,7 +815,7 @@ static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, str
       resources->netSendComm = comms->sendComm[resources->channelId];
       if (comms->sendComm[resources->channelId]) comms->sendRefCount[resources->channelId]++;
     } else {
-      if (rccl_anp) {
+      if (rcclAinicRoce) {
         ncclNetCtxt.chId = resources->channelId;
         ret = proxyState->ncclNet->connect(resources->netDev, &commConfig, req->handle, &resources->netSendComm, (ncclNetDeviceHandle_t **)&ncclNetCtxt);
       } else {
@@ -825,7 +824,7 @@ static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, str
     }
   } else {
     // Connect to remote peer
-    if (rccl_anp) {
+    if (rcclAinicRoce) {
       ncclNetCtxt.chId = resources->channelId;
       ret = proxyState->ncclNet->connect(resources->netDev, &commConfig, req->handle, &resources->netSendComm, (ncclNetDeviceHandle_t **)&ncclNetCtxt);
     } else {
@@ -979,8 +978,8 @@ static ncclResult_t recvProxyConnect(struct ncclProxyConnection* connection, str
   resources->tpRemoteProxyRank = req->proxyRank;
   ncclResult_t ret = ncclSuccess;
   ncclNet_ctxt_t ncclNetCtxt = {};
+  bool rcclAinicRoce = ((rcclParamAinicRoce() == 1) ? true : false);
 
-  bool rccl_anp = !(strcmp(proxyState->ncclNet->name, RCCL_ANP_PLUGIN_STR));
   NCCLCHECK(ncclNetGetDeviceHandle(resources->netDeviceType, resources->netDeviceVersion, true /*isRecv*/, &resources->netDeviceHandle));
   // Finish connection establishment from remote peer
   if (resources->shared) {
@@ -1007,7 +1006,7 @@ static ncclResult_t recvProxyConnect(struct ncclProxyConnection* connection, str
       //try connecting while comm is null
       if (comms->recvComm[resources->channelId] == NULL
          && comms->activeAccept[resources->channelId] == (resources->tpLocalRank + 1)) {
-        if (rccl_anp) {
+        if (rcclAinicRoce) {
           ncclNetCtxt.chId = resources->channelId;
           ret = proxyState->ncclNet->accept(resources->netListenComm,
             comms->recvComm+resources->channelId, (ncclNetDeviceHandle_t **)&ncclNetCtxt);
@@ -1019,7 +1018,7 @@ static ncclResult_t recvProxyConnect(struct ncclProxyConnection* connection, str
       resources->netRecvComm = comms->recvComm[resources->channelId];
       if (comms->recvComm[resources->channelId]) comms->recvRefCount[resources->channelId]++;
     } else {
-      if (rccl_anp) {
+      if (rcclAinicRoce) {
         ncclNetCtxt.chId = resources->channelId;
         ret = proxyState->ncclNet->accept(resources->netListenComm, &resources->netRecvComm, (ncclNetDeviceHandle_t **)&ncclNetCtxt);
       } else {
@@ -1028,7 +1027,7 @@ static ncclResult_t recvProxyConnect(struct ncclProxyConnection* connection, str
     }
   } else {
     // Connect to remote peer
-    if (rccl_anp) {
+    if (rcclAinicRoce) {
       ncclNetCtxt.chId = resources->channelId;
       ret = proxyState->ncclNet->accept(resources->netListenComm, &resources->netRecvComm, (ncclNetDeviceHandle_t **)&ncclNetCtxt);
     } else {