fx_optimal UT Test-pytorch-AtomGit

Atest/_inductor/_pass_test_utils.py+92-0

@@ -0,0 +1,92 @@
	1	+"""
	2	+为 ascend_custom_passes 中各个 fx pass 提供共享测试工具。
	3	+
	4	+这些 pass 大多在 POST 阶段运行（即 AOT autograd 之后），其输入是一个含有
	5	+``meta['val']`` (FakeTensor) 的 aten-level FX 图。本模块提供一个轻量的
	6	+``GraphBuilder``：每调用一次 ``call`` 既追加一个 ``call_function`` 节点，
	7	+又在 ``FakeTensorMode`` 下重放该算子，把结果 FakeTensor 写入节点的
	8	+``meta['val']``，从而拼出一个与真实 pass 输入等价的 FX 图。
	9	+
	10	+测试只检查图结构变换是否符合预期，不依赖 NPU 硬件，可在 CPU 上运行。
	11	+"""
	12	+import torch
	13	+import torch.fx as fx
	14	+from torch._subclasses.fake_tensor import FakeTensorMode
	15	+
	16	+
	17	+def new_fake_mode():
	18	+ """构造一个新的 FakeTensorMode；与图节点共享同一 fake_mode 才能复用 FakeTensor。"""
	19	+ return FakeTensorMode()
	20	+
	21	+
	22	+def fake_from_real(fake_mode, real_tensor):
	23	+ """将真实 CPU tensor 转换为 fake_mode 下的 FakeTensor。"""
	24	+ return fake_mode.from_tensor(real_tensor)
	25	+
	26	+
	27	+def make_empty_fake(fake_mode, shape, dtype=torch.float32, device="cpu"):
	28	+ """直接在 fake_mode 中创建一个指定 shape/dtype 的 FakeTensor。"""
	29	+ with fake_mode:
	30	+ return torch.empty(shape, dtype=dtype, device=device)
	31	+
	32	+
	33	+class GraphBuilder:
	34	+ """轻量构图工具：自动把每个 call_function 节点的 ``meta['val']`` 设置为
	35	+ 在 ``fake_mode`` 下执行该算子的输出 FakeTensor。"""
	36	+
	37	+ def __init__(self, fake_mode):
	38	+ self.graph = fx.Graph()
	39	+ self.fake_mode = fake_mode
	40	+ self._placeholders = {}
	41	+
	42	+ def placeholder(self, name, fake_tensor):
	43	+ node = self.graph.placeholder(name)
	44	+ node.meta["val"] = fake_tensor
	45	+ self._placeholders[name] = node
	46	+ return node
	47	+
	48	+ def call(self, target, args=(), kwargs=None):
	49	+ kwargs = kwargs or {}
	50	+ node = self.graph.call_function(target, args=args, kwargs=kwargs)
	51	+
	52	+ def resolve(a):
	53	+ if isinstance(a, fx.Node):
	54	+ return a.meta.get("val", a)
	55	+ if isinstance(a, (list, tuple)):
	56	+ return type(a)(resolve(x) for x in a)
	57	+ return a
	58	+
	59	+ try:
	60	+ with self.fake_mode:
	61	+ node.meta["val"] = target(
	62	+ *[resolve(a) for a in args],
	63	+ **{k: resolve(v) for k, v in kwargs.items()},
	64	+ )
	65	+ except Exception:
	66	+ # 某些算子（如 prims.iota）需要走显式 kwargs，调用方可手动补 meta。
	67	+ pass
	68	+ return node
	69	+
	70	+ def output(self, value):
	71	+ self.graph.output(value)
	72	+
	73	+ def to_module(self):
	74	+ gm = fx.GraphModule(torch.nn.Module(), self.graph)
	75	+ return gm
	76	+
	77	+
	78	+def count_target(graph, target):
	79	+ """统计图中 target == ``target`` 的 call_function 节点数。"""
	80	+ return sum(
	81	+ 1 for n in graph.nodes
	82	+ if n.op == "call_function" and n.target is target
	83	+ )
	84	+
	85	+
	86	+def count_any_of(graph, targets):
	87	+ """统计图中 target 在 ``targets`` 集合内的 call_function 节点数。"""
	88	+ targets = tuple(targets)
	89	+ return sum(
	90	+ 1 for n in graph.nodes
	91	+ if n.op == "call_function" and n.target in targets
	92	+ )

Atest/_inductor/test_batch_embedding_fusion_pass.py+201-0

@@ -0,0 +1,201 @@
	1	+import torch
	2	+from torch.testing._internal.common_utils import TestCase, run_tests
	3	+
	4	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	5	+ batch_embedding_fusion_pass,
	6	+ _has_default_embedding_args,
	7	+ _symbolic_shape_key,
	8	+ _weight_node_key,
	9	+ _reduce_call_args,
	10	+ _detect_reduce_pattern,
	11	+ _detect_indices_parent,
	12	+)
	13	+
	14	+from _pass_test_utils import (
	15	+ GraphBuilder,
	16	+ count_target,
	17	+ new_fake_mode,
	18	+)
	19	+
	20	+
	21	+_EMBED = torch.ops.aten.embedding.default
	22	+_SLICE = torch.ops.aten.slice.Tensor
	23	+_SUM = torch.ops.aten.sum.dim_IntList
	24	+_PROD = torch.ops.aten.prod.dim_int
	25	+_RESHAPE = torch.ops.aten.reshape.default
	26	+
	27	+
	28	+def _build_batch_embedding(
	29	+ n_emb=2,
	30	+ V=10,
	31	+ D=4,
	32	+ L=3,
	33	+ keepdim=False,
	34	+ use_mixed_reduce=False,
	35	+ padding_idx=-1,
	36	+ indices_from_single_parent=True,
	37	+):
	38	+ """构建 n_emb 个 embedding + slice + sum 的图。"""
	39	+ fm = new_fake_mode()
	40	+ gb = GraphBuilder(fm)
	41	+ parent_len = n_emb * L
	42	+ with fm:
	43	+ w_fake = torch.empty((V, D), dtype=torch.float32)
	44	+ idx_fake = torch.empty((parent_len,), dtype=torch.int64).fill_(0)
	45	+ weight = gb.placeholder("weight", w_fake)
	46	+ parent = gb.placeholder("idx", idx_fake)
	47	+ reduces = []
	48	+ for i in range(n_emb):
	49	+ start = i * L
	50	+ end = start + L
	51	+ if indices_from_single_parent:
	52	+ src = parent
	53	+ else:
	54	+ # 第二个 embedding 使用一个独立的 placeholder 作为索引父节点
	55	+ with fm:
	56	+ other_idx = torch.empty((parent_len,), dtype=torch.int64)
	57	+ src = gb.placeholder(f"idx_alt_{i}", other_idx)
	58	+ sl = gb.call(_SLICE, args=(src, 0, start, end, 1))
	59	+ emb_args = (weight, sl)
	60	+ if padding_idx != -1:
	61	+ emb_args = (weight, sl, padding_idx)
	62	+ emb = gb.call(_EMBED, args=emb_args)
	63	+ if use_mixed_reduce and i == 1:
	64	+ red = gb.call(_PROD, args=(emb, 0, keepdim))
	65	+ else:
	66	+ red = gb.call(_SUM, args=(emb, [0], keepdim))
	67	+ reduces.append(red)
	68	+ gb.output(tuple(reduces))
	69	+ return gb.to_module(), weight, parent, reduces
	70	+
	71	+
	72	+class TestBatchEmbeddingFusionPass(TestCase):
	73	+ # ===== _has_default_embedding_args =====
	74	+ def test_has_default_embedding_args_default(self):
	75	+ gm, *_ = _build_batch_embedding(n_emb=1)
	76	+ emb = next(
	77	+ n for n in gm.graph.nodes
	78	+ if n.op == "call_function" and n.target is _EMBED
	79	+ )
	80	+ self.assertTrue(_has_default_embedding_args(emb))
	81	+
	82	+ def test_has_default_embedding_args_padding_idx(self):
	83	+ gm, *_ = _build_batch_embedding(n_emb=1, padding_idx=0)
	84	+ emb = next(
	85	+ n for n in gm.graph.nodes
	86	+ if n.op == "call_function" and n.target is _EMBED
	87	+ )
	88	+ self.assertFalse(_has_default_embedding_args(emb))
	89	+
	90	+ def test_has_default_embedding_args_kwargs(self):
	91	+ fm = new_fake_mode()
	92	+ gb = GraphBuilder(fm)
	93	+ with fm:
	94	+ w_fake = torch.empty((4, 4), dtype=torch.float32)
	95	+ idx_fake = torch.empty((4,), dtype=torch.int64)
	96	+ w = gb.placeholder("w", w_fake)
	97	+ idx = gb.placeholder("idx", idx_fake)
	98	+ emb = gb.graph.call_function(
	99	+ _EMBED, args=(w, idx), kwargs={"scale_grad_by_freq": True}
	100	+ )
	101	+ self.assertFalse(_has_default_embedding_args(emb))
	102	+
	103	+ # ===== helpers =====
	104	+ def test_symbolic_shape_key(self):
	105	+ self.assertEqual(_symbolic_shape_key([4, 8]), (4, 8))
	106	+
	107	+ def test_weight_node_key_placeholder(self):
	108	+ fm = new_fake_mode()
	109	+ gb = GraphBuilder(fm)
	110	+ ph = gb.placeholder("w", torch.empty((2,), dtype=torch.float32))
	111	+ key = _weight_node_key(ph)
	112	+ self.assertEqual(key, ("placeholder", "w"))
	113	+
	114	+ def test_weight_node_key_call_function(self):
	115	+ fm = new_fake_mode()
	116	+ gb = GraphBuilder(fm)
	117	+ x = gb.placeholder("x", torch.empty((2,), dtype=torch.float32))
	118	+ c = gb.call(torch.ops.aten.relu.default, args=(x,))
	119	+ self.assertEqual(_weight_node_key(c), id(c))
	120	+
	121	+ def test_reduce_call_args_dim_list(self):
	122	+ self.assertEqual(_reduce_call_args(_SUM, "x", 1), ("x", [1]))
	123	+
	124	+ def test_reduce_call_args_dim_int(self):
	125	+ self.assertEqual(_reduce_call_args(_PROD, "x", 1), ("x", 1))
	126	+
	127	+ def test_detect_reduce_pattern_consistent_sum(self):
	128	+ gm, *_ = _build_batch_embedding(n_emb=2)
	129	+ embs = [
	130	+ n for n in gm.graph.nodes
	131	+ if n.op == "call_function" and n.target is _EMBED
	132	+ ]
	133	+ result = _detect_reduce_pattern(embs, cat_dim=0)
	134	+ self.assertIsNotNone(result)
	135	+ target, _ = result
	136	+ self.assertIs(target, _SUM)
	137	+
	138	+ def test_detect_reduce_pattern_mixed_skipped(self):
	139	+ gm, *_ = _build_batch_embedding(n_emb=2, use_mixed_reduce=True)
	140	+ embs = [
	141	+ n for n in gm.graph.nodes
	142	+ if n.op == "call_function" and n.target is _EMBED
	143	+ ]
	144	+ self.assertIsNone(_detect_reduce_pattern(embs, cat_dim=0))
	145	+
	146	+ def test_detect_reduce_pattern_keepdim_skipped(self):
	147	+ gm, *_ = _build_batch_embedding(n_emb=2, keepdim=True)
	148	+ embs = [
	149	+ n for n in gm.graph.nodes
	150	+ if n.op == "call_function" and n.target is _EMBED
	151	+ ]
	152	+ self.assertIsNone(_detect_reduce_pattern(embs, cat_dim=0))
	153	+
	154	+ def test_detect_indices_parent_same_parent(self):
	155	+ gm, _, parent, _ = _build_batch_embedding(n_emb=2)
	156	+ embs = [
	157	+ n for n in gm.graph.nodes
	158	+ if n.op == "call_function" and n.target is _EMBED
	159	+ ]
	160	+ p, dim = _detect_indices_parent(embs)
	161	+ self.assertIs(p, parent)
	162	+ self.assertEqual(dim, 0)
	163	+
	164	+ def test_detect_indices_parent_different_parents(self):
	165	+ gm, *_ = _build_batch_embedding(n_emb=2, indices_from_single_parent=False)
	166	+ embs = [
	167	+ n for n in gm.graph.nodes
	168	+ if n.op == "call_function" and n.target is _EMBED
	169	+ ]
	170	+ p, dim = _detect_indices_parent(embs)
	171	+ self.assertIsNone(p)
	172	+
	173	+ # ===== full pass =====
	174	+ def test_two_embeddings_fused(self):
	175	+ gm, weight, parent, reduces = _build_batch_embedding(n_emb=2, V=10, D=4, L=3)
	176	+ n_embed_before = count_target(gm.graph, _EMBED)
	177	+ batch_embedding_fusion_pass(gm.graph)
	178	+ n_embed_after = count_target(gm.graph, _EMBED)
	179	+ self.assertLess(n_embed_after, n_embed_before)
	180	+
	181	+ def test_single_embedding_unchanged(self):
	182	+ gm, *_ = _build_batch_embedding(n_emb=1)
	183	+ orig = str(gm.graph)
	184	+ batch_embedding_fusion_pass(gm.graph)
	185	+ self.assertEqual(orig, str(gm.graph))
	186	+
	187	+ def test_non_default_args_skipped(self):
	188	+ gm, *_ = _build_batch_embedding(n_emb=2, padding_idx=0)
	189	+ n_before = count_target(gm.graph, _EMBED)
	190	+ batch_embedding_fusion_pass(gm.graph)
	191	+ self.assertEqual(count_target(gm.graph, _EMBED), n_before)
	192	+
	193	+ def test_mixed_reduce_skipped(self):
	194	+ gm, *_ = _build_batch_embedding(n_emb=2, use_mixed_reduce=True)
	195	+ n_before = count_target(gm.graph, _EMBED)
	196	+ batch_embedding_fusion_pass(gm.graph)
	197	+ self.assertEqual(count_target(gm.graph, _EMBED), n_before)
	198	+
	199	+
	200	+if __name__ == "__main__":
	201	+ run_tests()

