tvm_tirx.backend.cuda.transform.SpecializeEntryHelpers

tvm_tirx.backend.cuda.transform.SpecializeEntryHelpers()

Specialize helpers for each CUDA device entry’s resolved settings.

Run after BindBackendConfig and before TileDispatch and SplitHostDevice. Clone helpers when entry settings differ, propagating them through the call graph; remove originals that become unreachable.

Examples

import tvm_tirx as tvm
from tvm_tirx.script import tirx as T
from tvm_tirx.backend.cuda.transform import SpecializeEntryHelpers

def settings(arch):
    return {"target": tvm.target.Target({"kind": "cuda", "arch": arch}),
            "backend_config": '{"cuda":{"arch":"' + arch + '"}}'}

host_attrs, entry_attrs = settings("sm_100a"), settings("sm_90a")

@T.ir_module
class Before:
    @T.function(private=True)
    def helper(x: T.int32) -> T.int32:
        T.func_attr(host_attrs)
        return x + 1

    @T.function
    def main(A: T.Tensor((1,), "int32")):
        T.func_attr(host_attrs)
        with T.device_entry(attrs=entry_attrs):
            A[0] = Before.helper(A[0])

@T.ir_module
class Expected:
    @T.function(private=True)
    def helper_cuda(x: T.int32) -> T.int32:
        T.func_attr(entry_attrs)
        return x + 1

    @T.function
    def main(A: T.Tensor((1,), "int32")):
        T.func_attr(host_attrs)
        with T.device_entry(attrs=entry_attrs):
            A[0] = Expected.helper_cuda(A[0])

after = SpecializeEntryHelpers()(Before)
tvm.ir.assert_structural_equal(after, Expected)