You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "demoncoder-crypto (via GitHub)" <gi...@apache.org> on 2025/04/26 17:46:51 UTC

[PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

demoncoder-crypto opened a new pull request, #17898:
URL: https://github.com/apache/tvm/pull/17898

   Addresses #17818


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "mshr-h (via GitHub)" <gi...@apache.org>.
mshr-h commented on code in PR #17898:
URL: https://github.com/apache/tvm/pull/17898#discussion_r2065386482


##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -519,13 +521,18 @@ 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
-            ]
+            # UPDATED: Create SizeVars and map SymInts (removed original shape creation)

Review Comment:
   Why is this update needed? Looks like there's no functional change.



##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -534,7 +541,47 @@ def create_input_vars(
             else:
                 parameters_buffers_constants[name_hint] = relax_var
 
-        return parameters_buffers_constants, user_inputs
+        # NEW: Process range constraints (basic support for simple SymInt keys)
+        if hasattr(exported_program, "range_constraints"):
+            for torch_sym_expr, value_range in exported_program.range_constraints.items():
+                # Basic support: Only handle constraints where the key is a simple SymInt
+                if isinstance(torch_sym_expr, torch.SymInt):
+                    s_str = str(torch_sym_expr)
+                    if s_str in torch_symbol_to_relax_var:
+                        relax_tir_var = torch_symbol_to_relax_var[s_str]
+
+                        # Extract bounds, using None for infinity
+                        min_val = int(value_range.lower) if value_range.lower != -sympy.oo else None
+                        max_val = int(value_range.upper) if value_range.upper != sympy.oo else None
+
+                        if relax_tir_var not in relax_range_constraints:
+                            relax_range_constraints[relax_tir_var] = (min_val, max_val)
+                        else:
+                            # Refine existing constraints if the new one is tighter

Review Comment:
   Please also include a testcase for this condition.



##########
tests/python/relax/test_frontend_from_exported_program.py:
##########
@@ -4625,6 +4625,62 @@ def main(
     dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}}
 
     verify_model(DynamicModel(), example_args, {}, Expected, dynamic_shapes=dynamic_shapes)
+    
+def test_dynamic_shape_with_constraints():
+        B = torch.export.Dim("B", min=2, max=10)
+        S = torch.export.Dim("S", min=1)
+        # Use a tuple for args
+        example_args = (torch.randn(3, 4, dtype=torch.float32),)
+        # Dynamic shapes dict maps arg index to shape spec {dim_index: Dim obj}
+        dynamic_shapes = {0: {0: B, 1: S}}
+
+        class SimpleDynamic(torch.nn.Module):
+            def forward(self, x):
+                return torch.relu(x)
+
+        # Explicit export and import

Review Comment:
   Any reason to manually check structural equality? If no, please use the `verify_model()` utility function.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "mshr-h (via GitHub)" <gi...@apache.org>.
mshr-h commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2849069029

   I think we need multiple testcases so that we can cover all the possible cases:
   1. lower constraints only
   2. upper constraints only
   3. both the lower and upper constraints
   4. the constraints are overridden by another constraints (this is already in the test)


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "mshr-h (via GitHub)" <gi...@apache.org>.
mshr-h commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2849066379

   Adding helper functions is totally fine as long as it's readable and maintainable. @demoncoder-crypto 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "guan404ming (via GitHub)" <gi...@apache.org>.
guan404ming commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-3503910788

   Hi @mshr-h, it seems like this PR has been inactive for a while. Could I help with this? Thanks~


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "mshr-h (via GitHub)" <gi...@apache.org>.
mshr-h closed pull request #17898: feat(relax/frontend/torch): Add basic range constraint support 
URL: https://github.com/apache/tvm/pull/17898


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "mshr-h (via GitHub)" <gi...@apache.org>.
mshr-h commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-3520362712

   Closing as #18429 merged


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "demoncoder-crypto (via GitHub)" <gi...@apache.org>.
demoncoder-crypto commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2840322694

   Understood @mshr-h I will fix it. 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "demoncoder-crypto (via GitHub)" <gi...@apache.org>.
demoncoder-crypto commented on code in PR #17898:
URL: https://github.com/apache/tvm/pull/17898#discussion_r2071552597


##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -519,13 +521,18 @@ 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
-            ]
+            # UPDATED: Create SizeVars and map SymInts (removed original shape creation)

Review Comment:
   I thought List comprehensions can sometimes be a more compact way to express list creation based on iteration. 



##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -534,7 +541,47 @@ def create_input_vars(
             else:
                 parameters_buffers_constants[name_hint] = relax_var
 
-        return parameters_buffers_constants, user_inputs
+        # NEW: Process range constraints (basic support for simple SymInt keys)
+        if hasattr(exported_program, "range_constraints"):
+            for torch_sym_expr, value_range in exported_program.range_constraints.items():
+                # Basic support: Only handle constraints where the key is a simple SymInt
+                if isinstance(torch_sym_expr, torch.SymInt):
+                    s_str = str(torch_sym_expr)
+                    if s_str in torch_symbol_to_relax_var:
+                        relax_tir_var = torch_symbol_to_relax_var[s_str]
+
+                        # Extract bounds, using None for infinity
+                        min_val = int(value_range.lower) if value_range.lower != -sympy.oo else None
+                        max_val = int(value_range.upper) if value_range.upper != sympy.oo else None
+
+                        if relax_tir_var not in relax_range_constraints:
+                            relax_range_constraints[relax_tir_var] = (min_val, max_val)
+                        else:
+                            # Refine existing constraints if the new one is tighter

Review Comment:
   Done
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "mshr-h (via GitHub)" <gi...@apache.org>.
mshr-h commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-3507911045

   @guan404ming Feel free to take it over!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "demoncoder-crypto (via GitHub)" <gi...@apache.org>.
demoncoder-crypto commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2848604837

   Tried to fix the errors do let me know if it works now thanks @mshr-h 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "demoncoder-crypto (via GitHub)" <gi...@apache.org>.
demoncoder-crypto commented on code in PR #17898:
URL: https://github.com/apache/tvm/pull/17898#discussion_r2071554507


##########
tests/python/relax/test_frontend_from_exported_program.py:
##########
@@ -4625,6 +4625,62 @@ def main(
     dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}}
 
     verify_model(DynamicModel(), example_args, {}, Expected, dynamic_shapes=dynamic_shapes)
+    
+def test_dynamic_shape_with_constraints():
+        B = torch.export.Dim("B", min=2, max=10)
+        S = torch.export.Dim("S", min=1)
+        # Use a tuple for args
+        example_args = (torch.randn(3, 4, dtype=torch.float32),)
+        # Dynamic shapes dict maps arg index to shape spec {dim_index: Dim obj}
+        dynamic_shapes = {0: {0: B, 1: S}}
+
+        class SimpleDynamic(torch.nn.Module):
+            def forward(self, x):
+                return torch.relu(x)
+
+        # Explicit export and import

Review Comment:
   I thought it makes it extremely explicit exactly which part of the new feature is being tested, the presence and correctness of the constraint attributes. This was helpful during development to immediately pinpoint if the constraint propagation itself is failing. I forgot to update it after I did the validation. Sorry



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "mshr-h (via GitHub)" <gi...@apache.org>.
mshr-h commented on code in PR #17898:
URL: https://github.com/apache/tvm/pull/17898#discussion_r2072546077


##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -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

Review Comment:
   Please remove the unnecessary comment.



##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -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):

Review Comment:
   `@staticmethod` would be better since it doesn't access the instance variable or method.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "demoncoder-crypto (via GitHub)" <gi...@apache.org>.
demoncoder-crypto commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2834884777

   done @mshr-h Please check if my implementation is up to expectation


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "mshr-h (via GitHub)" <gi...@apache.org>.
mshr-h commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2833027362

   Please include test. @demoncoder-crypto 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "demoncoder-crypto (via GitHub)" <gi...@apache.org>.
demoncoder-crypto commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2847103694

   @mshr-h I have made the changes requested by you


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "demoncoder-crypto (via GitHub)" <gi...@apache.org>.
demoncoder-crypto commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2848617998

   I cannot fix the lint/pr-head its expecting less nesting which would require me to make a helper method should I proceed and make it? @mshr-h 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] feat(relax/frontend/torch): Add basic range constraint support [tvm]

Posted by "yongwww (via GitHub)" <gi...@apache.org>.
yongwww commented on PR #17898:
URL: https://github.com/apache/tvm/pull/17898#issuecomment-2892683144

   @demoncoder-crypto pls fix the conflicts and ci


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org