From fb9938cffad4560f456f797285b110dc07cfa31c Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Wed, 30 Nov 2022 08:53:03 -0800
Subject: [PATCH] Query DMABuf support through HSA runtime API (#654)

---
 CMakeLists.txt       |  9 +++++++++
 src/misc/rocmwrap.cc | 14 ++++++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 8d5bef6525..24e08d983d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -102,6 +102,15 @@ find_package(hip REQUIRED)
 message(STATUS "HIP compiler: ${HIP_COMPILER}")
 message(STATUS "HIP runtime: ${HIP_RUNTIME}")
 
+find_package(hsa-runtime64 REQUIRED)
+get_target_property(HSA_INCLUDE_PATH hsa-runtime64::hsa-runtime64 INTERFACE_INCLUDE_DIRECTORIES)
+message(STATUS "HSA runtime: ${HSA_INCLUDE_PATH}")
+set(CMAKE_REQUIRED_INCLUDES ${HSA_INCLUDE_PATH})
+check_symbol_exists("HSA_AMD_SYSTEM_INFO_DMABUF_SUPPORTED" "hsa.h" HAS_HSA_AMD_SYSTEM_INFO_DMABUF_SUPPORTED)
+if(${HAS_HSA_AMD_SYSTEM_INFO_DMABUF_SUPPORTED})
+  add_definitions(-DQUERY_HSA_DMABUF)
+endif()
+
 if(BUILD_STATIC)
   option(BUILD_SHARED_LIBS "Build as a shared library" OFF)
 else()
diff --git a/src/misc/rocmwrap.cc b/src/misc/rocmwrap.cc
index f452f107cc..e2545ad2c9 100644
--- a/src/misc/rocmwrap.cc
+++ b/src/misc/rocmwrap.cc
@@ -96,11 +96,17 @@ ncclResult_t rocmLibraryInit(void) {
     //goto error;
   //}
 
-  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;
+#if defined(QUERY_HSA_DMABUF)
+  bool supported;
+  res = pfn_hsa_system_get_info(HSA_AMD_SYSTEM_INFO_DMABUF_SUPPORTED, &supported);
+  if (res == 0 && supported) {
+    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;
+    }
   }
+#endif
   /*
    * Required to initialize the ROCr Driver.
    * Multiple calls of hsa_init() will return immediately