diff --git a/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py b/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py index 15ae608e52a..a3f11a00904 100644 --- a/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py +++ b/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py @@ -2,6 +2,7 @@ import logging from typing import Any, Dict, List +import sympy import torch from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule @@ -48,26 +49,166 @@ def _apply_symbolic_shape_expressions( ) return outputs - # Build a mapping from compile-time symbolic expressions to runtime SymInts - # by aligning captured input info with actual runtime input tensors - symbol_to_symint = {} - symbol_to_concrete = {} - shape_env = None + # Shape symbols are local to a ShapeEnv. The expressions in shape_info came + # from the ShapeEnv used to compile the engine, so they must not be inserted + # verbatim into the unrelated ShapeEnv doing this fake execution. Both + # environments can contain a symbol called u0 for different quantities. + # + # Give every compile-time symbol a private identity before relating input + # expressions to their runtime counterparts. Sympy symbols compare by name + # and assumptions, whereas Dummy objects are unique even when their printed + # names match a runtime symbol. + compile_symbols = { + symbol + for info in (*input_info, *output_info) + for expr in info["shape_exprs"] + if not isinstance(expr, int) + for symbol in expr.free_symbols + } + compile_symbol_namespace = { + symbol: sympy.Dummy(symbol.name, integer=True) for symbol in compile_symbols + } + + def in_compile_namespace(expr: sympy.Expr) -> sympy.Expr: + return expr.xreplace(compile_symbol_namespace) + + # Prefer the ShapeEnv owned by the active FakeTensorMode. Inputs may all be + # statically shaped even when a data-dependent engine output is symbolic. + shape_env = getattr(fake_mode, "shape_env", None) + compile_input_symbols = set() + compile_to_runtime: Dict[sympy.Expr, sympy.Expr] = {} + runtime_expr_to_symint = {} + composite_input_equations = [] # Align inputs: for each captured input, match it with the corresponding runtime input - for idx, (inp_tensor, inp_info) in enumerate(zip(inputs, input_info)): - for d, s in zip(inp_tensor.shape, inp_info["shape_exprs"]): + for inp_tensor, inp_info in zip(inputs, input_info): + for d, compile_expr in zip(inp_tensor.shape, inp_info["shape_exprs"]): + if isinstance(compile_expr, int): + continue + + compile_expr = in_compile_namespace(compile_expr) + compile_input_symbols.update(compile_expr.free_symbols) if isinstance(d, torch.SymInt): - symbol_to_symint[s] = d + runtime_expr = d.node.expr + runtime_expr_to_symint[runtime_expr] = d if shape_env is None: shape_env = d.node.shape_env - - elif isinstance(d, int): - symbol_to_concrete[s] = d + else: + runtime_expr = sympy.Integer(d) + + if compile_expr.is_Symbol: + if compile_expr in compile_to_runtime: + residual = compile_to_runtime[compile_expr] - runtime_expr + # Two runtime SymInts can be tied to the same compile-time + # symbol yet be distinct symbol objects here -- e.g. a + # shared torch.export.Dim gets reallocated as separate, + # SymInts on retrace/re-export. Preserve the compile-time + # equality by installing it as a guard in the active + # ShapeEnv instead of requiring the symbols to have + # already been merged. + error_message = ( + "[torch.ops.tensorrt.execute_engine]: Runtime input shapes " + f"disagree on compile-time symbol {compile_expr}: already " + f"mapped to {compile_to_runtime[compile_expr]} from an earlier " + f"input, but this input maps it to {runtime_expr}" + ) + if shape_env is not None: + residual = shape_env.simplify(residual) + residual = sympy.simplify(residual) + if residual != 0 and ( + shape_env is None + or not shape_env.guard_or_defer_runtime_assert( + sympy.Eq(residual, 0), error_message + ) + ): + raise RuntimeError(error_message) + compile_to_runtime[compile_expr] = runtime_expr + else: + # Store the difference: sympy.Eq can auto-collapse to a bare + # Boolean with no .lhs/.rhs (e.g. Eq(2*d, 7) -> False). + composite_input_equations.append(compile_expr - runtime_expr) logger.debug( - f"[torch.ops.tensorrt.execute_engine]: Meta kernel captured and mapped symbol from input {inp_tensor} (compile time symbol: {s}, new symbol: {d})" + f"[torch.ops.tensorrt.execute_engine]: Meta kernel captured input shape mapping from {compile_expr} to {runtime_expr}" + ) + + # A constrained input can be recorded as an expression such as 2*s0. + # Solve those equations for compile-time symbols not mapped by a direct + # symbolic dimension. + unresolved_input_symbols = compile_input_symbols - compile_to_runtime.keys() + if composite_input_equations and unresolved_input_symbols: + equations = [ + equation.xreplace(compile_to_runtime) + for equation in composite_input_equations + ] + solutions = sympy.solve( + equations, + tuple(sorted(unresolved_input_symbols, key=str)), + dict=True, + ) + if len(solutions) == 1: + # Underdetermined systems still return one dict, but a value can + # contain other unresolved compile-time symbols, e.g. solve(s0+s1-10, + # (s0,s1)) -> {s0: 10-s1}. Only accept fully resolved values. + for symbol, value in solutions[0].items(): + if not (value.free_symbols & compile_input_symbols): + compile_to_runtime[symbol] = value + + # Validate every composite equation, even if unresolved_input_symbols was + # empty above (a symbol already mapped elsewhere doesn't mean this + # relationship was actually satisfied by these runtime inputs). + for diff in composite_input_equations: + residual = sympy.simplify(diff.xreplace(compile_to_runtime)) + unresolved_symbols = residual.free_symbols & compile_input_symbols + if unresolved_symbols: + raise RuntimeError( + "[torch.ops.tensorrt.execute_engine]: Could not verify the compile-time " + f"input relationship {diff} == 0 against these runtime shapes " + f"(unresolved residual: {residual})" ) + if shape_env is not None: + residual = sympy.simplify(shape_env.simplify(residual)) + if residual == 0: + continue + error_message = ( + "[torch.ops.tensorrt.execute_engine]: Runtime input shapes violate " + f"a relationship captured at compile time: {diff} == 0 does not hold " + f"for these inputs (residual {residual} != 0)" + ) + if shape_env is not None and shape_env.guard_or_defer_runtime_assert( + sympy.Eq(residual, 0), error_message + ): + continue + if residual.is_number: + raise RuntimeError(error_message) + # Still symbolic: can't prove it holds, so fail closed. + raise RuntimeError( + "[torch.ops.tensorrt.execute_engine]: Could not verify the compile-time " + f"input relationship {diff} == 0 against these runtime shapes " + f"(unresolved residual: {residual})" + ) + + # Symbols which occur only in engine outputs represent quantities created + # by the engine, such as the row count of nonzero. Allocate fresh runtime + # symbols once per fake invocation, preserve sharing between output + # expressions, and mark them as valid tensor sizes. + output_only_symbols = { + symbol + for info in output_info + for expr in info["shape_exprs"] + if not isinstance(expr, int) + for symbol in in_compile_namespace(expr).free_symbols + if symbol not in compile_input_symbols + } + if output_only_symbols and shape_env is None: + raise RuntimeError( + "[torch.ops.tensorrt.execute_engine]: No shape_env available during meta kernel execution" + ) + for symbol in sorted(output_only_symbols, key=str): + runtime_symint = shape_env.create_unbacked_symint() + shape_env._constrain_range_for_size(runtime_symint.node.expr) + compile_to_runtime[symbol] = runtime_symint.node.expr + runtime_expr_to_symint[runtime_symint.node.expr] = runtime_symint # Create output fake tensors with symbolic shapes logger.debug(f"Deserialized output shape expressions: {output_info}") @@ -80,75 +221,37 @@ def _apply_symbolic_shape_expressions( # Concrete dimension output_shape.append(expr) else: - logger.debug(f"Symbolic expression: {expr}") - # Symbolic expression (sympy expr) - - # Check if this expression uses any symbols that are now concrete - has_concrete_symbols = any( - sym in symbol_to_concrete for sym in expr.free_symbols + compile_expr = in_compile_namespace(expr) + missing_input_symbols = ( + compile_expr.free_symbols + & compile_input_symbols - compile_to_runtime.keys() ) - - if has_concrete_symbols: - # Case 2: Some compile-time symbols are now concrete ints - # Evaluate the expression to a concrete value - try: - # Build substitution dict with concrete values - subs_dict = {} - for sym in expr.free_symbols: - if sym in symbol_to_concrete: - subs_dict[sym] = symbol_to_concrete[sym] - elif sym in symbol_to_symint: - subs_dict[sym] = symbol_to_symint[sym].node.hint - else: - subs_dict[sym] = sym - - val = expr.subs(subs_dict) - concrete_dim = int(val) - output_shape.append(concrete_dim) - logger.debug( - f"Evaluated {expr} to concrete value {concrete_dim} using concrete mappings" - ) - except Exception as e: - raise RuntimeError( - f"[torch.ops.tensorrt.execute_engine]: Failed to evaluate symbolic expression {expr} " - f"with concrete values. Free symbols: {expr.free_symbols}, " - f"Concrete mappings: {symbol_to_concrete}, " - f"SymInt mappings: {list(symbol_to_symint.keys())}. Error: {e}" - ) - elif expr in symbol_to_symint: - # Case 1a: Direct mapping - compile-time symbol is represented by runtime SymInt - output_shape.append(symbol_to_symint[expr]) - logger.debug( - f"Reused SymInt from input: {expr} -> {symbol_to_symint[expr]}" + if missing_input_symbols: + raise RuntimeError( + "[torch.ops.tensorrt.execute_engine]: Could not map " + f"compile-time input symbols {missing_input_symbols} " + f"while applying output expression {expr}" ) - elif shape_env is not None: - # Case 1b: Create new SymInt from expression using existing SymInts - try: - # Calculate hint by substituting known values - hint_val = expr.subs( - { - sym: symbol_to_symint[sym].node.hint - for sym in expr.free_symbols - if sym in symbol_to_symint - } - ) - hint = int(hint_val) if hint_val.is_number else None - # Create new SymInt from the expression - output_symint = shape_env.create_symintnode(expr, hint=hint) - output_shape.append(output_symint) - logger.debug( - f"Created new SymInt for {expr} with hint {hint}" + runtime_expr = sympy.simplify( + compile_expr.xreplace(compile_to_runtime) + ) + logger.debug( + f"Remapped symbolic output expression {expr} to {runtime_expr}" + ) + if runtime_expr.is_number: + output_shape.append(int(runtime_expr)) + elif runtime_expr in runtime_expr_to_symint: + output_shape.append(runtime_expr_to_symint[runtime_expr]) + else: + try: + output_shape.append( + shape_env.create_symintnode(runtime_expr, hint=None) ) except Exception as e: raise RuntimeError( - f"[torch.ops.tensorrt.execute_engine]: Failed to create SymInt for expression {expr}. " - f"Error: {e}" - ) - else: - raise RuntimeError( - "[torch.ops.tensorrt.execute_engine]: No shape_env available during meta kernel execution" - ) + f"[torch.ops.tensorrt.execute_engine]: Failed to create SymInt for remapped expression {runtime_expr} (captured as {expr}). Error: {e}" + ) from e outputs.append( torch.empty(output_shape, dtype=info["dtype"], device=inputs[0].device) diff --git a/tests/py/dynamo/models/test_meta_kernel_shape_inference.py b/tests/py/dynamo/models/test_meta_kernel_shape_inference.py index 173a2e366fe..eb562142275 100644 --- a/tests/py/dynamo/models/test_meta_kernel_shape_inference.py +++ b/tests/py/dynamo/models/test_meta_kernel_shape_inference.py @@ -17,10 +17,16 @@ """ import pytest +import sympy import torch import torch_tensorrt +from torch._dynamo.source import LocalSource from torch._subclasses.fake_tensor import FakeTensorMode from torch.export import Dim +from torch.fx.experimental.symbolic_shapes import DimDynamic, ShapeEnv +from torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops import ( + _apply_symbolic_shape_expressions, +) class TestMetaKernelShapeInference: @@ -294,5 +300,265 @@ def forward(self, x): assert exported_output.shape == trt_output.shape +class TestApplySymbolicShapeExpressions: + """CPU-only tests for remapping serialized expressions into a fresh ShapeEnv.""" + + @staticmethod + def _shape_info(input_expr, output_expr): + return { + "inputs": [ + { + "shape_exprs": [input_expr], + "dtype": torch.float32, + "name": "x", + } + ], + "outputs": [ + { + "shape_exprs": [output_expr], + "dtype": torch.float32, + } + ], + } + + def test_output_symbol_does_not_alias_same_named_runtime_symbol(self): + compile_input = sympy.Symbol("s0", integer=True) + compile_output = sympy.Symbol("u0", integer=True) + shape_env = ShapeEnv() + + with FakeTensorMode(shape_env=shape_env): + runtime_input = shape_env.create_unbacked_symint() + shape_env._constrain_range_for_size(runtime_input.node.expr) + fake_input = torch.empty(runtime_input) + output = _apply_symbolic_shape_expressions( + [fake_input], self._shape_info(compile_input, compile_output) + )[0] + + assert output.shape[0].node.expr != runtime_input.node.expr + + def test_output_only_symbol_uses_fake_mode_shape_env_and_is_size_like(self): + compile_output = sympy.Symbol("u0", integer=True) + shape_env = ShapeEnv() + shape_info = { + "inputs": [ + { + "shape_exprs": [16], + "dtype": torch.float32, + "name": "x", + } + ], + "outputs": [ + { + "shape_exprs": [compile_output], + "dtype": torch.float32, + } + ], + } + + with FakeTensorMode(shape_env=shape_env): + output = _apply_symbolic_shape_expressions([torch.empty(16)], shape_info)[0] + output_dim = output.shape[0] + query = 1 > output_dim + torch.sym_max(0, 1 - output_dim) + + assert output_dim.node.shape_env is shape_env + assert output_dim.node.expr in shape_env.size_like + assert not bool(shape_env.evaluate_expr(query.node.expr, size_oblivious=True)) + + def test_derived_expression_uses_runtime_input_symbol(self): + compile_input = sympy.Symbol("s0", integer=True) + shape_env = ShapeEnv() + + with FakeTensorMode(shape_env=shape_env): + runtime_input = shape_env.create_unbacked_symint() + shape_env._constrain_range_for_size(runtime_input.node.expr) + fake_input = torch.empty(runtime_input) + output = _apply_symbolic_shape_expressions( + [fake_input], + self._shape_info(compile_input, 2 * compile_input + 1), + )[0] + + assert ( + sympy.simplify( + output.shape[0].node.expr - (2 * runtime_input.node.expr + 1) + ) + == 0 + ) + + def test_composite_input_expression_is_solved_in_runtime_namespace(self): + compile_base = sympy.Symbol("s0", integer=True) + shape_env = ShapeEnv() + + with FakeTensorMode(shape_env=shape_env): + runtime_base = shape_env.create_unbacked_symint() + shape_env._constrain_range_for_size(runtime_base.node.expr) + fake_input = torch.empty(2 * runtime_base) + output = _apply_symbolic_shape_expressions( + [fake_input], + self._shape_info(2 * compile_base, compile_base), + )[0] + + assert output.shape[0].node.expr == runtime_base.node.expr + + def test_rejects_underdetermined_composite_mapping(self): + """One input dim s0+s1=10 cannot pin down s0 alone, so the output + (which is just s0) cannot be inferred. sympy.solve still returns one + solution ({s0: 10-s1}), so this must be rejected explicitly rather + than accepted because len(solutions) == 1.""" + compile_s0 = sympy.Symbol("s0", integer=True) + compile_s1 = sympy.Symbol("s1", integer=True) + + shape_info = self._shape_info( + input_expr=compile_s0 + compile_s1, + output_expr=compile_s0, + ) + + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + fake_input = torch.empty(10) + + with pytest.raises( + RuntimeError, + match="Could not verify the compile-time input relationship", + ): + _apply_symbolic_shape_expressions( + [fake_input], + shape_info, + ) + + @staticmethod + def _two_input_shape_info(x_expr, y_expr, output_expr): + return { + "inputs": [ + {"shape_exprs": [x_expr], "dtype": torch.float32, "name": "x"}, + {"shape_exprs": [y_expr], "dtype": torch.float32, "name": "y"}, + ], + "outputs": [ + {"shape_exprs": [output_expr], "dtype": torch.float32}, + ], + } + + def test_rejects_inconsistent_composite_input_relationship(self): + """s0 is mapped directly from x, so composite_input_equations from y's + `2*s0` is never solved (nothing is left unresolved) -- it still has to + be checked against the runtime shapes, or an inconsistent pairing like + x=4, y=7 (requires y=8) is silently accepted.""" + compile_base = sympy.Symbol("s0", integer=True) + shape_info = self._two_input_shape_info( + compile_base, 2 * compile_base, compile_base + ) + + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + fake_x = torch.empty(4) + fake_y = torch.empty(7) + with pytest.raises(RuntimeError): + _apply_symbolic_shape_expressions([fake_x, fake_y], shape_info) + + def test_accepts_consistent_composite_input_relationship(self): + compile_base = sympy.Symbol("s0", integer=True) + shape_info = self._two_input_shape_info( + compile_base, 2 * compile_base, compile_base + ) + + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + runtime_base = shape_env.create_unbacked_symint() + shape_env._constrain_range_for_size(runtime_base.node.expr) + fake_x = torch.empty(runtime_base) + fake_y = torch.empty(2 * runtime_base) + output = _apply_symbolic_shape_expressions([fake_x, fake_y], shape_info)[0] + + assert output.shape[0].node.expr == runtime_base.node.expr + + def test_rejects_inconsistent_repeated_direct_mapping(self): + """The same compile-time symbol s0 as a bare dimension on two inputs + must resolve to the same runtime value; x=4, z=5 disagree.""" + compile_base = sympy.Symbol("s0", integer=True) + shape_info = self._two_input_shape_info( + compile_base, compile_base, compile_base + ) + + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + fake_x = torch.empty(4) + fake_z = torch.empty(5) + with pytest.raises(RuntimeError): + _apply_symbolic_shape_expressions([fake_x, fake_z], shape_info) + + def test_guards_distinct_runtime_symbols_for_repeated_direct_mapping(self): + """A re-export can allocate distinct backed symbols for dimensions + that the engine metadata records as equal. Preserve that relationship + as a ShapeEnv guard rather than rejecting the fake execution.""" + compile_base = sympy.Symbol("s0", integer=True) + shape_info = self._two_input_shape_info( + compile_base, compile_base, compile_base + ) + + shape_env = ShapeEnv() + runtime_x_expr = shape_env.create_symbol( + 4, LocalSource("x"), dynamic_dim=DimDynamic.DYNAMIC + ) + runtime_y_expr = shape_env.create_symbol( + 4, LocalSource("y"), dynamic_dim=DimDynamic.DYNAMIC + ) + runtime_x = shape_env.create_symintnode(runtime_x_expr, hint=4) + runtime_y = shape_env.create_symintnode(runtime_y_expr, hint=4) + + with FakeTensorMode(shape_env=shape_env): + fake_x = torch.empty(runtime_x) + fake_y = torch.empty(runtime_y) + _apply_symbolic_shape_expressions([fake_x, fake_y], shape_info) + + assert runtime_x_expr != runtime_y_expr + assert shape_env.simplify(runtime_x_expr - runtime_y_expr) == 0 + + def test_guards_distinct_runtime_symbols_for_composite_mapping(self): + """The same guard mechanism preserves derived engine-input + relationships such as y.shape[0] == 2 * x.shape[0].""" + compile_base = sympy.Symbol("s0", integer=True) + shape_info = self._two_input_shape_info( + compile_base, 2 * compile_base, compile_base + ) + + shape_env = ShapeEnv() + runtime_x_expr = shape_env.create_symbol( + 4, LocalSource("x"), dynamic_dim=DimDynamic.DYNAMIC + ) + runtime_y_expr = shape_env.create_symbol( + 8, LocalSource("y"), dynamic_dim=DimDynamic.DYNAMIC + ) + runtime_x = shape_env.create_symintnode(runtime_x_expr, hint=4) + runtime_y = shape_env.create_symintnode(runtime_y_expr, hint=8) + + with FakeTensorMode(shape_env=shape_env): + fake_x = torch.empty(runtime_x) + fake_y = torch.empty(runtime_y) + _apply_symbolic_shape_expressions([fake_x, fake_y], shape_info) + + assert shape_env.simplify(2 * runtime_x_expr - runtime_y_expr) == 0 + + def test_defers_repeated_direct_mapping_for_unbacked_symbols(self): + """Distinct data-dependent symbols have no hints, so their required + equality must be recorded as a deferred runtime assertion.""" + compile_base = sympy.Symbol("s0", integer=True) + shape_info = self._two_input_shape_info( + compile_base, compile_base, compile_base + ) + + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + runtime_x = shape_env.create_unbacked_symint() + runtime_y = shape_env.create_unbacked_symint() + shape_env._constrain_range_for_size(runtime_x.node.expr) + shape_env._constrain_range_for_size(runtime_y.node.expr) + fake_x = torch.empty(runtime_x) + fake_y = torch.empty(runtime_y) + deferred_asserts_before = shape_env.num_deferred_runtime_asserts + + _apply_symbolic_shape_expressions([fake_x, fake_y], shape_info) + + assert shape_env.num_deferred_runtime_asserts == deferred_asserts_before + 1 + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"])