已合并
fix(inductor): keep autotune from nesting a profiler session #44117
dezheng889创建于 28 天前
fix(inductor): keep autotune from nesting a profiler session #44117
已合并
共 4 个文件变更+317-2
| @@ -0,0 +1,153 @@ | |||
| 1 | +"""Autotune must not open a profiler inside one the caller already opened. | ||
| 2 | + | ||
| 3 | +The torch_npu profiler backend is a process-wide singleton: a session started | ||
| 4 | +inside another one finalizes the shared trace when it exits and clears | ||
| 5 | +ProfPathCreator, so the outer session is torn down early and reports "Incorrect | ||
| 6 | +schedule: Stop profiler while current state is RECORD". With aggresive_autotune | ||
| 7 | +on, both the batch benchmark and the bandwidth benchmark would open one, so each | ||
| 8 | +has to probe the global state and fall back to the event timer instead. | ||
| 9 | +""" | ||
| 10 | + | ||
| 11 | +from types import SimpleNamespace | ||
| 12 | +from unittest.mock import patch | ||
| 13 | + | ||
| 14 | +from torch.testing._internal.common_utils import ( | ||
| 15 | + instantiate_parametrized_tests, | ||
| 16 | + parametrize, | ||
| 17 | + run_tests, | ||
| 18 | + TestCase, | ||
| 19 | +) | ||
| 20 | + | ||
| 21 | +import torch_npu | ||
| 22 | +import torch_npu._inductor | ||
| 23 | +from torch_npu._inductor.runtime import triton_heuristics | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +class TestAutotuneProfilerSession(TestCase): | ||
| 27 | + | ||
| 28 | + def _autotuner(profile_bandwidth=True): | ||
| 29 | + """An autotuner carrying just enough state to reach the benchmark call.""" | ||
| 30 | + autotuner = object.__new__(triton_heuristics.NPUCachingAutotuner) | ||
| 31 | + autotuner.inductor_meta = { | ||
| 32 | + "profile_bandwidth_with_do_bench_using_profiling": profile_bandwidth, | ||
| 33 | + } | ||
| 34 | + autotuner.get_device_interface = lambda: SimpleNamespace( | ||
| 35 | + current_device=lambda: 0, | ||
| 36 | + get_raw_stream=lambda _device: None, | ||
| 37 | + ) | ||
| 38 | + autotuner.clone_args = lambda *args, **kwargs: ((), {}) | ||
| 39 | + autotuner.reset_to_zero_args = lambda *args, **kwargs: None | ||
| 40 | + return autotuner | ||
| 41 | + | ||
| 42 | + | ||
| 43 | + def test_probe_follows_the_global_profiler_state(self, is_prof_inited): | ||
| 44 | + from torch_npu.profiler._profiler_path_creator import ProfPathCreator | ||
| 45 | + | ||
| 46 | + with patch.object(ProfPathCreator(), "is_prof_inited", is_prof_inited): | ||
| 47 | + self.assertEqual( | ||
| 48 | + triton_heuristics.npu_profiler_session_active(), | ||
| 49 | + is_prof_inited, | ||
| 50 | + ) | ||
| 51 | + | ||
| 52 | + def test_probe_never_breaks_a_kernel_launch(self): | ||
| 53 | + # The probe runs on the launch path, so a profiler backend it cannot | ||
| 54 | + # reach has to read as "no session" rather than raise. | ||
| 55 | + import torch_npu.profiler._profiler_path_creator as path_creator | ||
| 56 | + | ||
| 57 | + with patch.object( | ||
| 58 | + path_creator, | ||
| 59 | + "ProfPathCreator", | ||
| 60 | + side_effect=RuntimeError("no profiler backend"), | ||
| 61 | + ): | ||
| 62 | + self.assertFalse(triton_heuristics.npu_profiler_session_active()) | ||
| 63 | + | ||
| 64 | + | ||
| 65 | + def test_batch_benchmark_is_skipped_while_the_caller_profiles( | ||
| 66 | + self, session_active | ||
| 67 | + ): | ||
| 68 | + autotuner = object.__new__(triton_heuristics.NPUCachingAutotuner) | ||
| 69 | + kernel_funcs = [lambda: None] | ||
| 70 | + | ||
| 71 | + with ( | ||
| 72 | + patch.object(triton_heuristics.npu_config, "aggresive_autotune", True), | ||
| 73 | + patch.object( | ||
| 74 | + triton_heuristics, | ||
| 75 | + "npu_profiler_session_active", | ||
| 76 | + return_value=session_active, | ||
| 77 | + ), | ||
| 78 | + patch.object( | ||
| 79 | + triton_heuristics, | ||
| 80 | + "mspti_batch_benchmark", | ||
| 81 | + return_value=[1.0], | ||
| 82 | + ) as batch_benchmark, | ||
| 83 | + ): | ||
| 84 | + timings = autotuner._benchmark_kernel_funcs_batch(kernel_funcs, "grouped") | ||
| 85 | + | ||
| 86 | + if session_active: | ||
| 87 | + # None sends the caller back to the per-config timer. | ||
| 88 | + self.assertIsNone(timings) | ||
| 89 | + batch_benchmark.assert_not_called() | ||
| 90 | + else: | ||
| 91 | + self.assertEqual(timings, (1.0,)) | ||
| 92 | + batch_benchmark.assert_called_once() | ||
| 93 | + | ||
| 94 | + | ||
| 95 | + def test_bandwidth_benchmark_falls_back_to_the_timer(self, session_active): | ||
| 96 | + autotuner = self._autotuner() | ||
| 97 | + | ||
| 98 | + with ( | ||
| 99 | + patch.object( | ||
| 100 | + triton_heuristics, | ||
| 101 | + "npu_profiler_session_active", | ||
| 102 | + return_value=session_active, | ||
| 103 | + ), | ||
| 104 | + patch.object( | ||
| 105 | + triton_heuristics, | ||
| 106 | + "do_bench_using_profiling_npu", | ||
| 107 | + return_value=2.5, | ||
| 108 | + ) as profiling_benchmark, | ||
| 109 | + patch.object( | ||
| 110 | + triton_heuristics.benchmarker, | ||
| 111 | + "benchmark_gpu", | ||
| 112 | + return_value=1.5, | ||
| 113 | + ) as event_timer, | ||
| 114 | + ): | ||
| 115 | + timing = autotuner._bench_with_launch_args(lambda **kwargs: None, (), ()) | ||
| 116 | + | ||
| 117 | + if session_active: | ||
| 118 | + self.assertEqual(timing, 1.5) | ||
| 119 | + profiling_benchmark.assert_not_called() | ||
| 120 | + event_timer.assert_called_once() | ||
| 121 | + else: | ||
| 122 | + self.assertEqual(timing, 2.5) | ||
| 123 | + profiling_benchmark.assert_called_once() | ||
| 124 | + event_timer.assert_not_called() | ||
| 125 | + | ||
| 126 | + def test_probe_stays_behind_the_bandwidth_config_check(self): | ||
| 127 | + autotuner = self._autotuner(profile_bandwidth=False) | ||
| 128 | + | ||
| 129 | + with ( | ||
| 130 | + patch.object( | ||
| 131 | + triton_heuristics, | ||
| 132 | + "npu_profiler_session_active", | ||
| 133 | + return_value=False, | ||
| 134 | + ) as probe, | ||
| 135 | + patch.object( | ||
| 136 | + triton_heuristics.benchmarker, | ||
| 137 | + "benchmark_gpu", | ||
| 138 | + return_value=1.5, | ||
| 139 | + ), | ||
| 140 | + ): | ||
| 141 | + timing = autotuner._bench_with_launch_args(lambda **kwargs: None, (), ()) | ||
| 142 | + | ||
| 143 | + self.assertEqual(timing, 1.5) | ||
| 144 | + # The probe imports from torch_npu.profiler, and autotune benchmarks | ||
| 145 | + # every config, so it must not run when there is nothing to guard. | ||
| 146 | + probe.assert_not_called() | ||
| 147 | + | ||
| 148 | + | ||
| 149 | +instantiate_parametrized_tests(TestAutotuneProfilerSession) | ||
| 150 | + | ||
| 151 | + | ||
| 152 | +if __name__ == "__main__": | ||
| 153 | + run_tests() | ||
| @@ -0,0 +1,125 @@ | |||
| 1 | +"""filter_masks must keep the mask guarding a load over a dynamically sized cat. | ||
| 2 | + | ||
| 3 | +ops.masked ands one <axis>_mask into the body per loop axis of the masked | ||
| 4 | +subblock and skips size symbols, since "y1 < s0" constrains y1 and not the | ||
| 5 | +symbol itself. A size symbol left in current_subblock_axis therefore contributed | ||
| 6 | +no mask while still failing the subset test in filter_masks for every load whose | ||
| 7 | +index did not spell it out, and the tmp mask guarding that load was dropped. The | ||
| 8 | +first slice of a cat over a dynamic dimension then read out of bounds. | ||
| 9 | +""" | ||
| 10 | + | ||
| 11 | +import unittest | ||
| 12 | +from types import SimpleNamespace | ||
| 13 | +from unittest.mock import patch | ||
| 14 | + | ||
| 15 | +import torch | ||
| 16 | +from torch.testing._internal.common_utils import ( | ||
| 17 | + instantiate_parametrized_tests, | ||
| 18 | + parametrize, | ||
| 19 | + run_tests, | ||
| 20 | + TestCase, | ||
| 21 | +) | ||
| 22 | + | ||
| 23 | +import torch_npu | ||
| 24 | +import torch_npu._inductor | ||
| 25 | +from torch_npu._inductor.codegen import triton as triton_codegen | ||
| 26 | + | ||
| 27 | +if not torch.npu.is_available(): | ||
| 28 | + raise unittest.SkipTest("NPU is not available") | ||
| 29 | + | ||
| 30 | +device = "npu" | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +class TestDynamicCatMask(TestCase): | ||
| 34 | + """filter_masks may drop a redundant axis mask, never a semantic one.""" | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + def _axis(name): | ||
| 38 | + return SimpleNamespace( | ||
| 39 | + name=name, | ||
| 40 | + is_vectorized_split=False, | ||
| 41 | + is_tiling_axis=True, | ||
| 42 | + is_reduction=False, | ||
| 43 | + is_no_loop_axis=False, | ||
| 44 | + ) | ||
| 45 | + | ||
| 46 | + def _filter_masks(self, subblock_axis, index_vars): | ||
| 47 | + """Run filter_masks over a two axis kernel and report what survived. | ||
| 48 | + | ||
| 49 | + The mask set holds the two axis masks plus one tmp mask standing in for | ||
| 50 | + the guard a masked subblock puts on its load. | ||
| 51 | + """ | ||
| 52 | + kernel = object.__new__(triton_codegen.NPUIndexTritonKernel) | ||
| 53 | + kernel.sorted_axis = [self._axis("x0"), self._axis("y1")] | ||
| 54 | + kernel.persistent_reduction = False | ||
| 55 | + kernel.npu_kernel_type = triton_codegen.NPUKernelType.SIMD | ||
| 56 | + | ||
| 57 | + guard = object.__new__(triton_codegen.TritonCSEVariable) | ||
| 58 | + guard.name = "tmp3" | ||
| 59 | + mask_vars = {guard, "x0_mask", "y1_mask"} # noqa: set_linter | ||
| 60 | + | ||
| 61 | + virtualized = SimpleNamespace( | ||
| 62 | + kernel=SimpleNamespace(current_subblock_axis=set(subblock_axis)) | ||
| 63 | + ) | ||
| 64 | + with ( | ||
| 65 | + patch.object(triton_codegen, "V", virtualized), | ||
| 66 | + patch.object(triton_codegen, "get_allow_dynamic", return_value=True), | ||
| 67 | + ): | ||
| 68 | + kernel.filter_masks(mask_vars, index_vars) | ||
| 69 | + return {str(mask_var) for mask_var in mask_vars} | ||
| 70 | + | ||
| 71 | + | ||
| 72 | + def test_size_symbol_alone_keeps_the_guarding_mask(self, size_symbol): | ||
| 73 | + # What a cat over a dynamic dimension records: the masked subblock is | ||
| 74 | + # indexed by the size symbol, while the load reads the first slice and | ||
| 75 | + # never mentions it. | ||
| 76 | + kept = self._filter_masks({size_symbol}, ["y1"]) | ||
| 77 | + | ||
| 78 | + self.assertIn("tmp3", kept) | ||
| 79 | + | ||
| 80 | + def test_size_symbol_beside_a_loop_axis_keeps_the_guarding_mask(self): | ||
| 81 | + kept = self._filter_masks({"y1", "s0"}, ["y1"]) | ||
| 82 | + | ||
| 83 | + self.assertIn("tmp3", kept) | ||
| 84 | + # y1 is a subblock axis, so ops.masked already anded y1_mask into the | ||
| 85 | + # body and the axis mask on the load is the redundant kind. | ||
| 86 | + self.assertEqual(kept, {"tmp3", "x0_mask"}) | ||
| 87 | + | ||
| 88 | + def test_loop_axis_outside_the_index_still_drops_the_guarding_mask(self): | ||
| 89 | + # x0 is a real loop axis that this load does not index, so the guard was | ||
| 90 | + # built for a different shape and must not travel with it. | ||
| 91 | + kept = self._filter_masks({"x0"}, ["y1"]) | ||
| 92 | + | ||
| 93 | + self.assertEqual(kept, {"x0_mask", "y1_mask"}) | ||
| 94 | + | ||
| 95 | + def test_no_subblock_leaves_every_mask_alone(self): | ||
| 96 | + kept = self._filter_masks(set(), ["y1"]) | ||
| 97 | + | ||
| 98 | + self.assertEqual(kept, {"tmp3", "x0_mask", "y1_mask"}) | ||
| 99 | + | ||
| 100 | + def test_dynamic_cat_slice_stays_in_bounds(self): | ||
| 101 | + def fn(head, tail): | ||
| 102 | + return torch.cat([head, tail], dim=1) + 1.0 | ||
| 103 | + | ||
| 104 | + head = torch.randn(32, 200, device=device) | ||
| 105 | + tail = torch.randn(32, 56, device=device) | ||
| 106 | + torch._dynamo.mark_dynamic(head, 1) | ||
| 107 | + expected = fn(head, tail) | ||
| 108 | + | ||
| 109 | + try: | ||
| 110 | + compiled = torch.compile(fn, backend="inductor", dynamic=True) | ||
| 111 | + torch.testing.assert_close(compiled(head, tail), expected) | ||
| 112 | + | ||
| 113 | + # A second extent reuses the compiled kernel, so the guard has to | ||
| 114 | + # follow the symbol rather than the size that was traced. | ||
| 115 | + wider = torch.randn(32, 328, device=device) | ||
| 116 | + torch.testing.assert_close(compiled(wider, tail), fn(wider, tail)) | ||
| 117 | + finally: | ||
| 118 | + torch._dynamo.reset() | ||
| 119 | + | ||
| 120 | + | ||
| 121 | +instantiate_parametrized_tests(TestDynamicCatMask) | ||
| 122 | + | ||
| 123 | + | ||
| 124 | +if __name__ == "__main__": | ||
| 125 | + run_tests() | ||
| @@ -4168,7 +4168,16 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4168 | # loads/stores, indirect indexing validity masks, etc.), otherwise we | 4168 | # loads/stores, indirect indexing validity masks, etc.), otherwise we |
| 4169 | # can generate incorrect tl.load/tl.store. | 4169 | # can generate incorrect tl.load/tl.store. |
| 4170 | masked_axis_name = [] | 4170 | masked_axis_name = [] |
| 4171 | - subblock_axis = V.kernel.current_subblock_axis | 4171 | + # Size symbols are not loop axes: "y0 < s0" only constrains y0, and |
| 4172 | + # ops.masked skips them when it ands the axis masks in. Left in, they | ||
| 4173 | + # fail the subset test below for every load whose index does not spell | ||
| 4174 | + # the size symbol out, e.g. the first slice of a cat over a dynamic | ||
| 4175 | + # dimension, whose guarding mask would then be dropped. | ||
| 4176 | + subblock_axis = { # noqa: set_linter | ||
| 4177 | + axis | ||
| 4178 | + for axis in V.kernel.current_subblock_axis | ||
| 4179 | + if not str(axis).startswith(("s", "ps", "i")) | ||
| 4180 | + } | ||
| 4172 | save_variable_mask = True | 4181 | save_variable_mask = True |
| 4173 | if index_vars and subblock_axis: | 4182 | if index_vars and subblock_axis: |
| 4174 | save_variable_mask = subblock_axis.issubset( | 4183 | save_variable_mask = subblock_axis.issubset( |
| @@ -131,6 +131,22 @@ class CompileThreadPool: | |||
| 131 | compile_thread_pool = CompileThreadPool() | 131 | compile_thread_pool = CompileThreadPool() |
| 132 | 132 | ||
| 133 | 133 | ||
| 134 | +def npu_profiler_session_active(): | ||
| 135 | + """Whether a torch_npu profiler session is already tracing in this process. | ||
| 136 | + | ||
| 137 | + The NPU profiler backend is a process-wide singleton. A session opened | ||
| 138 | + inside another one finalizes the shared trace when it exits, so the outer | ||
| 139 | + session loses its data and reports "Profiler is not initialized". Autotune | ||
| 140 | + must therefore never open a profiler while the caller is profiling. | ||
| 141 | + """ | ||
| 142 | + try: | ||
| 143 | + from torch_npu.profiler._profiler_path_creator import ProfPathCreator | ||
| 144 | + return bool(ProfPathCreator().is_prof_inited) | ||
| 145 | + except Exception as exc: # the probe must never break a kernel launch | ||
| 146 | + log.debug("could not probe the torch_npu profiler state: %s", exc) | ||
| 147 | + return False | ||
| 148 | + | ||
| 149 | + | ||
| 134 | 150 | ||
| 135 | def create_profiler(torch_path, wait=0, warmup=1, active=1, repeat=1, skip_first=1): | 151 | def create_profiler(torch_path, wait=0, warmup=1, active=1, repeat=1, skip_first=1): |
| 136 | experimental_config = torch_npu.profiler._ExperimentalConfig( | 152 | experimental_config = torch_npu.profiler._ExperimentalConfig( |
| @@ -1531,7 +1547,10 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1531 | stream=stream, | 1547 | stream=stream, |
| 1532 | ) | 1548 | ) |
| 1533 | 1549 | ||
| 1534 | - if self.inductor_meta.get("profile_bandwidth_with_do_bench_using_profiling", False): | 1550 | + if ( |
| 1551 | + self.inductor_meta.get("profile_bandwidth_with_do_bench_using_profiling", False) | ||
| 1552 | + and not npu_profiler_session_active() | ||
| 1553 | + ): | ||
| 1535 | return do_bench_using_profiling_npu(kernel_call, rep=1) | 1554 | return do_bench_using_profiling_npu(kernel_call, rep=1) |
| 1536 | 1555 | ||
| 1537 | return benchmarker.benchmark_gpu(kernel_call, rep=1) | 1556 | return benchmarker.benchmark_gpu(kernel_call, rep=1) |
| @@ -1594,6 +1613,15 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1594 | if not kernel_funcs or not npu_config.aggresive_autotune: | 1613 | if not kernel_funcs or not npu_config.aggresive_autotune: |
| 1595 | return None | 1614 | return None |
| 1596 | 1615 | ||
| 1616 | + if npu_profiler_session_active(): | ||
| 1617 | + warning_once( | ||
| 1618 | + log, | ||
| 1619 | + "autotune ran while the caller was profiling, so the batch benchmark " | ||
| 1620 | + "was skipped in favour of the timer; warm the model up before " | ||
| 1621 | + "starting the profiler to keep autotune out of the profiled region", | ||
| 1622 | + ) | ||
| 1623 | + return None | ||
| 1624 | + | ||
| 1597 | try: | 1625 | try: |
| 1598 | batch_timings = tuple( | 1626 | batch_timings = tuple( |
| 1599 | mspti_batch_benchmark(kernel_funcs, filter_list=["triton"]) | 1627 | mspti_batch_benchmark(kernel_funcs, filter_list=["triton"]) |