tvm_tirx.tirx.transform.NarrowDataType

Contents

tvm_tirx.tirx.transform.NarrowDataType#

tvm_tirx.tirx.transform.NarrowDataType(target_bits)#

Narrow integer indices and loop variables when bounds prove it safe.

Keeps wider types when the analyzed range does not fit target_bits.

Parameters:

target_bits (int) – The target bit configuration.

Returns:

fpass (tvm_tirx.ir.transform.Pass) – The result pass

Note

Run this pass after FlattenTensor.

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)):
        for i in range(T.int64(0), T.int64(4)):
            A[i] = T.float32(1.0)

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(A: T.Tensor((4,), T.float32)):
        for i in range(4):
            A[i] = T.float32(1.0)

def run_pass():
    return tr.NarrowDataType(32)(Before)

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