tvm_tirx.tirx.transform.LowerTensorLayouts

tvm_tirx.tirx.transform.LowerTensorLayouts()

Apply tensor layouts to physical storage and accesses.

Run after TileDispatch() for generic tile compilation, or independently on already dispatched IR. Backends may pair TileDispatch with their own layout-lowering pass.

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():
        A = T.alloc_tensor(
            ty_args=[
                T.Tensor(
                    (3, 4),
                    T.uint32,
                    scope="local",
                    layout=T.TileLayout(T.S[(3, 4) : (16, 2)]),
                )
            ]
        )
        A[2, 3] = T.uint32(7)

@I.ir_module
class Expected:
    @T.function(private=True)
    def main():
        A = T.alloc_tensor(ty_args=[T.Tensor((39,), T.uint32, scope="local")])
        A[38] = T.uint32(7)

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

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