Tensor Instruction Lowering

Tensor Instruction Lowering#

Backend tensor operations are ordinary tvm.ir.Call nodes with opaque side effects and void return types. CUDA and Trainium register their contracts in python/tvm/backend/{cuda,trn}/tile/instructions.py using the shared python/tvm/tirx/tile/instruction.py machinery. C++ reflected attribute types and the native Call validator are in src/tirx/tile/op.cc.

The public C++ tile interfaces are grouped under include/tvm/tirx/tile/:

  • op.h declares the instruction registration and validation contract.

  • dispatch.h declares the dispatcher entry point.

  • exec_scope.h defines backend execution hierarchy nodes used by layouts and tile calls. The shared tirx::ExecScope type remains unchanged.

  • context.h defines the backend-neutral tile::DispatchContext and its allocation and initialization callbacks.

Their implementations live in src/tirx/tile/. CUDA-specific scope bindings, active-thread execution contexts, and predicate analysis remain private to src/backend/cuda/tile/. Generic tensor layout types remain in include/tvm/tirx/layout.h.

tirx.TileDispatch remains the first phase of LowerTIRx. The pass recognizes Evaluate(Call) by the operator’s TIRxOpCategory: tile_primitive or tile_composite. It validates the call, resolves the scope against the active thread set, constructs a DispatchContext, and invokes the instruction’s registered lowerer. It replaces the Evaluate with the returned function body and lowers any nested tensor calls.

This preserves launch parameters, inter/intra scope maps, value ranges, descriptor caches, host initialization, and allocation/initialization callbacks. Lowering occurs at a statement boundary because one instruction may require address preparation, layout expansion, and multiple occurrences of the same core instruction. A residual-call verifier rejects either category after expansion.

There is one lowerer per instruction identity. Unsupported layouts and scopes produce an error that includes the instruction, target, scope, and original cause. There is no cross-instruction priority search or forced variant bag.

The transient Python TensorCall view supplies semantic operand names to existing backend emission code. It is not a registered IR node and is never serialized. Its decoded options are private lowering data. Every expression in the original IR resides in Call.args; shared traversal does not inspect attributes for hidden expressions.

The TVMScript printer validates tensor Calls and prints canonical backend names, tensor regions, nested tuple operands, and static qualifiers. Parsing reconstructs the fixed signature. Private Trainium workspace allocation rewrites Call argument slots before instruction lowering.

See Tensor Instructions for instruction contracts and caller-owned algorithms, and Tensor Instruction Registration for registration.