Transform Ops

transform.air.broadcast_before_unary (transform::BroadcastBeforeUnaryOp)

Move vector.broadcast before element-wise unary operations to enable vector execution

Syntax:

operation ::= `transform.air.broadcast_before_unary` $target prop-dict attr-dict `:` functional-type(operands, results)

This transform identifies patterns where an element-wise unary operation operates on a single-element vector or scalar and its result is immediately broadcast to a larger vector. It rearranges the operations to broadcast first, then apply the unary operation, allowing the operation to execute on the full vector which can be more efficient on vector engines.

Pattern matched (vector<1xT>):

%unary_result = unary_op %x : vector<1xT>
%result = vector.broadcast %unary_result : vector<1xT> to vector<NxT>

Pattern matched (scalar):

%unary_result = unary_op %x : T
%result = vector.broadcast %unary_result : T to vector<NxT>

Transformed to:

%broadcast = vector.broadcast %x : vector<1xT> (or T) to vector<NxT>
%result = unary_op %broadcast : vector<NxT>

This is mathematically valid for element-wise operations where op(broadcast(x)) == broadcast(op(x)).

By default (when op_name is not specified), the transform uses trait-based checking to automatically support all Pure, single-operand, single-result, element-wise operations in math/arith dialects. If op_name is specified, only operations with that exact name are transformed.

Safety conditions checked: 1. The unary op must operate on a single-element vector (e.g., vector<1xf32>) or a scalar (e.g., f32) 2. The unary op result must have exactly one use (the broadcast) 3. The unary op must have exactly one operand and one result 4. Both operations must work on the same element type 5. If op_name is specified, the operation name must match 6. If op_name is empty, the operation must pass trait-based validation

This optimization is particularly beneficial for hardware accelerators like AMD AIEs that can only execute certain operations on vector engines, not on scalar units. Common in layer normalization and other neural network operations.

Example usage (all qualifying unary ops):

%func = transform.structured.match ops{["func.func"]} in %arg0
  : (!transform.any_op) -> !transform.any_op
transform.air.broadcast_before_unary %func

Example usage (specific operation only):

%func = transform.structured.match ops{["func.func"]} in %arg0
  : (!transform.any_op) -> !transform.any_op
transform.air.broadcast_before_unary %func {op_name = "math.rsqrt"}

Returns a handle to the transformed operation.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
op_name	::mlir::StringAttr	string attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.bubble_slice_through_collapse (transform::BubbleSliceThroughCollapseOp)

Move extract_slice ops above their producers, including through collapse_shape at non-constant offsets

Syntax:

operation ::= `transform.air.bubble_slice_through_collapse` $target attr-dict `:` functional-type(operands, results)

Applies, within the target, the patterns of transform.apply_patterns.tensor.bubble_up_extract_slice together with one for a tensor.extract_slice of a tensor.collapse_shape whose offsets are not all constants. The upstream patterns handle constant offsets only. The added pattern moves the slice above the collapse when, in each collapsed dim, the slice takes whole rows of the group's inner dims: its size is a multiple of their product and its offset provably is one (a constant, an scf.for induction variable whose lower bound and step are, or sums and products of such).

A loop over the packets of a collapsed dim then slices the tensor before the collapse, so the slice can be fused into the loop with its producer.

Returns the target.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.coalesce_slices (transform::CoalesceSlicesOp)

Read sibling slices of one tensor through a single copy

Syntax:

operation ::= `transform.air.coalesce_slices` $target (`full_dims` `=` $full_dims^)? attr-dict `:` functional-type(operands, results)

Groups the tensor.extract_slice ops that produce the inputs of the target ops, directly or through tensor.expand_shape and tensor.collapse_shape ops, by source. A group of two or more unit-stride slices in one block, whose offsets in each dimension are either the same value or constants, is rewritten as one linalg.copy of the slices' bounding region into a tensor.empty, with each slice taken from that copy. An offset counts as a constant when it simplifies to one through affine.apply, affine.delinearize_index and arith.addi / arith.muli by constants. The region may not hold more elements than the slices it replaces.

In each dimension listed in full_dims, the region instead spans the whole source, which must have a static size there, and the slices' offsets must be constants. A group may then be a single slice. This is for data whose transfer must cover the whole extent even though only part of it is read, such as a packet with an unused tail.

Promoting the returned copies (with transform.structured.bufferize_to_allocation) then moves the region with one DMA instead of one per slice.

Returns the new linalg.copy ops.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
full_dims	::mlir::DenseI64ArrayAttr	i64 dense array attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.convert_divf_sqrt_to_rsqrt (transform::ConvertDivfSqrtToRsqrtOp)

Convert arith.divf(1.0, math.sqrt(x)) pattern to math.rsqrt(x)

Syntax:

operation ::= `transform.air.convert_divf_sqrt_to_rsqrt` $target attr-dict `:` functional-type(operands, results)

This transform walks through operations with the IsolatedFromAbove trait (such as func.func, air.herd, air.segment, etc.) and identifies the pattern of dividing a constant 1.0 by the result of math.sqrt, converting it to a single math.rsqrt operation.

The transformation looks for:

%sqrt = math.sqrt %x : vector<NxT>
%rsqrt = arith.divf %cst_1.0, %sqrt : vector<NxT>

And replaces it with:

%rsqrt = math.rsqrt %x : vector<NxT>

Safety conditions checked: 1. The numerator must be exactly 1.0 (constant or constant splat for vectors) 2. The sqrt result should have exactly one use (the divf) 3. Both operations must work on the same floating-point type

This optimization is beneficial for hardware accelerators that have native rsqrt instructions, as it's typically faster and more accurate than separate sqrt + division operations. This is a common pattern in layer normalization and other neural network operations.

Example usage:

%herd = transform.structured.match ops{["air.herd"]} in %arg0 
  : (!transform.any_op) -> !transform.any_op
transform.air.convert_divf_sqrt_to_rsqrt %herd

Returns a handle to the transformed operation.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.convert_memref_copy_to_linalg_copy (transform::ConvertMemrefCopyToLinalgCopyOp)

Convert memref.copy operations to linalg.copy operations

Syntax:

operation ::= `transform.air.convert_memref_copy_to_linalg_copy` $target attr-dict `:` functional-type(operands, results)

This transform converts memref.copy operations to linalg.copy operations. This can be useful for enabling further linalg-based optimizations and transformations.

The transformation replaces:

memref.copy %source, %dest : memref<...> to memref<...>

With:

linalg.copy ins(%source : memref<...>) outs(%dest : memref<...>)

Returns a handle to the modified operation containing the transformed copies.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.convert_size1_vector_to_scalar (transform::ConvertSize1VectorToScalarOp)

Convert size-1 vector operations to scalar operations following LLVM's approach

Syntax:

operation ::= `transform.air.convert_size1_vector_to_scalar` $target attr-dict `:` functional-type(operands, results)

This transform converts operations on size-1 vectors (e.g., vector<1xf32>, vector<1x1xi32>) to use scalar operations instead. Following LLVM's VectorDropLeadUnitDim pattern, it uses rewrite patterns (NOT a TypeConverter) and inserts extract/broadcast operations at boundaries.

Key transformations: 1. vector.transfer_read with vector<1xT> → memref.load (scalar) + vector.broadcast 2. vector.transfer_write with vector<1xT> → vector.extract + memref.store (scalar) 3. vector.load with vector<1xT> → memref.load (scalar) + vector.broadcast 4. vector.store with vector<1xT> → vector.extract + memref.store (scalar) 5. Elementwise operations on vector<1xT> → vector.extract + scalar op + vector.broadcast

Important: Function signatures and block argument types remain as vector<1xT>. Only the operations inside are converted to scalars, with extract/broadcast at boundaries. Canonicalization then folds redundant extract(broadcast(x)) → x pairs.

Example with memory operations:

// Before:
%v = vector.transfer_read %mem[%i, %j], %pad : memref<8x8xf32>, vector<1xf32>
%result = arith.mulf %v, %v : vector<1xf32>
vector.transfer_write %result, %mem[%i, %j] : vector<1xf32>, memref<8x8xf32>

// After:
%scalar = memref.load %mem[%i, %j] : memref<8x8xf32>
%result = arith.mulf %scalar, %scalar : f32
memref.store %result, %mem[%i, %j] : memref<8x8xf32>

Example with function arguments (extract/broadcast at boundaries):

// Before:
func.func @foo(%a: vector<1xf32>, %b: vector<1xf32>) -> vector<1xf32> {
  %add = arith.addf %a, %b : vector<1xf32>
  %mul = arith.mulf %add, %a : vector<1xf32>
  return %mul : vector<1xf32>
}

