已合并
test(fx): add Interpreter internal API test cases for NPU [v2.7.1] #43909
zkx创建于 18 天前
test(fx): add Interpreter internal API test cases for NPU [v2.7.1] #43909
已合并
共 1 个文件变更+181-0
| @@ -0,0 +1,181 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | + | ||
| 16 | +""" | ||
| 17 | +Add validation cases for torch.fx.Interpreter APIs on NPU: | ||
| 18 | +1. PyTorch community lacks dedicated direct test cases for | ||
| 19 | + Interpreter.boxed_run, Interpreter.fetch_attr, | ||
| 20 | + Interpreter.map_nodes_to_values and | ||
| 21 | + Interpreter.fetch_args_kwargs_from_env, so this file is added. | ||
| 22 | +2. This file validates these internal methods on NPU. | ||
| 23 | +""" | ||
| 24 | + | ||
| 25 | +import torch | ||
| 26 | + | ||
| 27 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 28 | +from torch.fx import Interpreter, symbolic_trace | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +class TestInterpreterBoxedRun(TestCase): | ||
| 35 | + """Test Interpreter.boxed_run method.""" | ||
| 36 | + | ||
| 37 | + def test_boxed_run_basic(self): | ||
| 38 | + class AddModule(torch.nn.Module): | ||
| 39 | + def forward(self, lhs, rhs): | ||
| 40 | + return lhs + rhs | ||
| 41 | + | ||
| 42 | + gm = symbolic_trace(AddModule()) | ||
| 43 | + interpreter = Interpreter(gm) | ||
| 44 | + lhs = torch.tensor(1.0, device=device_type) | ||
| 45 | + rhs = torch.tensor(2.0, device=device_type) | ||
| 46 | + result = interpreter.boxed_run([lhs.clone(), rhs.clone()]) | ||
| 47 | + self.assertTrue(torch.equal(result, lhs + rhs)) | ||
| 48 | + | ||
| 49 | + def test_boxed_run_clears_args(self): | ||
| 50 | + class AddModule(torch.nn.Module): | ||
| 51 | + def forward(self, lhs, rhs): | ||
| 52 | + return lhs + rhs | ||
| 53 | + | ||
| 54 | + gm = symbolic_trace(AddModule()) | ||
| 55 | + interpreter = Interpreter(gm) | ||
| 56 | + lhs = torch.tensor(1.0, device=device_type) | ||
| 57 | + rhs = torch.tensor(2.0, device=device_type) | ||
| 58 | + args_list = [lhs.clone(), rhs.clone()] | ||
| 59 | + interpreter.boxed_run(args_list) | ||
| 60 | + self.assertEqual(args_list, []) | ||
| 61 | + | ||
| 62 | + | ||
| 63 | +class TestInterpreterFetchAttr(TestCase): | ||
| 64 | + """Test Interpreter.fetch_attr method.""" | ||
| 65 | + | ||
| 66 | + def test_fetch_attr_parameter(self): | ||
| 67 | + class M(torch.nn.Module): | ||
| 68 | + def __init__(self): | ||
| 69 | + super().__init__() | ||
| 70 | + self.param = torch.nn.Parameter( | ||
| 71 | + torch.ones(2, 2, device=device_type)) | ||
| 72 | + | ||
| 73 | + def forward(self, x): | ||
| 74 | + return x + self.param | ||
| 75 | + | ||
| 76 | + m = M() | ||
| 77 | + gm = symbolic_trace(m) | ||
| 78 | + interp = Interpreter(gm) | ||
| 79 | + param = interp.fetch_attr("param") | ||
| 80 | + self.assertTrue(torch.equal(param, torch.ones(2, 2, device=device_type))) | ||
| 81 | + | ||
| 82 | + def test_fetch_attr_submodule(self): | ||
| 83 | + class M(torch.nn.Module): | ||
| 84 | + def __init__(self): | ||
| 85 | + super().__init__() | ||
| 86 | + self.sub = torch.nn.Linear(3, 3).to(device_type) | ||
| 87 | + | ||
| 88 | + def forward(self, x): | ||
| 89 | + return self.sub(x) | ||
| 90 | + | ||
| 91 | + m = M() | ||
| 92 | + gm = symbolic_trace(m) | ||
| 93 | + interp = Interpreter(gm) | ||
| 94 | + sub = interp.fetch_attr("sub") | ||
| 95 | + self.assertIsInstance(sub, torch.nn.Module) | ||
| 96 | + | ||
| 97 | + | ||
| 98 | +class TestInterpreterMapNodesToValues(TestCase): | ||
| 99 | + """Test Interpreter.map_nodes_to_values method.""" | ||
| 100 | + | ||
| 101 | + def test_map_nodes_to_values_args(self): | ||
| 102 | + class M(torch.nn.Module): | ||
| 103 | + def forward(self, x, y): | ||
| 104 | + return x + y | ||
| 105 | + | ||
| 106 | + gm = symbolic_trace(M()) | ||
| 107 | + interp = Interpreter(gm) | ||
| 108 | + x = torch.ones(2, 2, device=device_type) | ||
| 109 | + y = torch.zeros(2, 2, device=device_type) | ||
| 110 | + interp.args_iter = iter([x, y]) | ||
| 111 | + add_node = [n for n in gm.graph.nodes if n.op == "call_function"][0] | ||
| 112 | + # fill env first so map_nodes_to_values can replace Node with values | ||
| 113 | + placeholder_nodes = [n for n in gm.graph.nodes if n.op == "placeholder"] | ||
| 114 | + for n in placeholder_nodes: | ||
| 115 | + interp.env[n] = next(interp.args_iter) | ||
| 116 | + mapped = interp.map_nodes_to_values(add_node.args, add_node) | ||
| 117 | + self.assertIsInstance(mapped, tuple) | ||
| 118 | + self.assertTrue(torch.equal(mapped[0], x)) | ||
| 119 | + self.assertTrue(torch.equal(mapped[1], y)) | ||
| 120 | + | ||
| 121 | + def test_map_nodes_to_values_kwargs(self): | ||
群 | |||
| 122 | + class M(torch.nn.Module): | ||
| 123 | + def forward(self, x): | ||
| 124 | + return torch.add(x, other=1) | ||
| 125 | + | ||
| 126 | + gm = symbolic_trace(M()) | ||
| 127 | + interp = Interpreter(gm) | ||
| 128 | + x = torch.ones(2, 2, device=device_type) | ||
| 129 | + interp.args_iter = iter([x]) | ||
| 130 | + add_node = [n for n in gm.graph.nodes if n.op == "call_function"][0] | ||
| 131 | + self.assertEqual(add_node.kwargs, {"other": 1}) | ||
| 132 | + mapped = interp.map_nodes_to_values(add_node.kwargs, add_node) | ||
| 133 | + self.assertIsInstance(mapped, dict) | ||
| 134 | + self.assertEqual(mapped, {"other": 1}) | ||
| 135 | + | ||
| 136 | + | ||
| 137 | +class TestInterpreterFetchArgsKwargsFromEnv(TestCase): | ||
| 138 | + """Test Interpreter.fetch_args_kwargs_from_env method.""" | ||
| 139 | + | ||
| 140 | + def test_fetch_args_kwargs_from_env(self): | ||
| 141 | + class M(torch.nn.Module): | ||
| 142 | + def forward(self, x, y): | ||
| 143 | + return torch.add(x, y) | ||
| 144 | + | ||
| 145 | + gm = symbolic_trace(M()) | ||
| 146 | + interp = Interpreter(gm) | ||
| 147 | + x = torch.ones(2, 2, device=device_type) | ||
| 148 | + y = torch.zeros(2, 2, device=device_type) | ||
| 149 | + interp.args_iter = iter([x, y]) | ||
| 150 | + placeholder_nodes = [n for n in gm.graph.nodes if n.op == "placeholder"] | ||
| 151 | + add_node = [n for n in gm.graph.nodes if n.op == "call_function"][0] | ||
| 152 | + interp.env = {} | ||
| 153 | + for n in placeholder_nodes: | ||
| 154 | + interp.env[n] = next(interp.args_iter) | ||
| 155 | + args, kwargs = interp.fetch_args_kwargs_from_env(add_node) | ||
| 156 | + self.assertEqual(len(args), 2) | ||
| 157 | + self.assertIsInstance(kwargs, dict) | ||
| 158 | + self.assertTrue(torch.equal(args[0], x)) | ||
| 159 | + self.assertTrue(torch.equal(args[1], y)) | ||
| 160 | + | ||
| 161 | + def test_fetch_args_kwargs_from_env_non_empty_kwargs(self): | ||
| 162 | + class M(torch.nn.Module): | ||
| 163 | + def forward(self, x): | ||
| 164 | + return torch.add(x, other=1) | ||
| 165 | + | ||
| 166 | + gm = symbolic_trace(M()) | ||
| 167 | + interp = Interpreter(gm) | ||
| 168 | + x = torch.ones(2, 2, device=device_type) | ||
| 169 | + interp.args_iter = iter([x]) | ||
| 170 | + placeholder_nodes = [n for n in gm.graph.nodes if n.op == "placeholder"] | ||
| 171 | + add_node = [n for n in gm.graph.nodes if n.op == "call_function"][0] | ||
| 172 | + interp.env = {} | ||
| 173 | + for n in placeholder_nodes: | ||
| 174 | + interp.env[n] = next(interp.args_iter) | ||
| 175 | + args, kwargs = interp.fetch_args_kwargs_from_env(add_node) | ||
| 176 | + self.assertEqual(args, (x,)) | ||
| 177 | + self.assertEqual(kwargs, {"other": 1}) | ||
| 178 | + | ||
| 179 | + | ||
| 180 | +if __name__ == "__main__": | ||
| 181 | + run_tests() | ||
请修改
test_map_nodes_to_values_kwargs，构造真实的关键字参数图节点并断言映射结果；当前使用torch.add(x, 1)生成的节点kwargs为空，测试始终只验证空字典类型。