[TVMScript] B6/B7: Symbolic shape and var shadowing (apache#245)

This PR features symbolic shape support and var shadowing in relax

Co-authored-by: Ruihang Lai <[email protected]>

Author: Siyuan Feng

python/tvm/script/ir_builder/relax/ir.py

@@ -28,6 +28,7 @@
 
 from . import _ffi_api
 from . import frame
+from ..tir import var as _tir_var
 
 ############################### Operators ###############################
 from tvm.relax.op import shape_of, make_closure, invoke_closure
@@ -44,15 +45,15 @@ class TensorType(Object):
 
 
 def tensor(
-    shape: Optional[List[PrimExpr]] = None,
+    shape: Optional[List[Union[PrimExpr, str]]] = None,
     dtype: Optional[str] = None,
     ndim: int = -1,
 ):
     """Helper function for `R.Tensor` in parser
 
     Parameters
     ----------
-    shape: Optional[List[PrimExpr]]
+    shape: Optional[List[Union[PrimExpr, str]]]
         The shape of the tensor. It's runtime dependent if `shape` is None.
 
     dtype: Optional[str]
@@ -66,6 +67,12 @@ def tensor(
     tensor_type: TensorType
         The TensorType that is only used in ir_builder.
     """
+    if isinstance(shape, (tuple, list)):
+        shape = list(shape)
+        for i, s in enumerate(shape):
+            if isinstance(s, str):
+                shape[i] = _tir_var("int64", s)
+
     return _ffi_api.Tensor(shape, dtype, ndim)  # pylint: disable=no-member # type: ignore
 
 

python/tvm/script/parser/core/parser.py

@@ -80,9 +80,16 @@ def pop_frame():
         self.frames.append(VarTableFrame())
         return _deferred(pop_frame)
 
-    def add(self, var: str, value: Any):
-        self.frames[-1].add(var)
-        self.name2value[var].append(value)
+    def add(self, var: str, value: Any, allow_shadowing: bool = False):
+        # Skip if the key and value are equal to those in the var_table
+        if self.name2value[var] and self.name2value[var][-1] == value:
+            return
+        if allow_shadowing and var in self.frames[-1].vars:
+            # Shadowing
+            self.name2value[var][-1] = value
+        else:
+            self.frames[-1].add(var)
+            self.name2value[var].append(value)
 
     def get(self) -> Dict[str, Any]:
         return {key: values[-1] for key, values in self.name2value.items() if values}
@@ -177,13 +184,14 @@ def eval_assign(
         target: doc.expr,
         source: Any,
         bind_value: Callable[["Parser", doc.expr, str, Any], Any],
+        allow_shadowing: bool = False,
     ) -> Dict[str, Any]:
         if self._duplicate_lhs_check(target) is True:
             self.report_error(target, "Duplicate vars assigned.")
         var_values = eval_assign(self, target, source)
         for k, v in var_values.items():
             var = bind_value(self, target, k, v)
-            self.var_table.add(k, var)
+            self.var_table.add(k, var, allow_shadowing)
         return var_values
 
     def report_error(

python/tvm/script/parser/relax/parser.py

@@ -18,16 +18,41 @@
 
 from typing import Any
 
-from tvm import relax
+from tvm import tir, relax
 
 from ...ir_builder import relax as R
 from ...ir_builder.base import IRBuilder
 from .._core import Parser, dispatch, doc
 
 
 def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
-    # pylint: disable=unused-argument
-    if isinstance(value, relax.Expr):
+    var_table = self.var_table.get()
+
+    if isinstance(value, tir.Var):
+        if value.name and var_name != value.name:
+            self.report_error(
+                node,
+                "Cannot define TIR variables with different names. The LHS of binding should has "
+                "the same name provided in RHS.",
+            )
+        if var_name in var_table:
+            prev_value = var_table[var_name]
+            if not isinstance(prev_value, tir.Var):
+                self.report_error(
+                    node,
+                    "Cannot redefine a non-TIR-variable object to a TIR variable. Please define "
+                    "the TIR variable with another name.",
+                )
+            if prev_value.dtype != value.dtype:
+                self.report_error(
+                    node,
+                    "Expected the same dtype for TIR vars "
+                    f"but got {value.dtype} vs {prev_value.dtype}",
+                )
+            return prev_value
+        IRBuilder.name(var_name, value)
+        return value
+    elif isinstance(value, relax.Expr):
         var = R.emit(value)
         IRBuilder.name(var_name, var)
         return var
@@ -60,6 +85,11 @@ def visit_arguments(self: Parser, node: doc.arguments) -> None:
             self.report_error(arg, "Type annotation is required for function parameters.")
         param_type = self.visit_tvm_annotation(arg.annotation)
         param = R.arg(arg.arg, param_type)
+        # Define the symbolic shape var
+        if param_type.shape is not None:
+            for shape_expr in param_type.shape:
+                if isinstance(shape_expr, tir.Var):
+                    self.var_table.add(shape_expr.name, shape_expr)
 
         self.var_table.add(arg.arg, param)
 
@@ -78,7 +108,7 @@ def visit_assign(self: Parser, node: doc.Assign) -> None:
         self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
     lhs = node.targets[0]
     rhs = self.eval_expr(node.value)
-    self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value)
+    self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value, allow_shadowing=True)
 
 
 @dispatch.register(token="relax", type_name="Return")

tests/python/relax/test_tvmscript_parser.py

@@ -20,7 +20,7 @@
 import tvm
 import tvm.testing
 
-from tvm import relax
+from tvm import relax, tir
 from tvm import IRModule
 from tvm.script.parser import ir as I, tir as T, relax as R
 
@@ -118,5 +118,67 @@ def foo(x: R.Tensor((4, 4), "float32")):
     _check(foo, bb.get()["foo"])
 
 
+def test_symbolic_shape():
+    @R.function
+    def foo(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(None, "float32", ndim=2):
+        m = T.var("int64", "m")
+        n = T.var("int64", "n")
+        gv0 = R.call_tir("extern_func", x, (m, n), dtype="float32")
+        return gv0
+
+    @R.function
+    def bar(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(None, "float32", ndim=2):
+        m = T.var("int64")
+        n = T.var("int64")
+        gv0 = R.call_tir("extern_func", x, (m, n), dtype="float32")
+        return gv0
+
+    with pytest.raises(tvm.error.DiagnosticError):
+
+        @R.function
+        def mismatch_dtype(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(None, "float32", ndim=2):
+            m = T.var("int64")
+            n = T.var("int32")  # The shape dtype should be int64
+            gv0 = R.call_tir("extern_func", x, (m, n), dtype="float32")
+            return gv0
+
+    def _expected(name: str):
+        n, m = tir.Var("n", "int64"), tir.Var("m", "int64")
+        x = relax.Var("x", [m, n], relax.DynTensorType(2, "float32"))
+        bb = relax.BlockBuilder()
+        with bb.function(name, (x,)):
+            out = bb.emit(relax.call_tir("extern_func", x, (m, n), dtype="float32"))
+            bb.emit_func_output(out)
+        return bb.get()[name]
+
+    _check(foo, _expected("foo"))
+    _check(bar, _expected("bar"))
+
+
+def test_shadowing():
+    @R.function
+    def foo(x: R.Tensor((4, 4), "float32")):
+        y = R.add(x, x)
+        z = R.multiply(x, y)
+        y = R.add(x, y)
+        y = z
+        y = R.multiply(y, x)
+        z = y
+        return z
+
+    x = relax.Var("x", [4, 4], relax.DynTensorType(2, "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", (x,)):
+        y = bb.emit(relax.op.add(x, x))
+        z = bb.emit(relax.op.multiply(x, y))
+        y = bb.emit(relax.op.add(x, y))
+        y = bb.emit(z)
+        y = bb.emit(relax.op.multiply(y, x))
+        z = bb.emit(y)
+        bb.emit_func_output(z)
+
+    _check(foo, bb.get()["foo"])
+
+
 if __name__ == "__main__":
     tvm.testing.main()