已合并
test(fx): add NPU test cases for torch.fx.GraphModule APIs [v2.10.0] #39608
test(fx): add NPU test cases for torch.fx.GraphModule APIs [v2.10.0] #39608
已合并
冬阳创建于 6月30日
1 个文件变更+463-0
@@ -0,0 +1,463 @@
1+"""
2+Add validation cases for torch.fx.GraphModule APIs on NPU:
3+1. PyTorch community lacks direct and independent test cases for
4+ torch.fx.GraphModule.code, torch.fx.GraphModule.graph, and
5+ several submodule management APIs, so this file is added.
6+2. This file validates torch.fx.GraphModule.__init__,
7+ torch.fx.GraphModule.code, torch.fx.GraphModule.graph,
8+ torch.fx.GraphModule.add_submodule,
9+ torch.fx.GraphModule.delete_submodule,
10+ torch.fx.GraphModule.delete_all_unused_submodules,
11+ torch.fx.GraphModule.print_readable,
12+ torch.fx.GraphModule.recompile,
13+ torch.fx.GraphModule.to_folder (extendable).
14+"""
15+ 
16+import os
17+import tempfile
18+ 
19+import torch
20+import torch.nn as nn
21+from torch.fx import symbolic_trace, GraphModule, Graph
22+from torch.fx.graph import PythonCode
23+from torch.testing._internal.common_utils import run_tests, TestCase
24+ 
25+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
26+ 
27+ 
28+class TestFxGraphModuleInit(TestCase):
29+ 
30+ def test_init_from_module(self):
31+ """GraphModule.__init__ with nn.Module root copies submodules."""
32+ class TestMod(nn.Module):
33+ def __init__(self):
34+ super().__init__()
35+ self.lin = nn.Linear(4, 3)
36+ 
37+ def forward(self, x):
38+ return self.lin(x)
39+ 
40+ m = TestMod()
41+ gm = symbolic_trace(m)
42+ self.assertTrue(hasattr(gm, "lin"))
43+ self.assertIsInstance(gm.lin, nn.Linear)
44+ 
45+ def test_init_from_dict(self):
46+ """GraphModule.__init__ with dict root assigns attributes."""
47+ graph = Graph()
48+ x = graph.placeholder("x")
49+ lin = graph.call_module("lin", args=(x,))
50+ graph.output(lin)
51+ 
52+ gm = GraphModule({"lin": nn.Linear(4, 3)}, graph)
53+ x_in = torch.randn(2, 4)
54+ out = gm(x_in)
55+ self.assertEqual(out.shape, (2, 3))
56+ 
57+ def test_init_sets_class_name(self):
58+ """GraphModule.__init__ with custom class_name."""
59+ graph = Graph()
60+ x = graph.placeholder("x")
61+ graph.output(x)
62+ 
63+ gm = GraphModule(torch.nn.Module(), graph, class_name="MyGM")
64+ self.assertEqual(gm.__class__.__name__, "MyGM")
65+ 
66+ def test_init_raises_on_bad_type(self):
67+ """GraphModule.__init__ rejects non-Module/non-dict root."""
68+ graph = Graph()
69+ x = graph.placeholder("x")
70+ graph.output(x)
71+ 
72+ with self.assertRaises(RuntimeError):
73+ GraphModule("bad_root", graph)
74+ 
75+ 
76+class TestFxGraphModuleCode(TestCase):
77+ 
78+ def setUp(self):
79+ super().setUp()
80+ 
81+ class SimpleModule(nn.Module):
82+ def __init__(self):
83+ super().__init__()
84+ self.linear = nn.Linear(4, 3)
85+ 
86+ def forward(self, x):
87+ return torch.relu(self.linear(x))
88+ 
89+ self.gm = symbolic_trace(SimpleModule())
90+ 
91+ def test_code_returns_string(self):
92+ """code property returns a non-empty string."""
93+ code = self.gm.code
94+ self.assertIsInstance(code, str)
95+ self.assertGreater(len(code), 0)
96+ 
97+ def test_code_contains_forward(self):
98+ """code contains 'def forward'."""
99+ self.assertIn("def forward", self.gm.code)
100+ self.assertIn("self", self.gm.code)
101+ 
102+ def test_code_contains_op_names(self):
103+ """code contains operator names from traced model."""
104+ code = self.gm.code
105+ self.assertIn("relu", code)
106+ self.assertIn("linear", code)
107+ 
108+ def test_code_consistent_after_recompile(self):
109+ """code is consistent after multiple recompiles."""
110+ code1 = self.gm.code
111+ self.gm.recompile()
112+ code2 = self.gm.code
113+ self.assertEqual(code1, code2)
114+ 
115+ 
116+class TestFxGraphModuleGraph(TestCase):
117+ 
118+ def setUp(self):
119+ super().setUp()
120+ 
121+ class SimpleModule(nn.Module):
122+ def __init__(self):
123+ super().__init__()
124+ self.linear = nn.Linear(4, 3)
125+ 
126+ def forward(self, x):
127+ return torch.relu(self.linear(x))
128+ 
129+ self.module = SimpleModule()
130+ self.gm = symbolic_trace(self.module)
131+ 
132+ def test_graph_getter_returns_graph(self):
133+ """graph getter returns a Graph instance."""
134+ g = self.gm.graph
135+ self.assertIsInstance(g, Graph)
136+ 
137+ def test_graph_getter_has_nodes(self):
138+ """graph getter returns graph with placeholder and output."""
139+ g = self.gm.graph
140+ nodes = list(g.nodes)
141+ self.assertGreater(len(nodes), 0)
142+ ops = {node.op for node in nodes}
143+ self.assertIn("placeholder", ops)
144+ self.assertIn("output", ops)
145+ 
146+ def test_graph_getter_is_consistent(self):
147+ """graph getter returns same object on repeated access."""
148+ g1 = self.gm.graph
149+ g2 = self.gm.graph
150+ self.assertIs(g1, g2)
151+ 
152+ def test_graph_setter_reassigns_graph(self):
153+ """Setting graph reassigns internal reference."""
154+ gm_new = symbolic_trace(self.module)
155+ new_graph = gm_new.graph
156+ self.gm.graph = new_graph
157+ self.assertIs(self.gm.graph, new_graph)
158+ 
159+ def test_graph_setter_triggers_recompile(self):
160+ """Setting graph triggers recompile and produces valid code."""
161+ gm_new = symbolic_trace(self.module)
162+ self.gm.graph = gm_new.graph
163+ code = self.gm.code
164+ self.assertIsInstance(code, str)
165+ self.assertGreater(len(code), 0)
166+ 
167+ def test_graph_setter_forward_works(self):
168+ """After graph set, forward produces correct output."""
169+ gm_new = symbolic_trace(self.module)
170+ self.gm.graph = gm_new.graph
171+ x = torch.randn(2, 4)
172+ expected = self.module(x)
173+ actual = self.gm(x)
174+ torch.testing.assert_close(actual, expected)
175+ 
176+ def test_graph_setter_raises_on_non_graph(self):
177+ """Setting graph to non-Graph raises AssertionError."""
178+ with self.assertRaises(AssertionError):
179+ self.gm.graph = "not_a_graph"
180+ with self.assertRaises(AssertionError):
181+ self.gm.graph = 42
182+ 
183+ def test_graph_setter_preserves_lint(self):
184+ """After setting a valid graph, lint() does not raise."""
185+ gm_new = symbolic_trace(self.module)
186+ self.gm.graph = gm_new.graph
187+ self.gm.graph.lint()
188+ 
189+ 
190+class TestFxGraphModuleSubmodule(TestCase):
191+ 
192+ def setUp(self):
193+ super().setUp()
194+ graph = Graph()
195+ x = graph.placeholder("x")
196+ lin = graph.call_module("lin", args=(x,))
197+ graph.output(lin)
198+ self.gm = GraphModule({"lin": nn.Linear(4, 3)}, graph)
199+ 
200+ def test_add_submodule_root_level(self):
201+ """add_submodule at root level adds the module."""
202+ new_mod = nn.ReLU()
203+ result = self.gm.add_submodule("relu", new_mod)
204+ self.assertTrue(result)
205+ self.assertIs(self.gm.relu, new_mod)
206+ 
207+ def test_add_submodule_nested(self):
208+ """add_submodule creates intermediate modules for nested path."""
209+ new_mod = nn.ReLU()
210+ result = self.gm.add_submodule("a.b.c", new_mod)
211+ self.assertTrue(result)
212+ self.assertIsInstance(self.gm.a, nn.Module)
213+ self.assertIsInstance(self.gm.a.b, nn.Module)
214+ self.assertIs(self.gm.a.b.c, new_mod)
215+ 
216+ def test_add_submodule_overwrite_fails_on_non_module(self):
217+ """add_submodule returns False when path blocked by non-Module."""
218+ self.gm.add_submodule("blocker", nn.ReLU())
219+ # Install a non-Module attribute to block the path
220+ self.gm.blocker.some_attr = "not_a_module"
221+ result = self.gm.add_submodule("blocker.some_attr.sub", nn.ReLU())
222+ self.assertFalse(result)
223+ 
224+ def test_delete_submodule_existing(self):
225+ """delete_submodule removes an existing submodule."""
226+ self.gm.add_submodule("temp", nn.ReLU())
227+ self.assertTrue(hasattr(self.gm, "temp"))
228+ result = self.gm.delete_submodule("temp")
229+ self.assertTrue(result)
230+ self.assertFalse(hasattr(self.gm, "temp"))
231+ 
232+ def test_delete_submodule_nested(self):
233+ """delete_submodule removes a nested submodule."""
234+ self.gm.add_submodule("outer.inner", nn.ReLU())
235+ self.assertTrue(hasattr(self.gm.outer, "inner"))
236+ result = self.gm.delete_submodule("outer.inner")
237+ self.assertTrue(result)
238+ self.assertFalse(hasattr(self.gm.outer, "inner"))
239+ 
240+ def test_delete_submodule_nonexistent(self):
241+ """delete_submodule on nonexistent path returns False."""
242+ result = self.gm.delete_submodule("nonexistent.path")
243+ self.assertFalse(result)
244+ 
245+ def test_delete_submodule_non_module(self):
246+ """delete_submodule on non-Module attribute returns False."""
247+ self.gm.some_param = nn.Parameter(torch.randn(2, 2))
248+ result = self.gm.delete_submodule("some_param")
249+ self.assertFalse(result)
250+ 
251+ def test_delete_all_unused_removes_orphans(self):
252+ """delete_all_unused_submodules removes modules not in graph."""
253+ graph = Graph()
254+ x = graph.placeholder("x")
255+ lin = graph.call_module("lin", args=(x,))
256+ graph.output(lin)
257+ gm = GraphModule({"lin": nn.Linear(4, 3)}, graph)
258+ gm.add_submodule("orphan", nn.ReLU())
259+ self.assertTrue(hasattr(gm, "orphan"))
260+ gm.delete_all_unused_submodules()
261+ self.assertFalse(hasattr(gm, "orphan"))
262+ self.assertTrue(hasattr(gm, "lin"))
263+ 
264+ def test_delete_all_unused_preserves_used(self):
265+ """delete_all_unused_submodules preserves modules in graph."""
266+ graph = Graph()
267+ x = graph.placeholder("x")
268+ a = graph.call_module("a", args=(x,))
269+ b = graph.call_module("b", args=(a,))
270+ graph.output(b)
271+ gm = GraphModule({
272+ "a": nn.Linear(4, 4),
273+ "b": nn.Linear(4, 3),
274+ }, graph)
275+ gm.add_submodule("orphan", nn.ReLU())
276+ self.assertTrue(hasattr(gm, "orphan"))
277+ gm.delete_all_unused_submodules()
278+ self.assertFalse(hasattr(gm, "orphan"))
279+ self.assertTrue(hasattr(gm, "a"))
280+ self.assertTrue(hasattr(gm, "b"))
281+ 
282+ 
283+class TestFxGraphModulePrintReadable(TestCase):
284+ 
285+ def test_print_readable_returns_string(self):
286+ """print_readable returns a non-empty string."""
287+ class SimpleMod(nn.Module):
288+ def __init__(self):
289+ super().__init__()
290+ self.lin = nn.Linear(4, 3)
291+ 
292+ def forward(self, x):
293+ return torch.relu(self.lin(x))
294+ 
295+ gm = symbolic_trace(SimpleMod())
296+ output = gm.print_readable(print_output=False)
297+ self.assertIsInstance(output, str)
298+ self.assertGreater(len(output), 0)
299+ self.assertIn("class", output)
300+ self.assertIn("def forward", output)
301+ 
302+ def test_print_readable_contains_child_code(self):
303+ """print_readable includes code from child GraphModules."""
304+ class ChildMod(nn.Module):
305+ def __init__(self):
306+ super().__init__()
307+ self.w = nn.Parameter(torch.randn(3, 4))
308+ 
309+ def forward(self, x):
310+ return x + self.w
311+ 
312+ class ParentMod(nn.Module):
313+ def __init__(self):
314+ super().__init__()
315+ self.child = symbolic_trace(ChildMod())
316+ 
317+ def forward(self, x):
318+ return self.child(x)
319+ 
320+ gm = symbolic_trace(ParentMod())
321+ output = gm.print_readable(print_output=False)
322+ self.assertIsInstance(output, str)
323+ self.assertIn("class", output)
324+ 
325+ 
326+class TestFxGraphModuleRecompile(TestCase):
327+ 
328+ def test_recompile_returns_python_code(self):
329+ """recompile returns a PythonCode object."""
330+ class SimpleMod(nn.Module):
331+ def __init__(self):
332+ super().__init__()
333+ self.lin = nn.Linear(4, 3)
334+ 
335+ def forward(self, x):
336+ return self.lin(x)
337+ 
338+ gm = symbolic_trace(SimpleMod())
339+ pc = gm.recompile()
340+ self.assertIsInstance(pc, PythonCode)
341+ self.assertGreater(len(pc.src), 0)
342+ 
343+ def test_recompile_preserves_forward(self):
344+ """After recompile, forward still works correctly."""
345+ class SimpleMod(nn.Module):
346+ def __init__(self):
347+ super().__init__()
348+ self.lin = nn.Linear(4, 3)
349+ 
350+ def forward(self, x):
351+ return self.lin(x)
352+ 
353+ m = SimpleMod()
354+ gm = symbolic_trace(m)
355+ x = torch.randn(2, 4)
356+ expected = m(x)
357+ gm.recompile()
358+ actual = gm(x)
359+ torch.testing.assert_close(actual, expected)
360+ 
361+ 
362+class TestFxGraphModuleToFolder(TestCase):
363+ 
364+ def test_to_folder_creates_files(self):
365+ """to_folder creates module.py and __init__.py in folder."""
366+ class SimpleMod(nn.Module):
367+ def __init__(self):
368+ super().__init__()
369+ self.lin = nn.Linear(4, 3)
370+ 
371+ def forward(self, x):
372+ return self.lin(x)
373+ 
374+ gm = symbolic_trace(SimpleMod())
375+ with tempfile.TemporaryDirectory() as tmpdir:
376+ gm.to_folder(tmpdir, "TestMod")
377+ self.assertTrue(os.path.isfile(os.path.join(tmpdir, "module.py")))
378+ self.assertTrue(os.path.isfile(
379+ os.path.join(tmpdir, "__init__.py")))
380+ 
381+ def test_to_folder_module_file_content(self):
382+ """to_folder output can be imported."""
383+ class SimpleMod(nn.Module):
384+ def __init__(self):
385+ super().__init__()
386+ self.lin = nn.Linear(4, 3)
387+ 
388+ def forward(self, x):
389+ return self.lin(x)
390+ 
391+ gm = symbolic_trace(SimpleMod())
392+ with tempfile.TemporaryDirectory() as tmpdir:
393+ gm.to_folder(tmpdir, "TestMod")
394+ # Verify the module file content is valid Python
395+ with open(os.path.join(tmpdir, "module.py")) as f:
396+ content = f.read()
397+ self.assertIn("class TestMod(torch.nn.Module)", content)
398+ self.assertIn("def forward", content)
399+ 
400+ 
401+class TestFxGraphModuleOnNpu(TestCase):
402+ """Verify GraphModule APIs work with tensors on NPU device."""
403+ 
404+ def setUp(self):
405+ super().setUp()
406+ if device_type == "cpu":
407+ self.skipTest("Test requires NPU device")
408+ 
409+ def test_code_and_graph_on_npu(self):
410+ """code and graph properties work after moving module to NPU."""
411+ class SimpleMod(nn.Module):
412+ def __init__(self):
413+ super().__init__()
414+ self.lin = nn.Linear(4, 3)
415+ 
416+ def forward(self, x):
417+ return torch.relu(self.lin(x))
418+ 
419+ m = SimpleMod().to(device_type)
420+ gm = symbolic_trace(m)
421+ self.assertIsInstance(gm.code, str)
422+ self.assertGreater(len(gm.code), 0)
423+ self.assertIsInstance(gm.graph, Graph)
424+ self.assertGreater(len(list(gm.graph.nodes)), 0)
425+ 
426+ def test_forward_on_npu(self):
427+ """Generated forward works with NPU tensors."""
428+ class SimpleMod(nn.Module):
429+ def __init__(self):
430+ super().__init__()
431+ self.lin = nn.Linear(4, 3)
432+ 
433+ def forward(self, x):
434+ return self.lin(x)
435+ 
436+ m = SimpleMod().to(device_type)
437+ gm = symbolic_trace(m)
438+ x = torch.randn(2, 4).to(device_type)
439+ expected = m(x)
440+ actual = gm(x)
441+ torch.testing.assert_close(actual, expected)
442+ 
443+ def test_recompile_on_npu(self):
444+ """recompile works after module is on NPU."""
445+ class SimpleMod(nn.Module):
446+ def __init__(self):
447+ super().__init__()
448+ self.lin = nn.Linear(4, 3)
449+ 
450+ def forward(self, x):
451+ return self.lin(x)
452+ 
453+ m = SimpleMod().to(device_type)
454+ gm = symbolic_trace(m)
455+ x = torch.randn(2, 4).to(device_type)
456+ gm.recompile()
457+ actual = gm(x)
458+ expected = m(x)
459+ torch.testing.assert_close(actual, expected)
460+ 
461+ 
462+if __name__ == "__main__":
463+ run_tests()