已合并
refactor(inductor): isolate NPU integrations by device #42564
Xuan Peng创建于 7月23日
refactor(inductor): isolate NPU integrations by device #42564
已合并
共 27 个文件变更+1237-948
| @@ -0,0 +1,108 @@ | |||
| 1 | +import contextlib | ||
| 2 | +import inspect | ||
| 3 | +from types import SimpleNamespace | ||
| 4 | +from unittest import mock | ||
| 5 | + | ||
| 6 | +from torch._inductor.select_algorithm import AlgorithmSelectorCache | ||
| 7 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 8 | + | ||
| 9 | + | ||
| 10 | +_UPSTREAM_CALL = AlgorithmSelectorCache.__call__ | ||
| 11 | +_UPSTREAM_MAKE_BENCHMARK_FN = AlgorithmSelectorCache.__dict__[ | ||
| 12 | + "make_benchmark_fn" | ||
| 13 | +].__func__ | ||
| 14 | + | ||
| 15 | +import torch_npu._inductor.select_algorithm as npu_select_algorithm # noqa: E402 | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +def _layout(device_type): | ||
| 19 | + return SimpleNamespace(device=SimpleNamespace(type=device_type)) | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | +def _install_with_parents(parent_call, parent_make_benchmark_fn): | ||
| 24 | + with ( | ||
| 25 | + mock.patch.object(AlgorithmSelectorCache, "__call__", parent_call), | ||
| 26 | + mock.patch.object( | ||
| 27 | + AlgorithmSelectorCache, | ||
| 28 | + "make_benchmark_fn", | ||
| 29 | + classmethod(parent_make_benchmark_fn), | ||
| 30 | + ), | ||
| 31 | + ): | ||
| 32 | + npu_select_algorithm.patch_algorithm_selector() | ||
| 33 | + yield ( | ||
| 34 | + AlgorithmSelectorCache.__call__, | ||
| 35 | + AlgorithmSelectorCache.__dict__["make_benchmark_fn"].__func__, | ||
| 36 | + ) | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +class TestAlgorithmSelectorDispatch(TestCase): | ||
| 40 | + def test_wrappers_preserve_v213_signatures(self): | ||
| 41 | + with _install_with_parents( | ||
| 42 | + _UPSTREAM_CALL, _UPSTREAM_MAKE_BENCHMARK_FN | ||
| 43 | + ) as (patched_call, patched_make_benchmark_fn): | ||
| 44 | + self.assertIs(patched_call.__wrapped__, _UPSTREAM_CALL) | ||
| 45 | + self.assertIs( | ||
| 46 | + patched_make_benchmark_fn.__wrapped__, | ||
| 47 | + _UPSTREAM_MAKE_BENCHMARK_FN, | ||
| 48 | + ) | ||
| 49 | + self.assertEqual( | ||
| 50 | + inspect.signature(patched_call), | ||
| 51 | + inspect.signature(_UPSTREAM_CALL), | ||
| 52 | + ) | ||
| 53 | + self.assertEqual( | ||
| 54 | + inspect.signature(patched_make_benchmark_fn), | ||
| 55 | + inspect.signature(_UPSTREAM_MAKE_BENCHMARK_FN), | ||
| 56 | + ) | ||
| 57 | + | ||
| 58 | + def test_non_npu_call_delegates_and_preserves_parent_result(self): | ||
| 59 | + choice = SimpleNamespace(output_node=mock.Mock()) | ||
| 60 | + parent_result = (object(), object()) | ||
| 61 | + | ||
| 62 | + for device_type in ("cpu", "cuda", "xpu"): | ||
| 63 | + with self.subTest(device_type=device_type): | ||
| 64 | + parent_call = mock.Mock(return_value=parent_result) | ||
| 65 | + layout = _layout(device_type) | ||
| 66 | + | ||
| 67 | + with _install_with_parents(parent_call, mock.Mock()) as ( | ||
| 68 | + patched_call, | ||
| 69 | + _, | ||
| 70 | + ): | ||
| 71 | + result = patched_call( | ||
| 72 | + object(), | ||
| 73 | + "test", | ||
| 74 | + [choice], | ||
| 75 | + [], | ||
| 76 | + layout, | ||
| 77 | + best_config_future=object(), | ||
| 78 | + ) | ||
| 79 | + | ||
| 80 | + self.assertIs(result, parent_result) | ||
| 81 | + parent_call.assert_called_once() | ||
| 82 | + choice.output_node.assert_not_called() | ||
| 83 | + | ||
| 84 | + def test_plain_npu_choice_returns_tuple_abi(self): | ||
| 85 | + selected_node = object() | ||
| 86 | + choice = SimpleNamespace(output_node=mock.Mock(return_value=selected_node)) | ||
| 87 | + parent_call = mock.Mock() | ||
| 88 | + | ||
| 89 | + with _install_with_parents(parent_call, mock.Mock()) as (patched_call, _): | ||
| 90 | + result = patched_call( | ||
| 91 | + object(), | ||
| 92 | + "test", | ||
| 93 | + [choice], | ||
| 94 | + [], | ||
| 95 | + _layout("npu"), | ||
| 96 | + best_config_future=object(), | ||
| 97 | + is_collective=True, | ||
| 98 | + min_speedup_threshold=1.5, | ||
| 99 | + benchmark_with_cudagraphs=True, | ||
| 100 | + ) | ||
| 101 | + | ||
| 102 | + self.assertIs(result[0], selected_node) | ||
| 103 | + self.assertIs(result[1], choice) | ||
| 104 | + parent_call.assert_not_called() | ||
| 105 | + | ||
| 106 | + | ||
| 107 | +if __name__ == "__main__": | ||
| 108 | + run_tests() | ||
| @@ -0,0 +1,128 @@ | |||
| 1 | +import os | ||
| 2 | +from unittest import mock | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | +from torch._inductor import ( | ||
| 6 | + autotune_process as inductor_autotune_process, | ||
| 7 | + config as inductor_config, | ||
| 8 | + utils as inductor_utils, | ||
| 9 | +) | ||
| 10 | +from torch._inductor.autotune_process import TuningProcess, TuningProcessPool | ||
| 11 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +COMMUNITY_GET_DEVICE_LIST = TuningProcessPool.get_device_list | ||
| 15 | + | ||
| 16 | +from torch_npu._inductor import utils as npu_utils | ||
| 17 | +from torch_npu._inductor.autotune_process import ( | ||
| 18 | + ASCEND_VISIBLE_DEVICES, | ||
| 19 | + patch_tuning_process, | ||
| 20 | +) | ||
| 21 | + | ||
| 22 | + | ||
| 23 | +class TestAutotuneProcessAdapter(TestCase): | ||
| 24 | + def setUp(self): | ||
| 25 | + super().setUp() | ||
| 26 | + self.visible_devices_key = inductor_autotune_process.CUDA_VISIBLE_DEVICES | ||
| 27 | + self.pool = object.__new__(TuningProcessPool) | ||
| 28 | + | ||
| 29 | + def tearDown(self): | ||
| 30 | + inductor_autotune_process.CUDA_VISIBLE_DEVICES = self.visible_devices_key | ||
| 31 | + super().tearDown() | ||
| 32 | + | ||
| 33 | + def test_pool_get_device_list_is_community_method(self): | ||
| 34 | + self.assertEqual( | ||
| 35 | + COMMUNITY_GET_DEVICE_LIST.__module__, | ||
| 36 | + "torch._inductor.autotune_process", | ||
| 37 | + ) | ||
| 38 | + self.assertIs(TuningProcessPool.get_device_list, COMMUNITY_GET_DEVICE_LIST) | ||
| 39 | + | ||
| 40 | + def test_single_device_mode(self): | ||
| 41 | + with inductor_config.patch("autotune_multi_device", False): | ||
| 42 | + self.assertEqual(self.pool.get_device_list(), [None]) | ||
| 43 | + | ||
| 44 | + def test_multi_device_uses_npu_interface_and_visible_key(self): | ||
| 45 | + interface = mock.Mock() | ||
| 46 | + interface.device_count.return_value = 4 | ||
| 47 | + patch_tuning_process() | ||
| 48 | + with ( | ||
| 49 | + inductor_config.patch("autotune_multi_device", True), | ||
| 50 | + mock.patch.object( | ||
| 51 | + inductor_autotune_process, "get_gpu_type", return_value="npu" | ||
| 52 | + ), | ||
| 53 | + mock.patch.object( | ||
| 54 | + inductor_autotune_process, | ||
| 55 | + "get_interface_for_device", | ||
| 56 | + return_value=interface, | ||
| 57 | + ), | ||
| 58 | + mock.patch.dict( | ||
| 59 | + os.environ, {ASCEND_VISIBLE_DEVICES: "3,1"}, clear=True | ||
| 60 | + ), | ||
| 61 | + ): | ||
| 62 | + self.assertEqual(self.pool.get_device_list(), [3, 1]) | ||
| 63 | + | ||
| 64 | + with ( | ||
| 65 | + inductor_config.patch("autotune_multi_device", True), | ||
| 66 | + mock.patch.object( | ||
| 67 | + inductor_autotune_process, "get_gpu_type", return_value="npu" | ||
| 68 | + ), | ||
| 69 | + mock.patch.object( | ||
| 70 | + inductor_autotune_process, | ||
| 71 | + "get_interface_for_device", | ||
| 72 | + return_value=interface, | ||
| 73 | + ), | ||
| 74 | + mock.patch.dict(os.environ, {}, clear=True), | ||
| 75 | + ): | ||
| 76 | + self.assertEqual(self.pool.get_device_list(), [0, 1, 2, 3]) | ||
| 77 | + | ||
| 78 | + def test_tuning_process_scopes_visible_device_to_child(self): | ||
| 79 | + patch_tuning_process() | ||
| 80 | + with ( | ||
| 81 | + mock.patch.dict( | ||
| 82 | + os.environ, {ASCEND_VISIBLE_DEVICES: "2,3"}, clear=True | ||
| 83 | + ), | ||
| 84 | + mock.patch.object( | ||
| 85 | + inductor_autotune_process.subprocess, "Popen" | ||
| 86 | + ) as popen_mock, | ||
| 87 | + ): | ||
| 88 | + process = TuningProcess(1) | ||
| 89 | + self.addCleanup(process.write_pipe.close) | ||
| 90 | + self.addCleanup(process.read_pipe.close) | ||
| 91 | + self.addCleanup(process.selector.close) | ||
| 92 | + | ||
| 93 | + child_env = popen_mock.call_args.kwargs["env"] | ||
| 94 | + self.assertEqual(child_env[ASCEND_VISIBLE_DEVICES], "1") | ||
| 95 | + self.assertEqual(os.environ[ASCEND_VISIBLE_DEVICES], "2,3") | ||
| 96 | + | ||
| 97 | + | ||
| 98 | +class TestPatchIsGpu(TestCase): | ||
| 99 | + def setUp(self): | ||
| 100 | + super().setUp() | ||
| 101 | + self.gpu_types = list(inductor_utils.GPU_TYPES) | ||
| 102 | + inductor_utils.get_gpu_type.cache_clear() | ||
| 103 | + | ||
| 104 | + def tearDown(self): | ||
| 105 | + inductor_utils.GPU_TYPES[:] = self.gpu_types | ||
| 106 | + inductor_utils.get_gpu_type.cache_clear() | ||
| 107 | + super().tearDown() | ||
| 108 | + | ||
| 109 | + def test_patch_is_gpu_is_idempotent_and_clears_cached_device(self): | ||
| 110 | + inductor_utils.GPU_TYPES[:] = ["cuda"] | ||
| 111 | + with mock.patch.object(torch.cuda, "is_available", return_value=False): | ||
| 112 | + self.assertEqual(inductor_utils.get_gpu_type(), "cuda") | ||
| 113 | + | ||
| 114 | + with ( | ||
| 115 | + mock.patch.object(torch.cuda, "is_available", return_value=False), | ||
| 116 | + mock.patch.object(torch.npu, "is_available", return_value=True), | ||
| 117 | + ): | ||
| 118 | + npu_utils.patch_is_gpu() | ||
| 119 | + self.assertEqual(inductor_utils.get_gpu_type(), "npu") | ||
| 120 | + self.assertEqual(inductor_utils.get_gpu_type.cache_info().currsize, 1) | ||
| 121 | + npu_utils.patch_is_gpu() | ||
| 122 | + self.assertEqual(inductor_utils.get_gpu_type.cache_info().currsize, 0) | ||
| 123 | + | ||
| 124 | + self.assertEqual(inductor_utils.GPU_TYPES.count("npu"), 1) | ||
| 125 | + | ||
| 126 | + | ||
| 127 | +if __name__ == "__main__": | ||
| 128 | + run_tests() | ||
| @@ -1,6 +1,4 @@ | |||
| 1 | -import os | ||
| 2 | import torch | 1 | import torch |
| 3 | -import numpy as np | ||
| 4 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | 2 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests |
| 5 | from testutils import TestUtils | 3 | from testutils import TestUtils |
| 6 | 4 | ||
| @@ -40,8 +38,51 @@ class TestNetworkCompile(TestUtils): | |||
| 40 | 38 | ||
| 41 | print(cpu_out) | 39 | print(cpu_out) |
| 42 | 40 | ||
| 41 | + def test_gemm_lowerings_are_device_dispatched(self): | ||
| 42 | + from torch._inductor import lowering | ||
| 43 | + | ||
| 44 | + for target in ( | ||
| 45 | + torch.ops.aten.mm.default, | ||
| 46 | + torch.ops.aten.addmm.default, | ||
| 47 | + torch.ops.aten.bmm.default, | ||
| 48 | + ): | ||
| 49 | + with self.subTest(target=target): | ||
| 50 | + handler = lowering.lowerings[target] | ||
| 51 | + self.assertTrue( | ||
| 52 | + getattr( | ||
| 53 | + handler, | ||
| 54 | + "_torch_npu_device_lowering_dispatch", | ||
| 55 | + False, | ||
| 56 | + ) | ||
| 57 | + ) | ||
| 58 | + | ||
| 59 | + def test_cpu_gemm_compile_after_npu_registration(self): | ||
| 60 | + cases = ( | ||
| 61 | + ( | ||
| 62 | + "mm", | ||
| 63 | + lambda mat1, mat2: torch.mm(mat1, mat2), | ||
| 64 | + (torch.randn(4, 5), torch.randn(5, 6)), | ||
| 65 | + ), | ||
| 66 | + ( | ||
| 67 | + "addmm", | ||
| 68 | + lambda bias, mat1, mat2: torch.addmm(bias, mat1, mat2), | ||
| 69 | + (torch.randn(4, 6), torch.randn(4, 5), torch.randn(5, 6)), | ||
| 70 | + ), | ||
| 71 | + ( | ||
| 72 | + "bmm", | ||
| 73 | + lambda mat1, mat2: torch.bmm(mat1, mat2), | ||
| 74 | + (torch.randn(2, 4, 5), torch.randn(2, 5, 6)), | ||
| 75 | + ), | ||
| 76 | + ) | ||
| 77 | + | ||
| 78 | + for name, fn, args in cases: | ||
| 79 | + with self.subTest(name=name): | ||
| 80 | + torch._dynamo.reset() | ||
| 81 | + compiled = torch.compile(fn, fullgraph=True) | ||
| 82 | + self.assertEqual(compiled(*args), fn(*args)) | ||
 | |||
