[SYCL][PI][OpenCL] Generalizing interop handler getters

This commit makes get_mem and get_queue of interop_handler return
types based on a specified backend. The backend defaults to OpenCL to
avoid breakages.

Signed-off-by: Steffen Larsen <[email protected]>

Author: Steffen Larsen

sycl/include/CL/sycl.hpp

@@ -11,6 +11,7 @@
 #include <CL/sycl/accessor.hpp>
 #include <CL/sycl/atomic.hpp>
 #include <CL/sycl/backend.hpp>
+#include <CL/sycl/backend/opencl.hpp>
 #include <CL/sycl/buffer.hpp>
 #include <CL/sycl/builtins.hpp>
 #include <CL/sycl/context.hpp>

sycl/include/CL/sycl/backend/opencl.hpp

@@ -0,0 +1,31 @@
+
+//==---------------- opencl.hpp - SYCL OpenCL backend ----------------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <CL/cl.h>
+#include <CL/sycl/accessor.hpp>
+#include <CL/sycl/backend_types.hpp>
+
+__SYCL_INLINE_NAMESPACE(cl) {
+namespace sycl {
+
+template <> struct interop<backend::opencl, queue> {
+  using type = cl_command_queue;
+};
+
+template <typename DataT, int Dimensions, access::mode AccessMode>
+struct interop<backend::opencl, accessor<DataT, Dimensions, AccessMode,
+                                         access::target::global_buffer,
+                                         access::placeholder::false_t>> {
+  using type = cl_mem;
+};
+
+} // namespace sycl
+} // __SYCL_INLINE_NAMESPACE(cl)

sycl/include/CL/sycl/detail/cg.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include <CL/sycl/backend_types.hpp>
 #include <CL/sycl/detail/accessor_impl.hpp>
 #include <CL/sycl/detail/common.hpp>
 #include <CL/sycl/detail/export.hpp>
@@ -29,6 +30,12 @@
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
 
