[MetaSchedule] TensorRT BYOC (#518)

* Meta-schedule w/ TRT works for the network-scale

Extend meta-schedule builder to allow graph-level execution

Code review for BYOC-TensorRT in meta schedule

Reflect comments on PR#518

* update

Co-authored-by: Junru Shao <junrushao1994@gmail.com>

include/tvm/meta_schedule/builder.h

@@ -32,10 +32,13 @@ class BuilderInputNode : public runtime::Object {
   IRModule mod;
   /*! \brief The target to be built for. */
   Target target;
+  /*! \brief Parameters for Relay build module. */
+  Optional<Map<String, runtime::NDArray>> params;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("mod", &mod);
     v->Visit("target", &target);
+    v->Visit("params", &params);
   }
 
   static constexpr const char* _type_key = "meta_schedule.BuilderInput";
@@ -52,8 +55,9 @@ class BuilderInput : public runtime::ObjectRef {
    * \brief Constructor of BuilderInput.
    * \param mod The IRModule to be built.
    * \param target The target to be built for.
+   * \param params Parameters for Relay build module.
    */
-  TVM_DLL explicit BuilderInput(IRModule mod, Target target);
+  TVM_DLL explicit BuilderInput(IRModule mod, Target target, Optional<Map<String, runtime::NDArray>> params = NullOpt);
   TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderInput, runtime::ObjectRef, BuilderInputNode);
 };
 

python/tvm/meta_schedule/builder/builder.py

@@ -15,8 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 """Meta Schedule builders that translate IRModule to runtime.Module, and then export"""
-from typing import List, Optional
+from typing import List, Optional, Dict
 
+from tvm.runtime import NDArray
 from tvm._ffi import register_object
 from tvm.ir import IRModule
 from tvm.runtime import Object
@@ -36,12 +37,15 @@ class BuilderInput(Object):
         The IRModule to be built.
     target : Target
         The target to be built for.
+    params: Optional[Dict[str, NDArray]]
+        The parameters for Relay build module
     """
 
     mod: IRModule
     target: Target
+    params: Optional[Dict[str, NDArray]]
 
-    def __init__(self, mod: IRModule, target: Target) -> None:
+    def __init__(self, mod: IRModule, target: Target, params: Optional[Dict[str, NDArray]]) -> None:
         """Constructor.
 
         Parameters
@@ -50,11 +54,14 @@ def __init__(self, mod: IRModule, target: Target) -> None:
             The IRModule to be built.
         target : Target
             The target to be built for.
+        params: Optional[Dict[str, NDArray]]
+            The parameters for Relay build module
         """
         self.__init_handle_by_constructor__(
             _ffi_api.BuilderInput,  # type: ignore # pylint: disable=no-member
             mod,
             target,
+            params,
         )
 
 

python/tvm/meta_schedule/builder/local_builder.py

@@ -18,11 +18,13 @@
 import logging
 import os
 import tempfile
-from typing import Callable, List, Optional, Union
+from typing import Callable, List, Optional, Union, Dict
+from numpy import byte
+from tvm.runtime import NDArray
 
 from tvm._ffi import register_func
 from tvm.ir import IRModule
-from tvm.runtime import Module
+from tvm.runtime import Module, save_param_dict, load_param_dict
 from tvm.target import Target
 
 from ...contrib.popen_pool import MapResult, PopenPoolExecutor, StatusKind
@@ -32,6 +34,18 @@
 logger = logging.getLogger(__name__)
 
 
+def _serialize_params(params: Optional[Dict[str, NDArray]]) -> Optional[bytearray]:
+    if params is None:
+        return None
+    return save_param_dict(params)
+
+
+def _deserialize_params(params: Optional[bytearray]) -> Optional[Dict[str, NDArray]]:
+    if params is None:
+        return None
+    return load_param_dict(params)
+
+
 class LocalBuilder(PyBuilder):
     """A builder that builds the given input on local host.
 
@@ -138,6 +152,7 @@ def build(self, build_inputs: List[BuilderInput]) -> List[BuilderResult]:
                     self.f_export,
                     build_input.mod,
                     build_input.target,
+                    _serialize_params(build_input.params),
                 )
                 for build_input in build_inputs
             ],
