[SYCL] Store stream offset in flush buffer's memory (#2767)

Flush buffer's offset is not shared between copies of stream object.
It may lead to a loss or double printing of work items' statements
written in one of copied stream objects.
Use two first bytes of work item's flush buffer for storing buffer's
offset to share it between copies of stream object.

Signed-off-by: Mikhail Lychkov <[email protected]>

Author: mlychkov

sycl/include/CL/sycl/handler.hpp

@@ -65,6 +65,9 @@ template <typename T_Src, int Dims_Src, cl::sycl::access::mode AccessMode_Src,
           cl::sycl::access::placeholder IsPlaceholder_Dst>
 class __copyAcc2Acc;
 
+// For unit testing purposes
+class MockHandler;
+
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
 
@@ -1951,6 +1954,8 @@ class __SYCL_EXPORT handler {
   friend void detail::associateWithHandler(handler &,
                                            detail::AccessorBaseHost *,
                                            access::target);
+
+  friend class ::MockHandler;
 };
 } // namespace sycl
 } // __SYCL_INLINE_NAMESPACE(cl)

sycl/include/CL/sycl/stream.hpp

@@ -134,6 +134,7 @@ set(SYCL_SOURCES
     "detail/scheduler/commands.cpp"
     "detail/scheduler/leaves_collection.cpp"
     "detail/scheduler/scheduler.cpp"
+    "detail/scheduler/scheduler_helpers.cpp"
     "detail/scheduler/graph_processor.cpp"
     "detail/scheduler/graph_builder.cpp"
     "detail/spec_constant_impl.cpp"

sycl/source/CMakeLists.txt

@@ -11,6 +11,7 @@
 #include <detail/global_handler.hpp>
 #include <detail/queue_impl.hpp>
 #include <detail/scheduler/scheduler.hpp>
+#include <detail/scheduler/scheduler_helpers.hpp>
 #include <detail/stream_impl.hpp>
 
 #include <chrono>
@@ -72,6 +73,19 @@ EventImplPtr Scheduler::addCG(std::unique_ptr<detail::CG> CommandGroup,
   const bool IsKernel = CommandGroup->getType() == CG::KERNEL;
   const bool IsHostKernel = CommandGroup->getType() == CG::RUN_ON_HOST_INTEL;
   vector_class<StreamImplPtr> Streams;
+
+  if (IsKernel) {
+    Streams = ((CGExecKernel *)CommandGroup.get())->getStreams();
+    // Stream's flush buffer memory is mainly initialized in stream's __init
+    // method. However, this method is not available on host device.
+    // Initializing stream's flush buffer on the host side in a separate task.
+    if (Queue->is_host()) {
+      for (const StreamImplPtr &Stream : Streams) {
+        initStream(Stream, Queue);
+      }
+    }
+  }
+
   {
     std::unique_lock<std::shared_timed_mutex> Lock(MGraphLock, std::defer_lock);
     lockSharedTimedMutex(Lock);
@@ -102,9 +116,6 @@ EventImplPtr Scheduler::addCG(std::unique_ptr<detail::CG> CommandGroup,
       if (!Enqueued && EnqueueResultT::SyclEnqueueFailed == Res.MResult)
         throw runtime_error("Enqueue process failed.", PI_INVALID_OPERATION);
 
-      if (IsKernel)
-        Streams = ((ExecCGCommand *)NewCmd)->getStreams();
-
       // If there are no memory dependencies decouple and free the command.
       // Though, dismiss ownership of native kernel command group as it's
       // resources may be in use by backend and synchronization point here is

sycl/source/detail/scheduler/scheduler.cpp

@@ -764,6 +764,7 @@ class Scheduler {
   };
 
   friend class stream_impl;
+  friend void initStream(StreamImplPtr, QueueImplPtr);
 
   // Protects stream buffers pool
   std::recursive_mutex StreamBuffersPoolMutex;

sycl/source/detail/scheduler/scheduler.hpp

@@ -0,0 +1,55 @@
+//==-- scheduler_helpers.cpp - SYCL Scheduler helper functions --*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include <CL/sycl/queue.hpp>
+#include <detail/queue_impl.hpp>
+#include <detail/scheduler/scheduler.hpp>
+#include <detail/scheduler/scheduler_helpers.hpp>
+#include <detail/stream_impl.hpp>
+
+__SYCL_INLINE_NAMESPACE(cl) {
+namespace sycl {
+namespace detail {
+
+void initStream(StreamImplPtr Stream, QueueImplPtr Queue) {
+  Scheduler::StreamBuffers *StrBufs{};
+
+  {
+    std::lock_guard<std::recursive_mutex> lock(
+        Scheduler::getInstance().StreamBuffersPoolMutex);
+
+    auto StreamBuf =
+        Scheduler::getInstance().StreamBuffersPool.find(Stream.get());
+    assert((StreamBuf != Scheduler::getInstance().StreamBuffersPool.end()) &&
+           "Stream is unexpectedly not found in pool.");
+
+    StrBufs = StreamBuf->second;
+  }
+
+  assert(StrBufs && "No buffers for a stream.");
+
+  // Real size of full flush buffer is saved only in buffer_impl field of
+  // FlushBuf object.
+  size_t FlushBufSize = getSyclObjImpl(StrBufs->FlushBuf)->get_count();
+
+  auto Q = createSyclObjFromImpl<queue>(Queue);
+  Q.submit([&](handler &cgh) {
+    auto FlushBufAcc =
+        StrBufs->FlushBuf.get_access<access::mode::discard_write,
+                                     access::target::host_buffer>(
+            cgh, range<1>(FlushBufSize), id<1>(0));
+    cgh.codeplay_host_task([=] {
+      char *FlushBufPtr = FlushBufAcc.get_pointer();
+      std::memset(FlushBufPtr, 0, FlushBufAcc.get_size());
+    });
+  });
+}
+
+} // namespace detail
+} // namespace sycl
+} // __SYCL_INLINE_NAMESPACE(cl)

sycl/source/detail/scheduler/scheduler_helpers.cpp

@@ -0,0 +1,29 @@
+//==---------- scheduler_helpers.hpp - SYCL standard header file -----------==//
+//
+// 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/sycl/detail/defines_elementary.hpp>
+
+#include <memory>
+
+__SYCL_INLINE_NAMESPACE(cl) {
+namespace sycl {
+namespace detail {
+
+class stream_impl;
+class queue_impl;
+
+using StreamImplPtr = std::shared_ptr<detail::stream_impl>;
+using QueueImplPtr = std::shared_ptr<detail::queue_impl>;
+
+void initStream(StreamImplPtr Stream, QueueImplPtr Queue);
+
+} // namespace detail
+} // namespace sycl
+} // __SYCL_INLINE_NAMESPACE(cl)

sycl/source/detail/scheduler/scheduler_helpers.hpp

@@ -27,7 +27,7 @@ stream_impl::stream_impl(size_t BufferSize, size_t MaxStatementSize,
   // the end of line symbol.
   detail::Scheduler::getInstance().allocateStreamBuffers(
       this, BufferSize + OffsetSize + 1 /* size of the stream buffer */,
-      MaxStatementSize /* size of the flush buffer */);
+      MaxStatementSize + FLUSH_BUF_OFFSET_SIZE /* size of the flush buffer */);
 }
 
 // Method to provide an access to the global stream buffer
@@ -43,7 +43,7 @@ GlobalBufAccessorT stream_impl::accessGlobalFlushBuf(handler &CGH) {
   return detail::Scheduler::getInstance()
       .StreamBuffersPool.find(this)
       ->second->FlushBuf.get_access<cl::sycl::access::mode::read_write>(
-          CGH, range<1>(MaxStatementSize_), id<1>(0));
+          CGH, range<1>(MaxStatementSize_ + FLUSH_BUF_OFFSET_SIZE), id<1>(0));
 }
 
 // Method to provide an atomic access to the offset in the global stream
@@ -76,7 +76,7 @@ void stream_impl::flush() {
                 cgh, range<1>(BufferSize_), id<1>(OffsetSize));
     // Create accessor to the flush buffer even if not using it yet. Otherwise
     // kernel will be a leaf for the flush buffer and scheduler will not be able
-    // to cleanup the kernel. TODO: git rid of finalize method by using host
+    // to cleanup the kernel. TODO: get rid of finalize method by using host
     // accessor to the flush buffer.
     auto FlushBufHostAcc =
         detail::Scheduler::getInstance()

sycl/source/detail/stream_impl.cpp

@@ -23,6 +23,8 @@ namespace detail {
 
 class __SYCL_EXPORT stream_impl {
 public:
+  // TODO: Handler argument is not used in constructor.
+  // To be removed when API/ABI changes are allowed.
   stream_impl(size_t BufferSize, size_t MaxStatementSize, handler &CGH);
 
   // Method to provide an access to the global stream buffer

sycl/source/detail/stream_impl.hpp

@@ -13,14 +13,24 @@
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
 
+// Maximum possible size of a flush buffer statement in bytes
+static constexpr size_t MAX_STATEMENT_SIZE =
+    (1 << (CHAR_BIT * detail::FLUSH_BUF_OFFSET_SIZE)) - 1;
+
 stream::stream(size_t BufferSize, size_t MaxStatementSize, handler &CGH)
     : impl(std::make_shared<detail::stream_impl>(BufferSize, MaxStatementSize,
                                                  CGH)),
       GlobalBuf(impl->accessGlobalBuf(CGH)),
       GlobalOffset(impl->accessGlobalOffset(CGH)),
       // Allocate the flush buffer, which contains space for each work item
       GlobalFlushBuf(impl->accessGlobalFlushBuf(CGH)),
-      FlushBufferSize(MaxStatementSize) {
+      FlushBufferSize(MaxStatementSize + detail::FLUSH_BUF_OFFSET_SIZE) {
+  if (MaxStatementSize > MAX_STATEMENT_SIZE) {
+    throw sycl::invalid_parameter_error(
+        "Maximum statement size exceeds limit of " +
+            std::to_string(MAX_STATEMENT_SIZE) + " bytes.",
+        PI_INVALID_VALUE);
+  }
 
   // Save stream implementation in the handler so that stream will be alive
   // during kernel execution