TensorIR v0.1.0.

TensorIR-to-CUDA-Tile Release Notes (v0.1.0)

The Tensor-to-CUDA-Tile pipeline works best on statically shaped, memory-bound programs that the compiler can fuse into one kernel. A TensorIR operation having a CUDA Tile lowering does not guarantee that every graph containing that operation can be lowered.

Programs that fit the current pipeline

A good input is a single nv_tensor_ir.graph with static tensor shapes, static strides, and a mostly uniform iteration space. Chains of pointwise operations are the best fit. Broadcasts and straightforward view-like transformations, such as transposes or slices, can also fit when layout propagation can preserve one iteration space.

The pipeline also lowers reductions and contractions, but these operations need explicit scheduling information as described below.

The pipeline determines support from the whole graph. The tiling analysis may reject an otherwise supported operation when it cannot schedule the surrounding graph.

One graph, one kernel

TensorIR attempts to produce one kernel from each nv_tensor_ir.graph. The tiling analysis must therefore find a schedule that covers the graph without materializing an intermediate tensor in global memory.

Operations that change or combine iteration spaces can prevent this.

Concatenations

TensorIR currently relies on lowering memory accesses to tile-based loads, which limits the concatenation patterns it can compile. Pointer-based loads and broader support for concatenations are in progress.

Reshapes

A reshape changes the relation between tensor dimensions and tile dimensions. Reshapes inside a graph can block tiling analysis, especially when operations before and after the reshape use different iteration spaces.

For now, avoid concatenations and reshapes in the interior of a graph. Place them at graph boundaries when possible, or split the program into separate graphs if producing more than one kernel is acceptable.

Dynamic shapes

TensorIR has partial support for dynamic shapes. We do not currently recommend that users explore dynamic-shape programs.

Reductions and contractions

Early milestones are focused on memory-bound kernels. The pipeline does not yet have heuristics for choosing tile sizes for reduction or contraction iterators. The user must select the reduction tile size and pass it through reduction-tile-size. The value must be a positive power of two.

For example:

tensor_ir-opt \
  "-layout-propagation-pipeline=tile-size=4 tile-size=1024 reduction-tile-size=128" \
  input.mlir

The tile-size values select the tile for the graph's non-reduction iteration space. reduction-tile-size selects the tile used along reduction or contracting dimensions. The value 128 above is an example, not a general performance recommendation. Choose and benchmark it for the operation shapes and target GPU.

Reduction-heavy and contraction-heavy programs may compile, but the compiler does not currently tune them. Expect to provide tile sizes and measure the generated kernels.