[Unity] Validate struct info in relax::Call constructor (#16311)

* [Unity] Validate struct info in relax::Call constructor

All operations called by a `relax::Call` node must have a
`FuncStructInfo`.  Prior to this commit, an invalid struct info would
be caught by the `BlockBuilder` during normalization.  This delay
between the invalid `relax::Call` being constructed and the invalid
`relax::Call` being detected makes debugging difficult.
This commit adds an additional check during the `relax::Call`
constructor, to provide earlier error detection.

* Updated unit test to avoid using Tensor as callable function

Author: Lunderberg

src/relax/ir/expr.cc

@@ -36,6 +36,12 @@ Id::Id(String name_hint) {
 }
 
 Call::Call(Expr op, Array<Expr> args, Attrs attrs, Array<StructInfo> sinfo_args, Span span) {
+  CHECK(!op->struct_info_.defined() || op->struct_info_->IsInstance<FuncStructInfoNode>())
+      << "ValueError: "
+      << "Call expects its operator to have FuncStructInfo, "
+      << "but operator " << op << ", which was called with arguments " << args
+      << ", has struct info " << op->struct_info_;
+
   ObjectPtr<CallNode> n = make_object<CallNode>();
   n->op = std::move(op);
   n->args = std::move(args);

tests/python/relax/test_expr.py

@@ -271,5 +271,25 @@ def test_datatype_imm():
     _check_json_roundtrip(d0)
 
 
+def test_call():
+    dtype = rx.PrimStructInfo("int32")
+    func = rx.Var("func", rx.FuncStructInfo([dtype], dtype))
+    arg = rx.Var("arg", dtype)
+    call = rx.Call(func, [arg])
+    assert call.op.same_as(func)
+    assert len(call.args) == 1
+    assert call.args[0].same_as(arg)
+
+
+def test_call_raises_error_for_invalid_function():
+    """relax::Call requires the function to have FuncStructInfo"""
+    dtype = rx.PrimStructInfo("int32")
+    func = rx.Var("func", dtype)
+    arg = rx.Var("arg", dtype)
+
+    with pytest.raises(ValueError):
+        rx.Call(func, [arg])
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/python/relax/test_op_misc.py

@@ -66,6 +66,16 @@ def test_implicit_op():
     m, n = tvm.tir.Var("m", "int64"), tvm.tir.Var("n", "int64")
     x = rx.Var("x", R.Tensor([m, n], "float32"))
     y = rx.Var("y", R.Tensor([m, n], "float32"))
+    func = rx.Var(
+        "func",
+        R.Callable(
+            [R.Tensor([m, n], "float32")],
+            R.Callable(
+                [R.Tensor([m, n], "float32")],
+                R.Tuple,
+            ),
+        ),
+    )
 
     def _check_call(expr, op_name: str):
         assert isinstance(expr, rx.Call)
@@ -94,9 +104,9 @@ def _check_call(expr, op_name: str):
     _check_call(x.astype("float32"), "astype")
 
     # Call
-    call_expr = x(y)(y)
+    call_expr = func(y)(y)
     assert isinstance(call_expr.op, rx.Call)
-    assert call_expr.op.op == x
+    assert call_expr.op.op == func
 
     # GetTupleItem
     ## Eager get item for tuple