tvm_tirx.tirx.transform.StorageRewrite

Contents

tvm_tirx.tirx.transform.StorageRewrite#

tvm_tirx.tirx.transform.StorageRewrite()#

Plan storage by moving allocations and reusing compatible storage.

Hoists allocations where valid and shares storage when lifetimes and memory-scope allocation properties permit it.

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(output: T.Tensor((2,), "float32")):
        A = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32")])
        for i in T.serial(4):
            A[i] = T.float32(i)
        output[0] = A[3]
        B = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32")])
        for i in T.serial(4):
            B[i] = T.float32(i + 1)
        output[1] = B[3]

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(output: T.Tensor((2,), T.float32)):
        A = T.alloc_tensor(ty_args=[T.Tensor((4,), T.float32)])
        for i in range(4):
            A[i] = T.Cast("float32", i)
        output[0] = A[3]
        A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((4,), T.float32)])
        for i_1 in range(4):
            A_1[i_1] = T.Cast("float32", i_1 + 1)
        output[1] = A_1[3]

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

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