tvm_tirx.tirx.transform.MakePackedAPI

Contents

tvm_tirx.tirx.transform.MakePackedAPI#

tvm_tirx.tirx.transform.MakePackedAPI()#

Expose exported functions through the TVM FFI packed-call ABI.

Requires a target attribute with a host target on functions with a global_symbol and the default calling convention. Private functions, functions with other calling conventions, and functions whose target has no host are left unchanged.

Tensor parameter types generate runtime checks for shape, dtype and layout. Symbolic dimensions not defined by other arguments are bound from the incoming DLTensor metadata. For example, a T.Tensor((16, 32), "float32") argument is checked against that shape, while symbolic dimensions are read from the argument’s shape. The resulting function uses the host target and an exported __tvm_ffi_ symbol.

Returns:

fpass (tvm_tirx.ir.transform.Pass) – The result pass

Examples

An exported no-argument function illustrates the packed-call ABI.

import tvm_tirx as tvm
from tvm_tirx.script import ir as I, tirx as T
from tvm_tirx.tirx import transform as tr

@I.ir_module
class Before:
    @T.function
    def main():
        T.func_attr({"target": T.target("llvm", host="llvm")})
        T.evaluate(0)

@I.ir_module
class Expected:
    @T.function
    def main(
        self_handle: T.Ptr(T.void),
        args: T.Ptr(T.void),
        num_args: T.int32,
        result: T.Ptr(T.void),
    ) -> T.int32:
        T.func_attr({"global_symbol": "__tvm_ffi_main"})
        T.func_attr({"calling_conv": 1, "target": T.target("llvm")})
        assert num_args == 0, (
            "TypeError",
            ["Expected ", "0", " arguments", " when calling:\n  `", "main()", "`"],
        )
        with T.compute_scope("main_compute_"):
            pass
        return 0

def run_pass():
    return tr.MakePackedAPI()(Before)

tvm.ir.assert_structural_equal(run_pass(), Expected)
Contents