@@ -176,6 +191,7 @@ def _worker_func(
         _f_export: Union[None, str, T_EXPORT],
         mod: IRModule,
         target: Target,
+        params: Optional[bytearray],
     ) -> str:
         # Step 0. Get the registered functions
         f_build: LocalBuilder.T_BUILD = get_global_func_with_default_on_worker(
@@ -187,7 +203,7 @@ def _worker_func(
             default_export,
         )
         # Step 1. Build the IRModule
-        rt_mod: Module = f_build(mod, target)
+        rt_mod: Module = f_build(mod, target, _deserialize_params(params))
         # Step 2. Export the Module
         artifact_path: str = f_export(rt_mod)
         return artifact_path

python/tvm/meta_schedule/testing/byoc_trt.py

@@ -0,0 +1,139 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""TensorRT-MetaSchedule integration"""
+# pylint: disable=import-outside-toplevel
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from tvm.ir import IRModule
+    from tvm.target import Target
+    from tvm.runtime import NDArray, Module, Device
+    from tvm.meta_schedule.runner import EvaluatorConfig
+
+
+def build_relay(
+    mod: "IRModule",
+    target: "Target",
+    params: Dict[str, "NDArray"],
+) -> "Module":
+    """Build a Relay IRModule
+
+    Parameters
+    ----------
+    mod : IRModule
+        The Relay IRModule to build.
+    target : Target
+        The target to build the module for.
+    params : Dict[str, NDArray]
+        The parameter dict to build the module with.
+
+    Returns
+    -------
+    mod : runtime.Module
+        The built module.
+    """
+    from tvm.relay.build_module import _build_module_no_factory as relay_build
+    from tvm.runtime import Module
+
+    result = relay_build(mod, target=target, target_host=None, params=params)
+    assert isinstance(result, Module)
+    return result
+
+
+def build_relay_with_tensorrt(
+    mod: "IRModule",
+    target: "Target",
+    params: Dict[str, "NDArray"],
+) -> "Module":
+    """Build a Relay IRModule with TensorRT BYOC
+
+    Parameters
+    ----------
+    mod : IRModule
+        The Relay IRModule to build.
+
+    target : Target
+        The target to build the module for.
+
+    params : Dict[str, NDArray]
+        The parameter dict to build the module with.
+
+    Returns
+    -------
+    mod : runtime.Module
+        The built module.
+    """
+    from tvm.ir.transform import PassContext
+    from tvm.relay.op.contrib import tensorrt
+    from tvm.relay.build_module import _build_module_no_factory as relay_build
+    from tvm.runtime import Module
+
+    mod, config = tensorrt.partition_for_tensorrt(mod, params)
+    with PassContext(
+        opt_level=3,
+        config={"relay.ext.tensorrt.options": config},
+    ):
+        result = relay_build(mod, target=target, target_host=None, params=params)
+    assert isinstance(result, Module)
+    return result
+
+
+def run_with_graph_executor(
+    rt_mod: "Module",
+    device: "Device",
+    evaluator_config: "EvaluatorConfig",
+    repeated_args: List["NDArray"],
+) -> List[float]:
+    """Run a Relay module with GraphExecutor
+
+    Parameters
+    ----------
+    rt_mod : Module
+        The Relay module to run.
+    device : Device
+        The device to run the module on.
+    evaluator_config : EvaluatorConfig
+        The evaluator configuration to run the module with.
+    repeated_args : List[NDArray]
+        The list of repeated arguments to run the module with.
+
+    Returns
+    -------
+    results : List[float]
+        The list of results.
+    """
+    import itertools
+    from tvm.contrib.graph_executor import GraphModule
+
+    graph_mod = GraphModule(rt_mod["default"](device))
+    evaluator = graph_mod.module.time_evaluator(
+        func_name="run",
+        dev=device,
+        number=evaluator_config.number,
+        repeat=evaluator_config.repeat,
+        min_repeat_ms=evaluator_config.min_repeat_ms,
+        f_preproc="cache_flush_cpu_non_first_arg"
+        if evaluator_config.enable_cpu_cache_flush
+        else "",
+    )
+    repeated_costs = []
+    for args in repeated_args:
+        profile_result = evaluator(*args)
+        repeated_costs.append(profile_result.results)
+    costs = [float(cost) for cost in itertools.chain.from_iterable(repeated_costs)]
+    return costs

python/tvm/relay/build_module.py

@@ -271,13 +271,17 @@ def _module_export(module, file_name):  # fcompile, addons, kwargs?
 
 
 @register_func("tvm.relay.build")
+def _build_module_no_factory_impl(mod, target, target_host, params, mod_name):
+    target, target_host = Target.check_and_update_host_consist(target, target_host)
+    return build(mod, target, params=params, mod_name=mod_name).module
+
+
 def _build_module_no_factory(mod, target=None, target_host=None, params=None, mod_name="default"):
     """A wrapper around build which discards the Python GraphFactoryRuntime.
     This wrapper is suitable to be used from other programming languages as
     the runtime::Module can be freely passed between language boundaries.
     """
-    target, target_host = Target.check_and_update_host_consist(target, target_host)
-    return build(mod, target, params=params, mod_name=mod_name).module
+    return _build_module_no_factory_impl(mod, target, target_host, params, mod_name)
 
 
 def _reconstruct_from_deprecated_options(deprecated_params_target):

src/meta_schedule/builder/builder.cc

@@ -23,10 +23,11 @@ namespace meta_schedule {
 
 /******** Constructors ********/
 
-BuilderInput::BuilderInput(IRModule mod, Target target) {
+BuilderInput::BuilderInput(IRModule mod, Target target, Optional<Map<String, runtime::NDArray>> params) {
   ObjectPtr<BuilderInputNode> n = make_object<BuilderInputNode>();
   n->mod = std::move(mod);
   n->target = std::move(target);
+  n->params = std::move(params);
   data_ = std::move(n);
 }
 
@@ -51,8 +52,8 @@ TVM_REGISTER_OBJECT_TYPE(BuilderNode);
 TVM_REGISTER_NODE_TYPE(PyBuilderNode);
 
 TVM_REGISTER_GLOBAL("meta_schedule.BuilderInput")
-    .set_body_typed([](IRModule mod, Target target) -> BuilderInput {
-      return BuilderInput(mod, target);
+    .set_body_typed([](IRModule mod, Target target, Optional<Map<String, runtime::NDArray>> params) -> BuilderInput {
+      return BuilderInput(mod, target, params);
     });
 
 TVM_REGISTER_GLOBAL("meta_schedule.BuilderResult")