tvm_tirx.tirx.transform.SplitHostDevice

tvm_tirx.tirx.transform.SplitHostDevice()

Annotate, split, and lower host/device functions.

This pass first annotates device regions within host functions, then splits them into host and device-side Functions, and finally lowers host-to-device calls into the device kernel launch ABI.

Returns:

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

Examples

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(private=True)
    def main(n: T.int32):
        T.func_attr({"target": T.target("cuda", host="llvm")})
        with T.region("tirx.device_scope", []):
            T.evaluate(n)

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(n: T.int32):
        T.func_attr(
            {
                "target": T.target(
                    {
                        "host": {
                            "keys": ["cpu"],
                            "kind": "llvm",
                            "mtriple": "",
                            "tag": "",
                        },
                        "keys": ["cuda", "gpu"],
                        "kind": "cuda",
                        "max_num_threads": 1024,
                        "tag": "",
                        "thread_warp_size": 32,
                    }
                )
            }
        )
        T.call_ffi_kernel("main_kernel", n)

    @T.function
    def main_kernel(n: T.int32):
        T.func_attr(
            {
                "calling_conv": 2,
                "target": T.target("cuda"),
                "tirx.is_global_func": True,
                "tirx.kernel_launch_params": [],
                "tirx.noalias": True,
            }
        )
        T.evaluate(n)

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

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