From 597831dd66c0d71d9bb580c67c86df68fabe44b7 Mon Sep 17 00:00:00 2001
From: Sarbojit Sarkar <Sarbojit.Sarkar@amd.com>
Date: Thu, 22 Jul 2021 13:08:52 +0000
Subject: [PATCH] SWDEV-292547 - hipStreamPerThread support

Change-Id: Id621ce073b0fee9eac03c59ffb78b197fda4ddb5


[ROCm/clr commit: c3ca1faee78871d8c69606ae8a68b21f87393709]
---
 .../nvidia_detail/nvidia_hip_runtime_api.h    |  3 ++
 projects/clr/hipamd/src/hip_internal.hpp      |  4 +-
 projects/clr/hipamd/src/hip_stream.cpp        | 39 ++++++++++++++++++-
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/projects/clr/hipamd/include/hip/nvidia_detail/nvidia_hip_runtime_api.h b/projects/clr/hipamd/include/hip/nvidia_detail/nvidia_hip_runtime_api.h
index a53d68e72e..fbe70bab8a 100644
--- a/projects/clr/hipamd/include/hip/nvidia_detail/nvidia_hip_runtime_api.h
+++ b/projects/clr/hipamd/include/hip/nvidia_detail/nvidia_hip_runtime_api.h
@@ -198,6 +198,9 @@ inline static CUresourcetype hipResourcetype_enumToCUresourcetype(
     }
 }
 
+// hipStreamPerThread
+#define hipStreamPerThread ((cudaStream_t)2)
+
 #define hipTexRef CUtexref
 #define hiparray CUarray
 
diff --git a/projects/clr/hipamd/src/hip_internal.hpp b/projects/clr/hipamd/src/hip_internal.hpp
index 082975635b..3ceb61d3a3 100755
--- a/projects/clr/hipamd/src/hip_internal.hpp
+++ b/projects/clr/hipamd/src/hip_internal.hpp
@@ -127,6 +127,7 @@ typedef struct ihipIpcEventHandle_st {
   } while (0);
 
 #define STREAM_CAPTURE(name, stream, ...)                                                          \
+  getStreamPerThread(stream);                                                                      \
   if (stream != nullptr &&                                                                         \
       reinterpret_cast<hip::Stream*>(stream)->GetCaptureStatus() ==                                \
           hipStreamCaptureStatusActive) {                                                          \
@@ -325,7 +326,7 @@ namespace hip {
   /// Get default stream of the thread
   extern amd::HostQueue* getNullStream();
   /// Check if stream is valid
-  extern bool isValid(hipStream_t stream);
+  extern bool isValid(hipStream_t& stream);
 };
 
 struct ihipExec_t {
@@ -347,6 +348,7 @@ extern int ihipGetDevice();
 extern hipError_t ihipMalloc(void** ptr, size_t sizeBytes, unsigned int flags);
 extern amd::Memory* getMemoryObject(const void* ptr, size_t& offset);
 extern amd::Memory* getMemoryObjectWithOffset(const void* ptr, const size_t size);
+extern hipError_t getStreamPerThread(hipStream_t& stream);
 
 constexpr bool kOptionChangeable = true;
 constexpr bool kNewDevProg = false;
diff --git a/projects/clr/hipamd/src/hip_stream.cpp b/projects/clr/hipamd/src/hip_stream.cpp
index 3fbcc73d5d..cbf3e27105 100755
--- a/projects/clr/hipamd/src/hip_stream.cpp
+++ b/projects/clr/hipamd/src/hip_stream.cpp
@@ -172,12 +172,16 @@ void Stream::destroyAllStreams(int deviceId) {
 }
 
 // ================================================================================================
-bool isValid(hipStream_t stream) {
+bool isValid(hipStream_t& stream) {
   // NULL stream is always valid
   if (stream == nullptr) {
     return true;
   }
 
+  if (hipStreamPerThread == stream) {
+    getStreamPerThread(stream);
+  }
+
   hip::Stream* s = reinterpret_cast<hip::Stream*>(stream);
   amd::ScopedLock lock(streamSetLock);
   if (streamSet.find(s) == streamSet.end()) {
@@ -266,6 +270,36 @@ static hipError_t ihipStreamCreate(hipStream_t* stream,
   return hipSuccess;
 }
 
+// ================================================================================================
+class stream_per_thread {
+private:
+  hipStream_t m_stream;
+public:
+  stream_per_thread():m_stream(nullptr) {}
+  stream_per_thread(const stream_per_thread& ) = delete;
+  void operator=(const stream_per_thread& ) = delete;
+  ~stream_per_thread() {
+    if (m_stream != nullptr && hip::isValid(m_stream)) {
+      delete reinterpret_cast<hip::Stream*>(m_stream);
+    }
+  }
+
+  hipStream_t& get() {
+    if (m_stream == nullptr) {
+      ihipStreamCreate(&m_stream, hipStreamDefault, hip::Stream::Priority::Normal);
+    }
+    return m_stream;
+  }
+};
+thread_local stream_per_thread streamPerThreadObj;
+
+// ================================================================================================
+hipError_t getStreamPerThread(hipStream_t& stream) {
+  if (stream == hipStreamPerThread) {
+    stream = streamPerThreadObj.get();
+  }
+}
+
 // ================================================================================================
 hipError_t hipStreamCreateWithFlags(hipStream_t *stream, unsigned int flags) {
   HIP_INIT_API(hipStreamCreateWithFlags, stream, flags);
@@ -520,6 +554,7 @@ hipError_t hipExtStreamCreateWithCUMask(hipStream_t* stream, uint32_t cuMaskSize
 // ================================================================================================
 hipError_t hipStreamGetPriority(hipStream_t stream, int* priority) {
   HIP_INIT_API(hipStreamGetPriority, stream, priority);
+
   if ((priority != nullptr) && (stream != nullptr)) {
     if (!hip::isValid(stream)) {
       return HIP_RETURN(hipErrorContextIsDestroyed);
@@ -576,7 +611,7 @@ hipError_t hipExtStreamGetCUMask(hipStream_t stream, uint32_t cuMaskSize, uint32
 
   // if the stream is null then either return globalCUMask_ (if it is defined)
   // or return defaultCUMask
-  if (stream == nullptr) {
+  if (stream == nullptr || stream == hipStreamPerThread) {
     if (info.globalCUMask_.size() != 0) {
       std::copy(info.globalCUMask_.begin(), info.globalCUMask_.end(), cuMask);
     } else {