vm external codegen (#4544)

src/relay/backend/vm/compiler.cc

@@ -476,24 +476,32 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
       argument_registers.push_back(reg->second);
     }
 
-    // Next generate the invoke instruction.
     Target target;
+
+    if (!func->UseDefaultCompiler()) {
+      target = tvm::target::ext_dev();
+    } else {
+      // Next generate the invoke instruction.
       if (targets_.size() == 1) {
         // homogeneous execution.
-      for (auto kv : targets_) {
-        target = kv.second;
-      }
+        const auto& it = targets_.begin();
+        target = (*it).second;
       } else {
         // heterogeneous execution.
         LOG(FATAL) << "Currently VM compiler doesn't support heterogeneous compilation";
       }
+    }
 
     auto key = CCacheKeyNode::make(func, target);
     auto cfunc = engine_->Lower(key);
 
+    auto op_index = -1;
+    if (!func->UseDefaultCompiler()) {
+      op_index = context_->cached_funcs.size();
+      context_->cached_funcs.push_back(cfunc);
+    } else {
       // TODO(jroesch): support lowered funcs for multiple targets
       CHECK_EQ(cfunc->funcs.size(), 1);
-    auto op_index = -1;
       if (context_->seen_funcs.find(cfunc->funcs[0]) == context_->seen_funcs.end()) {
         op_index = context_->cached_funcs.size();
         context_->cached_funcs.push_back(cfunc);
@@ -501,6 +509,7 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
       } else {
         op_index = context_->seen_funcs[cfunc->funcs[0]];
       }
+    }
 
     Emit(Instruction::InvokePacked(op_index,
       argument_registers.size(),
@@ -950,33 +959,47 @@ void VMCompiler::LibraryCodegen() {
   if (cached_funcs.size() == 0) {
     return;
   }
-  std::unordered_map<std::string, Array<LoweredFunc>> tgt_funcs;
-  for (auto &cfunc : cached_funcs) {
+  std::unordered_map<std::string, Array<LoweredFunc>> funcs;
+  for (auto& cfunc : cached_funcs) {
     std::string target_str = cfunc->target->str();
-    if (tgt_funcs.count(target_str) == 0) {
-      tgt_funcs.emplace(target_str, Array<LoweredFunc>{cfunc->funcs[0]});
+    if (target_str == "ext_dev") {
+      continue;
+    } else if (funcs.count(target_str) == 0) {
+      funcs.emplace(target_str, Array<LoweredFunc>{cfunc->funcs[0]});
     } else {
-      tgt_funcs[target_str].push_back(cfunc->funcs[0]);
-    }
+      funcs[target_str].push_back(cfunc->funcs[0]);
     }
-  Map<Target, Array<LoweredFunc>> funcs;
-  for (auto &it : tgt_funcs) {
-    funcs.Set(Target::Create(it.first), it.second);
   }
 
-  if (const auto *f = runtime::Registry::Get("relay.backend.build")) {
-    // The target is just a dummy arg because funcs already contains corresponding target
-    // therefore target won't be used in the build function
-    runtime::Module mod = (*f)(funcs, Target(), target_host_);
+  auto compile_engine = CompileEngine::Global();
+  auto ext_mods = compile_engine->LowerExternalFunctions();
+  runtime::Module mod;
+  if (funcs.size() > 0) {
+    mod = tvm::build(funcs, target_host_, tvm::BuildConfig::Current());
     CHECK(mod.operator->());
-    exec_->lib = mod;
   } else {
-    LOG(FATAL) << "relay.backend.build is not registered";
+    CHECK_EQ(ext_mods.size(), 1U)
+        << "Expect to have a TVM DSOModule when multiple runtime modules exist";
+  }
+  if (!ext_mods.empty()) {
+    if (funcs.size() == 0) {
+      mod = ext_mods[0];
+    } else {
+      // Import all external runtime modules.
+      for (auto it : ext_mods) {
+        mod.Import(it);
       }
+    }
+  }
+  exec_->lib = mod;
   size_t primitive_index = 0;
   for (auto cfunc : cached_funcs) {
+    if (cfunc->target->str() == "ext_dev") {
+      exec_->primitive_map.insert({cfunc->func_name, primitive_index++});
+    } else {
       exec_->primitive_map.insert({cfunc->funcs[0]->name, primitive_index++});
     }
+  }
 }
 
 runtime::Module CreateVMCompiler() {

src/runtime/vm/vm.cc

@@ -800,7 +800,9 @@ void VirtualMachine::LoadExecutable(const Executable* exec) {
     if (packed_funcs_.size() <= packed_index) {
       packed_funcs_.resize(packed_index + 1);
     }
-    packed_funcs_[packed_index] = lib.GetFunction(packed_name);
+    tvm::runtime::PackedFunc pf = lib.GetFunction(packed_name, true);
+    CHECK(pf != nullptr) << "Cannot find function in module: " << packed_name;
+    packed_funcs_[packed_index] = pf;
   }
 }
 

tests/python/relay/test_external_codegen.py

@@ -26,13 +26,13 @@ import tvm.relay.transform
 from tvm import relay
 from tvm.contrib import util
 
-def check_result(mod, map_inputs, out_shape, result, tol=1e-5):
+def check_result(mod, map_inputs, out_shape, result, tol=1e-5, target="llvm",
+                 ctx=tvm.cpu()):
     if sys.platform == "win32":
         print("Skip test on Windows for now")
         return
 
-    with relay.build_config(opt_level=3, disabled_pass=["AlterOpLayout"]):
-        json, lib, _ = relay.build(mod, "llvm")
+    def update_lib(lib):
         test_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
         source_dir = os.path.join(test_dir, "..", "..", "..")
         contrib_path = os.path.join(source_dir, "src", "runtime", "contrib")
@@ -45,18 +45,36 @@ def check_result(mod, map_inputs, out_shape, result, tol=1e-5):
         lib.export_library(lib_path, fcompile=False, **kwargs)
         lib = tvm.module.load(lib_path)
 
-    ctx = tvm.cpu()
+        return lib
+
+    def check_vm_result():
+        with relay.build_config(opt_level=3, disabled_pass=["AlterOpLayout"]):
+            exe = relay.vm.compile(mod, target=target)
+        code, lib = exe.save()
+        lib = update_lib(lib)
+        exe = relay.vm.Executable.load_exec(code, lib)
+        vm = relay.vm.VirtualMachine(exe)
+        vm.init(ctx)
+        out = vm.run(**map_inputs)
+        tvm.testing.assert_allclose(out.asnumpy(), result, rtol=tol, atol=tol)
+
+    def check_graph_runtime_result():
+        with relay.build_config(opt_level=3, disabled_pass=["AlterOpLayout"]):
+            json, lib, _ = relay.build(mod, target=target)
+        lib = update_lib(lib)
         rt_mod = tvm.contrib.graph_runtime.create(json, lib, ctx)
 
         for name, data in map_inputs.items():
             rt_mod.set_input(name, data)
-
         rt_mod.run()
         out = tvm.nd.empty(out_shape, ctx=ctx)
         out = rt_mod.get_output(0, out)
 
         tvm.testing.assert_allclose(out.asnumpy(), result, rtol=tol, atol=tol)
 
+    check_vm_result()
+    check_graph_runtime_result()
+
 
 def set_external_func_attr(func, compiler, ext_symbol):
     func = func.set_attribute("Primitive", tvm.expr.IntImm("int32", 1))