已合并
perf(inductor): refine pointwise default grouping #44926
luqichao创建于 5 天前
perf(inductor): refine pointwise default grouping #44926
已合并
共 4 个文件变更+228-37
| @@ -889,23 +889,46 @@ class TestPointwiseSymbolicGrouping(TestCase): | |||
| 889 | self.assertEqual(combined[0].name, "pointwise_broadcast_axis") | 889 | self.assertEqual(combined[0].name, "pointwise_broadcast_axis") |
| 890 | self.assertEqual(combined[0].buckets, (16, 64, 256, 1024, 4096)) | 890 | self.assertEqual(combined[0].buckets, (16, 64, 256, 1024, 4096)) |
| 891 | 891 | ||
| 892 | - def test_existing_pointwise_group_feature_is_unchanged(self): | 892 | + def test_default_pointwise_group_feature(self): |
| 893 | primary_axis = make_axis("x0", sympy.Symbol("s0", positive=True)) | 893 | primary_axis = make_axis("x0", sympy.Symbol("s0", positive=True)) |
| 894 | kernel = SimpleNamespace( | 894 | kernel = SimpleNamespace( |
| 895 | persistent_reduction=False, | 895 | persistent_reduction=False, |
| 896 | inside_reduction=False, | 896 | inside_reduction=False, |
| 897 | sorted_axis=[primary_axis], | 897 | sorted_axis=[primary_axis], |
| 898 | + split_axis=[primary_axis], | ||
| 899 | + tiling_axis=[], | ||
| 900 | + get_axis_dtype=lambda axis: torch.float32, | ||
| 898 | ) | 901 | ) |
| 899 | split_tiling = object.__new__(SplitTiling) | 902 | split_tiling = object.__new__(SplitTiling) |
| 900 | split_tiling.kernel = kernel | 903 | split_tiling.kernel = kernel |
| 901 | 904 | ||
| 902 | - features = split_tiling._build_group_features(None, primary_axis) | 905 | + features = split_tiling._build_group_features( |
| 906 | + None, | ||
| 907 | + primary_axis, | ||
| 908 | + dynamic_split_axes=(primary_axis,), | ||
| 909 | + ) | ||
| 910 | + | ||
| 911 | + vector_core = int(num_vector_core) | ||
| 912 | + lower = max(1024, split_tiling_module.next_power_of_2(2 * vector_core)) | ||
| 913 | + upper = max( | ||
| 914 | + split_tiling_module.next_power_of_2(8 * vector_core), | ||
| 915 | + 4096 * vector_core, | ||
| 916 | + ) | ||
| 917 | + expected_buckets = [lower, upper] | ||
| 918 | + if upper // lower > 8: | ||
| 919 | + expected_buckets.append( | ||
| 920 | + min( | ||
| 921 | + split_tiling_module.next_power_of_2(lower * 8), | ||
| 922 | + split_tiling_module.next_power_of_2((upper + 1) // 2), | ||
| 923 | + ) | ||
| 924 | + ) | ||
| 925 | + expected_buckets = tuple(sorted(set(expected_buckets))) | ||
| 903 | 926 | ||
| 904 | self.assertEqual(len(features), 1) | 927 | self.assertEqual(len(features), 1) |
| 905 | self.assertEqual(features[0].name, "pointwise") | 928 | self.assertEqual(features[0].name, "pointwise") |
| 906 | self.assertEqual(features[0].source, "outer_product") | 929 | self.assertEqual(features[0].source, "outer_product") |
| 907 | self.assertEqual(features[0].axis_names, ("x0",)) | 930 | self.assertEqual(features[0].axis_names, ("x0",)) |
| 908 | - self.assertEqual(features[0].buckets, (num_vector_core * 4096,)) | 931 | + self.assertEqual(features[0].buckets, expected_buckets) |
| 909 | 932 | ||
| 910 | def test_plain_pointwise_does_not_use_tiling_fallback(self): | 933 | def test_plain_pointwise_does_not_use_tiling_fallback(self): |
| 911 | split_axis = make_axis("x0", sympy.Integer(128)) | 934 | split_axis = make_axis("x0", sympy.Integer(128)) |
| @@ -18,6 +18,20 @@ from ..runtime.symbolic_grouping import GroupFeatureSpec, GroupedKernelMeta | |||
| 18 | 18 | ||
| 19 | _ELEMENTWISE_UNSUPPORTED_OPS = ("masked", "scan", "sort", "rand", "randn", "load_seed") | 19 | _ELEMENTWISE_UNSUPPORTED_OPS = ("masked", "scan", "sort", "rand", "randn", "load_seed") |
| 20 | _NEUTRAL_CONSTANT_OPS = frozenset(("constant", "store", "output")) | 20 | _NEUTRAL_CONSTANT_OPS = frozenset(("constant", "store", "output")) |
| 21 | +_LARGE_GROUP_SYMBOLIC_AXIS = 128 | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +def _add_wide_group_midpoints(boundaries): | ||
| 25 | + boundaries = sorted(set(int(boundary) for boundary in boundaries)) | ||
| 26 | + expanded = [] | ||
| 27 | + for index, boundary in enumerate(boundaries): | ||
| 28 | + if index and boundary // boundaries[index - 1] > 8: | ||
| 29 | + lower = boundaries[index - 1] | ||
| 30 | + upper_midpoint = next_power_of_2((boundary + 1) // 2) | ||
| 31 | + lower_limit = next_power_of_2(lower * 8) | ||
| 32 | + expanded.append(min(lower_limit, upper_midpoint)) | ||
| 33 | + expanded.append(boundary) | ||
| 34 | + return tuple(expanded) | ||
| 21 | 35 | ||
| 22 | 36 | ||
| 23 | # split and tiling axis selector | 37 | # split and tiling axis selector |
| @@ -423,7 +437,101 @@ class SplitTiling: | |||
| 423 | _REDUCTION_BUCKETS = (8192,) | 437 | _REDUCTION_BUCKETS = (8192,) |
| 424 | _OUTER_BUCKETS = (256,) | 438 | _OUTER_BUCKETS = (256,) |
| 425 | 439 | ||
| 426 | - def _build_group_features(self, workload, primary_axis, pointwise_layout=None): | 440 | + def _axis_static_length(self, axis): |
| 441 | + try: | ||
| 442 | + return int(self.get_length_val(axis)) | ||
| 443 | + except (TypeError, ValueError): | ||
| 444 | + try: | ||
| 445 | + return int(V.graph.sizevars.size_hint(axis.length)) | ||
| 446 | + except (AttributeError, KeyError, TypeError, ValueError): | ||
| 447 | + return 1 | ||
| 448 | + | ||
| 449 | + def _default_dynamic_axis_feature( | ||
| 450 | + self, | ||
| 451 | + dynamic_split_axes, | ||
| 452 | + static_split_axes, | ||
| 453 | + feature_name="pointwise", | ||
| 454 | + fallback_dynamic_axis=None, | ||
| 455 | + ): | ||
| 456 | + # The default policy is defined for one dynamic axis. Keep unsupported | ||
| 457 | + # multi-symbol split cases out of this feature construction. | ||
| 458 | + if len(dynamic_split_axes) > 1: | ||
| 459 | + return None | ||
| 460 | + | ||
| 461 | + dynamic_axis = ( | ||
| 462 | + dynamic_split_axes[0] | ||
| 463 | + if dynamic_split_axes | ||
| 464 | + else fallback_dynamic_axis | ||
| 465 | + ) | ||
| 466 | + if dynamic_axis is None or isinstance(dynamic_axis.length, sympy.Integer): | ||
| 467 | + return None | ||
| 468 | + axis_order = { | ||
| 469 | + axis.name: index for index, axis in enumerate(self.kernel.sorted_axis) | ||
| 470 | + } | ||
| 471 | + dynamic_order = axis_order[dynamic_axis.name] | ||
| 472 | + prefix_axes = [ | ||
| 473 | + axis | ||
| 474 | + for axis in static_split_axes | ||
| 475 | + if axis_order[axis.name] < dynamic_order | ||
| 476 | + ] | ||
| 477 | + suffix_axes = [ | ||
| 478 | + axis | ||
| 479 | + for axis in static_split_axes | ||
| 480 | + if axis_order[axis.name] > dynamic_order | ||
| 481 | + ] | ||
| 482 | + | ||
| 483 | + # Static split axes retain their split role. Axes before the dynamic | ||
| 484 | + # axis are part of the grouped outer-product feature. All split axes | ||
| 485 | + # are excluded from the tiling product used to classify the dynamic | ||
| 486 | + # workload, so the resulting bounds already describe that product. | ||
| 487 | + split_axis_names = {axis.name for axis in self.kernel.split_axis} | ||
| 488 | + tiling_product = 1 | ||
| 489 | + for axis in self.kernel.tiling_axis: | ||
| 490 | + if axis.name not in split_axis_names and axis.name != dynamic_axis.name: | ||
| 491 | + tiling_product *= max(1, self._axis_static_length(axis)) | ||
| 492 | + | ||
| 493 | + vector_core = int(num_vector_core) | ||
| 494 | + dtype_axis = ( | ||
| 495 | + self.kernel.split_axis[0] | ||
| 496 | + if self.kernel.split_axis | ||
| 497 | + else dynamic_axis | ||
| 498 | + ) | ||
| 499 | + axis_dtype = self.kernel.get_axis_dtype(dtype_axis) | ||
| 500 | + dtype_bytes = max(1, get_byte_per_numel(axis_dtype)) | ||
| 501 | + base = max(1, (4096 * vector_core + tiling_product - 1) // tiling_product) | ||
| 502 | + lower = max( | ||
| 503 | + (4 * 1024) // dtype_bytes // max(1, tiling_product), | ||
| 504 | + next_power_of_2(2 * vector_core), | ||
| 505 | + ) | ||
| 506 | + prefix_product = 1 | ||
| 507 | + for axis in prefix_axes: | ||
| 508 | + prefix_product *= max(1, self._axis_static_length(axis)) | ||
| 509 | + upper = max( | ||
| 510 | + prefix_product * _LARGE_GROUP_SYMBOLIC_AXIS, | ||
| 511 | + next_power_of_2(8 * vector_core), | ||
| 512 | + base, | ||
| 513 | + ) | ||
| 514 | + boundaries = [lower, upper] | ||
| 515 | + if suffix_axes: | ||
| 516 | + boundaries.insert(0, next_power_of_2(max(1, vector_core // 2))) | ||
| 517 | + boundaries = _add_wide_group_midpoints(boundaries) | ||
| 518 | + | ||
| 519 | + feature_axis_names = tuple(axis.name for axis in (*prefix_axes, dynamic_axis)) | ||
| 520 | + return GroupFeatureSpec( | ||
| 521 | + feature_name, | ||
| 522 | + "outer_product", | ||
| 523 | + feature_axis_names, | ||
| 524 | + boundaries, | ||
| 525 | + ) | ||
| 526 | + | ||
| 527 | + def _build_group_features( | ||
| 528 | + self, | ||
| 529 | + workload, | ||
| 530 | + primary_axis, | ||
| 531 | + pointwise_layout=None, | ||
| 532 | + dynamic_split_axes=(), | ||
| 533 | + static_split_axes=(), | ||
| 534 | + ): | ||
| 427 | if self.kernel.persistent_reduction or self.kernel.inside_reduction: | 535 | if self.kernel.persistent_reduction or self.kernel.inside_reduction: |
| 428 | outer_names = self.non_reduction_axis_names() | 536 | outer_names = self.non_reduction_axis_names() |
| 429 | reduction_names = self.reduction_axis_names() | 537 | reduction_names = self.reduction_axis_names() |
| @@ -449,6 +557,14 @@ class SplitTiling: | |||
| 449 | ) | 557 | ) |
| 450 | return tuple(features) | 558 | return tuple(features) |
| 451 | if workload == "elementwise": | 559 | if workload == "elementwise": |
| 560 | + default_feature = self._default_dynamic_axis_feature( | ||
| 561 | + dynamic_split_axes, | ||
| 562 | + static_split_axes, | ||
| 563 | + feature_name="elementwise_numel", | ||
| 564 | + fallback_dynamic_axis=primary_axis, | ||
| 565 | + ) | ||
| 566 | + if default_feature is not None: | ||
| 567 | + return (default_feature,) | ||
| 452 | return ( | 568 | return ( |
| 453 | GroupFeatureSpec( | 569 | GroupFeatureSpec( |
| 454 | "elementwise_numel", | 570 | "elementwise_numel", |
| @@ -475,14 +591,10 @@ class SplitTiling: | |||
| 475 | (64, 128, 256, 512), | 591 | (64, 128, 256, 512), |
| 476 | ), | 592 | ), |
| 477 | ) | 593 | ) |
| 478 | - return ( | 594 | + default_feature = self._default_dynamic_axis_feature( |
| 479 | - GroupFeatureSpec( | 595 | + dynamic_split_axes, static_split_axes |
| 480 | - "pointwise", | ||
| 481 | - "outer_product", | ||
| 482 | - self.all_axis_names(), | ||
| 483 | - (num_vector_core * 4096,), | ||
| 484 | - ), | ||
| 485 | ) | 596 | ) |
| 597 | + return (default_feature,) if default_feature is not None else () | ||
| 486 | 598 | ||
| 487 | 599 | ||
| 488 | def _alpha_rename_access_vars(dep): | 600 | def _alpha_rename_access_vars(dep): |
| @@ -643,7 +755,11 @@ class SplitTiling: | |||
| 643 | if primary_axis is None: | 755 | if primary_axis is None: |
| 644 | return None | 756 | return None |
| 645 | secondary_axes = [axis for axis in dynamic_split_axes if axis is not primary_axis] | 757 | secondary_axes = [axis for axis in dynamic_split_axes if axis is not primary_axis] |
| 646 | - self._downgrade_secondary_runtime_split_axes(secondary_axes) | 758 | + self._downgrade_split_axes(secondary_axes) |
| 759 | + if template == "pointwise" and len(dynamic_split_axes) == 1: | ||
| 760 | + static_split_axes = self._downgrade_suffix_static_split_axes( | ||
| 761 | + primary_axis, static_split_axes | ||
| 762 | + ) | ||
| 647 | static_names = tuple(axis.name for axis in static_split_axes) | 763 | static_names = tuple(axis.name for axis in static_split_axes) |
| 648 | secondary_names = tuple(axis.name for axis in secondary_axes) | 764 | secondary_names = tuple(axis.name for axis in secondary_axes) |
| 649 | runtime_block_arg_names = tuple( | 765 | runtime_block_arg_names = tuple( |
| @@ -663,7 +779,11 @@ class SplitTiling: | |||
| 663 | f"{axis.name.upper()}BLOCK" for axis in self.kernel.split_axis | 779 | f"{axis.name.upper()}BLOCK" for axis in self.kernel.split_axis |
| 664 | ) | 780 | ) |
| 665 | feature_specs = self._build_group_features( | 781 | feature_specs = self._build_group_features( |
| 666 | - workload, primary_axis, pointwise_layout | 782 | + workload, |
| 783 | + primary_axis, | ||
| 784 | + pointwise_layout, | ||
| 785 | + dynamic_split_axes, | ||
| 786 | + static_split_axes, | ||
| 667 | ) | 787 | ) |
| 668 | if not feature_specs: | 788 | if not feature_specs: |
| 669 | return None | 789 | return None |
| @@ -678,12 +798,12 @@ class SplitTiling: | |||
| 678 | runtime_block_arg_names=runtime_block_arg_names, | 798 | runtime_block_arg_names=runtime_block_arg_names, |
| 679 | ) | 799 | ) |
| 680 | 800 | ||
| 681 | - def _downgrade_secondary_runtime_split_axes(self, secondary_axes): | 801 | + def _downgrade_split_axes(self, axes): |
| 682 | - if not secondary_axes: | 802 | + if not axes: |
| 683 | return | 803 | return |
| 684 | retained = [] | 804 | retained = [] |
| 685 | for axis in self.kernel.split_axis: | 805 | for axis in self.kernel.split_axis: |
| 686 | - if axis in secondary_axes: | 806 | + if axis in axes: |
| 687 | axis.is_split_axis = False | 807 | axis.is_split_axis = False |
| 688 | continue | 808 | continue |
| 689 | retained.append(axis) | 809 | retained.append(axis) |
| @@ -692,6 +812,36 @@ class SplitTiling: | |||
| 692 | for i, axis in enumerate(self.kernel.split_axis): | 812 | for i, axis in enumerate(self.kernel.split_axis): |
| 693 | axis.split_order = i | 813 | axis.split_order = i |
| 694 | 814 | ||
| 815 | + def _downgrade_suffix_static_split_axes(self, dynamic_axis, static_split_axes): | ||
| 816 | + axis_order = { | ||
| 817 | + axis.name: index for index, axis in enumerate(self.kernel.sorted_axis) | ||
| 818 | + } | ||
| 819 | + dynamic_order = axis_order[dynamic_axis.name] | ||
| 820 | + split_axes_through_dynamic = [ | ||
| 821 | + axis | ||
| 822 | + for axis in self.kernel.split_axis | ||
| 823 | + if axis_order[axis.name] <= dynamic_order | ||
| 824 | + ] | ||
| 825 | + try: | ||
| 826 | + split_size_hint = V.graph.sizevars.size_hint( | ||
| 827 | + self.total_split_numels(split_axes_through_dynamic) | ||
| 828 | + ) | ||
| 829 | + except TypeError: | ||
| 830 | + return static_split_axes | ||
| 831 | + if split_size_hint < num_vector_core: | ||
| 832 | + return static_split_axes | ||
| 833 | + | ||
| 834 | + suffix_axes = [ | ||
| 835 | + axis | ||
| 836 | + for axis in static_split_axes | ||
| 837 | + if axis_order[axis.name] > dynamic_order | ||
| 838 | + ] | ||
| 839 | + self._downgrade_split_axes(suffix_axes) | ||
| 840 | + suffix_names = {axis.name for axis in suffix_axes} | ||
| 841 | + return [ | ||
| 842 | + axis for axis in static_split_axes if axis.name not in suffix_names | ||
| 843 | + ] | ||
| 844 | + | ||
| 695 | # the below logic doesn't work when there're two reduction axis, but only one need outer reduction | 845 | # the below logic doesn't work when there're two reduction axis, but only one need outer reduction |
| 696 | def should_outer_reduce_me(self, x): | 846 | def should_outer_reduce_me(self, x): |
| 697 | should_outer = ( | 847 | should_outer = ( |
| @@ -237,17 +237,6 @@ def build_group_representatives( | |||
| 237 | return lower * 2 | 237 | return lower * 2 |
| 238 | return upper | 238 | return upper |
| 239 | 239 | ||
| 240 | - def choose_symbolic_axis_value(static_factor: int, bucket_idx: int, feature_spec): | ||
| 241 | - lower, upper = bucket_bounds(feature_spec, bucket_idx) | ||
| 242 | - dyn_min = lower // static_factor + 1 | ||
| 243 | - if upper is None: | ||
| 244 | - target = max(lower + 1, lower * 2) | ||
| 245 | - return max(1, (target + static_factor - 1) // static_factor) | ||
| 246 | - dyn_max = upper // static_factor | ||
| 247 | - if dyn_min > dyn_max: | ||
| 248 | - return None | ||
| 249 | - return max(1, dyn_max) | ||
| 250 | - | ||
| 251 | def representative_for_feature(feature_spec, bucket_idx: int): | 240 | def representative_for_feature(feature_spec, bucket_idx: int): |
| 252 | feature_axis_names = tuple(_feature_field(feature_spec, "axis_names")) | 241 | feature_axis_names = tuple(_feature_field(feature_spec, "axis_names")) |
| 253 | source = _feature_field(feature_spec, "source") | 242 | source = _feature_field(feature_spec, "source") |
| @@ -291,11 +280,9 @@ def build_group_representatives( | |||
| 291 | return None | 280 | return None |
| 292 | return feature_value, tuple(static_axis_values) | 281 | return feature_value, tuple(static_axis_values) |
| 293 | 282 | ||
| 294 | - symbolic_axis_value = choose_symbolic_axis_value( | 283 | + lower, upper = bucket_bounds(feature_spec, bucket_idx) |
| 295 | - static_factor, bucket_idx, feature_spec | 284 | + target = upper if upper is not None else max(lower + 1, lower * 2) |
| 296 | - ) | 285 | + symbolic_axis_value = max(1, (target + static_factor - 1) // static_factor) |
| 297 | - if symbolic_axis_value is None: | ||
| 298 | - return None | ||
| 299 | axis_values = [ | 286 | axis_values = [ |
| 300 | (axis_name, symbolic_axis_value) | 287 | (axis_name, symbolic_axis_value) |
| 301 | if axis_name == symbolic_axis_names[0] | 288 | if axis_name == symbolic_axis_names[0] |
| @@ -2283,6 +2283,8 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner): | |||
| 2283 | axis_numel=axis_numel, | 2283 | axis_numel=axis_numel, |
| 2284 | expected_grid=grid_target, | 2284 | expected_grid=grid_target, |
| 2285 | block_sub=int(rule["block_sub"]), | 2285 | block_sub=int(rule["block_sub"]), |
| 2286 | + representative_numel=rule.get("representative_numel"), | ||
| 2287 | + max_grid=rule.get("max_grid"), | ||
| 2286 | ) | 2288 | ) |
| 2287 | missing_runtime_block_names = [ | 2289 | missing_runtime_block_names = [ |
| 2288 | name for name in runtime_block_names if name not in resolved_blocks | 2290 | name for name in runtime_block_names if name not in resolved_blocks |
| @@ -3209,17 +3211,31 @@ def build_grouped_launch_policy( | |||
| 3209 | // primary_static_block, | 3211 | // primary_static_block, |
| 3210 | ) | 3212 | ) |
| 3211 | 3213 | ||
| 3214 | + primary_feature = group_features[primary_feature_index] | ||
| 3215 | + if isinstance(primary_feature, dict): | ||
| 3216 | + primary_feature_name = primary_feature.get("name") | ||
| 3217 | + else: | ||
| 3218 | + primary_feature_name = getattr(primary_feature, "name", None) | ||
| 3219 | + runtime_rule = [ | ||
| 3220 | + ("op", "ceildiv"), | ||
| 3221 | + ("axis_name", primary_group_axis), | ||
| 3222 | + ("block_sub", primary_block_sub), | ||
| 3223 | + ] | ||
| 3224 | + if primary_feature_name in ("pointwise", "elementwise_numel"): | ||
| 3225 | + runtime_rule.extend( | ||
| 3226 | + ( | ||
| 3227 | + ("representative_numel", int(axis_env[primary_group_axis])), | ||
| 3228 | + ("max_grid", max(1, npu_num_vector_core // prior_programs)), | ||
| 3229 | + ) | ||
| 3230 | + ) | ||
| 3231 | + | ||
| 3212 | return { | 3232 | return { |
| 3213 | "group_id": group_id, | 3233 | "group_id": group_id, |
| 3214 | "static_blocks": tuple(static_blocks), | 3234 | "static_blocks": tuple(static_blocks), |
| 3215 | "runtime_block_rules": ( | 3235 | "runtime_block_rules": ( |
| 3216 | ( | 3236 | ( |
| 3217 | primary_block_name, | 3237 | primary_block_name, |
| 3218 | - ( | 3238 | + tuple(runtime_rule), |
| 3219 | - ("op", "ceildiv"), | ||
| 3220 | - ("axis_name", primary_group_axis), | ||
| 3221 | - ("block_sub", primary_block_sub), | ||
| 3222 | - ), | ||
| 3223 | ), | 3239 | ), |
| 3224 | ), | 3240 | ), |
| 3225 | "grid_target": primary_program_target, | 3241 | "grid_target": primary_program_target, |
| @@ -3230,6 +3246,8 @@ def resolve_grouped_runtime_block( | |||
| 3230 | axis_numel: int, | 3246 | axis_numel: int, |
| 3231 | expected_grid: int, | 3247 | expected_grid: int, |
| 3232 | block_sub: int, | 3248 | block_sub: int, |
| 3249 | + representative_numel: int | None = None, | ||
| 3250 | + max_grid: int | None = None, | ||
| 3233 | ) -> int: | 3251 | ) -> int: |
| 3234 | axis_numel = int(axis_numel) | 3252 | axis_numel = int(axis_numel) |
| 3235 | expected_grid = int(expected_grid) | 3253 | expected_grid = int(expected_grid) |
| @@ -3242,6 +3260,19 @@ def resolve_grouped_runtime_block( | |||
| 3242 | return 1 | 3260 | return 1 |
| 3243 | 3261 | ||
| 3244 | total_subblock_num = (axis_numel + block_sub - 1) // block_sub | 3262 | total_subblock_num = (axis_numel + block_sub - 1) // block_sub |
| 3263 | + if ( | ||
| 3264 | + representative_numel is not None | ||
| 3265 | + and max_grid is not None | ||
| 3266 | + and axis_numel < int(representative_numel) | ||
| 3267 | + ): | ||
| 3268 | + # Smaller shapes may use more programs than the representative while | ||
| 3269 | + # keeping BLOCK aligned to the compile-time BLOCK_SUB. | ||
| 3270 | + max_grid = max(1, int(max_grid)) | ||
| 3271 | + program_subblock_num = ( | ||
| 3272 | + total_subblock_num + max_grid - 1 | ||
| 3273 | + ) // max_grid | ||
| 3274 | + return program_subblock_num * block_sub | ||
| 3275 | + | ||
| 3245 | program_subblock_num = ( | 3276 | program_subblock_num = ( |
| 3246 | total_subblock_num + expected_grid - 1 | 3277 | total_subblock_num + expected_grid - 1 |
| 3247 | ) // expected_grid | 3278 | ) // expected_grid |