feat(inductor): add experimental planned fast launch for NPU Python Wrapper-pytorch-AtomGit

Atest/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch.py+705-0

@@ -0,0 +1,705 @@
	1	+import importlib
	2	+import os
	3	+import sys
	4	+import types
	5	+import unittest
	6	+from contextlib import contextmanager
	7	+from unittest import mock
	8	+
	9	+
	10	+REPO_ROOT = os.path.abspath(
	11	+ os.path.join(os.path.dirname(__file__), "../../../..")
	12	+)
	13	+PACKAGE = "torch_npu._inductor.experimental.python_wrapper_fast_launch"
	14	+
	15	+
	16	+@contextmanager
	17	+def isolated_fast_launch(c_extension=None, , fast_launch=True, *config_values):
	18	+ for name in list(sys.modules):
	19	+ if name == PACKAGE or name.startswith(PACKAGE + "."):
	20	+ sys.modules.pop(name, None)
	21	+
	22	+ torch = types.ModuleType("torch")
	23	+ torch.__path__ = []
	24	+ autograd = types.ModuleType("torch.autograd")
	25	+ profiler = types.ModuleType("torch.autograd.profiler")
	26	+ profiler._is_profiler_enabled = False
	27	+ autograd.profiler = profiler
	28	+ torch.autograd = autograd
	29	+
	30	+ torch_npu = types.ModuleType("torch_npu")
	31	+ torch_npu.__path__ = [os.path.join(REPO_ROOT, "torch_npu")]
	32	+ inductor = types.ModuleType("torch_npu._inductor")
	33	+ inductor.__path__ = [os.path.join(REPO_ROOT, "torch_npu", "_inductor")]
	34	+ config = types.ModuleType("torch_npu._inductor.config")
	35	+ defaults = {
	36	+ "dump_fx_graph": False,
	37	+ "check_accuracy": False,
	38	+ "enable_fast_launch": fast_launch,
	39	+ }
	40	+ defaults.update(config_values)
	41	+ for name, value in defaults.items():
	42	+ setattr(config, name, value)
	43	+ inductor.config = config
	44	+
	45	+ modules = {
	46	+ "torch": torch,
	47	+ "torch.autograd": autograd,
	48	+ "torch.autograd.profiler": profiler,
	49	+ "torch_npu": torch_npu,
	50	+ "torch_npu._inductor": inductor,
	51	+ "torch_npu._inductor.config": config,
	52	+ }
	53	+ if c_extension is not None:
	54	+ modules["torch_npu._C"] = c_extension
	55	+ torch_npu._C = c_extension
	56	+
	57	+ with mock.patch.dict(sys.modules, modules):
	58	+ try:
	59	+ yield profiler
	60	+ finally:
	61	+ for name in list(sys.modules):
	62	+ if name == PACKAGE or name.startswith(PACKAGE + "."):
	63	+ sys.modules.pop(name, None)
	64	+
	65	+
	66	+class FakeLauncher:
	67	+ def __init__(self, *, arg_kinds=("tensor", "i32")):
	68	+ self.config = types.SimpleNamespace(found_by_coordesc=True)
	69	+ self.store_cubin = False
	70	+ self._npu_fast_launch_kernel_name = "triton_poi_fused_0"
	71	+ self._npu_fast_launch_kernel_stub = 1234
	72	+ self._npu_fast_launch_kernel_stub_owner = object()
	73	+ self._npu_fast_launch_arg_kinds = arg_kinds
	74	+ self._npu_fast_launch_grid_exprs = ("xnumel", "1", "1")
	75	+ self._npu_fast_launch_get_grid = lambda *args: (2, 1, 1)
	76	+ self._npu_fast_launch_enable_simt = False
	77	+ self._npu_fast_launch_shared_mem_dynamic_size = 0
	78	+ self._npu_fast_launch_force_simt_only = False
	79	+ self._npu_fast_launch_target_support_ffts = False
	80	+ self._npu_fast_launch_workspace_size = -1
	81	+ self._npu_fast_launch_lock_num = -1
	82	+ self._npu_fast_launch_device_print_enabled = False
	83	+ self._npu_fast_launch_enter_hook = None
	84	+ self._npu_fast_launch_exit_hook = None
	85	+ self._npu_fast_launch_enter_hook_callbacks = ()
	86	+ self._npu_fast_launch_exit_hook_callbacks = ()
	87	+ self.calls = []
	88	+
	89	+ def __call__(self, *args, stream):
	90	+ metadata = {"args": args, "stream": stream}
	91	+ if self._npu_fast_launch_enter_hook is not None:
	92	+ self._npu_fast_launch_enter_hook(metadata)
	93	+ self.calls.append((args, stream))
	94	+ if self._npu_fast_launch_exit_hook is not None:
	95	+ self._npu_fast_launch_exit_hook(metadata)
	96	+ return "launcher"
	97	+
	98	+
	99	+class HookChain:
	100	+ def __init__(self, *, reversed=False):
	101	+ self.calls = []
	102	+ self.reversed = reversed
	103	+
	104	+ def add(self, callback):
	105	+ if callback not in self.calls:
	106	+ self.calls.append(callback)
	107	+
	108	+ def remove(self, callback):
	109	+ if callback in self.calls:
	110	+ self.calls.remove(callback)
	111	+
	112	+ def __call__(self, args, *kwargs):
	113	+ callbacks = reversed(self.calls) if self.reversed else self.calls
	114	+ for callback in callbacks:
	115	+ callback(args, *kwargs)
	116	+
	117	+
	118	+def set_launcher_hook_chains(launcher):
	119	+ launcher._npu_fast_launch_enter_hook = HookChain()
	120	+ launcher._npu_fast_launch_exit_hook = HookChain(reversed=True)
	121	+ launcher._npu_fast_launch_enter_hook_callbacks = (
	122	+ launcher._npu_fast_launch_enter_hook.calls
	123	+ )
	124	+ launcher._npu_fast_launch_exit_hook_callbacks = (
	125	+ launcher._npu_fast_launch_exit_hook.calls
	126	+ )
	127	+
	128	+
	129	+class FakeTensor:
	130	+ def __init__(self, pointer=1):
	131	+ self.pointer = pointer
	132	+
	133	+ def data_ptr(self):
	134	+ return self.pointer
	135	+
	136	+
	137	+class FakeAutotuner:
	138	+ def __init__(self, launcher):
	139	+ self.launchers = [launcher]
	140	+ self.best_launcher = None
	141	+ self.best_runtime_blocks = ()
	142	+ self.inductor_meta = {}
	143	+ self.triton_interpret = False
	144	+ self.dump_launch_params = False
	145	+ self.run_calls = []
	146	+
	147	+ def _build_runtime_launch_args(self, args, runtime_blocks):
	148	+ return (args, runtime_blocks)
	149	+
	150	+ def run(self, args, stream, benchmark_run=False, *kwargs):
	151	+ self.run_calls.append((args, stream, benchmark_run, kwargs))
	152	+ self.best_launcher = self.launchers[0]
	153	+ return "fallback"
	154	+
	155	+
	156	+def metadata(
	157	+ *,
	158	+ eligible=True,
	159	+ schema_state=None,
	160	+ arg_kinds=("tensor", "i32"),
	161	+ runtime_arg_count=2,
	162	+):
	163	+ result = {
	164	+ "graph_id": "graph-0",
	165	+ "callsite_id": "graph-0:0",
	166	+ "kernel_name": "triton_poi_fused_0",
	167	+ "schema_hash": "schema-0",
	168	+ "arg_kinds": arg_kinds,
	169	+ "runtime_arg_count": runtime_arg_count,
	170	+ "eligible": eligible,
	171	+ "fallback_reason": None if eligible else "unsupported",
	172	+ }
	173	+ if schema_state is not None:
	174	+ result["schema_state"] = schema_state
	175	+ return result
	176	+
	177	+
	178	+class TestNPUFastLaunch(unittest.TestCase):
	179	+ def test_disabled_environment_forces_cached_wrapper_to_full_entry(self):
	180	+ extension = types.SimpleNamespace(
	181	+ _npu_inductor_make_fast_launch_plan=mock.Mock(),
	182	+ _npu_inductor_fast_launch_with_plan=mock.Mock(),
	183	+ )
	184	+ with isolated_fast_launch(extension, fast_launch=False):
	185	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	186	+ launcher = FakeLauncher()
	187	+ autotuner = FakeAutotuner(launcher)
	188	+ autotuner.best_launcher = launcher
	189	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	190	+ tensor = FakeTensor()
	191	+
	192	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	193	+ self.assertEqual(bound(tensor, 4, stream=99), "fallback")
	194	+
	195	+ self.assertEqual(len(autotuner.run_calls), 2)
	196	+ self.assertEqual(launcher.calls, [])
	197	+ extension._npu_inductor_make_fast_launch_plan.assert_not_called()
	198	+ extension._npu_inductor_fast_launch_with_plan.assert_not_called()
	199	+
	200	+ def test_codegen_missing_signature_is_promotable_incomplete_schema(self):
	201	+ with isolated_fast_launch():
	202	+ codegen = importlib.import_module(f"{PACKAGE}.codegen")
	203	+ result = codegen.build_callsite_metadata(
	204	+ kernel_name="triton_poi_fused_0",
	205	+ call_args=("buf0", "xnumel"),
	206	+ triton_meta={},
	207	+ graph_id="graph-0",
	208	+ callsite_index=0,
	209	+ )
	210	+
	211	+ self.assertEqual(result["schema_state"], "incomplete")
	212	+ self.assertEqual(result["schema_reason"], "codegen_signature_missing")
	213	+ self.assertEqual(result["runtime_arg_count"], 2)
	214	+ self.assertTrue(result["eligible"])
	215	+
	216	+ def test_codegen_complete_schema_keeps_ordered_arg_kinds(self):
	217	+ with isolated_fast_launch():
	218	+ codegen = importlib.import_module(f"{PACKAGE}.codegen")
	219	+ result = codegen.build_callsite_metadata(
	220	+ kernel_name="triton_poi_fused_0",
	221	+ call_args=("buf0", "xnumel"),
	222	+ triton_meta={
	223	+ "signature": {"in_ptr0": "*fp32", "xnumel": "i32"},
	224	+ "constants": {},
	225	+ },
	226	+ graph_id="graph-0",
	227	+ callsite_index=0,
	228	+ )
	229	+
	230	+ self.assertEqual(result["schema_state"], "complete")
	231	+ self.assertEqual(result["arg_kinds"], ("tensor", "i32"))
	232	+
	233	+ def test_plan_owns_binary_and_launches_with_tuple_args(self):
	234	+ calls = []
	235	+
	236	+ class Plan:
	237	+ pass
	238	+
	239	+ def make_plan(*args):
	240	+ calls.append(("make", args))
	241	+ return Plan()
	242	+
	243	+ def launch(*args):
	244	+ calls.append(("launch", args))
	245	+
	246	+ extension = types.SimpleNamespace(
	247	+ _npu_inductor_make_fast_launch_plan=make_plan,
	248	+ _npu_inductor_fast_launch_with_plan=launch,
	249	+ )
	250	+ with isolated_fast_launch(extension):
	251	+ backend = importlib.import_module(f"{PACKAGE}.backend")
	252	+ launcher = FakeLauncher()
	253	+ tensor = FakeTensor()
	254	+ planned = backend.build_planned_fast_launch(
	255	+ launcher,
	256	+ metadata(),
	257	+ canonical_args=(tensor, 3),
	258	+ runtime_arg_count=2,
	259	+ )
	260	+ planned((tensor, 3), stream=99)
	261	+
	262	+ self.assertIs(planned.plan._owner, launcher._npu_fast_launch_kernel_stub_owner)
	263	+ self.assertEqual(calls[1][0], "launch")
	264	+ self.assertEqual(calls[1][1][1:5], (99, 2, 1, 1))
	265	+ self.assertEqual(calls[1][1][5], (tensor, 3))
	266	+ self.assertFalse(calls[0][1][-1])
	267	+
	268	+ def test_plan_forwards_ffts_abi_requirement(self):
	269	+ calls = []
	270	+
	271	+ class Plan:
	272	+ pass
	273	+
	274	+ extension = types.SimpleNamespace(
	275	+ _npu_inductor_make_fast_launch_plan=lambda *args: (
	276	+ calls.append(args) or Plan()
	277	+ ),
	278	+ _npu_inductor_fast_launch_with_plan=lambda *args: None,
	279	+ )
	280	+ with isolated_fast_launch(extension):
	281	+ backend = importlib.import_module(f"{PACKAGE}.backend")
	282	+ launcher = FakeLauncher()
	283	+ launcher._npu_fast_launch_target_support_ffts = True
	284	+ backend.build_planned_fast_launch(
	285	+ launcher,
	286	+ metadata(),
	287	+ canonical_args=(FakeTensor(), 3),
	288	+ runtime_arg_count=2,
	289	+ )
	290	+
	291	+ self.assertTrue(calls[0][-1])
	292	+
	293	+ def test_launcher_requiring_hidden_resources_is_negative(self):
	294	+ extension = types.SimpleNamespace(
	295	+ _npu_inductor_make_fast_launch_plan=mock.Mock(),
	296	+ _npu_inductor_fast_launch_with_plan=mock.Mock(),
	297	+ )
	298	+ with isolated_fast_launch(extension):
	299	+ backend = importlib.import_module(f"{PACKAGE}.backend")
	300	+ types_module = importlib.import_module(f"{PACKAGE}.types")
	301	+ for attribute, reason in (
	302	+ ("_npu_fast_launch_workspace_size", "launcher_workspace_required"),
	303	+ ("_npu_fast_launch_lock_num", "launcher_sync_block_lock_required"),
	304	+ ):
	305	+ launcher = FakeLauncher()
	306	+ setattr(launcher, attribute, 1)
	307	+ with self.subTest(attribute=attribute), self.assertRaises(
	308	+ types_module.FastLaunchPlanUnavailable
	309	+ ) as error:
	310	+ backend.build_planned_fast_launch(
	311	+ launcher,
	312	+ metadata(),
	313	+ canonical_args=(FakeTensor(), 3),
	314	+ runtime_arg_count=2,
	315	+ )
	316	+ self.assertEqual(str(error.exception), reason)
	317	+
	318	+ extension._npu_inductor_make_fast_launch_plan.assert_not_called()
	319	+
	320	+ def test_warmup_then_promotes_to_planned_backend(self):
	321	+ launches = []
	322	+
	323	+ class Plan:
	324	+ pass
	325	+
	326	+ extension = types.SimpleNamespace(
	327	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	328	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	329	+ )
	330	+ with isolated_fast_launch(extension):
	331	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	332	+ launcher = FakeLauncher()
	333	+ autotuner = FakeAutotuner(launcher)
	334	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	335	+ tensor = FakeTensor()
	336	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	337	+ self.assertIsNone(bound(tensor, 4, stream=99))
	338	+
	339	+ self.assertEqual(len(autotuner.run_calls), 1)
	340	+ self.assertEqual(len(launches), 1)
	341	+ self.assertEqual(launches[0][5], (tensor, 4))
	342	+
	343	+ def test_cold_coordinate_descent_promotes_after_autotune(self):
	344	+ launches = []
	345	+
	346	+ class Plan:
	347	+ pass
	348	+
	349	+ extension = types.SimpleNamespace(
	350	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	351	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	352	+ )
	353	+ with isolated_fast_launch(extension):
	354	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	355	+ launcher = FakeLauncher()
	356	+ launcher.config.found_by_coordesc = False
	357	+ autotuner = FakeAutotuner(launcher)
	358	+ autotuner.inductor_meta["coordinate_descent_tuning"] = True
	359	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	360	+ tensor = FakeTensor()
	361	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	362	+ self.assertIsNone(bound(tensor, 4, stream=99))
	363	+
	364	+ self.assertEqual(len(autotuner.run_calls), 1)
	365	+ self.assertEqual(len(launches), 1)
	366	+ self.assertEqual(launches[0][5], (tensor, 4))
	367	+
	368	+ def test_grouped_autotuner_keeps_original_entry(self):
	369	+ extension = types.SimpleNamespace(
	370	+ _npu_inductor_make_fast_launch_plan=mock.Mock(),
	371	+ _npu_inductor_fast_launch_with_plan=mock.Mock(),
	372	+ )
	373	+ with isolated_fast_launch(extension):
	374	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	375	+ launcher = FakeLauncher()
	376	+ autotuner = FakeAutotuner(launcher)
	377	+ autotuner.inductor_meta["group_enabled"] = True
	378	+ autotuner.best_launcher_map = {}
	379	+ call_slot = [None]
	380	+ call = bind.bind_python_wrapper_kernel_fast(
	381	+ metadata(),
	382	+ autotuner,
	383	+ call_slot=call_slot,
	384	+ )
	385	+ tensor = FakeTensor()
	386	+ self.assertIs(call_slot[0], call)
	387	+ self.assertNotIsInstance(call, bind.BoundFastLaunch)
	388	+ self.assertEqual(call(tensor, 3, stream=99), "fallback")
	389	+ self.assertEqual(call(tensor, 4, stream=99), "fallback")
	390	+
	391	+ self.assertEqual(len(autotuner.run_calls), 2)
	392	+ extension._npu_inductor_make_fast_launch_plan.assert_not_called()
	393	+ extension._npu_inductor_fast_launch_with_plan.assert_not_called()
	394	+
	395	+ def test_negative_cache_calls_stable_original_launcher(self):
	396	+ extension = types.SimpleNamespace()
	397	+ with isolated_fast_launch(extension):
	398	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	399	+ launcher = FakeLauncher()
	400	+ autotuner = FakeAutotuner(launcher)
	401	+ bound = bind.BoundFastLaunch(autotuner, metadata(eligible=False))
	402	+ tensor = FakeTensor()
	403	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	404	+ self.assertEqual(bound(tensor, 4, stream=99), "launcher")
	405	+
	406	+ self.assertEqual(len(autotuner.run_calls), 1)
	407	+ self.assertEqual(launcher.calls, [((tensor, 4), 99)])
	408	+
	409	+ def test_profiler_forces_full_entry(self):
	410	+ extension = types.SimpleNamespace()
	411	+ with isolated_fast_launch(extension) as profiler:
	412	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	413	+ launcher = FakeLauncher()
	414	+ autotuner = FakeAutotuner(launcher)
	415	+ autotuner.best_launcher = launcher
	416	+ profiler._is_profiler_enabled = True
	417	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	418	+ self.assertEqual(bound(FakeTensor(), 3, stream=99), "fallback")
	419	+
	420	+ self.assertEqual(len(autotuner.run_calls), 1)
	421	+
	422	+ def test_launch_hook_callback_source_tracks_hook_chain_mutation(self):
	423	+ with isolated_fast_launch():
	424	+ launcher_module = importlib.import_module(f"{PACKAGE}.launcher")
	425	+ chain = HookChain()
	426	+ callback = lambda metadata: None
	427	+
	428	+ callbacks = launcher_module._launch_hook_callbacks(chain)
	429	+ self.assertIs(callbacks, chain.calls)
	430	+ self.assertFalse(callbacks)
	431	+
	432	+ chain.add(callback)
	433	+ self.assertTrue(callbacks)
	434	+ chain.remove(callback)
	435	+ self.assertFalse(callbacks)
	436	+
	437	+ legacy_hook = lambda metadata: None
	438	+ self.assertEqual(
	439	+ launcher_module._launch_hook_callbacks(legacy_hook),
	440	+ (legacy_hook,),
	441	+ )
	442	+
	443	+ def test_empty_triton_launch_hook_chain_does_not_block_fast_path(self):
	444	+ launches = []
	445	+
	446	+ class Plan:
	447	+ pass
	448	+
	449	+ extension = types.SimpleNamespace(
	450	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	451	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	452	+ )
	453	+ with isolated_fast_launch(extension):
	454	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	455	+ launcher = FakeLauncher()
	456	+ set_launcher_hook_chains(launcher)
	457	+ autotuner = FakeAutotuner(launcher)
	458	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	459	+ tensor = FakeTensor()
	460	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	461	+ self.assertIsNone(bound(tensor, 4, stream=99))
	462	+
	463	+ self.assertEqual(len(autotuner.run_calls), 1)
	464	+ self.assertEqual(len(launches), 1)
	465	+ self.assertEqual(launches[0][5], (tensor, 4))
	466	+
	467	+ def test_active_triton_launch_hooks_use_original_launcher(self):
	468	+ launches = []
	469	+ hook_events = []
	470	+
	471	+ class Plan:
	472	+ pass
	473	+
	474	+ extension = types.SimpleNamespace(
	475	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	476	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	477	+ )
	478	+ with isolated_fast_launch(extension):
	479	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	480	+ launcher = FakeLauncher()
	481	+ set_launcher_hook_chains(launcher)
	482	+ launcher._npu_fast_launch_enter_hook.add(
	483	+ lambda metadata: hook_events.append(("enter", metadata))
	484	+ )
	485	+ launcher._npu_fast_launch_exit_hook.add(
	486	+ lambda metadata: hook_events.append(("exit", metadata))
	487	+ )
	488	+ autotuner = FakeAutotuner(launcher)
	489	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	490	+ tensor = FakeTensor()
	491	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	492	+ self.assertEqual(bound(tensor, 4, stream=100), "launcher")
	493	+
	494	+ self.assertEqual(len(autotuner.run_calls), 1)
	495	+ self.assertEqual(launches, [])
	496	+ self.assertEqual(launcher.calls, [((tensor, 4), 100)])
	497	+ self.assertEqual([event[0] for event in hook_events], ["enter", "exit"])
	498	+ self.assertEqual(hook_events[0][1]["stream"], 100)
	499	+
	500	+ def test_triton_launch_hook_activation_temporarily_falls_back(self):
	501	+ launches = []
	502	+ hook_events = []
	503	+
	504	+ class Plan:
	505	+ pass
	506	+
	507	+ extension = types.SimpleNamespace(
	508	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	509	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	510	+ )
	511	+ with isolated_fast_launch(extension):
	512	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	513	+ launcher = FakeLauncher()
	514	+ set_launcher_hook_chains(launcher)
	515	+ autotuner = FakeAutotuner(launcher)
	516	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	517	+ tensor = FakeTensor()
	518	+
	519	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	520	+ self.assertIsNone(bound(tensor, 4, stream=100))
	521	+
	522	+ enter = lambda metadata: hook_events.append(("enter", metadata))
	523	+ exit = lambda metadata: hook_events.append(("exit", metadata))
	524	+ launcher._npu_fast_launch_enter_hook.add(enter)
	525	+ launcher._npu_fast_launch_exit_hook.add(exit)
	526	+ self.assertEqual(bound(tensor, 5, stream=101), "launcher")
	527	+
	528	+ launcher._npu_fast_launch_enter_hook.remove(enter)
	529	+ launcher._npu_fast_launch_exit_hook.remove(exit)
	530	+ self.assertIsNone(bound(tensor, 6, stream=102))
	531	+
	532	+ self.assertEqual(len(autotuner.run_calls), 1)
	533	+ self.assertEqual(
	534	+ [launch[5] for launch in launches],
	535	+ [(tensor, 4), (tensor, 6)],
	536	+ )
	537	+ self.assertEqual(launcher.calls, [((tensor, 5), 101)])
	538	+ self.assertEqual([event[0] for event in hook_events], ["enter", "exit"])
	539	+
	540	+ def test_dynamic_grid_failure_does_not_poison_plan(self):
	541	+ launches = []
	542	+
	543	+ class Plan:
	544	+ pass
	545	+
	546	+ extension = types.SimpleNamespace(
	547	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	548	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	549	+ )
	550	+ with isolated_fast_launch(extension):
	551	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	552	+ launcher = FakeLauncher()
	553	+ launcher._npu_fast_launch_get_grid = lambda tensor, size: (size, 1, 1)
	554	+ autotuner = FakeAutotuner(launcher)
	555	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	556	+ tensor = FakeTensor()
	557	+ self.assertEqual(bound(tensor, 2, stream=99), "fallback")
	558	+ self.assertEqual(bound(tensor, 70_000, stream=99), "fallback")
	559	+ self.assertIsNone(bound(tensor, 3, stream=100))
	560	+
	561	+ self.assertEqual(len(autotuner.run_calls), 2)
	562	+ self.assertEqual(len(launches), 1)
	563	+ self.assertEqual(launches[0][1:5], (100, 3, 1, 1))
	564	+
	565	+ def test_backend_error_never_replays_fallback(self):
	566	+ class Plan:
	567	+ pass
	568	+
	569	+ def fail_after_boundary(*args):
	570	+ raise RuntimeError("rtKernelLaunch failed")
	571	+
	572	+ extension = types.SimpleNamespace(
	573	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	574	+ _npu_inductor_fast_launch_with_plan=fail_after_boundary,
	575	+ )
	576	+ with isolated_fast_launch(extension):
	577	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	578	+ types_module = importlib.import_module(f"{PACKAGE}.types")
	579	+ launcher = FakeLauncher()
	580	+ autotuner = FakeAutotuner(launcher)
	581	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	582	+ tensor = FakeTensor()
	583	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	584	+ with self.assertRaises(types_module.FastLaunchError) as error:
	585	+ bound(tensor, 4, stream=99)
	586	+ self.assertTrue(error.exception.backend_submitted)
	587	+
	588	+ self.assertEqual(len(autotuner.run_calls), 1)
	589	+
	590	+ def test_incomplete_schema_is_completed_from_launcher_abi(self):
	591	+ launches = []
	592	+
	593	+ class Plan:
	594	+ pass
	595	+
	596	+ extension = types.SimpleNamespace(
	597	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	598	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	599	+ )
	600	+ with isolated_fast_launch(extension):
	601	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	602	+ launcher = FakeLauncher(
	603	+ arg_kinds=("tensor", "tensor", "i32", "bool", "i32")
	604	+ )
	605	+ autotuner = FakeAutotuner(launcher)
	606	+ autotuner.best_runtime_blocks = (16,)
	607	+ bound = bind.BoundFastLaunch(
	608	+ autotuner,
	609	+ metadata(
	610	+ schema_state="incomplete",
	611	+ arg_kinds=(),
	612	+ runtime_arg_count=4,
	613	+ ),
	614	+ )
	615	+ first = FakeTensor(1)
	616	+ second = FakeTensor(2)
	617	+ self.assertEqual(bound(first, second, 7, True, stream=99), "fallback")
	618	+ self.assertIsNone(bound(first, second, 8, False, stream=100))
	619	+
	620	+ self.assertEqual(len(launches), 1)
	621	+ self.assertEqual(launches[0][5], (first, second, 8, False, 16))
	622	+
	623	+ def test_complete_codegen_launcher_schema_conflict_is_negative(self):
	624	+ launches = []
	625	+
	626	+ class Plan:
	627	+ pass
	628	+
	629	+ extension = types.SimpleNamespace(
	630	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	631	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	632	+ )
	633	+ with isolated_fast_launch(extension):
	634	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	635	+ launcher = FakeLauncher(arg_kinds=("tensor", "i64"))
	636	+ autotuner = FakeAutotuner(launcher)
	637	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	638	+ tensor = FakeTensor()
	639	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	640	+ self.assertEqual(bound(tensor, 4, stream=99), "launcher")
	641	+
	642	+ self.assertEqual(launches, [])
	643	+ self.assertEqual(len(autotuner.run_calls), 1)
	644	+
	645	+ def test_negative_hit_skips_launcher_stability_recheck(self):
	646	+ extension = types.SimpleNamespace()
	647	+ with isolated_fast_launch(extension):
	648	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	649	+ launcher = FakeLauncher()
	650	+ autotuner = FakeAutotuner(launcher)
	651	+ bound = bind.BoundFastLaunch(autotuner, metadata(eligible=False))
	652	+ tensor = FakeTensor()
	653	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	654	+ with mock.patch.object(
	655	+ bind.BoundFastLaunch,
	656	+ "_stable_launcher",
	657	+ side_effect=AssertionError("stable launcher was rechecked"),
	658	+ ):
	659	+ self.assertEqual(bound(tensor, 4, stream=99), "launcher")
	660	+
	661	+ def test_negative_cache_is_invalidated_when_launcher_changes(self):
	662	+ extension = types.SimpleNamespace()
	663	+ with isolated_fast_launch(extension):
	664	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	665	+ first_launcher = FakeLauncher()
	666	+ second_launcher = FakeLauncher()
	667	+ autotuner = FakeAutotuner(first_launcher)
	668	+ bound = bind.BoundFastLaunch(autotuner, metadata(eligible=False))
	669	+ tensor = FakeTensor()
	670	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	671	+ autotuner.launchers = [second_launcher]
	672	+ autotuner.best_launcher = second_launcher
	673	+ self.assertEqual(bound(tensor, 4, stream=99), "fallback")
	674	+
	675	+ self.assertEqual(first_launcher.calls, [])
	676	+ self.assertEqual(len(autotuner.run_calls), 2)
	677	+
	678	+ def test_no_runtime_blocks_reuses_args_without_builder_call(self):
	679	+ launches = []
	680	+
	681	+ class Plan:
	682	+ pass
	683	+
	684	+ extension = types.SimpleNamespace(
	685	+ _npu_inductor_make_fast_launch_plan=lambda *args: Plan(),
	686	+ _npu_inductor_fast_launch_with_plan=lambda *args: launches.append(args),
	687	+ )
	688	+ with isolated_fast_launch(extension):
	689	+ bind = importlib.import_module(f"{PACKAGE}.bind")
	690	+ launcher = FakeLauncher()
	691	+ autotuner = FakeAutotuner(launcher)
	692	+ autotuner._build_runtime_launch_args = mock.Mock(
	693	+ side_effect=AssertionError("runtime arg builder should be cold")
	694	+ )
	695	+ bound = bind.BoundFastLaunch(autotuner, metadata())
	696	+ tensor = FakeTensor()
	697	+ self.assertEqual(bound(tensor, 3, stream=99), "fallback")
	698	+ self.assertIsNone(bound(tensor, 4, stream=99))
	699	+
	700	+ self.assertEqual(len(launches), 1)
	701	+ autotuner._build_runtime_launch_args.assert_not_called()
	702	+
	703	+
	704	+if __name__ == "__main__":
	705	+ unittest.main()

