已合并
use torch_npu._C instead of option #43704
huangyunlong创建于 23 天前
use torch_npu._C instead of option #43704
已合并
共 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 |
| @@ -2533,6 +2533,27 @@ PyObject* THNPModule_get_deterministic_level(PyObject* self, PyObject* noargs) { | |||
| 2533 | END_HANDLE_TH_ERRORS | 2533 | END_HANDLE_TH_ERRORS |
| 2534 | } | 2534 | } |
| 2535 | 2535 | ||
| 2536 | +static PyObject* THNPModule_setNpuFillUninitializedMemory( | ||
| 2537 | + PyObject* _unused, | ||
| 2538 | + PyObject* arg) { | ||
| 2539 | + HANDLE_TH_ERRORS | ||
| 2540 | + TORCH_CHECK( | ||
| 2541 | + PyBool_Check(arg), "expected a bool, but got ", THPUtils_typename(arg)); | ||
| 2542 | + at_npu::native::env::globalNpuContext().setNpuFillUninitializedMemory( | ||
| 2543 | + arg == Py_True); | ||
| 2544 | + Py_RETURN_NONE; | ||
| 2545 | + END_HANDLE_TH_ERRORS | ||
| 2546 | +} | ||
| 2547 | + | ||
| 2548 | +static PyObject* THNPModule_npuFillUninitializedMemory( | ||
| 2549 | + PyObject* _unused, | ||
| 2550 | + PyObject* noargs) { | ||
| 2551 | + if (at_npu::native::env::globalNpuContext().npuFillUninitializedMemory()) | ||
| 2552 | + Py_RETURN_TRUE; | ||
| 2553 | + else | ||
| 2554 | + Py_RETURN_FALSE; | ||
| 2555 | +} | ||
| 2556 | + | ||
| 2536 | PyObject* THNPModule_hasPrimaryContext_wrap(PyObject* self, PyObject* arg) { | 2557 | PyObject* THNPModule_hasPrimaryContext_wrap(PyObject* self, PyObject* arg) { |
| 2537 | HANDLE_TH_ERRORS | 2558 | HANDLE_TH_ERRORS |
| 2538 | TORCH_CHECK( | 2559 | TORCH_CHECK( |
| @@ -2863,6 +2884,14 @@ static struct PyMethodDef THNPModule_methods[] = { | |||
| 2863 | (PyCFunction)THNPModule_get_deterministic_level, | 2884 | (PyCFunction)THNPModule_get_deterministic_level, |
| 2864 | METH_NOARGS, | 2885 | METH_NOARGS, |
| 2865 | nullptr}, | 2886 | nullptr}, |
| 2887 | + {"_npu_set_fill_uninitialized_memory", | ||
| 2888 | + (PyCFunction)THNPModule_setNpuFillUninitializedMemory, | ||
| 2889 | + METH_O, | ||
| 2890 | + nullptr}, | ||
| 2891 | + {"_npu_get_fill_uninitialized_memory", | ||
| 2892 | + (PyCFunction)THNPModule_npuFillUninitializedMemory, | ||
| 2893 | + METH_NOARGS, | ||
| 2894 | + nullptr}, | ||
| 2866 | {"_npu_sleep", | 2895 | {"_npu_sleep", |
| 2867 | (PyCFunction)THNPModule_npuSleep, | 2896 | (PyCFunction)THNPModule_npuSleep, |
| 2868 | METH_O, | 2897 | 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 第 2548-2555 行新增的
THNPModule_npuFillUninitializedMemory函数。该函数是 CPython 扩展方法(通过
PyMethodDef注册为_npu_get_fill_uninitialized_memory,METH_NOARGS),但其函数体中缺少HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS异常处理宏。问题链条:
THNPModule_get_deterministic_level第 2529-2534 行、THNPModule_npu_get_sync_debug_mode第 1975 行、以及配对的 setterTHNPModule_setNpuFillUninitializedMemory第 2539-2545 行)均使用HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS包裹函数体。建议:在函数体首尾添加
HANDLE_TH_ERRORS和END_HANDLE_TH_ERRORS宏,与代码库中所有其他 CPython 导出函数保持一致。