已合并
skip npu sync func to master #27578
lihui创建于 2025年12月6日
skip npu sync func to master #27578
已合并
共 2 个文件变更+39-0
| @@ -0,0 +1,36 @@ | |||
| 1 | +# Owner(s): ["module: tests"] | ||
| 2 | + | ||
| 3 | +import unittest | ||
| 4 | +import os | ||
| 5 | +import torch | ||
| 6 | +from torch.testing._internal.common_utils import run_tests, TestCase, load_tests | ||
| 7 | +from torch._inductor.utils import run_and_get_code | ||
| 8 | +import torch_npu | ||
| 9 | +import torch_npu.testing | ||
| 10 | + | ||
| 11 | +# load_tests from torch.testing._internal.common_utils is used to automatically filter tests for | ||
| 12 | +# sharding on sandcastle. This line silences flake warnings | ||
| 13 | +load_tests = load_tests | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +class TestSynchronizeSkip(TestCase): | ||
| 17 | + | ||
| 18 | + def test_synchronize_not_in_compiled_graph(self): | ||
| 19 | + | ||
| 20 | + def func_with_synchronize(x): | ||
| 21 | + y = x + 1.0 | ||
| 22 | + torch_npu.npu.utils.synchronize() | ||
| 23 | + return y * 2.0 | ||
| 24 | + | ||
| 25 | + x = torch.randn(32, 16, device="npu", dtype=torch.float32) | ||
| 26 | + expected = (x + 1.0) * 2.0 | ||
| 27 | + compiled_func = torch.compile(func_with_synchronize, backend="inductor", dynamic=False) | ||
| 28 | + result, inductor_code_list = run_and_get_code(compiled_func, x) | ||
| 29 | + torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5) | ||
| 30 | + full_code = "\n".join(inductor_code_list) | ||
| 31 | + self.assertNotIn("synchronize", full_code) | ||
| 32 | + self.assertIn("async_compile.triton", full_code) | ||
| 33 | + | ||
| 34 | + | ||
| 35 | +if __name__ == "__main__": | ||
| 36 | + run_tests() | ||
| @@ -1,5 +1,7 @@ | |||
| 1 | import torch | 1 | import torch |
| 2 | from torch._dynamo.variables import TorchInGraphFunctionVariable | 2 | from torch._dynamo.variables import TorchInGraphFunctionVariable |
| 3 | +from torch._dynamo.trace_rules import manual_torch_name_rule_map, SkipFunctionVariable | ||
| 4 | +import torch_npu | ||
| 3 | 5 | ||
| 4 | __all__ = [] | 6 | __all__ = [] |
| 5 | 7 | ||
| @@ -17,3 +19,4 @@ torch_c_binding_in_graph_functions_npu = dict.fromkeys( | |||
| 17 | def _patch_npu_trace_rules(): | 19 | def _patch_npu_trace_rules(): |
| 18 | torch._dynamo.trace_rules.clear_lru_cache() | 20 | torch._dynamo.trace_rules.clear_lru_cache() |
| 19 | torch._dynamo.trace_rules.torch_name_rule_map.append(torch_c_binding_in_graph_functions_npu) | 21 | torch._dynamo.trace_rules.torch_name_rule_map.append(torch_c_binding_in_graph_functions_npu) |
| 22 | + manual_torch_name_rule_map["torch_npu.npu.utils.synchronize"] = SkipFunctionVariable | ||