tvm_tirx.tirx.transform.HoistIf

Contents

tvm_tirx.tirx.transform.HoistIf#

tvm_tirx.tirx.transform.HoistIf()#

Hoist loop-invariant If statements outside eligible loops.

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.Tensor((4,), T.int32), n: T.int32):
        for i in range(4):
            if n > 0:
                A[i] = 1

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(A: T.Tensor((4,), T.int32), n: T.int32):
        if 0 < n:
            for i in range(4):
                A[i] = 1

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

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