Function build

python/tvm/tirx/build.py:155–240

Build a function with a signature, generating code for devices coupled with target information. Parameters ---------- mod : Union[PrimFunc, IRModule] The input to be built. target : Optional[Union[str, Target]] The target for compilation. pipeline : Union[Non

(
    mod: PrimFunc | IRModule,
    target: str | Target | None = None,
    pipeline: None | str | tvm.transform.Pass = "default",
)

def build(
    mod: PrimFunc | IRModule,
    target: str | Target | None = None,
    pipeline: None | str | tvm.transform.Pass = "default",
):
    """Build a function with a signature, generating code for devices
    coupled with target information.

    Parameters
    ----------
    mod : Union[PrimFunc, IRModule]
        The input to be built.
    target : Optional[Union[str, Target]]
        The target for compilation.
    pipeline : Union[None, str, tvm.transform.Pass]
        The pipeline to use for compilation.

    Returns
    -------
    tvm.runtime.Module
        A module combining both host and device code.
    """
    # Convert PrimFunc to IRModule
    if isinstance(mod, PrimFunc):
        mod = tvm.IRModule.from_expr(mod)
    else:
        assert isinstance(mod, tvm.IRModule)

    # Step 0: Determine the target in environment
    # It's used to bind the PrimFunc without target attr to serve as a default target
    target_to_bind = Target.current() if target is None else target
    if target_to_bind is None:
        target_to_bind = "llvm"
    assert target_to_bind is not None
    target_to_bind = Target(target_to_bind)

    # Step 1: Determine the target to search for tirx pipeline
    target = Target.current() if target is None else target
    if target is None:
        for func in mod.functions.values():
            f_target = func.attrs.get("target", None)
            if f_target is not None:
                target = f_target
                break
    if target is not None:
        target = Target(target)

    # Step 2: Determine the host target
    target_host = "llvm" if tvm.runtime.enabled("llvm") else "c"
    if target is not None:
        if target.host is not None:
            target_host = target.host
        elif (
            tvm.device(target.kind.name, 0).dlpack_device_type() == tvm.cpu(0).dlpack_device_type()
        ):
            target_host = target
    target_host = Target(target_host)
    target_to_bind = target_to_bind.with_host(target_host)