Skip to content

Conversation

@sunggg
Copy link
Contributor

@sunggg sunggg commented Jul 4, 2023

This PR introduces an utility pass that binds symbolic variables to user-provided integer values.
For example, say we have a following IRModule.

@tvm.script.ir_module
class Before:
  @R.function
  def main(
    x: R.Tensor(("m", "n")),
    y: R.Tensor(("m", "n"))
   ) -> R.Tensor(("m", "n")):
      m = T.Var("m", "int64")
      n = T.Var("m", "int64")
      with R.dataflow():
        out = R.matmul(x, y)
        R.output(out)
      return out

We can conveniently bind the symbolic variable by applying After = relax.transform.BindSymVars("main", {"m": 10, "n": 10})(Before).

@tvm.script.ir_module
class After:
  @R.function
  def main(
    x: R.Tensor((10, 10)),
    y: R.Tensor((10, 10))
   ) -> R.Tensor((10, 10)):
      with R.dataflow():
        out = R.matmul(x, y)
        R.output(out)
      return out

This would be useful when specializing shape and providing compile-time shape info (e.g., model params or batch sizes) by eliminating the need to rewrite the model.
cc. @tqchen @psrivas2

@tvm-bot
Copy link
Collaborator

tvm-bot commented Jul 4, 2023

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 by @-ing them in a comment.

Generated by tvm-bot

@sunggg sunggg closed this Jul 5, 2023
@sunggg sunggg deleted the BindSymVars branch July 5, 2023 17:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants