[Core][Streams] New API to wrap native backend streams (#525)

* [Core][Streams] New API to wrap native backend streams * [Core][Streams] Small changes based on review feedback * [OpenCL][Streams] Adding polyfill entry for clRetainCommandQueue Co-authored-by: Kris Rowe <kris.rowe@anl.gov>
libocca · Dec 6, 2021 · 1a4c374 · 1a4c374
1 parent 5f5ec0f
commit 1a4c374
Show file tree

Hide file tree

Showing 23 changed files with 111 additions and 16 deletions.
diff --git a/include/occa/core/device.hpp b/include/occa/core/device.hpp
@@ -390,6 +390,22 @@ namespace occa {
      */
     stream createStream(const occa::json &props = occa::json());
 
+    /**
+     * @startDoc{wrapStream}
+     *
+     * Description:
+     *   Wrap a native backend stream object inside a [[stream]] for the device.
+     *   The simplest example would be on a `CUDA` device, where a pointer to a cuStream_t, created via cudaStreamCreate, is passed in.
+     *
+     *   > Note that automatic garbage collection is not set for wrapped stream objects.
+     *
+     * Returns:
+     *   The wrapped [[stream]]
+     *
+     * @endDoc
+     */
+    stream wrapStream(void* ptr, const occa::json &props = occa::json());
+
     /**
      * @startDoc{getStream}
      *

diff --git a/src/core/device.cpp b/src/core/device.cpp
@@ -256,6 +256,14 @@ namespace occa {
     return modeDevice->createStream(streamProperties(props));
   }
 
+  stream device::wrapStream(void* ptr, const occa::json &props) {
+    assertInitialized();
+
+    occa::json streamProps = streamProperties(props);
+
+    return modeDevice->wrapStream(ptr, streamProps);
+  }
+
   stream device::getStream() {
     assertInitialized();
     return modeDevice->currentStream;

diff --git a/src/occa/internal/core/device.hpp b/src/occa/internal/core/device.hpp
@@ -71,6 +71,7 @@ namespace occa {
 
     //  |---[ Stream ]------------------
     virtual modeStream_t* createStream(const occa::json &props) = 0;
+    virtual modeStream_t* wrapStream(void *ptr, const occa::json &props) = 0;
 
     virtual streamTag tagStream() = 0;
     virtual void waitFor(streamTag tag) = 0;

diff --git a/src/occa/internal/modes/cuda/device.cpp b/src/occa/internal/modes/cuda/device.cpp
@@ -149,6 +149,12 @@ namespace occa {
       return new stream(this, props, cuStream);
     }
 
+    modeStream_t* device::wrapStream(void* ptr, const occa::json &props) {
+      OCCA_ERROR("A nullptr was passed to cuda::device::wrapStream",nullptr != ptr);
+      CUstream cuStream = *static_cast<CUstream*>(ptr);
+      return new stream(this, props, cuStream, true);
+    }
+
     occa::streamTag device::tagStream() {
       CUevent cuEvent = NULL;
 

diff --git a/src/occa/internal/modes/cuda/device.hpp b/src/occa/internal/modes/cuda/device.hpp
@@ -46,6 +46,7 @@ namespace occa {
 
       //---[ Stream ]-------------------
       virtual modeStream_t* createStream(const occa::json &props);
+      virtual modeStream_t* wrapStream(void* ptr, const occa::json &props);
 
       virtual streamTag tagStream();
       virtual void waitFor(streamTag tag);

diff --git a/src/occa/internal/modes/cuda/stream.cpp b/src/occa/internal/modes/cuda/stream.cpp
@@ -5,15 +5,18 @@ namespace occa {
   namespace cuda {
     stream::stream(modeDevice_t *modeDevice_,
                    const occa::json &properties_,
-                   CUstream cuStream_) :
+                   CUstream cuStream_, bool isWrapped_) :
       modeStream_t(modeDevice_, properties_),
-      cuStream(cuStream_) {}
+      cuStream(cuStream_),
+      isWrapped(isWrapped_) {}
 
     stream::~stream() {
-      OCCA_CUDA_DESTRUCTOR_ERROR(
-        "Device: freeStream",
-        cuStreamDestroy(cuStream)
-      );
+      if (!isWrapped) {
+        OCCA_CUDA_DESTRUCTOR_ERROR(
+          "Device: freeStream",
+          cuStreamDestroy(cuStream)
+        );
+      }
     }
   }
 }
diff --git a/src/occa/internal/modes/cuda/stream.hpp b/src/occa/internal/modes/cuda/stream.hpp
@@ -12,9 +12,12 @@ namespace occa {
     public:
       CUstream cuStream;
 
+      bool isWrapped;
+
       stream(modeDevice_t *modeDevice_,
              const occa::json &properties_,
-             CUstream cuStream_);
+             CUstream cuStream_,
+             bool isWrapped_=false);
 
       virtual ~stream();
     };

diff --git a/src/occa/internal/modes/dpcpp/device.cpp b/src/occa/internal/modes/dpcpp/device.cpp
@@ -88,6 +88,12 @@ namespace occa
       return new occa::dpcpp::stream(this, props, q);
     }
 
+    modeStream_t* device::wrapStream(void* ptr, const occa::json &props) {
+      OCCA_ERROR("A nullptr was passed to dpcpp::device::wrapStream",nullptr != ptr);
+      ::sycl::queue q = *static_cast<::sycl::queue*>(ptr);
+      return new stream(this, props, q);
+    }
+
     occa::streamTag device::tagStream()
     {
       //@note: This creates a host event which will return immediately.

diff --git a/src/occa/internal/modes/dpcpp/device.hpp b/src/occa/internal/modes/dpcpp/device.hpp
@@ -37,6 +37,7 @@ namespace occa
 
       //---[ Stream ]-------------------
       virtual modeStream_t *createStream(const occa::json &props) override;
+      virtual modeStream_t* wrapStream(void* ptr, const occa::json &props) override;
 
       virtual occa::streamTag tagStream() override;
       virtual void waitFor(occa::streamTag tag) override;

diff --git a/src/occa/internal/modes/hip/device.cpp b/src/occa/internal/modes/hip/device.cpp
@@ -133,6 +133,12 @@ namespace occa {
       return new stream(this, props, hipStream);
     }
 
+    modeStream_t* device::wrapStream(void* ptr, const occa::json &props) {
+      OCCA_ERROR("A nullptr was passed to hip::device::wrapStream",nullptr != ptr);
+      hipStream_t hipStream = *static_cast<hipStream_t*>(ptr);
+      return new stream(this, props, hipStream);
+    }
+
     occa::streamTag device::tagStream() {
       hipEvent_t hipEvent = NULL;
 

diff --git a/src/occa/internal/modes/hip/device.hpp b/src/occa/internal/modes/hip/device.hpp
@@ -37,6 +37,7 @@ namespace occa {
 
       //---[ Stream ]-------------------
       virtual modeStream_t* createStream(const occa::json &props);
+      virtual modeStream_t* wrapStream(void* ptr, const occa::json &props);
 
       virtual streamTag tagStream();
       virtual void waitFor(streamTag tag);

diff --git a/src/occa/internal/modes/hip/stream.cpp b/src/occa/internal/modes/hip/stream.cpp
@@ -5,13 +5,17 @@ namespace occa {
   namespace hip {
     stream::stream(modeDevice_t *modeDevice_,
                    const occa::json &properties_,
-                   hipStream_t hipStream_) :
+                   hipStream_t hipStream_,
+                   bool isWrapped_) :
       modeStream_t(modeDevice_, properties_),
-      hipStream(hipStream_) {}
+      hipStream(hipStream_),
+      isWrapped(isWrapped_) {}
 
     stream::~stream() {
-      OCCA_HIP_ERROR("Device: freeStream",
-                      hipStreamDestroy(hipStream));
+      if (!isWrapped) {
+        OCCA_HIP_ERROR("Device: freeStream",
+                        hipStreamDestroy(hipStream));
+      }
     }
   }
 }
diff --git a/src/occa/internal/modes/hip/stream.hpp b/src/occa/internal/modes/hip/stream.hpp
@@ -10,9 +10,12 @@ namespace occa {
     public:
       hipStream_t hipStream;
 
+      bool isWrapped;
+
       stream(modeDevice_t *modeDevice_,
              const occa::json &properties_,
-             hipStream_t hipStream_);
+             hipStream_t hipStream_,
+             bool isWrapped_=false);
 
       virtual ~stream();
     };

diff --git a/src/occa/internal/modes/metal/device.cpp b/src/occa/internal/modes/metal/device.cpp
@@ -74,6 +74,12 @@ namespace occa {
       return new stream(this, props, metalCommandQueue);
     }
 
+    modeStream_t* device::wrapStream(void* ptr, const occa::json &props) {
+      OCCA_ERROR("A nullptr was passed to metal::device::wrapStream",nullptr != ptr);
+      api::metal::commandQueue_t q = *static_cast<api::metal::commandQueue_t*>(ptr);
+      return new stream(this, props, q, true);
+    }
+
     occa::streamTag device::tagStream() {
       metal::stream &stream = (
         *((metal::stream*) (currentStream.getModeStream()))

diff --git a/src/occa/internal/modes/metal/device.hpp b/src/occa/internal/modes/metal/device.hpp
@@ -34,6 +34,7 @@ namespace occa {
 
       //---[ Stream ]-------------------
       virtual modeStream_t* createStream(const occa::json &props);
+      virtual modeStream_t* wrapStream(void* ptr, const occa::json &props);
 
       virtual streamTag tagStream();
       virtual void waitFor(streamTag tag);

diff --git a/src/occa/internal/modes/metal/stream.cpp b/src/occa/internal/modes/metal/stream.cpp
@@ -4,12 +4,16 @@ namespace occa {
   namespace metal {
     stream::stream(modeDevice_t *modeDevice_,
                    const occa::json &properties_,
-                   api::metal::commandQueue_t metalCommandQueue_) :
+                   api::metal::commandQueue_t metalCommandQueue_,
+                   bool isWrapped_) :
         modeStream_t(modeDevice_, properties_),
-        metalCommandQueue(metalCommandQueue_) {}
+        metalCommandQueue(metalCommandQueue_),
+        isWrapped(isWrapped_) {}
 
     stream::~stream() {
-      metalCommandQueue.free();
+      if (!isWrapped) {
+        metalCommandQueue.free();
+      }
     }
   }
 }
diff --git a/src/occa/internal/modes/metal/stream.hpp b/src/occa/internal/modes/metal/stream.hpp
@@ -10,9 +10,12 @@ namespace occa {
     public:
       api::metal::commandQueue_t metalCommandQueue;
 
+      bool isWrapped;
+
       stream(modeDevice_t *modeDevice_,
              const occa::json &properties_,
-             api::metal::commandQueue_t metalCommandQueue_);
+             api::metal::commandQueue_t metalCommandQueue_,
+             bool isWrapped_=false);
 
       virtual ~stream();
     };

diff --git a/src/occa/internal/modes/opencl/device.cpp b/src/occa/internal/modes/opencl/device.cpp
@@ -107,6 +107,16 @@ namespace occa {
       return new stream(this, props, commandQueue);
     }
 
+    modeStream_t* device::wrapStream(void* ptr, const occa::json &props) {
+      OCCA_ERROR("A nullptr was passed to opencl::device::wrapStream",nullptr != ptr);
+
+      cl_command_queue commandQueue = *static_cast<cl_command_queue*>(ptr);
+      OCCA_OPENCL_ERROR("Device: Retaining Command Queue",
+                        clRetainCommandQueue(commandQueue));
+
+      return new stream(this, props, commandQueue);
+    }
+
     occa::streamTag device::tagStream() {
       cl_event clEvent;
 

diff --git a/src/occa/internal/modes/opencl/device.hpp b/src/occa/internal/modes/opencl/device.hpp
@@ -35,6 +35,7 @@ namespace occa {
 
       //---[ Stream ]-------------------
       virtual modeStream_t* createStream(const occa::json &props);
+      virtual modeStream_t* wrapStream(void* ptr, const occa::json &props);
 
       virtual streamTag tagStream();
       virtual void waitFor(streamTag tag);

diff --git a/src/occa/internal/modes/opencl/polyfill.hpp b/src/occa/internal/modes/opencl/polyfill.hpp
@@ -428,6 +428,10 @@ namespace occa {
     return NULL;
   }
 
+  inline cl_int clRetainCommandQueue(cl_command_queue command_queue) {
+    return OCCA_OPENCL_IS_NOT_ENABLED;
+  }
+
   inline cl_int clReleaseCommandQueue(cl_command_queue command_queue) {
     return OCCA_OPENCL_IS_NOT_ENABLED;
   }

diff --git a/src/occa/internal/modes/opencl/stream.hpp b/src/occa/internal/modes/opencl/stream.hpp
@@ -10,6 +10,8 @@ namespace occa {
     public:
       cl_command_queue commandQueue;
 
+      bool isWrapped;
+
       stream(modeDevice_t *modeDevice_,
              const occa::json &properties_,
              cl_command_queue commandQueue_);

diff --git a/src/occa/internal/modes/serial/device.cpp b/src/occa/internal/modes/serial/device.cpp
@@ -47,6 +47,10 @@ namespace occa {
       return new stream(this, props);
     }
 
+    modeStream_t* device::wrapStream(void* ptr, const occa::json &props) {
+      return new stream(this, props);
+    }
+
     occa::streamTag device::tagStream() {
       return new occa::serial::streamTag(this, sys::currentTime());
     }

diff --git a/src/occa/internal/modes/serial/device.hpp b/src/occa/internal/modes/serial/device.hpp
@@ -23,6 +23,7 @@ namespace occa {
 
       //---[ Stream ]-------------------
       virtual modeStream_t* createStream(const occa::json &props);
+      virtual modeStream_t* wrapStream(void* ptr, const occa::json &props);
 
       virtual streamTag tagStream();
       virtual void waitFor(streamTag tag);