mbs-octoml commented on a change in pull request #10352:
URL: https://github.com/apache/tvm/pull/10352#discussion_r814394621
##########
File path: src/relay/ir/expr_functor.cc
##########
@@ -472,45 +472,38 @@ class ExprBinder : public MixedModeMutator,
PatternMutator {
const tvm::Map<Var, Expr>& args_map_;
};
+// This function should be called SubstAndBind, since it assumes any variables
introduced
+// in the substitution right hand side should be implicitly bound in the
function.
Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>& args_map) {
if (const FunctionNode* func = expr.as<FunctionNode>()) {
Expr new_body = ExprBinder(args_map).VisitExpr(func->body);
Array<Var> new_params;
- std::vector<VirtualDevice> new_param_virtual_devices;
for (size_t i = 0; i < func->params.size(); ++i) {
if (!args_map.count(func->params[i])) {
new_params.push_back(func->params[i]);
-
new_param_virtual_devices.push_back(GetFunctionParamVirtualDevice(func, i));
}
}
if (new_body.same_as(func->body) && new_params.size() ==
func->params.size()) {
return expr;
}
+
auto ret =
Function(new_params, new_body, func->ret_type, func->type_params,
func->attrs, func->span);
- ret =
- MaybeFunctionOnDevice(ret, new_param_virtual_devices,
GetFunctionResultVirtualDevice(func));
+ ret->virtual_device_ = func->virtual_device();
+
std::unordered_set<Var, ObjectPtrHash, ObjectPtrEqual> set;
for (const auto& v : FreeVars(expr)) {
set.insert(v);
}
for (const auto& v : FreeVars(ret)) {
if (set.count(v) == 0) {
new_params.push_back(v);
- if (!GetFunctionResultVirtualDevice(func)->IsFullyUnconstrained()) {
- // TODO(mbs): The function has been annotated with a device, which
means we are supposed
- // to be preserving device annotations on every transformation.
However there's no
- // such context for the free vars in args_map.
- LOG(WARNING) << "introduced free var '" << PrettyPrint(v)
- << "' into function body but no device is known for it";
- }
-
new_param_virtual_devices.push_back(VirtualDevice::FullyUnconstrained());
}
}
- ret =
Review comment:
you still need to bind the new_params right?
##########
File path: src/relay/ir/expr_functor.cc
##########
@@ -528,6 +521,31 @@ TVM_REGISTER_GLOBAL("relay.ir.Bind").set_body([](TVMArgs
args, TVMRetValue* ret)
}
});
+Expr SubstituteVars(const Expr& expr, const tvm::Map<Var, Expr>& args_map) {
Review comment:
How about 'SubstituteBoundVars' and make it only accept a Function?
##########
File path: src/relay/transforms/device_planner.cc
##########
@@ -877,7 +876,7 @@ class DeviceDefaulter : public ExprVisitor {
};
/* =============== Phase 3 =============== */
-
+// TODO(@electriclilies): rewrite this comment
Review comment:
I think just function attributes bullet needs to be updated, the rest
still stands.
--
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]
