Keep track of only shared USM allocations and prefetch only for those

Author: Hugh Delaney

sycl/plugins/unified_runtime/ur/adapters/hip/context.hpp

@@ -8,7 +8,6 @@
 #pragma once
 
 #include <set>
-#include <unordered_map>
 
 #include "common.hpp"
 #include "device.hpp"
@@ -104,57 +103,6 @@ struct ur_context_handle_t_ {
 
   ur_usm_pool_handle_t getOwningURPool(umf_memory_pool_t *UMFPool);
 
-  /// We need to keep track of USM mappings in AMD HIP, as certain extra
-  /// synchronization *is* actually required for correctness.
-  /// During kernel enqueue we must dispatch a prefetch for each kernel argument
-  /// that points to a USM mapping to ensure the mapping is correctly
-  /// populated on the device (https://github.com/intel/llvm/issues/7252). Thus,
-  /// we keep track of mappings in the context, and then check against them just
-  /// before the kernel is launched. The stream against which the kernel is
-  /// launched is not known until enqueue time, but the USM mappings can happen
-  /// at any time. Thus, they are tracked on the context used for the urUSM*
-  /// mapping.
-  ///
-  /// The three utility function are simple wrappers around a mapping from a
-  /// pointer to a size.
-  void addUSMMapping(void *Ptr, size_t Size) {
-    std::lock_guard<std::mutex> Guard(Mutex);
-    assert(USMMappings.find(Ptr) == USMMappings.end() &&
-           "mapping already exists");
-    USMMappings[Ptr] = Size;
-  }
-
-  void removeUSMMapping(const void *Ptr) {
-    std::lock_guard<std::mutex> guard(Mutex);
-    auto It = USMMappings.find(Ptr);
-    if (It != USMMappings.end())
-      USMMappings.erase(It);
-  }
-
-  std::pair<const void *, size_t> getUSMMapping(const void *Ptr) {
-    std::lock_guard<std::mutex> Guard(Mutex);
-    auto It = USMMappings.find(Ptr);
-    // The simple case is the fast case...
-    if (It != USMMappings.end())
-      return *It;
-
-    // ... but in the failure case we have to fall back to a full scan to search
-    // for "offset" pointers in case the user passes in the middle of an
-    // allocation. We have to do some not-so-ordained-by-the-standard ordered
-    // comparisons of pointers here, but it'll work on all platforms we support.
-    uintptr_t PtrVal = (uintptr_t)Ptr;
-    for (std::pair<const void *, size_t> Pair : USMMappings) {
-      uintptr_t BaseAddr = (uintptr_t)Pair.first;
-      uintptr_t EndAddr = BaseAddr + Pair.second;
-      if (PtrVal > BaseAddr && PtrVal < EndAddr) {
-        // If we've found something now, offset *must* be nonzero
-        assert(Pair.second);
-        return Pair;
-      }
-    }
-    return {nullptr, 0};
-  }
-
 private:
   std::mutex Mutex;
   std::vector<deleter_data> ExtendedDeleters;

sycl/plugins/unified_runtime/ur/adapters/hip/enqueue.cpp

@@ -254,7 +254,6 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   try {
     ur_device_handle_t Dev = hQueue->getDevice();
     ScopedContext Active(Dev);
-    ur_context_handle_t Ctx = hQueue->getContext();
 
     uint32_t StreamToken;
     ur_stream_quard Guard;
@@ -263,15 +262,15 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
     hipFunction_t HIPFunc = hKernel->get();
 
     hipDevice_t HIPDev = Dev->get();
-    for (const void *P : hKernel->getPtrArgs()) {
-      auto [Addr, Size] = Ctx->getUSMMapping(P);
-      if (!Addr)
-        continue;
-      if (hipMemPrefetchAsync(Addr, Size, HIPDev, HIPStream) != hipSuccess)
-        return UR_RESULT_ERROR_INVALID_KERNEL_ARGS;
+
+    // Some args using shared USM require prefetch
+    for (auto [Ptr, Size] : hKernel->Args.PtrArgsRequiringPrefetch) {
+      if (Ptr && Size) {
+        UR_CHECK_ERROR(hipMemPrefetchAsync(Ptr, Size, HIPDev, HIPStream));
+      }
     }
-    Result = enqueueEventsWait(hQueue, HIPStream, numEventsInWaitList,
-                               phEventWaitList);
+    UR_CHECK_ERROR(enqueueEventsWait(hQueue, HIPStream, numEventsInWaitList,
+                                     phEventWaitList));
 
     // Set the implicit global offset parameter if kernel has offset variant
     if (hKernel->getWithOffsetParameter()) {

sycl/plugins/unified_runtime/ur/adapters/hip/kernel.hpp

@@ -57,6 +57,8 @@ struct ur_kernel_handle_t_ {
     args_index_t Indices;
     args_size_t OffsetPerIndex;
     std::set<const void *> PtrArgs;
+    // Ptr args needing prefetch arranged [Ptr, Size of alloca]
+    std::set<std::pair<const void *, size_t>> PtrArgsRequiringPrefetch;
 
     std::uint32_t ImplicitOffsetArgs[3] = {0, 0, 0};
 
@@ -177,11 +179,20 @@ struct ur_kernel_handle_t_ {
     Args.addArg(Index, Size, Arg);
   }
 
-  /// We track all pointer arguments to be able to issue prefetches at enqueue
-  /// time
   void setKernelPtrArg(int Index, size_t Size, const void *PtrArg) {
     Args.PtrArgs.insert(*static_cast<void *const *>(PtrArg));
     setKernelArg(Index, Size, PtrArg);
+    // Ptr args using managed memory may require prefetch
+    hipPointerAttribute_t Attribs;
+    // We are only using hipPointerGetAttributes to check if the ptr refers to
+    // a managed memory location, meaning the Ptr may require a prefetch at
+    // kernel launch. If this call fails then it means that the Ptr may have
+    // been a host Ptr, which is not a problem
+    if (hipPointerGetAttributes(&Attribs, PtrArg) == hipSuccess &&
+        Attribs.isManaged) {
+      Args.PtrArgsRequiringPrefetch.insert(
+          {*static_cast<void *const *>(PtrArg), Size});
+    }
   }
 
   bool isPtrArg(const void *ptr) {