inductor/fx_passes: dynamic graph optimization e2e-pytorch-AtomGit

已合并

inductor/fx_passes: dynamic graph optimization e2e #44486

dezheng889创建于 8月13日从crazyDannyBoy/pytorch_dezheng： v2.10.0_graph_opt_clean 合入到 Ascend/pytorch： v2.10.0

inductor/fx_passes: dynamic graph optimization e2e #44486

已合并

dezheng889创建于 8月13日

从crazyDannyBoy/pytorch_dezheng v2.10.0_graph_opt_clean合入到Ascend/pytorch v2.10.0

共 15 个文件变更+3176-221

Atest/_inductor/test_grouped_matmul_pass.py+337-0

@@ -0,0 +1,337 @@
	1	+import operator
	2	+import unittest
	3	+
	4	+import torch
	5	+from torch.export import Dim, export
	6	+from torch.fx.experimental.proxy_tensor import make_fx
	7	+from torch.testing._internal.common_utils import (
	8	+ instantiate_parametrized_tests,
	9	+ parametrize,
	10	+ run_tests,
	11	+)
	12	+from testutils import TestUtils
	13	+import torch_npu
	14	+import torch_npu._inductor
	15	+from torch_npu._inductor import config as npu_config
	16	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	17	+ _gmm_plan_batch_size,
	18	+ grouped_matmul_fusion_pass,
	19	+)
	20	+
	21	+
	22	+_CAT = torch.ops.aten.cat.default
	23	+_ADDMM = torch.ops.aten.addmm.default
	24	+_MM = torch.ops.aten.mm.default
	25	+_RELU = torch.ops.aten.relu.default
	26	+_GMM = torch.ops.npu.npu_grouped_matmul.default
	27	+
	28	+_ROWS = 200
	29	+_WIDTH = 64
	30	+# Tower widths from a real model: three orders of magnitude, most far below a cube's capacity.
	31	+_TOWER_K = (16, 32, 48, 64, 96, 128, 160, 192, 256, 320, 384, 512, 640, 768, 1024)
	32	+
	33	+_ENABLED = npu_config.enable_grouped_matmul_fusion
	34	+_NEEDS_FLAG = "requires TORCHINDUCTOR_ENABLE_GROUPED_MATMUL_FUSION=1 (the pass registers on the switch)"
	35	+
	36	+
	37	+def _count(graph, target):
	38	+ return len([n for n in graph.nodes if n.op == "call_function" and n.target == target])
	39	+
	40	+
	41	+def _grouped_nodes(graph):
	42	+ return [n for n in graph.nodes if n.op == "call_function" and n.target == _GMM]
	43	+
	44	+
	45	+def _cat_node(graph):
	46	+ return [n for n in graph.nodes if n.op == "call_function" and n.target == _CAT][-1]
	47	+
	48	+
	49	+def _batch_sizes(graph):
	50	+ return [len(n.args[0]) for n in _grouped_nodes(graph)]
	51	+
	52	+
	53	+def _member_order(graph):
	54	+ """Map each cat input back to (batch index, group index) to confirm concat order is intact.
	55	+
	56	+ After the rewrite every cat input is a getitem, so node counts cannot catch a permutation,
	57	+ which would silently interleave the features of different towers.
	58	+ """
	59	+ grouped = _grouped_nodes(graph)
	60	+ order = []
	61	+ for inp in _cat_node(graph).args[0]:
	62	+ if not (isinstance(inp, torch.fx.Node) and inp.target is operator.getitem):
	63	+ order.append(None)
	64	+ continue
	65	+ order.append((grouped.index(inp.args[0]), inp.args[1]))
	66	+ return order
	67	+
	68	+
	69	+class _Captured:
	70	+ """A captured graph plus a callable that runs it."""
	71	+
	72	+ def __init__(self, gm, call):
	73	+ self.gm = gm
	74	+ self.graph = gm.graph
	75	+ self._call = call
	76	+
	77	+ def __call__(self, *args):
	78	+ return self._call(*args)
	79	+
	80	+
	81	+class _FnModule(torch.nn.Module):
	82	+ """torch.export only accepts nn.Module, and the cases under test are functions.
	83	+
	84	+ Inputs arrive as one tuple so dynamic_shapes stays a single positional spec no
	85	+ matter how many tensors a case passes; a tower of n projections passes 2n or 3n
	86	+ of them. Export flattens it back to one placeholder per tensor.
	87	+ """
	88	+
	89	+ def __init__(self, fn):
	90	+ super().__init__()
	91	+ self._fn = fn
	92	+
	93	+ def forward(self, tensors):
	94	+ return self._fn(*tensors)
	95	+
	96	+
	97	+def _capture(fn, *tensors, dynamic=False):
	98	+ """Capture the graph the pass actually operates on.
	99	+
	100	+ Static mode uses ``tracing_mode="fake"`` so all placeholders share one FakeTensorMode, as in
	101	+ production. Dynamic mode uses ``torch.export`` with only dim 0 dynamic, because batch is the
	102	+ only symbolic shape in production and each branch's K is a compile-time constant.
	103	+ """
	104	+ if not dynamic:
	105	+ gm = make_fx(fn, tracing_mode="fake")(*tensors)
	106	+ return _Captured(gm, gm)
	107	+
	108	+ batch = Dim("batch", min=2)
	109	+ dynamic_shapes = tuple(
	110	+ {0: batch} if t.dim() == 2 and t.shape[0] == _ROWS else None for t in tensors
	111	+ )
	112	+ program = export(_FnModule(fn), (tuple(tensors),), dynamic_shapes=(dynamic_shapes,))
	113	+ return _Captured(program.graph_module, lambda *a: program.module()(tuple(a)))
	114	+
	115	+
	116	+class TestGroupedMatmulPass(TestUtils):
	117	+ def _tower_inputs(self, n, rows=_ROWS, width=_WIDTH, dtype="float16", bias=True):
	118	+ """Build n projections with distinct K but equal output width, plus the final cat inputs."""
	119	+ torch_dtype = eval("torch." + dtype)
	120	+ device = torch.device("npu")
	121	+ ks = [_TOWER_K[i % len(_TOWER_K)] for i in range(n)]
	122	+ xs = [torch.randn((rows, k), dtype=torch_dtype, device=device) for k in ks]
	123	+ ws = [torch.randn((k, width), dtype=torch_dtype, device=device) for k in ks]
	124	+ if not bias:
	125	+ return xs + ws
	126	+ bs = [torch.randn((width,), dtype=torch_dtype, device=device) for _ in ks]
	127	+ return xs + ws + bs
	128	+
	129	+ def _tower_fn(self, n, bias=True):
	130	+ def fn(*tensors):
	131	+ xs, ws = tensors[:n], tensors[n:2 * n]
	132	+ if bias:
	133	+ bs = tensors[2 * n:3 * n]
	134	+ parts = [_ADDMM(b, x, w) for b, x, w in zip(bs, xs, ws)]
	135	+ else:
	136	+ parts = [_MM(x, w) for x, w in zip(xs, ws)]
	137	+ return _CAT(parts, 1)
	138	+
	139	+ return fn
	140	+
	141	+ def _run_pass(self, fn, tensors, dynamic=False, *limits):
	142	+ captured = _capture(fn, *tensors, dynamic=dynamic)
	143	+ grouped_matmul_fusion_pass(captured.graph, **limits)
	144	+ captured.gm.recompile()
	145	+ return captured
	146	+
	147	+ # ------------------------------------------------------------------
	148	+ # Batching rule. The cliff sits between 32 and 40 groups: minimize batches, then split evenly.
	149	+ # ------------------------------------------------------------------
	150	+ @parametrize(
	151	+ "total,expected",
	152	+ [
	153	+ (8, 8), # fits in one batch
	154	+ (32, 32), # exactly at the cap
	155	+ (33, 17), # just over the cap: two even batches, not 32 + 1
	156	+ (64, 32),
	157	+ (80, 27), # measured optimum: 3 batches of 27/27/26, 14% faster than 32/32/16
	158	+ (100, 25),
	159	+ ],
	160	+ )
	161	+ def test_batch_size_is_balanced(self, total, expected):
	162	+ self.assertEqual(_gmm_plan_batch_size(total, 32), expected)
	163	+
	164	+ def test_batch_size_never_exceeds_cap(self):
	165	+ for cap in (8, 16, 32):
	166	+ for total in range(2, 200):
	167	+ self.assertLessEqual(_gmm_plan_batch_size(total, cap), cap,
	168	+ msg=f"cap={cap} total={total}")
	169	+
	170	+ # ------------------------------------------------------------------
	171	+ # Graph rewriting
	172	+ # ------------------------------------------------------------------
	173	+ def test_folds_tower_into_balanced_batches(self):
	174	+ n = 80
	175	+ captured = self._run_pass(self._tower_fn(n), *self._tower_inputs(n),
	176	+ max_groups_per_call=32)
	177	+ self.assertEqual(_batch_sizes(captured.graph), [27, 27, 26])
	178	+ self.assertEqual(_count(captured.graph, _ADDMM), 0)
	179	+ # Dispatch count is the whole point of this rewrite: 81 (80 matmul + 1 cat) down to 4.
	180	+ self.assertEqual(_count(captured.graph, _GMM) + _count(captured.graph, _CAT), 4)
	181	+
	182	+ def test_preserves_concat_order(self):
	183	+ n = 80
	184	+ captured = self._run_pass(self._tower_fn(n), *self._tower_inputs(n),
	185	+ max_groups_per_call=32)
	186	+ expected = ([(0, i) for i in range(27)]
	187	+ + [(1, i) for i in range(27)]
	188	+ + [(2, i) for i in range(26)])
	189	+ self.assertEqual(_member_order(captured.graph), expected)
	190	+
	191	+ def test_passes_group_semantics(self):
	192	+ """split_item=0 emits one output per group; group_type=-1 gives each group its own K."""
	193	+ n = 16
	194	+ captured = self._run_pass(self._tower_fn(n), *self._tower_inputs(n))
	195	+ for node in _grouped_nodes(captured.graph):
	196	+ self.assertEqual(node.kwargs["split_item"], 0)
	197	+ self.assertEqual(node.kwargs["group_type"], -1)
	198	+ self.assertEqual(len(node.kwargs["bias"]), len(node.args[0]))
	199	+ self.assertEqual(len(node.args[0]), len(node.args[1]))
	200	+
	201	+ def test_folds_without_bias(self):
	202	+ n = 16
	203	+ captured = self._run_pass(self._tower_fn(n, bias=False),
	204	+ *self._tower_inputs(n, bias=False))
	205	+ self.assertEqual(_count(captured.graph, _MM), 0)
	206	+ for node in _grouped_nodes(captured.graph):
	207	+ self.assertNotIn("bias", node.kwargs)
	208	+
	209	+ def test_folds_under_dynamic_batch(self):
	210	+ """Row count is symbolic in production graphs; a symbolic dim must not block matching."""
	211	+ n = 16
	212	+ captured = self._run_pass(self._tower_fn(n), *self._tower_inputs(n),
	213	+ dynamic=True)
	214	+ self.assertEqual(_count(captured.graph, _GMM), 1)
	215	+ self.assertEqual(_count(captured.graph, _ADDMM), 0)
	216	+
	217	+ def test_splits_biased_and_plain_into_separate_calls(self):
	218	+ """Biased and plain matmuls cannot share a batch: bias is a whole-call argument."""
	219	+ n = 20
	220	+ tensors = self._tower_inputs(n)
	221	+
	222	+ def fn(*t):
	223	+ xs, ws, bs = t[:n], t[n:2 * n], t[2 * n:]
	224	+ parts = [_ADDMM(b, x, w) for b, x, w in zip(bs[:10], xs[:10], ws[:10])]
	225	+ parts += [_MM(x, w) for x, w in zip(xs[10:], ws[10:])]
	226	+ return _CAT(parts, 1)
	227	+
	228	+ captured = self._run_pass(fn, *tensors, min_groups=8)
	229	+ self.assertEqual(_count(captured.graph, _GMM), 2)
	230	+ with_bias = [n_ for n_ in _grouped_nodes(captured.graph) if "bias" in n_.kwargs]
	231	+ self.assertEqual(len(with_bias), 1)
	232	+
	233	+ # ------------------------------------------------------------------
	234	+ # Rejection conditions
	235	+ # ------------------------------------------------------------------
	236	+ def test_skips_few_branches(self):
	237	+ """Few wide GEMMs like QKV already fill the cube; grouping them costs 12% device time."""
	238	+ n = 3
	239	+ captured = self._run_pass(self._tower_fn(n), *self._tower_inputs(n),
	240	+ min_groups=8)
	241	+ self.assertEqual(_count(captured.graph, _GMM), 0)
	242	+ self.assertEqual(_count(captured.graph, _ADDMM), n)
	243	+
	244	+ def test_skips_wide_rows_when_static(self):
	245	+ n = 16
	246	+ captured = self._run_pass(self._tower_fn(n),
	247	+ *self._tower_inputs(n, rows=_ROWS),
	248	+ max_rows=_ROWS - 1)
	249	+ self.assertEqual(_count(captured.graph, _GMM), 0)
	250	+
	251	+ def test_skips_row_concat(self):
	252	+ n = 16
	253	+ tensors = self._tower_inputs(n, width=_WIDTH)
	254	+
	255	+ def fn(*t):
	256	+ xs, ws, bs = t[:n], t[n:2 * n], t[2 * n:]
	257	+ return _CAT([_ADDMM(b, x, w) for b, x, w in zip(bs, xs, ws)], 0)
	258	+
	259	+ captured = self._run_pass(fn, *tensors)
	260	+ self.assertEqual(_count(captured.graph, _GMM), 0)
	261	+
	262	+ def test_skips_mismatched_width(self):
	263	+ """Differing output widths form separate buckets, neither meeting the minimum group."""
	264	+ n = 16
	265	+ device = torch.device("npu")
	266	+ xs = [torch.randn((_ROWS, 64), dtype=torch.float16, device=device)
	267	+ for _ in range(n)]
	268	+ ws = [torch.randn((64, _WIDTH if i < n // 2 else _WIDTH * 2),
	269	+ dtype=torch.float16, device=device) for i in range(n)]
	270	+ bs = [torch.randn((_WIDTH if i < n // 2 else _WIDTH * 2,),
	271	+ dtype=torch.float16, device=device) for i in range(n)]
	272	+
	273	+ def fn(*t):
	274	+ a, b, c = t[:n], t[n:2 * n], t[2 * n:]
	275	+ return _CAT([_ADDMM(z, x, w) for z, x, w in zip(c, a, b)], 1)
	276	+
	277	+ captured = self._run_pass(fn, xs, ws, *bs, min_groups=12)
	278	+ self.assertEqual(_count(captured.graph, _GMM), 0)
	279	+
	280	+ def test_skips_scaled_addmm(self):
	281	+ """addmm with beta/alpha is not plain x @ w + b; the grouped operator cannot express it."""
	282	+ n = 16
	283	+ tensors = self._tower_inputs(n)
	284	+
	285	+ def fn(*t):
	286	+ xs, ws, bs = t[:n], t[n:2 * n], t[2 * n:]
	287	+ parts = [_ADDMM(b, x, w, beta=2.0) for b, x, w in zip(bs, xs, ws)]
	288	+ return _CAT(parts, 1)
	289	+
	290	+ captured = self._run_pass(fn, *tensors)
	291	+ self.assertEqual(_count(captured.graph, _GMM), 0)
	292	+ self.assertEqual(_count(captured.graph, _ADDMM), n)
	293	+
	294	+ def test_skips_matmul_with_other_users(self):
	295	+ """Extra users force the grouped call before the earliest user, so the gain is uncertain."""
	296	+ n = 16
	297	+ tensors = self._tower_inputs(n)
	298	+
	299	+ def fn(*t):
	300	+ xs, ws, bs = t[:n], t[n:2 * n], t[2 * n:]
	301	+ parts = [_ADDMM(b, x, w) for b, x, w in zip(bs, xs, ws)]
	302	+ return _CAT(parts, 1), _RELU(parts[0])
	303	+
	304	+ captured = self._run_pass(fn, *tensors, min_groups=8)
	305	+ # Only the first branch has an extra user; the remaining 15 still meet the minimum.
	306	+ self.assertEqual(_batch_sizes(captured.graph), [15])
	307	+ self.assertEqual(_count(captured.graph, _ADDMM), 1)
	308	+
	309	+ # ------------------------------------------------------------------
	310	+ # Numerics. Not bitwise-identical: the operator accumulates in a different order than separate
	311	+ # matmuls; measured relative RMS error in fp16 is 6.6e-06, matmul's own noise level.
	312	+ # ------------------------------------------------------------------
	313	+ @unittest.skipUnless(_ENABLED, _NEEDS_FLAG)
	314	+ @parametrize("n", [16, 80])
	315	+ def test_matches_separate_matmuls(self, n):
	316	+ tensors = self._tower_inputs(n)
	317	+ fn = self._tower_fn(n)
	318	+ expected = fn(*tensors)
	319	+ captured = self._run_pass(fn, *tensors,
	320	+ max_groups_per_call=32)
	321	+ self.assertGreater(_count(captured.graph, _GMM), 0, "pass did not fire, so this case proves nothing")
	322	+ actual = captured(*tensors)
	323	+ self.assertEqual(expected.shape, actual.shape)
	324	+
	325	+ # Element-wise comparison is meaningless here: accumulating 1024 fp16 products puts one ULP
	326	+ # at 3e-2, so any element-wise threshold is either too loose to catch bugs or bound to
	327	+ # flake. Relative RMS error is stable: 6.6e-06 measured at harder tower widths (K to 5036).
	328	+ diff = (actual.float() - expected.float())
	329	+ rms = diff.pow(2).mean().sqrt() / expected.float().pow(2).mean().sqrt()
	330	+ self.assertLess(rms.item(), 1e-4, f"relative RMS error {rms.item():.3e} is too large")
	331	+
	332	+
	333	+instantiate_parametrized_tests(TestGroupedMatmulPass)
	334	+
	335	+
	336	+if __name__ == "__main__":
	337	+ run_tests()

Atest/_inductor/test_multi_slice_concat_pass.py+1014-0

@@ -0,0 +1,1014 @@
	1	+import contextlib
	2	+import types
	3	+import unittest
	4	+
	5	+import torch
	6	+from torch.export import Dim, export
	7	+from torch.fx.experimental.proxy_tensor import make_fx
	8	+from torch.testing._internal.common_utils import (
	9	+ instantiate_parametrized_tests,
	10	+ parametrize,
	11	+ run_tests,
	12	+)
	13	+from testutils import TestUtils
	14	+import torch_npu
	15	+import torch_npu._inductor
	16	+from torch_npu._inductor import config as npu_config
	17	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	18	+ MULTI_SLICE_CONCAT_TARGET,
	19	+ multi_slice_concat_pass,
	20	+)
	21	+from torch_npu._inductor.kernel.multi_slice_concat import _dedup_inputs
	22	+
	23	+
	24	+_CAT = torch.ops.aten.cat.default
	25	+_SLICE = torch.ops.aten.slice.Tensor
	26	+_WHERE = torch.ops.aten.where.self
	27	+_FULL = torch.ops.aten.full.default
	28	+_MSC = torch.ops.npu_ext.multi_slice_concat
	29	+
	30	+# operator arg positions, so node.args is not indexed by bare numbers
	31	+_A_SRCS, _A_MASKS, _A_SRC_IDX, _A_OFFSETS, _A_WIDTHS, _A_MASK_IDX = range(6)
	32	+
	33	+# wide feature table and offsets from real model output code, all compile-time constants.
	34	+_WIDE_COLS = 69876
	35	+# L24837: 15 segments, width=128.
	36	+_W128_OFFSETS = (55834, 56452, 45421, 44203, 44744, 45962, 60741, 61328,
	37	+ 62267, 65602, 66495, 66956, 67445, 68021, 69060)
	38	+# L24895: 12 segments, width=16.
	39	+_W16_OFFSETS = (55818, 56436, 45405, 44187, 44728, 45946, 61312, 62251,
	40	+ 69044, 68199, 68221, 69822)
	41	+# L25021: 45 mixed-width segments, the longest column concat in the model.
	42	+_MIXED_OFFSETS = (55818, 56436, 45405, 44187, 44728, 45946, 61312, 62251, 69044,
	43	+ 45421, 44203, 44744, 45962, 60741, 61328, 62267, 65602, 66495,
	44	+ 66956, 67445, 69060, 45549, 44331, 44872, 46090, 61456, 62395,
	45	+ 69188, 68199, 68221, 69822, 55580, 59302, 1693, 1749, 1807,
	46	+ 2033, 3276, 3350, 2463, 58205, 58980, 62964, 2159, 64802)
	47	+_MIXED_WIDTHS = (16, 16, 16, 16, 16, 16, 16, 16, 16, 128, 128, 128, 128, 128, 128,
	48	+ 128, 128, 128, 128, 128, 128, 32, 32, 32, 32, 32, 32, 32, 16, 16,
	49	+ 16, 128, 128, 32, 32, 32, 32, 32, 32, 32, 16, 16, 16, 16, 16)
	50	+# L27947: 28 segments, width=8, the narrowest ones in the model.
	51	+_W8_OFFSETS = (48025, 47595, 46868, 47890, 48779, 49346, 47769, 49082, 49180,
	52	+ 48337, 48467, 48597, 48402, 48532, 48714, 49240, 48070, 48204,
	53	+ 48119, 11076, 14257, 18643, 11922, 15103, 19489, 12752, 15933, 20319)
	54	+# cat_124: masked and bare slices strictly alternating; nothing merged before masks.
	55	+_ALT_OFFSETS = (35098, 39296, 35470, 39913, 34648, 38647, 35842, 40530)
	56	+_ALT_WIDTHS = (80, 80, 80, 80, 112, 112, 128, 128)
	57	+
	58	+_ENABLED = npu_config.enable_multi_slice_concat
	59	+_NEEDS_FLAG = "needs TORCHINDUCTOR_ENABLE_MULTI_SLICE_CONCAT=1 (pass is registered by the flag)"
	60	+
	61	+
	62	+def _count(graph, target):
	63	+ return len([n for n in graph.nodes if n.op == "call_function" and n.target == target])
	64	+
	65	+
	66	+def _op_nodes(graph):
	67	+ return [n for n in graph.nodes
	68	+ if n.op == "call_function" and n.target == MULTI_SLICE_CONCAT_TARGET]
	69	+
	70	+
	71	+@contextlib.contextmanager
	72	+def _pass_spy():
	73	+ """Count the op nodes the pass produces during compilation.
	74	+
	75	+ When the pass does not fire the result falls back to aten.cat and is still
	76	+ bitwise correct, so a numeric-only end-to-end check passes vacuously. The
	77	+ pass is registered as a function object, so swap in a counting wrapper.
	78	+ """
	79	+ from torch_npu._inductor.fx_passes.ascend_custom_passes.register_custom_pass import (
	80	+ ASCEND_CUSTOME_PASS_REGISTER,
	81	+ )
	82	+
	83	+ stats = {"nodes": 0}
	84	+ original = []
	85	+
	86	+ def spy(graph):
	87	+ multi_slice_concat_pass(graph)
	88	+ stats["nodes"] += len(_op_nodes(graph))
	89	+
	90	+ spy.__name__ = multi_slice_concat_pass.__name__
	91	+ for pass_type, levels in ASCEND_CUSTOME_PASS_REGISTER.items():
	92	+ for level, fns in levels.items():
	93	+ if any(getattr(f, "__name__", "") == spy.__name__ for f in fns):
	94	+ original.append((pass_type, level, list(fns)))
	95	+ levels[level] = [
	96	+ spy if getattr(f, "__name__", "") == spy.__name__ else f
	97	+ for f in fns
	98	+ ]
	99	+ try:
	100	+ yield stats
	101	+ finally:
	102	+ for pass_type, level, fns in original:
	103	+ ASCEND_CUSTOME_PASS_REGISTER[pass_type][level] = fns
	104	+
	105	+
	106	+class _Captured:
	107	+ def __init__(self, gm, call):
	108	+ self.gm = gm
	109	+ self.graph = gm.graph
	110	+ self._call = call
	111	+
	112	+ def __call__(self, *args):
	113	+ return self._call(*args)
	114	+
	115	+
	116	+class _FnModule(torch.nn.Module):
	117	+ """torch.export only accepts nn.Module, and the cases under test are functions.
	118	+
	119	+ Inputs arrive as one tuple so dynamic_shapes stays a single positional spec no
	120	+ matter how many tensors a case passes; export flattens it back to one
	121	+ placeholder per tensor, which is what the pass sees.
	122	+ """
	123	+
	124	+ def __init__(self, fn):
	125	+ super().__init__()
	126	+ self._fn = fn
	127	+
	128	+ def forward(self, tensors):
	129	+ return self._fn(*tensors)
	130	+
	131	+
	132	+def _capture(fn, *tensors, dynamic=False):
	133	+ """Capture the graph the pass actually sees.
	134	+
	135	+ Static uses ``tracing_mode="fake"``, not ``"real"``: real builds one
	136	+ FakeTensorMode per input, so inferring output meta for multi-source segments
	137	+ hits "Mixing fake modes" and silently skips the rewrite.
	138	+
	139	+ Dynamic uses ``torch.export`` with only dim 0 dynamic, like production:
	140	+ symbolic batch, static last dim. ``tracing_mode="symbolic"`` would symbolize
	141	+ the last dim too, leaving the pass without constant offsets so it silently
	142	+ stops matching.
	143	+ """
	144	+ if not dynamic:
	145	+ gm = make_fx(fn, tracing_mode="fake")(*tensors)
	146	+ return _Captured(gm, gm)
	147	+
	148	+ batch = Dim("batch", min=2)
	149	+ program = export(_FnModule(fn), (tuple(tensors),),
	150	+ dynamic_shapes=(tuple({0: batch} for _ in tensors),))
	151	+ return _Captured(program.graph_module, lambda *a: program.module()(tuple(a)))
	152	+
	153	+
	154	+def _col_concat(wide, offsets, widths):
	155	+ """Reproduce the [aten.slice] -> cat(dim=-1) shape from the output code."""
	156	+ parts = [_SLICE(wide, -1, off, off + w) for off, w in zip(offsets, widths)]
	157	+ return _CAT(parts, -1)
	158	+
	159	+
	160	+def _col_concat_ref(wide, offsets, widths):
	161	+ return torch.cat([wide[..., off:off + w] for off, w in zip(offsets, widths)], dim=-1)
	162	+
	163	+
	164	+def _masked_slice(wide, mask, off, width):
	165	+ """Reproduce the where(row mask, zeros, column slice) shape from output code."""
	166	+ part = _SLICE(wide, -1, off, off + width)
	167	+ zero = _FULL([wide.shape[0], width], 0, dtype=wide.dtype, device=wide.device)
	168	+ return _WHERE(mask, zero, part)
	169	+
	170	+
	171	+def _masked_slice_ref(wide, mask, off, width):
	172	+ return torch.where(mask, torch.zeros_like(wide[..., off:off + width]),
	173	+ wide[..., off:off + width])
	174	+
	175	+
	176	+class _StubInput:
	177	+ """_dedup_inputs only reads buffer name and layout, so no real IR nodes needed."""
	178	+
	179	+ def __init__(self, name, size=(2, 1), stride=(1, 1), offset=0):
	180	+ self._name = name
	181	+ self._layout = types.SimpleNamespace(
	182	+ size=list(size), stride=list(stride), offset=offset
	183	+ )
	184	+
	185	+ def get_name(self):
	186	+ return self._name
	187	+
	188	+ def get_layout(self):
	189	+ return self._layout
	190	+
	191	+
	192	+class TestMultiSliceConcatPass(TestUtils):
	193	+ # pure data movement plus mask select, so output must be bitwise-identical to cat.
	194	+ def _assert_bitwise_equal(self, expected, actual, note=""):
	195	+ self.assertEqual(expected, actual, atol=0, rtol=0,
	196	+ msg=f"{note}: a pure-copy rewrite must be bitwise identical")
	197	+
	198	+ def _wide(self, rows, dtype='float16', cols=_WIDE_COLS):
	199	+ return torch.randn((rows, cols), dtype=eval('torch.' + dtype),
	200	+ device=torch.device("npu"))
	201	+
	202	+ def _mask(self, rows, cols=1):
	203	+ """Row-wise bool mask [rows, 1], matching logical_not in the model."""
	204	+ return torch.randint(0, 2, (rows, cols), device=torch.device("npu"),
	205	+ dtype=torch.bool)
	206	+
	207	+ def _run_pass(self, fn, tensors, dynamic=False, *limits):
	208	+ captured = _capture(fn, *tensors, dynamic=dynamic)
	209	+ multi_slice_concat_pass(captured.graph, **limits)
	210	+ captured.gm.recompile()
	211	+ return captured
	212	+
	213	+ # ------------------------------------------------------------------
	214	+ # operator contract: eager runs the CompositeExplicitAutograd reference impl
	215	+ # (per-segment copy + cat), not the fused kernel, which only exists after
	216	+ # lowering. These pin semantics and arg validation; perf in bench_multi_slice_concat.py.
	217	+ # ------------------------------------------------------------------
	218	+ @parametrize('rows', [1, 32, 200])
	219	+ @parametrize('dtype', ['float16'])
	220	+ def test_op_matches_cat_of_slices(self, rows, dtype):
	221	+ wide = self._wide(rows, dtype)
	222	+ widths = (128,) * len(_W128_OFFSETS)
	223	+ n = len(widths)
	224	+ out = _MSC([wide], [], [0] * n, list(_W128_OFFSETS), list(widths), [-1] * n)
	225	+ self._assert_bitwise_equal(_col_concat_ref(wide, _W128_OFFSETS, widths), out,
	226	+ "op reference impl")
	227	+ self.assertEqual(out.shape, (rows, sum(widths)))
	228	+
	229	+ @parametrize('rows', [200])
	230	+ @parametrize('dtype', ['float16'])
	231	+ def test_op_handles_mixed_widths(self, rows, dtype):
	232	+ wide = self._wide(rows, dtype)
	233	+ n = len(_MIXED_WIDTHS)
	234	+ out = _MSC([wide], [], [0] * n, list(_MIXED_OFFSETS), list(_MIXED_WIDTHS),
	235	+ [-1] * n)
	236	+ self._assert_bitwise_equal(
	237	+ _col_concat_ref(wide, _MIXED_OFFSETS, _MIXED_WIDTHS), out, "45 segments mixed widths"
	238	+ )
	239	+
	240	+ @parametrize('rows', [200])
	241	+ @parametrize('dtype', ['float16'])
	242	+ def test_op_handles_non_pow2_width(self, rows, dtype):
	243	+ """Model widths are all powers of two, but the template padding branch must work."""
	244	+ wide = self._wide(rows, dtype)
	245	+ offsets, widths = (100, 500, 900), (24, 48, 12)
	246	+ out = _MSC([wide], [], [0] * 3, list(offsets), list(widths), [-1] * 3)
	247	+ self._assert_bitwise_equal(_col_concat_ref(wide, offsets, widths), out,
	248	+ "non power-of-2 widths")
	249	+
	250	+ @parametrize('rows', [32, 200])
	251	+ @parametrize('dtype', ['float16'])
	252	+ def test_op_masked_segment(self, rows, dtype):
	253	+ """A masked segment is equivalent to where(row mask, 0, slice)."""
	254	+ wide = self._wide(rows, dtype)
	255	+ mask = self._mask(rows)
	256	+ offsets, widths = (100, 500, 900), (16, 32, 64)
	257	+ out = _MSC([wide], [mask], [0] * 3, list(offsets), list(widths), [0, -1, 0])
	258	+
	259	+ expected = torch.cat([
	260	+ _masked_slice_ref(wide, mask, 100, 16),
	261	+ wide[..., 500:532],
	262	+ _masked_slice_ref(wide, mask, 900, 964 - 900),
	263	+ ], dim=-1)
	264	+ self._assert_bitwise_equal(expected, out, "masked and bare segments mixed")
	265	+
	266	+ @parametrize('rows', [32])
	267	+ @parametrize('dtype', ['float16'])
	268	+ def test_op_multi_source(self, rows, dtype):
	269	+ """src_idx picks which source tensor each segment reads from."""
	270	+ a = self._wide(rows, dtype, cols=1024)
	271	+ b = self._wide(rows, dtype, cols=512)
	272	+ out = _MSC([a, b], [], [0, 1, 0], [100, 64, 700], [32, 16, 48], [-1] * 3)
	273	+
	274	+ expected = torch.cat([a[..., 100:132], b[..., 64:80], a[..., 700:748]], dim=-1)
	275	+ self._assert_bitwise_equal(expected, out, "two-source segment plan")
	276	+
	277	+ @parametrize('rows', [32])
	278	+ @parametrize('dtype', ['float16'])
	279	+ def test_op_rejects_out_of_range_segment(self, rows, dtype):
	280	+ wide = self._wide(rows, dtype)
	281	+ with self.assertRaises(RuntimeError):
	282	+ _MSC([wide], [], [0], [_WIDE_COLS - 8], [16], [-1])
	283	+
	284	+ @parametrize('rows', [32])
	285	+ @parametrize('dtype', ['float16'])
	286	+ def test_op_rejects_bad_plan(self, rows, dtype):
	287	+ """Mismatched plan lengths or out-of-range indices must raise, not miscompute."""
	288	+ wide = self._wide(rows, dtype)
	289	+ with self.assertRaises(RuntimeError):
	290	+ _MSC([wide], [], [0, 0], [100], [16], [-1])
	291	+ with self.assertRaises(RuntimeError):
	292	+ _MSC([wide], [], [3], [100], [16], [-1])
	293	+ with self.assertRaises(RuntimeError):
	294	+ _MSC([wide], [], [0], [100], [16], [0])
	295	+
	296	+ @parametrize('rows', [32])
	297	+ @parametrize('dtype', ['float16'])
	298	+ def test_op_rejects_non_row_mask(self, rows, dtype):
	299	+ """A non-[rows, 1] mask broadcasts differently, so it must raise, not guess."""
	300	+ wide = self._wide(rows, dtype)
	301	+ bad = torch.randint(0, 2, (rows, 16), device=torch.device("npu"),
	302	+ dtype=torch.bool)
	303	+ with self.assertRaises(RuntimeError):
	304	+ _MSC([wide], [bad], [0], [100], [16], [0])
	305	+
	306	+ # ------------------------------------------------------------------
	307	+ # rewrite under static and dynamic shapes
	308	+ # ------------------------------------------------------------------
	309	+ @parametrize('rows', [32])
	310	+ @parametrize('dtype', ['float16'])
	311	+ @parametrize('dynamic', [False, True])
	312	+ def test_whole_cat_collapses(self, rows, dtype, dynamic):
	313	+ wide = self._wide(rows, dtype)
	314	+ widths = (128,) * len(_W128_OFFSETS)
	315	+
	316	+ def fn(x):
	317	+ return _col_concat(x, _W128_OFFSETS, widths)
	318	+
	319	+ gm = self._run_pass(fn, wide, dynamic=dynamic)
	320	+
	321	+ self.assertEqual(_count(gm.graph, _CAT), 0, "slices should be replaced by the op, cat gone")
	322	+ self.assertEqual(len(_op_nodes(gm.graph)), 1)
	323	+ self.assertEqual(_count(gm.graph, _SLICE), 0, "slices should be eliminated too")
	324	+ self._assert_bitwise_equal(_col_concat_ref(wide, _W128_OFFSETS, widths),
	325	+ gm(wide), f"15 segments width=128 (dynamic={dynamic})")
	326	+
	327	+ @parametrize('rows', [32])
	328	+ @parametrize('dtype', ['float16'])
	329	+ def test_mixed_widths_from_real_model(self, rows, dtype):
	330	+ """L25021 prototype: 45 segments with widths mixed 16/32/128."""
	331	+ wide = self._wide(rows, dtype)
	332	+
	333	+ def fn(x):
	334	+ return _col_concat(x, _MIXED_OFFSETS, _MIXED_WIDTHS)
	335	+
	336	+ gm = self._run_pass(fn, wide)
	337	+ ops = _op_nodes(gm.graph)
	338	+
	339	+ self.assertEqual(_count(gm.graph, _CAT), 0)
	340	+ self.assertEqual(len(ops), 1, "45 segments of one base should collapse into one op")
	341	+ self.assertEqual(len(ops[0].args[_A_OFFSETS]), 45, "all segments must be preserved")
	342	+ self.assertEqual(list(ops[0].args[_A_OFFSETS]), list(_MIXED_OFFSETS),
	343	+ "offset order must not change")
	344	+ self._assert_bitwise_equal(
	345	+ _col_concat_ref(wide, _MIXED_OFFSETS, _MIXED_WIDTHS), gm(wide), "45 segments mixed widths"
	346	+ )
	347	+
	348	+ @parametrize('rows', [32])
	349	+ @parametrize('dtype', ['float16'])
	350	+ def test_narrow_segments(self, rows, dtype):
	351	+ """L27947 prototype: 28 segments of width=8, the narrowest case."""
	352	+ wide = self._wide(rows, dtype)
	353	+ widths = (8,) * len(_W8_OFFSETS)
	354	+
	355	+ def fn(x):
	356	+ return _col_concat(x, _W8_OFFSETS, widths)
	357	+
	358	+ gm = self._run_pass(fn, wide)
	359	+ self.assertEqual(len(_op_nodes(gm.graph)), 1)
	360	+ self._assert_bitwise_equal(_col_concat_ref(wide, _W8_OFFSETS, widths),
	361	+ gm(wide), "28 segments width=8")
	362	+
	363	+ @parametrize('rows', [32])
	364	+ @parametrize('dtype', ['float16'])
	365	+ def test_dynamic_batch_reruns_with_other_shape(self, rows, dtype):
	366	+ wide = self._wide(rows, dtype)
	367	+ widths = (16,) * len(_W16_OFFSETS)
	368	+
	369	+ def fn(x):
	370	+ return _col_concat(x, _W16_OFFSETS, widths)
	371	+
	372	+ gm = self._run_pass(fn, wide, dynamic=True)
	373	+ self.assertEqual(len(_op_nodes(gm.graph)), 1)
	374	+
	375	+ # batch is never 1: dims 0/1 are hard-specialized, so symbols range over [2, inf).
	376	+ for other_rows in (2, rows * 3):
	377	+ other = self._wide(other_rows, dtype)
	378	+ self._assert_bitwise_equal(_col_concat_ref(other, _W16_OFFSETS, widths),
	379	+ gm(other), f"dynamic batch rows={other_rows}")
	380	+
	381	+ @parametrize('rows', [32])
	382	+ @parametrize('dtype', ['float16'])
	383	+ def test_new_node_carries_fake_meta(self, rows, dtype):
	384	+ wide = self._wide(rows, dtype)
	385	+ widths = (16,) * len(_W16_OFFSETS)
	386	+
	387	+ def fn(x):
	388	+ return _col_concat(x, _W16_OFFSETS, widths)
	389	+
	390	+ node = _op_nodes(self._run_pass(fn, wide).graph)[0]
	391	+
	392	+ self.assertIn('val', node.meta, "new node must carry meta['val'] or lowering has no shape")
	393	+ self.assertEqual(node.meta['val'].dtype, torch.float16)
	394	+ self.assertEqual(int(node.meta['val'].shape[-1]), sum(widths))
	395	+
	396	+ # ------------------------------------------------------------------
	397	+ # masked segments: more common than bare slices in production, and what breaks up runs
	398	+ # ------------------------------------------------------------------
	399	+ @parametrize('rows', [32])
	400	+ @parametrize('dtype', ['float16'])
	401	+ @parametrize('dynamic', [False, True])
	402	+ def test_masked_slices_collapse(self, rows, dtype, dynamic):
	403	+ wide = self._wide(rows, dtype)
	404	+ mask = self._mask(rows)
	405	+ offsets, widths = (100, 500, 900, 1300), (16, 32, 64, 16)
	406	+
	407	+ def fn(x, m):
	408	+ return _CAT([_masked_slice(x, m, o, w)
	409	+ for o, w in zip(offsets, widths)], -1)
	410	+
	411	+ gm = self._run_pass(fn, wide, mask, dynamic=dynamic)
	412	+ ops = _op_nodes(gm.graph)
	413	+
	414	+ self.assertEqual(len(ops), 1, "the whole masked slice run should collapse into one op")
	415	+ self.assertEqual(_count(gm.graph, _CAT), 0)
	416	+ self.assertEqual(_count(gm.graph, _WHERE), 0, "where should be eliminated too")
	417	+ self.assertEqual(len(ops[0].args[_A_MASKS]), 1, "one mask must be registered only once")
	418	+ self.assertEqual(list(ops[0].args[_A_MASK_IDX]), [0] * 4)
	419	+
	420	+ expected = torch.cat([_masked_slice_ref(wide, mask, o, w)
	421	+ for o, w in zip(offsets, widths)], dim=-1)
	422	+ self._assert_bitwise_equal(expected, gm(wide, mask),
	423	+ f"masked slices (dynamic={dynamic})")
	424	+
	425	+ @parametrize('rows', [32])
	426	+ @parametrize('dtype', ['float16'])
	427	+ def test_interleaved_masked_and_bare_collapse(self, rows, dtype):
	428	+ """cat_124 prototype: masked and bare slices strictly alternating.
	429	+
	430	+ Bare-slice-only matching merges nothing here, since no two bare slices are
	431	+ adjacent. Covering masked segments joins the run and is the main gain.
	432	+ """
	433	+ wide = self._wide(rows, dtype)
	434	+ mask = self._mask(rows)
	435	+
	436	+ def fn(x, m):
	437	+ parts = []
	438	+ for i, (off, w) in enumerate(zip(_ALT_OFFSETS, _ALT_WIDTHS)):
	439	+ parts.append(_masked_slice(x, m, off, w) if i % 2 == 0
	440	+ else _SLICE(x, -1, off, off + w))
	441	+ return _CAT(parts, -1)
	442	+
	443	+ gm = self._run_pass(fn, wide, mask)
	444	+ ops = _op_nodes(gm.graph)
	445	+
	446	+ self.assertEqual(len(ops), 1, "an alternating layout should form one run")
	447	+ self.assertEqual(len(ops[0].args[_A_OFFSETS]), len(_ALT_OFFSETS))
	448	+ self.assertEqual(list(ops[0].args[_A_MASK_IDX]), [0, -1, 0, -1, 0, -1, 0, -1],
	449	+ "masked and bare segments must each be recorded correctly")
	450	+
	451	+ expected = torch.cat([
	452	+ _masked_slice_ref(wide, mask, off, w) if i % 2 == 0
	453	+ else wide[..., off:off + w]
	454	+ for i, (off, w) in enumerate(zip(_ALT_OFFSETS, _ALT_WIDTHS))
	455	+ ], dim=-1)
	456	+ self._assert_bitwise_equal(expected, gm(wide, mask), "masked and bare slices alternating")
	457	+
	458	+ @parametrize('rows', [32])
	459	+ @parametrize('dtype', ['float16'])
	460	+ def test_mask_reshaped_before_where_is_matched(self, rows, dtype):
	461	+ """Broadcasts of the condition (reshape/expand to width) must be stripped."""
	462	+ wide = self._wide(rows, dtype)
	463	+ mask = self._mask(rows)
	464	+
	465	+ def fn(x, m):
	466	+ parts = []
	467	+ for off, w in ((100, 16), (500, 32)):
	468	+ cond = torch.ops.aten.expand.default(
	469	+ torch.ops.aten.reshape.default(m, [-1, 1]), [x.shape[0], w]
	470	+ )
	471	+ zero = _FULL([x.shape[0], w], 0, dtype=x.dtype, device=x.device)
	472	+ parts.append(_WHERE(cond, zero, _SLICE(x, -1, off, off + w)))
	473	+ return _CAT(parts, -1)
	474	+
	475	+ gm = self._run_pass(fn, wide, mask)
	476	+ ops = _op_nodes(gm.graph)
	477	+
	478	+ self.assertEqual(len(ops), 1, "a broadcast condition must still count as a row mask")
	479	+ self.assertEqual(ops[0].args[_A_MASKS][0].op, "placeholder",
	480	+ "the registered mask must be the one before broadcast")
	481	+ expected = torch.cat([_masked_slice_ref(wide, mask, 100, 16),
	482	+ _masked_slice_ref(wide, mask, 500, 32)], dim=-1)
	483	+ self._assert_bitwise_equal(expected, gm(wide, mask), "broadcast condition")
	484	+
	485	+ @parametrize('rows', [32])
	486	+ @parametrize('dtype', ['float16'])
	487	+ def test_mask_chain_stops_at_deepest_row_mask(self, rows, dtype):
	488	+ """Broadcast stripping must stop at the deepest [rows, 1].
	489	+
	490	+ The producer here is [1, 1], so stripping all the way down lands on a node
	491	+ whose row count does not match and drops the segment; the intermediate
	492	+ [rows, 1] is the real row mask.
	493	+ """
	494	+ wide = self._wide(rows, dtype)
	495	+ seed = self._mask(1)
	496	+
	497	+ def fn(x, s):
	498	+ row = torch.ops.aten.expand.default(s, [x.shape[0], 1])
	499	+ parts = []
	500	+ for off, w in ((100, 16), (500, 32)):
	501	+ cond = torch.ops.aten.expand.default(row, [x.shape[0], w])
	502	+ zero = _FULL([x.shape[0], w], 0, dtype=x.dtype, device=x.device)
	503	+ parts.append(_WHERE(cond, zero, _SLICE(x, -1, off, off + w)))
	504	+ return _CAT(parts, -1)
	505	+
	506	+ gm = self._run_pass(fn, wide, seed)
	507	+ ops = _op_nodes(gm.graph)
	508	+
	509	+ self.assertEqual(len(ops), 1, "a qualifying row mask mid-chain must not drop the segment")
	510	+ self.assertEqual(len(ops[0].args[_A_MASKS]), 1, "one mask must be registered only once")
	511	+ recorded = ops[0].args[_A_MASKS][0]
	512	+ self.assertEqual(tuple(recorded.meta["val"].shape), (rows, 1),
	513	+ "the registered mask must be [rows, 1]")
	514	+ self.assertNotEqual(recorded.op, "placeholder",
	515	+ "a [1, 1] graph input is not a row mask and must not be registered")
	516	+
	517	+ row = seed.expand(rows, 1)
	518	+ expected = torch.cat([_masked_slice_ref(wide, row, 100, 16),
	519	+ _masked_slice_ref(wide, row, 500, 32)], dim=-1)
	520	+ self._assert_bitwise_equal(expected, gm(wide, seed), "scalar broadcast mask")
	521	+
	522	+ @parametrize('rows', [32])
	523	+ @parametrize('dtype', ['float16'])
	524	+ def test_two_masks_registered_separately(self, rows, dtype):
	525	+ wide = self._wide(rows, dtype)
	526	+ m0, m1 = self._mask(rows), self._mask(rows)
	527	+
	528	+ def fn(x, a, b):
	529	+ return _CAT([_masked_slice(x, a, 100, 16), _masked_slice(x, b, 500, 32),
	530	+ _masked_slice(x, a, 900, 16)], -1)
	531	+
	532	+ gm = self._run_pass(fn, wide, m0, m1)
	533	+ ops = _op_nodes(gm.graph)
	534	+
	535	+ self.assertEqual(len(ops), 1)
	536	+ self.assertEqual(len(ops[0].args[_A_MASKS]), 2, "two distinct masks must each be registered")
	537	+ self.assertEqual(list(ops[0].args[_A_MASK_IDX]), [0, 1, 0], "mask indices must be reused")
	538	+ expected = torch.cat([_masked_slice_ref(wide, m0, 100, 16),
	539	+ _masked_slice_ref(wide, m1, 500, 32),
	540	+ _masked_slice_ref(wide, m0, 900, 16)], dim=-1)
	541	+ self._assert_bitwise_equal(expected, gm(wide, m0, m1), "two masks")
	542	+
	543	+ # ------------------------------------------------------------------
	544	+ # multi-source: one concat often mixes several model inputs
	545	+ # ------------------------------------------------------------------
	546	+ @parametrize('rows', [32])
	547	+ @parametrize('dtype', ['float16'])
	548	+ def test_two_bases_merge_into_one_run(self, rows, dtype):
	549	+ """Slices of different bases can share one op; each source costs one pointer arg."""
	550	+ a, b = self._wide(rows, dtype), self._wide(rows, dtype)
	551	+
	552	+ def fn(x, y):
	553	+ return _CAT([
	554	+ _SLICE(x, -1, 100, 116), _SLICE(x, -1, 900, 916),
	555	+ _SLICE(y, -1, 500, 532), _SLICE(y, -1, 700, 732),
	556	+ ], -1)
	557	+
	558	+ gm = self._run_pass(fn, a, b)
	559	+ ops = _op_nodes(gm.graph)
	560	+
	561	+ self.assertEqual(len(ops), 1, "a contiguous multi-source run should collapse into one op")
	562	+ self.assertEqual(len(ops[0].args[_A_SRCS]), 2, "both sources must be registered")
	563	+ self.assertEqual(list(ops[0].args[_A_SRC_IDX]), [0, 0, 1, 1], "source indices must match")
	564	+ expected = torch.cat([a[..., 100:116], a[..., 900:916],
	565	+ b[..., 500:532], b[..., 700:732]], dim=-1)
	566	+ self._assert_bitwise_equal(expected, gm(a, b), "two base tensors")
	567	+
	568	+ @parametrize('rows', [32])
	569	+ @parametrize('dtype', ['float16'])
	570	+ def test_source_cap_splits_run(self, rows, dtype):
	571	+ a, b = self._wide(rows, dtype), self._wide(rows, dtype)
	572	+
	573	+ def fn(x, y):
	574	+ return _CAT([
	575	+ _SLICE(x, -1, 100, 116), _SLICE(x, -1, 900, 916),
	576	+ _SLICE(y, -1, 500, 532), _SLICE(y, -1, 700, 732),
	577	+ ], -1)
	578	+
	579	+ gm = self._run_pass(fn, a, b, max_sources=1)
	580	+ ops = _op_nodes(gm.graph)
	581	+
	582	+ self.assertEqual(len(ops), 2, "with a source cap of 1 the run must split per base")
	583	+ for node in ops:
	584	+ self.assertEqual(len(node.args[_A_SRCS]), 1)
	585	+ expected = torch.cat([a[..., 100:116], a[..., 900:916],
	586	+ b[..., 500:532], b[..., 700:732]], dim=-1)
	587	+ self._assert_bitwise_equal(expected, gm(a, b), "source cap split")
	588	+
	589	+ @parametrize('rows', [32])
	590	+ @parametrize('dtype', ['float16'])
	591	+ def test_mask_cap_splits_run(self, rows, dtype):
	592	+ wide = self._wide(rows, dtype)
	593	+ m0, m1 = self._mask(rows), self._mask(rows)
	594	+
	595	+ def fn(x, a, b):
	596	+ return _CAT([_masked_slice(x, a, 100, 16), _masked_slice(x, a, 300, 16),
	597	+ _masked_slice(x, b, 500, 32), _masked_slice(x, b, 700, 32)], -1)
	598	+
	599	+ gm = self._run_pass(fn, wide, m0, m1, max_masks=1)
	600	+ ops = _op_nodes(gm.graph)
	601	+
	602	+ self.assertEqual(len(ops), 2, "with a mask cap of 1 the run must split per mask")
	603	+ for node in ops:
	604	+ self.assertEqual(len(node.args[_A_MASKS]), 1)
	605	+ expected = torch.cat([_masked_slice_ref(wide, m0, 100, 16),
	606	+ _masked_slice_ref(wide, m0, 300, 16),
	607	+ _masked_slice_ref(wide, m1, 500, 32),
	608	+ _masked_slice_ref(wide, m1, 700, 32)], dim=-1)
	609	+ self._assert_bitwise_equal(expected, gm(wide, m0, m1), "mask cap split")
	610	+
	611	+ @parametrize('rows', [32])
	612	+ @parametrize('dtype', ['float16'])
	613	+ def test_masks_disabled_falls_back_to_bare_slices(self, rows, dtype):
	614	+ wide = self._wide(rows, dtype)
	615	+ mask = self._mask(rows)
	616	+
	617	+ def fn(x, m):
	618	+ return _CAT([_masked_slice(x, m, 100, 16), _masked_slice(x, m, 500, 32)], -1)
	619	+
	620	+ gm = self._run_pass(fn, wide, mask, max_masks=0)
	621	+ self.assertEqual(len(_op_nodes(gm.graph)), 0, "no rewrite when masks are disabled")
	622	+
	623	+ # ------------------------------------------------------------------
	624	+ # slices mixed with other inputs: only contiguous runs merge, order is preserved
	625	+ # ------------------------------------------------------------------
	626	+ @parametrize('rows', [32])
	627	+ @parametrize('dtype', ['float16'])
	628	+ def test_slices_mixed_with_dense_input(self, rows, dtype):
	629	+ wide = self._wide(rows, dtype)
	630	+ dense = torch.randn((rows, 64), dtype=torch.float16, device=torch.device("npu"))
	631	+ widths = (16,) * len(_W16_OFFSETS)
	632	+
	633	+ def fn(x, d):
	634	+ parts = [_SLICE(x, -1, off, off + w) for off, w in zip(_W16_OFFSETS, widths)]
	635	+ return _CAT(parts + [torch.ops.aten.relu.default(d)], -1)
	636	+
	637	+ gm = self._run_pass(fn, wide, dense)
	638	+
	639	+ self.assertEqual(len(_op_nodes(gm.graph)), 1, "the slice run should collapse into one op")
	640	+ self.assertEqual(_count(gm.graph, _CAT), 1, "outer cat must stay to join the dense input")
	641	+ expected = torch.cat(
	642	+ [_col_concat_ref(wide, _W16_OFFSETS, widths), torch.relu(dense)], dim=-1
	643	+ )
	644	+ self._assert_bitwise_equal(expected, gm(wide, dense), "slice run + dense input")
	645	+
	646	+ @parametrize('rows', [32])
	647	+ @parametrize('dtype', ['float16'])
	648	+ def test_interleaved_slices_split_into_runs(self, rows, dtype):
	649	+ wide = self._wide(rows, dtype)
	650	+ dense = torch.randn((rows, 8), dtype=torch.float16, device=torch.device("npu"))
	651	+
	652	+ def fn(x, d):
	653	+ return _CAT([
	654	+ _SLICE(x, -1, 100, 116),
	655	+ _SLICE(x, -1, 500, 516),
	656	+ torch.ops.aten.relu.default(d),
	657	+ _SLICE(x, -1, 900, 916),
	658	+ _SLICE(x, -1, 1300, 1316),
	659	+ ], -1)
	660	+
	661	+ gm = self._run_pass(fn, wide, dense)
	662	+
	663	+ self.assertEqual(len(_op_nodes(gm.graph)), 2, "each slice run collapses into its own op")
	664	+ expected = torch.cat([
	665	+ wide[..., 100:116], wide[..., 500:516], torch.relu(dense),
	666	+ wide[..., 900:916], wide[..., 1300:1316],
	667	+ ], dim=-1)
	668	+ self._assert_bitwise_equal(expected, gm(wide, dense), "slices split apart")
	669	+
	670	+ @parametrize('rows', [32])
	671	+ @parametrize('dtype', ['float16'])
	672	+ def test_segment_cap_splits_run(self, rows, dtype):
	673	+ wide = self._wide(rows, dtype)
	674	+ widths = (128,) * len(_W128_OFFSETS)
	675	+
	676	+ def fn(x):
	677	+ return _col_concat(x, _W128_OFFSETS, widths)
	678	+
	679	+ gm = self._run_pass(fn, wide, max_segments=4)
	680	+ ops = _op_nodes(gm.graph)
	681	+
	682	+ self.assertEqual(len(ops), 4, "15 segments capped at 4 per group make 4 groups")
	683	+ for node in ops:
	684	+ self.assertLessEqual(len(node.args[_A_OFFSETS]), 4, "group exceeds the segment cap")
	685	+ self._assert_bitwise_equal(_col_concat_ref(wide, _W128_OFFSETS, widths),
	686	+ gm(wide), "segment cap split")
	687	+
	688	+ # ------------------------------------------------------------------
	689	+ # profitability: merging costs one dispatch, so a segment paying less than two loses
	690	+ # ------------------------------------------------------------------
	691	+ @parametrize('rows', [32])
	692	+ @parametrize('dtype', ['float16'])
	693	+ def test_whole_inputs_alone_not_rewritten(self, rows, dtype):
	694	+ """aclnnCat reads contiguous inputs directly; folding them in adds a dispatch."""
	695	+ a = self._wide(rows, dtype, cols=64)
	696	+ b = self._wide(rows, dtype, cols=32)
	697	+
	698	+ def fn(x, y):
	699	+ return _CAT([x, y], -1)
	700	+
	701	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, a, b).graph)), 0,
	702	+ "a plain concat of two dense inputs must not be rewritten")
	703	+
	704	+ @parametrize('rows', [32])
	705	+ @parametrize('dtype', ['float16'])
	706	+ def test_whole_input_joins_paying_run(self, rows, dtype):
	707	+ """A whole input saves nothing alone, but rides along a run that already pays."""
	708	+ wide = self._wide(rows, dtype)
	709	+ dense = self._wide(rows, dtype, cols=64)
	710	+
	711	+ def fn(x, d):
	712	+ return _CAT([_SLICE(x, -1, 100, 116), d, _SLICE(x, -1, 900, 916)], -1)
	713	+
	714	+ gm = self._run_pass(fn, wide, dense)
	715	+ ops = _op_nodes(gm.graph)
	716	+
	717	+ self.assertEqual(len(ops), 1, "two slices already pay off, so the whole input can join")
	718	+ self.assertEqual(_count(gm.graph, _CAT), 0)
	719	+ expected = torch.cat([wide[..., 100:116], dense, wide[..., 900:916]], dim=-1)
	720	+ self._assert_bitwise_equal(expected, gm(wide, dense), "whole input between slices")
	721	+
	722	+ @parametrize('rows', [32])
	723	+ @parametrize('dtype', ['float16'])
	724	+ def test_single_masked_slice_with_whole_input_not_rewritten(self, rows, dtype):
	725	+ """Only one segment costs a dispatch, so the rewrite gains nothing."""
	726	+ wide = self._wide(rows, dtype)
	727	+ mask = self._mask(rows)
	728	+ dense = self._wide(rows, dtype, cols=64)
	729	+
	730	+ def fn(x, m, d):
	731	+ return _CAT([_masked_slice(x, m, 100, 16), d], -1)
	732	+
	733	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, wide, mask, dense).graph)), 0)
	734	+
	735	+ # ------------------------------------------------------------------
	736	+ # must reject: a false match silently miscomputes, far worse than not optimizing
	737	+ # ------------------------------------------------------------------
	738	+ @parametrize('rows', [32])
	739	+ @parametrize('dtype', ['float16'])
	740	+ def test_single_slice_not_rewritten(self, rows, dtype):
	741	+ """One slice means one dispatch before and after, so there is no gain."""
	742	+ wide = self._wide(rows, dtype)
	743	+ dense = torch.randn((rows, 16), dtype=torch.float16, device=torch.device("npu"))
	744	+
	745	+ def fn(x, d):
	746	+ return _CAT([_SLICE(x, -1, 100, 116), torch.ops.aten.relu.default(d)], -1)
	747	+
	748	+ gm = self._run_pass(fn, wide, dense)
	749	+ self.assertEqual(len(_op_nodes(gm.graph)), 0)
	750	+
	751	+ @parametrize('rows', [32])
	752	+ @parametrize('dtype', ['float16'])
	753	+ def test_row_direction_cat_not_rewritten(self, rows, dtype):
	754	+ wide = self._wide(rows, dtype)
	755	+
	756	+ def fn(x):
	757	+ return _CAT([_SLICE(x, -1, 0, 16), _SLICE(x, -1, 64, 80)], 0)
	758	+
	759	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, wide).graph)), 0)
	760	+
	761	+ @parametrize('rows', [32])
	762	+ @parametrize('dtype', ['float16'])
	763	+ def test_slice_on_other_dim_not_rewritten(self, rows, dtype):
	764	+ """The slice axis differs from the concat axis, so these are not column segments."""
	765	+ wide = self._wide(rows, dtype)
	766	+
	767	+ def fn(x):
	768	+ return _CAT([_SLICE(x, 0, 0, 8), _SLICE(x, 0, 8, 16)], -1)
	769	+
	770	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, wide).graph)), 0)
	771	+
	772	+ @parametrize('rows', [32])
	773	+ @parametrize('dtype', ['float16'])
	774	+ def test_strided_slice_not_rewritten(self, rows, dtype):
	775	+ wide = self._wide(rows, dtype)
	776	+
	777	+ def fn(x):
	778	+ return _CAT([_SLICE(x, -1, 0, 32, 2), _SLICE(x, -1, 64, 96, 2)], -1)
	779	+
	780	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, wide).graph)), 0,
	781	+ "strided slices are not fixed contiguous segments, must be rejected")
	782	+
	783	+ @parametrize('rows', [32])
	784	+ @parametrize('dtype', ['float16'])
	785	+ def test_mixed_dtype_not_rewritten(self, rows, dtype):
	786	+ """cat promotes mismatched dtypes, the op does not, so leave it to the original."""
	787	+ a, b = self._wide(rows, 'float16'), self._wide(rows, 'float32')
	788	+
	789	+ def fn(x, y):
	790	+ return _CAT([_SLICE(x, -1, 0, 16), _SLICE(y, -1, 0, 16)], -1)
	791	+
	792	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, a, b).graph)), 0)
	793	+
	794	+ @parametrize('rows', [8])
	795	+ @parametrize('dtype', ['float16'])
	796	+ def test_three_dim_not_rewritten(self, rows, dtype):
	797	+ """The op only supports 2D; higher rank goes back to the original cat."""
	798	+ x = torch.randn((rows, 4, 256), dtype=torch.float16, device=torch.device("npu"))
	799	+
	800	+ def fn(t):
	801	+ return _CAT([_SLICE(t, -1, 0, 16), _SLICE(t, -1, 64, 80)], -1)
	802	+
	803	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, x).graph)), 0)
	804	+
	805	+ @parametrize('rows', [32])
	806	+ @parametrize('dtype', ['float16'])
	807	+ def test_computed_base_not_rewritten(self, rows, dtype):
	808	+ """Skip computed bases so they are not forced to materialize early.
	809	+
	810	+ In the target case the wide table is a model input and already landed. For a
	811	+ computed tensor, realize_input in lowering would demand a contiguous buffer
	812	+ first, which can block fusion on the producer side.
	813	+ """
	814	+ wide = self._wide(rows, dtype)
	815	+
	816	+ def fn(x):
	817	+ base = torch.ops.aten.mul.Tensor(x, 2.0)
	818	+ return _CAT([_SLICE(base, -1, 0, 16), _SLICE(base, -1, 64, 80)], -1)
	819	+
	820	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, wide).graph)), 0,
	821	+ "a base that is not a graph input must fall back to the original cat")
	822	+
	823	+ @parametrize('rows', [32])
	824	+ @parametrize('dtype', ['float16'])
	825	+ def test_masked_computed_base_not_rewritten(self, rows, dtype):
	826	+ """The selected value of a masked segment must also come from a graph input."""
	827	+ wide = self._wide(rows, dtype)
	828	+ mask = self._mask(rows)
	829	+
	830	+ def fn(x, m):
	831	+ base = torch.ops.aten.mul.Tensor(x, 2.0)
	832	+ return _CAT([_masked_slice(base, m, 100, 16),
	833	+ _masked_slice(base, m, 500, 32)], -1)
	834	+
	835	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, wide, mask).graph)), 0)
	836	+
	837	+ @parametrize('rows', [32])
	838	+ @parametrize('dtype', ['float16'])
	839	+ def test_elementwise_mask_not_rewritten(self, rows, dtype):
	840	+ """An elementwise mask is not row-wise, so the op cannot express it."""
	841	+ wide = self._wide(rows, dtype)
	842	+
	843	+ def fn(x, m0, m1):
	844	+ zero0 = _FULL([x.shape[0], 16], 0, dtype=x.dtype, device=x.device)
	845	+ zero1 = _FULL([x.shape[0], 16], 0, dtype=x.dtype, device=x.device)
	846	+ return _CAT([_WHERE(m0, zero0, _SLICE(x, -1, 100, 116)),
	847	+ _WHERE(m1, zero1, _SLICE(x, -1, 500, 516))], -1)
	848	+
	849	+ wide_mask0 = torch.randint(0, 2, (rows, 16), device=torch.device("npu"),
	850	+ dtype=torch.bool)
	851	+ wide_mask1 = torch.randint(0, 2, (rows, 16), device=torch.device("npu"),
	852	+ dtype=torch.bool)
	853	+ gm = self._run_pass(fn, wide, wide_mask0, wide_mask1)
	854	+ self.assertEqual(len(_op_nodes(gm.graph)), 0, "an elementwise mask must be rejected")
	855	+
	856	+ @parametrize('rows', [32])
	857	+ @parametrize('dtype', ['float16'])
	858	+ def test_nonzero_fill_not_rewritten(self, rows, dtype):
	859	+ wide = self._wide(rows, dtype)
	860	+ mask = self._mask(rows)
	861	+
	862	+ def fn(x, m):
	863	+ parts = []
	864	+ for off, w in ((100, 16), (500, 32)):
	865	+ one = _FULL([x.shape[0], w], 1, dtype=x.dtype, device=x.device)
	866	+ parts.append(_WHERE(m, one, _SLICE(x, -1, off, off + w)))
	867	+ return _CAT(parts, -1)
	868	+
	869	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, wide, mask).graph)), 0)
	870	+
	871	+ @parametrize('rows', [32])
	872	+ @parametrize('dtype', ['float16'])
	873	+ def test_where_between_two_slices_not_rewritten(self, rows, dtype):
	874	+ """Selecting between two slices reads two places, not zero-fill; unsupported."""
	875	+ wide = self._wide(rows, dtype)
	876	+ mask = self._mask(rows)
	877	+
	878	+ def fn(x, m):
	879	+ return _CAT([
	880	+ _WHERE(m, _SLICE(x, -1, 100, 116), _SLICE(x, -1, 200, 216)),
	881	+ _WHERE(m, _SLICE(x, -1, 500, 516), _SLICE(x, -1, 600, 616)),
	882	+ ], -1)
	883	+
	884	+ self.assertEqual(len(_op_nodes(self._run_pass(fn, wide, mask).graph)), 0)
	885	+
	886	+ @parametrize('rows', [32])
	887	+ @parametrize('dtype', ['float16'])
	888	+ def test_slice_shared_with_other_consumer_is_kept(self, rows, dtype):
	889	+ wide = self._wide(rows, dtype)
	890	+
	891	+ def fn(x):
	892	+ a = _SLICE(x, -1, 55818, 55834)
	893	+ b = _SLICE(x, -1, 56436, 56452)
	894	+ # a feeds both the cat and a pointwise op, the form Inductor fuses.
	895	+ return _CAT([a, b], -1), torch.ops.aten.mul.Tensor(a, 2.0)
	896	+
	897	+ gm = self._run_pass(fn, wide)
	898	+ self.assertEqual(len(_op_nodes(gm.graph)), 1, "the cat path should still be rewritten")
	899	+ self.assertGreaterEqual(_count(gm.graph, _SLICE), 1,
	900	+ "a slice used elsewhere must not be deleted")
	901	+ expected = fn(wide)
	902	+ actual = gm(wide)
	903	+ self._assert_bitwise_equal(expected[0], actual[0], "cat result with a shared slice")
	904	+ self._assert_bitwise_equal(expected[1], actual[1], "other consumer of a shared slice")
	905	+
	906	+ # ------------------------------------------------------------------
	907	+ # end to end; needs the flag on, otherwise pass and lowering are not registered
	908	+ # ------------------------------------------------------------------
	909	+ @unittest.skipUnless(_ENABLED, _NEEDS_FLAG)
	910	+ @parametrize('rows', [200])
	911	+ @parametrize('dtype', ['float16'])
	912	+ @parametrize('dynamic', [False, True])
	913	+ def test_compile_end_to_end(self, rows, dtype, dynamic):
	914	+ """The dynamic case marks only batch as symbolic and cannot use
	915	+ ``torch.compile(dynamic=True)``: that also symbolizes the column count, and
	916	+ the pass needs a static last dim to read constant offsets, so it silently
	917	+ falls back to aten.cat. Values stay correct but the pass is not covered.
	918	+ """
	919	+ widths = (16,) * len(_W16_OFFSETS)
	920	+ wide = self._wide(rows, dtype)
	921	+
	922	+ def fn(x):
	923	+ return _col_concat_ref(x, _W16_OFFSETS, widths)
	924	+
	925	+ with torch.no_grad(), _pass_spy() as stats:
	926	+ compiled = torch.compile(fn, backend="inductor", dynamic=False)
	927	+ if dynamic:
	928	+ torch._dynamo.mark_dynamic(wide, 0)
	929	+ self._assert_bitwise_equal(fn(wide), compiled(wide), "compiled artifact")
	930	+ self.assertGreater(stats["nodes"], 0,
	931	+ f"pass did not fire (dynamic={dynamic}), still aten.cat")
	932	+ if dynamic:
	933	+ # another batch must reuse the same artifact. never 1: dims 0/1 are
	934	+ # hard-specialized, backed symbols are [2, inf), so 1 recompiles.
	935	+ other = self._wide(rows * 2, dtype)
	936	+ torch._dynamo.mark_dynamic(other, 0)
	937	+ self._assert_bitwise_equal(fn(other), compiled(other), "dynamic batch artifact")
	938	+
	939	+ @unittest.skipUnless(_ENABLED, _NEEDS_FLAG)
	940	+ @parametrize('rows', [200])
	941	+ @parametrize('dtype', ['float16'])
	942	+ def test_compile_masked_end_to_end(self, rows, dtype):
	943	+ wide = self._wide(rows, dtype)
	944	+ mask = self._mask(rows)
	945	+ plan = ((100, 16), (500, 32), (900, 64), (1300, 16))
	946	+
	947	+ def fn(x, m):
	948	+ return torch.cat([_masked_slice_ref(x, m, off, w) for off, w in plan], dim=-1)
	949	+
	950	+ with torch.no_grad(), _pass_spy() as stats:
	951	+ compiled = torch.compile(fn, backend="inductor", dynamic=False)
	952	+ self._assert_bitwise_equal(fn(wide, mask), compiled(wide, mask), "masked artifact")
	953	+ self.assertGreater(stats["nodes"], 0, "masked segments were not rewritten")
	954	+
	955	+ @unittest.skipUnless(_ENABLED, _NEEDS_FLAG)
	956	+ @parametrize('rows', [200])
	957	+ @parametrize('dtype', ['float16'])
	958	+ def test_compile_aliased_masks_end_to_end(self, rows, dtype):
	959	+ """Two mask nodes landing on the same buffer must still compile.
	960	+
	961	+ One row mask is squeezed and reshaped twice; squeeze is not a stripped
	962	+ broadcast, so both chains keep a reshape node and the pass registers two
	963	+ masks. In lowering both views point at the same ComputedBuffer, and
	964	+ def_kernel registers args by buffer name, so the entries overwrite each
	965	+ other, the first arg is left undefined, and the rendered kernel raises
	966	+ NameError. is_live_type in the production model has this exact shape.
	967	+ """
	968	+ wide = self._wide(rows, dtype)
	969	+ flag = self._mask(rows)
	970	+ plan = ((100, 16), (500, 32), (900, 64), (1300, 16))
	971	+
	972	+ def fn(x, f):
	973	+ cond = torch.logical_not(f)
	974	+ m0 = cond.squeeze().reshape(-1, 1)
	975	+ m1 = cond.squeeze().reshape(-1, 1)
	976	+ picks = (m0, m1, m1, m0)
	977	+ return torch.cat(
	978	+ [_masked_slice_ref(x, m, off, w)
	979	+ for m, (off, w) in zip(picks, plan)], dim=-1)
	980	+
	981	+ with torch.no_grad(), _pass_spy() as stats:
	982	+ compiled = torch.compile(fn, backend="inductor", dynamic=False)
	983	+ self._assert_bitwise_equal(fn(wide, flag), compiled(wide, flag),
	984	+ "aliased masks artifact")
	985	+ self.assertGreater(stats["nodes"], 0, "aliased mask segments were not rewritten")
	986	+
	987	+ # ------------------------------------------------------------------
	988	+ # input dedup in lowering, pinned here since the case above needs the whole chain
	989	+ # ------------------------------------------------------------------
	990	+ def test_dedup_inputs_merges_aliased_buffers(self):
	991	+ """Duplicate inputs on one buffer merge, indices are remapped, and -1 stays."""
	992	+ first, second = _StubInput("buf5"), _StubInput("buf5")
	993	+ nodes, indices = _dedup_inputs([first, second], [0, -1, 1, 1, 0])
	994	+ self.assertEqual(nodes, [first])
	995	+ self.assertEqual(indices, [0, -1, 0, 0, 0])
	996	+
	997	+ def test_dedup_inputs_keeps_distinct_buffers(self):
	998	+ first, second = _StubInput("buf5"), _StubInput("buf7")
	999	+ nodes, indices = _dedup_inputs([first, second], [1, 0, -1])
	1000	+ self.assertEqual(nodes, [first, second])
	1001	+ self.assertEqual(indices, [1, 0, -1])
	1002	+
	1003	+ def test_dedup_inputs_rejects_same_buffer_different_layout(self):
	1004	+ """Same name, different layouts fight over one arg slot; caller must fall back."""
	1005	+ broadcast = _StubInput("buf5", stride=(0, 1))
	1006	+ contiguous = _StubInput("buf5", stride=(1, 1))
	1007	+ self.assertEqual(_dedup_inputs([broadcast, contiguous], [0, 1]), (None, None))
	1008	+
	1009	+
	1010	+instantiate_parametrized_tests(TestMultiSliceConcatPass)
	1011	+
	1012	+
	1013	+if __name__ == "__main__":
	1014	+ run_tests()

