已合并
fix(inductor): guard grouped benchmark memory footprint #43884
Xuan Peng创建于 8月6日
fix(inductor): guard grouped benchmark memory footprint #43884
已合并
共 5 个文件变更+533-0
| @@ -36,6 +36,7 @@ from torch.testing._internal.common_utils import ( | |||
| 36 | 36 | ||
| 37 | import torch_npu | 37 | import torch_npu |
| 38 | import torch_npu._inductor | 38 | import torch_npu._inductor |
| 39 | +from torch_npu._inductor.codegen import triton as triton_codegen | ||
| 39 | from torch_npu._inductor.codegen import split_tiling as split_tiling_module | 40 | from torch_npu._inductor.codegen import split_tiling as split_tiling_module |
| 40 | from torch_npu._inductor.codegen.split_tiling import SplitTiling | 41 | from torch_npu._inductor.codegen.split_tiling import SplitTiling |
| 41 | from torch_npu._inductor.config import num_vector_core | 42 | from torch_npu._inductor.config import num_vector_core |
| @@ -59,6 +60,11 @@ def make_axis(name, length): | |||
| 59 | ) | 60 | ) |
| 60 | 61 | ||
| 61 | 62 | ||
| 63 | +class _FixedMemoryDeviceProperties(SimpleNamespace): | ||
| 64 | + def __getattr__(self, name): | ||
| 65 | + return getattr(self._base_properties, name) | ||
| 66 | + | ||
| 67 | + | ||
| 62 | # ============================================================ | 68 | # ============================================================ |
| 63 | # Base class: provides compile counting + dynamic=True helpers | 69 | # Base class: provides compile counting + dynamic=True helpers |
| 64 | # ============================================================ | 70 | # ============================================================ |
| @@ -653,6 +659,168 @@ class TestSymbolicGroupElementwise(TestCase): | |||
| 653 | # ============================================================ | 659 | # ============================================================ |
| 654 | 660 | ||
| 655 | class TestPointwiseSymbolicGrouping(TestCase): | 661 | class TestPointwiseSymbolicGrouping(TestCase): |
| 662 | + | ||
| 663 | + def _benchmark_guard_meta(width): | ||
| 664 | + return { | ||
| 665 | + "kernel_name": "pointwise_footprint_guard_test", | ||
| 666 | + "group_enabled": True, | ||
| 667 | + "group_template": "pointwise", | ||
| 668 | + "group_workload": None, | ||
| 669 | + "primary_group_axis": "x0", | ||
| 670 | + "axis_names": ("x0",), | ||
| 671 | + "axis_static_values": (), | ||
| 672 | + "group_features": ({ | ||
| 673 | + "name": "pointwise", | ||
| 674 | + "source": "outer_product", | ||
| 675 | + "axis_names": ("x0",), | ||
| 676 | + "buckets": (229376,), | ||
| 677 | + },), | ||
| 678 | + "mutated_arg_names": (), | ||
| 679 | + "ordered_arg_specs": ({ | ||
| 680 | + "kind": "tensor", | ||
| 681 | + "source": "buffer", | ||
| 682 | + "name": "in_ptr0", | ||
| 683 | + "dtype": "torch.float16", | ||
| 684 | + "device": "npu:0", | ||
| 685 | + "size_exprs": ( | ||
| 686 | + {"axis_name": "x0"}, | ||
| 687 | + {"const": width}, | ||
| 688 | + ), | ||
| 689 | + "stride_exprs": ( | ||
| 690 | + {"const": width}, | ||
| 691 | + {"const": 1}, | ||
| 692 | + ), | ||
| 693 | + },), | ||
| 694 | + } | ||
| 695 | + | ||
| 696 | + | ||
| 697 | + def _run_benchmark_guard(meta, total_memory=108 * 1024**3): | ||
| 698 | + kernel = object.__new__(triton_codegen.NPUIndexTritonKernel) | ||
| 699 | + properties = SimpleNamespace(total_memory=total_memory) | ||
| 700 | + with ( | ||
| 701 | + patch.object(torch.npu, "get_device_properties", return_value=properties), | ||
| 702 | + patch.object( | ||
| 703 | + triton_codegen.npu_config, | ||
| 704 | + "symbolic_group_max_benchmark_memory_ratio", | ||
| 705 | + 0.25, | ||
| 706 | + ), | ||
| 707 | + ): | ||
| 708 | + kernel._disable_grouped_autotune_if_benchmark_too_large( | ||
| 709 | + meta, | ||
| 710 | + torch.device("npu", 0), | ||
| 711 | + ) | ||
| 712 | + return kernel | ||
| 713 | + | ||
| 714 | + | ||
| 715 | + "width, expected_group_enabled", | ||
| 716 | + ((69876, False), (1, True)), | ||
| 717 | + ) | ||
| 718 | + def test_benchmark_footprint_guard_updates_pointwise_group( | ||
| 719 | + self, | ||
| 720 | + width, | ||
| 721 | + expected_group_enabled, | ||
| 722 | + ): | ||
| 723 | + meta = self._benchmark_guard_meta(width) | ||
| 724 | + self._run_benchmark_guard(meta) | ||
| 725 | + | ||
| 726 | + self.assertEqual(meta["group_enabled"], expected_group_enabled) | ||
| 727 | + if not expected_group_enabled: | ||
| 728 | + self.assertIsNone(meta["group_template"]) | ||
| 729 | + | ||
| 730 | + def test_benchmark_footprint_budget_boundary_is_inclusive(self): | ||
| 731 | + from torch_npu._inductor.runtime.symbolic_grouping import ( | ||
| 732 | + build_group_representatives, | ||
| 733 | + estimate_grouped_benchmark_footprint, | ||
| 734 | + ) | ||
| 735 | + | ||
| 736 | + meta = self._benchmark_guard_meta(1) | ||
| 737 | + representatives = build_group_representatives( | ||
| 738 | + meta["group_features"], | ||
| 739 | + meta["axis_names"], | ||
| 740 | + meta["axis_static_values"], | ||
| 741 | + ) | ||
| 742 | + footprint = estimate_grouped_benchmark_footprint( | ||
| 743 | + representatives, | ||
| 744 | + meta["ordered_arg_specs"], | ||
| 745 | + ) | ||
| 746 | + self._run_benchmark_guard(meta, footprint.total_bytes * 4) | ||
| 747 | + | ||
| 748 | + self.assertTrue(meta["group_enabled"]) | ||
| 749 | + | ||
| 750 | + def test_benchmark_footprint_rejection_enables_auto_blockify(self): | ||
| 751 | + meta = self._benchmark_guard_meta(69876) | ||
| 752 | + kernel = self._run_benchmark_guard(meta) | ||
| 753 | + with ( | ||
| 754 | + patch.object(kernel, "_has_dynamic_shape_axis", return_value=True), | ||
| 755 | + patch.object( | ||
| 756 | + triton_codegen.npu_config, | ||
| 757 | + "enable_symbolic_shape_group_autotune", | ||
| 758 | + True, | ||
| 759 | + ), | ||
| 760 | + ): | ||
| 761 | + kernel._enable_auto_blockify_for_grouped_fallback_if_needed(meta) | ||
| 762 | + | ||
| 763 | + self.assertFalse(meta["group_enabled"]) | ||
| 764 | + self.assertTrue(meta["enable_auto_blockify"]) | ||
| 765 | + | ||
| 766 | + def test_wide_backing_storage_falls_back_before_grouped_benchmark(self): | ||
| 767 | + import torch_npu._inductor.config as npu_config | ||
| 768 | + | ||
| 769 | + def fn(values): | ||
| 770 | + return ( | ||
| 771 | + values[:, 6826] | ||
| 772 | + + values[:, 42400] | ||
| 773 | + + values[:, 43912] | ||
| 774 | + ) | ||
| 775 | + | ||
| 776 | + values = torch.randn( | ||
| 777 | + (200, 69876), | ||
| 778 | + device=device, | ||
| 779 | + dtype=torch.float16, | ||
| 780 | + ) | ||
| 781 | + torch._dynamo.mark_dynamic(values, 0) | ||
| 782 | + expected = fn(values) | ||
| 783 | + fixed_properties = _FixedMemoryDeviceProperties( | ||
| 784 | + total_memory=108 * 1024**3, | ||
| 785 | + _base_properties=torch.npu.get_device_properties(torch.device("npu", 0)), | ||
| 786 | + ) | ||
| 787 | + | ||
| 788 | + try: | ||
| 789 | + with ( | ||
| 790 | + patch.object( | ||
| 791 | + torch.npu, | ||
| 792 | + "get_device_properties", | ||
| 793 | + return_value=fixed_properties, | ||
| 794 | + ), | ||
| 795 | + patch.object( | ||
| 796 | + npu_config, | ||
| 797 | + "enable_symbolic_shape_group_autotune", | ||
| 798 | + True, | ||
| 799 | + ), | ||
| 800 | + patch.object( | ||
| 801 | + npu_config, | ||
| 802 | + "symbolic_group_max_benchmark_memory_ratio", | ||
| 803 | + 0.25, | ||
| 804 | + ), | ||
| 805 | + ): | ||
| 806 | + compiled = torch.compile(fn, backend="inductor", dynamic=True) | ||
| 807 | + actual, codes = run_and_get_code(compiled, values) | ||
 | |||
| 808 | + finally: | ||
| 809 | + torch._dynamo.reset() | ||
| 810 | + | ||
| 811 | + torch.testing.assert_close(actual, expected) | ||
| 812 | + matching_codes = [ | ||
| 813 | + code | ||
| 814 | + for code in codes | ||
| 815 | + if "69876" in code | ||
| 816 | + and "'group_enabled': False" in code | ||
| 817 | + and "'enable_auto_blockify': True" in code | ||
| 818 | + ] | ||
| 819 | + self.assertTrue( | ||
| 820 | + matching_codes, | ||
| 821 | + f"Expected wide pointwise grouped fallback, got:\n{codes}", | ||
| 822 | + ) | ||
| 823 | + | ||
| 656 | def test_static_split_dynamic_tiling_group(self): | 824 | def test_static_split_dynamic_tiling_group(self): |
| 657 | # Transpose and broadcast commonly leave the dynamic inner axis tiled | 825 | # Transpose and broadcast commonly leave the dynamic inner axis tiled |
| 658 | # while a static outer axis supplies grid parallelism. | 826 | # while a static outer axis supplies grid parallelism. |
| @@ -0,0 +1,112 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 3 | + | ||
| 4 | +from torch_npu._inductor.runtime.symbolic_grouping import ( | ||
| 5 | + UnsupportedGroupedPlan, | ||
| 6 | + build_group_representatives, | ||
| 7 | + estimate_grouped_benchmark_footprint, | ||
| 8 | + evaluate_grouped_benchmark_expr, | ||
| 9 | + required_storage_numel, | ||
| 10 | +) | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +class TestGroupedBenchmarkFootprint(TestCase): | ||
| 14 | + def test_required_storage_numel_matches_rand_strided_layout(self): | ||
| 15 | + self.assertEqual(required_storage_numel((8, 10), (10, 1)), 80) | ||
| 16 | + self.assertEqual(required_storage_numel((3, 4), (8, 1)), 20) | ||
| 17 | + self.assertEqual(required_storage_numel((8, 10), (0, 1)), 10) | ||
| 18 | + self.assertEqual(required_storage_numel((0, 10), (10, 1)), 0) | ||
| 19 | + | ||
| 20 | + def test_invalid_storage_layout_fails_closed(self): | ||
| 21 | + with self.assertRaisesRegex( | ||
| 22 | + UnsupportedGroupedPlan, | ||
| 23 | + "rank mismatch", | ||
| 24 | + ): | ||
| 25 | + required_storage_numel((8, 10), (1,)) | ||
| 26 | + with self.assertRaisesRegex( | ||
| 27 | + UnsupportedGroupedPlan, | ||
| 28 | + "negative benchmark tensor stride", | ||
| 29 | + ): | ||
| 30 | + required_storage_numel((8,), (-1,)) | ||
| 31 | + | ||
| 32 | + def test_expression_evaluation_uses_group_axis_environment(self): | ||
| 33 | + expr = { | ||
| 34 | + "add": ( | ||
| 35 | + {"mul": ({"axis_name": "x0"}, {"const": 10})}, | ||
| 36 | + {"const": 3}, | ||
| 37 | + ) | ||
| 38 | + } | ||
| 39 | + self.assertEqual(evaluate_grouped_benchmark_expr(expr, {"x0": 8}), 83) | ||
| 40 | + | ||
| 41 | + def test_footprint_sums_groups_and_one_mutated_clone_peak(self): | ||
| 42 | + features = ({ | ||
| 43 | + "name": "pointwise", | ||
| 44 | + "source": "outer_product", | ||
| 45 | + "axis_names": ("x0",), | ||
| 46 | + "buckets": (8,), | ||
| 47 | + },) | ||
| 48 | + representatives = build_group_representatives(features, ("x0",), ()) | ||
| 49 | + specs = ( | ||
| 50 | + { | ||
| 51 | + "kind": "tensor", | ||
| 52 | + "source": "buffer", | ||
| 53 | + "name": "out_ptr0", | ||
| 54 | + "dtype": "torch.float16", | ||
| 55 | + "device": "npu:0", | ||
| 56 | + "size_exprs": ({"axis_name": "x0"},), | ||
| 57 | + "stride_exprs": ({"const": 1},), | ||
| 58 | + }, | ||
| 59 | + { | ||
| 60 | + "kind": "tensor", | ||
| 61 | + "source": "buffer", | ||
| 62 | + "name": "in_ptr0", | ||
| 63 | + "dtype": "torch.float16", | ||
| 64 | + "device": "npu:0", | ||
| 65 | + "size_exprs": ({"axis_name": "x0"}, {"const": 10}), | ||
| 66 | + "stride_exprs": ({"const": 10}, {"const": 1}), | ||
| 67 | + }, | ||
| 68 | + { | ||
| 69 | + "kind": "size", | ||
| 70 | + "source": "axis_expr", | ||
| 71 | + "name": "x0_numel", | ||
| 72 | + "expr": {"axis_name": "x0"}, | ||
| 73 | + }, | ||
| 74 | + ) | ||
| 75 | + | ||
| 76 | + footprint = estimate_grouped_benchmark_footprint( | ||
| 77 | + representatives, | ||
| 78 | + specs, | ||
| 79 | + mutated_arg_names=("out_ptr0",), | ||
| 80 | + ) | ||
| 81 | + | ||
| 82 | + self.assertEqual(footprint.group_bytes, ((0, 176), (1, 352))) | ||
| 83 | + self.assertEqual(footprint.synthetic_bytes, 528) | ||
| 84 | + self.assertEqual(footprint.mutated_clone_bytes, 32) | ||
| 85 | + self.assertEqual(footprint.total_bytes, 560) | ||
| 86 | + self.assertEqual(footprint.largest_group_id, 1) | ||
| 87 | + self.assertEqual(footprint.dominant_arg.name, "in_ptr0") | ||
| 88 | + self.assertEqual(footprint.dominant_arg.num_bytes, 320) | ||
| 89 | + | ||
| 90 | + def test_runtime_dependent_expression_fails_closed(self): | ||
| 91 | + representatives = { | ||
| 92 | + "reachable_group_ids": (0,), | ||
| 93 | + "benchmark_axis_values_by_group": ((('x0', 8),),), | ||
| 94 | + } | ||
| 95 | + specs = ({ | ||
| 96 | + "kind": "tensor", | ||
| 97 | + "source": "buffer", | ||
| 98 | + "name": "in_ptr0", | ||
| 99 | + "dtype": "torch.float16", | ||
| 100 | + "device": "npu:0", | ||
| 101 | + "size_exprs": ({"runtime_arg_index": 1},), | ||
| 102 | + "stride_exprs": ({"const": 1},), | ||
| 103 | + },) | ||
| 104 | + with self.assertRaisesRegex( | ||
| 105 | + UnsupportedGroupedPlan, | ||
| 106 | + "runtime_arg_index cannot be bounded", | ||
| 107 | + ): | ||
| 108 | + estimate_grouped_benchmark_footprint(representatives, specs) | ||
| 109 | + | ||
| 110 | + | ||
| 111 | +if __name__ == "__main__": | ||
| 112 | + run_tests() | ||
| @@ -81,6 +81,7 @@ from ..runtime.symbolic_grouping import ( | |||
| 81 | GroupedKernelMeta, | 81 | GroupedKernelMeta, |
| 82 | UnsupportedGroupedPlan, | 82 | UnsupportedGroupedPlan, |
| 83 | build_group_representatives, | 83 | build_group_representatives, |
| 84 | + estimate_grouped_benchmark_footprint, | ||
| 84 | ) | 85 | ) |
| 85 | from .kernel_analysis import ( | 86 | from .kernel_analysis import ( |
| 86 | collect_stride_sorted_vars_from_indexings, | 87 | collect_stride_sorted_vars_from_indexings, |
| @@ -2401,6 +2402,62 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2401 | except UnsupportedGroupedPlan as exc: | 2402 | except UnsupportedGroupedPlan as exc: |
| 2402 | self._disable_grouped_autotune(inductor_meta, str(exc)) | 2403 | self._disable_grouped_autotune(inductor_meta, str(exc)) |
| 2403 | 2404 | ||
| 2405 | + def _disable_grouped_autotune_if_benchmark_too_large( | ||
| 2406 | + self, | ||
| 2407 | + inductor_meta, | ||
| 2408 | + device, | ||
| 2409 | + ): | ||
| 2410 | + if not inductor_meta.get("group_enabled", False): | ||
| 2411 | + return | ||
| 2412 | + | ||
| 2413 | + try: | ||
| 2414 | + representatives = build_group_representatives( | ||
| 2415 | + inductor_meta.get("group_features", ()), | ||
| 2416 | + inductor_meta.get("axis_names", ()), | ||
| 2417 | + inductor_meta.get("axis_static_values", ()), | ||
| 2418 | + ) | ||
| 2419 | + footprint = estimate_grouped_benchmark_footprint( | ||
| 2420 | + representatives, | ||
| 2421 | + inductor_meta.get("ordered_arg_specs", ()), | ||
| 2422 | + inductor_meta.get("mutated_arg_names", ()), | ||
| 2423 | + ) | ||
| 2424 | + total_memory = int( | ||
| 2425 | + torch.npu.get_device_properties(device).total_memory | ||
| 2426 | + ) | ||
| 2427 | + if total_memory <= 0: | ||
| 2428 | + raise UnsupportedGroupedPlan( | ||
| 2429 | + f"invalid NPU total memory: {total_memory}" | ||
| 2430 | + ) | ||
| 2431 | + except (AttributeError, KeyError, RuntimeError, TypeError, ValueError) as exc: | ||
| 2432 | + self._disable_grouped_autotune( | ||
| 2433 | + inductor_meta, | ||
| 2434 | + f"grouped benchmark footprint is not bounded: {exc}", | ||
| 2435 | + ) | ||
| 2436 | + return | ||
| 2437 | + | ||
| 2438 | + ratio = npu_config.symbolic_group_max_benchmark_memory_ratio | ||
| 2439 | + budget_bytes = int(total_memory * ratio) | ||
| 2440 | + if footprint.total_bytes <= budget_bytes: | ||
| 2441 | + return | ||
| 2442 | + | ||
| 2443 | + dominant = footprint.dominant_arg | ||
| 2444 | + dominant_text = "none" | ||
| 2445 | + if dominant is not None: | ||
| 2446 | + dominant_text = ( | ||
| 2447 | + f"group_id={dominant.group_id} arg={dominant.name} " | ||
| 2448 | + f"kind={dominant.kind} size={dominant.size} " | ||
| 2449 | + f"stride={dominant.stride} dtype={dominant.dtype} " | ||
| 2450 | + f"bytes={dominant.num_bytes}" | ||
| 2451 | + ) | ||
| 2452 | + self._disable_grouped_autotune( | ||
| 2453 | + inductor_meta, | ||
| 2454 | + "grouped benchmark footprint exceeds budget: " | ||
| 2455 | + f"estimated_bytes={footprint.total_bytes} " | ||
| 2456 | + f"budget_bytes={budget_bytes} ratio={ratio} " | ||
| 2457 | + f"largest_group_id={footprint.largest_group_id} " | ||
| 2458 | + f"dominant=({dominant_text})", | ||
| 2459 | + ) | ||
| 2460 | + | ||
| 2404 | def _has_dynamic_shape_axis(self): | 2461 | def _has_dynamic_shape_axis(self): |
| 2405 | for axis in self.sorted_axis: | 2462 | for axis in self.sorted_axis: |
| 2406 | length = V.graph.sizevars.simplify(axis.length) | 2463 | length = V.graph.sizevars.simplify(axis.length) |
| @@ -2668,6 +2725,10 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2668 | self.build_grouped_benchmark_arg_specs(argdefs, signature) | 2725 | self.build_grouped_benchmark_arg_specs(argdefs, signature) |
| 2669 | ) | 2726 | ) |
| 2670 | inductor_meta.setdefault("extra_launcher_arg_specs", ()) | 2727 | inductor_meta.setdefault("extra_launcher_arg_specs", ()) |
| 2728 | + self._disable_grouped_autotune_if_benchmark_too_large( | ||
| 2729 | + inductor_meta, | ||
| 2730 | + V.graph.get_current_device_or_throw(), | ||
| 2731 | + ) | ||
| 2671 | except RuntimeError as exc: | 2732 | except RuntimeError as exc: |
| 2672 | self._disable_grouped_autotune(inductor_meta, str(exc)) | 2733 | self._disable_grouped_autotune(inductor_meta, str(exc)) |
| 2673 | 2734 | ||
| @@ -423,6 +423,12 @@ enable_fast_launch = _parse_bool_env( | |||
| 423 | ) | 423 | ) |
| 424 | enable_costmodel_backend = _parse_bool_env("INDUCTOR_ASCEND_ENABLE_COSTMODEL", False) | 424 | enable_costmodel_backend = _parse_bool_env("INDUCTOR_ASCEND_ENABLE_COSTMODEL", False) |
| 425 | costmodel_ratio = _parse_float_env("INDUCTOR_ASCEND_COSTMODEL_RATIO", 0.25, 0.0, 1.0) | 425 | costmodel_ratio = _parse_float_env("INDUCTOR_ASCEND_COSTMODEL_RATIO", 0.25, 0.0, 1.0) |
| 426 | +symbolic_group_max_benchmark_memory_ratio = _parse_float_env( | ||
| 427 | + "INDUCTOR_ASCEND_SYMBOLIC_GROUP_MAX_BENCHMARK_MEMORY_RATIO", | ||
| 428 | + 0.25, | ||
| 429 | + 0.0, | ||
| 430 | + 1.0, | ||
| 431 | +) | ||
| 426 | 432 | ||
| 427 | 433 | ||
| 428 | lowering_axis_count = None | 434 | lowering_axis_count = None |
| @@ -5,6 +5,7 @@ from collections.abc import Mapping, Sequence | |||
| 5 | from typing import Literal | 5 | from typing import Literal |
| 6 | 6 | ||
| 7 | import sympy | 7 | import sympy |
| 8 | +import torch | ||
| 8 | 9 | ||
| 9 | 10 | ||
| 10 | JsonScalar = int | float | bool | str | None | 11 | JsonScalar = int | float | bool | str | None |
| @@ -110,6 +111,27 @@ class GroupedKernelMeta: | |||
| 110 | ) | 111 | ) |
| 111 | 112 | ||
| 112 | 113 | ||
| 114 | + | ||
| 115 | +class GroupedBenchmarkArgFootprint: | ||
| 116 | + group_id: int | ||
| 117 | + name: str | ||
| 118 | + kind: str | ||
| 119 | + dtype: str | ||
| 120 | + size: tuple[int, ...] | ||
| 121 | + stride: tuple[int, ...] | ||
| 122 | + num_bytes: int | ||
| 123 | + | ||
| 124 | + | ||
| 125 | + | ||
| 126 | +class GroupedBenchmarkFootprint: | ||
| 127 | + synthetic_bytes: int | ||
| 128 | + mutated_clone_bytes: int | ||
| 129 | + total_bytes: int | ||
| 130 | + group_bytes: tuple[tuple[int, int], ...] | ||
| 131 | + largest_group_id: int | None | ||
| 132 | + dominant_arg: GroupedBenchmarkArgFootprint | None | ||
| 133 | + | ||
| 134 | + | ||
| 113 | def is_runtime_symbolic_length(length) -> bool: | 135 | def is_runtime_symbolic_length(length) -> bool: |
| 114 | return not isinstance(length, (int, sympy.Integer)) | 136 | return not isinstance(length, (int, sympy.Integer)) |
| 115 | 137 | ||
| @@ -343,6 +365,170 @@ def build_group_representatives( | |||
| 343 | ), | 365 | ), |
| 344 | } | 366 | } |
| 345 | return plan | 367 | return plan |
| 368 | + | ||
| 369 | + | ||
| 370 | +def evaluate_grouped_benchmark_expr(expr, axis_env) -> int: | ||
| 371 | + if isinstance(expr, int): | ||
| 372 | + return int(expr) | ||
| 373 | + if not isinstance(expr, Mapping): | ||
| 374 | + raise UnsupportedGroupedPlan( | ||
| 375 | + f"unsupported grouped benchmark expression: {expr}" | ||
| 376 | + ) | ||
| 377 | + if "const" in expr: | ||
| 378 | + return int(expr["const"]) | ||
| 379 | + if "axis_name" in expr: | ||
| 380 | + axis_name = expr["axis_name"] | ||
| 381 | + if axis_name not in axis_env: | ||
| 382 | + raise UnsupportedGroupedPlan( | ||
| 383 | + f"benchmark axis environment is missing {axis_name}" | ||
| 384 | + ) | ||
| 385 | + return int(axis_env[axis_name]) | ||
| 386 | + if "runtime_arg_index" in expr: | ||
| 387 | + raise UnsupportedGroupedPlan( | ||
| 388 | + "runtime_arg_index cannot be bounded during grouped codegen" | ||
| 389 | + ) | ||
| 390 | + if "mul" in expr: | ||
| 391 | + product = 1 | ||
| 392 | + for operand in expr["mul"]: | ||
| 393 | + product *= evaluate_grouped_benchmark_expr(operand, axis_env) | ||
| 394 | + return product | ||
| 395 | + if "add" in expr: | ||
| 396 | + return sum( | ||
| 397 | + evaluate_grouped_benchmark_expr(operand, axis_env) | ||
| 398 | + for operand in expr["add"] | ||
| 399 | + ) | ||
| 400 | + if "floordiv" in expr: | ||
| 401 | + operands = tuple( | ||
| 402 | + evaluate_grouped_benchmark_expr(operand, axis_env) | ||
| 403 | + for operand in expr["floordiv"] | ||
| 404 | + ) | ||
| 405 | + if len(operands) != 2 or operands[1] == 0: | ||
| 406 | + raise UnsupportedGroupedPlan( | ||
| 407 | + f"invalid grouped benchmark floordiv expression: {expr}" | ||
| 408 | + ) | ||
| 409 | + return operands[0] // operands[1] | ||
| 410 | + raise UnsupportedGroupedPlan( | ||
| 411 | + f"unsupported grouped benchmark expression: {expr}" | ||
| 412 | + ) | ||
| 413 | + | ||
| 414 | + | ||
| 415 | +def required_storage_numel(size, stride) -> int: | ||
| 416 | + size = tuple(int(dim) for dim in size) | ||
| 417 | + stride = tuple(int(dim_stride) for dim_stride in stride) | ||
| 418 | + if len(size) != len(stride): | ||
| 419 | + raise UnsupportedGroupedPlan( | ||
| 420 | + f"benchmark size/stride rank mismatch: {size} vs {stride}" | ||
| 421 | + ) | ||
| 422 | + if any(dim < 0 for dim in size): | ||
| 423 | + raise UnsupportedGroupedPlan(f"negative benchmark tensor size: {size}") | ||
| 424 | + if any(dim_stride < 0 for dim_stride in stride): | ||
| 425 | + raise UnsupportedGroupedPlan( | ||
| 426 | + f"negative benchmark tensor stride: {stride}" | ||
| 427 | + ) | ||
| 428 | + if any(dim == 0 for dim in size): | ||
| 429 | + return 0 | ||
| 430 | + return 1 + sum( | ||
| 431 | + (dim - 1) * dim_stride | ||
| 432 | + for dim, dim_stride in zip(size, stride) | ||
| 433 | + ) | ||
| 434 | + | ||
| 435 | + | ||
| 436 | +def _grouped_benchmark_dtype_itemsize(dtype) -> tuple[str, int]: | ||
| 437 | + if isinstance(dtype, str): | ||
| 438 | + dtype = getattr(torch, dtype.removeprefix("torch."), None) | ||
| 439 | + if not isinstance(dtype, torch.dtype): | ||
| 440 | + raise UnsupportedGroupedPlan(f"unknown grouped benchmark dtype: {dtype}") | ||
| 441 | + return str(dtype), int(dtype.itemsize) | ||
| 442 | + | ||
| 443 | + | ||
| 444 | +def estimate_grouped_benchmark_footprint( | ||
| 445 | + group_representatives, | ||
| 446 | + ordered_arg_specs, | ||
| 447 | + mutated_arg_names=(), | ||
| 448 | +) -> GroupedBenchmarkFootprint: | ||
| 449 | + reachable = tuple(group_representatives["reachable_group_ids"]) | ||
| 450 | + axis_values = tuple( | ||
| 451 | + group_representatives["benchmark_axis_values_by_group"] | ||
| 452 | + ) | ||
| 453 | + ordered_arg_specs = tuple(ordered_arg_specs) | ||
| 454 | + mutated_names = set(mutated_arg_names) | ||
| 455 | + synthetic_bytes = 0 | ||
| 456 | + mutated_clone_bytes = 0 | ||
| 457 | + group_bytes = [] | ||
| 458 | + dominant_arg = None | ||
| 459 | + | ||
| 460 | + for group_id in reachable: | ||
| 461 | + axis_env = dict(axis_values[group_id]) | ||
| 462 | + current_bytes = 0 | ||
| 463 | + current_mutated_bytes = 0 | ||
| 464 | + for spec in ordered_arg_specs: | ||
| 465 | + kind = spec.get("kind") | ||
| 466 | + source = spec.get("source") | ||
| 467 | + if source == "runtime_arg" or kind == "size": | ||
| 468 | + continue | ||
| 469 | + if kind == "tensor" and source in ("buffer", "constant"): | ||
| 470 | + size = tuple( | ||
| 471 | + evaluate_grouped_benchmark_expr(expr, axis_env) | ||
| 472 | + for expr in spec["size_exprs"] | ||
| 473 | + ) | ||
| 474 | + stride = tuple( | ||
| 475 | + evaluate_grouped_benchmark_expr(expr, axis_env) | ||
| 476 | + for expr in spec["stride_exprs"] | ||
| 477 | + ) | ||
| 478 | + dtype_name, itemsize = _grouped_benchmark_dtype_itemsize( | ||
| 479 | + spec["dtype"] | ||
| 480 | + ) | ||
| 481 | + num_bytes = required_storage_numel(size, stride) * itemsize | ||
| 482 | + elif kind == "workspace" and source == "workspace": | ||
| 483 | + count = evaluate_grouped_benchmark_expr( | ||
| 484 | + spec["count_expr"], axis_env | ||
| 485 | + ) | ||
| 486 | + if count < 0: | ||
| 487 | + raise UnsupportedGroupedPlan( | ||
| 488 | + f"negative grouped benchmark workspace count: {count}" | ||
| 489 | + ) | ||
| 490 | + dtype_name, itemsize = _grouped_benchmark_dtype_itemsize( | ||
| 491 | + spec["dtype"] | ||
| 492 | + ) | ||
| 493 | + size, stride, num_bytes = (count,), (1,), count * itemsize | ||
| 494 | + else: | ||
| 495 | + raise UnsupportedGroupedPlan( | ||
| 496 | + f"unsupported grouped benchmark arg: kind={kind}, source={source}" | ||
| 497 | + ) | ||
| 498 | + | ||
| 499 | + current_bytes += num_bytes | ||
| 500 | + if spec.get("name") in mutated_names: | ||
| 501 | + current_mutated_bytes += num_bytes | ||
| 502 | + if dominant_arg is None or num_bytes > dominant_arg.num_bytes: | ||
| 503 | + dominant_arg = GroupedBenchmarkArgFootprint( | ||
| 504 | + group_id=int(group_id), | ||
| 505 | + name=str(spec.get("name", "unknown")), | ||
| 506 | + kind=str(kind), | ||
| 507 | + dtype=dtype_name, | ||
| 508 | + size=size, | ||
| 509 | + stride=stride, | ||
| 510 | + num_bytes=int(num_bytes), | ||
| 511 | + ) | ||
| 512 | + | ||
| 513 | + group_bytes.append((int(group_id), int(current_bytes))) | ||
| 514 | + synthetic_bytes += current_bytes | ||
| 515 | + mutated_clone_bytes = max(mutated_clone_bytes, current_mutated_bytes) | ||
| 516 | + | ||
| 517 | + largest_group_id = ( | ||
| 518 | + max(group_bytes, key=lambda item: item[1])[0] | ||
| 519 | + if group_bytes | ||
| 520 | + else None | ||
| 521 | + ) | ||
| 522 | + return GroupedBenchmarkFootprint( | ||
| 523 | + int(synthetic_bytes), | ||
| 524 | + int(mutated_clone_bytes), | ||
| 525 | + int(synthetic_bytes + mutated_clone_bytes), | ||
| 526 | + tuple(group_bytes), | ||
| 527 | + largest_group_id, | ||
| 528 | + dominant_arg, | ||
| 529 | + ) | ||
| 530 | + | ||
| 531 | + | ||
| 346 | def serialize_grouped_plan(plan: GroupedKernelMeta) -> dict[str, object]: | 532 | def serialize_grouped_plan(plan: GroupedKernelMeta) -> dict[str, object]: |
| 347 | return plan.to_payload() | 533 | return plan.to_payload() |
| 348 | 534 | ||
🟡 Medium Priority
变更行：
test_wide_backing_storage_falls_back_before_grouped_benchmark（第 761–808 行）仅 patch 了enable_symbolic_shape_group_autotune和symbolic_group_max_benchmark_memory_ratio，但未 mocktorch.npu.get_device_properties。受影响的逻辑：
_disable_grouped_autotune_if_benchmark_too_large使用torch.npu.get_device_properties(device).total_memory计算预算（total_memory * ratio）。在未 mock 的情况下，该值取决于实际 NPU 设备内存大小。失效模式：在具有大容量 NPU 内存（如 512 GB 以上）的设备上，预算（128 GB）可能超过分组 benchmark 预估的内存占用，导致
group_enabled保持为True。测试断言要求'group_enabled': False出现在生成代码中，因此测试将失败。而在小容量 NPU 上测试可能通过。这使得测试在不同硬件环境下结果不一致。建议：在
with上下文管理器中添加patch.object(torch.npu, "get_device_properties", return_value=SimpleNamespace(total_memory=108 * 1024**3))，使其与_run_benchmark_guard中的做法一致，确保测试在不同 NPU 硬件上行为确定。