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.build_module import _build_module_no_factory as relay_build from tvm.relay.op.contrib import tensorrt 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