Commit 8d50312f by 雾雨魔理沙 Committed by masahi

[Relay] Fix Fuse (#3035)

* save

* fix

* Update fuse_ops.cc
parent d5f7224d
......@@ -865,9 +865,17 @@ class FuseMutator : private ExprMutator {
}
Expr MakeNewFunction(GraphPartitioner::Group* group, Type ret_type, Expr body) {
// If the function has no call, it is not a primitive function.
struct HasCallVisitor : ExprVisitor {
bool has_call = false;
void VisitExpr_(const CallNode* op) final {
has_call = true;
}
} visitor;
visitor(body);
const GroupInfo& ginfo = ginfo_[group];
auto func = FunctionNode::make(ginfo.params, body, ret_type, {});
func = FunctionSetAttr(func, "Primitive", tvm::Integer(1));
func = FunctionSetAttr(func, "Primitive", tvm::Integer(visitor.has_call));
return CallNode::make(func, ginfo.arguments, Attrs());
}
......
......@@ -51,6 +51,12 @@ def test_tuple_value():
np.testing.assert_allclose(tv[2].asnumpy(), 3)
def test_tuple_getitem():
two = relay.add(relay.const(1), relay.const(1))
func = relay.Function([], relay.TupleGetItem(relay.Tuple([relay.const(1), relay.const(2)]), 0))
check_eval(func, [], 1)
def test_id():
x = relay.var('x', 'float32')
ident = relay.Function([x], x)
......@@ -223,4 +229,6 @@ if __name__ == "__main__":
test_kwargs_params()
test_ref()
test_tensor_value()
test_function_taking_adt_ref_tuple()
test_tuple_value()
test_tuple_getitem()
test_function_taking_adt_ref_tuple()
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment