已合并
[Inductor] Symbolic dynamic-shape grouping for reductions #42292
dezheng889创建于 7月21日
[Inductor] Symbolic dynamic-shape grouping for reductions #42292
已合并
共 5 个文件变更+225-44
| @@ -61,6 +61,41 @@ class TestSum(TestUtils): | |||
| 61 | 61 | ||
| 62 | self.assertEqual(std_sum, inductor_sum, atol=1e-1, rtol=1e-1) | 62 | self.assertEqual(std_sum, inductor_sum, atol=1e-1, rtol=1e-1) |
| 63 | 63 | ||
| 64 | + # Sizes straddle bucket boundaries plus non-power-of-2 / large, to exercise | ||
| 65 | + # the runtime loop bound + tail mask. | ||
| 66 | + _sum_1d_dynamic_sizes = [8, 63, 64, 65, 255, 257, 1024, 1025, 8191, 8193, 50000, 999983] | ||
| 67 | + | ||
| 68 | + | ||
| 69 | + def test_sum_1d_dynamic_shape(self, dtype): | ||
| 70 | + # One dynamic kernel must stay correct across all sizes without recompiling. | ||
| 71 | + torch_dtype = eval('torch.' + dtype) | ||
| 72 | + torch._dynamo.reset() | ||
| 73 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor", dynamic=True) | ||
| 74 | + for n in self._sum_1d_dynamic_sizes: | ||
| 75 | + x = torch.ones((n,), dtype=torch_dtype, device=torch.device("npu")) | ||
| 76 | + std_sum = self.op_calc(x, None) | ||
| 77 | + inductor_sum = compiled_op_calc(x, None) | ||
| 78 | + self.assertEqual(std_sum, inductor_sum) | ||
| 79 | + | ||
| 80 | + | ||
| 81 | + def test_sum_1d_dynamic_shape_group_autotune(self, dtype): | ||
| 82 | + # Same, with symbolic group-autotune (bucketed path) enabled. | ||
| 83 | + import torch_npu._inductor.config as npu_config | ||
| 84 | + torch_dtype = eval('torch.' + dtype) | ||
| 85 | + prev = npu_config.enable_symbolic_shape_group_autotune | ||
| 86 | + npu_config.enable_symbolic_shape_group_autotune = True | ||
| 87 | + try: | ||
| 88 | + torch._dynamo.reset() | ||
| 89 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor", dynamic=True) | ||
| 90 | + for n in self._sum_1d_dynamic_sizes: | ||
| 91 | + x = torch.ones((n,), dtype=torch_dtype, device=torch.device("npu")) | ||
| 92 | + std_sum = self.op_calc(x, None) | ||
| 93 | + inductor_sum = compiled_op_calc(x, None) | ||
| 94 | + self.assertEqual(std_sum, inductor_sum) | ||
| 95 | + finally: | ||
| 96 | + npu_config.enable_symbolic_shape_group_autotune = prev | ||
| 97 | + torch._dynamo.reset() | ||
| 98 | + | ||
| 64 | 99 | ||
| 65 | instantiate_parametrized_tests(TestSum) | 100 | instantiate_parametrized_tests(TestSum) |
| 66 | 101 | ||
| @@ -1132,6 +1132,54 @@ def should_skip_linearization_on_a5(var_ranges, indexing): | |||
| 1132 | return False | 1132 | return False |
| 1133 | 1133 | ||
| 1134 | 1134 | ||
| 1135 | +# Reduction types that produce an index or are multi-output/Welford. They have | ||
| 1136 | +# never been validated on simt_template and miscompute there (e.g. wrong argmax | ||
| 1137 | +# index for a dynamic-length reduction axis), so they must fall back to the | ||
| 1138 | +# community SIMT_ONLY path (same as static SIMD and the community dynamic | ||
| 1139 | +# fallback). reduction_type is recorded as a bare string arg on the reduction | ||
| 1140 | +# node, so match by value rather than a fragile positional index. | ||
| 1141 | +_SIMT_TEMPLATE_UNSAFE_REDUCTION_TYPES = { | ||
| 1142 | + "argmax", | ||
| 1143 | + "argmin", | ||
| 1144 | + "welford_reduce", | ||
| 1145 | + "welford_combine", | ||
| 1146 | +} | ||
| 1147 | + | ||
| 1148 | + | ||
| 1149 | +def _loop_body_has_unsafe_reduction(loop_body): | ||
| 1150 | + graphs = [loop_body.root_block.graph] | ||
| 1151 | + graphs.extend(sub.graph for sub in getattr(loop_body, "subblocks", {}).values()) | ||
| 1152 | + for graph in graphs: | ||
| 1153 | + for node in graph.nodes: | ||
| 1154 | + if "reduction" not in node.name: | ||
| 1155 | + continue | ||
| 1156 | + for arg in node.args: | ||
| 1157 | + if isinstance(arg, str) and arg in _SIMT_TEMPLATE_UNSAFE_REDUCTION_TYPES: | ||
| 1158 | + return True | ||
| 1159 | + return False | ||
| 1160 | + | ||
| 1161 | + | ||
| 1162 | +def is_linear_dynamic_reduction(loop_body): | ||
| 1163 | + """Keep a dynamic reduction on simt_template only when it is safe to do so. | ||
| 1164 | + Bail (fall back to community SIMT_ONLY) when an index carries a SYMBOLIC bound | ||
| 1165 | + in ModularIndexing/FloorDiv (breaks linearization), or when the reduction | ||
| 1166 | + produces an index / is Welford (argmax/argmin from max.dim/min.dim, var/std). | ||
| 1167 | + Static-bound modular/floordiv from folding static dims are fine; the step-3 | ||
| 1168 | + can_split_all check is the final fallback.""" | ||
| 1169 | + kernel = V.kernel | ||
| 1170 | + if kernel is None or not getattr(kernel, "inside_reduction", False): | ||
| 1171 | + return False | ||
| 1172 | + indexing = loop_body.indexing | ||
| 1173 | + if not indexing: | ||
| 1174 | + return False | ||
| 1175 | + for index_expr in indexing.values(): | ||
| 1176 | + if check_subexpr_for_dynamic_symbols(index_expr): | ||
| 1177 | + return False | ||
| 1178 | + if _loop_body_has_unsafe_reduction(loop_body): | ||
| 1179 | + return False | ||
| 1180 | + return True | ||
| 1181 | + | ||
| 1182 | + | ||
| 1135 | def transform_dims_in_indexing(self, indices): | 1183 | def transform_dims_in_indexing(self, indices): |
| 1136 | # Step 1: Generate basic indexing (always executed) | 1184 | # Step 1: Generate basic indexing (always executed) |
| 1137 | if self.indexing is None: | 1185 | if self.indexing is None: |
| @@ -1140,12 +1188,16 @@ def transform_dims_in_indexing(self, indices): | |||
| 1140 | 1188 | ||
| 1141 | # Step 2: Check for dynamic shapes (only skip on A5 platform) | 1189 | # Step 2: Check for dynamic shapes (only skip on A5 platform) |
| 1142 | if should_skip_linearization_on_a5(self.var_ranges, self.indexing): | 1190 | if should_skip_linearization_on_a5(self.var_ranges, self.indexing): |
| 1143 | - log.info( | 1191 | + # A linearizable plain-value dynamic reduction stays on simt_template |
| 1144 | - "Skip memory access linearization due to dynamic shape on A5, var_ranges: %s", | 1192 | + # (full UB tile); indexed/Welford reductions and everything else fall |
| 1145 | - self.var_ranges, | 1193 | + # back below. |
| 1146 | - ) | 1194 | + if not is_linear_dynamic_reduction(self): |
| 1147 | - # Set SIMT_ONLY compile option for dynamic shapes on A5 | 1195 | + log.info( |
| 1148 | - raise ValueError(f"fallback to community simt for SIMT_ONLY compile option for dynamic shapes on A5") | 1196 | + "Skip memory access linearization due to dynamic shape on A5, var_ranges: %s", |
| 1197 | + self.var_ranges, | ||
| 1198 | + ) | ||
| 1199 | + # Set SIMT_ONLY compile option for dynamic shapes on A5 | ||
| 1200 | + raise ValueError(f"fallback to community simt for SIMT_ONLY compile option for dynamic shapes on A5") | ||
| 1149 | 1201 | ||
| 1150 | # Step 3: Perform memory access linearization | 1202 | # Step 3: Perform memory access linearization |
| 1151 | log.debug( | 1203 | log.debug( |
| @@ -128,8 +128,24 @@ class SplitTiling: | |||
| 128 | break | 128 | break |
| 129 | 129 | ||
| 130 | if not self.kernel.split_axis and self.kernel.sorted_axis: | 130 | if not self.kernel.split_axis and self.kernel.sorted_axis: |
| 131 | - self.kernel.split_axis.append(self.kernel.sorted_axis[0]) | 131 | + only_axis = self.kernel.sorted_axis[0] |
| 132 | - self.kernel.sorted_axis[0].is_split_axis = True | 132 | + # A dynamic full reduction (every axis is a reduction axis, at least one |
| 133 | + # dynamic) must keep grid==1: making a reduction axis a grid split axis | ||
| 134 | + # gives grid>1 with no cross-core combine (A5), which overwrites the | ||
| 135 | + # scalar output. Leave split_axis empty (loop the runtime numel with a | ||
| 136 | + # tail mask). Covers pure 1D ([n]) and full reduce with extra static | ||
| 137 | + # reduction dims ([n, 1024]). | ||
| 138 | + all_reduction = all( | ||
| 139 | + axis.prefix == "r" for axis in self.kernel.sorted_axis | ||
| 140 | + ) | ||
| 141 | + any_dynamic = any( | ||
| 142 | + not isinstance(axis.length, sympy.Integer) | ||
| 143 | + for axis in self.kernel.sorted_axis | ||
| 144 | + ) | ||
| 145 | + is_dynamic_full_reduction = all_reduction and any_dynamic | ||
| 146 | + if not is_dynamic_full_reduction: | ||
| 147 | + self.kernel.split_axis.append(only_axis) | ||
| 148 | + only_axis.is_split_axis = True | ||
| 133 | 149 | ||
| 134 | self.kernel.split_axis.sort(reverse=True, key=self.key) | 150 | self.kernel.split_axis.sort(reverse=True, key=self.key) |
| 135 | for i, x in enumerate(self.kernel.split_axis): | 151 | for i, x in enumerate(self.kernel.split_axis): |
| @@ -277,6 +293,23 @@ class SplitTiling: | |||
| 277 | return axis | 293 | return axis |
| 278 | return dynamic_split_axes[0] | 294 | return dynamic_split_axes[0] |
| 279 | 295 | ||
| 296 | + def _dynamic_reduction_tiling_axis(self): | ||
| 297 | + """Sole dynamic reduction axis to bucket by runtime size while it stays a | ||
| 298 | + tiling axis (grid==1 over it). The reduction axis must NOT itself be a grid | ||
| 299 | + split axis (A5 has no cross-core combine -> grid>1 would overwrite the | ||
| 300 | + output); static non-reduction split axes are fine and drive the grid.""" | ||
| 301 | + reduction_dynamic = [ | ||
| 302 | + axis | ||
| 303 | + for axis in self.kernel.sorted_axis | ||
| 304 | + if axis.prefix == "r" and not isinstance(axis.length, sympy.Integer) | ||
| 305 | + ] | ||
| 306 | + if len(reduction_dynamic) != 1: | ||
| 307 | + return None | ||
| 308 | + axis = reduction_dynamic[0] | ||
| 309 | + if axis in self.kernel.split_axis: | ||
| 310 | + return None | ||
| 311 | + return axis | ||
| 312 | + | ||
| 280 | def non_reduction_axis_names(self): | 313 | def non_reduction_axis_names(self): |
| 281 | return tuple(axis.name for axis in self.kernel.sorted_axis if axis.prefix != "r") | 314 | return tuple(axis.name for axis in self.kernel.sorted_axis if axis.prefix != "r") |
| 282 | 315 | ||
| @@ -286,22 +319,49 @@ class SplitTiling: | |||
| 286 | def reduction_axis_names(self): | 319 | def reduction_axis_names(self): |
| 287 | return tuple(axis.name for axis in self.kernel.sorted_axis if axis.prefix == "r") | 320 | return tuple(axis.name for axis in self.kernel.sorted_axis if axis.prefix == "r") |
| 288 | 321 | ||
| 322 | + def _has_dynamic_axis(self, axis_names): | ||
| 323 | + names = set(axis_names) | ||
| 324 | + for axis in self.kernel.sorted_axis: | ||
| 325 | + if axis.name in names and not isinstance(axis.length, sympy.Integer): | ||
| 326 | + return True | ||
| 327 | + return False | ||
| 328 | + | ||
| 329 | + # Bucket boundaries on the group PRODUCT (reduction_product / outer_product). | ||
| 330 | + # A closed bucket tunes at its upper bound and the open tail at (max boundary * | ||
| 331 | + # 2). Kept to a single coarse boundary each. Adding closed buckets to the OUTER | ||
| 332 | + # (grid) axis backfires: the outer axis bakes XBLOCK from the representative, so | ||
| 333 | + # a mid-of-bucket runtime size (e.g. 257 in (256, 4096]) gets an XBLOCK tuned | ||
| 334 | + # for 4096 -> too few programs -> core underutilization. The single (256,) tail | ||
| 335 | + # (representative 512) is a better all-round compromise. Reduction likewise | ||
| 336 | + # stays a single boundary; the ~1e6 1D corner is out of scope. | ||
| 337 | + _REDUCTION_BUCKETS = (8192,) | ||
| 338 | + _OUTER_BUCKETS = (256,) | ||
| 339 | + | ||
| 289 | def _build_group_features(self, primary_axis): | 340 | def _build_group_features(self, primary_axis): |
| 290 | if self.kernel.persistent_reduction or self.kernel.inside_reduction: | 341 | if self.kernel.persistent_reduction or self.kernel.inside_reduction: |
| 291 | - return ( | 342 | + outer_names = self.non_reduction_axis_names() |
| 292 | - GroupFeatureSpec( | 343 | + reduction_names = self.reduction_axis_names() |
| 293 | - "outer", | 344 | + # Emit a feature only for a group with a dynamic axis: a static group |
| 294 | - "outer_product", | 345 | + # has a constant product (one reachable bucket), so bucketing it just |
| 295 | - self.non_reduction_axis_names(), | 346 | + # adds unreachable group ids. Works for the symbolic dim on either |
| 296 | - (256,), | 347 | + # side. Reduction axis stays grid==1; grid comes from outer split axes. |
| 297 | - ), | 348 | + features = [] |
| 298 | - GroupFeatureSpec( | 349 | + if outer_names and self._has_dynamic_axis(outer_names): |
| 299 | - "reduction", | 350 | + features.append( |
| 300 | - "reduction_product", | 351 | + GroupFeatureSpec( |
| 301 | - self.reduction_axis_names(), | 352 | + "outer", "outer_product", outer_names, self._OUTER_BUCKETS |
| 302 | - (8192,), | 353 | + ) |
| 303 | - ), | 354 | + ) |
| 304 | - ) | 355 | + if reduction_names and self._has_dynamic_axis(reduction_names): |
| 356 | + features.append( | ||
| 357 | + GroupFeatureSpec( | ||
| 358 | + "reduction", | ||
| 359 | + "reduction_product", | ||
| 360 | + reduction_names, | ||
| 361 | + self._REDUCTION_BUCKETS, | ||
| 362 | + ) | ||
| 363 | + ) | ||
| 364 | + return tuple(features) | ||
| 305 | return ( | 365 | return ( |
| 306 | GroupFeatureSpec( | 366 | GroupFeatureSpec( |
| 307 | "pointwise", | 367 | "pointwise", |
| @@ -313,24 +373,42 @@ class SplitTiling: | |||
| 313 | 373 | ||
| 314 | def _build_grouped_meta(self): | 374 | def _build_grouped_meta(self): |
| 315 | dynamic_split_axes, static_split_axes = self._classify_split_axes() | 375 | dynamic_split_axes, static_split_axes = self._classify_split_axes() |
| 316 | - if not dynamic_split_axes: | 376 | + if dynamic_split_axes: |
| 317 | - return None | 377 | + primary_axis = self._select_primary_group_axis(dynamic_split_axes) |
| 318 | - primary_axis = self._select_primary_group_axis(dynamic_split_axes) | 378 | + if primary_axis is None: |
| 319 | - if primary_axis is None: | 379 | + return None |
| 320 | - return None | 380 | + secondary_axes = [axis for axis in dynamic_split_axes if axis is not primary_axis] |
| 321 | - secondary_axes = [axis for axis in dynamic_split_axes if axis is not primary_axis] | 381 | + self._downgrade_secondary_runtime_split_axes(secondary_axes) |
| 322 | - self._downgrade_secondary_runtime_split_axes(secondary_axes) | 382 | + static_names = tuple(axis.name for axis in static_split_axes) |
| 383 | + secondary_names = tuple(axis.name for axis in secondary_axes) | ||
| 384 | + runtime_block_arg_names = tuple( | ||
| 385 | + f"{axis.name.upper()}BLOCK" for axis in self.kernel.split_axis | ||
| 386 | + ) | ||
| 387 | + else: | ||
| 388 | + # No dynamic grid split axis: bucket the reduction tiling axis by its | ||
| 389 | + # runtime size (grid==1 over it). Grid parallelism, if any, comes from | ||
| 390 | + # the STATIC non-reduction split axes; their blocks are passed so the | ||
| 391 | + # grid is computed from them (build_grouped_launch_policy treats a | ||
| 392 | + # reduction-tiling primary specially -- no runtime rule for it). | ||
| 393 | + primary_axis = self._dynamic_reduction_tiling_axis() | ||
| 394 | + if primary_axis is None: | ||
| 395 | + return None | ||
| 396 | + static_names = tuple(axis.name for axis in static_split_axes) | ||
| 397 | + secondary_names = () | ||
| 398 | + runtime_block_arg_names = tuple( | ||
| 399 | + f"{axis.name.upper()}BLOCK" for axis in self.kernel.split_axis | ||
| 400 | + ) | ||
| 323 | feature_specs = self._build_group_features(primary_axis) | 401 | feature_specs = self._build_group_features(primary_axis) |
| 402 | + if not feature_specs: | ||
| 403 | + return None | ||
| 324 | return GroupedKernelMeta( | 404 | return GroupedKernelMeta( |
| 325 | enabled=True, | 405 | enabled=True, |
| 326 | template=self._grouped_template_name(), | 406 | template=self._grouped_template_name(), |
| 327 | primary_group_axis=primary_axis.name, | 407 | primary_group_axis=primary_axis.name, |
| 328 | - static_split_axes=tuple(axis.name for axis in static_split_axes), | 408 | + static_split_axes=static_names, |
| 329 | - secondary_runtime_symbolic_axes=tuple(axis.name for axis in secondary_axes), | 409 | + secondary_runtime_symbolic_axes=secondary_names, |
| 330 | group_features=tuple(feature_specs), | 410 | group_features=tuple(feature_specs), |
| 331 | - runtime_block_arg_names=tuple( | 411 | + runtime_block_arg_names=runtime_block_arg_names, |
| 332 | - f"{axis.name.upper()}BLOCK" for axis in self.kernel.split_axis | ||
| 333 | - ), | ||
| 334 | ) | 412 | ) |
| 335 | 413 | ||
| 336 | def _downgrade_secondary_runtime_split_axes(self, secondary_axes): | 414 | def _downgrade_secondary_runtime_split_axes(self, secondary_axes): |
| @@ -119,6 +119,15 @@ def num_splits( | |||
| 119 | return 1 | 119 | return 1 |
| 120 | 120 | ||
| 121 | if numel_hint == 1: | 121 | if numel_hint == 1: |
| 122 | + # A reduction whose reduction group has exactly ONE dynamic axis (any other | ||
| 123 | + # reduction axes static): keep a single Reduction node (split==1) instead of | ||
| 124 | + # a degenerate multilayer split; grid is forced to 1 in select_split_axis so | ||
| 125 | + # it loops the runtime numel with a tail mask. Covers pure 1D | ||
| 126 | + # (reduction_ranges == [n]) and full reduce with extra static reduction dims | ||
| 127 | + # (e.g. [n, 1024] -> reduction_ranges == [n, 1024]). | ||
| 128 | + num_dynamic_reduction = sum(1 for r in reduction_ranges if not _is_static(r)) | ||
| 129 | + if num_dynamic_reduction == 1: | ||
| 130 | + return ReductionHint.INNER, 1 | ||
| 122 | split = inner_reduction_splits(reduction_ranges) | 131 | split = inner_reduction_splits(reduction_ranges) |
| 123 | return ReductionHint.INNER, split | 132 | return ReductionHint.INNER, split |
| 124 | return ReductionHint.DEFAULT, 1 | 133 | return ReductionHint.DEFAULT, 1 |
| @@ -2752,10 +2752,11 @@ def _triton_config_npu_index_grouped( | |||
| 2752 | if primary_group_axis is None: | 2752 | if primary_group_axis is None: |
| 2753 | raise RuntimeError("grouped autotune plan is missing primary_group_axis") | 2753 | raise RuntimeError("grouped autotune plan is missing primary_group_axis") |
| 2754 | primary_block_name = f"{primary_group_axis.upper()}BLOCK" | 2754 | primary_block_name = f"{primary_group_axis.upper()}BLOCK" |
| 2755 | - if runtime_block_arg_names and primary_block_name not in runtime_block_arg_names: | 2755 | + # primary_block_name may legitimately be absent from runtime_block_arg_names |
| 2756 | - raise RuntimeError( | 2756 | + # when the primary group axis is a reduction TILING axis (bucketed via |
| 2757 | - f"runtime_block_arg_names is missing primary block {primary_block_name}" | 2757 | + # R0BLOCK_SUB per variant, no grid block); in that case any grid comes from |
| 2758 | - ) | 2758 | + # the static non-reduction split axes. build_grouped_launch_policy handles |
| 2759 | + # both shapes, so no guard is needed here. | ||
| 2759 | 2760 | ||
| 2760 | def benchmark_axis_env(group_id: int) -> dict[str, int]: | 2761 | def benchmark_axis_env(group_id: int) -> dict[str, int]: |
| 2761 | axis_values = group_representatives["benchmark_axis_values_by_group"][group_id] | 2762 | axis_values = group_representatives["benchmark_axis_values_by_group"][group_id] |
| @@ -2989,14 +2990,20 @@ def build_grouped_launch_policy( | |||
| 2989 | ) -> dict[str, object]: | 2990 | ) -> dict[str, object]: |
| 2990 | primary_block_name = f"{primary_group_axis.upper()}BLOCK" | 2991 | primary_block_name = f"{primary_group_axis.upper()}BLOCK" |
| 2991 | runtime_blocks = extract_runtime_blocks_from_cfg(cfg, runtime_block_arg_names) | 2992 | runtime_blocks = extract_runtime_blocks_from_cfg(cfg, runtime_block_arg_names) |
| 2992 | - if primary_block_name not in runtime_blocks: | 2993 | + if primary_block_name not in runtime_block_arg_names: |
| 2993 | - if runtime_block_arg_names: | 2994 | + # Primary is a reduction TILING axis (bucketed via R0BLOCK_SUB per variant), |
| 2994 | - raise RuntimeError( | 2995 | + # not a grid split axis, so it has no grid block. Grid==1 over it; any grid |
| 2995 | - f"legacy grouped config is missing primary block {primary_block_name}" | 2996 | + # parallelism comes from the STATIC non-reduction split axes, whose blocks |
| 2996 | - ) | 2997 | + # are emitted as static_blocks (no runtime rule for the absent primary |
| 2998 | + # block). Covers both the pure scalar reduction (no split axes -> empty | ||
| 2999 | + # static_blocks) and the partial reduction (static split axes present). | ||
| 2997 | return { | 3000 | return { |
| 2998 | "group_id": group_id, | 3001 | "group_id": group_id, |
| 2999 | - "static_blocks": (), | 3002 | + "static_blocks": tuple( |
| 3003 | + (block_name, runtime_blocks[block_name]) | ||
| 3004 | + for block_name in runtime_block_arg_names | ||
| 3005 | + if block_name in runtime_blocks | ||
| 3006 | + ), | ||
| 3000 | "runtime_block_rules": (), | 3007 | "runtime_block_rules": (), |
| 3001 | "grid_target": 1, | 3008 | "grid_target": 1, |
| 3002 | } | 3009 | } |