[Codegen] Support more operations in tile size analysis (#23971)
Add support for `inner_tiled`, `linalg.pack` and `linalg.unpack`
operations to the vector tile size analysis.
This is part of https://github.com/iree-org/iree/issues/23415.
Assisted-by: Claude Code
---------
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp
index a03d1e5..dc99cbe 100644
--- a/compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp
@@ -6,6 +6,7 @@
#include "iree/compiler/Codegen/Common/Passes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
#include "llvm/Support/DebugLog.h"
@@ -157,6 +158,73 @@
return result;
}
+ /// Map tile sizes from pack source space (rank N) to pack dest space
+ /// (rank N+K). Divides packed dims by inner tile sizes, applies
+ /// outer_dims_perm, and appends inner tile sizes as new dimensions.
+ /// The static inner tiles are assumed to not be dynamic values.
+ TileSizes mapPackSourceToDest(ArrayRef<int64_t> innerDimsPos,
+ ArrayRef<int64_t> staticInnerTiles,
+ ArrayRef<int64_t> outerDimsPerm) const {
+ if (empty()) {
+ return {};
+ }
+ TileSizes result = *this;
+ for (auto [dimPos, tileSize] :
+ llvm::zip_equal(innerDimsPos, staticInnerTiles)) {
+ assert(ShapedType::isStatic(tileSize) &&
+ "expecting static inner tile size");
+ if (result.dims[dimPos] == kUninitialized ||
+ result.dims[dimPos] == kOverdefined) {
+ continue;
+ }
+ result.dims[dimPos] =
+ llvm::divideCeilSigned(result.dims[dimPos], tileSize);
+ }
+ if (!outerDimsPerm.empty()) {
+ applyPermutationToVector(result.dims, outerDimsPerm);
+ }
+ llvm::append_range(result.dims, staticInnerTiles);
+ return result;
+ }
+
+ /// Map tile sizes from pack dest space (rank N+K) back to source space
+ /// (rank N). Truncates to outer dims, applies inverse permutation, and
+ /// multiplies packed dims by inner tile sizes.
+ TileSizes mapPackDestToSource(unsigned sourceRank,
+ ArrayRef<int64_t> innerDimsPos,
+ ArrayRef<int64_t> staticInnerTiles,
+ ArrayRef<int64_t> outerDimsPerm) const {
+ if (empty()) {
+ return {};
+ }
+ assert(sourceRank <= rank() && "sourceRank exceeds dest tile size rank");
+ TileSizes result(sourceRank);
+ for (unsigned i = 0; i < sourceRank; ++i) {
+ result.dims[i] = dims[i];
+ }
+ if (!outerDimsPerm.empty()) {
+ applyPermutationToVector(result.dims,
+ invertPermutationVector(outerDimsPerm));
+ }
+ for (auto [dimPos, tileSize] :
+ llvm::zip_equal(innerDimsPos, staticInnerTiles)) {
+ if (result.dims[dimPos] == kUninitialized ||
+ result.dims[dimPos] == kOverdefined) {
+ // Uninitialized and overdefined are preserved.
+ continue;
+ }
+ result.dims[dimPos] *= tileSize;
+ }
+ return result;
+ }
+
+ /// Append dimensions from `suffix` to produce a higher-rank TileSizes.
+ TileSizes append(ArrayRef<int64_t> suffix) const {
+ SmallVector<int64_t> fullDims(dims);
+ fullDims.append(suffix.begin(), suffix.end());
+ return TileSizes(fullDims);
+ }
+
/// Lattice join: per-dimension merge. Uninitialized is identity.
static TileSizes join(const TileSizes &lhs, const TileSizes &rhs) {
TileSizes result = lhs;
@@ -304,6 +372,67 @@
return success();
}
+ // InnerTiledOp: Propagate the outer dimensions through the indexing maps
+ // and then append the static inner dimensions.
+ if (auto innerTiledOp = dyn_cast<IREE::Codegen::InnerTiledOp>(op)) {
+ SmallVector<AffineMap> indexingMaps = innerTiledOp.getIndexingMapsArray();
+ unsigned numLoops = indexingMaps[0].getNumDims();
+ TileSizes iterTileSizes(numLoops);
+
+ // Merge tile sizes from all operands into iteration space.
+ // mapToIterationSpace reads only the first getNumResults() elements
+ // from the operand TileSizes, naturally skipping inner dims.
+ for (OpOperand &operand : op->getOpOperands()) {
+ TileSizes opTileSizes = getTileSizesFor(
+ operand.get(), operands[operand.getOperandNumber()]);
+ AffineMap map = indexingMaps[operand.getOperandNumber()];
+ iterTileSizes.merge(opTileSizes.mapToIterationSpace(map));
+ }
+
+ // Propagate to results. Results correspond to outputs.
+ unsigned numInputs = innerTiledOp.getNumInputs();
+ for (unsigned i = 0; i < innerTiledOp.getNumOutputs(); ++i) {
+ AffineMap map = indexingMaps[numInputs + i];
+ TileSizes outerTileSizes = iterTileSizes.mapFromIterationSpace(map);
+ if (outerTileSizes.empty()) {
+ continue;
+ }
+ // Extend with static inner dims to match full operand rank.
+ ArrayRef<int64_t> innerShape =
+ innerTiledOp.getOperandInnerShape(numInputs + i);
+ TileSizes fullTileSizes = outerTileSizes.append(innerShape);
+ propagateIfChanged(results[i], results[i]->join(fullTileSizes));
+ }
+ return success();
+ }
+
+ // Pack ops: map source tile sizes to dest space.
+ if (auto packOp = dyn_cast<linalg::PackOp>(op)) {
+ if (ShapedType::isDynamicShape(packOp.getStaticInnerTiles())) {
+ return success();
+ }
+ TileSizes srcTileSizes = getTileSizesFor(packOp.getSource(), operands[0]);
+ TileSizes destTileSizes = srcTileSizes.mapPackSourceToDest(
+ packOp.getInnerDimsPos(), packOp.getStaticInnerTiles(),
+ packOp.getOuterDimsPerm());
+ propagateIfChanged(results[0], results[0]->join(destTileSizes));
+ return success();
+ }
+
+ // Unpack ops: map source (packed) tile sizes to dest (unpacked) space.
+ if (auto unpackOp = dyn_cast<linalg::UnPackOp>(op)) {
+ if (ShapedType::isDynamicShape(unpackOp.getStaticInnerTiles())) {
+ return success();
+ }
+ TileSizes srcTileSizes =
+ getTileSizesFor(unpackOp.getSource(), operands[0]);
+ TileSizes destTileSizes = srcTileSizes.mapPackDestToSource(
+ unpackOp.getDestRank(), unpackOp.getInnerDimsPos(),
+ unpackOp.getStaticInnerTiles(), unpackOp.getOuterDimsPerm());
+ propagateIfChanged(results[0], results[0]->join(destTileSizes));
+ return success();
+ }
+
return success();
}
};
@@ -365,6 +494,83 @@
return success();
}
+ // InnerTiledOp: propagate backward through indexing maps.
+ if (auto innerTiledOp = dyn_cast<IREE::Codegen::InnerTiledOp>(op)) {
+ SmallVector<AffineMap> indexingMaps = innerTiledOp.getIndexingMapsArray();
+ unsigned numLoops = indexingMaps[0].getNumDims();
+ unsigned numInputs = innerTiledOp.getNumInputs();
+ TileSizes iterTileSizes(numLoops);
+
+ // Gather from results (full-rank lattice, mapToIterationSpace skips
+ // inner dims naturally).
+ for (auto [result, resultLattice] :
+ llvm::zip_equal(op->getResults(), results)) {
+ unsigned resultIdx = cast<OpResult>(result).getResultNumber();
+ AffineMap map = indexingMaps[numInputs + resultIdx];
+ TileSizes tileSizes = getTileSizesFor(result, resultLattice);
+ iterTileSizes.merge(tileSizes.mapToIterationSpace(map));
+ }
+
+ // Gather from operands.
+ for (OpOperand &operand : op->getOpOperands()) {
+ TileSizes opTileSizes = getTileSizesFor(
+ operand.get(), operands[operand.getOperandNumber()]);
+ AffineMap map = indexingMaps[operand.getOperandNumber()];
+ iterTileSizes.merge(opTileSizes.mapToIterationSpace(map));
+ }
+
+ // Propagate back to each operand. Extend outer-rank result with inner
+ // dims to match full operand rank.
+ for (OpOperand &operand : op->getOpOperands()) {
+ unsigned idx = operand.getOperandNumber();
+ AffineMap map = indexingMaps[idx];
+ TileSizes outerTileSizes = iterTileSizes.mapFromIterationSpace(map);
+ if (outerTileSizes.empty()) {
+ continue;
+ }
+ ArrayRef<int64_t> innerShape = innerTiledOp.getOperandInnerShape(idx);
+ TileSizes fullTileSizes = outerTileSizes.append(innerShape);
+ TileSizeLattice *opLattice = operands[idx];
+ propagateIfChanged(opLattice, opLattice->meet(fullTileSizes));
+ }
+ return success();
+ }
+
+ // Pack ops: result tile sizes → source tile sizes (backward).
+ if (auto packOp = dyn_cast<linalg::PackOp>(op)) {
+ if (ShapedType::isDynamicShape(packOp.getStaticInnerTiles())) {
+ return success();
+ }
+ if (packOp.getPaddingValue()) {
+ // We do not backward propagate for pack with padding, as it would
+ // potentially propagate too large tile sizes.
+ return success();
+ }
+ TileSizes resultTileSizes =
+ getTileSizesFor(packOp->getResult(0), results[0]);
+ TileSizes srcTileSizes = resultTileSizes.mapPackDestToSource(
+ packOp.getSourceRank(), packOp.getInnerDimsPos(),
+ packOp.getStaticInnerTiles(), packOp.getOuterDimsPerm());
+ TileSizeLattice *srcLattice = operands[0];
+ propagateIfChanged(srcLattice, srcLattice->meet(srcTileSizes));
+ return success();
+ }
+
+ // Unpack ops: result tile sizes → source tile sizes (backward).
+ if (auto unpackOp = dyn_cast<linalg::UnPackOp>(op)) {
+ if (ShapedType::isDynamicShape(unpackOp.getStaticInnerTiles())) {
+ return success();
+ }
+ TileSizes resultTileSizes =
+ getTileSizesFor(unpackOp->getResult(0), results[0]);
+ TileSizes srcTileSizes = resultTileSizes.mapPackSourceToDest(
+ unpackOp.getInnerDimsPos(), unpackOp.getStaticInnerTiles(),
+ unpackOp.getOuterDimsPerm());
+ TileSizeLattice *srcLattice = operands[0];
+ propagateIfChanged(srcLattice, srcLattice->meet(srcTileSizes));
+ return success();
+ }
+
return success();
}
@@ -383,20 +589,20 @@
// Result querying
//===----------------------------------------------------------------------===//
-/// Gather tile sizes into the iteration space of a linalg op by looking up each
+/// Gather tile sizes into the iteration space of an op by looking up each
/// operand's lattice state in the solver.
-static TileSizes getIterationSpaceTileSizes(linalg::LinalgOp linalgOp,
+static TileSizes getIterationSpaceTileSizes(Operation *op, unsigned numLoops,
+ ArrayRef<AffineMap> indexingMaps,
const DataFlowSolver &solver) {
- unsigned numLoops = linalgOp.getNumLoops();
TileSizes iterTileSizes(numLoops);
- for (OpOperand &operand : linalgOp->getOpOperands()) {
+ for (OpOperand &operand : op->getOpOperands()) {
Value val = operand.get();
const TileSizeLattice *lattice = solver.lookupState<TileSizeLattice>(val);
TileSizes tileSize = getTileSizesFor(val, lattice);
if (tileSize.empty()) {
continue;
}
- AffineMap map = linalgOp.getMatchingIndexingMap(&operand);
+ AffineMap map = indexingMaps[operand.getOperandNumber()];
assert(map.getNumDims() == numLoops);
iterTileSizes.merge(tileSize.mapToIterationSpace(map));
}
@@ -426,23 +632,63 @@
return signalPassFailure();
}
- funcOp->walk([&](linalg::LinalgOp linalgOp) {
- TileSizes tileSizes = getIterationSpaceTileSizes(linalgOp, solver);
+ auto materialize = [](Operation *op, TileSizes tileSizes) -> LogicalResult {
if (tileSizes.isOverdefined()) {
- linalgOp.emitOpError()
+ op->emitOpError()
<< "tile size analysis did not determine a valid tile size";
- return;
+ return failure();
}
if (!tileSizes.isDefined()) {
- LDBG() << "Analysis did not determine tile size for " << *linalgOp;
- return;
+ LDBG() << "Analysis did not determine tile size for " << *op;
+ return success();
}
- assert(tileSizes.rank() == linalgOp.getNumLoops());
-
- linalgOp->setAttr(
+ op->setAttr(
kVectorTileSizesAttrName,
- DenseI64ArrayAttr::get(linalgOp->getContext(), tileSizes.getDims()));
+ DenseI64ArrayAttr::get(op->getContext(), tileSizes.getDims()));
+ return success();
+ };
+
+ auto result = funcOp->walk([&](Operation *op) -> WalkResult {
+ if (isa<linalg::PackOp>(op) || isa<linalg::UnPackOp>(op)) {
+ // linalg.pack and linalg.unpack have an unpacked (rank N) and a packed
+ // (rank N + K) domain. linalg.pack converts from the unpacked domain to
+ // the packed domain, linalg.unpack works the other way round.
+ // Vectorization of the operations expects vector sizes in the packed
+ // domain. After analysis, these are available on operand of linalg.pack
+ // and the result of linalg.unpack, respectively.
+ Value packedVal;
+ if (auto packOp = dyn_cast<linalg::PackOp>(op)) {
+ packedVal = packOp.getResult();
+ } else if (auto unpackOp = dyn_cast<linalg::UnPackOp>(op)) {
+ packedVal = unpackOp.getSource();
+ }
+ const TileSizeLattice *lattice =
+ solver.lookupState<TileSizeLattice>(packedVal);
+ TileSizes tileSizes = getTileSizesFor(packedVal, lattice);
+ return WalkResult(materialize(op, tileSizes));
+ }
+ if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {
+ SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
+ TileSizes tileSizes = getIterationSpaceTileSizes(
+ op, linalgOp.getNumLoops(), indexingMaps, solver);
+ assert(!tileSizes.isDefined() ||
+ tileSizes.rank() == linalgOp.getNumLoops());
+ return WalkResult(materialize(op, tileSizes));
+ }
+
+ if (auto innerTiledOp = dyn_cast<IREE::Codegen::InnerTiledOp>(op)) {
+ SmallVector<AffineMap> indexingMaps =
+ innerTiledOp.getIndexingMapsArray();
+ unsigned numLoops = indexingMaps[0].getNumDims();
+ TileSizes tileSizes =
+ getIterationSpaceTileSizes(op, numLoops, indexingMaps, solver);
+ return WalkResult(materialize(op, tileSizes));
+ }
+ return WalkResult::advance();
});
+ if (result.wasInterrupted()) {
+ signalPassFailure();
+ }
}
};
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_vector_tile_sizes.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_vector_tile_sizes.mlir
index 514546b..98f3d33 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/materialize_vector_tile_sizes.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_vector_tile_sizes.mlir
@@ -247,36 +247,288 @@
// -----
-// Conflicting tile sizes: two to_layout ops with different tile sizes for the
-// same dimension feed into one generic. The dimension becomes overdefined, so
-// no tile size attribute should be materialized.
+#layout_pack = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1], batch_tile = [1, 32], outer_tile = [1, 1],
+ thread_tile = [1, 1], element_tile = [8, 8],
+ subgroup_strides = [0, 0], thread_strides = [0, 0]>
-#layout_32 = #iree_vector_ext.nested_layout<
- subgroup_tile = [1], batch_tile = [4], outer_tile = [1],
- thread_tile = [1], element_tile = [8],
- subgroup_strides = [0], thread_strides = [0]>
+// CHECK-LABEL: @pack_forward_propagation
+func.func @pack_forward_propagation(%arg0: tensor<8x256xf16>) -> tensor<8x8x32xf16> {
+ %laid_out = iree_vector_ext.to_layout %arg0 to layout(#layout_pack) : tensor<8x256xf16>
+ %empty_gen = tensor.empty() : tensor<8x256xf16>
+ // CHECK: linalg.generic
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 256>
+ %gen = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%laid_out : tensor<8x256xf16>) outs(%empty_gen : tensor<8x256xf16>) {
+ ^bb0(%in: f16, %out: f16):
+ %neg = arith.negf %in : f16
+ linalg.yield %neg : f16
+ } -> tensor<8x256xf16>
+ %empty_pack = tensor.empty() : tensor<8x8x32xf16>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 8, 32>
+ %pack = linalg.pack %gen inner_dims_pos = [1] inner_tiles = [32]
+ into %empty_pack : tensor<8x256xf16> -> tensor<8x8x32xf16>
+ return %pack : tensor<8x8x32xf16>
+}
-#layout_64 = #iree_vector_ext.nested_layout<
- subgroup_tile = [1], batch_tile = [8], outer_tile = [1],
- thread_tile = [1], element_tile = [8],
- subgroup_strides = [0], thread_strides = [0]>
+// -----
-// CHECK-LABEL: @conflicting_tile_sizes
-func.func @conflicting_tile_sizes(%arg0: tensor<64xf16>, %arg1: tensor<64xf16>) -> tensor<64xf16> {
- %a = iree_vector_ext.to_layout %arg0 to layout(#layout_32) : tensor<64xf16>
- %b = iree_vector_ext.to_layout %arg1 to layout(#layout_64) : tensor<64xf16>
- %empty = tensor.empty() : tensor<64xf16>
+#layout_pack_perm = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1, 1], batch_tile = [1, 1, 32], outer_tile = [1, 1, 1],
+ thread_tile = [1, 1, 1], element_tile = [4, 16, 8],
+ subgroup_strides = [0, 0, 0], thread_strides = [0, 0, 0]>
+
+// CHECK-LABEL: @pack_forward_propagation_outer_perm
+func.func @pack_forward_propagation_outer_perm(%arg0: tensor<4x16x256xf16>) -> tensor<8x4x16x32xf16> {
+ %laid_out = iree_vector_ext.to_layout %arg0 to layout(#layout_pack_perm) : tensor<4x16x256xf16>
+ %empty_gen = tensor.empty() : tensor<4x16x256xf16>
+ // CHECK: linalg.generic
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 4, 16, 256>
+ %gen = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ } ins(%laid_out : tensor<4x16x256xf16>) outs(%empty_gen : tensor<4x16x256xf16>) {
+ ^bb0(%in: f16, %out: f16):
+ %neg = arith.negf %in : f16
+ linalg.yield %neg : f16
+ } -> tensor<4x16x256xf16>
+ %empty_pack = tensor.empty() : tensor<8x4x16x32xf16>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 4, 16, 32>
+ %pack = linalg.pack %gen outer_dims_perm = [2, 0, 1] inner_dims_pos = [2] inner_tiles = [32]
+ into %empty_pack : tensor<4x16x256xf16> -> tensor<8x4x16x32xf16>
+ return %pack : tensor<8x4x16x32xf16>
+}
+
+// -----
+
+// Pack with padding value: backward propagation through the pack should be
+// blocked, so the generic producing the pack source should not get tile sizes.
+
+#layout_pack_pad = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1, 1], batch_tile = [1, 1, 1], outer_tile = [1, 1, 1],
+ thread_tile = [1, 1, 1], element_tile = [8, 8, 32],
+ subgroup_strides = [0, 0, 0], thread_strides = [0, 0, 0]>
+
+// CHECK-LABEL: @pack_no_backward_propagation_with_padding
+func.func @pack_no_backward_propagation_with_padding(
+ %arg0: tensor<8x256xf16>) -> tensor<8x8x32xf16> {
+ %cst = arith.constant 0.0 : f16
+ %empty_gen = tensor.empty() : tensor<8x256xf16>
// CHECK: linalg.generic
// CHECK-NOT: iree_codegen.vector_tile_sizes
- %result = linalg.generic {
- indexing_maps = [affine_map<(d0) -> (d0)>,
- affine_map<(d0) -> (d0)>,
- affine_map<(d0) -> (d0)>],
- iterator_types = ["parallel"]
- } ins(%a, %b : tensor<64xf16>, tensor<64xf16>) outs(%empty : tensor<64xf16>) {
- ^bb0(%in0: f16, %in1: f16, %out: f16):
- %add = arith.addf %in0, %in1 : f16
- linalg.yield %add : f16
- } -> tensor<64xf16>
- return %result : tensor<64xf16>
+ %gen = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%arg0 : tensor<8x256xf16>) outs(%empty_gen : tensor<8x256xf16>) {
+ ^bb0(%in: f16, %out: f16):
+ %neg = arith.negf %in : f16
+ linalg.yield %neg : f16
+ } -> tensor<8x256xf16>
+ %empty_pack = tensor.empty() : tensor<8x8x32xf16>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 8, 32>
+ %pack = linalg.pack %gen padding_value(%cst : f16)
+ inner_dims_pos = [1] inner_tiles = [32]
+ into %empty_pack : tensor<8x256xf16> -> tensor<8x8x32xf16>
+ %result = iree_vector_ext.to_layout %pack to layout(#layout_pack_pad) : tensor<8x8x32xf16>
+ return %result : tensor<8x8x32xf16>
+}
+
+// -----
+
+#layout_unpack = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1], batch_tile = [1, 32], outer_tile = [1, 1],
+ thread_tile = [1, 1], element_tile = [8, 8],
+ subgroup_strides = [0, 0], thread_strides = [0, 0]>
+
+// CHECK-LABEL: @unpack_backward_propagation
+func.func @unpack_backward_propagation(%arg0: tensor<8x8x32xf16>) -> tensor<8x256xf16> {
+ %empty_unpack = tensor.empty() : tensor<8x256xf16>
+ // CHECK: linalg.unpack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 8, 32>
+ %unpack = linalg.unpack %arg0 inner_dims_pos = [1] inner_tiles = [32]
+ into %empty_unpack : tensor<8x8x32xf16> -> tensor<8x256xf16>
+ %empty_gen = tensor.empty() : tensor<8x256xf16>
+ // CHECK: linalg.generic
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 256>
+ %gen = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%unpack : tensor<8x256xf16>) outs(%empty_gen : tensor<8x256xf16>) {
+ ^bb0(%in: f16, %out: f16):
+ %neg = arith.negf %in : f16
+ linalg.yield %neg : f16
+ } -> tensor<8x256xf16>
+ %result = iree_vector_ext.to_layout %gen to layout(#layout_unpack) : tensor<8x256xf16>
+ return %result : tensor<8x256xf16>
+}
+
+// -----
+
+#layout_chain = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1], batch_tile = [1, 32], outer_tile = [1, 1],
+ thread_tile = [1, 1], element_tile = [8, 8],
+ subgroup_strides = [0, 0], thread_strides = [0, 0]>
+
+// CHECK-LABEL: @pack_unpack_chain
+func.func @pack_unpack_chain(%arg0: tensor<8x256xf16>) -> tensor<8x256xf16> {
+ %laid_out = iree_vector_ext.to_layout %arg0 to layout(#layout_chain) : tensor<8x256xf16>
+ %empty_pack = tensor.empty() : tensor<8x8x32xf16>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 8, 32>
+ %pack = linalg.pack %laid_out inner_dims_pos = [1] inner_tiles = [32]
+ into %empty_pack : tensor<8x256xf16> -> tensor<8x8x32xf16>
+ %empty_gen = tensor.empty() : tensor<8x8x32xf16>
+ // CHECK: linalg.generic
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 8, 32>
+ %gen = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ } ins(%pack : tensor<8x8x32xf16>) outs(%empty_gen : tensor<8x8x32xf16>) {
+ ^bb0(%in: f16, %out: f16):
+ %neg = arith.negf %in : f16
+ linalg.yield %neg : f16
+ } -> tensor<8x8x32xf16>
+ %empty_unpack = tensor.empty() : tensor<8x256xf16>
+ // CHECK: linalg.unpack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 8, 8, 32>
+ %unpack = linalg.unpack %gen inner_dims_pos = [1] inner_tiles = [32]
+ into %empty_unpack : tensor<8x8x32xf16> -> tensor<8x256xf16>
+ return %unpack : tensor<8x256xf16>
+}
+
+// -----
+
+#layout_pv_lhs = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1, 1], batch_tile = [1, 1, 4], outer_tile = [1, 1, 1],
+ thread_tile = [1, 16, 4], element_tile = [1, 1, 4],
+ subgroup_strides = [0, 0, 0], thread_strides = [0, 1, 16]>
+
+#layout_pv_rhs = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1, 1], batch_tile = [1, 4, 4], outer_tile = [1, 1, 1],
+ thread_tile = [1, 4, 16], element_tile = [1, 4, 1],
+ subgroup_strides = [0, 0, 0], thread_strides = [0, 16, 1]>
+
+#layout_pv_acc = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1, 1], batch_tile = [1, 1, 4], outer_tile = [1, 1, 1],
+ thread_tile = [1, 4, 16], element_tile = [1, 4, 1],
+ subgroup_strides = [0, 0, 0], thread_strides = [0, 16, 1]>
+
+// CHECK-LABEL: @inner_tiled_attention_pv
+func.func @inner_tiled_attention_pv(
+ %lhs: tensor<1x16x64xf16>,
+ %rhs: tensor<1x64x64xf16>,
+ %acc: tensor<1x16x64xf32>) -> tensor<1x16x64xf32> {
+ %lhs_l = iree_vector_ext.to_layout %lhs to layout(#layout_pv_lhs) : tensor<1x16x64xf16>
+ %rhs_l = iree_vector_ext.to_layout %rhs to layout(#layout_pv_rhs) : tensor<1x64x64xf16>
+ %acc_l = iree_vector_ext.to_layout %acc to layout(#layout_pv_acc) : tensor<1x16x64xf32>
+ %empty_lhs = tensor.empty() : tensor<1x1x4x16x16xf16>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4, 16, 16>
+ %pack_lhs = linalg.pack %lhs_l inner_dims_pos = [1, 2] inner_tiles = [16, 16]
+ into %empty_lhs : tensor<1x16x64xf16> -> tensor<1x1x4x16x16xf16>
+ %empty_rhs = tensor.empty() : tensor<1x4x4x16x16xf16>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 4, 4, 16, 16>
+ %pack_rhs = linalg.pack %rhs_l inner_dims_pos = [1, 2] inner_tiles = [16, 16]
+ into %empty_rhs : tensor<1x64x64xf16> -> tensor<1x4x4x16x16xf16>
+ %empty_acc = tensor.empty() : tensor<1x1x4x16x16xf32>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4, 16, 16>
+ %pack_acc = linalg.pack %acc_l inner_dims_pos = [1, 2] inner_tiles = [16, 16]
+ into %empty_acc : tensor<1x16x64xf32> -> tensor<1x1x4x16x16xf32>
+ // CHECK: iree_codegen.inner_tiled
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4, 4>
+ %result_packed = iree_codegen.inner_tiled ins(%pack_lhs, %pack_rhs) outs(%pack_acc) {
+ indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>,
+ affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>],
+ iterator_types = [#linalg.iterator_type<parallel>,
+ #linalg.iterator_type<parallel>,
+ #linalg.iterator_type<reduction>,
+ #linalg.iterator_type<parallel>],
+ kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>,
+ semantics = #iree_gpu.mma_semantics<distributed = false, opaque = true>
+ } : tensor<1x1x4x16x16xf16>, tensor<1x4x4x16x16xf16> into tensor<1x1x4x16x16xf32>
+ %empty_result = tensor.empty() : tensor<1x16x64xf32>
+ // CHECK: linalg.unpack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4, 16, 16>
+ %result = linalg.unpack %result_packed inner_dims_pos = [1, 2] inner_tiles = [16, 16]
+ into %empty_result : tensor<1x1x4x16x16xf32> -> tensor<1x16x64xf32>
+ return %result : tensor<1x16x64xf32>
+}
+
+// -----
+
+#layout_pv_lhs = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1, 1], batch_tile = [1, 1, 4], outer_tile = [1, 1, 1],
+ thread_tile = [1, 16, 4], element_tile = [1, 1, 4],
+ subgroup_strides = [0, 0, 0], thread_strides = [0, 1, 16]>
+
+#layout_pv_rhs = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1, 1], batch_tile = [1, 4, 4], outer_tile = [1, 1, 1],
+ thread_tile = [1, 4, 16], element_tile = [1, 4, 1],
+ subgroup_strides = [0, 0, 0], thread_strides = [0, 16, 1]>
+
+#layout_pv_acc = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1, 1], batch_tile = [1, 1, 4], outer_tile = [1, 1, 1],
+ thread_tile = [1, 4, 16], element_tile = [1, 4, 1],
+ subgroup_strides = [0, 0, 0], thread_strides = [0, 16, 1]>
+
+// CHECK-LABEL: @inner_tiled_dynamic
+func.func @inner_tiled_dynamic(
+ %lhs: tensor<1x?x?xf16>,
+ %rhs: tensor<1x?x?xf16>,
+ %acc: tensor<1x?x?xf32>) -> tensor<1x?x?xf32> {
+ %lhs_l = iree_vector_ext.to_layout %lhs to layout(#layout_pv_lhs) : tensor<1x?x?xf16>
+ %rhs_l = iree_vector_ext.to_layout %rhs to layout(#layout_pv_rhs) : tensor<1x?x?xf16>
+ %acc_l = iree_vector_ext.to_layout %acc to layout(#layout_pv_acc) : tensor<1x?x?xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %m = tensor.dim %lhs, %c0 : tensor<1x?x?xf16>
+ %n = tensor.dim %rhs, %c0 : tensor<1x?x?xf16>
+ %k = tensor.dim %lhs, %c1 : tensor<1x?x?xf16>
+ %empty_lhs = tensor.empty(%m, %k) : tensor<1x?x?x16x16xf16>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4, 16, 16>
+ %pack_lhs = linalg.pack %lhs_l inner_dims_pos = [1, 2] inner_tiles = [16, 16]
+ into %empty_lhs : tensor<1x?x?xf16> -> tensor<1x?x?x16x16xf16>
+ %empty_rhs = tensor.empty(%n, %k) : tensor<1x?x?x16x16xf16>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 4, 4, 16, 16>
+ %pack_rhs = linalg.pack %rhs_l inner_dims_pos = [1, 2] inner_tiles = [16, 16]
+ into %empty_rhs : tensor<1x?x?xf16> -> tensor<1x?x?x16x16xf16>
+ %empty_acc = tensor.empty(%m, %n) : tensor<1x?x?x16x16xf32>
+ // CHECK: linalg.pack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4, 16, 16>
+ %pack_acc = linalg.pack %acc_l inner_dims_pos = [1, 2] inner_tiles = [16, 16]
+ into %empty_acc : tensor<1x?x?xf32> -> tensor<1x?x?x16x16xf32>
+ // CHECK: iree_codegen.inner_tiled
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4, 4>
+ %result_packed = iree_codegen.inner_tiled ins(%pack_lhs, %pack_rhs) outs(%pack_acc) {
+ indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>,
+ affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>],
+ iterator_types = [#linalg.iterator_type<parallel>,
+ #linalg.iterator_type<parallel>,
+ #linalg.iterator_type<reduction>,
+ #linalg.iterator_type<parallel>],
+ kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>,
+ semantics = #iree_gpu.mma_semantics<distributed = false, opaque = true>
+ } : tensor<1x?x?x16x16xf16>, tensor<1x?x?x16x16xf16> into tensor<1x?x?x16x16xf32>
+ %empty_result = tensor.empty(%m, %n) : tensor<1x?x?xf32>
+ // CHECK: linalg.unpack
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4, 16, 16>
+ %result = linalg.unpack %result_packed inner_dims_pos = [1, 2] inner_tiles = [16, 16]
+ into %empty_result : tensor<1x?x?x16x16xf32> -> tensor<1x?x?xf32>
+ return %result : tensor<1x?x?xf32>
}