Atest/_inductor/test_bool_cast_mul_to_where_pass.py+141-0

@@ -0,0 +1,141 @@
	1	+import torch
	2	+from torch.testing._internal.common_utils import TestCase, run_tests
	3	+
	4	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	5	+ bool_cast_mul_to_where_pass,
	6	+ _walk_back_view_chain_to_cast,
	7	+)
	8	+
	9	+from _pass_test_utils import (
	10	+ GraphBuilder,
	11	+ count_target,
	12	+ new_fake_mode,
	13	+)
	14	+
	15	+
	16	+_MUL = torch.ops.aten.mul.Tensor
	17	+_WHERE = torch.ops.aten.where.self
	18	+_CAST = torch.ops.prims.convert_element_type.default
	19	+_VIEW = torch.ops.aten.view.default
	20	+_UNSQUEEZE = torch.ops.aten.unsqueeze.default
	21	+
	22	+
	23	+def _build_bool_cast_mul(
	24	+ shape=(4, 4),
	25	+ target_dtype=torch.float32,
	26	+ other_dtype=None,
	27	+ insert_view=False,
	28	+ bool_src=True,
	29	+):
	30	+ other_dtype = other_dtype or target_dtype
	31	+ fm = new_fake_mode()
	32	+ gb = GraphBuilder(fm)
	33	+ with fm:
	34	+ mask_fake = torch.empty(shape, dtype=torch.bool if bool_src else torch.int32)
	35	+ other_fake = torch.empty(shape, dtype=other_dtype)
	36	+ mask = gb.placeholder("mask", mask_fake)
	37	+ other = gb.placeholder("other", other_fake)
	38	+ cast = gb.call(_CAST, args=(mask, target_dtype))
	39	+ after_cast = cast
	40	+ if insert_view:
	41	+ after_cast = gb.call(_VIEW, args=(cast, list(shape)))
	42	+ mul = gb.call(_MUL, args=(after_cast, other))
	43	+ gb.output(mul)
	44	+ return gb.to_module(), mask, other, cast, mul
	45	+
	46	+
	47	+class TestBoolCastMulToWherePass(TestCase):
	48	+ # ===== _walk_back_view_chain_to_cast =====
	49	+ def test_walk_back_direct_cast(self):
	50	+ gm, mask, other, cast, mul = _build_bool_cast_mul()
	51	+ chain, found = _walk_back_view_chain_to_cast(cast)
	52	+ self.assertEqual(chain, [])
	53	+ self.assertIs(found, cast)
	54	+
	55	+ def test_walk_back_through_view(self):
	56	+ gm, mask, other, cast, mul = _build_bool_cast_mul(insert_view=True)
	57	+ view = next(
	58	+ n for n in gm.graph.nodes
	59	+ if n.op == "call_function" and n.target is _VIEW
	60	+ )
	61	+ chain, found = _walk_back_view_chain_to_cast(view)
	62	+ self.assertEqual(len(chain), 1)
	63	+ self.assertIs(found, cast)
	64	+
	65	+ def test_walk_back_no_cast(self):
	66	+ fm = new_fake_mode()
	67	+ gb = GraphBuilder(fm)
	68	+ with fm:
	69	+ x_fake = torch.empty((4,), dtype=torch.float32)
	70	+ x = gb.placeholder("x", x_fake)
	71	+ v = gb.call(_VIEW, args=(x, [4]))
	72	+ chain, found = _walk_back_view_chain_to_cast(v)
	73	+ self.assertIsNone(found)
	74	+
	75	+ # ===== full pass =====
	76	+ def test_basic_rewrite(self):
	77	+ gm, mask, other, cast, mul = _build_bool_cast_mul()
	78	+ bool_cast_mul_to_where_pass(gm.graph)
	79	+ self.assertEqual(count_target(gm.graph, _MUL), 0)
	80	+ self.assertEqual(count_target(gm.graph, _WHERE), 1)
	81	+ w = next(
	82	+ n for n in gm.graph.nodes
	83	+ if n.op == "call_function" and n.target is _WHERE
	84	+ )
	85	+ self.assertIs(w.args[0], mask)
	86	+ self.assertIs(w.args[1], other)
	87	+
	88	+ def test_with_view_chain(self):
	89	+ gm, mask, other, cast, mul = _build_bool_cast_mul(insert_view=True)
	90	+ bool_cast_mul_to_where_pass(gm.graph)
	91	+ self.assertEqual(count_target(gm.graph, _MUL), 0)
	92	+ self.assertEqual(count_target(gm.graph, _WHERE), 1)
	93	+
	94	+ def test_non_bool_source_skipped(self):
	95	+ gm, *_ = _build_bool_cast_mul(bool_src=False)
	96	+ bool_cast_mul_to_where_pass(gm.graph)
	97	+ self.assertEqual(count_target(gm.graph, _MUL), 1)
	98	+ self.assertEqual(count_target(gm.graph, _WHERE), 0)
	99	+
	100	+ def test_dtype_mismatch_skipped(self):
	101	+ """cast 的目标 dtype 与 other 的 dtype 不同 → 不应重写。"""
	102	+ gm, *_ = _build_bool_cast_mul(target_dtype=torch.float16, other_dtype=torch.float32)
	103	+ bool_cast_mul_to_where_pass(gm.graph)
	104	+ self.assertEqual(count_target(gm.graph, _MUL), 1)
	105	+ self.assertEqual(count_target(gm.graph, _WHERE), 0)
	106	+
	107	+ def test_cast_node_with_extra_user_skipped(self):
	108	+ """cast 节点有不只一个用户 → 不应重写。"""
	109	+ fm = new_fake_mode()
	110	+ gb = GraphBuilder(fm)
	111	+ with fm:
	112	+ mask_fake = torch.empty((4,), dtype=torch.bool)
	113	+ other_fake = torch.empty((4,), dtype=torch.float32)
	114	+ mask = gb.placeholder("mask", mask_fake)
	115	+ other = gb.placeholder("other", other_fake)
	116	+ cast = gb.call(_CAST, args=(mask, torch.float32))
	117	+ mul = gb.call(_MUL, args=(cast, other))
	118	+ # 让 cast 同时被另一个算子使用
	119	+ relu = gb.call(torch.ops.aten.relu.default, args=(cast,))
	120	+ gb.output((mul, relu))
	121	+ gm = gb.to_module()
	122	+ bool_cast_mul_to_where_pass(gm.graph)
	123	+ self.assertEqual(count_target(gm.graph, _MUL), 1)
	124	+ self.assertEqual(count_target(gm.graph, _WHERE), 0)
	125	+
	126	+ def test_mul_with_scalar_skipped(self):
	127	+ """mul 一侧为标量 → 找不到 cast→bool 模式，应跳过。"""
	128	+ fm = new_fake_mode()
	129	+ gb = GraphBuilder(fm)
	130	+ with fm:
	131	+ x_fake = torch.empty((4,), dtype=torch.float32)
	132	+ x = gb.placeholder("x", x_fake)
	133	+ mul = gb.call(_MUL, args=(x, 2.0))
	134	+ gb.output(mul)
	135	+ gm = gb.to_module()
	136	+ bool_cast_mul_to_where_pass(gm.graph)
	137	+ self.assertEqual(count_target(gm.graph, _MUL), 1)
	138	+
	139	+
	140	+if __name__ == "__main__":
	141	+ run_tests()

