Star850
1331
代码
代码
Issues1518
Pull Requests2251
流水线
项目讨论
Wiki
分析
Star850
1331
已合并
test(nn): add test for fx.tracer.call_module_2.9.0 #38070
咿咿呀呀创建于 6月10日
test(nn): add test for fx.tracer.call_module_2.9.0 #38070
已合并
咿咿呀呀创建于 6月10日
从gcw_cRSfl0he/v2.9.0test-fx-tracer-call-module-2.9.0合入到Ascend/pytorchv2.9.0
共 1 个文件变更+270-1
Mtest/fx/test_fx_tracer_transformer_api.py+270-1
@@ -15,6 +15,7 @@ Add validation cases for torch.fx.Tracer/Transformer APIs on NPU:
15 - torch.fx.Transformer.get_attr15 - torch.fx.Transformer.get_attr
16 - torch.fx.Transformer.placeholder16 - torch.fx.Transformer.placeholder
17 - torch.fx.Tracer.getattr17 - torch.fx.Tracer.getattr
18+ - torch.fx.Tracer.call_module
18 19 
192. This file validates the core functionality of these APIs on NPU environment.202. This file validates the core functionality of these APIs on NPU environment.
203. This file also validates Tracer class-level behavior with module state,213. This file also validates Tracer class-level behavior with module state,
@@ -23,6 +24,8 @@ Add validation cases for torch.fx.Tracer/Transformer APIs on NPU:
23 24 
24import torch25import torch
25import torch_npu26import torch_npu
27+import torch.nn as nn
28+import torch.fx as fx
26 29 
27from torch.fx import Tracer, symbolic_trace, Transformer, GraphModule30from torch.fx import Tracer, symbolic_trace, Transformer, GraphModule
28from torch.fx.proxy import Proxy, TraceError31from torch.fx.proxy import Proxy, TraceError
@@ -32,6 +35,7 @@ import torch_npu
32 35 
33 36 
34torch_npu.npu.set_compile_mode(jit_compile=False)37torch_npu.npu.set_compile_mode(jit_compile=False)
38+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
35 39 
36 40 
37class TestFxTracerApi(TestCase):41class TestFxTracerApi(TestCase):
@@ -403,5 +407,270 @@ class TestTracerGetFreshQualname(TestCase):
403 self.assertEqual(tracer.get_fresh_qualname("new_attr"), "new_attr0")407 self.assertEqual(tracer.get_fresh_qualname("new_attr"), "new_attr0")
404 408 
405 409 
410+class TestTracerCallModule(TestCase):
411+ """
412+ DIRECT test suite for fx.Tracer.call_module API.
413+ Does NOT use trace() - validates call_module API behavior in isolation.
414+ """
415+ 
416+ def setUp(self):
417+ """Setup fresh tracer and empty graph for each test"""
418+ super().setUp()
419+ self.tracer = fx.Tracer()
420+ self.graph = fx.Graph()
421+ self.tracer.graph = self.graph
422+ self.tracer.root = nn.Module()
423+ 
424+ def _create_placeholder(self, name: str, value: torch.Tensor = None) -> fx.Proxy:
425+ """Helper to create a placeholder node and return its Proxy"""
426+ node = self.graph.placeholder(name)
427+ return fx.Proxy(node, self.tracer)
428+ 
429+ def test_call_module_creates_node_in_graph(self):
430+ """[CORE] Verify call_module creates a call_module node in the graph"""
431+ linear = nn.Linear(10, 5).to(device_type)
432+ self.tracer.root.add_module("linear", linear)
433+ 
434+ x_proxy = self._create_placeholder("x")
435+ result = self.tracer.call_module(linear, linear.forward, (x_proxy,), {})
436+ 
437+ nodes = list(self.graph.nodes)
438+ self.assertEqual(len(nodes), 2)
439+ self.assertEqual(nodes[0].op, "placeholder")
440+ self.assertEqual(nodes[1].op, "call_module")
441+ self.assertEqual(nodes[1].target, "linear")
442+ self.assertEqual(nodes[1].args, (x_proxy.node,))
443+ self.assertIsInstance(result, fx.Proxy)
444+ 
445+ def test_call_module_returns_proxy_with_correct_node(self):
446+ """[CORE] Verify call_module returns a Proxy whose node has correct attributes"""
447+ relu = nn.ReLU().to(device_type)
448+ self.tracer.root.add_module("relu", relu)
449+ 
450+ x_proxy = self._create_placeholder("x")
451+ result = self.tracer.call_module(relu, relu.forward, (x_proxy,), {})
452+ 
453+ self.assertIsInstance(result, fx.Proxy)
454+ self.assertEqual(result.node.op, "call_module")
455+ self.assertEqual(result.node.target, "relu")
456+ self.assertEqual(result.node.args, (x_proxy.node,))
457+ self.assertEqual(result.node.kwargs, {})
458+ 
459+ def test_call_module_preserves_positional_args(self):
460+ """[CORE] Verify multiple positional arguments are preserved"""
461+ class MultiArgModule(nn.Module):
462+ def forward(self, x, y, z):
463+ return x + y + z
464+ 
465+ mod = MultiArgModule().to(device_type)
466+ self.tracer.root.add_module("mod", mod)
467+ 
468+ x_proxy = self._create_placeholder("x")
469+ y_proxy = self._create_placeholder("y")
470+ z_proxy = self._create_placeholder("z")
471+ 
472+ result = self.tracer.call_module(mod, mod.forward, (x_proxy, y_proxy, z_proxy), {})
473+ 
474+ self.assertIsInstance(result, fx.Proxy)
475+ self.assertIn(result.node.op, ["call_module", "call_function"])
476+ 
477+ def test_call_module_preserves_kwargs(self):
478+ """[CORE] Verify keyword arguments are preserved"""
479+ class KwargModule(nn.Module):
480+ def forward(self, x, bias=None, scale=None):
481+ out = x
482+ if scale is not None:
483+ out = out * scale
484+ if bias is not None:
485+ out = out + bias
486+ return out
487+ 
488+ mod = KwargModule().to(device_type)
489+ self.tracer.root.add_module("mod", mod)
490+ 
491+ x_proxy = self._create_placeholder("x")
492+ bias_proxy = self._create_placeholder("bias")
493+ scale_proxy = self._create_placeholder("scale")
494+ 
495+ result = self.tracer.call_module(
496+ mod, mod.forward,
497+ (x_proxy,),
498+ {"bias": bias_proxy, "scale": scale_proxy}
499+ )
500+ self.assertIsNotNone(result.node.kwargs)
501+ 
502+ def test_call_module_chains_multiple_calls(self):
503+ """[CORE] Verify chained call_module calls produce correct dataflow"""
504+ linear1 = nn.Linear(10, 8).to(device_type)
505+ linear2 = nn.Linear(8, 5).to(device_type)
506+ self.tracer.root.add_module("linear1", linear1)
507+ self.tracer.root.add_module("linear2", linear2)
508+ 
509+ x_proxy = self._create_placeholder("x")
510+ intermediate = self.tracer.call_module(linear1, linear1.forward, (x_proxy,), {})
511+ result = self.tracer.call_module(linear2, linear2.forward, (intermediate,), {})
512+ 
513+ self.assertEqual(result.node.args[0], intermediate.node)
514+ nodes = list(self.graph.nodes)
515+ self.assertEqual(len(nodes), 3)
516+ self.assertEqual(nodes[0].op, "placeholder")
517+ self.assertEqual(nodes[1].op, "call_module")
518+ self.assertEqual(nodes[1].target, "linear1")
519+ self.assertEqual(nodes[2].op, "call_module")
520+ self.assertEqual(nodes[2].target, "linear2")
521+ self.assertEqual(nodes[2].args[0], nodes[1])
522+ 
523+ def test_call_module_same_module_multiple_times(self):
524+ """[CORE] Verify calling same module multiple times creates distinct nodes"""
525+ linear = nn.Linear(10, 5).to(device_type)
526+ self.tracer.root.add_module("linear", linear)
527+ 
528+ x_proxy = self._create_placeholder("x")
529+ y_proxy = self._create_placeholder("y")
530+ 
531+ result1 = self.tracer.call_module(linear, linear.forward, (x_proxy,), {})
532+ result2 = self.tracer.call_module(linear, linear.forward, (y_proxy,), {})
533+ 
534+ self.assertNotEqual(result1.node, result2.node)
535+ self.assertEqual(result1.node.target, result2.node.target)
536+ self.assertEqual(result1.node.target, "linear")
537+ 
538+ nodes = list(self.graph.nodes)
539+ self.assertEqual(len(nodes), 4)
540+ self.assertEqual(nodes[2].op, "call_module")
541+ self.assertEqual(nodes[3].op, "call_module")
542+ 
543+ def test_call_module_with_nested_module_path(self):
544+ """[CORE] Verify nested module target path resolution"""
545+ class Inner(nn.Module):
546+ def __init__(self):
547+ super().__init__()
548+ self.linear = nn.Linear(5, 5).to(device_type)
549+ 
550+ class Outer(nn.Module):
551+ def __init__(self):
552+ super().__init__()
553+ self.inner = Inner()
554+ 
555+ outer = Outer().to(device_type)
556+ self.tracer.root.add_module("outer", outer)
557+ 
558+ x_proxy = self._create_placeholder("x")
559+ result = self.tracer.call_module(
560+ outer.inner.linear,
561+ outer.inner.linear.forward,
562+ (x_proxy,),
563+ {}
564+ )
565+ 
566+ self.assertEqual(result.node.target, "outer.inner.linear")
567+ self.assertEqual(result.node.op, "call_module")
568+ 
569+ def test_call_module_with_sequential_indexing(self):
570+ """[CORE] Verify Sequential submodule indexing works"""
571+ seq = nn.Sequential(
572+ nn.Linear(5, 10).to(device_type),
573+ nn.ReLU().to(device_type),
574+ nn.Linear(10, 5).to(device_type)
575+ )
576+ self.tracer.root.add_module("seq", seq)
577+ 
578+ x_proxy = self._create_placeholder("x")
579+ result = self.tracer.call_module(seq[0], seq[0].forward, (x_proxy,), {})
580+ self.assertEqual(result.node.target, "seq.0")
581+ 
582+ def test_call_module_result_can_be_used_in_operations(self):
583+ """[CORE] Verify call_module result can be used in arithmetic operations"""
584+ linear = nn.Linear(10, 5).to(device_type)
585+ self.tracer.root.add_module("linear", linear)
586+ 
587+ x_proxy = self._create_placeholder("x")
588+ linear_out = self.tracer.call_module(linear, linear.forward, (x_proxy,), {})
589+ add_result = linear_out + 1.0
590+ 
591+ self.assertIsInstance(add_result, fx.Proxy)
592+ self.assertEqual(add_result.node.op, "call_function")
593+ self.assertEqual(add_result.node.args[0], linear_out.node)
594+ 
595+ def test_call_module_with_different_module_types(self):
596+ """[CORE] Verify call_module works with various module types"""
597+ modules = {
598+ "linear": nn.Linear(10, 5).to(device_type),
599+ "conv2d": nn.Conv2d(3, 16, 3).to(device_type),
600+ "relu": nn.ReLU().to(device_type),
601+ "dropout": nn.Dropout(0.5).to(device_type),
602+ "batchnorm": nn.BatchNorm2d(16).to(device_type),
603+ }
604+ 
605+ for name, module in modules.items():
606+ with self.subTest(module_type=name):
607+ self.graph = fx.Graph()
608+ self.tracer.graph = self.graph
609+ self.tracer.root = nn.Module()
610+ x_proxy = self._create_placeholder("x")
611+ 
612+ self.tracer.root.add_module(name, module)
613+ result = self.tracer.call_module(module, module.forward, (x_proxy,), {})
614+ 
615+ self.assertIsInstance(result, fx.Proxy)
616+ self.assertEqual(result.node.op, "call_module")
617+ self.assertEqual(result.node.target, name)
618+ 
619+ def test_call_module_graph_contains_only_call_module_nodes(self):
620+ """[CORE] Verify graph only contains nodes created by call_module (no trace artifacts)"""
621+ linear = nn.Linear(10, 5).to(device_type)
622+ self.tracer.root.add_module("linear", linear)
623+ 
624+ x_proxy = self._create_placeholder("x")
625+ self.tracer.call_module(linear, linear.forward, (x_proxy,), {})
626+ 
627+ nodes = list(self.graph.nodes)
628+ self.assertEqual(len(nodes), 2)
629+ self.assertEqual(nodes[0].op, "placeholder")
630+ self.assertEqual(nodes[1].op, "call_module")
631+ self.assertNotIn("output", [n.op for n in nodes])
632+ 
633+ def test_call_module_module_must_be_registered(self):
634+ """[CORE] Verify module must be registered in root to get string target"""
635+ linear = nn.Linear(10, 5).to(device_type)
636+ x_proxy = self._create_placeholder("x")
637+ 
638+ with self.assertRaises(NameError) as context:
639+ self.tracer.call_module(linear, linear.forward, (x_proxy,), {})
640+ 
641+ self.assertIn("not installed as a submodule", str(context.exception))
642+ 
643+ def test_call_module_with_single_arg(self):
644+ """[CORE] Verify call_module works with single argument"""
645+ linear = nn.Linear(10, 5).to(device_type)
646+ self.tracer.root.add_module("linear", linear)
647+ 
648+ x_proxy = self._create_placeholder("x")
649+ result = self.tracer.call_module(linear, linear.forward, (x_proxy,), {})
650+ 
651+ self.assertIsInstance(result, fx.Proxy)
652+ self.assertEqual(result.node.op, "call_module")
653+ self.assertEqual(len(result.node.args), 1)
654+ self.assertEqual(result.node.args[0], x_proxy.node)
655+ 
656+ def test_call_module_preserves_output_for_further_tracing(self):
657+ """[CORE] Verify the graph built by call_module can be traced/executed"""
658+ linear = nn.Linear(10, 5).to(device_type)
659+ self.tracer.root.add_module("linear", linear)
660+ 
661+ x_proxy = self._create_placeholder("x")
662+ result = self.tracer.call_module(linear, linear.forward, (x_proxy,), {})
663+ 
664+ output_node = self.graph.output(result.node)
665+ graph_module = fx.GraphModule(self.tracer.root, self.graph)
666+ 
667+ x = torch.randn(3, 10).to(device_type)
668+ output = graph_module(x)
669+ 
670+ self.assertEqual(output.shape, (3, 5))
671+ expected = linear(x)
672+ torch.testing.assert_close(output, expected)
673+ 
674+ 
406if __name__ == "__main__":675if __name__ == "__main__":
407- run_tests()676+ run_tests()