Atest/_inductor/test_stack_sum_pass.py+394-0

@@ -0,0 +1,394 @@
	1	+import torch
	2	+import torch.fx as fx
	3	+from torch.fx.experimental.proxy_tensor import make_fx
	4	+from torch.fx.passes.shape_prop import ShapeProp
	5	+from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests
	6	+from testutils import TestUtils
	7	+import torch_npu
	8	+import torch_npu._inductor
	9	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	10	+ stack_sum_to_add_chain_pass,
	11	+)
	12	+
	13	+
	14	+# From real model output code: a [B, 69876] fp16 wide feature table, cut into fixed-width
	15	+# column blocks that are stacked and summed. All offsets are compile-time constants.
	16	+_WIDE_COLS = 69876
	17	+_GROUP_W128_OFFSETS = (55834, 56452, 45421, 44203, 44744, 45962, 60741, 61328,
	18	+ 62267, 65602, 66495, 66956, 67445, 68021, 69060)
	19	+_USER_W128_OFFSETS = (55580, 59302)
	20	+_GROUP_W16_OFFSETS = (55818, 56436, 45405, 44187, 44728, 45946, 61312, 62251,
	21	+ 69044, 68199, 68221, 69822)
	22	+# Widest stack pooling in the model (28-way); used to stress accumulation error.
	23	+_GROUP_W8_OFFSETS = (48025, 47595, 46868, 47890, 48779, 49346, 47769, 49082,
	24	+ 49180, 48337, 48467, 48597, 48402, 48532, 48714, 49240,
	25	+ 48070, 48204, 48119, 11076, 14257, 18643, 11922, 15103,
	26	+ 19489, 12752, 15933, 20319)
	27	+# fp16 relative precision, used as the magnitude of one ulp.
	28	+_FP16_EPS = 2 ** -10
	29	+
	30	+
	31	+def _count(graph, target):
	32	+ return len([n for n in graph.nodes if n.op == "call_function" and n.target == target])
	33	+
	34	+
	35	+def _fm_pool(wide, offsets, width):
	36	+ """Reproduce the [aten.slice, aten.stack] -> sum combination from output code."""
	37	+ slices = [
	38	+ torch.ops.aten.slice.Tensor(wide, 1, off, off + width) for off in offsets
	39	+ ]
	40	+ cat = torch.ops.aten.cat.default(slices, 0)
	41	+ view = torch.ops.aten.reshape.default(cat, [len(offsets), wide.shape[0], width])
	42	+ return torch.ops.aten.sum.dim_IntList(view, [0])
	43	+
	44	+
	45	+def _fm_pool_ref(wide, offsets, width):
	46	+ return torch.stack([wide[:, off:off + width] for off in offsets]).sum(0)
	47	+
	48	+
	49	+class ViewCatSumModel(torch.nn.Module):
	50	+ """Stack form after Inductor simplification: cat(dim=0) -> reshape -> sum(dim=0)."""
	51	+
	52	+ def forward(self, t1, t2, t3):
	53	+ cat = torch.ops.aten.cat.default([t1, t2, t3], 0)
	54	+ view = torch.ops.aten.reshape.default(cat, [3, t1.shape[0], t1.shape[1]])
	55	+ return torch.ops.aten.sum.dim_IntList(view, [0])
	56	+
	57	+
	58	+class UnsqueezeCatSumModel(torch.nn.Module):
	59	+ """Original stack form: unsqueeze -> cat -> sum."""
	60	+
	61	+ def forward(self, t1, t2, t3):
	62	+ cat = torch.ops.aten.cat.default(
	63	+ [
	64	+ torch.ops.aten.unsqueeze.default(t1, 0),
	65	+ torch.ops.aten.unsqueeze.default(t2, 0),
	66	+ torch.ops.aten.unsqueeze.default(t3, 0),
	67	+ ],
	68	+ 0,
	69	+ )
	70	+ return torch.ops.aten.sum.dim_IntList(cat, [0])
	71	+
	72	+
	73	+class TestStackSumPass(TestUtils):
	74	+ def _assert_close(self, expected, actual, dtype):
	75	+ # At magnitude ~6000 one fp16 ulp is already several units, so compare relative error.
	76	+ tol = 1e-2 if dtype in ('float16', 'bfloat16') else 1e-4
	77	+ self.assertEqual(expected, actual, atol=tol, rtol=tol)
	78	+
	79	+ def _trace(self, model, *tensors):
	80	+ gm = fx.symbolic_trace(model)
	81	+ ShapeProp(gm).propagate(*tensors)
	82	+ return gm
	83	+
	84	+ def _run_pass(self, model, *tensors):
	85	+ gm = self._trace(model, *tensors)
	86	+ stack_sum_to_add_chain_pass(gm.graph)
	87	+ gm.recompile()
	88	+ return gm
	89	+
	90	+ @parametrize('shape', [(8, 16)])
	91	+ @parametrize('dtype', ['float32', 'float16'])
	92	+ def test_view_cat_sum_rewritten(self, shape, dtype):
	93	+ tensors = [self._generate_tensor(shape, dtype) for _ in range(3)]
	94	+ model = ViewCatSumModel()
	95	+ gm = self._run_pass(model, *tensors)
	96	+
	97	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 0, "cat should be eliminated")
	98	+ self.assertEqual(_count(gm.graph, torch.ops.aten.sum.dim_IntList), 0, "sum should be eliminated")
	99	+ self.assertEqual(_count(gm.graph, torch.ops.aten.add.Tensor), 2, "3 inputs should give 2 adds")
	100	+ self._assert_close(model(tensors), gm(tensors), dtype)
	101	+
	102	+ @parametrize('shape', [(8, 16)])
	103	+ @parametrize('dtype', ['float16'])
	104	+ def test_low_precision_accumulates_in_fp32(self, shape, dtype):
	105	+ tensors = [self._generate_tensor(shape, dtype) for _ in range(3)]
	106	+ gm = self._run_pass(ViewCatSumModel(), *tensors)
	107	+
	108	+ casts = [
	109	+ n for n in gm.graph.nodes
	110	+ if n.op == "call_function" and n.target == torch.ops.prims.convert_element_type.default
	111	+ ]
	112	+ self.assertEqual(len(casts), 4, "3 upcasts plus 1 downcast back to the original dtype")
	113	+ self.assertEqual(gm(*tensors).dtype, torch.float16, "output dtype should be unchanged")
	114	+
	115	+ @parametrize('shape', [(8, 16)])
	116	+ @parametrize('dtype', ['float32'])
	117	+ def test_unsqueeze_cat_sum_rewritten(self, shape, dtype):
	118	+ tensors = [self._generate_tensor(shape, dtype) for _ in range(3)]
	119	+ model = UnsqueezeCatSumModel()
	120	+ gm = self._run_pass(model, *tensors)
	121	+
	122	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 0, "cat should be eliminated")
	123	+ self.assertEqual(_count(gm.graph, torch.ops.aten.add.Tensor), 2)
	124	+ self._assert_close(model(tensors), gm(tensors), dtype)
	125	+
	126	+ @parametrize('shape', [(8, 16)])
	127	+ @parametrize('dtype', ['float32'])
	128	+ def test_keepdim_preserves_shape(self, shape, dtype):
	129	+ class M(torch.nn.Module):
	130	+ def forward(self, t1, t2):
	131	+ cat = torch.ops.aten.cat.default([t1, t2], 0)
	132	+ view = torch.ops.aten.reshape.default(cat, [2, t1.shape[0], t1.shape[1]])
	133	+ return torch.ops.aten.sum.dim_IntList(view, [0], True)
	134	+
	135	+ tensors = [self._generate_tensor(shape, dtype) for _ in range(2)]
	136	+ model = M()
	137	+ gm = self._run_pass(model, *tensors)
	138	+
	139	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 0)
	140	+ self._assert_close(model(tensors), gm(tensors), dtype)
	141	+
	142	+ @parametrize('shape', [(8, 16)])
	143	+ @parametrize('dtype', ['float32'])
	144	+ def test_slices_of_one_base(self, shape, dtype):
	145	+ """Real model form: column slices of a single wide table, stacked and summed."""
	146	+ class M(torch.nn.Module):
	147	+ def forward(self, wide):
	148	+ slices = [
	149	+ torch.ops.aten.slice.Tensor(wide, 1, off, off + 4)
	150	+ for off in (0, 8, 20, 36)
	151	+ ]
	152	+ cat = torch.ops.aten.cat.default(slices, 0)
	153	+ view = torch.ops.aten.reshape.default(cat, [4, wide.shape[0], 4])
	154	+ return torch.ops.aten.sum.dim_IntList(view, [0])
	155	+
	156	+ wide = self._generate_tensor((8, 64), dtype)
	157	+ model = M()
	158	+ gm = self._run_pass(model, wide)
	159	+
	160	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 0, "cat should be eliminated")
	161	+ self.assertEqual(_count(gm.graph, torch.ops.aten.slice.Tensor), 4, "slices should survive as views")
	162	+ self._assert_close(model(wide), gm(wide), dtype)
	163	+
	164	+ # ------------------------------------------------------------------
	165	+ # Real model scenario: symbolic batch, fp16, real offsets, traced by make_fx so nodes carry
	166	+ # meta['val']. Covers the propagate_fake_tensor path that ShapeProp never reaches.
	167	+ # ------------------------------------------------------------------
	168	+ def _wide_tensor(self, rows, dtype, scale=0.1):
	169	+ return torch.randn(
	170	+ (rows, _WIDE_COLS), dtype=eval('torch.' + dtype), device=torch.device("npu")
	171	+ ) * scale
	172	+
	173	+ @staticmethod
	174	+ def _fp32_golden(wide, offsets, width):
	175	+ return _fm_pool_ref(wide.float(), offsets, width)
	176	+
	177	+ def _assert_not_worse_than(self, golden, baseline, candidate, note):
	178	+ """Candidate error against the fp32 golden must not exceed the baseline, plus one fp16 ulp.
	179	+
	180	+ Both paths accumulate in fp32 and round once on write-back, so results are normally
	181	+ bitwise-identical; the extra ulp tolerates last-bit differences from accumulation order.
	182	+ """
	183	+ err_base = (baseline.float() - golden).abs().max().item()
	184	+ err_cand = (candidate.float() - golden).abs().max().item()
	185	+ one_ulp = golden.abs().max().item() * _FP16_EPS
	186	+ self.assertLessEqual(
	187	+ err_cand,
	188	+ err_base + one_ulp,
	189	+ f"{note}: rewritten error {err_cand:.3e} exceeds the original {err_base:.3e} plus one ulp {one_ulp:.3e}",
	190	+ )
	191	+ self.assertEqual(baseline, candidate, atol=one_ulp, rtol=_FP16_EPS)
	192	+
	193	+ @parametrize('rows', [64])
	194	+ @parametrize('dtype', ['float16'])
	195	+ def test_accuracy_28way_not_worse_than_unrewritten(self, rows, dtype):
	196	+ """Graph-level A/B against the fp32 result: both runs execute the same operators, so any
	197	+ error difference can only come from this pass. Comparing against eager cannot isolate it.
	198	+ """
	199	+ wide = self._wide_tensor(rows, dtype, scale=1.0)
	200	+
	201	+ def fn(x):
	202	+ return _fm_pool(x, _GROUP_W8_OFFSETS, 8)
	203	+
	204	+ baseline_gm = make_fx(fn, tracing_mode="symbolic")(wide)
	205	+ rewritten_gm = make_fx(fn, tracing_mode="symbolic")(wide)
	206	+ stack_sum_to_add_chain_pass(rewritten_gm.graph)
	207	+ rewritten_gm.recompile()
	208	+
	209	+ self.assertEqual(_count(baseline_gm.graph, torch.ops.aten.cat.default), 1, "baseline graph should keep the cat")
	210	+ self.assertEqual(_count(rewritten_gm.graph, torch.ops.aten.cat.default), 0, "rewritten graph should eliminate the cat")
	211	+
	212	+ golden = self._fp32_golden(wide, _GROUP_W8_OFFSETS, 8)
	213	+ self._assert_not_worse_than(golden, baseline_gm(wide), rewritten_gm(wide), "28-way stack pooling")
	214	+
	215	+ def _make_fx_symbolic(self, fn, *tensors):
	216	+ gm = make_fx(fn, tracing_mode="symbolic")(*tensors)
	217	+ stack_sum_to_add_chain_pass(gm.graph)
	218	+ gm.recompile()
	219	+ return gm
	220	+
	221	+ @parametrize('rows', [32])
	222	+ @parametrize('dtype', ['float16'])
	223	+ def test_real_model_fm_pooling(self, rows, dtype):
	224	+ """Prototype of output code L20616: 15-way width=128 column slice stack sum."""
	225	+ wide = self._wide_tensor(rows, dtype)
	226	+
	227	+ def fn(x):
	228	+ return _fm_pool(x, _GROUP_W128_OFFSETS, 128)
	229	+
	230	+ gm = self._make_fx_symbolic(fn, wide)
	231	+
	232	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 0, "cat should be eliminated")
	233	+ self.assertEqual(_count(gm.graph, torch.ops.aten.sum.dim_IntList), 0, "sum should be eliminated")
	234	+ self.assertEqual(_count(gm.graph, torch.ops.aten.add.Tensor), 14, "15 inputs should give 14 adds")
	235	+ self._assert_close(_fm_pool_ref(wide, _GROUP_W128_OFFSETS, 128), gm(wide), dtype)
	236	+
	237	+ @parametrize('rows', [32])
	238	+ @parametrize('dtype', ['float16'])
	239	+ def test_real_model_mul_of_two_pools(self, rows, dtype):
	240	+ """Prototype of triton_per_fused_mul_stack_sum_16: two pooling results multiplied."""
	241	+ wide = self._wide_tensor(rows, dtype)
	242	+
	243	+ def fn(x):
	244	+ user = _fm_pool(x, _USER_W128_OFFSETS, 128)
	245	+ group = _fm_pool(x, _GROUP_W128_OFFSETS, 128)
	246	+ return torch.ops.aten.mul.Tensor(user, group)
	247	+
	248	+ gm = self._make_fx_symbolic(fn, wide)
	249	+
	250	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 0, "both cats should be eliminated")
	251	+ self.assertEqual(_count(gm.graph, torch.ops.aten.add.Tensor), 15, "(2-1) + (15-1)")
	252	+ expected = (
	253	+ _fm_pool_ref(wide, _USER_W128_OFFSETS, 128)
	254	+ * _fm_pool_ref(wide, _GROUP_W128_OFFSETS, 128)
	255	+ )
	256	+ self._assert_close(expected, gm(wide), dtype)
	257	+
	258	+ @parametrize('rows', [32])
	259	+ @parametrize('dtype', ['float16'])
	260	+ def test_rewritten_nodes_carry_fake_meta(self, rows, dtype):
	261	+ """Inserted nodes must carry meta['val'], otherwise Inductor lowering gets no shape."""
	262	+ wide = self._wide_tensor(rows, dtype)
	263	+
	264	+ def fn(x):
	265	+ return _fm_pool(x, _GROUP_W16_OFFSETS, 16)
	266	+
	267	+ gm = self._make_fx_symbolic(fn, wide)
	268	+
	269	+ inserted = [
	270	+ n for n in gm.graph.nodes
	271	+ if n.op == "call_function"
	272	+ and n.target in (torch.ops.aten.add.Tensor,
	273	+ torch.ops.prims.convert_element_type.default)
	274	+ ]
	275	+ self.assertTrue(inserted, "add chain and cast nodes should be inserted")
	276	+ for node in inserted:
	277	+ self.assertIn('val', node.meta, f"{node.name} is missing meta['val']")
	278	+ output_node = [n for n in gm.graph.nodes if n.op == "output"][0]
	279	+ result = output_node.args[0]
	280	+ while isinstance(result, (list, tuple)):
	281	+ result = result[0]
	282	+ self.assertEqual(result.meta['val'].dtype, torch.float16, "output dtype should stay fp16")
	283	+
	284	+ @parametrize('rows', [32])
	285	+ @parametrize('dtype', ['float16'])
	286	+ def test_symbolic_batch_is_matched(self, rows, dtype):
	287	+ """Shape comparison must hold for a symbolic batch, or the pass never fires when dynamic."""
	288	+ wide = self._wide_tensor(rows, dtype)
	289	+
	290	+ def fn(x):
	291	+ return _fm_pool(x, _GROUP_W16_OFFSETS, 16)
	292	+
	293	+ traced = make_fx(fn, tracing_mode="symbolic")(wide)
	294	+ # SymInt is neither hashable nor directly comparable, so only check per-dim types.
	295	+ has_symbolic_dim = any(
	296	+ isinstance(dim, torch.SymInt)
	297	+ for node in traced.graph.nodes
	298	+ if 'val' in node.meta and hasattr(node.meta['val'], 'shape')
	299	+ for dim in node.meta['val'].shape
	300	+ )
	301	+ self.assertTrue(
	302	+ has_symbolic_dim,
	303	+ "trace must carry symbolic shapes, otherwise this case does not exercise dynamic shape",
	304	+ )
	305	+
	306	+ stack_sum_to_add_chain_pass(traced.graph)
	307	+ traced.recompile()
	308	+ self.assertEqual(_count(traced.graph, torch.ops.aten.cat.default), 0)
	309	+
	310	+ other = self._wide_tensor(rows * 2, dtype)
	311	+ self._assert_close(_fm_pool_ref(other, _GROUP_W16_OFFSETS, 16), traced(other), dtype)
	312	+
	313	+ @parametrize('shape', [(8, 16)])
	314	+ @parametrize('dtype', ['float32'])
	315	+ def test_multi_user_cat_not_rewritten(self, shape, dtype):
	316	+ class M(torch.nn.Module):
	317	+ def forward(self, t1, t2):
	318	+ cat = torch.ops.aten.cat.default([t1, t2], 0)
	319	+ view = torch.ops.aten.reshape.default(cat, [2, t1.shape[0], t1.shape[1]])
	320	+ return torch.ops.aten.sum.dim_IntList(view, [0]), cat
	321	+
	322	+ tensors = [self._generate_tensor(shape, dtype) for _ in range(2)]
	323	+ gm = self._run_pass(M(), *tensors)
	324	+
	325	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 1, "cat has another user, so it must not be rewritten")
	326	+ self.assertEqual(_count(gm.graph, torch.ops.aten.add.Tensor), 0)
	327	+
	328	+ @parametrize('shape', [(8, 16)])
	329	+ @parametrize('dtype', ['float32'])
	330	+ def test_sum_on_other_dim_not_rewritten(self, shape, dtype):
	331	+ class M(torch.nn.Module):
	332	+ def forward(self, t1, t2):
	333	+ cat = torch.ops.aten.cat.default([t1, t2], 0)
	334	+ view = torch.ops.aten.reshape.default(cat, [2, t1.shape[0], t1.shape[1]])
	335	+ return torch.ops.aten.sum.dim_IntList(view, [1])
	336	+
	337	+ tensors = [self._generate_tensor(shape, dtype) for _ in range(2)]
	338	+ gm = self._run_pass(M(), *tensors)
	339	+
	340	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 1, "a sum over a non-stack axis must not be rewritten")
	341	+
	342	+ @parametrize('shape', [(8, 16)])
	343	+ @parametrize('dtype', ['int32'])
	344	+ def test_integer_dtype_not_rewritten(self, shape, dtype):
	345	+ """Integer sum promotes to int64, which an add chain does not, so leave it untouched."""
	346	+ tensors = [self._generate_tensor(shape, dtype) for _ in range(3)]
	347	+ gm = self._run_pass(ViewCatSumModel(), *tensors)
	348	+
	349	+ self.assertEqual(_count(gm.graph, torch.ops.aten.cat.default), 1)
	350	+
	351	+ @parametrize('shape', [(8, 16)])
	352	+ @parametrize('dtype', ['float32'])
	353	+ def test_compile_cases(self, shape, dtype):
	354	+ def op_calc(t1, t2, t3):
	355	+ return torch.stack([t1, t2, t3]).sum(0)
	356	+
	357	+ tensors = [self._generate_tensor(shape, dtype) for _ in range(3)]
	358	+ std_result = op_calc(*tensors)
	359	+ with torch.no_grad():
	360	+ compiled = torch.compile(op_calc, backend="inductor")
	361	+ self._assert_close(std_result, compiled(*tensors), dtype)
	362	+
	363	+ @parametrize('rows', [32])
	364	+ @parametrize('dtype', ['float16'])
	365	+ def test_compile_real_pattern_dynamic_batch(self, rows, dtype):
	366	+ def op_calc(x):
	367	+ return _fm_pool_ref(x, _GROUP_W16_OFFSETS, 16)
	368	+
	369	+ wide = self._wide_tensor(rows, dtype)
	370	+ with torch.no_grad():
	371	+ compiled = torch.compile(op_calc, backend="inductor", dynamic=True)
	372	+ self._assert_close(op_calc(wide), compiled(wide), dtype)
	373	+ other = self._wide_tensor(rows * 2, dtype)
	374	+ self._assert_close(op_calc(other), compiled(other), dtype)
	375	+
	376	+ @parametrize('rows', [64])
	377	+ @parametrize('dtype', ['float16'])
	378	+ def test_compile_accuracy_against_fp32_golden(self, rows, dtype):
	379	+ """Eager fp16 error is the yardstick, so the bound holds across data magnitudes."""
	380	+ def op_calc(x):
	381	+ return _fm_pool_ref(x, _GROUP_W8_OFFSETS, 8)
	382	+
	383	+ wide = self._wide_tensor(rows, dtype, scale=1.0)
	384	+ golden = self._fp32_golden(wide, _GROUP_W8_OFFSETS, 8)
	385	+ with torch.no_grad():
	386	+ compiled = torch.compile(op_calc, backend="inductor")
	387	+ self._assert_not_worse_than(golden, op_calc(wide), compiled(wide), "end-to-end 28-way pooling")
	388	+
	389	+
	390	+instantiate_parametrized_tests(TestStackSumPass)
	391	+
	392	+
	393	+if __name__ == "__main__":
	394	+ run_tests()

