feat(inductor): support dynamic shapes in ascend custom graph passes-pytorch-AtomGit

已合并

feat(inductor): support dynamic shapes in ascend custom graph passes #42443

dezheng889创建于 7月22日从crazyDannyBoy/pytorch_dezheng： dynamic-graph-dev 合入到 Ascend/pytorch： v2.10.0

feat(inductor): support dynamic shapes in ascend custom graph passes #42443

已合并

dezheng889创建于 7月22日

从crazyDannyBoy/pytorch_dezheng dynamic-graph-dev合入到Ascend/pytorch v2.10.0

共 4 个文件变更+1165-171

Atest/_inductor/test_dynamic_shape_fx_passes.py+471-0

@@ -0,0 +1,471 @@
	1	+# Owner(s): ["module: inductor"]
	2	+"""Device-independent unit tests for dynamic-shape graph optimization.
	3	+
	4	+These cases validate the symbolic-shape utilities and passes purely at the FX
	5	+level (build graph + run pass + assert transform); they run no NPU/CUDA kernel and
	6	+thus work on CPU-only machines. Coverage focuses on:
	7	+1. symbolic_shape_util: three-valued checks / normalization / materialization / ranges;
	8	+2. per-pass "optimize only when provable" behavior on dynamic shapes;
	9	+3. boundary safety: undecidable, rank mismatch, distinct symbols, switch-off, etc.
	10	+ must not mis-fold;
	11	+4. static-shape regression: existing optimizations still fire on static inputs.
	12	+"""
	13	+
	14	+import collections.abc
	15	+import importlib
	16	+import logging
	17	+import os
	18	+import sys
	19	+import types
	20	+import unittest
	21	+
	22	+
	23	+def _shim_missing_torch_internals():
	24	+ """Shim an internal util present in the target torch (2.10) but missing on older local torch (e.g. 2.2).
	25	+
	26	+ Injected only when genuinely absent; the real target environment is untouched."""
	27	+ try:
	28	+ importlib.import_module("torch.utils._ordered_set")
	29	+ return
	30	+ except Exception:
	31	+ pass
	32	+
	33	+ module = types.ModuleType("torch.utils._ordered_set")
	34	+
	35	+ class OrderedSet(collections.abc.MutableSet):
	36	+ def __init__(self, iterable=()):
	37	+ self._data = dict.fromkeys(iterable)
	38	+
	39	+ def __contains__(self, value):
	40	+ return value in self._data
	41	+
	42	+ def __iter__(self):
	43	+ return iter(self._data)
	44	+
	45	+ def __len__(self):
	46	+ return len(self._data)
	47	+
	48	+ def add(self, value):
	49	+ self._data[value] = None
	50	+
	51	+ def discard(self, value):
	52	+ self._data.pop(value, None)
	53	+
	54	+ def pop(self):
	55	+ key = next(iter(self._data))
	56	+ del self._data[key]
	57	+ return key
	58	+
	59	+ module.OrderedSet = OrderedSet
	60	+ sys.modules["torch.utils._ordered_set"] = module
	61	+
	62	+
	63	+_NPU_TEST_LIB = None
	64	+
	65	+
	66	+def _shim_missing_npu_ops():
	67	+ """Register a schema stub for the npu custom op referenced at import time when the torch_npu C++ ext is absent."""
	68	+ global _NPU_TEST_LIB
	69	+ try:
	70	+ torch.ops.npu._npu_dtype_cast.default
	71	+ return
	72	+ except Exception:
	73	+ pass
	74	+ _NPU_TEST_LIB = torch.library.Library("npu", "FRAGMENT")
	75	+ try:
	76	+ _NPU_TEST_LIB.define("_npu_dtype_cast(Tensor self, ScalarType dtype) -> Tensor")
	77	+ except Exception:
	78	+ pass
	79	+
	80	+
	81	+def _ensure_torch_npu_importable():
	82	+ """Stub the torch_npu package root on NPU-less machines so the pure graph-pass modules import.
	83	+
	84	+ On a real NPU box ``import torch_npu`` works and we return early; otherwise we
	85	+ only stub the necessary package nodes and ``torch_npu._inductor.config.log``,
	86	+ while other submodules still load from the real source files.
	87	+ """
	88	+ try:
	89	+ import torch_npu # noqa: F401
	90	+
	91	+ return
	92	+ except Exception:
	93	+ pass
	94	+
	95	+ repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
	96	+
	97	+ def _pkg(name, rel):
	98	+ module = types.ModuleType(name)
	99	+ module.__path__ = [os.path.join(repo_root, *rel)]
	100	+ sys.modules[name] = module
	101	+
	102	+ _pkg("torch_npu", ["torch_npu"])
	103	+ _pkg("torch_npu._inductor", ["torch_npu", "_inductor"])
	104	+ config_stub = types.ModuleType("torch_npu._inductor.config")
	105	+ config_stub.log = logging.getLogger("dynamic_fx_test")
	106	+ sys.modules["torch_npu._inductor.config"] = config_stub
	107	+ _pkg("torch_npu._inductor.fx_passes", ["torch_npu", "_inductor", "fx_passes"])
	108	+ _pkg(
	109	+ "torch_npu._inductor.fx_passes.utils",
	110	+ ["torch_npu", "_inductor", "fx_passes", "utils"],
	111	+ )
	112	+ _pkg(
	113	+ "torch_npu._inductor.fx_passes.ascend_custom_passes",
	114	+ ["torch_npu", "_inductor", "fx_passes", "ascend_custom_passes"],
	115	+ )
	116	+
	117	+
	118	+import torch
	119	+
	120	+_shim_missing_torch_internals()
	121	+_shim_missing_npu_ops()
	122	+_ensure_torch_npu_importable()
	123	+
	124	+from torch.fx.experimental.proxy_tensor import make_fx
	125	+
	126	+from torch_npu._inductor.fx_passes.ascend_custom_passes import ascend_graph_pass as agp
	127	+from torch_npu._inductor.fx_passes.utils import symbolic_shape_util as ssu
	128	+
	129	+
	130	+def _symbolic_gm(fn, *example_inputs):
	131	+ """Trace an aten-level fx graph in symbolic mode; nodes carry symbolic FakeTensor meta."""
	132	+ return make_fx(fn, tracing_mode="symbolic")(*example_inputs)
	133	+
	134	+
	135	+def _first_symint(gm):
	136	+ """Grab one symbolic dim (SymInt) from the graph's first placeholder, for util tests."""
	137	+ for node in gm.graph.nodes:
	138	+ if node.op == "placeholder":
	139	+ val = node.meta.get("val")
	140	+ if val is not None and hasattr(val, "shape"):
	141	+ for dim in val.shape:
	142	+ if isinstance(dim, torch.SymInt):
	143	+ return dim
	144	+ raise AssertionError("no symbolic dim found in traced graph")
	145	+
	146	+
	147	+def _count_target(gm, target):
	148	+ return sum(
	149	+ 1
	150	+ for n in gm.graph.nodes
	151	+ if n.op == "call_function" and n.target is target
	152	+ )
	153	+
	154	+
	155	+class TestSymbolicShapeUtil(unittest.TestCase):
	156	+ """Three-valued logic, normalization, materialization and ranges of symbolic_shape_util."""
	157	+
	158	+ def setUp(self):
	159	+ os.environ.pop("NPU_INDUCTOR_DYNAMIC_FX_PASS", None)
	160	+ ssu.reset_stats()
	161	+ # Build two independent symbols s0, s1 plus derived expressions
	162	+ self.gm = _symbolic_gm(lambda a, b: a + b.sum(), torch.randn(6, 5), torch.randn(7))
	163	+ self.s0 = None
	164	+ self.s1 = None
	165	+ for node in self.gm.graph.nodes:
	166	+ if node.op == "placeholder":
	167	+ val = node.meta.get("val")
	168	+ if val is None or not hasattr(val, "shape") or not val.shape:
	169	+ continue
	170	+ dim0 = val.shape[0]
	171	+ if isinstance(dim0, torch.SymInt):
	172	+ if self.s0 is None:
	173	+ self.s0 = dim0
	174	+ elif self.s1 is None:
	175	+ self.s1 = dim0
	176	+ self.assertIsNotNone(self.s0)
	177	+ self.assertIsNotNone(self.s1)
	178	+
	179	+ # ---- three-valued checks -----------------------------------------
	180	+ def test_statically_known_eq_int(self):
	181	+ self.assertTrue(ssu.statically_known_eq(3, 3))
	182	+ self.assertFalse(ssu.statically_known_eq(3, 4))
	183	+
	184	+ def test_statically_known_eq_same_symbol(self):
	185	+ self.assertTrue(ssu.statically_known_eq(self.s0, self.s0))
	186	+ self.assertTrue(ssu.statically_known_eq(self.s0 + 1, self.s0 + 1))
	187	+
	188	+ def test_statically_known_eq_distinct_symbols_unprovable(self):
	189	+ # s0 == s1 is undecidable -> False (neither mis-proven true nor adds a guard)
	190	+ self.assertFalse(ssu.statically_known_eq(self.s0, self.s1))
	191	+
	192	+ def test_statically_known_eq_symbol_vs_const_unprovable(self):
	193	+ self.assertFalse(ssu.statically_known_eq(self.s0, 5))
	194	+
	195	+ def test_statically_known_geq_gt_leq(self):
	196	+ self.assertTrue(ssu.statically_known_geq(self.s0 + 1, self.s0))
	197	+ self.assertTrue(ssu.statically_known_gt(self.s0 + 1, self.s0))
	198	+ self.assertTrue(ssu.statically_known_leq(self.s0, self.s0 + 1))
	199	+ self.assertFalse(ssu.statically_known_gt(self.s0, self.s0 + 1))
	200	+
	201	+ def test_is_statically_one(self):
	202	+ self.assertTrue(ssu.is_statically_one(1))
	203	+ self.assertFalse(ssu.is_statically_one(2))
	204	+ self.assertFalse(ssu.is_statically_one(self.s0))
	205	+
	206	+ def test_shapes_statically_equal(self):
	207	+ self.assertTrue(ssu.shapes_statically_equal([self.s0, 3], [self.s0, 3]))
	208	+ self.assertFalse(ssu.shapes_statically_equal([self.s0, 3], [self.s1, 3]))
	209	+ self.assertFalse(ssu.shapes_statically_equal([self.s0], [self.s0, 1]))
	210	+ self.assertFalse(ssu.shapes_statically_equal(None, [1]))
	211	+
	212	+ def test_has_free_symbols(self):
	213	+ self.assertTrue(ssu.has_free_symbols([self.s0, 2]))
	214	+ self.assertFalse(ssu.has_free_symbols([2, 3]))
	215	+ self.assertFalse(ssu.has_free_symbols(None))
	216	+
	217	+ # ---- normalization -----------------------------------------------
	218	+ def test_resolve_size_arg_scalar(self):
	219	+ self.assertEqual(ssu.resolve_size_arg(4), 4)
	220	+ self.assertIs(ssu.resolve_size_arg(self.s0), self.s0)
	221	+ self.assertIsNone(ssu.resolve_size_arg(True))
	222	+ self.assertIsNone(ssu.resolve_size_arg(1.5))
	223	+
	224	+ def test_resolve_size_arg_node(self):
	225	+ # sym_size node: meta['val'] is a SymInt
	226	+ sym_size_node = None
	227	+ for n in self.gm.graph.nodes:
	228	+ if n.op == "call_function" and isinstance(n.meta.get("val"), torch.SymInt):
	229	+ sym_size_node = n
	230	+ break
	231	+ if sym_size_node is not None:
	232	+ self.assertIsInstance(ssu.resolve_size_arg(sym_size_node), torch.SymInt)
	233	+
	234	+ def test_resolve_size_list(self):
	235	+ self.assertEqual(ssu.resolve_size_list([1, 2, 3]), [1, 2, 3])
	236	+ self.assertIsNone(ssu.resolve_size_list([1, object()]))
	237	+ self.assertIsNone(ssu.resolve_size_list(5))
	238	+
	239	+ # ---- value ranges ------------------------------------------------
	240	+ def test_statically_fits_int32_int(self):
	241	+ self.assertTrue(ssu.statically_fits_int32(0, 100, -100))
	242	+ self.assertFalse(ssu.statically_fits_int32(2**31))
	243	+ self.assertFalse(ssu.statically_fits_int32())
	244	+
	245	+ def test_statically_fits_int32_unbounded_symbol(self):
	246	+ # An unbounded symbol cannot be proven to fit int32
	247	+ self.assertFalse(ssu.statically_fits_int32(self.s0))
	248	+
	249	+ # ---- materialization ---------------------------------------------
	250	+ def test_materialize_shape_from_anchor(self):
	251	+ gm = _symbolic_gm(lambda x: x + 1, torch.randn(6, 4))
	252	+ ph = [n for n in gm.graph.nodes if n.op == "placeholder"][0]
	253	+ s0 = ph.meta["val"].shape[0]
	254	+ anchor_user = [
	255	+ n for n in gm.graph.nodes if n.target is torch.ops.aten.add.Tensor
	256	+ ][0]
	257	+ with gm.graph.inserting_before(anchor_user):
	258	+ mat = ssu.materialize_shape(gm.graph, [s0, 8], ph)
	259	+ self.assertIsNotNone(mat)
	260	+ self.assertEqual(mat[1], 8)
	261	+ self.assertIsInstance(mat[0], torch.fx.Node)
	262	+ self.assertIs(mat[0].target, torch.ops.aten.sym_size.int)
	263	+
	264	+ def test_materialize_shape_unresolvable_returns_none(self):
	265	+ gm = _symbolic_gm(lambda x: x + 1, torch.randn(6, 4))
	266	+ ph = [n for n in gm.graph.nodes if n.op == "placeholder"][0]
	267	+ s0 = ph.meta["val"].shape[0]
	268	+ anchor_user = [
	269	+ n for n in gm.graph.nodes if n.target is torch.ops.aten.add.Tensor
	270	+ ][0]
	271	+ # s0*s0 is not any anchor dim length -> cannot materialize
	272	+ with gm.graph.inserting_before(anchor_user):
	273	+ mat = ssu.materialize_shape(gm.graph, [s0 * s0], ph)
	274	+ self.assertIsNone(mat)
	275	+
	276	+ # ---- switch -------------------------------------------------------
	277	+ def test_switch_off_degrades_to_static(self):
	278	+ os.environ["NPU_INDUCTOR_DYNAMIC_FX_PASS"] = "0"
	279	+ try:
	280	+ # When off, symbolic comparisons are all undecidable, but pure ints still decide
	281	+ self.assertFalse(ssu.statically_known_eq(self.s0, self.s0))
	282	+ self.assertTrue(ssu.statically_known_eq(3, 3))
	283	+ self.assertIsNone(ssu.resolve_size_arg(self.s0))
	284	+ self.assertEqual(ssu.resolve_size_arg(4), 4)
	285	+ finally:
	286	+ os.environ.pop("NPU_INDUCTOR_DYNAMIC_FX_PASS", None)
	287	+
	288	+ def test_materialize_shape_switch_off_symbolic_returns_none(self):
	289	+ # Switch off is a hard kill-switch: symbolic dims cannot be materialized,
	290	+ # but a fully-static shape still passes through unchanged.
	291	+ gm = _symbolic_gm(lambda x: x + 1, torch.randn(6, 4))
	292	+ ph = [n for n in gm.graph.nodes if n.op == "placeholder"][0]
	293	+ s0 = ph.meta["val"].shape[0]
	294	+ anchor_user = [
	295	+ n for n in gm.graph.nodes if n.target is torch.ops.aten.add.Tensor
	296	+ ][0]
	297	+ os.environ["NPU_INDUCTOR_DYNAMIC_FX_PASS"] = "0"
	298	+ try:
	299	+ with gm.graph.inserting_before(anchor_user):
	300	+ self.assertIsNone(ssu.materialize_shape(gm.graph, [s0, 8], ph))
	301	+ self.assertEqual(
	302	+ ssu.materialize_shape(gm.graph, [2, 8], ph), [2, 8]
	303	+ )
	304	+ finally:
	305	+ os.environ.pop("NPU_INDUCTOR_DYNAMIC_FX_PASS", None)
	306	+
	307	+
	308	+class TestDynamicShapePasses(unittest.TestCase):
	309	+ """Per-pass behavior and boundary safety on dynamic shapes."""
	310	+
	311	+ def setUp(self):
	312	+ os.environ.pop("NPU_INDUCTOR_DYNAMIC_FX_PASS", None)
	313	+
	314	+ def test_fold_expand_identity_symbolic(self):
	315	+ def fn(x):
	316	+ return torch.ops.aten.expand.default(x, [x.size(0), x.size(1)]).relu()
	317	+
	318	+ gm = _symbolic_gm(fn, torch.randn(6, 4))
	319	+ before = _count_target(gm, torch.ops.aten.expand.default)
	320	+ self.assertGreaterEqual(before, 1)
	321	+ agp.fold_expand(gm.graph)
	322	+ self.assertEqual(_count_target(gm, torch.ops.aten.expand.default), 0)
	323	+
	324	+ def test_view_fold_identity_symbolic(self):
	325	+ def fn(x):
	326	+ return torch.ops.aten.view.default(x, [x.size(0), x.size(1)]).relu()
	327	+
	328	+ gm = _symbolic_gm(fn, torch.randn(6, 4))
	329	+ agp.view_fold_pass(gm.graph)
	330	+ self.assertEqual(_count_target(gm, torch.ops.aten.view.default), 0)
	331	+
	332	+ def test_fold_reduce_static_one_dim(self):
	333	+ # Middle dim is statically 1 (0/1 specialization); symbolic batch dim is kept
	334	+ def fn(x):
	335	+ return torch.ops.aten.sum.dim_IntList(x, [1], True)
	336	+
	337	+ gm = _symbolic_gm(fn, torch.randn(6, 1, 4))
	338	+ before = _count_target(gm, torch.ops.aten.sum.dim_IntList)
	339	+ self.assertEqual(before, 1)
	340	+ agp.fold_reduce(gm.graph)
	341	+ self.assertEqual(_count_target(gm, torch.ops.aten.sum.dim_IntList), 0)
	342	+
	343	+ def test_fold_reduce_symbolic_dim_not_folded(self):
	344	+ # Reducing a symbolic dim that cannot be proven 1 -> must be kept
	345	+ def fn(x):
	346	+ return torch.ops.aten.sum.dim_IntList(x, [0], True)
	347	+
	348	+ gm = _symbolic_gm(fn, torch.randn(6, 4))
	349	+ agp.fold_reduce(gm.graph)
	350	+ self.assertEqual(_count_target(gm, torch.ops.aten.sum.dim_IntList), 1)
	351	+
	352	+ def test_fold_slice_full_symbolic(self):
	353	+ def fn(x):
	354	+ return torch.ops.aten.slice.Tensor(x, 0, 0, x.size(0)).relu()
	355	+
	356	+ gm = _symbolic_gm(fn, torch.randn(6, 4))
	357	+ agp.fold_slice(gm.graph)
	358	+ self.assertEqual(_count_target(gm, torch.ops.aten.slice.Tensor), 0)
	359	+
	360	+ def test_fold_slice_partial_symbolic_not_folded(self):
	361	+ # Slicing to s0-1 cannot be proven to cover the full dim -> kept
	362	+ def fn(x):
	363	+ return torch.ops.aten.slice.Tensor(x, 0, 0, x.size(0) - 1).relu()
	364	+
	365	+ gm = _symbolic_gm(fn, torch.randn(6, 4))
	366	+ agp.fold_slice(gm.graph)
	367	+ self.assertEqual(_count_target(gm, torch.ops.aten.slice.Tensor), 1)
	368	+
	369	+ def test_repeat_to_expand_symbolic(self):
	370	+ # x:[s0,1] repeat(1,3) -> pure broadcast, can become expand (mul consumer is broadcast-friendly)
	371	+ def fn(x):
	372	+ r = torch.ops.aten.repeat.default(x, [1, 3])
	373	+ return torch.ops.aten.mul.Tensor(r, r)
	374	+
	375	+ gm = _symbolic_gm(fn, torch.randn(6, 1))
	376	+ agp.repeat_to_expand_pass(gm.graph)
	377	+ self.assertEqual(_count_target(gm, torch.ops.aten.repeat.default), 0)
	378	+ self.assertGreaterEqual(_count_target(gm, torch.ops.aten.expand.default), 1)
	379	+
	380	+ def test_repeat_physical_copy_symbolic_not_rewritten(self):
	381	+ # x:[s0,4] repeat(1,3): 2nd dim is neither 1 nor kept -> needs a physical copy, keep repeat
	382	+ def fn(x):
	383	+ r = torch.ops.aten.repeat.default(x, [1, 3])
	384	+ return torch.ops.aten.mul.Tensor(r, r)
	385	+
	386	+ gm = _symbolic_gm(fn, torch.randn(6, 4))
	387	+ agp.repeat_to_expand_pass(gm.graph)
	388	+ self.assertEqual(_count_target(gm, torch.ops.aten.repeat.default), 1)
	389	+
	390	+ def test_fold_four_op_add_zeros_symbolic(self):
	391	+ def fn(x):
	392	+ return x + torch.zeros_like(x)
	393	+
	394	+ gm = _symbolic_gm(fn, torch.randn(6, 4))
	395	+ before = _count_target(gm, torch.ops.aten.add.Tensor)
	396	+ self.assertGreaterEqual(before, 1)
	397	+ agp.fold_four_op_pass(gm.graph)
	398	+ self.assertEqual(_count_target(gm, torch.ops.aten.add.Tensor), 0)
	399	+
	400	+ def test_cat_to_view_identity_symbolic(self):
	401	+ # cat([x[:, 0:2], x[:, 2:s1]], dim=1) covers all of dim1 -> identity view
	402	+ def fn(x):
	403	+ a = torch.ops.aten.slice.Tensor(x, 1, 0, 2)
	404	+ b = torch.ops.aten.slice.Tensor(x, 1, 2, x.size(1))
	405	+ return torch.ops.aten.cat.default([a, b], 1)
	406	+
	407	+ gm = _symbolic_gm(fn, torch.randn(6, 5))
	408	+ self.assertEqual(_count_target(gm, torch.ops.aten.cat.default), 1)
	409	+ agp.cat_to_view_pass(gm.graph)
	410	+ self.assertEqual(_count_target(gm, torch.ops.aten.cat.default), 0)
	411	+
	412	+ def test_cat_to_view_partial_not_folded(self):
	413	+ # Slices do not cover the whole dim (missing tail) -> must not fold
	414	+ def fn(x):
	415	+ a = torch.ops.aten.slice.Tensor(x, 1, 0, 2)
	416	+ b = torch.ops.aten.slice.Tensor(x, 1, 2, 4)
	417	+ return torch.ops.aten.cat.default([a, b], 1)
	418	+
	419	+ gm = _symbolic_gm(fn, torch.randn(6, 5))
	420	+ agp.cat_to_view_pass(gm.graph)
	421	+ self.assertEqual(_count_target(gm, torch.ops.aten.cat.default), 1)
	422	+
	423	+
	424	+class TestSwitchAndStaticRegression(unittest.TestCase):
	425	+ """Switch-off degrades to static behavior; existing optimizations still fire on static shapes."""
	426	+
	427	+ def tearDown(self):
	428	+ os.environ.pop("NPU_INDUCTOR_DYNAMIC_FX_PASS", None)
	429	+
	430	+ def test_switch_off_symbolic_slice_not_folded(self):
	431	+ os.environ["NPU_INDUCTOR_DYNAMIC_FX_PASS"] = "0"
	432	+
	433	+ def fn(x):
	434	+ return torch.ops.aten.slice.Tensor(x, 0, 0, x.size(0)).relu()
	435	+
	436	+ gm = _symbolic_gm(fn, torch.randn(6, 4))
	437	+ agp.fold_slice(gm.graph)
	438	+ # Switch off: symbolic full-slice is undecidable -> kept
	439	+ self.assertEqual(_count_target(gm, torch.ops.aten.slice.Tensor), 1)
	440	+
	441	+ def test_switch_off_symbolic_repeat_not_rewritten(self):
	442	+ os.environ["NPU_INDUCTOR_DYNAMIC_FX_PASS"] = "0"
	443	+
	444	+ def fn(x):
	445	+ r = torch.ops.aten.repeat.default(x, [1, 3])
	446	+ return torch.ops.aten.mul.Tensor(r, r)
	447	+
	448	+ gm = _symbolic_gm(fn, torch.randn(6, 1))
	449	+ agp.repeat_to_expand_pass(gm.graph)
	450	+ # Switch off: symbolic broadcast is not rewritten (materialize refuses sym dims)
	451	+ self.assertEqual(_count_target(gm, torch.ops.aten.repeat.default), 1)
	452	+
	453	+ def test_static_expand_still_folded(self):
	454	+ def fn(x):
	455	+ return torch.ops.aten.expand.default(x, [4, 5]).relu()
	456	+
	457	+ gm = make_fx(fn, tracing_mode="fake")(torch.randn(4, 5))
	458	+ agp.fold_expand(gm.graph)
	459	+ self.assertEqual(_count_target(gm, torch.ops.aten.expand.default), 0)
	460	+
	461	+ def test_static_slice_still_folded(self):
	462	+ def fn(x):
	463	+ return torch.ops.aten.slice.Tensor(x, 0, 0, 4).relu()
	464	+
	465	+ gm = make_fx(fn, tracing_mode="fake")(torch.randn(4, 5))
	466	+ agp.fold_slice(gm.graph)
	467	+ self.assertEqual(_count_target(gm, torch.ops.aten.slice.Tensor), 0)
	468	+
	469	+
	470	+if __name__ == "__main__":
	471	+ unittest.main(verbosity=2)

