Dma buf support optin (#905) (#907)

* Dma buf support optin (#905) * dmaBufSupport Optin added on every part of the code that should invoke it * added to the original commit
ROCm · Oct 6, 2023 · 3d014cc · 3d014cc
1 parent f9012b6
commit 3d014cc
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 11 deletions.
diff --git a/src/init.cc b/src/init.cc
@@ -425,11 +425,10 @@ NCCL_PARAM(GdrCopyFifoEnable, "GDRCOPY_FIFO_ENABLE", 1);
 NCCL_PARAM(WorkFifoDepth, "WORK_FIFO_DEPTH", 64<<10);
 enum ncclLaunchMode ncclParamLaunchMode;
 
-NCCL_PARAM(DmaBufEnable, "DMABUF_ENABLE", 0);
 
 // Detect DMA-BUF support
 static ncclResult_t dmaBufSupported(struct ncclComm* comm) {
- if (ncclParamDmaBufEnable() == 0 || comm->ncclNet->regMrDmaBuf == NULL || rocmLibraryInit() != ncclSuccess) return ncclInternalError;
+ if (comm->ncclNet->regMrDmaBuf == NULL || rocmLibraryInit() != ncclSuccess) return ncclInternalError;
 #if CUDA_VERSION >= 11070
  int flag = 0;
  CUdevice dev;
@@ -1727,7 +1726,7 @@ constexpr nvtxPayloadSchemaEntry_t CommInitRankSchema[] = {
 NCCL_API(ncclResult_t, ncclCommInitRank, ncclComm_t* newcomm, int nranks, ncclUniqueId commId, int myrank);
 ncclResult_t ncclCommInitRank(ncclComm_t* newcomm, int nranks, ncclUniqueId commId, int myrank) {
  // Load the CUDA driver and dlsym hooks (can fail on old drivers)
- if (ncclParamDmaBufEnable()) rocmLibraryInit();
+ rocmLibraryInit();
 
  int cudaDev;
  ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
@@ -1743,7 +1742,7 @@ ncclResult_t ncclCommInitRank(ncclComm_t* newcomm, int nranks, ncclUniqueId comm
 NCCL_API(ncclResult_t, ncclCommInitRankMulti, ncclComm_t* newcomm, int nranks, ncclUniqueId commId, int myrank, int virtualId);
 ncclResult_t ncclCommInitRankMulti(ncclComm_t* newcomm, int nranks, ncclUniqueId commId, int myrank, int virtualId) {
  // Load the CUDA driver and dlsym hooks (can fail on old drivers)
- if (ncclParamDmaBufEnable()) rocmLibraryInit();
+ rocmLibraryInit();
 
  int cudaDev;
  ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
@@ -1770,7 +1769,7 @@ ncclResult_t ncclCommInitAll(ncclComm_t* comms, int ndev, const int* devlist) {
  NVTX3_FUNC_WITH_PARAMS(CommInitAll, CommInitAllSchema, ndev)
 
  // Load the CUDA driver and dlsym hooks (can fail on old drivers)
- if (ncclParamDmaBufEnable()) (void) rocmLibraryInit();
+ rocmLibraryInit();
 
  NCCLCHECKGOTO(PtrCheck(comms, "CommInitAll", "comms"), ret, fail);
  if (ndev < 0) {
@@ -1834,7 +1833,7 @@ ncclResult_t ncclCommInitRankConfig(ncclComm_t *newcomm, int nranks, ncclUniqueI
  ncclConfig_t *internalConfigPtr = NULL;
  NCCLCHECK(ncclGroupStartInternal());
 
- if (ncclParamDmaBufEnable()) (void) rocmLibraryInit();
+ rocmLibraryInit();
  CUDACHECKGOTO(cudaGetDevice(&cudaDev), ret, fail);
 
  if (config == NULL)

diff --git a/src/misc/rocmwrap.cc b/src/misc/rocmwrap.cc
@@ -9,6 +9,7 @@
 #include "debug.h"
 #include "rocmwrap.h"
 #include "hsa/hsa.h"
+#include "param.h"
 
 #include <dlfcn.h>
 #include <sys/utsname.h>
@@ -17,7 +18,7 @@
 #define DECLARE_ROCM_PFN(symbol) PFN_##symbol pfn_##symbol = nullptr
 
 DECLARE_ROCM_PFN(hsa_amd_portable_export_dmabuf); // DMA-BUF support
-
+NCCL_PARAM(DmaBufEnable, "DMABUF_ENABLE", 0);
 /* ROCr Driver functions loaded with dlsym() */
 DECLARE_ROCM_PFN(hsa_init);
 DECLARE_ROCM_PFN(hsa_system_get_info);
@@ -28,14 +29,14 @@ static enum { hsaUninitialized, hsaInitializing, hsaInitialized, hsaError } hsaS
 static void *hsaLib;
 static uint16_t version_major, version_minor;
 bool ncclCudaLaunchBlocking = false;
-bool dmaBufSupport = false;
 
 ncclResult_t rocmLibraryInit(void) {
  do {
  char* val = getenv("CUDA_LAUNCH_BLOCKING");
  ncclCudaLaunchBlocking = val!=nullptr && val[0]!=0 && !(val[0]=='0' && val[1]==0);
  } while (0);
 
+ bool dmaBufSupport = false;
  hsa_status_t res;
 
  if (hsaState == hsaInitialized)
@@ -108,14 +109,21 @@ ncclResult_t rocmLibraryInit(void) {
 
  /* DMA-BUF support */
  //ROCm support
+ if (ncclParamDmaBufEnable() == 0 ) {
+ INFO(NCCL_INIT, "Dmabuf feature disabled without NCCL_ENABLE_DMABUF_SUPPORT=1");
+ goto error;
+ }
  res = pfn_hsa_system_get_info((hsa_system_info_t) 0x204, &dmaBufSupport);
- if (res != HSA_STATUS_SUCCESS || !dmaBufSupport) INFO(NCCL_INIT, "Current version of ROCm does not support dmabuf feature.");
+ if (res != HSA_STATUS_SUCCESS || !dmaBufSupport) {
+ INFO(NCCL_INIT, "Current version of ROCm does not support dmabuf feature.");
+ goto error;
+ }
  else {
  pfn_hsa_amd_portable_export_dmabuf = (PFN_hsa_amd_portable_export_dmabuf) dlsym(hsaLib, "hsa_amd_portable_export_dmabuf");
  if (pfn_hsa_amd_portable_export_dmabuf == NULL) {
  WARN("Failed to load ROCr missing symbol hsa_amd_portable_export_dmabuf");
  goto error;
- } 
+ }
  else {
  //check OS kernel support
  struct utsname utsname;
@@ -126,7 +134,7 @@ ncclResult_t rocmLibraryInit(void) {
  char buf[256];
  int found_opt1 = 0;
  int found_opt2 = 0;
- 
+
  //check for kernel name exists
  if (uname(&utsname) == -1) INFO(NCCL_INIT,"Could not get kernel name");
  //format and store the kernel conf file location

diff --git a/src/transport/net_ib.cc b/src/transport/net_ib.cc
@@ -319,6 +319,7 @@ ncclResult_t ncclIbDmaBufSupport(int dev) {
  static int dmaBufSupported = -1;
  if (dmaBufSupported == -1) {
  ncclResult_t res;
+ NCCLCHECKGOTO(rocmLibraryInit(), res, failure);
  struct ibv_pd* pd;
  struct ibv_context* ctx;
  ctx = ncclIbDevs[dev].context;