This is an automated email from the ASF dual-hosted git repository.
lunderberg pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 1af82ad666 [Unity] Validate struct info in relax::Call constructor
(#16311)
1af82ad666 is described below
commit 1af82ad6660314f9f42f8cde2a4696af6ed94ef7
Author: Eric Lunderberg <[email protected]>
AuthorDate: Wed Jan 3 09:50:36 2024 -0600
[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
---
src/relax/ir/expr.cc | 6 ++++++
tests/python/relax/test_expr.py | 20 ++++++++++++++++++++
tests/python/relax/test_op_misc.py | 14 ++++++++++++--
3 files changed, 38 insertions(+), 2 deletions(-)
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index 00ad252ec4..1bc7267af6 100644
--- a/src/relax/ir/expr.cc
+++ b/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);
diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py
index fbd37b307e..af1bc851be 100644
--- a/tests/python/relax/test_expr.py
+++ b/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()
diff --git a/tests/python/relax/test_op_misc.py
b/tests/python/relax/test_op_misc.py
index 2bbc18af6b..f786b89043 100644
--- a/tests/python/relax/test_op_misc.py
+++ b/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 test_implicit_op():
_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
