tvm_tirx.tirx.transform.VectorizeLoop

Contents

tvm_tirx.tirx.transform.VectorizeLoop#

tvm_tirx.tirx.transform.VectorizeLoop(enable_vectorize=True)#

Lower vectorization loops.

Target-dependent vectorization uses the Function’s target attribute, not an ambient target context or nested target attributes. Target-independent fixed-width loops do not require a target. Code needing different vectorization targets must be separated into functions before this pass.

Unsupported loops retain scalar execution. Loops containing assertions always remain scalar. Accessing multiple tensors requires the function’s tirx.noalias attribute unless they are views of the same storage.

Parameters:

enable_vectorize (bool) – Whether vectorization is enabled. Will lower to scalar loop when it is turned off.

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.float32), B: T.Tensor((4,), T.float32)):
        T.func_attr({"tirx.noalias": True})
        for i in T.vectorized(0, 4):
            B[i] = A[i]

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(A: T.Tensor((4,), T.float32), B: T.Tensor((4,), T.float32)):
        T.func_attr({"tirx.noalias": True})
        B[0:4] = A[T.Ramp(0, 1, 4)]

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

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