!4 merge inductor-dynamic-reduction-group-autotune into v2.10.0

Mtorch_npu/_inductor/codegen/split_tiling.py+79-2

	@@ -104,9 +104,40 @@ class SplitTiling:
104	self.kernel.split_axis.clear()	104	self.kernel.split_axis.clear()
105		105
106	# total numel exceed aicore or total split axis exceed 3	106	# total numel exceed aicore or total split axis exceed 3
		107	+ def has_dynamic_split_axis():
		108	+ return any(
		109	+ not isinstance(axis.length, sympy.Integer)
		110	+ for axis in self.kernel.split_axis
		111	+ )
		112	+
		113	+ def is_reduction_kernel():
		114	+ return any(axis.prefix == "r" for axis in self.kernel.sorted_axis)
		115	+
107	def meet_stop_condition():	116	def meet_stop_condition():
108	sv = V.graph.sizevars	117	sv = V.graph.sizevars
109	current_numels = self.total_split_numels(self.kernel.split_axis)	118	current_numels = self.total_split_numels(self.kernel.split_axis)
		119	+ # A dynamic split axis already supplies the grid on its own: the size
		120	+ # hint seen here is one example shape, while the launch derives the
		121	+ # grid from the runtime numel. Adding a second split axis to reach the
		122	+ # core count is therefore speculative, and the cost is permanent: a
		123	+ # split axis can never become a no_loop_axis (see select_no_loop_axis),
		124	+ # so a static low dim pulled in here loses whole-axis residency and its
		125	+ # BLOCK_SUB can no longer span the axis. That fragments the innermost
		126	+ # contiguous dimension for every runtime shape, not just the one whose
		127	+ # hint was used. Reductions only: their grid comes solely from the
		128	+ # non-reduction axes, so the dynamic one is the parallelism.
		129	+ #
		130	+ # Confined to grouped autotune, which is what supplies the per-bucket
		131	+ # runtime block that makes one split axis sufficient. Without it the
		132	+ # single axis would keep the hint's block at every shape, so the old
		133	+ # speculative second axis stays the better bet.
		134	+ if (
		135	+ npu_config.enable_symbolic_shape_group_autotune
		136	+ and self.kernel.split_axis
		137	+ and is_reduction_kernel()
		138	+ and has_dynamic_split_axis()
		139	+ ):
		140	+ return True
110	try:	141	try:
111	val = sv.size_hint(current_numels)	142	val = sv.size_hint(current_numels)
112	except TypeError:	143	except TypeError:
	@@ -436,6 +467,20 @@ class SplitTiling:
436	# stays a single boundary; the ~1e6 1D corner is out of scope.	467	# stays a single boundary; the ~1e6 1D corner is out of scope.
437	_REDUCTION_BUCKETS = (8192,)	468	_REDUCTION_BUCKETS = (8192,)
438	_OUTER_BUCKETS = (256,)	469	_OUTER_BUCKETS = (256,)
		470	+ # Boundaries for an outer group that also carries static axes, in units of the
		471	+ # dynamic axis (see _scale_buckets_by_static_factor). Such a group spans every
		472	+ # runtime size from a grid-starved shape up to a device-saturating one, which
		473	+ # a single bucket cannot serve: the representative bakes the compile-time
		474	+ # sub-blocks, and those decide how many rows each program walks.
		475	+ #
		476	+ # The boundaries are multiples of the core count because that is what makes a
		477	+ # bucket internally uniform: every shape in (k-1)cores .. kcores wants the
		478	+ # same k rows per program, so the representative's sub-blocks are right for
		479	+ # all of them. Aligning the boundaries anywhere else leaves shapes that want
		480	+ # fewer rows sharing a variant tuned for more.
		481	+ _OUTER_BUCKETS_WITH_STATIC_AXES = tuple(
		482	+ num_vector_core * k for k in range(1, 6)
		483	+ )
439		484
440	def _axis_static_length(self, axis):	485	def _axis_static_length(self, axis):
441	try:	486	try:
	@@ -446,6 +491,29 @@ class SplitTiling:
446	except (AttributeError, KeyError, TypeError, ValueError):	491	except (AttributeError, KeyError, TypeError, ValueError):
447	return 1	492	return 1
448		493
		494	+ def _static_group_factor(self, axis_names):
		495	+ names = set(axis_names)
		496	+ factor = 1
		497	+ for axis in self.kernel.sorted_axis:
		498	+ if axis.name in names and isinstance(axis.length, sympy.Integer):
		499	+ factor *= max(1, self._axis_static_length(axis))
		500	+ return factor
		501	+
		502	+ def _scale_buckets_by_static_factor(self, buckets, axis_names, wide_buckets=None):
		503	+ # The boundaries describe the group PRODUCT, so the static axes of the group
		504	+ # are already baked into every runtime value. A group whose static factor
		505	+ # exceeds a boundary would otherwise put every reachable shape in the open
		506	+ # tail, and the tail representative (2 * boundary) divided by that factor
		507	+ # collapses the dynamic axis to a single element: the group then gets tuned
		508	+ # on a workload that cannot fill the grid. Scaling keeps the boundaries
		509	+ # expressed in the dynamic axis, which is what the buckets mean to separate.
		510	+ static_factor = self._static_group_factor(axis_names)
		511	+ if static_factor <= 1:
		512	+ return buckets
		513	+ if wide_buckets is not None:
		514	+ buckets = wide_buckets
		515	+ return tuple(boundary * static_factor for boundary in buckets)
		516	+
449	def _default_dynamic_axis_feature(	517	def _default_dynamic_axis_feature(
450	self,	518	self,
451	dynamic_split_axes,	519	dynamic_split_axes,
	@@ -543,7 +611,14 @@ class SplitTiling:
543	if outer_names and self._has_dynamic_axis(outer_names):	611	if outer_names and self._has_dynamic_axis(outer_names):
544	features.append(	612	features.append(
545	GroupFeatureSpec(	613	GroupFeatureSpec(
546	- "outer", "outer_product", outer_names, self._OUTER_BUCKETS	614	+ "outer",
		615	+ "outer_product",
		616	+ outer_names,
		617	+ self._scale_buckets_by_static_factor(
		618	+ self._OUTER_BUCKETS,
		619	+ outer_names,
		620	+ self._OUTER_BUCKETS_WITH_STATIC_AXES,
		621	+ ),
547	)	622	)
548	)	623	)
549	if reduction_names and self._has_dynamic_axis(reduction_names):	624	if reduction_names and self._has_dynamic_axis(reduction_names):
	@@ -552,7 +627,9 @@ class SplitTiling:
552	"reduction",	627	"reduction",
553	"reduction_product",	628	"reduction_product",
554	reduction_names,	629	reduction_names,
555	- self._REDUCTION_BUCKETS,	630	+ self._scale_buckets_by_static_factor(
		631	+ self._REDUCTION_BUCKETS, reduction_names
		632	+ ),
556	)	633	)
557	)	634	)
558	return tuple(features)	635	return tuple(features)

Mtorch_npu/_inductor/config.py+3-0

	@@ -166,6 +166,9 @@ inductor_config.triton.mix_order_reduction = False
166	enable_fast_gelu = os.getenv("TORCHINDUCTOR_ENABLE_FAST_GELU", "0") == "1"	166	enable_fast_gelu = os.getenv("TORCHINDUCTOR_ENABLE_FAST_GELU", "0") == "1"
167	allow_pad_lowering = os.getenv("TORCHINDUCTOR_ENABLE_PAD_LOWERING", "0") == "1"	167	allow_pad_lowering = os.getenv("TORCHINDUCTOR_ENABLE_PAD_LOWERING", "0") == "1"
168	allow_searchsorted_lowering = os.getenv("TORCHINDUCTOR_ENABLE_SEARCHSORTED_LOWERING", "0") == "1"	168	allow_searchsorted_lowering = os.getenv("TORCHINDUCTOR_ENABLE_SEARCHSORTED_LOWERING", "0") == "1"
		169	+allow_embedding_dense_backward_lowering = os.getenv(
		170	+ "TORCHINDUCTOR_ENABLE_EMBEDDING_DENSE_BACKWARD_LOWERING", "0"
		171	+) == "1"
169	enable_flex_attention_dq_before_scale_materialize = os.environ.get(	172	enable_flex_attention_dq_before_scale_materialize = os.environ.get(
170	"FLEX_ATTENTION_DQ_BEFORE_SCALE_MATERIALIZE", "1"	173	"FLEX_ATTENTION_DQ_BEFORE_SCALE_MATERIALIZE", "1"
171	).lower() in ("1", "true", "yes")	174	).lower() in ("1", "true", "yes")

Mtorch_npu/_inductor/decomposition.py+10-2

	@@ -168,7 +168,11 @@ def _matmul_backward_inductor(grad, self, other, mask):
168		168
169		169
170	def _register_triton_decompositions():	170	def _register_triton_decompositions():
171	- from .config import is_ascend950, enable_fast_gelu	171	+ from .config import (
		172	+ is_ascend950,
		173	+ enable_fast_gelu,
		174	+ allow_embedding_dense_backward_lowering,
		175	+ )
172	from .lowering import _add_overload	176	from .lowering import _add_overload
173		177
174	DECOMPOSITION_OVERLOAD_OP = [	178	DECOMPOSITION_OVERLOAD_OP = [
	@@ -181,7 +185,11 @@ def _register_triton_decompositions():
181	aten.expm1,	185	aten.expm1,
182	aten.native_layer_norm,	186	aten.native_layer_norm,
183	aten.repeat_interleave.Tensor, # perf issue	187	aten.repeat_interleave.Tensor, # perf issue
184	- aten.embedding_dense_backward,	188	+ *(
		189	+ [aten.embedding_dense_backward]
		190	+ if not allow_embedding_dense_backward_lowering
		191	+ else []
		192	+ ),
185	aten.matmul.default,	193	aten.matmul.default,
186	aten.matmul_backward.default,	194	aten.matmul_backward.default,
187	]	195	]

Mtorch_npu/_inductor/lowering_fallback_list.py+2-1

	@@ -36,6 +36,7 @@ from .config import (
36	inductor_ascend_linear_mode,	36	inductor_ascend_linear_mode,
37	allow_pad_lowering,	37	allow_pad_lowering,
38	allow_searchsorted_lowering,	38	allow_searchsorted_lowering,
		39	+ allow_embedding_dense_backward_lowering,
39	)	40	)
40		41
41	aten = torch.ops.aten	42	aten = torch.ops.aten
	@@ -606,7 +607,7 @@ TORCH_NATIVE_FALLBACK_LIST = [
606	aten._embedding_bag_forward_only.out,	607	aten._embedding_bag_forward_only.out,
607	aten._embedding_bag_per_sample_weights_backward.default,	608	aten._embedding_bag_per_sample_weights_backward.default,
608	aten._embedding_bag_per_sample_weights_backward.out,	609	aten._embedding_bag_per_sample_weights_backward.out,
609	- aten.embedding_dense_backward.default,	610	+ *([aten.embedding_dense_backward.default] if not allow_embedding_dense_backward_lowering else []),
610	aten._fft_r2c.default,	611	aten._fft_r2c.default,
611	aten._fft_r2c.out,	612	aten._fft_r2c.out,
612	aten._flash_attention_backward.default,	613	aten._flash_attention_backward.default,

Mtorch_npu/_inductor/runtime/triton_heuristics.py+127-5

	@@ -2577,7 +2577,37 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner):
2577	return torch.zeros(count, device=device, dtype=dtype)	2577	return torch.zeros(count, device=device, dtype=dtype)
2578	return torch.empty(count, device=device, dtype=dtype)	2578	return torch.empty(count, device=device, dtype=dtype)
2579		2579
2580	- def materialize_spec(spec, kind):	2580	+ def integer_benchmark_range(spec_index):
		2581	+ """Value range the real argument at this position actually spans.
		2582	+
		2583	+ rand_strided zero-fills integer tensors, so every indirect load in
		2584	+ the benchmark resolves to the same element and the gather costs
		2585	+ nothing. A gather kernel is then ranked on memory behaviour it will
		2586	+ never see: for an embedding reduction the candidates that handle
		2587	+ scattered rows well and the ones that do not land within a few
		2588	+ percent of each other, while differing by an order of magnitude on
		2589	+ real data. Sampling the range the real argument spans restores the
		2590	+ divergence, and cannot go out of bounds because the real launch
		2591	+ already indexes with exactly these values.
		2592	+
		2593	+ Restricted to reduction kernels so that every other grouped
		2594	+ heuristic keeps benchmarking on byte-identical inputs to today.
		2595	+ """
		2596	+ if self.heuristic_type not in (
		2597	+ HeuristicType.REDUCTION,
		2598	+ HeuristicType.PERSISTENT_REDUCTION,
		2599	+ ):
		2600	+ return None
		2601	+ if spec_index >= len(self._grouped_runtime_args_snapshot):
		2602	+ return None
		2603	+ real_arg = self._grouped_runtime_args_snapshot[spec_index]
		2604	+ if not isinstance(real_arg, torch.Tensor) or real_arg.numel() == 0:
		2605	+ return None
		2606	+ if real_arg.dtype.is_floating_point or real_arg.dtype == torch.bool:
		2607	+ return None
		2608	+ return int(real_arg.min()), int(real_arg.max())
		2609	+
		2610	+ def materialize_spec(spec, kind, spec_index=-1):
2581	spec_kind = spec.get("kind")	2611	spec_kind = spec.get("kind")
2582	source = spec.get("source")	2612	source = spec.get("source")
2583	if source == "runtime_arg":	2613	if source == "runtime_arg":
	@@ -2596,7 +2626,12 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner):
2596	tensor_factory = getattr(self, "_benchmark_tensor_factory", None)	2626	tensor_factory = getattr(self, "_benchmark_tensor_factory", None)
2597	if tensor_factory is None:	2627	if tensor_factory is None:
2598	tensor_factory = rand_strided	2628	tensor_factory = rand_strided
2599	- return tensor_factory(size, stride, dtype, spec["device"])	2629	+ tensor = tensor_factory(size, stride, dtype, spec["device"])
		2630	+ if not dtype.is_floating_point and dtype != torch.bool:
		2631	+ value_range = integer_benchmark_range(spec_index)
		2632	+ if value_range is not None:
		2633	+ tensor.random_(value_range[0], value_range[1] + 1)
		2634	+ return tensor
2600	if spec_kind == "workspace" and source == "workspace":	2635	if spec_kind == "workspace" and source == "workspace":
2601	count = eval_expr(spec["count_expr"])	2636	count = eval_expr(spec["count_expr"])
2602	return workspace_factory(	2637	return workspace_factory(
	@@ -2608,8 +2643,10 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner):
2608	raise RuntimeError(f"{kind} spec source {source} is not supported")	2643	raise RuntimeError(f"{kind} spec source {source} is not supported")
2609		2644
2610	return tuple(	2645	return tuple(
2611	- materialize_spec(spec, f"{spec.get('kind', 'unknown')} arg")	2646	+ materialize_spec(spec, f"{spec.get('kind', 'unknown')} arg", spec_index)
2612	- for spec in self.candidate_plan.get("ordered_arg_specs", ())	2647	+ for spec_index, spec in enumerate(
		2648	+ self.candidate_plan.get("ordered_arg_specs", ())
		2649	+ )
2613	)	2650	)
2614		2651
2615	def _build_grouped_benchmark_entries(self):	2652	def _build_grouped_benchmark_entries(self):
	@@ -2701,6 +2738,12 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner):
2701	timings[entry_idx] = timing	2738	timings[entry_idx] = timing
2702	return tuple(timings)	2739	return tuple(timings)
2703		2740
		2741	+ # Leaves only the candidates that can still win this group as None, so the
		2742	+ # full measurement below runs on those alone.
		2743	+ self._screen_grouped_entries(
		2744	+ entries, kernel_funcs, kernel_entry_indices, timings
		2745	+ )
		2746	+
2704	return tuple(	2747	return tuple(
2705	timing	2748	timing
2706	if timing is not None	2749	if timing is not None
	@@ -2713,6 +2756,76 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner):
2713	for entry, timing in zip(entries, timings)	2756	for entry, timing in zip(entries, timings)
2714	)	2757	)
2715		2758
		2759	+ # Screening pays for itself only above the selector's own ceiling: below it the
		2760	+ # screen would time every candidate and then keep them all anyway.
		2761	+ _GROUPED_SCREEN_MIN_CANDIDATES = 200
		2762	+
		2763	+ def _screen_grouped_entries(
		2764	+ self, entries, kernel_funcs, kernel_entry_indices, timings
		2765	+ ):
		2766	+ """Rank a group's candidates by one launch each, keeping only contenders.
		2767	+
		2768	+ The non-grouped sweep already works this way (see the prerun block in
		2769	+ _benchmark_candidate_entries); the grouped sweep never picked it up and so
		2770	+ pays the full repeated measurement for every candidate. A specialized
		2771	+ compile can afford that, but a grouped reduction benchmarks a bucket
		2772	+ representative rather than a real shape: the pool is an order of magnitude
		2773	+ larger, and it holds configs whose loop count explodes at the
		2774	+ representative's reduction length. Measuring those to full precision only
		2775	+ establishes how badly they lose.
		2776	+
		2777	+ Candidates that are screened out keep their single-launch time as their
		2778	+ timing. That is a fair ranking for them because they are not going to be
		2779	+ the group's winner, and the winner itself is still chosen by the same full
		2780	+ measurement as before.
		2781	+
		2782	+ Mutates `timings` in place; entries left at None are the ones to measure.
		2783	+ """
		2784	+ if self.heuristic_type not in (
		2785	+ HeuristicType.REDUCTION,
		2786	+ HeuristicType.PERSISTENT_REDUCTION,
		2787	+ ):
		2788	+ return
		2789	+ if len(kernel_entry_indices) <= self._GROUPED_SCREEN_MIN_CANDIDATES:
		2790	+ return
		2791	+
		2792	+ screened_ms = {}
		2793	+ by_group = {}
		2794	+ for kernel_func, entry_idx in zip(kernel_funcs, kernel_entry_indices):
		2795	+ try:
		2796	+ elapsed_ms = _measure_prerun_ms(kernel_func)
		2797	+ except Exception as exc:
		2798	+ screen_log_str = (
		2799	+ f"PreRun [{self.fn.__name__}], grouped candidate: "
		2800	+ f"{entries[entry_idx]['candidate']['candidate_id']}\n err: {exc}"
		2801	+ )
		2802	+ if not autotune_continue_on_failure:
		2803	+ raise RuntimeError(screen_log_str) from exc
		2804	+ log.warning(screen_log_str)
		2805	+ timings[entry_idx] = float("inf")
		2806	+ continue
		2807	+ screened_ms[entry_idx] = elapsed_ms
		2808	+ by_group.setdefault(entries[entry_idx]["group_id"], []).append(
		2809	+ (entry_idx, elapsed_ms)
		2810	+ )
		2811	+
		2812	+ # Each group picks its own winner, so each group keeps its own contenders.
		2813	+ selected = set()
		2814	+ for group_timings in by_group.values():
		2815	+ selected \|= set(_select_prerun_top_candidates(group_timings))
		2816	+
		2817	+ for entry_idx, elapsed_ms in screened_ms.items():
		2818	+ if entry_idx not in selected:
		2819	+ timings[entry_idx] = elapsed_ms
		2820	+
		2821	+ log.info(
		2822	+ "%s grouped screen kept %s of %s candidates across %s groups",
		2823	+ self.get_fn_name(),
		2824	+ len(selected),
		2825	+ len(screened_ms),
		2826	+ len(by_group),
		2827	+ )
		2828	+
2716	def _autotune_all_groups(self, args, *kwargs):	2829	def _autotune_all_groups(self, args, *kwargs):
2717	self._grouped_runtime_args_snapshot = args	2830	self._grouped_runtime_args_snapshot = args
2718	try:	2831	try:
	@@ -3428,7 +3541,16 @@ def build_grouped_launch_policy(
3428	("axis_name", primary_group_axis),	3541	("axis_name", primary_group_axis),
3429	("block_sub", primary_block_sub),	3542	("block_sub", primary_block_sub),
3430	]	3543	]
3431	- if primary_feature_name in ("pointwise", "elementwise_numel"):	3544	+ if primary_feature_name in (
		3545	+ "pointwise",
		3546	+ "elementwise_numel",
		3547	+ # Reduction grouping needs this just as much: the primary axis is the
		3548	+ # only source of grid parallelism, so without it every shape below the
		3549	+ # representative inherits the representative's grid and each program
		3550	+ # loops over the representative's rows instead of its own.
		3551	+ "outer",
		3552	+ "reduction",
		3553	+ ):
3432	runtime_rule.extend(	3554	runtime_rule.extend(
3433	(	3555	(
3434	("representative_numel", int(axis_env[primary_group_axis])),	3556	("representative_numel", int(axis_env[primary_group_axis])),