Esempio n. 1
0
def default_build(mod: IRModule, target: Target) -> Module:
    """Default build function.

    Parameters
    ----------
    mod : IRModule
        The IRModule to be built.
    target : Target
        The target to be built.

    Returns
    -------
    rt_mod : Module
        The built Module.
    """
    # pylint: disable=import-outside-toplevel
    from tvm.autotvm.measure.measure_methods import set_cuda_target_arch
    from tvm.driver import build as tvm_build

    # pylint: enable=import-outside-toplevel

    if target.kind.name == "cuda":
        set_cuda_target_arch(target.attrs["arch"])

    return tvm_build(mod, target=target)
Esempio n. 2
0
def default_build(mod: IRModule, target: Target,
                  _params: Optional[Dict[str, NDArray]]) -> Module:
    """Default build function.

    Parameters
    ----------
    mod : IRModule
        The IRModule to be built.
    target : Target
        The target to be built.
    _params : Optional[Dict[str, NDArray]]
        The parameters to be used for the build. Must be None.

    Returns
    -------
    rt_mod : Module
        The built Module.
    """
    # pylint: disable=import-outside-toplevel
    from tvm.driver import build as tvm_build
    from tvm.tir.transform import RemoveWeightLayoutRewriteBlock

    # pylint: enable=import-outside-toplevel
    mod = RemoveWeightLayoutRewriteBlock()(mod)
    return tvm_build(mod, target=target)
Esempio n. 3
0
def default_build(mod: IRModule, target: Target,
                  _params: Optional[Dict[str, NDArray]]) -> Module:
    """Default build function.

    Parameters
    ----------
    mod : IRModule
        The IRModule to be built.
    target : Target
        The target to be built.
    _params : Optional[Dict[str, NDArray]]
        The parameters to be used for the build. Must be None.

    Returns
    -------
    rt_mod : Module
        The built Module.
    """
    # pylint: disable=import-outside-toplevel
    from tvm.driver import build as tvm_build
    from tvm.ir.transform import PassContext

    # pylint: enable=import-outside-toplevel
    with PassContext(disabled_pass=["tir.CommonSubexprElimTIR"]):
        return tvm_build(mod, target=target)
Esempio n. 4
0
def default_build(mod: IRModule, target: Target) -> Module:
    """Default build function.

    Parameters
    ----------
    mod : IRModule
        The IRModule to be built.
    target : Target
        The target to be built.

    Returns
    -------
    rt_mod : Module
        The built Module.
    """
    # pylint: disable=import-outside-toplevel
    from tvm.driver import build as tvm_build

    # pylint: enable=import-outside-toplevel

    return tvm_build(mod, target=target)