Mtorch_npu/_inductor/__init__.py+3-0

	@@ -136,6 +136,7 @@ def _load_triton_backend():
136	_register_npu_inductor_flex_attention,	136	_register_npu_inductor_flex_attention,
137	_register_npu_inductor_grouped_mm,	137	_register_npu_inductor_grouped_mm,
138	_register_npu_inductor_mm,	138	_register_npu_inductor_mm,
		139	+ _register_npu_inductor_multi_slice_concat,
139	_validate_device,	140	_validate_device,
140	patch_flex_attention,	141	patch_flex_attention,
141	)	142	)
	@@ -216,6 +217,8 @@ def _load_triton_backend():
216	_register_npu_inductor_addmm()	217	_register_npu_inductor_addmm()
217	_register_npu_inductor_bmm()	218	_register_npu_inductor_bmm()
218	_register_npu_inductor_grouped_mm()	219	_register_npu_inductor_grouped_mm()
		220	+ if npu_config.enable_multi_slice_concat:
		221	+ _register_npu_inductor_multi_slice_concat()
219		222
220	_register_npu_inductor_flex_attention()	223	_register_npu_inductor_flex_attention()
221		224

Mtorch_npu/_inductor/config.py+16-0

	@@ -430,6 +430,22 @@ enable_fused_matmul_relu = _parse_bool_env(
430	"TORCHINDUCTOR_ENABLE_FUSED_MATMUL_RELU", False	430	"TORCHINDUCTOR_ENABLE_FUSED_MATMUL_RELU", False
431	)	431	)
432		432
		433	+# multi_slice_concat_pass: rewrite a run of constant-offset column slices feeding
		434	+# one cat into npu_ext::multi_slice_concat, so a single kernel replaces the N Slice
		435	+# copies aclnnCat needs. Gated off until validated on hardware; the rewrite is pure
		436	+# data movement and cannot affect numerics.
		437	+enable_multi_slice_concat = _parse_bool_env(
		438	+ "TORCHINDUCTOR_ENABLE_MULTI_SLICE_CONCAT", False
		439	+)
		440	+
		441	+# grouped_matmul_fusion_pass: merge the independent small GEMMs feeding one cat into
		442	+# npu_grouped_matmul. Gated off until validated on the target model; the rewrite keeps
		443	+# every GEMM's operands intact but the kernel accumulates differently, so results are
		444	+# close but not bit-exact.
		445	+enable_grouped_matmul_fusion = _parse_bool_env(
		446	+ "TORCHINDUCTOR_ENABLE_GROUPED_MATMUL_FUSION", False
		447	+)
		448	+
