fix(inductor): keep multi_slice_concat out of epilogue fusion-pytorch-AtomGit

Mtest/_inductor/test_multi_slice_concat_pass.py+54-0

	@@ -3,6 +3,7 @@ import types
3	import unittest	3	import unittest
4		4
5	import torch	5	import torch
		6	+from torch._inductor.utils import run_and_get_code
6	from torch.export import Dim, export	7	from torch.export import Dim, export
7	from torch.fx.experimental.proxy_tensor import make_fx	8	from torch.fx.experimental.proxy_tensor import make_fx
8	from torch.testing._internal.common_utils import (	9	from torch.testing._internal.common_utils import (
	@@ -68,6 +69,22 @@ def _op_nodes(graph):
68	if n.op == "call_function" and n.target == MULTI_SLICE_CONCAT_TARGET]	69	if n.op == "call_function" and n.target == MULTI_SLICE_CONCAT_TARGET]
69		70
70		71
		72	+def _kernel_source(code, marker):
		73	+ """Source of the one compiled kernel whose body contains marker, or None.
		74	+
		75	+ Every kernel in the output code is its own ``async_compile.triton()`` string, so
		76	+ splitting on that keeps kernels apart; requiring a def line drops the call section.
		77	+ Kernels are matched on a rendered argument name, not on the kernel name, which
		78	+ define_kernel builds from node origins rather than from the template.
		79	+ """
		80	+ for chunk in code.split("async_compile.triton("):
		81	+ if marker in chunk and any(
		82	+ line.startswith("def ") for line in chunk.splitlines()
		83	+ ):
		84	+ return chunk
		85	+ return None
		86	+
		87	+
71	@contextlib.contextmanager	88	@contextlib.contextmanager
72	def _pass_spy():	89	def _pass_spy():
73	"""Count the op nodes the pass produces during compilation.	90	"""Count the op nodes the pass produces during compilation.
	@@ -984,6 +1001,43 @@ class TestMultiSliceConcatPass(TestUtils):
984	"aliased masks artifact")	1001	"aliased masks artifact")
985	self.assertGreater(stats["nodes"], 0, "aliased mask segments were not rewritten")	1002	self.assertGreater(stats["nodes"], 0, "aliased mask segments were not rewritten")
986		1003
		1004	+ @unittest.skipUnless(_ENABLED, _NEEDS_FLAG)
		1005	+ @parametrize('rows', [2, 200])
		1006	+ @parametrize('dtype', ['float16'])
		1007	+ def test_compile_epilogue_is_not_fused_away(self, rows, dtype):
		1008	+ """A pointwise consumer of the concat must not be fused into the template.
		1009	+
		1010	+ The template stores every segment with its own tl.store and renders no
		1011	+ store_output hook, so an epilogue fused into it has nowhere to be codegened and
		1012	+ is skipped: the kernel returns the bare concat and the consumer disappears. In
		1013	+ the model this was a clamp between the concat and an mm, which left the mm reading
		1014	+ unclamped values. The lowering marks the output no-fuse, so the clamp stays a
		1015	+ kernel of its own while the rewrite still happens.
		1016	+ """
		1017	+ widths = (16,) * len(_W16_OFFSETS)
		1018	+ wide = self._wide(rows, dtype)
		1019	+ # exactly representable in fp16, so it survives into the kernel source as written
		1020	+ bound = 0.375
		1021	+
		1022	+ def fn(x):
		1023	+ return torch.clamp(_col_concat_ref(x, _W16_OFFSETS, widths), -bound, bound)
		1024	+
		1025	+ with torch.no_grad(), _pass_spy() as stats:
		1026	+ compiled = torch.compile(fn, backend="inductor", dynamic=False)
		1027	+ actual, codes = run_and_get_code(compiled, wide)
		1028	+ self.assertGreater(stats["nodes"], 0, "pass did not fire, still aten.cat")
		1029	+ self._assert_bitwise_equal(fn(wide), actual, "clamp after the concat")
		1030	+
		1031	+ # arg_SRC0 is the first source pointer the template renders
		1032	+ source = _kernel_source(codes[0], "arg_SRC0")
		1033	+ self.assertIsNotNone(
		1034	+ source, "no template kernel in the output code, lowering fell back")
		1035	+ # whichever form the clamp takes, none of it belongs in the template
		1036	+ for trace in (str(bound), "maximum", "minimum", "clamp("):
		1037	+ self.assertNotIn(trace, source,
		1038	+ f"clamp was fused into the template ({trace})")
		1039	+
		1040	+
