From 181dbd8e94222b1dad5da4c3f15b8c63facc3582 Mon Sep 17 00:00:00 2001 From: Bing Xu Date: Wed, 8 May 2019 21:06:33 -0700 Subject: [PATCH] [BuildModule] Fix AlterLayout Pass (#3155) --- src/relay/backend/build_module.cc | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/relay/backend/build_module.cc b/src/relay/backend/build_module.cc index b60a048..67ab750 100644 --- a/src/relay/backend/build_module.cc +++ b/src/relay/backend/build_module.cc @@ -504,7 +504,14 @@ class RelayBuildModule : public runtime::ModuleNode { if (cfg.pass_enabled("AlterOpLayout")) { if (targets.size() == 1) { func = CallPackedFunc("relay._ir_pass.infer_type", func, nullptr); - func = CallPackedFunc("relay._ir_pass.AlterOpLayout", func); + auto enter_pf = GetPackedFunc("_EnterTargetScope"); + auto exit_pf = GetPackedFunc("_ExitTargetScope"); + for (const auto& kv : targets) { + auto target = Target::create(kv.second); + (*enter_pf)(target); + func = CallPackedFunc("relay._ir_pass.AlterOpLayout", func); + (*exit_pf)(); + } } else { LOG(WARNING) << "AlterOpLayout pass is not enabled for heterogeneous" << " execution yet."; -- 2.7.4