feat(relax/frontend/torch): Add basic range constraint support

python/tvm/relax/frontend/torch/exported_program_translator.py

@@ -20,12 +20,12 @@
 """PyTorch ExportedProgram of Relax."""
 from collections import ChainMap, OrderedDict
 from functools import partial
-from typing import Callable, Dict, List, Tuple
+from typing import Callable, Dict, List, Tuple, Optional
 
 import torch
 import tvm
 from tvm import relax
-
+import tvm.tir as tir  # pylint: disable=unused-import, consider-using-from-import
 from .base_fx_graph_translator import BaseFXGraphImporter
 
 
@@ -497,11 +497,16 @@ def create_convert_map(
 
     def create_input_vars(
         self, exported_program: torch.export.ExportedProgram
-    ) -> Tuple[Dict[str, relax.Var], Dict[str, relax.Var]]:
+    ) -> Tuple[
+        Dict[str, relax.Var],
+        Dict[str, relax.Var],
+        Dict[tvm.tir.Var, Tuple[Optional[int], Optional[int]]],
+    ]:
         """Create relax input vars."""
         parameters_buffers_constants = OrderedDict()
         user_inputs = OrderedDict()
         torch_symbol_to_relax_var: Dict[str, tvm.tir.Var] = {}
+        relax_range_constraints: Dict[tvm.tir.Var, Tuple[Optional[int], Optional[int]]] = {}
 
         for spec in exported_program.graph_signature.input_specs:
             name_hint = spec.arg.name
