PR openxla#15311: [ROCm] GPU/CPU unified memory for rocm

Imported from GitHub PR openxla#15311 @xla-rotation Copybara import of the project: -- 2c4cee2 by Chao Chen <[email protected]>: unified memory for rocm Merging this change closes openxla#15311 COPYBARA_INTEGRATE_REVIEW=openxla#15311 from ROCm:ci_rocm_unify_mem 2c4cee2 PiperOrigin-RevId: 657168704
ROCm · Aug 5, 2024 · 4d88619 · 4d88619
1 parent dc89176
commit 4d88619
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/xla/stream_executor/rocm/rocm_driver.cc b/xla/stream_executor/rocm/rocm_driver.cc
@@ -1350,16 +1350,32 @@ struct BitPatternToValue {
 /* static */ void* GpuDriver::UnifiedMemoryAllocate(GpuContext* context,
                                                     uint64_t bytes) {
   ScopedActivateContext activated{context};
-
-  LOG(ERROR)
-      << "Feature not supported on ROCm platform (UnifiedMemoryAllocate)";
-  return nullptr;
+  hipDeviceptr_t result = 0;
+  // "managed" memory is visible to both CPU and GPU.
+  hipError_t res = wrap::hipMallocManaged(&result, bytes, hipMemAttachGlobal);
+  if (res != hipSuccess) {
+    LOG(ERROR) << "failed to alloc " << bytes
+               << " bytes unified memory; result: " << ToString(res);
+    return nullptr;
+  }
+  void* ptr = reinterpret_cast<void*>(result);
+  VLOG(2) << "allocated " << ptr << " for context " << context->context()
+          << " of " << bytes << " bytes in unified memory";
+  return ptr;
 }
 
 /* static */ void GpuDriver::UnifiedMemoryDeallocate(GpuContext* context,
                                                      void* location) {
-  LOG(ERROR)
-      << "Feature not supported on ROCm platform (UnifiedMemoryDeallocate)";
+  ScopedActivateContext activation(context);
+  hipDeviceptr_t pointer = absl::bit_cast<hipDeviceptr_t>(location);
+  hipError_t res = wrap::hipFree(pointer);
+  if (res != hipSuccess) {
+    LOG(ERROR) << "failed to free unified memory at " << location
+               << "; result: " << ToString(res);
+  } else {
+    VLOG(2) << "deallocated unified memory at " << location << " for context "
+            << context->context();
+  }
 }
 
 /* static */ void* GpuDriver::HostAllocate(GpuContext* context,

diff --git a/xla/stream_executor/rocm/rocm_driver_wrapper.h b/xla/stream_executor/rocm/rocm_driver_wrapper.h
@@ -134,6 +134,7 @@ namespace wrap {
   __macro(hipLaunchHostFunc)                        \
   __macro(hipLaunchKernel)                          \
   __macro(hipMalloc)                                \
+  __macro(hipMallocManaged)                         \
   __macro(hipMemGetAddressRange)                    \
   __macro(hipMemGetInfo)                            \
   __macro(hipMemcpyDtoD)                            \