// After:
func.func @foo(%a: vector<1xf32>, %b: vector<1xf32>) -> vector<1xf32> {
  %0 = vector.extract %a[0] : f32 from vector<1xf32>
  %1 = vector.extract %b[0] : f32 from vector<1xf32>
  %2 = arith.addf %0, %1 : f32
  %3 = vector.extract %a[0] : f32 from vector<1xf32>
  %4 = arith.mulf %2, %3 : f32
  %5 = vector.broadcast %4 : f32 to vector<1xf32>
  return %5 : vector<1xf32>
}

Example with loops (boundaries preserved):

// Before:
%result = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg = %init) -> (vector<1xf32>) {
  %add = arith.addf %arg, %cst : vector<1xf32>
  scf.yield %add : vector<1xf32>
}

// After:
%result = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg = %init) -> (vector<1xf32>) {
  %0 = vector.extract %arg[0] : f32 from vector<1xf32>
  %1 = arith.addf %0, %cst_scalar : f32
  %2 = vector.broadcast %1 : f32 to vector<1xf32>
  scf.yield %2 : vector<1xf32>
}

Result: Clean scalar operations on memref.load/store with extract/broadcast only at region boundaries. Canonicalization folds adjacent extract/broadcast pairs automatically.

Usage:

%func = transform.structured.match ops{["func.func"]} in %arg0 
  : (!transform.any_op) -> !transform.any_op
transform.air.convert_size1_vector_to_scalar %func

Returns a handle to the transformed operation.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.copy_to_dma (transform::CopyToDmaOp)

Syntax:

operation ::= `transform.air.copy_to_dma` $target attr-dict `:` functional-type(operands, results)

Transform a memref.copy operation into a air.dma_memcpy_nd operation. Returns the new air.dma_memcpy_nd operation.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.eliminate_cascade_memcpy (transform::EliminateCascadeMemcpyOp)

Eliminate intermediate memref buffers in cascaded DMA operations

Syntax:

operation ::= `transform.air.eliminate_cascade_memcpy` $target attr-dict `:` functional-type(operands, results)

This transform identifies and eliminates intermediate memref buffers in cascaded air.dma_memcpy_nd operations. It looks for the pattern where an intermediate buffer is used exactly twice: once as the destination of a DMA operation and once as the source of another DMA operation, with both operations using default access patterns (empty offsets, sizes, and strides).

The transformation replaces:

air.dma_memcpy_nd (%intermediate[] [] [], %source[] [] []) : (memref<...>, memref<...>)
air.dma_memcpy_nd (%dest[] [] [], %intermediate[] [] []) : (memref<...>, memref<...>)

With:

air.dma_memcpy_nd (%dest[] [] [], %source[] [] []) : (memref<...>, memref<...>)

This optimization eliminates unnecessary intermediate memory allocations and reduces memory traffic, which is particularly beneficial for cascade patterns in AIR programs.

Returns a handle to the modified operation.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.eliminate_redundant_vector_transfers (transform::EliminateRedundantVectorTransfersOp)

Eliminate redundant vector.transfer_read operations

Syntax:

operation ::= `transform.air.eliminate_redundant_vector_transfers` $target attr-dict `:` functional-type(operands, results)

This transform identifies and eliminates redundant vector.transfer_read operations within the target operation. Two vector.transfer_read operations are considered redundant when: 1. They read from the same memref source 2. They use identical indices for the read 3. They have the same result type 4. No write operations to the source memref occur between them

The transformation walks through all vector.transfer_read operations in the target, compares each pair, and when a redundant read is found, replaces all uses of the second read with the result of the first read, then erases the redundant operation.

This optimization is particularly useful after loop unrolling or other transformations that may duplicate read operations unnecessarily, reducing memory traffic and register pressure.

Example:

// Before:
%0 = vector.transfer_read %memref[%i, %j], %pad : memref<8x8xi32>, vector<4xi32>
%1 = vector.add %0, %cst : vector<4xi32>
%2 = vector.transfer_read %memref[%i, %j], %pad : memref<8x8xi32>, vector<4xi32>  // Redundant!
%3 = vector.mul %2, %other : vector<4xi32>

// After:
%0 = vector.transfer_read %memref[%i, %j], %pad : memref<8x8xi32>, vector<4xi32>
%1 = vector.add %0, %cst : vector<4xi32>
%3 = vector.mul %0, %other : vector<4xi32>  // Uses %0 instead of redundant %2

Returns a handle to the transformed operation.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.flatten_for_iter_args (transform::FlattenForIterArgsOp)

Flatten vector-typed iter_args of an scf.for loop using vector.shape_cast

Syntax:

operation ::= `transform.air.flatten_for_iter_args` $target attr-dict `:` functional-type(operands, results)

This transform takes a handle to an scf.for loop and flattens all vector-typed iter_args by inserting vector.shape_cast operations. The transformation: 1. Identifies all iter_args with vector types 2. For each vector iter_arg, inserts a vector.shape_cast before the loop to flatten it 3. Updates the iter_arg type to the flattened vector type 4. Inserts vector.shape_cast operations inside the loop body to convert back to the original shape 5. Updates the scf.yield to flatten yielded values with vector.shape_cast

This is useful for ensuring that loop-carried dependencies use flattened vector types, which can be required by certain backend lowerings or optimization passes.

Example:

// Before:
%result:4 = scf.for %i = %c0 to %c4 step %c1 
    iter_args(%arg0 = %v0, %arg1 = %v1, %arg2 = %v2, %arg3 = %v3) 
    -> (vector<1x1x8x8xi16>, vector<1x1x8x8xi16>, vector<1x1x8x8xi16>, vector<1x1x8x8xi16>) {
  // ... computation ...
  scf.yield %r0, %r1, %r2, %r3 : vector<1x1x8x8xi16>, vector<1x1x8x8xi16>, vector<1x1x8x8xi16>, vector<1x1x8x8xi16>
}

// After:
%v0_flat = vector.shape_cast %v0 : vector<1x1x8x8xi16> to vector<64xi16>
%v1_flat = vector.shape_cast %v1 : vector<1x1x8x8xi16> to vector<64xi16>
%v2_flat = vector.shape_cast %v2 : vector<1x1x8x8xi16> to vector<64xi16>
%v3_flat = vector.shape_cast %v3 : vector<1x1x8x8xi16> to vector<64xi16>
%result:4 = scf.for %i = %c0 to %c4 step %c1 
    iter_args(%arg0 = %v0_flat, %arg1 = %v1_flat, %arg2 = %v2_flat, %arg3 = %v3_flat) 
    -> (vector<64xi16>, vector<64xi16>, vector<64xi16>, vector<64xi16>) {
  %arg0_shaped = vector.shape_cast %arg0 : vector<64xi16> to vector<1x1x8x8xi16>
  %arg1_shaped = vector.shape_cast %arg1 : vector<64xi16> to vector<1x1x8x8xi16>
  %arg2_shaped = vector.shape_cast %arg2 : vector<64xi16> to vector<1x1x8x8xi16>
  %arg3_shaped = vector.shape_cast %arg3 : vector<64xi16> to vector<1x1x8x8xi16>
  // ... computation using %arg0_shaped, %arg1_shaped, etc. ...
  %r0_flat = vector.shape_cast %r0 : vector<1x1x8x8xi16> to vector<64xi16>
  %r1_flat = vector.shape_cast %r1 : vector<1x1x8x8xi16> to vector<64xi16>
  %r2_flat = vector.shape_cast %r2 : vector<1x1x8x8xi16> to vector<64xi16>
  %r3_flat = vector.shape_cast %r3 : vector<1x1x8x8xi16> to vector<64xi16>
  scf.yield %r0_flat, %r1_flat, %r2_flat, %r3_flat : vector<64xi16>, vector<64xi16>, vector<64xi16>, vector<64xi16>
}

Returns a handle to the transformed loop.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.fold_extract_slice_of_insert_slice (transform::FoldExtractSliceOfInsertSliceOp)

Read slices of an assembled tensor from the parts it was built from

Syntax:

operation ::= `transform.air.fold_extract_slice_of_insert_slice` $target attr-dict `:` functional-type(operands, results)

Applies, within the target, patterns that read a tensor.extract_slice of a tensor.insert_slice from the operand that holds the extracted box. The boxes need static sizes and unit strides; the extract's offsets may be values with known bounds, such as an induction variable.

• Disjoint boxes: the extract reads the insert's destination.
• The inserted box contains the extracted one, at constant offsets: the extract reads the insert's source. A rank-reducing insert is undone by a tensor.expand_shape of that read.
• insert(x at S, insert(y at T, insert(w at S, z))) with S and T disjoint becomes insert(y at T, insert(x at S, z)).
• A slice of a tensor.expand_shape that only adds unit dims reads the expand's source.

The same box read back with the inserted value's type is left to the tensor.extract_slice folder. A value assembled from parts then reads each part where it was made, and a part no one reads becomes dead.

Returns the target.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.fold_iter_args_into_induction_var (transform::FoldIterArgsIntoInductionVarOp)

Compute values advanced by a fixed step from the induction variable

Syntax:

operation ::= `transform.air.fold_iter_args_into_induction_var` $target attr-dict `:` functional-type(operands, results)

Each target is an scf.for. An integer or index iteration argument whose yielded value is itself plus a value defined outside the loop is a linear function of the iteration count. Its uses in the body read init + i * delta instead, where i is the iteration number (iv - lb) / step, the loop's result reads init + n * delta for the trip count n, and the argument is then passed through unchanged, which canonicalization removes.

Loops written with pointer increments carry their offsets this way. An offset computed from the induction variable can be folded into a DMA's strides; one carried by the loop cannot.

Returns the loops.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.fold_pack_into_generic (transform::FoldPackIntoGenericOp)

Read a packed input through the generic's indexing map

Syntax:

operation ::= `transform.air.fold_pack_into_generic` $target attr-dict `:` functional-type(operands, results)

For each input of the target linalg.generic produced by a linalg.pack without padding (every tiled dimension divides evenly), replaces the operand by the pack's source and composes the pack's relayout into the operand's indexing map: a packed index (..., o, ..., i) reads source element o * tile + i. Packs left without uses are erased. The generic then reads the source directly, and no relayout copy (on AIE, a DMA between two L1 buffers) is made.

The composed maps are not projected permutations: tile the generic so the outer packed dimensions are 1 and fold unit dimensions before vectorizing it. structured.lower_pack, generalizing its transpose and fusing it into the generic gives an equivalent read through a tensor.expand_shape view with a projected-permutation map; this op is for packs that have already been tiled and fused into a loop.

Returns the updated generics.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.fold_unit_extent_dims (transform::FoldUnitExtentDimsOp)

Fold unit-extent dimensions with bounded greedy iterations

Syntax:

operation ::= `transform.air.fold_unit_extent_dims` $target attr-dict `:` functional-type(operands, results)

Applies linalg fold_unit_extent_dims_via_reshapes patterns to the target function with a bounded number of greedy rewrite iterations. This is needed because LLVM 23's populateFoldUnitExtentDimsPatterns doesn't converge in the greedy driver on IR with air.herd ops containing unit-extent memref dimensions. This op runs the patterns with a limited iteration count and ignores non-convergence, preserving the partial results.

Also applies canonicalization, tiling canonicalization, and scf.for loop canonicalization patterns alongside fold_unit_extent to ensure proper pattern interactions.

Returns a handle to the modified function.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.forall_with_reduce_to_parallel (transform::ForallWithReduceToParallelOp)

Converts a pattern of scf.forall and linalg.reduce to scf.parallel

Syntax:

operation ::= `transform.air.forall_with_reduce_to_parallel` $target attr-dict `:` functional-type(operands, results)

.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
transformed	variadic of TransformHandleTypeInterface instance

transform.air.fuse_elementwise_linalg (transform::FuseElementwiseLinalgOp)

Apply linalg elementwise fusion patterns to a func.func operation

Syntax:

operation ::= `transform.air.fuse_elementwise_linalg` $target (`fold_reshapes` $fold_reshapes^)? attr-dict `:` functional-type(operands, results)

This transform walks the body of a func.func operation and applies the linalg elementwise fusion patterns, which include: - FuseElementwiseOps: Fuses producer elementwise operations into consumer operations - FoldFillWithGenericOp: Folds fill operations into generic operations - FoldScalarOrSplatConstant: Folds scalar or splat constants - RemoveOutsDependency: Removes unnecessary output dependencies - EraseUnusedOperandsAndResults: Cleans up dead operands and results

This is the transform dialect equivalent of running the -linalg-fuse-elementwise-ops pass on the target function.

Example:

// Before fusion:
%0 = linalg.generic {
  indexing_maps = [#map0, #map1],
  iterator_types = ["parallel"]
} ins(%input : tensor<16xf32>) outs(%temp1 : tensor<16xf32>) {
  ^bb0(%arg0: f32, %arg1: f32):
    %1 = arith.mulf %arg0, %cst : f32
    linalg.yield %1 : f32
} -> tensor<16xf32>

%result = linalg.generic {
  indexing_maps = [#map0, #map2, #map1],
  iterator_types = ["parallel"]
} ins(%0, %input2 : tensor<16xf32>, tensor<16xf32>) outs(%temp2 : tensor<16xf32>) {
  ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
    %2 = arith.addf %arg0, %arg1 : f32
    linalg.yield %2 : f32
} -> tensor<16xf32>

// After fusion:
%result = linalg.generic {
  indexing_maps = [#map0, #map2, #map1],
  iterator_types = ["parallel"]
} ins(%input, %input2 : tensor<16xf32>, tensor<16xf32>) outs(%temp2 : tensor<16xf32>) {
  ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
    %1 = arith.mulf %arg0, %cst : f32
    %2 = arith.addf %1, %arg1 : f32
    linalg.yield %2 : f32
} -> tensor<16xf32>

With fold_reshapes, also folds tensor.expand_shape and tensor.collapse_shape ops between two all-parallel linalg.generic ops into them by expanding the iteration space (upstream fold-reshape-by-expansion), so that elementwise chains split by reshapes fuse too. Reshapes next to any other op, a contraction included, are left alone.

Returns a handle to the modified function.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
fold_reshapes	::mlir::UnitAttr	unit attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.fuse_into_containing_op (transform::FuseIntoContainingMemrefOp)

Fuse a producer into a containing operation.

Syntax:

operation ::= `transform.air.fuse_into_containing_op` $producer_op `into` $containing_op attr-dict `:` functional-type(operands, results)

Fuses the producer_op into the containing_op. Returns a handle to the fused ops.

The producer is a subview slice of a tiled op. This transform computes the accessed producer slice inside of the containing op ("tile and fuse").

The containing op handle must be associated with exactly one payload op. The producer op handle may be associated with multiple payload ops. This transform fuses exactly one producer.

Return modes

If the producer could not be fused, this operation fails silently. This is the case when tiling fails or when the producer op has zero uses within the containing op. I.e., "producers" that are not consumed within the containing op are rejected by this operation.

This operation reads and frees the producer handle. This operation reads the containing op handle.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Operands:

Operand	Description
producer_op	TransformHandleTypeInterface instance
containing_op	TransformHandleTypeInterface instance

Results:

Result	Description
fused_op	TransformHandleTypeInterface instance

transform.air.fuse_multi_op_linalg (transform::FuseMultiOpLinalgOp)

Fuse a linalg operation containing multiple element-wise ops with its consumer

Syntax:

operation ::= `transform.air.fuse_multi_op_linalg` $first_op `,` $second_op attr-dict `:` functional-type(operands, results)

This transform fuses two linalg operations where: 1. The first operation contains one or more element-wise pure operations in its body 2. The first operation has only parallel iterator types (no reduction dimensions) 3. The second operation directly consumes the result of the first operation

The second operation may have any iterator types (parallel or reduction).

This is a generalization of FuseExtfLinalgOp that supports multiple operations. The fusion is performed by: 1. Validating that the first op contains only element-wise pure operations 2. Validating that the first op has no reduction dimensions 3. Cloning the entire body of the first op into the second op at the consumption point 4. Updating block argument mappings to maintain SSA form 5. Erasing the first operation

All operations in the first op's body must be: - Pure (no side effects) - Element-wise (operate independently on each element)

Example with reduction in second op:

// Before fusion:
%0 = linalg.generic {
  indexing_maps = [#map0, #map1],
  iterator_types = ["parallel", "parallel"]
} ins(%input : tensor<16x8xf16>) outs(%temp : tensor<16x8xf32>) {
  ^bb0(%arg0: f16):
    %1 = arith.extf %arg0 : f16 to f32
    %2 = arith.mulf %1, %cst : f32
    linalg.yield %2 : f32
} -> tensor<16x8xf32>

%result = linalg.generic {
  indexing_maps = [#map2, #map3],
  iterator_types = ["parallel", "reduction"]
} ins(%0 : tensor<16x8xf32>) outs(%output : tensor<16xf32>) {
  ^bb0(%arg0: f32, %arg1: f32):
    %3 = arith.addf %arg0, %arg1 : f32
    linalg.yield %3 : f32
} -> tensor<16xf32>

// After fusion:
%result = linalg.generic {
  indexing_maps = [#map4, #map3],
  iterator_types = ["parallel", "reduction"]
} ins(%input : tensor<16x8xf16>) outs(%output : tensor<16xf32>) {
  ^bb0(%arg0: f16, %arg1: f32):
    %1 = arith.extf %arg0 : f16 to f32
    %2 = arith.mulf %1, %cst : f32
    %3 = arith.addf %2, %arg1 : f32
    linalg.yield %3 : f32
} -> tensor<16xf32>

Returns a handle to the fused operation (the second operation after modification).

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Operands:

Operand	Description
first_op	TransformHandleTypeInterface instance
second_op	TransformHandleTypeInterface instance

Results:

Result	Description
fused_op	TransformHandleTypeInterface instance

transform.air.fuse_truncf_linalg (transform::FuseTruncfLinalgOp)

Fuse a linalg operation containing only arith.truncf into its producer

Syntax:

operation ::= `transform.air.fuse_truncf_linalg` $truncf_op `,` $producer_op attr-dict `:` functional-type(operands, results)

This transform fuses two linalg operations where: 1. The truncf operation contains only an arith.truncf operation in its body (apart from terminator) 2. The producer operation produces a result that is consumed by the truncf operation

The fusion is performed by: 1. Taking the producer operation's body 2. Adding arith.truncf operation before the terminator 3. Updating the output type to use the truncated (narrower) type 4. Erasing both the original truncf operation and producer operation

This optimization folds the arithmetic truncations into the producer linalg ops, enabling the use of native intrinsics on narrower datatypes, such as AMD AIEs, and reducing intermediate memory storage requirements.

Example:

// Before fusion:
%0 = linalg.generic {
  ^bb0(%arg0: f32, %arg1: f32):
    %1 = arith.addf %arg0, %arg1 : f32
    linalg.yield %1 : f32
} ins(%input1, %input2 : tensor<16xf32>, tensor<16xf32>) outs(%temp : tensor<16xf32>)

%result = linalg.generic {
  ^bb0(%arg0: f32):
    %2 = arith.truncf %arg0 : f32 to f16
    linalg.yield %2 : f16
} ins(%0 : tensor<16xf32>) outs(%output : tensor<16xf16>)

// After fusion:
%result = linalg.generic {
  ^bb0(%arg0: f32, %arg1: f32):
    %1 = arith.addf %arg0, %arg1 : f32
    %2 = arith.truncf %1 : f32 to f16
    linalg.yield %2 : f16
} ins(%input1, %input2 : tensor<16xf32>, tensor<16xf32>) outs(%output : tensor<16xf16>)

Returns a handle to the fused operation (the producer operation after modification).

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Operands:

Operand	Description
truncf_op	TransformHandleTypeInterface instance
producer_op	TransformHandleTypeInterface instance

Results:

Result	Description
fused_op	TransformHandleTypeInterface instance

transform.air.get_segment_for (transform::GetSegmentForOp)

Gets a handle to the parent 'air.segment' of the given operation

Syntax:

operation ::= `transform.air.get_segment_for` $target attr-dict `:` functional-type(operands, results)

Produces a handle to the parent air.segment op for each payload IR operation associated with the operand. Fails if a segment cannot be found. The list of operations associated with the handle contains parent operations in the same order as the list associated with the operand, except for operations that are parents to more than one input which are only present once.

Traits: NavigationTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
parent	TransformHandleTypeInterface instance

transform.air.herd_vectorize (transform::AIRHerdVectorizeOp)

Vectorize operations inside air.herd operations

Syntax:

operation ::= `transform.air.herd_vectorize` $target prop-dict attr-dict `:` functional-type(operands, results)

This transform takes a handle to air.herd operations and vectorizes the operations inside their bodies using the same logic as the AIRHerdVectorizePass. It walks the body of each herd operation and applies vectorization patterns to linalg operations and other vectorizable operations.

The transform supports the same options as the AIRHerdVectorizePass: - vectorize_nd_extract: Controls whether to vectorize tensor.extract when the input tensor is rank >= 2 - flatten_1d_depthwise_conv: Controls whether to "flatten" the channel dimension when vectorizing 1D depthwise convolutions - disable_transfer_permutation_map_lowering_patterns: Disables vector transfer permutation map lowering patterns - disable_multi_reduction_to_contract_patterns: Disables multi-reduction to contract patterns - vectorize_padding: Enables vectorization of padding operations

Example:

%herd = transform.structured.match ops{["air.herd"]} in %f : (!transform.any_op) -> !transform.any_op
%vectorized = transform.air.herd_vectorize %herd {
  vectorize_nd_extract = false,
  flatten_1d_depthwise_conv = false,
  vectorize_padding = true
} : (!transform.any_op) -> !transform.any_op

Returns a handle to the transformed air.herd operations.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
vectorize_nd_extract	::mlir::BoolAttr	bool attribute
flatten_1d_depthwise_conv	::mlir::BoolAttr	bool attribute
disable_transfer_permutation_map_lowering_patterns	::mlir::BoolAttr	bool attribute
disable_multi_reduction_to_contract_patterns	::mlir::BoolAttr	bool attribute
vectorize_padding	::mlir::BoolAttr	bool attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.hoist_cast_pair (transform::HoistCastPairOp)

Hoist extension/truncation operation pairs out of a loop

Syntax:

operation ::= `transform.air.hoist_cast_pair` $extension_op `,` $truncation_op `,` $loop_op attr-dict `:` functional-type(operands, results)

This transform takes handles to an extension operation (arith.extsi, arith.extui, or arith.extf), a truncation operation (arith.trunci or arith.truncf), and their parent scf.for loop. It hoists the extension/truncation pair out of the loop by: 1. Moving the extension operation before the loop to extend the initial iter_arg value 2. Changing the loop's iter_arg type from the narrow type to the wide type 3. Removing the extension from inside the loop (now using the iter_arg directly) 4. Removing the truncation from inside the loop 5. Adding a truncation after the loop to convert the result back to the narrow type

Supports the following extension/truncation pairs: - Integer signed extension: arith.extsi + arith.trunci - Integer unsigned extension: arith.extui + arith.trunci - Floating-point extension: arith.extf + arith.truncf

This optimization is beneficial when accumulator values are repeatedly extended to a wider type for computation and then truncated back to a narrow type at each iteration. By keeping the accumulator in the wide type throughout all loop iterations, we eliminate redundant extend/truncate operations.

Example (Integer):

// Before:
%init = ... : vector<64xi16>
%result = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg = %init) -> (vector<64xi16>) {
  %arg_shaped = vector.shape_cast %arg : vector<64xi16> to vector<1x1x8x8xi16>
  %arg_ext = arith.extsi %arg_shaped : vector<1x1x8x8xi16> to vector<1x1x8x8xi32>
  // ... computation using %arg_ext ...
  %result_i32 = vector.contract ... : ... into vector<1x1x8x8xi32>
  %result_i16 = arith.trunci %result_i32 : vector<1x1x8x8xi32> to vector<1x1x8x8xi16>
  %result_flat = vector.shape_cast %result_i16 : vector<1x1x8x8xi16> to vector<64xi16>
  scf.yield %result_flat : vector<64xi16>
}

// After:
%init = ... : vector<64xi16>
%init_shaped = vector.shape_cast %init : vector<64xi16> to vector<1x1x8x8xi16>
%init_ext = arith.extsi %init_shaped : vector<1x1x8x8xi16> to vector<1x1x8x8xi32>
%init_flat = vector.shape_cast %init_ext : vector<1x1x8x8xi32> to vector<64xi32>
%result_i32 = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg = %init_flat) -> (vector<64xi32>) {
  %arg_shaped = vector.shape_cast %arg : vector<64xi32> to vector<1x1x8x8xi32>
  // ... computation using %arg_shaped directly (no extsi needed) ...
  %result_i32 = vector.contract ... : ... into vector<1x1x8x8xi32>
  %result_flat = vector.shape_cast %result_i32 : vector<1x1x8x8xi32> to vector<64xi32>
  scf.yield %result_flat : vector<64xi32>
}
%result_shaped = vector.shape_cast %result_i32 : vector<64xi32> to vector<1x1x8x8xi32>
%result_i16 = arith.trunci %result_shaped : vector<1x1x8x8xi32> to vector<1x1x8x8xi16>
%result = vector.shape_cast %result_i16 : vector<1x1x8x8xi16> to vector<64xi16>

Example (Floating-point):

// Before:
%init = ... : vector<64xbf16>
%result = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg = %init) -> (vector<64xbf16>) {
  %arg_ext = arith.extf %arg : vector<64xbf16> to vector<64xf32>
  // ... computation using %arg_ext ...
  %result_f32 = vector.fma ... : vector<64xf32>
  %result_bf16 = arith.truncf %result_f32 : vector<64xf32> to vector<64xbf16>
  scf.yield %result_bf16 : vector<64xbf16>
}

// After:
%init = ... : vector<64xbf16>
%init_ext = arith.extf %init : vector<64xbf16> to vector<64xf32>
%result_f32 = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg = %init_ext) -> (vector<64xf32>) {
  // ... computation using %arg directly (no extf needed) ...
  %result_f32 = vector.fma ... : vector<64xf32>
  scf.yield %result_f32 : vector<64xf32>
}
%result = arith.truncf %result_f32 : vector<64xf32> to vector<64xbf16>

Requirements: - Extension and truncation operations must be in the same scf.for loop - Extension must extend an iter_arg (or value derived from iter_arg via shape_cast) - Truncation result must be yielded (directly or indirectly via shape_cast) - Extension and truncation operations must form a valid pair (matching types)

Returns a handle to the transformed loop.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectOpInterface, MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
extension_op	TransformHandleTypeInterface instance
truncation_op	TransformHandleTypeInterface instance
loop_op	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.hoist_loop_carried_pack (transform::HoistLoopCarriedPackOp)

Carry a packed value through an scf.for instead of packing it every iteration

Syntax:

operation ::= `transform.air.hoist_loop_carried_pack` $target attr-dict `:` functional-type(operands, results)

Each target is an scf.for. An iteration argument qualifies when its only uses in the body are as the source of one linalg.pack and as the destination of one linalg.unpack whose result is yielded back to it, and the two use the same layout and the pack has no padding. That is the shape packing a loop-carried accumulator leaves:

%r = scf.for ... iter_args(%acc = %init) {
  %p = linalg.pack %acc ... into %empty
  %x = <compute on %p>
  %u = linalg.unpack %x ... into %acc
  scf.yield %u
}

The loop then carries the packed value instead. The pack moves before the loop, applied to the initial value, and the unpack after it:

%p0 = linalg.pack %init ... into %empty
%r:2 = scf.for ... iter_args(%acc = %init, %pa = %p0) {
  %x = <compute on %pa>
  scf.yield %acc, %x
}
%u = linalg.unpack %r#1 ... into %r#0

The pack's destination must be a tensor.empty or defined outside the loop. Iteration arguments that do not qualify are left alone; a target with none is a silenceable failure.

Returns the new loops and the hoisted packs.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
loops	TransformHandleTypeInterface instance
packs	TransformHandleTypeInterface instance

transform.air.hoist_loop_invariant_transfers (transform::HoistLoopInvariantTransfersOp)

Discover and hoist all loop-invariant vector transfer read/write pairs

Syntax:

operation ::= `transform.air.hoist_loop_invariant_transfers` $scope_op `,` $loop_op attr-dict `:` functional-type(operands, results)

This transform takes handles to a scope operation and an scf.for loop inside it. It automatically discovers all vector.transfer_read/write pairs in the loop that: 1. Have loop-invariant indices (don't depend on the loop induction variable) 2. Access the same memref with equivalent indices (forming a load-modify-store pair)

Each discovered pair is hoisted out of the loop: the read is moved before the loop (with an iter_arg), and the write is moved after the loop, keeping the order the writes had in the loop body. All necessary operand-producing operations (constants, affine.apply, etc.) are also hoisted to maintain SSA dominance.

Index equivalence is checked using areEquivalentIndices(), which handles direct SSA value equality, affine.apply ops with the same map and operands, and constant index equality.

This eliminates the need for fragile split_handle patterns that depend on the exact number and ordering of transfer operations, which can change with different unroll factors, tile sizes, or data types.

The op works across all matmul variants (BF16, I8, I16) and any unroll factor.

Example usage:

%herd = transform.structured.match ops{["air.herd"]} attributes{compute_herd}
  in %arg0 : (!transform.any_op) -> !transform.any_op
%loop = ...  // innermost scf.for loop
%updated_loop = transform.air.hoist_loop_invariant_transfers %herd, %loop
  : (!transform.any_op, !transform.any_op) -> !transform.any_op

Requirements: - The loop must be inside the scope operation - Transfer operations to be hoisted must have loop-invariant indices

Returns a handle to the transformed loop.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectOpInterface, MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
scope_op	TransformHandleTypeInterface instance
loop_op	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.hoist_static_alloc (transform::AIRHoistStaticAllocOp)

Hoist static allocations.

Syntax:

operation ::= `transform.air.hoist_static_alloc` $target attr-dict `:` functional-type(operands, results)

Moves certain statically-sized memref.alloc operations from inner blocks to the entry block of the target function. This shortens and unifies buffer lifetimes, which can unlock reuse and downstream optimizations.

Notes / limitations

• Currently targets memref.alloc buffers with static shapes.
• Uses that require exact type equality across region boundaries (e.g., scf.yield, func.return) are not rewritten; such allocations are skipped.
• Hoisting increases the buffer's lifetime; apply with care on large buffers.

Example

Before:

func.func @foo(%arg0: memref<64xi32>) {
  scf.for %i = %c0 to %c4 step %c1 {
    %tmp = memref.alloc() : memref<64xi32>
    linalg.fill ins(%cst : i32) outs(%tmp : memref<64xi32>)
    memref.dealloc %tmp : memref<64xi32>
  }
  return
}

After:

func.func @foo(%arg0: memref<64xi32>) {
  %tmp.hoisted = memref.alloc() : memref<64xi32>
  scf.for %i = %c0 to %c4 step %c1 {
    linalg.fill ins(%cst : i32) outs(%tmp.hoisted : memref<64xi32>)
  }
  memref.dealloc %tmp.hoisted : memref<64xi32>
  return
}

Usage (Transform dialect)

transform.sequence %arg0 : !transform.any_op failures(propagate) {
^bb0(%f: !transform.any_op):
  transform.air.hoist_static_alloc %f
    : (!transform.any_op) -> ()
}

Traits: ReportTrackingListenerFailuresOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

transform.air.hoist_vector_transfer_pointers (transform::HoistVectorTransferPointersOp)

Optimize vector transfers by hoisting pointer computations out of loops

Syntax:

operation ::= `transform.air.hoist_vector_transfer_pointers` $target attr-dict `:` functional-type(operands, results)

This transform takes a handle to an scf.for loop and optimizes vector transfer operations (vector.transfer_read and vector.transfer_write) inside the loop by: 1. Flattening the vector types to 1D using vector.shape_cast before and after the transfer 2. Flattening multi-dimensional memrefs to 1D using memref.collapse_shape 3. Computing a linearized base pointer from the operation's indices using affine.apply 4. Hoisting the base pointer computation out of the loop 5. For IV-dependent indices: - Passing base pointers as iter_args to the loop - Using the pointer iter_arg directly in the loop body - Incrementing the pointer by a constant stride at each iteration - Yielding the updated pointer for the next iteration

This optimization converts expensive multi-dimensional address calculations inside loops into simple "pointer + constant" arithmetic with iter_args, which is particularly beneficial for hardware accelerators with limited address computation capabilities.

Example with IV-dependent indices:

// Before:
scf.for %i = %c0 to %c8 step %c1 {
  %val = vector.transfer_read %mem[%c0, %i], %pad 
    : memref<32x32xi16>, vector<8x8xi16>
  // ... computation ...
  vector.transfer_write %result, %mem[%c0, %i] 
    : vector<8x8xi16>, memref<32x32xi16>
}

// After:
%flat_mem = memref.collapse_shape %mem [[0, 1]] : memref<32x32xi16> into memref<1024xi16>
%base_ptr = affine.apply affine_map<(d0, d1) -> (d0 * 32 + d1)>(%c0, %c0)
%stride = arith.constant 1 : index
scf.for %i = %c0 to %c8 step %c1 iter_args(%ptr = %base_ptr) -> (index) {
  %val_1d = vector.transfer_read %flat_mem[%ptr], %pad : memref<1024xi16>, vector<64xi16>
  %val = vector.shape_cast %val_1d : vector<64xi16> to vector<8x8xi16>
  // ... computation ...
  %result_1d = vector.shape_cast %result : vector<8x8xi16> to vector<64xi16>
  vector.transfer_write %result_1d, %flat_mem[%ptr] : vector<64xi16>, memref<1024xi16>
  %next_ptr = arith.addi %ptr, %stride : index
  scf.yield %next_ptr : index
}

Requirements: - Target must be an scf.for operation - Vector transfer operations must have concrete vector types (no dynamic dimensions) - The memref must have static shapes for proper stride calculation

Returns a handle to the transformed loop.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.linalg_promote (transform::LinalgPromoteOp)

Syntax:

operation ::= `transform.air.linalg_promote` $target prop-dict attr-dict `:` functional-type(operands, results)

Promotes the specified operands of the target into a separate memory buffer using the mlir::linalg::promoteSubViews utility.

This operation applies to Linalg ops that satisfy the mlir::linalg::promoteSubviewsPrecondition, otherwise it fails.

When successful, several optimization passes are run on the resulting IR. The return handle points to the target operation that was modified inplace.

The operation accepts as attributes the fields in mlir::linalg::LinalgPromotionOptions. In addition the memory space in allocated buffers can be specified with with the memory_space attribute as "L1", "L2" or "L3". The default memory space is L1.

example:

%0 = transform.structured.match ops{["linalg.matmul"]} in %code  : (!transform.any_op) -> !transform.any_op
%1 = transform.air.linalg_promote %0 {memory_space="L2", operands_to_promote=[0]}

The group_size attribute is used to apply promotion to multiple linalg ops. When group_size=N, the operands_to_promote attribute refers to N payload operations at a time and the operand indices apply to the operands of the N operations in the order they appear in the target handle.

For example,

%m = transform.structured.match ops{["linalg.matmul"]} in %f : (!transform.any_op) -> !transform.any_op
%f = transform.structured.match ops{["linalg.fill"]} in %f : (!transform.any_op) -> !transform.any_op
%h = transform.merge_handles %f, %m : !transform.any_op
// promote the input of the fill operation and the output of the matmul operation to L1 memory
transform.air.linalg_promote %h {"group_size"=2, "operands_to_promote"=[1,4], "memory_space"="L1"}

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
operands_to_promote	::mlir::ArrayAttr	64-bit integer array attribute
group_size	::mlir::IntegerAttr	64-bit signless integer attribute
use_full_tile_buffers	::mlir::ArrayAttr	1-bit boolean array attribute
use_full_tiles_by_default	::mlir::UnitAttr	unit attribute
use_alloca	::mlir::UnitAttr	unit attribute
alignment	::mlir::IntegerAttr	64-bit signless integer attribute
memory_space	::mlir::StringAttr	string attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
transformed	TransformHandleTypeInterface instance

transform.air.linalg_tile (transform::LinalgTileOp)

Tile a linalg operation with the given sizes. The new linalg operantion and the generated loop are returned. Tiling is performed with the transform::tileToForallOpImpl so that an scf.forall loop is generated whenever possible.

This is a variant of transform.structured.tile_using_forall.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
static_sizes	::mlir::DenseI64ArrayAttr	i64 dense array attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance
dynamic_sizes	variadic of TransformHandleTypeInterface instance

Results:

Result	Description
tiled_linalg_op	TransformHandleTypeInterface instance
loops	TransformHandleTypeInterface instance

transform.air.linalg_to_library_call (transform::LinalgToLibraryCallOp)

Convert a linalg op to a function call (library call)

Syntax:

operation ::= `transform.air.linalg_to_library_call` $target prop-dict attr-dict `:` functional-type(operands, results)

Replaces a linalg op with a call to a function. If the function_name attribute is provided, it is used as the function name. Otherwise, the linalg op's library_call attribute is used. The function is created if it does not exist. If the link_with attribute is provided, it is used to link the function call to a prebuilt object that contains the implementation of the function. If the linalg op is inside a herd, the link_with attribute is propagated to the herd.

Example:

%matmul = transform.structured.match ops{["linalg.matmul"]} in %f : (!transform.any_op) -> !transform.any_op
%call = transform.air.linalg_to_library_call %matmul { function_name = "my_matmul", link_with = "extern_func.o" } : (!transform.any_op) -> !transform.any_op

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
function_name	::mlir::StringAttr	string attribute
link_with	::mlir::StringAttr	string attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.linearize_vectors (transform::LinearizeVectorsOp)

Rewrite n-D elementwise vector computation as rank-1 vectors

Syntax:

operation ::= `transform.air.linearize_vectors` $target (`arch` `=` $arch^)? (`f32_lanes` `=` $f32_lanes^)? (`contract` $contract^)? attr-dict `:` functional-type(operands, results)

Rewrites the elementwise vector computation nested in target so every value it produces is a rank-1 vector, the form the AIE vector lowerings handle. A rank-2+ elementwise op would otherwise be unrolled along its leading dims during LLVM lowering into one narrow op per innermost row. target need not be isolated from above.

The rewrite runs in phases, each on the ops the previous ones leave:

1. Transfer reads the AIE core lowering cannot take become ones it can:
2. a read broadcasting along a memref dim of extent 1 (a map like (d0, d1, d2) -> (d0, 0, d2)) becomes an identity-map read and a vector.broadcast;
3. a read whose vector is not one contiguous run of its memref (a 4x8 block of rows 64 wide) becomes one read per innermost row and vector.inserts; the rows must be in bounds;
4. a rank-1 read whose map selects one dim followed only by unit dims reads the collapsed memref with a minor identity map;
5. a read of the dims just above the innermost one, with the innermost index a constant k and that dim's extent s between 2 and 8, becomes a read that includes the innermost dim and a vector.shuffle taking every s-th lane from k.
6. Upstream vector linearization of the elementwise ops, constants, vector.bitcast, vector.insert and vector.broadcast (the last as a vector.shuffle). Values crossing into ops left alone are reshaped with vector.shape_cast; transfers, contractions and loops keep their n-D types.
7. Rank-1 rewrites:
8. (x >> [0, 4, 8, ...]) & 15 on 16- or 32-bit words, where lane l of x is word l / k of one vector (k fields per word), is unpacked as vector.bitcast to i4 and arith.extui, or with aievec.unpack on AIE2P when it is exactly 64 or 128 fields;
9. an f32 addf(mulf(a, b), c) with a and b widened from bf16 becomes vector.fma, when both ops carry the contract fast-math flag or contract is set here: contraction drops the product's rounding;
10. elementwise float ops wider than the native width are halved until they fit: 32 lanes for bf16, f32_lanes for ops on f32;
11. an f32 add, sub or mul on fewer than 16 lanes whose result is only replicated (a value per row, broadcast across the row) is done on replicated operands instead;
12. a strided rank-1 read that is in bounds is read element by element.
13. On AIE2P, with f32_lanes 32: a 32-lane bf16 x bf16 + f32 vector.fma becomes aievec.mac_elem and a 32-lane f32 to bf16 arith.truncf (default rounding) aievec.srs. With f32_lanes 64: the 64-lane fma becomes the llvm.aie2p.I1024.I1024.ACC2048.bf.mac.conf intrinsic, subf(c, mulf(a, b)) (with a widened from bf16 and b widened from bf16 or a splat that bf16 represents exactly) the matching bf.msc.conf under the same condition as the fma, the 64-lane truncf two aievec.srs (stored as two halves when its only use is a shape_cast to an in-bounds memref write), and bf16(0x4300 | zext(q)) for unpacked 4-bit q two llvm.aie2p.vshuffle byte interleaves of q with 0x43. Intrinsics are called through private func.func declarations of their names.

arch is "aie2" (default) or "aie2p". f32_lanes is 16 (default), or 32 or 64 on AIE2P.

Run it after vectorization, then canonicalize so that adjacent vector.shape_cast pairs fold.

Example:

%s = vector.transfer_read %scale[%c0, %c0, %c0], %pad
    {in_bounds = [true, true, true],
     permutation_map = affine_map<(d0, d1, d2) -> (d0, 0, d2)>}
    : memref<1x1x8xbf16>, vector<1x4x8xbf16>
%w = arith.mulf %q, %s : vector<1x4x8xbf16>

becomes

%r = vector.transfer_read %scale[%c0, %c0, %c0], %pad
    {in_bounds = [true, true, true]} : memref<1x1x8xbf16>, vector<1x1x8xbf16>
%r1 = vector.shape_cast %r : vector<1x1x8xbf16> to vector<8xbf16>
%s1 = vector.shuffle %r1, %r1 [0, 1, ..., 7, 0, 1, ..., 7, ...]
    : vector<8xbf16>, vector<8xbf16>
%q1 = vector.shape_cast %q : vector<1x4x8xbf16> to vector<32xbf16>
%w1 = arith.mulf %q1, %s1 : vector<32xbf16>

Returns a handle to target.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
arch	::mlir::StringAttr	string attribute
f32_lanes	::mlir::IntegerAttr	64-bit signless integer attribute
contract	::mlir::UnitAttr	unit attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.merge_sibling_copies (transform::MergeSiblingCopiesOp)

Copy several views of one buffer as one copy of their bounding region

Syntax:

operation ::= `transform.air.merge_sibling_copies` $target attr-dict `:` functional-type(operands, results)

Finds the memref.copy ops under target that copy a memref.reinterpret_cast of a common base into a fresh memref.alloc read only through bufferization.to_tensor, and merges those whose views differ only by a constant offset: same base, element type, strides and destination memory space, in the same block, with no write between them. Each group becomes one alloc and one copy of the views' bounding region, and each original buffer's bufferization.to_tensor a tensor.extract_slice of the merged buffer's tensor. Several fields of one packed buffer then move as one DMA stream.

The offsets must differ by a provable constant, which must decompose along the strides into a position on the common view's grid. The region may not hold more elements than the copies it replaces, nor end past the last element any of them reads. Copies that do not qualify are left alone.

Returns the merged copies.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.normalize_for_bounds (transform::NormalizeForBoundsOp)

Normalize scf.for loop bounds by folding affine.apply on induction variable

Syntax:

operation ::= `transform.air.normalize_for_bounds` $target attr-dict `:` functional-type(operands, results)

This transform normalizes an scf.for loop by folding affine.apply operations that multiply the induction variable by a constant factor into the loop bounds.

The transformation looks for patterns where the induction variable is multiplied by a constant via affine.apply, and folds this multiplication into the loop bounds to eliminate the affine.apply operation.

For example, if a loop iterates with bounds (0, 64, step=8) and has an affine.apply that scales the induction variable by 8 (i.e., affine_map<(d0) -> (d0 * 8)>), the transformation will: 1. Scale the loop bounds to (0, 512, step=64) 2. Replace all uses of the affine.apply result with the induction variable directly 3. Erase the affine.apply operation

Example:

// Before:
scf.for %i = %c0 to %c64 step %c8 {
  %scaled = affine.apply affine_map<(d0) -> (d0 * 8)>(%i)
  // ... uses of %scaled ...
}

// After:
scf.for %i = %c0 to %c512 step %c64 {
  // ... uses of %i directly (scaled bounds) ...
}

The transform supports: - Affine maps of the form (d0) -> (d0 * constant) - Both positive and negative scaling constants - Multiple affine.apply operations on the same induction variable (applies the first found)

Returns a handle to the transformed loop.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.normalize_loop_to_unit_step (transform::NormalizeLoopToUnitStepOp)

Rewrite scf.for to count from 0 by 1, the IV's value an affine.apply

Syntax:

operation ::= `transform.air.normalize_loop_to_unit_step` $target attr-dict `:` functional-type(operands, results)

Rewrites each target scf.for %i = %lb to %ub step %s in place as scf.for %j = 0 to ceildiv(%ub - %lb, %s) step 1, replacing the uses of %i with affine.apply (%j * %s + %lb). The inverse of air.normalize_for_bounds.

Offsets computed from the induction variable then show their stride, so that tensor.bubble_up_extract_slice can prove a slice of a collapsed dim aligned to the inner dims: a slice at %j * 8 of a dim collapsed from 64x4 is, one at %i (step 8) cannot be shown to be.

The targets must be signed loops over index; otherwise this is a silenceable failure and no loop is rewritten.

Returns the rewritten loops.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.override_memref_memory_space (transform::OverrideMemRefMemorySpaceOp)

Override memref memory spaces within a target operation scope

Syntax:

operation ::= `transform.air.override_memref_memory_space` $target prop-dict attr-dict `:` functional-type(operands, results)

Overrides the memory space of all memref.alloc operations within the target operation to the specified memory_space value. The scope is inferred from the target operation type:

• air.herd -> overrides allocs inside the herd (L1)
• air.segment -> overrides allocs inside the segment but NOT inside herds (L2)
• air.launch -> overrides allocs inside the launch but NOT inside segments/herds
• func.func -> overrides allocs inside the function but NOT inside launch/segment/herds

This exclusive scoping allows assigning different memory spaces at different hierarchy levels by invoking the op multiple times with different targets.

After overriding alloc types, the op also: 1. Propagates updated types through AIR hierarchy block arguments 2. Corrects memory spaces of view-like operations (subview, reinterpret_cast, etc.)

This is the transform dialect equivalent of the air-override-memref-memory-space pass.

Example:

// Override herd allocs to L1 (memory_space 2)
%herd = transform.structured.match ops{["air.herd"]} in %arg1
  : (!transform.any_op) -> !transform.any_op
transform.air.override_memref_memory_space %herd {memory_space = 2}
  : (!transform.any_op) -> !transform.any_op

// Override func-level allocs to L2 (memory_space 1), excluding herds
%func = transform.structured.match ops{["func.func"]} in %arg1
  : (!transform.any_op) -> !transform.any_op
transform.air.override_memref_memory_space %func {memory_space = 1}
  : (!transform.any_op) -> !transform.any_op

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
memory_space	::mlir::IntegerAttr	32-bit signless integer attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.par_to_herd (transform::ParToHerdOp)

Syntax:

operation ::= `transform.air.par_to_herd` $target prop-dict attr-dict `:` functional-type(operands, results)

Transform a scf.parallel operation into a air.herd operation. If the scf.parallel operation has more than two dimensions, then only the last two are used and a new scf.parallel is created outside of the herd. Returns the new air.herd operation.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
first_dim	::mlir::IntegerAttr	64-bit signless integer attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.par_to_launch (transform::ParToLaunchOp)

Syntax:

operation ::= `transform.air.par_to_launch` $target prop-dict attr-dict `:` functional-type(operands, results)

Transform a scf.parallel operation into a air.launch operation. Returns the new air.launch operation.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
has_air_segment	::mlir::BoolAttr	bool attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.par_to_segment (transform::ParToSegmentOp)

Syntax:

operation ::= `transform.air.par_to_segment` $target prop-dict attr-dict `:` functional-type(operands, results)

Transform a scf.parallel operation into a air.segment operation. Returns the new air.segment operation.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
has_air_segment	::mlir::BoolAttr	bool attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.pipeline_reduce (transform::PipelineReduceOp)

Syntax:

operation ::= `transform.air.pipeline_reduce` $target prop-dict attr-dict `:` functional-type(operands, results)

Experimental

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
tile_size	::mlir::ArrayAttr	64-bit integer array attribute
pipeline_depth	::mlir::IntegerAttr	64-bit signless integer attribute
direction	::mlir::StringAttr	string attribute
promote	::mlir::UnitAttr	unit attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.push_unpack_through_slices (transform::PushUnpackThroughSlicesOp)

Move a linalg.unpack below an expand_shape and the single-index slices of it

Syntax:

operation ::= `transform.air.push_unpack_through_slices` $target attr-dict `:` functional-type(operands, results)

Each target is a linalg.unpack whose only user is a tensor.expand_shape splitting exactly one packed dim d into [n, s], whose users are all rank-reducing tensor.extract_slices taking one index k of the s dim and all of every other dim, and whose inner tile t on d is a multiple of s. Each slice is rewritten as an unpack, with inner tile t / s on d, of the packed source's tile dim split into [t / s, s] and sliced at k.

Then an all-parallel linalg.generic with identity maps whose inputs are all such unpacks, of one layout, that does not use linalg.index and does not read its output, runs on the packed values with one unpack after it.

Every target is matched before any is rewritten; a target that does not match is a silenceable failure.

Returns the new unpacks.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.raise_load_buffer_to_tensor (transform::RaiseLoadBufferToTensorOp)

Rewrite a buffer built by a fill and box copies as tensor ops

Syntax:

operation ::= `transform.air.raise_load_buffer_to_tensor` $target attr-dict `:` functional-type(operands, results)

Looks for memref.alloc ops in the target that are written only by an optional linalg.fill of the whole buffer (possibly alone in an scf.if) and by memref.copy ops into boxes of the buffer, and are then read once, by a bufferization.to_tensor with restrict. The tensor read is rewritten as the fill value (or a tensor.empty when there is no fill) with each copied box inserted by tensor.insert_slice, in copy order. Each box is copied into a buffer of its own size and read through bufferization.to_tensor.

A masked Triton load lowers to such a buffer: the other value is filled in when the mask cuts the tile, and the in-mask box is copied over it. As tensor values, a box read back out of the result folds to the copied data, and a fill nothing reads becomes dead.

A fill under a condition is raised as if it always ran. Where it does not run, the elements no copy writes are uninitialized, and the fill value is one of the values they may hold.

A buffer with any other use, a write after the read, or a box with non-unit strides is left unchanged.

Returns the target.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.remove_uninitialized_copy (transform::RemoveUninitializedCopyOp)

Remove copy operations that copy from uninitialized memrefs

Syntax:

operation ::= `transform.air.remove_uninitialized_copy` $target attr-dict `:` functional-type(operands, results)

This transform walks through a func.func operation and identifies memref.copy and linalg.copy operations where the source is an uninitialized memref (allocated but not written to). Such copy operations are erased as they copy undefined data.

The transform detects the pattern where: 1. A memref is allocated with memref.alloc 2. A subview of that memref is created (optional) 3. The memref/subview is used as source in memref.copy or linalg.copy before any write operations

Returns a handle to the modified function.

Examples:

// memref.copy case
%alloc = memref.alloc() : memref<2x16x8xi32, 1>
%subview = memref.subview %alloc[0, 0, 0] [1, 16, 8] [1, 1, 1] : ...
%target = memref.alloc() : memref<1x16x8xi32, 2>
memref.copy %subview, %target  // <- This copy will be erased

// linalg.copy case
%alloc2 = memref.alloc() : memref<16x8xi32, 1>
%target2 = memref.alloc() : memref<16x8xi32, 2>
linalg.copy ins(%alloc2 : memref<16x8xi32, 1>) outs(%target2 : memref<16x8xi32, 2>)  // <- This copy will be erased

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.segment_to_aie (transform::SegmentToAIEOp)

Syntax:

operation ::= `transform.air.segment_to_aie` $target attr-dict `:` functional-type(operands, results)

Lower air.segment operations to mlir-aie modules.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
transformed	TransformHandleTypeInterface instance

transform.air.share_alloc (transform::ShareAllocOp)

Use one buffer for several allocations that are never live together

Syntax:

operation ::= `transform.air.share_alloc` $target attr-dict `:` functional-type(operands, results)

The targets are memref.alloc ops in one block, with static identity layouts, one memory space, and the same element type and element count. The first target is kept, moved ahead of the others in the block, and the others are replaced by it: an alloc of the same type directly, one of a different shape through a memref.collapse_shape of the kept buffer to one dimension and a memref.expand_shape to its own shape. Their memref.dealloc ops are erased.

The transform does not check lifetimes. The caller asserts that no two of the buffers hold live values at the same time, as for the buffers of two loops that run one after the other.

Returns the kept alloc.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.swap_transpose_with_collapse (transform::SwapTransposeWithCollapseOp)

Move transposes above the collapse_shape ops they read

Syntax:

operation ::= `transform.air.swap_transpose_with_collapse` $target attr-dict `:` functional-type(operands, results)

Applies, within the target, a pattern that rewrites a linalg.transpose of a tensor.collapse_shape as a tensor.collapse_shape of a transpose of the uncollapsed tensor, when every dim the transpose moves is a whole reassociation group. The new transpose moves each group's dims together. A transpose that the collapse separated from another one then reads it directly, and the two can compose.

Returns the target.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.transpose_reduce (transform::TransposeReduceOp)

Transpose inputs of linalg.reduce ops to make reduction dimensions innermost

Syntax:

operation ::= `transform.air.transpose_reduce` $target attr-dict `:` functional-type(operands, results)

This transform takes a handle to linalg.reduce operations and checks if the reduction dimensions are at the innermost (last/lowest) dimensions. If any reduction dimension has non-reduction dimensions to the right, it transposes the corresponding inputs to ensure all reduction dimensions are innermost.

For example, if a linalg.reduce operation reduces along dimension 1 in a 3D tensor (shape [M, N, K] reducing along N), this transform will transpose the input to [M, K, N] so that the reduction dimension N becomes innermost.

This optimization is beneficial for hardware accelerators that perform more efficient reductions when the reduction dimensions are contiguous and innermost.

The transformation: 1. Analyzes each linalg.reduce operation's reduction dimensions 2. Determines if any reduction dimension has non-reduction dimensions to its right 3. If so, creates a transpose operation to move reduction dimensions to the end 4. Updates the linalg.reduce operation to work with the transposed input 5. Optionally transposes the output back to the original layout if needed

Returns a handle to the transformed linalg.reduce operations.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.unswitch_loop (transform::UnswitchLoopOp)

Hoist a loop-invariant select condition out of a loop

Syntax:

operation ::= `transform.air.unswitch_loop` $target attr-dict `:` functional-type(operands, results)

Each target is an scf.for. When an arith.select in the loop is decided by an i1 defined outside it (directly, or through a vector.broadcast of it), the loop is replaced by an scf.if on that condition with a copy of the loop in each branch. In the first copy every select on the condition is replaced by its true operand, in the second by its false operand. The operand a copy no longer uses then becomes dead, so work done for only one value of a run-time flag (an optional activation, say) costs nothing when the flag has the other value.

Loops without such a condition are left unchanged.

Returns the new scf.if ops and the unchanged loops.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance

transform.air.vector_type_cast (transform::VectorTypeCastOp)

Cast vector operands and results of vector operations to a user-provided datatype

Syntax:

operation ::= `transform.air.vector_type_cast` $target prop-dict attr-dict `:` functional-type(operands, results)

This transform takes a handle to vector dialect operations and casts input operands and/or results of vector type to a user-provided datatype. By default, if none of input_indices or output_indices are specified, all vector operands and results are cast.

The transformation works by: 1. Finding vector dialect operations within the target 2. For each vector operation, examining its operands and results 3. Skipping the operation entirely if ALL vector operands and results contain exactly one element (e.g., vector<1xf32>, vector<1x1xf16>) 4. Creating cast operations to convert selected vector operands to the target element type 5. Updating the operation to work with the new vector types 6. Creating cast operations to convert selected results back to the original types

This optimization is useful for hardware accelerators that can perform vector operations natively on specific data types (e.g., bf16, f16) while maintaining compatibility with the original precision through selective casting.

Special handling for single-element vectors: Operations where ALL vector operands and results have exactly one element are skipped to avoid unnecessary type conversions that provide no performance benefit. Note that operations with mixed vector sizes (e.g., vector.multi_reduction with a multi-element input and single-element accumulator) are still transformed, as they contain at least one multi-element vector.

Example 1 - Cast all inputs and outputs (default behavior):

// Before:
%result = vector.fma %a, %b, %c : vector<8xf32>

// After (with target_element_type = f16):
%a_cast = arith.truncf %a : vector<8xf32> to vector<8xf16>
%b_cast = arith.truncf %b : vector<8xf32> to vector<8xf16>
%c_cast = arith.truncf %c : vector<8xf32> to vector<8xf16>
%result_f16 = vector.fma %a_cast, %b_cast, %c_cast : vector<8xf16>
%result = arith.extf %result_f16 : vector<8xf16> to vector<8xf32>

Example 2 - Cast only specific inputs:

// Before:
%result = vector.fma %a, %b, %c : vector<8xf32>

// After (with target_element_type = f16, input_indices = [0, 1]):
%a_cast = arith.truncf %a : vector<8xf32> to vector<8xf16>
%b_cast = arith.truncf %b : vector<8xf32> to vector<8xf16>
%result_f16 = vector.fma %a_cast, %b_cast, %c : vector<8xf16, f32, f32>
%result = arith.extf %result_f16 : vector<8xf16> to vector<8xf32>

Example 3 - Cast only outputs:

// Transform only the output
transform.air.vector_type_cast %op {
  target_element_type = f16,
  output_indices = [0]
}

Attributes: - target_element_type: The element type to cast to (required). Supported types include f16, bf16, f32, f64, i8, i16, i32, i64. - input_indices: Optional array of input operand indices to cast. If empty, all vector inputs are cast. - output_indices: Optional array of output result indices to cast. If empty, all vector results are cast.

Returns a handle to the modified operations containing the transformed vector operations.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes:

Attribute	MLIR Type	Description
target_element_type	::mlir::TypeAttr	any type attribute
input_indices	::mlir::ArrayAttr	64-bit integer array attribute
output_indices	::mlir::ArrayAttr	64-bit integer array attribute

Operands:

Operand	Description
target	TransformHandleTypeInterface instance

Results:

Result	Description
result	TransformHandleTypeInterface instance