From 304bf368dfdcd75f07786f2ed0486c62e50af7f8 Mon Sep 17 00:00:00 2001
From: Chao <cchen104@amd.com>
Date: Mon, 29 Jul 2024 05:40:31 -0700
Subject: [PATCH] PR #15311: [ROCm] GPU/CPU unified memory for rocm

Imported from GitHub PR https://github.com/openxla/xla/pull/15311

@xla-rotation
Copybara import of the project:

--
2c4cee2bc335c72538261c41f485d49f1eb7c08f by Chao Chen <cchen104@amd.com>:

unified memory for rocm

Merging this change closes #15311

COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/15311 from ROCm:ci_rocm_unify_mem 2c4cee2bc335c72538261c41f485d49f1eb7c08f
PiperOrigin-RevId: 657168704
---
 xla/stream_executor/rocm/rocm_driver.cc       | 28 +++++++++++++++----
 .../rocm/rocm_driver_wrapper.h                |  1 +
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/xla/stream_executor/rocm/rocm_driver.cc b/xla/stream_executor/rocm/rocm_driver.cc
index ec4404a215a60..77e9884b47cb9 100644
--- a/xla/stream_executor/rocm/rocm_driver.cc
+++ b/xla/stream_executor/rocm/rocm_driver.cc
@@ -1338,16 +1338,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
index 2a8274b527eca..66c7dbb119ace 100644
--- 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)                            \