You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "Lunderberg (via GitHub)" <gi...@apache.org> on 2023/08/08 20:44:26 UTC

[GitHub] [tvm] Lunderberg opened a new pull request, #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Lunderberg opened a new pull request, #15509:
URL: https://github.com/apache/tvm/pull/15509

   If a function has dynamic shape parameters, it can be useful to replace them with static parameters (e.g. when producing several models within the same family).  This commit introduces a utility function `relax.Function.bind_symbolic_vars`, which allows symbolic variables to be replaced with static values.
   
   This is a related to the parameter binding done in `relax.transform.BindParam`, but does not require the bound parameter to be fully static data array.


-- 
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


[GitHub] [tvm] sunggg commented on pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

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

   Thanks, @Lunderberg! We might have implemented the same feature haha
   Here is my PR https://github.com/apache/tvm/pull/15246 under review. 


-- 
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


[GitHub] [tvm] Lunderberg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1294638372


##########
src/relax/analysis/struct_info_analysis.cc:
##########
@@ -1014,14 +1034,14 @@ class SymbolicVarCollector : public relax::ExprVisitor,
   }
 
   void VisitStructInfoExprField(const PrimExpr& expr) final {
-    if (mode_ == VisitMode::kMatchVarDef && expr->IsInstance<tir::VarNode>()) {
-      // populate symbolic var in first occurrence
-      const auto& var = Downcast<tir::Var>(expr);
-      if (defined_symbolic_var_.count(var) == 0) {
-        defined_symbolic_var_.insert(var);
+    if (mode_ >= VisitMode::kMatchVarDef) {

Review Comment:
   The pattern of GT/LT enums is used a couple of places, most commonly in the `OpPatternKind` when checking if an operation is as permissible or more permissible than required for a transformation.  That said, I agree that this isn't the most readable, and it probably would be better to have separate visitors to collect defined variables and to collect free variables.



-- 
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


[GitHub] [tvm] slyubomirsky commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1296454252


##########
src/relax/analysis/struct_info_analysis.cc:
##########
@@ -994,6 +1003,17 @@ class SymbolicVarCollector : public relax::ExprVisitor,
 
   void VisitStructInfo_(const FuncStructInfoNode* op) final {
     if (op->params.defined()) {
+      // Visit the parameters once to collect bindings, and another
+      // time to collect usages.  Otherwise, a symbolic variable
+      // defined by a later parameter may be treated as undefined when
+      // used by an earlier parameter.

Review Comment:
   Aha, I hadn't thought of that case. Very good.



-- 
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


[GitHub] [tvm] Lunderberg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1294748776


##########
src/relax/analysis/struct_info_analysis.cc:
##########
@@ -1014,14 +1034,14 @@ class SymbolicVarCollector : public relax::ExprVisitor,
   }
 
   void VisitStructInfoExprField(const PrimExpr& expr) final {
-    if (mode_ == VisitMode::kMatchVarDef && expr->IsInstance<tir::VarNode>()) {
-      // populate symbolic var in first occurrence
-      const auto& var = Downcast<tir::Var>(expr);
-      if (defined_symbolic_var_.count(var) == 0) {
-        defined_symbolic_var_.insert(var);
+    if (mode_ >= VisitMode::kMatchVarDef) {

Review Comment:
   Making two separate visitors ended up with too much duplicate code, so I instead made `VisitMode` explicitly be a bitflag, with `kProvideDefinition` and `kRequireDefinition`.  I think this cleans up the implementation, allowing the behavior to be specified with more granularity, and avoids the use of `>=` and `<=` for enums.



-- 
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


[GitHub] [tvm] Lunderberg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1294748776


##########
src/relax/analysis/struct_info_analysis.cc:
##########
@@ -1014,14 +1034,14 @@ class SymbolicVarCollector : public relax::ExprVisitor,
   }
 
   void VisitStructInfoExprField(const PrimExpr& expr) final {
-    if (mode_ == VisitMode::kMatchVarDef && expr->IsInstance<tir::VarNode>()) {
-      // populate symbolic var in first occurrence
-      const auto& var = Downcast<tir::Var>(expr);
-      if (defined_symbolic_var_.count(var) == 0) {
-        defined_symbolic_var_.insert(var);
+    if (mode_ >= VisitMode::kMatchVarDef) {

Review Comment:
   Making two separate visitors ended up with too much duplicate code, so I instead made `VisitMode` explicitly be a bitflag, with `kProvideDefinition` and `kRequireDefinition`.  I think this cleans up the implementation, allowing the behavior to be specified with more granularity.



-- 
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


[GitHub] [tvm] Lunderberg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1290609732


##########
tests/python/relax/test_bind_symbolic_vars.py:
##########
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import relax as R, tir as T
+
+replace_by_tir_var = tvm.testing.parameter(
+    by_dict={"replace-by-string": False, "replace-by-tir-var": True}
+)
+
+
+def test_bind_static_value(replace_by_tir_var):
+    """Symbolic vars may be replaced
+
+    The replaced variables may be given either as strings, or as TIR variables
+    """
+
+    @R.function(private=True)
+    def before(A: R.Tensor(("M", "K")), B: R.Tensor(("K", "N"))) -> R.Tensor(("M", "N")):
+        return R.matmul(A, B)
+
+    @R.function(private=True)
+    def expected(A: R.Tensor((128, 64)), B: R.Tensor((64, 32))) -> R.Tensor((128, 32)):
+        return R.matmul(A, B)
+
+    if replace_by_tir_var:
+        M, K = before.params[0].struct_info.shape
+        _, N = before.params[1].struct_info.shape
+        symbolic_var_map = {M: 128, K: 64, N: 32}
+    else:
+        symbolic_var_map = {"M": 128, "K": 64, "N": 32}
+
+    after = before.bind_symbolic_vars(symbolic_var_map)
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_error_with_duplicate_var_names():
+    """Duplicate variable names may not be replaced by string

Review Comment:
   On the C++ side, both TIR variables and Relax variables have reference identity, and duplicate names can exist within the same function.  (Technically, the variables only have a `name_hint`, not a name, which I believe is to avoid the assumption of unique names.  (Though, double parenthetically, some of the python APIs expose `name_hint` as `name`.))  TVMScript has a tighter restriction, where variables tracked by the parser may not repeat within the same scope.  In this test, I define the variables with the same TIR variable name outside of TVMScript, to produce the ambiguous test case anyways.
   
   As a bit of fun, the function can be inspected, and even though TVMScript modifies the variable names to avoid printing duplicate variables, the underlying objects still have duplicate names.
   
   ```python
   M = tir.Var("M", "int64")
   N1 = tir.Var("N", "int64")
   N2 = tir.Var("N", "int64")
   
   dtype = "float32"
   
   
   @R.function
   def func(
       A: R.Tensor([M, N2], dtype), B: R.Tensor([N1, N2], dtype)
   ) -> R.Tensor([M, N1], dtype):
       A: R.Tensor([M, N2], dtype) = A
       B: R.Tensor([N1, N2], dtype) = B
       out = R.linear(A, B)
       return out
   
   
   # Prints  B: R.Tensor(("N", "N_1"), dtype="float32")
   func.show()
   # But the true name can still be extracted
   assert func.params[1].struct_info.shape[0].name == "N"
   assert func.params[1].struct_info.shape[1].name == "N"
   # And the function is well-formed, despite the duplicate variable
   # names.
   assert tvm.relax.analysis.well_formed(tvm.IRModule.from_expr(func))
   ```



-- 
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


[GitHub] [tvm] slyubomirsky commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1293837242


##########
src/relax/analysis/struct_info_analysis.cc:
##########
@@ -994,6 +1003,17 @@ class SymbolicVarCollector : public relax::ExprVisitor,
 
   void VisitStructInfo_(const FuncStructInfoNode* op) final {
     if (op->params.defined()) {
+      // Visit the parameters once to collect bindings, and another
+      // time to collect usages.  Otherwise, a symbolic variable
+      // defined by a later parameter may be treated as undefined when
+      // used by an earlier parameter.

Review Comment:
   When does this happen? I think by definition, the symbolic variables are defined at their first appearance in the signature (going left to right)



-- 
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


[GitHub] [tvm] Lunderberg commented on pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

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

   @sunggg Also, I've added you as a co-author for this PR.


-- 
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


[GitHub] [tvm] Lunderberg commented on pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

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

   @sunggg How does the IRModule pass in the latest commit look?  
   
   * Swapped order of `binding_map` and `func_name`, to allow a default value of `func_name = None`.  If no function name is specified, all functions within the module with be updated.
   * Pulled in the error checking from the `relax.Function.bind_symbolic_vars` for the `IRModule` pass.  If a single function is updated, each value in the `binding_map` must be used in that function.  If multiple functions are updated, each values in the `binding_map` must occur as a replacement in at least one function.
   * Allows replacement either by TIR variable or by name.  If multiple functions are updated, a replacement by name may result in updates to a different TIR variable in each function, so long as it uniquely maps to a TIR variable within a given function.  (e.g. Replacing `"batch_size"` in all functions in a module.)
   
   The error checking was why I ended up going with a module pass instead of just wrapping the utility with a function pass.  Since the string names would be very convenient to hand-write, they're also the greatest risk of a user typo, which I wanted to guard against.  Having a function pass would require either disabling the "all replacements are used", or restricting the updates to a single 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


[GitHub] [tvm] Lunderberg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1294622521


##########
src/relax/utils.cc:
##########
@@ -61,6 +65,49 @@ class ExprBinder : public ExprMutator {
     }
   }
 
+  Expr VisitExpr_(const CallNode* op) final {
+    auto call_node = Downcast<Call>(ExprMutator::VisitExpr_(op));
+
+    // Special case for strided_slice
+    //
+    // The strided_slice operator currently stores the begins/ends in
+    // the CallNode::attrs.  Because the CallNode::attrs is only
+    // intended to store static information, any PrimExpr members in
+    // the attributes are not visited by `ExprMutator::VisitPrimExpr`.
+    // Therefore, these must be explicitly visited.
+    //
+    // When the strided_slice operator is updated to store begins/ends
+    // as a tuple of `relax::PrimValue` in the arguments, this special
+    // case can be removed.

Review Comment:
   That's true, and I'd like to update `strided_slice` at some point as well.  The comments in the `strided_slice` implementation state that it should be updated when `relax.PrimValue` is more widely supported, so that is the long-term goal.  However, I think that's a larger change, as it would be a change to the Relax IR, and not just a transformation between Relax functions, and should probably be left for a later PR.



-- 
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


[GitHub] [tvm] sunggg commented on pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

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

   Thank you for the great summary! 
   In fact, I like your PR in the sense that it offers broader supports for TIR and it accepts both strings and `tir.Var`.
   Since I think a pass would be convenient when we are dealing with IRModule with multiple functions inside, how about creating a function pass that wraps your utility function? I agree that this would be easier path to merge our PRs. 


-- 
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


[GitHub] [tvm] tvm-bot commented on pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

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

   <!---bot-comment-->
   
   Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @-ing them in a comment.
   
   <!--bot-comment-ccs-start-->
    * cc @quic-sanirudh <sub>See [#10317](https://github.com/apache/tvm/issues/10317) for details</sub><!--bot-comment-ccs-end-->
   
   <sub>Generated by [tvm-bot](https://github.com/apache/tvm/blob/main/ci/README.md#github-actions)</sub>


-- 
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


[GitHub] [tvm] Lunderberg merged pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg merged PR #15509:
URL: https://github.com/apache/tvm/pull/15509


-- 
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


[GitHub] [tvm] sunggg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "sunggg (via GitHub)" <gi...@apache.org>.
sunggg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1290717965


##########
tests/python/relax/test_bind_symbolic_vars.py:
##########
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import relax as R, tir as T
+
+replace_by_tir_var = tvm.testing.parameter(
+    by_dict={"replace-by-string": False, "replace-by-tir-var": True}
+)
+
+
+def test_bind_static_value(replace_by_tir_var):
+    """Symbolic vars may be replaced
+
+    The replaced variables may be given either as strings, or as TIR variables
+    """
+
+    @R.function(private=True)
+    def before(A: R.Tensor(("M", "K")), B: R.Tensor(("K", "N"))) -> R.Tensor(("M", "N")):
+        return R.matmul(A, B)
+
+    @R.function(private=True)
+    def expected(A: R.Tensor((128, 64)), B: R.Tensor((64, 32))) -> R.Tensor((128, 32)):
+        return R.matmul(A, B)
+
+    if replace_by_tir_var:
+        M, K = before.params[0].struct_info.shape
+        _, N = before.params[1].struct_info.shape
+        symbolic_var_map = {M: 128, K: 64, N: 32}
+    else:
+        symbolic_var_map = {"M": 128, "K": 64, "N": 32}
+
+    after = before.bind_symbolic_vars(symbolic_var_map)
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_error_with_duplicate_var_names():
+    """Duplicate variable names may not be replaced by string

Review Comment:
   Ah, my apology. I missed the part that it already shows renamed varname when printing functions in your example. Just realized by trying by my own. My concern was printing the same name for different variables, which is turned out to be a false alarm. So, I agree with you, this seems okay. Thank you for clarifying!



-- 
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


[GitHub] [tvm] sunggg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "sunggg (via GitHub)" <gi...@apache.org>.
sunggg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1289500660


##########
tests/python/relax/test_bind_symbolic_vars.py:
##########
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import relax as R, tir as T
+
+replace_by_tir_var = tvm.testing.parameter(
+    by_dict={"replace-by-string": False, "replace-by-tir-var": True}
+)
+
+
+def test_bind_static_value(replace_by_tir_var):
+    """Symbolic vars may be replaced
+
+    The replaced variables may be given either as strings, or as TIR variables
+    """
+
+    @R.function(private=True)
+    def before(A: R.Tensor(("M", "K")), B: R.Tensor(("K", "N"))) -> R.Tensor(("M", "N")):
+        return R.matmul(A, B)
+
+    @R.function(private=True)
+    def expected(A: R.Tensor((128, 64)), B: R.Tensor((64, 32))) -> R.Tensor((128, 32)):
+        return R.matmul(A, B)
+
+    if replace_by_tir_var:
+        M, K = before.params[0].struct_info.shape
+        _, N = before.params[1].struct_info.shape
+        symbolic_var_map = {M: 128, K: 64, N: 32}
+    else:
+        symbolic_var_map = {"M": 128, "K": 64, "N": 32}
+
+    after = before.bind_symbolic_vars(symbolic_var_map)
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_error_with_duplicate_var_names():
+    """Duplicate variable names may not be replaced by string

Review Comment:
   I'm surprised to see this case passing the well_formed check. I think this should be filtered before any pass is applied, because this seems incorrect to me. 



-- 
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


[GitHub] [tvm] slyubomirsky commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1293840657


##########
src/relax/analysis/struct_info_analysis.cc:
##########
@@ -1014,14 +1034,14 @@ class SymbolicVarCollector : public relax::ExprVisitor,
   }
 
   void VisitStructInfoExprField(const PrimExpr& expr) final {
-    if (mode_ == VisitMode::kMatchVarDef && expr->IsInstance<tir::VarNode>()) {
-      // populate symbolic var in first occurrence
-      const auto& var = Downcast<tir::Var>(expr);
-      if (defined_symbolic_var_.count(var) == 0) {
-        defined_symbolic_var_.insert(var);
+    if (mode_ >= VisitMode::kMatchVarDef) {

Review Comment:
   Do we use this pattern elsewhere? I'm not personally very fond of it because it could create problems for adding more modes later, but I'm not against adding it if we do this already.



-- 
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


[GitHub] [tvm] Lunderberg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1290661105


##########
tests/python/relax/test_bind_symbolic_vars.py:
##########
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import relax as R, tir as T
+
+replace_by_tir_var = tvm.testing.parameter(
+    by_dict={"replace-by-string": False, "replace-by-tir-var": True}
+)
+
+
+def test_bind_static_value(replace_by_tir_var):
+    """Symbolic vars may be replaced
+
+    The replaced variables may be given either as strings, or as TIR variables
+    """
+
+    @R.function(private=True)
+    def before(A: R.Tensor(("M", "K")), B: R.Tensor(("K", "N"))) -> R.Tensor(("M", "N")):
+        return R.matmul(A, B)
+
+    @R.function(private=True)
+    def expected(A: R.Tensor((128, 64)), B: R.Tensor((64, 32))) -> R.Tensor((128, 32)):
+        return R.matmul(A, B)
+
+    if replace_by_tir_var:
+        M, K = before.params[0].struct_info.shape
+        _, N = before.params[1].struct_info.shape
+        symbolic_var_map = {M: 128, K: 64, N: 32}
+    else:
+        symbolic_var_map = {"M": 128, "K": 64, "N": 32}
+
+    after = before.bind_symbolic_vars(symbolic_var_map)
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_error_with_duplicate_var_names():
+    """Duplicate variable names may not be replaced by string

Review Comment:
   Hmm.  It is really useful if you're merging parts of TIR/Relax that were generated independently.  Those are guaranteed to have different `tir.Var` and `relax.Var` objects, but can very easily have duplicate names.  I think the only place we de-duplicate the names themselves is in an `IRModule`, and even that gets kind of hairy with both `String` and `GlobalVar` acting as unrelated unique keys.
   
   I think that re-use of the same name should be well-formed, but there should be a utility pass (potentially called several times during lowering) that de-duplicates the name within a function.  That way, it provides the better debugging experience, but doesn't require renaming variables when constructing TIR.  It also avoids promising that variable names will remain stable, or that they should be relied upon, since that isn't a stable assumption at the moment.



-- 
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


[GitHub] [tvm] Lunderberg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1294635288


##########
src/relax/analysis/struct_info_analysis.cc:
##########
@@ -994,6 +1003,17 @@ class SymbolicVarCollector : public relax::ExprVisitor,
 
   void VisitStructInfo_(const FuncStructInfoNode* op) final {
     if (op->params.defined()) {
+      // Visit the parameters once to collect bindings, and another
+      // time to collect usages.  Otherwise, a symbolic variable
+      // defined by a later parameter may be treated as undefined when
+      // used by an earlier parameter.

Review Comment:
   The symbolic variables are defined in their first appearance in the signature, but only if that appearance is sufficient to provide a definition.  If the symbolic variable can be matched to a specific expression in the caller (e.g. ), then it is bound.  If the symbolic variable cannot be matched to a specific expression in the caller, then it is a usage-site instead.  So, `R.Tensor(shape=[16])` binding to `R.Tensor(shape=[sym_var])` implies a binding of `{sym_var: 16}`, but `R.Tensor(shape=[16])` binding to `R.Tensor(shape=[sym_var1 * sym_var2])` cannot provide any definitions.
   
   So, this edge case occurs in cases like the one shown below, where the first occurrence is a usage, and the definition doesn't occur until a later parameter.
   
   ```python
   def func(A: R.Tensor(shape=[sym_var1*sym_var2]), B: R.Tensor(shape=[sym_var1, sym_var2])):
       #                       ^^^^^^^^                                ^^^^^^^^
       #                       Usage of sym_var1                       Definition of sym_var1
       ...
   ```



-- 
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


[GitHub] [tvm] sunggg commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "sunggg (via GitHub)" <gi...@apache.org>.
sunggg commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1290636588


##########
tests/python/relax/test_bind_symbolic_vars.py:
##########
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import relax as R, tir as T
+
+replace_by_tir_var = tvm.testing.parameter(
+    by_dict={"replace-by-string": False, "replace-by-tir-var": True}
+)
+
+
+def test_bind_static_value(replace_by_tir_var):
+    """Symbolic vars may be replaced
+
+    The replaced variables may be given either as strings, or as TIR variables
+    """
+
+    @R.function(private=True)
+    def before(A: R.Tensor(("M", "K")), B: R.Tensor(("K", "N"))) -> R.Tensor(("M", "N")):
+        return R.matmul(A, B)
+
+    @R.function(private=True)
+    def expected(A: R.Tensor((128, 64)), B: R.Tensor((64, 32))) -> R.Tensor((128, 32)):
+        return R.matmul(A, B)
+
+    if replace_by_tir_var:
+        M, K = before.params[0].struct_info.shape
+        _, N = before.params[1].struct_info.shape
+        symbolic_var_map = {M: 128, K: 64, N: 32}
+    else:
+        symbolic_var_map = {"M": 128, "K": 64, "N": 32}
+
+    after = before.bind_symbolic_vars(symbolic_var_map)
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_error_with_duplicate_var_names():
+    """Duplicate variable names may not be replaced by string

Review Comment:
   Thanks for the explanation, and yeah I see why it currently passes the well-formed checker. But I'm wondering if this is what we want or what it is supposed to be. I feel like this duplicated name would be very confusing if we are debugging things, especially when we print out IRs. Is there any necessity for this feature? 



-- 
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


[GitHub] [tvm] slyubomirsky commented on a diff in pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #15509:
URL: https://github.com/apache/tvm/pull/15509#discussion_r1293832517


##########
src/relax/utils.cc:
##########
@@ -61,6 +65,49 @@ class ExprBinder : public ExprMutator {
     }
   }
 
+  Expr VisitExpr_(const CallNode* op) final {
+    auto call_node = Downcast<Call>(ExprMutator::VisitExpr_(op));
+
+    // Special case for strided_slice
+    //
+    // The strided_slice operator currently stores the begins/ends in
+    // the CallNode::attrs.  Because the CallNode::attrs is only
+    // intended to store static information, any PrimExpr members in
+    // the attributes are not visited by `ExprMutator::VisitPrimExpr`.
+    // Therefore, these must be explicitly visited.
+    //
+    // When the strided_slice operator is updated to store begins/ends
+    // as a tuple of `relax::PrimValue` in the arguments, this special
+    // case can be removed.

Review Comment:
   I wonder if this is worth changing to avoid needing a special case here. Namely, we can use PrimExpr arguments for strided_slice (the reason it's implemented this way is probably because we didn't have PrimExpr at the time we added this op).



-- 
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


[GitHub] [tvm] Lunderberg commented on pull request #15509: [Unity] Implement relax.Function.bind_symbolic_vars

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

   Ooh, thank you for pointing it out!  That would explain why I didn't see the functionality anywhere in the `unity` branch itself.
   
   Looking through it, I think the main differences are:
   
   * This PR accepts either `tir.Var` (less ambiguous, easier to generate programmatically), or strings (more ambiguous, much easier to hand-write), where #15246 only accepts strings.
   
   * This PR operates on a relax function, where #15246 operates on an IRModule.
   
   * This PR would ignore functions that are defined within a local function definition, where #15246 would update their symbolic variables.
   
   * This PR would replace any TIR usage supported by `tir::ExprMutator`, where #15246 checks for specific TIR nodes.
   
   * This PR raises an error if no symbolic variable is known by that name, where #15246 would silently ignore it.
   
   Thoughts on making a merged PR with the best of both implementations:
   
   * Providing methods to update either a IRModule transform or a relax Function would provide the most flexibility.
   * Implementing a IRModule transform on top of a relax Function update would be require constructing fewer intermediates than implementing a relax Function update on top of an IRModule transform.
   
   Any preference on which of the two PRs should be updated overall? Since both provide similar functionality, I don't have a huge preference.


-- 
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