tvm_tirx.tirx.transform.FlattenTensor

Contents

tvm_tirx.tirx.transform.FlattenTensor#

tvm_tirx.tirx.transform.FlattenTensor()#

Flatten tensor allocations and accesses to one-dimensional storage.

Uses shapes, strides and offsets to compute linear indices. Tensor parameters retain their interface types and gain flat views as needed.

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((2, 4), T.float32, scope="local")])
        A[1, 2] = T.float32(1.0)

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

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

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