diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index e42e354a8c003dae43fca375a97ec5a11e8a7db4..a9da9e90ca25ec915c653ae94cc9885fb1ac3a8d 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -3023,6 +3023,8 @@ private: void insert_events(Block *block) { + int pre_device = -1; + NPU_CHECK_ERROR(c10_npu::GetDevice(&pre_device)); aclrtContext compiler_ctx = aclrtContext(); aclError ret_ctx = aclrtGetCurrentContext(&compiler_ctx); NPU_CHECK_ERROR(aclrtSetCurrentContext(c10_npu::GetDeviceContext(block->device))); @@ -3040,7 +3042,9 @@ private: npu_events[stream].emplace_back(std::move(event), block); } if (ret_ctx == ACL_ERROR_NONE) { - NPU_CHECK_ERROR(aclrtSetCurrentContext(compiler_ctx)); + NPU_CHECK_ERROR(aclrtSetCurrentContext(compiler_ctx)); + // Setting context will exchange device implicitly, so we need to reset the cached device here to ensure consistency. + NPU_CHECK_ERROR(c10_npu::SetDevice(pre_device)); } } diff --git a/torch_npu/csrc/core/npu/NPUQueue.cpp b/torch_npu/csrc/core/npu/NPUQueue.cpp index abc993f37b43424ab36396a0410b83f64a85fd4f..20bee1e650f7aa24acc694be417166552922748c 100644 --- a/torch_npu/csrc/core/npu/NPUQueue.cpp +++ b/torch_npu/csrc/core/npu/NPUQueue.cpp @@ -5,6 +5,7 @@ #include "torch_npu/csrc/framework/utils/NpuUtils.h" #include "torch_npu/csrc/core/npu/NPUFunctions.h" #include "torch_npu/csrc/framework/OpParamMaker.h" +#include "torch_npu/csrc/framework/OpCommand.h" #include "torch_npu/csrc/core/npu/register/OptionsManager.h" #include "torch_npu/csrc/core/npu/NPUEventManager.h" @@ -249,7 +250,7 @@ NPUStatus Repository::MakeSureQueueEmpty(bool check_error) // occur. #ifndef BUILD_LIBTORCH PyThreadState *gilState = nullptr; - if (PyGILState_Check() != 0) { + if (PyGILState_Check() != 0 && g_used_aclop) { gilState = PyEval_SaveThread(); } #endif @@ -531,7 +532,14 @@ void Repository::Enqueue(void *cur_paras) if (IsFullQueue()) { #ifndef BUILD_LIBTORCH // double check the current thread hold a Gil lock - if (PyGILState_Check() != 0) { + // and release the GIL to TE op compiler in case the acl thread deadlock. + // However, this operator could produce another form of deadlock. + // When thread A deconstract a tensor, it will hold the mutex of deviceCachingAllocator and insert an event into the taskqueue. + // If the taskqueue is full, thead A will run into here and release the GIL. + // Once another thread B get GIL and trigger GC, it may deconstract another tensor + // and try to get deviceCachingAllocator's mutex, which would cause another form of deadlock. + // Since the aclop will be deprecated soon, we just add a using-aclop check here to aviod the second case of deadlock. + if (PyGILState_Check() != 0 && g_used_aclop) { Py_BEGIN_ALLOW_THREADS s = eventfd_read(efd_write, &u); Py_END_ALLOW_THREADS } else { diff --git a/torch_npu/csrc/framework/OpCommand.cpp b/torch_npu/csrc/framework/OpCommand.cpp index 787fdb7bfd1f7e1345a97592b6a353f51a907cf4..383dcb8ecf3e2c5aa9f0664082242e4889a735c3 100644 --- a/torch_npu/csrc/framework/OpCommand.cpp +++ b/torch_npu/csrc/framework/OpCommand.cpp @@ -35,6 +35,8 @@ static std::unordered_map> integral_limits_map {at::ScalarType::Short, {std::numeric_limits::max(), std::numeric_limits::min()}}}; } // namespace +std::atomic g_used_aclop{false}; + namespace at_npu { namespace native { @@ -126,6 +128,7 @@ void OpCommand::Run() { // Check for npu graph if (aclCmd->CheckCustomHandlerNull()) { + g_used_aclop = true; c10_npu::assertNotCapturingAclop(aclCmd->GetName()); } diff --git a/torch_npu/csrc/framework/OpCommand.h b/torch_npu/csrc/framework/OpCommand.h index e60617077c976b0109a37b72c33254a25333a095..f30d9fb4988bfc8e32902b8b2fab783f820b4d54 100644 --- a/torch_npu/csrc/framework/OpCommand.h +++ b/torch_npu/csrc/framework/OpCommand.h @@ -10,6 +10,8 @@ #include "torch_npu/csrc/framework/utils/NpuUtils.h" #include "torch_npu/csrc/framework/utils/NPUDefinition.h" +extern std::atomic g_used_aclop; + namespace at_npu { namespace native {