+// Forward declaration
+class queue;
+namespace detail {
+class queue_impl;
+} // namespace detail
+
 // Interoperability handler
 //
 class interop_handler {
@@ -37,26 +44,51 @@ class interop_handler {
             access::target AccTarget, access::placeholder isPlaceholder>
   friend class accessor;
 public:
+  using QueueImplPtr = std::shared_ptr<detail::queue_impl>;
   using ReqToMem = std::pair<detail::Requirement*, pi_mem>;
 
-  interop_handler(std::vector<ReqToMem> MemObjs, cl_command_queue PiQueue) :
-    MQueue(PiQueue), MMemObjs(MemObjs) {}
+  interop_handler(std::vector<ReqToMem> MemObjs, QueueImplPtr Queue)
+      : MQueue(std::move(Queue)), MMemObjs(std::move(MemObjs)) {}
 
-  cl_command_queue get_queue() const noexcept { return MQueue; };
+  template <backend BackendName = backend::opencl>
+  auto get_queue() const -> typename interop<BackendName, queue>::type {
+    return reinterpret_cast<typename interop<BackendName, queue>::type>(
+        GetNativeQueue());
+  }
 
-  template <typename DataT, int Dims, access::mode AccessMode,
-            access::target AccessTarget,
+  template <backend BackendName = backend::opencl, typename DataT, int Dims,
+            access::mode AccessMode, access::target AccessTarget,
             access::placeholder IsPlaceholder = access::placeholder::false_t>
-  cl_mem get_mem(accessor<DataT, Dims, AccessMode, AccessTarget,
-                          access::placeholder::false_t>
-                     Acc) const {
+  auto get_mem(accessor<DataT, Dims, AccessMode, AccessTarget,
+                        access::placeholder::false_t>
+                   Acc) const ->
+      typename interop<BackendName,
+                       accessor<DataT, Dims, AccessMode, AccessTarget,
+                                access::placeholder::false_t>>::type {
     detail::AccessorBaseHost *AccBase = (detail::AccessorBaseHost *)&Acc;
-    return getMemImpl(detail::getSyclObjImpl(*AccBase).get());
+    return getMemImpl<BackendName, DataT, Dims, AccessMode, AccessTarget,
+                      access::placeholder::false_t>(
+        detail::getSyclObjImpl(*AccBase).get());
   }
+
 private:
-  cl_command_queue MQueue;
+  QueueImplPtr MQueue;
   std::vector<ReqToMem> MMemObjs;
-  __SYCL_EXPORT cl_mem getMemImpl(detail::Requirement *Req) const;
+
+  template <backend BackendName, typename DataT, int Dims,
+            access::mode AccessMode, access::target AccessTarget,
+            access::placeholder IsPlaceholder>
+  __SYCL_EXPORT auto
+  getMemImpl(detail::Requirement *Req) const -> typename interop<
+      BackendName,
+      accessor<DataT, Dims, AccessMode, AccessTarget, IsPlaceholder>>::type {
+    return (typename interop<BackendName,
+                             accessor<DataT, Dims, AccessMode, AccessTarget,
+                                      IsPlaceholder>>::type)GetNativeMem(Req);
+  }
+
+  __SYCL_EXPORT pi_native_handle GetNativeMem(detail::Requirement *Req) const;
+  __SYCL_EXPORT pi_native_handle GetNativeQueue() const;
 };
 
 namespace detail {

sycl/plugins/opencl/pi_opencl.cpp

@@ -1043,6 +1043,23 @@ static pi_result OCL(piextProgramSetSpecializationConstantImpl)(
   return cast<pi_result>(Res);
 }
 
+/// API to get the native handle of a PI object
+///
+/// \param handleType is an identifier representing the type of the handle
+/// \param piObject is the PI object to get the handle of
+/// \param nativeHandle is the native handle of piObject
+pi_result OCL(piGetNativeHandle)(pi_handle_type handleType, void *piObject,
+                                 pi_native_handle *nativeHandle) {
+  switch (handleType) {
+  case pi_handle_type::PI_NATIVE_HANDLE_MEM:
+  case pi_handle_type::PI_NATIVE_HANDLE_QUEUE:
+    *nativeHandle = reinterpret_cast<pi_native_handle>(piObject);
+    return PI_SUCCESS;
+  default:
+    return PI_INVALID_VALUE;
+  }
+}
+
 pi_result piPluginInit(pi_plugin *PluginInit) {
   int CompareVersions = strcmp(PluginInit->PiVersion, SupportedVersion);
   if (CompareVersions < 0) {
@@ -1154,6 +1171,8 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piextUSMEnqueuePrefetch, OCL(piextUSMEnqueuePrefetch))
   _PI_CL(piextUSMEnqueueMemAdvise, OCL(piextUSMEnqueueMemAdvise))
   _PI_CL(piextUSMGetMemAllocInfo, OCL(piextUSMGetMemAllocInfo))
+  // Native
+  _PI_CL(piGetNativeHandle, OCL(piGetNativeHandle))
 
   _PI_CL(piextKernelSetArgMemObj, OCL(piextKernelSetArgMemObj))
 

sycl/source/detail/cg.cpp

@@ -8,11 +8,11 @@
 
 #include "CL/sycl/detail/cg.hpp"
 #include <CL/sycl/detail/memory_manager.hpp>
+#include <CL/sycl/detail/pi.hpp>
 #include <detail/queue_impl.hpp>
 #include <detail/scheduler/commands.hpp>
 #include <detail/scheduler/scheduler.hpp>
 
-
 #include <memory>
 #include <string>
 #include <type_traits>
@@ -21,17 +21,24 @@
 namespace cl {
 namespace sycl {
 
-cl_mem interop_handler::getMemImpl(detail::Requirement* Req) const {
-    auto Iter = std::find_if(std::begin(MMemObjs), std::end(MMemObjs),
-      [=](ReqToMem Elem) {
-        return (Elem.first == Req);
-    });
+pi_native_handle interop_handler::GetNativeQueue() const {
+  return MQueue->getNative();
+}
+
+pi_native_handle interop_handler::GetNativeMem(detail::Requirement *Req) const {
+  auto Iter = std::find_if(std::begin(MMemObjs), std::end(MMemObjs),
+                           [=](ReqToMem Elem) { return (Elem.first == Req); });
 
-    if (Iter == std::end(MMemObjs)) {
-        throw("Invalid memory object used inside interop");
-    }
-    return detail::pi::cast<cl_mem>(Iter->second);
+  if (Iter == std::end(MMemObjs)) {
+    throw("Invalid memory object used inside interop");
   }
 
+  auto Plugin = MQueue->getPlugin();
+  pi_native_handle Handle;
+  Plugin.call<detail::PiApiKind::piGetNativeHandle>(
+      pi_handle_type::PI_NATIVE_HANDLE_MEM, Iter->second, &Handle);
+  return Handle;
+}
+
 }  // sycl
 }  // cl

sycl/source/detail/scheduler/commands.cpp

@@ -1772,12 +1772,13 @@ cl_int ExecCGCommand::enqueueImp() {
       ReqMemObjs.emplace_back(ReqToMem);
     });
 
-    auto interop_queue = MQueue->get();
     std::sort(std::begin(ReqMemObjs), std::end(ReqMemObjs));
-    interop_handler InteropHandler(std::move(ReqMemObjs), interop_queue);
+    interop_handler InteropHandler(std::move(ReqMemObjs), MQueue);
     ExecInterop->MInteropTask->call(InteropHandler);
-    Plugin.call<PiApiKind::piEnqueueEventsWait>(MQueue->getHandleRef(), 0, nullptr, &Event);
-    Plugin.call<PiApiKind::piQueueRelease>(reinterpret_cast<pi_queue>(interop_queue));
+    Plugin.call<PiApiKind::piEnqueueEventsWait>(MQueue->getHandleRef(), 0,
+                                                nullptr, &Event);
+    Plugin.call<PiApiKind::piQueueRelease>(
+        reinterpret_cast<pi_queue>(MQueue->get()));
     return CL_SUCCESS;
   }
   case CG::CGTYPE::NONE:

sycl/unittests/pi/cuda/test_interop_get_native.cpp

@@ -1,8 +1,4 @@
-// REQUIRES: cuda
-// RUN: %clangxx -fsycl -fsycl-targets=%sycl_triple %s -I%opencl_include_dir -I%cuda_toolkit_include -o %t.out -lcuda -lsycl
-// RUN: env SYCL_DEVICE_TYPE=GPU %t.out
-
-//==---------- interop_get_native.cpp - SYCL cuda get_native tests ---------==//
+//==------- test_interop_get_native.cpp - SYCL CUDA get_native tests -------==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -74,4 +70,28 @@ TEST_F(DISABLED_CudaInteropGetNativeTests, getNativeQueue) {
 
   CUcontext cudaContext = get_native<backend::cuda>(syclContext_);
   ASSERT_EQ(streamContext, cudaContext);
-}
+}
+
+TEST_F(DISABLED_CudaInteropGetNativeTests, interopTaskGetMem) {
+  buffer<int, 1> syclBuffer(range<1>{1});
+  syclQueue_.submit([&](cl::sycl::handler &cgh) {
+    auto syclAccessor = syclBuffer.get_access<access::mode::read>(cgh);
+    cgh.interop_task([=](sycl::interop_handler ih) {
+      CUdeviceptr cudaPtr = ih.get_mem<backend::cuda>(syclAccessor);
+      CUdeviceptr cudaPtrBase;
+      size_t cudaPtrSize = 0;
+      cuMemGetAddressRange(&cudaPtrBase, &cudaPtrSize, cudaPtr);
+      ASSERT_EQ(cudaPtrSize, sizeof(int));
+    });
+  });
+}
+
+TEST_F(DISABLED_CudaInteropGetNativeTests, interopTaskGetBufferMem) {
+  CUstream cudaStream = get_native<backend::cuda>(syclQueue_);
+  syclQueue_.submit([&](cl::sycl::handler &cgh) {
+    cgh.interop_task([=](sycl::interop_handler ih) {
+      CUstream cudaInteropStream = ih.get_queue<backend::cuda>();
+      ASSERT_EQ(cudaInteropStream, cudaStream);
+    });
+  });
+}