From df11e5480bfcc4abbbb488c0aec48429ad7ffbb8 Mon Sep 17 00:00:00 2001
From: chen_liqing <jiangyiwen5@huawei.com>
Date: Mon, 18 Aug 2025 15:38:11 +0800
Subject: [PATCH] [bugfix]set jit compile while LazyInitAclops

---
 torch_npu/csrc/framework/LazyInitAclops.cpp | 19 +++++++++++++++++++
 torch_npu/csrc/framework/LazyInitAclops.h   |  2 ++
 torch_npu/csrc/npu/Module.cpp               | 13 +++----------
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/torch_npu/csrc/framework/LazyInitAclops.cpp b/torch_npu/csrc/framework/LazyInitAclops.cpp
index 440d237fef3..a3960bc9c9a 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<c10::IValue>({}));
+        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 b842b786522..fccd8858f4a 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 ab0218f5853..12fc060feaa 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
 }
-- 
Gitee