diff --git a/torch_npu/csrc/framework/LazyInitAclops.cpp b/torch_npu/csrc/framework/LazyInitAclops.cpp index 440d237fef303f45fcd0883e43298df053492e70..a3960bc9c9ac884a7d3f0f97a7bc06f9688ddfa7 100644 --- a/torch_npu/csrc/framework/LazyInitAclops.cpp +++ b/torch_npu/csrc/framework/LazyInitAclops.cpp @@ -174,6 +174,22 @@ void LazyInitAclopsCore() #endif } +bool IsJitCompileDisable() +{ + static const std::string jit_compile_option_name = "jitCompile"; + auto option_value = c10_npu::option::GetOption(jit_compile_option_name); + if (option_value.has_value() && (option_value.value() == "disable")) { + return true; + } else { + static const std::string jit_compile_init_option_name = "jitCompileInit"; + auto init_option_value = c10_npu::option::GetOption(jit_compile_init_option_name); + if (init_option_value.has_value() && (init_option_value.value() == "disable")) { + return true; + } + } + return false; +} + void LazyInitAclops() { static auto acl_op_init_mode = c10_npu::option::OptionsManager::GetAclOpInitMode(); @@ -186,6 +202,9 @@ void LazyInitAclops() if (!encounteredAclops.exchange(true) && c10_npu::NpuSysCtrl::GetInstance().GetInitFlag()) { RECORD_FUNCTION("LazyInitAclops", std::vector({})); + std::string val = IsJitCompileDisable() ? "disable" : "enable"; + NPU_CHECK_ERROR(at_npu::native::AclSetCompileopt(aclCompileOpt::ACL_OP_JIT_COMPILE, val.c_str())); + ASCEND_LOGI("Set jitCompileInit option to %s", val.c_str()); LazyInitAclopsCore(); ASCEND_LOGI("Lazy init for aclops finished.") } diff --git a/torch_npu/csrc/framework/LazyInitAclops.h b/torch_npu/csrc/framework/LazyInitAclops.h index b842b786522309ec08b548aea5488e6f93d52cc3..fccd8858f4a6555b74b6d9664ca87547db3f987c 100644 --- a/torch_npu/csrc/framework/LazyInitAclops.h +++ b/torch_npu/csrc/framework/LazyInitAclops.h @@ -8,6 +8,8 @@ void InitAclops(); void LazyInitAclops(); void InitializeJitCompilationMode(); +bool IsJitCompileDisable(); + } // namespace aclops } // namespace at_npu diff --git a/torch_npu/csrc/npu/Module.cpp b/torch_npu/csrc/npu/Module.cpp index ab0218f58536127b92a2c3959e41267f7e291b00..12fc060feaa752727c0508203145173ab98c731c 100644 --- a/torch_npu/csrc/npu/Module.cpp +++ b/torch_npu/csrc/npu/Module.cpp @@ -34,6 +34,7 @@ #include "torch_npu/csrc/core/npu/register/OptionRegister.h" #include "torch_npu/csrc/core/OverflowUtils.h" #include "torch_npu/csrc/framework/StorageDescHelper.h" +#include "torch_npu/csrc/framework/LazyInitAclops.h" #include "torch_npu/csrc/npu/DataParallelComm.h" #include "torch_npu/csrc/npu/NPUPluggableAllocator.h" #include "torch_npu/csrc/npu/Stream.h" @@ -945,18 +946,10 @@ PyObject *THNPModule_is_jit_compile_false_wrap(PyObject *self, PyObject *noargs) { HANDLE_TH_ERRORS pybind11::gil_scoped_release no_gil; - static const std::string jit_compile_option_name = "jitCompile"; - auto option_value = c10_npu::option::GetOption(jit_compile_option_name); - if (option_value.has_value() && (option_value.value() == "disable")) { + if (at_npu::aclops::IsJitCompileDisable()) { Py_RETURN_TRUE; } else { - static const std::string jit_compile_init_option_name = "jitCompileInit"; - auto init_option_value = c10_npu::option::GetOption(jit_compile_init_option_name); - if (init_option_value.has_value() && (init_option_value.value() == "disable")) { - Py_RETURN_TRUE; - } else { - Py_RETURN_FALSE; - } + Py_RETURN_FALSE; } END_HANDLE_TH_ERRORS }