Atest/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch_patch.py+283-0

@@ -0,0 +1,283 @@
	1	+import importlib
	2	+import os
	3	+import sys
	4	+import types
	5	+import unittest
	6	+from contextlib import contextmanager
	7	+from unittest import mock
	8	+
	9	+
	10	+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))
	11	+PACKAGE = "torch_npu._inductor.experimental.python_wrapper_fast_launch"
	12	+
	13	+
	14	+@contextmanager
	15	+def isolated_patch_imports(*, fast_launch=False):
	16	+ for name in list(sys.modules):
	17	+ if name == PACKAGE or name.startswith(PACKAGE + "."):
	18	+ sys.modules.pop(name, None)
	19	+
	20	+ torch_npu = types.ModuleType("torch_npu")
	21	+ torch_npu.__path__ = [os.path.join(REPO_ROOT, "torch_npu")]
	22	+ inductor = types.ModuleType("torch_npu._inductor")
	23	+ inductor.__path__ = [os.path.join(REPO_ROOT, "torch_npu", "_inductor")]
	24	+ config = types.ModuleType("torch_npu._inductor.config")
	25	+ config.enable_fast_launch = fast_launch
	26	+ inductor.config = config
	27	+
	28	+ with mock.patch.dict(
	29	+ sys.modules,
	30	+ {
	31	+ "torch_npu": torch_npu,
	32	+ "torch_npu._inductor": inductor,
	33	+ "torch_npu._inductor.config": config,
	34	+ },
	35	+ ):
	36	+ try:
	37	+ yield
	38	+ finally:
	39	+ for name in list(sys.modules):
	40	+ if name == PACKAGE or name.startswith(PACKAGE + "."):
	41	+ sys.modules.pop(name, None)
	42	+
	43	+
	44	+class _DebugPrinter:
	45	+ def __enter__(self):
	46	+ return self
	47	+
	48	+ def __exit__(self, exc_type, exc_value, traceback):
	49	+ return False
	50	+
	51	+ def set_printer_args(self, *args):
	52	+ self.args = args
	53	+
	54	+
	55	+def _fake_wrapper_module():
	56	+ events = []
	57	+
	58	+ class FakeWrapper:
	59	+ def __init__(self):
	60	+ self.imports = types.SimpleNamespace(
	61	+ writeline=lambda line: events.append(("import", line))
	62	+ )
	63	+ self.kernel_autotune_names = set()
	64	+
	65	+ def write_triton_header_once(self):
	66	+ events.append(("header",))
	67	+ return "header"
	68	+
	69	+ def prepare_triton_kernel_call(self, call_args):
	70	+ return tuple(call_args)
	71	+
	72	+ def generate_kernel_call(self, args, *kwargs):
	73	+ events.append(("original", args, kwargs))
	74	+ return "original"
	75	+
	76	+ graph = types.SimpleNamespace(
	77	+ cpp_wrapper=False,
	78	+ get_current_device_or_throw=lambda: types.SimpleNamespace(index=0),
	79	+ wrapper_code=types.SimpleNamespace(debug_printer=_DebugPrinter()),
	80	+ )
	81	+ module = types.SimpleNamespace(
	82	+ NPUPythonWrapperCodeGen=FakeWrapper,
	83	+ PythonWrapperCodegen=types.SimpleNamespace(
	84	+ write_get_raw_stream=lambda owner, index, graph: f"stream{index}"
	85	+ ),
	86	+ V=types.SimpleNamespace(graph=graph),
	87	+ config=types.SimpleNamespace(
	88	+ triton=types.SimpleNamespace(autotune_at_compile_time=False)
	89	+ ),
	90	+ is_multi_stream=lambda: False,
	91	+ _is_codegen_graph_partition_subgraph=lambda owner: False,
	92	+ )
	93	+ return module, FakeWrapper, events
	94	+
	95	+
	96	+def _fake_triton_module():
	97	+ class FakeGrid:
	98	+ prefix = ()
	99	+ x_grid = "xnumel"
	100	+ y_grid = "1"
	101	+ z_grid = "1"
	102	+
	103	+ class FakeGridExpr:
	104	+ @staticmethod
	105	+ def from_meta(inductor_meta, cfg):
	106	+ return FakeGrid()
	107	+
	108	+ class FakeCompileResult:
	109	+ def __init__(self):
	110	+ self.config = types.SimpleNamespace(kwargs={})
	111	+ self.compile_meta = {
	112	+ "constants": {},
	113	+ "signature": {"xnumel": "i64"},
	114	+ }
	115	+ fn = types.SimpleNamespace(
	116	+ __name__="triton_poi_fused_0",
	117	+ arg_names=("xnumel",),
	118	+ constexprs=set(),
	119	+ )
	120	+ self.kernel = types.SimpleNamespace(
	121	+ src=types.SimpleNamespace(fn=fn),
	122	+ metadata=types.SimpleNamespace(),
	123	+ )
	124	+ self.inductor_meta = {"kernel_name": "triton_poi_fused_0"}
	125	+
	126	+ def make_launcher(self):
	127	+ scope = {
	128	+ "function": 1234,
	129	+ "launch_enter_hook": None,
	130	+ "launch_exit_hook": None,
	131	+ }
	132	+ exec(
	133	+ "def launcher(xnumel, stream):\n return (xnumel, stream)\n",
	134	+ scope,
	135	+ )
	136	+ return scope["launcher"]
	137	+
	138	+ module = types.SimpleNamespace(
	139	+ TritonCompileResultNpu=FakeCompileResult,
	140	+ OrderedSet=set,
	141	+ triton_version_uses_attrs_dict=lambda: False,
	142	+ config_to_dict=lambda cfg: dict(cfg.kwargs),
	143	+ filter_launcher_def_args=lambda names, cfg, none, runtime: [
	144	+ name for name in names if name not in cfg and name not in none
	145	+ ],
	146	+ GridExpr=FakeGridExpr,
	147	+ GridExprNpu=types.SimpleNamespace(),
	148	+ )
	149	+ return module, FakeCompileResult
	150	+
	151	+
	152	+class TestNPUFastLaunchPatch(unittest.TestCase):
	153	+ def test_default_off_does_not_resolve_patch_targets(self):
	154	+ with isolated_patch_imports():
	155	+ fast_patch = importlib.import_module(f"{PACKAGE}.patch")
	156	+ fast_patch._resolve_wrapper_module = mock.Mock(side_effect=AssertionError)
	157	+ fast_patch._resolve_triton_heuristics_module = mock.Mock(
	158	+ side_effect=AssertionError
	159	+ )
	160	+
	161	+ self.assertFalse(fast_patch.patch_fast_launch())
	162	+
	163	+ fast_patch._resolve_wrapper_module.assert_not_called()
	164	+ fast_patch._resolve_triton_heuristics_module.assert_not_called()
	165	+
	166	+ def test_enabled_patch_installs_once_and_routes_codegen(self):
	167	+ wrapper_module, wrapper_cls, events = _fake_wrapper_module()
	168	+ triton_module, triton_cls = _fake_triton_module()
	169	+ emitted = []
	170	+
	171	+ class FakeEmitter:
	172	+ def __init__(self, owner):
	173	+ self.owner = owner
	174	+
	175	+ def emit_triton_call(self, **kwargs):
	176	+ emitted.append(kwargs)
	177	+
	178	+ attach_metadata = mock.Mock()
	179	+ with (
	180	+ isolated_patch_imports(fast_launch=True),
	181	+ mock.patch.dict(
	182	+ sys.modules,
	183	+ {
	184	+ f"{PACKAGE}.wrapper_codegen": types.SimpleNamespace(
	185	+ FastLaunchWrapperEmitter=FakeEmitter
	186	+ ),
	187	+ f"{PACKAGE}.launcher": types.SimpleNamespace(
	188	+ attach_python_wrapper_launcher_metadata=attach_metadata
	189	+ ),
	190	+ },
	191	+ ),
	192	+ ):
	193	+ fast_patch = importlib.import_module(f"{PACKAGE}.patch")
	194	+ fast_patch._resolve_wrapper_module = mock.Mock(return_value=wrapper_module)
	195	+ fast_patch._resolve_triton_heuristics_module = mock.Mock(
	196	+ return_value=triton_module
	197	+ )
	198	+ original_generate = wrapper_cls.generate_kernel_call
	199	+ original_make_launcher = triton_cls.make_launcher
	200	+
	201	+ self.assertTrue(fast_patch.patch_fast_launch())
	202	+ self.assertTrue(fast_patch.patch_fast_launch())
	203	+
	204	+ self.assertIsNot(wrapper_cls.generate_kernel_call, original_generate)
	205	+ self.assertIsNot(triton_cls.make_launcher, original_make_launcher)
	206	+ self.assertEqual(fast_patch._resolve_wrapper_module.call_count, 1)
	207	+ self.assertEqual(
	208	+ fast_patch._resolve_triton_heuristics_module.call_count,
	209	+ 1,
	210	+ )
	211	+
	212	+ wrapper = wrapper_cls()
	213	+ wrapper.generate_kernel_call(
	214	+ "triton_poi_fused_0",
	215	+ ("xnumel",),
	216	+ triton_meta={"signature": {"xnumel": "i64"}},
	217	+ )
	218	+ self.assertEqual(len(emitted), 1)
	219	+ self.assertIn(
	220	+ "python_wrapper_fast_launch.bind",
	221	+ events[1][1],
	222	+ )
	223	+
	224	+ wrapper_module.config.triton.autotune_at_compile_time = True
	225	+ self.assertEqual(
	226	+ wrapper.generate_kernel_call(
	227	+ "triton_poi_fused_compile_time_autotune",
	228	+ ("buf0",),
	229	+ arg_types=(object(),),
	230	+ raw_args=None,
	231	+ ),
	232	+ "original",
	233	+ )
	234	+ self.assertEqual(len(emitted), 1)
	235	+ self.assertEqual(events[-1][0], "original")
	236	+ wrapper_module.config.triton.autotune_at_compile_time = False
	237	+
	238	+ wrapper_module.V.graph.cpp_wrapper = True
	239	+ self.assertEqual(
	240	+ wrapper.generate_kernel_call(
	241	+ "triton_poi_fused_1",
	242	+ ("xnumel",),
	243	+ ),
	244	+ "original",
	245	+ )
	246	+
	247	+ launcher = triton_cls().make_launcher()
	248	+ attach_metadata.assert_called_once()
	249	+ get_grid = attach_metadata.call_args.kwargs["get_grid"]
	250	+ self.assertEqual(get_grid(7), (7, 1, 1))
	251	+ self.assertEqual(launcher(3, "stream"), (3, "stream"))
	252	+
	253	+ grouped_result = triton_cls()
	254	+ grouped_result.inductor_meta["group_enabled"] = True
	255	+ grouped_launcher = grouped_result.make_launcher()
	256	+ attach_metadata.assert_called_once()
	257	+ self.assertEqual(grouped_launcher(5, "stream"), (5, "stream"))
	258	+
	259	+ def test_patch_operation_failure_rolls_back_previous_changes(self):
	260	+ class RejectingMeta(type):
	261	+ def __setattr__(cls, name, value):
	262	+ if name == "reject":
	263	+ raise RuntimeError("reject")
	264	+ return super().__setattr__(name, value)
	265	+
	266	+ class Target(metaclass=RejectingMeta):
	267	+ first = "original"
	268	+
	269	+ with isolated_patch_imports():
	270	+ fast_patch = importlib.import_module(f"{PACKAGE}.patch")
	271	+ with self.assertRaisesRegex(RuntimeError, "reject"):
	272	+ fast_patch._apply_patch_operations(
	273	+ [
	274	+ (Target, "first", "patched"),
	275	+ (Target, "reject", True),
	276	+ ]
	277	+ )
	278	+ self.assertEqual(Target.first, "original")
	279	+ self.assertFalse(hasattr(Target, "reject"))
	280	+
	281	+
	282	+if __name__ == "__main__":
	283	+ unittest.main()

