tvm_tirx.tirx.transform.ForceNarrowIndexToInt32

tvm_tirx.tirx.transform.ForceNarrowIndexToInt32()

Force index expressions to int32 without range proofs.

Use only when the caller can guarantee representability. Unlike NarrowDataType(), this pass may change overflowing computations. Rejects signed integer tensor parameters and ordinary allocations wider than int32. Explicitly typed local scalars retain their storage type and are cast at narrowed index uses.

Returns:

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

Note

This pass should not be used in default cases.

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.ForceNarrowIndexToInt32()(Before)

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