已合并
test(fx): add NPU test cases and patch for torch.fx.node APIs #37787
hantao55_创建于 6月6日
test(fx): add NPU test cases and patch for torch.fx.node APIs #37787
已合并
共 2 个文件变更+232-0
| @@ -0,0 +1,208 @@ | |||
| 1 | +""" | ||
| 2 | +Add validation cases for torch.fx.node APIs on NPU: | ||
| 3 | +1. PyTorch community lacks direct test cases for the following APIs, | ||
| 4 | + so this file is added. | ||
| 5 | +2. This file validates torch.fx.node._type_repr, torch.fx.Node.kwargs, | ||
| 6 | + torch.fx.node.map_arg, torch.fx.Node.next, torch.fx.Node.prev (extendable). | ||
| 7 | +""" | ||
| 8 | + | ||
| 9 | +import operator | ||
| 10 | +import types | ||
| 11 | + | ||
| 12 | +import torch | ||
| 13 | +import torch.fx | ||
| 14 | +from torch.fx.node import _type_repr, map_arg | ||
| 15 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +class TestFxNodeExtraApis(TestCase): | ||
| 19 | + | ||
| 20 | + def test_type_repr_builtin_type(self): | ||
| 21 | + # builtin types return qualname only, no module prefix | ||
| 22 | + self.assertEqual(_type_repr(int), "int") | ||
| 23 | + self.assertEqual(_type_repr(float), "float") | ||
| 24 | + self.assertEqual(_type_repr(str), "str") | ||
| 25 | + self.assertEqual(_type_repr(bool), "bool") | ||
| 26 | + self.assertEqual(_type_repr(type(None)), "NoneType") | ||
| 27 | + | ||
| 28 | + def test_type_repr_non_builtin_type(self): | ||
| 29 | + # non-builtin types return module.qualname | ||
| 30 | + self.assertEqual(_type_repr(torch.Tensor), "torch.Tensor") | ||
| 31 | + self.assertEqual(_type_repr(torch.nn.Linear), "torch.nn.modules.linear.Linear") | ||
| 32 | + | ||
| 33 | + def test_type_repr_ellipsis(self): | ||
| 34 | + # Ellipsis object returns "..." | ||
| 35 | + self.assertEqual(_type_repr(...), "...") | ||
| 36 | + | ||
| 37 | + def test_type_repr_function(self): | ||
| 38 | + # FunctionType returns function.__name__ | ||
| 39 | + def my_func(): | ||
| 40 | + pass | ||
| 41 | + | ||
| 42 | + self.assertEqual(_type_repr(my_func), "my_func") | ||
| 43 | + | ||
| 44 | + def test_type_repr_generic_alias(self): | ||
| 45 | + # GenericAlias (e.g. list[int]) falls through to repr() | ||
| 46 | + # not treated as a plain type, so module.qualname logic is skipped | ||
| 47 | + result = _type_repr(list[int]) | ||
| 48 | + self.assertEqual(result, repr(list[int])) | ||
| 49 | + | ||
| 50 | + def test_type_repr_other(self): | ||
| 51 | + # Non-type objects fall back to repr() | ||
| 52 | + self.assertEqual(_type_repr(42), "42") | ||
| 53 | + self.assertEqual(_type_repr("hello"), "'hello'") | ||
| 54 | + | ||
| 55 | + def test_node_kwargs_getter(self): | ||
| 56 | + graph = torch.fx.Graph() | ||
| 57 | + x = graph.placeholder("x") | ||
| 58 | + node = graph.call_function(torch.relu, kwargs={"input": x}) | ||
| 59 | + graph.output(node) | ||
| 60 | + | ||
| 61 | + # getter returns the kwargs dict | ||
| 62 | + self.assertEqual(node.kwargs, {"input": x}) | ||
| 63 | + self.assertIsInstance(node.kwargs, dict) | ||
| 64 | + | ||
| 65 | + def test_node_kwargs_setter_updates_use_def(self): | ||
| 66 | + graph = torch.fx.Graph() | ||
| 67 | + x = graph.placeholder("x") | ||
| 68 | + y = graph.placeholder("y") | ||
| 69 | + node = graph.call_function(torch.relu, kwargs={"input": x}) | ||
| 70 | + graph.output(node) | ||
| 71 | + | ||
| 72 | + # x should be in node's users before reassignment | ||
| 73 | + self.assertIn(node, x.users) | ||
| 74 | + self.assertNotIn(node, y.users) | ||
| 75 | + | ||
| 76 | + # reassign kwargs: x is removed, y is added | ||
| 77 | + node.kwargs = {"input": y} | ||
| 78 | + self.assertEqual(node.kwargs, {"input": y}) | ||
| 79 | + self.assertNotIn(node, x.users) | ||
| 80 | + self.assertIn(node, y.users) | ||
| 81 | + | ||
| 82 | + def test_node_kwargs_setter_empty(self): | ||
| 83 | + graph = torch.fx.Graph() | ||
| 84 | + x = graph.placeholder("x") | ||
| 85 | + node = graph.call_function(torch.relu, kwargs={"input": x}) | ||
| 86 | + graph.output(node) | ||
| 87 | + | ||
| 88 | + node.kwargs = {} | ||
| 89 | + self.assertEqual(node.kwargs, {}) | ||
| 90 | + # x no longer used by node | ||
| 91 | + self.assertNotIn(node, x.users) | ||
| 92 | + | ||
| 93 | + def test_map_arg_applies_fn_to_nodes(self): | ||
| 94 | + graph = torch.fx.Graph() | ||
| 95 | + x = graph.placeholder("x") | ||
| 96 | + y = graph.placeholder("y") | ||
| 97 | + add = graph.call_function(operator.add, args=(x, y)) | ||
| 98 | + graph.output(add) | ||
| 99 | + | ||
| 100 | + visited = [] | ||
| 101 | + | ||
| 102 | + def collect(node): | ||
| 103 | + visited.append(node.name) | ||
| 104 | + return node | ||
| 105 | + | ||
| 106 | + map_arg(add.args, collect) | ||
| 107 | + self.assertEqual(visited, ["x", "y"]) | ||
| 108 | + | ||
| 109 | + def test_map_arg_non_node_passthrough(self): | ||
| 110 | + graph = torch.fx.Graph() | ||
| 111 | + x = graph.placeholder("x") | ||
| 112 | + | ||
| 113 | + # non-Node elements pass through unchanged | ||
| 114 | + result = map_arg((x, 42, "const", 3.14), lambda n: n) | ||
| 115 | + self.assertIs(result[0], x) | ||
| 116 | + self.assertEqual(result[1], 42) | ||
| 117 | + self.assertEqual(result[2], "const") | ||
| 118 | + self.assertAlmostEqual(result[3], 3.14) | ||
| 119 | + | ||
| 120 | + def test_map_arg_nested_structure(self): | ||
| 121 | + graph = torch.fx.Graph() | ||
| 122 | + x = graph.placeholder("x") | ||
| 123 | + y = graph.placeholder("y") | ||
| 124 | + | ||
| 125 | + visited = [] | ||
| 126 | + | ||
| 127 | + def collect(node): | ||
| 128 | + visited.append(node.name) | ||
| 129 | + return node | ||
| 130 | + | ||
| 131 | + # nested list containing nodes | ||
| 132 | + map_arg([x, [y]], collect) | ||
| 133 | + self.assertIn("x", visited) | ||
| 134 | + self.assertIn("y", visited) | ||
| 135 | + | ||
| 136 | + def test_map_arg_dict(self): | ||
| 137 | + graph = torch.fx.Graph() | ||
| 138 | + x = graph.placeholder("x") | ||
| 139 | + | ||
| 140 | + visited = [] | ||
| 141 | + | ||
| 142 | + def collect(node): | ||
| 143 | + visited.append(node.name) | ||
| 144 | + return node | ||
| 145 | + | ||
| 146 | + map_arg({"input": x, "scale": 2.0}, collect) | ||
| 147 | + self.assertEqual(visited, ["x"]) | ||
| 148 | + | ||
| 149 | + def test_map_arg_requires_callable(self): | ||
| 150 | + graph = torch.fx.Graph() | ||
| 151 | + x = graph.placeholder("x") | ||
| 152 | + | ||
| 153 | + with self.assertRaises(AssertionError): | ||
| 154 | + map_arg(x, "not_a_callable") | ||
| 155 | + | ||
| 156 | + def test_node_next(self): | ||
| 157 | + graph = torch.fx.Graph() | ||
| 158 | + x = graph.placeholder("x") | ||
| 159 | + relu = graph.call_function(torch.relu, args=(x,)) | ||
| 160 | + out = graph.output(relu) | ||
| 161 | + | ||
| 162 | + self.assertIs(x.next, relu) | ||
| 163 | + self.assertIs(relu.next, out) | ||
| 164 | + | ||
| 165 | + def test_node_next_after_append(self): | ||
| 166 | + graph = torch.fx.Graph() | ||
| 167 | + x = graph.placeholder("x") | ||
| 168 | + relu = graph.call_function(torch.relu, args=(x,)) | ||
| 169 | + neg = graph.call_function(torch.neg, args=(x,)) | ||
| 170 | + graph.output(relu) | ||
| 171 | + | ||
| 172 | + # move neg to after relu | ||
| 173 | + relu.append(neg) | ||
| 174 | + self.assertIs(relu.next, neg) | ||
| 175 | + | ||
| 176 | + def test_node_prev(self): | ||
| 177 | + graph = torch.fx.Graph() | ||
| 178 | + x = graph.placeholder("x") | ||
| 179 | + relu = graph.call_function(torch.relu, args=(x,)) | ||
| 180 | + out = graph.output(relu) | ||
| 181 | + | ||
| 182 | + self.assertIs(relu.prev, x) | ||
| 183 | + self.assertIs(out.prev, relu) | ||
| 184 | + | ||
| 185 | + def test_node_prev_after_prepend(self): | ||
| 186 | + graph = torch.fx.Graph() | ||
| 187 | + x = graph.placeholder("x") | ||
| 188 | + relu = graph.call_function(torch.relu, args=(x,)) | ||
| 189 | + neg = graph.call_function(torch.neg, args=(x,)) | ||
| 190 | + graph.output(relu) | ||
| 191 | + | ||
| 192 | + # move neg to before relu | ||
| 193 | + relu.prepend(neg) | ||
| 194 | + self.assertIs(relu.prev, neg) | ||
| 195 | + | ||
| 196 | + def test_node_next_prev_consistent(self): | ||
| 197 | + # next and prev are inverses of each other | ||
| 198 | + graph = torch.fx.Graph() | ||
| 199 | + x = graph.placeholder("x") | ||
| 200 | + relu = graph.call_function(torch.relu, args=(x,)) | ||
| 201 | + graph.output(relu) | ||
| 202 | + | ||
| 203 | + self.assertIs(x.next.prev, x) | ||
| 204 | + self.assertIs(relu.prev.next, relu) | ||
| 205 | + | ||
| 206 | + | ||
| 207 | +if __name__ == "__main__": | ||
| 208 | + run_tests() | ||
| @@ -0,0 +1,24 @@ | |||
| 1 | +diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py | ||
| 2 | +index 5b8aa5c..295bbd5 100644 | ||
| 3 | +--- a/test/dynamo/test_functions.py | ||
| 4 | ++++ b/test/dynamo/test_functions.py | ||
| 5 | + class DefaultsTests(torch._dynamo.test_case.TestCase): | ||
| 6 | + | ||
| 7 | + opt_fn = torch.compile(fn, backend="eager", fullgraph=True) | ||
| 8 | + | ||
| 9 | +- x = [torch.randn(4), [torch.randn(4), torch.randn(4)]] | ||
| 10 | ++ x = [torch.randn(4).npu(), [torch.randn(4).npu(), torch.randn(4).npu()]] | ||
| 11 | + | ||
| 12 | + def f(y): | ||
| 13 | + return y * 2 | ||
| 14 | + class DefaultsTests(torch._dynamo.test_case.TestCase): | ||
| 15 | + self.assertTrue(type(ref) is type(res)) | ||
| 16 | + | ||
| 17 | + x = { | ||
| 18 | +- "a": torch.randn(4), | ||
| 19 | +- "b": [torch.randn(4), torch.randn(4)], | ||
| 20 | ++ "a": torch.randn(4).npu(), | ||
| 21 | ++ "b": [torch.randn(4).npu(), torch.randn(4).npu()], | ||
| 22 | + } | ||
| 23 | + ref = fn(x, f) | ||
| 24 | + res = opt_fn(x, f) | ||