已合并
test(fx): add DimConstraints API tests #38825
Goko创建于 6月18日
test(fx): add DimConstraints API tests #38825
已合并
共 1 个文件变更+85-0
| @@ -9,6 +9,10 @@ Add validation cases for torch.fx symbolic_shapes related APIs on NPU: | |||
| 9 | 3. Current covered APIs / behaviors include: | 9 | 3. Current covered APIs / behaviors include: |
| 10 | - symbolic_shapes.DimConstraints.add | 10 | - symbolic_shapes.DimConstraints.add |
| 11 | - symbolic_shapes.DimConstraints.add_equality | 11 | - symbolic_shapes.DimConstraints.add_equality |
| 12 | + - symbolic_shapes.DimConstraints.rewrite_with_congruences | ||
| 13 | + - symbolic_shapes.DimConstraints.solve | ||
| 14 | + - symbolic_shapes.DimConstraints.forced_specializations | ||
| 15 | + - symbolic_shapes.DimConstraints.prettify_results | ||
| 12 | - torch.fx.experimental.symbolic_shapes.ShapeEnv.size_hint | 16 | - torch.fx.experimental.symbolic_shapes.ShapeEnv.size_hint |
| 13 | - torch.fx.experimental.symbolic_shapes.ShapeEnv.suppress_guards | 17 | - torch.fx.experimental.symbolic_shapes.ShapeEnv.suppress_guards |
| 14 | - torch.fx.experimental.symbolic_shapes.ShapeEnvSettings | 18 | - torch.fx.experimental.symbolic_shapes.ShapeEnvSettings |
| @@ -33,8 +37,10 @@ import torch | |||
| 33 | 37 | ||
| 34 | import torch_npu | 38 | import torch_npu |
| 35 | from torch._dynamo.source import ConstantSource | 39 | from torch._dynamo.source import ConstantSource |
| 40 | +from torch.export import Dim | ||
| 36 | from torch.fx.experimental import symbolic_shapes | 41 | from torch.fx.experimental import symbolic_shapes |
| 37 | from torch.testing._internal.common_utils import TestCase, run_tests | 42 | from torch.testing._internal.common_utils import TestCase, run_tests |
| 43 | +from torch.utils._sympy.functions import FloorDiv, Mod | ||
| 38 | from torch.utils._sympy.value_ranges import ValueRanges | 44 | from torch.utils._sympy.value_ranges import ValueRanges |
| 39 | 45 | ||
| 40 | 46 | ||
| @@ -97,6 +103,85 @@ class TestSymbolicShapesAPI(TestCase): | |||
| 97 | constraints.add_equality(source, symbolic_expr) | 103 | constraints.add_equality(source, symbolic_expr) |
| 98 | self.assertEqual(constraints._symbolic_equivalences, [(source, symbolic_expr)]) | 104 | self.assertEqual(constraints._symbolic_equivalences, [(source, symbolic_expr)]) |
| 99 | 105 | ||
| 106 | + def test_dim_constraints_rewrite_with_congruences_records_mod_guard(self): | ||
| 107 | + # Verify that congruence guards rewrite floor division and record | ||
| 108 | + # the corresponding modular relationship for the symbol. | ||
| 109 | + symbol = sympy.Symbol("s0", positive=True, integer=True) | ||
| 110 | + constraints = symbolic_shapes.DimConstraints( | ||
| 111 | + {}, | ||
| 112 | + {symbol: sympy.Integer(5)}, | ||
| 113 | + set(), | ||
| 114 | + {}, | ||
| 115 | + ) | ||
| 116 | + | ||
| 117 | + rewritten = constraints.rewrite_with_congruences(symbol, FloorDiv(symbol, 2)) | ||
| 118 | + | ||
| 119 | + self.assertEqual(rewritten, symbol / 2 - sympy.Rational(1, 2)) | ||
| 120 | + self.assertEqual( | ||
| 121 | + {str(expr) for expr in constraints._congruences[symbol]}, | ||
| 122 | + {"Mod(s0 + 1, 2)"}, | ||
| 123 | + ) | ||
| 124 | + | ||
| 125 | + def test_dim_constraints_solve_records_dynamic_results(self): | ||
| 126 | + # Verify that `solve` classifies a satisfiable lower-bound guard | ||
| 127 | + # as a dynamic result rather than a static specialization. | ||
| 128 | + symbol = sympy.Symbol("s0", positive=True, integer=True) | ||
| 129 | + constraints = symbolic_shapes.DimConstraints( | ||
| 130 | + {symbol: [ConstantSource("x")]}, | ||
| 131 | + {symbol: sympy.Integer(4)}, | ||
| 132 | + {symbol}, | ||
| 133 | + {}, | ||
| 134 | + ) | ||
| 135 | + | ||
| 136 | + constraints.add(symbol >= 2) | ||
| 137 | + constraints.solve() | ||
| 138 | + | ||
| 139 | + self.assertEqual(constraints._static_results, set()) | ||
| 140 | + self.assertEqual(constraints._dynamic_results, {"2 <= x"}) | ||
| 141 | + | ||
| 142 | + def test_dim_constraints_forced_specializations_reports_marked_dynamic_equalities(self): | ||
| 143 | + # Verify that marked dynamic equalities are reported as forced | ||
| 144 | + # specializations using the configured debug name. | ||
| 145 | + symbol = sympy.Symbol("s0", positive=True, integer=True) | ||
| 146 | + constraints = symbolic_shapes.DimConstraints( | ||
| 147 | + {symbol: [ConstantSource("x")]}, | ||
| 148 | + {symbol: sympy.Integer(4)}, | ||
| 149 | + {symbol}, | ||
| 150 | + {"x": "dx"}, | ||
| 151 | + ) | ||
| 152 | + | ||
| 153 | + constraints.add(sympy.Eq(symbol, 4)) | ||
| 154 | + constraints.solve() | ||
| 155 | + | ||
| 156 | + self.assertEqual(constraints.forced_specializations(), {"dx = x": 4}) | ||
| 157 | + | ||
| 158 | + def test_dim_constraints_prettify_results_reports_forced_specialization(self): | ||
| 159 | + # Verify that `prettify_results` explains a forced specialization | ||
| 160 | + # and includes the suggested concrete dimension value. | ||
| 161 | + def fn(x): | ||
| 162 | + return x | ||
| 163 | + | ||
| 164 | + symbol = sympy.Symbol("s0", positive=True, integer=True) | ||
| 165 | + constraints = symbolic_shapes.DimConstraints( | ||
| 166 | + {symbol: [ConstantSource("x")]}, | ||
| 167 | + {symbol: sympy.Integer(4)}, | ||
| 168 | + {symbol}, | ||
| 169 | + {"x": "dx"}, | ||
| 170 | + ) | ||
| 171 | + | ||
| 172 | + constraints.add(sympy.Eq(symbol, 4)) | ||
| 173 | + constraints.solve() | ||
| 174 | + message = constraints.prettify_results( | ||
| 175 | + inspect.signature(fn), | ||
| 176 | + {"x": Dim("dx")}, | ||
| 177 | + ValueError("dummy constraint violation"), | ||
| 178 | + constraints.forced_specializations(), | ||
| 179 | + ) | ||
| 180 | + | ||
| 181 | + self.assertIn("Specializations unexpectedly required (dx)!", message) | ||
| 182 | + self.assertIn("dx = x", message) | ||
| 183 | + self.assertIn("dx = 4", message) | ||
| 184 | + | ||
| 100 | def test_strict_min_max_constraint_records_warn_only_and_value_range(self): | 185 | def test_strict_min_max_constraint_records_warn_only_and_value_range(self): |
| 101 | # Verify that StrictMinMaxConstraint stores warn_only and ValueRanges | 186 | # Verify that StrictMinMaxConstraint stores warn_only and ValueRanges |
| 102 | # according to its actual constructor signature. | 187 | # according to its actual constructor signature. |