[Call TIR] Fix bug when invoking call_tir with scalar values. (#254)

This small PR changes a check in the tvmscript parser to support empty shape tuples which are used to represent scalars. I added a scalar addition test to make sure it works properly.

Author: Josh Fromm

python/tvm/script/relax/parser.py

@@ -1399,7 +1399,8 @@ def transform_expr(self, expr: ast.Expr) -> relax.Expr:
         elif isinstance(expr, ast.Tuple):
             fields = [self.transform_expr(field) for field in expr.values]
 
-            if all([isinstance(f, str) for f in fields]):
+            # Empty shape tuples should be treated as shape expressions.
+            if all([isinstance(f, str) for f in fields]) and len(fields) != 0:
                 return tuple(fields)
 
             # TODO(@altanh): this check might be too weak; we really only accept integral PrimExprs

tests/python/relax/test_parser.py

@@ -721,6 +721,34 @@ def f(x: Tensor) -> Tensor:
     )
 
 
+def test_empty_shape():
+    @R.function
+    def f(x: Tensor((), "float32"), y: Tensor((), "float32")):
+        @T.prim_func
+        def scalar_add(a: T.handle, b: T.handle, c: T.handle) -> None:
+            A = T.match_buffer(a, ())
+            B = T.match_buffer(b, ())
+            C = T.match_buffer(c, ())
+
+            with T.block("add"):
+                C[()] = A[()] + B[()]
+
+        z = relax.call_tir(scalar_add, (x, y), (), dtype="float32")
+        return z
+
+    x, y = f.params
+    add_bind, z_bind = f.body.blocks[0].bindings
+
+    assert add_bind.var.name_hint == "scalar_add"
+    assert isinstance(add_bind.value, tir.PrimFunc)
+
+    check_call(
+        z_bind.value,
+        "relax.call_tir",
+        [add_bind.var, relax.Tuple([x, y]), relax.ShapeExpr([])],
+    )
+
+
 def test_class_irmodule():
     @tvm.script.ir_module
     class MyModule: