feat(inductor): adapt group autotune for cpp wrapper-pytorch-AtomGit

已合并

feat(inductor): adapt group autotune for cpp wrapper #44686

Xuan Peng创建于 20 天前从HinPeng/pytorch： feat/group-autotune-cppwrapper-v2.10.0 合入到 Ascend/pytorch： v2.10.0

feat(inductor): adapt group autotune for cpp wrapper #44686

已合并

Xuan Peng创建于 20 天前

从HinPeng/pytorch feat/group-autotune-cppwrapper-v2.10.0合入到Ascend/pytorch v2.10.0

共 3 个文件变更+746-5

Atest/_inductor/test_grouped_cpp_wrapper.py+347-0

@@ -0,0 +1,347 @@
	1	+import unittest
	2	+from types import SimpleNamespace
	3	+
	4	+import torch
	5	+from torch._inductor import config
	6	+from torch._inductor.codecache import CudaKernelParamCache
	7	+from torch._inductor.utils import IndentedBuffer, run_and_get_cpp_code
	8	+from torch._inductor.virtualized import V
	9	+from torch.testing._internal.common_utils import run_tests, TestCase
	10	+
	11	+import torch_npu
	12	+import torch_npu._inductor
	13	+from torch_npu._inductor.codegen.cpp_wrapper_npu import (
	14	+ CppWrapperNpu,
	15	+ DeferredNpuTritonCallWrapper,
	16	+)
	17	+
	18	+GROUP_COUNT = 32
	19	+HIDDEN_SIZE = 1600
	20	+GATE_SIZE = 256
	21	+POST_SIZE = 64
	22	+SIDE_SIZE = 32
	23	+COMPILE_SEQUENCE_LENGTH = 200
	24	+
	25	+
	26	+def _variant_load_meta(variant_id):
	27	+ return {
	28	+ "cubin_path": f"/tmp/triton_kernel_{variant_id}.cubin",
	29	+ "mangled_name": f"mangled_{variant_id}",
	30	+ "mix_mode": "aiv",
	31	+ "shared_mem": 64,
	32	+ "parallel_mode": "vector",
	33	+ "force_simt_only": False,
	34	+ "shared_mem_dynamic_size": 0,
	35	+ "has_auto_blockify_blacklist_op": False,
	36	+ }
	37	+
	38	+
	39	+def _grouped_plan():
	40	+ return {
	41	+ "variants": {
	42	+ "v0": {"config": {}, "load_meta": _variant_load_meta("v0")},
	43	+ "v1": {"config": {}, "load_meta": _variant_load_meta("v1")},
	44	+ },
	45	+ "variant_order": ("v0", "v1"),
	46	+ "best_by_group": {
	47	+ "0": {"variant_id": "v0", "policy_id": "p0"},
	48	+ "1": {"variant_id": "v1", "policy_id": "p1"},
	49	+ },
	50	+ "policies": {
	51	+ "p0": {
	52	+ "grid_target": 8,
	53	+ "static_blocks": (),
	54	+ "runtime_block_rules": (
	55	+ (
	56	+ "XBLOCK",
	57	+ (
	58	+ ("op", "ceildiv"),
	59	+ ("axis_name", "x"),
	60	+ ("block_sub", 4),
	61	+ ),
	62	+ ),
	63	+ ),
	64	+ },
	65	+ "p1": {
	66	+ "grid_target": 16,
	67	+ "static_blocks": (),
	68	+ "runtime_block_rules": (
	69	+ (
	70	+ "XBLOCK",
	71	+ (
	72	+ ("op", "ceildiv"),
	73	+ ("axis_name", "x"),
	74	+ ("block_sub", 8),
	75	+ ),
	76	+ ),
	77	+ ),
	78	+ },
	79	+ },
	80	+ "runtime_block_append_order": ("XBLOCK",),
	81	+ "group_id_count": 2,
	82	+ "reachable_group_ids": (0, 1),
	83	+ "group_features": (
	84	+ {
	85	+ "name": "x",
	86	+ "source": "axis",
	87	+ "axis_names": ("x",),
	88	+ "buckets": (128,),
	89	+ },
	90	+ ),
	91	+ "axis_arg_indices": {"x": 1},
	92	+ "feature_arg_indices": ((1,),),
	93	+ "feature_sources": (
	94	+ {
	95	+ "name": "x",
	96	+ "source": "axis",
	97	+ "axis_names": ("x",),
	98	+ },
	99	+ ),
	100	+ }
	101	+
	102	+
	103	+def _render_grouped_wrapper(grouped_plan=None, *, return_files=False):
	104	+ grouped_plan = grouped_plan or _grouped_plan()
	105	+ inductor_meta = {
	106	+ "group_enabled": True,
	107	+ "grouped_candidate_plan": grouped_plan,
	108	+ "grid_type": "GridNpu",
	109	+ "axis_names": ("x",),
	110	+ "runtime_block_arg_names": ("XBLOCK",),
	111	+ "primary_group_axis": "x",
	112	+ }
	113	+ params = {
	114	+ "def_args": ["in_ptr", "x_numel", "XBLOCK"],
	115	+ "call_args": ["in_ptr", "x_numel", "XBLOCK"],
	116	+ "config": {"split_axis": (0,), "split_blocks": (128,)},
	117	+ "inductor_meta": inductor_meta,
	118	+ "triton_meta": {
	119	+ "signature": {
	120	+ "in_ptr": "fp16",
	121	+ "x_numel": "i64",
	122	+ "XBLOCK": "i64",
	123	+ },
	124	+ "constants": {},
	125	+ },
	126	+ "mangled_name": "unused",
	127	+ "shared_mem": 64,
	128	+ "cubin_path": "/tmp/unused.cubin",
	129	+ "mix_mode": "aiv",
	130	+ "parallel_mode": "vector",
	131	+ "force_simt_only": False,
	132	+ }
	133	+ graph = SimpleNamespace(
	134	+ cpp_wrapper=True,
	135	+ aot_mode=False,
	136	+ is_const_graph=False,
	137	+ constant_reprs={},
	138	+ inputs_to_check=[],
	139	+ graph_input_names=[],
	140	+ graph_inputs={},
	141	+ device_types={"npu"},
	142	+ wrapper_code=None,
	143	+ )
	144	+ with V.set_graph_handler(graph):
	145	+ wrapper = CppWrapperNpu()
	146	+ graph.wrapper_code = wrapper
	147	+ wrapper.prefix = IndentedBuffer()
	148	+ CudaKernelParamCache.cache_clear()
	149	+ CudaKernelParamCache.cache["triton_kernel"] = params
	150	+ deferred = DeferredNpuTritonCallWrapper(
	151	+ wrapper_name="call_triton_kernel",
	152	+ kernel_name="triton_kernel",
	153	+ kernel_name_to_body={},
	154	+ arg_types=[torch.float16, int],
	155	+ kernel_id=0,
	156	+ )
	157	+ deferred.generate(wrapper)
	158	+ source = wrapper.prefix.getvalue()
	159	+ if return_files:
	160	+ return source, tuple(wrapper.additional_files)
	161	+ return source
	162	+
	163	+
	164	+class GatedTransposeBmmModel(torch.nn.Module):
	165	+ def forward(
	166	+ self,
	167	+ source,
	168	+ gate_down_rhs,
	169	+ gate_down_bias,
	170	+ gate_up_rhs,
	171	+ bias,
	172	+ post_bmm_rhs,
	173	+ side_bmm_rhs,
	174	+ ):
	175	+ transposed = source.transpose(0, 1)
	176	+ gate_down = torch.bmm(transposed, gate_down_rhs)
	177	+ gate_down = torch.tanh(gate_down + gate_down_bias)
	178	+ bmm_result = torch.bmm(gate_down, gate_up_rhs)
	179	+ gated = transposed * torch.tanh(bmm_result + bias)
	180	+ post = torch.bmm(gated, post_bmm_rhs)
	181	+ side = torch.bmm(transposed, side_bmm_rhs)
	182	+ return post, side
	183	+
	184	+
	185	+def _make_gated_transpose_inputs(sequence_length):
	186	+ def rand(*shape):
	187	+ return torch.randn(shape, device="npu:0", dtype=torch.float16)
	188	+
	189	+ return (
	190	+ rand(sequence_length, GROUP_COUNT, HIDDEN_SIZE),
	191	+ rand(GROUP_COUNT, HIDDEN_SIZE, GATE_SIZE),
	192	+ rand(GROUP_COUNT, 1, GATE_SIZE),
	193	+ rand(GROUP_COUNT, GATE_SIZE, HIDDEN_SIZE),
	194	+ rand(GROUP_COUNT, 1, HIDDEN_SIZE),
	195	+ rand(GROUP_COUNT, HIDDEN_SIZE, POST_SIZE) * 0.01,
	196	+ rand(GROUP_COUNT, HIDDEN_SIZE, SIDE_SIZE),
	197	+ )
	198	+
	199	+
	200	+def _mark_sequence_length_dynamic(inputs):
	201	+ torch._dynamo.mark_dynamic(
	202	+ inputs[0],
	203	+ 0,
	204	+ hint_override=COMPILE_SEQUENCE_LENGTH,
	205	+ )
	206	+
	207	+
	208	+class TestGroupedCppWrapper(TestCase):
	209	+ def tearDown(self):
	210	+ CudaKernelParamCache.cache_clear()
	211	+ torch._dynamo.reset()
	212	+ super().tearDown()
	213	+
	214	+ def test_grouped_wrapper_emits_bucket_dispatch_and_variants(self):
	215	+ source = _render_grouped_wrapper()
	216	+
	217	+ self.assertIn("switch (grouped_group_id)", source)
	218	+ self.assertIn("case 0: {", source)
	219	+ self.assertIn("case 1: {", source)
	220	+ self.assertIn("grouped_kernel_v0", source)
	221	+ self.assertIn("grouped_kernel_v1", source)
	222	+ self.assertIn('"mangled_v0"', source)
	223	+ self.assertIn('"mangled_v1"', source)
	224	+
	225	+ def test_grouped_wrapper_materializes_block_sub_aligned_runtime_block(self):
	226	+ source = _render_grouped_wrapper()
	227	+
	228	+ self.assertIn(
	229	+ "auto resolve_grouped_runtime_block",
	230	+ source,
	231	+ )
	232	+ self.assertIn(
	233	+ "total_subblocks = ceildiv(axis_numel, block_sub)",
	234	+ source,
	235	+ )
	236	+ self.assertIn(
	237	+ "program_subblocks = ceildiv(",
	238	+ source,
	239	+ )
	240	+ self.assertIn(
	241	+ "effective_grid = ceildiv(",
	242	+ source,
	243	+ )
	244	+ self.assertIn(
	245	+ "XBLOCK = resolve_grouped_runtime_block(x_numel, 8, 4)",
	246	+ source,
	247	+ )
	248	+ self.assertIn(
	249	+ "XBLOCK = resolve_grouped_runtime_block(x_numel, 16, 8)",
	250	+ source,
	251	+ )
	252	+ self.assertLess(source.index("int64_t XBLOCK"), source.index("uint32_t grid_0"))
	253	+
	254	+ def test_grouped_wrapper_omits_unselected_variant(self):
	255	+ grouped_plan = _grouped_plan()
	256	+ grouped_plan["best_by_group"]["1"] = {
	257	+ "variant_id": "v0",
	258	+ "policy_id": "p0",
	259	+ }
	260	+ source, additional_files = _render_grouped_wrapper(
	261	+ grouped_plan,
	262	+ return_files=True,
	263	+ )
	264	+
	265	+ self.assertIn("grouped_kernel_v0", source)
	266	+ self.assertNotIn("grouped_kernel_v1", source)
	267	+ self.assertIn("/tmp/triton_kernel_v0.cubin", additional_files)
	268	+ self.assertNotIn("/tmp/triton_kernel_v1.cubin", additional_files)
	269	+ self.assertNotIn("/tmp/unused.cubin", additional_files)
	270	+
	271	+ @unittest.skipIf(not torch.npu.is_available(), "NPU is not available")
	272	+ def test_gated_transpose_dynamic_shapes_functionality_and_accuracy(self):
	273	+ import torch_npu._inductor.config as npu_config
	274	+
	275	+ previous_group_autotune = (
	276	+ npu_config.enable_symbolic_shape_group_autotune
	277	+ )
	278	+ npu_config.enable_symbolic_shape_group_autotune = True
	279	+ try:
	280	+ with config.patch(
	281	+ {
	282	+ "cpp_wrapper": True,
	283	+ "compile_threads": 1,
	284	+ "force_disable_caches": True,
	285	+ }
	286	+ ):
	287	+ model = GatedTransposeBmmModel().eval()
	288	+ compile_inputs = _make_gated_transpose_inputs(
	289	+ COMPILE_SEQUENCE_LENGTH
	290	+ )
	291	+ _mark_sequence_length_dynamic(compile_inputs)
	292	+ compiled = torch.compile(
	293	+ model,
	294	+ backend="inductor",
	295	+ fullgraph=True,
	296	+ dynamic=None,
	297	+ )
	298	+
	299	+ with torch.no_grad():
	300	+ expected = model(*compile_inputs)
	301	+ actual, cpp_code = run_and_get_cpp_code(
	302	+ compiled, *compile_inputs
	303	+ )
	304	+ torch.npu.synchronize()
	305	+ torch.testing.assert_close(
	306	+ actual,
	307	+ expected,
	308	+ rtol=0.02,
	309	+ atol=0.03,
	310	+ msg=(
	311	+ "grouped cpp wrapper mismatch for "
	312	+ f"sequence_length={COMPILE_SEQUENCE_LENGTH}"
	313	+ ),
	314	+ )
	315	+ self.assertIn("'group_enabled': True", cpp_code)
	316	+ self.assertIn("switch (grouped_group_id)", cpp_code)
	317	+ self.assertGreaterEqual(
	318	+ cpp_code.count("static void* grouped_kernel_v"), 2
	319	+ )
	320	+
	321	+ for sequence_length, inputs in (
	322	+ (4, _make_gated_transpose_inputs(4)),
	323	+ (256, _make_gated_transpose_inputs(256)),
	324	+ ):
	325	+ _mark_sequence_length_dynamic(inputs)
	326	+ with torch.no_grad():
	327	+ expected = model(*inputs)
	328	+ actual = compiled(*inputs)
	329	+ torch.npu.synchronize()
	330	+ torch.testing.assert_close(
	331	+ actual,
	332	+ expected,
	333	+ rtol=0.02,
	334	+ atol=0.03,
	335	+ msg=(
	336	+ "grouped cpp wrapper mismatch for "
	337	+ f"sequence_length={sequence_length}"
	338	+ ),
	339	+ )
	340	+ finally:
	341	+ npu_config.enable_symbolic_shape_group_autotune = (
	342	+ previous_group_autotune
	343	+ )
	344	+
	345	+
	346	+if __name__ == "__main__":
	347	+ run_tests()

Mtorch_npu/_inductor/codegen/cpp_wrapper_npu.py+320-5

	@@ -30,6 +30,7 @@ from torch._inductor.ir import GraphPartitionSignature
30	from torch._inductor.runtime.runtime_utils import dynamo_timed	30	from torch._inductor.runtime.runtime_utils import dynamo_timed
31	from torch._inductor.utils import IndentedBuffer	31	from torch._inductor.utils import IndentedBuffer
32	from torch._inductor.virtualized import V	32	from torch._inductor.virtualized import V
		33	+from torch.utils._ordered_set import OrderedSet
33		34
34	from .. import config as npu_config	35	from .. import config as npu_config
35	from ..runtime.triton_heuristics import GridExprNpu	36	from ..runtime.triton_heuristics import GridExprNpu
	@@ -100,11 +101,280 @@ class DeferredNpuTritonCallWrapper(DeferredTritonCallWrapper):
100	self.arg_types = original_arg_types	101	self.arg_types = original_arg_types
101		102
102	def generate(self, wrapper: CppWrapperGpu):	103	def generate(self, wrapper: CppWrapperGpu):
		104	+ additional_files = V.graph.wrapper_code.additional_files
		105	+ existing_files = OrderedSet(additional_files)
103	with self._patch_runtime_block_params() as params:	106	with self._patch_runtime_block_params() as params:
104	super().generate(wrapper)	107	super().generate(wrapper)
105	- cubin_path = params[get_cpp_wrapper_cubin_path_name()]	108	+ inductor_meta = params["inductor_meta"]
106	- if cubin_path not in V.graph.wrapper_code.additional_files:	109	+ if inductor_meta.get("group_enabled", False):
107	- V.graph.wrapper_code.additional_files.append(cubin_path)	110	+ grouped_plan = inductor_meta["grouped_candidate_plan"]
		111	+ cubin_paths = tuple(
		112	+ self._grouped_load_meta(grouped_plan, variant_id)[
		113	+ "cubin_path"
		114	+ ]
		115	+ for variant_id in self._grouped_active_variants(grouped_plan)
		116	+ )
		117	+ default_cubin_path = params[get_cpp_wrapper_cubin_path_name()]
		118	+ if (
		119	+ default_cubin_path not in existing_files
		120	+ and default_cubin_path not in cubin_paths
		121	+ and default_cubin_path in additional_files
		122	+ ):
		123	+ additional_files.remove(default_cubin_path)
		124	+ else:
		125	+ cubin_paths = (params[get_cpp_wrapper_cubin_path_name()],)
		126	+ for cubin_path in cubin_paths:
		127	+ if cubin_path not in additional_files:
		128	+ additional_files.append(cubin_path)
		129	+
		130	+ @staticmethod
		131	+ def _grouped_runtime_block_names(
		132	+ grouped_plan: dict[str, Any],
		133	+ ) -> tuple[str, ...]:
		134	+ return tuple(grouped_plan.get("runtime_block_append_order", ()))
		135	+
		136	+ @staticmethod
		137	+ def _grouped_active_variants(
		138	+ grouped_plan: dict[str, Any],
		139	+ ) -> tuple[str, ...]:
		140	+ best_by_group = grouped_plan.get("best_by_group", {})
		141	+ if not best_by_group:
		142	+ raise RuntimeError(
		143	+ "grouped cpp wrapper expects best_by_group to be populated"
		144	+ )
		145	+ selected_variant_ids = OrderedSet(
		146	+ selected["variant_id"] for selected in best_by_group.values()
		147	+ )
		148	+ active_variants = tuple(
		149	+ variant_id
		150	+ for variant_id in grouped_plan.get("variant_order", ())
		151	+ if variant_id in selected_variant_ids
		152	+ )
		153	+ if not active_variants:
		154	+ raise RuntimeError(
		155	+ "grouped cpp wrapper has no active compiled variants"
		156	+ )
		157	+ return active_variants
		158	+
		159	+ @staticmethod
		160	+ def _grouped_load_meta(
		161	+ grouped_plan: dict[str, Any], variant_id: str
		162	+ ) -> dict[str, Any]:
		163	+ load_meta = dict(
		164	+ grouped_plan["variants"][variant_id].get("load_meta", {})
		165	+ )
		166	+ if not load_meta:
		167	+ raise RuntimeError(
		168	+ f"grouped cpp wrapper expects load_meta for variant {variant_id}"
		169	+ )
		170	+ return load_meta
		171	+
		172	+ def _generate_grouped_feature_inputs(
		173	+ self,
		174	+ prefix: IndentedBuffer,
		175	+ grouped_plan: dict[str, Any],
		176	+ def_args: list[str],
		177	+ ) -> None:
		178	+ feature_arg_indices = tuple(
		179	+ grouped_plan.get("feature_arg_indices", ())
		180	+ )
		181	+ feature_sources = tuple(grouped_plan.get("feature_sources", ()))
		182	+ if len(feature_arg_indices) != len(feature_sources):
		183	+ raise RuntimeError(
		184	+ "grouped cpp wrapper feature inputs and sources do not match"
		185	+ )
		186	+ for feature_idx, (arg_indices, feature_source) in enumerate(
		187	+ zip(feature_arg_indices, feature_sources)
		188	+ ):
		189	+ arg_names = [def_args[arg_index] for arg_index in arg_indices]
		190	+ source = feature_source.get("source")
		191	+ if source in ("outer_product", "reduction_product"):
		192	+ feature_expr = " * ".join(arg_names)
		193	+ elif len(arg_names) == 1:
		194	+ feature_expr = arg_names[0]
		195	+ else:
		196	+ raise RuntimeError(
		197	+ f"grouped cpp wrapper feature {source} expects one axis"
		198	+ )
		199	+ prefix.writeline(
		200	+ f"int64_t grouped_feature_{feature_idx} = {feature_expr};"
		201	+ )
		202	+
		203	+ def _generate_grouped_group_id(
		204	+ self, prefix: IndentedBuffer, grouped_plan: dict[str, Any]
		205	+ ) -> None:
		206	+ group_features = tuple(grouped_plan.get("group_features", ()))
		207	+ prefix.writeline("int64_t grouped_group_id = 0;")
		208	+ prefix.writeline("int64_t grouped_group_stride = 1;")
		209	+ for feature_idx, feature_spec in enumerate(group_features):
		210	+ buckets = tuple(feature_spec.get("buckets", ()))
		211	+ prefix.writeline(f"int64_t grouped_bucket_{feature_idx} = 0;")
		212	+ for bucket_idx, upper_bound in enumerate(buckets):
		213	+ keyword = "if" if bucket_idx == 0 else "else if"
		214	+ prefix.writeline(
		215	+ f"{keyword} (grouped_feature_{feature_idx} <= "
		216	+ f"{int(upper_bound)}) grouped_bucket_{feature_idx} = "
		217	+ f"{bucket_idx};"
		218	+ )
		219	+ if buckets:
		220	+ prefix.writeline(
		221	+ f"else grouped_bucket_{feature_idx} = {len(buckets)};"
		222	+ )
		223	+ prefix.writeline(
		224	+ f"grouped_group_id += grouped_bucket_{feature_idx} * "
		225	+ "grouped_group_stride;"
		226	+ )
		227	+ prefix.writeline(
		228	+ f"grouped_group_stride *= {len(buckets) + 1};"
		229	+ )
		230	+ group_id_count = int(grouped_plan.get("group_id_count", 0))
		231	+ if group_id_count:
		232	+ prefix.writeline(
		233	+ f"if (grouped_group_id >= {group_id_count}) "
		234	+ 'throw std::runtime_error("grouped cpp wrapper resolved group '
		235	+ 'id out of range");'
		236	+ )
		237	+
		238	+ def _generate_grouped_selection(
		239	+ self,
		240	+ prefix: IndentedBuffer,
		241	+ grouped_plan: dict[str, Any],
		242	+ def_args: list[str],
		243	+ active_variants: tuple[str, ...],
		244	+ ) -> None:
		245	+ runtime_block_names = self._grouped_runtime_block_names(grouped_plan)
		246	+ axis_arg_indices = dict(grouped_plan.get("axis_arg_indices", {}))
		247	+ best_by_group = grouped_plan["best_by_group"]
		248	+ prefix.writeline("int64_t grouped_variant_index = -1;")
		249	+ for block_name in runtime_block_names:
		250	+ prefix.writeline(f"int64_t {block_name} = 0;")
		251	+ if runtime_block_names:
		252	+ prefix.splice(
		253	+ """
		254	+ auto resolve_grouped_runtime_block = [](
		255	+ int64_t axis_numel,
		256	+ int64_t expected_grid,
		257	+ int64_t block_sub
		258	+ ) -> int64_t {
		259	+ if (axis_numel <= 0) return 1;
		260	+ auto ceildiv = [](int64_t value, int64_t divisor) {
		261	+ return (value + divisor - 1) / divisor;
		262	+ };
		263	+ int64_t total_subblocks = ceildiv(axis_numel, block_sub);
		264	+ int64_t program_subblocks = ceildiv(
		265	+ total_subblocks, expected_grid
		266	+ );
		267	+ int64_t effective_grid = ceildiv(
		268	+ total_subblocks, program_subblocks
		269	+ );
		270	+ return ceildiv(axis_numel, effective_grid);
		271	+ };
		272	+ """
		273	+ )
		274	+ prefix.writeline("switch (grouped_group_id) {")
		275	+ with prefix.indent():
		276	+ for group_id in grouped_plan.get("reachable_group_ids", ()):
		277	+ selected = best_by_group.get(str(group_id))
		278	+ if selected is None:
		279	+ raise RuntimeError(
		280	+ "grouped cpp wrapper is missing winner for reachable "
		281	+ f"group {group_id}"
		282	+ )
		283	+ variant_id = selected["variant_id"]
		284	+ policy = grouped_plan["policies"][selected["policy_id"]]
		285	+ prefix.writeline(f"case {group_id}: {{")
		286	+ with prefix.indent():
		287	+ prefix.writeline(
		288	+ f"grouped_variant_index = "
		289	+ f"{active_variants.index(variant_id)};"
		290	+ )
		291	+ assigned_blocks = OrderedSet()
		292	+ for block_name, block_value in policy.get(
		293	+ "static_blocks", ()
		294	+ ):
		295	+ prefix.writeline(
		296	+ f"{block_name} = {int(block_value)};"
		297	+ )
		298	+ assigned_blocks.add(block_name)
		299	+ for block_name, rule_items in policy.get(
		300	+ "runtime_block_rules", ()
		301	+ ):
		302	+ rule = dict(rule_items)
		303	+ if rule.get("op") != "ceildiv":
		304	+ raise RuntimeError(
		305	+ "grouped cpp wrapper only supports ceildiv "
		306	+ "runtime block rules"
		307	+ )
		308	+ axis_name = rule["axis_name"]
		309	+ if axis_name not in axis_arg_indices:
		310	+ raise RuntimeError(
		311	+ "grouped cpp wrapper is missing axis argument "
		312	+ f"for {axis_name}"
		313	+ )
		314	+ axis_arg = def_args[axis_arg_indices[axis_name]]
		315	+ prefix.writeline(
		316	+ f"{block_name} = resolve_grouped_runtime_block("
		317	+ f"{axis_arg}, {int(policy['grid_target'])}, "
		318	+ f"{int(rule['block_sub'])});"
		319	+ )
		320	+ assigned_blocks.add(block_name)
		321	+ missing_blocks = [
		322	+ name
		323	+ for name in runtime_block_names
		324	+ if name not in assigned_blocks
		325	+ ]
		326	+ if missing_blocks:
		327	+ raise RuntimeError(
		328	+ "grouped cpp wrapper policy is missing runtime "
		329	+ f"blocks {missing_blocks}"
		330	+ )
		331	+ prefix.writeline("break;")
		332	+ prefix.writeline("}")
		333	+ prefix.writeline("default:")
		334	+ with prefix.indent():
		335	+ prefix.writeline(
		336	+ 'throw std::runtime_error("grouped cpp wrapper resolved '
		337	+ 'an unavailable group");'
		338	+ )
		339	+ prefix.writeline("}")
		340	+
		341	+ def _generate_grouped_variant_launch(
		342	+ self,
		343	+ prefix: IndentedBuffer,
		344	+ wrapper: CppWrapperGpu,
		345	+ params: dict[str, Any],
		346	+ grouped_plan: dict[str, Any],
		347	+ active_variants: tuple[str, ...],
		348	+ ) -> None:
		349	+ prefix.writeline("switch (grouped_variant_index) {")
		350	+ with prefix.indent():
		351	+ for variant_index, variant_id in enumerate(active_variants):
		352	+ load_meta = self._grouped_load_meta(grouped_plan, variant_id)
		353	+ prefix.writeline(f"case {variant_index}: {{")
		354	+ with prefix.indent():
		355	+ kernel_var_name = f"grouped_kernel_{variant_id}"
		356	+ variant_params = {
		357	+ **params,
		358	+ **load_meta,
		359	+ }
		360	+ self._generate_single_kernel_load(
		361	+ prefix, kernel_var_name, variant_params
		362	+ )
		363	+ self._generate_single_kernel_launch(
		364	+ prefix,
		365	+ wrapper,
		366	+ kernel_var_name,
		367	+ variant_params,
		368	+ )
		369	+ prefix.writeline("break;")
		370	+ prefix.writeline("}")
		371	+ prefix.writeline("default:")
		372	+ with prefix.indent():
		373	+ prefix.writeline(
		374	+ 'throw std::runtime_error("grouped cpp wrapper could not '
		375	+ 'launch selected variant");'
		376	+ )
		377	+ prefix.writeline("}")
108		378
109	def generate_grid(	379	def generate_grid(
110	self,	380	self,
	@@ -112,9 +382,25 @@ class DeferredNpuTritonCallWrapper(DeferredTritonCallWrapper):
112	inductor_meta: dict[str, Any],	382	inductor_meta: dict[str, Any],
113	params: dict[str, Any],	383	params: dict[str, Any],
114	):	384	):
		385	+ if inductor_meta.get("group_enabled", False):
		386	+ grouped_plan = inductor_meta["grouped_candidate_plan"]
		387	+ active_variants = self._grouped_active_variants(grouped_plan)
		388	+ self._generate_grouped_feature_inputs(
		389	+ prefix, grouped_plan, params["def_args"]
		390	+ )
		391	+ self._generate_grouped_group_id(prefix, grouped_plan)
		392	+ self._generate_grouped_selection(
		393	+ prefix,
		394	+ grouped_plan,
		395	+ params["def_args"],
		396	+ active_variants,
		397	+ )
115	numels = [arg for arg in params["def_args"] if "_numel" in arg]	398	numels = [arg for arg in params["def_args"] if "_numel" in arg]
116	- for block_name, block_value in dict(params.get("runtime_blocks", {})).items():	399	+ if not inductor_meta.get("group_enabled", False):
117	- prefix.writeline(f"int64_t {block_name} = {block_value};")	400	+ for block_name, block_value in dict(
		401	+ params.get("runtime_blocks", {})
		402	+ ).items():
		403	+ prefix.writeline(f"int64_t {block_name} = {block_value};")
118	grid = GridExprNpu.from_meta_and_set_numel(	404	grid = GridExprNpu.from_meta_and_set_numel(
119	inductor_meta, params["config"], numels, "cpp"	405	inductor_meta, params["config"], numels, "cpp"
120	)	406	)
	@@ -130,6 +416,17 @@ class DeferredNpuTritonCallWrapper(DeferredTritonCallWrapper):
130	prefix.writeline("if (grid_0 == 0 \|\| grid_1 == 0 \|\| grid_2 == 0) return;")	416	prefix.writeline("if (grid_0 == 0 \|\| grid_1 == 0 \|\| grid_2 == 0) return;")
131		417
132	def generate_load_kernel(self, prefix, kernel_var_name, params):	418	def generate_load_kernel(self, prefix, kernel_var_name, params):
		419	+ inductor_meta = params["inductor_meta"]
		420	+ if inductor_meta.get("group_enabled", False):
		421	+ grouped_plan = inductor_meta["grouped_candidate_plan"]
		422	+ for variant_id in self._grouped_active_variants(grouped_plan):
		423	+ prefix.writeline(
		424	+ f"static void* grouped_kernel_{variant_id} = nullptr;"
		425	+ )
		426	+ return
		427	+ self._generate_single_kernel_load(prefix, kernel_var_name, params)
		428	+
		429	+ def _generate_single_kernel_load(self, prefix, kernel_var_name, params):
133	prefix.writeline(f"if ({kernel_var_name} == nullptr) {{")	430	prefix.writeline(f"if ({kernel_var_name} == nullptr) {{")
134	with prefix.indent():	431	with prefix.indent():
135	load_kernel_args = [	432	load_kernel_args = [
	@@ -146,6 +443,24 @@ class DeferredNpuTritonCallWrapper(DeferredTritonCallWrapper):
146	prefix.writeline("}")	443	prefix.writeline("}")
147		444
148	def generate_launch_kernel(self, prefix, wrapper, kernel_var_name, params):	445	def generate_launch_kernel(self, prefix, wrapper, kernel_var_name, params):
		446	+ inductor_meta = params["inductor_meta"]
		447	+ if inductor_meta.get("group_enabled", False):
		448	+ grouped_plan = inductor_meta["grouped_candidate_plan"]
		449	+ self._generate_grouped_variant_launch(
		450	+ prefix,
		451	+ wrapper,
		452	+ params,
		453	+ grouped_plan,
		454	+ self._grouped_active_variants(grouped_plan),
		455	+ )
		456	+ return
		457	+ self._generate_single_kernel_launch(
		458	+ prefix, wrapper, kernel_var_name, params
		459	+ )
		460	+
		461	+ def _generate_single_kernel_launch(
		462	+ self, prefix, wrapper, kernel_var_name, params
		463	+ ):
149	triton_meta = params["triton_meta"]	464	triton_meta = params["triton_meta"]
150	arg_type_lookup = dict(zip(params["def_args"], self.arg_types))	465	arg_type_lookup = dict(zip(params["def_args"], self.arg_types))
151	for block_name in params["inductor_meta"].get("runtime_block_arg_names", ()):	466	for block_name in params["inductor_meta"].get("runtime_block_arg_names", ()):

Mtorch_npu/_inductor/runtime/triton_heuristics.py+79-0

	@@ -26,6 +26,7 @@ import triton
26	from torch._dynamo.testing import rand_strided	26	from torch._dynamo.testing import rand_strided
27	from torch._dynamo.utils import dynamo_timed	27	from torch._dynamo.utils import dynamo_timed
28	from torch._inductor import config	28	from torch._inductor import config
		29	+from torch._inductor.codecache import split_aot_inductor_output_path, write
29	from torch._inductor.compile_fx import clone_preserve_strides	30	from torch._inductor.compile_fx import clone_preserve_strides
30	from torch._inductor.runtime.autotune_cache import AutotuneCache	31	from torch._inductor.runtime.autotune_cache import AutotuneCache
31	from torch._inductor.runtime.benchmarking import benchmarker	32	from torch._inductor.runtime.benchmarking import benchmarker
	@@ -2184,6 +2185,80 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner):
2184	self._grouped_runtime_args_snapshot = ()	2185	self._grouped_runtime_args_snapshot = ()
2185	self._grouped_variant_launchers_initialized = False	2186	self._grouped_variant_launchers_initialized = False
2186		2187
		2188	+ @staticmethod
		2189	+ def _grouped_variant_load_meta(launcher) -> dict[str, Any]:
		2190	+ if not hasattr(launcher, "bin") or not hasattr(launcher.bin, "asm"):
		2191	+ raise RuntimeError("grouped cpp wrapper variant has no compiled binary")
		2192	+ metadata = launcher.bin.metadata
		2193	+ binary_path = launcher.bin.asm.get("cubin_path")
		2194	+ if not binary_path:
		2195	+ binary = launcher.bin.asm.get("npubin")
		2196	+ if binary is None:
		2197	+ raise RuntimeError("grouped cpp wrapper variant is missing npubin")
		2198	+ _, binary_path = write(
		2199	+ binary,
		2200	+ "cubin",
		2201	+ hash_type="cubin",
		2202	+ specified_dir=split_aot_inductor_output_path(
		2203	+ config.aot_inductor.output_path
		2204	+ )[0],
		2205	+ )
		2206	+ return {
		2207	+ "mangled_name": (
		2208	+ metadata.name
		2209	+ if hasattr(metadata, "name")
		2210	+ else metadata["name"]
		2211	+ ),
		2212	+ "shared_mem": (
		2213	+ launcher.bin.shared
		2214	+ if hasattr(launcher.bin, "shared")
		2215	+ else metadata.shared
		2216	+ ),
		2217	+ "mix_mode": metadata.mix_mode,
		2218	+ "parallel_mode": metadata.parallel_mode,
		2219	+ "force_simt_only": metadata.force_simt_only,
		2220	+ "shared_mem_dynamic_size": getattr(
		2221	+ metadata, "shared_mem_dynamic_size", 0
		2222	+ ),
		2223	+ "has_auto_blockify_blacklist_op": getattr(
		2224	+ metadata, "has_auto_blockify_blacklist_op", False
		2225	+ ),
		2226	+ "cubin_path": binary_path,
		2227	+ }
		2228	+
		2229	+ def _record_grouped_cpp_wrapper_plan(self) -> None:
		2230	+ best_by_group = {}
		2231	+ selected_variant_ids = OrderedSet()
		2232	+ for group_id in self.reachable_selection_keys:
		2233	+ candidate = self.best_candidate_map.get(group_id)
		2234	+ launcher = self.best_launcher_map.get(group_id)
		2235	+ if candidate is None or launcher is None:
		2236	+ raise RuntimeError(
		2237	+ f"grouped cpp wrapper is missing winner for group {group_id}"
		2238	+ )
		2239	+ variant_id = candidate["variant_id"]
		2240	+ best_by_group[str(group_id)] = {
		2241	+ "variant_id": variant_id,
		2242	+ "policy_id": candidate["policy_id"],
		2243	+ }
		2244	+ selected_variant_ids.add(variant_id)
		2245	+
		2246	+ for variant_id in selected_variant_ids:
		2247	+ launcher = self.variant_launcher_map.get(variant_id)
		2248	+ if launcher is None:
		2249	+ raise RuntimeError(
		2250	+ f"grouped cpp wrapper is missing launcher for variant {variant_id}"
		2251	+ )
		2252	+ self.candidate_plan["variants"][variant_id]["load_meta"] = (
		2253	+ self._grouped_variant_load_meta(launcher)
		2254	+ )
		2255	+
		2256	+ self.candidate_plan["best_by_group"] = best_by_group
		2257	+
		2258	+ def save_npu_kernel(self, input_stream, input_launcher):
		2259	+ self._record_grouped_cpp_wrapper_plan()
		2260	+ super().save_npu_kernel(input_stream, input_launcher)
		2261	+
2187	def _set_group_best_candidate(self, group_id, candidate, launcher):	2262	def _set_group_best_candidate(self, group_id, candidate, launcher):
2188	self.best_candidate_map[group_id] = candidate	2263	self.best_candidate_map[group_id] = candidate
2189	self.best_launcher_map[group_id] = launcher	2264	self.best_launcher_map[group_id] = launcher
	@@ -2600,6 +2675,10 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner):
2600	selected_config,	2675	selected_config,
2601	runtime_blocks,	2676	runtime_blocks,
2602	)	2677	)
		2678	+ if launcher.store_cubin and (
		2679	+ not benchmark_run or not self.cuda_kernel_saved
		2680	+ ):
		2681	+ self.save_gpu_kernel(stream, launcher)
2603	return launcher(	2682	return launcher(
2604	*self._build_runtime_launch_args(args, runtime_blocks),	2683	*self._build_runtime_launch_args(args, runtime_blocks),
2605	stream=stream,	2684	stream=stream,