tvm_tirx.tirx.transform

tvm_tirx.tirx.transform#

Transformation passes for lowering and optimizing tensor-level IR.

Pass factories return callable passes: after = PassName()(before), where before is an IRModule. Each pass example includes a complete input module, expected output and invocation.

Pass decorator helpers#

FunctionPass(pass_info)

A pass that transforms each tvm_tirx.tirx.Function in a module.

function_pass([pass_func, opt_level, name])

Create a TIRx function pass from a callback or class.

Passes#

DecorateDeviceScope()

Wrap each function body in a tirx.device_scope region.

Apply(ftransform)

Apply ftransform to each function in the Module.

VectorizeLoop([enable_vectorize])

Lower vectorization loops.

StorageRewrite()

Plan storage by moving allocations and reusing compatible storage.

PointerValueTypeRewrite()

Rewrite pointer and tensor element types to match their accesses.

UnrollLoop()

Unroll loops with constant extents marked T.unroll.

RemoveNoOp()

Remove effect-free statements and provably empty control flow.

HoistIf()

Hoist loop-invariant If statements outside eligible loops.

BF16ComputeLegalize()

Promote unsupported BF16 computations and internal temporaries to float32.

FP8ComputeLegalize([promote_dtype])

Promote unsupported FP8 computations and internal temporaries.

BF16StorageLegalize()

Represent BF16 storage as uint16 while preserving its bit patterns.

FP8StorageLegalize()

Represent FP8 storage as uint8 while preserving its bit patterns.

CommonSubexprElim()

Replace redundant computations by new variables.

RenormalizeSplitPattern()

Renormalize the split pattern from floordiv(floormod()) to floormod(floordiv())

StmtSimplify()

Run statement-level arithmetic simplifications on the TIR Function.

LowerThreadAllreduce()

Lower gpu_thread_allreduce to target-specific communication.

ConvertSSA()

Give every variable definition in an IRModule a distinct identity.

MakePackedAPI()

Expose exported functions through the TVM FFI packed-call ABI.

SplitHostDevice()

Annotate, split, and lower host/device functions.

SkipAssert()

Replace assertion statements with no-ops.

LowerTVMBuiltin()

Lower runtime builtins to explicit stack storage and ABI operations.

LowerIntrin()

Lower intrinsic calls using rules registered for the function target.

NarrowDataType(target_bits)

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

ForceNarrowIndexToInt32()

Force index expressions to int32 without range proofs.

VerifyMemory()

Reject illegal host-side access to device memory.

FlattenTensor()

Flatten tensor allocations and accesses to one-dimensional storage.

BindTarget(target)

Bind targets and classify host/device helper functions.

AnnotateEntryFunc()

Infer and mark an unambiguous TIRx entry function.

Filter(fcond)

Remove TIRx functions for which fcond(function) is false.

TileDispatch()

Lower TIRx tile calls using the active backend dispatch table.

LowerTensorLayouts()

Apply tensor layouts to physical storage and accesses.