433	# permute_continous_reduction: when enabled, detects the "permute contiguous reduction"	449	# permute_continous_reduction: when enabled, detects the "permute contiguous reduction"
434	# pattern (a non-reduction axis sitting between two reduction axes in stride order)	450	# pattern (a non-reduction axis sitting between two reduction axes in stride order)
435	# and applies special handling: selects the permute axis as a tiling axis, uses NDDMA	451	# and applies special handling: selects the permute axis as a tiling axis, uses NDDMA

Mtorch_npu/_inductor/fx_passes/ascend_custom_passes/__init__.py+1-1

	@@ -32,7 +32,7 @@ def run_register_post_custom_passes(gm):
32	stable_topological_sort(gm)	32	stable_topological_sort(gm)
33	for level in sorted(FxPassLevel):	33	for level in sorted(FxPassLevel):
34	for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.POST][level]:	34	for fn in ASCEND_CUSTOME_PASS_REGISTER[PassType.POST][level]:
35	- # 标记了 ignore_inference_check 的 pass 在训练图下也执行。	35	+ # a pass marked ignore_inference_check also runs on training graphs.
36	if inference or getattr(fn, "ignore_inference_check", False):	36	if inference or getattr(fn, "ignore_inference_check", False):
37	fn(gm)	37	fn(gm)
38		38

Mtorch_npu/_inductor/fx_passes/ascend_custom_passes/ascend_graph_pass.py+1073-180

Mtorch_npu/_inductor/fx_passes/ascend_custom_passes/register_custom_pass.py+22-14

Mtorch_npu/_inductor/fx_passes/parallelism_strategy_cv.py+5-5

Mtorch_npu/_inductor/fx_passes/parallelism_strategy_default.py+10-10

Mtorch_npu/_inductor/fx_passes/pattern_match/npu_fusion_attention_graph.py+6-6

Mtorch_npu/_inductor/fx_passes/utils/check_mode.py+2-2

Mtorch_npu/_inductor/fx_passes/utils/check_op_util.py+25-3

Mtorch_npu/_inductor/kernel/__init__.py+1-0

Atorch_npu/_inductor/kernel/multi_slice_concat.py+267-0