Atest/_inductor/experimental/python_wrapper_fast_launch/test_npu_fast_launch_static.py+109-0

@@ -0,0 +1,109 @@
	1	+import unittest
	2	+from pathlib import Path
	3	+
	4	+
	5	+REPO_ROOT = Path(__file__).resolve().parents[4]
	6	+
	7	+
	8	+class TestNPUFastLaunchStatic(unittest.TestCase):
	9	+ def test_codegen_integration_is_installed_only_by_patch(self):
	10	+ wrapper = (REPO_ROOT / "torch_npu/_inductor/codegen/wrapper.py").read_text(
	11	+ encoding="utf-8"
	12	+ )
	13	+ triton_runtime = (
	14	+ REPO_ROOT / "torch_npu/_inductor/runtime/triton_heuristics.py"
	15	+ ).read_text(encoding="utf-8")
	16	+ config = (REPO_ROOT / "torch_npu/_inductor/config.py").read_text(
	17	+ encoding="utf-8"
	18	+ )
	19	+ inductor_init = (REPO_ROOT / "torch_npu/_inductor/__init__.py").read_text(
	20	+ encoding="utf-8"
	21	+ )
	22	+ patch = (
	23	+ REPO_ROOT
	24	+ / "torch_npu/_inductor/experimental/python_wrapper_fast_launch/patch.py"
	25	+ ).read_text(encoding="utf-8")
	26	+ bind = (
	27	+ REPO_ROOT
	28	+ / "torch_npu/_inductor/experimental/python_wrapper_fast_launch/bind.py"
	29	+ ).read_text(encoding="utf-8")
	30	+ emitter = (
	31	+ REPO_ROOT
	32	+ / "torch_npu/_inductor/experimental/python_wrapper_fast_launch/wrapper_codegen.py"
	33	+ ).read_text(encoding="utf-8")
	34	+
	35	+ self.assertNotIn("python_wrapper_fast_launch", wrapper)
	36	+ self.assertNotIn("_npu_fast_launch", triton_runtime)
	37	+ self.assertRegex(
	38	+ config,
	39	+ r'enable_fast_launch\s=\s_parse_bool_env\(\s*'
	40	+ r'"TORCHINDUCTOR_NPU_FAST_LAUNCH",\sFalse,\s\)',
	41	+ )
	42	+ self.assertIn("patch_fast_launch()", inductor_init)
	43	+ self.assertIn("npu_config.enable_fast_launch", patch)
	44	+ self.assertIn("npu_config.enable_fast_launch", bind)
	45	+ self.assertIn("NPUPythonWrapperCodeGen", patch)
	46	+ self.assertIn("TritonCompileResultNpu", patch)
	47	+ self.assertIn("bind_python_wrapper_kernel_fast", patch)
	48	+ self.assertIn("call_slot=", emitter)
	49	+ self.assertIn("[None]", emitter)
	50	+
	51	+ def test_package_import_does_not_eagerly_load_runtime_bindings(self):
	52	+ package_init = (
	53	+ REPO_ROOT
	54	+ / "torch_npu/_inductor/experimental/python_wrapper_fast_launch/__init__.py"
	55	+ ).read_text(encoding="utf-8")
	56	+ eager_imports = package_init.split("def __getattr__", 1)[0]
	57	+ self.assertNotIn("from .bind import", eager_imports)
	58	+ self.assertIn("def __getattr__(name)", package_init)
	59	+
	60	+ def test_cpp_backend_keeps_opcommand_and_runtime_guards(self):
	61	+ source = (
	62	+ REPO_ROOT
	63	+ / "torch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.cpp"
	64	+ ).read_text(
	65	+ encoding="utf-8"
	66	+ )
	67	+ self.assertIn("struct FastLaunchPlan", source)
	68	+ self.assertIn("rtKernelLaunch(", source)
	69	+ self.assertIn("rtKernelLaunchWithFlagV2", source)
	70	+ self.assertIn("SetCustomHandler", source)
	71	+ self.assertIn("grid product exceeds uint16 max", source)
	72	+ self.assertIn("args and arg_kinds size mismatch", source)
	73	+
	74	+ def test_no_unplanned_or_operator_fast_launch_is_added(self):
	75	+ source = (
	76	+ REPO_ROOT
	77	+ / "torch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.cpp"
	78	+ ).read_text(
	79	+ encoding="utf-8"
	80	+ )
	81	+ self.assertNotIn('"_npu_inductor_fast_launch"', source)
	82	+ self.assertNotIn("OperatorFastLaunch", source)
	83	+
	84	+ def test_fast_launch_plan_precomputes_packed_argument_layout(self):
	85	+ source = (
	86	+ REPO_ROOT
	87	+ / "torch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.cpp"
	88	+ ).read_text(
	89	+ encoding="utf-8"
	90	+ )
	91	+ self.assertIn("struct FastLaunchArgLayout", source)
	92	+ self.assertIn("BuildPackedLayout(*plan)", source)
	93	+ self.assertIn("plan.argLayouts[index]", source)
	94	+ self.assertIn("plan.gridOffsets[index]", source)
	95	+ self.assertIn("packed.args.resize(plan.packedArgsSize, 0)", source)
	96	+ self.assertIn("if (!plan.forceSimtOnly)", source)
	97	+ self.assertIn("if (plan.targetSupportFfts)", source)
	98	+ self.assertIn("rtGetC2cCtrlAddr(&fftsAddress, &fftsLength)", source)
	99	+ self.assertIn(
	100	+ "plan.packedArgsSize = AlignOffset(offset, packedAlignment)", source
	101	+ )
	102	+ self.assertIn(
	103	+ 'TORCH_CHECK(alignment != 0, "alignment must be non-zero")', source
	104	+ )
	105	+ self.assertNotIn("void AppendBytes(", source)
	106	+
	107	+
	108	+if __name__ == "__main__":
	109	+ unittest.main()

