已合并
api_trace_rule_changes #30533
Ambi创建于 2月5日
api_trace_rule_changes #30533
已合并
从已删除 :v2.7.1合入到Ascend/pytorchv2.7.1
共 2 个文件变更+76-1
| @@ -0,0 +1,39 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch._inductor.utils import run_and_get_code | ||
| 3 | +from torch.testing._internal.common_utils import ( | ||
| 4 | + run_tests, | ||
| 5 | + parametrize, | ||
| 6 | + instantiate_parametrized_tests, | ||
| 7 | +) | ||
| 8 | +from testutils import TestUtils | ||
| 9 | +import torch_npu | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +class TestNpuSetDeviceInGraph(TestUtils): | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + def test_npu_set_device_in_graph(self, shape, dtype): | ||
| 17 | + def op_calc(x, y): | ||
| 18 | + x = x * y | ||
| 19 | + # 保留你原来 graph 内的 NPU 调用 | ||
| 20 | + torch_npu._C._npu_getDefaultStream(0) | ||
| 21 | + out = x + y | ||
| 22 | + return out | ||
| 23 | + | ||
| 24 | + x = self._generate_tensor(shape, dtype) | ||
| 25 | + y = self._generate_tensor(shape, dtype) | ||
| 26 | + | ||
| 27 | + compiled_fn = torch.compile(op_calc, backend="inductor") | ||
| 28 | + out, codes = run_and_get_code(compiled_fn, x, y) | ||
| 29 | + | ||
| 30 | + # 基本 sanity check：确保 Inductor kernel 生成 | ||
| 31 | + self.assertTrue(len(codes) > 0) | ||
| 32 | + self.assertIn("triton", codes[0]) | ||
| 33 | + | ||
| 34 | + | ||
| 35 | +instantiate_parametrized_tests(TestNpuSetDeviceInGraph) | ||
| 36 | + | ||
| 37 | + | ||
| 38 | +if __name__ == "__main__": | ||
| 39 | + run_tests() | ||
| @@ -7,7 +7,7 @@ import torch_npu | |||
| 7 | 7 | ||
| 8 | __all__ = [] | 8 | __all__ = [] |
| 9 | 9 | ||
| 10 | -torch_c_binding_in_graph_functions_npu = dict.fromkeys( | 10 | +torch_non_c_binding_in_graph_functions_npu = dict.fromkeys( |
| 11 | [ | 11 | [ |
| 12 | "torch.npu.current_stream", | 12 | "torch.npu.current_stream", |
| 13 | "torch.npu.default_stream", | 13 | "torch.npu.default_stream", |
| @@ -26,6 +26,41 @@ torch_c_binding_in_graph_functions_npu = dict.fromkeys( | |||
| 26 | "torch.npu.memory._set_allocator_settings", | 26 | "torch.npu.memory._set_allocator_settings", |
| 27 | "torch.npu.memory.empty_cache", | 27 | "torch.npu.memory.empty_cache", |
| 28 | "torch.npu.mem_get_info", | 28 | "torch.npu.mem_get_info", |
| 29 | + "torch.npu.memory.reset_accumulated_host_memory_stats", | ||
| 30 | + "torch.npu.memory.reset_accumulated_memory_stats", | ||
| 31 | + "torch.npu.memory.reset_max_memory_allocated", | ||
| 32 | + "torch.npu.memory.reset_max_memory_cached", | ||
| 33 | + "torch.npu.memory.reset_peak_host_memory_stats", | ||
| 34 | + "torch.npu.memory.reset_peak_memory_stats", | ||
| 35 | + "torch.npu.memory.set_per_process_memory_fraction", | ||
| 36 | + "torch.npu.random.manual_seed_all", | ||
| 37 | + "torch.npu.random.manual_seed", | ||
| 38 | + "torch.npu.random.seed_all", | ||
| 39 | + "torch.npu.random.seed", | ||
| 40 | + "torch.npu.set_sync_debug_mode", | ||
| 41 | + "torch.npu._set_rng_state_offset", | ||
| 42 | + ], | ||
| 43 | + TorchInGraphFunctionVariable, | ||
| 44 | +) | ||
| 45 | + | ||
| 46 | +torch_c_binding_in_graph_functions_npu = dict.fromkeys( | ||
| 47 | + [ | ||
| 48 | + "torch_npu._C._npu_changeCurrentAllocator", | ||
| 49 | + "torch_npu._C._npu_npuCachingAllocator_set_allocator_settings", | ||
| 50 | + "torch_npu._C._npu_emptyCache", | ||
| 51 | + "torch_npu._C._npu_getAllocator", | ||
| 52 | + "torch_npu._C._npu_getCheckpointState", | ||
| 53 | + "torch_npu._C._npu_getCurrentStream", | ||
| 54 | + "torch_npu._C._npu_getDefaultStream", | ||
| 55 | + "torch_npu._C._npu_init", | ||
| 56 | + "torch_npu._C._npu_ipc_collect", | ||
| 57 | + "torch_npu._C._npu_resetAccumulatedHostMemoryStats", | ||
| 58 | + "torch_npu._C._npu_resetPeakHostMemoryStats", | ||
| 59 | + "torch_npu._C._npu_resetPeakMemoryStats", | ||
| 60 | + "torch_npu._C._npu_set_sync_debug_mode", | ||
| 61 | + "torch_npu._C._npu_setDevice", | ||
| 62 | + "torch_npu._C._npu_setMemoryFraction", | ||
| 63 | + "torch_npu._C._npu_synchronize", | ||
| 29 | ], | 64 | ], |
| 30 | TorchInGraphFunctionVariable, | 65 | TorchInGraphFunctionVariable, |
| 31 | ) | 66 | ) |
| @@ -47,6 +82,7 @@ skip_functions_npu = dict.fromkeys( | |||
| 47 | 82 | ||
| 48 | def _patch_npu_trace_rules(): | 83 | def _patch_npu_trace_rules(): |
| 49 | torch._dynamo.trace_rules.clear_lru_cache() | 84 | torch._dynamo.trace_rules.clear_lru_cache() |
| 85 | + torch._dynamo.trace_rules.torch_name_rule_map.append(torch_non_c_binding_in_graph_functions_npu) | ||
| 50 | torch._dynamo.trace_rules.torch_name_rule_map.append(torch_c_binding_in_graph_functions_npu) | 86 | torch._dynamo.trace_rules.torch_name_rule_map.append(torch_c_binding_in_graph_functions_npu) |
| 51 | torch._dynamo.trace_rules.torch_name_rule_map.append(skip_functions_npu) | 87 | torch._dynamo.trace_rules.torch_name_rule_map.append(skip_functions_npu) |
| 52 | torch_module.constant_fold_functions[torch.npu.current_device] = True | 88 | torch_module.constant_fold_functions[torch.npu.current_device] = True |