[SYCL] Add implementation of host-interop-task and test. (#1748)

This patch is number two in series of patches for interop part of host task.
This patch introduces an API to enqueue host-task with `interop_handle` argument
See the proposal https://github.com/codeplaysoftware/standards-proposals/blob/master/host_task/host_task.md

Signed-off-by: Sergey Kanaev <[email protected]>

Author: s-kanaev

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

@@ -305,9 +305,16 @@ class CGInteropTask : public CG {
 class CGHostTask : public CG {
 public:
   std::unique_ptr<HostTask> MHostTask;
+  // queue for host-interop task
+  shared_ptr_class<detail::queue_impl> MQueue;
+  // context for host-interop task
+  shared_ptr_class<detail::context_impl> MContext;
   vector_class<ArgDesc> MArgs;
 
-  CGHostTask(std::unique_ptr<HostTask> HostTask, vector_class<ArgDesc> Args,
+  CGHostTask(std::unique_ptr<HostTask> HostTask,
+             std::shared_ptr<detail::queue_impl> Queue,
+             std::shared_ptr<detail::context_impl> Context,
+             vector_class<ArgDesc> Args,
              std::vector<std::vector<char>> ArgsStorage,
              std::vector<detail::AccessorImplPtr> AccStorage,
              std::vector<std::shared_ptr<const void>> SharedPtrStorage,
@@ -317,7 +324,8 @@ class CGHostTask : public CG {
       : CG(Type, std::move(ArgsStorage), std::move(AccStorage),
            std::move(SharedPtrStorage), std::move(Requirements),
            std::move(Events), std::move(loc)),
-        MHostTask(std::move(HostTask)), MArgs(std::move(Args)) {}
+        MHostTask(std::move(HostTask)), MQueue(Queue), MContext(Context),
+        MArgs(std::move(Args)) {}
 };
 
 class CGBarrier : public CG {

sycl/include/CL/sycl/detail/cg_types.hpp

@@ -12,6 +12,7 @@
 #include <CL/sycl/detail/kernel_desc.hpp>
 #include <CL/sycl/group.hpp>
 #include <CL/sycl/id.hpp>
+#include <CL/sycl/interop_handle.hpp>
 #include <CL/sycl/interop_handler.hpp>
 #include <CL/sycl/kernel.hpp>
 #include <CL/sycl/nd_item.hpp>
@@ -143,12 +144,17 @@ class InteropTask {
 
 class HostTask {
   std::function<void()> MHostTask;
+  std::function<void(interop_handle)> MInteropTask;
 
 public:
   HostTask() : MHostTask([]() {}) {}
   HostTask(std::function<void()> &&Func) : MHostTask(Func) {}
+  HostTask(std::function<void(interop_handle)> &&Func) : MInteropTask(Func) {}
+
+  bool isInteropTask() const { return !!MInteropTask; }
 
   void call() { MHostTask(); }
+  void call(interop_handle handle) { MInteropTask(handle); }
 };
 
 // Class which stores specific lambda object.

sycl/include/CL/sycl/detail/sycl_mem_obj_i.hpp

@@ -69,6 +69,7 @@ class SYCLMemObjI {
   // which is unavailable.
   shared_ptr_class<MemObjRecord> MRecord;
   friend class Scheduler;
+  friend class ExecCGCommand;
 };
 
 } // namespace detail

sycl/include/CL/sycl/handler.hpp

@@ -19,6 +19,7 @@
 #include <CL/sycl/detail/os_util.hpp>
 #include <CL/sycl/event.hpp>
 #include <CL/sycl/id.hpp>
+#include <CL/sycl/interop_handle.hpp>
 #include <CL/sycl/item.hpp>
 #include <CL/sycl/kernel.hpp>
 #include <CL/sycl/nd_item.hpp>
@@ -856,8 +857,22 @@ class __SYCL_EXPORT handler {
   }
 
   template <typename FuncT>
-  typename std::enable_if<detail::check_fn_signature<
-      typename std::remove_reference<FuncT>::type, void()>::value>::type
+  detail::enable_if_t<detail::check_fn_signature<
+      detail::remove_reference_t<FuncT>, void()>::value>
+  codeplay_host_task(FuncT Func) {
+    throwIfActionIsCreated();
+
+    MNDRDesc.set(range<1>(1));
+    MArgs = std::move(MAssociatedAccesors);
+
+    MHostTask.reset(new detail::HostTask(std::move(Func)));
+
+    MCGType = detail::CG::CODEPLAY_HOST_TASK;
+  }
+
+  template <typename FuncT>
+  detail::enable_if_t<detail::check_fn_signature<
+      detail::remove_reference_t<FuncT>, void(interop_handle)>::value>
   codeplay_host_task(FuncT Func) {
     throwIfActionIsCreated();
 

sycl/include/CL/sycl/interop_handle.hpp

@@ -84,8 +84,13 @@ class interop_handle {
   template <backend BackendName = backend::opencl>
   auto get_native_queue() const noexcept ->
       typename interop<BackendName, queue>::type {
+#ifndef __SYCL_DEVICE_ONLY__
     return reinterpret_cast<typename interop<BackendName, queue>::type>(
         getNativeQueue());
+#else
+    // we believe this won't be ever called on device side
+    return nullptr;
+#endif
   }
 
   /// Returns an underlying OpenCL device associated with the SYCL queue used
@@ -94,8 +99,13 @@ class interop_handle {
   template <backend BackendName = backend::opencl>
   auto get_native_device() const noexcept ->
       typename interop<BackendName, device>::type {
+#ifndef __SYCL_DEVICE_ONLY__
     return reinterpret_cast<typename interop<BackendName, device>::type>(
         getNativeDevice());
+#else
+    // we believe this won't be ever called on device side
+    return nullptr;
+#endif
   }
 
   /// Returns an underlying OpenCL context associated with the SYCL queue used
@@ -104,14 +114,20 @@ class interop_handle {
   template <backend BackendName = backend::opencl>
   auto get_native_context() const noexcept ->
       typename interop<BackendName, context>::type {
+#ifndef __SYCL_DEVICE_ONLY__
     return reinterpret_cast<typename interop<BackendName, context>::type>(
         getNativeContext());
+#else
+    // we believe this won't be ever called on device side
+    return nullptr;
+#endif
   }
 
 private:
+  friend class detail::ExecCGCommand;
+  friend class detail::DispatchHostTask;
   using ReqToMem = std::pair<detail::Requirement *, pi_mem>;
 
-public:
   // TODO set c-tor private
   interop_handle(std::vector<ReqToMem> MemObjs,
                  const std::shared_ptr<detail::queue_impl> &Queue,
@@ -131,10 +147,10 @@ class interop_handle {
         getNativeMem(Req));
   }
 
-  pi_native_handle getNativeMem(detail::Requirement *Req) const;
-  pi_native_handle getNativeQueue() const;
-  pi_native_handle getNativeDevice() const;
-  pi_native_handle getNativeContext() const;
+  __SYCL_EXPORT pi_native_handle getNativeMem(detail::Requirement *Req) const;
+  __SYCL_EXPORT pi_native_handle getNativeQueue() const;
+  __SYCL_EXPORT pi_native_handle getNativeDevice() const;
+  __SYCL_EXPORT pi_native_handle getNativeContext() const;
 
   std::shared_ptr<detail::queue_impl> MQueue;
   std::shared_ptr<detail::device_impl> MDevice;

sycl/source/detail/scheduler/commands.cpp

@@ -8,8 +8,9 @@
 
 #include <detail/error_handling/error_handling.hpp>
 
-#include "CL/sycl/access/access.hpp"
+#include <CL/sycl/access/access.hpp>
 #include <CL/sycl/backend_types.hpp>
+#include <CL/sycl/detail/cg_types.hpp>
 #include <CL/sycl/detail/cl.h>
 #include <CL/sycl/detail/kernel_desc.hpp>
 #include <CL/sycl/detail/memory_manager.hpp>
@@ -159,6 +160,7 @@ getPiEvents(const std::vector<EventImplPtr> &EventImpls) {
 
 class DispatchHostTask {
   ExecCGCommand *MThisCmd;
+  std::vector<interop_handle::ReqToMem> MReqToMem;
 
   void waitForEvents() const {
     std::map<const detail::plugin *, std::vector<EventImplPtr>>
@@ -187,7 +189,9 @@ class DispatchHostTask {
   }
 
 public:
-  DispatchHostTask(ExecCGCommand *ThisCmd) : MThisCmd{ThisCmd} {}
+  DispatchHostTask(ExecCGCommand *ThisCmd,
+                   std::vector<interop_handle::ReqToMem> ReqToMem)
+      : MThisCmd{ThisCmd}, MReqToMem(std::move(ReqToMem)) {}
 
   void operator()() const {
     waitForEvents();
@@ -197,7 +201,15 @@ class DispatchHostTask {
     CGHostTask &HostTask = static_cast<CGHostTask &>(MThisCmd->getCG());
 
     // we're ready to call the user-defined lambda now
-    HostTask.MHostTask->call();
+    if (HostTask.MHostTask->isInteropTask()) {
+      interop_handle IH{MReqToMem, HostTask.MQueue,
+                        getSyclObjImpl(HostTask.MQueue->get_device()),
+                        HostTask.MQueue->getContextImplPtr()};
+
+      HostTask.MHostTask->call(IH);
+    } else
+      HostTask.MHostTask->call();
+
     HostTask.MHostTask.reset();
 
     // unblock user empty command here
@@ -1943,7 +1955,38 @@ cl_int ExecCGCommand::enqueueImp() {
       }
     }
 
-    MQueue->getThreadPool().submit<DispatchHostTask>(DispatchHostTask(this));
+    std::vector<interop_handle::ReqToMem> ReqToMem;
+
+    if (HostTask->MHostTask->isInteropTask()) {
+      // Extract the Mem Objects for all Requirements, to ensure they are
+      // available if a user asks for them inside the interop task scope
+      const std::vector<Requirement *> &HandlerReq = HostTask->MRequirements;
+      auto ReqToMemConv = [&ReqToMem, HostTask](Requirement *Req) {
+        const std::vector<AllocaCommandBase *> &AllocaCmds =
+            Req->MSYCLMemObj->MRecord->MAllocaCommands;
+
+        for (AllocaCommandBase *AllocaCmd : AllocaCmds)
+          if (HostTask->MQueue == AllocaCmd->getQueue()) {
+            auto MemArg =
+                reinterpret_cast<pi_mem>(AllocaCmd->getMemAllocation());
+            ReqToMem.emplace_back(std::make_pair(Req, MemArg));
+
+            return;
+          }
+
+        assert(false &&
+               "Can't get memory object due to no allocation available");
+
+        throw runtime_error(
+            "Can't get memory object due to no allocation available",
+            PI_INVALID_MEM_OBJECT);
+      };
+      std::for_each(std::begin(HandlerReq), std::end(HandlerReq), ReqToMemConv);
+      std::sort(std::begin(ReqToMem), std::end(ReqToMem));
+    }
+
+    MQueue->getThreadPool().submit<DispatchHostTask>(
+        DispatchHostTask(this, std::move(ReqToMem)));
 
     MShouldCompleteEventIfPossible = false;
 

sycl/source/detail/scheduler/graph_builder.cpp

@@ -679,6 +679,16 @@ Scheduler::GraphBuilder::addEmptyCmd(Command *Cmd, const std::vector<T *> &Reqs,
   return EmptyCmd;
 }
 
+static bool isInteropHostTask(const std::unique_ptr<ExecCGCommand> &Cmd) {
+  if (Cmd->getCG().getType() != CG::CGTYPE::CODEPLAY_HOST_TASK)
+    return false;
+
+  const detail::CGHostTask &HT =
+      static_cast<detail::CGHostTask &>(Cmd->getCG());
+
+  return HT.MHostTask->isInteropTask();
+}
+
 Command *
 Scheduler::GraphBuilder::addCG(std::unique_ptr<detail::CG> CommandGroup,
                                QueueImplPtr Queue) {
@@ -695,13 +705,29 @@ Scheduler::GraphBuilder::addCG(std::unique_ptr<detail::CG> CommandGroup,
     printGraphAsDot("before_addCG");
 
   for (Requirement *Req : Reqs) {
-    MemObjRecord *Record = getOrInsertMemObjRecord(Queue, Req);
-    markModifiedIfWrite(Record, Req);
+    MemObjRecord *Record = nullptr;
+    AllocaCommandBase *AllocaCmd = nullptr;
+
+    bool isSameCtx = false;
+
+    {
+      const QueueImplPtr &QueueForAlloca =
+          isInteropHostTask(NewCmd)
+              ? static_cast<detail::CGHostTask &>(NewCmd->getCG()).MQueue
+              : Queue;
+
+      Record = getOrInsertMemObjRecord(QueueForAlloca, Req);
+      markModifiedIfWrite(Record, Req);
+
+      AllocaCmd = getOrCreateAllocaForReq(Record, Req, QueueForAlloca);
+
+      isSameCtx =
+          sameCtx(QueueForAlloca->getContextImplPtr(), Record->MCurContext);
+    }
 
-    AllocaCommandBase *AllocaCmd = getOrCreateAllocaForReq(Record, Req, Queue);
     // If there is alloca command we need to check if the latest memory is in
     // required context.
-    if (sameCtx(Queue->getContextImplPtr(), Record->MCurContext)) {
+    if (isSameCtx) {
       // If the memory is already in the required host context, check if the
       // required access mode is valid, remap if not.
       if (Record->MCurContext->is_host() &&
@@ -710,10 +736,24 @@ Scheduler::GraphBuilder::addCG(std::unique_ptr<detail::CG> CommandGroup,
     } else {
       // Cannot directly copy memory from OpenCL device to OpenCL device -
       // create two copies: device->host and host->device.
-      if (!Queue->is_host() && !Record->MCurContext->is_host())
+      bool NeedMemMoveToHost = false;
+      auto MemMoveTargetQueue = Queue;
+
+      if (isInteropHostTask(NewCmd)) {
+        const detail::CGHostTask &HT =
+            static_cast<detail::CGHostTask &>(NewCmd->getCG());
+
+        if (HT.MQueue->getContextImplPtr() != Record->MCurContext) {
+          NeedMemMoveToHost = true;
+          MemMoveTargetQueue = HT.MQueue;
+        }
+      } else if (!Queue->is_host() && !Record->MCurContext->is_host())
+        NeedMemMoveToHost = true;
+
+      if (NeedMemMoveToHost)
         insertMemoryMove(Record, Req,
                          Scheduler::getInstance().getDefaultHostQueue());
-      insertMemoryMove(Record, Req, Queue);
+      insertMemoryMove(Record, Req, MemMoveTargetQueue);
     }
     std::set<Command *> Deps =
         findDepsForReq(Record, Req, Queue->getContextImplPtr());
@@ -927,10 +967,11 @@ void Scheduler::GraphBuilder::connectDepEvent(Command *const Cmd,
   {
     std::unique_ptr<detail::HostTask> HT(new detail::HostTask);
     std::unique_ptr<detail::CG> ConnectCG(new detail::CGHostTask(
-        std::move(HT), /* Args = */ {}, /* ArgsStorage = */ {},
-        /* AccStorage = */ {}, /* SharedPtrStorage = */ {},
-        /* Requirements = */ {}, /* DepEvents = */ {DepEvent},
-        CG::CODEPLAY_HOST_TASK, /* Payload */ {}));
+        std::move(HT), /* Queue = */ {}, /* Context = */ {}, /* Args = */ {},
+        /* ArgsStorage = */ {}, /* AccStorage = */ {},
+        /* SharedPtrStorage = */ {}, /* Requirements = */ {},
+        /* DepEvents = */ {DepEvent}, CG::CODEPLAY_HOST_TASK,
+        /* Payload */ {}));
     ConnectCmd = new ExecCGCommand(
         std::move(ConnectCG), Scheduler::getInstance().getDefaultHostQueue());
   }

sycl/source/handler.cpp

@@ -85,9 +85,10 @@ event handler::finalize() {
     break;
   case detail::CG::CODEPLAY_HOST_TASK:
     CommandGroup.reset(new detail::CGHostTask(
-        std::move(MHostTask), std::move(MArgs), std::move(MArgsStorage),
-        std::move(MAccStorage), std::move(MSharedPtrStorage),
-        std::move(MRequirements), std::move(MEvents), MCGType, MCodeLoc));
+        std::move(MHostTask), MQueue, MQueue->getContextImplPtr(),
+        std::move(MArgs), std::move(MArgsStorage), std::move(MAccStorage),
+        std::move(MSharedPtrStorage), std::move(MRequirements),
+        std::move(MEvents), MCGType, MCodeLoc));
     break;
   case detail::CG::BARRIER:
   case detail::CG::BARRIER_WAITLIST:

sycl/test/abi/sycl_symbols_linux.dump

@@ -3320,6 +3320,10 @@ _ZNK2cl4sycl13host_selectorclERKNS0_6deviceE
 _ZNK2cl4sycl14exception_list3endEv
 _ZNK2cl4sycl14exception_list4sizeEv
 _ZNK2cl4sycl14exception_list5beginEv
+_ZNK2cl4sycl14interop_handle12getNativeMemEPNS0_6detail16AccessorImplHostE
+_ZNK2cl4sycl14interop_handle14getNativeQueueEv
+_ZNK2cl4sycl14interop_handle15getNativeDeviceEv
+_ZNK2cl4sycl14interop_handle16getNativeContextEv
 _ZNK2cl4sycl15device_selector13select_deviceEv
 _ZNK2cl4sycl15interop_handler12GetNativeMemEPNS0_6detail16AccessorImplHostE
 _ZNK2cl4sycl15interop_handler14GetNativeQueueEv