[llvm] [Offload][Lang] Add internal StreamTy (PR #216374)

Sophia Herrmann via llvm-commits llvm-commits at lists.llvm.org
Tue Aug 18 10:18:34 PDT 2026


https://github.com/jellytabby updated https://github.com/llvm/llvm-project/pull/216374

>From a83c926fd9f214be99a65bdce52b9ff3413a6257 Mon Sep 17 00:00:00 2001
From: Sophia Herrmann <herrmann15 at llnl.gov>
Date: Thu, 13 Aug 2026 16:03:12 -0700
Subject: [PATCH] add interal StreamTy

---
 offload/languages/kernel/CMakeLists.txt       |   1 +
 .../languages/kernel/include/LanguageUtils.h  |  17 +-
 offload/languages/kernel/include/State.h      |  50 +++-
 offload/languages/kernel/include/Stream.h     |  33 +++
 .../languages/kernel/src/LanguageLaunch.cpp   |   4 +-
 .../languages/kernel/src/LanguageRuntime.cpp  |  16 +-
 offload/languages/kernel/src/State.cpp        | 273 ++++++++++++++++--
 offload/test/offloading/CUDA/stream_api.cu    |  43 ++-
 offload/test/offloading/HIP/stream_api.hip    |  41 ++-
 9 files changed, 412 insertions(+), 66 deletions(-)
 create mode 100644 offload/languages/kernel/include/Stream.h

diff --git a/offload/languages/kernel/CMakeLists.txt b/offload/languages/kernel/CMakeLists.txt
index ff0bb81c7f017..92a051f06729e 100644
--- a/offload/languages/kernel/CMakeLists.txt
+++ b/offload/languages/kernel/CMakeLists.txt
@@ -81,5 +81,6 @@ install(FILES
         ${CMAKE_CURRENT_SOURCE_DIR}/../include/kernel/LanguageRuntime.h
         ${CMAKE_CURRENT_SOURCE_DIR}/../include/kernel/UndefineLanguageNames.inc
         ${CMAKE_CURRENT_SOURCE_DIR}/include/LanguageLaunch.h
+        ${CMAKE_CURRENT_SOURCE_DIR}/include/Stream.h
         ${CMAKE_CURRENT_SOURCE_DIR}/include/Types.h
         DESTINATION ${CMAKE_INSTALL_PREFIX}/include/offload/kernel/)
diff --git a/offload/languages/kernel/include/LanguageUtils.h b/offload/languages/kernel/include/LanguageUtils.h
index 9726d0599a9f4..730857e00461f 100644
--- a/offload/languages/kernel/include/LanguageUtils.h
+++ b/offload/languages/kernel/include/LanguageUtils.h
@@ -12,6 +12,7 @@
 #include "LanguageRuntime.h"
 #include "OffloadAPI.h"
 #include "State.h"
+#include "Stream.h"
 
 /// Convert an ol_result_t to the active language's Error_t.
 static inline Error_t convertResult(ol_result_t Result) {
@@ -49,12 +50,26 @@ static inline Error_t convertAndSetLastError(ol_result_t Result) {
   return setLastError(convertResult(Result));
 }
 
+/// Convert between the language-facing opaque stream and the internal stream.
+static inline Stream_t makeLanguageStream(llvm::offload::StreamTy *Stream) {
+  return reinterpret_cast<Stream_t>(Stream);
+}
+
+static inline llvm::offload::StreamTy *getInternalStream(Stream_t Stream) {
+  return reinterpret_cast<llvm::offload::StreamTy *>(Stream);
+}
+
 /// Convert a Stream_t to an ol_queue_handle_t.
 static inline Error_t getQueueFromStream(Stream_t Stream,
                                          ol_queue_handle_t *Queue) {
   if (!Stream)
     return ErrorInvalidValue;
-  *Queue = reinterpret_cast<ol_queue_handle_t>(Stream);
+
+  llvm::offload::StreamTy *InternalStream = getInternalStream(Stream);
+  if (!llvm::offload::StateTy::isStreamRegistered(InternalStream))
+    return ErrorInvalidResourceHandle;
+
+  *Queue = InternalStream->Queue;
   return Success;
 }
 
diff --git a/offload/languages/kernel/include/State.h b/offload/languages/kernel/include/State.h
index 7315b2e710df5..bb5422b5d3bb6 100644
--- a/offload/languages/kernel/include/State.h
+++ b/offload/languages/kernel/include/State.h
@@ -10,10 +10,12 @@
 #define LLVM_OFFLOAD_LANGUAGES_KERNEL_INCLUDE_STATE_H
 
 #include "OffloadAPI.h"
+#include "Stream.h"
 #include "Types.h"
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cstdint>
@@ -51,9 +53,12 @@ using KernelIDTy = const void *;
 struct ThreadStateTy {
   ~ThreadStateTy();
 
-  /// Return the default queue for the current host thread
+  /// Return the default queue for the current host thread.
   static ol_queue_handle_t getDefaultQueue();
 
+  /// Return the default stream for the current host thread and device.
+  static StreamTy *getDefaultStream();
+
   /// Return the thread-local default device, or the first discovered device.
   static ol_device_handle_t getDefaultDevice();
 
@@ -75,17 +80,16 @@ struct ThreadStateTy {
   /// Return the pending kernel launch configuration for this thread.
   static CallConfigurationTy &getCallConfiguration();
 
-  /// Set the thread-local default device to \p Device and recreate its queue.
-  static void setDefaultDevice(ol_device_handle_t Device);
-
 private:
   static ThreadStateTy &get();
 
-  void createDefaultQueue(ol_device_handle_t Device);
+  StreamTy *getOrCreateDefaultStream(ol_device_handle_t Device);
+  void destroyDefaultStreams();
 
   int DefaultDevice = 0;
   uint32_t LastError = 0;
-  ol_queue_handle_t DefaultQueue = nullptr;
+  llvm::DenseMap<ol_device_handle_t, StreamTy *>
+      PerThreadDeviceDefaultStreamMap;
 
   CallConfigurationTy CC = {};
 
@@ -136,6 +140,27 @@ struct StateTy {
   /// Return the loaded program handle for binary image key \p ID.
   static ol_program_handle_t getProgram(const void *ID);
 
+  /// Return all streams currently known for \p Device.
+  static llvm::SmallPtrSet<StreamTy *, 8>
+  getDeviceStreams(ol_device_handle_t Device);
+
+  /// Return all explicitly created blocking streams for \p Device.
+  static llvm::SmallPtrSet<StreamTy *, 8>
+  getBlockingStreams(ol_device_handle_t Device);
+
+  /// Return true if \p Device has an existing legacy default stream.
+  static bool hasLegacyDefaultStream(ol_device_handle_t Device);
+
+  /// Create a stream for \p Device and register it with the process state.
+  static ol_result_t createStream(ol_device_handle_t Device, QueueKind Kind,
+                                  StreamTy **Stream);
+
+  /// Destroy \p Stream after removing it from the process state.
+  static ol_result_t destroyStream(StreamTy *Stream);
+
+  /// Return true if \p Stream is currently registered with the process state.
+  static bool isStreamRegistered(StreamTy *Stream);
+
 private:
   static StateTy &get();
   static StateTy *tryGet();
@@ -146,6 +171,12 @@ struct StateTy {
   void addDevice(ol_device_handle_t Device);
   void setHostDevice(ol_device_handle_t Device);
 
+  StreamTy *getOrCreateDefaultStream(ol_device_handle_t Device);
+  void destroyDefaultStreams();
+
+  void addStream(StreamTy *Stream);
+  void removeStream(StreamTy *Stream);
+
   void addKernel(KernelIDTy KernelID, ol_symbol_handle_t Kernel);
   void removeKernel(KernelIDTy KernelID);
   ol_symbol_handle_t lookupKernel(KernelIDTy KernelID);
@@ -154,14 +185,19 @@ struct StateTy {
   ol_program_handle_t removeProgram(const void *Binary);
   ol_program_handle_t lookupProgram(const void *Binary);
 
+  void destroyRegisteredStreams();
   void destroyRegisteredPrograms();
 
   llvm::DenseMap<const void *, ol_program_handle_t> BinaryRegisterMap;
   llvm::DenseMap<KernelIDTy, ol_symbol_handle_t> KernelMap;
   llvm::SmallVector<ol_device_handle_t, 8> Devices;
+  llvm::DenseMap<ol_device_handle_t, StreamTy *> DeviceDefaultStreamsMap;
+  llvm::DenseMap<ol_device_handle_t, llvm::SmallPtrSet<StreamTy *, 8>>
+      DeviceStreamsMap;
+  llvm::DenseMap<ol_device_handle_t, llvm::SmallPtrSet<StreamTy *, 8>>
+      DeviceBlockingStreamsMap;
 
   ol_context_handle_t Context = nullptr;
-  ol_queue_handle_t DefaultQueue = nullptr;
   ol_device_handle_t HostDevice = nullptr;
 
   StateTy();
diff --git a/offload/languages/kernel/include/Stream.h b/offload/languages/kernel/include/Stream.h
new file mode 100644
index 0000000000000..e66fc7000281d
--- /dev/null
+++ b/offload/languages/kernel/include/Stream.h
@@ -0,0 +1,33 @@
+//===-- Stream.h - Kernel language stream state ---------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_OFFLOAD_LANGUAGES_KERNEL_INCLUDE_STREAM_H
+#define LLVM_OFFLOAD_LANGUAGES_KERNEL_INCLUDE_STREAM_H
+
+#include "OffloadAPI.h"
+
+namespace llvm {
+namespace offload {
+
+enum class QueueKind {
+  LegacyDefault,
+  PerThreadDefault,
+  ExplicitBlocking,
+  ExplicitNonBlocking,
+};
+
+struct StreamTy {
+  ol_queue_handle_t Queue = nullptr;
+  ol_device_handle_t Device = nullptr;
+  QueueKind Kind = QueueKind::ExplicitBlocking;
+};
+
+} // namespace offload
+} // namespace llvm
+
+#endif // LLVM_OFFLOAD_LANGUAGES_KERNEL_INCLUDE_STREAM_H
diff --git a/offload/languages/kernel/src/LanguageLaunch.cpp b/offload/languages/kernel/src/LanguageLaunch.cpp
index 7199e3a3f10f0..cbbe9e41291b8 100644
--- a/offload/languages/kernel/src/LanguageLaunch.cpp
+++ b/offload/languages/kernel/src/LanguageLaunch.cpp
@@ -9,10 +9,12 @@
 #include "LanguageLaunch.h"
 #include "LanguageUtils.h"
 #include "State.h"
+#include "Stream.h"
 #include <cstdio>
 
 using RuntimeState = llvm::offload::StateTy;
 using ThreadState = llvm::offload::ThreadStateTy;
+using StreamTy = llvm::offload::StreamTy;
 
 static constexpr ol_error_struct_t InvalidKernelError = {
     OL_ERRC_INVALID_NULL_HANDLE, "kernel is not registered"};
@@ -54,7 +56,7 @@ ol_result_t __llvmLaunchKernelImpl(const char *KernelID, dim3 GridDim,
   LaunchSizeArgs.GroupSize.z = BlockDim.z;
   LaunchSizeArgs.DynSharedMemory = DynamicSharedMem;
 
-  ol_queue_handle_t Queue = Stream ? reinterpret_cast<ol_queue_handle_t>(Stream)
+  ol_queue_handle_t Queue = Stream ? reinterpret_cast<StreamTy *>(Stream)->Queue
                                    : ThreadState::getDefaultQueue();
 
   struct OffloadKernelArgs {
diff --git a/offload/languages/kernel/src/LanguageRuntime.cpp b/offload/languages/kernel/src/LanguageRuntime.cpp
index d294412d37e8f..c4a2cd6d18cb9 100644
--- a/offload/languages/kernel/src/LanguageRuntime.cpp
+++ b/offload/languages/kernel/src/LanguageRuntime.cpp
@@ -19,6 +19,7 @@
 
 #include "LanguageUtils.h"
 #include "State.h"
+#include "Stream.h"
 #include "Types.h"
 
 #include "OffloadAPI.h"
@@ -144,20 +145,17 @@ Error_t GetDeviceProperties(DeviceProp_t *DeviceProp, int DeviceNo) {
 }
 
 Error_t StreamCreate(Stream_t *Stream) {
-  ol_queue_handle_t Queue;
-  ol_result_t Result = olCreateQueue(RuntimeState::getContext(),
-                                     ThreadState::getDefaultDevice(), &Queue);
+  llvm::offload::StreamTy *StreamObj = nullptr;
+  ol_result_t Result = RuntimeState::createStream(
+      ThreadState::getDefaultDevice(),
+      llvm::offload::QueueKind::ExplicitBlocking, &StreamObj);
   if (Result == OL_SUCCESS)
-    *Stream = reinterpret_cast<Stream_t>(Queue);
+    *Stream = makeLanguageStream(StreamObj);
   return convertAndSetLastError(Result);
 }
 
 Error_t StreamDestroy(Stream_t Stream) {
-  ol_queue_handle_t Queue;
-  Error_t Err = getQueueFromStream(Stream, &Queue);
-  if (Err != Success)
-    return setLastError(Err);
-  ol_result_t Result = olDestroyQueue(Queue);
+  ol_result_t Result = RuntimeState::destroyStream(getInternalStream(Stream));
   return convertAndSetLastError(Result);
 }
 
diff --git a/offload/languages/kernel/src/State.cpp b/offload/languages/kernel/src/State.cpp
index baaceeb377bc7..eed2e53dc7607 100644
--- a/offload/languages/kernel/src/State.cpp
+++ b/offload/languages/kernel/src/State.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "State.h"
+#include "Stream.h"
 #include "Types.h"
 
 #include "OffloadAPI.h"
@@ -27,6 +28,15 @@ using namespace offload;
 // Weak so another runtime object can override the default stream mode.
 __attribute__((weak)) uint32_t PerThreadQueue = 0;
 
+static constexpr ol_error_struct_t InvalidNullPointerError = {
+    OL_ERRC_INVALID_NULL_POINTER, "invalid null stream pointer"};
+
+static constexpr ol_error_struct_t InvalidDeviceError = {OL_ERRC_INVALID_DEVICE,
+                                                         "invalid device"};
+
+static constexpr ol_error_struct_t InvalidStreamError = {OL_ERRC_INVALID_QUEUE,
+                                                         "invalid stream"};
+
 // Process-wide singleton and thread-state registry.
 static std::mutex &getStateLock() {
   static std::mutex StateLock;
@@ -43,6 +53,21 @@ static std::mutex &getThreadStatesLock() {
 using ThreadStatesTy = SmallVector<ThreadStateTy *, 64>;
 static ThreadStatesTy *ThreadStatesPtr = nullptr;
 
+static std::mutex &getDeviceDefaultStreamsMapLock() {
+  static std::mutex DeviceDefaultStreamsMapLock;
+  return DeviceDefaultStreamsMapLock;
+}
+
+static std::mutex &getDeviceStreamsMapLock() {
+  static std::mutex DeviceStreamsMapLock;
+  return DeviceStreamsMapLock;
+}
+
+static std::mutex &getDeviceBlockingStreamsMapLock() {
+  static std::mutex DeviceBlockingStreamsMapLock;
+  return DeviceBlockingStreamsMapLock;
+}
+
 static void deleteThreadStates() {
   // Detach the registry before deletion because deleteThreadState may be called
   // more than once via atexit and StateTy teardown.
@@ -61,31 +86,47 @@ static void deleteThreadStates() {
 
 static void deleteState() {
   StateTy *ST = StatePtr.load();
-  StatePtr.store(nullptr);
+  if (!ST)
+    return;
   delete ST;
-  StatePtr = nullptr;
+  StatePtr.store(nullptr);
 }
 
-static void destroyQueue(ol_queue_handle_t &Queue) {
-  if (!Queue)
+static void destroyStreamHandle(StreamTy *&Stream) {
+  if (!Stream)
     return;
 
-  olSyncQueue(Queue);
-  olDestroyQueue(Queue);
-  Queue = nullptr;
+  olSyncQueue(Stream->Queue);
+  olDestroyQueue(Stream->Queue);
+  delete Stream;
+  Stream = nullptr;
 }
 
 namespace llvm {
 namespace offload {
 
-// ThreadStateTy implementation.
+static bool removeStreamFromMap(
+    DenseMap<ol_device_handle_t, SmallPtrSet<StreamTy *, 8>> &StreamsMap,
+    StreamTy *Stream) {
+  bool Removed = false;
+  SmallVector<ol_device_handle_t, 8> EmptyDevices;
 
-ThreadStateTy::ThreadStateTy() {
-  if (PerThreadQueue) [[unlikely]]
-    createDefaultQueue(getDefaultDevice());
-  atexit(deleteThreadStates);
+  for (auto &It : StreamsMap) {
+    Removed |= It.second.erase(Stream);
+    if (It.second.empty())
+      EmptyDevices.push_back(It.first);
+  }
+
+  for (ol_device_handle_t Device : EmptyDevices)
+    StreamsMap.erase(Device);
+
+  return Removed;
 }
-ThreadStateTy::~ThreadStateTy() { destroyQueue(DefaultQueue); }
+
+// ThreadStateTy implementation.
+
+ThreadStateTy::ThreadStateTy() { atexit(deleteThreadStates); }
+ThreadStateTy::~ThreadStateTy() { destroyDefaultStreams(); }
 
 ThreadStateTy &ThreadStateTy::get() {
   auto *&TS = ThreadState;
@@ -100,14 +141,28 @@ ThreadStateTy &ThreadStateTy::get() {
 }
 
 ol_device_handle_t ThreadStateTy::getDefaultDevice() {
+  ArrayRef<ol_device_handle_t> Devices = StateTy::get().getDevices();
   int DD = ThreadStateTy::get().DefaultDevice;
-  return StateTy::get().getDevices()[DD];
+  if (DD < 0 || DD >= static_cast<int>(Devices.size()))
+    return nullptr;
+  return Devices[DD];
 }
 
-ol_queue_handle_t ThreadStateTy::getDefaultQueue() {
+StreamTy *ThreadStateTy::getDefaultStream() {
+  ol_device_handle_t Device = getDefaultDevice();
+  if (!Device)
+    return nullptr;
+
   if (!PerThreadQueue) [[likely]]
-    return StateTy::get().DefaultQueue;
-  return ThreadStateTy::get().DefaultQueue;
+    return StateTy::get().getOrCreateDefaultStream(Device);
+
+  return ThreadStateTy::get().getOrCreateDefaultStream(Device);
+}
+
+ol_queue_handle_t ThreadStateTy::getDefaultQueue() {
+  if (StreamTy *Stream = getDefaultStream())
+    return Stream->Queue;
+  return nullptr;
 }
 
 CallConfigurationTy &ThreadStateTy::getCallConfiguration() {
@@ -119,9 +174,7 @@ ol_device_handle_t ThreadStateTy::setDefaultDevice(int DeviceNo) {
   if (DeviceNo < 0 || DeviceNo >= static_cast<int>(Devices.size()))
     return nullptr;
   ThreadStateTy::get().DefaultDevice = DeviceNo;
-  ol_device_handle_t DD = Devices[DeviceNo];
-  ThreadStateTy::get().createDefaultQueue(DD);
-  return DD;
+  return Devices[DeviceNo];
 }
 
 ol_device_handle_t ThreadStateTy::getDevice(int *DeviceNo) {
@@ -137,11 +190,32 @@ uint32_t ThreadStateTy::setLastError(uint32_t Error) {
   return ThreadStateTy::get().LastError = Error;
 }
 
-void ThreadStateTy::createDefaultQueue(ol_device_handle_t Device) {
-  if (DefaultQueue)
-    olDestroyQueue(DefaultQueue);
-  CHECK_FATAL(olCreateQueue(StateTy::getContext(), Device, &DefaultQueue),
-              "Failed to create per-thread default queue");
+StreamTy *ThreadStateTy::getOrCreateDefaultStream(ol_device_handle_t Device) {
+  if (!Device)
+    return nullptr;
+
+  ol_context_handle_t Context = StateTy::getContext();
+  if (!Context)
+    return nullptr;
+
+  StreamTy *&Stream = PerThreadDeviceDefaultStreamMap[Device];
+  if (!Stream) {
+    ol_queue_handle_t Queue = nullptr;
+    CHECK_FATAL(olCreateQueue(Context, Device, &Queue),
+                "Failed to create per-thread default queue for device");
+    Stream = new StreamTy{Queue, Device, QueueKind::PerThreadDefault};
+    StateTy::get().addStream(Stream);
+  }
+  return Stream;
+}
+
+void ThreadStateTy::destroyDefaultStreams() {
+  for (auto &It : PerThreadDeviceDefaultStreamMap) {
+    if (StateTy *State = StateTy::tryGet())
+      State->removeStream(It.second);
+    destroyStreamHandle(It.second);
+  }
+  PerThreadDeviceDefaultStreamMap.clear();
 }
 
 // StateTy implementation.
@@ -161,6 +235,130 @@ StateTy &StateTy::get() {
 
 StateTy *StateTy::tryGet() { return StatePtr.load(); }
 
+StreamTy *StateTy::getOrCreateDefaultStream(ol_device_handle_t Device) {
+  if (!Device)
+    return nullptr;
+
+  ol_context_handle_t RuntimeContext = StateTy::getContext();
+  if (!RuntimeContext)
+    return nullptr;
+
+  std::lock_guard<std::mutex> LG(getDeviceDefaultStreamsMapLock());
+  StreamTy *&Stream = DeviceDefaultStreamsMap[Device];
+  if (!Stream) {
+    ol_queue_handle_t Queue = nullptr;
+    CHECK_FATAL(olCreateQueue(RuntimeContext, Device, &Queue),
+                "Failed to create default queue for device");
+    Stream = new StreamTy{Queue, Device, QueueKind::LegacyDefault};
+    addStream(Stream);
+  }
+  return Stream;
+}
+
+void StateTy::destroyDefaultStreams() {
+  std::lock_guard<std::mutex> LG(getDeviceDefaultStreamsMapLock());
+  for (auto &It : DeviceDefaultStreamsMap) {
+    removeStream(It.second);
+    destroyStreamHandle(It.second);
+  }
+  DeviceDefaultStreamsMap.clear();
+}
+
+void StateTy::addStream(StreamTy *Stream) {
+  std::lock_guard<std::mutex> LG(getDeviceStreamsMapLock());
+  DeviceStreamsMap[Stream->Device].insert(Stream);
+}
+
+void StateTy::removeStream(StreamTy *Stream) {
+  if (!Stream)
+    return;
+
+  {
+    std::lock_guard<std::mutex> LG(getDeviceStreamsMapLock());
+    if (!removeStreamFromMap(DeviceStreamsMap, Stream))
+      return;
+  }
+
+  std::lock_guard<std::mutex> LG(getDeviceBlockingStreamsMapLock());
+  removeStreamFromMap(DeviceBlockingStreamsMap, Stream);
+}
+
+SmallPtrSet<StreamTy *, 8>
+StateTy::getDeviceStreams(ol_device_handle_t Device) {
+  StateTy &State = get();
+  std::lock_guard<std::mutex> LG(getDeviceStreamsMapLock());
+  auto It = State.DeviceStreamsMap.find(Device);
+  if (It == State.DeviceStreamsMap.end())
+    return {};
+  return It->second;
+}
+
+SmallPtrSet<StreamTy *, 8>
+StateTy::getBlockingStreams(ol_device_handle_t Device) {
+  StateTy &State = get();
+  std::lock_guard<std::mutex> LG(getDeviceBlockingStreamsMapLock());
+  auto It = State.DeviceBlockingStreamsMap.find(Device);
+  if (It == State.DeviceBlockingStreamsMap.end())
+    return {};
+  return It->second;
+}
+
+bool StateTy::hasLegacyDefaultStream(ol_device_handle_t Device) {
+  StateTy &State = get();
+  std::lock_guard<std::mutex> LG(getDeviceDefaultStreamsMapLock());
+  auto It = State.DeviceDefaultStreamsMap.find(Device);
+  return It != State.DeviceDefaultStreamsMap.end() && It->second;
+}
+
+ol_result_t StateTy::createStream(ol_device_handle_t Device, QueueKind Kind,
+                                  StreamTy **Stream) {
+  if (!Stream)
+    return &InvalidNullPointerError;
+  *Stream = nullptr;
+
+  ol_context_handle_t RuntimeContext = getContext();
+  if (!Device || !RuntimeContext)
+    return &InvalidDeviceError;
+
+  ol_queue_handle_t Queue = nullptr;
+  ol_result_t Result = olCreateQueue(RuntimeContext, Device, &Queue);
+  if (Result == OL_SUCCESS) {
+    *Stream = new StreamTy{Queue, Device, Kind};
+    get().addStream(*Stream);
+    if (Kind == QueueKind::ExplicitBlocking) {
+      StateTy &State = get();
+      std::lock_guard<std::mutex> LG(getDeviceBlockingStreamsMapLock());
+      State.DeviceBlockingStreamsMap[Device].insert(*Stream);
+    }
+  }
+  return Result;
+}
+
+ol_result_t StateTy::destroyStream(StreamTy *Stream) {
+  if (!Stream)
+    return &InvalidNullPointerError;
+
+  if (!isStreamRegistered(Stream))
+    return &InvalidStreamError;
+
+  get().removeStream(Stream);
+  ol_result_t Result = olDestroyQueue(Stream->Queue);
+  delete Stream;
+  return Result;
+}
+
+bool StateTy::isStreamRegistered(StreamTy *Stream) {
+  if (!Stream)
+    return false;
+
+  StateTy &State = get();
+  std::lock_guard<std::mutex> LG(getDeviceStreamsMapLock());
+  for (auto &It : State.DeviceStreamsMap)
+    if (It.second.contains(Stream))
+      return true;
+  return false;
+}
+
 ol_device_handle_t StateTy::getHostDevice() { return get().HostDevice; }
 
 ol_context_handle_t StateTy::getContext() { return get().Context; }
@@ -269,23 +467,36 @@ StateTy::StateTy() {
     CHECK_FATAL(olCreateContext(Devices.size(), Devices.data(), &Context),
                 "Failed to create default context");
 
-  if (!PerThreadQueue) [[likely]]
-    if (!Devices.empty()) [[likely]]
-      CHECK_FATAL(olCreateQueue(Context, Devices.front(), &DefaultQueue),
-                  "Failed to create default queue");
-
   atexit(deleteState);
 }
 
 StateTy::~StateTy() {
   deleteThreadStates();
-  destroyQueue(DefaultQueue);
+  destroyDefaultStreams();
+  destroyRegisteredStreams();
   destroyRegisteredPrograms();
   if (Context)
     olDestroyContext(Context);
   olShutDown();
 }
 
+void StateTy::destroyRegisteredStreams() {
+  SmallVector<StreamTy *, 16> Streams;
+  {
+    std::lock_guard<std::mutex> LG(getDeviceStreamsMapLock());
+    for (auto &It : DeviceStreamsMap)
+      Streams.append(It.second.begin(), It.second.end());
+    DeviceStreamsMap.clear();
+  }
+  {
+    std::lock_guard<std::mutex> LG(getDeviceBlockingStreamsMapLock());
+    DeviceBlockingStreamsMap.clear();
+  }
+
+  for (StreamTy *&Stream : Streams)
+    destroyStreamHandle(Stream);
+}
+
 void StateTy::destroyRegisteredPrograms() {
   SmallPtrSet<ol_program_handle_t, 8> Programs;
   for (auto &It : BinaryRegisterMap)
diff --git a/offload/test/offloading/CUDA/stream_api.cu b/offload/test/offloading/CUDA/stream_api.cu
index 7202751f8207e..885ef3b300069 100644
--- a/offload/test/offloading/CUDA/stream_api.cu
+++ b/offload/test/offloading/CUDA/stream_api.cu
@@ -14,7 +14,12 @@
 
 #include <stdio.h>
 
-__global__ void setValue(int *Out) { *Out = 42; }
+static void print_error(const char *Label, cudaError_t Error) {
+  printf("%s value: %u\n", Label, static_cast<unsigned>(Error));
+  printf("%s name: %s\n", Label, cudaGetErrorName(Error));
+}
+
+__global__ void setValue(int *Out, int Value) { *Out = Value; }
 
 int main(int argc, char **argv) {
   cudaStream_t Stream = nullptr;
@@ -24,23 +29,43 @@ int main(int argc, char **argv) {
   printf("stream created: %d\n", Stream != nullptr);
   // CHECK: stream created: 1
 
-  int *DevPtr = nullptr;
-  int Result = 0;
-  if (cudaMalloc(&DevPtr, sizeof(int)) != cudaSuccess)
+  int *StreamPtr = nullptr;
+  int *DefaultPtr = nullptr;
+  int StreamResult = 0;
+  int DefaultResult = 0;
+  if (cudaMalloc(&StreamPtr, sizeof(int)) != cudaSuccess)
+    return 1;
+  if (cudaMalloc(&DefaultPtr, sizeof(int)) != cudaSuccess)
     return 1;
 
-  setValue<<<1, 1, 0, Stream>>>(DevPtr);
+  setValue<<<1, 1, 0, Stream>>>(StreamPtr, 42);
+  setValue<<<1, 1>>>(DefaultPtr, 17);
 
   if (cudaStreamSynchronize(Stream) != cudaSuccess)
     return 1;
-  if (cudaMemcpy(&Result, DevPtr, sizeof(int), cudaMemcpyDeviceToHost) !=
-      cudaSuccess)
+  if (cudaDeviceSynchronize() != cudaSuccess)
+    return 1;
+  if (cudaMemcpy(&StreamResult, StreamPtr, sizeof(int),
+                 cudaMemcpyDeviceToHost) != cudaSuccess)
+    return 1;
+  if (cudaMemcpy(&DefaultResult, DefaultPtr, sizeof(int),
+                 cudaMemcpyDeviceToHost) != cudaSuccess)
     return 1;
 
-  printf("stream result: %d\n", Result);
+  printf("stream result: %d\n", StreamResult);
   // CHECK: stream result: 42
+  printf("default result: %d\n", DefaultResult);
+  // CHECK: default result: 17
 
   if (cudaStreamDestroy(Stream) != cudaSuccess)
     return 1;
-  cudaFree(DevPtr);
+  print_error("destroyed stream destroy", cudaStreamDestroy(Stream));
+  // CHECK: destroyed stream destroy value: 4
+  // CHECK: destroyed stream destroy name: cudaErrorInvalidResourceHandle
+  print_error("destroyed stream synchronize", cudaStreamSynchronize(Stream));
+  // CHECK: destroyed stream synchronize value: 4
+  // CHECK: destroyed stream synchronize name: cudaErrorInvalidResourceHandle
+
+  cudaFree(StreamPtr);
+  cudaFree(DefaultPtr);
 }
diff --git a/offload/test/offloading/HIP/stream_api.hip b/offload/test/offloading/HIP/stream_api.hip
index c0e2699822814..6d16c2546245f 100644
--- a/offload/test/offloading/HIP/stream_api.hip
+++ b/offload/test/offloading/HIP/stream_api.hip
@@ -14,7 +14,12 @@
 
 #include <stdio.h>
 
-__global__ void setValue(int *Out) { *Out = 42; }
+static void print_error(const char *Label, hipError_t Error) {
+  printf("%s value: %u\n", Label, static_cast<unsigned>(Error));
+  printf("%s name: %s\n", Label, hipGetErrorName(Error));
+}
+
+__global__ void setValue(int *Out, int Value) { *Out = Value; }
 
 int main(int argc, char **argv) {
   hipStream_t Stream = nullptr;
@@ -24,23 +29,43 @@ int main(int argc, char **argv) {
   printf("stream created: %d\n", Stream != nullptr);
   // CHECK: stream created: 1
 
-  int *DevPtr = nullptr;
-  int Result = 0;
-  if (hipMalloc(&DevPtr, sizeof(int)) != hipSuccess)
+  int *StreamPtr = nullptr;
+  int *DefaultPtr = nullptr;
+  int StreamResult = 0;
+  int DefaultResult = 0;
+  if (hipMalloc(&StreamPtr, sizeof(int)) != hipSuccess)
+    return 1;
+  if (hipMalloc(&DefaultPtr, sizeof(int)) != hipSuccess)
     return 1;
 
-  setValue<<<1, 1, 0, Stream>>>(DevPtr);
+  setValue<<<1, 1, 0, Stream>>>(StreamPtr, 42);
+  setValue<<<1, 1>>>(DefaultPtr, 17);
 
   if (hipStreamSynchronize(Stream) != hipSuccess)
     return 1;
-  if (hipMemcpy(&Result, DevPtr, sizeof(int), hipMemcpyDeviceToHost) !=
+  if (hipDeviceSynchronize() != hipSuccess)
+    return 1;
+  if (hipMemcpy(&StreamResult, StreamPtr, sizeof(int), hipMemcpyDeviceToHost) !=
       hipSuccess)
     return 1;
+  if (hipMemcpy(&DefaultResult, DefaultPtr, sizeof(int),
+                hipMemcpyDeviceToHost) != hipSuccess)
+    return 1;
 
-  printf("stream result: %d\n", Result);
+  printf("stream result: %d\n", StreamResult);
   // CHECK: stream result: 42
+  printf("default result: %d\n", DefaultResult);
+  // CHECK: default result: 17
 
   if (hipStreamDestroy(Stream) != hipSuccess)
     return 1;
-  hipFree(DevPtr);
+  print_error("destroyed stream destroy", hipStreamDestroy(Stream));
+  // CHECK: destroyed stream destroy value: 4
+  // CHECK: destroyed stream destroy name: hipErrorInvalidResourceHandle
+  print_error("destroyed stream synchronize", hipStreamSynchronize(Stream));
+  // CHECK: destroyed stream synchronize value: 4
+  // CHECK: destroyed stream synchronize name: hipErrorInvalidResourceHandle
+
+  hipFree(StreamPtr);
+  hipFree(DefaultPtr);
 }



More information about the llvm-commits mailing list