slyubomirsky commented on code in PR #14394:
URL: https://github.com/apache/tvm/pull/14394#discussion_r1183173385
##########
src/relax/utils.cc:
##########
@@ -111,6 +113,75 @@ bool IsLeafOrTuple(const Expr& expr) {
expr.as<OpNode>() || expr.as<TupleNode>();
}
+bool IsImpureCall(const Call& call) {
+ if (auto op_ptr = call->op.as<OpNode>()) {
+ auto op = GetRef<Op>(op_ptr);
+ static auto purity_map = Op::GetAttrMap<Bool>("FPurity");
+ ICHECK(purity_map.count(op)) << "Cannot find the registered purity of this
op: " << op->name;
+ return !(purity_map[op]->value);
+ }
+ // the StructInfo must be FuncStructInfo
+ auto func_struct_info = GetStructInfoAs<FuncStructInfoNode>(call->op);
+ return !func_struct_info->purity;
+}
+
+Call WrapCallPure(const Call& call) {
+ static const Op& call_pure_packed_op = Op::Get("relax.call_pure_packed");
+ static const Op& call_pure_dps_packed_op =
Op::Get("relax.call_pure_dps_packed");
+ static const Op& call_dps_packed_op = Op::Get("relax.call_dps_packed");
+ static const Op& invoke_closure_op = Op::Get("relax.invoke_closure");
+ static const Op& invoke_pure_closure_op =
Op::Get("relax.invoke_pure_closure");
+
+ Call ret;
+ if (call->op == call_dps_packed_op) {
+ ret = std::move(Call(call_pure_dps_packed_op, call->args, call->attrs,
call->sinfo_args));
+ } else if (call->op == invoke_closure_op) {
+ ret = std::move(Call(invoke_pure_closure_op, call->args, call->attrs,
call->sinfo_args));
+ } else {
+ Array<Expr> call_args = {call->op};
+ for (auto arg : call->args) {
+ call_args.push_back(arg);
+ }
+ ret = std::move(Call(call_pure_packed_op, call_args, call->attrs,
call->sinfo_args));
+ }
+
+ // transfer over struct info if we can
+ if (call->struct_info_) {
+ UpdateStructInfo(ret, GetStructInfo(call));
+ }
+ return ret;
+}
+
+Call UnwrapCallPure(const Call& call) {
+ static const Op& call_pure_packed_op = Op::Get("relax.call_pure_packed");
Review Comment:
I could split this to handle the different op variants. I kept them together
for convenience, but there's no inherent reason they need to be kept together.
It's convenient in passes to unwrap and rewrap, but I'll see how it will look
if I get rid of these helpers.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