Atest/_inductor/test_broadcast_const_mask_compress_pass.py+131-0

@@ -0,0 +1,131 @@
	1	+import torch
	2	+from torch.testing._internal.common_utils import TestCase, run_tests
	3	+
	4	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	5	+ broadcast_const_mask_compress,
	6	+ _extract_const_full_scalar,
	7	+)
	8	+
	9	+from _pass_test_utils import (
	10	+ GraphBuilder,
	11	+ count_target,
	12	+ new_fake_mode,
	13	+)
	14	+
	15	+
	16	+_WHERE = torch.ops.aten.where.self
	17	+_FULL = torch.ops.aten.full.default
	18	+_CAST = torch.ops.prims.convert_element_type.default
	19	+_NOT = torch.ops.aten.logical_not.default
	20	+_ADD = torch.ops.aten.add.Tensor
	21	+
	22	+
	23	+def _build_cast_where_full_full(
	24	+ shape=(4, 4),
	25	+ t_val=1,
	26	+ f_val=0,
	27	+ target_dtype=torch.float32,
	28	+ mask_is_bool=True,
	29	+):
	30	+ fm = new_fake_mode()
	31	+ gb = GraphBuilder(fm)
	32	+ # 对 bool dtype 的 full，使用 Python bool 字面量，避免某些版本不接受 0/1。
	33	+ fill_dtype = torch.float32 if target_dtype == torch.bool else target_dtype
	34	+ t_payload = bool(t_val) if target_dtype == torch.bool else t_val
	35	+ f_payload = bool(f_val) if target_dtype == torch.bool else f_val
	36	+ with fm:
	37	+ cond_fake = torch.empty(shape, dtype=torch.bool if mask_is_bool else torch.int32)
	38	+ full_t_fake = torch.ops.aten.full.default(list(shape), t_payload, dtype=fill_dtype)
	39	+ full_f_fake = torch.ops.aten.full.default(list(shape), f_payload, dtype=fill_dtype)
	40	+ cond = gb.placeholder("cond", cond_fake)
	41	+ full_t = gb.graph.call_function(
	42	+ _FULL, args=(list(shape), t_val), kwargs={"dtype": fill_dtype}
	43	+ )
	44	+ full_t.meta["val"] = full_t_fake
	45	+ full_f = gb.graph.call_function(
	46	+ _FULL, args=(list(shape), f_val), kwargs={"dtype": fill_dtype}
	47	+ )
	48	+ full_f.meta["val"] = full_f_fake
	49	+ w = gb.call(_WHERE, args=(cond, full_t, full_f))
	50	+ cast = gb.call(_CAST, args=(w, target_dtype))
	51	+ out = gb.call(_ADD, args=(cast, cast))
	52	+ gb.output(out)
	53	+ return gb.to_module(), cond, w, cast
	54	+
	55	+
	56	+class TestBroadcastConstMaskCompress(TestCase):
	57	+ def test_mask_full1_full0_replaces_with_mask(self):
	58	+ gm, cond, w, cast = _build_cast_where_full_full(t_val=1, f_val=0)
	59	+ broadcast_const_mask_compress(gm.graph)
	60	+ self.assertEqual(count_target(gm.graph, _NOT), 0)
	61	+ self.assertIs(cast.args[0], cond)
	62	+
	63	+ def test_mask_full0_full1_inserts_logical_not(self):
	64	+ gm, cond, w, cast = _build_cast_where_full_full(t_val=0, f_val=1)
	65	+ broadcast_const_mask_compress(gm.graph)
	66	+ self.assertEqual(count_target(gm.graph, _NOT), 1)
	67	+ new_cond = next(
	68	+ n for n in gm.graph.nodes
	69	+ if n.op == "call_function" and n.target is _NOT
	70	+ )
	71	+ self.assertIs(cast.args[0], new_cond)
	72	+
	73	+ def test_target_dtype_bool_drops_cast(self):
	74	+ gm, cond, w, cast = _build_cast_where_full_full(
	75	+ t_val=1, f_val=0, target_dtype=torch.bool
	76	+ )
	77	+ broadcast_const_mask_compress(gm.graph)
	78	+ self.assertEqual(count_target(gm.graph, _CAST), 0)
	79	+ out_consumer = next(
	80	+ n for n in gm.graph.nodes
	81	+ if n.op == "call_function" and n.target is _ADD
	82	+ )
	83	+ self.assertIs(out_consumer.args[0], cond)
	84	+
	85	+ def test_target_dtype_non_bool_non_01_skipped(self):
	86	+ gm, cond, w, cast = _build_cast_where_full_full(t_val=5, f_val=3)
	87	+ broadcast_const_mask_compress(gm.graph)
	88	+ self.assertEqual(count_target(gm.graph, _WHERE), 1)
	89	+
	90	+ def test_equal_constants_skipped(self):
	91	+ gm, cond, w, cast = _build_cast_where_full_full(t_val=1, f_val=1)
	92	+ broadcast_const_mask_compress(gm.graph)
	93	+ self.assertEqual(count_target(gm.graph, _WHERE), 1)
	94	+
	95	+ def test_mask_not_bool_skipped(self):
	96	+ gm, cond, w, cast = _build_cast_where_full_full(
	97	+ t_val=1, f_val=0, mask_is_bool=False
	98	+ )
	99	+ broadcast_const_mask_compress(gm.graph)
	100	+ self.assertEqual(count_target(gm.graph, _WHERE), 1)
	101	+
	102	+ def test_where_with_non_full_branch_skipped(self):
	103	+ fm = new_fake_mode()
	104	+ gb = GraphBuilder(fm)
	105	+ with fm:
	106	+ cond_fake = torch.empty((4,), dtype=torch.bool)
	107	+ x_fake = torch.empty((4,), dtype=torch.float32)
	108	+ cond = gb.placeholder("cond", cond_fake)
	109	+ x = gb.placeholder("x", x_fake)
	110	+ full_f = gb.graph.call_function(_FULL, args=([4], 0), kwargs={"dtype": torch.float32})
	111	+ with fm:
	112	+ full_f.meta["val"] = torch.ops.aten.full.default([4], 0, dtype=torch.float32)
	113	+ w = gb.call(_WHERE, args=(cond, x, full_f))
	114	+ cast = gb.call(_CAST, args=(w, torch.float32))
	115	+ gb.output(cast)
	116	+ gm = gb.to_module()
	117	+ broadcast_const_mask_compress(gm.graph)
	118	+ self.assertEqual(count_target(gm.graph, _WHERE), 1)
	119	+
	120	+ def test_extract_const_full_scalar_helpers(self):
	121	+ fm = new_fake_mode()
	122	+ gb = GraphBuilder(fm)
	123	+ f = gb.graph.call_function(_FULL, args=([4], 7), kwargs={"dtype": torch.float32})
	124	+ self.assertEqual(_extract_const_full_scalar(f), 7)
	125	+ x = gb.placeholder("x", torch.empty((4,), dtype=torch.float32))
	126	+ self.assertIsNone(_extract_const_full_scalar(x))
	127	+ self.assertIsNone(_extract_const_full_scalar(42))
	128	+
	129	+
	130	+if __name__ == "__main__":
	131	+ run_tests()

Atest/_inductor/test_cat_to_view_pass.py+176-0