| 83 | + | ||
| 43 | 84 | ||
| 44 | instantiate_parametrized_tests(TestNetworkCompile) | 85 | instantiate_parametrized_tests(TestNetworkCompile) |
| 45 | 86 | ||
| 46 | if __name__ == "__main__": | 87 | if __name__ == "__main__": |
| 47 | - run_tests() | 88 | + run_tests() |
| @@ -9,7 +9,6 @@ import torch_npu # noqa: F401 | |||
| 9 | src_code_1 = ''' | 9 | src_code_1 = ''' |
| 10 | import triton | 10 | import triton |
| 11 | import triton.language as tl | 11 | import triton.language as tl |
| 12 | -from triton.compiler.compiler import AttrsDescriptor | ||
| 13 | 12 | ||
| 14 | from torch._inductor.runtime import triton_helpers, triton_heuristics | 13 | from torch._inductor.runtime import triton_helpers, triton_heuristics |
| 15 | from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math | 14 | from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math |
| @@ -28,11 +27,12 @@ import torch_npu | |||
| 28 | 'device': DeviceProperties(type='npu', index=0, multi_processor_count=40, cc='Ascend910B3', | 27 | 'device': DeviceProperties(type='npu', index=0, multi_processor_count=40, cc='Ascend910B3', |
| 29 | major=None, regs_per_multiprocessor=None, max_threads_per_multi_processor=None, warp_size=32), | 28 | major=None, regs_per_multiprocessor=None, max_threads_per_multi_processor=None, warp_size=32), |
| 30 | 'constants': {}, 'mix_mode': 'aiv'}, | 29 | 'constants': {}, 'mix_mode': 'aiv'}, |
| 31 | - inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_unk_fused_add_0', 'mutated_arg_names': [], | 30 | + inductor_meta={'grid_type': 'GridNpu', 'autotune_hints': set(), 'kernel_name': 'triton_unk_fused_add_0', 'mutated_arg_names': [], |
| 32 | 'backend_hash': 'bc71dba4086164e7ac2b0779fa861dbf7467f0265d4a57b8f48cf6dda02b150f', 'split_axis': [0], | 31 | 'backend_hash': 'bc71dba4086164e7ac2b0779fa861dbf7467f0265d4a57b8f48cf6dda02b150f', 'split_axis': [0], |
| 33 | - 'tiling_axis': [0, 1], 'no_loop_axis': [1], 'axis_names': ['y0', 'x1'], 'low_dims': {1}, 'numof_reduction_axis': 0, 'inductor_ascend_linear_mode': 'linear', | 32 | + 'tiling_axis': [0, 1], 'no_loop_axis': [1], 'axis_names': ['y0', 'x1'], |
| 33 | + 'axis_static_values': (('y0', 16384), ('x1', 32)), 'low_dims': {1}, 'numof_reduction_axis': 0, | ||
| 34 | 'split_axis_dtype': torch.float16, 'dual_reduction': False, 'traced_graph_hash': 'TRACED_GRAPH_HASH', | 34 | 'split_axis_dtype': torch.float16, 'dual_reduction': False, 'traced_graph_hash': 'TRACED_GRAPH_HASH', |
| 35 | - 'traced_graph_dir': 'TRACED_GRAPH_DIR'}, | 35 | + 'traced_graph_dir': 'TRACED_GRAPH_DIR', 'runtime_block_arg_names': ('Y0BLOCK',)}, |
| 36 | min_elem_per_thread=0 | 36 | min_elem_per_thread=0 |
| 37 | ) | 37 | ) |
| 38 | 38 | ||
| @@ -1,5 +1,6 @@ | |||
| 1 | import os | 1 | import os |
| 2 | 2 | ||
| 3 | + | ||
| 3 | # Some ops (e.g. aten.matmul_backward) only get their NPU meta registration when | 4 | # Some ops (e.g. aten.matmul_backward) only get their NPU meta registration when |
| 4 | # compatible impl mode is enabled. This env var is read at torch_npu import time, | 5 | # compatible impl mode is enabled. This env var is read at torch_npu import time, |
| 5 | # so it must be set before importing torch_npu. | 6 | # so it must be set before importing torch_npu. |
| @@ -14,14 +15,20 @@ import unittest | |||
| 14 | import warnings | 15 | import warnings |
| 15 | import weakref | 16 | import weakref |
| 16 | from io import StringIO | 17 | from io import StringIO |
| 18 | +from types import SimpleNamespace | ||
| 19 | +from unittest.mock import Mock | ||
| 17 | 20 | ||
| 18 | import torch | 21 | import torch |
| 19 | -import torch.nn as nn # noqa: F401 | ||
| 20 | import torch._dynamo.config as dynamo_config | 22 | import torch._dynamo.config as dynamo_config |
| 23 | +import torch.nn as nn # noqa: F401 | ||
| 21 | from torch._inductor import config | 24 | from torch._inductor import config |
| 22 | from torch._inductor.compile_fx import compile_fx_inner | 25 | from torch._inductor.compile_fx import compile_fx_inner |
| 23 | -from torch._inductor.utils import run_and_get_code | 26 | +from torch._inductor.dependencies import ReadWrites, StarDep, WeakDep |
| 27 | +from torch._inductor.ir import GraphPartitionSignature, NoneLayout | ||
| 28 | +from torch._inductor.scheduler import Scheduler | ||
| 24 | from torch._inductor.test_case import TestCase as InductorTestCase | 29 | from torch._inductor.test_case import TestCase as InductorTestCase |
| 30 | +from torch._inductor.utils import run_and_get_code | ||
| 31 | +from torch._inductor.virtualized import V | ||
| 25 | from torch.fx.experimental.proxy_tensor import make_fx | 32 | from torch.fx.experimental.proxy_tensor import make_fx |
| 26 | from torch.testing import FileCheck | 33 | from torch.testing import FileCheck |
| 27 | from torch.testing._internal.common_utils import ( # noqa: F401 | 34 | from torch.testing._internal.common_utils import ( # noqa: F401 |
| @@ -29,11 +36,17 @@ from torch.testing._internal.common_utils import ( # noqa: F401 | |||
| 29 | parametrize, | 36 | parametrize, |
| 30 | ) | 37 | ) |
| 31 | from torch.testing._internal.logging_utils import logs_to_string | 38 | from torch.testing._internal.logging_utils import logs_to_string |
| 39 | +from torch.utils._ordered_set import OrderedSet | ||
| 32 | from torch.utils._python_dispatch import TorchDispatchMode | 40 | from torch.utils._python_dispatch import TorchDispatchMode |
| 33 | 41 | ||
| 42 | + | ||
| 43 | +PARENT_GET_GRAPH_PARTITION_SIGNATURE = Scheduler.get_graph_partition_signature | ||
| 44 | + | ||
| 34 | import torch_npu # noqa: F401 | 45 | import torch_npu # noqa: F401 |
| 46 | +import torch_npu._inductor # noqa: F401 | ||
| 35 | from torch_npu.npu._graph_tree import get_container | 47 | from torch_npu.npu._graph_tree import get_container |
| 36 | 48 | ||
| 49 | + | ||
| 37 | TEST_NPU = torch.npu.is_available() | 50 | TEST_NPU = torch.npu.is_available() |
| 38 | aten = torch.ops.aten | 51 | aten = torch.ops.aten |
| 39 | 52 | ||
| @@ -103,6 +116,111 @@ class TestCase(InductorTestCase): | |||
| 103 | torch._dynamo.reset() | 116 | torch._dynamo.reset() |
| 104 | 117 | ||
| 105 | 118 | ||
| 119 | +class TestGraphPartitionSchedulerContract(TestCase): | ||
| 120 | + """Locks NPU Inductor to the v2.13 graph-partition implementation.""" | ||
| 121 | + | ||
| 122 | + def test_uses_parent_signature_with_extra_input_weakdep_and_mutation_alias(self): | ||
| 123 | + self.assertIs( | ||
| 124 | + Scheduler.get_graph_partition_signature, | ||
| 125 | + PARENT_GET_GRAPH_PARTITION_SIGNATURE, | ||
| 126 | + ) | ||
| 127 | + | ||
| 128 | + scheduler = Scheduler.__new__(Scheduler) | ||
| 129 | + scheduler.mutation_real_name = {"mutation_alias": "mutation_real"} | ||
| 130 | + scheduler.name_to_buf = { | ||
| 131 | + "mutation_alias": SimpleNamespace( | ||
| 132 | + node=SimpleNamespace(layout=NoneLayout(device=None)) | ||
| 133 | + ), | ||
| 134 | + "mutation_real": SimpleNamespace( | ||
| 135 | + node=SimpleNamespace(layout=object()) | ||
| 136 | + ), | ||
| 137 | + } | ||
| 138 | + input_nodes = { | ||
| 139 | + name: SimpleNamespace(name=name) | ||
| 140 | + for name in ("cross_partition", "mutation_real", "weak_dependency") | ||
| 141 | + } | ||
| 142 | + scheduler.get_name_to_nodes = lambda: input_nodes | ||
| 143 | + scheduler.get_graph_partition_symbol_inputs = ( | ||
| 144 | + lambda partition, inputs: OrderedSet() | ||
| 145 | + ) | ||
| 146 | + | ||
| 147 | + partition_node = SimpleNamespace( | ||
| 148 | + outputs_by_name={"produced_here": object()}, | ||
| 149 | + read_writes=ReadWrites( | ||
| 150 | + OrderedSet( | ||
| 151 | + [ | ||
| 152 | + StarDep("mutation_alias"), | ||
| 153 | + WeakDep("weak_dependency", mutating_buf="produced_here"), | ||
| 154 | + ] | ||
| 155 | + ), | ||
| 156 | + OrderedSet(), | ||
| 157 | + OrderedSet(), | ||
| 158 | + ), | ||
| 159 | + last_usage=OrderedSet(["cross_partition"]), | ||
| 160 | + ) | ||
| 161 | + graph = SimpleNamespace( | ||
| 162 | + get_output_names=lambda: OrderedSet(), constants=OrderedSet() | ||
| 163 | + ) | ||
| 164 | + | ||
| 165 | + with V.set_graph_handler(graph): | ||
| 166 | + signature = scheduler.get_graph_partition_signature( | ||
| 167 | + [[partition_node]], [False] | ||
| 168 | + )[0] | ||
| 169 | + | ||
| 170 | + self.assertEqual( | ||
| 171 | + list(signature.input_nodes), ["mutation_real", "cross_partition"] | ||
| 172 | + ) | ||
| 173 | + self.assertNotIn("weak_dependency", signature.input_nodes) | ||
| 174 | + self.assertEqual( | ||
| 175 | + signature.input_deallocation, | ||
| 176 | + {"mutation_real": False, "cross_partition": True}, | ||
| 177 | + ) | ||
| 178 | + | ||
| 179 | + def test_cleans_only_buffers_removed_during_codegen(self): | ||
| 180 | + scheduler = Scheduler.__new__(Scheduler) | ||
| 181 | + | ||
| 182 | + live_output = Mock() | ||
| 183 | + live_output.maybe_get_name.return_value = "live_output" | ||
| 184 | + prior_output = Mock() | ||
| 185 | + prior_output.maybe_get_name.return_value = "prior_output" | ||
| 186 | + codegen_output = Mock() | ||
| 187 | + codegen_output.maybe_get_name.return_value = "codegen_output" | ||
| 188 | + signature = GraphPartitionSignature( | ||
| 189 | + symbol_inputs=OrderedSet(), | ||
| 190 | + input_nodes={ | ||
| 191 | + "live_input": Mock(), | ||
| 192 | + "prior_input": Mock(), | ||
| 193 | + "codegen_input": Mock(), | ||
| 194 | + }, | ||
| 195 | + output_nodes=[live_output, prior_output, codegen_output], | ||
| 196 | + input_deallocation={ | ||
| 197 | + "live_input": False, | ||
| 198 | + "prior_input": True, | ||
| 199 | + "codegen_input": False, | ||
| 200 | + }, | ||
| 201 | + skip_cudagraph=False, | ||
| 202 | + constant_names=["live_constant", "prior_constant", "codegen_constant"], | ||
| 203 | + ) | ||
| 204 | + removed_before_codegen = OrderedSet( | ||
| 205 | + ["prior_input", "prior_output", "prior_constant"] | ||
| 206 | + ) | ||
| 207 | + removed_after_codegen = removed_before_codegen | OrderedSet( | ||
| 208 | + ["codegen_input", "codegen_output", "codegen_constant"] | ||
| 209 | + ) | ||
| 210 | + | ||
| 211 | + cleaned = scheduler.clean_removed_buffer_from_partition_signatures( | ||
| 212 | + signature, removed_after_codegen - removed_before_codegen | ||
| 213 | + ) | ||
| 214 | + | ||
| 215 | + self.assertEqual(list(cleaned.input_nodes), ["live_input", "prior_input"]) | ||
| 216 | + self.assertEqual( | ||
| 217 | + cleaned.input_deallocation, | ||
| 218 | + {"live_input": False, "prior_input": True}, | ||
| 219 | + ) | ||
| 220 | + self.assertEqual(cleaned.output_nodes, [live_output, prior_output]) | ||
| 221 | + self.assertEqual(cleaned.constant_names, ["live_constant", "prior_constant"]) | ||
| 222 | + | ||
| 223 | + | ||
| 106 | # =========================================================================== | 224 | # =========================================================================== |
| 107 | # Graph Partition Tests — Codegen correctness | 225 | # Graph Partition Tests — Codegen correctness |
| 108 | # (ported from test_torchinductor.py, device-agnostic via self.device) | 226 | # (ported from test_torchinductor.py, device-agnostic via self.device) |
| @@ -326,7 +444,7 @@ class TestGraphPartitionNPU(TestCase): | |||
| 326 | inp = torch.rand([20, 20], device="npu", requires_grad=True) | 444 | inp = torch.rand([20, 20], device="npu", requires_grad=True) |
| 327 | out = foo(inp) | 445 | out = foo(inp) |
| 328 | 446 | ||
| 329 | - with config.patch(always_complex_memory_overlap_TESTING_ONLY=True): | 447 | + with config.patch(force_disable_cudagraph_TESTING_ONLY=True): |
| 330 | back_inp = torch.empty_strided([20, 20], [0, 1], device="npu") | 448 | back_inp = torch.empty_strided([20, 20], [0, 1], device="npu") |
| 331 | out.backward(back_inp) | 449 | out.backward(back_inp) |
| 332 | 450 | ||
| @@ -0,0 +1,236 @@ | |||
| 1 | +from types import SimpleNamespace | ||
| 2 | +from unittest import mock | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 6 | + | ||
| 7 | +import torch_npu._inductor.lowering_patch as lowering_patch | ||
| 8 | +from torch_npu._inductor.lowering_common import LOWERING_REGISTRY_ATTRS | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class _IRValue: | ||
| 12 | + def __init__(self, device_type): | ||
| 13 | + self.device = torch.device(device_type) | ||
| 14 | + | ||
| 15 | + def get_device(self): | ||
| 16 | + return self.device | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +class TestLoweringDeviceDispatch(TestCase): | ||
| 20 | + def setUp(self): | ||
| 21 | + super().setUp() | ||
| 22 | + self.target = "target" | ||
| 23 | + self.upstream_call = mock.Mock(return_value="upstream") | ||
| 24 | + self.device_call = mock.Mock(return_value="npu") | ||
| 25 | + | ||
| 26 | + def upstream(*args, **kwargs): | ||
| 27 | + return self.upstream_call(*args, **kwargs) | ||
| 28 | + | ||
| 29 | + def device_handler(*args, **kwargs): | ||
| 30 | + return self.device_call(*args, **kwargs) | ||
| 31 | + | ||
| 32 | + self.upstream = upstream | ||
| 33 | + self.device_handler = device_handler | ||
| 34 | + self.registry = {self.target: self.device_handler} | ||
| 35 | + self.make_reduction = object() | ||
| 36 | + self.lowering = SimpleNamespace( | ||
| 37 | + lowerings=self.registry, | ||
| 38 | + make_reduction=self.make_reduction, | ||
| 39 | + ) | ||
| 40 | + registry_copies = {} | ||
| 41 | + for attr in LOWERING_REGISTRY_ATTRS: | ||
| 42 | + if attr == "lowerings": | ||
| 43 | + registry_copies[attr] = {self.target: self.upstream} | ||
| 44 | + continue | ||
| 45 | + value = {} | ||
| 46 | + setattr(self.lowering, attr, value) | ||
| 47 | + registry_copies[attr] = {} | ||
| 48 | + | ||
| 49 | + self.baseline = lowering_patch.LoweringSnapshot( | ||
| 50 | + functions={}, | ||
| 51 | + lowerings_ref=self.registry, | ||
| 52 | + lowerings_copy={self.target: self.upstream}, | ||
| 53 | + registry_copies=registry_copies, | ||
| 54 | + make_reduction=self.make_reduction, | ||
| 55 | + ) | ||
| 56 | + self.get_lowering_patch = mock.patch.object( | ||
| 57 | + lowering_patch, | ||
| 58 | + "_get_inductor_lowering", | ||
| 59 | + return_value=self.lowering, | ||
| 60 | + ) | ||
| 61 | + self.capture_patch = mock.patch.object( | ||
| 62 | + lowering_patch, | ||
| 63 | + "capture_lowering_baseline", | ||
| 64 | + return_value=self.baseline, | ||
| 65 | + ) | ||
| 66 | + self.get_lowering_patch.start() | ||
| 67 | + self.capture_patch.start() | ||
| 68 | + | ||
| 69 | + def tearDown(self): | ||
| 70 | + self.capture_patch.stop() | ||
| 71 | + self.get_lowering_patch.stop() | ||
| 72 | + super().tearDown() | ||
| 73 | + | ||
| 74 | + def install(self, targets=None): | ||
| 75 | + registry_id = id(self.registry) | ||
| 76 | + lowering_patch.install_device_lowering_dispatch( | ||
| 77 | + targets or [self.target] | ||
| 78 | + ) | ||
| 79 | + self.assertEqual(id(self.registry), registry_id) | ||
| 80 | + return self.registry[self.target] | ||
| 81 | + | ||
| 82 | + def test_non_npu_layout_uses_upstream_handler(self): | ||
| 83 | + dispatcher = self.install() | ||
| 84 | + | ||
| 85 | + for device_type in ("cpu", "cuda", "xpu"): | ||
| 86 | + with self.subTest(device_type=device_type): | ||
| 87 | + result = dispatcher( | ||
| 88 | + _IRValue(device_type), | ||
| 89 | + layout=SimpleNamespace(device=torch.device(device_type)), | ||
| 90 | + ) | ||
| 91 | + | ||
| 92 | + self.assertEqual(result, "upstream") | ||
| 93 | + | ||
| 94 | + self.assertEqual(self.upstream_call.call_count, 3) | ||
| 95 | + self.device_call.assert_not_called() | ||
| 96 | + | ||
| 97 | + def test_npu_layout_wins_over_cpu_metadata(self): | ||
| 98 | + dispatcher = self.install() | ||
| 99 | + | ||
| 100 | + result = dispatcher( | ||
| 101 | + _IRValue("cpu"), | ||
| 102 | + layout=SimpleNamespace(device=torch.device("npu")), | ||
| 103 | + ) | ||
| 104 | + | ||
| 105 | + self.assertEqual(result, "npu") | ||
| 106 | + self.device_call.assert_called_once() | ||
| 107 | + self.upstream_call.assert_not_called() | ||
| 108 | + | ||
| 109 | + def test_non_npu_layout_wins_over_npu_input(self): | ||
| 110 | + dispatcher = self.install() | ||
| 111 | + | ||
| 112 | + result = dispatcher( | ||
| 113 | + _IRValue("npu"), | ||
| 114 | + layout=SimpleNamespace(device=torch.device("cpu")), | ||
| 115 | + ) | ||
| 116 | + | ||
| 117 | + self.assertEqual(result, "upstream") | ||
| 118 | + self.upstream_call.assert_called_once() | ||
| 119 | + self.device_call.assert_not_called() | ||
| 120 | + | ||
| 121 | + def test_npu_input_without_layout_uses_device_handler(self): | ||
| 122 | + dispatcher = self.install() | ||
| 123 | + | ||
| 124 | + result = dispatcher({"nested": [_IRValue("npu")]}) | ||
| 125 | + | ||
| 126 | + self.assertEqual(result, "npu") | ||
| 127 | + self.device_call.assert_called_once() | ||
| 128 | + self.upstream_call.assert_not_called() | ||
| 129 | + | ||
| 130 | + def test_missing_device_information_uses_upstream(self): | ||
| 131 | + dispatcher = self.install() | ||
| 132 | + | ||
| 133 | + result = dispatcher(1, flag=True) | ||
| 134 | + | ||
| 135 | + self.assertEqual(result, "upstream") | ||
| 136 | + self.upstream_call.assert_called_once() | ||
| 137 | + self.device_call.assert_not_called() | ||
| 138 | + | ||
| 139 | + def test_missing_upstream_handler_leaves_target_unchanged(self): | ||
| 140 | + self.baseline.lowerings_copy.clear() | ||
| 141 | + | ||
| 142 | + dispatcher = self.install() | ||
| 143 | + | ||
| 144 | + self.assertIs(dispatcher, self.device_handler) | ||
| 145 | + | ||
| 146 | + def test_repeated_install_is_idempotent(self): | ||
| 147 | + first = self.install() | ||
| 148 | + | ||
| 149 | + second = self.install() | ||
| 150 | + | ||
| 151 | + self.assertIs(second, first) | ||
| 152 | + self.assertTrue( | ||
| 153 | + getattr(second, "_torch_npu_device_lowering_dispatch", False) | ||
| 154 | + ) | ||
| 155 | + | ||
| 156 | + def test_multiple_targets_capture_their_own_handlers(self): | ||
| 157 | + second_target = "second_target" | ||
| 158 | + second_upstream_call = mock.Mock(return_value="second_upstream") | ||
| 159 | + second_device_call = mock.Mock(return_value="second_npu") | ||
| 160 | + | ||
| 161 | + def second_upstream(*args, **kwargs): | ||
| 162 | + return second_upstream_call(*args, **kwargs) | ||
| 163 | + | ||
| 164 | + def second_device(*args, **kwargs): | ||
| 165 | + return second_device_call(*args, **kwargs) | ||
| 166 | + | ||
| 167 | + self.baseline.lowerings_copy[second_target] = second_upstream | ||
| 168 | + self.registry[second_target] = second_device | ||
| 169 | + | ||
| 170 | + self.install([self.target, second_target]) | ||
| 171 | + | ||
| 172 | + self.assertEqual(self.registry[self.target](_IRValue("cpu")), "upstream") | ||
| 173 | + self.assertEqual( | ||
| 174 | + self.registry[second_target](_IRValue("npu")), "second_npu" | ||
| 175 | + ) | ||
| 176 | + self.upstream_call.assert_called_once() | ||
| 177 | + second_device_call.assert_called_once() | ||
| 178 | + | ||
| 179 | + def test_op_packet_expands_registered_overloads(self): | ||
| 180 | + packet = torch.ops.aten.mm | ||
| 181 | + overload = packet.default | ||
| 182 | + | ||
| 183 | + def packet_upstream(*args, **kwargs): | ||
| 184 | + return "packet_upstream" | ||
| 185 | + | ||
| 186 | + def packet_device(*args, **kwargs): | ||
| 187 | + return "packet_npu" | ||
| 188 | + | ||
| 189 | + def overload_upstream(*args, **kwargs): | ||
| 190 | + return "overload_upstream" | ||
| 191 | + | ||
| 192 | + def overload_device(*args, **kwargs): | ||
| 193 | + return "overload_npu" | ||
| 194 | + | ||
| 195 | + self.baseline.lowerings_copy.update( | ||
| 196 | + { | ||
| 197 | + packet: packet_upstream, | ||
| 198 | + overload: overload_upstream, | ||
| 199 | + } | ||
| 200 | + ) | ||
| 201 | + self.registry.update( | ||
| 202 | + { | ||
| 203 | + packet: packet_device, | ||
| 204 | + overload: overload_device, | ||
| 205 | + } | ||
| 206 | + ) | ||
| 207 | + | ||
| 208 | + lowering_patch.install_device_lowering_dispatch([packet]) | ||
| 209 | + | ||
| 210 | + self.assertTrue( | ||
| 211 | + getattr( | ||
| 212 | + self.registry[packet], | ||
| 213 | + "_torch_npu_device_lowering_dispatch", | ||
| 214 | + False, | ||
| 215 | + ) | ||
| 216 | + ) | ||
| 217 | + self.assertTrue( | ||
| 218 | + getattr( | ||
| 219 | + self.registry[overload], | ||
| 220 | + "_torch_npu_device_lowering_dispatch", | ||
| 221 | + False, | ||
| 222 | + ) | ||
| 223 | + ) | ||
| 224 | + | ||
| 225 | + def test_restore_removes_dispatcher_in_place(self): | ||
| 226 | + registry_id = id(self.registry) | ||
| 227 | + self.install() | ||
| 228 | + | ||
| 229 | + lowering_patch.restore_lowering_baseline() | ||
| 230 | + | ||
| 231 | + self.assertEqual(id(self.registry), registry_id) | ||
| 232 | + self.assertIs(self.registry[self.target], self.upstream) | ||
| 233 | + | ||
| 234 | + | ||
| 235 | +if __name__ == "__main__": | ||
| 236 | + run_tests() | ||
| @@ -0,0 +1,103 @@ | |||
| 1 | +import inspect | ||
| 2 | +from types import SimpleNamespace | ||
| 3 | +from unittest import mock | ||
| 4 | + | ||
| 5 | +from torch._inductor.codegen.cuda_combined_scheduling import CUDACombinedScheduling | ||
| 6 | +from torch._inductor.codegen.triton import TritonScheduling | ||
| 7 | +from torch._inductor.runtime.triton_heuristics import Grid1D | ||
| 8 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 9 | + | ||
| 10 | +from torch_npu._inductor.codegen.npu_combined_scheduling import ( | ||
| 11 | + NPUCombinedScheduling, | ||
| 12 | +) | ||
| 13 | +from torch_npu._inductor.codegen.scheduling import ( | ||
| 14 | + NPUNoLinearTritonScheduling, | ||
| 15 | + NPUTritonScheduling, | ||
| 16 | +) | ||
| 17 | +from torch_npu._inductor.codegen.triton import ( | ||
| 18 | + NPUIndexTritonKernel, | ||
| 19 | + NPUTritonKernel, | ||
| 20 | +) | ||
| 21 | +from torch_npu._inductor.runtime.triton_heuristics import ( | ||
| 22 | + _create_launcher_grid, | ||
| 23 | + _remap_fallback_block_subs, | ||
| 24 | +) | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +class TestSchedulingContract(TestCase): | ||
| 28 | + def test_combined_scheduling_keeps_pytorch_213_protocol(self): | ||
| 29 | + self.assertTrue(issubclass(NPUCombinedScheduling, CUDACombinedScheduling)) | ||
| 30 | + self.assertTrue(issubclass(NPUCombinedScheduling, TritonScheduling)) | ||
| 31 | + signature = inspect.signature( | ||
| 32 | + NPUCombinedScheduling.generate_kernel_code_from_nodes | ||
| 33 | + ) | ||
| 34 | + self.assertIn("hint_override", signature.parameters) | ||
| 35 | + | ||
| 36 | + def test_index_codegen_is_always_first(self): | ||
| 37 | + scheduling = object.__new__(NPUCombinedScheduling) | ||
| 38 | + scheduling._triton_scheduling = mock.Mock() | ||
| 39 | + scheduling._triton_scheduling.codegen_node.return_value = "index" | ||
| 40 | + scheduling._nolinear_triton_scheduling = mock.Mock() | ||
| 41 | + node = mock.Mock() | ||
| 42 | + | ||
| 43 | + self.assertEqual(scheduling.codegen_node(node), "index") | ||
| 44 | + scheduling._nolinear_triton_scheduling.codegen_node.assert_not_called() | ||
| 45 | + | ||
| 46 | + def test_index_failure_regroups_then_falls_back(self): | ||
| 47 | + scheduling = object.__new__(NPUCombinedScheduling) | ||
| 48 | + scheduling._triton_scheduling = mock.Mock() | ||
| 49 | + scheduling._triton_scheduling.codegen_node.side_effect = RuntimeError("index") | ||
| 50 | + scheduling._nolinear_triton_scheduling = mock.Mock() | ||
| 51 | + scheduling._nolinear_triton_scheduling.group_fn.return_value = (64, 1) | ||
| 52 | + scheduling._nolinear_triton_scheduling.codegen_node.return_value = "fallback" | ||
| 53 | + snode = SimpleNamespace(group=("npu", "old"), _sizes=[[64], [1]]) | ||
| 54 | + node = mock.Mock() | ||
| 55 | + node.get_nodes.return_value = [snode] | ||
| 56 | + | ||
| 57 | + self.assertEqual(scheduling.codegen_node(node), "fallback") | ||
| 58 | + self.assertEqual(snode.group, ("npu", (64, 1))) | ||
| 59 | + | ||
| 60 | + def test_scheduling_kernel_types_are_fixed(self): | ||
| 61 | + self.assertIs(NPUTritonScheduling.kernel_type, NPUIndexTritonKernel) | ||
| 62 | + self.assertIs(NPUNoLinearTritonScheduling.kernel_type, NPUTritonKernel) | ||
| 63 | + | ||
| 64 | + def test_only_fallback_configs_are_remapped(self): | ||
| 65 | + untouched = SimpleNamespace(kwargs={"XBLOCK_SUB": 32}) | ||
| 66 | + fallback = SimpleNamespace(kwargs={"XBLOCK_SUB": 32}) | ||
| 67 | + _remap_fallback_block_subs([untouched], {}) | ||
| 68 | + _remap_fallback_block_subs( | ||
| 69 | + [fallback], {"requires_no_linear_block_remap": True} | ||
| 70 | + ) | ||
| 71 | + self.assertEqual(untouched.kwargs, {"XBLOCK_SUB": 32}) | ||
| 72 | + self.assertEqual(fallback.kwargs, {"XBLOCK": 32}) | ||
| 73 | + | ||
| 74 | + def test_standard_grid_uses_upstream_grid_factory(self): | ||
| 75 | + grid = _create_launcher_grid( | ||
| 76 | + {"grid_type": "Grid1D"}, | ||
| 77 | + {"XBLOCK": 32}, | ||
| 78 | + ["xnumel"], | ||
| 79 | + (), | ||
| 80 | + ) | ||
| 81 | + | ||
| 82 | + self.assertIsInstance(grid, Grid1D) | ||
| 83 | + self.assertEqual(grid.eval_slow({"xnumel": 33}), (2, 1, 1)) | ||
| 84 | + | ||
| 85 | + def test_fallback_kernel_marks_block_remap_capability(self): | ||
| 86 | + kernel = object.__new__(NPUTritonKernel) | ||
| 87 | + kernel.range_trees = [ | ||
| 88 | + SimpleNamespace(is_reduction=False, tensor_dim=0, prefix="x") | ||
| 89 | + ] | ||
| 90 | + kernel.inside_reduction = False | ||
| 91 | + | ||
| 92 | + metadata = kernel.add_npu_inductor_meta({}) | ||
| 93 | + | ||
| 94 | + self.assertIs(metadata["requires_no_linear_block_remap"], True) | ||
| 95 | + | ||
| 96 | + def test_index_kernel_does_not_emit_fallback_capability(self): | ||
| 97 | + source = inspect.getsource(NPUIndexTritonKernel.create_inductor_meta) | ||
| 98 | + self.assertNotIn("requires_no_linear_block_remap", source) | ||
| 99 | + self.assertNotIn("inductor_" + "ascend_linear_mode", source) | ||
| 100 | + | ||
| 101 | + | ||
| 102 | +if __name__ == "__main__": | ||
| 103 | + run_tests() | ||
| @@ -632,7 +632,7 @@ class TestPublicBindings(TestCase): | |||
| 632 | "torch_npu._inductor.codegen.npu_kernel_features", | 632 | "torch_npu._inductor.codegen.npu_kernel_features", |
| 633 | "torch_npu._inductor.codegen.scheduling", | 633 | "torch_npu._inductor.codegen.scheduling", |
| 634 | "torch_npu._inductor.codegen.split_tiling", | 634 | "torch_npu._inductor.codegen.split_tiling", |
| 635 | - "torch_npu._inductor.codegen.tile_generator", | 635 | + "torch_npu._inductor.runtime.tile_generator", |
| 636 | "torch_npu._inductor.codegen.triton", | 636 | "torch_npu._inductor.codegen.triton", |
| 637 | "torch_npu._inductor.codegen.triton_utils", | 637 | "torch_npu._inductor.codegen.triton_utils", |
| 638 | "torch_npu._inductor.codegen.cpp_utils", | 638 | "torch_npu._inductor.codegen.cpp_utils", |
| @@ -737,7 +737,7 @@ class TestPublicBindings(TestCase): | |||
| 737 | "torch.ao.quantization.experimental.linear", | 737 | "torch.ao.quantization.experimental.linear", |
| 738 | "torch.ao.quantization.experimental.observer", | 738 | "torch.ao.quantization.experimental.observer", |
| 739 | "torch.ao.quantization.experimental.qconfig", | 739 | "torch.ao.quantization.experimental.qconfig", |
| 740 | - "torch_npu._inductor.fasta_autotune", | 740 | + "torch_npu._inductor.runtime.fasta_autotune", |
| 741 | "torch_npu._inductor.profiler", | 741 | "torch_npu._inductor.profiler", |
| 742 | "torch_npu._inductor.kernel.flex_attention", | 742 | "torch_npu._inductor.kernel.flex_attention", |
| 743 | "torch_npu._inductor.fx_passes.parallel_scheduler_pass", | 743 | "torch_npu._inductor.fx_passes.parallel_scheduler_pass", |
| @@ -7,11 +7,10 @@ from torch._inductor.async_compile import AsyncCompile | |||
| 7 | AsyncCompile.warm_pool() | 7 | AsyncCompile.warm_pool() |
| 8 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD | 8 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD |
| 9 | 9 | ||
| 10 | -import os | 10 | +from torch_npu.utils._dynamo import _dynamo_register_interface_for_device |
| 11 | -from torch_npu.utils._dynamo import _dynamo_register_interface_for_device, patch_SkipFunctionVariable, patch_TensorVariable_call_method # noqa: B950 | ||
| 12 | # all backends need register npu/cpu/mps device_op_overrides | 11 | # all backends need register npu/cpu/mps device_op_overrides |
| 13 | from .graph import patch_codegen_with_cpp_wrapper | 12 | from .graph import patch_codegen_with_cpp_wrapper |
| 14 | -from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, get_current_raw_stream | 13 | +from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu |
| 15 | # All backends need npu/cpu/mps device_op_overrides. | 14 | # All backends need npu/cpu/mps device_op_overrides. |
| 16 | from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system | 15 | from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system |
| 17 | from ._npu_meta_registration import npu_patch_meta | 16 | from ._npu_meta_registration import npu_patch_meta |
| @@ -39,13 +38,12 @@ def _load_ascendc_backend(): | |||
| 39 | 38 | ||
| 40 | def _load_mlir_backend(): | 39 | def _load_mlir_backend(): |
| 41 | _apply_common_patches() | 40 | _apply_common_patches() |
| 42 | - import torch | ||
| 43 | try: | 41 | try: |
| 44 | import torch_mlir | 42 | import torch_mlir |
| 45 | from torch_mlir import ir | 43 | from torch_mlir import ir |
| 46 | except ImportError as e: | 44 | except ImportError as e: |
| 47 | raise ImportError("torch_mlir is not installed, install it first.") from e | 45 | raise ImportError("torch_mlir is not installed, install it first.") from e |
| 48 | - from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin, torch_mlir_patch | 46 | + from .ascend_npu_ir.ascend_npu_ir.npu import torch_mlir_patch |
| 49 | from .lowering_patch import apply_mlir_inductor_patch | 47 | from .lowering_patch import apply_mlir_inductor_patch |
| 50 | from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import ( | 48 | from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import ( |
| 51 | register_mlir_codegen_backend, | 49 | register_mlir_codegen_backend, |
| @@ -57,7 +55,6 @@ def _load_mlir_backend(): | |||
| 57 | def _load_dvm_backend(): | 55 | def _load_dvm_backend(): |
| 58 | _apply_common_patches() | 56 | _apply_common_patches() |
| 59 | import torch | 57 | import torch |
| 60 | - from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin | ||
| 61 | from .lowering_patch import apply_mlir_inductor_patch | 58 | from .lowering_patch import apply_mlir_inductor_patch |
| 62 | from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import ( | 59 | from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import ( |
| 63 | register_mlir_codegen_backend, | 60 | register_mlir_codegen_backend, |
| @@ -67,14 +64,11 @@ def _load_dvm_backend(): | |||
| 67 | from .dvm import mlir_fusion | 64 | from .dvm import mlir_fusion |
| 68 | has_triton = torch.utils._triton.has_triton() | 65 | has_triton = torch.utils._triton.has_triton() |
| 69 | if has_triton: | 66 | if has_triton: |
| 70 | - from .codegen.triton import patch_triton_scheduling | ||
| 71 | from .runtime import patch_triton_heuristics_cached_autotune | 67 | from .runtime import patch_triton_heuristics_cached_autotune |
| 72 | - patch_triton_scheduling() | ||
| 73 | patch_triton_heuristics_cached_autotune() | 68 | patch_triton_heuristics_cached_autotune() |
| 74 | 69 | ||
| 75 | def _load_triton_backend(): | 70 | def _load_triton_backend(): |
| 76 | _apply_common_patches() | 71 | _apply_common_patches() |
| 77 | - import os | ||
| 78 | import torch | 72 | import torch |
| 79 | has_triton = torch.utils._triton.has_triton() | 73 | has_triton = torch.utils._triton.has_triton() |
| 80 | if not has_triton: | 74 | if not has_triton: |
| @@ -84,34 +78,20 @@ def _load_triton_backend(): | |||
| 84 | import logging | 78 | import logging |
| 85 | log = logging.getLogger(__name__) | 79 | log = logging.getLogger(__name__) |
| 86 | 80 | ||
| 87 | - import torch | ||
| 88 | from torch._dynamo.device_interface import get_interface_for_device | 81 | from torch._dynamo.device_interface import get_interface_for_device |
| 89 | from torch._inductor import lowering as inductor_lowering | 82 | from torch._inductor import lowering as inductor_lowering |
| 90 | - from torch._inductor.codegen.common import ( | 83 | + from torch._inductor.codegen.common import register_backend_for_device |
| 91 | - register_backend_for_device, | ||
| 92 | - register_device_op_overrides, | ||
| 93 | - ) | ||
| 94 | from torch.nn.attention import flex_attention | 84 | from torch.nn.attention import flex_attention |
| 95 | 85 | ||
| 96 | - from . import codegen, config as npu_config | 86 | + from . import config as npu_config |
| 97 | from .async_compile import patch_async_compile | 87 | from .async_compile import patch_async_compile |
| 98 | from .codecache import patch_get_cpp_wrapper_header | 88 | from .codecache import patch_get_cpp_wrapper_header |
| 99 | from .codegen._sizevars import patch_simplify | 89 | from .codegen._sizevars import patch_simplify |
| 100 | from .codegen.ir import patch_indexing, patch_loop_body | 90 | from .codegen.ir import patch_indexing, patch_loop_body |
| 101 | - from .codegen.triton import ( | ||
| 102 | - patch_triton_scheduling, | ||
| 103 | - ) | ||
| 104 | - from .config import ( | ||
| 105 | - aggresive_autotune, | ||
| 106 | - log as npulog, | ||
| 107 | - max_precompiled_thread_num, | ||
| 108 | - num_vector_core, | ||
| 109 | - ) | ||
| 110 | from .cpp_builder import ( | 91 | from .cpp_builder import ( |
| 111 | patch_get_cpp_torch_device_options, | 92 | patch_get_cpp_torch_device_options, |
| 112 | patch_get_optimization_cflags, | 93 | patch_get_optimization_cflags, |
| 113 | ) | 94 | ) |
| 114 | - from .codegen.cpp_utils import patch_device_to_aten | ||
| 115 | from .decomposition import _register_triton_decompositions | 95 | from .decomposition import _register_triton_decompositions |
| 116 | from .dependencies import patch_extract_read_writes | 96 | from .dependencies import patch_extract_read_writes |
| 117 | from .fx_passes import patch_pattern_mm_plus_mm | 97 | from .fx_passes import patch_pattern_mm_plus_mm |
| @@ -135,12 +115,12 @@ def _load_triton_backend(): | |||
| 135 | patch_load_cached_autotuning, | 115 | patch_load_cached_autotuning, |
| 136 | patch_triton_heuristics_cached_autotune, | 116 | patch_triton_heuristics_cached_autotune, |
| 137 | ) | 117 | ) |
| 138 | - from .scheduler import patch_scheduler, patch_get_graph_partition_signature | 118 | + from .scheduler import patch_scheduler |
| 139 | from .select_algorithm import patch_algorithm_selector | 119 | from .select_algorithm import patch_algorithm_selector |
| 140 | from .utils import patch_get_first_incompatible_cudagraph_node | 120 | from .utils import patch_get_first_incompatible_cudagraph_node |
| 141 | 121 | ||
| 142 | from .graph import patch_count_bytes | 122 | from .graph import patch_count_bytes |
| 143 | - from .autotune_process import patch_tuning_process, patch_tuning_process_pool | 123 | + from .autotune_process import patch_tuning_process |
| 144 | flex_attention._validate_device = _validate_device | 124 | flex_attention._validate_device = _validate_device |
| 145 | 125 | ||
| 146 | def _inductor_register_backend_for_device(): | 126 | def _inductor_register_backend_for_device(): |
| @@ -154,14 +134,13 @@ def _load_triton_backend(): | |||
| 154 | 134 | ||
| 155 | _inductor_register_backend_for_device() | 135 | _inductor_register_backend_for_device() |
| 156 | 136 | ||
| 157 | - device = get_interface_for_device("npu") | 137 | + get_interface_for_device("npu") |
| 158 | 138 | ||
| 159 | inductor_lowering.make_reduction = make_reduction | 139 | inductor_lowering.make_reduction = make_reduction |
| 160 | 140 | ||
| 161 | patch_get_cpp_wrapper_header() | 141 | patch_get_cpp_wrapper_header() |
| 162 | patch_get_cpp_torch_device_options() | 142 | patch_get_cpp_torch_device_options() |
| 163 | patch_constant_fold_uniform_value() | 143 | patch_constant_fold_uniform_value() |
| 164 | - patch_device_to_aten() | ||
| 165 | 144 | ||
| 166 | if npu_config.dump_fx_graph: | 145 | if npu_config.dump_fx_graph: |
| 167 | from .codegen.ir_fx import _patch_npu_inductor_ir | 146 | from .codegen.ir_fx import _patch_npu_inductor_ir |
| @@ -169,9 +148,11 @@ def _load_triton_backend(): | |||
| 169 | _patch_npu_inductor_ir() | 148 | _patch_npu_inductor_ir() |
| 170 | 149 | ||
| 171 | from .lowering import ( | 150 | from .lowering import ( |
| 151 | + LOWERING_OVERRIDE_OP, | ||
| 172 | _enable_full_lowering_fallback, | 152 | _enable_full_lowering_fallback, |
| 173 | _register_npu_inductor_fallbacks, | 153 | _register_npu_inductor_fallbacks, |
| 174 | ) | 154 | ) |
| 155 | + from .lowering_patch import install_device_lowering_dispatch | ||
| 175 | 156 | ||
| 176 | _register_triton_decompositions() | 157 | _register_triton_decompositions() |
| 177 | 158 | ||
| @@ -185,6 +166,7 @@ def _load_triton_backend(): | |||
| 185 | _register_npu_inductor_grouped_mm() | 166 | _register_npu_inductor_grouped_mm() |
| 186 | 167 | ||
| 187 | _register_npu_inductor_flex_attention() | 168 | _register_npu_inductor_flex_attention() |
| 169 | + install_device_lowering_dispatch(LOWERING_OVERRIDE_OP) | ||
| 188 | 170 | ||
| 189 | patch_pattern_mm_plus_mm() | 171 | patch_pattern_mm_plus_mm() |
| 190 | patch_algorithm_selector() | 172 | patch_algorithm_selector() |
| @@ -194,8 +176,6 @@ def _load_triton_backend(): | |||
| 194 | patch_num_splits() | 176 | patch_num_splits() |
| 195 | patch_loop_body() | 177 | patch_loop_body() |
| 196 | patch_indexing() | 178 | patch_indexing() |
| 197 | - patch_triton_scheduling() | ||
| 198 | - | ||
| 199 | patch_create_device_properties() | 179 | patch_create_device_properties() |
| 200 | patch_load_cached_autotuning() | 180 | patch_load_cached_autotuning() |
| 201 | patch_triton_heuristics_cached_autotune() | 181 | patch_triton_heuristics_cached_autotune() |
| @@ -207,36 +187,11 @@ def _load_triton_backend(): | |||
| 207 | 187 | ||
| 208 | parallel_scheduler() | 188 | parallel_scheduler() |
| 209 | 189 | ||
| 210 | - # register fx_pass should be put behind of _register_triton_decompositions | ||
| 211 | - def _replace_benchmark_all_configs(): | ||
| 212 | - from torch._inductor.runtime.triton_heuristics import CachingAutotuner | ||
| 213 | - | ||
| 214 | - from .runtime.triton_heuristics import ( | ||
| 215 | - _benchmark_all_configs, | ||
| 216 | - benchmark_all_configs, | ||
| 217 | - ) | ||
| 218 | - | ||
| 219 | - CachingAutotuner._benchmark_all_configs = _benchmark_all_configs | ||
| 220 | - CachingAutotuner.benchmark_all_configs = benchmark_all_configs | ||
| 221 | - | ||
| 222 | - def _replace_precompile(): | ||
| 223 | - from .runtime.triton_heuristics import NPUCachingAutotuner, precompile_parallel | ||
| 224 | - | ||
| 225 | - NPUCachingAutotuner.precompile = precompile_parallel | ||
| 226 | - | ||
| 227 | - if aggresive_autotune: | ||
| 228 | - _replace_benchmark_all_configs() | ||
| 229 | - | ||
| 230 | - if max_precompiled_thread_num > 1: | ||
| 231 | - _replace_precompile() | ||
| 232 | - | ||
| 233 | patch_get_first_incompatible_cudagraph_node() | 190 | patch_get_first_incompatible_cudagraph_node() |
| 234 | - patch_get_graph_partition_signature() | ||
| 235 | patch_get_optimization_cflags() | 191 | patch_get_optimization_cflags() |
| 236 | patch_extract_read_writes() | 192 | patch_extract_read_writes() |
| 237 | patch_count_bytes() | 193 | patch_count_bytes() |
| 238 | patch_tuning_process() | 194 | patch_tuning_process() |
| 239 | - patch_tuning_process_pool() | ||
| 240 | 195 | ||
| 241 | def add_additional_op(): | 196 | def add_additional_op(): |
| 242 | from torch._inductor.ops_handler import OpsHandler | 197 | from torch._inductor.ops_handler import OpsHandler |
| @@ -3,14 +3,12 @@ from __future__ import annotations | |||
| 3 | import copy | 3 | import copy |
| 4 | import functools | 4 | import functools |
| 5 | import logging | 5 | import logging |
| 6 | -import os | ||
| 7 | from ctypes import byref, c_size_t, c_void_p | 6 | from ctypes import byref, c_size_t, c_void_p |
| 8 | from typing import (Any, Callable, Iterable, List, | 7 | from typing import (Any, Callable, Iterable, List, |
| 9 | - Optional, Sequence, Union) | 8 | + Optional, Union) |
| 10 | 9 | ||
| 11 | import torch | 10 | import torch |
| 12 | import torch._inductor.async_compile # noqa: F401 required to warm up AsyncCompile pools | 11 | import torch._inductor.async_compile # noqa: F401 required to warm up AsyncCompile pools |
| 13 | -from torch._inductor import config | ||
| 14 | from torch._inductor.autotune_process import ( | 12 | from torch._inductor.autotune_process import ( |
| 15 | BenchmarkRequest, NonzeroWorkspaceNotSupportedError, TensorMeta) | 13 | BenchmarkRequest, NonzeroWorkspaceNotSupportedError, TensorMeta) |
| 16 | from torch._inductor.codecache import DLLWrapper | 14 | from torch._inductor.codecache import DLLWrapper |
| @@ -31,31 +29,6 @@ def patch_tuning_process(): | |||
| 31 | autotune_process.CUDA_VISIBLE_DEVICES = ASCEND_VISIBLE_DEVICES | 29 | autotune_process.CUDA_VISIBLE_DEVICES = ASCEND_VISIBLE_DEVICES |
| 32 | 30 | ||
| 33 | 31 | ||
| 34 | -def patch_tuning_process_pool(): | ||
| 35 | - from torch._inductor.autotune_process import TuningProcessPool | ||
| 36 | - | ||
| 37 | - def get_device_list(self) -> Sequence[Optional[int]]: | ||
| 38 | - """ | ||
| 39 | - Gather the list of devices to be used in the pool. | ||
| 40 | - """ | ||
| 41 | - if not config.autotune_multi_device: | ||
| 42 | - # Don't use multiple devices | ||
| 43 | - return [None] | ||
| 44 | - | ||
| 45 | - count = torch.npu.device_count() | ||
| 46 | - | ||
| 47 | - # If the user specified the visible devices in the env, use those. | ||
| 48 | - if ASCEND_VISIBLE_DEVICES in os.environ: | ||
| 49 | - devices = [int(d) for d in os.environ[ASCEND_VISIBLE_DEVICES].split(",")] | ||
| 50 | - if len(devices) > count: | ||
| 51 | - raise ValueError(f"Specified visible devices exceed the number of total devices: {devices}") | ||
| 52 | - return devices | ||
| 53 | - | ||
| 54 | - return list(range(count)) | ||
| 55 | - | ||
| 56 | - TuningProcessPool.get_device_list = get_device_list | ||
| 57 | - | ||
| 58 | - | ||
| 59 | class NPUDeviceBenchmarkMixin: | 32 | class NPUDeviceBenchmarkMixin: |
| 60 | def do_bench( | 33 | def do_bench( |
| 61 | self, | 34 | self, |
| @@ -13,7 +13,7 @@ from torch._inductor.scheduler import ( | |||
| 13 | ) | 13 | ) |
| 14 | 14 | ||
| 15 | from ..autotune_process import FusedCATLASSBenchmarkRequest | 15 | from ..autotune_process import FusedCATLASSBenchmarkRequest |
| 16 | -from ..config import is_ascend950, log | 16 | +from ..config import log |
| 17 | from .catlass.catlass_scheduling import CATLASSScheduling | 17 | from .catlass.catlass_scheduling import CATLASSScheduling |
| 18 | from .scheduling import NPUNoLinearTritonScheduling, NPUTritonScheduling | 18 | from .scheduling import NPUNoLinearTritonScheduling, NPUTritonScheduling |
| 19 | 19 | ||
| @@ -74,31 +74,20 @@ class NPUCombinedScheduling(CUDACombinedScheduling, TritonScheduling): | |||
| 74 | template_node, epilogue_nodes, prologue_nodes | 74 | template_node, epilogue_nodes, prologue_nodes |
| 75 | ) | 75 | ) |
| 76 | 76 | ||
| 77 | - def node_can_linear(self): | ||
| 78 | - # user config use linear scheduling for ascend. | ||
| 79 | - from ..config import inductor_ascend_linear_mode | ||
| 80 | - | ||
| 81 | - if inductor_ascend_linear_mode != "linear": | ||
| 82 | - return False | ||
| 83 | - return True | ||
| 84 | - | ||
| 85 | def codegen_node(self, node: FusedSchedulerNode | SchedulerNode): | 77 | def codegen_node(self, node: FusedSchedulerNode | SchedulerNode): |
| 86 | - if not is_ascend950: | 78 | + try: |
| 87 | return self._triton_scheduling.codegen_node(node) | 79 | return self._triton_scheduling.codegen_node(node) |
| 88 | - | 80 | + except Exception: |
| 89 | - if self.node_can_linear(): | 81 | + log.debug( |
| 90 | - try: | 82 | + "index codegen for node %s failed, falling back to no-linear codegen", |
| 91 | - return self._triton_scheduling.codegen_node(node) | 83 | + node, |
| 92 | - except Exception: | 84 | + exc_info=True, |
| 93 | - log.debug( | 85 | + ) |
| 94 | - "linear codegen for node %s raise error, fallback to origin codegen", | ||
| 95 | - node, | ||
| 96 | - exc_info=True, | ||
| 97 | - ) | ||
| 98 | - # regroup snode | ||
| 99 | for snode in node.get_nodes(): | 86 | for snode in node.get_nodes(): |
| 100 | - group_fn = self._nolinear_triton_scheduling.group_fn | 87 | + snode.group = ( |
| 101 | - snode.group = (snode.group[0], group_fn(snode._sizes)) | 88 | + snode.group[0], |
| 89 | + self._nolinear_triton_scheduling.group_fn(snode._sizes), | ||
| 90 | + ) | ||
| 102 | return self._nolinear_triton_scheduling.codegen_node(node) | 91 | return self._nolinear_triton_scheduling.codegen_node(node) |
| 103 | 92 | ||
| 104 | def benchmark_codegened_module(self, module): | 93 | def benchmark_codegened_module(self, module): |
| @@ -27,7 +27,7 @@ from torch.utils._ordered_set import OrderedSet | |||
| 27 | from torch._inductor.codegen.simd import CandidateTiling | 27 | from torch._inductor.codegen.simd import CandidateTiling |
| 28 | 28 | ||
| 29 | from .npu_kernel_features import NumelList, NPUKernelFeatures | 29 | from .npu_kernel_features import NumelList, NPUKernelFeatures |
| 30 | -from .triton import NPUIndexTritonKernel, NPUTritonKernel, NPUTritonKernelWithLoop, flatten | 30 | +from .triton import NPUIndexTritonKernel, NPUTritonKernel, flatten |
| 31 | from .triton_combo_kernel import NPUComboKernel | 31 | from .triton_combo_kernel import NPUComboKernel |
| 32 | from .. import config as npu_config | 32 | from .. import config as npu_config |
| 33 | from ..lowering_fx import ( | 33 | from ..lowering_fx import ( |
| @@ -51,12 +51,7 @@ def flatten_groups(nums): | |||
| 51 | return res | 51 | return res |
| 52 | 52 | ||
| 53 | class NPUNoLinearTritonScheduling(TritonScheduling): | 53 | class NPUNoLinearTritonScheduling(TritonScheduling): |
| 54 | - def __init__(self, input_scheduler): | 54 | + kernel_type = NPUTritonKernel |
| 55 | - super().__init__(input_scheduler) | ||
| 56 | - from ..config import inductor_ascend_linear_mode | ||
| 57 | - self.kernel_type = NPUTritonKernel | ||
| 58 | - if inductor_ascend_linear_mode == 'no_linear_loop': | ||
| 59 | - self.kernel_type = NPUTritonKernelWithLoop | ||
| 60 | 55 | ||
| 61 | class NPUTritonScheduling(TritonScheduling): | 56 | class NPUTritonScheduling(TritonScheduling): |
| 62 | kernel_type = NPUIndexTritonKernel | 57 | kernel_type = NPUIndexTritonKernel |
| @@ -39,7 +39,6 @@ from torch._inductor.codegen.triton import ( | |||
| 39 | TritonCSEVariable, | 39 | TritonCSEVariable, |
| 40 | TritonKernel, | 40 | TritonKernel, |
| 41 | TritonKernelOverrides, | 41 | TritonKernelOverrides, |
| 42 | - TritonScheduling, | ||
| 43 | upcast_acc_dtype, | 42 | upcast_acc_dtype, |
| 44 | TMACompatibilityChecker, | 43 | TMACompatibilityChecker, |
| 45 | ) | 44 | ) |
| @@ -433,11 +432,6 @@ def get_allow_dynamic(): | |||
| 433 | return False | 432 | return False |
| 434 | 433 | ||
| 435 | 434 | ||
| 436 | - | ||
| 437 | -def select_index_dtype(node_schedule, numel, reduction_numel): | ||
| 438 | - return "tl.int32" | ||
| 439 | - | ||
| 440 | - | ||
| 441 | def flatten_groups(nums): | 435 | def flatten_groups(nums): |
| 442 | res = [] | 436 | res = [] |
| 443 | for i in nums: | 437 | for i in nums: |
| @@ -449,12 +443,6 @@ def flatten_groups(nums): | |||
| 449 | return res | 443 | return res |
| 450 | 444 | ||
| 451 | 445 | ||
| 452 | -def patch_triton_scheduling(): | ||
| 453 | - # need to enable this to speedup attn_cp_test | ||
| 454 | - # triton scheduling | ||
| 455 | - TritonScheduling.select_index_dtype = select_index_dtype | ||
| 456 | - | ||
| 457 | - | ||
| 458 | class IterationRangesEntryNPUIndex(IterationRangesEntry): | 446 | class IterationRangesEntryNPUIndex(IterationRangesEntry): |
| 459 | """ | 447 | """ |
| 460 | NPU entry index implement | 448 | NPU entry index implement |
| @@ -827,13 +815,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 827 | tiling_axis += [tree.tensor_dim for tree in reduction_range_trees] | 815 | tiling_axis += [tree.tensor_dim for tree in reduction_range_trees] |
| 828 | axis_names += [tree.prefix for tree in reduction_range_trees] | 816 | axis_names += [tree.prefix for tree in reduction_range_trees] |
| 829 | 817 | ||
| 830 | - # NPU tiling inductor meta | 818 | + inductor_meta["requires_no_linear_block_remap"] = True |
| 831 | - from ..config import inductor_ascend_linear_mode | ||
| 832 | - if inductor_ascend_linear_mode == "linear": | ||
| 833 | - # Linear fallback to no_linear_loop | ||
| 834 | - inductor_meta["inductor_ascend_linear_mode"] = "no_linear" | ||
| 835 | - else: | ||
| 836 | - inductor_meta["inductor_ascend_linear_mode"] = inductor_ascend_linear_mode | ||
| 837 | inductor_meta["npu_kernel_type"] = str(NPUKernelType.SIMD_SIMT_MIX) | 819 | inductor_meta["npu_kernel_type"] = str(NPUKernelType.SIMD_SIMT_MIX) |
| 838 | inductor_meta["split_axis"] = split_axis | 820 | inductor_meta["split_axis"] = split_axis |
| 839 | inductor_meta["tiling_axis"] = tiling_axis | 821 | inductor_meta["tiling_axis"] = tiling_axis |
| @@ -862,529 +844,6 @@ class NPUTritonKernel(TritonKernel): | |||
| 862 | return patched_kernel | 844 | return patched_kernel |
| 863 | 845 | ||
| 864 | 846 | ||
| 865 | -class NPUTritonKernelWithLoop(NPUTritonKernel): | ||
| 866 | - """ | ||
| 867 | - NPU triton kernel without linear, but keep loop | ||
| 868 | - """ | ||
| 869 | - | ||
| 870 | - def __init__( | ||
| 871 | - self, | ||
| 872 | - tiling: dict[str, sympy.Expr], | ||
| 873 | - min_elem_per_thread=0, | ||
| 874 | - optimize_mask=True, | ||
| 875 | - fixed_config: FixedTritonConfig | None = None, | ||
| 876 | - **kwargs, | ||
| 877 | - ): | ||
| 878 | - self.loop_header = IndentedBuffer() # reduction loops prefix info | ||
| 879 | - self.body_header = IndentedBuffer() # all index prefix info | ||
| 880 | - self.normal_loop_body = IndentedBuffer() | ||
| 881 | - self.normal_loop_body_rendered = False | ||
| 882 | - super().__init__( | ||
| 883 | - tiling=tiling, | ||
| 884 | - min_elem_per_thread=min_elem_per_thread, | ||
| 885 | - optimize_mask=optimize_mask, | ||
| 886 | - fixed_config=fixed_config, | ||
| 887 | - **kwargs, | ||
| 888 | - ) | ||
| 889 | - for tree in self.range_trees: | ||
| 890 | - tree.indexing_code = IndentedBuffer() | ||
| 891 | - self.iteration_ranges_split_tiling() | ||
| 892 | - | ||
| 893 | - def dense_size_list(self) -> list[str]: | ||
| 894 | - sizes = ["1"] * self.triton_tensor_ndim() | ||
| 895 | - for tree in self.range_trees: | ||
| 896 | - if tree.tensor_dim is None: | ||
| 897 | - continue | ||
| 898 | - | ||
| 899 | - if not tree.is_reduction or self.inside_reduction: | ||
| 900 | - sizes[tree.tensor_dim] = ( | ||
| 901 | - f"{tree.prefix.upper()}BLOCK" | ||
| 902 | - if tree.is_reduction | ||
| 903 | - else f"{tree.prefix.upper()}BLOCK_SUB" | ||
| 904 | - ) | ||
| 905 | - return sizes | ||
| 906 | - | ||
| 907 | - def get_load_buffer(self, indexing): | ||
| 908 | - if indexing.has_indirect() or indexing.has_tmpmask(): | ||
| 909 | - # Masked loads must come after the mask is computed | ||
| 910 | - return self.compute | ||
| 911 | - | ||
| 912 | - return self.loads | ||
| 913 | - | ||
| 914 | - def iteration_ranges_split_tiling(self): | ||
| 915 | - for entry in self.range_trees: | ||
| 916 | - if entry.is_reduction: | ||
| 917 | - continue | ||
| 918 | - entry.is_loop = True | ||
| 919 | - | ||
| 920 | - def iteration_ranges_ranges_code(self, entry: IterationRangesRoot) -> str: | ||
| 921 | - if entry.is_reduction: | ||
| 922 | - return super().iteration_ranges_ranges_code(entry) | ||
| 923 | - size = self.indexing_size_str(entry.tensor_dim) | ||
| 924 | - index_dtype = self.index_dtype | ||
| 925 | - suffix = f".to({index_dtype})" if index_dtype != "tl.int32" else "" | ||
| 926 | - return f"tl.arange(0, {entry.prefix.upper()}BLOCK_SUB){size}{suffix}" | ||
| 927 | - | ||
| 928 | - def iteration_ranges_codegen_header( | ||
| 929 | - self, entry: IterationRangesRoot, code: IndentedBuffer | ||
| 930 | - ) -> None: | ||
| 931 | - if entry.is_reduction: | ||
| 932 | - return super().iteration_ranges_codegen_header(entry, code) | ||
| 933 | - x = entry.prefix | ||
| 934 | - code.writeline(f"{entry.name} = {x}offset + {x}base") | ||
| 935 | - if entry.grid_dim is not None: | ||
| 936 | - # split to different vec core | ||
| 937 | - code.writeline( | ||
| 938 | - f"{x}mask = {entry.name} < min({x}numel, {x}offset_start + {x.upper()}BLOCK)" | ||
| 939 | - ) | ||
| 940 | - else: | ||
| 941 | - code.writeline(f"{x}mask = {entry.name} < {x}numel") | ||
| 942 | - | ||
| 943 | - def codegen_range_tree(self): | ||
| 944 | - for tree in self.range_trees: | ||
| 945 | - if tree.is_reduction: | ||
| 946 | - if not tree.is_loop: | ||
| 947 | - self.iteration_ranges_codegen_header(tree, self.body) | ||
| 948 | - elif self.inside_reduction: | ||
| 949 | - # workaround for this issue: | ||
| 950 | - # https://gist.github.com/jansel/6527126f781559095c5531f98a4235a7 | ||
| 951 | - self.body.writeline( | ||
| 952 | - f"{tree.prefix}base = {self.iteration_ranges_ranges_code(tree)}" | ||
| 953 | - ) | ||
| 954 | - # scalar don't need loop | ||
| 955 | - elif tree.tensor_dim is not None: | ||
| 956 | - self.body.writeline( | ||
| 957 | - f"{tree.prefix}base = {self.iteration_ranges_ranges_code(tree)}" | ||
| 958 | - ) | ||
| 959 | - self.body.writeline( | ||
| 960 | - f"{tree.prefix}offset_start = {self.iteration_ranges_get_pid(tree)} * {tree.prefix.upper()}BLOCK", | ||
| 961 | - ) | ||
| 962 | - | ||
| 963 | - if self.inside_reduction: | ||
| 964 | - if any(tree.is_loop and tree.is_reduction for tree in self.range_trees): | ||
| 965 | - # If the kernel contains loops, compute rbase. | ||
| 966 | - rn_bases = self._get_reduction_symbols( | ||
| 967 | - "base", integer=True, nonnegative=True | ||
| 968 | - ) | ||
| 969 | - rbase = self._flatten_reduction_indices(rn_bases) | ||
| 970 | - self.body.splice(f"rbase = {self.index_to_str(rbase)}") | ||
| 971 | - else: | ||
| 972 | - # For looped reductions, indexing is deferred to the innermost loop. | ||
| 973 | - self.codegen_reduction_indices(self.body) | ||
| 974 | - | ||
| 975 | - self.body_header.splice(self.body) | ||
| 976 | - self.body.clear() | ||
| 977 | - | ||
| 978 | - def codegen_iteration_ranges_entry(self, entry: IterationRangesEntry): | ||
| 979 | - if entry.is_reduction: | ||
| 980 | - return super().codegen_iteration_ranges_entry(entry) | ||
| 981 | - line = f"{entry.name} = {self.kexpr(self.rename_indexing(entry.expr))}" | ||
| 982 | - entry.root.indexing_code.writeline(line) | ||
| 983 | - | ||
| 984 | - def codegen_body(self): | ||
| 985 | - """ | ||
| 986 | - Concat output code from index_code, loads, compute, stores, | ||
| 987 | - suffix into self.body. | ||
| 988 | - | ||
| 989 | - For pointwise kernels, this is called just once at the end. | ||
| 990 | - | ||
| 991 | - For reduction kernels, this generates a loop over the reduction | ||
| 992 | - axis. | ||
| 993 | - """ | ||
| 994 | - from torch._inductor.codegen.triton import TritonSymbols # noqa: F401 | ||
| 995 | - from torch.utils._sympy.functions import CeilDiv | ||
| 996 | - | ||
| 997 | - if not ( | ||
| 998 | - self.indexing_code | ||
| 999 | - or self.loads | ||
| 1000 | - or self.stores | ||
| 1001 | - or self.compute | ||
| 1002 | - or self.post_loop_combine | ||
| 1003 | - or self.post_loop_store | ||
| 1004 | - ): | ||
| 1005 | - return | ||
| 1006 | - | ||
| 1007 | - # loop not scalar entry | ||
| 1008 | - normal_loop_trees = [ | ||
| 1009 | - tree | ||
| 1010 | - for tree in self.range_trees | ||
| 1011 | - if not tree.is_reduction and tree.tensor_dim is not None | ||
| 1012 | - ] | ||
| 1013 | - normal_axis_base = len(normal_loop_trees) | ||
| 1014 | - reduction_loop_trees = [ | ||
| 1015 | - tree for tree in self.range_trees if tree.is_loop and tree.is_reduction | ||
| 1016 | - ] | ||
| 1017 | - body_chunk = IndentedBuffer() | ||
| 1018 | - | ||
| 1019 | - # self.body is the rendered kernel body. codegen_body() can be called | ||
| 1020 | - # multiple times by disable_reduction(), so rebuild it from accumulated | ||
| 1021 | - # loop-inner chunks instead of wrapping the previous render again. | ||
| 1022 | - if self.body and not self.normal_loop_body_rendered: | ||
| 1023 | - self.normal_loop_body.splice(self.body) | ||
| 1024 | - self.body.clear() | ||
| 1025 | - | ||
| 1026 | - self.body.splice(self.body_header) | ||
| 1027 | - if len(normal_loop_trees) > 0: | ||
| 1028 | - # Write the loop headers. | ||
| 1029 | - for level, tree in enumerate(normal_loop_trees): | ||
| 1030 | - with self.body.indent(offset=level): | ||
| 1031 | - prefix = tree.prefix | ||
| 1032 | - loop_start = "0" | ||
| 1033 | - if tree.grid_dim is not None: | ||
| 1034 | - loop_start = f"{prefix}offset_start" | ||
| 1035 | - | ||
| 1036 | - loop_end = f"{loop_start} + {prefix.upper()}BLOCK" | ||
| 1037 | - self.body.writeline( | ||
| 1038 | - f"for {prefix}offset in range({loop_start}, {loop_end}, {prefix.upper()}BLOCK_SUB):" | ||
| 1039 | - ) | ||
| 1040 | - with self.body.indent(offset=level + 1): | ||
| 1041 | - self.iteration_ranges_codegen_header(tree, self.body) | ||
| 1042 | - self.body.splice(tree.indexing_code) | ||
| 1043 | - | ||
| 1044 | - if self.inside_reduction and len(reduction_loop_trees) > 0: | ||
| 1045 | - # Write the loop headers. | ||
| 1046 | - for level, tree in enumerate(reduction_loop_trees): | ||
| 1047 | - with body_chunk.indent(offset=level): | ||
| 1048 | - prefix = tree.prefix | ||
| 1049 | - loop_start = "rsplit_start" if self.cooperative_reduction else "0" | ||
| 1050 | - loop_end = ( | ||
| 1051 | - "rsplit_end" if self.cooperative_reduction else f"{prefix}numel" | ||
| 1052 | - ) | ||
| 1053 | - body_chunk.writeline( | ||
| 1054 | - f"for {prefix}offset in range({loop_start}, {loop_end}, {prefix.upper()}BLOCK):" | ||
| 1055 | - ) | ||
| 1056 | - with body_chunk.indent(offset=level + 1): | ||
| 1057 | - self.iteration_ranges_codegen_header(tree, body_chunk) | ||
| 1058 | - | ||
| 1059 | - # The innermost loop performs the reduction_loop_trees. | ||
| 1060 | - with body_chunk.indent(offset=len(reduction_loop_trees)): | ||
| 1061 | - self.codegen_reduction_indices(body_chunk) | ||
| 1062 | - body_chunk.splice(self.indexing_code) | ||
| 1063 | - body_chunk.splice(self.loads) | ||
| 1064 | - body_chunk.splice(self.compute) | ||
| 1065 | - body_chunk.splice(self.stores) | ||
| 1066 | - | ||
| 1067 | - # Write loop suffixes. | ||
| 1068 | - for level, tree in reversed([*enumerate(reduction_loop_trees)]): | ||
| 1069 | - # persistent reduction doesn't need split loop | ||
| 1070 | - if not tree.is_reduction: | ||
| 1071 | - continue | ||
| 1072 | - with body_chunk.indent(offset=level + 1): | ||
| 1073 | - # Advance pointers at the end of each loop. | ||
| 1074 | - for block_ptr, advancement in self.pointer_advancements[ | ||
| 1075 | - tree.symt | ||
| 1076 | - ].items(): | ||
| 1077 | - # Subtract any advancements made in the previous loop level. | ||
| 1078 | - if level < len(reduction_loop_trees) - 1: | ||
| 1079 | - prev_tree = reduction_loop_trees[level + 1] | ||
| 1080 | - prev_advancement = self.pointer_advancements[ | ||
| 1081 | - prev_tree.symt | ||
| 1082 | - ][block_ptr] | ||
| 1083 | - prev_block = NPUTritonSymbols.get_block_size(prev_tree) | ||
| 1084 | - prev_num_iter = CeilDiv(prev_tree.numel, prev_block) | ||
| 1085 | - advancement = [ | ||
| 1086 | - cur - prev * prev_num_iter | ||
| 1087 | - for cur, prev in zip(advancement, prev_advancement) | ||
| 1088 | - ] | ||
| 1089 | - | ||
| 1090 | - body_chunk.writeline( | ||
| 1091 | - DeferredLine( | ||
| 1092 | - self.block_ptr_to_buffer[block_ptr], | ||
| 1093 | - f"{block_ptr} = tl.advance({block_ptr}, {V.kernel.index_to_str(advancement)})", | ||
| 1094 | - ) | ||
| 1095 | - ) | ||
| 1096 | - | ||
| 1097 | - # Invalidate any cache entries that came from inside the loop. | ||
| 1098 | - self.cse.invalidate(self.outside_loop_vars) | ||
| 1099 | - tree.cache_clear() | ||
| 1100 | - else: | ||
| 1101 | - body_chunk.splice(self.indexing_code) | ||
| 1102 | - body_chunk.splice(self.loads) | ||
| 1103 | - body_chunk.splice(self.compute) | ||
| 1104 | - body_chunk.splice(self.stores) | ||
| 1105 | - body_chunk.splice(self.post_loop_combine) | ||
| 1106 | - if self.cooperative_reduction and ( | ||
| 1107 | - self.post_loop_combine or self.post_loop_store | ||
| 1108 | - ): | ||
| 1109 | - sem_ptr = f"{self.semaphores_name} + tl.program_id(1)" | ||
| 1110 | - body_chunk.splice( | ||
| 1111 | - f""" | ||
| 1112 | - if HAS_RSPLIT: | ||
| 1113 | - triton_helpers.x_grid_barrier({sem_ptr}) | ||
| 1114 | - """, | ||
| 1115 | - strip=True, | ||
| 1116 | - ) | ||
| 1117 | - self.cooperative_reduction_workspace_cache.on_loop_end() | ||
| 1118 | - | ||
| 1119 | - body_chunk.splice(self.post_loop_store) | ||
| 1120 | - | ||
| 1121 | - self.normal_loop_body.splice(body_chunk) | ||
| 1122 | - with self.body.indent(offset=normal_axis_base): | ||
| 1123 | - self.body.splice(self.normal_loop_body) | ||
| 1124 | - self.normal_loop_body_rendered = True | ||
| 1125 | - | ||
| 1126 | - self.loop_header.clear() | ||
| 1127 | - self.indexing_code.clear() | ||
| 1128 | - self.loads.clear() | ||
| 1129 | - self.compute.clear() | ||
| 1130 | - self.stores.clear() | ||
| 1131 | - self.post_loop_combine.clear() | ||
| 1132 | - self.post_loop_store.clear() | ||
| 1133 | - | ||
| 1134 | - def codegen_kernel(self, name=None): | ||
| 1135 | - """ | ||
| 1136 | - codegen triton kernel with loop | ||
| 1137 | - """ | ||
| 1138 | - | ||
| 1139 | - from torch._inductor.codegen.common import ( | ||
| 1140 | - ArgName, | ||
| 1141 | - ConstexprArg, | ||
| 1142 | - IndentedBuffer, | ||
| 1143 | - InplacedBuffer, | ||
| 1144 | - RemovedArg, | ||
| 1145 | - SizeArg, | ||
| 1146 | - WorkspaceArg, | ||
| 1147 | - WorkspaceZeroMode, | ||
| 1148 | - ) | ||
| 1149 | - from torch._inductor.codegen.triton_utils import ( | ||
| 1150 | - config_of, | ||
| 1151 | - equal_1_arg_indices, | ||
| 1152 | - non_constexpr_signature, | ||
| 1153 | - signature_to_meta, | ||
| 1154 | - ) | ||
| 1155 | - from torch._inductor.utils import ( | ||
| 1156 | - Placeholder, | ||
| 1157 | - prefix_is_reduction, | ||
| 1158 | - triton_version_uses_attrs_dict, | ||
| 1159 | - ) | ||
| 1160 | - | ||
| 1161 | - code = IndentedBuffer() | ||
| 1162 | - | ||
| 1163 | - size_hints = {} | ||
| 1164 | - for prefix, numel in self.numels.items(): | ||
| 1165 | - if prefix_is_reduction(prefix) and not self.inside_reduction: | ||
| 1166 | - continue | ||
| 1167 | - | ||
| 1168 | - numel_hint = V.graph.sizevars.optimization_hint(numel) | ||
| 1169 | - if not isinstance(numel_hint, (int, sympy.Integer)): | ||
| 1170 | - # This default heuristic hint was picked carefully: it is | ||
| 1171 | - # large, to ensure that we don't shrink the block size (since | ||
| 1172 | - # if you don't have many elements, it'd be wasteful to pick a | ||
| 1173 | - # large block size). Since we don't know how many elements we | ||
| 1174 | - # might have, we should be OK with some inefficiency to make | ||
| 1175 | - # sure we handle the large case well. 8192 is the largest | ||
| 1176 | - # block size we support, so we pick that. | ||
| 1177 | - # | ||
| 1178 | - # If we have a better hint for unbacked SymInts (e.g., because | ||
| 1179 | - # a user told us, or we are tracking upper bounds) we could | ||
| 1180 | - # use that here. | ||
| 1181 | - size_hint = 8192 | ||
| 1182 | - else: | ||
| 1183 | - size_hint = next_power_of_2(int(numel_hint)) | ||
| 1184 | - size_hints[prefix] = size_hint | ||
| 1185 | - | ||
| 1186 | - if name is None: | ||
| 1187 | - code.splice(self.gen_common_triton_imports()) | ||
| 1188 | - device_type = V.graph.get_current_device_or_throw().type | ||
| 1189 | - if device_type == "cpu": | ||
| 1190 | - code.splice("triton_helpers.set_driver_to_cpu()") | ||
| 1191 | - else: | ||
| 1192 | - code.splice("triton_helpers.set_driver_to_gpu()") | ||
| 1193 | - | ||
| 1194 | - if config.benchmark_kernel: | ||
| 1195 | - code.splice(self.imports_for_benchmark_kernel()) | ||
| 1196 | - | ||
| 1197 | - argdefs, _, signature, _ = self.args.python_argdefs() | ||
| 1198 | - # maps actual expression to SizeArg if it is in sizevars replacements | ||
| 1199 | - for i, arg in enumerate(signature): | ||
| 1200 | - if isinstance(arg, SizeArg): | ||
| 1201 | - # mypy is unhappy about the sympy.Expr | ||
| 1202 | - # type for the key of the dict below | ||
| 1203 | - symbol = cast(sympy.Symbol, arg.expr) | ||
| 1204 | - if symbol in V.graph.sizevars.inv_precomputed_replacements: | ||
| 1205 | - signature[i] = SizeArg( | ||
| 1206 | - arg.name, V.graph.sizevars.inv_precomputed_replacements[symbol] | ||
| 1207 | - ) | ||
| 1208 | - | ||
| 1209 | - mutated_args = OrderedSet[str]() | ||
| 1210 | - for mutation in self.mutations: | ||
| 1211 | - if mutation in self.args.input_buffers: | ||
| 1212 | - mutated_args.add(self.args.input_buffers[mutation]) | ||
| 1213 | - if ( | ||
| 1214 | - mutation in self.args.inplace_buffers | ||
| 1215 | - and mutation not in V.graph.removed_buffers | ||
| 1216 | - and mutation not in self.removed_buffers | ||
| 1217 | - ): | ||
| 1218 | - mutated_args.add( | ||
| 1219 | - cast(InplacedBuffer, self.args.inplace_buffers[mutation]).inner_name | ||
| 1220 | - ) | ||
| 1221 | - if mutation in self.args.output_buffers: | ||
| 1222 | - mutation_arg = self.args.output_buffers[mutation] | ||
| 1223 | - assert not isinstance(mutation_arg, RemovedArg) # noqa: S101 | ||
| 1224 | - mutated_args.add(mutation_arg) | ||
| 1225 | - | ||
| 1226 | - # Note: [Workspace Mutation] | ||
| 1227 | - # workspace arguments are mutated, but are not marked as mutations in self.mutations | ||
| 1228 | - # because their buffers are added during codegen, and aren't tracked during | ||
| 1229 | - # lowering/scheduling. So we add them as mutated_args explicitly below. | ||
| 1230 | - # | ||
| 1231 | - # In the logic below, we only mark the workspaces a mutated if they are marked with | ||
| 1232 | - # zero_fill: that's because, if we don't expect the buffer to be pre-filled with | ||
| 1233 | - # zeros, then, although we still mutate the data, we don't care about those | ||
| 1234 | - # mutations because we don't make any assumptions about the contents of the | ||
| 1235 | - # workspace buffer. Similarly, ZERO_PER_GRAPH requires the kernel to return | ||
| 1236 | - # the buffer back to its original state. | ||
| 1237 | - for argname, arg in zip(argdefs, signature): | ||
| 1238 | - if ( | ||
| 1239 | - isinstance(arg, WorkspaceArg) | ||
| 1240 | - and arg.zero_mode == WorkspaceZeroMode.ZERO_ON_CALL | ||
| 1241 | - ): | ||
| 1242 | - mutated_args.add(argname.name) | ||
| 1243 | - | ||
| 1244 | - mutated_args = sorted(mutated_args) | ||
| 1245 | - | ||
| 1246 | - for tree in self.active_range_trees(): | ||
| 1247 | - sizearg = SizeArg(f"{tree.prefix}numel", tree.numel) | ||
| 1248 | - signature.append(sizearg) | ||
| 1249 | - argdefs.append(ArgName(sizearg.name)) | ||
| 1250 | - # constexpr version causes issues, see | ||
| 1251 | - # https://github.com/pytorch/torchdynamo/pull/1362 | ||
| 1252 | - # triton_meta["constants"][len(argdefs)] = V.graph.sizevars.optimization_hint( | ||
| 1253 | - # tree.numel | ||
| 1254 | - # ) | ||
| 1255 | - # argdefs.append(f"{tree.prefix}numel: tl.constexpr") | ||
| 1256 | - | ||
| 1257 | - def add_constexpr_arg(arg_name): | ||
| 1258 | - # new versions (but not old versions) of Triton need constexprs included in the signature | ||
| 1259 | - if triton_version_uses_attrs_dict(): | ||
| 1260 | - signature.append(ConstexprArg(arg_name)) | ||
| 1261 | - argdefs.append(ArgName(arg_name, is_constexpr=True)) | ||
| 1262 | - | ||
| 1263 | - for tree in self.range_trees: | ||
| 1264 | - if tree.is_reduction and self.persistent_reduction: | ||
| 1265 | - # Rn_BLOCK for persistent_reduction is defined in codegen_static_numels | ||
| 1266 | - continue | ||
| 1267 | - if tree.tensor_dim is None: | ||
| 1268 | - continue | ||
| 1269 | - | ||
| 1270 | - add_constexpr_arg(f"{tree.prefix.upper()}BLOCK") | ||
| 1271 | - # This is used for split axis IterationRange to single vector call | ||
| 1272 | - if not tree.is_reduction: | ||
| 1273 | - add_constexpr_arg(f"{tree.prefix.upper()}BLOCK_SUB") | ||
| 1274 | - | ||
| 1275 | - if self.cooperative_reduction: | ||
| 1276 | - add_constexpr_arg("RSPLIT") | ||
| 1277 | - | ||
| 1278 | - triton_meta_signature = signature_to_meta( | ||
| 1279 | - signature, size_dtype=self.index_dtype, argdefs=argdefs | ||
| 1280 | - ) | ||
| 1281 | - triton_meta: dict[str, Any] = { | ||
| 1282 | - "signature": triton_meta_signature, | ||
| 1283 | - "device": DeviceProperties.create(V.graph.get_current_device_or_throw()), | ||
| 1284 | - "constants": {}, | ||
| 1285 | - } | ||
| 1286 | - | ||
| 1287 | - # Skip memory optimization for forward of the training loop where we expect | ||
| 1288 | - # every new node will increase the peak memory and our greedy approach would | ||
| 1289 | - # introduce a lot of unnecessary cpu copies. | ||
| 1290 | - optimize_mem = V.graph.is_inference or V.graph.is_backward | ||
| 1291 | - | ||
| 1292 | - inductor_meta = { | ||
| 1293 | - # Triton will not accept an OrderedSet for autotune_hints | ||
| 1294 | - "grid_type": self._get_grid_type().__name__, | ||
| 1295 | - "autotune_hints": set(self.autotune_hints), # noqa: set_linter | ||
| 1296 | - "kernel_name": str(Placeholder.DESCRIPTIVE_NAME), | ||
| 1297 | - "mutated_arg_names": mutated_args, | ||
| 1298 | - "optimize_mem": optimize_mem, | ||
| 1299 | - "no_x_dim": self.no_x_dim, | ||
| 1300 | - "num_load": self.num_load, | ||
| 1301 | - "num_reduction": self.num_reduction, | ||
| 1302 | - **self.inductor_meta_common(), | ||
| 1303 | - } | ||
| 1304 | - if self.cooperative_reduction: | ||
| 1305 | - inductor_meta["persistent_reduction"] = self.persistent_reduction | ||
| 1306 | - self.add_npu_inductor_meta(inductor_meta) | ||
| 1307 | - | ||
| 1308 | - num_gb = None | ||
| 1309 | - if config.benchmark_kernel or config.profile_bandwidth: | ||
| 1310 | - num_gb = self.estimate_kernel_num_bytes() / 1e9 | ||
| 1311 | - inductor_meta["kernel_num_gb"] = num_gb | ||
| 1312 | - | ||
| 1313 | - triton_meta["configs"] = [config_of(signature)] | ||
| 1314 | - | ||
| 1315 | - # Triton compiler includes equal_to_1 args into constants even | ||
| 1316 | - # when they are not constexpr. otherwise there may be a segfault | ||
| 1317 | - # during launching the Inductor-compiled Triton kernel. | ||
| 1318 | - # https://github.com/pytorch/pytorch/issues/120478#issuecomment-1962822307 | ||
| 1319 | - # https://github.com/openai/triton/blob/231efe9ed2d200be0f69a07c298e4342b08efe3d/python/triton/runtime/jit.py#L384 | ||
| 1320 | - for arg_num in equal_1_arg_indices(signature): # type: ignore[index] | ||
| 1321 | - triton_meta["constants"][signature[arg_num].name] = 1 # type: ignore[index,union-attr] | ||
| 1322 | - | ||
| 1323 | - self.triton_meta = triton_meta | ||
| 1324 | - | ||
| 1325 | - self.codegen_body() | ||
| 1326 | - | ||
| 1327 | - for helper in self.helper_functions: | ||
| 1328 | - code.writeline("") | ||
| 1329 | - code.splice(helper) | ||
| 1330 | - | ||
| 1331 | - if self.fixed_config: | ||
| 1332 | - heuristics_line = f""" | ||
| 1333 | - | ||
| 1334 | - config={self.fixed_config.config!r}, | ||
| 1335 | - filename=__file__, | ||
| 1336 | - triton_meta={triton_meta!r}, | ||
| 1337 | - inductor_meta={inductor_meta!r} | ||
| 1338 | - ) | ||
| 1339 | - | ||
| 1340 | - """ | ||
| 1341 | - elif self.inside_reduction: | ||
| 1342 | - reduction_hint = self.features.get_reduction_hint() | ||
| 1343 | - heuristics_line = f""" | ||
| 1344 | - | ||
| 1345 | - size_hints={size_hints!r}, | ||
| 1346 | - reduction_hint={reduction_hint}, | ||
| 1347 | - filename=__file__, | ||
| 1348 | - triton_meta={triton_meta!r}, | ||
| 1349 | - inductor_meta={inductor_meta!r} | ||
| 1350 | - ) | ||
| 1351 | - | ||
| 1352 | - """ | ||
| 1353 | - else: | ||
| 1354 | - tile_hint = "" | ||
| 1355 | - if len(size_hints) == 2: | ||
| 1356 | - if ( | ||
| 1357 | - len(non_constexpr_signature(signature)) == 4 | ||
| 1358 | - ): # input, output and 2 args | ||
| 1359 | - tile_hint = "tile_hint=TileHint.SQUARE," | ||
| 1360 | - else: | ||
| 1361 | - tile_hint = "tile_hint=TileHint.DEFAULT," | ||
| 1362 | - heuristics_line = f""" | ||
| 1363 | - | ||
| 1364 | - size_hints={size_hints!r}, {tile_hint} | ||
| 1365 | - filename=__file__, | ||
| 1366 | - triton_meta={triton_meta!r}, | ||
| 1367 | - inductor_meta={inductor_meta!r}, | ||
| 1368 | - min_elem_per_thread={self.min_elem_per_thread} | ||
| 1369 | - ) | ||
| 1370 | - | ||
| 1371 | - """ | ||
| 1372 | - code.splice(heuristics_line) | ||
| 1373 | - code.writeline( | ||
| 1374 | - f"def {name or str(Placeholder.KERNEL_NAME)}({', '.join(x.full_name() for x in argdefs)}):" | ||
| 1375 | - ) | ||
| 1376 | - with code.indent(): | ||
| 1377 | - self.codegen_static_numels(code) | ||
| 1378 | - for old, new in self.args.aliases(): | ||
| 1379 | - code.writeline(f"{old} = {new}") | ||
| 1380 | - code.splice(self.body) | ||
| 1381 | - | ||
| 1382 | - if config.benchmark_kernel: | ||
| 1383 | - code.splice(self.codegen_kernel_benchmark(num_gb)) | ||
| 1384 | - | ||
| 1385 | - return code.getvalue() | ||
| 1386 | - | ||
| 1387 | - | ||
| 1388 | class NPUIndexTritonKernel(TritonKernel): | 847 | class NPUIndexTritonKernel(TritonKernel): |
| 1389 | """ | 848 | """ |
| 1390 | NPU triton kernel with linear and block_sub loop | 849 | NPU triton kernel with linear and block_sub loop |
| @@ -1980,8 +1439,6 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1980 | f"{axis.name.upper()}BLOCK" for axis in self.split_axis | 1439 | f"{axis.name.upper()}BLOCK" for axis in self.split_axis |
| 1981 | ) | 1440 | ) |
| 1982 | 1441 | ||
| 1983 | - from ..config import inductor_ascend_linear_mode | ||
| 1984 | - | ||
| 1985 | inductor_meta = { | 1442 | inductor_meta = { |
| 1986 | "grid_type": self._get_grid_type().__name__, | 1443 | "grid_type": self._get_grid_type().__name__, |
| 1987 | "autotune_hints": set(self.autotune_hints), # noqa: set_linter | 1444 | "autotune_hints": set(self.autotune_hints), # noqa: set_linter |
| @@ -2002,7 +1459,6 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2002 | "traced_graph_hash": "TRACED_GRAPH_HASH", | 1459 | "traced_graph_hash": "TRACED_GRAPH_HASH", |
| 2003 | "traced_graph_dir": "TRACED_GRAPH_DIR", | 1460 | "traced_graph_dir": "TRACED_GRAPH_DIR", |
| 2004 | "are_deterministic_algorithms_enabled": torch.are_deterministic_algorithms_enabled(), | 1461 | "are_deterministic_algorithms_enabled": torch.are_deterministic_algorithms_enabled(), |
| 2005 | - "inductor_ascend_linear_mode": inductor_ascend_linear_mode, | ||
| 2006 | "runtime_block_arg_names": runtime_block_arg_names, | 1462 | "runtime_block_arg_names": runtime_block_arg_names, |
| 2007 | **TritonKernel.inductor_meta_common(), | 1463 | **TritonKernel.inductor_meta_common(), |
| 2008 | } | 1464 | } |
| @@ -284,6 +284,5 @@ if "TORCHNPU_PRECOMPILE_THREADS" in os.environ: | |||
| 284 | max_precompiled_thread_num = int(os.environ["TORCHNPU_PRECOMPILE_THREADS"]) | 284 | max_precompiled_thread_num = int(os.environ["TORCHNPU_PRECOMPILE_THREADS"]) |
| 285 | 285 | ||
| 286 | lowering_axis_count = None | 286 | lowering_axis_count = None |
| 287 | -inductor_ascend_linear_mode = "linear" | ||
| 288 | 287 | ||
| 289 | autotune_continue_on_failure = os.environ.get('TORCHINDUCTOR_NPU_BACKEND') == "default" | 288 | autotune_continue_on_failure = os.environ.get('TORCHINDUCTOR_NPU_BACKEND') == "default" |
| @@ -1,5 +1,19 @@ | |||
| 1 | +import functools | ||
| 2 | + | ||
| 1 | import torch | 3 | import torch |
| 2 | import torch._inductor.fx_passes.joint_graph as joint_graph | 4 | import torch._inductor.fx_passes.joint_graph as joint_graph |
| 5 | +import torch.utils._pytree as pytree | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def _is_npu_graph(gm): | ||
| 9 | + if not isinstance(gm, torch.fx.GraphModule): | ||
| 10 | + return False | ||
| 11 | + | ||
| 12 | + return any( | ||
| 13 | + getattr(getattr(value, "device", None), "type", None) == "npu" | ||
| 14 | + for node in gm.graph.nodes | ||
| 15 | + for value in pytree.tree_leaves(node.meta.get("val")) | ||
| 16 | + ) | ||
| 3 | 17 | ||
| 4 | 18 | ||
| 5 | def patch_constant_fold_uniform_value(): | 19 | def patch_constant_fold_uniform_value(): |
| @@ -7,9 +21,12 @@ def patch_constant_fold_uniform_value(): | |||
| 7 | # Eliminate dead-nodes to remove extra constants generated by torch.tensor. | 21 | # Eliminate dead-nodes to remove extra constants generated by torch.tensor. |
| 8 | src_func = joint_graph.constant_fold_uniform_value | 22 | src_func = joint_graph.constant_fold_uniform_value |
| 9 | 23 | ||
| 24 | + | ||
| 10 | def new_constant_fold_uniform_value(gm): | 25 | def new_constant_fold_uniform_value(gm): |
| 11 | - src_func(gm) | 26 | + is_npu_graph = _is_npu_graph(gm) |
| 12 | - if isinstance(gm, torch.fx.GraphModule): | 27 | + result = src_func(gm) |
| 28 | + if is_npu_graph: | ||
| 13 | gm.graph.eliminate_dead_code() | 29 | gm.graph.eliminate_dead_code() |
| 30 | + return result | ||
| 14 | 31 | ||
| 15 | - joint_graph.constant_fold_uniform_value = new_constant_fold_uniform_value | 32 | + joint_graph.constant_fold_uniform_value = new_constant_fold_uniform_value |
| @@ -1,29 +1,36 @@ | |||
| 1 | -import torch | 1 | +import functools |
| 2 | -from torch._inductor.pattern_matcher import CallFunction, KeywordArg, LoweringPatternEntry | ||
| 3 | -from torch._inductor.fx_passes.post_grad import pass_patterns, is_valid_mm_plus_mm | ||
| 4 | 2 | ||
| 5 | -aten = torch.ops.aten | 3 | +import torch.utils._pytree as pytree |
| 4 | +from torch._inductor.fx_passes import post_grad | ||
| 5 | +from torch._inductor.pattern_matcher import LoweringPatternEntry | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def _is_npu_match(match): | ||
| 9 | + return any( | ||
| 10 | + getattr(getattr(value, "device", None), "type", None) == "npu" | ||
| 11 | + for node in pytree.tree_leaves((match.args, match.kwargs)) | ||
| 12 | + for value in pytree.tree_leaves(getattr(node, "meta", {}).get("val")) | ||
| 13 | + ) | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +def _npu_aware_extra_check(src_check): | ||
| 17 | + | ||
| 18 | + def extra_check(match): | ||
| 19 | + return not _is_npu_match(match) and src_check(match) | ||
| 20 | + | ||
| 21 | + return extra_check | ||
| 6 | 22 | ||
| 7 | 23 | ||
| 8 | def patch_pattern_mm_plus_mm(): | 24 | def patch_pattern_mm_plus_mm(): |
| 9 | - | 25 | + # Keep the shared pattern registered for other devices. NPU does not yet |
| 10 | - def is_mm_plus_mm(entry) -> bool: | 26 | + # support this lowering, so only narrow its applicability predicate. |
| 11 | - if isinstance(entry, LoweringPatternEntry): | 27 | + seen_entries = set() |
| 12 | - handler_name = getattr(entry.handler, '__name__', '') | 28 | + for entries in post_grad.pass_patterns[1].patterns.values(): |
| 13 | - return handler_name == 'mm_plus_mm' | 29 | + for entry in entries: |
| 14 | - return False | 30 | + if ( |
| 15 | - | 31 | + id(entry) not in seen_entries |
| 16 | - pattern = CallFunction( | 32 | + and isinstance(entry, LoweringPatternEntry) |
| 17 | - aten.add, | 33 | + and entry.handler is post_grad.mm_plus_mm |
| 18 | - CallFunction(aten.mm, KeywordArg("mat1"), KeywordArg("mat2")), | 34 | + ): |
| 19 | - CallFunction(aten.mm, KeywordArg("mat3"), KeywordArg("mat4")), | 35 | + seen_entries.add(id(entry)) |
| 20 | - extra_check=is_valid_mm_plus_mm | 36 | + entry.extra_check = _npu_aware_extra_check(entry.extra_check) |
| 21 | - ) | ||
| 22 | - | ||
| 23 | - # currently, torch_npu does not support mm_plus_mm fusion | ||
| 24 | - for fn in pattern.fns: | ||
| 25 | - index = None | ||
| 26 | - for i, pattern_entry in enumerate(pass_patterns[1].patterns[(pattern.op, fn)]): | ||
| 27 | - if is_mm_plus_mm(pattern_entry): | ||
| 28 | - pass_patterns[1].patterns[(pattern.op, fn)].pop(i) | ||
| 29 | - break | ||
| @@ -122,4 +122,7 @@ def _register_npu_inductor_bmm(): | |||
| 122 | log.warning("No choices for GEMM, using ATen backend as fallback") | 122 | log.warning("No choices for GEMM, using ATen backend as fallback") |
| 123 | choices.append(aten_bmm.bind((mat1, mat2), layout)) | 123 | choices.append(aten_bmm.bind((mat1, mat2), layout)) |
| 124 | 124 | ||
| 125 | - return autotune_select_algorithm("bmm", choices, [mat1, mat2], layout) | 125 | + node, _ = autotune_select_algorithm( |
| 126 | + "bmm", choices, [mat1, mat2], layout | ||
| 127 | + ) | ||
| 128 | + return node | ||
| @@ -1586,16 +1586,14 @@ def _register_npu_inductor_flex_attention(): | |||
| 1586 | 6: create_num_blocks_fake_generator(full_kv_indices), | 1586 | 6: create_num_blocks_fake_generator(full_kv_indices), |
| 1587 | 7: create_indices_fake, | 1587 | 7: create_indices_fake, |
| 1588 | } | 1588 | } |
| 1589 | - return ( | 1589 | + selected, _ = autotune_select_algorithm( |
| 1590 | - autotune_select_algorithm( | 1590 | + "flex_attention", |
| 1591 | - "flex_attention", | 1591 | + choices, |
| 1592 | - choices, | 1592 | + inputs_for_autotuning, |
| 1593 | - inputs_for_autotuning, | 1593 | + layout, |
| 1594 | - layout, | 1594 | + input_gen_fns=input_gen_fns, |
| 1595 | - input_gen_fns=input_gen_fns, | ||
| 1596 | - ), | ||
| 1597 | - logsumexp, | ||
| 1598 | ) | 1595 | ) |
| 1596 | + return selected, logsumexp | ||
| 1599 | 1597 | ||
| 1600 | 1598 | ||
| 1601 | def flex_attention_backward(*args, **kwargs): | 1599 | def flex_attention_backward(*args, **kwargs): |
| @@ -1908,7 +1906,7 @@ def _register_npu_inductor_flex_attention(): | |||
| 1908 | 15: create_indices_fake, | 1906 | 15: create_indices_fake, |
| 1909 | } | 1907 | } |
| 1910 | 1908 | ||
| 1911 | - broadcasted_grad_key = autotune_select_algorithm( | 1909 | + broadcasted_grad_key, _ = autotune_select_algorithm( |
| 1912 | "flex_attention_backward", | 1910 | "flex_attention_backward", |
| 1913 | choices, | 1911 | choices, |
| 1914 | inputs_for_autotuning, | 1912 | inputs_for_autotuning, |
| @@ -111,7 +111,10 @@ def _register_npu_inductor_mm(): | |||
| 111 | choices.append(lazy_register_extern_choice(k).bind((mat1, mat2), layout)) | 111 | choices.append(lazy_register_extern_choice(k).bind((mat1, mat2), layout)) |
| 112 | 112 | ||
| 113 | try: | 113 | try: |
| 114 | - return autotune_select_algorithm(name, choices, [mat1, mat2], layout) | 114 | + node, _ = autotune_select_algorithm( |
| 115 | + name, choices, [mat1, mat2], layout | ||
| 116 | + ) | ||
| 117 | + return node | ||
| 115 | except NoValidChoicesError: | 118 | except NoValidChoicesError: |
| 116 | if not inductor_config.autotune_fallback_to_aten: | 119 | if not inductor_config.autotune_fallback_to_aten: |
| 117 | raise | 120 | raise |
| @@ -182,9 +185,10 @@ def _register_npu_inductor_addmm(): | |||
| 182 | if use_aten_gemm_kernels() | 185 | if use_aten_gemm_kernels() |
| 183 | else [] | 186 | else [] |
| 184 | ) | 187 | ) |
| 185 | - return autotune_select_algorithm( | 188 | + node, _ = autotune_select_algorithm( |
| 186 | "addmm", choices, [inp, mat1, mat2], layout | 189 | "addmm", choices, [inp, mat1, mat2], layout |
| 187 | ) | 190 | ) |
| 191 | + return node | ||
| 188 | 192 | ||
| 189 | choices = ( | 193 | choices = ( |
| 190 | [ | 194 | [ |
| @@ -261,9 +265,10 @@ def _register_npu_inductor_addmm(): | |||
| 261 | ) | 265 | ) |
| 262 | 266 | ||
| 263 | try: | 267 | try: |
| 264 | - return autotune_select_algorithm( | 268 | + node, _ = autotune_select_algorithm( |
| 265 | "addmm", choices, [inp_expanded, mat1, mat2], layout | 269 | "addmm", choices, [inp_expanded, mat1, mat2], layout |
| 266 | ) | 270 | ) |
| 271 | + return node | ||
| 267 | except NoValidChoicesError: | 272 | except NoValidChoicesError: |
| 268 | if not inductor_config.autotune_fallback_to_aten: | 273 | if not inductor_config.autotune_fallback_to_aten: |
| 269 | raise | 274 | raise |
| @@ -266,7 +266,7 @@ def _tuned_grouped_mm_common( | |||
| 266 | ), | 266 | ), |
| 267 | } | 267 | } |
| 268 | 268 | ||
| 269 | - tb = autotune_select_algorithm( | 269 | + tb, _ = autotune_select_algorithm( |
| 270 | algorithm_name, choices, input_nodes, layout, input_gen_fns=input_gen_fns | 270 | algorithm_name, choices, input_nodes, layout, input_gen_fns=input_gen_fns |
| 271 | ) | 271 | ) |
| 272 | return [tb] | 272 | return [tb] |
| @@ -9,14 +9,22 @@ | |||
| 9 | from __future__ import annotations | 9 | from __future__ import annotations |
| 10 | 10 | ||
| 11 | import copy | 11 | import copy |
| 12 | +import functools | ||
| 12 | import importlib | 13 | import importlib |
| 14 | +from collections.abc import Iterable | ||
| 13 | from dataclasses import dataclass, field | 15 | from dataclasses import dataclass, field |
| 14 | from typing import Any, Callable, Optional | 16 | from typing import Any, Callable, Optional |
| 15 | 17 | ||
| 18 | +import torch | ||
| 19 | +import torch.utils._pytree as pytree | ||
| 20 | + | ||
| 16 | from .lowering_common import LOWERING_REGISTRY_ATTRS, get_module_functions | 21 | from .lowering_common import LOWERING_REGISTRY_ATTRS, get_module_functions |
| 17 | 22 | ||
| 18 | _BASELINE: Optional["LoweringSnapshot"] = None | 23 | _BASELINE: Optional["LoweringSnapshot"] = None |
| 19 | _INDUCTOR_ATTR_BASELINE = None | 24 | _INDUCTOR_ATTR_BASELINE = None |
| 25 | +_DEVICE_DISPATCH_MARKER = "_torch_npu_device_lowering_dispatch" | ||
| 26 | +_UPSTREAM_HANDLER_ATTR = "_torch_npu_upstream_handler" | ||
| 27 | +_DEVICE_HANDLER_ATTR = "_torch_npu_device_handler" | ||
| 20 | 28 | ||
| 21 | 29 | ||
| 22 | 30 | ||
| @@ -74,6 +82,91 @@ def capture_lowering_baseline() -> LoweringSnapshot: | |||
| 74 | return _BASELINE | 82 | return _BASELINE |
| 75 | 83 | ||
| 76 | 84 | ||
| 85 | +def _expand_lowering_targets(ops: Iterable[Any]) -> list[Any]: | ||
| 86 | + targets = [] | ||
| 87 | + seen = set() | ||
| 88 | + for op in ops: | ||
| 89 | + candidates = [op] | ||
| 90 | + if isinstance(op, torch._ops.OpOverloadPacket): | ||
| 91 | + candidates.extend(op.op_overloads()) | ||
| 92 | + for target in candidates: | ||
| 93 | + if target not in seen: | ||
| 94 | + seen.add(target) | ||
| 95 | + targets.append(target) | ||
| 96 | + return targets | ||
| 97 | + | ||
| 98 | + | ||
| 99 | +def _iter_ir_device_types(value: Any): | ||
| 100 | + for leaf in pytree.tree_leaves(value): | ||
| 101 | + get_device = getattr(leaf, "get_device", None) | ||
| 102 | + if not callable(get_device): | ||
| 103 | + continue | ||
| 104 | + try: | ||
| 105 | + device = get_device() | ||
| 106 | + except NotImplementedError: | ||
| 107 | + continue | ||
| 108 | + device_type = getattr(device, "type", None) | ||
| 109 | + if device_type is not None: | ||
| 110 | + yield device_type | ||
| 111 | + | ||
| 112 | + | ||
| 113 | +def _uses_device_lowering( | ||
| 114 | + args: tuple[Any, ...], | ||
| 115 | + kwargs: dict[str, Any], | ||
| 116 | + device_type: str, | ||
| 117 | +) -> bool: | ||
| 118 | + layout = kwargs.get("layout") | ||
| 119 | + layout_device_type = getattr( | ||
| 120 | + getattr(layout, "device", None), "type", None | ||
| 121 | + ) | ||
| 122 | + if layout_device_type is not None: | ||
| 123 | + return layout_device_type == device_type | ||
| 124 | + return device_type in _iter_ir_device_types((args, kwargs)) | ||
| 125 | + | ||
| 126 | + | ||
| 127 | +def _make_device_lowering_dispatcher( | ||
| 128 | + upstream_handler: Callable[..., Any], | ||
| 129 | + device_handler: Callable[..., Any], | ||
| 130 | + device_type: str, | ||
| 131 | +) -> Callable[..., Any]: | ||
| 132 | + | ||
| 133 | + def dispatcher(*args, **kwargs): | ||
| 134 | + handler = ( | ||
| 135 | + device_handler | ||
| 136 | + if _uses_device_lowering(args, kwargs, device_type) | ||
| 137 | + else upstream_handler | ||
| 138 | + ) | ||
| 139 | + return handler(*args, **kwargs) | ||
| 140 | + | ||
| 141 | + setattr(dispatcher, _DEVICE_DISPATCH_MARKER, True) | ||
| 142 | + setattr(dispatcher, _UPSTREAM_HANDLER_ATTR, upstream_handler) | ||
| 143 | + setattr(dispatcher, _DEVICE_HANDLER_ATTR, device_handler) | ||
| 144 | + return dispatcher | ||
| 145 | + | ||
| 146 | + | ||
| 147 | +def install_device_lowering_dispatch( | ||
| 148 | + ops: Iterable[Any], device_type: str = "npu" | ||
| 149 | +) -> None: | ||
| 150 | + baseline = capture_lowering_baseline() | ||
| 151 | + lowering = _get_inductor_lowering() | ||
| 152 | + | ||
| 153 | + for target in _expand_lowering_targets(ops): | ||
| 154 | + upstream_handler = baseline.lowerings_copy.get(target) | ||
| 155 | + device_handler = lowering.lowerings.get(target) | ||
| 156 | + if ( | ||
| 157 | + upstream_handler is None | ||
| 158 | + or device_handler is None | ||
| 159 | + or device_handler is upstream_handler | ||
| 160 | + or getattr(device_handler, _DEVICE_DISPATCH_MARKER, False) | ||
| 161 | + ): | ||
| 162 | + continue | ||
| 163 | + lowering.lowerings[target] = _make_device_lowering_dispatcher( | ||
| 164 | + upstream_handler, | ||
| 165 | + device_handler, | ||
| 166 | + device_type, | ||
| 167 | + ) | ||
| 168 | + | ||
| 169 | + | ||
| 77 | def restore_lowering_baseline() -> None: | 170 | def restore_lowering_baseline() -> None: |
| 78 | """Reset torch._inductor.lowering to the captured PT baseline.""" | 171 | """Reset torch._inductor.lowering to the captured PT baseline.""" |
| 79 | baseline = capture_lowering_baseline() | 172 | baseline = capture_lowering_baseline() |
| @@ -1,6 +1,7 @@ | |||
| 1 | import os | 1 | import os |
| 2 | import shutil | 2 | import shutil |
| 3 | import subprocess | 3 | import subprocess |
| 4 | +import sys | ||
| 4 | from datetime import datetime, timezone | 5 | from datetime import datetime, timezone |
| 5 | import glob | 6 | import glob |
| 6 | from typing import List, Callable | 7 | from typing import List, Callable |
| @@ -50,7 +51,7 @@ class SimpleProfilingAnalyzer: | |||
| 50 | mindstudio_profiler_output_dir = os.path.join(prof_dir, "mindstudio_profiler_output") | 51 | mindstudio_profiler_output_dir = os.path.join(prof_dir, "mindstudio_profiler_output") |
| 51 | if os.path.exists(mindstudio_profiler_output_dir): | 52 | if os.path.exists(mindstudio_profiler_output_dir): |
| 52 | shutil.rmtree(mindstudio_profiler_output_dir) | 53 | shutil.rmtree(mindstudio_profiler_output_dir) |
| 53 | - export_cmd = ['python', msprof_py_script_path, "export", "summary", "-dir", prof_dir] | 54 | + export_cmd = [sys.executable, msprof_py_script_path, "export", "summary", "-dir", prof_dir] |
| 54 | completed_analysis = subprocess.run(export_cmd, capture_output=True) | 55 | completed_analysis = subprocess.run(export_cmd, capture_output=True) |
| 55 | if completed_analysis.returncode != 0: | 56 | if completed_analysis.returncode != 0: |
| 56 | raise RuntimeError("subprocess return code is not 0.") | 57 | raise RuntimeError("subprocess return code is not 0.") |
| @@ -22,12 +22,12 @@ from torch._inductor.runtime.runtime_utils import next_power_of_2 | |||
| 22 | from torch._inductor import config | 22 | from torch._inductor import config |
| 23 | 23 | ||
| 24 | import torch_npu | 24 | import torch_npu |
| 25 | -from .codegen.tile_generator import TileGenerator | 25 | +from .tile_generator import TileGenerator |
| 26 | -from .config import log | 26 | +from .triton_heuristics import NPUCachingAutotuner |
| 27 | -from .runtime.triton_heuristics import NPUCachingAutotuner | 27 | +from .. import config as npu_config |
| 28 | -from . import config as npu_config | 28 | +from ..codegen.triton_utils import get_byte_per_numel, NPUKernelType |
| 29 | -from .codegen.triton_utils import get_byte_per_numel, NPUKernelType | 29 | +from ..config import log |
| 30 | -from .profiler import simple_trace_handler | 30 | +from ..profiler import simple_trace_handler |
| 31 | 31 | ||
| 32 | 32 | ||
| 33 | def fast_a_log_message(content, tag='autotuner', level='debug'): | 33 | def fast_a_log_message(content, tag='autotuner', level='debug'): |
| @@ -4,8 +4,8 @@ import sys | |||
| 4 | from torch._inductor.runtime.runtime_utils import next_power_of_2 | 4 | from torch._inductor.runtime.runtime_utils import next_power_of_2 |
| 5 | from torch._inductor.runtime.triton_heuristics import Config | 5 | from torch._inductor.runtime.triton_heuristics import Config |
| 6 | 6 | ||
| 7 | -from .triton_utils import get_byte_per_numel, NPUKernelType | ||
| 8 | from .. import config | 7 | from .. import config |
| 8 | +from ..codegen.triton_utils import get_byte_per_numel, NPUKernelType | ||
| 9 | 9 | ||
| 10 | 10 | ||
| 11 | def aligned_numel_32byte(numel, dtype_bytes): | 11 | def aligned_numel_32byte(numel, dtype_bytes): |
| @@ -14,13 +14,12 @@ import csv | |||
| 14 | import uuid | 14 | import uuid |
| 15 | import threading | 15 | import threading |
| 16 | from itertools import count | 16 | from itertools import count |
| 17 | -from typing import Any, Callable, Dict, Literal, Optional, Union, List | 17 | +from typing import Any, Dict, Literal, Optional, Union, List |
| 18 | from contextlib import contextmanager | 18 | from contextlib import contextmanager |
| 19 | -from concurrent.futures import ThreadPoolExecutor, as_completed | 19 | +from concurrent.futures import ThreadPoolExecutor |
| 20 | import torch | 20 | import torch |
| 21 | import triton | 21 | import triton |
| 22 | from torch._dynamo.testing import rand_strided | 22 | from torch._dynamo.testing import rand_strided |
| 23 | -from torch._dynamo.utils import dynamo_timed | ||
| 24 | from torch._inductor import config | 23 | from torch._inductor import config |
| 25 | from torch._inductor.runtime.autotune_cache import AutotuneCache | 24 | from torch._inductor.runtime.autotune_cache import AutotuneCache |
| 26 | from torch._inductor.runtime.benchmarking import TritonBenchmarker | 25 | from torch._inductor.runtime.benchmarking import TritonBenchmarker |
| @@ -32,6 +31,8 @@ from torch._inductor.utils import triton_version_uses_attrs_dict | |||
| 32 | from torch.utils._ordered_set import OrderedSet | 31 | from torch.utils._ordered_set import OrderedSet |
| 33 | from torch._inductor.runtime.triton_heuristics import ( | 32 | from torch._inductor.runtime.triton_heuristics import ( |
| 34 | CachingAutotuner, | 33 | CachingAutotuner, |
| 34 | + CachingAutotunerPlugin, | ||
| 35 | + DEFER, | ||
| 35 | HeuristicType, | 36 | HeuristicType, |
| 36 | unique_configs, | 37 | unique_configs, |
| 37 | hash_configs, | 38 | hash_configs, |
| @@ -82,7 +83,7 @@ import torch_npu | |||
| 82 | 83 | ||
| 83 | from torch_npu._inductor.npu_compare import check_accuracy_triton | 84 | from torch_npu._inductor.npu_compare import check_accuracy_triton |
| 84 | 85 | ||
| 85 | -from ..codegen.tile_generator import TileGenerator | 86 | +from .tile_generator import TileGenerator |
| 86 | from ..codegen.triton_utils import NPUKernelType | 87 | from ..codegen.triton_utils import NPUKernelType |
| 87 | from ..config import log, autotune_continue_on_failure | 88 | from ..config import log, autotune_continue_on_failure |
| 88 | from .. import config as npu_config | 89 | from .. import config as npu_config |
| @@ -236,6 +237,8 @@ class GridExprNpu(GridExpr): | |||
| 236 | grid_type = inductor_meta["grid_type"] | 237 | grid_type = inductor_meta["grid_type"] |
| 237 | 238 | ||
| 238 | grid_cls = globals().get(grid_type) | 239 | grid_cls = globals().get(grid_type) |
| 240 | + if grid_cls is None: | ||
| 241 | + return GridExpr.from_meta(inductor_meta, cfg, mode) | ||
| 239 | if not issubclass(grid_cls, GridNpu): | 242 | if not issubclass(grid_cls, GridNpu): |
| 240 | grid = grid_cls(inductor_meta=inductor_meta, mode=mode) | 243 | grid = grid_cls(inductor_meta=inductor_meta, mode=mode) |
| 241 | if isinstance(cfg, Config): | 244 | if isinstance(cfg, Config): |
| @@ -318,6 +321,23 @@ class GridExprNpu(GridExpr): | |||
| 318 | grid.z_grid = grouped_grid_fn(2) | 321 | grid.z_grid = grouped_grid_fn(2) |
| 319 | return grid | 322 | return grid |
| 320 | 323 | ||
| 324 | + | ||
| 325 | +def _create_launcher_grid( | ||
| 326 | + inductor_meta, | ||
| 327 | + cfg, | ||
| 328 | + numels, | ||
| 329 | + runtime_block_names, | ||
| 330 | +): | ||
| 331 | + if inductor_meta.get("group_enabled", False): | ||
| 332 | + return GridExprNpu.from_grouped_meta_and_numel( | ||
| 333 | + inductor_meta, | ||
| 334 | + cfg, | ||
| 335 | + numels, | ||
| 336 | + runtime_block_names=runtime_block_names, | ||
| 337 | + ) | ||
| 338 | + return GridExprNpu.from_meta_and_set_numel(inductor_meta, cfg, numels) | ||
| 339 | + | ||
| 340 | + | ||
| 321 | class TritonCompileResultNpu(TritonCompileResult): | 341 | class TritonCompileResultNpu(TritonCompileResult): |
| 322 | def make_launcher(self): | 342 | def make_launcher(self): |
| 323 | cfg = self.config | 343 | cfg = self.config |
| @@ -490,22 +510,15 @@ class TritonCompileResultNpu(TritonCompileResult): | |||
| 490 | for arg in fn.arg_names | 510 | for arg in fn.arg_names |
| 491 | if "_numel" in arg | 511 | if "_numel" in arg |
| 492 | ] | 512 | ] |
| 493 | - linear_mode = self.inductor_meta.get('inductor_ascend_linear_mode', 'no_linear') | ||
| 494 | runtime_block_names = tuple( | 513 | runtime_block_names = tuple( |
| 495 | self.inductor_meta.get("runtime_block_arg_names", ()) | 514 | self.inductor_meta.get("runtime_block_arg_names", ()) |
| 496 | ) | 515 | ) |
| 497 | - grid = None | 516 | + grid = _create_launcher_grid( |
| 498 | - if self.inductor_meta.get("group_enabled", False): | 517 | + self.inductor_meta, |
| 499 | - grid = GridExprNpu.from_grouped_meta_and_numel( | 518 | + cfg, |
| 500 | - self.inductor_meta, | 519 | + numels, |
| 501 | - cfg, | 520 | + runtime_block_names, |
| 502 | - numels, | 521 | + ) |
| 503 | - runtime_block_names=runtime_block_names, | ||
| 504 | - ) | ||
| 505 | - elif linear_mode == 'no_linear' and not runtime_block_names: | ||
| 506 | - grid = GridExpr.from_meta(self.inductor_meta, cfg) | ||
| 507 | - else: | ||
| 508 | - grid = GridExprNpu.from_meta_and_set_numel(self.inductor_meta, cfg, numels) | ||
| 509 | # grid.prefix is usually empty, grid.x_grid is something like `-(xnumel//-1024)` | 522 | # grid.prefix is usually empty, grid.x_grid is something like `-(xnumel//-1024)` |
| 510 | lines = [ | 523 | lines = [ |
| 511 | f"def launcher({', '.join(def_args)}, stream):", | 524 | f"def launcher({', '.join(def_args)}, stream):", |
| @@ -542,6 +555,14 @@ class TritonCompileResultNpu(TritonCompileResult): | |||
| 542 | launcher.call_args = call_args | 555 | launcher.call_args = call_args |
| 543 | return launcher | 556 | return launcher |
| 544 | 557 | ||
| 558 | + | ||
| 559 | +class _NPUSkipPrecompilePlugin(CachingAutotunerPlugin): | ||
| 560 | + def pre_compile(self, autotuner): | ||
| 561 | + if getattr(autotuner, "skip_precompile", False): | ||
| 562 | + return None | ||
| 563 | + return DEFER | ||
| 564 | + | ||
| 565 | + | ||
| 545 | class NPUCachingAutotuner(CachingAutotuner): | 566 | class NPUCachingAutotuner(CachingAutotuner): |
| 546 | def __init__( | 567 | def __init__( |
| 547 | self, | 568 | self, |
| @@ -562,6 +583,7 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 562 | super().__init__(fn, triton_meta, configs, save_cache_hook, mutated_arg_names, optimize_mem, heuristic_type, | 583 | super().__init__(fn, triton_meta, configs, save_cache_hook, mutated_arg_names, optimize_mem, heuristic_type, |
| 563 | size_hints, inductor_meta, custom_kernel, filename, reset_to_zero_arg_names, | 584 | size_hints, inductor_meta, custom_kernel, filename, reset_to_zero_arg_names, |
| 564 | autotune_cache_info) | 585 | autotune_cache_info) |
| 586 | + self._plugins.insert(0, _NPUSkipPrecompilePlugin()) | ||
| 565 | 587 | ||
| 566 | self.exceptions = [] | 588 | self.exceptions = [] |
| 567 | self.fn_name = None | 589 | self.fn_name = None |
| @@ -601,6 +623,11 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 601 | def _build_runtime_launch_args(self, args, runtime_blocks: tuple[int, ...]): | 623 | def _build_runtime_launch_args(self, args, runtime_blocks: tuple[int, ...]): |
| 602 | return (*args, *runtime_blocks) | 624 | return (*args, *runtime_blocks) |
| 603 | 625 | ||
| 626 | + def _diagnostic_runtime_blocks(self, candidate, args) -> tuple[int, ...]: | ||
| 627 | + return tuple( | ||
| 628 | + value for _, value in candidate.get("runtime_blocks", ()) | ||
| 629 | + ) | ||
| 630 | + | ||
| 604 | def _first_compiled_candidate_entry(self): | 631 | def _first_compiled_candidate_entry(self): |
| 605 | if not self.compiled_candidate_entries: | 632 | if not self.compiled_candidate_entries: |
| 606 | raise RuntimeError( | 633 | raise RuntimeError( |
| @@ -668,32 +695,24 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 668 | if candidate["variant_id"] in self.variant_launcher_map | 695 | if candidate["variant_id"] in self.variant_launcher_map |
| 669 | ) | 696 | ) |
| 670 | 697 | ||
| 671 | - def precompile( | 698 | + def _build_candidate_plan(self, configs): |
| 672 | - self, | 699 | + return build_candidate_plan(configs, self.runtime_block_arg_names) |
| 673 | - warm_cache_only=False, | 700 | + |
| 674 | - reload_kernel: Optional[Callable[[], CachingAutotuner]] = None, | 701 | + def _prepare_precompile(self): |
| 675 | - static_triton_bundle_key: Optional[str] = None, | ||
| 676 | - ): | ||
| 677 | runtime_args, runtime_kwargs = self._resolve_costmodel_runtime_inputs() | 702 | runtime_args, runtime_kwargs = self._resolve_costmodel_runtime_inputs() |
| 678 | self._apply_costmodel_to_configs(*runtime_args, **runtime_kwargs) | 703 | self._apply_costmodel_to_configs(*runtime_args, **runtime_kwargs) |
| 679 | - if self.candidate_plan is None: | 704 | + self.candidate_plan = self._build_candidate_plan(self.configs) |
| 680 | - self.candidate_plan = build_candidate_plan( | 705 | + |
| 681 | - self.configs, self.runtime_block_arg_names | 706 | + def _precompile_worker(self): |
| 682 | - ) | 707 | + if getattr(self, "skip_precompile", False): |
| 683 | - if warm_cache_only: | ||
| 684 | - self.kernel_name = self.get_fn_name() | ||
| 685 | - self._precompile_worker() | ||
| 686 | return | 708 | return |
| 687 | - with self.lock: | 709 | + if self.compile_results: |
| 688 | - # Helper function for reloading a kernel generated in a worker | 710 | + return self._precompile_worker_serial() |
| 689 | - # in the parent class. Normally we don't need to reload the kernel | 711 | + self._prepare_precompile() |
| 690 | - # in the parent process, but in certain cases (coordesc tuning, dynamic_scale_rblock), | 712 | + if npu_config.max_precompiled_thread_num > 1: |
| 691 | - # we need to actually run compilation on the parent process | 713 | + self._precompile_worker_parallel() |
| 692 | - if reload_kernel is not None: | 714 | + else: |
| 693 | - self._reload_kernel = reload_kernel | 715 | + self._precompile_worker_serial() |
| 694 | - self._precompile_worker() | ||
| 695 | - self._make_launchers() | ||
| 696 | - self._refresh_variant_launchers() | ||
| 697 | 716 | ||
| 698 | def _make_ttir_module_from_cfg(self, cfg): | 717 | def _make_ttir_module_from_cfg(self, cfg): |
| 699 | """Compile one config to TTIR module only (no backend lowering/launch).""" | 718 | """Compile one config to TTIR module only (no backend lowering/launch).""" |
| @@ -1041,27 +1060,70 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1041 | self.configs = ranked_cfgs[:selected_count] | 1060 | self.configs = ranked_cfgs[:selected_count] |
| 1042 | self._costmodel_fallback_configs = ranked_cfgs[selected_count:] or None | 1061 | self._costmodel_fallback_configs = ranked_cfgs[selected_count:] or None |
| 1043 | 1062 | ||
| 1044 | - def _precompile_configs(self, configs): | 1063 | + def _precompile_configs_once(self, configs): |
| 1064 | + compile_results = [None] * len(configs) | ||
| 1065 | + compile_exceptions = [None] * len(configs) | ||
| 1066 | + compile_exception_stacks = [""] * len(configs) | ||
| 1067 | + for index, candidate_config in enumerate(configs): | ||
| 1068 | + try: | ||
| 1069 | + compile_results[index] = self._precompile_config(candidate_config) | ||
| 1070 | + except Exception as exc: | ||
| 1071 | + import traceback | ||
| 1072 | + compile_exception_stacks[index] = traceback.format_exc() | ||
| 1073 | + compile_exceptions[index] = exc | ||
| 1074 | + return compile_results, compile_exceptions, compile_exception_stacks | ||
| 1075 | + | ||
| 1076 | + def _configs_with_vf_fusion(self, configs): | ||
| 1077 | + retry_configs = [] | ||
| 1078 | + for candidate_config in configs: | ||
| 1079 | + retry_config = copy.deepcopy(candidate_config) | ||
| 1080 | + retry_config.kwargs["enable_vf_fusion"] = True | ||
| 1081 | + retry_configs.append(retry_config) | ||
| 1082 | + return retry_configs | ||
| 1083 | + | ||
| 1084 | + def _precompile_configs_with_vf_retry(self, configs, compile_once): | ||
| 1045 | if not configs: | 1085 | if not configs: |
| 1046 | raise NoTritonConfigsError("No triton configs are available") | 1086 | raise NoTritonConfigsError("No triton configs are available") |
| 1047 | 1087 | ||
| 1048 | - compile_results = [] | ||
| 1049 | - exc = None | ||
| 1050 | - exc_stack = "" | ||
| 1051 | compile_start_time = time.perf_counter() | 1088 | compile_start_time = time.perf_counter() |
| 1052 | - for c in configs: | 1089 | + compile_results, exceptions, exception_stacks = compile_once(configs) |
| 1053 | - try: | 1090 | + successful_results = [ |
| 1054 | - compile_results.append(self._precompile_config(c)) | 1091 | + result for result in compile_results if result is not None |
| 1055 | - except Exception as e: | 1092 | + ] |
| 1056 | - import traceback | 1093 | + if not successful_results: |
| 1057 | - exc_stack = traceback.format_exc() | 1094 | + retry_results, retry_exceptions, retry_exception_stacks = compile_once( |
| 1058 | - exc = e | 1095 | + self._configs_with_vf_fusion(configs) |
| 1059 | - if len(compile_results) == 0: | ||
| 1060 | - raise NoTritonConfigsError( | ||
| 1061 | - f"No valid triton configs. {type(exc).__name__}: {exc} \nStack trace:{exc_stack}" | ||
| 1062 | ) | 1096 | ) |
| 1063 | - log.info("kernel: %s compile cost time: %ss", self.get_fn_name(), time.perf_counter() - compile_start_time) | 1097 | + successful_results = [ |
| 1064 | - return compile_results | 1098 | + result for result in retry_results if result is not None |
| 1099 | + ] | ||
| 1100 | + exceptions = retry_exceptions + exceptions | ||
| 1101 | + exception_stacks = retry_exception_stacks + exception_stacks | ||
| 1102 | + | ||
| 1103 | + if not successful_results: | ||
| 1104 | + exc, exc_stack = next( | ||
| 1105 | + ( | ||
| 1106 | + (exc, stack) | ||
| 1107 | + for exc, stack in zip(exceptions, exception_stacks) | ||
| 1108 | + if exc is not None | ||
| 1109 | + ), | ||
| 1110 | + (None, ""), | ||
| 1111 | + ) | ||
| 1112 | + raise NoTritonConfigsError( | ||
| 1113 | + f"No valid triton configs for kernel {self.get_fn_name()}. " | ||
| 1114 | + f"{type(exc).__name__}: {exc} \nStack trace:{exc_stack}" | ||
| 1115 | + ) | ||
| 1116 | + log.debug( | ||
| 1117 | + "kernel: %s config compile elapsed time: %ss", | ||
| 1118 | + self.get_fn_name(), | ||
| 1119 | + time.perf_counter() - compile_start_time, | ||
| 1120 | + ) | ||
| 1121 | + return successful_results | ||
| 1122 | + | ||
| 1123 | + def _precompile_configs(self, configs): | ||
| 1124 | + return self._precompile_configs_with_vf_retry( | ||
| 1125 | + configs, self._precompile_configs_once | ||
| 1126 | + ) | ||
| 1065 | 1127 | ||
| 1066 | def _precompile_with_costmodel_fallback(self, compile_fn): | 1128 | def _precompile_with_costmodel_fallback(self, compile_fn): |
| 1067 | primary_configs = self.configs | 1129 | primary_configs = self.configs |
| @@ -1103,7 +1165,7 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1103 | f"Primary error: {primary_exc}. Fallback error: {fallback_exc}" | 1165 | f"Primary error: {primary_exc}. Fallback error: {fallback_exc}" |
| 1104 | ) from fallback_exc | 1166 | ) from fallback_exc |
| 1105 | 1167 | ||
| 1106 | - def _precompile_worker(self): | 1168 | + def _precompile_worker_with(self, compile_fn): |
| 1107 | if self.compile_results: | 1169 | if self.compile_results: |
| 1108 | for result in self.compile_results: | 1170 | for result in self.compile_results: |
| 1109 | TritonBundler.put( | 1171 | TritonBundler.put( |
| @@ -1114,9 +1176,12 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1114 | if self.launchers: | 1176 | if self.launchers: |
| 1115 | raise AssertionError("Before _precompile_worker, launchers must bt empty") | 1177 | raise AssertionError("Before _precompile_worker, launchers must bt empty") |
| 1116 | 1178 | ||
| 1117 | - self._precompile_with_costmodel_fallback(self._precompile_configs) | 1179 | + self._precompile_with_costmodel_fallback(compile_fn) |
| 1118 | self._costmodel_fallback_configs = None | 1180 | self._costmodel_fallback_configs = None |
| 1119 | 1181 | ||
| 1182 | + def _precompile_worker_serial(self): | ||
| 1183 | + self._precompile_worker_with(self._precompile_configs) | ||
| 1184 | + | ||
| 1120 | def parse_triton_ascend_options(self, tiling_kwargs, options): | 1185 | def parse_triton_ascend_options(self, tiling_kwargs, options): |
| 1121 | from triton.backends.ascend.compiler import NPUOptions | 1186 | from triton.backends.ascend.compiler import NPUOptions |
| 1122 | for k in NPUOptions.__dataclass_fields__.keys(): | 1187 | for k in NPUOptions.__dataclass_fields__.keys(): |
| @@ -1192,15 +1257,15 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1192 | } | 1257 | } |
| 1193 | 1258 | ||
| 1194 | options = self.parse_triton_ascend_options(cfg_kwargs, options) | 1259 | options = self.parse_triton_ascend_options(cfg_kwargs, options) |
| 1195 | - if self.inductor_meta.get("enable_auto_blockify", False): | 1260 | + if ( |
| 1261 | + bool(self.inductor_meta.get("enable_auto_blockify", False)) | ||
| 1262 | + or self.inductor_meta.get("requires_no_linear_block_remap") is True | ||
| 1263 | + ): | ||
| 1196 | options["enable_auto_blockify"] = True | 1264 | options["enable_auto_blockify"] = True |
| 1197 | # pure simt stack overflow check | 1265 | # pure simt stack overflow check |
| 1198 | if compile_meta['compile_mode'] == NPUKernelType.SIMT_ONLY.compile_mode(): | 1266 | if compile_meta['compile_mode'] == NPUKernelType.SIMT_ONLY.compile_mode(): |
| 1199 | options['simt_stack_limit'] = npu_config.simt_default_warp_stacksize | 1267 | options['simt_stack_limit'] = npu_config.simt_default_warp_stacksize |
| 1200 | 1268 | ||
| 1201 | - if self.inductor_meta.get("inductor_ascend_linear_mode", "no_linear") == "no_linear": | ||
| 1202 | - options['enable_auto_blockify'] = True | ||
| 1203 | - | ||
| 1204 | compile_kwargs = { | 1269 | compile_kwargs = { |
| 1205 | "target": target, | 1270 | "target": target, |
| 1206 | "options": options, | 1271 | "options": options, |
| @@ -1268,6 +1333,7 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1268 | raise RuntimeError(f"No valid triton configs. {type(exc).__name__}: {exc}\n" | 1333 | raise RuntimeError(f"No valid triton configs. {type(exc).__name__}: {exc}\n" |
| 1269 | f"Stack trace: {exc_stack}") | 1334 | f"Stack trace: {exc_stack}") |
| 1270 | self.launchers = launchers | 1335 | self.launchers = launchers |
| 1336 | + self._refresh_variant_launchers() | ||
| 1271 | 1337 | ||
| 1272 | def save_gpu_kernel(self, stream, launcher): | 1338 | def save_gpu_kernel(self, stream, launcher): |
| 1273 | self.save_npu_kernel(stream, launcher) | 1339 | self.save_npu_kernel(stream, launcher) |
| @@ -1321,85 +1387,41 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1321 | 1387 | ||
| 1322 | self.cuda_kernel_saved = True | 1388 | self.cuda_kernel_saved = True |
| 1323 | 1389 | ||
| 1324 | - def _precompile_configs_parallel(self, configs): | 1390 | + def _precompile_configs_parallel_once(self, configs): |
| 1325 | - if not configs: | 1391 | + compile_results = [None] * len(configs) |
| 1326 | - raise NoTritonConfigsError("No triton configs are available") | 1392 | + compile_exceptions = [None] * len(configs) |
| 1393 | + compile_exception_stacks = [""] * len(configs) | ||
| 1327 | 1394 | ||
| 1328 | - config_len = len(configs) | 1395 | + def worker(index, config): |
| 1329 | - compile_exc_results = [None for _ in range(config_len)] | ||
| 1330 | - compile_exc_stack_results = ["" for _ in range(config_len)] | ||
| 1331 | - | ||
| 1332 | - def worker(i, kernel_config): | ||
| 1333 | try: | 1396 | try: |
| 1334 | - return self._precompile_config(kernel_config) | 1397 | + return self._precompile_config(config) |
| 1335 | - except Exception as e: | 1398 | + except Exception as exc: |
| 1336 | import traceback | 1399 | import traceback |
| 1337 | - compile_exc_stack_results[i] = traceback.format_exc() | 1400 | + compile_exception_stacks[index] = traceback.format_exc() |
| 1338 | - compile_exc_results[i] = e | 1401 | + compile_exceptions[index] = exc |
| 1339 | return None | 1402 | return None |
| 1340 | 1403 | ||
| 1341 | - tasks = [] | 1404 | + tasks = [ |
| 1342 | - for i, c in enumerate(configs): | 1405 | + compile_thread_pool.submit(worker, index, config) |
| 1343 | - task_handler = compile_thread_pool.submit(worker, i, c) | 1406 | + for index, config in enumerate(configs) |
| 1344 | - tasks.append(task_handler) | 1407 | + ] |
| 1345 | 1408 | ||
| 1346 | from torch._dynamo.device_interface import DeviceGuard | 1409 | from torch._dynamo.device_interface import DeviceGuard |
| 1347 | device_interface = self.get_device_interface() | 1410 | device_interface = self.get_device_interface() |
| 1348 | - # load binary to the correct device | ||
| 1349 | - compile_results = [] | ||
| 1350 | with DeviceGuard(device_interface, self.triton_meta["device"]): | 1411 | with DeviceGuard(device_interface, self.triton_meta["device"]): |
| 1351 | - # need to initialize context | ||
| 1352 | device_interface.synchronize(device_interface.current_device()) | 1412 | device_interface.synchronize(device_interface.current_device()) |
| 1353 | - for future in as_completed(tasks): | 1413 | + for index, future in enumerate(tasks): |
| 1354 | - compiled_kernel = future.result() | 1414 | + compile_results[index] = future.result() |
| 1355 | - if compiled_kernel is None: | ||
| 1356 | - continue | ||
| 1357 | - compile_results.append(compiled_kernel) | ||
| 1358 | 1415 | ||
| 1359 | - # first try but return no valid configs | 1416 | + return compile_results, compile_exceptions, compile_exception_stacks |
| 1360 | - # so we try tuning more options | ||
| 1361 | - if len(compile_results) == 0: | ||
| 1362 | - # set up new configs | ||
| 1363 | - for i in range(len(configs)): | ||
| 1364 | - # in future, adjust more options | ||
| 1365 | - configs[i].kwargs["enable_vf_fusion"] = True | ||
| 1366 | - # start compilation tasks | ||
| 1367 | - tasks = [] | ||
| 1368 | - for i, c in enumerate(configs): | ||
| 1369 | - task_handler = compile_thread_pool.submit(worker, i, c) | ||
| 1370 | - tasks.append(task_handler) | ||
| 1371 | - # collect compiled results | ||
| 1372 | - with DeviceGuard(device_interface, self.triton_meta["device"]): | ||
| 1373 | - # need to initialize context | ||
| 1374 | - device_interface.synchronize(device_interface.current_device()) | ||
| 1375 | - for future in as_completed(tasks): | ||
| 1376 | - compiled_kernel = future.result() | ||
| 1377 | - if compiled_kernel is None: | ||
| 1378 | - continue | ||
| 1379 | - compile_results.append(compiled_kernel) | ||
| 1380 | 1417 | ||
| 1381 | - if len(compile_results) == 0: | 1418 | + def _precompile_configs_parallel(self, configs): |
| 1382 | - raise NoTritonConfigsError( | 1419 | + return self._precompile_configs_with_vf_retry( |
| 1383 | - f"No valid triton configs for kernel {self.get_fn_name()}. " | 1420 | + configs, self._precompile_configs_parallel_once |
| 1384 | - f"{type(compile_exc_results[0]).__name__}: {compile_exc_results[0]} " | 1421 | + ) |
| 1385 | - f"\nStack trace:{compile_exc_stack_results[0]}" | ||
| 1386 | - ) | ||
| 1387 | - return compile_results | ||
| 1388 | 1422 | ||
| 1389 | def _precompile_worker_parallel(self): | 1423 | def _precompile_worker_parallel(self): |
| 1390 | - if self.compile_results: | 1424 | + self._precompile_worker_with(self._precompile_configs_parallel) |
| 1391 | - for result in self.compile_results: | ||
| 1392 | - TritonBundler.put( | ||
| 1393 | - triton_hash_to_path_key(result.kernel.hash), | ||
| 1394 | - self.triton_meta.get("device", 0), | ||
| 1395 | - ) | ||
| 1396 | - return | ||
| 1397 | - | ||
| 1398 | - if self.launchers: | ||
| 1399 | - raise AssertionError("Before _precompile_worker, launchers must bt empty") | ||
| 1400 | - | ||
| 1401 | - self._precompile_with_costmodel_fallback(self._precompile_configs_parallel) | ||
| 1402 | - self._costmodel_fallback_configs = None | ||
| 1403 | 1425 | ||
| 1404 | # bench method is called by torch, grid can not be modified | 1426 | # bench method is called by torch, grid can not be modified |
| 1405 | def bench(self, launcher, *args, with_profiler=False, runtime_blocks=None, **kwargs): | 1427 | def bench(self, launcher, *args, with_profiler=False, runtime_blocks=None, **kwargs): |
| @@ -1411,9 +1433,16 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1411 | return float("inf") | 1433 | return float("inf") |
| 1412 | 1434 | ||
| 1413 | if runtime_blocks is None: | 1435 | if runtime_blocks is None: |
| 1414 | - runtime_blocks = ( | 1436 | + if launcher is self.best_launcher: |
| 1415 | - self.best_runtime_blocks if launcher is self.best_launcher else () | 1437 | + runtime_blocks = self.best_runtime_blocks |
| 1416 | - ) | 1438 | + else: |
| 1439 | + runtime_blocks = () | ||
| 1440 | + for entry in self.compiled_candidate_entries: | ||
| 1441 | + if entry["launcher"] is launcher: | ||
| 1442 | + runtime_blocks = self._diagnostic_runtime_blocks( | ||
| 1443 | + entry["candidate"], args | ||
| 1444 | + ) | ||
| 1445 | + break | ||
| 1417 | launch_args = self._build_runtime_launch_args(args, tuple(runtime_blocks)) | 1446 | launch_args = self._build_runtime_launch_args(args, tuple(runtime_blocks)) |
| 1418 | return self._bench_with_launch_args( | 1447 | return self._bench_with_launch_args( |
| 1419 | launcher, | 1448 | launcher, |
| @@ -2032,6 +2061,9 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner): | |||
| 2032 | self._grouped_runtime_args_snapshot = () | 2061 | self._grouped_runtime_args_snapshot = () |
| 2033 | self._grouped_variant_launchers_initialized = False | 2062 | self._grouped_variant_launchers_initialized = False |
| 2034 | 2063 | ||
| 2064 | + def _build_candidate_plan(self, configs): | ||
| 2065 | + return self.candidate_plan | ||
| 2066 | + | ||
| 2035 | def _set_group_best_candidate(self, group_id, candidate, launcher): | 2067 | def _set_group_best_candidate(self, group_id, candidate, launcher): |
| 2036 | self.best_candidate_map[group_id] = candidate | 2068 | self.best_candidate_map[group_id] = candidate |
| 2037 | self.best_launcher_map[group_id] = launcher | 2069 | self.best_launcher_map[group_id] = launcher |
| @@ -2039,17 +2071,18 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner): | |||
| 2039 | def _precompile_variants(self): | 2071 | def _precompile_variants(self): |
| 2040 | if self._grouped_variant_launchers_initialized: | 2072 | if self._grouped_variant_launchers_initialized: |
| 2041 | return | 2073 | return |
| 2042 | - if len(self.launchers) == 0: | 2074 | + with self.lock: |
| 2043 | - start_time = time.time_ns() | 2075 | + if self._grouped_variant_launchers_initialized: |
| 2044 | - self.kernel_name = self.get_fn_name() | 2076 | + return |
| 2045 | - if getattr(self, "_precompile_worker_parallel", None) is not None: | 2077 | + if len(self.launchers) == 0: |
| 2046 | - self._precompile_worker_parallel() | 2078 | + start_time = time.time_ns() |
| 2047 | - else: | 2079 | + self.kernel_name = self.get_fn_name() |
| 2048 | self._precompile_worker() | 2080 | self._precompile_worker() |
| 2049 | - self._make_launchers() | 2081 | + self._make_launchers() |
| 2050 | - self.precompile_time_taken_ns = time.time_ns() - start_time | 2082 | + self.precompile_time_taken_ns = time.time_ns() - start_time |
| 2051 | - self._refresh_variant_launchers() | 2083 | + else: |
| 2052 | - self._grouped_variant_launchers_initialized = True | 2084 | + self._refresh_variant_launchers() |
| 2085 | + self._grouped_variant_launchers_initialized = True | ||
| 2053 | 2086 | ||
| 2054 | def _runtime_feature_inputs(self, args) -> tuple[int, ...]: | 2087 | def _runtime_feature_inputs(self, args) -> tuple[int, ...]: |
| 2055 | if "feature_arg_indices" not in self.candidate_plan: | 2088 | if "feature_arg_indices" not in self.candidate_plan: |
| @@ -2138,6 +2171,9 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner): | |||
| 2138 | ) | 2171 | ) |
| 2139 | return tuple(resolved_blocks[name] for name in runtime_block_names) | 2172 | return tuple(resolved_blocks[name] for name in runtime_block_names) |
| 2140 | 2173 | ||
| 2174 | + def _diagnostic_runtime_blocks(self, candidate, args) -> tuple[int, ...]: | ||
| 2175 | + return self._materialize_runtime_blocks(candidate, args) | ||
| 2176 | + | ||
| 2141 | def _benchmark_feature_inputs_for_group(self, group_id: int) -> tuple[int, ...]: | 2177 | def _benchmark_feature_inputs_for_group(self, group_id: int) -> tuple[int, ...]: |
| 2142 | if "benchmark_feature_inputs_by_group" not in self.candidate_plan: | 2178 | if "benchmark_feature_inputs_by_group" not in self.candidate_plan: |
| 2143 | raise RuntimeError( | 2179 | raise RuntimeError( |
| @@ -2406,10 +2442,10 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner): | |||
| 2406 | ) | 2442 | ) |
| 2407 | 2443 | ||
| 2408 | def ensure_grouped_autotune_ready(self, *args, **kwargs): | 2444 | def ensure_grouped_autotune_ready(self, *args, **kwargs): |
| 2445 | + self._precompile_variants() | ||
| 2409 | with self.lock: | 2446 | with self.lock: |
| 2410 | if self._all_reachable_groups_tuned(): | 2447 | if self._all_reachable_groups_tuned(): |
| 2411 | return | 2448 | return |
| 2412 | - self._precompile_variants() | ||
| 2413 | self._autotune_all_groups(*args, **kwargs) | 2449 | self._autotune_all_groups(*args, **kwargs) |
| 2414 | 2450 | ||
| 2415 | def run( | 2451 | def run( |
| @@ -3084,23 +3120,7 @@ def _triton_config_npu_index_legacy( | |||
| 3084 | cfg.kwargs["split_axis"] = tuple(split_axis) | 3120 | cfg.kwargs["split_axis"] = tuple(split_axis) |
| 3085 | cfg.kwargs["split_blocks"] = tuple(split_blocks) | 3121 | cfg.kwargs["split_blocks"] = tuple(split_blocks) |
| 3086 | 3122 | ||
| 3087 | - inductor_ascend_linear_mode = inductor_meta.get( | 3123 | + _remap_fallback_block_subs(configs, inductor_meta) |
| 3088 | - "inductor_ascend_linear_mode", "no_linear" | ||
| 3089 | - ) | ||
| 3090 | - if inductor_ascend_linear_mode == "no_linear": | ||
| 3091 | - for tiling_cfg in configs: | ||
| 3092 | - tiling_kwargs = copy.deepcopy(tiling_cfg.kwargs) | ||
| 3093 | - for tiling, tling_value in tiling_kwargs.items(): | ||
| 3094 | - if isinstance(tiling, str) and tiling.endswith("SUB"): | ||
| 3095 | - tiling_cfg.kwargs[tiling.rstrip("_SUB")] = tling_value | ||
| 3096 | - tiling_cfg.kwargs.pop(tiling) | ||
| 3097 | - elif inductor_ascend_linear_mode == "no_linear_loop": | ||
| 3098 | - for tiling_cfg in configs: | ||
| 3099 | - tiling_kwargs = copy.deepcopy(tiling_cfg.kwargs) | ||
| 3100 | - for tiling, tling_value in tiling_kwargs.items(): | ||
| 3101 | - if isinstance(tiling, str) and tiling.endswith( | ||
| 3102 | - "SUB") and tiling.startswith("R"): | ||
| 3103 | - tiling_cfg.kwargs[tiling.rstrip("_SUB")] = tling_value | ||
| 3104 | 3124 | ||
| 3105 | set_reduction_runtime_blocks_to_numel(configs, split_axis, axis_names, size_hints, | 3125 | set_reduction_runtime_blocks_to_numel(configs, split_axis, axis_names, size_hints, |
| 3106 | inductor_meta.get("runtime_block_arg_names", ())) | 3126 | inductor_meta.get("runtime_block_arg_names", ())) |
| @@ -3113,6 +3133,16 @@ def _triton_config_npu_index_legacy( | |||
| 3113 | configs = brutal_prune_tiling_configs_if_fast_run(configs, inductor_meta) | 3133 | configs = brutal_prune_tiling_configs_if_fast_run(configs, inductor_meta) |
| 3114 | return configs | 3134 | return configs |
| 3115 | 3135 | ||
| 3136 | +def _remap_fallback_block_subs(configs, inductor_meta): | ||
| 3137 | + if inductor_meta.get("requires_no_linear_block_remap") is not True: | ||
| 3138 | + return | ||
| 3139 | + for tiling_cfg in configs: | ||
| 3140 | + for name, value in list(tiling_cfg.kwargs.items()): | ||
| 3141 | + if isinstance(name, str) and name.endswith("_SUB"): | ||
| 3142 | + tiling_cfg.kwargs[name.removesuffix("_SUB")] = value | ||
| 3143 | + tiling_cfg.kwargs.pop(name) | ||
| 3144 | + | ||
| 3145 | + | ||
| 3116 | def strip_runtime_blocks_from_cfg( | 3146 | def strip_runtime_blocks_from_cfg( |
| 3117 | cfg: Config, | 3147 | cfg: Config, |
| 3118 | runtime_block_arg_names: tuple[str, ...], | 3148 | runtime_block_arg_names: tuple[str, ...], |
| @@ -3233,10 +3263,6 @@ def foreach(size_hints, triton_meta, num_warps, filename=None, inductor_meta=Non | |||
| 3233 | filename=filename, | 3263 | filename=filename, |
| 3234 | ) | 3264 | ) |
| 3235 | 3265 | ||
| 3236 | -def benchmark_all_configs(self, *args, **kwargs): | ||
| 3237 | - with dynamo_timed("benchmark_all_configs"): | ||
| 3238 | - return self._benchmark_all_configs(*args, **kwargs) | ||
| 3239 | - | ||
| 3240 | def _measure_prerun_ms(kernel_call_fn): | 3266 | def _measure_prerun_ms(kernel_call_fn): |
| 3241 | start_event = torch.npu.Event(enable_timing=True) | 3267 | start_event = torch.npu.Event(enable_timing=True) |
| 3242 | end_event = torch.npu.Event(enable_timing=True) | 3268 | end_event = torch.npu.Event(enable_timing=True) |
| @@ -3268,47 +3294,3 @@ def _select_prerun_top_candidates( | |||
| 3268 | total_ms += cost_ms | 3294 | total_ms += cost_ms |
| 3269 | 3295 | ||
| 3270 | return selected | 3296 | return selected |
| 3271 | - | ||
| 3272 | -def _benchmark_all_configs(self, *args, **kwargs): | ||
| 3273 | - if getattr(self, "candidate_plan", None) is None: | ||
| 3274 | - self.candidate_plan = build_candidate_plan( | ||
| 3275 | - self.configs, getattr(self, "runtime_block_arg_names", ()) | ||
| 3276 | - ) | ||
| 3277 | - return self._benchmark_candidate_entries(*args, **kwargs) | ||
| 3278 | - | ||
| 3279 | -def precompile_parallel( | ||
| 3280 | - self, | ||
| 3281 | - warm_cache_only=False, | ||
| 3282 | - reload_kernel: Optional[Callable[[], CachingAutotuner]] = None, | ||
| 3283 | - static_triton_bundle_key: Optional[str] = None, | ||
| 3284 | -): | ||
| 3285 | - if reload_kernel is not None: | ||
| 3286 | - self._reload_kernel = reload_kernel | ||
| 3287 | - start_time = time.perf_counter() | ||
| 3288 | - if hasattr(self, "skip_precompile"): | ||
| 3289 | - if self.skip_precompile: | ||
| 3290 | - return | ||
| 3291 | - | ||
| 3292 | - runtime_args, runtime_kwargs = self._resolve_costmodel_runtime_inputs() | ||
| 3293 | - self._apply_costmodel_to_configs(*runtime_args, **runtime_kwargs) | ||
| 3294 | - | ||
| 3295 | - if warm_cache_only: | ||
| 3296 | - self.kernel_name = self.get_fn_name() | ||
| 3297 | - self._precompile_worker_parallel() | ||
| 3298 | - log.info("kernel: %s precompile elapsed time: %ss", self.get_fn_name(), time.perf_counter() - start_time) | ||
| 3299 | - return | ||
| 3300 | - | ||
| 3301 | - if self.compile_results: | ||
| 3302 | - for result in self.compile_results: | ||
| 3303 | - TritonBundler.put( | ||
| 3304 | - triton_hash_to_path_key(result.kernel.hash), | ||
| 3305 | - self.triton_meta.get("device", 0), | ||
| 3306 | - ) | ||
| 3307 | - self._make_launchers() | ||
| 3308 | - self._refresh_variant_launchers() | ||
| 3309 | - return | ||
| 3310 | - | ||
| 3311 | - self._precompile_worker_parallel() | ||
| 3312 | - self._make_launchers() | ||
| 3313 | - self._refresh_variant_launchers() | ||
| 3314 | - log.info("kernel: %s precompile elapsed time: %ss", self.get_fn_name(), time.perf_counter() - start_time) | ||
| @@ -414,7 +414,14 @@ def patch_algorithm_selector() -> None: | |||
| 414 | specific to NPU hardware. | 414 | specific to NPU hardware. |
| 415 | """ | 415 | """ |
| 416 | 416 | ||
| 417 | - def __call__( | 417 | + from torch._inductor.select_algorithm import AlgorithmSelectorCache |
| 418 | + | ||
| 419 | + original_call = AlgorithmSelectorCache.__call__ | ||
| 420 | + original_make_benchmark_fn = AlgorithmSelectorCache.__dict__[ | ||
| 421 | + "make_benchmark_fn" | ||
| 422 | + ].__func__ | ||
| 423 | + | ||
| 424 | + def npu_call( | ||
| 418 | self, | 425 | self, |
| 419 | name: str, | 426 | name: str, |
| 420 | choices: List[ChoiceCaller], | 427 | choices: List[ChoiceCaller], |
| @@ -454,7 +461,7 @@ def patch_algorithm_selector() -> None: | |||
| 454 | if len(choices) == 1: | 461 | if len(choices) == 1: |
| 455 | if not isinstance(choices[0], CATLASSTemplateCaller): | 462 | if not isinstance(choices[0], CATLASSTemplateCaller): |
| 456 | # CATLASSTemplateCaller still needs to go through autotuning process to retrieve workspace size. | 463 | # CATLASSTemplateCaller still needs to go through autotuning process to retrieve workspace size. |
| 457 | - return choices[0].output_node() | 464 | + return choices[0].output_node(), choices[0] |
| 458 | 465 | ||
| 459 | 466 | ||
| 460 | def make_benchmark_fn(): | 467 | def make_benchmark_fn(): |
| @@ -674,27 +681,29 @@ def patch_algorithm_selector() -> None: | |||
| 674 | if isinstance(c, TritonTemplateCaller): | 681 | if isinstance(c, TritonTemplateCaller): |
| 675 | allowed_prologue_inps |= c.allowed_prologue_inps | 682 | allowed_prologue_inps |= c.allowed_prologue_inps |
| 676 | 683 | ||
| 677 | - return torch._inductor.ir.TensorBox.create( | 684 | + return ( |
| 678 | - torch._inductor.ir.MultiTemplateBuffer( | 685 | + torch._inductor.ir.TensorBox.create( |
| 679 | - layout, | 686 | + torch._inductor.ir.MultiTemplateBuffer( |
| 680 | - input_nodes, | 687 | + layout, |
| 681 | - get_timings, | 688 | + input_nodes, |
| 682 | - choices, | 689 | + get_timings, |
| 683 | - allowed_prologue_inps, | 690 | + choices, |
| 684 | - ) | 691 | + allowed_prologue_inps, |
| 692 | + ) | ||
| 693 | + ), | ||
| 694 | + None, | ||
| 685 | ) | 695 | ) |
| 686 | 696 | ||
| 687 | timings = do_autotuning(precompile_fn) | 697 | timings = do_autotuning(precompile_fn) |
| 688 | if timings == {} or choices[0] not in timings: | 698 | if timings == {} or choices[0] not in timings: |
| 689 | - return choices[0].output_node() | 699 | + return choices[0].output_node(), choices[0] |
| 690 | 700 | ||
| 691 | selected_key = builtins.min(timings, key=timings.__getitem__) | 701 | selected_key = builtins.min(timings, key=timings.__getitem__) |
| 692 | - selected_choice = selected_key.output_node() | 702 | + selected_node = selected_key.output_node() |
| 693 | - log.debug("selected choice: %s", str(selected_choice)) | 703 | + log.debug("selected choice: %s", str(selected_node)) |
| 694 | - return selected_choice | 704 | + return selected_node, selected_key |
| 695 | 705 | ||
| 696 | - @classmethod | 706 | + def npu_make_benchmark_fn( |
| 697 | - def make_benchmark_fn( | ||
| 698 | cls, | 707 | cls, |
| 699 | choices: List[ChoiceCaller], | 708 | choices: List[ChoiceCaller], |
| 700 | input_nodes: list[ir.IRNode], | 709 | input_nodes: list[ir.IRNode], |
| @@ -978,7 +987,78 @@ def patch_algorithm_selector() -> None: | |||
| 978 | 987 | ||
| 979 | return benchmark | 988 | return benchmark |
| 980 | 989 | ||
| 981 | - from torch._inductor.select_algorithm import AlgorithmSelectorCache | 990 | + @functools.wraps(original_call, assigned=(), updated=()) |
| 991 | + def __call__( | ||
| 992 | + self, | ||
| 993 | + name, | ||
| 994 | + choices, | ||
| 995 | + input_nodes, | ||
| 996 | + layout, | ||
| 997 | + input_gen_fns=None, | ||
| 998 | + precompilation_timeout_seconds=60 * 60, | ||
| 999 | + return_multi_template=False, | ||
| 1000 | + best_config_future=None, | ||
| 1001 | + is_collective=False, | ||
| 1002 | + min_speedup_threshold=1.0, | ||
| 1003 | + benchmark_with_cudagraphs=False, | ||
| 1004 | + ): | ||
| 1005 | + args = ( | ||
| 1006 | + self, | ||
| 1007 | + name, | ||
| 1008 | + choices, | ||
| 1009 | + input_nodes, | ||
| 1010 | + layout, | ||
| 1011 | + input_gen_fns, | ||
| 1012 | + precompilation_timeout_seconds, | ||
| 1013 | + return_multi_template, | ||
| 1014 | + best_config_future, | ||
| 1015 | + is_collective, | ||
| 1016 | + min_speedup_threshold, | ||
| 1017 | + benchmark_with_cudagraphs, | ||
| 1018 | + ) | ||
| 1019 | + if layout.device.type != "npu": | ||
| 1020 | + return original_call(*args) | ||
| 1021 | + return npu_call( | ||
| 1022 | + self, | ||
| 1023 | + name, | ||
| 1024 | + choices, | ||
| 1025 | + input_nodes, | ||
| 1026 | + layout, | ||
| 1027 | + input_gen_fns, | ||
| 1028 | + precompilation_timeout_seconds, | ||
| 1029 | + return_multi_template, | ||
| 1030 | + ) | ||
| 1031 | + | ||
| 1032 | + | ||
| 1033 | + | ||
| 1034 | + original_make_benchmark_fn, assigned=(), updated=() | ||
| 1035 | + ) | ||
| 1036 | + def make_benchmark_fn( | ||
| 1037 | + cls, | ||
| 1038 | + choices, | ||
| 1039 | + input_nodes, | ||
| 1040 | + layout, | ||
| 1041 | + input_gen_fns, | ||
| 1042 | + hint_override=None, | ||
| 1043 | + is_collective=False, | ||
| 1044 | + ): | ||
| 1045 | + if layout.device.type != "npu": | ||
| 1046 | + return original_make_benchmark_fn( | ||
| 1047 | + cls, | ||
| 1048 | + choices, | ||
| 1049 | + input_nodes, | ||
| 1050 | + layout, | ||
| 1051 | + input_gen_fns, | ||
| 1052 | + hint_override, | ||
| 1053 | + is_collective, | ||
| 1054 | + ) | ||
| 1055 | + return npu_make_benchmark_fn( | ||
| 1056 | + cls, | ||
| 1057 | + choices, | ||
| 1058 | + input_nodes, | ||
| 1059 | + layout, | ||
| 1060 | + input_gen_fns, | ||
| 1061 | + ) | ||
| 982 | 1062 | ||
| 983 | AlgorithmSelectorCache.__call__ = __call__ | 1063 | AlgorithmSelectorCache.__call__ = __call__ |
| 984 | AlgorithmSelectorCache.make_benchmark_fn = make_benchmark_fn | 1064 | AlgorithmSelectorCache.make_benchmark_fn = make_benchmark_fn |
🟡 Medium Priority 在 然而， 但如果外部代码（如用户自定义 lowering）直接调用 建议：在 PR 描述或 CHANGELOG 中明确标注   不准确？ | |||
| @@ -14,9 +14,11 @@ def get_current_raw_stream(device): | |||
| 14 | 14 | ||
| 15 | 15 | ||
| 16 | def patch_is_gpu(): | 16 | def patch_is_gpu(): |
| 17 | - from torch._inductor.utils import GPU_TYPES | 17 | + from torch._inductor.utils import GPU_TYPES, get_gpu_type |
| 18 | 18 | ||
| 19 | - GPU_TYPES.append("npu") | 19 | + if "npu" not in GPU_TYPES: |
| 20 | + GPU_TYPES.append("npu") | ||
| 21 | + get_gpu_type.cache_clear() | ||
| 20 | 22 | ||
| 21 | 23 | ||
| 22 | def patch_has_triton(): | 24 | def patch_has_triton(): |
🟠 High Priority
变更后的
test_gemm_lowerings_are_device_dispatched直接检查torch._inductor.lowering.lowerings[target]上是否存在_torch_npu_device_lowering_dispatch属性。该属性仅在torch_npu/_inductor/__init__.py的_load_triton_backend()调用install_device_lowering_dispatch()时设置。然而，test_cpu_compile.py仅通过testutils.py间接import torch_npu，torch_npu/__init__.py并不导入_inductor子包，因此 lowering dispatch 从未安装。test_gemm_lowerings_are_device_dispatched中的self.assertTrue(getattr(handler, "_torch_npu_device_lowering_dispatch", False))将始终失败。对比同批次新增的
test_inductor_graph_partition.py（第 46 行显式import torch_npu._inductor # noqa: F401），此处明显遗漏。test_cpu_gemm_compile_after_npu_registration虽然可能因torch.compile的延迟加载而侥幸通过，但其测试意图（"在 NPU 注册后 CPU GEMM 仍可编译"）也无法真正验证。建议：在测试文件顶部（import 区域）添加
import torch_npu._inductor # noqa: F401，确保 lowering dispatch 在测试运行前已安装。