user provided bound for torchtrt compile when export dimension is unbounded

Author: apbose

py/torch_tensorrt/dynamo/_compiler.py

@@ -5,8 +5,9 @@
 import os
 import platform
 import warnings
-from typing import Any, Collection, List, Optional, Sequence, Union
+from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union
 
+import sympy
 import torch
 from torch.export import ExportedProgram
 from torch.fx.node import Target
@@ -791,7 +792,7 @@ def _insert_complex_io_adapters(
       Outputs: insert view_as_complex before the output node for each originally-complex
                output that comes from a TRT block.
 
-    Leverages metadata that was captued when the complex rewriter pass was run
+    Leverages metadata that was captured when the complex rewriter pass was run
     """
     complex_input_names = gm.meta.get("complex_input_names", [])
     complex_input_dtypes = gm.meta.get("complex_input_dtypes", {})
@@ -875,6 +876,75 @@ def _insert_complex_io_adapters(
         partitioned_module.recompile()
 
 
+def _build_user_symbol_bounds(
+    gm: torch.fx.GraphModule,
+    sample_arg_inputs: Sequence[Input],
+    sample_kwarg_inputs: dict[Any, Any],
+) -> Dict[sympy.Symbol, Tuple[int, int]]:
+    """Collect ``{sympy.Symbol: (min, max)}`` from user-supplied ``Input``s.
+
+    This is a *read-only* bridge between ``torch_tensorrt.Input`` (where the
+    user declares ``min_shape``/``max_shape``) and the partitioner's shape
+    reader (``extract_var_range_info``). Each sympy symbol that appears in a
+    top-level placeholder's ``meta["val"].shape`` is recorded once with the
+    corresponding ``(min_shape[d], max_shape[d])`` from the user-provided
+    dynamic ``Input``.
+
+    The map is consulted only to *fill* missing upper bounds (``int_oo`` /
+    unbounded) left by ``Dim.DYNAMIC``; the parent ``ShapeEnv`` is never
+    mutated. As a result, downstream consumers such as
+    :func:`torch_tensorrt.save(..., output_format="exported_program")` see
+    the original ``range_constraints`` from the exporter.
+    """
+    placeholders = [n for n in gm.graph.nodes if n.op == "placeholder"]
+
+    sample_by_name: dict[str, Input] = {}
+    for i, node in enumerate(placeholders):
+        if i < len(sample_arg_inputs):
+            inp = sample_arg_inputs[i]
+            if isinstance(inp, Input) and inp.shape_mode == Input._ShapeMode.DYNAMIC:
+                sample_by_name[node.target] = inp
+    for name, inp in sample_kwarg_inputs.items():
+        if isinstance(inp, Input) and inp.shape_mode == Input._ShapeMode.DYNAMIC:
+            sample_by_name[name] = inp
+
+    user_symbol_bounds: Dict[sympy.Symbol, Tuple[int, int]] = {}
+    if not sample_by_name:
+        return user_symbol_bounds
+
+    for node in placeholders:
+        if node.target not in sample_by_name:
+            continue
+        sample_input = sample_by_name[node.target]
+        fake_val = node.meta.get("val")
+        if not isinstance(fake_val, torch.Tensor):
+            continue
+
+        min_shape = sample_input.shape["min_shape"]
+        max_shape = sample_input.shape["max_shape"]
+
+        for d, dim in enumerate(fake_val.size()):
+            if not isinstance(dim, torch.SymInt) or d >= len(min_shape):
+                continue
+            expr = dim.node.expr
+            # Only record bounds for plain symbols. Composite expressions
+            # (e.g. ``2*s0``) are reconstructed by ShapeEnv.bound_sympy and
+            # would be incorrect to override directly.
+            if not isinstance(expr, sympy.Symbol):
+                continue
+            if expr in user_symbol_bounds:
+                continue
+            user_symbol_bounds[expr] = (int(min_shape[d]), int(max_shape[d]))
+            logger.debug(
+                "Recorded user-supplied bounds for %s: [%d, %d]",
+                expr,
+                int(min_shape[d]),
+                int(max_shape[d]),
+            )
+
+    return user_symbol_bounds
+
+
 @fn_supports_debugger  # type: ignore[misc]
 def compile_module(
     gm: torch.fx.GraphModule,
@@ -906,6 +976,16 @@ def compile_module(
     if sample_kwarg_inputs is None:
         sample_kwarg_inputs = {}
 
+    # Build a read-only ``{sympy.Symbol: (min, max)}`` map from the user's
+    # sample ``Input`` objects. This is forwarded to the partitioner so that
+    # symbols whose upper bound is left unbounded by ``Dim.DYNAMIC`` get
+    # filled in with the user's declared bounds, *without* mutating the
+    # exporter's ``ShapeEnv.var_to_range`` (which preserves the original
+    # ``range_constraints`` on save / re-export).
+    user_symbol_bounds = _build_user_symbol_bounds(
+        gm, sample_arg_inputs, sample_kwarg_inputs
+    )
+
     # Configure user compilation settings to converters.
     CONVERTERS.set_compilation_settings(settings)
 
@@ -1087,7 +1167,9 @@ def preserve_module_specs(
         )
 
         # Get the submodule inputs for min, opt, max shapes of the graph inputs
-        submodule_inputs = partitioning.construct_submodule_inputs(submodule)
+        submodule_inputs = partitioning.construct_submodule_inputs(
+            submodule, user_symbol_bounds=user_symbol_bounds
+        )
 
         assert submodule_inputs is not None
 

py/torch_tensorrt/dynamo/partitioning/common.py

@@ -1,10 +1,10 @@
 import logging
 from typing import Any, Dict, Optional, Sequence, Set, Tuple
 
+import sympy
 import torch
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
-
 from torch_tensorrt._Input import Input
 from torch_tensorrt.dynamo.utils import (
     COMPLEX_TO_REAL_DTYPE,
@@ -20,11 +20,16 @@ def construct_dynamic_input(
     input_dtype: torch.dtype,
     name: str = "",
     is_shape_tensor: bool = False,
+    user_symbol_bounds: Optional[Dict[sympy.Symbol, Tuple[int, int]]] = None,
 ) -> Input:
     """
     Constructs a torch_tensorrt.Input based on a symbolic input
     Args:
         input_shape: A symbolic shape / regular shape of a tensor (which can have a  mix of SymInt nodes and static values)
+        user_symbol_bounds: Optional read-only ``{sym: (min, max)}`` map, used
+            only to fill missing upper bounds when the exporter's ``ShapeEnv``
+            reports them as unbounded. See
+            :func:`torch_tensorrt.dynamo.utils.extract_var_range_info`.
     Returns:
         A dynamic shaped torch_tensorrt.Input which has the properties of the symbolic shaped input.
     """
@@ -33,7 +38,9 @@ def construct_dynamic_input(
     max_shape = []
     for d, dim in enumerate(input_shape):
         if isinstance(dim, torch.SymInt):
-            min_max_opt = extract_var_range_info(dim)
+            min_max_opt = extract_var_range_info(
+                dim, user_symbol_bounds=user_symbol_bounds
+            )
             unwrapped_min_max_opt: Dict[str, int] = {}
             if "min" not in min_max_opt or min_max_opt["min"] is None:
                 logger.warning(
@@ -85,9 +92,14 @@ def get_input(
     dtype: torch.dtype,
     name: str = "",
     is_shape_tensor: bool = False,
+    user_symbol_bounds: Optional[Dict[sympy.Symbol, Tuple[int, int]]] = None,
 ) -> Input:
     """
     Based on type of dimensions in the input_shape, construct regular or dynamic shaped inputs
+
+    ``user_symbol_bounds`` is forwarded to :func:`construct_dynamic_input` and
+    is only consulted when an exporter-supplied symbolic dimension has no
+    upper bound (e.g. ``Dim.DYNAMIC`` without an explicit ``max``).
     """
     if dtype in COMPLEX_TO_REAL_DTYPE:
         real_dtype = COMPLEX_TO_REAL_DTYPE[dtype]
@@ -106,19 +118,27 @@ def get_input(
             dtype,
             name=name,
             is_shape_tensor=is_shape_tensor,
+            user_symbol_bounds=user_symbol_bounds,
         )
     else:
         return Input(
             shape=input_shape, dtype=dtype, name=name, is_shape_tensor=is_shape_tensor
         )
 
 
-def construct_submodule_inputs(module: torch.fx.GraphModule) -> Sequence[Input]:
+def construct_submodule_inputs(
+    module: torch.fx.GraphModule,
+    user_symbol_bounds: Optional[Dict[sympy.Symbol, Tuple[int, int]]] = None,
+) -> Sequence[Input]:
     """
     Construct torch_tensorrt Inputs based on the module inputs.
     The module inputs will have meta data which has the shape and dtype info
     Args:
         module: Input FX GraphModule
+        user_symbol_bounds: Optional read-only ``{sym: (min, max)}`` map built
+            from user-supplied ``torch_tensorrt.Input`` objects. Used only to
+            fill missing upper bounds left by ``Dim.DYNAMIC`` in the exporter;
+            never mutates the parent ``ShapeEnv``.
     Returns:
         Sequence of torch_tensorrt.Input's representing inputs to given module
     """
@@ -134,7 +154,12 @@ def construct_submodule_inputs(module: torch.fx.GraphModule) -> Sequence[Input]:
                     if isinstance(input_meta, (FakeTensor, torch.Tensor)):
                         input_shape = input_meta.size()
                         torchtrt_inputs.append(
-                            get_input(input_shape, input_meta.dtype, name=input.name)
+                            get_input(
+                                input_shape,
+                                input_meta.dtype,
+                                name=input.name,
+                                user_symbol_bounds=user_symbol_bounds,
+                            )
                         )
                     elif isinstance(input_meta, torch.SymInt):
                         # Assuming sym_integers | shape inputs always have torch.int64 dtype
@@ -144,6 +169,7 @@ def construct_submodule_inputs(module: torch.fx.GraphModule) -> Sequence[Input]:
                                 torch.int64,
                                 name=input.name,
                                 is_shape_tensor=True,
+                                user_symbol_bounds=user_symbol_bounds,
                             )
                         )
                     elif isinstance(input_meta, torch.SymFloat):
@@ -153,6 +179,7 @@ def construct_submodule_inputs(module: torch.fx.GraphModule) -> Sequence[Input]:
                                 torch.float32,
                                 name=input.name,
                                 is_shape_tensor=False,  # Only SymInt inputs are treated as shape tensors
+                                user_symbol_bounds=user_symbol_bounds,
                             )
                         )
                     else:
@@ -164,7 +191,12 @@ def construct_submodule_inputs(module: torch.fx.GraphModule) -> Sequence[Input]:
                     input_meta = input.meta["tensor_meta"]
                     input_shape = input_meta.shape
                     torchtrt_inputs.append(
-                        get_input(input_shape, input_meta.dtype, name=input.name)
+                        get_input(
+                            input_shape,
+                            input_meta.dtype,
+                            name=input.name,
+                            user_symbol_bounds=user_symbol_bounds,
+                        )
                     )
                 else:
                     raise AssertionError(

py/torch_tensorrt/dynamo/utils.py

@@ -24,12 +24,9 @@
 import sympy
 import tensorrt as trt
 import torch
-from torch._subclasses.fake_tensor import FakeTensor
-from torch._subclasses.fake_tensor import FakeScriptObject
+from torch._subclasses.fake_tensor import FakeScriptObject, FakeTensor
 from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
 from torch.utils._sympy.numbers import int_oo
-
-from packaging import version
 from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import dtype
 from torch_tensorrt._features import ENABLED_FEATURES
@@ -40,6 +37,8 @@
 from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
 from torch_tensorrt.dynamo._settings import CompilationSettings
 
+from packaging import version
+
 from .types import TRTDataType
 
 logger = logging.getLogger(__name__)
@@ -105,7 +104,7 @@ class Frameworks(Enum):
     torch.complex128: torch.float64,
 }
 
-COMPLEX_DTYPES: frozenset = frozenset(COMPLEX_TO_REAL_DTYPE)
+COMPLEX_DTYPES: frozenset[torch.dtype] = frozenset(COMPLEX_TO_REAL_DTYPE)
 
 
 def unified_dtype_converter(
@@ -438,9 +437,24 @@ def contains_sym_int(tensor: torch.Tensor) -> bool:
     return any(isinstance(dim, torch.SymInt) for dim in tensor)
 
 
-def extract_var_range_info(symbolic_integer: torch.SymInt) -> Dict[str, Optional[int]]:
+def extract_var_range_info(
+    symbolic_integer: torch.SymInt,
+    user_symbol_bounds: Optional[Dict[sympy.Symbol, Tuple[int, int]]] = None,
+) -> Dict[str, Optional[int]]:
     """
     This function returns the min, max, opt values of a symbolic integer.
+
+    Args:
+        symbolic_integer: The ``torch.SymInt`` whose range is being queried.
+        user_symbol_bounds: Optional read-only map from a top-level sympy symbol
+            to ``(min, max)`` bounds supplied by the user via
+            ``torch_tensorrt.Input``. These are used **only** to fill the gap
+            when the exporter's ``ShapeEnv`` reports an unbounded upper range
+            (``int_oo``). The exporter's bounds always win when they are
+            finite; the lower bound is intersected (``max(exporter_lower,
+            user_lower)``) so we never widen the exporter's 0/1 specialization
+            (e.g. ``lower == 2 -> 1``) to ``0`` when the user passes
+            ``min_shape=0``. ``ShapeEnv`` itself is never mutated.
     """
     node = symbolic_integer.node
     expr = node.expr
@@ -474,6 +488,23 @@ def extract_var_range_info(symbolic_integer: torch.SymInt) -> Dict[str, Optional
 
     # Torchdynamo 0/1 specialization outlier
     min_val = 1 if min_val == 2 else min_val
+
+    # If the exporter left this symbol with an unbounded upper range (i.e.
+    # the user used ``Dim.DYNAMIC`` without an explicit upper), fall back to
+    # the bounds the user supplied via ``torch_tensorrt.Input(min_shape=...,
+    # max_shape=...)``. Only fills the gap; never overrides a finite exporter
+    # max. The lower bound is intersected so the exporter's specialization
+    # (e.g. lower == 1) is preserved.
+    if (
+        max_val is None
+        and user_symbol_bounds
+        and isinstance(expr, sympy.Symbol)
+        and expr in user_symbol_bounds
+    ):
+        user_min, user_max = user_symbol_bounds[expr]
+        min_val = max(min_val, int(user_min))
+        max_val = int(user_max)
+
     min_max_opt: Dict[str, Optional[int]] = {}
     min_max_opt["min"] = min_val
     min_max_opt["max"] = max_val