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