已合并
use torch_npu._C instead of option #43708
huangyunlong创建于 8月4日
use torch_npu._C instead of option #43708
已合并
共 7 个文件变更+59-10
| @@ -122,7 +122,7 @@ class TorchFillUninitializedMemoryTestCase(TestCase): | |||
| 122 | self.assertTrue(res.eq(max_val).all()) | 122 | self.assertTrue(res.eq(max_val).all()) |
| 123 | 123 | ||
| 124 | 124 | ||
| 125 | -instantiate_device_type_tests(TorchFillUninitializedMemoryTestCase, globals(), only_for='privateuse1') | 125 | +instantiate_device_type_tests(TorchFillUninitializedMemoryTestCase, globals(), only_for=('privateuse1',)) |
| 126 | 126 | ||
| 127 | 127 | ||
| 128 | if __name__ == '__main__': | 128 | if __name__ == '__main__': |
| @@ -3,6 +3,7 @@ | |||
| 3 | 3 | ||
| 4 | 4 | ||
| 5 | 5 | ||
| 6 | + | ||
| 6 | 7 | ||
| 7 | 8 | ||
| 8 | 9 | ||
| @@ -62,7 +63,7 @@ const at::Tensor& NPUNativeFunctions::resize_( | |||
| 62 | // fill newly added storage section with NaN/MAX_INT when requested. | 63 | // fill newly added storage section with NaN/MAX_INT when requested. |
| 63 | if (C10_UNLIKELY(at::globalContext().deterministicAlgorithms() && | 64 | if (C10_UNLIKELY(at::globalContext().deterministicAlgorithms() && |
| 64 | at::globalContext().deterministicFillUninitializedMemory() && | 65 | at::globalContext().deterministicFillUninitializedMemory() && |
| 65 | - at_npu::native::env::CheckFillUninitializedMemory())) { | 66 | + at_npu::native::env::globalNpuContext().npuFillUninitializedMemory())) { |
| 66 | at::native::fill_resize_deterministic_(self, static_cast<int64_t>(old_storage_nbytes)); | 67 | at::native::fill_resize_deterministic_(self, static_cast<int64_t>(old_storage_nbytes)); |
| 67 | } | 68 | } |
| 68 | return self; | 69 | return self; |
| @@ -120,7 +120,7 @@ inline bool should_fill_empty_deterministic() | |||
| 120 | { | 120 | { |
| 121 | return at::globalContext().deterministicAlgorithms() && | 121 | return at::globalContext().deterministicAlgorithms() && |
| 122 | at::globalContext().deterministicFillUninitializedMemory() && | 122 | at::globalContext().deterministicFillUninitializedMemory() && |
| 123 | - at_npu::native::env::CheckFillUninitializedMemory(); | 123 | + at_npu::native::env::globalNpuContext().npuFillUninitializedMemory(); |
| 124 | } | 124 | } |
| 125 | 125 | ||
| 126 | void window_function_checks( | 126 | void window_function_checks( |
| @@ -320,9 +320,6 @@ bool CheckCompatibleImplFor(CompatibleKey compatible_key) | |||
| 320 | return CheckCompatibleImpl(); | 320 | return CheckCompatibleImpl(); |
| 321 | } | 321 | } |
| 322 | 322 | ||
| 323 | -TORCH_NPU_REGISTER_OPTION(TORCH_NPU_FILL_UNINITIALIZED_MEMORY) | ||
| 324 | -REGISTER_OPTION_BOOL_FUNCTION(CheckFillUninitializedMemory, TORCH_NPU_FILL_UNINITIALIZED_MEMORY, "0", "1") | ||
| 325 | - | ||
| 326 | REGISTER_OPTION_HOOK(ALLOW_CONV_HF32, [](const std::string &val) { | 323 | REGISTER_OPTION_HOOK(ALLOW_CONV_HF32, [](const std::string &val) { |
| 327 | static const std::string mm_hf32_option_name = "ALLOW_MATMUL_HF32"; | 324 | static const std::string mm_hf32_option_name = "ALLOW_MATMUL_HF32"; |
| 328 | auto mm_hf32_val = c10_npu::option::GetOption(mm_hf32_option_name); | 325 | auto mm_hf32_val = c10_npu::option::GetOption(mm_hf32_option_name); |
| @@ -363,6 +360,19 @@ REGISTER_OPTION_HOOK(ACL_OP_DEBUG_OPTION, [](const std::string &val) { | |||
| 363 | 360 | ||
| 364 | TORCH_NPU_REGISTER_OPTION(CUBE_MATH_TYPE) | 361 | TORCH_NPU_REGISTER_OPTION(CUBE_MATH_TYPE) |
| 365 | 362 | ||
| 363 | +NpuContext& globalNpuContext() { | ||
| 364 | + static NpuContext globalNpuContext_; | ||
| 365 | + return globalNpuContext_; | ||
| 366 | +} | ||
| 367 | + | ||
| 368 | +bool NpuContext::npuFillUninitializedMemory() const { | ||
| 369 | + return _npu_fill_uninitialized_memory; | ||
| 370 | +} | ||
| 371 | + | ||
| 372 | +void NpuContext::setNpuFillUninitializedMemory(bool b) { | ||
| 373 | + _npu_fill_uninitialized_memory = b; | ||
| 374 | +} | ||
| 375 | + | ||
| 366 | } // namespace env | 376 | } // namespace env |
| 367 | } // namespace native | 377 | } // namespace native |
| 368 | } // namespace at_npu | 378 | } // namespace at_npu |
| @@ -32,12 +32,22 @@ enum class CompatibleKey { | |||
| 32 | bool CheckCompatibleImpl(); | 32 | bool CheckCompatibleImpl(); |
| 33 | bool CheckCompatibleImplFor(CompatibleKey compatible_key); | 33 | bool CheckCompatibleImplFor(CompatibleKey compatible_key); |
| 34 | bool CheckCompatibleImplBlackListFor(CompatibleKey compatible_key); | 34 | bool CheckCompatibleImplBlackListFor(CompatibleKey compatible_key); |
| 35 | -bool CheckFillUninitializedMemory(); | ||
| 36 | 35 | ||
| 37 | bool IsAllowFP32ToFP16(); | 36 | bool IsAllowFP32ToFP16(); |
| 38 | bool IsAllowConvHF32(); | 37 | bool IsAllowConvHF32(); |
| 39 | bool IsAllowMatmulHF32(); | 38 | bool IsAllowMatmulHF32(); |
| 40 | 39 | ||
| 40 | +class NpuContext { | ||
| 41 | + public: | ||
| 42 | + bool npuFillUninitializedMemory() const; | ||
| 43 | + void setNpuFillUninitializedMemory(bool /*b*/); | ||
| 44 | + | ||
| 45 | + private: | ||
| 46 | + bool _npu_fill_uninitialized_memory = false; | ||
| 47 | +}; | ||
| 48 | + | ||
| 49 | +NpuContext& globalNpuContext(); | ||
| 50 | + | ||
| 41 | } // namespace env | 51 | } // namespace env |
| 42 | } // namespace native | 52 | } // namespace native |
| 43 | } // namespace at_npu | 53 | } // namespace at_npu |
| @@ -2525,6 +2525,27 @@ PyObject* THNPModule_get_deterministic_level(PyObject* self, PyObject* noargs) { | |||
| 2525 | END_HANDLE_TH_ERRORS | 2525 | END_HANDLE_TH_ERRORS |
| 2526 | } | 2526 | } |
| 2527 | 2527 | ||
| 2528 | +static PyObject* THNPModule_setNpuFillUninitializedMemory( | ||
| 2529 | + PyObject* _unused, | ||
| 2530 | + PyObject* arg) { | ||
| 2531 | + HANDLE_TH_ERRORS | ||
| 2532 | + TORCH_CHECK( | ||
| 2533 | + PyBool_Check(arg), "expected a bool, but got ", THPUtils_typename(arg)); | ||
| 2534 | + at_npu::native::env::globalNpuContext().setNpuFillUninitializedMemory( | ||
| 2535 | + arg == Py_True); | ||
| 2536 | + Py_RETURN_NONE; | ||
| 2537 | + END_HANDLE_TH_ERRORS | ||
| 2538 | +} | ||
| 2539 | + | ||
| 2540 | +static PyObject* THNPModule_npuFillUninitializedMemory( | ||
| 2541 | + PyObject* _unused, | ||
| 2542 | + PyObject* noargs) { | ||
| 2543 | + if (at_npu::native::env::globalNpuContext().npuFillUninitializedMemory()) | ||
| 2544 | + Py_RETURN_TRUE; | ||
| 2545 | + else | ||
| 2546 | + Py_RETURN_FALSE; | ||
| 2547 | +} | ||
 | |||
| 2548 | + | ||
| 2528 | PyObject* THNPModule_hasPrimaryContext_wrap(PyObject* self, PyObject* arg) { | 2549 | PyObject* THNPModule_hasPrimaryContext_wrap(PyObject* self, PyObject* arg) { |
| 2529 | HANDLE_TH_ERRORS | 2550 | HANDLE_TH_ERRORS |
| 2530 | TORCH_CHECK( | 2551 | TORCH_CHECK( |
| @@ -2855,6 +2876,14 @@ static struct PyMethodDef THNPModule_methods[] = { | |||
| 2855 | (PyCFunction)THNPModule_get_deterministic_level, | 2876 | (PyCFunction)THNPModule_get_deterministic_level, |
| 2856 | METH_NOARGS, | 2877 | METH_NOARGS, |
| 2857 | nullptr}, | 2878 | nullptr}, |
| 2879 | + {"_npu_set_fill_uninitialized_memory", | ||
| 2880 | + (PyCFunction)THNPModule_setNpuFillUninitializedMemory, | ||
| 2881 | + METH_O, | ||
| 2882 | + nullptr}, | ||
| 2883 | + {"_npu_get_fill_uninitialized_memory", | ||
| 2884 | + (PyCFunction)THNPModule_npuFillUninitializedMemory, | ||
| 2885 | + METH_NOARGS, | ||
| 2886 | + nullptr}, | ||
| 2858 | {"_npu_sleep", | 2887 | {"_npu_sleep", |
| 2859 | (PyCFunction)THNPModule_npuSleep, | 2888 | (PyCFunction)THNPModule_npuSleep, |
| 2860 | METH_O, | 2889 | METH_O, |
| @@ -2,7 +2,7 @@ def _add_deterministic_patch(): | |||
| 2 | """ | 2 | """ |
| 3 | Patch torch.utils.deterministic.fill_uninitialized_memory setter so that | 3 | Patch torch.utils.deterministic.fill_uninitialized_memory setter so that |
| 4 | setting it also synchronizes the NPU-side option | 4 | setting it also synchronizes the NPU-side option |
| 5 | - TORCH_NPU_FILL_UNINITIALIZED_MEMORY via torch_npu._C._npu_setOption. | 5 | + _npu_fill_uninitialized_memory via torch_npu._C._npu_set_fill_uninitialized_memory. |
| 6 | """ | 6 | """ |
| 7 | import torch_npu | 7 | import torch_npu |
| 8 | import torch.utils.deterministic as det_mod | 8 | import torch.utils.deterministic as det_mod |
| @@ -19,7 +19,6 @@ def _add_deterministic_patch(): | |||
| 19 | # 1. call original setter: torch._C._set_deterministic_fill_uninitialized_memory(mode) | 19 | # 1. call original setter: torch._C._set_deterministic_fill_uninitialized_memory(mode) |
| 20 | _orig_fset(self, mode) | 20 | _orig_fset(self, mode) |
| 21 | # 2. sync NPU-side option | 21 | # 2. sync NPU-side option |
| 22 | - option = {"TORCH_NPU_FILL_UNINITIALIZED_MEMORY": "1" if mode else "0"} | 22 | + torch_npu._C._npu_set_fill_uninitialized_memory(mode) |
| 23 | - torch_npu._C._npu_setOption(option) | ||
| 24 | 23 | ||
| 25 | _Deterministic.fill_uninitialized_memory = property(_orig_fget, _new_fset) | 24 | _Deterministic.fill_uninitialized_memory = property(_orig_fget, _new_fset) |
🟡 Medium Priority
变更行：Module.cpp 第 2540–2547 行新增的
THNPModule_npuFillUninitializedMemory函数（_npu_get_fill_uninitialized_memory的 C 实现）没有HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS宏包裹。受影响的调用链：Python 调用
torch_npu._C._npu_get_fill_uninitialized_memory()→THNPModule_npuFillUninitializedMemory→globalNpuContext().npuFillUninitializedMemory()。失败模式：该文件内所有其他 Python 可调用函数（包括其兄弟 setter
THNPModule_setNpuFillUninitializedMemory、THNPModule_get_deterministic_level等）均使用HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS将 C++ 异常转换为 Python 异常。缺少该保护时，若内部代码（如globalNpuContext()的初始化、或未来对NpuContext的扩展）抛出 C++ 异常，异常将直接穿透 Python 解释器边界导致进程崩溃（UB）。虽然当前
NpuContext::npuFillUninitializedMemory()仅返回一个 bool、几乎不可能抛出异常，但这是一个一致性/维护性缺陷：后续任何对NpuContext的修改都可能无意中引入异常风险，而此处缺乏保护不会被编译器或审查人注意到。建议：为
THNPModule_npuFillUninitializedMemory添加HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS宏包裹，与文件中所有其他 Python C API 函数保持一致。