已合并
[feat]: A5 only supports INF_NAN mode #35335
wanglijun55创建于 5月11日
[feat]: A5 only supports INF_NAN mode #35335
已合并
wanglijun55创建于 5月11日
2 个文件变更+52-0
Atest/npu/test_inf_nan_mode.py+47-0
@@ -0,0 +1,47 @@
1+import os
2+import subprocess
3+import sys
4+ 
5+from torch_npu.testing.common_utils import SupportedDevices
6+from torch_npu.testing.testcase import TestCase, run_tests
7+ 
8+ 
9+class TestInfNanMode(TestCase):
10+ 
11+ @SupportedDevices(['Ascend950'])
12+ def test_is_support_inf_nan_always_enabled(self):
13+ import torch_npu.npu.utils as utils
14+ self.assertTrue(utils.is_support_inf_nan())
15+ 
16+ @SupportedDevices(['Ascend950'])
17+ def test_env_vars_ignored(self):
18+ code = (
19+ "import torch_npu.npu.utils as utils;"
20+ "assert utils.is_support_inf_nan()"
21+ )
22+ env_cases = [
23+ ({'INF_NAN_MODE_ENABLE': '1'}, "INF_NAN_MODE_ENABLE=1"),
24+ ({'INF_NAN_MODE_ENABLE': '0'}, "INF_NAN_MODE_ENABLE=0"),
25+ ({'INF_NAN_MODE_FORCE_DISABLE': '0'}, "INF_NAN_MODE_FORCE_DISABLE=0"),
26+ ({'INF_NAN_MODE_FORCE_DISABLE': '1'}, "INF_NAN_MODE_FORCE_DISABLE=1"),
27+ ({'INF_NAN_MODE_ENABLE': '0', 'INF_NAN_MODE_FORCE_DISABLE': '1'},
28+ "INF_NAN_MODE_ENABLE=0 + INF_NAN_MODE_FORCE_DISABLE=1"),
29+ ]
30+ for env_vars, desc in env_cases:
31+ env = os.environ.copy()
32+ env.update(env_vars)
33+ result = subprocess.run(
34+ [sys.executable, '-c', code],
35+ env=env,
36+ capture_output=True,
37+ text=True,
38+ )
39+ self.assertEqual(
40+ result.returncode, 0,
41+ f"{desc} should be ignored on Ascend950.\n"
42+ f"stdout: {result.stdout}\nstderr: {result.stderr}"
43+ )
44+ 
45+ 
46+if __name__ == "__main__":
47+ run_tests()
Mtorch_npu/csrc/core/npu/NpuVariables.cpp+5-0
@@ -76,6 +76,11 @@ const SocVersion& GetSocVersion()
76 76 
77bool IsSupportInfNan()77bool IsSupportInfNan()
78{78{
79+ // Ascend950 only supports INF_NAN mode, saturation mode not available,
80+ // env vars INF_NAN_MODE_ENABLE and INF_NAN_MODE_FORCE_DISABLE are ignored.
81+ if (GetSocVersion() >= SocVersion::Ascend950) {
82+ return true;
83+ }
79 static bool default_support_inf_nan = ((GetSocVersion() >= SocVersion::Ascend910B1) &&84 static bool default_support_inf_nan = ((GetSocVersion() >= SocVersion::Ascend910B1) &&
80 (GetSocVersion() < SocVersion::Ascend310B1)) ||85 (GetSocVersion() < SocVersion::Ascend310B1)) ||
81 (GetSocVersion() >= SocVersion::Ascend910_9391);86 (GetSocVersion() >= SocVersion::Ascend910_9391);