@@ -519,13 +524,17 @@ def create_input_vars(
                 torch_shape = exported_program.state_dict[spec.target].shape
                 torch_dtype = exported_program.state_dict[spec.target].dtype
 
-            # TODO(mshr-h): Support range constraints
-            relax_shape = [
-                torch_symbol_to_relax_var.setdefault(str(s), tvm.tir.SizeVar(str(s), "int64"))
-                if isinstance(s, torch.SymInt)
-                else s
-                for s in torch_shape
-            ]
+            # Create SizeVars and map SymInts
+            relax_shape = []
+            for s in torch_shape:
+                if isinstance(s, torch.SymInt):
+                    s_str = str(s)
+                    if s_str not in torch_symbol_to_relax_var:
+                        torch_symbol_to_relax_var[s_str] = tvm.tir.SizeVar(s_str, "int64")
+                    relax_shape.append(torch_symbol_to_relax_var[s_str])
+                else:
+                    relax_shape.append(s)
+
             dtype = self._convert_data_type(torch_dtype)
 
             relax_var = relax.Var(name_hint, relax.TensorStructInfo(relax_shape, dtype))
@@ -534,7 +543,55 @@ def create_input_vars(
             else:
                 parameters_buffers_constants[name_hint] = relax_var
 
-        return parameters_buffers_constants, user_inputs
+        # Extract range constraints for TIR vars
+        if hasattr(exported_program, "range_constraints") and exported_program.range_constraints:
+            for torch_sym_expr, constraint in exported_program.range_constraints.items():
+                # Convert sympy expression to string for mapping
+                torch_sym_expr_str = str(torch_sym_expr)
+
+                if torch_sym_expr_str in torch_symbol_to_relax_var:
+                    relax_tir_var = torch_symbol_to_relax_var[torch_sym_expr_str]
+                    # TODO(sjt): Handle SymFloat, SymBool cases as well.
+                    # Note: min / max could be int or SymInt objects.
+                    # Need to handle symbolic shapes as well.
+                    min_val = constraint.min
+                    max_val = constraint.max
+                    # Call helper to add/refine constraint
+                    self._add_range_constraint(
+                        relax_range_constraints, relax_tir_var, min_val, max_val
+                    )
+                # else:
+                # FIXED Indentation for Black:
+                # TODO: Handle complex expressions (e.g., s0 + 1) for advanced support
+                # print(f"Skipping complex constraint expression: {torch_sym_expr}")
+
+        return parameters_buffers_constants, user_inputs, relax_range_constraints
+
+    # NEW HELPER METHOD
+    def _add_range_constraint(self, constraints_dict, relax_tir_var, min_val, max_val):
+        """Adds or refines a range constraint for a TIR variable."""
+        if relax_tir_var not in constraints_dict:
+            constraints_dict[relax_tir_var] = (min_val, max_val)
+        else:
+            # Refine existing constraints if the new one is tighter
+            existing_min, existing_max = constraints_dict[relax_tir_var]
+            # Merge lower bounds (take the max)
+            if existing_min is None:
+                new_min = min_val
+            elif min_val is None:
+                new_min = existing_min
+            else:
+                new_min = max(existing_min, min_val)
+
+            # Merge upper bounds (take the min)
+            if existing_max is None:
+                new_max = max_val
+            elif max_val is None:
+                new_max = existing_max
+            else:
+                new_max = min(existing_max, max_val)
+
+            constraints_dict[relax_tir_var] = (new_min, new_max)
 
     def from_exported_program(
         self,
@@ -546,23 +603,47 @@ def from_exported_program(
         """Convert a PyTorch ExportedProgram to a Relax program."""
         from torch import fx  # type: ignore
 
-        # Create input variables.
-        parameter_buffer_constant_vars, user_input_vars = self.create_input_vars(exported_program)
+        # Create input variables and get range constraints.
+        (
+            parameter_buffer_constant_vars,
+            user_input_vars,
+            relax_range_constraints,
+        ) = self.create_input_vars(exported_program)
         inputs_vars = user_input_vars.copy()
         inputs_vars.update(parameter_buffer_constant_vars)
 
         # Initialize the block builder with a function and a dataflow block.
         self.block_builder = relax.BlockBuilder()
         func_name = "main"
-        func_attrs = {"num_input": len(user_input_vars)} if keep_params_as_input else None
+
+        # Prepare function attributes
+        func_attrs = {"num_input": len(user_input_vars)} if keep_params_as_input else {}
+
+        # NEW: Add range constraints to function attributes if they exist
+        if relax_range_constraints:
+            lower_bounds = {}
+            upper_bounds = {}
+            for tir_var, (min_val, max_val) in relax_range_constraints.items():
+                if min_val is not None:
+                    lower_bounds[tir_var] = tvm.tir.IntImm("int64", min_val)
+                if max_val is not None:
+                    upper_bounds[tir_var] = tvm.tir.IntImm("int64", max_val)
+
+            if lower_bounds:
+                func_attrs["tir_var_lower_bound"] = lower_bounds
+            if upper_bounds:
+                func_attrs["tir_var_upper_bound"] = upper_bounds
+
+        # Use None if func_attrs is empty, otherwise use the dictionary
+        final_func_attrs = func_attrs if func_attrs else None
 
         nodes: List[fx.Node] = exported_program.graph.nodes
 
         # Find all the missing function types
         self._check_unsupported_func_type(nodes)
 
         with self.block_builder.function(
-            name=func_name, params=list(inputs_vars.values()).copy(), attrs=func_attrs
+            name=func_name, params=list(inputs_vars.values()).copy(), attrs=final_func_attrs
         ):
             output = None
             with self.block_builder.dataflow():

tests/python/relax/test_frontend_from_exported_program.py

@@ -20,6 +20,8 @@
 from torch.nn import Module
 from torch.export import export
 
+import tvm.tir as tir  # pylint: disable=unused-import, consider-using-from-import
+
 import tvm
 from tvm import relax
 import tvm.testing
@@ -4627,6 +4629,60 @@ def main(
     verify_model(DynamicModel(), example_args, {}, Expected, dynamic_shapes=dynamic_shapes)
 
 
+def test_dynamic_shape_with_constraints():
+    # Define SymInts with constraints
+    B = torch.export.Dim("B", min=2, max=10)
+    # Use B again for another dimension to test refinement (max(10, 15) -> 15)
+    B_refined = torch.export.Dim("B", min=3, max=15)
+    S = torch.export.Dim("S", min=1)  # Test min constraint only (-> (1, None))
+
+    # Example args matching initial B dim (max=10)
+    example_args = (torch.randn(3, 4, dtype=torch.float32), torch.randn(5, 2, dtype=torch.float32))
+
+    # Dynamic shapes using the Dim objects
+    # Input 0: Dim 0 uses B (min=2, max=10), Dim 1 uses S (min=1)
+    # Input 1: Dim 0 uses B_refined (min=3, max=15)
+    # The final constraint for tir.Var("B") should be max(2,3) to min(10,15) => min=3, max=10
+    dynamic_shapes = {0: {0: B, 1: S}, 1: {0: B_refined}}
+
+    class SimpleDynamic(torch.nn.Module):
+        # Simple op, the main thing is testing the input signature and constraints
+        def forward(self, x, y):
+            # Add tensors with different shapes requires broadcasting,
+            # but we only care about the input signature here.
+            # Use an op that doesn't depend on exact shapes matching.
+            return torch.relu(x)  # Return just one to simplify output signature
+
+    # NEW: Define TIR Vars for TVMScript parsing
+    B = tir.Var("B", "int64")
+    S = tir.Var("S", "int64")
+
+    # Define the expected Relax IRModule
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(
+            # Note: B has refined constraints: min=3, max=10
+            # Note: S has constraints: min=1
+            x: R.Tensor((B, S), dtype="float32"),
+            y: R.Tensor((B, 2), dtype="float32"),
+        ) -> R.Tuple(R.Tensor((B, S), dtype="float32")):
+            B = T.int64()
+            S = T.int64()
+            # tell TIR about the constraints via function attributes
+            T.func_attr({"tir_var_lower_bound": {B: 3, S: 1}, "tir_var_upper_bound": {B: 10}})
+            with R.dataflow():
+                # The actual body isn't the focus, just the signature
+                lv: R.Tensor((B, S), dtype="float32") = R.relu(x)
+                # Output must be a tuple
+                gv: R.Tuple(R.Tensor((B, S), dtype="float32")) = (lv,)
+                R.output(gv)
+            return gv
+
+    # Use verify_model utility
+    verify_model(SimpleDynamic(), example_args, {}, Expected, dynamic_shapes=dynamic_shapes)
+
+
 def test_broadcast_to():
     class BroadcastTo(Module):
         def forward(self, x):