987	# ------------------------------------------------------------------	1041	# ------------------------------------------------------------------
988	# input dedup in lowering, pinned here since the case above needs the whole chain	1042	# input dedup in lowering, pinned here since the case above needs the whole chain
989	# ------------------------------------------------------------------	1043	# ------------------------------------------------------------------

Mtest/dynamo/test_compile_trigger.py+13-30

	@@ -473,11 +473,10 @@ class TorchCompileTriggerTests(unittest.TestCase):
473	"""	473	"""
474	)	474	)
475		475
476	- # Verify shape handling is installed only after the selected backend scope.	476	+ # Verify shape handling is installed after selecting the requested NPU backend.
477	def test_shape_handling_initializes_after_backend_selection(self):	477	def test_shape_handling_initializes_after_backend_selection(self):
478	self.run_in_subprocess(	478	self.run_in_subprocess(
479	"""	479	"""
480	- import os
481	import types	480	import types
482	from unittest import mock	481	from unittest import mock
483		482
	@@ -485,6 +484,11 @@ class TorchCompileTriggerTests(unittest.TestCase):
485	import torch_npu	484	import torch_npu
486	from torch_npu.utils import _dynamo	485	from torch_npu.utils import _dynamo
487		486
		487	+ def fake_setup(actual_options):
		488	+ actual_options = dict(actual_options)
		489	+ events.append(("setup", actual_options))
		490	+ return actual_options["npu_backend"]
		491	+
488	options = {	492	options = {
489	"npu_backend": "mlir",	493	"npu_backend": "mlir",
490	"enable_shape_handling": True,	494	"enable_shape_handling": True,
	@@ -492,9 +496,7 @@ class TorchCompileTriggerTests(unittest.TestCase):
492	events = []	496	events = []
493		497
494	def scope_register():	498	def scope_register():
495	- events.append(	499	+ events.append(("scope_register", None))
496	- ("scope_register", os.environ.get("TORCHINDUCTOR_NPU_BACKEND"))
497	- )
498		500
499	fake_inductor = types.SimpleNamespace(	501	fake_inductor = types.SimpleNamespace(
500	patch_shape_handling=lambda: events.append(	502	patch_shape_handling=lambda: events.append(
	@@ -502,9 +504,7 @@ class TorchCompileTriggerTests(unittest.TestCase):
502	)	504	)
503	)	505	)
504	with mock.patch.object(	506	with mock.patch.object(
505	- _dynamo, "_lazy_dynamo_setup", lambda: None	507	+ _dynamo, "_setup_inductor_for_compile", fake_setup
506	- ), mock.patch.object(
507	- _dynamo, "_lazy_inductor_setup", lambda: None
508	), mock.patch.object(	508	), mock.patch.object(
509	_dynamo, "register_inductor_npu", scope_register	509	_dynamo, "register_inductor_npu", scope_register
510	), mock.patch.object(	510	), mock.patch.object(
	@@ -514,7 +514,11 @@ class TorchCompileTriggerTests(unittest.TestCase):
514		514
515	assert wrapper.config["npu_backend"] == "mlir"	515	assert wrapper.config["npu_backend"] == "mlir"
516	assert wrapper.config["enable_shape_handling"] is True	516	assert wrapper.config["enable_shape_handling"] is True
517	- assert events == [("scope_register", "mlir")], events	517	+ assert events == [
		518	+ ("setup", options),
		519	+ ("shape_handling", None),
		520	+ ("scope_register", None),
		521	+ ], events
518	"""	522	"""
519	)	523	)
520		524
	@@ -664,27 +668,6 @@ class TorchCompileTriggerTests(unittest.TestCase):
664	"""	668	"""
665	)	669	)
666		670
667	- # Creating an Inductor wrapper must not load the Triton backend yet.
668	- def test_inductor_backend_load_is_deferred_until_first_call(self):
669	- self.run_in_subprocess(
670	- """
671	- import sys
672	- import torch
673	- import torch_npu
674	- from torch_npu.utils import _dynamo
675	-
676	- torch.compile(
677	- lambda x: x + 1,
678	- backend="inductor",
679	- options={"enable_shape_handling": True},
680	- )
681	-
682	- assert _dynamo._lazy_dynamo_setup.has_run
683	- assert not _dynamo._lazy_inductor_setup.has_run
684	- assert "torch_npu._inductor" not in sys.modules
685	- """
686	- )
687	-
688	# Verify lazy setup completes before compile backend lookup.	671	# Verify lazy setup completes before compile backend lookup.
689	def test_compile_triggers_setup_before_backend_lookup(self):	672	def test_compile_triggers_setup_before_backend_lookup(self):
690	self.run_in_subprocess(	673	self.run_in_subprocess(

Mtorch_npu/_inductor/kernel/multi_slice_concat.py+15-3

	@@ -10,8 +10,11 @@ front and reused across segments.
10		10
11	Output goes through ``manual_output_buffer``: the full output can be thousands of	11	Output goes through ``manual_output_buffer``: the full output can be thousands of
12	columns wide, too much to stage on chip for a single ``store_output``, so the template	12	columns wide, too much to stage on chip for a single ``store_output``, so the template
13	-stores per segment instead. That gives up epilogue fusion, but these concats feed	13	+stores per segment instead. That gives up epilogue fusion: with no ``store_output`` there
14	-extern kernels like ``mm`` / ``npu_fused_matmul``, which would not fuse anyway.	14	+is no hook to render an epilogue into, and one fused in anyway would be skipped by codegen
		15	+rather than emitted, silently dropping the consumer's computation. The lowering therefore
		16	+marks the output in ``V.graph.no_fuse_buffer_names``, which the scheduler honours before
		17	+it asks the backend, so the concat is fused into nothing at all.
15		18
16	The segment plan is baked into the template source and the instance cached by plan, so	19	The segment plan is baked into the template source and the instance cached by plan, so
17	autotune kwargs stay down to ``BLOCK_ROWS`` and the kernel name does not grow with the	20	autotune kwargs stay down to ``BLOCK_ROWS`` and the kernel name does not grow with the
	@@ -33,6 +36,7 @@ from torch._inductor.select_algorithm import (
33	autotune_select_algorithm,	36	autotune_select_algorithm,
34	SymbolicGridFn,	37	SymbolicGridFn,
35	)	38	)
		39	+from torch._inductor.virtualized import V
36		40
37	from ..select_algorithm import NPUTritonTemplate	41	from ..select_algorithm import NPUTritonTemplate
38		42
	@@ -262,6 +266,14 @@ def _register_npu_inductor_multi_slice_concat():
262	)	266	)
263	return per_slice_copies(srcs, masks, src_idx, mask_idx)	267	return per_slice_copies(srcs, masks, src_idx, mask_idx)
264		268
265	- return autotune_select_algorithm(	269	+ out = autotune_select_algorithm(
266	"multi_slice_concat", choices, input_nodes, layout	270	"multi_slice_concat", choices, input_nodes, layout
267	)	271	)
		272	+ # Nothing may be fused into this kernel: it stores its own output, so it renders
		273	+ # no store_output hook, and an epilogue the scheduler fused in would be skipped by
		274	+ # codegen rather than emitted -- the kernel would return the bare concat with the
		275	+ # consumer's computation gone. ``Scheduler.can_fuse`` tests this set before it
		276	+ # consults the backend, so one name here refuses epilogue and horizontal fusion
		277	+ # alike.
		278	+ V.graph.no_fuse_buffer_names.add(out.get_name())
		279	+ return out

Mtorch_npu/utils/_dynamo.py+28-6

	@@ -168,7 +168,6 @@ class _NpuBackendScope:
168	try:	168	try:
169	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend	169	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend
170	register_inductor_npu()	170	register_inductor_npu()
171	- _lazy_inductor_setup()
172	if self.backend == "ascendc":	171	if self.backend == "ascendc":
173	from torch_npu._inductor.deterministic_cache import (	172	from torch_npu._inductor.deterministic_cache import (
174	patch_npu_deterministic_level_cache_keys,	173	patch_npu_deterministic_level_cache_keys,
	@@ -210,8 +209,10 @@ def patch_inductor_wrapper():
210	if shape_handling_requested:	209	if shape_handling_requested:
211	if getattr(self, "_npu_defer_shape_handling", False):	210	if getattr(self, "_npu_defer_shape_handling", False):
212	self._npu_shape_handling_requested = True	211	self._npu_shape_handling_requested = True
213	- # Shape handling is installed in new_call, after the selected	212	+ return
214	- # backend scope has loaded the matching NPU Inductor backend.	213	+ if not is_inductor_npu_initialized():
		214	+ register_inductor_npu()
		215	+ torch_npu._inductor.patch_shape_handling()
215		216
216	def new_get_config_copy(self) -> dict[str, Any]:	217	def new_get_config_copy(self) -> dict[str, Any]:
217	ori_dict = src_get_config_copy(self)	218	ori_dict = src_get_config_copy(self)
	@@ -249,10 +250,13 @@ def patch_inductor_wrapper():
249	self._npu_shape_handling_requested = False	250	self._npu_shape_handling_requested = False
250	try:	251	try:
251	src_init(self, mode, options, dynamic)	252	src_init(self, mode, options, dynamic)
		253	+ shape_handling_requested = self._npu_shape_handling_requested
252	finally:	254	finally:
253	del self._npu_defer_shape_handling	255	del self._npu_defer_shape_handling
254	del self._npu_shape_handling_requested	256	del self._npu_shape_handling_requested
255	- _lazy_dynamo_setup()	257	+ _setup_inductor_for_compile(self.config)
		258	+ if shape_handling_requested:
		259	+ torch_npu._inductor.patch_shape_handling()
256	backend = _resolve_npu_backend_from_wrapper(self)	260	backend = _resolve_npu_backend_from_wrapper(self)
257	if backend=="mlir":	261	if backend=="mlir":
258	with _NpuBackendScope(backend):	262	with _NpuBackendScope(backend):
	@@ -266,8 +270,6 @@ def patch_inductor_wrapper():
266	def new_call(self, model_, inputs_):	270	def new_call(self, model_, inputs_):
267	backend = _resolve_npu_backend_from_wrapper(self)	271	backend = _resolve_npu_backend_from_wrapper(self)
268	with _NpuBackendScope(backend):	272	with _NpuBackendScope(backend):
269	- if self.config.get("enable_shape_handling", False):
270	- torch_npu._inductor.patch_shape_handling()
271	if backend == "ascendc":	273	if backend == "ascendc":
272	from torch_npu.dynamo._deterministic_guard import (	274	from torch_npu.dynamo._deterministic_guard import (
273	install_npu_deterministic_level_guard,	275	install_npu_deterministic_level_guard,
	@@ -708,6 +710,26 @@ def _lazy_inductor_setup():
708	_inject_inductor_npu_backend_config()	710	_inject_inductor_npu_backend_config()
709		711
710		712
		713	+def _setup_inductor_for_compile(options=None):
		714	+ """Initialize the NPU Inductor backend selected for this compile call."""
		715	+ _lazy_dynamo_setup()
		716	+
		717	+ option_backend = options.get("npu_backend") if isinstance(options, dict) else None
		718	+ selected_backend = _resolve_npu_backend(option_backend)
		719	+
		720	+ old_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
		721	+ if selected_backend not in (None, "", "default"):
		722	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = selected_backend
		723	+ try:
		724	+ _lazy_inductor_setup()
		725	+ finally:
		726	+ if old_backend is None:
		727	+ os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
		728	+ else:
		729	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = old_backend
		730	+ return selected_backend
		731	+
		732	+
711	@run_once	733	@run_once
712	def install_npugraph_mark_step_trigger():	734	def install_npugraph_mark_step_trigger():
713	"""Expose the public NPUGraph step API without importing compiler internals."""	735	"""Expose the public NPUGraph step API without importing compiler internals."""