@@ -0,0 +1,176 @@
	1	+import torch
	2	+import torch.fx as fx
	3	+from torch.testing._internal.common_utils import TestCase, run_tests
	4	+
	5	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	6	+ cat_to_view_pass,
	7	+)
	8	+
	9	+from _pass_test_utils import (
	10	+ GraphBuilder,
	11	+ count_target,
	12	+ new_fake_mode,
	13	+)
	14	+
	15	+
	16	+_CAT = torch.ops.aten.cat.default
	17	+_SLICE = torch.ops.aten.slice.Tensor
	18	+_ROLL = torch.ops.aten.roll.default
	19	+
	20	+
	21	+def _build_cat_of_slices(intervals, dim=1, shape=(2, 6, 4), other_parent_idx=None):
	22	+ """构造一个 cat([slice_0, slice_1, ...]) 图。
	23	+
	24	+ intervals: List[(start, end)]，按列出顺序填进 cat 的输入。
	25	+ other_parent_idx: 若非 None，则该索引位置的 slice 的 parent 替换成另一个 placeholder。
	26	+ """
	27	+ fm = new_fake_mode()
	28	+ gb = GraphBuilder(fm)
	29	+ with fm:
	30	+ parent_fake = torch.empty(shape, dtype=torch.float32)
	31	+ other_fake = torch.empty(shape, dtype=torch.float32)
	32	+ parent = gb.placeholder("parent", parent_fake)
	33	+ other = gb.placeholder("other", other_fake)
	34	+ slice_nodes = []
	35	+ for i, (s, e) in enumerate(intervals):
	36	+ src = other if other_parent_idx == i else parent
	37	+ n = gb.call(_SLICE, args=(src, dim, s, e, 1))
	38	+ slice_nodes.append(n)
	39	+ cat = gb.call(_CAT, args=(slice_nodes, dim))
	40	+ gb.output(cat)
	41	+ return gb.to_module(), parent, cat
	42	+
	43	+
	44	+class TestCatToViewPass(TestCase):
	45	+ def test_identity_full_cover(self):
	46	+ """连续 [0,2)+[2,4)+[4,6) 完整覆盖 dim=1 → cat 应被替换为 parent。"""
	47	+ gm, parent, cat = _build_cat_of_slices(
	48	+ [(0, 2), (2, 4), (4, 6)], dim=1, shape=(2, 6, 4)
	49	+ )
	50	+ cat_to_view_pass(gm.graph)
	51	+ self.assertEqual(count_target(gm.graph, _CAT), 0)
	52	+ self.assertEqual(count_target(gm.graph, _ROLL), 0)
	53	+ out_node = next(n for n in gm.graph.nodes if n.op == "output")
	54	+ self.assertIs(out_node.args[0], parent)
	55	+
	56	+ def test_rotation_full_cover(self):
	57	+ """[2,4)+[4,6)+[0,2) 完整覆盖但是循环位移 → 应插入 roll 节点。"""
	58	+ gm, _, _ = _build_cat_of_slices(
	59	+ [(2, 4), (4, 6), (0, 2)], dim=1, shape=(2, 6, 4)
	60	+ )
	61	+ cat_to_view_pass(gm.graph)
	62	+ self.assertEqual(count_target(gm.graph, _CAT), 0)
	63	+ self.assertEqual(count_target(gm.graph, _ROLL), 1)
	64	+ roll = next(
	65	+ n for n in gm.graph.nodes
	66	+ if n.op == "call_function" and n.target is _ROLL
	67	+ )
	68	+ self.assertEqual(list(roll.args[1]), [-2])
	69	+ self.assertEqual(list(roll.args[2]), [1])
	70	+
	71	+ def test_partial_cover_skipped(self):
	72	+ """[0,2)+[2,4) 仅部分覆盖 dim=1（size=6）→ 不应折叠。"""
	73	+ gm, _, _ = _build_cat_of_slices(
	74	+ [(0, 2), (2, 4)], dim=1, shape=(2, 6, 4)
	75	+ )
	76	+ cat_to_view_pass(gm.graph)
	77	+ self.assertEqual(count_target(gm.graph, _CAT), 1)
	78	+
	79	+ def test_different_parents_skipped(self):
	80	+ """不同 parent 的 slice → 不应折叠。"""
	81	+ gm, _, _ = _build_cat_of_slices(
	82	+ [(0, 3), (3, 6)], dim=1, shape=(2, 6, 4), other_parent_idx=1
	83	+ )
	84	+ cat_to_view_pass(gm.graph)
	85	+ self.assertEqual(count_target(gm.graph, _CAT), 1)
	86	+
	87	+ def test_step_not_one_skipped(self):
	88	+ """slice step != 1 → 不应折叠。"""
	89	+ fm = new_fake_mode()
	90	+ gb = GraphBuilder(fm)
	91	+ with fm:
	92	+ parent_fake = torch.empty((2, 6), dtype=torch.float32)
	93	+ parent = gb.placeholder("parent", parent_fake)
	94	+ s0 = gb.call(_SLICE, args=(parent, 1, 0, 6, 2))
	95	+ s1 = gb.call(_SLICE, args=(parent, 1, 1, 6, 2))
	96	+ cat = gb.call(_CAT, args=([s0, s1], 1))
	97	+ gb.output(cat)
	98	+ gm = gb.to_module()
	99	+ cat_to_view_pass(gm.graph)
	100	+ self.assertEqual(count_target(gm.graph, _CAT), 1)
	101	+
	102	+ def test_non_int_dim_skipped(self):
	103	+ """cat dim 非 int（例如 SymInt 之类）→ 直接跳过。"""
	104	+ gm, _, cat = _build_cat_of_slices(
	105	+ [(0, 3), (3, 6)], dim=1, shape=(2, 6)
	106	+ )
	107	+ cat.args = (cat.args[0], "bad")
	108	+ cat_to_view_pass(gm.graph)
	109	+ self.assertEqual(count_target(gm.graph, _CAT), 1)
	110	+
	111	+ def test_non_cat_node_untouched(self):
	112	+ """图中没有 cat 节点 → pass 不修改图。"""
	113	+ fm = new_fake_mode()
	114	+ gb = GraphBuilder(fm)
	115	+ with fm:
	116	+ x_fake = torch.empty((4, 4), dtype=torch.float32)
	117	+ x = gb.placeholder("x", x_fake)
	118	+ add = gb.call(torch.ops.aten.add.Tensor, args=(x, x))
	119	+ gb.output(add)
	120	+ gm = gb.to_module()
	121	+ orig = str(gm.graph)
	122	+ cat_to_view_pass(gm.graph)
	123	+ self.assertEqual(orig, str(gm.graph))
	124	+
	125	+ def test_negative_dim_normalised(self):
	126	+ """dim=-1 应被规范化为 rank-1 并正确进入 identity 路径。"""
	127	+ fm = new_fake_mode()
	128	+ gb = GraphBuilder(fm)
	129	+ with fm:
	130	+ parent_fake = torch.empty((2, 6), dtype=torch.float32)
	131	+ parent = gb.placeholder("parent", parent_fake)
	132	+ s0 = gb.call(_SLICE, args=(parent, -1, 0, 3, 1))
	133	+ s1 = gb.call(_SLICE, args=(parent, -1, 3, 6, 1))
	134	+ cat = gb.call(_CAT, args=([s0, s1], -1))
	135	+ gb.output(cat)
	136	+ gm = gb.to_module()
	137	+ cat_to_view_pass(gm.graph)
	138	+ self.assertEqual(count_target(gm.graph, _CAT), 0)
	139	+ out_node = next(n for n in gm.graph.nodes if n.op == "output")
	140	+ self.assertIs(out_node.args[0], parent)
	141	+
	142	+ def test_dim_via_kwargs(self):
	143	+ """cat 用 kwargs 传 dim 也应能正确处理（默认 dim=0）。"""
	144	+ fm = new_fake_mode()
	145	+ gb = GraphBuilder(fm)
	146	+ with fm:
	147	+ parent_fake = torch.empty((6, 2), dtype=torch.float32)
	148	+ parent = gb.placeholder("parent", parent_fake)
	149	+ s0 = gb.call(_SLICE, args=(parent, 0, 0, 3, 1))
	150	+ s1 = gb.call(_SLICE, args=(parent, 0, 3, 6, 1))
	151	+ # 仅传 inputs 列表，不传 dim → 走 kwargs.get("dim", 0) 分支
	152	+ cat = gb.call(_CAT, args=([s0, s1],))
	153	+ gb.output(cat)
	154	+ gm = gb.to_module()
	155	+ cat_to_view_pass(gm.graph)
	156	+ self.assertEqual(count_target(gm.graph, _CAT), 0)
	157	+
	158	+ def test_cat_with_non_slice_input_skipped(self):
	159	+ """cat 输入不全是 slice → 不应折叠。"""
	160	+ fm = new_fake_mode()
	161	+ gb = GraphBuilder(fm)
	162	+ with fm:
	163	+ parent_fake = torch.empty((2, 6), dtype=torch.float32)
	164	+ parent = gb.placeholder("parent", parent_fake)
	165	+ s0 = gb.call(_SLICE, args=(parent, 1, 0, 3, 1))
	166	+ not_a_slice = gb.call(torch.ops.aten.relu.default, args=(parent,))
	167	+ s1 = gb.call(_SLICE, args=(not_a_slice, 1, 3, 6, 1))
	168	+ cat = gb.call(_CAT, args=([s0, s1], 1))
	169	+ gb.output(cat)
	170	+ gm = gb.to_module()
	171	+ cat_to_view_pass(gm.graph)
	172	+ self.assertEqual(count_target(gm.graph, _CAT), 1)
	173	+
	174	+
	175	+if __name__ == "__main__":
	176	+ run_tests()

Atest/_inductor/test_fold_iota_arithmetic_pass.py+270-0

