| @@ -0,0 +1,204 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch.fx as fx | ||
| 3 | +from torch.fx.passes.shape_prop import ShapeProp | ||
| 4 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 5 | +from testutils import TestUtils | ||
| 6 | +import torch_npu | ||
| 7 | +import torch_npu._inductor | ||
| 8 | +from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_four_op_pass | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class FoldAddModel(torch.nn.Module): | ||
| 12 | + def forward(self, first_element): | ||
| 13 | + add = torch.ops.aten.add.Tensor(first_element, 0) | ||
| 14 | + add_output = torch.ops.aten.relu.default(add) | ||
| 15 | + return add_output | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +class TestFoldAddPass(TestUtils): | ||
| 19 | + def op_calc(self, first_element): | ||
| 20 | + add = torch.add(first_element, 0) | ||
| 21 | + add_output = torch.relu(add) | ||
| 22 | + return add_output | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + def test_compile_cases(self, shape, dtype): | ||
| 28 | + first_element = self._generate_tensor(shape, dtype) | ||
| 29 | + std_result = self.op_calc(first_element) | ||
| 30 | + with torch.no_grad(): | ||
| 31 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor") | ||
| 32 | + inductor_result = compiled_op_calc(first_element) | ||
| 33 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 34 | + | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + | ||
| 38 | + def test_ut_cases(self, shape, dtype): | ||
| 39 | + first_element = self._generate_tensor(shape, dtype) | ||
| 40 | + model = FoldAddModel() | ||
| 41 | + graph_module = fx.symbolic_trace(model) | ||
| 42 | + ShapeProp(graph_module).propagate(first_element) | ||
| 43 | + | ||
| 44 | + # 应用优化 Pass | ||
| 45 | + fold_four_op_pass(graph_module.graph) | ||
| 46 | + graph_module.recompile() | ||
| 47 | + | ||
| 48 | + # 验证输出是否一致 | ||
| 49 | + std_result = model(first_element) | ||
| 50 | + inductor_result = graph_module(first_element) | ||
| 51 | + | ||
| 52 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 53 | + | ||
| 54 | + | ||
| 55 | + | ||
| 56 | + | ||
| 57 | + def test_add_zero_left_fold(self, shape, dtype): | ||
| 58 | + """ 0 + x 应该折叠成 x """ | ||
| 59 | + class M(torch.nn.Module): | ||
| 60 | + def forward(self, x): | ||
| 61 | + y = torch.ops.aten.add.Tensor(0, x) # 左边是 0 (scalar) | ||
| 62 | + return torch.relu(y + 1.0) | ||
| 63 | + | ||
| 64 | + t = self._generate_tensor(shape, dtype) | ||
| 65 | + m = M() | ||
| 66 | + gm = fx.symbolic_trace(m) | ||
| 67 | + ShapeProp(gm).propagate(t) | ||
| 68 | + | ||
| 69 | + fold_four_op_pass(gm.graph) | ||
| 70 | + gm.recompile() | ||
| 71 | + | ||
| 72 | + # 检查 pass 是否删除了 add 节点 | ||
| 73 | + self.assertFalse( | ||
| 74 | + any(n.target in (torch.add, torch.ops.aten.add.Tensor, torch.ops.aten.add.Scalar) | ||
| 75 | + for n in gm.graph.nodes), | ||
| 76 | + "add + 0 should be folded away" | ||
| 77 | + ) | ||
| 78 | + | ||
| 79 | + self.assertEqual(m(t), gm(t)) | ||
| 80 | + | ||
| 81 | + | ||
| 82 | + | ||
| 83 | + def test_add_scalar_zero_right_fold(self, shape): | ||
| 84 | + """ x + 0.0 (scalar) 应该折叠 """ | ||
| 85 | + class M(torch.nn.Module): | ||
| 86 | + def forward(self, x): | ||
| 87 | + y = torch.add(x, 0.0) # 使用 python torch.add → 应转为 add.Scalar | ||
| 88 | + return y * 2.0 | ||
| 89 | + | ||
| 90 | + t = torch.randn(shape) | ||
| 91 | + m = M() | ||
| 92 | + gm = fx.symbolic_trace(m) | ||
| 93 | + ShapeProp(gm).propagate(t) | ||
| 94 | + | ||
| 95 | + fold_four_op_pass(gm.graph) | ||
| 96 | + gm.recompile() | ||
| 97 | + | ||
| 98 | + self.assertFalse( | ||
| 99 | + any(n.target in (torch.add, torch.ops.aten.add.Tensor, torch.ops.aten.add.Scalar) | ||
| 100 | + for n in gm.graph.nodes) | ||
| 101 | + ) | ||
| 102 | + torch.testing.assert_close(m(t), gm(t), atol=1e-5, rtol=1e-5) | ||
| 103 | + | ||
| 104 | + | ||
| 105 | + | ||
| 106 | + def test_add_int_zero_right_fold(self, shape): | ||
| 107 | + """ x + 0 (int literal) """ | ||
| 108 | + class M(torch.nn.Module): | ||
| 109 | + def forward(self, x): | ||
| 110 | + return torch.ops.aten.add.Tensor(x, 0) + x | ||
| 111 | + | ||
| 112 | + t = self._generate_tensor(shape, 'float32') | ||
| 113 | + m = M() | ||
| 114 | + gm = fx.symbolic_trace(m) | ||
| 115 | + ShapeProp(gm).propagate(t) | ||
| 116 | + | ||
| 117 | + fold_four_op_pass(gm.graph) | ||
| 118 | + gm.recompile() | ||
| 119 | + | ||
| 120 | + self.assertEqual(m(t), gm(t)) | ||
| 121 | + | ||
| 122 | + | ||
| 123 | + | ||
| 124 | + def test_add_zero_not_fold_nonzero(self, shape): | ||
| 125 | + """ x + 1.5 不应该折叠 """ | ||
| 126 | + class M(torch.nn.Module): | ||
| 127 | + def forward(self, x): | ||
| 128 | + y = torch.add(x, 1.5) | ||
| 129 | + return torch.sigmoid(y) | ||
| 130 | + | ||
| 131 | + t = self._generate_tensor(shape, 'float32') | ||
| 132 | + m = M() | ||
| 133 | + gm = fx.symbolic_trace(m) | ||
| 134 | + ShapeProp(gm).propagate(t) | ||
| 135 | + | ||
| 136 | + before_nodes = len([n for n in gm.graph.nodes if n.op == 'call_function']) | ||
| 137 | + | ||
| 138 | + fold_four_op_pass(gm.graph) | ||
| 139 | + gm.recompile() | ||
| 140 | + | ||
| 141 | + after_nodes = len([n for n in gm.graph.nodes if n.op == 'call_function']) | ||
| 142 | + self.assertEqual(before_nodes, after_nodes, "不应折叠非零常量") | ||
| 143 | + | ||
| 144 | + | ||
| 145 | + | ||
| 146 | + def test_add_zero_chain_fold(self, shape): | ||
| 147 | + """ 多个连续的 +0 折叠 """ | ||
| 148 | + class M(torch.nn.Module): | ||
| 149 | + def forward(self, x): | ||
| 150 | + a = torch.ops.aten.add.Tensor(x, 0) | ||
| 151 | + b = torch.ops.aten.add.Tensor(a, 0.0) | ||
| 152 | + c = torch.ops.aten.add.Tensor(b, torch.zeros_like(x)) | ||
| 153 | + return torch.ops.aten.relu.default(c) | ||
| 154 | + | ||
| 155 | + t = self._generate_tensor(shape, 'float32') | ||
| 156 | + m = M() | ||
| 157 | + gm = fx.symbolic_trace(m) | ||
| 158 | + ShapeProp(gm).propagate(t) | ||
| 159 | + fold_four_op_pass(gm.graph) | ||
| 160 | + gm.recompile() | ||
| 161 | + | ||
| 162 | + add_nodes = [n for n in gm.graph.nodes if 'add' in str(n.target)] | ||
| 163 | + self.assertEqual(len(add_nodes), 0, "所有 +0 都应该被折叠掉") | ||
| 164 | + | ||
| 165 | + | ||
| 166 | + def test_no_add_op_no_change(self): | ||
| 167 | + """ 图中没有 add 节点,pass 应安全通过 """ | ||
| 168 | + class M(torch.nn.Module): | ||
| 169 | + def forward(self, x): | ||
| 170 | + return torch.mul(x, 2.0) + torch.pow(x, 2) | ||
| 171 | + | ||
| 172 | + t = torch.randn(8) | ||
| 173 | + gm = fx.symbolic_trace(M()) | ||
| 174 | + ShapeProp(gm).propagate(t) | ||
| 175 | + | ||
| 176 | + orig_graph_str = str(gm.graph) | ||
| 177 | + fold_four_op_pass(gm.graph) | ||
| 178 | + self.assertEqual(orig_graph_str, str(gm.graph), "无 add 节点不应修改图") | ||
| 179 | + | ||
| 180 | + | ||
| 181 | + | ||
| 182 | + def test_add_zero_with_downstream_users(self, shape): | ||
| 183 | + """ +0 后有多个下游用户,都应正确替换 """ | ||
| 184 | + class M(torch.nn.Module): | ||
| 185 | + def forward(self, x): | ||
| 186 | + y = torch.add(x, 0) | ||
| 187 | + return (y * 2, y + 3, torch.relu(y)) | ||
| 188 | + | ||
| 189 | + t = self._generate_tensor(shape, 'float32') | ||
| 190 | + m = M() | ||
| 191 | + gm = fx.symbolic_trace(m) | ||
| 192 | + ShapeProp(gm).propagate(t) | ||
| 193 | + | ||
| 194 | + fold_four_op_pass(gm.graph) | ||
| 195 | + gm.recompile() | ||
| 196 | + | ||
| 197 | + self.assertEqual(m(t), gm(t)) | ||
| 198 | + | ||
| 199 | + | ||
| 200 | +instantiate_parametrized_tests(TestFoldAddPass) | ||
| 201 | + | ||
| 202 | + | ||
| 203 | +if __name__ == "__main__": | ||
| 204 | + run_tests() | ||
| @@ -0,0 +1,204 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch.fx as fx | ||
| 3 | +from torch.fx.passes.shape_prop import ShapeProp | ||
| 4 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 5 | +from testutils import TestUtils | ||
| 6 | +import torch_npu | ||
| 7 | +import torch_npu._inductor | ||
| 8 | +from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_cat | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class FoldCatModel(torch.nn.Module): | ||
| 12 | + def forward(self, t1, t2, t3, t4, t5): | ||
| 13 | + cat1 = torch.ops.aten.cat.default([t1, t2], 1) | ||
| 14 | + cat2 = torch.ops.aten.cat.default([cat1, t3], 1) | ||
| 15 | + cat3 = torch.ops.aten.cat.default([cat2, t4], 1) | ||
| 16 | + cat4 = torch.ops.aten.cat.default([cat3, t5], 1) | ||
| 17 | + return cat4 | ||
| 18 | + | ||
| 19 | + | ||
| 20 | +class TestFoldCatPass(TestUtils): | ||
| 21 | + def op_calc(self, t1, t2, t3, t4, t5): | ||
| 22 | + cat1 = torch.cat([t1, t2], dim=1) | ||
| 23 | + cat2 = torch.cat([cat1, t3], dim=1) | ||
| 24 | + cat3 = torch.cat([cat2, t4], dim=1) | ||
| 25 | + cat4 = torch.cat([cat3, t5], dim=1) | ||
| 26 | + return cat4 | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + def test_compile_cases(self, shape, dtype): | ||
| 32 | + t1 = self._generate_tensor(shape, dtype) | ||
| 33 | + t2 = self._generate_tensor(shape, dtype) | ||
| 34 | + t3 = self._generate_tensor(shape, dtype) | ||
| 35 | + t4 = self._generate_tensor(shape, dtype) | ||
| 36 | + t5 = self._generate_tensor(shape, dtype) | ||
| 37 | + std_result = self.op_calc(t1, t2, t3, t4, t5) | ||
| 38 | + with torch.no_grad(): | ||
| 39 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor") | ||
| 40 | + inductor_result = compiled_op_calc(t1, t2, t3, t4, t5) | ||
| 41 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 42 | + | ||
| 43 | + | ||
| 44 | + | ||
| 45 | + | ||
| 46 | + def test_ut_cases(self, shape, dtype): | ||
| 47 | + t1 = self._generate_tensor(shape, dtype) | ||
| 48 | + t2 = self._generate_tensor(shape, dtype) | ||
| 49 | + t3 = self._generate_tensor(shape, dtype) | ||
| 50 | + t4 = self._generate_tensor(shape, dtype) | ||
| 51 | + t5 = self._generate_tensor(shape, dtype) | ||
| 52 | + model = FoldCatModel() | ||
| 53 | + graph_module = fx.symbolic_trace(model) | ||
| 54 | + ShapeProp(graph_module).propagate(t1, t2, t3, t4, t5) | ||
| 55 | + | ||
| 56 | + # 应用优化 Pass | ||
| 57 | + fold_cat(graph_module.graph) | ||
| 58 | + graph_module.recompile() | ||
| 59 | + | ||
| 60 | + # 验证输出是否一致 | ||
| 61 | + std_result = model(t1, t2, t3, t4, t5) | ||
| 62 | + inductor_result = graph_module(t1, t2, t3, t4, t5) | ||
| 63 | + | ||
| 64 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 65 | + | ||
| 66 | + | ||
| 67 | + | ||
| 68 | + | ||
| 69 | + def test_deep_nested_same_axis_fold(self, shape, dtype): | ||
| 70 | + """多层嵌套同轴 cat,应完全折叠为单层 cat""" | ||
| 71 | + class M(torch.nn.Module): | ||
| 72 | + def forward(self, a, b, c, d, e): | ||
| 73 | + c1 = torch.ops.aten.cat.default([a, b], dim=1) | ||
| 74 | + c2 = torch.ops.aten.cat.default([c1, c], dim=1) | ||
| 75 | + c3 = torch.ops.aten.cat.default([c2, d], dim=1) | ||
| 76 | + c4 = torch.ops.aten.cat.default([c3, e], dim=1) | ||
| 77 | + return c4 | ||
| 78 | + | ||
| 79 | + tensors = [self._generate_tensor(shape, dtype) for _ in range(5)] | ||
| 80 | + model = M() | ||
| 81 | + gm = fx.symbolic_trace(model) | ||
| 82 | + ShapeProp(gm).propagate(*tensors) | ||
| 83 | + gm.graph.print_tabular() | ||
| 84 | + fold_cat(gm.graph) | ||
| 85 | + gm.graph.print_tabular() | ||
| 86 | + gm.recompile() | ||
| 87 | + | ||
| 88 | + cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default] | ||
| 89 | + self.assertEqual(len(cat_nodes), 1, "应折叠为单层 cat") | ||
| 90 | + self.assertEqual(len(cat_nodes[0].args[0]), 5, "输入应包含所有 5 个原始 tensor") | ||
| 91 | + | ||
| 92 | + self.assertEqual(model(*tensors), gm(*tensors)) | ||
| 93 | + | ||
| 94 | + | ||
| 95 | + | ||
| 96 | + | ||
| 97 | + def test_different_axis_no_fold(self, shape, dtype): | ||
| 98 | + """不同轴的 cat 不应折叠""" | ||
| 99 | + class M(torch.nn.Module): | ||
| 100 | + def forward(self, a, b, c): | ||
| 101 | + c1 = torch.ops.aten.cat.default([a, b], 0) # dim=0 | ||
| 102 | + c2 = torch.ops.aten.cat.default([c1, c], 1) # dim=1 → 不同 | ||
| 103 | + return c2 | ||
| 104 | + | ||
| 105 | + t1, t2 = [self._generate_tensor(shape, dtype) for _ in range(2)] | ||
| 106 | + shape1 = (shape[0] * 2, shape[1]) | ||
| 107 | + t3 = self._generate_tensor(shape1, dtype) | ||
| 108 | + model = M() | ||
| 109 | + gm = fx.symbolic_trace(model) | ||
| 110 | + ShapeProp(gm).propagate(t1, t2, t3) | ||
| 111 | + | ||
| 112 | + fold_cat(gm.graph) | ||
| 113 | + | ||
| 114 | + cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default] | ||
| 115 | + self.assertEqual(len(cat_nodes), 2, "不同轴不应折叠") | ||
| 116 | + | ||
| 117 | + | ||
| 118 | + | ||
| 119 | + | ||
| 120 | + def test_input_cat_has_multiple_users_no_fold(self, shape, dtype): | ||
| 121 | + """被多个下游使用的 cat 输入不应折叠""" | ||
| 122 | + class M(torch.nn.Module): | ||
| 123 | + def forward(self, a, b, c): | ||
| 124 | + inner = torch.ops.aten.cat.default([a, b], 1) | ||
| 125 | + out1 = inner + 1.0 | ||
| 126 | + out2 = inner * 2.0 | ||
| 127 | + out3 = torch.ops.aten.cat.default([inner, c], 1) | ||
| 128 | + return out1, out2, out3 | ||
| 129 | + | ||
| 130 | + t1, t2, t3 = [self._generate_tensor(shape, dtype) for _ in range(3)] | ||
| 131 | + model = M() | ||
| 132 | + gm = fx.symbolic_trace(model) | ||
| 133 | + ShapeProp(gm).propagate(t1, t2, t3) | ||
| 134 | + | ||
| 135 | + fold_cat(gm.graph) | ||
| 136 | + | ||
| 137 | + cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default] | ||
| 138 | + self.assertEqual(len(cat_nodes), 2, "有多个用户,不应折叠 inner cat") | ||
| 139 | + | ||
| 140 | + | ||
| 141 | + | ||
| 142 | + | ||
| 143 | + def test_last_dim_negative_and_positive_axis(self, shape, dtype): | ||
| 144 | + """最后一维使用 -1 和正数轴应视为相同""" | ||
| 145 | + class M(torch.nn.Module): | ||
| 146 | + def forward(self, a, b, c): | ||
| 147 | + inner = torch.ops.aten.cat.default([a, b], -1) # -1 | ||
| 148 | + outer = torch.ops.aten.cat.default([inner, c], 2) # 2 | ||
| 149 | + return outer | ||
| 150 | + | ||
| 151 | + t1, t2, t3 = [self._generate_tensor(shape, dtype) for _ in range(3)] | ||
| 152 | + model = M() | ||
| 153 | + gm = fx.symbolic_trace(model) | ||
| 154 | + ShapeProp(gm).propagate(t1, t2, t3) | ||
| 155 | + fold_cat(gm.graph) | ||
| 156 | + cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default] | ||
| 157 | + self.assertEqual(len(cat_nodes), 1, "最后一维 -1 和 2 应折叠") | ||
| 158 | + self.assertEqual(cat_nodes[0].args[1], -1) | ||
| 159 | + | ||
| 160 | + | ||
| 161 | + | ||
| 162 | + | ||
| 163 | + def test_mixed_cat_and_non_cat_inputs(self, shape, dtype): | ||
| 164 | + """cat 输入中混有非 cat 节点""" | ||
| 165 | + class M(torch.nn.Module): | ||
| 166 | + def forward(self, a, b, c): | ||
| 167 | + inner = torch.ops.aten.cat.default([a, b], 0) | ||
| 168 | + outer = torch.ops.aten.cat.default([inner, c * 1.5, torch.ops.aten.relu.default(c)], 0) | ||
| 169 | + return outer | ||
| 170 | + | ||
| 171 | + t1, t2, t3 = [self._generate_tensor(shape, dtype) for _ in range(3)] | ||
| 172 | + model = M() | ||
| 173 | + gm = fx.symbolic_trace(model) | ||
| 174 | + ShapeProp(gm).propagate(t1, t2, t3) | ||
| 175 | + | ||
| 176 | + fold_cat(gm.graph) | ||
| 177 | + | ||
| 178 | + cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default] | ||
| 179 | + self.assertEqual(len(cat_nodes), 1, "混有非 cat 输入,应保留一层 cat") | ||
| 180 | + self.assertEqual(len(cat_nodes[0].args[0]), 4, "应包含 inner + c*1.5 + relu(c)") | ||
| 181 | + | ||
| 182 | + | ||
| 183 | + def test_single_input_cat_no_fold(self): | ||
| 184 | + """只有一个输入的 cat 不应折叠""" | ||
| 185 | + class M(torch.nn.Module): | ||
| 186 | + def forward(self, x): | ||
| 187 | + return torch.ops.aten.cat.default([x], dim=0) | ||
| 188 | + | ||
| 189 | + t = self._generate_tensor((4, 5), 'float32') | ||
| 190 | + model = M() | ||
| 191 | + gm = fx.symbolic_trace(model) | ||
| 192 | + ShapeProp(gm).propagate(t) | ||
| 193 | + | ||
| 194 | + fold_cat(gm.graph) | ||
| 195 | + | ||
| 196 | + cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default] | ||
| 197 | + self.assertEqual(len(cat_nodes), 1, "单输入 cat 不应被移除") | ||
| 198 | + | ||
| 199 | + | ||
| 200 | +instantiate_parametrized_tests(TestFoldCatPass) | ||
| 201 | + | ||
| 202 | + | ||
| 203 | +if __name__ == "__main__": | ||
| 204 | + run_tests() | ||
| @@ -0,0 +1,201 @@ | |||||||||||||||||||||||||||||||||||||||||||
| 1 | +import torch | ||||||||||||||||||||||||||||||||||||||||||
| 2 | +import torch.fx as fx | ||||||||||||||||||||||||||||||||||||||||||
| 3 | +from torch.fx.passes.shape_prop import ShapeProp | ||||||||||||||||||||||||||||||||||||||||||
| 4 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||||||||||||||||||||||||||||||||||||||||||
| 5 | +from testutils import TestUtils | ||||||||||||||||||||||||||||||||||||||||||
| 6 | +import torch_npu | ||||||||||||||||||||||||||||||||||||||||||
| 7 | +import torch_npu._inductor | ||||||||||||||||||||||||||||||||||||||||||
| 8 | +from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_clone | ||||||||||||||||||||||||||||||||||||||||||
| 9 | + | ||||||||||||||||||||||||||||||||||||||||||
| 10 | + | ||||||||||||||||||||||||||||||||||||||||||
| 11 | +class FoldCloneModel(torch.nn.Module): | ||||||||||||||||||||||||||||||||||||||||||
| 12 | + def forward(self, t1): | ||||||||||||||||||||||||||||||||||||||||||
| 13 | + clone_1 = torch.ops.aten.clone.default(t1) | ||||||||||||||||||||||||||||||||||||||||||
| 14 | + relu_1 = torch.ops.aten.relu.default(clone_1) | ||||||||||||||||||||||||||||||||||||||||||
| 15 | + return relu_1 | ||||||||||||||||||||||||||||||||||||||||||
| 16 | + | ||||||||||||||||||||||||||||||||||||||||||
| 17 | + | ||||||||||||||||||||||||||||||||||||||||||
| 18 | +class TestFoldClonePass(TestUtils): | ||||||||||||||||||||||||||||||||||||||||||
| 19 | + def op_calc(self, t1): | ||||||||||||||||||||||||||||||||||||||||||
| 20 | + clone_1 = torch.clone(t1) | ||||||||||||||||||||||||||||||||||||||||||
| 21 | + output = torch.relu(clone_1) | ||||||||||||||||||||||||||||||||||||||||||
| 22 | + return output | ||||||||||||||||||||||||||||||||||||||||||
| 23 | + | ||||||||||||||||||||||||||||||||||||||||||
| 24 | + | ||||||||||||||||||||||||||||||||||||||||||
| 25 | + | ||||||||||||||||||||||||||||||||||||||||||
| 26 | + | ||||||||||||||||||||||||||||||||||||||||||
| 27 | + def test_compile_cases(self, shape, dtype): | ||||||||||||||||||||||||||||||||||||||||||
| 28 | + t1 = self._generate_tensor(shape, dtype) | ||||||||||||||||||||||||||||||||||||||||||
| 29 | + std_result = self.op_calc(t1) | ||||||||||||||||||||||||||||||||||||||||||
| 30 | + with torch.no_grad(): | ||||||||||||||||||||||||||||||||||||||||||
| 31 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor") | ||||||||||||||||||||||||||||||||||||||||||
| 32 | + inductor_result = compiled_op_calc(t1) | ||||||||||||||||||||||||||||||||||||||||||
| 33 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||||||||||||||||||||||||||||||||||||||||||
| 34 | + | ||||||||||||||||||||||||||||||||||||||||||
| 35 | + | ||||||||||||||||||||||||||||||||||||||||||
| 36 | + | ||||||||||||||||||||||||||||||||||||||||||
| 37 | + | ||||||||||||||||||||||||||||||||||||||||||
| 38 | + def test_ut_cases(self, shape, dtype): | ||||||||||||||||||||||||||||||||||||||||||
| 39 | + t1 = self._generate_tensor(shape, dtype) | ||||||||||||||||||||||||||||||||||||||||||
| 40 | + model = FoldCloneModel() | ||||||||||||||||||||||||||||||||||||||||||
| 41 | + graph_module = fx.symbolic_trace(model) | ||||||||||||||||||||||||||||||||||||||||||
| 42 | + ShapeProp(graph_module).propagate(t1) | ||||||||||||||||||||||||||||||||||||||||||
| 43 | + | ||||||||||||||||||||||||||||||||||||||||||
| 44 | + # 应用优化 Pass | ||||||||||||||||||||||||||||||||||||||||||
| 45 | + fold_clone(graph_module.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 46 | + graph_module.recompile() | ||||||||||||||||||||||||||||||||||||||||||
| 47 | + | ||||||||||||||||||||||||||||||||||||||||||
| 48 | + # 验证输出是否一致 | ||||||||||||||||||||||||||||||||||||||||||
| 49 | + std_result = model(t1) | ||||||||||||||||||||||||||||||||||||||||||
| 50 | + inductor_result = graph_module(t1) | ||||||||||||||||||||||||||||||||||||||||||
| 51 | + | ||||||||||||||||||||||||||||||||||||||||||
| 52 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||||||||||||||||||||||||||||||||||||||||||
| 53 | + | ||||||||||||||||||||||||||||||||||||||||||
| 54 | + | ||||||||||||||||||||||||||||||||||||||||||
| 55 | + | ||||||||||||||||||||||||||||||||||||||||||
| 56 | + | ||||||||||||||||||||||||||||||||||||||||||
| 57 | + def test_clone_same_memory_format_should_fold(self, shape, dtype): | ||||||||||||||||||||||||||||||||||||||||||
| 58 | + """memory_format 相同 → 应折叠删除 clone""" | ||||||||||||||||||||||||||||||||||||||||||
| 59 | + class M(torch.nn.Module): | ||||||||||||||||||||||||||||||||||||||||||
| 60 | + def forward(self, x): | ||||||||||||||||||||||||||||||||||||||||||
| 61 | + cloned = torch.ops.aten.clone.default(x) # 默认 memory_format 相同 | ||||||||||||||||||||||||||||||||||||||||||
| 62 | + return cloned + 1.0 | ||||||||||||||||||||||||||||||||||||||||||
| 63 | + | ||||||||||||||||||||||||||||||||||||||||||
| 64 | + x = self._generate_tensor(shape, dtype) | ||||||||||||||||||||||||||||||||||||||||||
| 65 | + model = M() | ||||||||||||||||||||||||||||||||||||||||||
| 66 | + gm = fx.symbolic_trace(model) | ||||||||||||||||||||||||||||||||||||||||||
| 67 | + ShapeProp(gm).propagate(x) | ||||||||||||||||||||||||||||||||||||||||||
| 68 | + fold_clone(gm.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 69 | + gm.recompile() | ||||||||||||||||||||||||||||||||||||||||||
| 70 | + | ||||||||||||||||||||||||||||||||||||||||||
| 71 | + clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default] | ||||||||||||||||||||||||||||||||||||||||||
| 72 | + self.assertEqual(len(clone_nodes), 0, "相同 memory_format 的 clone 应被删除") | ||||||||||||||||||||||||||||||||||||||||||
| 73 | + | ||||||||||||||||||||||||||||||||||||||||||
| 74 | + self.assertEqual(model(x), gm(x)) | ||||||||||||||||||||||||||||||||||||||||||
| 75 | + | ||||||||||||||||||||||||||||||||||||||||||
| 76 | + | ||||||||||||||||||||||||||||||||||||||||||
| 77 | + | ||||||||||||||||||||||||||||||||||||||||||
| 78 | + | ||||||||||||||||||||||||||||||||||||||||||
| 79 | + def test_clone_different_memory_format_no_fold(self, shape, dtype): | ||||||||||||||||||||||||||||||||||||||||||
| 80 | + """memory_format 不同 → 不折叠""" | ||||||||||||||||||||||||||||||||||||||||||
| 81 | + class M(torch.nn.Module): | ||||||||||||||||||||||||||||||||||||||||||
| 82 | + def forward(self, x): | ||||||||||||||||||||||||||||||||||||||||||
| 83 | + # 显式指定不同 memory_format(假设原 x 是 contiguous) | ||||||||||||||||||||||||||||||||||||||||||
| 84 | + cloned = torch.ops.aten.clone.default(x, memory_format=torch.preserve_format) | ||||||||||||||||||||||||||||||||||||||||||
| 85 | + # 如果想强制不同,可以用 channels_last,但需确保 dtype 支持 | ||||||||||||||||||||||||||||||||||||||||||
| 86 | + return cloned | ||||||||||||||||||||||||||||||||||||||||||
| 87 | + | ||||||||||||||||||||||||||||||||||||||||||
| 88 | + x = self._generate_tensor(shape, dtype) | ||||||||||||||||||||||||||||||||||||||||||
| 89 | + model = M() | ||||||||||||||||||||||||||||||||||||||||||
| 90 | + gm = fx.symbolic_trace(model) | ||||||||||||||||||||||||||||||||||||||||||
| 91 | + ShapeProp(gm).propagate(x) | ||||||||||||||||||||||||||||||||||||||||||
| 92 | + | ||||||||||||||||||||||||||||||||||||||||||
| 93 | + fold_clone(gm.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 94 | + | ||||||||||||||||||||||||||||||||||||||||||
| 95 | + clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default] | ||||||||||||||||||||||||||||||||||||||||||
| 96 | + self.assertEqual(len(clone_nodes), 1, "memory_format 不同不应折叠") | ||||||||||||||||||||||||||||||||||||||||||
🟡 Medium Priority 变更行:第 79-96 行( 证据链:
失败模式:如果后续有人误修改了 建议:将模型改为 clone 不直接作为输出,让 clone 进入 fold_clone 的候选列表,真正触发 memory_format 比较路径。例如: ![]() ![]() | |||||||||||||||||||||||||||||||||||||||||||
| 97 | + | ||||||||||||||||||||||||||||||||||||||||||
| 98 | + | ||||||||||||||||||||||||||||||||||||||||||
| 99 | + | ||||||||||||||||||||||||||||||||||||||||||
| 100 | + def test_clone_in_output_path_no_fold(self, shape): | ||||||||||||||||||||||||||||||||||||||||||
| 101 | + """clone 是 output 的直接/间接输入 → 不折叠""" | ||||||||||||||||||||||||||||||||||||||||||
| 102 | + class M(torch.nn.Module): | ||||||||||||||||||||||||||||||||||||||||||
| 103 | + def forward(self, x): | ||||||||||||||||||||||||||||||||||||||||||
| 104 | + cloned = torch.ops.aten.clone.default(x) | ||||||||||||||||||||||||||||||||||||||||||
| 105 | + return cloned # 直接输出 clone | ||||||||||||||||||||||||||||||||||||||||||
| 106 | + | ||||||||||||||||||||||||||||||||||||||||||
| 107 | + x = self._generate_tensor(shape, 'float32') | ||||||||||||||||||||||||||||||||||||||||||
| 108 | + model = M() | ||||||||||||||||||||||||||||||||||||||||||
| 109 | + gm = fx.symbolic_trace(model) | ||||||||||||||||||||||||||||||||||||||||||
| 110 | + ShapeProp(gm).propagate(x) | ||||||||||||||||||||||||||||||||||||||||||
| 111 | + fold_clone(gm.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 112 | + clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default] | ||||||||||||||||||||||||||||||||||||||||||
| 113 | + self.assertEqual(len(clone_nodes), 1, "clone 在 output 路径上不应折叠") | ||||||||||||||||||||||||||||||||||||||||||
| 114 | + | ||||||||||||||||||||||||||||||||||||||||||
| 115 | + | ||||||||||||||||||||||||||||||||||||||||||
| 116 | + | ||||||||||||||||||||||||||||||||||||||||||
| 117 | + def test_multi_clone_some_fold_some_not(self, shape): | ||||||||||||||||||||||||||||||||||||||||||
| 118 | + """多个 clone,只有部分可折叠""" | ||||||||||||||||||||||||||||||||||||||||||
| 119 | + class M(torch.nn.Module): | ||||||||||||||||||||||||||||||||||||||||||
| 120 | + def forward(self, x, y): | ||||||||||||||||||||||||||||||||||||||||||
| 121 | + c1 = torch.ops.aten.clone.default(x) | ||||||||||||||||||||||||||||||||||||||||||
| 122 | + c2 = torch.ops.aten.clone.default(y, memory_format=torch.channels_last) | ||||||||||||||||||||||||||||||||||||||||||
| 123 | + c3 = torch.ops.aten.clone.default(c1) | ||||||||||||||||||||||||||||||||||||||||||
| 124 | + return c1 + c2 + c3 | ||||||||||||||||||||||||||||||||||||||||||
| 125 | + | ||||||||||||||||||||||||||||||||||||||||||
| 126 | + x = self._generate_tensor(shape, 'float32') | ||||||||||||||||||||||||||||||||||||||||||
| 127 | + y = self._generate_tensor(shape, 'float32') | ||||||||||||||||||||||||||||||||||||||||||
| 128 | + model = M() | ||||||||||||||||||||||||||||||||||||||||||
| 129 | + gm = fx.symbolic_trace(model) | ||||||||||||||||||||||||||||||||||||||||||
| 130 | + ShapeProp(gm).propagate(x, y) | ||||||||||||||||||||||||||||||||||||||||||
| 131 | + fold_clone(gm.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 132 | + clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default] | ||||||||||||||||||||||||||||||||||||||||||
| 133 | + # 预期:c1 和 c3 折叠,c2 保留(或根据你的 memory_format 判断) | ||||||||||||||||||||||||||||||||||||||||||
| 134 | + self.assertLessEqual(len(clone_nodes), 1) | ||||||||||||||||||||||||||||||||||||||||||
| 135 | + | ||||||||||||||||||||||||||||||||||||||||||
| 136 | + | ||||||||||||||||||||||||||||||||||||||||||
| 137 | + | ||||||||||||||||||||||||||||||||||||||||||
| 138 | + def test_clone_with_downstream_users(self, shape): | ||||||||||||||||||||||||||||||||||||||||||
| 139 | + """clone 被折叠后,下游多个用户应正确指向原输入""" | ||||||||||||||||||||||||||||||||||||||||||
| 140 | + class M(torch.nn.Module): | ||||||||||||||||||||||||||||||||||||||||||
| 141 | + def forward(self, x): | ||||||||||||||||||||||||||||||||||||||||||
| 142 | + cloned = torch.ops.aten.clone.default(x) | ||||||||||||||||||||||||||||||||||||||||||
| 143 | + a = cloned * 2 | ||||||||||||||||||||||||||||||||||||||||||
| 144 | + b = cloned + 3 | ||||||||||||||||||||||||||||||||||||||||||
| 145 | + return a, b | ||||||||||||||||||||||||||||||||||||||||||
| 146 | + | ||||||||||||||||||||||||||||||||||||||||||
| 147 | + x = self._generate_tensor(shape, 'float32') | ||||||||||||||||||||||||||||||||||||||||||
| 148 | + model = M() | ||||||||||||||||||||||||||||||||||||||||||
| 149 | + gm = fx.symbolic_trace(model) | ||||||||||||||||||||||||||||||||||||||||||
| 150 | + ShapeProp(gm).propagate(x) | ||||||||||||||||||||||||||||||||||||||||||
| 151 | + | ||||||||||||||||||||||||||||||||||||||||||
| 152 | + fold_clone(gm.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 153 | + gm.recompile() | ||||||||||||||||||||||||||||||||||||||||||
| 154 | + | ||||||||||||||||||||||||||||||||||||||||||
| 155 | + clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default] | ||||||||||||||||||||||||||||||||||||||||||
| 156 | + self.assertEqual(len(clone_nodes), 0) | ||||||||||||||||||||||||||||||||||||||||||
| 157 | + | ||||||||||||||||||||||||||||||||||||||||||
| 158 | + self.assertEqual(model(x), gm(x)) | ||||||||||||||||||||||||||||||||||||||||||
| 159 | + | ||||||||||||||||||||||||||||||||||||||||||
| 160 | + | ||||||||||||||||||||||||||||||||||||||||||
| 161 | + def test_no_clone_node_no_change(self): | ||||||||||||||||||||||||||||||||||||||||||
| 162 | + """图中无 clone 节点,pass 不应修改图""" | ||||||||||||||||||||||||||||||||||||||||||
| 163 | + class M(torch.nn.Module): | ||||||||||||||||||||||||||||||||||||||||||
| 164 | + def forward(self, x): | ||||||||||||||||||||||||||||||||||||||||||
| 165 | + return x + x * 2.0 | ||||||||||||||||||||||||||||||||||||||||||
| 166 | + | ||||||||||||||||||||||||||||||||||||||||||
| 167 | + x = torch.randn(8, 8) | ||||||||||||||||||||||||||||||||||||||||||
| 168 | + gm = fx.symbolic_trace(M()) | ||||||||||||||||||||||||||||||||||||||||||
| 169 | + ShapeProp(gm).propagate(x) | ||||||||||||||||||||||||||||||||||||||||||
| 170 | + | ||||||||||||||||||||||||||||||||||||||||||
| 171 | + orig_graph = str(gm.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 172 | + fold_clone(gm.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 173 | + self.assertEqual(orig_graph, str(gm.graph)) | ||||||||||||||||||||||||||||||||||||||||||
| 174 | + | ||||||||||||||||||||||||||||||||||||||||||
| 175 | + | ||||||||||||||||||||||||||||||||||||||||||
| 176 | + | ||||||||||||||||||||||||||||||||||||||||||
| 177 | + def test_clone_without_tensor_meta_no_fold(self, shape): | ||||||||||||||||||||||||||||||||||||||||||
| 178 | + """输入没有 tensor_meta → 不折叠(安全检查)""" | ||||||||||||||||||||||||||||||||||||||||||
| 179 | + class M(torch.nn.Module): | ||||||||||||||||||||||||||||||||||||||||||
| 180 | + def forward(self, x): | ||||||||||||||||||||||||||||||||||||||||||
| 181 | + cloned = torch.ops.aten.clone.default(x) | ||||||||||||||||||||||||||||||||||||||||||
| 182 | + return cloned | ||||||||||||||||||||||||||||||||||||||||||
| 183 | + | ||||||||||||||||||||||||||||||||||||||||||
| 184 | + x = self._generate_tensor(shape, 'float32') | ||||||||||||||||||||||||||||||||||||||||||
| 185 | + gm = fx.symbolic_trace(M()) | ||||||||||||||||||||||||||||||||||||||||||
| 186 | + # 故意移除 tensor_meta | ||||||||||||||||||||||||||||||||||||||||||
| 187 | + for node in gm.graph.nodes: | ||||||||||||||||||||||||||||||||||||||||||
| 188 | + if "tensor_meta" in node.meta: | ||||||||||||||||||||||||||||||||||||||||||
| 189 | + del node.meta["tensor_meta"] | ||||||||||||||||||||||||||||||||||||||||||
| 190 | + | ||||||||||||||||||||||||||||||||||||||||||
| 191 | + fold_clone(gm.graph) | ||||||||||||||||||||||||||||||||||||||||||
| 192 | + | ||||||||||||||||||||||||||||||||||||||||||
| 193 | + clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default] | ||||||||||||||||||||||||||||||||||||||||||
| 194 | + self.assertEqual(len(clone_nodes), 1, "缺少 tensor_meta 不应折叠") | ||||||||||||||||||||||||||||||||||||||||||
🟡 Medium Priority 变更行:第 176-194 行( 证据链:
失败模式:如果后续有人误修改了 建议:将模型改为 clone 不直接作为输出(如 改动建议
![]() ![]() | |||||||||||||||||||||||||||||||||||||||||||
| 195 | + | ||||||||||||||||||||||||||||||||||||||||||
| 196 | + | ||||||||||||||||||||||||||||||||||||||||||
| 197 | +instantiate_parametrized_tests(TestFoldClonePass) | ||||||||||||||||||||||||||||||||||||||||||
| 198 | + | ||||||||||||||||||||||||||||||||||||||||||
| 199 | + | ||||||||||||||||||||||||||||||||||||||||||
| 200 | +if __name__ == "__main__": | ||||||||||||||||||||||||||||||||||||||||||
| 201 | + run_tests() | ||||||||||||||||||||||||||||||||||||||||||
| @@ -0,0 +1,201 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch.fx as fx | ||
| 3 | +from torch.fx.passes.shape_prop import ShapeProp | ||
| 4 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 5 | +from testutils import TestUtils | ||
| 6 | +import torch_npu | ||
| 7 | +import torch_npu._inductor | ||
| 8 | +from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_four_op_pass | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class FoldDivModel(torch.nn.Module): | ||
| 12 | + def forward(self, t1, t2, t3): | ||
| 13 | + div_1 = torch.ops.aten.div(t1, 1) | ||
| 14 | + div_output = torch.relu(div_1) | ||
| 15 | + return div_output | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +class TestFoldDivPass(TestUtils): | ||
| 19 | + def op_calc(self, t1): | ||
| 20 | + div_1 = torch.ops.aten.div(t1, 1) | ||
| 21 | + div_output = torch.relu(div_1) | ||
| 22 | + return div_output | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + def test_compile_cases(self, shape, dtype): | ||
| 28 | + t1 = self._generate_tensor(shape, dtype) | ||
| 29 | + std_result = self.op_calc(t1) | ||
| 30 | + with torch.no_grad(): | ||
| 31 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor") | ||
| 32 | + inductor_result = compiled_op_calc(t1) | ||
| 33 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 34 | + | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + | ||
| 38 | + def test_ut_cases(self, shape, dtype): | ||
| 39 | + t1 = self._generate_tensor(shape, dtype) | ||
| 40 | + t2 = self._generate_tensor(shape, dtype) | ||
| 41 | + t3 = self._generate_tensor(shape, dtype) | ||
| 42 | + model = FoldDivModel() | ||
| 43 | + graph_module = fx.symbolic_trace(model) | ||
| 44 | + ShapeProp(graph_module).propagate(t1, t2, t3) | ||
| 45 | + | ||
| 46 | + # 应用优化 Pass | ||
| 47 | + fold_four_op_pass(graph_module.graph) | ||
| 48 | + graph_module.recompile() | ||
| 49 | + | ||
| 50 | + # 验证输出是否一致 | ||
| 51 | + std_result = model(t1, t2, t3) | ||
| 52 | + inductor_result = graph_module(t1, t2, t3) | ||
| 53 | + | ||
| 54 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 55 | + | ||
| 56 | + | ||
| 57 | + | ||
| 58 | + | ||
| 59 | + def test_div_by_one_right_fold(self, shape, dtype): | ||
| 60 | + """ x / 1.0 → should be x """ | ||
| 61 | + class M(torch.nn.Module): | ||
| 62 | + def forward(self, x): | ||
| 63 | + y = torch.div(x, 1.0) | ||
| 64 | + return torch.relu(y * 2.0) | ||
| 65 | + | ||
| 66 | + t = self._generate_tensor(shape, dtype) | ||
| 67 | + m = M() | ||
| 68 | + gm = fx.symbolic_trace(m) | ||
| 69 | + ShapeProp(gm).propagate(t) | ||
| 70 | + | ||
| 71 | + fold_four_op_pass(gm.graph) | ||
| 72 | + gm.recompile() | ||
| 73 | + | ||
| 74 | + self.assertFalse( | ||
| 75 | + any(n.target in (torch.ops.aten.div.Tensor, torch.ops.aten.div.Scalar) | ||
| 76 | + for n in gm.graph.nodes), | ||
| 77 | + "x / 1.0 should be folded to x" | ||
| 78 | + ) | ||
| 79 | + torch.testing.assert_close(m(t), gm(t), atol=1e-5, rtol=1e-5) | ||
| 80 | + | ||
| 81 | + | ||
| 82 | + | ||
| 83 | + def test_div_by_one_int_fold(self, shape): | ||
| 84 | + """ x / 1 (int) should be folded """ | ||
| 85 | + class M(torch.nn.Module): | ||
| 86 | + def forward(self, x): | ||
| 87 | + y = torch.div(x, 1) | ||
| 88 | + return y + x.mean() | ||
| 89 | + | ||
| 90 | + t = self._generate_tensor(shape, 'float32') | ||
| 91 | + m = M() | ||
| 92 | + gm = fx.symbolic_trace(m) | ||
| 93 | + ShapeProp(gm).propagate(t) | ||
| 94 | + | ||
| 95 | + fold_four_op_pass(gm.graph) | ||
| 96 | + gm.recompile() | ||
| 97 | + | ||
| 98 | + self.assertFalse(any('div' in str(n.target) for n in gm.graph.nodes if n.op == 'call_function')) | ||
| 99 | + self.assertEqual(m(t), gm(t), "x / 1 (int) should be folded") | ||
| 100 | + | ||
| 101 | + | ||
| 102 | + | ||
| 103 | + def test_rdiv_one_left_fold(self, shape): | ||
| 104 | + """ 1.0 / x should be folded """ | ||
| 105 | + class M(torch.nn.Module): | ||
| 106 | + def forward(self, x): | ||
| 107 | + y = torch.div(1.0, x) # div.Scalar(1.0, x) 或 rdiv | ||
| 108 | + return torch.sigmoid(y) | ||
| 109 | + | ||
| 110 | + t = self._generate_tensor(shape, 'float32') + 0.1 # 避免除零 | ||
| 111 | + m = M() | ||
| 112 | + gm = fx.symbolic_trace(m) | ||
| 113 | + ShapeProp(gm).propagate(t) | ||
| 114 | + | ||
| 115 | + fold_four_op_pass(gm.graph) | ||
| 116 | + gm.recompile() | ||
| 117 | + | ||
| 118 | + torch.testing.assert_close(m(t), gm(t), atol=1e-5, rtol=1e-5) | ||
🟡 Medium Priority 测试名和 docstring 都声称 "1.0 / x should be folded",但测试体中仅通过 查看 后果:测试会通过(pass 未破坏正确性),但并未验证其声称的折叠行为,造成虚假的测试覆盖信心。 修复方向:要么在 pass 中补充 rdiv 折叠逻辑并加入对应的节点断言;要么将测试改名/改 docstring 为验证"不对 1.0/x 做错误折叠"。 建议:要么为 fold_four_op_pass 补充 rdiv 折叠能力并在测试中加入节点计数断言;要么将测试名和 docstring 改为表示验证"1.0/x 不被错误折叠"。 ![]() ![]() | |||
| 119 | + | ||
| 120 | + | ||
| 121 | + | ||
| 122 | + def test_div_one_chain_fold(self, shape): | ||
| 123 | + """ all / 1.0 should be folded """ | ||
| 124 | + class M(torch.nn.Module): | ||
| 125 | + def forward(self, x): | ||
| 126 | + a = torch.div(x, 1.0) | ||
| 127 | + b = torch.div(a, 1) | ||
| 128 | + c = torch.div(b, torch.ones_like(x)) | ||
| 129 | + return torch.relu(c + x) | ||
| 130 | + | ||
| 131 | + t = self._generate_tensor(shape, 'float32') | ||
| 132 | + m = M() | ||
| 133 | + gm = fx.symbolic_trace(m) | ||
| 134 | + ShapeProp(gm).propagate(t) | ||
| 135 | + | ||
| 136 | + fold_four_op_pass(gm.graph) | ||
| 137 | + gm.recompile() | ||
| 138 | + | ||
| 139 | + div_nodes = [n for n in gm.graph.nodes if 'div' in str(n.target)] | ||
| 140 | + self.assertEqual(len(div_nodes), 0, "all / 1.0 should be folded") | ||
| 141 | + | ||
| 142 | + | ||
| 143 | + | ||
| 144 | + def test_div_non_one_no_fold(self, shape): | ||
| 145 | + """ should not fold since div not 1 """ | ||
| 146 | + class M(torch.nn.Module): | ||
| 147 | + def forward(self, x): | ||
| 148 | + y = torch.div(x, 2.0) | ||
| 149 | + return torch.tanh(y) | ||
| 150 | + | ||
| 151 | + t = self._generate_tensor(shape, 'float32') | ||
| 152 | + m = M() | ||
| 153 | + gm = fx.symbolic_trace(m) | ||
| 154 | + ShapeProp(gm).propagate(t) | ||
| 155 | + | ||
| 156 | + before = len([n for n in gm.graph.nodes if n.op == 'call_function']) | ||
| 157 | + fold_four_op_pass(gm.graph) | ||
| 158 | + after = len([n for n in gm.graph.nodes if n.op == 'call_function']) | ||
| 159 | + | ||
| 160 | + self.assertEqual(before, after, "should not fold since div not 1") | ||
| 161 | + | ||
| 162 | + | ||
| 163 | + | ||
| 164 | + def test_div_one_multi_users(self, shape): | ||
| 165 | + """ there are multiple usrs after div /1 all of them should be replaced """ | ||
| 166 | + class M(torch.nn.Module): | ||
| 167 | + def forward(self, x): | ||
| 168 | + y = torch.div(x, 1.0) | ||
| 169 | + return y * 3.0, torch.relu(y), y + x | ||
| 170 | + | ||
| 171 | + t = self._generate_tensor(shape, 'float32') | ||
| 172 | + m = M() | ||
| 173 | + gm = fx.symbolic_trace(m) | ||
| 174 | + ShapeProp(gm).propagate(t) | ||
| 175 | + | ||
| 176 | + fold_four_op_pass(gm.graph) | ||
| 177 | + gm.recompile() | ||
| 178 | + | ||
| 179 | + self.assertEqual(m(t), gm(t), "there are multiple usrs after div /1 all of them should be replaced") | ||
| 180 | + | ||
| 181 | + | ||
| 182 | + def test_no_div_op_no_change(self): | ||
| 183 | + """ no div node, pass should not modify graph """ | ||
| 184 | + class M(torch.nn.Module): | ||
| 185 | + def forward(self, x): | ||
| 186 | + return torch.mul(x, x) + torch.add(x, 1.0) | ||
| 187 | + | ||
| 188 | + t = torch.randn(4, 4) | ||
| 189 | + gm = fx.symbolic_trace(M()) | ||
| 190 | + ShapeProp(gm).propagate(t) | ||
| 191 | + | ||
| 192 | + orig_str = str(gm.graph) | ||
| 193 | + fold_four_op_pass(gm.graph) | ||
| 194 | + self.assertEqual(orig_str, str(gm.graph), "no div node, pass should not modify graph") | ||
| 195 | + | ||
| 196 | + | ||
| 197 | +instantiate_parametrized_tests(TestFoldDivPass) | ||
| 198 | + | ||
| 199 | + | ||
| 200 | +if __name__ == "__main__": | ||
| 201 | + run_tests() | ||


🟠 High Priority
测试
test_add_zero_left_fold在第 73-77 行断言图中不应存在任何 add 节点:self.assertFalse(any(n.target in (torch.add, torch.ops.aten.add.Tensor, torch.ops.aten.add.Scalar) for n in gm.graph.nodes))。但模型定义(第 59-62 行)中:
y = torch.ops.aten.add.Tensor(0, x)折叠后 y→x,而torch.relu(y + 1.0)中的y + 1.0仍会生成一个新的 add 节点(add.Tensor(x, 1.0)或add.Scalar(x, 1.0),因为 1.0 ≠ 0,不会被 fold_four_op_pass 折叠)。因此图中必然残留一个 add 节点,assertFalse将失败。触发条件:
fold_four_op_pass正确折叠0 + x后,x + 1.0的 add 节点仍然存在。修复方向:将断言改为验证折叠后的图中仅剩余
x + 1.0这一个 add 节点(即len(add_nodes) == 1),或修改模型使y + 1.0不产生 add 节点(例如改为torch.mul(y, 2.0))。建议:将断言从 assertFalse(零 add 节点)改为 assertEqual(len(add_nodes), 1),只验证
0 + x被折叠,而x + 1.0正常保留。