已合并
test(fx): add tests for symbolic shapes APIs #37573
chenlan114514创建于 6月4日
test(fx): add tests for symbolic shapes APIs #37573
已合并
共 1 个文件变更+102-1
| @@ -1,7 +1,10 @@ | |||
| 1 | """ | 1 | """ |
| 2 | Add validation cases for torch.fx.experimental.symbolic_shapes APIs on NPU: | 2 | Add validation cases for torch.fx.experimental.symbolic_shapes APIs on NPU: |
| 3 | 1. PyTorch community lacks sufficient and direct API validations for these APIs, so this file is added. | 3 | 1. PyTorch community lacks sufficient and direct API validations for these APIs, so this file is added. |
| 4 | -2. This file validates RelaxedUnspecConstraint, resolve_unbacked_bindings, safe_expand, ShapeEnv.add_var_to_val, ShapeEnv.create_symboolnode and ShapeEnv.create_symfloatnode. | 4 | +2. This file validates RelaxedUnspecConstraint, resolve_unbacked_bindings, safe_expand, |
| 5 | +-ShapeEnv.add_var_to_val, ShapeEnv.create_symboolnode, ShapeEnv.create_symfloatnode, | ||
| 6 | +-ShapeEnv.create_symbol, ShapeEnv.bound_sympy, ShapeEnv.check_equal, | ||
| 7 | +-ShapeEnv.cleanup, ShapeEnv.bind_symbols. | ||
| 5 | 3. This file is extendable for other torch.fx.experimental.symbolic_shapes APIs. | 8 | 3. This file is extendable for other torch.fx.experimental.symbolic_shapes APIs. |
| 6 | """ | 9 | """ |
| 7 | 10 | ||
| @@ -120,5 +123,103 @@ class TestShapeEnvSymbolicShapes(TestCase): | |||
| 120 | self.assertEqual(dummy_tensor.device.type, "npu") | 123 | self.assertEqual(dummy_tensor.device.type, "npu") |
| 121 | 124 | ||
| 122 | 125 | ||
| 126 | +class TestShapeEnvCoreMethods(TestCase): | ||
| 127 | + """Unit tests for ShapeEnv core methods: create_symbol, bound_sympy, check_equal, cleanup, bind_symbols.""" | ||
| 128 | + | ||
| 129 | + def setUp(self): | ||
| 130 | + self.env = ShapeEnv() | ||
| 131 | + self.source = ConstantSource("x") | ||
| 132 | + | ||
| 133 | + def test_create_symbol(self): | ||
| 134 | + """Test create_symbol returns unique sympy.Symbol instances usable in expressions.""" | ||
| 135 | + sym1 = self.env.create_symbol(5, self.source) | ||
| 136 | + sym2 = self.env.create_symbol(10, self.source) | ||
| 137 | + | ||
| 138 | + self.assertIsInstance(sym1, sympy.Symbol) | ||
| 139 | + self.assertIsInstance(sym2, sympy.Symbol) | ||
| 140 | + self.assertTrue(sym1.name.startswith('s')) | ||
| 141 | + self.assertTrue(sym2.name.startswith('s')) | ||
| 142 | + self.assertNotEqual(sym1.name, sym2.name, "Symbols must have distinct names") | ||
| 143 | + expr = sym1 + sym2 | ||
| 144 | + self.assertIsInstance(expr, sympy.Expr) | ||
| 145 | + | ||
| 146 | + def test_bound_sympy(self): | ||
| 147 | + """Test bound_sympy returns correct lower bound for 's0+2' (at least 4).""" | ||
| 148 | + s0 = self.env.create_symbol(5, self.source) | ||
| 149 | + expr = s0 + 2 | ||
| 150 | + bounds = self.env.bound_sympy(expr) | ||
| 151 | + | ||
| 152 | + self.assertTrue(hasattr(bounds, 'lower')) | ||
| 153 | + self.assertTrue(hasattr(bounds, 'upper')) | ||
| 154 | + self.assertLessEqual(bounds.lower, bounds.upper) | ||
| 155 | + self.assertGreaterEqual(bounds.lower, 4) | ||
| 156 | + | ||
| 157 | + def test_check_equal(self): | ||
| 158 | + """Test check_equal passes for self-equality and fails for different environment states.""" | ||
| 159 | + self.env.check_equal(self.env) | ||
| 160 | + | ||
| 161 | + other = ShapeEnv() | ||
| 162 | + other.create_symbol(5, self.source) | ||
| 163 | + # NotEqualError is defined in torch.fx.experimental.recording (PyTorch 2.9+) | ||
| 164 | + try: | ||
| 165 | + from torch.fx.experimental.recording import NotEqualError | ||
| 166 | + expected_exceptions = (AssertionError, NotEqualError) | ||
| 167 | + except ImportError: | ||
| 168 | + expected_exceptions = AssertionError | ||
| 169 | + with self.assertRaises(expected_exceptions): | ||
| 170 | + self.env.check_equal(other) | ||
| 171 | + | ||
| 172 | + def test_cleanup(self): | ||
| 173 | + """Test cleanup does not break future symbol creation.""" | ||
| 174 | + self.env.create_symbol(5, self.source) | ||
| 175 | + self.env.cleanup() | ||
| 176 | + new_sym = self.env.create_symbol(10, self.source) | ||
| 177 | + self.assertIsInstance(new_sym, sympy.Symbol) | ||
| 178 | + | ||
| 179 | + def test_bind_symbols(self): | ||
| 180 | + """Test bind_symbols correctly maps symbolic placeholders to concrete tensor dimensions.""" | ||
| 181 | + try: | ||
| 182 | + from torch.fx.experimental.proxy_tensor import make_fake_tensor | ||
| 183 | + has_fake = True | ||
| 184 | + except ImportError: | ||
| 185 | + has_fake = False | ||
| 186 | + | ||
| 187 | + if not has_fake: | ||
| 188 | + self.skipTest("make_fake_tensor not available in this environment") | ||
| 189 | + | ||
| 190 | + # Single tensor binding | ||
| 191 | + s0 = self.env.create_symbol(5, self.source) | ||
| 192 | + s1 = self.env.create_symbol(2, self.source) | ||
| 193 | + fake_input = make_fake_tensor(torch.empty(s0, s1), self.env, self.source) | ||
| 194 | + real_input = torch.randn(5, 2) | ||
| 195 | + bindings = self.env.bind_symbols([fake_input], [real_input]) | ||
| 196 | + self.assertIsInstance(bindings, dict) | ||
| 197 | + self.assertIn(s0, bindings) | ||
| 198 | + self.assertIn(s1, bindings) | ||
| 199 | + self.assertEqual(bindings[s0], 5) | ||
| 200 | + self.assertEqual(bindings[s1], 2) | ||
| 201 | + | ||
| 202 | + # Multiple tensor batch binding | ||
| 203 | + s2 = self.env.create_symbol(3, self.source) | ||
| 204 | + s3 = self.env.create_symbol(4, self.source) | ||
| 205 | + fake_input_2 = make_fake_tensor(torch.empty(s2, s3), self.env, self.source) | ||
| 206 | + real_input_2 = torch.randn(3, 4) | ||
| 207 | + bindings_batch = self.env.bind_symbols([fake_input, fake_input_2], [real_input, real_input_2]) | ||
| 208 | + self.assertEqual(len(bindings_batch), 4) | ||
| 209 | + # Verify both old and new symbols are correctly mapped | ||
| 210 | + self.assertIn(s0, bindings_batch) | ||
| 211 | + self.assertIn(s1, bindings_batch) | ||
| 212 | + self.assertIn(s2, bindings_batch) | ||
| 213 | + self.assertIn(s3, bindings_batch) | ||
| 214 | + self.assertEqual(bindings_batch[s0], 5) | ||
| 215 | + self.assertEqual(bindings_batch[s1], 2) | ||
| 216 | + self.assertEqual(bindings_batch[s2], 3) | ||
| 217 | + self.assertEqual(bindings_batch[s3], 4) | ||
| 218 | + | ||
| 219 | + # Negative case: mismatched argument count | ||
| 220 | + with self.assertRaises(ValueError): | ||
| 221 | + self.env.bind_symbols([fake_input], [real_input, torch.randn(3, 4)]) | ||
| 222 | + | ||
| 223 | + | ||
| 123 | if __name__ == "__main__": | 224 | if __name__ == "__main__": |
| 124 | run_tests() | 225 | run_tests() |