已合并
use torch_npu._C instead of option #43705
huangyunlong创建于 25 天前
use torch_npu._C instead of option #43705
已合并
共 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 |
| @@ -2534,6 +2534,27 @@ PyObject* THNPModule_get_deterministic_level(PyObject* self, PyObject* noargs) { | |||
| 2534 | END_HANDLE_TH_ERRORS | 2534 | END_HANDLE_TH_ERRORS |
| 2535 | } | 2535 | } |
| 2536 | 2536 | ||
| 2537 | +static PyObject* THNPModule_setNpuFillUninitializedMemory( | ||
| 2538 | + PyObject* _unused, | ||
| 2539 | + PyObject* arg) { | ||
| 2540 | + HANDLE_TH_ERRORS | ||
| 2541 | + TORCH_CHECK( | ||
| 2542 | + PyBool_Check(arg), "expected a bool, but got ", THPUtils_typename(arg)); | ||
| 2543 | + at_npu::native::env::globalNpuContext().setNpuFillUninitializedMemory( | ||
| 2544 | + arg == Py_True); | ||
| 2545 | + Py_RETURN_NONE; | ||
| 2546 | + END_HANDLE_TH_ERRORS | ||
| 2547 | +} | ||
| 2548 | + | ||
| 2549 | +static PyObject* THNPModule_npuFillUninitializedMemory( | ||
| 2550 | + PyObject* _unused, | ||
| 2551 | + PyObject* noargs) { | ||
| 2552 | + if (at_npu::native::env::globalNpuContext().npuFillUninitializedMemory()) | ||
| 2553 | + Py_RETURN_TRUE; | ||
| 2554 | + else | ||
| 2555 | + Py_RETURN_FALSE; | ||
| 2556 | +} | ||
 | |||
| 2557 | + | ||
| 2537 | PyObject* THNPModule_hasPrimaryContext_wrap(PyObject* self, PyObject* arg) { | 2558 | PyObject* THNPModule_hasPrimaryContext_wrap(PyObject* self, PyObject* arg) { |
| 2538 | HANDLE_TH_ERRORS | 2559 | HANDLE_TH_ERRORS |
| 2539 | TORCH_CHECK( | 2560 | TORCH_CHECK( |
| @@ -2864,6 +2885,14 @@ static struct PyMethodDef THNPModule_methods[] = { | |||
| 2864 | (PyCFunction)THNPModule_get_deterministic_level, | 2885 | (PyCFunction)THNPModule_get_deterministic_level, |
| 2865 | METH_NOARGS, | 2886 | METH_NOARGS, |
| 2866 | nullptr}, | 2887 | nullptr}, |
| 2888 | + {"_npu_set_fill_uninitialized_memory", | ||
| 2889 | + (PyCFunction)THNPModule_setNpuFillUninitializedMemory, | ||
| 2890 | + METH_O, | ||
| 2891 | + nullptr}, | ||
| 2892 | + {"_npu_get_fill_uninitialized_memory", | ||
| 2893 | + (PyCFunction)THNPModule_npuFillUninitializedMemory, | ||
| 2894 | + METH_NOARGS, | ||
| 2895 | + nullptr}, | ||
| 2867 | {"_npu_sleep", | 2896 | {"_npu_sleep", |
| 2868 | (PyCFunction)THNPModule_npuSleep, | 2897 | (PyCFunction)THNPModule_npuSleep, |
| 2869 | METH_O, | 2898 | 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
THNPModule_npuFillUninitializedMemory（第 2549-2556 行）是一个 Python C API 函数，注册为METH_NOARGS，但其函数体缺少HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS宏包裹。证据链：
HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS。METH_NOARGS函数都正确使用了这对宏。特别地，紧邻的THNPModule_get_deterministic_level（第 2530-2535 行）和同批新增的THNPModule_setNpuFillUninitializedMemory（第 2537-2547 行）都正确使用了HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS。这是整个Module.cpp中唯一缺少这对宏的METH_NOARGS函数。建议：在函数体首尾添加 HANDLE_TH_ERRORS / END_HANDLE_TH_ERRORS 宏，与项目中所有其他 METH_NOARGS 函数保持一致。