[Codegen] Masked vectorization of inner_tiled (#23935)
So far, vectorization of `inner_tiled` did not support dynamic tensor
shapes. This PR enables vectorization of `inner_tiled` with dynamic
tensor shapes in the outer dimensions through masking for the case that
vector sizes are provided.
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/test/generic_vectorization_masked_inferred.mlir b/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_masked_inferred.mlir
index 5449920..1badf10 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_masked_inferred.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_masked_inferred.mlir
@@ -873,3 +873,45 @@
// CHECK: %[[MMA:.+]] = iree_codegen.inner_tiled ins(%[[LHS]], %[[RHS]], %[[LHS_SCALE]], %[[RHS_SCALE]]) outs(%[[ACC]])
// CHECK-SAME: : vector<3x5x1x32xf4E2M1FN>, vector<5x1x7x32xf8E4M3FN>, vector<3x5x1xf8E8M0FNU>, vector<5x7x1xf8E8M0FNU> into vector<3x7x4xf32>
// CHECK: vector.transfer_write %[[MMA]], %arg4[%c0, %c0, %c0] {{.*}} : vector<3x7x4xf32>, tensor<3x7x4xf32>
+
+// -----
+
+// Masked vectorization of inner_tiled with dynamic outer dimensions.
+// Vector sizes come from the iree_codegen.vector_tile_sizes attribute.
+// The LHS has shape <?x?x4xf16> (outer dims dynamic), iteration space is
+// [i=2, j=5, k=3], so reads should be masked.
+#contraction_accesses = [
+ affine_map<(i, j, k) -> (i, k)>,
+ affine_map<(i, j, k) -> (k, j)>,
+ affine_map<(i, j, k) -> (i, j)>
+]
+func.func @masked_tensor_multi_mma(%lhs: tensor<?x?x4xf16>, %rhs: tensor<?x?x4xf16>, %acc: tensor<?x?x4xf32>) -> tensor<?x?x4xf32> {
+ %0 = iree_codegen.inner_tiled ins(%lhs, %rhs) outs(%acc) {
+ iree_codegen.vector_tile_sizes = array<i64: 2, 5, 3>,
+ indexing_maps = #contraction_accesses,
+ iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>],
+ kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>,
+ semantics = #iree_gpu.mma_semantics<distributed = true, opaque = false>
+ } : tensor<?x?x4xf16>, tensor<?x?x4xf16> into tensor<?x?x4xf32>
+ return %0 : tensor<?x?x4xf32>
+}
+
+// CHECK-LABEL: func @masked_tensor_multi_mma
+
+// With vectorSizes, reads use create_mask for dynamic dims.
+// LHS: outer (i=2, k=3) + inner (4) → vector<2x3x4xf16>
+// RHS: outer (k=3, j=5) + inner (4) → vector<3x5x4xf16>
+// ACC: outer (i=2, j=5) + inner (4) → vector<2x5x4xf32>
+// CHECK-DAG: %[[CSTF16:.+]] = arith.constant 0.000000e+00 : f16
+// CHECK-DAG: %[[CSTF32:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[LHS_MASK:.+]] = vector.create_mask {{.*}} : vector<2x3x4xi1>
+// CHECK: %[[LHS:.+]] = vector.transfer_read %arg0{{.*}}, %[[CSTF16]], %[[LHS_MASK]]{{.*}} : tensor<?x?x4xf16>, vector<2x3x4xf16>
+// CHECK: %[[RHS_MASK:.+]] = vector.create_mask {{.*}} : vector<3x5x4xi1>
+// CHECK: %[[RHS:.+]] = vector.transfer_read %arg1{{.*}}, %[[CSTF16]], %[[RHS_MASK]]{{.*}} : tensor<?x?x4xf16>, vector<3x5x4xf16>
+// CHECK: %[[ACC_MASK:.+]] = vector.create_mask {{.*}} : vector<2x5x4xi1>
+// CHECK: %[[ACC:.+]] = vector.transfer_read %arg2{{.*}}, %[[CSTF32]], %[[ACC_MASK]]{{.*}} : tensor<?x?x4xf32>, vector<2x5x4xf32>
+// CHECK: %[[MMA:.+]] = iree_codegen.inner_tiled ins(%[[LHS]], %[[RHS]]) outs(%[[ACC]])
+// CHECK-SAME: : vector<2x3x4xf16>, vector<3x5x4xf16> into vector<2x5x4xf32>
+// CHECK: %[[WRITE_MASK:.+]] = vector.create_mask {{.*}} : vector<2x5x4xi1>
+// CHECK: vector.transfer_write %[[MMA]], %arg2{{.*}}, %[[WRITE_MASK]] {in_bounds = [false, false, true]}
+// CHECK-SAME: : vector<2x5x4xf32>, tensor<?x?x4xf32>
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel
index 7d54aa3..5771116 100644
--- a/compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel
@@ -247,6 +247,7 @@
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:UBDialect",
"@llvm-project//mlir:VectorDialect",
+ "@llvm-project//mlir:VectorUtils",
],
)
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt
index 5afd252..e4e0d8b 100644
--- a/compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt
@@ -181,6 +181,7 @@
MLIRTensorDialect
MLIRUBDialect
MLIRVectorDialect
+ MLIRVectorUtils
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
iree::compiler::Dialect::LinalgExt::IR
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp b/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
index a07d3d8..be4e33e 100644
--- a/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
+++ b/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
@@ -20,6 +20,7 @@
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/IRMapping.h"
@@ -1055,9 +1056,17 @@
if (!tiledOp.hasTensorSemantics()) {
return false;
}
- SmallVector<ShapedType> argTypes = tiledOp.getOperandShapedTypes();
- return llvm::all_of(argTypes,
- [](ShapedType st) { return st.hasStaticShape(); });
+ SmallVector<int64_t> loopRanges;
+ tiledOp.getIterationBounds(loopRanges);
+ // If vector sizes are provided (from tile size analysis or config),
+ // dynamic outer shapes are fine - they'll be masked during vectorization.
+ // However, vector sizes must be >= the static outer dimension sizes.
+ if (!vectorSizes.empty()) {
+ return succeeded(
+ vector::isValidMaskedInputVector(loopRanges, vectorSizes));
+ }
+ // Without vector sizes, require static outer shapes.
+ return ShapedType::isStaticShape(loopRanges);
}
FailureOr<SmallVector<Value>> vectorize(Operation *op, RewriterBase &rewriter,
@@ -1070,6 +1079,19 @@
Location loc = tiledOp.getLoc();
SmallVector<ShapedType> argTypes = tiledOp.getOperandShapedTypes();
+ SmallVector<AffineMap> indexingMaps = tiledOp.getIndexingMapsArray();
+
+ // If no vector sizes are provided, use static loop ranges and the inBounds
+ // attribute instead of masking.
+ bool needsMasking = true;
+ if (vectorSizes.empty()) {
+ SmallVector<int64_t> loopRanges;
+ tiledOp.getIterationBounds(loopRanges);
+ assert(ShapedType::isStaticShape(loopRanges) &&
+ "unable to infer vector sizes");
+ vectorSizes = loopRanges;
+ needsMasking = false;
+ }
// Construct the zero padding value for each operand. Ideally, we'd need the
// InnerTile interface to return the padding value to use. If it is not
@@ -1081,13 +1103,35 @@
rewriter, loc, rewriter.getZeroAttr(argType.getElementType()));
});
- SmallVector<Value> newOperands = tiledOp.getOperands();
- for (auto [operand, type, padValue] :
- llvm::zip_equal(newOperands, argTypes, padValues)) {
- operand = vector::createReadOrMaskedRead(
- rewriter, loc, operand, type.getShape(), padValue,
- /*useInBoundsInsteadOfMasking=*/true);
+ // Compute the read shape for each operand.
+ SmallVector<SmallVector<int64_t>> readShapes;
+ for (auto [i, argType] : llvm::enumerate(argTypes)) {
+ if (!needsMasking) {
+ readShapes.push_back(llvm::to_vector(argType.getShape()));
+ continue;
+ }
+ // Outer dimensions come from vector sizes via the indexing map, inner
+ // dimensions are static.
+ SmallVector<int64_t> readShape;
+ AffineMap map = indexingMaps[i];
+ for (AffineExpr expr : map.getResults()) {
+ auto dimExpr = cast<AffineDimExpr>(expr);
+ readShape.push_back(vectorSizes[dimExpr.getPosition()]);
+ }
+ ArrayRef<int64_t> innerShape = tiledOp.getOperandInnerShape(i);
+ readShape.append(innerShape.begin(), innerShape.end());
+ readShapes.push_back(std::move(readShape));
}
+
+ // Read each operand into a vector, with masking if needed.
+ SmallVector<Value> newOperands(tiledOp->getOperands());
+ for (auto [operand, readShape, padValue] :
+ llvm::zip_equal(newOperands, readShapes, padValues)) {
+ operand = vector::createReadOrMaskedRead(
+ rewriter, loc, operand, readShape, padValue,
+ /*useInBoundsInsteadOfMasking=*/!needsMasking);
+ }
+
auto newTiledOp = IREE::Codegen::InnerTiledOp::create(
rewriter, loc,
ValueRange{newOperands}.take_front(tiledOp.getNumInputs()),
@@ -1095,16 +1139,40 @@
tiledOp.getIndexingMaps(), tiledOp.getIteratorTypes(),
tiledOp.getKind(), tiledOp.getSemantics());
+ // Write results back to tensor, with masking if needed.
+ // TODO: Use createWriteOrMaskedWrite once it is promoted to a public
+ // utility in mlir/Dialect/Vector/Utils/VectorUtils.h.
auto zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
SmallVector<Value> results;
- for (auto [result, tensorAcc] :
- llvm::zip_equal(newTiledOp.getResults(), tiledOp.getOutputs())) {
- int64_t rank = cast<RankedTensorType>(tensorAcc.getType()).getRank();
- auto write = vector::TransferWriteOp::create(
- rewriter, loc, result, tensorAcc,
- /*indices=*/SmallVector<Value>(rank, zero),
- /*inBounds=*/SmallVector<bool>(rank, true));
- results.push_back(write.getResults().front());
+ unsigned numInputs = tiledOp.getNumInputs();
+ for (auto [i, result, dest] :
+ llvm::enumerate(newTiledOp.getResults(), tiledOp.getOutputs())) {
+ auto destType = cast<ShapedType>(dest.getType());
+ int64_t rank = destType.getRank();
+ SmallVector<Value> indices(rank, zero);
+
+ ArrayRef<int64_t> writeShape = readShapes[numInputs + i];
+ SmallVector<bool> inBounds(rank);
+ for (int64_t d = 0; d < rank; ++d) {
+ inBounds[d] =
+ destType.isStaticDim(d) && destType.getDimSize(d) >= writeShape[d];
+ }
+
+ auto write = vector::TransferWriteOp::create(rewriter, loc, result, dest,
+ indices, inBounds);
+ if (!needsMasking) {
+ results.push_back(write.getResults().front());
+ continue;
+ }
+ auto vecType = cast<VectorType>(result.getType());
+ auto maskType =
+ vecType.cloneWith(/*shape=*/std::nullopt, rewriter.getI1Type());
+ SmallVector<OpFoldResult> mixedSizes =
+ tensor::getMixedSizes(rewriter, loc, dest);
+ Value mask =
+ vector::CreateMaskOp::create(rewriter, loc, maskType, mixedSizes);
+ results.push_back(
+ mlir::vector::maskOperation(rewriter, write, mask)->getResult(0));
}
return results;
}