Mtorch_npu/_inductor/__init__.py+2-0

	@@ -128,6 +128,7 @@ def _load_triton_backend():
128	pre_grad_custom_pass_fuc,	128	pre_grad_custom_pass_fuc,
129	)	129	)
130	from .fx_passes.joint_graph import patch_constant_fold_uniform_value	130	from .fx_passes.joint_graph import patch_constant_fold_uniform_value
		131	+ from .experimental.python_wrapper_fast_launch.patch import patch_fast_launch
131	from .ir import patch_num_splits	132	from .ir import patch_num_splits
132	from .kernel import (	133	from .kernel import (
133	_register_npu_inductor_addmm,	134	_register_npu_inductor_addmm,
	@@ -210,6 +211,7 @@ def _load_triton_backend():
210	patch_create_device_properties()	211	patch_create_device_properties()
211	patch_load_cached_autotuning()	212	patch_load_cached_autotuning()
212	patch_triton_heuristics_cached_autotune()	213	patch_triton_heuristics_cached_autotune()
		214	+ patch_fast_launch()
213		215
214	pre_grad_custom_pass_fuc()	216	pre_grad_custom_pass_fuc()
215	post_grad_custom_pass_fuc()	217	post_grad_custom_pass_fuc()

Mtorch_npu/_inductor/config.py+5-0

	@@ -399,8 +399,13 @@ def _parse_float_env(name: str, default: float = 0.25, min_value: float = 0.0, m
399		399
400		400
401	# Frontend -> inductor controls (env-driven)	401	# Frontend -> inductor controls (env-driven)
		402	+# - TORCHINDUCTOR_NPU_FAST_LAUNCH: enable planned fast launch for NPU Python Wrapper
402	# - INDUCTOR_ASCEND_ENABLE_COSTMODEL: whether to forward costmodel backend signal to triton-ascend	403	# - INDUCTOR_ASCEND_ENABLE_COSTMODEL: whether to forward costmodel backend signal to triton-ascend
403	# - INDUCTOR_ASCEND_COSTMODEL_RATIO: select the shortest-latency top ratio of configs	404	# - INDUCTOR_ASCEND_COSTMODEL_RATIO: select the shortest-latency top ratio of configs
		405	+enable_fast_launch = _parse_bool_env(
		406	+ "TORCHINDUCTOR_NPU_FAST_LAUNCH",
		407	+ False,
		408	+)
404	enable_costmodel_backend = _parse_bool_env("INDUCTOR_ASCEND_ENABLE_COSTMODEL", False)	409	enable_costmodel_backend = _parse_bool_env("INDUCTOR_ASCEND_ENABLE_COSTMODEL", False)
405	costmodel_ratio = _parse_float_env("INDUCTOR_ASCEND_COSTMODEL_RATIO", 0.25, 0.0, 1.0)	410	costmodel_ratio = _parse_float_env("INDUCTOR_ASCEND_COSTMODEL_RATIO", 0.25, 0.0, 1.0)
406		411

Atorch_npu/_inductor/docs/npu_inductor_fast_launch_zh.md+164-0

@@ -0,0 +1,164 @@
	1	+# NPU Inductor Planned Fast Launch 使用指南
	2	+
	3	+## 1. 功能说明
	4	+
	5	+Planned Fast Launch 用于降低 `torch.compile` 生成的 Python Wrapper 在稳态运行时
	6	+下发 NPU Triton 融合 kernel 的 Host 侧固定开销。launcher 稳定后，Wrapper 会创建
	7	+C++ `FastLaunchPlan`；后续调用复用计划，只更新 stream、grid 和参数。
	8	+
	9	+该功能有以下边界：
	10	+
	11	+- 仅覆盖 NPU Inductor Python Wrapper 生成的 Triton kernel callsite。
	12	+- 不覆盖 eager/ACLNN 算子、C++ Wrapper 和图分区子图。
	13	+- compile-time autotune 保留原 codegen；grouped dynamic-shape autotuner 不附加
	14	+ Fast Launch launcher 元数据，运行时直接绑定原 `run` 入口。
	15	+- PyTorch profiler、FX graph/launch-params dump 和准确性检查等状态不进入
	16	+ planned path，由原路径保留对应语义。
	17	+- Triton launcher scope 中空的 launch `HookChain` 对象本身不阻止 planned path；
	18	+ 注册真实回调后，当前调用临时使用原 launcher，回调移除后自动恢复 planned path。
	19	+- 默认关闭，需要显式启用。
	20	+
	21	+## 2. 前置条件
	22	+
	23	+- 使用 `torch_npu` 2.10 系列中包含本功能的构建产物。
	24	+- PyTorch、torch_npu、CANN 和 NPU 驱动版本相互匹配。
	25	+- 模型通过 `torch.compile(..., backend="inductor")` 执行，并实际生成 NPU Triton
	26	+ kernel。
	27	+
	28	+环境变量在 Python 进程导入 `torch_npu` 以及执行 `torch.compile` 之前设置。
	29	+运行中修改环境变量不会改写已经生成的 Wrapper；需要重启进程并重新编译。
	30	+
	31	+## 3. 启用方式
	32	+
	33	+Linux Shell：
	34	+
	35	+```bash
	36	+export TORCHINDUCTOR_NPU_FAST_LAUNCH=1
	37	+python run_model.py
	38	+```
	39	+
	40	+Python 启动脚本也可以在导入 PyTorch/torch_npu 前设置：
	41	+
	42	+```python
	43	+import os
	44	+
	45	+os.environ["TORCHINDUCTOR_NPU_FAST_LAUNCH"] = "1"
	46	+
	47	+import torch
	48	+import torch_npu
	49	+```
	50	+
	51	+最小使用示例：
	52	+
	53	+```python
	54	+import torch
	55	+import torch_npu
	56	+
	57	+
	58	+def fn(x, y):
	59	+ return torch.sin(x + y) * 2
	60	+
	61	+
	62	+x = torch.randn(4096, device="npu")
	63	+y = torch.randn(4096, device="npu")
	64	+compiled_fn = torch.compile(fn, backend="inductor")
	65	+
	66	+result = compiled_fn(x, y)
	67	+```
	68	+
	69	+无需修改模型代码或显式调用 `_npu_inductor_fast_launch_with_plan`。计划创建和路由
	70	+由生成的 Wrapper 自动完成。初始调用可能包含编译、autotune 和计划创建开销；若启用
	71	+coordinate-descent tuning 或 kernel binary 保存，也会先完成相应生命周期工作。
	72	+
	73	+## 4. 生命周期、ABI 和回退
	74	+
	75	+生成的 Wrapper 会为每个 Triton callsite 惰性创建绑定对象。没有完整入口条件时，
	76	+如果选定 launcher 尚未稳定，本次调用使用原 `NPUCachingAutotuner.run`，调用结束后
	77	+再尝试创建 `FastLaunchPlan`；如果 launcher 已经稳定，则可以直接创建计划并进入
	78	+快路径。存在完整入口条件时只执行原入口，不尝试创建 plan。
	79	+
	80	+创建计划需要满足以下条件：
	81	+
	82	+- 已选定唯一且稳定的 launcher，且不是 fallback launcher。
	83	+- autotune 已完成。
	84	+- 启用 coordinate-descent tuning 时，调优已经完成。
	85	+- 启用 kernel binary 保存时，保存已经完成。
	86	+- 最终编译 ABI、C++ backend 和所需资源均受支持。
	87	+
	88	+codegen schema 不完整本身不会拒绝快路径；launcher 稳定后会使用最终编译 ABI 补全
	89	+参数类型。schema 明确冲突、最终 ABI 或资源不受支持，kernel、stub、grid resolver、
	90	+FFTS 状态或 backend 等必要信息缺失，以及 plan 创建失败等稳定问题，会为当前
	91	+launcher 安装负缓存并继续使用原 launcher。动态 grid 等单次调用错误不会污染已有
	92	+plan。
	93	+
	94	+当前计划 backend 支持以下 ABI 参数：
	95	+
	96	+- `torch.Tensor` 参数，提交时按其 `data_ptr` 打包。
	97	+- `i32`、`i64`、`u32`、`u64`。
	98	+- `f32`、`f64`。
	99	+- 按 `int32` 打包的 `bool`。
	100	+- 由最终 launcher ABI 提供、且类型属于上述范围的 runtime block 参数。
	101	+
	102	+Plan 保存 kernel name、kernel stub 及其 owner、参数布局、SIMT/shared memory 配置，
	103	+以及目标设备需要的 FFTS 地址。FFTS 地址在 plan 创建时查询一次。Plan 不记录或校验
	104	+device id；C++ backend 对传入 stream 只检查非空，不核对其所属设备，因此同一个 plan
	105	+不应跨设备复用。以下信息在每次调用时重新解析并打包：
	106	+
	107	+- 当前非空 NPU stream。
	108	+- 当前三维 grid。
	109	+- Tensor `data_ptr`。
	110	+- scalar 和 runtime block 参数值。
	111	+
	112	+Grid 的每个维度会先执行 `int(value)` 转换；转换结果必须恰好包含三个正整数，每维
	113	+不超过 `INT32_MAX`，三维乘积不超过 65535。plan 创建时会检查 SIMT 配置和 shared
	114	+memory 范围；每次调用会检查参数数量、stream 和 grid。提交继续经过
	115	+`OpCommand.Run`。
	116	+
	117	+以下状态会进入原 `NPUCachingAutotuner.run` 完整入口：
	118	+
	119	+- grouped dynamic-shape autotuner。
	120	+- launcher 生命周期尚未稳定，或当前为 fallback launcher。
	121	+- benchmark run、runtime kwargs 或 `TRITON_INTERPRET`。
	122	+- `INDUCTOR_ASCEND_DUMP_FX_GRAPH`、`INDUCTOR_ASCEND_CHECK_ACCURACY` 或
	123	+ launch-params dump。
	124	+- PyTorch profiler 正在运行。
	125	+
	126	+Triton launch `HookChain` 中存在已注册回调（或旧版 Triton 提供非空单 hook）时，
	127	+当前调用直接使用已选定的原 launcher，以保留 enter/exit hook 及 launch metadata
	128	+语义。该状态不安装负缓存，已有 plan 继续保留；hook 移除后的下一次调用会自动恢复
	129	+planned path。
	130	+
	131	+以下稳定问题会阻止 plan 创建并为当前 launcher 安装负缓存：
	132	+
	133	+- codegen schema 与最终 ABI 冲突，或最终 ABI 不受支持。
	134	+- launcher 元数据明确报告需要非零 workspace、sync block lock 或 device print
	135	+ 缓冲区。
	136	+- kernel、stub、grid resolver、FFTS 状态或 C++ backend 等必要信息缺失，或 plan
	137	+ 创建失败。
	138	+
	139	+负缓存首次安装前可能先进入完整入口；没有上述完整入口条件时，后续命中同一个
	140	+launcher 会直接调用已选定的原 launcher，不再重复进入 autotuner。launcher 变化后
	141	+负缓存失效并重新判断。
	142	+
	143	+Python 端发现 schema、backend 或动态 grid 等可恢复问题时可以使用原 launcher。一旦
	144	+调用 C++ fast-launch 入口，包括参数打包、Tensor/scalar 转换、stream 检查和实际提交
	145	+阶段发生的异常，均按“可能已经提交”处理，不再重放原 launcher，避免同一 kernel
	146	+重复执行。
	147	+
	148	+## 5. 关闭和故障处理
	149	+
	150	+关闭功能：
	151	+
	152	+```bash
	153	+export TORCHINDUCTOR_NPU_FAST_LAUNCH=0
	154	+```
	155	+
	156	+修改后重启 Python 进程并重新执行 `torch.compile`。
	157	+
	158	+常见问题：
	159	+
	160	+\| 现象 \| 检查项 \|
	161	+\| --- \| --- \|
	162	+\| 只有部分 callsite 使用 fast launch \| 其他 callsite 可能不是 NPU Triton Python Wrapper，或依赖 fallback launcher、不支持的 ABI 或额外资源 \|
	163	+\| PyTorch profiler、dump 或准确性检查下未使用 planned path \| 这是预期行为；这些状态保留原完整入口语义 \|
	164	+\| C++ fast-launch 调用失败后没有自动重试 \| 进入 C++ 入口后不会重放原 launcher，以避免 kernel 重复执行 \|

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/__init__.py+19-0

L linjiyuan7月28日

确认下是否放实验性目录？或者在资料中体现这个变量是个实验性特性

L

linjiyuan

7月29日评论：

@@ -0,0 +1,19 @@
	1	+from .types import FastLaunchError
	2	+
	3	+
	4	+def __getattr__(name):
	5	+ if name in ("BoundFastLaunch", "bind_python_wrapper_kernel_fast"):
	6	+ from .bind import BoundFastLaunch, bind_python_wrapper_kernel_fast
	7	+
	8	+ return {
	9	+ "BoundFastLaunch": BoundFastLaunch,
	10	+ "bind_python_wrapper_kernel_fast": bind_python_wrapper_kernel_fast,
	11	+ }[name]
	12	+ raise AttributeError(name)
	13	+
	14	+
	15	+__all__ = [
	16	+ "BoundFastLaunch",
	17	+ "FastLaunchError",
	18	+ "bind_python_wrapper_kernel_fast",
	19	+]

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/backend.py+342-0

@@ -0,0 +1,342 @@
	1	+from __future__ import annotations
	2	+
	3	+from importlib import import_module
	4	+from numbers import Real
	5	+from operator import index as operator_index
	6	+from typing import TYPE_CHECKING, Any
	7	+
	8	+if TYPE_CHECKING:
	9	+ from collections.abc import Callable
	10	+
	11	+from .types import FastLaunchError, FastLaunchPlanUnavailable
	12	+
	13	+
	14	+_SUPPORTED_ARG_KINDS = frozenset(
	15	+ (
	16	+ "tensor",
	17	+ "i32",
	18	+ "i64",
	19	+ "u32",
	20	+ "u64",
	21	+ "f32",
	22	+ "f64",
	23	+ "bool",
	24	+ )
	25	+)
	26	+_INT32_MAX = 2**31 - 1
	27	+_UINT16_MAX = 2**16 - 1
	28	+_INTEGER_ARG_KINDS = frozenset(("i32", "i64", "u32", "u64"))
	29	+_FLOAT_ARG_KINDS = frozenset(("f32", "f64"))
	30	+
	31	+
	32	+def arg_kind_from_abi_signature(signature: Any) -> str \| None:
	33	+ text = str(signature or "").strip().lower()
	34	+ if text.startswith(("*", "memref")) or "ptr" in text:
	35	+ return "tensor"
	36	+ if text in ("i1", "bool"):
	37	+ return "bool"
	38	+ if "u32" in text:
	39	+ return "u32"
	40	+ if "i32" in text or text == "int":
	41	+ return "i32"
	42	+ if "u64" in text:
	43	+ return "u64"
	44	+ if "i64" in text or text == "long":
	45	+ return "i64"
	46	+ if "fp32" in text or "f32" in text or text == "float":
	47	+ return "f32"
	48	+ if "fp64" in text or "f64" in text or text == "double":
	49	+ return "f64"
	50	+ return None
	51	+
	52	+
	53	+def _normalize_grid(grid: Any) -> tuple[int, int, int]:
	54	+ try:
	55	+ values = tuple(int(value) for value in grid)
	56	+ except (TypeError, ValueError) as exc:
	57	+ raise FastLaunchError(
	58	+ "invalid_grid",
	59	+ backend_submitted=False,
	60	+ ) from exc
	61	+ if len(values) != 3:
	62	+ raise FastLaunchError(
	63	+ f"grid_rank_mismatch:{len(values)}",
	64	+ backend_submitted=False,
	65	+ )
	66	+ product = 1
	67	+ for index, value in enumerate(values):
	68	+ if value <= 0:
	69	+ raise FastLaunchError(
	70	+ f"grid_dim_non_positive:{index}",
	71	+ backend_submitted=False,
	72	+ )
	73	+ if value > _INT32_MAX:
	74	+ raise FastLaunchError(
	75	+ f"grid_dim_exceeds_int32:{index}",
	76	+ backend_submitted=False,
	77	+ )
	78	+ product *= value
	79	+ if product > _UINT16_MAX:
	80	+ raise FastLaunchError(
	81	+ "grid_product_exceeds_uint16",
	82	+ backend_submitted=False,
	83	+ )
	84	+ return values
	85	+
	86	+
	87	+def _load_c_extension() -> Any:
	88	+ try:
	89	+ return import_module("torch_npu._C")
	90	+ except Exception as exc:
	91	+ raise FastLaunchPlanUnavailable("c_extension_unavailable") from exc
	92	+
	93	+
	94	+def _kernel_stub_supported(kernel_stub: Any) -> bool:
	95	+ if isinstance(kernel_stub, int):
	96	+ return kernel_stub != 0
	97	+ if type(kernel_stub).__name__ == "PyCapsule":
	98	+ return True
	99	+ if hasattr(kernel_stub, "value"):
	100	+ try:
	101	+ return int(kernel_stub.value) != 0
	102	+ except (TypeError, ValueError):
	103	+ return False
	104	+ try:
	105	+ return int(kernel_stub) != 0
	106	+ except (TypeError, ValueError):
	107	+ return False
	108	+
	109	+
	110	+def _callsite_schema_state(callsite_metadata: dict[str, Any]) -> str:
	111	+ state = callsite_metadata.get("schema_state")
	112	+ if state in ("complete", "incomplete", "conflict"):
	113	+ return str(state)
	114	+
	115	+ # Generated wrappers from the first fast-launch revision only carried an
	116	+ # eligible bit. Missing/unknown codegen signatures can be completed from
	117	+ # the selected launcher ABI; an explicit non-codegen rejection remains a
	118	+ # conflict for backward compatibility.
	119	+ if callsite_metadata.get("eligible", False):
	120	+ return "complete" if callsite_metadata.get("arg_kinds") else "incomplete"
	121	+ fallback_reason = str(callsite_metadata.get("fallback_reason") or "")
	122	+ if fallback_reason.startswith("codegen_"):
	123	+ return "incomplete"
	124	+ return "conflict"
	125	+
	126	+
	127	+def _validate_callsite_schema(
	128	+ callsite_metadata: dict[str, Any],
	129	+ launcher_arg_kinds: tuple[str, ...],
	130	+ runtime_arg_count: int,
	131	+) -> None:
	132	+ state = _callsite_schema_state(callsite_metadata)
	133	+ if state == "conflict":
	134	+ raise FastLaunchPlanUnavailable(
	135	+ str(callsite_metadata.get("fallback_reason") or "codegen_schema_conflict")
	136	+ )
	137	+ if runtime_arg_count > len(launcher_arg_kinds):
	138	+ raise FastLaunchPlanUnavailable("launcher_schema_too_short")
	139	+ if state != "complete":
	140	+ return
	141	+
	142	+ callsite_arg_kinds = tuple(callsite_metadata.get("arg_kinds", ()) or ())
	143	+ if len(callsite_arg_kinds) != runtime_arg_count:
	144	+ raise FastLaunchPlanUnavailable("codegen_schema_size_conflict")
	145	+ if callsite_arg_kinds != launcher_arg_kinds[:runtime_arg_count]:
	146	+ raise FastLaunchPlanUnavailable("codegen_launcher_schema_conflict")
	147	+
	148	+
	149	+def _runtime_arg_matches_kind(arg: Any, kind: str) -> bool:
	150	+ if kind == "tensor":
	151	+ return callable(getattr(arg, "data_ptr", None))
	152	+ if kind == "bool":
	153	+ return isinstance(arg, bool)
	154	+ if kind in _INTEGER_ARG_KINDS:
	155	+ if isinstance(arg, bool):
	156	+ return False
	157	+ try:
	158	+ operator_index(arg)
	159	+ return True
	160	+ except TypeError:
	161	+ return False
	162	+ if kind in _FLOAT_ARG_KINDS:
	163	+ return isinstance(arg, Real) and not isinstance(arg, bool)
	164	+ return False
	165	+
	166	+
	167	+def _validate_runtime_arg_categories(
	168	+ canonical_args: tuple[Any, ...],
	169	+ launcher_arg_kinds: tuple[str, ...],
	170	+) -> None:
	171	+ if len(canonical_args) != len(launcher_arg_kinds):
	172	+ raise FastLaunchPlanUnavailable(
	173	+ f"canonical_args_size_mismatch:{len(canonical_args)}:"
	174	+ f"{len(launcher_arg_kinds)}"
	175	+ )
	176	+ for arg_index, (arg, kind) in enumerate(zip(canonical_args, launcher_arg_kinds)):
	177	+ if not _runtime_arg_matches_kind(arg, kind):
	178	+ raise FastLaunchPlanUnavailable(
	179	+ f"runtime_arg_category_mismatch:{arg_index}:{kind}"
	180	+ )
	181	+
	182	+
	183	+class PlannedFastLaunch:
	184	+ __slots__ = (
	185	+ "arg_kinds",
	186	+ "get_grid",
	187	+ "launcher",
	188	+ "plan",
	189	+ "untimed_launch",
	190	+ )
	191	+
	192	+ def __init__(
	193	+ self,
	194	+ *,
	195	+ launcher: Any,
	196	+ plan: Any,
	197	+ arg_kinds: tuple[str, ...],
	198	+ get_grid: Callable[..., Any],
	199	+ untimed_launch: Callable[..., Any],
	200	+ ) -> None:
	201	+ self.launcher = launcher
	202	+ self.plan = plan
	203	+ self.arg_kinds = arg_kinds
	204	+ self.get_grid = get_grid
	205	+ self.untimed_launch = untimed_launch
	206	+
	207	+ def __call__(
	208	+ self,
	209	+ args: tuple[Any, ...],
	210	+ *,
	211	+ stream: Any,
	212	+ ) -> None:
	213	+ if len(args) != len(self.arg_kinds):
	214	+ raise FastLaunchError(
	215	+ f"args_size_mismatch:{len(args)}:{len(self.arg_kinds)}",
	216	+ backend_submitted=False,
	217	+ stable=True,
	218	+ )
	219	+ if stream is None:
	220	+ raise FastLaunchError(
	221	+ "stream_is_none",
	222	+ backend_submitted=False,
	223	+ )
	224	+ try:
	225	+ grid = _normalize_grid(self.get_grid(*args))
	226	+ except FastLaunchError:
	227	+ raise
	228	+ except Exception as exc:
	229	+ raise FastLaunchError(
	230	+ f"grid_resolve_error:{type(exc).__name__}",
	231	+ backend_submitted=False,
	232	+ ) from exc
	233	+
	234	+ try:
	235	+ self.untimed_launch(
	236	+ self.plan,
	237	+ stream,
	238	+ grid[0],
	239	+ grid[1],
	240	+ grid[2],
	241	+ args,
	242	+ )
	243	+ except Exception as exc:
	244	+ # All recoverable validation is completed before entering C++.
	245	+ # Treat errors after the boundary as submitted so fallback can never
	246	+ # replay an already queued kernel.
	247	+ raise FastLaunchError(
	248	+ f"backend_error:{type(exc).__name__}:{exc}",
	249	+ backend_submitted=True,
	250	+ ) from exc
	251	+ return None
	252	+
	253	+
	254	+def build_planned_fast_launch(
	255	+ launcher: Any,
	256	+ callsite_metadata: dict[str, Any],
	257	+ *,
	258	+ canonical_args: tuple[Any, ...] \| None = None,
	259	+ runtime_arg_count: int \| None = None,
	260	+) -> PlannedFastLaunch:
	261	+ kernel_name = str(getattr(launcher, "_npu_fast_launch_kernel_name", "") or "")
	262	+ kernel_stub = getattr(launcher, "_npu_fast_launch_kernel_stub", None)
	263	+ kernel_stub_owner = getattr(launcher, "_npu_fast_launch_kernel_stub_owner", None)
	264	+ get_grid = getattr(launcher, "_npu_fast_launch_get_grid", None)
	265	+ arg_kinds = tuple(getattr(launcher, "_npu_fast_launch_arg_kinds", ()) or ())
	266	+ if not kernel_name:
	267	+ raise FastLaunchPlanUnavailable("kernel_name_missing")
	268	+ if not _kernel_stub_supported(kernel_stub):
	269	+ raise FastLaunchPlanUnavailable("kernel_stub_unsupported")
	270	+ if kernel_stub_owner is None:
	271	+ raise FastLaunchPlanUnavailable("kernel_stub_owner_missing")
	272	+ if not callable(get_grid):
	273	+ raise FastLaunchPlanUnavailable("grid_resolver_missing")
	274	+ # Triton keeps launch hook objects in the generated launcher scope even
	275	+ # for launchers that the established planned fast path can invoke
	276	+ # directly. Do not make their mere presence a permanent negative cache.
	277	+ # BoundFastLaunch checks active hook callbacks on every direct call.
	278	+ if int(getattr(launcher, "_npu_fast_launch_workspace_size", 0) or 0) > 0:
	279	+ raise FastLaunchPlanUnavailable("launcher_workspace_required")
	280	+ if int(getattr(launcher, "_npu_fast_launch_lock_num", 0) or 0) > 0:
	281	+ raise FastLaunchPlanUnavailable("launcher_sync_block_lock_required")
	282	+ if getattr(launcher, "_npu_fast_launch_device_print_enabled", False):
	283	+ raise FastLaunchPlanUnavailable("launcher_device_print_required")
	284	+ target_support_ffts = getattr(
	285	+ launcher, "_npu_fast_launch_target_support_ffts", None
	286	+ )
	287	+ if target_support_ffts is None:
	288	+ raise FastLaunchPlanUnavailable("launcher_ffts_abi_unknown")
	289	+ if not arg_kinds or any(kind not in _SUPPORTED_ARG_KINDS for kind in arg_kinds):
	290	+ raise FastLaunchPlanUnavailable("arg_kinds_unsupported")
	291	+
	292	+ if runtime_arg_count is None:
	293	+ runtime_arg_count = int(
	294	+ callsite_metadata.get("runtime_arg_count", len(arg_kinds))
	295	+ )
	296	+ _validate_callsite_schema(callsite_metadata, arg_kinds, runtime_arg_count)
	297	+ if canonical_args is not None:
	298	+ _validate_runtime_arg_categories(canonical_args, arg_kinds)
	299	+
	300	+ extension = _load_c_extension()
	301	+ make_plan = getattr(extension, "_npu_inductor_make_fast_launch_plan", None)
	302	+ launch = getattr(extension, "_npu_inductor_fast_launch_with_plan", None)
	303	+ if not callable(make_plan) or not callable(launch):
	304	+ raise FastLaunchPlanUnavailable("planned_backend_unavailable")
	305	+
	306	+ enable_simt = bool(getattr(launcher, "_npu_fast_launch_enable_simt", False))
	307	+ shared_mem_dynamic_size = int(
	308	+ getattr(launcher, "_npu_fast_launch_shared_mem_dynamic_size", 0) or 0
	309	+ )
	310	+ force_simt_only = bool(getattr(launcher, "_npu_fast_launch_force_simt_only", False))
	311	+ try:
	312	+ plan = make_plan(
	313	+ kernel_name,
	314	+ kernel_stub,
	315	+ arg_kinds,
	316	+ enable_simt,
	317	+ shared_mem_dynamic_size,
	318	+ force_simt_only,
	319	+ bool(target_support_ffts),
	320	+ )
	321	+ # The C++ plan owns the stub object; this additional reference owns the
	322	+ # loaded binary that produced it.
	323	+ plan._owner = kernel_stub_owner
	324	+ except Exception as exc:
	325	+ raise FastLaunchPlanUnavailable(
	326	+ f"plan_creation_error:{type(exc).__name__}"
	327	+ ) from exc
	328	+
	329	+ return PlannedFastLaunch(
	330	+ launcher=launcher,
	331	+ plan=plan,
	332	+ arg_kinds=arg_kinds,
	333	+ get_grid=get_grid,
	334	+ untimed_launch=launch,
	335	+ )
	336	+
	337	+
	338	+__all__ = [
	339	+ "PlannedFastLaunch",
	340	+ "arg_kind_from_abi_signature",
	341	+ "build_planned_fast_launch",
	342	+]

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/bind.py+303-0

@@ -0,0 +1,303 @@
	1	+from __future__ import annotations
	2	+
	3	+from typing import Any
	4	+
	5	+import torch.autograd.profiler as autograd_profiler
	6	+
	7	+from torch_npu._inductor import config as npu_config
	8	+
	9	+from .backend import PlannedFastLaunch, build_planned_fast_launch
	10	+from .types import FastLaunchError, FastLaunchPlanUnavailable
	11	+
	12	+
	13	+_MISSING_HOOK_CALLBACKS = object()
	14	+
	15	+
	16	+def _static_full_entry_reason(autotuner: Any) -> str \| None:
	17	+ if not npu_config.enable_fast_launch:
	18	+ return "disabled"
	19	+ if getattr(npu_config, "dump_fx_graph", False):
	20	+ return "dump_fx_graph"
	21	+ if getattr(npu_config, "check_accuracy", False):
	22	+ return "check_accuracy"
	23	+ if getattr(autotuner, "triton_interpret", False):
	24	+ return "triton_interpret"
	25	+ if getattr(autotuner, "dump_launch_params", False):
	26	+ return "dump_launch_params"
	27	+ return None
	28	+
	29	+
	30	+def _is_grouped_autotuner(autotuner: Any) -> bool:
	31	+ inductor_meta = getattr(autotuner, "inductor_meta", {}) or {}
	32	+ return bool(
	33	+ inductor_meta.get("group_enabled", False)
	34	+ or hasattr(autotuner, "best_launcher_map")
	35	+ )
	36	+
	37	+
	38	+def _store_cubin_pending(autotuner: Any, launcher: Any) -> bool:
	39	+ return bool(
	40	+ getattr(launcher, "store_cubin", False)
	41	+ and not getattr(autotuner, "cuda_kernel_saved", False)
	42	+ )
	43	+
	44	+
	45	+def _launcher_has_active_launch_hooks(launcher: Any) -> bool:
	46	+ enter_callbacks = getattr(
	47	+ launcher,
	48	+ "_npu_fast_launch_enter_hook_callbacks",
	49	+ _MISSING_HOOK_CALLBACKS,
	50	+ )
	51	+ exit_callbacks = getattr(
	52	+ launcher,
	53	+ "_npu_fast_launch_exit_hook_callbacks",
	54	+ _MISSING_HOOK_CALLBACKS,
	55	+ )
	56	+ if (
	57	+ enter_callbacks is _MISSING_HOOK_CALLBACKS
	58	+ and exit_callbacks is _MISSING_HOOK_CALLBACKS
	59	+ ):
	60	+ # Conservatively preserve behavior for launchers carrying metadata
	61	+ # produced by the earlier boolean-only implementation.
	62	+ return bool(getattr(launcher, "_npu_fast_launch_has_launch_hooks", False))
	63	+ if (
	64	+ enter_callbacks is not _MISSING_HOOK_CALLBACKS
	65	+ and enter_callbacks
	66	+ ):
	67	+ return True
	68	+ return bool(
	69	+ exit_callbacks is not _MISSING_HOOK_CALLBACKS
	70	+ and exit_callbacks
	71	+ )
	72	+
	73	+
	74	+class BoundFastLaunch:
	75	+ __slots__ = (
	76	+ "autotuner",
	77	+ "metadata",
	78	+ "_direct",
	79	+ "_negative_callable",
	80	+ "_negative_launcher",
	81	+ "_static_full_entry_reason",
	82	+ )
	83	+
	84	+ def __init__(self, autotuner: Any, metadata: dict[str, Any]) -> None:
	85	+ self.autotuner = autotuner
	86	+ self.metadata = dict(metadata)
	87	+ self._direct: PlannedFastLaunch \| None = None
	88	+ self._static_full_entry_reason = _static_full_entry_reason(autotuner)
	89	+ self._negative_launcher: Any \| None = None
	90	+ self._negative_callable: Any \| None = None
	91	+
	92	+ def _stable_launcher(self) -> Any \| None:
	93	+ launchers = tuple(getattr(self.autotuner, "launchers", ()) or ())
	94	+ if len(launchers) != 1:
	95	+ return None
	96	+ launcher = launchers[0]
	97	+ if hasattr(launcher, "fallback"):
	98	+ return None
	99	+ best_launcher = getattr(self.autotuner, "best_launcher", None)
	100	+ if best_launcher is None or best_launcher is not launcher:
	101	+ return None
	102	+ if _store_cubin_pending(self.autotuner, launcher):
	103	+ return None
	104	+ return launcher
	105	+
	106	+ def _canonical_args(self, args: tuple[Any, ...]) -> tuple[Any, ...]:
	107	+ runtime_blocks = tuple(getattr(self.autotuner, "best_runtime_blocks", ()) or ())
	108	+ if not runtime_blocks:
	109	+ return args
	110	+ builder = getattr(self.autotuner, "_build_runtime_launch_args", None)
	111	+ if callable(builder):
	112	+ return tuple(builder(args, runtime_blocks))
	113	+ raise FastLaunchError(
	114	+ "runtime_block_builder_missing",
	115	+ backend_submitted=False,
	116	+ stable=True,
	117	+ )
	118	+
	119	+ def _clear_negative(self) -> None:
	120	+ self._negative_launcher = None
	121	+ self._negative_callable = None
	122	+
	123	+ def _install_negative(self, launcher: Any) -> None:
	124	+ self._direct = None
	125	+ self._negative_launcher = launcher
	126	+ self._negative_callable = launcher
	127	+
	128	+ def _try_promote(self, args: tuple[Any, ...]) -> bool:
	129	+ launcher = self._stable_launcher()
	130	+ if launcher is None:
	131	+ return False
	132	+ if self._negative_launcher is launcher:
	133	+ return False
	134	+ if self._direct is not None and self._direct.launcher is launcher:
	135	+ return True
	136	+ try:
	137	+ canonical_args = self._canonical_args(args)
	138	+ self._direct = build_planned_fast_launch(
	139	+ launcher,
	140	+ self.metadata,
	141	+ canonical_args=canonical_args,
	142	+ runtime_arg_count=len(args),
	143	+ )
	144	+ except FastLaunchPlanUnavailable as exc:
	145	+ if exc.stable:
	146	+ self._install_negative(launcher)
	147	+ return False
	148	+ except FastLaunchError as exc:
	149	+ if exc.stable:
	150	+ self._install_negative(launcher)
	151	+ return False
	152	+ self._clear_negative()
	153	+ return True
	154	+
	155	+ def _fallback(
	156	+ self,
	157	+ args: tuple[Any, ...],
	158	+ *,
	159	+ stream: Any,
	160	+ benchmark_run: bool,
	161	+ kwargs: dict[str, Any],
	162	+ ) -> Any:
	163	+ return self.autotuner.run(
	164	+ *args,
	165	+ stream=stream,
	166	+ benchmark_run=benchmark_run,
	167	+ **kwargs,
	168	+ )
	169	+
	170	+ def _call_launcher(
	171	+ self,
	172	+ launcher_call: Any,
	173	+ args: tuple[Any, ...],
	174	+ *,
	175	+ stream: Any,
	176	+ ) -> Any:
	177	+ canonical_args = self._canonical_args(args)
	178	+ return launcher_call(*canonical_args, stream=stream)
	179	+
	180	+ def _negative_fallback(
	181	+ self,
	182	+ launcher_call: Any,
	183	+ args: tuple[Any, ...],
	184	+ *,
	185	+ stream: Any,
	186	+ ) -> Any:
	187	+ return self._call_launcher(launcher_call, args, stream=stream)
	188	+
	189	+ def _hook_fallback(
	190	+ self,
	191	+ launcher_call: Any,
	192	+ args: tuple[Any, ...],
	193	+ *,
	194	+ stream: Any,
	195	+ ) -> Any:
	196	+ return self._call_launcher(launcher_call, args, stream=stream)
	197	+
	198	+ def _call_direct(
	199	+ self,
	200	+ direct: PlannedFastLaunch,
	201	+ args: tuple[Any, ...],
	202	+ *,
	203	+ stream: Any,
	204	+ ) -> Any:
	205	+ if _launcher_has_active_launch_hooks(direct.launcher):
	206	+ # Hook registration is dynamic. Use the selected original launcher
	207	+ # for this call only, keep the plan, and resume direct launch as
	208	+ # soon as the HookChain becomes empty again.
	209	+ return self._hook_fallback(
	210	+ direct.launcher,
	211	+ args,
	212	+ stream=stream,
	213	+ )
	214	+ canonical_args = self._canonical_args(args)
	215	+ try:
	216	+ direct(canonical_args, stream=stream)
	217	+ except FastLaunchError as exc:
	218	+ if exc.backend_submitted:
	219	+ raise
	220	+ if exc.stable:
	221	+ self._install_negative(direct.launcher)
	222	+ return self._fallback(
	223	+ args,
	224	+ stream=stream,
	225	+ benchmark_run=False,
	226	+ kwargs={},
	227	+ )
	228	+ return None
	229	+
	230	+ def __call__(
	231	+ self,
	232	+ *args: Any,
	233	+ stream: Any,
	234	+ benchmark_run: bool = False,
	235	+ **kwargs: Any,
	236	+ ) -> Any:
	237	+ reason = self._static_full_entry_reason
	238	+ if reason is None and benchmark_run:
	239	+ reason = "benchmark_run"
	240	+ if reason is None and kwargs:
	241	+ reason = "runtime_kwargs"
	242	+ if reason is None and autograd_profiler._is_profiler_enabled:
	243	+ reason = "profiler"
	244	+ if reason is not None:
	245	+ return self._fallback(
	246	+ args,
	247	+ stream=stream,
	248	+ benchmark_run=benchmark_run,
	249	+ kwargs=kwargs,
	250	+ )
	251	+
	252	+ direct = self._direct
	253	+ if direct is not None:
	254	+ # The plan is created only after all launcher lifecycle work is
	255	+ # complete. The selected launcher identity is the only mutable
	256	+ # state that must be rechecked on a steady-state hit.
	257	+ if getattr(self.autotuner, "best_launcher", None) is direct.launcher:
	258	+ return self._call_direct(direct, args, stream=stream)
	259	+ self._direct = None
	260	+
	261	+ negative_launcher = self._negative_launcher
	262	+ if negative_launcher is not None:
	263	+ if getattr(self.autotuner, "best_launcher", None) is negative_launcher:
	264	+ return self._negative_fallback(
	265	+ self._negative_callable,
	266	+ args,
	267	+ stream=stream,
	268	+ )
	269	+ self._clear_negative()
	270	+
	271	+ if self._try_promote(args) and self._direct is not None:
	272	+ return self._call_direct(self._direct, args, stream=stream)
	273	+
	274	+ result = self._fallback(
	275	+ args,
	276	+ stream=stream,
	277	+ benchmark_run=False,
	278	+ kwargs={},
	279	+ )
	280	+ # Promotion is deliberately after the original call: autotune,
	281	+ # coordinate descent, store-cubin, and debug semantics must complete
	282	+ # before the planned path may observe a launcher as stable.
	283	+ self._try_promote(args)
	284	+ return result
	285	+
	286	+
	287	+def bind_python_wrapper_kernel_fast(
	288	+ metadata: dict[str, Any],
	289	+ autotuner: Any,
	290	+ *,
	291	+ call_slot: list[Any] \| None = None,
	292	+) -> Any:
	293	+ bound = (
	294	+ autotuner.run
	295	+ if _is_grouped_autotuner(autotuner)
	296	+ else BoundFastLaunch(autotuner, metadata)
	297	+ )
	298	+ if call_slot is not None:
	299	+ call_slot[0] = bound
	300	+ return bound
	301	+
	302	+
	303	+__all__ = ["BoundFastLaunch", "bind_python_wrapper_kernel_fast"]

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/codegen.py+85-0

@@ -0,0 +1,85 @@
	1	+from __future__ import annotations
	2	+
	3	+import hashlib
	4	+import json
	5	+from collections.abc import Mapping
	6	+from typing import Any
	7	+
	8	+from .backend import arg_kind_from_abi_signature
	9	+
	10	+
	11	+def _runtime_signatures(triton_meta: Any, arg_count: int) -> tuple[str, ...]:
	12	+ if not isinstance(triton_meta, Mapping):
	13	+ return ()
	14	+ signature = triton_meta.get("signature")
	15	+ if not isinstance(signature, Mapping):
	16	+ return ()
	17	+ constants = triton_meta.get("constants", {}) or {}
	18	+ constant_names = {str(name) for name in constants}
	19	+ constant_indices = {
	20	+ int(index)
	21	+ for index in constants
	22	+ if isinstance(index, int) or str(index).isdigit()
	23	+ }
	24	+ values = tuple(
	25	+ str(value)
	26	+ for index, (name, value) in enumerate(signature.items())
	27	+ if str(name) not in constant_names
	28	+ and index not in constant_indices
	29	+ and str(value) != "constexpr"
	30	+ )
	31	+ if len(values) == arg_count:
	32	+ return values
	33	+ all_values = tuple(str(value) for value in signature.values())
	34	+ return all_values if len(all_values) == arg_count else ()
	35	+
	36	+
	37	+def build_callsite_metadata(
	38	+ *,
	39	+ kernel_name: str,
	40	+ call_args: Any,
	41	+ triton_meta: Any,
	42	+ graph_id: str,
	43	+ callsite_index: int,
	44	+) -> dict[str, Any]:
	45	+ arg_exprs = tuple(str(arg) for arg in call_args)
	46	+ signatures = _runtime_signatures(triton_meta, len(arg_exprs))
	47	+ arg_kinds = tuple(arg_kind_from_abi_signature(value) for value in signatures)
	48	+ schema_state = "complete"
	49	+ schema_reason = None
	50	+ if not signatures:
	51	+ schema_state = "incomplete"
	52	+ schema_reason = "codegen_signature_missing"
	53	+ elif any(kind is None for kind in arg_kinds):
	54	+ schema_state = "incomplete"
	55	+ schema_reason = "codegen_arg_kind_unsupported"
	56	+ schema_payload = {
	57	+ "kernel_name": kernel_name,
	58	+ "arg_exprs": arg_exprs,
	59	+ "arg_signatures": signatures,
	60	+ "arg_kinds": arg_kinds,
	61	+ "schema_state": schema_state,
	62	+ }
	63	+ schema_hash = hashlib.sha256(
	64	+ json.dumps(schema_payload, sort_keys=True).encode("utf-8")
	65	+ ).hexdigest()[:16]
	66	+ callsite_id = f"{graph_id}:{callsite_index}"
	67	+ return {
	68	+ "graph_id": graph_id,
	69	+ "callsite_id": callsite_id,
	70	+ "callsite_index": callsite_index,
	71	+ "kernel_name": kernel_name,
	72	+ "schema_hash": schema_hash,
	73	+ "runtime_arg_count": len(arg_exprs),
	74	+ "arg_kinds": tuple(kind for kind in arg_kinds if kind is not None),
	75	+ # An incomplete codegen schema is only a hint. The selected launcher
	76	+ # owns the final ABI, including runtime block arguments, so promotion
	77	+ # may safely complete the schema after autotuning has stabilized.
	78	+ "schema_state": schema_state,
	79	+ "schema_reason": schema_reason,
	80	+ "eligible": True,
	81	+ "fallback_reason": None,
	82	+ }
	83	+
	84	+
	85	+__all__ = ["build_callsite_metadata"]

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.cpp+498-0

@@ -0,0 +1,498 @@
	1	+#ifndef BUILD_LIBTORCH
	2	+
	3	+#include "torch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.h"
	4	+
	5	+#include <algorithm>
	6	+#include <cstddef>
	7	+#include <cstdint>
	8	+#include <cstring>
	9	+#include <limits>
	10	+#include <memory>
	11	+#include <string>
	12	+#include <utility>
	13	+#include <vector>
	14	+
	15	+#include <ATen/ATen.h>
	16	+#include <torch/csrc/Exceptions.h>
	17	+#include <torch/csrc/utils/pybind.h>
	18	+
	19	+#include <torch_npu/csrc/framework/OpCommand.h>
	20	+#include <torch_npu/csrc/inductor/mlir/hacl_rt.h>
	21	+
	22	+namespace py = pybind11;
	23	+
	24	+namespace {
	25	+
	26	+struct FastLaunchRtTaskCfgInfo {
	27	+ uint8_t qos = 0;
	28	+ uint8_t partId = 0;
	29	+ uint8_t schemMode = 0;
	30	+ bool d2dCrossFlag = false;
	31	+ uint32_t blockDimOffset = 0;
	32	+ uint8_t dumpflag = 0;
	33	+ uint8_t neverTimeout = 0;
	34	+ uint8_t rev[2] = {0, 0};
	35	+ uint32_t localMemorySize = 0;
	36	+};
	37	+
	38	+static_assert(offsetof(FastLaunchRtTaskCfgInfo, localMemorySize) == 12);
	39	+static_assert(sizeof(FastLaunchRtTaskCfgInfo) == 16);
	40	+
	41	+struct FastLaunchRtHostInputInfo {
	42	+ uint32_t addrOffset = 0;
	43	+ uint32_t dataOffset = 0;
	44	+};
	45	+
	46	+struct FastLaunchRtArgsExInfo {
	47	+ void* args = nullptr;
	48	+ FastLaunchRtHostInputInfo* hostInputInfoPtr = nullptr;
	49	+ uint32_t argsSize = 0;
	50	+ uint32_t tilingAddrOffset = 0;
	51	+ uint32_t tilingDataOffset = 0;
	52	+ uint16_t hostInputInfoNum = 0;
	53	+ uint8_t hasTiling = 0;
	54	+ uint8_t isNoNeedH2DCopy = 0;
	55	+ uint8_t reserved[4] = {0, 0, 0, 0};
	56	+};
	57	+
	58	+static_assert(offsetof(FastLaunchRtArgsExInfo, argsSize) == 16);
	59	+static_assert(offsetof(FastLaunchRtArgsExInfo, reserved) == 32);
	60	+static_assert(sizeof(FastLaunchRtArgsExInfo) == 40);
	61	+
	62	+extern "C" rtError_t rtKernelLaunchWithFlagV2(
	63	+ const void*,
	64	+ uint32_t,
	65	+ FastLaunchRtArgsExInfo*,
	66	+ rtSmDesc_t*,
	67	+ rtStream_t,
	68	+ uint32_t,
	69	+ const FastLaunchRtTaskCfgInfo*) __attribute__((weak));
	70	+
	71	+using FastLaunchRtKernelLaunchWithFlagV2 = decltype(&rtKernelLaunchWithFlagV2);
	72	+
	73	+enum class FastLaunchArgKind {
	74	+ Tensor,
	75	+ I32,
	76	+ I64,
	77	+ U32,
	78	+ U64,
	79	+ F32,
	80	+ F64,
	81	+ Bool,
	82	+};
	83	+
	84	+struct FastLaunchArgLayout {
	85	+ FastLaunchArgKind kind = FastLaunchArgKind::Tensor;
	86	+ size_t offset = 0;
	87	+};
	88	+
	89	+struct FastLaunchPlan {
	90	+ std::string kernelName;
	91	+ py::object kernelStubOwner;
	92	+ void* kernelStub = nullptr;
	93	+ std::vector<FastLaunchArgKind> argKinds;
	94	+ std::vector<FastLaunchArgLayout> argLayouts;
	95	+ size_t fftsOffset = 0;
	96	+ size_t gridOffsets[3] = {0, 0, 0};
	97	+ size_t packedArgsSize = 0;
	98	+ bool enableSimt = false;
	99	+ uint64_t sharedMemDynamicSize = 0;
	100	+ bool forceSimtOnly = false;
	101	+ bool targetSupportFfts = false;
	102	+ void* fftsAddress = nullptr;
	103	+};
	104	+
	105	+size_t AlignOffset(size_t offset, size_t alignment) {
	106	+ TORCH_CHECK(alignment != 0, "alignment must be non-zero");
	107	+ return (offset + alignment - 1) / alignment * alignment;
	108	+}
	109	+
	110	+void WriteBytesAt(
	111	+ std::vector<uint8_t>& buffer,
	112	+ size_t offset,
	113	+ const void* data,
	114	+ size_t size) {
	115	+ std::memcpy(buffer.data() + offset, data, size);
	116	+}
	117	+
	118	+void WritePointerAt(
	119	+ std::vector<uint8_t>& buffer,
	120	+ size_t offset,
	121	+ void* pointer) {
	122	+ WriteBytesAt(buffer, offset, &pointer, sizeof(void*));
	123	+}
	124	+
	125	+template <typename T>
	126	+void WriteScalarAt(
	127	+ std::vector<uint8_t>& buffer,
	128	+ size_t offset,
	129	+ py::handle arg) {
	130	+ T value = py::cast<T>(arg);
	131	+ WriteBytesAt(buffer, offset, &value, sizeof(T));
	132	+}
	133	+
	134	+FastLaunchArgKind ParseArgKind(const std::string& kind) {
	135	+ if (kind == "tensor") {
	136	+ return FastLaunchArgKind::Tensor;
	137	+ }
	138	+ if (kind == "i32") {
	139	+ return FastLaunchArgKind::I32;
	140	+ }
	141	+ if (kind == "i64") {
	142	+ return FastLaunchArgKind::I64;
	143	+ }
	144	+ if (kind == "u32") {
	145	+ return FastLaunchArgKind::U32;
	146	+ }
	147	+ if (kind == "u64") {
	148	+ return FastLaunchArgKind::U64;
	149	+ }
	150	+ if (kind == "f32") {
	151	+ return FastLaunchArgKind::F32;
	152	+ }
	153	+ if (kind == "f64") {
	154	+ return FastLaunchArgKind::F64;
	155	+ }
	156	+ if (kind == "bool") {
	157	+ return FastLaunchArgKind::Bool;
	158	+ }
	159	+ TORCH_CHECK(false, "unsupported fast launch arg kind: ", kind);
	160	+ return FastLaunchArgKind::Tensor;
	161	+}
	162	+
	163	+std::vector<FastLaunchArgKind> ParseArgKinds(
	164	+ const std::vector<std::string>& kinds) {
	165	+ std::vector<FastLaunchArgKind> parsed;
	166	+ parsed.reserve(kinds.size());
	167	+ for (const auto& kind : kinds) {
	168	+ parsed.emplace_back(ParseArgKind(kind));
	169	+ }
	170	+ return parsed;
	171	+}
	172	+
	173	+size_t ArgSize(FastLaunchArgKind kind) {
	174	+ switch (kind) {
	175	+ case FastLaunchArgKind::Tensor:
	176	+ return sizeof(void*);
	177	+ case FastLaunchArgKind::I32:
	178	+ return sizeof(int32_t);
	179	+ case FastLaunchArgKind::I64:
	180	+ return sizeof(int64_t);
	181	+ case FastLaunchArgKind::U32:
	182	+ return sizeof(uint32_t);
	183	+ case FastLaunchArgKind::U64:
	184	+ return sizeof(uint64_t);
	185	+ case FastLaunchArgKind::F32:
	186	+ return sizeof(float);
	187	+ case FastLaunchArgKind::F64:
	188	+ return sizeof(double);
	189	+ case FastLaunchArgKind::Bool:
	190	+ return sizeof(int32_t);
	191	+ }
	192	+ TORCH_INTERNAL_ASSERT(false, "unsupported fast launch arg kind");
	193	+ return 0;
	194	+}
	195	+
	196	+size_t ArgAlignment(FastLaunchArgKind kind) {
	197	+ switch (kind) {
	198	+ case FastLaunchArgKind::Tensor:
	199	+ return alignof(void*);
	200	+ case FastLaunchArgKind::I32:
	201	+ return alignof(int32_t);
	202	+ case FastLaunchArgKind::I64:
	203	+ return alignof(int64_t);
	204	+ case FastLaunchArgKind::U32:
	205	+ return alignof(uint32_t);
	206	+ case FastLaunchArgKind::U64:
	207	+ return alignof(uint64_t);
	208	+ case FastLaunchArgKind::F32:
	209	+ return alignof(float);
	210	+ case FastLaunchArgKind::F64:
	211	+ return alignof(double);
	212	+ case FastLaunchArgKind::Bool:
	213	+ return alignof(int32_t);
	214	+ }
	215	+ TORCH_INTERNAL_ASSERT(false, "unsupported fast launch arg kind");
	216	+ return 0;
	217	+}
	218	+
	219	+void BuildPackedLayout(FastLaunchPlan& plan) {
	220	+ size_t offset = 0;
	221	+ // The generated runner uses a packed struct whose individual fields carry
	222	+ // explicit alignment. Such a struct still has tail padding up to the
	223	+ // largest field alignment, so argsSize must be rounded up as well.
	224	+ size_t packedAlignment = alignof(int32_t);
	225	+ if (plan.targetSupportFfts) {
	226	+ packedAlignment = std::max(packedAlignment, alignof(void*));
	227	+ offset = AlignOffset(offset, alignof(void*));
	228	+ plan.fftsOffset = offset;
	229	+ offset += sizeof(void*);
	230	+ }
	231	+ // This is an ABI property, not a launch-API property. Ascend's generated
	232	+ // runner keeps the sync-lock and workspace slots for every kernel except a
	233	+ // force_simt_only binary, including SIMT-capable mixed-mode kernels.
	234	+ if (!plan.forceSimtOnly) {
	235	+ packedAlignment = std::max(packedAlignment, alignof(void*));
	236	+ for (int index = 0; index < 2; ++index) {
	237	+ offset = AlignOffset(offset, alignof(void*));
	238	+ offset += sizeof(void*);
	239	+ }
	240	+ }
	241	+
	242	+ plan.argLayouts.clear();
	243	+ plan.argLayouts.reserve(plan.argKinds.size());
	244	+ for (FastLaunchArgKind kind : plan.argKinds) {
	245	+ size_t alignment = ArgAlignment(kind);
	246	+ packedAlignment = std::max(packedAlignment, alignment);
	247	+ offset = AlignOffset(offset, alignment);
	248	+ plan.argLayouts.push_back({kind, offset});
	249	+ offset += ArgSize(kind);
	250	+ }
	251	+ for (size_t index = 0; index < 3; ++index) {
	252	+ offset = AlignOffset(offset, alignof(int32_t));
	253	+ plan.gridOffsets[index] = offset;
	254	+ offset += sizeof(int32_t);
	255	+ }
	256	+ plan.packedArgsSize = AlignOffset(offset, packedAlignment);
	257	+}
	258	+
	259	+void WriteArgAt(
	260	+ std::vector<uint8_t>& buffer,
	261	+ py::handle arg,
	262	+ const FastLaunchArgLayout& layout) {
	263	+ switch (layout.kind) {
	264	+ case FastLaunchArgKind::Tensor: {
	265	+ at::Tensor tensor = py::cast<at::Tensor>(arg);
	266	+ WritePointerAt(buffer, layout.offset, tensor.data_ptr());
	267	+ return;
	268	+ }
	269	+ case FastLaunchArgKind::I32:
	270	+ WriteScalarAt<int32_t>(buffer, layout.offset, arg);
	271	+ return;
	272	+ case FastLaunchArgKind::I64:
	273	+ WriteScalarAt<int64_t>(buffer, layout.offset, arg);
	274	+ return;
	275	+ case FastLaunchArgKind::U32:
	276	+ WriteScalarAt<uint32_t>(buffer, layout.offset, arg);
	277	+ return;
	278	+ case FastLaunchArgKind::U64:
	279	+ WriteScalarAt<uint64_t>(buffer, layout.offset, arg);
	280	+ return;
	281	+ case FastLaunchArgKind::F32:
	282	+ WriteScalarAt<float>(buffer, layout.offset, arg);
	283	+ return;
	284	+ case FastLaunchArgKind::F64:
	285	+ WriteScalarAt<double>(buffer, layout.offset, arg);
	286	+ return;
	287	+ case FastLaunchArgKind::Bool: {
	288	+ int32_t value = py::cast<bool>(arg) ? 1 : 0;
	289	+ WriteBytesAt(buffer, layout.offset, &value, sizeof(value));
	290	+ return;
	291	+ }
	292	+ }
	293	+ TORCH_CHECK(false, "unsupported fast launch arg kind");
	294	+}
	295	+
	296	+void* ExtractPointer(py::handle object, const char* name) {
	297	+ PyObject* raw = object.ptr();
	298	+ if (PyCapsule_CheckExact(raw)) {
	299	+ const char* capsuleName = PyCapsule_GetName(raw);
	300	+ if (PyErr_Occurred()) {
	301	+ PyErr_Clear();
	302	+ capsuleName = nullptr;
	303	+ }
	304	+ void* pointer = PyCapsule_GetPointer(raw, capsuleName);
	305	+ TORCH_CHECK(pointer != nullptr, name, " PyCapsule pointer is null");
	306	+ return pointer;
	307	+ }
	308	+ if (PyLong_Check(raw)) {
	309	+ void* pointer = PyLong_AsVoidPtr(raw);
	310	+ TORCH_CHECK(!PyErr_Occurred(), name, " cannot be converted to pointer");
	311	+ TORCH_CHECK(pointer != nullptr, name, " pointer is null");
	312	+ return pointer;
	313	+ }
	314	+ if (py::hasattr(object, "value")) {
	315	+ return ExtractPointer(object.attr("value"), name);
	316	+ }
	317	+ TORCH_CHECK(false, name, " must be an integer address or PyCapsule");
	318	+}
	319	+
	320	+struct PackedLaunch {
	321	+ std::vector<uint8_t> args;
	322	+ uint32_t blockNum = 0;
	323	+ rtStream_t stream = nullptr;
	324	+};
	325	+
	326	+PackedLaunch PackLaunch(
	327	+ const FastLaunchPlan& plan,
	328	+ uint64_t streamValue,
	329	+ uint32_t grid0,
	330	+ uint32_t grid1,
	331	+ uint32_t grid2,
	332	+ const py::sequence& args) {
	333	+ size_t argCount = static_cast<size_t>(py::len(args));
	334	+ TORCH_CHECK(
	335	+ argCount == plan.argKinds.size(),
	336	+ "fast launch args and arg_kinds size mismatch: ",
	337	+ argCount,
	338	+ " vs ",
	339	+ plan.argKinds.size());
	340	+ rtStream_t stream = reinterpret_cast<rtStream_t>(streamValue);
	341	+ TORCH_CHECK(stream != nullptr, "fast launch stream pointer is null");
	342	+
	343	+ const uint32_t grid[3] = {grid0, grid1, grid2};
	344	+ uint64_t blockNum = 1;
	345	+ for (size_t index = 0; index < 3; ++index) {
	346	+ TORCH_CHECK(grid[index] > 0, "fast launch grid dim must be positive");
	347	+ TORCH_CHECK(
	348	+ grid[index] <=
	349	+ static_cast<uint32_t>(std::numeric_limits<int32_t>::max()),
	350	+ "fast launch grid dim exceeds int32 max");
	351	+ blockNum *= grid[index];
	352	+ TORCH_CHECK(
	353	+ blockNum <= std::numeric_limits<uint16_t>::max(),
	354	+ "fast launch grid product exceeds uint16 max");
	355	+ }
	356	+
	357	+ PackedLaunch packed;
	358	+ packed.blockNum = static_cast<uint32_t>(blockNum);
	359	+ packed.stream = stream;
	360	+ TORCH_INTERNAL_ASSERT(plan.argLayouts.size() == argCount);
	361	+ packed.args.resize(plan.packedArgsSize, 0);
	362	+ if (plan.targetSupportFfts) {
	363	+ WritePointerAt(packed.args, plan.fftsOffset, plan.fftsAddress);
	364	+ }
	365	+ for (size_t index = 0; index < argCount; ++index) {
	366	+ WriteArgAt(packed.args, args[index], plan.argLayouts[index]);
	367	+ }
	368	+ int32_t signedGrid[3] = {
	369	+ static_cast<int32_t>(grid0),
	370	+ static_cast<int32_t>(grid1),
	371	+ static_cast<int32_t>(grid2),
	372	+ };
	373	+ for (size_t index = 0; index < 3; ++index) {
	374	+ WriteBytesAt(
	375	+ packed.args,
	376	+ plan.gridOffsets[index],
	377	+ &signedGrid[index],
	378	+ sizeof(signedGrid[index]));
	379	+ }
	380	+ return packed;
	381	+}
	382	+
	383	+void SubmitLaunch(const FastLaunchPlan& plan, PackedLaunch packed) {
	384	+ auto launchCall = [kernelStub = plan.kernelStub,
	385	+ enableSimt = plan.enableSimt,
	386	+ sharedMemDynamicSize = plan.sharedMemDynamicSize,
	387	+ packed = std::move(packed)]() mutable {
	388	+ void* args = packed.args.data();
	389	+ uint32_t argsSize = static_cast<uint32_t>(packed.args.size());
	390	+ rtError_t result;
	391	+ if (enableSimt) {
	392	+ FastLaunchRtKernelLaunchWithFlagV2 launchWithFlag =
	393	+ rtKernelLaunchWithFlagV2;
	394	+ TORCH_CHECK(
	395	+ launchWithFlag != nullptr,
	396	+ "rtKernelLaunchWithFlagV2 symbol not found");
	397	+ FastLaunchRtArgsExInfo argsInfo = {};
	398	+ argsInfo.args = args;
	399	+ argsInfo.argsSize = argsSize;
	400	+ FastLaunchRtTaskCfgInfo taskConfig = {};
	401	+ taskConfig.localMemorySize = static_cast<uint32_t>(sharedMemDynamicSize);
	402	+ result = launchWithFlag(
	403	+ kernelStub,
	404	+ packed.blockNum,
	405	+ &argsInfo,
	406	+ nullptr,
	407	+ packed.stream,
	408	+ 0,
	409	+ &taskConfig);
	410	+ } else {
	411	+ result = rtKernelLaunch(
	412	+ kernelStub, packed.blockNum, args, argsSize, nullptr, packed.stream);
	413	+ }
	414	+ return static_cast<int>(result);
	415	+ };
	416	+
	417	+ at_npu::native::OpCommand command;
	418	+ command.Name(plan.kernelName).SetCustomHandler(std::move(launchCall)).Run();
	419	+}
	420	+
	421	+std::shared_ptr<FastLaunchPlan> MakeFastLaunchPlan(
	422	+ const std::string& kernelName,
	423	+ py::object kernelStub,
	424	+ const std::vector<std::string>& argKinds,
	425	+ bool enableSimt,
	426	+ uint64_t sharedMemDynamicSize,
	427	+ bool forceSimtOnly,
	428	+ bool targetSupportFfts) {
	429	+ TORCH_CHECK(
	430	+ sharedMemDynamicSize <= std::numeric_limits<uint32_t>::max(),
	431	+ "shared_mem_dynamic_size exceeds uint32 max");
	432	+ TORCH_CHECK(
	433	+ !forceSimtOnly \|\| enableSimt, "force_simt_only requires enable_simt");
	434	+ auto plan = std::make_shared<FastLaunchPlan>();
	435	+ plan->kernelName = kernelName;
	436	+ plan->kernelStubOwner = kernelStub;
	437	+ plan->kernelStub = ExtractPointer(kernelStub, "kernel_stub");
	438	+ plan->argKinds = ParseArgKinds(argKinds);
	439	+ plan->enableSimt = enableSimt;
	440	+ plan->sharedMemDynamicSize = sharedMemDynamicSize;
	441	+ plan->forceSimtOnly = forceSimtOnly;
	442	+ plan->targetSupportFfts = targetSupportFfts;
	443	+ if (targetSupportFfts) {
	444	+ uint64_t fftsAddress = 0;
	445	+ uint32_t fftsLength = 0;
	446	+ rtError_t result = rtGetC2cCtrlAddr(&fftsAddress, &fftsLength);
	447	+ TORCH_CHECK(
	448	+ result == RT_ERROR_NONE,
	449	+ "rtGetC2cCtrlAddr failed while creating fast launch plan: ",
	450	+ static_cast<int>(result));
	451	+ TORCH_CHECK(
	452	+ fftsAddress != 0,
	453	+ "rtGetC2cCtrlAddr returned a null fast launch FFTS address");
	454	+ plan->fftsAddress = reinterpret_cast<void*>(fftsAddress);
	455	+ }
	456	+ BuildPackedLayout(*plan);
	457	+ return plan;
	458	+}
	459	+
	460	+void FastLaunchWithPlan(
	461	+ const std::shared_ptr<FastLaunchPlan>& plan,
	462	+ uint64_t stream,
	463	+ uint32_t grid0,
	464	+ uint32_t grid1,
	465	+ uint32_t grid2,
	466	+ const py::sequence& args) {
	467	+ TORCH_CHECK(plan != nullptr, "fast launch plan is null");
	468	+ SubmitLaunch(plan, PackLaunch(plan, stream, grid0, grid1, grid2, args));
	469	+}
	470	+
	471	+} // namespace
	472	+
	473	+void RegisterNPUFastLaunchBindings(PyObject* module) {
	474	+ auto m = py::handle(module).cast<py::module>();
	475	+ py::class_<FastLaunchPlan, std::shared_ptr<FastLaunchPlan>>(
	476	+ m, "_NPUInductorFastLaunchPlan", py::dynamic_attr());
	477	+ m.def(
	478	+ "_npu_inductor_make_fast_launch_plan",
	479	+ &MakeFastLaunchPlan,
	480	+ py::arg("kernel_name"),
	481	+ py::arg("kernel_stub"),
	482	+ py::arg("arg_kinds"),
	483	+ py::arg("enable_simt") = false,
	484	+ py::arg("shared_mem_dynamic_size") = 0,
	485	+ py::arg("force_simt_only") = false,
	486	+ py::arg("target_support_ffts") = false);
	487	+ m.def(
	488	+ "_npu_inductor_fast_launch_with_plan",
	489	+ &FastLaunchWithPlan,
	490	+ py::arg("plan"),
	491	+ py::arg("stream"),
	492	+ py::arg("grid_0"),
	493	+ py::arg("grid_1"),
	494	+ py::arg("grid_2"),
	495	+ py::arg("args"));
	496	+}
	497	+
	498	+#endif

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.h+11-0

@@ -0,0 +1,11 @@
	1	+#pragma once
	2	+
	3	+#ifndef BUILD_LIBTORCH
	4	+
	5	+#include <Python.h>
	6	+
	7	+#include "torch_npu/csrc/core/npu/NPUMacros.h"
	8	+
	9	+TORCH_NPU_API void RegisterNPUFastLaunchBindings(PyObject* module);
	10	+
	11	+#endif

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/launcher.py+109-0

@@ -0,0 +1,109 @@
	1	+from __future__ import annotations
	2	+
	3	+import os
	4	+from typing import Any
	5	+
	6	+from .backend import arg_kind_from_abi_signature
	7	+
	8	+
	9	+def _metadata_value(metadata: Any, name: str, default: Any) -> Any:
	10	+ if isinstance(metadata, dict):
	11	+ return metadata.get(name, default)
	12	+ return getattr(metadata, name, default)
	13	+
	14	+
	15	+def _launch_hook_callbacks(hook: Any) -> Any:
	16	+ if hook is None:
	17	+ return ()
	18	+
	19	+ calls = getattr(hook, "calls", None)
	20	+ if (
	21	+ type(hook).__name__ == "HookChain"
	22	+ and calls is not None
	23	+ and callable(getattr(hook, "add", None))
	24	+ and callable(getattr(hook, "remove", None))
	25	+ ):
	26	+ # Triton keeps one HookChain object in generated launcher scopes and
	27	+ # mutates its calls list when instrumentation starts or stops. Keep the
	28	+ # list itself so the steady-state check is both dynamic and cheap.
	29	+ return calls
	30	+
	31	+ # Older Triton versions expose a single hook callable instead of a
	32	+ # HookChain. Its presence means that the original launcher is required.
	33	+ return (hook,)
	34	+
	35	+
	36	+def attach_python_wrapper_launcher_metadata(
	37	+ launcher: Any,
	38	+ *,
	39	+ kernel_name: str,
	40	+ kernel_stub: Any,
	41	+ kernel_stub_owner: Any,
	42	+ get_grid: Any,
	43	+ grid: Any,
	44	+ def_args: Any,
	45	+ compile_meta: dict[str, Any],
	46	+ binary: Any,
	47	+ launcher_enter: Any,
	48	+ launcher_exit: Any,
	49	+) -> Any:
	50	+ signature = compile_meta.get("signature", {}) or {}
	51	+ arg_signatures = tuple(str(signature.get(name, "")) for name in def_args)
	52	+ launcher._npu_fast_launch_kernel_name = str(kernel_name)
	53	+ launcher._npu_fast_launch_kernel_stub = kernel_stub
	54	+ launcher._npu_fast_launch_kernel_stub_owner = kernel_stub_owner
	55	+ launcher._npu_fast_launch_get_grid = get_grid
	56	+ launcher._npu_fast_launch_grid_exprs = (
	57	+ str(grid.x_grid),
	58	+ str(grid.y_grid),
	59	+ str(grid.z_grid),
	60	+ )
	61	+ launcher._npu_fast_launch_def_args = tuple(str(name) for name in def_args)
	62	+ launcher._npu_fast_launch_arg_signatures = arg_signatures
	63	+ launcher._npu_fast_launch_arg_kinds = tuple(
	64	+ arg_kind_from_abi_signature(value) for value in arg_signatures
	65	+ )
	66	+
	67	+ metadata = getattr(binary, "metadata", None)
	68	+ parallel_mode = _metadata_value(metadata, "parallel_mode", "")
	69	+ force_simt_only = _metadata_value(metadata, "force_simt_only", False)
	70	+ shared_mem_dynamic_size = _metadata_value(
	71	+ metadata,
	72	+ "shared_mem_dynamic_size",
	73	+ 0,
	74	+ )
	75	+ workspace_size = _metadata_value(metadata, "workspace_size", -1)
	76	+ lock_num = _metadata_value(metadata, "lock_num", -1)
	77	+
	78	+ launcher._npu_fast_launch_enable_simt = "simt" in str(
	79	+ parallel_mode
	80	+ ).lower() or bool(force_simt_only)
	81	+ launcher._npu_fast_launch_force_simt_only = bool(force_simt_only)
	82	+ launcher._npu_fast_launch_shared_mem_dynamic_size = int(
	83	+ shared_mem_dynamic_size or 0
	84	+ )
	85	+
	86	+ # The Ascend runner prepends the FFTS synchronization address to the
	87	+ # packed ABI when the target supports FFTS. Preserve that compile-time
	88	+ # decision so the planned backend builds the same hidden-argument layout.
	89	+ from torch_npu._inductor.utils import triton_support_ffts
	90	+
	91	+ launcher._npu_fast_launch_target_support_ffts = bool(triton_support_ffts())
	92	+ launcher._npu_fast_launch_workspace_size = int(workspace_size or 0)
	93	+ launcher._npu_fast_launch_lock_num = int(lock_num or 0)
	94	+ launcher._npu_fast_launch_device_print_enabled = os.getenv(
	95	+ "TRITON_DEVICE_PRINT",
	96	+ "false",
	97	+ ).lower() in ("true", "1")
	98	+ launcher._npu_fast_launch_enter_hook = launcher_enter
	99	+ launcher._npu_fast_launch_exit_hook = launcher_exit
	100	+ launcher._npu_fast_launch_enter_hook_callbacks = _launch_hook_callbacks(
	101	+ launcher_enter
	102	+ )
	103	+ launcher._npu_fast_launch_exit_hook_callbacks = _launch_hook_callbacks(
	104	+ launcher_exit
	105	+ )
	106	+ return launcher
	107	+
	108	+
	109	+__all__ = ["attach_python_wrapper_launcher_metadata"]

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/patch.py+396-0

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/types.py+28-0

Atorch_npu/_inductor/experimental/python_wrapper_fast_launch/wrapper_codegen.py+55-0

Mtorch_npu/csrc/InitNpuBindings.cpp+4-0

Mtorch_npu/csrc/inductor/CMakeLists.txt+1-0