[draft][HIP][CUDA] Use new UR entry point to get native mem with device param

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

sycl/include/sycl/detail/pi.h

@@ -1402,6 +1402,15 @@ __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 dev is the PI device that the native allocation will be resident on.
+/// \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.
 ///

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) {

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) {

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) {

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) {

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) {

sycl/plugins/unified_runtime/CMakeLists.txt

@@ -56,13 +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 8d1486a607666763413fea3c70b74a3dd8f73c0d
-  # Author: aarongreig <[email protected]>
-  # Date: Mon Dec 18 13:15:43 2023 +0000
-  #     Merge pull request #939 from steffenlarsen/steffen/virtual_mem_adapters
-  #     [UR][CUDA][L0][HIP] Add virtual memory adapter implementations
-  set(UNIFIED_RUNTIME_TAG 8d1486a607666763413fea3c70b74a3dd8f73c0d)
+  # 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}")

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,

sycl/plugins/unified_runtime/pi_unified_runtime.cpp

@@ -240,6 +240,12 @@ 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 +1396,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)

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;
 }
 

sycl/test-e2e/HostInteropTask/interop-task-hip.cpp

@@ -0,0 +1,175 @@
+// 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);
+
+      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] = B1Init;
+      Acc2[Id] = B2Init;
+    });
+  });
+}
+
+// A test that uses HIP interop to copy data from buffer A to buffer B, by
+// getting HIP ptrs and calling the hipMemcpyWithStream. 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);
+}
+
+// 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);
+  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);
+  std::cout << "Test PASSED" << std::endl;
+  return 0;
+}

sycl/unittests/helpers/PiMockPlugin.hpp

@@ -619,6 +619,13 @@ inline pi_result mock_piextMemGetNativeHandle(pi_mem mem,
   return PI_SUCCESS;
 }
 
+inline pi_result
+mock_piextMemGetNativeHandleExp(pi_mem mem, pi_device dev,
+                                pi_native_handle *nativeHandle) {
+  *nativeHandle = reinterpret_cast<pi_native_handle>(mem);
+  return PI_SUCCESS;
+}
+
 inline pi_result
 mock_piextMemCreateWithNativeHandle(pi_native_handle nativeHandle,
                                     pi_context context, bool ownNativeHandle,