已合并
[test] add fx graph api coverage for find_nodes and graph_copy #36539
AZT创建于 5月23日
[test] add fx graph api coverage for find_nodes and graph_copy #36539
已合并
共 1 个文件变更+112-2
| @@ -1,10 +1,12 @@ | |||
| 1 | """ | 1 | """ |
| 2 | Add validation cases for torch.fx Graph APIs on NPU: | 2 | Add validation cases for torch.fx Graph APIs on NPU: |
| 3 | 3 | ||
| 4 | -1. This file adds lightweight direct validations for torch.fx Graph APIs on NPU. | 4 | +1. PyTorch community lacks sufficient and direct API validations for some APIs, so this file is added. |
| 5 | 2. This file validates torch.fx.Graph.inserting_after, | 5 | 2. This file validates torch.fx.Graph.inserting_after, |
| 6 | torch.fx.Graph.inserting_before, torch.fx.graph.magic_methods.format, | 6 | torch.fx.Graph.inserting_before, torch.fx.graph.magic_methods.format, |
| 7 | - and torch.fx.graph.inplace_methods.format. | 7 | + torch.fx.graph.inplace_methods.format, |
| 8 | + torch.fx.Graph.find_nodes, torch.fx.Graph.graph_copy, | ||
| 9 | + torch.fx.Graph.erase_node, and torch.fx.Graph.get_attr. | ||
| 8 | """ | 10 | """ |
| 9 | 11 | ||
| 10 | import operator | 12 | import operator |
| @@ -17,6 +19,7 @@ device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else | |||
| 17 | 19 | ||
| 18 | 20 | ||
| 19 | class TestFxGraphApiNPU(TestCase): | 21 | class TestFxGraphApiNPU(TestCase): |
| 22 | + | ||
| 20 | def test_graph_inserting_after(self): | 23 | def test_graph_inserting_after(self): |
| 21 | graph = torch.fx.Graph() | 24 | graph = torch.fx.Graph() |
| 22 | x = graph.placeholder("x") | 25 | x = graph.placeholder("x") |
| @@ -104,6 +107,113 @@ class TestFxGraphApiNPU(TestCase): | |||
| 104 | self.assertEqual(output, torch.full((2, 3), 3.0).to(device_type)) | 107 | self.assertEqual(output, torch.full((2, 3), 3.0).to(device_type)) |
| 105 | self.assertEqual(input_tensor, torch.full((2, 3), 3.0).to(device_type)) | 108 | self.assertEqual(input_tensor, torch.full((2, 3), 3.0).to(device_type)) |
| 106 | 109 | ||
| 110 | + # Test find_nodes: query nodes by op type, target, and sort parameter. | ||
| 111 | + def test_find_nodes(self): | ||
| 112 | + class M(torch.nn.Module): | ||
| 113 | + def __init__(self): | ||
| 114 | + super().__init__() | ||
| 115 | + self.p = torch.nn.Parameter(torch.tensor([1.0])) | ||
| 116 | + self.linear = torch.nn.Linear(3, 3) | ||
| 117 | + | ||
| 118 | + def forward(self, x): | ||
| 119 | + return torch.relu(x) + self.linear(x) * self.p | ||
| 120 | + | ||
| 121 | + m = M().to(device_type) | ||
| 122 | + gm = symbolic_trace(m) | ||
| 123 | + # call_function + target | ||
| 124 | + relu_nodes = gm.graph.find_nodes(op="call_function", target=torch.relu) | ||
| 125 | + self.assertEqual(len(relu_nodes), 1) | ||
| 126 | + self.assertEqual(relu_nodes[0].op, "call_function") | ||
| 127 | + # call_module | ||
| 128 | + module_nodes = gm.graph.find_nodes(op="call_module") | ||
| 129 | + self.assertEqual(len(module_nodes), 1) | ||
| 130 | + self.assertEqual(module_nodes[0].target, "linear") | ||
| 131 | + # get_attr | ||
| 132 | + attr_nodes = gm.graph.find_nodes(op="get_attr") | ||
| 133 | + self.assertEqual(len(attr_nodes), 1) | ||
| 134 | + self.assertEqual(attr_nodes[0].target, "p") | ||
| 135 | + # sort requires target | ||
| 136 | + sorted_relu = gm.graph.find_nodes(op="call_function", target=torch.relu, sort=True) | ||
| 137 | + unsorted_relu = gm.graph.find_nodes(op="call_function", target=torch.relu, sort=False) | ||
| 138 | + self.assertEqual(set(sorted_relu), set(unsorted_relu)) | ||
| 139 | + # empty result | ||
| 140 | + self.assertEqual(gm.graph.find_nodes(op="call_function", target=torch.sigmoid), []) | ||
| 141 | + | ||
| 142 | + input_tensor = torch.randn(2, 3).to(device_type) | ||
| 143 | + self.assertEqual(gm(input_tensor), m(input_tensor)) | ||
| 144 | + | ||
| 145 | + # Test graph_copy: copy nodes between graphs with val_map and return_output_node. | ||
| 146 | + def test_graph_copy(self): | ||
| 147 | + g_src = torch.fx.Graph() | ||
| 148 | + x = g_src.placeholder("x") | ||
| 149 | + neg = g_src.call_function(torch.neg, (x,)) | ||
| 150 | + relu = g_src.call_function(torch.relu, (neg,)) | ||
| 151 | + g_src.output((relu,)) | ||
| 152 | + | ||
| 153 | + # Basic copy with val_map populated. | ||
| 154 | + g_dst = torch.fx.Graph() | ||
| 155 | + val_map = {} | ||
| 156 | + g_dst.graph_copy(g_src, val_map) | ||
| 157 | + self.assertEqual(len(val_map), 3) | ||
| 158 | + x_dst, neg_dst, relu_dst = val_map[x], val_map[neg], val_map[relu] | ||
| 159 | + self.assertEqual(neg_dst.args, (x_dst,)) | ||
| 160 | + self.assertEqual(relu_dst.args, (neg_dst,)) | ||
| 161 | + | ||
| 162 | + # return_output_node returns the source output node. | ||
| 163 | + g_dst2 = torch.fx.Graph() | ||
| 164 | + rv = g_dst2.graph_copy(g_src, {}, return_output_node=True) | ||
| 165 | + self.assertIsNotNone(rv) | ||
| 166 | + _, src_output = rv | ||
| 167 | + self.assertEqual(src_output.op, "output") | ||
| 168 | + | ||
| 169 | + # val_map reuse: existing node replaces source node. | ||
| 170 | + g_dst3 = torch.fx.Graph() | ||
| 171 | + existing = g_dst3.placeholder("existing") | ||
| 172 | + val_map3 = {x: existing} | ||
| 173 | + g_dst3.graph_copy(g_src, val_map3) | ||
| 174 | + self.assertEqual(val_map3[neg].args[0], existing) | ||
| 175 | + | ||
| 176 | + # Test erase_node: remove a node and rewire its users. | ||
| 177 | + def test_erase_node(self): | ||
| 178 | + graph = torch.fx.Graph() | ||
| 179 | + x = graph.placeholder("x") | ||
| 180 | + neg = graph.call_function(torch.neg, (x,)) | ||
| 181 | + relu = graph.call_function(torch.relu, (neg,)) | ||
| 182 | + graph.output(relu) | ||
| 183 | + | ||
| 184 | + neg.replace_all_uses_with(x) | ||
| 185 | + graph.erase_node(neg) | ||
| 186 | + | ||
| 187 | + self.assertNotIn(neg, list(graph.nodes)) | ||
| 188 | + self.assertEqual(relu.args, (x,)) | ||
| 189 | + graph.lint() | ||
| 190 | + | ||
| 191 | + gm = GraphModule(torch.nn.Module(), graph) | ||
| 192 | + input_tensor = torch.randn(2, 3).to(device_type) | ||
| 193 | + self.assertEqual(gm(input_tensor), torch.relu(input_tensor)) | ||
| 194 | + | ||
| 195 | + # Test get_attr: retrieve module parameters as graph nodes via symbolic tracing. | ||
| 196 | + def test_get_attr(self): | ||
| 197 | + class M(torch.nn.Module): | ||
| 198 | + def __init__(self): | ||
| 199 | + super().__init__() | ||
| 200 | + self.weight = torch.nn.Parameter(torch.tensor([2.0])) | ||
| 201 | + self.bias = torch.nn.Parameter(torch.tensor([1.0])) | ||
| 202 | + | ||
| 203 | + def forward(self, x): | ||
| 204 | + return x * self.weight + self.bias | ||
| 205 | + | ||
| 206 | + gm = symbolic_trace(M()) | ||
| 207 | + attr_nodes = gm.graph.find_nodes(op="get_attr") | ||
| 208 | + self.assertEqual(len(attr_nodes), 2) | ||
| 209 | + targets = {n.target for n in attr_nodes} | ||
| 210 | + self.assertIn("weight", targets) | ||
| 211 | + self.assertIn("bias", targets) | ||
| 212 | + | ||
| 213 | + gm = gm.to(device_type) | ||
| 214 | + input_tensor = torch.randn(2, 3).to(device_type) | ||
| 215 | + self.assertEqual(gm(input_tensor), M().to(device_type)(input_tensor)) | ||
| 216 | + | ||
| 107 | 217 | ||
| 108 | if __name__ == "__main__": | 218 | if __name__ == "__main__": |
| 109 | run_tests() | 219 | run_tests() |