tvm_tirx.tirx.transform.TileDispatch

Contents

tvm_tirx.tirx.transform.TileDispatch#

tvm_tirx.tirx.transform.TileDispatch()#

Lower TIRx tile calls using the active backend dispatch table.

Resolves tile instructions and device launch/execution scopes. Select the backend with a target context; follow with LowerTensorLayouts() or the backend’s own layout-lowering pass.

Returns:

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

Examples

Dispatch a tile-wide move into a loop of scalar PTX moves.

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():
        A = T.alloc_local((4, 1), "float32", layout=T.TileLayout(T.S[(4, 1) : (1, 1)]))
        T.cuda.tile.apply(T.ptx["mov.b32"], A[:, :], T.float32(0), scope="thread")

@I.ir_module
class Expected:
    @T.function(private=True)
    def main():
        A = T.alloc_local((4, 1), "float32", layout=T.TileLayout(T.S[(4, 1) : (1, 1)]))
        tile_d = T.decl_tensor(
            A.data,
            ty_args=[
                T.Tensor(
                    (4,),
                    T.float32,
                    scope="local",
                    layout=T.TileLayout(T.S[(4,) : (1,)]),
                )
            ],
        )
        for repeat_0 in T.unroll(0, 4):
            T.ptx.mov(tile_d[repeat_0], T.float32(0.0), "b32", "")

def run_pass():
    with tvm.target.Target({"kind": "cuda", "arch": "sm_75"}):
        return tr.TileDispatch()(Before)

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