add stream lock to avoid use same steam in multithread at same time (DeepLink-org#793)

* add stream lock to avoid use same steam in multithread at same time

Author: zhaoguochun1995

impl/ascend/common/acloprunner.hpp

@@ -24,6 +24,7 @@
 #include "debug.hpp"
 #include "gil_scoped_release.hpp"
 #include "impl_functions.hpp"
+#include "stream_lock.hpp"
 #include "utils.hpp"
 
 namespace impl {
@@ -630,6 +631,7 @@ class AclOpRunner {
         diopiGetStream(context_, &stream);
         diopi::GilScopedRelease gilReleaeGuard;
         if (sync_) {
+            diopi::StreamLockGuard streamLockGuard(stream);
             CALL_ACLRT(aclopCompileAndExecuteV2(opname_.data(),
                                                 inputIndex_,
                                                 inputDescs_.data(),
@@ -643,6 +645,7 @@ class AclOpRunner {
                                                 nullptr,
                                                 stream));
         } else {
+            diopi::StreamLockGuard streamLockGuard(stream);
             CALL_ACLRT(aclopCompileAndExecute(opname_.data(),
                                               inputIndex_,
                                               inputDescs_.data(),

impl/ascend/common/stream_lock.cpp

@@ -0,0 +1,31 @@
+#include "stream_lock.hpp"
+
+#include <map>
+#include <mutex>
+#include <thread>
+
+namespace diopi {
+
+using MutexType = std::mutex;
+
+namespace {
+
+MutexType* getLockForStream(void* aclrtStreamHandle) {
+    static std::map<void*, MutexType> streamThreadMutexMap;
+    return &streamThreadMutexMap[aclrtStreamHandle];
+}
+
+}  // namespace
+
+StreamLockGuard::StreamLockGuard(void* aclrtStreamHandle) {
+    MutexType* mutexPtr = getLockForStream(aclrtStreamHandle);
+    mutex_ = mutexPtr;
+    mutexPtr->lock();
+}
+
+StreamLockGuard::~StreamLockGuard() {
+    MutexType* mutexPtr = static_cast<MutexType*>(mutex_);
+    mutexPtr->unlock();
+}
+
+}  // namespace diopi

impl/ascend/common/stream_lock.hpp

@@ -0,0 +1,14 @@
+#pragma once
+
+namespace diopi {
+
+class StreamLockGuard {
+private:
+    void* mutex_ = nullptr;
+
+public:
+    explicit StreamLockGuard(void* aclrtStreamHandle);
+    ~StreamLockGuard();
+};
+
+}  // namespace diopi

impl/ascend_npu/CMakeLists.txt

@@ -655,6 +655,7 @@ set(OLD_IMPL_SRC
     ${OLD_IMPL_DIR}/common/format_helper.cpp
     ${OLD_IMPL_DIR}/common/debug.cpp
     ${OLD_IMPL_DIR}/common/generator_helper.cpp
+    ${OLD_IMPL_DIR}/common/stream_lock.cpp
     ${OLD_IMPL_DIR}/env_vars.cpp
 )
 

impl/ascend_npu/torch_npu/csrc/DIOPIAdapter.cpp

@@ -7,6 +7,7 @@
 #include <torch/library.h>
 
 #include "../../../ascend/common/gil_scoped_release.hpp"
+#include "../../../ascend/common/stream_lock.hpp"
 #include "diopi_impl/helper.hpp"
 #include "op_plugin/AclOpsInterface.h"
 
@@ -2132,18 +2133,21 @@ aclError OpCommandImpl::InnerRun(const string& name, AclExecParam& params, bool
           }
         }
 #endif
-        ret = AclopCompileAndExecuteV2(name.c_str(),
-                                       inputSize,
-                                       const_cast<aclTensorDesc**>(params.inDesc.data()),
-                                       const_cast<aclDataBuffer**>(params.inBuffer.data()),
-                                       outputSize,
-                                       const_cast<aclTensorDesc**>(params.outDesc.data()),
-                                       params.outBuffer.data(),
-                                       params.attr,
-                                       ACL_ENGINE_SYS,
-                                       ACL_COMPILE_SYS,
-                                       NULL,
-                                       stream);
+        {
+            diopi::StreamLockGuard streamLockGuard(stream.stream());
+            ret = AclopCompileAndExecuteV2(name.c_str(),
+                                           inputSize,
+                                           const_cast<aclTensorDesc**>(params.inDesc.data()),
+                                           const_cast<aclDataBuffer**>(params.inBuffer.data()),
+                                           outputSize,
+                                           const_cast<aclTensorDesc**>(params.outDesc.data()),
+                                           params.outBuffer.data(),
+                                           params.attr,
+                                           ACL_ENGINE_SYS,
+                                           ACL_COMPILE_SYS,
+                                           NULL,
+                                           stream);
+        }
         NPU_CHECK_ERROR(ret);
         if (sync) {
             int64_t dimSize;