Use new UR entry point to get native mem with device param

Hugh Delaney · Hugh Delaney · commit 5f7b8308e45a · 2023-12-18T16:05:17.000Z
diff --git a/sycl/include/sycl/detail/pi.def b/sycl/include/sycl/detail/pi.def
@@ -60,6 +60,7 @@ _PI_API(piMemRetain)
 _PI_API(piMemRelease)
 _PI_API(piMemBufferPartition)
 _PI_API(piextMemGetNativeHandle)
+_PI_API(piextMemGetNativeHandleExp)
 _PI_API(piextMemCreateWithNativeHandle)
 _PI_API(piextMemImageCreateWithNativeHandle)
 // Program
diff --git a/sycl/include/sycl/detail/pi.h b/sycl/include/sycl/detail/pi.h
@@ -1402,6 +1402,13 @@ __SYCL_EXPORT pi_result piMemBufferPartition(
 __SYCL_EXPORT pi_result piextMemGetNativeHandle(pi_mem mem,
                                                 pi_native_handle *nativeHandle);
 
+/// Gets the native handle of a PI mem object on a particular device
+///
+/// \param mem is the PI mem to get the native handle of.
+/// \param nativeHandle is the native handle of mem.
+__SYCL_EXPORT pi_result piextMemGetNativeHandleExp(
+    pi_mem mem, pi_device dev, pi_native_handle *nativeHandle);
+
 /// Creates PI mem object from a native handle.
 /// NOTE: The created PI object takes ownership of the native handle.
 ///
diff --git a/sycl/plugins/cuda/pi_cuda.cpp b/sycl/plugins/cuda/pi_cuda.cpp
@@ -232,6 +232,11 @@ pi_result piextMemGetNativeHandle(pi_mem Mem, pi_native_handle *NativeHandle) {
   return pi2ur::piextMemGetNativeHandle(Mem, NativeHandle);
 }
 
+pi_result piextMemGetNativeHandleExp(pi_mem Mem, pi_device Dev,
+                                     pi_native_handle *NativeHandle) {
+  return pi2ur::piextMemGetNativeHandleExp(Mem, Dev, NativeHandle);
+}
+
 pi_result piextMemCreateWithNativeHandle(pi_native_handle NativeHandle,
                                          pi_context Context,
                                          bool ownNativeHandle, pi_mem *Mem) {
diff --git a/sycl/plugins/hip/pi_hip.cpp b/sycl/plugins/hip/pi_hip.cpp
@@ -240,6 +240,11 @@ pi_result piextMemGetNativeHandle(pi_mem Mem, pi_native_handle *NativeHandle) {
   return pi2ur::piextMemGetNativeHandle(Mem, NativeHandle);
 }
 
+pi_result piextMemGetNativeHandleExp(pi_mem Mem, pi_device Dev,
+                                     pi_native_handle *NativeHandle) {
+  return pi2ur::piextMemGetNativeHandleExp(Mem, Dev, NativeHandle);
+}
+
 pi_result piextMemCreateWithNativeHandle(pi_native_handle NativeHandle,
                                          pi_context Context,
                                          bool ownNativeHandle, pi_mem *Mem) {
diff --git a/sycl/plugins/level_zero/pi_level_zero.cpp b/sycl/plugins/level_zero/pi_level_zero.cpp
@@ -248,6 +248,11 @@ pi_result piextMemGetNativeHandle(pi_mem Mem, pi_native_handle *NativeHandle) {
   return pi2ur::piextMemGetNativeHandle(Mem, NativeHandle);
 }
 
+pi_result piextMemGetNativeHandleExp(pi_mem Mem, pi_device Dev,
+                                     pi_native_handle *NativeHandle) {
+  return pi2ur::piextMemGetNativeHandleExp(Mem, Dev, NativeHandle);
+}
+
 pi_result piextMemCreateWithNativeHandle(pi_native_handle NativeHandle,
                                          pi_context Context,
                                          bool ownNativeHandle, pi_mem *Mem) {
diff --git a/sycl/plugins/native_cpu/pi_native_cpu.cpp b/sycl/plugins/native_cpu/pi_native_cpu.cpp
@@ -244,6 +244,11 @@ pi_result piextMemGetNativeHandle(pi_mem Mem, pi_native_handle *NativeHandle) {
   return pi2ur::piextMemGetNativeHandle(Mem, NativeHandle);
 }
 
+pi_result piextMemGetNativeHandleExp(pi_mem Mem, pi_device Dev,
+                                     pi_native_handle *NativeHandle) {
+  return pi2ur::piextMemGetNativeHandleExp(Mem, Dev, NativeHandle);
+}
+
 pi_result piextMemCreateWithNativeHandle(pi_native_handle NativeHandle,
                                          pi_context Context,
                                          bool ownNativeHandle, pi_mem *Mem) {
diff --git a/sycl/plugins/opencl/pi_opencl.cpp b/sycl/plugins/opencl/pi_opencl.cpp
@@ -226,6 +226,11 @@ pi_result piextMemGetNativeHandle(pi_mem Mem, pi_native_handle *NativeHandle) {
   return pi2ur::piextMemGetNativeHandle(Mem, NativeHandle);
 }
 
+pi_result piextMemGetNativeHandleExp(pi_mem Mem, pi_device Dev,
+                                     pi_native_handle *NativeHandle) {
+  return pi2ur::piextMemGetNativeHandleExp(Mem, Dev, NativeHandle);
+}
+
 pi_result piextMemCreateWithNativeHandle(pi_native_handle NativeHandle,
                                          pi_context Context,
                                          bool ownNativeHandle, pi_mem *Mem) {
diff --git a/sycl/plugins/unified_runtime/CMakeLists.txt b/sycl/plugins/unified_runtime/CMakeLists.txt
@@ -56,14 +56,9 @@ endif()
 if(SYCL_PI_UR_USE_FETCH_CONTENT)
   include(FetchContent)
 
-  set(UNIFIED_RUNTIME_REPO "https://github.com/oneapi-src/unified-runtime.git")
-  # commit 69a56ea6d1369a6bde5fce97c85fc7dbda49252f
-  # Merge: b25bb64d b78f541d
-  # Author: Kenneth Benzie (Benie) <k.benzie@codeplay.com>
-  # Date:   Mon Dec 11 12:30:24 2023 +0000
-  #     Merge pull request #1123 from aarongreig/aaron/usmLocationProps
-  #     [OpenCL] Add ur_usm_alloc_location_desc struct and handle it in the CL adapter.
-  set(UNIFIED_RUNTIME_TAG 69a56ea6d1369a6bde5fce97c85fc7dbda49252f)
+  # DO NOT MERGE
+  set(UNIFIED_RUNTIME_REPO "https://github.com/hdelan/unified-runtime.git")
+  set(UNIFIED_RUNTIME_TAG get-native-mem-on-device)
 
   if(SYCL_PI_UR_OVERRIDE_FETCH_CONTENT_REPO)
     set(UNIFIED_RUNTIME_REPO "${SYCL_PI_UR_OVERRIDE_FETCH_CONTENT_REPO}")
diff --git a/sycl/plugins/unified_runtime/pi2ur.hpp b/sycl/plugins/unified_runtime/pi2ur.hpp
@@ -3090,6 +3090,20 @@ inline pi_result piextMemGetNativeHandle(pi_mem Mem,
   return PI_SUCCESS;
 }
 
+inline pi_result piextMemGetNativeHandleExp(pi_mem Mem, pi_device Device,
+                                            pi_native_handle *NativeHandle) {
+  PI_ASSERT(Mem, PI_ERROR_INVALID_MEM_OBJECT);
+
+  ur_mem_handle_t UrMem = reinterpret_cast<ur_mem_handle_t>(Mem);
+  ur_device_handle_t UrDevice = reinterpret_cast<ur_device_handle_t>(Device);
+  ur_native_handle_t NativeMem{};
+  HANDLE_ERRORS(urMemGetNativeHandleExp(UrMem, UrDevice, &NativeMem));
+
+  *NativeHandle = reinterpret_cast<pi_native_handle>(NativeMem);
+
+  return PI_SUCCESS;
+}
+
 inline pi_result
 piEnqueueMemImageCopy(pi_queue Queue, pi_mem SrcImage, pi_mem DstImage,
                       pi_image_offset SrcOrigin, pi_image_offset DstOrigin,
diff --git a/sycl/plugins/unified_runtime/pi_unified_runtime.cpp b/sycl/plugins/unified_runtime/pi_unified_runtime.cpp
@@ -240,6 +240,11 @@ piextMemGetNativeHandle(pi_mem Mem, pi_native_handle *NativeHandle) {
   return pi2ur::piextMemGetNativeHandle(Mem, NativeHandle);
 }
 
+__SYCL_EXPORT pi_result piextMemGetNativeHandleExp(
+    pi_mem Mem, pi_device Device, pi_native_handle *NativeHandle) {
+  return pi2ur::piextMemGetNativeHandleExp(Mem, Device, NativeHandle);
+}
+
 __SYCL_EXPORT pi_result
 piEnqueueMemImageCopy(pi_queue Queue, pi_mem SrcImage, pi_mem DstImage,
                       pi_image_offset SrcOrigin, pi_image_offset DstOrigin,
@@ -1390,6 +1395,7 @@ __SYCL_EXPORT pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_API(piMemBufferPartition)
   _PI_API(piEnqueueMemImageCopy)
   _PI_API(piextMemGetNativeHandle)
+  _PI_API(piextMemGetNativeHandleExp)
   _PI_API(piextMemCreateWithNativeHandle)
   _PI_API(piMemRetain)
   _PI_API(piextUSMGetMemAllocInfo)
diff --git a/sycl/source/interop_handle.cpp b/sycl/source/interop_handle.cpp
@@ -34,8 +34,13 @@ pi_native_handle interop_handle::getNativeMem(detail::Requirement *Req) const {
 
   auto Plugin = MQueue->getPlugin();
   pi_native_handle Handle;
-  Plugin->call<detail::PiApiKind::piextMemGetNativeHandle>(Iter->second,
-                                                           &Handle);
+  if (get_backend() == backend::ext_oneapi_cuda ||
+      get_backend() == backend::ext_oneapi_hip)
+    Plugin->call<detail::PiApiKind::piextMemGetNativeHandleExp>(
+        Iter->second, MDevice->getHandleRef(), &Handle);
+  else
+    Plugin->call<detail::PiApiKind::piextMemGetNativeHandle>(Iter->second,
+                                                             &Handle);
   return Handle;
 }
 
diff --git a/sycl/test-e2e/HostInteropTask/interop-task-hip.cpp b/sycl/test-e2e/HostInteropTask/interop-task-hip.cpp
@@ -0,0 +1,199 @@
+// FIXME: the rocm include path and link path are highly platform dependent,
+// we should set this with some variable instead
+// RUN: %{build} -o %t.out -I/opt/rocm/include -L/opt/rocm/lib -lamdhip64
+// RUN: %{run} %t.out
+// REQUIRES: hip
+
+#include <iostream>
+#include <sycl/sycl.hpp>
+
+#define __HIP_PLATFORM_AMD__
+
+#include <hip/hip_runtime.h>
+
+using namespace sycl;
+using namespace sycl::access;
+
+static constexpr size_t BUFFER_SIZE = 1024;
+
+template <typename T> class Modifier;
+
+template <typename T> class Init;
+
+template <typename BufferT, typename ValueT>
+void checkBufferValues(BufferT Buffer, ValueT Value) {
+  auto Acc = Buffer.get_host_access();
+  for (size_t Idx = 0; Idx < Acc.get_count(); ++Idx) {
+    if (Acc[Idx] != Value) {
+      std::cerr << "buffer[" << Idx << "] = " << Acc[Idx]
+                << ", expected val = " << Value << std::endl;
+      assert(0 && "Invalid data in the buffer");
+    }
+  }
+}
+
+template <typename DataT>
+void copy(buffer<DataT, 1> &Src, buffer<DataT, 1> &Dst, queue &Q) {
+  Q.submit([&](handler &CGH) {
+    auto SrcA = Src.template get_access<mode::read>(CGH);
+    auto DstA = Dst.template get_access<mode::write>(CGH);
+
+    auto Func = [=](interop_handle IH) {
+      auto HipStream = IH.get_native_queue<backend::ext_oneapi_hip>();
+      auto SrcMem = IH.get_native_mem<backend::ext_oneapi_hip>(SrcA);
+      auto DstMem = IH.get_native_mem<backend::ext_oneapi_hip>(DstA);
+      cl_event Event;
+
+      if (hipMemcpyWithStream(DstMem, SrcMem, sizeof(DataT) * SrcA.get_count(),
+                              hipMemcpyDefault, HipStream) != hipSuccess) {
+        throw;
+      }
+
+      if (hipStreamSynchronize(HipStream) != hipSuccess) {
+        throw;
+      }
+
+      if (Q.get_backend() != IH.get_backend())
+        throw;
+    };
+    CGH.host_task(Func);
+  });
+}
+
+template <typename DataT> void modify(buffer<DataT, 1> &B, queue &Q) {
+  Q.submit([&](handler &CGH) {
+    auto Acc = B.template get_access<mode::read_write>(CGH);
+
+    auto Kernel = [=](item<1> Id) { Acc[Id] += 1; };
+
+    CGH.parallel_for<Modifier<DataT>>(Acc.get_count(), Kernel);
+  });
+}
+
+template <typename DataT, DataT B1Init, DataT B2Init>
+void init(buffer<DataT, 1> &B1, buffer<DataT, 1> &B2, queue &Q) {
+  Q.submit([&](handler &CGH) {
+    auto Acc1 = B1.template get_access<mode::write>(CGH);
+    auto Acc2 = B2.template get_access<mode::write>(CGH);
+
+    CGH.parallel_for<Init<DataT>>(BUFFER_SIZE, [=](item<1> Id) {
+      Acc1[Id] = -1;
+      Acc2[Id] = -2;
+    });
+  });
+}
+
+// A test that uses OpenCL interop to copy data from buffer A to buffer B, by
+// getting cl_mem objects and calling the clEnqueueBufferCopy. Then run a SYCL
+// kernel that modifies the data in place for B, e.g. increment one, then copy
+// back to buffer A. Run it on a loop, to ensure the dependencies and the
+// reference counting of the objects is not leaked.
+void test1(queue &Q) {
+  static constexpr int COUNT = 4;
+  buffer<int, 1> Buffer1{BUFFER_SIZE};
+  buffer<int, 1> Buffer2{BUFFER_SIZE};
+
+  // init the buffer with a'priori invalid data
+  init<int, -1, -2>(Buffer1, Buffer2, Q);
+
+  // Repeat a couple of times
+  for (size_t Idx = 0; Idx < COUNT; ++Idx) {
+    copy(Buffer1, Buffer2, Q);
+    modify(Buffer2, Q);
+    copy(Buffer2, Buffer1, Q);
+  }
+
+  checkBufferValues(Buffer1, COUNT - 1);
+  checkBufferValues(Buffer2, COUNT - 1);
+}
+
+// Same as above, but performing each command group on a separate SYCL queue
+// (on the same or different devices). This ensures the dependency tracking
+// works well but also there is no accidental side effects on other queues.
+void test2(queue &Q) {
+  static constexpr int COUNT = 4;
+  buffer<int, 1> Buffer1{BUFFER_SIZE};
+  buffer<int, 1> Buffer2{BUFFER_SIZE};
+
+  // init the buffer with a'priori invalid data
+  init<int, -1, -2>(Buffer1, Buffer2, Q);
+
+  // Repeat a couple of times
+  for (size_t Idx = 0; Idx < COUNT; ++Idx) {
+    copy(Buffer1, Buffer2, Q);
+    modify(Buffer2, Q);
+    copy(Buffer2, Buffer1, Q);
+  }
+  checkBufferValues(Buffer1, COUNT - 1);
+  checkBufferValues(Buffer2, COUNT - 1);
+}
+
+// Same as above but with queue constructed out of context
+void test2_1(queue &Q) {
+  static constexpr int COUNT = 4;
+  buffer<int, 1> Buffer1{BUFFER_SIZE};
+  buffer<int, 1> Buffer2{BUFFER_SIZE};
+
+  device Device;
+  auto Context = context(Device);
+  // init the buffer with a'priori invalid data
+  init<int, -1, -2>(Buffer1, Buffer2, Q);
+
+  // Repeat a couple of times
+  for (size_t Idx = 0; Idx < COUNT; ++Idx) {
+    copy(Buffer1, Buffer2, Q);
+    modify(Buffer2, Q);
+    copy(Buffer2, Buffer1, Q);
+  }
+  checkBufferValues(Buffer1, COUNT - 1);
+  checkBufferValues(Buffer2, COUNT - 1);
+}
+
+// Check that a single host-interop-task with a buffer will work
+void test3(queue &Q) {
+  buffer<int, 1> Buffer{BUFFER_SIZE};
+
+  Q.submit([&](handler &CGH) {
+    auto Acc = Buffer.get_access<mode::write>(CGH);
+    auto Func = [=](interop_handle IH) { /*A no-op */ };
+    CGH.host_task(Func);
+  });
+}
+
+void test4(queue &Q) {
+  buffer<int, 1> Buffer1{BUFFER_SIZE};
+  buffer<int, 1> Buffer2{BUFFER_SIZE};
+
+  Q.submit([&](handler &CGH) {
+    auto Acc = Buffer1.template get_access<mode::write>(CGH);
+
+    auto Kernel = [=](item<1> Id) { Acc[Id] = 123; };
+    CGH.parallel_for<class Test5Init>(Acc.get_count(), Kernel);
+  });
+
+  copy(Buffer1, Buffer2, Q);
+
+  checkBufferValues(Buffer2, static_cast<int>(123));
+}
+
+void tests(queue &Q) {
+  test1(Q);
+  test2(Q);
+  test2_1(Q);
+  test3(Q);
+  test4(Q);
+}
+
+int main() {
+  queue Q([](sycl::exception_list ExceptionList) {
+    if (ExceptionList.size() != 1) {
+      std::cerr << "Should be one exception in exception list" << std::endl;
+      std::abort();
+    }
+    std::rethrow_exception(*ExceptionList.begin());
+  });
+  tests(Q);
+  tests(Q);
+  std::cout << "Test PASSED" << std::endl;
+  return 0;
+}
