已合并
cpu_compile #28486
Ambi创建于 2025年12月22日
cpu_compile #28486
已合并
从已删除 :v2.8.0合入到Ascend/pytorchv2.8.0
共 2 个文件变更+72-0
| @@ -0,0 +1,47 @@ | |||
| 1 | +import os | ||
| 2 | +import torch | ||
| 3 | +import numpy as np | ||
| 4 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 5 | +from testutils import TestUtils | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +class TestNetworkCompile(TestUtils): | ||
| 9 | + | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + def test_network_compile_inference(self, input_dim, reshape_shape, device): | ||
| 13 | + class Network(torch.nn.Module): | ||
| 14 | + def __init__(self): | ||
| 15 | + super().__init__() | ||
| 16 | + self.relu = torch.nn.ReLU() | ||
| 17 | + | ||
| 18 | + def forward(self, data1): | ||
| 19 | + relu_01 = self.relu(data1) | ||
| 20 | + reshape_01 = torch.reshape(relu_01, reshape_shape) | ||
| 21 | + softmax_01 = torch.nn.functional.softmax(reshape_01, dim=1) | ||
| 22 | + sqrt_01 = torch.sqrt(softmax_01) | ||
| 23 | + relu_02 = self.relu(sqrt_01) | ||
| 24 | + square_01 = torch.square(relu_02) | ||
| 25 | + add_01 = torch.add(square_01, square_01) | ||
| 26 | + return add_01 | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + torch.manual_seed(42) | ||
| 30 | + data1 = torch.randn(input_dim, device=device) | ||
| 31 | + | ||
| 32 | + model = Network().to(device) | ||
| 33 | + model.eval() | ||
| 34 | + | ||
| 35 | + compiled_model = torch.compile(model) | ||
| 36 | + | ||
| 37 | + with torch.no_grad(): | ||
| 38 | + output = compiled_model(data1) | ||
| 39 | + cpu_out = output.detach().cpu().numpy() | ||
| 40 | + | ||
| 41 | + print(cpu_out) | ||
| 42 | + | ||
| 43 | + | ||
| 44 | +instantiate_parametrized_tests(TestNetworkCompile) | ||
| 45 | + | ||
| 46 | +if __name__ == "__main__": | ||
| 47 | + run_tests() | ||
| @@ -102,10 +102,35 @@ if (aggresive_autotune): | |||
| 102 | InductorChoices.should_use_persistent_reduction = should_use_persistent_reduction | 102 | InductorChoices.should_use_persistent_reduction = should_use_persistent_reduction |
| 103 | autotune_cache._load_cached_autotuning = _load_cached_autotuning | 103 | autotune_cache._load_cached_autotuning = _load_cached_autotuning |
| 104 | 104 | ||
| 105 | + | ||
| 106 | +def patch_device_override_func(): | ||
| 107 | + def get_device_op_overrides_patch(device_name: str): | ||
| 108 | + def register_cpu_backend(): | ||
| 109 | + from torch._inductor.codegen import cpu_device_op_overrides | ||
| 110 | + | ||
| 111 | + return | ||
| 112 | + | ||
| 113 | + def register_mps_backend(): | ||
| 114 | + from torch._inductor.codegen import mps_device_op_overrides | ||
| 115 | + | ||
| 116 | + return | ||
| 117 | + | ||
| 118 | + backend_factory = {"cpu": register_cpu_backend, "mps": register_mps_backend} | ||
| 119 | + | ||
| 120 | + if device_name not in torch._inductor.codegen.common.device_op_overrides_dict: | ||
| 121 | + if device_name not in backend_factory: | ||
| 122 | + raise ValueError("backend not found: ", device_name) | ||
| 123 | + backend_factory[device_name]() | ||
| 124 | + | ||
| 125 | + return torch._inductor.codegen.common.device_op_overrides_dict[device_name] | ||
| 126 | + | ||
| 127 | + torch._inductor.graph.get_device_op_overrides = get_device_op_overrides_patch | ||
| 128 | + | ||
| 105 | register_fa_pass() | 129 | register_fa_pass() |
| 106 | patch_cache_base_get_system() | 130 | patch_cache_base_get_system() |
| 107 | patch_is_gpu() | 131 | patch_is_gpu() |
| 108 | patch_has_triton() | 132 | patch_has_triton() |
| 109 | disable_foreach() | 133 | disable_foreach() |
| 110 | patch_get_optimization_cflags() | 134 | patch_get_optimization_cflags() |
| 135 | +patch_device_override_func() | ||
| 111 | 136 | ||