@@ -0,0 +1,270 @@
	1	+import torch
	2	+from torch.testing._internal.common_utils import TestCase, run_tests
	3	+
	4	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	5	+ fold_iota_arithmetic_pass,
	6	+ _prims_iota_value_range,
	7	+ _collect_iota_downcast_closure,
	8	+ _hashable_const_key,
	9	+ _cse_constant_call,
	10	+ _refresh_fake_meta,
	11	+)
	12	+
	13	+from _pass_test_utils import (
	14	+ GraphBuilder,
	15	+ count_target,
	16	+ new_fake_mode,
	17	+)
	18	+
	19	+
	20	+_IOTA = torch.ops.prims.iota.default
	21	+_FULL = torch.ops.aten.full.default
	22	+_SUB = torch.ops.aten.sub.Tensor
	23	+_GE_S = torch.ops.aten.ge.Scalar
	24	+_GT_S = torch.ops.aten.gt.Scalar
	25	+_LT_S = torch.ops.aten.lt.Scalar
	26	+_LE_S = torch.ops.aten.le.Scalar
	27	+_EQ_S = torch.ops.aten.eq.Scalar
	28	+_NE_S = torch.ops.aten.ne.Scalar
	29	+_GE_T = torch.ops.aten.ge.Tensor
	30	+_GT_T = torch.ops.aten.gt.Tensor
	31	+
	32	+
	33	+def _build_iota_with_user(length, start=0, step=1, dtype=torch.int64,
	34	+ user_target=torch.ops.aten.add.Tensor, user_rhs_shape=None):
	35	+ """构造 iota → transparent_op → closing_op(_to_copy) 的图，
	36	+ 满足 fold_iota_arithmetic_pass 的闭包检测要求。"""
	37	+ fm = new_fake_mode()
	38	+ gb = GraphBuilder(fm)
	39	+ with fm:
	40	+ iota_val = torch.ops.prims.iota.default(
	41	+ length, start=start, step=step, dtype=dtype, device="cpu", requires_grad=False
	42	+ )
	43	+ rhs_shape = user_rhs_shape or (length,)
	44	+ rhs_val = torch.empty(rhs_shape, dtype=dtype)
	45	+ iota = gb.graph.call_function(
	46	+ _IOTA,
	47	+ args=(length,),
	48	+ kwargs={
	49	+ "start": start,
	50	+ "step": step,
	51	+ "dtype": dtype,
	52	+ "device": torch.device("cpu"),
	53	+ "requires_grad": False,
	54	+ },
	55	+ )
	56	+ iota.meta["val"] = iota_val
	57	+ rhs = gb.placeholder("rhs", rhs_val)
	58	+ middle = gb.call(user_target, args=(iota, rhs))
	59	+ # 最终下游必须是 closing op（如 convert_element_type），否则闭包检测会拒绝。
	60	+ closing = gb.call(
	61	+ torch.ops.prims.convert_element_type.default,
	62	+ args=(middle, torch.float32),
	63	+ )
	64	+ gb.output(closing)
	65	+ return gb.to_module(), iota
	66	+
	67	+
	68	+class TestFoldIotaArithmeticPass(TestCase):
	69	+ # ===== _prims_iota_value_range =====
	70	+ def test_iota_value_range_positive_step(self):
	71	+ fm = new_fake_mode()
	72	+ gb = GraphBuilder(fm)
	73	+ n = gb.graph.call_function(
	74	+ _IOTA, args=(5,),
	75	+ kwargs={"start": 3, "step": 2, "dtype": torch.int64,
	76	+ "device": torch.device("cpu"), "requires_grad": False},
	77	+ )
	78	+ self.assertEqual(_prims_iota_value_range(n), (3, 12))
	79	+
	80	+ def test_iota_value_range_negative_step(self):
	81	+ """负 step：序列 [10, 7, 4] → range=(4, 11)。"""
	82	+ fm = new_fake_mode()
	83	+ gb = GraphBuilder(fm)
	84	+ n = gb.graph.call_function(
	85	+ _IOTA, args=(3,),
	86	+ kwargs={"start": 10, "step": -3, "dtype": torch.int64,
	87	+ "device": torch.device("cpu"), "requires_grad": False},
	88	+ )
	89	+ self.assertEqual(_prims_iota_value_range(n), (4, 11))
	90	+
	91	+ def test_iota_value_range_zero_length(self):
	92	+ fm = new_fake_mode()
	93	+ gb = GraphBuilder(fm)
	94	+ n = gb.graph.call_function(
	95	+ _IOTA, args=(0,),
	96	+ kwargs={"start": 5, "step": 1, "dtype": torch.int64,
	97	+ "device": torch.device("cpu"), "requires_grad": False},
	98	+ )
	99	+ self.assertEqual(_prims_iota_value_range(n), (5, 5))
	100	+
	101	+ def test_iota_value_range_non_constant(self):
	102	+ fm = new_fake_mode()
	103	+ gb = GraphBuilder(fm)
	104	+ n = gb.graph.call_function(
	105	+ _IOTA, args=("nonint",),
	106	+ kwargs={"start": 0, "step": 1, "dtype": torch.int64},
	107	+ )
	108	+ self.assertIsNone(_prims_iota_value_range(n))
	109	+
	110	+ def test_iota_value_range_not_iota(self):
	111	+ fm = new_fake_mode()
	112	+ gb = GraphBuilder(fm)
	113	+ x = gb.placeholder("x", torch.empty((4,), dtype=torch.float32))
	114	+ self.assertIsNone(_prims_iota_value_range(x))
	115	+
	116	+ # ===== iota int64 → int32 downcast =====
	117	+ def test_iota_downcast_to_int32(self):
	118	+ gm, iota = _build_iota_with_user(length=128)
	119	+ fold_iota_arithmetic_pass(gm.graph)
	120	+ new_iota = next(
	121	+ n for n in gm.graph.nodes
	122	+ if n.op == "call_function" and n.target is _IOTA
	123	+ )
	124	+ self.assertIs(new_iota.kwargs["dtype"], torch.int32)
	125	+
	126	+ def test_iota_out_of_int32_range_no_downcast(self):
	127	+ gm, iota = _build_iota_with_user(length=2, start=2**31)
	128	+ fold_iota_arithmetic_pass(gm.graph)
	129	+ new_iota = next(
	130	+ n for n in gm.graph.nodes
	131	+ if n.op == "call_function" and n.target is _IOTA
	132	+ )
	133	+ self.assertIs(new_iota.kwargs["dtype"], torch.int64)
	134	+
	135	+ def test_iota_non_int64_dtype_skipped(self):
	136	+ gm, iota = _build_iota_with_user(length=10, dtype=torch.int32)
	137	+ fold_iota_arithmetic_pass(gm.graph)
	138	+ new_iota = next(
	139	+ n for n in gm.graph.nodes
	140	+ if n.op == "call_function" and n.target is _IOTA
	141	+ )
	142	+ self.assertIs(new_iota.kwargs["dtype"], torch.int32)
	143	+
	144	+ # ===== closure collection =====
	145	+ def test_collect_closure_only_transparent_then_closing(self):
	146	+ fm = new_fake_mode()
	147	+ gb = GraphBuilder(fm)
	148	+ with fm:
	149	+ iota_val = torch.ops.prims.iota.default(
	150	+ 4, start=0, step=1, dtype=torch.int64,
	151	+ device="cpu", requires_grad=False,
	152	+ )
	153	+ iota = gb.graph.call_function(
	154	+ _IOTA, args=(4,),
	155	+ kwargs={"start": 0, "step": 1, "dtype": torch.int64,
	156	+ "device": torch.device("cpu"), "requires_grad": False},
	157	+ )
	158	+ iota.meta["val"] = iota_val
	159	+ v = gb.call(torch.ops.aten.view.default, args=(iota, [2, 2]))
	160	+ cmp = gb.call(_GE_S, args=(v, 1))
	161	+ gb.output(cmp)
	162	+ ids = _collect_iota_downcast_closure(iota)
	163	+ self.assertIsNotNone(ids)
	164	+ self.assertIn(id(v), ids)
	165	+
	166	+ def test_collect_closure_rejects_non_transparent_user(self):
	167	+ fm = new_fake_mode()
	168	+ gb = GraphBuilder(fm)
	169	+ with fm:
	170	+ iota_val = torch.ops.prims.iota.default(
	171	+ 4, start=0, step=1, dtype=torch.int64,
	172	+ device="cpu", requires_grad=False,
	173	+ )
	174	+ iota = gb.graph.call_function(
	175	+ _IOTA, args=(4,),
	176	+ kwargs={"start": 0, "step": 1, "dtype": torch.int64,
	177	+ "device": torch.device("cpu"), "requires_grad": False},
	178	+ )
	179	+ iota.meta["val"] = iota_val
	180	+ non_trans = gb.call(torch.ops.aten.relu.default, args=(iota,))
	181	+ gb.output(non_trans)
	182	+ self.assertIsNone(_collect_iota_downcast_closure(iota))
	183	+
	184	+ # ===== cmp(sub(a,b), 0) → cmp(a,b) =====
	185	+ def _build_cmp_sub_zero(self, cmp_target, rhs=0, alpha=1):
	186	+ fm = new_fake_mode()
	187	+ gb = GraphBuilder(fm)
	188	+ with fm:
	189	+ a_fake = torch.empty((4,), dtype=torch.float32)
	190	+ b_fake = torch.empty((4,), dtype=torch.float32)
	191	+ a = gb.placeholder("a", a_fake)
	192	+ b = gb.placeholder("b", b_fake)
	193	+ sub_kwargs = {} if alpha == 1 else {"alpha": alpha}
	194	+ sub = gb.call(_SUB, args=(a, b), kwargs=sub_kwargs)
	195	+ cmp = gb.call(cmp_target, args=(sub, rhs))
	196	+ gb.output(cmp)
	197	+ return gb.to_module()
	198	+
	199	+ def test_cmp_sub_zero_simplified_ge(self):
	200	+ gm = self._build_cmp_sub_zero(_GE_S)
	201	+ fold_iota_arithmetic_pass(gm.graph)
	202	+ self.assertEqual(count_target(gm.graph, _GE_T), 1)
	203	+ self.assertEqual(count_target(gm.graph, _GE_S), 0)
	204	+
	205	+ def test_cmp_sub_zero_simplified_gt(self):
	206	+ gm = self._build_cmp_sub_zero(_GT_S)
	207	+ fold_iota_arithmetic_pass(gm.graph)
	208	+ self.assertEqual(count_target(gm.graph, _GT_T), 1)
	209	+
	210	+ def test_cmp_non_zero_rhs_skipped(self):
	211	+ gm = self._build_cmp_sub_zero(_GE_S, rhs=1)
	212	+ fold_iota_arithmetic_pass(gm.graph)
	213	+ self.assertEqual(count_target(gm.graph, _GE_S), 1)
	214	+ self.assertEqual(count_target(gm.graph, _GE_T), 0)
	215	+
	216	+ def test_cmp_sub_with_alpha_skipped(self):
	217	+ gm = self._build_cmp_sub_zero(_GE_S, alpha=2)
	218	+ fold_iota_arithmetic_pass(gm.graph)
	219	+ self.assertEqual(count_target(gm.graph, _GE_S), 1)
	220	+ self.assertEqual(count_target(gm.graph, _GE_T), 0)
	221	+
	222	+ def test_cmp_lhs_not_sub_skipped(self):
	223	+ fm = new_fake_mode()
	224	+ gb = GraphBuilder(fm)
	225	+ with fm:
	226	+ a_fake = torch.empty((4,), dtype=torch.float32)
	227	+ a = gb.placeholder("a", a_fake)
	228	+ cmp = gb.call(_GE_S, args=(a, 0))
	229	+ gb.output(cmp)
	230	+ gm = gb.to_module()
	231	+ fold_iota_arithmetic_pass(gm.graph)
	232	+ self.assertEqual(count_target(gm.graph, _GE_S), 1)
	233	+
	234	+ # ===== _hashable_const_key & _cse_constant_call =====
	235	+ def test_hashable_const_key_basic(self):
	236	+ self.assertEqual(_hashable_const_key(3), 3)
	237	+ self.assertEqual(
	238	+ _hashable_const_key([1, [2, 3]]),
	239	+ ("__list__", 1, ("__list__", 2, 3)),
	240	+ )
	241	+ self.assertEqual(
	242	+ _hashable_const_key((1, 2)),
	243	+ ("__tuple__", 1, 2),
	244	+ )
	245	+ self.assertEqual(
	246	+ _hashable_const_key({"a": 1}),
	247	+ ("__dict__", ("a", 1)),
	248	+ )
	249	+
	250	+ def test_cse_constant_call_dedups(self):
	251	+ """两个同参数 full 调用应被 CSE 合并。"""
	252	+ fm = new_fake_mode()
	253	+ gb = GraphBuilder(fm)
	254	+ with fm:
	255	+ f_fake = torch.ops.aten.full.default([4], 0.0, dtype=torch.float32)
	256	+ # 确保 placeholder 先于 call_function 节点出现，避免 FX 拓扑告警。
	257	+ _ = gb.placeholder("x", torch.empty((4,), dtype=torch.float32))
	258	+ f1 = gb.graph.call_function(_FULL, args=([4], 0.0), kwargs={"dtype": torch.float32})
	259	+ f1.meta["val"] = f_fake
	260	+ f2 = gb.graph.call_function(_FULL, args=([4], 0.0), kwargs={"dtype": torch.float32})
	261	+ f2.meta["val"] = f_fake
	262	+ add = gb.call(torch.ops.aten.add.Tensor, args=(f1, f2))
	263	+ gb.output(add)
	264	+ gm = gb.to_module()
	265	+ changed = _cse_constant_call(gm.graph, _FULL)
	266	+ self.assertTrue(changed)
	267	+
	268	+
	269	+if __name__ == "__main__":
	270	+ run_tests()

