implement a function for stream to wait for an event without blocking…

… the host (ex. cuStreamWaitEvent)

examples/cpp/CMakeLists.txt

@@ -14,6 +14,7 @@ add_subdirectory(14_cuda_interop)
 add_subdirectory(17_memory_pool)
 
 add_subdirectory(18_nonblocking_streams)
+add_subdirectory(19_stream_tags)
 add_subdirectory(20_native_dpcpp_kernel)
 add_subdirectory(30_device_function)
 

include/occa/core/base.hpp

@@ -42,6 +42,7 @@ namespace occa {
   stream createStream(const occa::json &props = occa::json());
   stream getStream();
   void setStream(stream s);
+  void streamWait(stream s, streamTag tag);
 
   streamTag tagStream();
 

include/occa/core/stream.hpp

@@ -12,6 +12,7 @@
 namespace occa {
   class modeStream_t; class stream;
   class modeDevice_t; class device;
+  class streamTag;
 
   /**
    * @startDoc{stream}
@@ -148,6 +149,17 @@ namespace occa {
      */
     void finish();
 
+     /**
+     * @startDoc{waitFor}
+     *
+     * Description:
+     *   Waits for any operations submitted before a given streamg tag recording
+     *   to complete where the stream tag may be created on a different stream.
+     *
+     * @endDoc
+     */
+    void waitFor(occa::streamTag tag);
+
     /**
      * @startDoc{unwrap}
      * 

src/core/base.cpp

@@ -94,6 +94,10 @@ namespace occa {
     return getDevice().tagStream();
   }
 
+  void streamWait(stream s, streamTag tag) {
+    s.waitFor(tag);
+  }
+
   memoryPool createMemoryPool(const occa::json &props) {
     return getDevice().createMemoryPool(props);
   }

src/core/stream.cpp

@@ -2,6 +2,7 @@
 #include <occa/core/device.hpp>
 #include <occa/internal/core/device.hpp>
 #include <occa/internal/core/stream.hpp>
+#include <occa/internal/core/streamTag.hpp>
 
 namespace occa {
   stream::stream() :
@@ -102,6 +103,10 @@ namespace occa {
     if(modeStream) modeStream->finish();
   }
 
+  void stream::waitFor(occa::streamTag tag) {
+    if(modeStream) modeStream->waitFor(tag);
+  }
+
   void* stream::unwrap() {
     OCCA_ERROR(
       "stream::unwrap: stream is uninitialized or has been free'd",

src/occa/internal/api/metal/commandQueue.hpp

@@ -34,6 +34,7 @@ namespace occa {
         void freeLastCommandBuffer();
 
         event_t createEvent() const;
+        void waitForEvent(const event_t &event);
 
         void clearCommandBuffer(void *commandBufferObj);
         void setLastCommandBuffer(void *commandBufferObj);

src/occa/internal/api/metal/commandQueue.mm

@@ -77,6 +77,16 @@
                        lastCommandBufferObj);
       }
 
+      void commandQueue_t::waitForEvent(const event_t &event) {
+        if (lastCommandBufferObj) {
+          id<MTLEvent> metalEvent = (__bridge id<MTLEvent>) event.eventObj;
+          id<MTLCommandBuffer> metalCommandBuffer = (
+            (__bridge id<MTLCommandBuffer>) lastCommandBufferObj
+          );
+          [metalCommandBuffer encodeWaitForEvent:metalEvent value:event.signalValue];
+        }
+      }
+
       void commandQueue_t::clearCommandBuffer(void *commandBufferObj) {
         if (commandBufferObj == lastCommandBufferObj) {
           freeLastCommandBuffer();

src/occa/internal/api/metal/event.hpp

@@ -13,6 +13,7 @@ namespace occa {
         int eventId;
         void *commandBufferObj;
         double eventTime;
+        int signalValue;
 
         event_t();
 

src/occa/internal/api/metal/event.mm

@@ -25,10 +25,17 @@
         eventObj(eventObj_),
         eventId(eventId_),
         commandBufferObj(commandBufferObj_),
-        eventTime(0) {
+        eventTime(0),
+        signalValue(1) {
           // If there are no active command buffers, use the current time
           if (!commandBufferObj) {
             eventTime = occa::sys::currentTime();
+          } else {
+            id<MTLEvent> metalEvent = (__bridge id<MTLEvent>) eventObj;
+            id<MTLCommandBuffer> metalCommandBuffer = (
+              (__bridge id<MTLCommandBuffer>) commandBufferObj
+            );
+            [metalCommandBuffer encodeSignalEvent:metalEvent value:signalValue];
           }
         }
 
@@ -37,14 +44,16 @@
         eventId(other.eventId),
         eventObj(other.eventObj),
         commandBufferObj(other.commandBufferObj),
-        eventTime(other.eventTime) {}
+        eventTime(other.eventTime),
+        signalValue(other.signalValue) {}
 
         event_t& event_t::operator = (const event_t &other) {
         commandQueue = other.commandQueue;
         eventId = other.eventId;
         eventObj = other.eventObj;
         commandBufferObj = other.commandBufferObj;
         eventTime = other.eventTime;
+        signalValue = other.signalValue;
         return *this;
       }
 

src/occa/internal/api/metal/polyfill.cpp

@@ -73,6 +73,7 @@ namespace occa {
       event_t commandQueue_t::createEvent() const {
         return event_t();
       }
+      void commandQueue_t::waitForEvent(const event_t &event) {}
 
       void commandQueue_t::clearCommandBuffer(void *commandBufferObj) {}
       void commandQueue_t::setLastCommandBuffer(void *commandBufferObj) {}

src/occa/internal/core/stream.hpp

@@ -25,6 +25,7 @@ namespace occa {
 
     //---[ Virtual Methods ]------------
     virtual void finish() = 0;
+    virtual void waitFor(streamTag tag) = 0;
 
     virtual void* unwrap() = 0;
     //==================================

src/occa/internal/modes/cuda/polyfill.hpp

@@ -275,6 +275,10 @@ namespace occa {
   inline CUresult cuStreamSynchronize(CUstream hStream) {
     return OCCA_CUDA_IS_NOT_ENABLED;
   }
+
+  inline CUresult cuStreamWaitEvent(CUstream hStream, CUevent hEvent, unsigned int  Flags) {
+    return OCCA_CUDA_IS_NOT_ENABLED;
+  }
 }
 
 #endif

src/occa/internal/modes/cuda/stream.cpp

@@ -1,4 +1,5 @@
 #include <occa/internal/modes/cuda/stream.hpp>
+#include <occa/internal/modes/cuda/streamTag.hpp>
 #include <occa/internal/modes/cuda/utils.hpp>
 
 namespace occa {
@@ -24,6 +25,14 @@ namespace occa {
                       cuStreamSynchronize(cuStream));
     }
 
+    void stream::waitFor(occa::streamTag tag) {
+      occa::cuda::streamTag *cuTag = (
+        dynamic_cast<occa::cuda::streamTag*>(tag.getModeStreamTag())
+      );
+      OCCA_CUDA_ERROR("Stream: waitFor",
+                      cuStreamWaitEvent(cuStream, cuTag->cuEvent, 0));
+    }
+
     void* stream::unwrap() {
       return static_cast<void*>(&cuStream);
     }

src/occa/internal/modes/cuda/stream.hpp

@@ -21,6 +21,7 @@ namespace occa {
 
       virtual ~stream();
       void finish() override;
+      void waitFor(occa::streamTag tag) override;
 
       void* unwrap() override;
     };

src/occa/internal/modes/dpcpp/polyfill.hpp

@@ -295,6 +295,7 @@ class queue {
   }
 
   sycl::event ext_oneapi_submit_barrier()  { return sycl::event();}
+  sycl::event ext_oneapi_submit_barrier( const std::vector<sycl::event> &waitList ) { return sycl::event(); }
 };
 
 inline void* malloc_device(size_t num_bytes, 

src/occa/internal/modes/dpcpp/stream.cpp

@@ -16,6 +16,13 @@ namespace occa {
                        commandQueue.wait_and_throw());
     }
 
+    void stream::waitFor(occa::streamTag tag)
+    {
+      auto& dpcppTag{getDpcppStreamTag(tag)};
+      OCCA_DPCPP_ERROR("stream::waitFor",
+                       commandQueue.ext_oneapi_submit_barrier( std::vector<sycl::event>{dpcppTag.dpcppEvent} ));
+    }
+
     occa::dpcpp::streamTag stream::memcpy(void * dest,const void* src, occa::udim_t num_bytes)
     {
       ::sycl::event e{commandQueue.memcpy(dest, src, num_bytes)};

src/occa/internal/modes/dpcpp/stream.hpp

@@ -19,6 +19,7 @@ namespace occa {
       virtual ~stream()=default;
 
       void finish() override;
+      void waitFor(occa::streamTag tag) override;
 
       void* unwrap() override;
 

src/occa/internal/modes/hip/polyfill.hpp

@@ -284,6 +284,10 @@ namespace occa {
   inline hipError_t hipStreamSynchronize(hipStream_t hStream) {
     return OCCA_HIP_IS_NOT_ENABLED;
   }
+
+  inline hipError_t hipStreamWaitEvent(hipStream_t stream, hipEvent_t event, unsigned int flags) {
+    return OCCA_HIP_IS_NOT_ENABLED;
+  }
 }
 
 #endif

src/occa/internal/modes/hip/stream.cpp

@@ -1,4 +1,5 @@
 #include <occa/internal/modes/hip/stream.hpp>
+#include <occa/internal/modes/hip/streamTag.hpp>
 #include <occa/internal/modes/hip/utils.hpp>
 
 namespace occa {
@@ -23,6 +24,14 @@ namespace occa {
                      hipStreamSynchronize(hipStream));
     }
 
+    void stream::waitFor(occa::streamTag tag) {
+      occa::hip::streamTag *hipTag = (
+        dynamic_cast<occa::hip::streamTag*>(tag.getModeStreamTag())
+      );
+      OCCA_HIP_ERROR("Stream: waitFor",
+                     hipStreamWaitEvent(hipStream, hipTag->hipEvent, 0));
+    }
+
     void* stream::unwrap() {
       return static_cast<void*>(&hipStream);
     }

src/occa/internal/modes/hip/stream.hpp

@@ -19,6 +19,7 @@ namespace occa {
 
       virtual ~stream();
       void finish() override;
+      void waitFor(occa::streamTag tag) override;
 
       void* unwrap() override;
     };

src/occa/internal/modes/metal/stream.cpp

@@ -1,4 +1,5 @@
 #include <occa/internal/modes/metal/stream.hpp>
+#include <occa/internal/modes/metal/streamTag.hpp>
 
 namespace occa {
   namespace metal {
@@ -20,6 +21,13 @@ namespace occa {
       metalCommandQueue.finish();
     }
 
+    void stream::waitFor(occa::streamTag tag) {
+      occa::metal::streamTag *metalTag = (
+        dynamic_cast<occa::metal::streamTag*>(tag.getModeStreamTag())
+      );
+      metalCommandQueue.waitForEvent(metalTag->metalEvent);
+    }
+
     void* stream::unwrap() {
       return static_cast<void*>(&metalCommandQueue);
     }

src/occa/internal/modes/metal/stream.hpp

@@ -19,6 +19,7 @@ namespace occa {
 
       virtual ~stream();
       void finish() override;
+      void waitFor(occa::streamTag tag) override;
 
       void* unwrap() override;
     };

src/occa/internal/modes/opencl/stream.cpp

@@ -1,4 +1,5 @@
 #include <occa/internal/modes/opencl/stream.hpp>
+#include <occa/internal/modes/opencl/streamTag.hpp>
 #include <occa/internal/modes/opencl/utils.hpp>
 
 namespace occa {
@@ -19,6 +20,15 @@ namespace occa {
                         clFinish(commandQueue));
     }
 
+    void stream::waitFor(occa::streamTag tag) {
+      occa::opencl::streamTag *clTag = (
+        dynamic_cast<occa::opencl::streamTag*>(tag.getModeStreamTag())
+      );
+      OCCA_OPENCL_ERROR("Stream: waitFor",
+                        clEnqueueBarrierWithWaitList(commandQueue,
+                                                     1, &(clTag->clEvent), NULL));
+    }
+
     void* stream::unwrap() {
       return static_cast<void*>(&commandQueue);
     }

src/occa/internal/modes/opencl/stream.hpp

@@ -18,6 +18,7 @@ namespace occa {
 
       virtual ~stream();
       void finish() override;
+      void waitFor(occa::streamTag tag) override;
 
       void* unwrap() override;
     };

src/occa/internal/modes/serial/stream.cpp

@@ -1,4 +1,5 @@
 #include <occa/internal/modes/serial/stream.hpp>
+#include <occa/internal/modes/serial/streamTag.hpp>
 
 namespace occa {
   namespace serial {
@@ -8,6 +9,7 @@ namespace occa {
 
     stream::~stream() {}
     void stream::finish() {}
+    void stream::waitFor(occa::streamTag tag) {}
 
     void* stream::unwrap() {
       OCCA_FORCE_ERROR("stream::unwrap is not defined for serial mode");

src/occa/internal/modes/serial/stream.hpp

@@ -13,6 +13,7 @@ namespace occa {
 
       virtual ~stream();
       void finish() override;
+      void waitFor(occa::streamTag tag) override;
 
       void* unwrap() override;
     };