[Relay][VM] Fix constant folding issue in VM compiler

python/tvm/relay/backend/profiler_vm.py

@@ -14,33 +14,53 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable, invalid-name
+# pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable, invalid-name, redefined-builtin
 """
 The Relay Virtual Machine profiler.
 
 Provides extra APIs for profiling vm execution.
 """
-import tvm
 from . import vm, _vm
 
-def _update_target(target):
-    target = target if target else tvm.target.current_target()
-    if target is None:
-        raise ValueError("Target is not set in env or passed as argument.")
+def compile(mod, target=None, target_host=None, params=None):
+    """
+    Parameters
+    ----------
+    mod : relay.Module
+        The Relay module to build.
 
-    tgts = {}
-    if isinstance(target, (str, tvm.target.Target)):
-        dev_type = tvm.expr.IntImm("int32", tvm.nd.context(str(target)).device_type)
-        tgts[dev_type] = tvm.target.create(target)
-    elif isinstance(target, dict):
-        for dev, tgt in target.items():
-            dev_type = tvm.expr.IntImm("int32", tvm.nd.context(dev).device_type)
-            tgts[dev_type] = tvm.target.create(tgt)
-    else:
-        raise TypeError("target is expected to be str, tvm.target.Target, " +
-                        "or dict of str to str/tvm.target.Target, but received " +
-                        "{}".format(type(target)))
-    return tgts
+    target : str, :any:`tvm.target.Target`, or dict of str(i.e.
+        device/context name) to str/tvm.target.Target, optional
+        For heterogeneous compilation, it is a dictionary indicating context
+        to target mapping. For homogeneous compilation, it is a build target.
+
+    target_host : str or :any:`tvm.target.Target`, optional
+        Host compilation target, if target is device.
+        When TVM compiles device specific program such as CUDA,
+        we also need host(CPU) side code to interact with the driver
+        to setup the dimensions and parameters correctly.
+        target_host is used to specify the host side codegen target.
+        By default, llvm is used if it is enabled,
+        otherwise a stackvm intepreter is used.
+
+    params : dict of str to NDArray
+        Input parameters to the graph that do not change
+        during inference time. Used for constant folding.
+
+    Returns
+    -------
+    vm : VirtualMachineProfiler
+        The profile VM runtime.
+    """
+    compiler = VMCompilerProfiler()
+    target = compiler.update_target(target)
+    target_host = compiler.update_target_host(target, target_host)
+    if params:
+        compiler.set_params(params)
+    tophub_context = compiler.tophub_context(target)
+    with tophub_context:
+        compiler._compile(mod, target, target_host)
+    return VirtualMachineProfiler(compiler._get_vm())
 
 class VMCompilerProfiler(vm.VMCompiler):
     """Build Relay module to run on VM runtime."""
@@ -49,36 +69,7 @@ def __init__(self):
         self.mod = _vm._VMCompilerProfiler()
         self._compile = self.mod["compile"]
         self._get_vm = self.mod["get_vm"]
-
-    def compile(self, mod, target=None, target_host=None):
-        """
-        Parameters
-        ----------
-        mod : relay.Module
-            The Relay module to build.
-
-        target : str, :any:`tvm.target.Target`, or dict of str(i.e.
-            device/context name) to str/tvm.target.Target, optional
-            For heterogeneous compilation, it is a dictionary indicating context
-            to target mapping. For homogeneous compilation, it is a build target.
-
-        target_host : str or :any:`tvm.target.Target`, optional
-            Host compilation target, if target is device.
-            When TVM compiles device specific program such as CUDA,
-            we also need host(CPU) side code to interact with the driver
-            to setup the dimensions and parameters correctly.
-            target_host is used to specify the host side codegen target.
-            By default, llvm is used if it is enabled,
-            otherwise a stackvm intepreter is used.
-
-        Returns
-        -------
-        vm : VirtualMachineProfiler
-            The profile VM runtime.
-        """
-        target = _update_target(target)
-        self._compile(mod, target, target_host)
-        return VirtualMachineProfiler(self._get_vm())
+        self._set_params_func = self.mod["set_params"]
 
 class VirtualMachineProfiler(vm.VirtualMachine):
     """Relay profile VM runtime."""

python/tvm/relay/backend/vm.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable, invalid-name
+# pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable, invalid-name, redefined-builtin
 """
 The Relay Virtual Machine.
 
@@ -25,30 +25,11 @@
 import tvm
 from tvm import autotvm
 from tvm._ffi.runtime_ctypes import TVMByteArray
+from tvm.relay import expr as _expr
 from . import _vm
 from . import vmobj as _obj
 from .interpreter import Executor
 
-
-def _update_target(target):
-    target = target if target else tvm.target.current_target()
-    if target is None:
-        raise ValueError("Target is not set in env or passed as argument.")
-
-    tgts = {}
-    if isinstance(target, (str, tvm.target.Target)):
-        dev_type = tvm.expr.IntImm("int32", tvm.nd.context(str(target)).device_type)
-        tgts[dev_type] = tvm.target.create(target)
-    elif isinstance(target, dict):
-        for dev, tgt in target.items():
-            dev_type = tvm.expr.IntImm("int32", tvm.nd.context(dev).device_type)
-            tgts[dev_type] = tvm.target.create(tgt)
-    else:
-        raise TypeError("target is expected to be str, tvm.target.Target, " +
-                        "or dict of str to str/tvm.target.Target, but received " +
-                        "{}".format(type(target)))
-    return tgts
-
 def _convert(arg, cargs):
     if isinstance(arg, (np.ndarray, tvm.nd.NDArray)):
         cargs.append(_obj.tensor_object(arg))
@@ -144,40 +125,85 @@ def module(self):
         return self.mod
 
 
+def compile(mod, target=None, target_host=None, params=None):
+    """
+    Parameters
+    ----------
+    mod : relay.Module
+        The Relay module to build.
+
+    target : str, :any:`tvm.target.Target`, or dict of str(i.e.
+        device/context name) to str/tvm.target.Target, optional
+        For heterogeneous compilation, it is a dictionary indicating context
+        to target mapping. For homogeneous compilation, it is a build target.
+
+    target_host : str or :any:`tvm.target.Target`, optional
+        Host compilation target, if target is device.
+        When TVM compiles device specific program such as CUDA,
+        we also need host(CPU) side code to interact with the driver
+        to setup the dimensions and parameters correctly.
+        target_host is used to specify the host side codegen target.
+        By default, llvm is used if it is enabled,
+        otherwise a stackvm intepreter is used.
+
+    params : dict of str to NDArray
+        Input parameters to the graph that do not change
+        during inference time. Used for constant folding.
+
+    Returns
+    -------
+    vm : VirtualMachine
+        The VM runtime.
+    """
+    compiler = VMCompiler()
+
+    target = compiler.update_target(target)
+    target_host = compiler.update_target_host(target, target_host)
+    if params:
+        compiler.set_params(params)
+    tophub_context = compiler.tophub_context(target)
+    with tophub_context:
+        compiler._compile(mod, target, target_host)
+    return VirtualMachine(compiler._get_vm())
+
 class VMCompiler(object):
     """Build Relay module to run on VM runtime."""
     def __init__(self):
         self.mod = _vm._VMCompiler()
         self._compile = self.mod["compile"]
         self._get_vm = self.mod["get_vm"]
+        self._set_params_func = self.mod["set_params"]
+
+    def set_params(self, params):
+        """Set constant parameters for the model"""
+        inputs = {}
+        for name, param in params.items():
+            if isinstance(param, np.ndarray):
+                param = _nd.array(param)
+            inputs[name] = _expr.const(param)
+        self._set_params_func(inputs)
+
+    def update_target(self, target):
+        """Update target"""
+        target = target if target else tvm.target.current_target()
+        if target is None:
+            raise ValueError("Target is not set in env or passed as argument.")
+        tgts = {}
+        if isinstance(target, (str, tvm.target.Target)):
+            dev_type = tvm.expr.IntImm("int32", tvm.nd.context(str(target)).device_type)
+            tgts[dev_type] = tvm.target.create(target)
+        elif isinstance(target, dict):
+            for dev, tgt in target.items():
+                dev_type = tvm.expr.IntImm("int32", tvm.nd.context(dev).device_type)
+                tgts[dev_type] = tvm.target.create(tgt)
+        else:
+            raise TypeError("target is expected to be str, tvm.target.Target, " +
+                            "or dict of str to str/tvm.target.Target, but received " +
+                            "{}".format(type(target)))
+        return tgts
 
-    def compile(self, mod, target=None, target_host=None):
-        """
-        Parameters
-        ----------
-        mod : relay.Module
-            The Relay module to build.
-
-        target : str, :any:`tvm.target.Target`, or dict of str(i.e.
-            device/context name) to str/tvm.target.Target, optional
-            For heterogeneous compilation, it is a dictionary indicating context
-            to target mapping. For homogeneous compilation, it is a build target.
-
-        target_host : str or :any:`tvm.target.Target`, optional
-            Host compilation target, if target is device.
-            When TVM compiles device specific program such as CUDA,
-            we also need host(CPU) side code to interact with the driver
-            to setup the dimensions and parameters correctly.
-            target_host is used to specify the host side codegen target.
-            By default, llvm is used if it is enabled,
-            otherwise a stackvm intepreter is used.
-
-        Returns
-        -------
-        vm : VirtualMachine
-            The VM runtime.
-        """
-        target = _update_target(target)
+    def update_target_host(self, target, target_host):
+        """Update target host"""
         target_host = None if target_host == "" else target_host
         if not target_host:
             for device_type, tgt in target.items():
@@ -186,19 +212,16 @@ def compile(self, mod, target=None, target_host=None):
                     break
         if not target_host:
             target_host = "llvm" if tvm.module.enabled("llvm") else "stackvm"
-        target_host = tvm.target.create(target_host)
+        return tvm.target.create(target_host)
 
+    def tophub_context(self, target):
         # If current dispatch context is fallback context (the default root context),
         # then load pre-tuned parameters from TopHub
         if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
             tophub_context = autotvm.tophub.context(list(target.values()))
         else:
             tophub_context = autotvm.util.EmptyContext()
-
-        with tophub_context:
-            self._compile(mod, target, target_host)
-        return VirtualMachine(self._get_vm())
-
+        return tophub_context
 
 class VMExecutor(Executor):
     """
