Set preferred location to the device for HIP Managed Memory

The semantics for specifying different kinds of advice is unclear so I set it in two stages.
nod-ai · Sep 23, 2023 · a03d9ae · a03d9ae
1 parent 25074d4
commit a03d9ae
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 0 deletions.
diff --git a/experimental/rocm/dynamic_symbol_tables.h b/experimental/rocm/dynamic_symbol_tables.h
@@ -24,6 +24,7 @@ RC_PFN_DECL(hipInit, unsigned int)
 RC_PFN_DECL(hipModuleLaunchKernel, hipFunction_t, unsigned int, unsigned int,
             unsigned int, unsigned int, unsigned int, unsigned int,
             unsigned int, hipStream_t, void **, void **)
+RC_PFN_DECL(hipMemAdvise, const void *, size_t, int, int)
 RC_PFN_DECL(hipMemset, void *, int, size_t)
 RC_PFN_DECL(hipMemsetAsync, void *, int, size_t, hipStream_t)
 RC_PFN_DECL(hipMemsetD32Async, void *, int, size_t, hipStream_t)

diff --git a/experimental/rocm/hip_headers.h b/experimental/rocm/hip_headers.h
@@ -234,6 +234,17 @@ typedef enum HIPmemAttach_flags_enum {
   HIP_MEM_ATTACH_SINGLE = 0x4,
 } HIPmemAttach_flags;
 
+typedef enum HIPmemoryAdvise_flags_enum {
+  hipMemAdviseSetReadMostly = 1,
+  hipMemAdviseUnsetReadMostly = 2,
+  hipMemAdviseSetPreferredLocation = 3,
+  hipMemAdviseUnsetPreferredLocation = 4,
+  hipMemAdviseSetAccessedBy = 5,
+  hipMemAdviseUnsetAccessedBy = 6,
+  hipMemAdviseSetCoarseGrain = 100,
+  hipMemAdviseUnsetCoarseGrain = 101,
+} HIPmemoryAdvise_flags;
+
 typedef enum HIPctx_flags_enum {
   hipDeviceScheduleAuto = 0x00,
   hipDeviceScheduleSpin = 0x01,

diff --git a/experimental/rocm/rocm_allocator.c b/experimental/rocm/rocm_allocator.c
@@ -285,6 +285,16 @@ static iree_status_t iree_hal_rocm_allocator_allocate_buffer(
       status = ROCM_RESULT_TO_STATUS(
           allocator->context->syms,
           hipMallocManaged(&device_ptr, allocation_size, hipMemAttachGlobal));
+      if (iree_status_is_ok(status)) {
+        status = ROCM_RESULT_TO_STATUS(
+            allocator->context->syms,
+            hipMemAdvise(device_ptr, allocation_size,
+                         hipMemAdviseSetPreferredLocation, allocator->device));
+        status = ROCM_RESULT_TO_STATUS(
+            allocator->context->syms,
+            hipMemAdvise(device_ptr, allocation_size,
+                         hipMemAdviseSetCoarseGrain, allocator->device));
+      }
       if (iree_status_is_ok(status) &&
           allocator->supports_concurrent_managed_access) {
         // Prefetch the buffer on the GPU device.