Mtorch_npu/_inductor/fx_passes/ascend_custom_passes/ascend_graph_pass.py+189-136

	@@ -1,4 +1,3 @@
1	-import math
2	import operator	1	import operator
3		2
4	import torch	3	import torch
	@@ -41,6 +40,18 @@ from ..utils.get_binary_fold_result import (
41	has_storage_or_layout,	40	has_storage_or_layout,
42	propagate_fake_tensor,	41	propagate_fake_tensor,
43	)	42	)
		43	+from ..utils.symbolic_shape_util import (
		44	+ is_statically_one,
		45	+ materialize_shape,
		46	+ refresh_fake_meta,
		47	+ resolve_size_arg,
		48	+ resolve_size_list,
		49	+ shapes_statically_equal,
		50	+ statically_fits_int32,
		51	+ statically_known_eq,
		52	+ statically_known_geq,
		53	+ statically_known_leq,
		54	+)
44	from .register_custom_pass import register_custom_pass	55	from .register_custom_pass import register_custom_pass
45		56
46		57
	@@ -60,7 +71,7 @@ def cat_slice_cat_fold_pass(graph: torch.fx.Graph) -> None:
60	cat2_node = node	71	cat2_node = node
61	cat2_inputs = cat2_node.args[0]	72	cat2_inputs = cat2_node.args[0]
62	cat2_dim = cat2_node.kwargs.get("dim", -1)	73	cat2_dim = cat2_node.kwargs.get("dim", -1)
63	- cat2_shape = get_node_shape(cat2_node)	74	+ cat2_shape = get_node_shape(cat2_node, allow_symbolic=True)
64	if not cat2_shape:	75	if not cat2_shape:
65	continue	76	continue
66	cat2_rank = len(cat2_shape)	77	cat2_rank = len(cat2_shape)
	@@ -95,26 +106,40 @@ def cat_slice_cat_fold_pass(graph: torch.fx.Graph) -> None:
95	continue	106	continue
96	cat1_inputs = cat1_node.args[0]	107	cat1_inputs = cat1_node.args[0]
97	cat1_dim = cat1_node.kwargs.get("dim", -1)	108	cat1_dim = cat1_node.kwargs.get("dim", -1)
98	- cat1_shape = get_node_shape(cat1_node)	109	+ cat1_shape = get_node_shape(cat1_node, allow_symbolic=True)
99	if not cat1_shape:	110	if not cat1_shape:
100	continue	111	continue
101	cat1_rank = len(cat1_shape)	112	cat1_rank = len(cat1_shape)
102	cat1_dim = cat1_dim + cat1_rank if cat1_dim == -1 else cat1_dim	113	cat1_dim = cat1_dim + cat1_rank if cat1_dim == -1 else cat1_dim
103	- if cat1_dim != cat2_dim or cat1_shape != cat2_shape:	114	+ if cat1_dim != cat2_dim or not shapes_statically_equal(cat1_shape, cat2_shape):
104	continue	115	continue
105	- sorted_ranges = [	116	+ # Normalize each slice's (start, stop), allowing symbolic bounds (step must be 1).
106	- (sl.start, sl.stop, sl.step)	117	+ resolved_ranges = []
107	- for sl in slice_ranges	118	+ valid_ranges = True
108	- if isinstance(sl.start, int) and isinstance(sl.stop, int)	119	+ for sl in slice_ranges:
109	- ]	120	+ if sl.step not in (1, None):
		121	+ valid_ranges = False
		122	+ break
		123	+ start = resolve_size_arg(0 if sl.start is None else sl.start)
		124	+ stop = resolve_size_arg(sl.stop) if sl.stop is not None else None
		125	+ if start is None or stop is None:
		126	+ valid_ranges = False
		127	+ break
		128	+ resolved_ranges.append((start, stop))
		129	+ if not valid_ranges or len(resolved_ranges) != len(cat1_inputs):
		130	+ continue
		131	+ # Chained coverage proof: start_0==0, start_i==stop_{i-1}, last stop covers the full dim length.
110	ranges_match = True	132	ranges_match = True
111	expected_start = 0	133	expected_start = 0
112	- for start, stop, step in sorted_ranges:	134	+ for start, stop in resolved_ranges:
113	- if start != expected_start or step not in (1, None):	135	+ if not statically_known_eq(start, expected_start):
114	ranges_match = False	136	ranges_match = False
115	break	137	break
116	expected_start = stop	138	expected_start = stop
117	- ranges_match = ranges_match and len(sorted_ranges) == len(cat1_inputs)	139	+ if ranges_match and not statically_known_eq(
		140	+ expected_start, cat2_shape[cat2_dim]
		141	+ ):
		142	+ ranges_match = False
118		143
119	if not ranges_match:	144	if not ranges_match:
120	continue	145	continue
	@@ -140,7 +165,7 @@ def pad_slice_fold(graph: torch.fx.Graph) -> None:
140	# 获取 pad 节点的输入和参数	165	# 获取 pad 节点的输入和参数
141	input_tensor = node.args[0]	166	input_tensor = node.args[0]
142	pad = node.args[1]	167	pad = node.args[1]
143	- input_shape = get_node_shape(input_tensor)	168	+ input_shape = get_node_shape(input_tensor, allow_symbolic=True)
144	if input_shape is None:	169	if input_shape is None:
145	continue	170	continue
146	pad_dim, _ = get_pad_dim_and_size(pad, input_shape)	171	pad_dim, _ = get_pad_dim_and_size(pad, input_shape)
	@@ -162,15 +187,16 @@ def pad_slice_fold(graph: torch.fx.Graph) -> None:
162	start = idx[pad_dim].start	187	start = idx[pad_dim].start
163	end = idx[pad_dim].stop	188	end = idx[pad_dim].stop
164	step = idx[pad_dim].step	189	step = idx[pad_dim].step
165	- slice_start = start if isinstance(start, int) else 0	190	+ slice_start = 0 if start is None else resolve_size_arg(start)
166	- slice_end = end if isinstance(end, int) else None	191	+ slice_end = None if end is None else resolve_size_arg(end)
167	slice_step = step if isinstance(step, int) else 1	192	slice_step = step if isinstance(step, int) else 1
168	- # 检查是否在维度上发生的切片，且切片范围不包含填充部分	193	+ # The slice upper bound must be provably within the valid pre-pad data region (not touching padding).
169	is_valid_prefix = (	194	is_valid_prefix = (
170	- isinstance(slice_end, int)	195	+ slice_start is not None
171	- and slice_end <= input_shape[pad_dim]	196	+ and slice_end is not None
172	and slice_step in (1, None)	197	and slice_step in (1, None)
173	- and slice_start <= slice_end	198	+ and statically_known_leq(slice_end, input_shape[pad_dim])
		199	+ and statically_known_leq(slice_start, slice_end)
174	)	200	)
175	if not is_valid_prefix:	201	if not is_valid_prefix:
176	all_slices_valid = False	202	all_slices_valid = False
	@@ -274,7 +300,7 @@ def fold_cat(graph: torch.fx.Graph) -> None:
274	is_cat, cat_axis = check_cat_op(node)	300	is_cat, cat_axis = check_cat_op(node)
275	if not is_cat:	301	if not is_cat:
276	continue	302	continue
277	- node_shape = get_node_shape(node)	303	+ node_shape = get_node_shape(node, allow_symbolic=True)
278	if not node_shape:	304	if not node_shape:
279	continue	305	continue
280	if cat_axis == len(node_shape) - 1:	306	if cat_axis == len(node_shape) - 1:
	@@ -285,7 +311,7 @@ def fold_cat(graph: torch.fx.Graph) -> None:
285	is_input_cat, input_cat_axis = check_cat_op(inp)	311	is_input_cat, input_cat_axis = check_cat_op(inp)
286	if is_input_cat:	312	if is_input_cat:
287	if len(inp.users) == 1:	313	if len(inp.users) == 1:
288	- inp_shape = get_node_shape(inp)	314	+ inp_shape = get_node_shape(inp, allow_symbolic=True)
289	effective_input_axis = input_cat_axis	315	effective_input_axis = input_cat_axis
290	if inp_shape and input_cat_axis == len(inp_shape) - 1:	316	if inp_shape and input_cat_axis == len(inp_shape) - 1:
291	effective_input_axis = -1	317	effective_input_axis = -1
	@@ -392,16 +418,21 @@ def fold_expand(graph: torch.fx.Graph) -> None:
392	if len(org_shape) != len(target_shape):	418	if len(org_shape) != len(target_shape):
393	return False	419	return False
394	for os, ts in zip(org_shape, target_shape):	420	for os, ts in zip(org_shape, target_shape):
395	- if os != ts and ts != -1:	421	+ if isinstance(ts, int) and ts == -1:
		422	+ continue
		423	+ if not statically_known_eq(os, ts):
396	return False	424	return False
397	return True	425	return True
398		426
399	for expand in candidates:	427	for expand in candidates:
400	inp = expand.args[0]	428	inp = expand.args[0]
401	target_shape = expand.args[1]	429	target_shape = expand.args[1]
402	- if not isinstance(target_shape, list):	430	+ if not isinstance(target_shape, (list, tuple)):
403	continue	431	continue
404	- inp_shape = get_node_shape(inp)	432	+ target_shape = resolve_size_list(target_shape)
		433	+ if target_shape is None:
		434	+ continue
		435	+ inp_shape = get_node_shape(inp, allow_symbolic=True)
405	if inp_shape is None:	436	if inp_shape is None:
406	continue	437	continue
407	org_shape = list(inp_shape)	438	org_shape = list(inp_shape)
	@@ -427,14 +458,14 @@ def fold_reduce(graph: torch.fx.Graph) -> None:
427		458
428	for reduce in reversed(candidates):	459	for reduce in reversed(candidates):
429	inp = get_input_node(reduce, 0)	460	inp = get_input_node(reduce, 0)
430	- shape = get_node_shape(inp)	461	+ shape = get_node_shape(inp, allow_symbolic=True)
431	if shape is None:	462	if shape is None:
432	continue	463	continue
433	dims = get_input_kw_node(reduce, "dim") or list(range(len(shape)))	464	dims = get_input_kw_node(reduce, "dim") or list(range(len(shape)))
434	if not isinstance(dims, list):	465	if not isinstance(dims, list):
435	dims = [dims]	466	dims = [dims]
436	keep_dim = get_input_kw_node(reduce, "keepdim") or False	467	keep_dim = get_input_kw_node(reduce, "keepdim") or False
437	- if all(shape[dim] == 1 for dim in dims):	468	+ if all(is_statically_one(shape[dim]) for dim in dims):
438	with graph.inserting_before(reduce):	469	with graph.inserting_before(reduce):
439	fold_res = _get_fold_result(graph, inp, dims, keep_dim)	470	fold_res = _get_fold_result(graph, inp, dims, keep_dim)
440	if fold_res:	471	if fold_res:
	@@ -454,7 +485,7 @@ def fold_sink_view(graph: torch.fx.Graph) -> None:
454	continue	485	continue
455	if len(node.users) != 1:	486	if len(node.users) != 1:
456	continue	487	continue
457	- view_shape = get_node_shape(node)	488	+ view_shape = get_node_shape(node, allow_symbolic=True)
458	if view_shape is None:	489	if view_shape is None:
459	continue	490	continue
460	user = next(iter(node.users))	491	user = next(iter(node.users))
	@@ -490,10 +521,10 @@ def fold_sink_view(graph: torch.fx.Graph) -> None:
490	other_shape = []	521	other_shape = []
491	other_val = other_node	522	other_val = other_node
492	else:	523	else:
493	- other_shape = get_node_shape(other_node)	524	+ other_shape = get_node_shape(other_node, allow_symbolic=True)
494	other_val = get_node_meta(other_node)	525	other_val = get_node_meta(other_node)
495	- result_shape = get_node_shape(user)	526	+ result_shape = get_node_shape(user, allow_symbolic=True)
496	- orig_shape = get_node_shape(node.args[0])	527	+ orig_shape = get_node_shape(node.args[0], allow_symbolic=True)
497	if (	528	if (
498	other_shape is not None	529	other_shape is not None
499	and result_shape is not None	530	and result_shape is not None
	@@ -501,10 +532,12 @@ def fold_sink_view(graph: torch.fx.Graph) -> None:
501	and orig_shape is not None	532	and orig_shape is not None
502	):	533	):
503	no_broadcast_dims = min(len(other_shape), len(orig_shape))	534	no_broadcast_dims = min(len(other_shape), len(orig_shape))
504	- if result_shape == view_shape and (	535	+ if shapes_statically_equal(result_shape, view_shape) and (
505	len(other_shape) == 0	536	len(other_shape) == 0
506	- or orig_shape[-no_broadcast_dims:]	537	+ or shapes_statically_equal(
507	- == view_shape[-no_broadcast_dims:]	538	+ orig_shape[-no_broadcast_dims:],
		539	+ view_shape[-no_broadcast_dims:],
		540	+ )
508	):	541	):
509	with graph.inserting_before(user):	542	with graph.inserting_before(user):
510	new_args = list(user.args)	543	new_args = list(user.args)
	@@ -700,11 +733,14 @@ def view_fold_pass(graph) -> None:
700	changed = True	733	changed = True
701	else:	734	else:
702	target_shape = view.args[1]	735	target_shape = view.args[1]
703	- if not isinstance(target_shape, list):	736	+ if not isinstance(target_shape, (list, tuple)):
704	continue	737	continue
705	- inp_shape = get_node_shape(inp)	738	+ target_shape = resolve_size_list(target_shape)
		739	+ if target_shape is None:
		740	+ continue
		741	+ inp_shape = get_node_shape(inp, allow_symbolic=True)
706	if inp_shape is not None:	742	if inp_shape is not None:
707	- if target_shape == list(inp_shape):	743	+ if shapes_statically_equal(target_shape, list(inp_shape)):
708	view.replace_all_uses_with(inp)	744	view.replace_all_uses_with(inp)
709	propagate_fake_tensor(inp, view, lambda x: x)	745	propagate_fake_tensor(inp, view, lambda x: x)
710	graph.erase_node(view)	746	graph.erase_node(view)
	@@ -775,13 +811,13 @@ def fold_redundant_ops(graph: torch.fx.Graph):
775	continue	811	continue
776	in_meta = _get_tensor_meta(first_arg)	812	in_meta = _get_tensor_meta(first_arg)
777	squeeze_out_meta = _get_tensor_meta(squeeze_node)	813	squeeze_out_meta = _get_tensor_meta(squeeze_node)
778	- in_shape = get_node_shape(first_arg)	814	+ in_shape = get_node_shape(first_arg, allow_symbolic=True)
779	- squeeze_out_shape = get_node_shape(squeeze_node)	815	+ squeeze_out_shape = get_node_shape(squeeze_node, allow_symbolic=True)
780	if in_meta is None or squeeze_out_meta is None:	816	if in_meta is None or squeeze_out_meta is None:
781	continue	817	continue
782	if in_shape is None or squeeze_out_shape is None:	818	if in_shape is None or squeeze_out_shape is None:
783	continue	819	continue
784	- if in_shape != squeeze_out_shape:	820	+ if not shapes_statically_equal(in_shape, squeeze_out_shape):
785	continue	821	continue
786	if in_meta.dtype != squeeze_out_meta.dtype:	822	if in_meta.dtype != squeeze_out_meta.dtype:
787	continue	823	continue
	@@ -807,7 +843,6 @@ def fold_redundant_ops(graph: torch.fx.Graph):
807	def dtype_optimal_pass(graph: torch.fx.Graph) -> None:	843	def dtype_optimal_pass(graph: torch.fx.Graph) -> None:
808	"""将不必要的 int64 优化为 int32：若 torch.arange 或 to(int64) 的取值	844	"""将不必要的 int64 优化为 int32：若 torch.arange 或 to(int64) 的取值
809	可被 int32 安全表示，则降级 dtype 以减少访存与计算开销。"""	845	可被 int32 安全表示，则降级 dtype 以减少访存与计算开销。"""
810	- int32_min, int32_max = -(231), 231 - 1
811	cast_dtype_limit = [torch.float32, torch.int32, torch.bool, torch.int16, torch.int8]	846	cast_dtype_limit = [torch.float32, torch.int32, torch.bool, torch.int16, torch.int8]
812	changed = False	847	changed = False
813	for node in list(graph.nodes): # 使用list避免修改时迭代问题	848	for node in list(graph.nodes): # 使用list避免修改时迭代问题
	@@ -837,24 +872,18 @@ def dtype_optimal_pass(graph: torch.fx.Graph) -> None:
837	# 如果 end 为 None，假设无限或跳过 (罕见，但安全)	872	# 如果 end 为 None，假设无限或跳过 (罕见，但安全)
838	if end is None:	873	if end is None:
839	continue	874	continue
840	- # 静态范围检查 (所有参数是常量)	875	+ # Normalize start/end/step (symbolic allowed); elements always lie in
841	- if all(isinstance(p, (int, float)) for p in [start, step, end]):	876	+ # [start, end), so downgrade is safe once both bounds provably fit int32; step must be a nonzero int.
842	- if step == 0:	877	+ r_start = resolve_size_arg(start)
843	- continue	878	+ r_end = resolve_size_arg(end)
844	- # 如果 step 非整数且 dtype 是 int，警告 (arange 会自动转为 float)	879	+ r_step = resolve_size_arg(step)
845	- if not isinstance(step, int):	880	+ if r_start is None or r_end is None or r_step is None:
846	- continue	881	+ continue
847	- # 计算序列长度和 min/max 值	882	+ if not isinstance(r_step, int) or r_step == 0:
848	- num_elements = (	883	+ continue
849	- math.ceil((end - start) / step)	884	+ if statically_fits_int32(r_start, r_end):
850	- if step > 0	885	+ node.kwargs = {**node.kwargs, "dtype": torch.int32}
851	- else math.ceil((start - end) / -step)	886	+ changed = True
852	- )
853	- seq_min = min(start, start + (num_elements - 1) * step)
854	- seq_max = max(start, start + (num_elements - 1) * step)
855	- if seq_min > int32_min and seq_max < int32_max:
856	- node.kwargs = {**node.kwargs, "dtype": torch.int32}
857	- changed = True
858	if node.op == "call_method":	887	if node.op == "call_method":
859	input_node = node.args[0]	888	input_node = node.args[0]
860	input_fake = (	889	input_fake = (
	@@ -922,7 +951,7 @@ def cat_to_view_pass(graph: torch.fx.Graph) -> None:
922	cat_inputs = cat.args[0]	951	cat_inputs = cat.args[0]
923	if not isinstance(cat_inputs, (list, tuple)) or len(cat_inputs) < 2:	952	if not isinstance(cat_inputs, (list, tuple)) or len(cat_inputs) < 2:
924	continue	953	continue
925	- cat_shape = get_node_shape(cat)	954	+ cat_shape = get_node_shape(cat, allow_symbolic=True)
926	if cat_shape is None:	955	if cat_shape is None:
927	continue	956	continue
928	rank = len(cat_shape)	957	rank = len(cat_shape)
	@@ -934,6 +963,7 @@ def cat_to_view_pass(graph: torch.fx.Graph) -> None:
934	parent = None	963	parent = None
935	intervals = []	964	intervals = []
936	valid = True	965	valid = True
		966	+ all_static = True
937	for inp in cat_inputs:	967	for inp in cat_inputs:
938	if not (	968	if not (
939	isinstance(inp, torch.fx.Node)	969	isinstance(inp, torch.fx.Node)
	@@ -958,30 +988,58 @@ def cat_to_view_pass(graph: torch.fx.Graph) -> None:
958	break	988	break
959	sl_start = inp.args[2] if len(inp.args) > 2 else 0	989	sl_start = inp.args[2] if len(inp.args) > 2 else 0
960	sl_end = inp.args[3] if len(inp.args) > 3 else None	990	sl_end = inp.args[3] if len(inp.args) > 3 else None
961	- if not isinstance(sl_start, int):	991	+ r_start = resolve_size_arg(0 if sl_start is None else sl_start)
962	- valid = False	992	+ r_end = None if sl_end is None else resolve_size_arg(sl_end)
963	- break	993	+ if r_start is None or (sl_end is not None and r_end is None):
964	- if sl_end is not None and not isinstance(sl_end, int):
965	valid = False	994	valid = False
966	break	995	break
		996	+ if not isinstance(r_start, int) or (
		997	+ r_end is not None and not isinstance(r_end, int)
		998	+ ):
		999	+ all_static = False
967	if parent is None:	1000	if parent is None:
968	parent = p	1001	parent = p
969	elif parent is not p:	1002	elif parent is not p:
970	valid = False	1003	valid = False
971	break	1004	break
972	- intervals.append((sl_start, sl_end))	1005	+ intervals.append((r_start, r_end))
973		1006
974	if not valid or parent is None:	1007	if not valid or parent is None:
975	continue	1008	continue
976	- parent_shape = get_node_shape(parent)	1009	+ parent_shape = get_node_shape(parent, allow_symbolic=True)
977	if parent_shape is None or len(parent_shape) != rank:	1010	if parent_shape is None or len(parent_shape) != rank:
978	continue	1011	continue
979	- if list(parent_shape) != list(cat_shape):	1012	+ if not shapes_statically_equal(parent_shape, cat_shape):
980	continue	1013	continue
981	- dim_size_raw = parent_shape[cat_dim]	1014	+ dim_size = parent_shape[cat_dim]
982	- if isinstance(dim_size_raw, torch.SymInt):	1015	+
		1016	+ if not (all_static and isinstance(dim_size, int)):
		1017	+ # Symbolic path: negative indices cannot be reliably normalized under
		1018	+ # symbols; only fold when the slices, in cat input order, contiguously
		1019	+ # cover the full dim length from 0 (identity view); skip rotation cases.
		1020	+ expected = 0
		1021	+ ok_sym = True
		1022	+ for s, e in intervals:
		1023	+ e_eff = dim_size if e is None else e
		1024	+ if not statically_known_geq(s, 0) or not statically_known_eq(
		1025	+ s, expected
		1026	+ ):
		1027	+ ok_sym = False
		1028	+ break
		1029	+ expected = e_eff
		1030	+ if ok_sym and statically_known_eq(expected, dim_size):
		1031	+ cat.replace_all_uses_with(parent)
		1032	+ changed = True
		1033	+ log.info(
		1034	+ "cat_to_view_pass: collapsed cat(%d slices, dim=%d) of %s "
		1035	+ "→ identity view (dynamic full cover)",
		1036	+ len(cat_inputs),
		1037	+ cat_dim,
		1038	+ parent.name,
		1039	+ )
983	continue	1040	continue
984	- dim_size = int(dim_size_raw)	1041	+
		1042	+ dim_size = int(dim_size)
985		1043
986	normalised = []	1044	normalised = []
987	ok = True	1045	ok = True
	@@ -1128,21 +1186,27 @@ def repeat_to_expand_pass(graph: torch.fx.Graph) -> None:
1128	if not isinstance(repeats, (list, tuple)):	1186	if not isinstance(repeats, (list, tuple)):
1129	continue	1187	continue
1130		1188
1131	- in_shape = get_node_shape(inp)	1189	+ in_shape = get_node_shape(inp, allow_symbolic=True)
1132	if in_shape is None:	1190	if in_shape is None:
1133	continue	1191	continue
1134	if len(repeats) != len(in_shape):	1192	if len(repeats) != len(in_shape):
1135	continue	1193	continue
1136		1194
		1195	+ # Only broadcast (no physical copy) can be rewritten: each dim either is
		1196	+ # not repeated (r==1, output keeps the original dim, may be symbolic) or the
		1197	+ # original dim is provably 1 (output dim is r). repeats must be int constants.
1137	valid = True	1198	valid = True
		1199	+ out_shape = []
1138	for r, s in zip(repeats, in_shape):	1200	for r, s in zip(repeats, in_shape):
1139	- if isinstance(r, torch.SymInt) or isinstance(s, torch.SymInt):	1201	+ r = resolve_size_arg(r)
		1202	+ if not isinstance(r, int):
1140	valid = False	1203	valid = False
1141	break	1204	break
1142	- if not (isinstance(r, int) and isinstance(s, int)):	1205	+ if r == 1:
1143	- valid = False	1206	+ out_shape.append(s)
1144	- break	1207	+ elif is_statically_one(s):
1145	- if r != 1 and s != 1:	1208	+ out_shape.append(r)
		1209	+ else:
1146	valid = False	1210	valid = False
1147	break	1211	break
1148	if not valid:	1212	if not valid:
	@@ -1155,17 +1219,19 @@ def repeat_to_expand_pass(graph: torch.fx.Graph) -> None:
1155	if not users_ok or not list(rpt.users):	1219	if not users_ok or not list(rpt.users):
1156	continue	1220	continue
1157		1221
1158	- out_shape = [int(r) * int(s) for r, s in zip(repeats, in_shape)]
1159	-
1160	inp_fake = inp.meta.get("val")	1222	inp_fake = inp.meta.get("val")
1161	fake_mode = (	1223	fake_mode = (
1162	getattr(inp_fake, "fake_mode", None) if inp_fake is not None else None	1224	getattr(inp_fake, "fake_mode", None) if inp_fake is not None else None
1163	)	1225	)
1164		1226
1165	with graph.inserting_before(rpt):	1227	with graph.inserting_before(rpt):
		1228	+ # Symbolic dims in the output shape all come from inp itself; materialize as sym_size refs.
		1229	+ expand_shape = materialize_shape(graph, out_shape, inp)
		1230	+ if expand_shape is None:
		1231	+ continue
1166	exp = graph.call_function(	1232	exp = graph.call_function(
1167	torch.ops.aten.expand.default,	1233	torch.ops.aten.expand.default,
1168	- args=(inp, list(out_shape)),	1234	+ args=(inp, expand_shape),
1169	)	1235	)
1170	if "val" in rpt.meta:	1236	if "val" in rpt.meta:
1171	exp.meta["val"] = rpt.meta["val"]	1237	exp.meta["val"] = rpt.meta["val"]
	@@ -1237,28 +1303,24 @@ _IOTA_DTYPE_CLOSING_OPS = frozenset(
1237	)	1303	)
1238		1304
1239		1305
1240	-def _prims_iota_value_range(node):	1306	+def _prims_iota_endpoints(node):
1241	- """计算 prims.iota 节点产生序列的 [lo, hi) 取值范围；	1307	+ """Return the two endpoints (start, last) of a prims.iota sequence (may be symbolic);
1242	- 若参数非常量整数则返回 None。"""	1308	+ all elements lie within [min(start,last), max(start,last)]. None if args are unresolvable."""
1243	if not (	1309	if not (
1244	node.op == "call_function"	1310	node.op == "call_function"
1245	and node.target is torch.ops.prims.iota.default	1311	and node.target is torch.ops.prims.iota.default
1246	and node.args	1312	and node.args
1247	):	1313	):
1248	return None	1314	return None
1249	- length = node.args[0]	1315	+ length = resolve_size_arg(node.args[0])
1250	- start = node.kwargs.get("start", 0)	1316	+ start = resolve_size_arg(node.kwargs.get("start", 0))
1251	- step = node.kwargs.get("step", 1)	1317	+ step = resolve_size_arg(node.kwargs.get("step", 1))
1252	- if not (	1318	+ if length is None or start is None or step is None:
1253	- isinstance(length, int) and isinstance(start, int) and isinstance(step, int)
1254	- ):
1255	return None	1319	return None
1256	- if length <= 0:	1320	+ if isinstance(length, int) and length <= 0:
1257	return (start, start)	1321	return (start, start)
1258	last = start + (length - 1) * step	1322	last = start + (length - 1) * step
1259	- lo = min(start, last)	1323	+ return (start, last)
1260	- hi = max(start, last)
1261	- return (lo, hi + 1)
1262		1324
1263		1325
1264	def _collect_iota_downcast_closure(iota_node):	1326	def _collect_iota_downcast_closure(iota_node):
	@@ -1284,31 +1346,6 @@ def _collect_iota_downcast_closure(iota_node):
1284	return middle_ids	1346	return middle_ids
1285		1347
1286		1348
1287	-def _refresh_fake_meta(node, fake_mode):
1288	- """基于当前 args/kwargs 在 fake_mode 下重新执行算子，刷新节点的 meta['val'] FakeTensor。"""
1289	-
1290	- def resolve(arg):
1291	- if isinstance(arg, torch.fx.Node):
1292	- return arg.meta.get("val", arg)
1293	- if isinstance(arg, (list, tuple)):
1294	- return type(arg)(resolve(x) for x in arg)
1295	- return arg
1296	-
1297	- try:
1298	- with fake_mode:
1299	- new_val = node.target(
1300	- *[resolve(a) for a in node.args],
1301	- **{k: resolve(v) for k, v in node.kwargs.items()},
1302	- )
1303	- node.meta["val"] = new_val
1304	- except Exception:
1305	- pass
1306	-
1307	-
1308	-_INT32_MIN = -(1 << 31)
1309	-_INT32_MAX = (1 << 31) - 1
1310	-
1311	-
1312	def _hashable_const_key(value):	1349	def _hashable_const_key(value):
1313	"""将常量参数（含嵌套 list/tuple/dict）转为可哈希的 key，便于做常量折叠的 CSE 比较。"""	1350	"""将常量参数（含嵌套 list/tuple/dict）转为可哈希的 key，便于做常量折叠的 CSE 比较。"""
1314	if isinstance(value, list):	1351	if isinstance(value, list):
	@@ -1412,11 +1449,10 @@ def fold_iota_arithmetic_pass(graph: torch.fx.Graph) -> None:
1412	continue	1449	continue
1413	if iota.kwargs.get("dtype") is not torch.int64:	1450	if iota.kwargs.get("dtype") is not torch.int64:
1414	continue	1451	continue
1415	- rng = _prims_iota_value_range(iota)	1452	+ endpoints = _prims_iota_endpoints(iota)
1416	- if rng is None:	1453	+ if endpoints is None:
1417	continue	1454	continue
1418	- lo, hi_exc = rng	1455	+ if not statically_fits_int32(*endpoints):
1419	- if lo < _INT32_MIN or hi_exc - 1 > _INT32_MAX:
1420	continue	1456	continue
1421		1457
1422	fake = iota.meta.get("val")	1458	fake = iota.meta.get("val")
	@@ -1431,30 +1467,28 @@ def fold_iota_arithmetic_pass(graph: torch.fx.Graph) -> None:
1431	new_kwargs = dict(iota.kwargs)	1467	new_kwargs = dict(iota.kwargs)
1432	new_kwargs["dtype"] = torch.int32	1468	new_kwargs["dtype"] = torch.int32
1433	iota.kwargs = new_kwargs	1469	iota.kwargs = new_kwargs
1434	- _refresh_fake_meta(iota, fake_mode)	1470	+ refresh_fake_meta(iota, fake_mode)
1435	if iota.meta.get("val") is None or iota.meta["val"].dtype is not torch.int32:	1471	if iota.meta.get("val") is None or iota.meta["val"].dtype is not torch.int32:
1436	new_kwargs["dtype"] = torch.int64	1472	new_kwargs["dtype"] = torch.int64
1437	iota.kwargs = new_kwargs	1473	iota.kwargs = new_kwargs
1438	- _refresh_fake_meta(iota, fake_mode)	1474	+ refresh_fake_meta(iota, fake_mode)
1439	continue	1475	continue
1440		1476
1441	middle_nodes_in_topo = [n for n in graph.nodes if id(n) in middle_ids]	1477	middle_nodes_in_topo = [n for n in graph.nodes if id(n) in middle_ids]
1442	for n in middle_nodes_in_topo:	1478	for n in middle_nodes_in_topo:
1443	- _refresh_fake_meta(n, fake_mode)	1479	+ refresh_fake_meta(n, fake_mode)
1444	for n in list(graph.nodes):	1480	for n in list(graph.nodes):
1445	if (	1481	if (
1446	n.op == "call_function"	1482	n.op == "call_function"
1447	and n.target in _IOTA_DTYPE_CLOSING_OPS	1483	and n.target in _IOTA_DTYPE_CLOSING_OPS
1448	and any(u is iota or id(u) in middle_ids for u in n.all_input_nodes)	1484	and any(u is iota or id(u) in middle_ids for u in n.all_input_nodes)
1449	):	1485	):
1450	- _refresh_fake_meta(n, fake_mode)	1486	+ refresh_fake_meta(n, fake_mode)
1451		1487
1452	changed = True	1488	changed = True
1453	log.info(	1489	log.info(
1454	- "fold_iota_arithmetic_pass: downcast iota[%d,%d) int64 → int32"	1490	+ "fold_iota_arithmetic_pass: downcast iota int64 → int32"
1455	" (%d transparent user%s in closure)",	1491	" (%d transparent user%s in closure)",
1456	- lo,
1457	- hi_exc,
1458	len(middle_ids),	1492	len(middle_ids),
1459	"" if len(middle_ids) == 1 else "s",	1493	"" if len(middle_ids) == 1 else "s",
1460	)	1494	)
	@@ -1620,7 +1654,7 @@ def broadcast_const_mask_compress(graph: torch.fx.Graph) -> None:
1620	f_val,	1654	f_val,
1621	replacement_kind,	1655	replacement_kind,
1622	action,	1656	action,
1623	- get_node_shape(w),	1657	+ get_node_shape(w, allow_symbolic=True),
1624	)	1658	)
1625		1659
1626	eliminate_dead_code(graph, changed, broadcast_const_mask_compress.__name__)	1660	eliminate_dead_code(graph, changed, broadcast_const_mask_compress.__name__)
	@@ -1938,7 +1972,7 @@ def bool_cast_mul_to_where_pass(graph: torch.fx.Graph) -> None:
1938	cast_target_dtype,	1972	cast_target_dtype,
1939	cast_src.name,	1973	cast_src.name,
1940	chain_desc,	1974	chain_desc,
1941	- get_node_shape(other),	1975	+ get_node_shape(other, allow_symbolic=True),
1942	)	1976	)
1943		1977
1944	eliminate_dead_code(graph, changed, bool_cast_mul_to_where_pass.__name__)	1978	eliminate_dead_code(graph, changed, bool_cast_mul_to_where_pass.__name__)
	@@ -2040,7 +2074,7 @@ def sign_diff_hamming_fuse_pass(graph: torch.fx.Graph) -> None:
2040		2074
2041	if fake_mode is not None:	2075	if fake_mode is not None:
2042	for n in (gt_x, gt_y, ne_node, new_sum):	2076	for n in (gt_x, gt_y, ne_node, new_sum):
2043	- _refresh_fake_meta(n, fake_mode)	2077	+ refresh_fake_meta(n, fake_mode)
2044	if "val" not in new_sum.meta and "val" in sum_node.meta:	2078	if "val" not in new_sum.meta and "val" in sum_node.meta:
2045	new_sum.meta["val"] = sum_node.meta["val"]	2079	new_sum.meta["val"] = sum_node.meta["val"]
2046		2080
	@@ -2079,8 +2113,18 @@ def _has_default_embedding_args(node):
2079		2113
2080		2114
2081	def _symbolic_shape_key(shape):	2115	def _symbolic_shape_key(shape):
2082	- """生成可哈希的 shape key：SymInt 维度转字符串，其它维度转 int，便于 embedding 分组。"""	2116	+ """Build a hashable shape key: SymInt dims use the canonical sympy expr string
2083	- return tuple(str(d) if isinstance(d, torch.SymInt) else int(d) for d in shape)	2117	+ (so s0+s0 and 2*s0 group together), other dims become int, for embedding grouping."""
		2118	+
		2119	+ def _dim_key(d):
		2120	+ if isinstance(d, torch.SymInt):
		2121	+ try:
		2122	+ return f"sym:{d.node.expr}"
		2123	+ except Exception:
		2124	+ return f"sym:{d}"
		2125	+ return int(d)
		2126	+
		2127	+ return tuple(_dim_key(d) for d in shape)
2084		2128
2085		2129
2086	def _weight_node_key(w):	2130	def _weight_node_key(w):
	@@ -2192,10 +2236,11 @@ def _detect_indices_parent(nodes):
2192	return None, None	2236	return None, None
2193	slices_info.append((int(start) if start is not None else 0, end))	2237	slices_info.append((int(start) if start is not None else 0, end))
2194		2238
2195	- parent_shape = get_node_shape(parent)	2239	+ parent_shape = get_node_shape(parent, allow_symbolic=True)
2196	if parent_shape is None or slice_dim >= len(parent_shape):	2240	if parent_shape is None or slice_dim >= len(parent_shape):
2197	return None, None	2241	return None, None
2198	dim_size = parent_shape[slice_dim]	2242	dim_size = parent_shape[slice_dim]
		2243	+ # The sliced field dim must be static (the N*L coverage proof needs a concrete length); batch dim may be symbolic.
2199	if isinstance(dim_size, torch.SymInt):	2244	if isinstance(dim_size, torch.SymInt):
2200	return None, None	2245	return None, None
2201	dim_size = int(dim_size)	2246	dim_size = int(dim_size)
	@@ -2351,7 +2396,7 @@ def _apply_pattern_c_reshape_first(
2351	再做一次 embedding + reduce，最后用 select（或 cat 折叠）接回原下游使用者。"""	2396	再做一次 embedding + reduce，最后用 select（或 cat 折叠）接回原下游使用者。"""
2352	N = len(nodes)	2397	N = len(nodes)
2353		2398
2354	- parent_shape = get_node_shape(parent_node)	2399	+ parent_shape = get_node_shape(parent_node, allow_symbolic=True)
2355	if parent_shape is None:	2400	if parent_shape is None:
2356	return False	2401	return False
2357	if slice_dim >= len(parent_shape):	2402	if slice_dim >= len(parent_shape):
	@@ -2405,9 +2450,16 @@ def _apply_pattern_c_reshape_first(
2405	return False	2450	return False
2406		2451
2407	with graph.inserting_before(ordered_emb_nodes[0]):	2452	with graph.inserting_before(ordered_emb_nodes[0]):
		2453	+ # Symbolic dims of new_idx_shape (e.g. batch dim) all come from parent;
		2454	+ # materialize as sym_size refs to parent instead of writing raw SymInt into node args.
		2455	+ materialized_idx_shape = materialize_shape(
		2456	+ graph, list(new_idx_shape), parent_node
		2457	+ )
		2458	+ if materialized_idx_shape is None:
		2459	+ return False
2408	reshaped_parent = graph.call_function(	2460	reshaped_parent = graph.call_function(
2409	torch.ops.aten.reshape.default,	2461	torch.ops.aten.reshape.default,
2410	- args=(parent_node, list(new_idx_shape)),	2462	+ args=(parent_node, materialized_idx_shape),
2411	)	2463	)
2412	reshaped_parent.meta["val"] = new_idx_fake	2464	reshaped_parent.meta["val"] = new_idx_fake
2413		2465
	@@ -2512,8 +2564,9 @@ def batch_embedding_fusion_pass(graph: torch.fx.Graph) -> None:
2512	weight = node.args[0]	2564	weight = node.args[0]
2513	indices = node.args[1]	2565	indices = node.args[1]
2514		2566
		2567	+ # The weight table (V, D) must be static; indices may contain symbolic dims (e.g. dynamic batch).
2515	w_shape = get_node_shape(weight)	2568	w_shape = get_node_shape(weight)
2516	- idx_shape = get_node_shape(indices)	2569	+ idx_shape = get_node_shape(indices, allow_symbolic=True)
2517	if w_shape is None or idx_shape is None or len(w_shape) != 2:	2570	if w_shape is None or idx_shape is None or len(w_shape) != 2:
2518	continue	2571	continue
2519	if len(idx_shape) < 1:	2572	if len(idx_shape) < 1:

Mtorch_npu/_inductor/fx_passes/utils/get_binary_fold_result.py+70-35

	@@ -3,6 +3,16 @@ import torch
3	from torch import fx	3	from torch import fx
4	from torch.multiprocessing.reductions import StorageWeakRef	4	from torch.multiprocessing.reductions import StorageWeakRef
5		5
		6	+from .symbolic_shape_util import (
		7	+ has_free_symbols,
		8	+ materialize_shape,
		9	+ resolve_size_arg,
		10	+ shapes_statically_equal,
		11	+ statically_known_eq,
		12	+ statically_known_geq,
		13	+ statically_known_gt,
		14	+)
		15	+
6		16
7	MAX_INT64 = 9223372036854775807	17	MAX_INT64 = 9223372036854775807
8		18
	@@ -22,7 +32,7 @@ def get_val_meta_info(target_meta: Dict[str, Any]):
22	return None	32	return None
23		33
24		34
25	-def get_node_shape(node: torch.fx.Node):	35	+def get_node_shape(node: torch.fx.Node, allow_symbolic: bool = False):
26	if (node.meta == {}) or ('example_value' not in node.meta and 'tensor_meta' not in node.meta and 'val' not in node.meta):	36	if (node.meta == {}) or ('example_value' not in node.meta and 'tensor_meta' not in node.meta and 'val' not in node.meta):
27	return None	37	return None
28	shape = None	38	shape = None
	@@ -34,7 +44,10 @@ def get_node_shape(node: torch.fx.Node):
34	shape = example_value.size()	44	shape = example_value.size()
35	elif 'tensor_meta' in node.meta:	45	elif 'tensor_meta' in node.meta:
36	shape = node.meta['tensor_meta'].shape	46	shape = node.meta['tensor_meta'].shape
37	- if any(isinstance(s, torch.SymInt) for s in shape):	47	+ if shape is None:
		48	+ return None
		49	+ # allow_symbolic=False keeps legacy behavior (None if any symbolic dim); True returns the symbolic shape as-is.
		50	+ if not allow_symbolic and any(isinstance(s, torch.SymInt) for s in shape):
38	return None	51	return None
39	return shape	52	return shape
40		53
	@@ -73,7 +86,7 @@ def get_binary_fold_result(
73	if isinstance(inp, fx.Node):	86	if isinstance(inp, fx.Node):
74	val_meta = get_val_meta_info(target_meta)	87	val_meta = get_val_meta_info(target_meta)
75	node_meta = get_node_meta(inp)	88	node_meta = get_node_meta(inp)
76	- node_meta_shape = get_node_shape(inp)	89	+ node_meta_shape = get_node_shape(inp, allow_symbolic=True)
77	if node_meta_shape is None:	90	if node_meta_shape is None:
78	return None	91	return None
79	if val_meta is None:	92	if val_meta is None:
	@@ -93,16 +106,22 @@ def get_binary_fold_result(
93	)	106	)
94	propagate_fake_tensor(new_node, inp, lambda fake: fake.full(val_meta.shape))	107	propagate_fake_tensor(new_node, inp, lambda fake: fake.full(val_meta.shape))
95	else:	108	else:
96	- if node_meta_shape == val_meta.shape:	109	+ if shapes_statically_equal(node_meta_shape, val_meta.shape):
97	expand = inp	110	expand = inp
98	else:	111	else:
99	- if any(isinstance(s, torch.SymInt) for s in val_meta.shape):	112	+ target_shape = list(val_meta.shape)
100	- return None	113	+ if has_free_symbols(target_shape):
		114	+ # Symbolic target shape must be materialized as sym_size refs on
		115	+ # inp's own dims; abandon fold if unmaterializable (e.g. a
		116	+ # broadcast dim originating from the other operand).
		117	+ target_shape = materialize_shape(graph, target_shape, inp)
		118	+ if target_shape is None:
		119	+ return None
101	expand = graph.call_function(	120	expand = graph.call_function(
102	torch.ops.aten.expand.default,	121	torch.ops.aten.expand.default,
103	args=(	122	args=(
104	inp,	123	inp,
105	- val_meta.shape	124	+ target_shape
106	)	125	)
107	)	126	)
108	propagate_fake_tensor(expand, inp, lambda fake: fake.expand(val_meta.shape))	127	propagate_fake_tensor(expand, inp, lambda fake: fake.expand(val_meta.shape))
	@@ -126,7 +145,7 @@ def get_binary_fold_result(
126		145
127		146
128	def _get_fold_result(graph: torch.fx.Graph, src, dims: List[int], keep_dim: bool) -> fx.Node:	147	def _get_fold_result(graph: torch.fx.Graph, src, dims: List[int], keep_dim: bool) -> fx.Node:
129	- # 检查 src 是否被原地操作引用	148	+ # Skip if src is referenced by an in-place op
130	for user in src.users:	149	for user in src.users:
131	if user.op == "call_function" and user.target in [torch.ops.aten.add_.Tensor, torch.ops.aten.add_.Scalar, torch.ops.aten.copy_.default]:	150	if user.op == "call_function" and user.target in [torch.ops.aten.add_.Tensor, torch.ops.aten.add_.Scalar, torch.ops.aten.copy_.default]:
132	return None	151	return None
	@@ -174,21 +193,25 @@ def _fold_slice(node: torch.fx.Node, graph: torch.fx.Graph) -> bool:
174	if len(node.args) < 4:	193	if len(node.args) < 4:
175	return False	194	return False
176	src_node, dim, start, end = node.args[0], node.args[1], node.args[2], node.args[3]	195	src_node, dim, start, end = node.args[0], node.args[1], node.args[2], node.args[3]
177	- if not isinstance(dim, int) or not isinstance(start, int):	196	+ if not isinstance(dim, int):
178	return False	197	return False
179	- if start != 0:	198	+ # start / end may be symbolic (e.g. slicing to a symbolic dim length); normalize to int/SymInt.
		199	+ start = resolve_size_arg(start)
		200	+ if start is None or not statically_known_eq(start, 0):
180	return False	201	return False
181	- if end is not None and not isinstance(end, int):	202	+ if end is not None:
182	- return False	203	+ end = resolve_size_arg(end)
183	- src_shape = get_node_shape(src_node)	204	+ if end is None:
		205	+ return False
		206	+ src_shape = get_node_shape(src_node, allow_symbolic=True)
184	if src_shape is None:	207	if src_shape is None:
185	return False	208	return False
186		209
187	- node_shape = get_node_shape(node)	210	+ node_shape = get_node_shape(node, allow_symbolic=True)
188	if node_shape is None:	211	if node_shape is None:
189	return False	212	return False
190		213
191	- if src_shape != node_shape:	214	+ if not shapes_statically_equal(src_shape, node_shape):
192	return False	215	return False
193		216
194	if dim >= len(src_shape):	217	if dim >= len(src_shape):
	@@ -197,8 +220,9 @@ def _fold_slice(node: torch.fx.Node, graph: torch.fx.Graph) -> bool:
197	dim_length = src_shape[dim]	220	dim_length = src_shape[dim]
198		221
199	is_full_slice = (	222	is_full_slice = (
200	- isinstance(start, int) and start == 0 and	223	+ end is None
201	- (end is None or end >= MAX_INT64 or end >= dim_length)	224	+ or statically_known_geq(end, MAX_INT64)
		225	+ or statically_known_geq(end, dim_length)
202	)	226	)
203		227
204	if is_full_slice:	228	if is_full_slice:
	@@ -214,21 +238,28 @@ def _fold_slice_scatter(node: torch.fx.Node, graph: torch.fx.Graph) -> bool:
214	return False	238	return False
215	base_node, view_node, dim, start, end = node.args[:5]	239	base_node, view_node, dim, start, end = node.args[:5]
216		240
217	- if not isinstance(dim, int) or not isinstance(start, int) or not isinstance(end, int):	241	+ if not isinstance(dim, int):
218	return False	242	return False
219	- base_shape = get_node_shape(base_node)	243	+ start = resolve_size_arg(start)
220	- view_shape = get_node_shape(view_node)	244	+ end = resolve_size_arg(end)
		245	+ if start is None or end is None:
		246	+ return False
		247	+ base_shape = get_node_shape(base_node, allow_symbolic=True)
		248	+ view_shape = get_node_shape(view_node, allow_symbolic=True)
221	if (base_shape is None) or (view_shape is None):	249	if (base_shape is None) or (view_shape is None):
222	return False	250	return False
223		251
224	- if base_shape != view_shape:	252	+ if not shapes_statically_equal(base_shape, view_shape):
225	return False	253	return False
226		254
227	- if start != 0:	255	+ if not statically_known_eq(start, 0):
		256	+ return False
		257	+
		258	+ if dim >= len(view_shape):
228	return False	259	return False
229		260
230	dim_length = view_shape[dim]	261	dim_length = view_shape[dim]
231	- if end != dim_length:	262	+ if not statically_known_eq(end, dim_length):
232	return False	263	return False
233		264
234	node.replace_all_uses_with(view_node)	265	node.replace_all_uses_with(view_node)
	@@ -238,28 +269,32 @@ def _fold_slice_scatter(node: torch.fx.Node, graph: torch.fx.Graph) -> bool:
238		269
239		270
240	def get_pad_dim_and_size(pad, input_shape):	271	def get_pad_dim_and_size(pad, input_shape):
241	- """	272	+ """Extract the padded dim and pad amount from pad args and input shape.
242	- 从 pad 参数和输入张量形状中提取填充的维度和填充量。	273	+
243	Args:	274	Args:
244	- pad: 填充参数列表，例如 [0, 0, 0, max_seq_len]。	275	+ pad: pad list, e.g. [0, 0, 0, max_seq_len].
245	- input_shape: 输入张量的形状，例如 [128, 50, 128]。	276	+ input_shape: input tensor shape, e.g. [128, 50, 128].
246	Returns:	277	Returns:
247	- pad_dim: 填充的维度索引（从 0 开始）。	278	+ pad_dim: padded dim index (0-based).
248	- pad_size: 右侧填充量。	279	+ pad_size: right-side pad amount.
249	"""	280	"""
250	- N = len(input_shape) # 张量维度数	281	+ N = len(input_shape) # tensor rank
251	pad_dim = None	282	pad_dim = None
252	pad_size = 0	283	pad_size = 0
253	for i in range(len(pad) // 2):	284	for i in range(len(pad) // 2):
254	left, right = pad[2 * i], pad[2 * i + 1]	285	left, right = pad[2 * i], pad[2 * i + 1]
255	- dim = N - 1 - i # 从后向前映射维度	286	+ dim = N - 1 - i # map back-to-front to dims
256	- if left == 0 and right > 0:	287	+ # left/right may be symbolic (e.g. pad to a fixed max_len - s0); use three-valued checks.
		288	+ left_zero = statically_known_eq(left, 0)
		289	+ if left_zero and statically_known_gt(right, 0):
257	if pad_dim is not None:	290	if pad_dim is not None:
258	- return None, 0 # 多个维度填充，复杂情况，跳过	291	+ return None, 0 # multiple padded dims, complex case -> skip
259	pad_dim = dim	292	pad_dim = dim
260	pad_size = right	293	pad_size = right
261	- elif left != 0 or right != 0:	294	+ elif left_zero and statically_known_eq(right, 0):
262	- return None, 0 # 复杂填充，跳过	295	+ continue # no padding on this dim
		296	+ else:
		297	+ return None, 0 # complex or undecidable padding -> skip
263	return pad_dim, pad_size	298	return pad_dim, pad_size
264		299
265		300

Atorch_npu/_inductor/fx_passes/utils/symbolic_shape_util.py+435-0

@@ -0,0 +1,435 @@
	1	+"""Symbolic-shape primitives for dynamic-shape graph optimization.
	2	+
	3	+Provides device-independent, guard-free symbolic reasoning helpers for the graph
	4	+passes under ``ascend_custom_passes``, so optimizations that used to fire only on
	5	+static shapes (integer dims) also work on dynamic shapes (``torch.SymInt`` dims).
	6	+
	7	+Design principles:
	8	+- Three-valued logic: a symbolic comparison is statically-true / statically-false
	9	+ / undecidable; optimize only when statically-true.
	10	+- Guard-free: rely on ``statically_known_true`` / ``bound_sympy`` style static
	11	+ reasoning, never adding guards to ``ShapeEnv`` nor changing recompile bounds.
	12	+- Pure-static inputs fall back to the original semantics (zero static regression).
	13	+"""
	14	+
	15	+import os
	16	+from collections import defaultdict
	17	+from typing import List, Optional, Tuple, Union
	18	+
	19	+import torch
	20	+
	21	+from ...config import log
	22	+
	23	+
	24	+try:
	25	+ from torch.fx.experimental.symbolic_shapes import statically_known_true as _skt
	26	+except Exception: # pragma: no cover - fallback for older torch
	27	+ _skt = None
	28	+
	29	+
	30	+Number = Union[int, "torch.SymInt"]
	31	+
	32	+
	33	+# ---------------------------------------------------------------------------
	34	+# Global switch
	35	+# ---------------------------------------------------------------------------
	36	+def dynamic_fx_pass_enabled() -> bool:
	37	+ """Master switch for dynamic-shape optimization (on by default).
	38	+
	39	+ When set to 0/false/off, passes fall back to the legacy static behavior.
	40	+ """
	41	+ return os.environ.get("NPU_INDUCTOR_DYNAMIC_FX_PASS", "1").lower() not in (
	42	+ "0",
	43	+ "false",
	44	+ "off",
	45	+ )
	46	+
	47	+
	48	+# ---------------------------------------------------------------------------
	49	+# Three-valued symbolic predicates
	50	+# ---------------------------------------------------------------------------
	51	+def statically_true(expr) -> bool:
	52	+ """Return True only if a bool/SymBool expression is provably true; never adds a guard.
	53	+
	54	+ Returns False for statically-false and undecidable cases. When the switch is
	55	+ off, only python bools are honored (symbols are treated as undecidable) so the
	56	+ whole layer degrades to static behavior. Prefers the official
	57	+ ``statically_known_true``; falls back to ShapeEnv static evaluation (also
	58	+ guard-free) on older torch.
	59	+ """
	60	+ if isinstance(expr, bool):
	61	+ return expr
	62	+ if not dynamic_fx_pass_enabled():
	63	+ return False
	64	+ if _skt is not None:
	65	+ try:
	66	+ return bool(_skt(expr))
	67	+ except Exception:
	68	+ pass
	69	+ return _static_eval_symbool(expr)
	70	+
	71	+
	72	+def _static_eval_symbool(expr) -> bool:
	73	+ """Fallback: statically evaluate a SymBool via its sympy expr + ShapeEnv, guard-free."""
	74	+ node = getattr(expr, "node", None)
	75	+ sym_expr = getattr(node, "expr", None)
	76	+ if sym_expr is None:
	77	+ return False
	78	+ try:
	79	+ import sympy
	80	+
	81	+ if sym_expr == sympy.true:
	82	+ return True
	83	+ shape_env = getattr(node, "shape_env", None)
	84	+ evaluator = getattr(shape_env, "_maybe_evaluate_static", None)
	85	+ if evaluator is None:
	86	+ return False
	87	+ result = evaluator(sym_expr)
	88	+ return result is not None and result == sympy.true
	89	+ except Exception:
	90	+ return False
	91	+
	92	+
	93	+def statically_known_eq(a: Number, b: Number) -> bool:
	94	+ """Whether a == b provably holds."""
	95	+ try:
	96	+ return statically_true(a == b)
	97	+ except Exception:
	98	+ return False
	99	+
	100	+
	101	+def statically_known_geq(a: Number, b: Number) -> bool:
	102	+ """Whether a >= b provably holds."""
	103	+ try:
	104	+ return statically_true(a >= b)
	105	+ except Exception:
	106	+ return False
	107	+
	108	+
	109	+def statically_known_gt(a: Number, b: Number) -> bool:
	110	+ """Whether a > b provably holds."""
	111	+ try:
	112	+ return statically_true(a > b)
	113	+ except Exception:
	114	+ return False
	115	+
	116	+
	117	+def statically_known_leq(a: Number, b: Number) -> bool:
	118	+ """Whether a <= b provably holds."""
	119	+ try:
	120	+ return statically_true(a <= b)
	121	+ except Exception:
	122	+ return False
	123	+
	124	+
	125	+def is_statically_one(d: Number) -> bool:
	126	+ """Whether a dim is provably 1."""
	127	+ return statically_known_eq(d, 1)
	128	+
	129	+
	130	+def has_free_symbols(shape) -> bool:
	131	+ """Whether a shape contains any symbolic dim."""
	132	+ if shape is None:
	133	+ return False
	134	+ return any(isinstance(d, torch.SymInt) for d in shape)
	135	+
	136	+
	137	+def shapes_statically_equal(s1, s2) -> bool:
	138	+ """Whether two shapes have equal rank and provably equal dims."""
	139	+ if s1 is None or s2 is None:
	140	+ return False
	141	+ if len(s1) != len(s2):
	142	+ return False
	143	+ return all(statically_known_eq(a, b) for a, b in zip(s1, s2))
	144	+
	145	+
	146	+# ---------------------------------------------------------------------------
	147	+# Symbol-aware shape reading
	148	+# ---------------------------------------------------------------------------
	149	+def get_symbolic_shape(node):
	150	+ """Read a node's shape keeping symbolic dims; None when no shape info.
	151	+
	152	+ Equivalent to ``get_binary_fold_result.get_node_shape(allow_symbolic=True)``,
	153	+ offered for call sites that do not import that module.
	154	+ """
	155	+ if not isinstance(node, torch.fx.Node):
	156	+ return None
	157	+ meta = node.meta
	158	+ val = meta.get("val", None)
	159	+ if val is not None and hasattr(val, "shape"):
	160	+ return val.shape
	161	+ example = meta.get("example_value", None)
	162	+ if example is not None and hasattr(example, "shape"):
	163	+ return example.shape
	164	+ tensor_meta = meta.get("tensor_meta", None)
	165	+ if tensor_meta is not None and hasattr(tensor_meta, "shape"):
	166	+ return tensor_meta.shape
	167	+ return None
	168	+
	169	+
	170	+# ---------------------------------------------------------------------------
	171	+# Size-argument normalization
	172	+# ---------------------------------------------------------------------------
	173	+def resolve_size_arg(arg) -> Optional[Number]:
	174	+ """Normalize one size argument to int / SymInt; None when unresolvable.
	175	+
	176	+ Handles four forms: python int, SymInt, ``sym_size``-like nodes (meta is
	177	+ SymInt/int), and PRE-graph derived nodes whose ``example_value`` is SymInt/int.
	178	+ """
	179	+ if isinstance(arg, bool):
	180	+ return None
	181	+ if isinstance(arg, int):
	182	+ return arg
	183	+ # Switch off: do not resolve symbolic sources (legacy "static int only").
	184	+ if not dynamic_fx_pass_enabled():
	185	+ return None
	186	+ if isinstance(arg, torch.SymInt):
	187	+ return arg
	188	+ if isinstance(arg, torch.fx.Node):
	189	+ for key in ("val", "example_value"):
	190	+ v = arg.meta.get(key, None)
	191	+ if isinstance(v, torch.SymInt):
	192	+ return v
	193	+ if isinstance(v, int) and not isinstance(v, bool):
	194	+ return v
	195	+ return None
	196	+
	197	+
	198	+def resolve_size_list(args) -> Optional[List[Number]]:
	199	+ """Normalize a size list element-wise; None if any element is unresolvable."""
	200	+ if not isinstance(args, (list, tuple)):
	201	+ return None
	202	+ resolved = []
	203	+ for a in args:
	204	+ r = resolve_size_arg(a)
	205	+ if r is None:
	206	+ return None
	207	+ resolved.append(r)
	208	+ return resolved
	209	+
	210	+
	211	+# ---------------------------------------------------------------------------
	212	+# Symbolic shape materialization
	213	+# ---------------------------------------------------------------------------
	214	+def _sym_key(value) -> Optional[str]:
	215	+ """Canonical string key for a size value (ints use their value), for equivalence."""
	216	+ if isinstance(value, bool):
	217	+ return None
	218	+ if isinstance(value, int):
	219	+ return f"int:{value}"
	220	+ if isinstance(value, torch.SymInt):
	221	+ try:
	222	+ return f"sym:{value.node.expr}"
	223	+ except Exception:
	224	+ return None
	225	+ return None
	226	+
	227	+
	228	+def _iter_anchor_nodes(anchors):
	229	+ if anchors is None:
	230	+ return
	231	+ if isinstance(anchors, torch.fx.Node):
	232	+ yield anchors
	233	+ return
	234	+ for a in anchors:
	235	+ if isinstance(a, torch.fx.Node):
	236	+ yield a
	237	+
	238	+
	239	+def materialize_shape(graph, shape, anchors) -> Optional[list]:
	240	+ """Turn a (possibly symbolic) target shape into a node-writable arg list.
	241	+
	242	+ - int dims are kept as-is;
	243	+ - SymInt dims: insert ``aten.sym_size.int`` referencing the same-valued dim of
	244	+ some anchor (an anchor is an input of the rewritten op, so it dominates the
	245	+ insertion point -> dominance-safe);
	246	+ - returns None if any SymInt dim has no source (rewrite is abandoned).
	247	+
	248	+ Caller must manage the insertion point via ``graph.inserting_before(node)``.
	249	+ """
	250	+ # Switch off: refuse to emit sym_size for symbolic dims -> degrade to legacy
	251	+ # static behavior (callers then skip the rewrite).
	252	+ if not dynamic_fx_pass_enabled() and has_free_symbols(shape):
	253	+ return None
	254	+ result = []
	255	+ created = {}
	256	+ for dim in shape:
	257	+ if isinstance(dim, bool):
	258	+ return None
	259	+ if isinstance(dim, int):
	260	+ result.append(dim)
	261	+ continue
	262	+ if not isinstance(dim, torch.SymInt):
	263	+ return None
	264	+ key = _sym_key(dim)
	265	+ if key is None:
	266	+ return None
	267	+ if key in created:
	268	+ result.append(created[key])
	269	+ continue
	270	+ node = _make_sym_size(graph, dim, key, anchors)
	271	+ if node is None:
	272	+ return None
	273	+ created[key] = node
	274	+ result.append(node)
	275	+ return result
	276	+
	277	+
	278	+def _make_sym_size(graph, sym, key, anchors):
	279	+ """Build a sym_size node for symbolic dim ``sym`` from a same-valued anchor dim."""
	280	+ for anchor in _iter_anchor_nodes(anchors):
	281	+ val = anchor.meta.get("val", None)
	282	+ if val is None or not hasattr(val, "shape"):
	283	+ continue
	284	+ for dim_idx, s in enumerate(val.shape):
	285	+ if isinstance(s, torch.SymInt) and _sym_key(s) == key:
	286	+ size_node = graph.call_function(
	287	+ torch.ops.aten.sym_size.int, args=(anchor, dim_idx)
	288	+ )
	289	+ size_node.meta["val"] = s
	290	+ return size_node
	291	+ return None
	292	+
	293	+
	294	+# ---------------------------------------------------------------------------
	295	+# Symbolic value-range analysis
	296	+# ---------------------------------------------------------------------------
	297	+_INT32_MIN = -(1 << 31)
	298	+_INT32_MAX = (1 << 31) - 1
	299	+
	300	+
	301	+def symbolic_value_range(v) -> Optional[Tuple]:
	302	+ """Provable compile-time range [lo, hi] of a value (ints give (v, v)); None if unknown.
	303	+
	304	+ Relies on ShapeEnv.bound_sympy (available on newer torch); returns None if absent.
	305	+ """
	306	+ if isinstance(v, bool):
	307	+ return None
	308	+ if isinstance(v, int):
	309	+ return (v, v)
	310	+ if isinstance(v, torch.SymInt):
	311	+ if not dynamic_fx_pass_enabled():
	312	+ return None
	313	+ try:
	314	+ node = v.node
	315	+ bound = getattr(node.shape_env, "bound_sympy", None)
	316	+ if bound is None:
	317	+ return None
	318	+ vr = bound(node.expr)
	319	+ return (vr.lower, vr.upper)
	320	+ except Exception:
	321	+ return None
	322	+ return None
	323	+
	324	+
	325	+def statically_fits_int32(*vals) -> bool:
	326	+ """Whether all given values provably lie within the int32 range.
	327	+
	328	+ Implemented with three-valued comparisons (no bound_sympy dependency); a
	329	+ symbolic value counts as fitting only when provably bounded.
	330	+ """
	331	+ if not vals:
	332	+ return False
	333	+ for v in vals:
	334	+ if isinstance(v, bool):
	335	+ return False
	336	+ if isinstance(v, int):
	337	+ if v < _INT32_MIN or v > _INT32_MAX:
	338	+ return False
	339	+ elif isinstance(v, torch.SymInt):
	340	+ if not (
	341	+ statically_known_leq(v, _INT32_MAX)
	342	+ and statically_known_geq(v, _INT32_MIN)
	343	+ ):
	344	+ return False
	345	+ else:
	346	+ return False
	347	+ return True
	348	+
	349	+
	350	+# ---------------------------------------------------------------------------
	351	+# Symbol-aware fake meta propagation
	352	+# ---------------------------------------------------------------------------
	353	+def _shape_env_from_val(val):
	354	+ if isinstance(val, torch.SymInt):
	355	+ try:
	356	+ return val.node.shape_env
	357	+ except Exception:
	358	+ return None
	359	+ fake_mode = getattr(val, "fake_mode", None)
	360	+ return getattr(fake_mode, "shape_env", None)
	361	+
	362	+
	363	+def get_shape_env(graph):
	364	+ """Resolve the ShapeEnv from any fake-carrying node in the graph; None if none."""
	365	+ for n in graph.nodes:
	366	+ for key in ("val", "example_value"):
	367	+ se = _shape_env_from_val(n.meta.get(key, None))
	368	+ if se is not None:
	369	+ return se
	370	+ return None
	371	+
	372	+
	373	+def get_fake_mode(graph):
	374	+ """Resolve the FakeTensorMode from any FakeTensor node in the graph; None if none."""
	375	+ for n in graph.nodes:
	376	+ val = n.meta.get("val", None)
	377	+ fake_mode = getattr(val, "fake_mode", None)
	378	+ if fake_mode is not None:
	379	+ return fake_mode
	380	+ return None
	381	+
	382	+
	383	+def _resolve_fake(arg):
	384	+ if isinstance(arg, torch.fx.Node):
	385	+ return arg.meta.get("val", arg)
	386	+ if isinstance(arg, (list, tuple)):
	387	+ return type(arg)(_resolve_fake(x) for x in arg)
	388	+ return arg
	389	+
	390	+
	391	+def refresh_fake_meta(node, fake_mode) -> bool:
	392	+ """Recompute node.meta['val'] from current args/kwargs under fake_mode; keep as-is on failure."""
	393	+ if fake_mode is None:
	394	+ return False
	395	+ try:
	396	+ with fake_mode:
	397	+ new_val = node.target(
	398	+ *[_resolve_fake(a) for a in node.args],
	399	+ **{k: _resolve_fake(v) for k, v in node.kwargs.items()},
	400	+ )
	401	+ node.meta["val"] = new_val
	402	+ return True
	403	+ except Exception:
	404	+ return False
	405	+
	406	+
	407	+# ---------------------------------------------------------------------------
	408	+# Lightweight hit/skip stats (debug-level only, avoids log spam)
	409	+# ---------------------------------------------------------------------------
	410	+_STATS = defaultdict(lambda: {"hit": 0, "skip": 0})
	411	+
	412	+
	413	+def note_hit(pass_name: str, n: int = 1) -> None:
	414	+ _STATS[pass_name]["hit"] += n
	415	+
	416	+
	417	+def note_skip(pass_name: str, n: int = 1) -> None:
	418	+ _STATS[pass_name]["skip"] += n
	419	+
	420	+
	421	+def dump_stats(pass_name: str) -> None:
	422	+ """Emit a single hit/skip stat line for a pass (debug level)."""
	423	+ stat = _STATS.get(pass_name)
	424	+ if stat and (stat["hit"] or stat["skip"]):
	425	+ log.debug(
	426	+ "[dynamic_fx] %s: hit=%d skip_undecidable=%d",
	427	+ pass_name,
	428	+ stat["hit"],
	429	+ stat["skip"],
	430	+ )
	431	+
	432	+
	433	+def reset_stats() -> None:
	434	+ """Clear stats, mainly for tests."""
	435	+ _STATS.clear()