[SYCL] Store stream buffers in the scheduler (#2416)

Stream buffers need to be alive after submitting a kernel because it is
executed by the scheduler asynchronously. For this reason currently
stream buffers are stored in an associated stream object. This stream
object is passed to the handler and then forwarded further to a command
group to keep stream buffers alive for the scheduler.

But there is a problem with this approach. A command group cannot be
destroyed while stream buffers (which are accessed in this command
group) are alive. Stream buffers are destroyed only if the stream
is destroyed. Stream object is destroyed only if command group is
destroyed. So, there is a loop dependency. Which results in memory leaks.

Solution is to store stream buffers in the scheduler for each stream.
With this approach resources are released properly.

Author: againull

sycl/source/detail/scheduler/scheduler.cpp

@@ -238,6 +238,19 @@ void Scheduler::enqueueLeavesOfReqUnlocked(const Requirement *const Req) {
   EnqueueLeaves(Record->MWriteLeaves);
 }
 
+void Scheduler::allocateStreamBuffers(stream_impl *Impl,
+                                      size_t StreamBufferSize,
+                                      size_t FlushBufferSize) {
+  std::lock_guard<std::mutex> lock(StreamBuffersPoolMutex);
+  StreamBuffersPool.insert(
+      {Impl, StreamBuffers(StreamBufferSize, FlushBufferSize)});
+}
+
+void Scheduler::deallocateStreamBuffers(stream_impl *Impl) {
+  std::lock_guard<std::mutex> lock(StreamBuffersPoolMutex);
+  StreamBuffersPool.erase(Impl);
+}
+
 Scheduler::Scheduler() {
   sycl::device HostDevice;
   DefaultHostQueue = QueueImplPtr(

sycl/source/detail/scheduler/scheduler.hpp

@@ -715,6 +715,43 @@ class Scheduler {
 
   friend class Command;
   friend class DispatchHostTask;
+
+  /// Stream buffers structure.
+  ///
+  /// The structure contains all buffers for a stream object.
+  struct StreamBuffers {
+    StreamBuffers(size_t StreamBufferSize, size_t FlushBufferSize)
+        // Initialize stream buffer with zeros, this is needed for two reasons:
+        // 1. We don't need to care about end of line when printing out
+        // streamed data.
+        // 2. Offset is properly initialized.
+        : Data(StreamBufferSize, 0),
+          Buf(Data.data(), range<1>(StreamBufferSize),
+              {property::buffer::use_host_ptr()}),
+          FlushBuf(range<1>(FlushBufferSize)) {}
+
+    // Vector on the host side which is used to initialize the stream
+    // buffer
+    std::vector<char> Data;
+
+    // Stream buffer
+    buffer<char, 1> Buf;
+
+    // Global flush buffer
+    buffer<char, 1> FlushBuf;
+  };
+
+  friend class stream_impl;
+
+  // Protects stream buffers pool
+  std::mutex StreamBuffersPoolMutex;
+  std::map<stream_impl *, StreamBuffers> StreamBuffersPool;
+
+  /// Allocate buffers in the pool for a provided stream
+  void allocateStreamBuffers(stream_impl *, size_t, size_t);
+
+  /// Deallocate buffers in the pool for a provided stream
+  void deallocateStreamBuffers(stream_impl *);
 };
 
 } // namespace detail

sycl/source/detail/stream_impl.cpp

@@ -17,30 +17,63 @@ namespace detail {
 
 stream_impl::stream_impl(size_t BufferSize, size_t MaxStatementSize,
                          handler &CGH)
-    : BufferSize_(BufferSize), MaxStatementSize_(MaxStatementSize),
-      // Allocate additional place for the offset variable and the end of line
-      // symbol. Initialize buffer with zeros, this is needed for two reasons:
-      // 1. We don't need to care about end of line when printing out streamed
-      // data.
-      // 2. Offset is properly initialized.
-      Data(BufferSize + OffsetSize + 1, 0),
-      Buf(Data.data(), range<1>(BufferSize + OffsetSize + 1),
-          {property::buffer::use_host_ptr()}),
-
-      FlushBuf(range<1>(MaxStatementSize)) {}
+    : BufferSize_(BufferSize), MaxStatementSize_(MaxStatementSize) {
+  // We need to store stream buffers in the scheduler because they need to be
+  // alive after submitting the kernel. They cannot be stored in the stream
+  // object because it causes loop dependency between objects and results in
+  // memory leak.
+  // Allocate additional place in the stream buffer for the offset variable and
+  // the end of line symbol.
+  detail::Scheduler::getInstance().allocateStreamBuffers(
+      this, BufferSize + OffsetSize + 1 /* size of the stream buffer */,
+      MaxStatementSize /* size of the flush buffer */);
+}
+
+// Method to provide an access to the global stream buffer
+GlobalBufAccessorT stream_impl::accessGlobalBuf(handler &CGH) {
+  return detail::Scheduler::getInstance()
+      .StreamBuffersPool.find(this)
+      ->second.Buf.get_access<cl::sycl::access::mode::read_write>(
+          CGH, range<1>(BufferSize_), id<1>(OffsetSize));
+}
 
+// Method to provide an accessor to the global flush buffer
+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));
+}
+
+// Method to provide an atomic access to the offset in the global stream
+// buffer and offset in the flush buffer
+GlobalOffsetAccessorT stream_impl::accessGlobalOffset(handler &CGH) {
+  auto OffsetSubBuf = buffer<char, 1>(
+      detail::Scheduler::getInstance().StreamBuffersPool.find(this)->second.Buf,
+      id<1>(0), range<1>(OffsetSize));
+  auto ReinterpretedBuf = OffsetSubBuf.reinterpret<unsigned, 1>(range<1>(2));
+  return ReinterpretedBuf.get_access<cl::sycl::access::mode::atomic>(
+      CGH, range<1>(2), id<1>(0));
+}
 size_t stream_impl::get_size() const { return BufferSize_; }
 
 size_t stream_impl::get_max_statement_size() const { return MaxStatementSize_; }
 
 void stream_impl::flush() {
   // Access the stream buffer on the host. This access guarantees that kernel is
   // executed and buffer contains streamed data.
-  auto HostAcc = Buf.get_access<cl::sycl::access::mode::read>(
-      range<1>(BufferSize_), id<1>(OffsetSize));
+  {
+    auto HostAcc = detail::Scheduler::getInstance()
+                       .StreamBuffersPool.find(this)
+                       ->second.Buf.get_access<cl::sycl::access::mode::read>(
+                           range<1>(BufferSize_), id<1>(OffsetSize));
+
+    printf("%s", HostAcc.get_pointer());
+    fflush(stdout);
+  }
 
-  printf("%s", HostAcc.get_pointer());
-  fflush(stdout);
+  // Flushed the stream, can deallocate the buffers now.
+  detail::Scheduler::getInstance().deallocateStreamBuffers(this);
 }
 } // namespace detail
 } // namespace sycl

