tvm_tirx.tirx.transform.PointerValueTypeRewrite

tvm_tirx.tirx.transform.PointerValueTypeRewrite()

Rewrite pointer and tensor element types to match their accesses.

Adjusts argument and allocation types, shapes and accesses where possible to reduce backend casts, including packing scalar storage into vectors.

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((16,), T.float32, scope="local")])
        A[0:4] = T.Broadcast(T.float32(1.0), 4)

@I.ir_module
class Expected:
    @T.function(private=True)
    def main():
        A = T.alloc_tensor(ty_args=[T.Tensor((4,), T.float32x4, scope="local")])
        A[0] = T.Broadcast(T.float32(1.0), 4)

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

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