Atest/_inductor/test_masked_add_compose_pass.py+219-0

@@ -0,0 +1,219 @@
	1	+import torch
	2	+from torch.testing._internal.common_utils import TestCase, run_tests
	3	+
	4	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	5	+ masked_add_compose_pass,
	6	+ _is_zero_tensor_source,
	7	+ _strip_logical_not,
	8	+ _are_logically_negated_masks,
	9	+ _match_masked_zero_where,
	10	+)
	11	+
	12	+from _pass_test_utils import (
	13	+ GraphBuilder,
	14	+ count_target,
	15	+ new_fake_mode,
	16	+)
	17	+
	18	+
	19	+_WHERE = torch.ops.aten.where.self
	20	+_ADD = torch.ops.aten.add.Tensor
	21	+_NOT = torch.ops.aten.logical_not.default
	22	+_FULL = torch.ops.aten.full.default
	23	+_ZEROS = torch.ops.aten.zeros.default
	24	+
	25	+
	26	+def _build_masked_add(shape=(4, 4), use_full_zero=True, rhs_other=None,
	27	+ negate_rhs_mask=True, alpha=1):
	28	+ fm = new_fake_mode()
	29	+ gb = GraphBuilder(fm)
	30	+ with fm:
	31	+ mask_fake = torch.empty(shape, dtype=torch.bool)
	32	+ a_fake = torch.empty(shape, dtype=torch.float32)
	33	+ b_fake = torch.empty(shape, dtype=torch.float32)
	34	+ mask = gb.placeholder("mask", mask_fake)
	35	+ a = gb.placeholder("a", a_fake)
	36	+ b = gb.placeholder("b", b_fake)
	37	+ if use_full_zero:
	38	+ zero = gb.graph.call_function(
	39	+ _FULL, args=(list(shape), 0), kwargs={"dtype": torch.float32}
	40	+ )
	41	+ with fm:
	42	+ zero.meta["val"] = torch.ops.aten.full.default(
	43	+ list(shape), 0, dtype=torch.float32
	44	+ )
	45	+ else:
	46	+ zero = gb.graph.call_function(
	47	+ _ZEROS, args=(list(shape),), kwargs={"dtype": torch.float32}
	48	+ )
	49	+ with fm:
	50	+ zero.meta["val"] = torch.ops.aten.zeros.default(
	51	+ list(shape), dtype=torch.float32
	52	+ )
	53	+
	54	+ not_mask = gb.call(_NOT, args=(mask,)) if negate_rhs_mask else mask
	55	+ w_lhs = gb.call(_WHERE, args=(mask, a, zero))
	56	+ rhs_other_node = rhs_other if rhs_other is not None else zero
	57	+ w_rhs = gb.call(_WHERE, args=(not_mask, b, rhs_other_node))
	58	+ if alpha == 1:
	59	+ add = gb.call(_ADD, args=(w_lhs, w_rhs))
	60	+ else:
	61	+ add = gb.call(_ADD, args=(w_lhs, w_rhs), kwargs={"alpha": alpha})
	62	+ gb.output(add)
	63	+ return gb.to_module(), mask, a, b
	64	+
	65	+
	66	+class TestMaskedAddComposePass(TestCase):
	67	+ # ===== helpers =====
	68	+ def test_is_zero_tensor_source_scalar(self):
	69	+ self.assertTrue(_is_zero_tensor_source(0))
	70	+ self.assertFalse(_is_zero_tensor_source(1))
	71	+
	72	+ def test_is_zero_tensor_source_zeros_call(self):
	73	+ fm = new_fake_mode()
	74	+ gb = GraphBuilder(fm)
	75	+ n = gb.graph.call_function(_ZEROS, args=([4],))
	76	+ self.assertTrue(_is_zero_tensor_source(n))
	77	+
	78	+ def test_is_zero_tensor_source_full_zero(self):
	79	+ fm = new_fake_mode()
	80	+ gb = GraphBuilder(fm)
	81	+ z = gb.graph.call_function(_FULL, args=([4], 0))
	82	+ self.assertTrue(_is_zero_tensor_source(z))
	83	+ nz = gb.graph.call_function(_FULL, args=([4], 5))
	84	+ self.assertFalse(_is_zero_tensor_source(nz))
	85	+
	86	+ def test_strip_logical_not_logical(self):
	87	+ fm = new_fake_mode()
	88	+ gb = GraphBuilder(fm)
	89	+ with fm:
	90	+ mask_fake = torch.empty((4,), dtype=torch.bool)
	91	+ m = gb.placeholder("m", mask_fake)
	92	+ nm = gb.call(_NOT, args=(m,))
	93	+ inner, neg = _strip_logical_not(nm)
	94	+ self.assertTrue(neg)
	95	+ self.assertIs(inner, m)
	96	+
	97	+ def test_strip_logical_not_bitwise_on_bool(self):
	98	+ """bitwise_not 作用在 bool 张量上时也应被识别为逻辑取反。"""
	99	+ fm = new_fake_mode()
	100	+ gb = GraphBuilder(fm)
	101	+ with fm:
	102	+ mask_fake = torch.empty((4,), dtype=torch.bool)
	103	+ m = gb.placeholder("m", mask_fake)
	104	+ nm = gb.call(torch.ops.aten.bitwise_not.default, args=(m,))
	105	+ inner, neg = _strip_logical_not(nm)
	106	+ self.assertTrue(neg)
	107	+ self.assertIs(inner, m)
	108	+
	109	+ def test_strip_logical_not_bitwise_on_int_passthrough(self):
	110	+ """bitwise_not 作用在非 bool 张量上时不视作逻辑取反。"""
	111	+ fm = new_fake_mode()
	112	+ gb = GraphBuilder(fm)
	113	+ with fm:
	114	+ int_fake = torch.empty((4,), dtype=torch.int32)
	115	+ x = gb.placeholder("x", int_fake)
	116	+ nx = gb.call(torch.ops.aten.bitwise_not.default, args=(x,))
	117	+ inner, neg = _strip_logical_not(nx)
	118	+ self.assertFalse(neg)
	119	+ self.assertIs(inner, nx)
	120	+
	121	+ def test_strip_logical_not_passthrough(self):
	122	+ fm = new_fake_mode()
	123	+ gb = GraphBuilder(fm)
	124	+ with fm:
	125	+ x_fake = torch.empty((4,), dtype=torch.float32)
	126	+ x = gb.placeholder("x", x_fake)
	127	+ inner, neg = _strip_logical_not(x)
	128	+ self.assertFalse(neg)
	129	+ self.assertIs(inner, x)
	130	+
	131	+ def test_are_logically_negated_masks(self):
	132	+ fm = new_fake_mode()
	133	+ gb = GraphBuilder(fm)
	134	+ with fm:
	135	+ mask_fake = torch.empty((4,), dtype=torch.bool)
	136	+ m = gb.placeholder("m", mask_fake)
	137	+ nm = gb.call(_NOT, args=(m,))
	138	+ self.assertTrue(_are_logically_negated_masks(m, nm))
	139	+ self.assertFalse(_are_logically_negated_masks(m, m))
	140	+
	141	+ def test_match_masked_zero_where_positive(self):
	142	+ fm = new_fake_mode()
	143	+ gb = GraphBuilder(fm)
	144	+ with fm:
	145	+ m_fake = torch.empty((4,), dtype=torch.bool)
	146	+ v_fake = torch.empty((4,), dtype=torch.float32)
	147	+ m = gb.placeholder("m", m_fake)
	148	+ v = gb.placeholder("v", v_fake)
	149	+ w = gb.call(_WHERE, args=(m, v, 0))
	150	+ match = _match_masked_zero_where(w)
	151	+ self.assertIsNotNone(match)
	152	+ self.assertIs(match[0], m)
	153	+ self.assertIs(match[1], v)
	154	+
	155	+ def test_match_masked_zero_where_negative(self):
	156	+ fm = new_fake_mode()
	157	+ gb = GraphBuilder(fm)
	158	+ with fm:
	159	+ m_fake = torch.empty((4,), dtype=torch.bool)
	160	+ v_fake = torch.empty((4,), dtype=torch.float32)
	161	+ m = gb.placeholder("m", m_fake)
	162	+ v = gb.placeholder("v", v_fake)
	163	+ # other 不是 0 → 不匹配
	164	+ w = gb.call(_WHERE, args=(m, v, v))
	165	+ self.assertIsNone(_match_masked_zero_where(w))
	166	+
	167	+ # ===== full pass =====
	168	+ def test_masked_add_composed_into_where(self):
	169	+ gm, mask, a, b = _build_masked_add()
	170	+ masked_add_compose_pass(gm.graph)
	171	+ self.assertEqual(count_target(gm.graph, _ADD), 0)
	172	+ wheres = [
	173	+ n for n in gm.graph.nodes
	174	+ if n.op == "call_function" and n.target is _WHERE
	175	+ ]
	176	+ self.assertEqual(len(wheres), 1)
	177	+ self.assertIs(wheres[0].args[0], mask)
	178	+ self.assertIs(wheres[0].args[1], a)
	179	+ self.assertIs(wheres[0].args[2], b)
	180	+
	181	+ def test_alpha_not_one_skipped(self):
	182	+ gm, *_ = _build_masked_add(alpha=2)
	183	+ masked_add_compose_pass(gm.graph)
	184	+ self.assertEqual(count_target(gm.graph, _ADD), 1)
	185	+
	186	+ def test_same_mask_skipped(self):
	187	+ """两个 where 用同一 mask（无逻辑取反）→ 不应合成。"""
	188	+ gm, *_ = _build_masked_add(negate_rhs_mask=False)
	189	+ masked_add_compose_pass(gm.graph)
	190	+ self.assertEqual(count_target(gm.graph, _ADD), 1)
	191	+
	192	+ def test_non_zero_other_skipped(self):
	193	+ """右侧 where 的 other 不是 0 → 不应合成。"""
	194	+ fm = new_fake_mode()
	195	+ gb = GraphBuilder(fm)
	196	+ with fm:
	197	+ mask_fake = torch.empty((4,), dtype=torch.bool)
	198	+ a_fake = torch.empty((4,), dtype=torch.float32)
	199	+ b_fake = torch.empty((4,), dtype=torch.float32)
	200	+ c_fake = torch.empty((4,), dtype=torch.float32)
	201	+ mask = gb.placeholder("mask", mask_fake)
	202	+ a = gb.placeholder("a", a_fake)
	203	+ b = gb.placeholder("b", b_fake)
	204	+ c = gb.placeholder("c", c_fake)
	205	+ z = gb.graph.call_function(_ZEROS, args=([4],))
	206	+ with fm:
	207	+ z.meta["val"] = torch.ops.aten.zeros.default([4])
	208	+ not_mask = gb.call(_NOT, args=(mask,))
	209	+ w_lhs = gb.call(_WHERE, args=(mask, a, z))
	210	+ w_rhs = gb.call(_WHERE, args=(not_mask, b, c)) # non-zero other
	211	+ add = gb.call(_ADD, args=(w_lhs, w_rhs))
	212	+ gb.output(add)
	213	+ gm = gb.to_module()
	214	+ masked_add_compose_pass(gm.graph)
	215	+ self.assertEqual(count_target(gm.graph, _ADD), 1)
	216	+
	217	+
	218	+if __name__ == "__main__":
	219	+ run_tests()