sycl/source/detail/stream_impl.hpp

@@ -26,25 +26,14 @@ class __SYCL_EXPORT stream_impl {
   stream_impl(size_t BufferSize, size_t MaxStatementSize, handler &CGH);
 
   // Method to provide an access to the global stream buffer
-  GlobalBufAccessorT accessGlobalBuf(handler &CGH) {
-    return Buf.get_access<cl::sycl::access::mode::read_write>(
-        CGH, range<1>(BufferSize_), id<1>(OffsetSize));
-  }
+  GlobalBufAccessorT accessGlobalBuf(handler &CGH);
 
   // Method to provide an accessor to the global flush buffer
-  GlobalBufAccessorT accessGlobalFlushBuf(handler &CGH) {
-    return FlushBuf.get_access<cl::sycl::access::mode::read_write>(
-        CGH, range<1>(MaxStatementSize_), id<1>(0));
-  }
+  GlobalBufAccessorT accessGlobalFlushBuf(handler &CGH);
 
   // Method to provide an atomic access to the offset in the global stream
   // buffer and offset in the flush buffer
-  GlobalOffsetAccessorT accessGlobalOffset(handler &CGH) {
-    auto OffsetSubBuf = buffer<char, 1>(Buf, id<1>(0), range<1>(OffsetSize));
-    auto ReinterpretedBuf = OffsetSubBuf.reinterpret<unsigned, 1>(range<1>(2));
-    return ReinterpretedBuf.get_access<cl::sycl::access::mode::atomic>(
-        CGH, range<1>(2), id<1>(0));
-  }
+  GlobalOffsetAccessorT accessGlobalOffset(handler &CGH);
 
   // Copy stream buffer to the host and print the contents
   void flush();
@@ -65,14 +54,6 @@ class __SYCL_EXPORT stream_impl {
   // 2 variables: offset in the stream buffer and offset in the flush buffer.
   static const size_t OffsetSize = 2 * sizeof(unsigned);
 
-  // Vector on the host side which is used to initialize the stream buffer
-  std::vector<char> Data;
-
-  // Stream buffer
-  buffer<char, 1> Buf;
-
-  // Global flush buffer
-  buffer<char, 1> FlushBuf;
 };
 
 } // namespace detail

sycl/test/abi/sycl_symbols_linux.dump

@@ -4161,3 +4161,6 @@ _ZNK2cl4sycl9exception11has_contextEv
 _ZNK2cl4sycl9exception4whatEv
 __sycl_register_lib
 __sycl_unregister_lib
+_ZN2cl4sycl6detail11stream_impl15accessGlobalBufERNS0_7handlerE
+_ZN2cl4sycl6detail11stream_impl20accessGlobalFlushBufERNS0_7handlerE
+_ZN2cl4sycl6detail11stream_impl18accessGlobalOffsetERNS0_7handlerE

sycl/test/basic_tests/stream/release_resources_test.cpp

@@ -0,0 +1,26 @@
+// RUN: %clangxx -fsycl -fsycl-targets=%sycl_triple %s -o %t.out
+// RUN: env SYCL_PI_TRACE=2 %CPU_RUN_PLACEHOLDER %t.out %CPU_CHECK_PLACEHOLDER
+// RUN: env SYCL_PI_TRACE=2 %GPU_RUN_PLACEHOLDER %t.out %GPU_CHECK_PLACEHOLDER
+// RUN: env SYCL_PI_TRACE=2 %ACC_RUN_PLACEHOLDER %t.out %ACC_CHECK_PLACEHOLDER
+
+// Check that buffer used by a stream object is released.
+
+#include <CL/sycl.hpp>
+
+using namespace cl::sycl;
+
+int main() {
+  {
+    queue Queue;
+
+    // CHECK:---> piMemRelease
+    Queue.submit([&](handler &CGH) {
+      stream Out(1024, 80, CGH);
+      CGH.parallel_for<class test_cleanup1>(
+          range<1>(2), [=](id<1> i) { Out << "Hello, World!" << endl; });
+    });
+    Queue.wait();
+  }
+
+  return 0;
+}