electriclilies commented on a change in pull request #10352:
URL: https://github.com/apache/tvm/pull/10352#discussion_r814402151
##########
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:
yup, good catch! thanks :)
--
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]