Atest/_inductor/test_repeat_to_expand_pass.py+112-0

@@ -0,0 +1,112 @@
	1	+import torch
	2	+from torch.testing._internal.common_utils import TestCase, run_tests
	3	+
	4	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	5	+ repeat_to_expand_pass,
	6	+)
	7	+
	8	+from _pass_test_utils import (
	9	+ GraphBuilder,
	10	+ count_target,
	11	+ new_fake_mode,
	12	+)
	13	+
	14	+
	15	+_REPEAT = torch.ops.aten.repeat.default
	16	+_EXPAND = torch.ops.aten.expand.default
	17	+_MUL = torch.ops.aten.mul.Tensor
	18	+_RELU = torch.ops.aten.relu.default
	19	+
	20	+
	21	+def _build(input_shape, repeats, user_target=_MUL, other_shape=None):
	22	+ fm = new_fake_mode()
	23	+ gb = GraphBuilder(fm)
	24	+ with fm:
	25	+ x_fake = torch.empty(input_shape, dtype=torch.float32)
	26	+ other_shape = other_shape or [int(s) * int(r) for s, r in zip(input_shape, repeats)]
	27	+ y_fake = torch.empty(other_shape, dtype=torch.float32)
	28	+ x = gb.placeholder("x", x_fake)
	29	+ y = gb.placeholder("y", y_fake)
	30	+ rpt = gb.call(_REPEAT, args=(x, list(repeats)))
	31	+ out = gb.call(user_target, args=(rpt, y)) if user_target is not _RELU else gb.call(user_target, args=(rpt,))
	32	+ gb.output(out)
	33	+ return gb.to_module()
	34	+
	35	+
	36	+class TestRepeatToExpandPass(TestCase):
	37	+ def test_pure_broadcast_repeat_rewritten(self):
	38	+ """input_shape=(1,4), repeats=(8,1) → 应替换为 expand."""
	39	+ gm = _build((1, 4), (8, 1))
	40	+ repeat_to_expand_pass(gm.graph)
	41	+ self.assertEqual(count_target(gm.graph, _REPEAT), 0)
	42	+ self.assertEqual(count_target(gm.graph, _EXPAND), 1)
	43	+ exp = next(
	44	+ n for n in gm.graph.nodes
	45	+ if n.op == "call_function" and n.target is _EXPAND
	46	+ )
	47	+ self.assertEqual(list(exp.args[1]), [8, 4])
	48	+
	49	+ def test_non_broadcast_repeat_skipped(self):
	50	+ """非 1 维度也有 repeat>1（实际产生新数据），不能替换。"""
	51	+ gm = _build((2, 4), (3, 1)) # 维度 0 上 size=2 且 repeat=3 → 真物理拷贝
	52	+ repeat_to_expand_pass(gm.graph)
	53	+ self.assertEqual(count_target(gm.graph, _REPEAT), 1)
	54	+ self.assertEqual(count_target(gm.graph, _EXPAND), 0)
	55	+
	56	+ def test_user_not_broadcast_friendly_skipped(self):
	57	+ """下游使用者不属于广播友好集合 → 不应替换。"""
	58	+ gm = _build((1, 4), (8, 1), user_target=_RELU)
	59	+ repeat_to_expand_pass(gm.graph)
	60	+ self.assertEqual(count_target(gm.graph, _REPEAT), 1)
	61	+ self.assertEqual(count_target(gm.graph, _EXPAND), 0)
	62	+
	63	+ def test_repeats_len_mismatch_skipped(self):
	64	+ """repeats 长度 != input 维度数 → 不替换（保留原 repeat）。"""
	65	+ fm = new_fake_mode()
	66	+ gb = GraphBuilder(fm)
	67	+ with fm:
	68	+ x_fake = torch.empty((1, 4), dtype=torch.float32)
	69	+ y_fake = torch.empty((2, 1, 4), dtype=torch.float32)
	70	+ x = gb.placeholder("x", x_fake)
	71	+ y = gb.placeholder("y", y_fake)
	72	+ rpt = gb.call(_REPEAT, args=(x, [2, 1, 1]))
	73	+ out = gb.call(_MUL, args=(rpt, y))
	74	+ gb.output(out)
	75	+ gm = gb.to_module()
	76	+ repeat_to_expand_pass(gm.graph)
	77	+ self.assertEqual(count_target(gm.graph, _REPEAT), 1)
	78	+ self.assertEqual(count_target(gm.graph, _EXPAND), 0)
	79	+
	80	+ def test_no_users_skipped(self):
	81	+ """repeat 节点没有 user（如它就是输出）→ 不替换。"""
	82	+ fm = new_fake_mode()
	83	+ gb = GraphBuilder(fm)
	84	+ with fm:
	85	+ x_fake = torch.empty((1, 4), dtype=torch.float32)
	86	+ x = gb.placeholder("x", x_fake)
	87	+ rpt = gb.call(_REPEAT, args=(x, [8, 1]))
	88	+ gb.output(rpt)
	89	+ gm = gb.to_module()
	90	+ repeat_to_expand_pass(gm.graph)
	91	+ # output 不算 call_function 用户，rpt.users 为空 → 不应替换
	92	+ self.assertEqual(count_target(gm.graph, _REPEAT), 1)
	93	+
	94	+ def test_repeats_not_list_skipped(self):
	95	+ """repeats 参数不是 list/tuple → 直接跳过。"""
	96	+ fm = new_fake_mode()
	97	+ gb = GraphBuilder(fm)
	98	+ with fm:
	99	+ x_fake = torch.empty((1, 4), dtype=torch.float32)
	100	+ y_fake = torch.empty((8, 4), dtype=torch.float32)
	101	+ x = gb.placeholder("x", x_fake)
	102	+ y = gb.placeholder("y", y_fake)
	103	+ rpt = gb.call(_REPEAT, args=(x, "not_a_list"))
	104	+ out = gb.call(_MUL, args=(rpt, y))
	105	+ gb.output(out)
	106	+ gm = gb.to_module()
	107	+ repeat_to_expand_pass(gm.graph)
	108	+ self.assertEqual(count_target(gm.graph, _REPEAT), 1)
	109	+
	110	+
	111	+if __name__ == "__main__":
	112	+ run_tests()

