已合并
API_trace_relu changes #30394
cuiduo创建于 2月2日
API_trace_relu changes #30394
已合并
共 4 个文件变更+72-2
| @@ -0,0 +1,37 @@ | |||
| 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 TestCurrentDevice(TestUtils): | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + def test_current_device(self, shape, dtype): | ||
| 17 | + def op_calc(x, y): | ||
| 18 | + x = x * y | ||
| 19 | + t = torch.empty( | ||
| 20 | + size=x.shape, | ||
| 21 | + dtype=x.dtype, | ||
| 22 | + device=torch.npu.current_device(), | ||
| 23 | + ) | ||
| 24 | + out = x + y | ||
| 25 | + return out | ||
| 26 | + | ||
| 27 | + x = self._generate_tensor(shape, dtype) | ||
| 28 | + y = self._generate_tensor(shape, dtype) | ||
| 29 | + compile_result, codes = run_and_get_code(torch.compile(op_calc, backend='inductor'), x, y) | ||
| 30 | + self.assertTrue('triton_unk_fused_add_mul_0.run' in codes[0]) | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +instantiate_parametrized_tests(TestCurrentDevice) | ||
| 34 | + | ||
| 35 | + | ||
| 36 | +if __name__ == "__main__": | ||
| 37 | + run_tests() | ||
| @@ -25,7 +25,6 @@ from torch._inductor.lowering import ( | |||
| 25 | square, | 25 | square, |
| 26 | sub, | 26 | sub, |
| 27 | fallback_handler, | 27 | fallback_handler, |
| 28 | - is_boolean_type, | ||
| 29 | logical_and, | 28 | logical_and, |
| 30 | make_pointwise, | 29 | make_pointwise, |
| 31 | _make_reduction_inner, | 30 | _make_reduction_inner, |
| @@ -1,5 +1,8 @@ | |||
| 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._dynamo.variables.torch as torch_module | ||
| 5 | +from torch._dynamo.utils import common_constant_types | ||
| 3 | import torch_npu | 6 | import torch_npu |
| 4 | 7 | ||
| 5 | __all__ = [] | 8 | __all__ = [] |
| @@ -11,12 +14,42 @@ torch_c_binding_in_graph_functions_npu = dict.fromkeys( | |||
| 11 | "torch.npu.stream", | 14 | "torch.npu.stream", |
| 12 | "torch.npu.set_stream", | 15 | "torch.npu.set_stream", |
| 13 | "torch_npu.npu.utils.synchronize", | 16 | "torch_npu.npu.utils.synchronize", |
| 17 | + "torch.npu.current_device", | ||
| 18 | + "torch.npu.get_device_capability", | ||
| 19 | + "torch.npu.get_device_properties", | ||
| 20 | + "torch.npu.graphs.graph_pool_handle", | ||
| 21 | + "torch.npu.ipc_collect", | ||
| 22 | + "torch.npu.is_available", | ||
| 23 | + "torch.npu.memory._dump_snapshot", | ||
| 24 | + "torch.npu.memory._free_mutex", | ||
| 25 | + "torch.npu.memory._record_memory_history_impl", | ||
| 26 | + "torch.npu.memory._set_allocator_settings", | ||
| 27 | + "torch.npu.memory.empty_cache", | ||
| 28 | + "torch.npu.mem_get_info", | ||
| 14 | ], | 29 | ], |
| 15 | TorchInGraphFunctionVariable, | 30 | TorchInGraphFunctionVariable, |
| 16 | ) | 31 | ) |
| 17 | 32 | ||
| 33 | +skip_functions_npu = dict.fromkeys( | ||
| 34 | + [ | ||
| 35 | + "torch.npu.set_device", | ||
| 36 | + "torch.npu._get_generator", | ||
| 37 | + "torch.npu._memory_viz._frames_fmt", | ||
| 38 | + "torch.npu._memory_viz._frame_fmt", | ||
| 39 | + "torch.npu.amp.autocast_mode.custom_bwd", | ||
| 40 | + "torch.npu.amp.autocast_mode.custom_fwd", | ||
| 41 | + | ||
| 42 | + ], | ||
| 43 | + SkipFunctionVariable | ||
| 44 | +) | ||
| 45 | + | ||
| 18 | 46 | ||
| 19 | def _patch_npu_trace_rules(): | 47 | def _patch_npu_trace_rules(): |
| 20 | torch._dynamo.trace_rules.clear_lru_cache() | 48 | torch._dynamo.trace_rules.clear_lru_cache() |
| 21 | torch._dynamo.trace_rules.torch_name_rule_map.append(torch_c_binding_in_graph_functions_npu) | 49 | torch._dynamo.trace_rules.torch_name_rule_map.append(torch_c_binding_in_graph_functions_npu) |
| 50 | + torch._dynamo.trace_rules.torch_name_rule_map.append(skip_functions_npu) | ||
| 51 | + torch_module.constant_fold_functions[torch.npu.current_device] = True | ||
| 52 | + torch_module.constant_fold_functions[torch.npu.get_device_properties] = True | ||
| 53 | + torch_module.constant_fold_functions[torch.npu.is_available] = True | ||
| 54 | + common_constant_types.add(torch_npu._C._NPUDeviceProperties) | ||
| 22 | 55 | ||
| @@ -49,7 +49,8 @@ def set_rng_state(new_state: torch.Tensor, device: Union[int, str, torch.device] | |||
| 49 | device (torch.device or int, optional): The device to set the RNG state. | 49 | device (torch.device or int, optional): The device to set the RNG state. |
| 50 | Default: ``'npu'`` (i.e., ``torch.device('npu')``, the current NPU device). | 50 | Default: ``'npu'`` (i.e., ``torch.device('npu')``, the current NPU device). |
| 51 | """ | 51 | """ |
| 52 | - new_state_copy = new_state.clone(memory_format=torch.contiguous_format) | 52 | + with torch._C._DisableFuncTorch(): |
| 53 | + new_state_copy = new_state.clone(memory_format=torch.contiguous_format) | ||
| 53 | if isinstance(device, str): | 54 | if isinstance(device, str): |
| 54 | device = torch.device(device) | 55 | device = torch.device(device) |
| 55 | elif isinstance(device, int): | 56 | elif isinstance(device, int): |