Dma buf support optin (#905)

* dmaBufSupport Optin added on every part of the code that should invoke it

[ROCm/rccl commit: 28d7fe5629]

projects/rccl/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,7 +29,6 @@ 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 {
@@ -36,6 +36,7 @@ ncclResult_t rocmLibraryInit(void) {
     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