Atest/_inductor/test_sign_diff_hamming_fuse_pass.py+135-0

@@ -0,0 +1,135 @@
	1	+import torch
	2	+from torch.testing._internal.common_utils import TestCase, run_tests
	3	+
	4	+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import (
	5	+ sign_diff_hamming_fuse_pass,
	6	+ _peel_single_user_relu_sign,
	7	+)
	8	+
	9	+from _pass_test_utils import (
	10	+ GraphBuilder,
	11	+ count_target,
	12	+ new_fake_mode,
	13	+)
	14	+
	15	+
	16	+_SIGN = torch.ops.aten.sign.default
	17	+_RELU = torch.ops.aten.relu.default
	18	+_SUB = torch.ops.aten.sub.Tensor
	19	+_ABS = torch.ops.aten.abs.default
	20	+_SUM = torch.ops.aten.sum.dim_IntList
	21	+_GT = torch.ops.aten.gt.Scalar
	22	+_NE = torch.ops.aten.ne.Tensor
	23	+
	24	+
	25	+def _build_full_pattern(shape=(8,), dim=[0], alpha=1, drop_abs=False, drop_sign_x=False):
	26	+ fm = new_fake_mode()
	27	+ gb = GraphBuilder(fm)
	28	+ with fm:
	29	+ x_fake = torch.empty(shape, dtype=torch.float32)
	30	+ y_fake = torch.empty(shape, dtype=torch.float32)
	31	+ x = gb.placeholder("x", x_fake)
	32	+ y = gb.placeholder("y", y_fake)
	33	+ sx = gb.call(_SIGN, args=(x,)) if not drop_sign_x else x
	34	+ sy = gb.call(_SIGN, args=(y,))
	35	+ rx = gb.call(_RELU, args=(sx,)) if not drop_sign_x else gb.call(_RELU, args=(sx,))
	36	+ ry = gb.call(_RELU, args=(sy,))
	37	+ sub_kwargs = {} if alpha == 1 else {"alpha": alpha}
	38	+ sub = gb.call(_SUB, args=(rx, ry), kwargs=sub_kwargs)
	39	+ upstream = sub if drop_abs else gb.call(_ABS, args=(sub,))
	40	+ s = gb.call(_SUM, args=(upstream, dim, False), kwargs={"dtype": torch.int64})
	41	+ gb.output(s)
	42	+ return gb.to_module(), x, y
	43	+
	44	+
	45	+class TestSignDiffHammingFusePass(TestCase):
	46	+ def test_peel_relu_sign(self):
	47	+ fm = new_fake_mode()
	48	+ gb = GraphBuilder(fm)
	49	+ with fm:
	50	+ x_fake = torch.empty((4,), dtype=torch.float32)
	51	+ x = gb.placeholder("x", x_fake)
	52	+ s = gb.call(_SIGN, args=(x,))
	53	+ r = gb.call(_RELU, args=(s,))
	54	+ gb.output(r)
	55	+ self.assertIs(_peel_single_user_relu_sign(r), x)
	56	+
	57	+ def test_peel_not_relu(self):
	58	+ fm = new_fake_mode()
	59	+ gb = GraphBuilder(fm)
	60	+ with fm:
	61	+ x_fake = torch.empty((4,), dtype=torch.float32)
	62	+ x = gb.placeholder("x", x_fake)
	63	+ self.assertIsNone(_peel_single_user_relu_sign(x))
	64	+
	65	+ def test_peel_sign_with_multiple_users(self):
	66	+ fm = new_fake_mode()
	67	+ gb = GraphBuilder(fm)
	68	+ with fm:
	69	+ x_fake = torch.empty((4,), dtype=torch.float32)
	70	+ x = gb.placeholder("x", x_fake)
	71	+ s = gb.call(_SIGN, args=(x,))
	72	+ r1 = gb.call(_RELU, args=(s,))
	73	+ r2 = gb.call(_RELU, args=(s,))
	74	+ gb.output((r1, r2))
	75	+ # sign 节点有两个使用者 → 不匹配
	76	+ self.assertIsNone(_peel_single_user_relu_sign(r1))
	77	+
	78	+ def test_full_pattern_fused(self):
	79	+ gm, x, y = _build_full_pattern()
	80	+ sign_diff_hamming_fuse_pass(gm.graph)
	81	+ self.assertEqual(count_target(gm.graph, _ABS), 0)
	82	+ self.assertEqual(count_target(gm.graph, _SUB), 0)
	83	+ self.assertEqual(count_target(gm.graph, _GT), 2)
	84	+ self.assertEqual(count_target(gm.graph, _NE), 1)
	85	+ new_sum = next(
	86	+ n for n in gm.graph.nodes
	87	+ if n.op == "call_function" and n.target is _SUM
	88	+ )
	89	+ self.assertEqual(new_sum.kwargs.get("dtype"), torch.int64)
	90	+
	91	+ def test_missing_abs_not_fused(self):
	92	+ gm, *_ = _build_full_pattern(drop_abs=True)
	93	+ sign_diff_hamming_fuse_pass(gm.graph)
	94	+ # 没有 abs，sum 的输入是 sub，模式不完整 → 不应改写
	95	+ self.assertEqual(count_target(gm.graph, _NE), 0)
	96	+
	97	+ def test_missing_sign_chain_not_fused(self):
	98	+ gm, *_ = _build_full_pattern(drop_sign_x=True)
	99	+ sign_diff_hamming_fuse_pass(gm.graph)
	100	+ self.assertEqual(count_target(gm.graph, _NE), 0)
	101	+
	102	+ def test_alpha_not_one_skipped(self):
	103	+ gm, *_ = _build_full_pattern(alpha=2)
	104	+ sign_diff_hamming_fuse_pass(gm.graph)
	105	+ self.assertEqual(count_target(gm.graph, _NE), 0)
	106	+
	107	+ def test_keepdim_true_still_fused(self):
	108	+ """keepdim=True 也应能识别（pass 直接透传 keepdim 给新 sum）。"""
	109	+ fm = new_fake_mode()
	110	+ gb = GraphBuilder(fm)
	111	+ with fm:
	112	+ x_fake = torch.empty((4, 8), dtype=torch.float32)
	113	+ y_fake = torch.empty((4, 8), dtype=torch.float32)
	114	+ x = gb.placeholder("x", x_fake)
	115	+ y = gb.placeholder("y", y_fake)
	116	+ sx = gb.call(_SIGN, args=(x,))
	117	+ sy = gb.call(_SIGN, args=(y,))
	118	+ rx = gb.call(_RELU, args=(sx,))
	119	+ ry = gb.call(_RELU, args=(sy,))
	120	+ sub = gb.call(_SUB, args=(rx, ry))
	121	+ a = gb.call(_ABS, args=(sub,))
	122	+ s = gb.call(_SUM, args=(a, [1], True), kwargs={"dtype": torch.int64})
	123	+ gb.output(s)
	124	+ gm = gb.to_module()
	125	+ sign_diff_hamming_fuse_pass(gm.graph)
	126	+ self.assertEqual(count_target(gm.graph, _NE), 1)
	127	+ new_sum = next(
	128	+ n for n in gm.graph.nodes
	129	+ if n.op == "call_function" and n.target is _SUM
	130	+ )
	131	+ self.assertEqual(new_sum.args[2], True)
	132	+
	133	+
	134	+if __name__ == "__main__":
	135	+ run_tests()