From 220066197aac32a8fb5c26b3057bbb23d30bb4f5 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Wed, 17 Apr 2024 16:56:38 -0700
Subject: [PATCH] Use hipExtMallocWithFlags to allocate host memory on APU
 (#1149)

Also use SM60 as CUDA compatibility level.
---
 src/graph/topo.cc   |  5 +++++
 src/include/alloc.h | 11 ++++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/graph/topo.cc b/src/graph/topo.cc
index cf99ff558d..2e91919a31 100644
--- a/src/graph/topo.cc
+++ b/src/graph/topo.cc
@@ -371,7 +371,12 @@ ncclResult_t ncclTopoAddNic(struct ncclXmlNode* xmlNic, struct ncclTopoSystem* s
 }
 
 ncclResult_t ncclTopoAddGpu(struct ncclXmlNode* xmlGpu, struct ncclTopoSystem* system, struct ncclTopoNode* gpu) {
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+  // There is no direct mapping between CUDA SM to HIP GFX. Use SM60 as compatibility level.
+  gpu->gpu.cudaCompCap = 60;
+#else
   NCCLCHECK(xmlGetAttrInt(xmlGpu, "sm", &gpu->gpu.cudaCompCap));
+#endif
   const char* gcnArch;
   const char* gcnArchName;
   NCCLCHECK(xmlGetAttr(xmlGpu, "gcn", &gcnArch));
diff --git a/src/include/alloc.h b/src/include/alloc.h
index 288433d0fc..256ba11b94 100644
--- a/src/include/alloc.h
+++ b/src/include/alloc.h
@@ -27,7 +27,16 @@ ncclResult_t ncclCudaHostCallocDebug(T** ptr, size_t nelem, const char *filefunc
   cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
   *ptr = nullptr;
   CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
-  CUDACHECKGOTO(hipHostMalloc(ptr, nelem*sizeof(T), cudaHostAllocMapped), result, finish);
+  int managed = 0;
+  CUDACHECK(hipDeviceGetAttribute(&managed, hipDeviceAttributeDirectManagedMemAccessFromHost, 0));
+  if (managed) {
+#if defined(HIP_UNCACHED_MEMORY)
+    CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocUncached), result, finish);
+#else
+    CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
+#endif
+  } else
+    CUDACHECKGOTO(hipHostMalloc(ptr, nelem*sizeof(T), cudaHostAllocMapped), result, finish);
   memset(*ptr, 0, nelem*sizeof(T));
 finish:
   CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));