SWDEV-546293 - hipMemPrefetchAsync_v2 and hipMemAdvise_v2 implementation (#869)

SWDEV-546293 - hipMemPrefetchAsync hipMemAdvise_v2

Please enter the commit message for your changes. Lines starting

CHANGELOG.md

@@ -23,6 +23,8 @@ Full documentation for HIP is available at [rocm.docs.amd.com](https://rocm.docs
     - `hipMemcpy3DPeerAsync`Copied memory between devices asynchronously
     - `hipMemsetD2D32Async` Used for setting 2D memory range with specified 32-bit values
       asynchronously
+    - `hipMemPrefetchAsync_v2`  prefetches memory to the specified location
+    - `hipMemAdvise_v2`         advise about the usage of a given memory range
 * Changed HIP APIs
     - `hipMemCreate`  now can take hipDeviceMallocUncached as a flag to allocate uncached memory
 
@@ -175,9 +177,9 @@ HIP runtime has the following functional improvements which greatly improve runt
 * Refactored memory validation, creates a unique function to validate a variety of memory copy operations.
 * Improved kernel logging using demangling shader names.
 * Advanced support for SPIRV, now kernel compilation caching is enabled by default. This feature is controlled by the environment variable `AMD_COMGR_CACHE`, for details, see [hip_rtc document](https://rocm.docs.amd.com/projects/HIP/en/latest/how-to/hip_rtc.html).
-* Programmatic support for scratch limits on MI300 and MI350 series up GPU devices. More enumeration values were added in `hipLimit_t` as following, 
-   - `hipExtLimitScratchMin`, minimum allowed value in bytes for scratch limit on the device. 
-   - `hipExtLimitScratchMax`, maximum allowed value in bytes for scratch limit on the device. 
+* Programmatic support for scratch limits on MI300 and MI350 series up GPU devices. More enumeration values were added in `hipLimit_t` as following,
+   - `hipExtLimitScratchMin`, minimum allowed value in bytes for scratch limit on the device.
+   - `hipExtLimitScratchMax`, maximum allowed value in bytes for scratch limit on the device.
    - `hipExtLimitScratchCurrent`, current scratch limit threshold in bytes on the device. Must be between the value `hipExtLimitScratchMin` and `hipExtLimitScratchMax`.
  Developers can now use the environment variable `HSA_SCRATCH_SINGLE_LIMIT_ASYNC` to change the default allocation size with expected scratch limit in ROCR runtime. On top of it, this value can also be overwritten programmatically in the application using the HIP API `hipDeviceSetLimit(hipExtLimitScratchCurrent, value)` to reset the scratch limit value.
 * HIP runtime now enables peer-to-peer (P2P) memory copies to utilize all available SDMA engines, rather than being limited to a single engine. It also selects the best engine first to give optimal bandwidth.
@@ -219,7 +221,7 @@ HIP runtime has the following functional improvements which greatly improve runt
 * The memory leak in virtual memory management (VMM). HIP runtime now uses the size of handle for allocated memory range instead of actual size for physical memory, which fixed the issue of address clash with VMM.
 * Large memory allocation issue. HIP runtime now checks GPU video RAM and system RAM properly and sets size limits during memory allocation either on the host or the GPU device.
 * Support of `hipDeviceMallocContiguous` flags in `hipExtMallocWithFlags()`. It now enables `HSA_AMD_MEMORY_POOL_CONTIGUOUS_FLAG` in the memory pool allocation on GPU device.
-* Radom memory segmentation fault in handling `GraphExec` object release and `hipDeviceSyncronization`. HIP runtime now uses internal device synchronize function in `__hipUnregisterFatBinary`. 
+* Radom memory segmentation fault in handling `GraphExec` object release and `hipDeviceSyncronization`. HIP runtime now uses internal device synchronize function in `__hipUnregisterFatBinary`.
 
 ## HIP 6.4.1 for ROCm 6.4.1
 
@@ -309,7 +311,7 @@ The following are the list of backwards incompatible changes planned for the upc
     - `hipModuleLoad`
     - `hipLaunchCooperativeKernelMultiDevice`
     - `hipExtLaunchCooperativeKernelMultiDevice`
- 
+
 * HIPRTC implementation, the compilation of hiprtc now uses  namespace ` __hip_internal`, instead of the standard headers `std`.
 * Stream capture mode update in the following hip APIs. Stream can only be captured in relax mode, to match the behavior of the corresponding CUDA APIs,
    - `hipMallocManaged`

hipamd/include/hip/amd_detail/hip_api_trace.hpp

@@ -484,6 +484,8 @@ typedef hipError_t (*t_hipMemAddressReserve)(void** ptr, size_t size, size_t ali
                                              unsigned long long flags);
 typedef hipError_t (*t_hipMemAdvise)(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
                                      int device);
+typedef hipError_t (*t_hipMemAdvise_v2)(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
+                                        hipMemLocation device);
 typedef hipError_t (*t_hipMemAllocHost)(void** ptr, size_t size);
 typedef hipError_t (*t_hipMemAllocPitch)(hipDeviceptr_t* dptr, size_t* pitch, size_t widthInBytes,
                                          size_t height, unsigned int elementSizeBytes);
@@ -535,6 +537,9 @@ typedef hipError_t (*t_hipMemPoolSetAttribute)(hipMemPool_t mem_pool, hipMemPool
 typedef hipError_t (*t_hipMemPoolTrimTo)(hipMemPool_t mem_pool, size_t min_bytes_to_hold);
 typedef hipError_t (*t_hipMemPrefetchAsync)(const void* dev_ptr, size_t count, int device,
                                             hipStream_t stream);
+typedef hipError_t (*t_hipMemPrefetchAsync_v2)(const void* dev_ptr, size_t count,
+                                               hipMemLocation location, unsigned int flags,
+                                               hipStream_t stream);
 typedef hipError_t (*t_hipMemPtrGetInfo)(void* ptr, size_t* size);
 typedef hipError_t (*t_hipMemRangeGetAttribute)(void* data, size_t data_size,
                                                 hipMemRangeAttribute attribute, const void* dev_ptr,
@@ -1638,7 +1643,8 @@ struct HipDispatchTable {
   t_hipMemcpy3DPeerAsync hipMemcpy3DPeerAsync_fn;
   t_hipGetDriverEntryPoint hipGetDriverEntryPoint_fn;
   t_hipGetDriverEntryPoint_spt hipGetDriverEntryPoint_spt_fn;
-
+  t_hipMemPrefetchAsync_v2 hipMemPrefetchAsync_v2_fn;
+  t_hipMemAdvise_v2 hipMemAdvise_v2_fn;
   // HIP_RUNTIME_API_TABLE_STEP_VERSION = 14
   // removed HIP_MEMSET_NODE_PARAMS replaced by hipMemsetParams
 

hipamd/include/hip/amd_detail/hip_prof_str.h

@@ -453,7 +453,9 @@ enum hip_api_id_t {
   HIP_API_ID_hipMemcpy3DPeerAsync = 433,
   HIP_API_ID_hipMemcpyBatchAsync = 434,
   HIP_API_ID_hipGetDriverEntryPoint = 435,
-  HIP_API_ID_LAST = 435,
+  HIP_API_ID_hipMemPrefetchAsync_v2 = 436,
+  HIP_API_ID_hipMemAdvise_v2 = 437,
+  HIP_API_ID_LAST = 437,
 
   HIP_API_ID_hipChooseDevice = HIP_API_ID_CONCAT(HIP_API_ID_,hipChooseDevice),
   HIP_API_ID_hipGetDeviceProperties = HIP_API_ID_CONCAT(HIP_API_ID_,hipGetDeviceProperties),
@@ -744,6 +746,7 @@ static inline const char* hip_api_name(const uint32_t id) {
     case HIP_API_ID_hipMemAddressFree: return "hipMemAddressFree";
     case HIP_API_ID_hipMemAddressReserve: return "hipMemAddressReserve";
     case HIP_API_ID_hipMemAdvise: return "hipMemAdvise";
+    case HIP_API_ID_hipMemAdvise_v2: return "hipMemAdvise_v2";
     case HIP_API_ID_hipMemAllocHost: return "hipMemAllocHost";
     case HIP_API_ID_hipMemAllocPitch: return "hipMemAllocPitch";
     case HIP_API_ID_hipMemCreate: return "hipMemCreate";
@@ -768,6 +771,7 @@ static inline const char* hip_api_name(const uint32_t id) {
     case HIP_API_ID_hipMemPoolSetAttribute: return "hipMemPoolSetAttribute";
     case HIP_API_ID_hipMemPoolTrimTo: return "hipMemPoolTrimTo";
     case HIP_API_ID_hipMemPrefetchAsync: return "hipMemPrefetchAsync";
+    case HIP_API_ID_hipMemPrefetchAsync_v2: return "hipMemPrefetchAsync_v2";
     case HIP_API_ID_hipMemPtrGetInfo: return "hipMemPtrGetInfo";
     case HIP_API_ID_hipMemRangeGetAttribute: return "hipMemRangeGetAttribute";
     case HIP_API_ID_hipMemRangeGetAttributes: return "hipMemRangeGetAttributes";
@@ -1174,6 +1178,7 @@ static inline uint32_t hipApiIdByName(const char* name) {
   if (strcmp("hipMemAddressFree", name) == 0) return HIP_API_ID_hipMemAddressFree;
   if (strcmp("hipMemAddressReserve", name) == 0) return HIP_API_ID_hipMemAddressReserve;
   if (strcmp("hipMemAdvise", name) == 0) return HIP_API_ID_hipMemAdvise;
+  if (strcmp("hipMemAdvise_v2", name) == 0) return HIP_API_ID_hipMemAdvise_v2;
   if (strcmp("hipMemAllocHost", name) == 0) return HIP_API_ID_hipMemAllocHost;
   if (strcmp("hipMemAllocPitch", name) == 0) return HIP_API_ID_hipMemAllocPitch;
   if (strcmp("hipMemCreate", name) == 0) return HIP_API_ID_hipMemCreate;
@@ -1198,6 +1203,7 @@ static inline uint32_t hipApiIdByName(const char* name) {
   if (strcmp("hipMemPoolSetAttribute", name) == 0) return HIP_API_ID_hipMemPoolSetAttribute;
   if (strcmp("hipMemPoolTrimTo", name) == 0) return HIP_API_ID_hipMemPoolTrimTo;
   if (strcmp("hipMemPrefetchAsync", name) == 0) return HIP_API_ID_hipMemPrefetchAsync;
+  if (strcmp("hipMemPrefetchAsync_v2", name) == 0) return HIP_API_ID_hipMemPrefetchAsync_v2;
   if (strcmp("hipMemPtrGetInfo", name) == 0) return HIP_API_ID_hipMemPtrGetInfo;
   if (strcmp("hipMemRangeGetAttribute", name) == 0) return HIP_API_ID_hipMemRangeGetAttribute;
   if (strcmp("hipMemRangeGetAttributes", name) == 0) return HIP_API_ID_hipMemRangeGetAttributes;
@@ -2814,6 +2820,12 @@ typedef struct hip_api_data_s {
       hipMemoryAdvise advice;
       int device;
     } hipMemAdvise;
+    struct {
+      const void* dev_ptr;
+      size_t count;
+      hipMemoryAdvise advice;
+      hipMemLocation location;
+    } hipMemAdvise_v2;
     struct {
       void** ptr;
       void* ptr__val;
@@ -2960,6 +2972,13 @@ typedef struct hip_api_data_s {
       int device;
       hipStream_t stream;
     } hipMemPrefetchAsync;
+    struct {
+      const void* dev_ptr;
+      size_t count;
+      hipMemLocation location;
+      unsigned int flags;
+      hipStream_t stream;
+    } hipMemPrefetchAsync_v2;
     struct {
       void* ptr;
       size_t* size;
@@ -5370,6 +5389,13 @@ typedef struct hip_api_data_s {
   cb_data.args.hipMemAdvise.advice = (hipMemoryAdvise)advice; \
   cb_data.args.hipMemAdvise.device = (int)device; \
 };
+// hipMemAdvise_v2[('const void*', 'dev_ptr'), ('size_t', 'count'), ('hipMemoryAdvise', 'advice'), ('hipMemLocation', 'location')]
+#define INIT_hipMemAdvise_v2_CB_ARGS_DATA(cb_data) { \
+  cb_data.args.hipMemAdvise_v2.dev_ptr = (const void*)dev_ptr; \
+  cb_data.args.hipMemAdvise_v2.count = (size_t)count; \
+  cb_data.args.hipMemAdvise_v2.advice = (hipMemoryAdvise)advice; \
+  cb_data.args.hipMemAdvise_v2.location = (hipMemLocation)location; \
+};
 // hipMemAllocHost[('void**', 'ptr'), ('size_t', 'size')]
 #define INIT_hipMemAllocHost_CB_ARGS_DATA(cb_data) { \
   cb_data.args.hipMemAllocHost.ptr = (void**)ptr; \
@@ -5515,6 +5541,14 @@ typedef struct hip_api_data_s {
   cb_data.args.hipMemPrefetchAsync.device = (int)device; \
   cb_data.args.hipMemPrefetchAsync.stream = (hipStream_t)stream; \
 };
+// hipMemPrefetchAsync_v2[('const void*', 'dev_ptr'), ('size_t', 'count'), ('hipMemLocation', 'location'), ('unsigned int', 'flags'), ('hipStream_t', 'stream')]
+#define INIT_hipMemPrefetchAsync_v2_CB_ARGS_DATA(cb_data) { \
+  cb_data.args.hipMemPrefetchAsync_v2.dev_ptr = (const void*)dev_ptr; \
+  cb_data.args.hipMemPrefetchAsync_v2.count = (size_t)count; \
+  cb_data.args.hipMemPrefetchAsync_v2.location = (hipMemLocation)location; \
+  cb_data.args.hipMemPrefetchAsync_v2.flags = (unsigned int)flags; \
+  cb_data.args.hipMemPrefetchAsync_v2.stream = (hipStream_t)stream; \
+};
 // hipMemPtrGetInfo[('void*', 'ptr'), ('size_t*', 'size')]
 #define INIT_hipMemPtrGetInfo_CB_ARGS_DATA(cb_data) { \
   cb_data.args.hipMemPtrGetInfo.ptr = (void*)ptr; \
@@ -7575,6 +7609,9 @@ static inline void hipApiArgsInit(hip_api_id_t id, hip_api_data_t* data) {
 // hipMemAdvise[('const void*', 'dev_ptr'), ('size_t', 'count'), ('hipMemoryAdvise', 'advice'), ('int', 'device')]
     case HIP_API_ID_hipMemAdvise:
       break;
+// hipMemAdvise_v2[('const void*', 'dev_ptr'), ('size_t', 'count'), ('hipMemoryAdvise', 'advice'), ('hipMemLocation', 'location')]
+    case HIP_API_ID_hipMemAdvise_v2:
+      break;
 // hipMemAllocHost[('void**', 'ptr'), ('size_t', 'size')]
     case HIP_API_ID_hipMemAllocHost:
       if (data->args.hipMemAllocHost.ptr) data->args.hipMemAllocHost.ptr__val = *(data->args.hipMemAllocHost.ptr);
@@ -7672,6 +7709,9 @@ static inline void hipApiArgsInit(hip_api_id_t id, hip_api_data_t* data) {
 // hipMemPrefetchAsync[('const void*', 'dev_ptr'), ('size_t', 'count'), ('int', 'device'), ('hipStream_t', 'stream')]
     case HIP_API_ID_hipMemPrefetchAsync:
       break;
+// hipMemPrefetchAsync_v2[('const void*', 'dev_ptr'), ('size_t', 'count'), ('hipMemLocation', 'location'), ('unsigned int', 'flags'), ('hipStream_t', 'stream')]
+    case HIP_API_ID_hipMemPrefetchAsync_v2:
+      break;
 // hipMemPtrGetInfo[('void*', 'ptr'), ('size_t*', 'size')]
     case HIP_API_ID_hipMemPtrGetInfo:
       if (data->args.hipMemPtrGetInfo.size) data->args.hipMemPtrGetInfo.size__val = *(data->args.hipMemPtrGetInfo.size);
@@ -10216,6 +10256,14 @@ static inline const char* hipApiString(hip_api_id_t id, const hip_api_data_t* da
       oss << ", device="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemAdvise.device);
       oss << ")";
     break;
+    case HIP_API_ID_hipMemAdvise_v2:
+      oss << "hipMemAdvise_v2(";
+      oss << "dev_ptr="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemAdvise_v2.dev_ptr);
+      oss << ", count="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemAdvise_v2.count);
+      oss << ", advice="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemAdvise_v2.advice);
+      oss << ", location="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemAdvise_v2.location);
+      oss << ")";
+    break;
     case HIP_API_ID_hipMemAllocHost:
       oss << "hipMemAllocHost(";
       if (data->args.hipMemAllocHost.ptr == NULL) oss << "ptr=NULL";
@@ -10410,6 +10458,15 @@ static inline const char* hipApiString(hip_api_id_t id, const hip_api_data_t* da
       oss << ", stream="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemPrefetchAsync.stream);
       oss << ")";
     break;
+    case HIP_API_ID_hipMemPrefetchAsync_v2:
+      oss << "hipMemPrefetchAsync_v2(";
+      oss << "dev_ptr="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemPrefetchAsync_v2.dev_ptr);
+      oss << ", count="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemPrefetchAsync_v2.count);
+      oss << ", location="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemPrefetchAsync_v2.location);
+      oss << ", flags="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemPrefetchAsync_v2.flags);
+      oss << ", stream="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemPrefetchAsync_v2.stream);
+      oss << ")";
+    break;
     case HIP_API_ID_hipMemPtrGetInfo:
       oss << "hipMemPtrGetInfo(";
       oss << "ptr="; roctracer::hip_support::detail::operator<<(oss, data->args.hipMemPtrGetInfo.ptr);

hipamd/src/amdhip.def

@@ -508,3 +508,5 @@ hipMemcpy3DPeer
 hipMemcpy3DPeerAsync
 hipGetDriverEntryPoint
 hipGetDriverEntryPoint_spt
+hipMemPrefetchAsync_v2
+hipMemAdvise_v2

hipamd/src/hip_api_trace.cpp

@@ -400,6 +400,8 @@ hipError_t hipMemAddressFree(void* devPtr, size_t size);
 hipError_t hipMemAddressReserve(void** ptr, size_t size, size_t alignment, void* addr,
                                 unsigned long long flags);
 hipError_t hipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advice, int device);
+hipError_t hipMemAdvise_v2(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
+                           hipMemLocation location);
 hipError_t hipMemAllocHost(void** ptr, size_t size);
 hipError_t hipMemAllocPitch(hipDeviceptr_t* dptr, size_t* pitch, size_t widthInBytes, size_t height,
                             unsigned int elementSizeBytes);
@@ -441,6 +443,8 @@ hipError_t hipMemPoolSetAccess(hipMemPool_t mem_pool, const hipMemAccessDesc* de
 hipError_t hipMemPoolSetAttribute(hipMemPool_t mem_pool, hipMemPoolAttr attr, void* value);
 hipError_t hipMemPoolTrimTo(hipMemPool_t mem_pool, size_t min_bytes_to_hold);
 hipError_t hipMemPrefetchAsync(const void* dev_ptr, size_t count, int device, hipStream_t stream);
+hipError_t hipMemPrefetchAsync_v2(const void* dev_ptr, size_t count, hipMemLocation location,
+                                  unsigned int flags, hipStream_t stream);
 hipError_t hipMemPtrGetInfo(void* ptr, size_t* size);
 hipError_t hipMemRangeGetAttribute(void* data, size_t data_size, hipMemRangeAttribute attribute,
                                    const void* dev_ptr, size_t count);
@@ -1124,6 +1128,7 @@ void UpdateDispatchTable(HipDispatchTable* ptrDispatchTable) {
   ptrDispatchTable->hipMemAddressFree_fn = hip::hipMemAddressFree;
   ptrDispatchTable->hipMemAddressReserve_fn = hip::hipMemAddressReserve;
   ptrDispatchTable->hipMemAdvise_fn = hip::hipMemAdvise;
+  ptrDispatchTable->hipMemAdvise_v2_fn = hip::hipMemAdvise_v2;
   ptrDispatchTable->hipMemAllocHost_fn = hip::hipMemAllocHost;
   ptrDispatchTable->hipMemAllocPitch_fn = hip::hipMemAllocPitch;
   ptrDispatchTable->hipMemCreate_fn = hip::hipMemCreate;
@@ -1150,6 +1155,7 @@ void UpdateDispatchTable(HipDispatchTable* ptrDispatchTable) {
   ptrDispatchTable->hipMemPoolSetAttribute_fn = hip::hipMemPoolSetAttribute;
   ptrDispatchTable->hipMemPoolTrimTo_fn = hip::hipMemPoolTrimTo;
   ptrDispatchTable->hipMemPrefetchAsync_fn = hip::hipMemPrefetchAsync;
+  ptrDispatchTable->hipMemPrefetchAsync_v2_fn = hip::hipMemPrefetchAsync_v2;
   ptrDispatchTable->hipMemPtrGetInfo_fn = hip::hipMemPtrGetInfo;
   ptrDispatchTable->hipMemRangeGetAttribute_fn = hip::hipMemRangeGetAttribute;
   ptrDispatchTable->hipMemRangeGetAttributes_fn = hip::hipMemRangeGetAttributes;
@@ -2050,13 +2056,15 @@ HIP_ENFORCE_ABI(HipDispatchTable, hipMemcpy3DPeer_fn, 489);
 HIP_ENFORCE_ABI(HipDispatchTable, hipMemcpy3DPeerAsync_fn, 490);
 HIP_ENFORCE_ABI(HipDispatchTable, hipGetDriverEntryPoint_fn, 491);
 HIP_ENFORCE_ABI(HipDispatchTable, hipGetDriverEntryPoint_spt_fn, 492);
+HIP_ENFORCE_ABI(HipDispatchTable, hipMemPrefetchAsync_v2_fn, 493);
+HIP_ENFORCE_ABI(HipDispatchTable, hipMemAdvise_v2_fn, 494);
 // if HIP_ENFORCE_ABI entries are added for each new function pointer in the table, the number below
 // will be +1 of the number in the last HIP_ENFORCE_ABI line. E.g.:
 //
 //  HIP_ENFORCE_ABI(<table>, <functor>, 8)
 //
 //  HIP_ENFORCE_ABI_VERSIONING(<table>, 9) <- 8 + 1 = 9
-HIP_ENFORCE_ABI_VERSIONING(HipDispatchTable, 493)
+HIP_ENFORCE_ABI_VERSIONING(HipDispatchTable, 495)
 
 static_assert(HIP_RUNTIME_API_TABLE_MAJOR_VERSION == 0 && HIP_RUNTIME_API_TABLE_STEP_VERSION == 14,
               "If you get this error, add new HIP_ENFORCE_ABI(...) code for the new function "

hipamd/src/hip_formatting.hpp

@@ -937,3 +937,29 @@ inline std::ostream& operator<<(std::ostream& os, const hipPitchedPtr* p) {
   }
   return os;
 }
+
+inline std::ostream& operator<<(std::ostream& os, const hipMemLocation& s) {
+  os << "{type=";
+  switch (s.type) {
+    case hipMemLocationTypeInvalid:
+      os << "hipMemLocationTypeInvalid";
+      break;
+    case hipMemLocationTypeDevice:
+      os << "hipMemLocationTypeDevice";
+      break;
+    case hipMemLocationTypeHost:
+      os << "hipMemLocationTypeHost";
+      break;
+    case hipMemLocationTypeHostNuma:
+      os << "hipMemLocationTypeHostNuma";
+      break;
+    case hipMemLocationTypeHostNumaCurrent:
+      os << "hipMemLocationTypeHostNumaCurrent";
+      break;
+    default:
+      os << static_cast<int>(s.type);
+      break;
+  }
+  os << ", id=" << s.id << "}";
+  return os;
+}

hipamd/src/hip_hcc.map.in

@@ -624,6 +624,8 @@ global:
     hipMemcpy3DPeerAsync;
     hipGetDriverEntryPoint;
     hipGetDriverEntryPoint_spt;
+    hipMemPrefetchAsync_v2;
+    hipMemAdvise_v2;
 local:
     *;
 } hip_6.5;

hipamd/src/hip_hmm.cpp

@@ -29,6 +29,10 @@ namespace hip {
 
 // Forward declaraiton of a function
 hipError_t ihipMallocManaged(void** ptr, size_t size, size_t align = 0, bool use_host_ptr = 0);
+hipError_t ihipMemPrefetchAsync(const void* dev_ptr, size_t count, hipMemLocation location,
+                                hipStream_t stream);
+hipError_t ihipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
+                         hipMemLocation location);
 
 // Make sure HIP defines match ROCclr to avoid double conversion
 static_assert(hipCpuDeviceId == amd::CpuDeviceId, "CPU device ID mismatch with ROCclr!");
@@ -83,98 +87,51 @@ hipError_t hipMallocManaged(void** dev_ptr, size_t size, unsigned int flags) {
 hipError_t hipMemPrefetchAsync(const void* dev_ptr, size_t count, int device,
                                hipStream_t stream) {
   HIP_INIT_API(hipMemPrefetchAsync, dev_ptr, count, device, stream);
-
-  if ((dev_ptr == nullptr) || (count == 0)) {
-    HIP_RETURN(hipErrorInvalidValue);
-  }
-
-  getStreamPerThread(stream);
-
-  size_t offset = 0;
-  amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
-
-  if ((memObj != nullptr) && (count  > (memObj->getSize() - offset))) {
-    HIP_RETURN(hipErrorInvalidValue);
-  }
-  if (device != hipCpuDeviceId && (static_cast<size_t>(device) >= g_devices.size())) {
-    HIP_RETURN(hipErrorInvalidDevice);
-  }
-
-  hip::Stream* hip_stream = nullptr;
-  amd::Device* dev = nullptr;
-  bool cpu_access = false;
-
-  if ((memObj == nullptr) && (device != hipCpuDeviceId) &&
-      (!g_devices[device]->devices()[0]->info().hmmCpuMemoryAccessible_)) {
-    HIP_RETURN(hipErrorNotSupported);
-  }
-
-  // Pick the specified stream or Null one from the provided device
+  CHECK_STREAM_CAPTURE_SUPPORTED();
+  hipMemLocation location;
   if (device == hipCpuDeviceId) {
-    cpu_access = true;
-    hip_stream = (stream == nullptr || stream == hipStreamLegacy) ?
-                  hip::getCurrentDevice()->NullStream() : hip::getStream(stream);
+    location.type = hipMemLocationTypeHost;
+    location.id = hipCpuDeviceId;
   } else {
-    dev = g_devices[device]->devices()[0];
-    hip_stream = (stream == nullptr || stream == hipStreamLegacy) ?
-                  g_devices[device]->NullStream() : hip::getStream(stream);
+    location.type = hipMemLocationTypeDevice;
+    location.id = device;
   }
+  HIP_RETURN(ihipMemPrefetchAsync(dev_ptr, count, location, stream));
+}
 
-  if (hip_stream == nullptr) {
+// ================================================================================================
+hipError_t hipMemPrefetchAsync_v2(const void* dev_ptr, size_t count, hipMemLocation location,
+                                  unsigned int flags, hipStream_t stream) {
+  HIP_INIT_API(hipMemPrefetchAsync_v2, dev_ptr, count, location, flags, stream);
+  CHECK_STREAM_CAPTURE_SUPPORTED();
+  if (flags != 0) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-
-  amd::Command::EventWaitList waitList;
-  amd::SvmPrefetchAsyncCommand* command =
-      new amd::SvmPrefetchAsyncCommand(*hip_stream, waitList, dev_ptr, count, dev, cpu_access);
-  if (command == nullptr) {
-    return hipErrorOutOfMemory;
-  }
-
-  command->enqueue();
-  command->release();
-
-  HIP_RETURN(hipSuccess);
+  HIP_RETURN(ihipMemPrefetchAsync(dev_ptr, count, location, stream));
 }
 
 // ================================================================================================
 hipError_t hipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advice, int device) {
   HIP_INIT_API(hipMemAdvise, dev_ptr, count, advice, device);
-
   CHECK_STREAM_CAPTURE_SUPPORTED();
-
-  bool isAdviseReadMostly = (advice == hipMemAdviseSetReadMostly) ||
-                            (advice == hipMemAdviseUnsetReadMostly);
-
-  if (!isAdviseReadMostly && ((device != hipCpuDeviceId) &&
-      (static_cast<size_t>(device) >= g_devices.size()))) {
-    HIP_RETURN(hipErrorInvalidDevice);
+  hipMemLocation location;
+  if (device == hipCpuDeviceId) {
+    location.type = hipMemLocationTypeHost;
+    location.id = hipCpuDeviceId;
+  } else {
+    location.type = hipMemLocationTypeDevice;
+    location.id = device;
   }
 
-  if ((dev_ptr == nullptr) || (count == 0)) {
-    HIP_RETURN(hipErrorInvalidValue);
-  }
+  HIP_RETURN(ihipMemAdvise(dev_ptr, count, advice, location));
+}
 
-  if (!hip::tls.capture_streams_.empty() || !g_captureStreams.empty()) {
-    HIP_RETURN(hipErrorStreamCaptureUnsupported);
-  }
-
-  size_t offset = 0;
-  amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
-  if (memObj && count > (memObj->getSize() - offset)) {
-    HIP_RETURN(hipErrorInvalidValue);
-  }
-
-  amd::Device* dev = (device == hipCpuDeviceId || isAdviseReadMostly) ?
-    g_devices[0]->devices()[0] : g_devices[device]->devices()[0];
-  bool use_cpu = (device == hipCpuDeviceId) ? true : false;
-
-  // Set the allocation attributes in AMD HMM
-  if (!dev->SetSvmAttributes(dev_ptr, count, static_cast<amd::MemoryAdvice>(advice), use_cpu)) {
-    HIP_RETURN(hipErrorInvalidValue);
-  }
-
-  HIP_RETURN(hipSuccess);
+// ================================================================================================
+hipError_t hipMemAdvise_v2(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
+                           hipMemLocation location) {
+  HIP_INIT_API(hipMemAdvise_v2, dev_ptr, count, advice, location);
+  CHECK_STREAM_CAPTURE_SUPPORTED();
+  HIP_RETURN(ihipMemAdvise(dev_ptr, count, advice, location));
 }
 
 // ================================================================================================
@@ -322,4 +279,122 @@ hipError_t ihipMallocManaged(void** ptr, size_t size, size_t align, bool use_hos
   ClPrint(amd::LOG_INFO, amd::LOG_API, "ihipMallocManaged ptr=0x%zx", *ptr);
   return hipSuccess;
 }
+// ================================================================================================
+hipError_t ihipMemPrefetchAsync(const void* dev_ptr, size_t count, hipMemLocation location,
+                                hipStream_t stream) {
+  if ((dev_ptr == nullptr) || (count == 0)) {
+    return hipErrorInvalidValue;
+  }
+
+  getStreamPerThread(stream);
+
+  size_t offset = 0;
+  amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
+  if ((memObj != nullptr) && (count > (memObj->getSize() - offset))) {
+    return hipErrorInvalidValue;
+  }
+  // Compute the type of prefetch
+  const bool isHost = (location.type == hipMemLocationTypeHost);
+  const bool isHostNuma = (location.type == hipMemLocationTypeHostNuma);
+  const bool isHostCurrent = (location.type == hipMemLocationTypeHostNumaCurrent);
+  const bool cpuAccess = isHost || isHostNuma || isHostCurrent;
+
+  // Determine the target device index:
+  //  - for host-prefetch and host-current, always use device 0
+  //  - for host-NUMA or device-prefetch, use the provided id
+  int targetDevice = (isHost || isHostCurrent) ? hipCpuDeviceId : location.id;
+
+  amd::Device* dev = nullptr;
+  if (cpuAccess == false) {
+    if (static_cast<size_t>(targetDevice) >= g_devices.size()) {
+      return hipErrorInvalidDevice;
+    }
+    dev = g_devices[targetDevice]->devices()[0];
+    if (memObj == nullptr && !dev->info().hmmCpuMemoryAccessible_) {
+      return hipErrorNotSupported;
+    }
+  }
+  hip::Stream* hip_stream = nullptr;
+  // Pick the specified stream or Null one from the provided target device
+  if (cpuAccess == true) {
+    hip_stream = (stream == nullptr || stream == hipStreamLegacy)
+        ? hip::getCurrentDevice()->NullStream()
+        : hip::getStream(stream);
+  } else {
+    dev = g_devices[targetDevice]->devices()[0];
+    hip_stream = (stream == nullptr || stream == hipStreamLegacy)
+        ? g_devices[targetDevice]->NullStream()
+        : hip::getStream(stream);
+  }
+
+  if (hip_stream == nullptr) {
+    return hipErrorInvalidValue;
+  }
+
+  amd::Command::EventWaitList waitList;
+  amd::SvmPrefetchAsyncCommand* command = new amd::SvmPrefetchAsyncCommand(
+      *hip_stream, waitList, dev_ptr, count, dev, cpuAccess, targetDevice);
+  if (command == nullptr) {
+    return hipErrorOutOfMemory;
+  }
+  command->enqueue();
+  command->release();
+  return hipSuccess;
+}
+// ================================================================================================
+hipError_t ihipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
+                         hipMemLocation location) {
+  if ((dev_ptr == nullptr) || (count == 0)) {
+    return hipErrorInvalidValue;
+  }
+
+  if (!hip::tls.capture_streams_.empty() || !g_captureStreams.empty()) {
+    return hipErrorStreamCaptureUnsupported;
+  }
+
+  // Determine device and CPU access from location
+  int targetDevice = hipCpuDeviceId;
+  bool use_cpu = true;
+  bool isAdviseReadMostly =
+      (advice == hipMemAdviseSetReadMostly) || (advice == hipMemAdviseUnsetReadMostly);
+
+  switch (location.type) {
+    case hipMemLocationTypeDevice:
+      targetDevice = location.id;
+      use_cpu = false;
+      break;
+    case hipMemLocationTypeHostNuma:
+      targetDevice = location.id;  // NUMA node ID
+      use_cpu = true;
+      break;
+    case hipMemLocationTypeHost:
+    case hipMemLocationTypeHostNumaCurrent:
+      targetDevice = hipCpuDeviceId;
+      use_cpu = true;
+      break;
+    default:
+      return hipErrorInvalidValue;
+  }
+
+  if (!isAdviseReadMostly && !use_cpu && (static_cast<size_t>(targetDevice) >= g_devices.size())) {
+    return hipErrorInvalidDevice;
+  }
+
+  size_t offset = 0;
+  amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
+  if (memObj && count > (memObj->getSize() - offset)) {
+    return hipErrorInvalidValue;
+  }
+
+  amd::Device* dev = (use_cpu || isAdviseReadMostly) ? g_devices[0]->devices()[0]
+                                                     : g_devices[targetDevice]->devices()[0];
+
+  // Set the allocation attributes in AMD HMM
+  if (!dev->SetSvmAttributes(dev_ptr, count, static_cast<amd::MemoryAdvice>(advice), use_cpu,
+                             targetDevice)) {
+    return hipErrorInvalidValue;
+  }
+
+  return hipSuccess;
+}
 } //namespace hip

hipamd/src/hip_table_interface.cpp

@@ -936,6 +936,10 @@ hipError_t hipMemAddressReserve(void** ptr, size_t size, size_t alignment, void*
 hipError_t hipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advice, int device) {
   return hip::GetHipDispatchTable()->hipMemAdvise_fn(dev_ptr, count, advice, device);
 }
+hipError_t hipMemAdvise_v2(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
+                           hipMemLocation location) {
+  return hip::GetHipDispatchTable()->hipMemAdvise_v2_fn(dev_ptr, count, advice, location);
+}
 hipError_t hipMemAllocHost(void** ptr, size_t size) {
   return hip::GetHipDispatchTable()->hipMemAllocHost_fn(ptr, size);
 }
@@ -1030,6 +1034,11 @@ hipError_t hipMemPoolTrimTo(hipMemPool_t mem_pool, size_t min_bytes_to_hold) {
 hipError_t hipMemPrefetchAsync(const void* dev_ptr, size_t count, int device, hipStream_t stream) {
   return hip::GetHipDispatchTable()->hipMemPrefetchAsync_fn(dev_ptr, count, device, stream);
 }
+hipError_t hipMemPrefetchAsync_v2(const void* dev_ptr, size_t count, hipMemLocation location,
+                                  unsigned int flags, hipStream_t stream) {
+  return hip::GetHipDispatchTable()->hipMemPrefetchAsync_v2_fn(dev_ptr, count, location, flags,
+                                                               stream);
+}
 hipError_t hipMemPtrGetInfo(void* ptr, size_t* size) {
   return hip::GetHipDispatchTable()->hipMemPtrGetInfo_fn(ptr, size);
 }

rocclr/device/device.hpp

@@ -2023,8 +2023,8 @@ class Device : public RuntimeObject {
   /**
    * @return True if the device successfully applied the SVM attributes in HMM for device memory
    */
-  virtual bool SetSvmAttributes(const void* dev_ptr, size_t count,
-                                amd::MemoryAdvice advice, bool use_cpu = false) const {
+  virtual bool SetSvmAttributes(const void* dev_ptr, size_t count, amd::MemoryAdvice advice,
+                                bool use_cpu = false, int numa_id = kDefaultNumaNode) const {
     ShouldNotCallThis();
     return false;
   }
@@ -2286,8 +2286,8 @@ class Device : public RuntimeObject {
   static amd::Monitor lockP2P_;
   Monitor* vaCacheAccess_;                            //!< Lock to serialize VA caching access
   std::map<uintptr_t, device::Memory*>* vaCacheMap_;  //!< VA cache map
-  uint32_t index_;  //!< Unique device index
-
+  uint32_t index_;                                    //!< Unique device index
+  static constexpr int kDefaultNumaNode = -1;         //! Default NUMA node value for SVM operations
   // Tracks all amd::Memory objects allocated via hostcall for this device.
   std::vector<amd::Memory*> hostcall_allocated_memories_;
 };

rocclr/device/rocm/rocdevice.cpp

@@ -2288,22 +2288,22 @@ void Device::updateFreeMemory(size_t size, bool free) {
 void* Device::svmAlloc(amd::Context& context, size_t size, size_t alignment, cl_svm_mem_flags flags,
                        void* svmPtr) const {
   amd::Memory* mem = nullptr;
-  void* svmPtrUsed = reinterpret_cast<void*>(amd::Memory::MemoryType::kSvmMemoryPtr); 
-                  
+  void* svmPtrUsed = reinterpret_cast<void*>(amd::Memory::MemoryType::kSvmMemoryPtr);
+
   if (nullptr != svmPtr) {
     // Find the existing amd::mem object
     mem = amd::MemObjMap::FindMemObj(svmPtr);
     if (mem != nullptr) {
       return mem->getSvmPtr();
-    } 
+    }
     if (flags & CL_MEM_USE_HOST_PTR ) {
       svmPtrUsed = svmPtr;
     } else {
       DevLogPrintfError("Cannot find svm_ptr: 0x%x \n", svmPtr);
       return nullptr;
     }
-  }     
-  
+  }
+
   // create a hidden buffer, which will allocated on the device later
   mem = new (context) amd::Buffer(context, flags, size, svmPtrUsed);
   if (mem == nullptr) {
@@ -2468,8 +2468,8 @@ amd::Memory* Device::ImportShareableVMMHandle(void* osHandle) {
 }
 
 // ================================================================================================
-bool Device::SetSvmAttributesInt(const void* dev_ptr, size_t count,
-                              amd::MemoryAdvice advice, bool first_alloc, bool use_cpu) const {
+bool Device::SetSvmAttributesInt(const void* dev_ptr, size_t count, amd::MemoryAdvice advice,
+                                 bool first_alloc, bool use_cpu, int numa_id) const {
   if ((settings().hmmFlags_ & Settings::Hmm::EnableSvmTracking) && !first_alloc) {
     amd::Memory* svm_mem = amd::MemObjMap::FindMemObj(dev_ptr);
     if ((nullptr == svm_mem) || ((svm_mem->getMemFlags() & CL_MEM_ALLOC_HOST_PTR) == 0) ||
@@ -2492,7 +2492,7 @@ bool Device::SetSvmAttributesInt(const void* dev_ptr, size_t count,
         break;
       case amd::MemoryAdvice::SetPreferredLocation:
         if (use_cpu) {
-          attr.push_back({HSA_AMD_SVM_ATTRIB_PREFERRED_LOCATION, getCpuAgent().handle});
+          attr.push_back({HSA_AMD_SVM_ATTRIB_PREFERRED_LOCATION, getCpuAgent(numa_id).handle});
         } else {
           attr.push_back({HSA_AMD_SVM_ATTRIB_PREFERRED_LOCATION, getBackendDevice().handle});
         }
@@ -2552,7 +2552,7 @@ bool Device::SetSvmAttributesInt(const void* dev_ptr, size_t count,
 
 // ================================================================================================
 bool Device::SetSvmAttributes(const void* dev_ptr, size_t count,
-                              amd::MemoryAdvice advice, bool use_cpu) const {
+                              amd::MemoryAdvice advice, bool use_cpu, int numa_id) const {
   constexpr bool kFirstAlloc = false;
   return SetSvmAttributesInt(dev_ptr, count, advice, kFirstAlloc, use_cpu);
 }

rocclr/device/rocm/rocdevice.hpp

@@ -345,6 +345,15 @@ class Device : public NullDevice {
   //! Get the CPU agent with the least NUMA distance to this GPU
   const hsa_agent_t &getCpuAgent() const { return cpu_agent_info_->agent; }
 
+  //! Get the CPU agent that is in a 'index' NUMA node
+  const hsa_agent_t getCpuAgent(int index) const {
+    if ((index < 0) || (index >= cpu_agents_.size())) {
+      // Return default CPU agent
+      return cpu_agent_info_->agent;
+    }
+    return cpu_agents_[index].agent;
+  }
+
   void setupCpuAgent(); // Setup the CPU agent which has the least NUMA distance to this GPU
 
   void checkAtomicSupport(); //!< Check the support for pcie atomics
@@ -425,8 +434,8 @@ class Device : public NullDevice {
 
   virtual void svmFree(void* ptr) const;
 
-  virtual bool SetSvmAttributes(const void* dev_ptr, size_t count,
-                                amd::MemoryAdvice advice, bool use_cpu = false) const;
+  virtual bool SetSvmAttributes(const void* dev_ptr, size_t count, amd::MemoryAdvice advice,
+                                bool use_cpu = false, int numa_id = kDefaultNumaNode) const;
   virtual bool GetSvmAttributes(void** data, size_t* data_sizes, int* attributes,
                                 size_t num_attributes, const void* dev_ptr, size_t count) const;
   virtual size_t ScratchLimitCurrent() const final;
@@ -578,8 +587,11 @@ class Device : public NullDevice {
   //! Construct a new physical HSA device
   Device(hsa_agent_t bkendDevice);
 
+  static constexpr int kDefaultNumaNode = -1;
+
   bool SetSvmAttributesInt(const void* dev_ptr, size_t count, amd::MemoryAdvice advice,
-                           bool first_alloc = false, bool use_cpu = false) const;
+                           bool first_alloc = false, bool use_cpu = false,
+                           int numa_id = kDefaultNumaNode) const;
   static constexpr hsa_signal_value_t InitSignalValue = 1;
 
   static hsa_ven_amd_loader_1_00_pfn_t amd_loader_ext_table;

rocclr/device/rocm/rocvirtual.cpp

@@ -2109,9 +2109,10 @@ void VirtualGPU::submitSvmPrefetchAsync(amd::SvmPrefetchAsyncCommand& cmd) {
     hsa_signal_t active = Barriers().ActiveSignal(kInitSignalValueOne, timestamp_);
 
     // Find the requested agent for the transfer
-    hsa_agent_t agent = (cmd.cpu_access() ||
-        (dev().settings().hmmFlags_ & Settings::Hmm::EnableSystemMemory)) ?
-        dev().getCpuAgent() : (static_cast<const roc::Device*>(cmd.device()))->getBackendDevice();
+    hsa_agent_t agent =
+        (cmd.cpu_access() || (dev().settings().hmmFlags_ & Settings::Hmm::EnableSystemMemory))
+        ? dev().getCpuAgent(cmd.numa_id())
+        : (static_cast<const roc::Device*>(cmd.device()))->getBackendDevice();
 
     // Initiate a prefetch command
     hsa_status_t status = hsa_amd_svm_prefetch_async(

rocclr/platform/command.hpp

@@ -1856,13 +1856,15 @@ class SvmPrefetchAsyncCommand : public Command {
   const void* dev_ptr_;   //!< Device pointer to memory for prefetch
   size_t count_;          //!< the size for prefetch
   bool cpu_access_;       //!< Prefetch data into CPU location
+  int numa_id_;           //!< Host NUMA node id
   amd::Device* dev_;      //!< Destination device to prefetch to
 
  public:
   SvmPrefetchAsyncCommand(HostQueue& queue, const EventWaitList& eventWaitList,
-                          const void* dev_ptr, size_t count, amd::Device* dev, bool cpu_access)
+                          const void* dev_ptr, size_t count, amd::Device* dev,
+                          bool cpu_access, int numa_id)
       : Command(queue, 1, eventWaitList), dev_ptr_(dev_ptr), count_(count),
-        cpu_access_(cpu_access), dev_(dev) {}
+        cpu_access_(cpu_access), dev_(dev), numa_id_(numa_id) {}
 
   virtual void submit(device::VirtualDevice& device) { device.submitSvmPrefetchAsync(*this); }
 
@@ -1872,6 +1874,7 @@ class SvmPrefetchAsyncCommand : public Command {
   size_t count() const { return count_; }
   amd::Device* device() const { return dev_; }
   size_t cpu_access() const { return cpu_access_; }
+  int numa_id() const { return numa_id_; }
 };
 
 /*! \brief  A virtual map memory command.