@@ -226,8 +249,7 @@ def __init__(self, mod, ctx, target):
         self.mod = mod
         self.ctx = ctx
         self.target = target
-        compiler = VMCompiler()
-        self.vm = compiler.compile(mod, target)
+        self.vm = compile(mod, target)
         self.vm.init(ctx)
 
     def _make_executor(self, expr=None):

src/relay/backend/vm/compiler.cc

@@ -780,31 +780,81 @@ PackedFunc VMCompiler::GetFunction(const std::string& name,
   if (name == "compile") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
       CHECK_EQ(args.num_args, 3);
-      this->Compile(args[0], args[1], args[2]);
+      Module mod = args[0];
+      this->Compile(mod, args[1], args[2]);
     });
   } else if (name == "get_vm") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
       *rv = runtime::Module(vm_);
     });
+  } else if (name == "set_params") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      Map<std::string, Constant> params = args[0];
+      for (const auto& kv : params) {
+        this->SetParam(kv.first, kv.second->data);
+      }
+    });
   } else {
     LOG(FATAL) << "Unknown packed function: " << name;
     return PackedFunc([sptr_to_self, name](TVMArgs args, TVMRetValue* rv) {});
   }
 }
 
-void VMCompiler::Compile(const Module& mod_ref,
+void VMCompiler::SetParam(const std::string& name, runtime::NDArray data_in) {
+  params_[name] = data_in;
+}
+
+relay::Function VMCompiler::BindParamsByName(
+    relay::Function func,
+    const std::unordered_map<std::string, runtime::NDArray>& params) {
+  std::unordered_map<std::string, relay::Var> name_dict;
+  std::unordered_set<relay::Var, NodeHash, NodeEqual> repeat_var;
+  for (auto arg : func->params) {
+    const auto &name = arg->name_hint();
+    if (name_dict.count(name)) {
+      repeat_var.insert(arg);
+    } else {
+      name_dict[name] = arg;
+    }
+  }
+  std::unordered_map<relay::Var, Expr, NodeHash, NodeEqual> bind_dict;
+  for (auto &kv : params) {
+    if (name_dict.count(kv.first) == 0) {
+      continue;
+    }
+    auto arg = name_dict.at(kv.first);
+    if (repeat_var.count(arg)) {
+      LOG(FATAL) << "Multiple args in the function have name " << kv.first;
+    }
+    bind_dict[arg] = ConstantNode::make(kv.second);
+  }
+  Expr bound_expr = relay::Bind(func, bind_dict);
+  Function ret = Downcast<Function>(bound_expr);
+  CHECK(ret.defined())
+      << "The returning type is expected to be a Relay Function."
+      << "\n";
+  return ret;
+}
+
+
+void VMCompiler::Compile(Module mod,
                          const TargetsMap& targets,
                          const tvm::Target& target_host) {
   CHECK_EQ(targets.size(), 1)
     << "Currently VM compiler doesn't support heterogeneous compilation";
+  if (params_.size()) {
+    auto f = BindParamsByName(mod->Lookup("main"), params_);
+    auto gvar = mod->GetGlobalVar("main");
+    mod->Add(gvar, f);
+  }
 
   InitVM();
   targets_ = targets;
   target_host_ = target_host;
 
   // Run some optimizations first, this code should
   // be moved to pass manager.
-  context_.module = OptimizeModule(mod_ref, targets_);
+  context_.module = OptimizeModule(mod, targets_);
 
   // Populate the global map.
   //