[Mlir-commits] [mlir] [mlir][vector] Pad the target shape when unrolling masks (PR #234448)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 19:00:09 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Hamza Qureshi (hamzaqureshi5)
<details>
<summary>Changes</summary>
`getTargetShape` only needs to divide the trailing dimensions, so it can have
a smaller rank than the op. Both mask unroll patterns index the offsets by the
op's rank, so a rank 3 mask with a rank 2 target shape reads past the end and
asserts. Pad with leading unit dims, the same way `UnrollElementwisePattern`
does.
With that, `constant_mask` hits a second issue: a tile past the mask ends up
with a zero in one dimension and its full extent in the others, which the op
rejects since it only accepts all zeros for an empty mask. Normalize those
tiles. This one also happens without a rank mismatch, e.g.
`vector.constant_mask [8, 8] : vector<16x16xi1>`.
Fixes #<!-- -->217175
---
Full diff: https://github.com/llvm/llvm-project/pull/234448.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+30-9)
- (modified) mlir/test/Dialect/Vector/vector-unroll-options.mlir (+41)
``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 19abe44081eb4..14bcc27f1fe96 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -1091,16 +1091,23 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
*createMaskOp.getIterationShapeToUnroll();
Location loc = createMaskOp.getLoc();
+ // `getTargetShape` is only required to divide the trailing dimensions, so
+ // it may have a smaller rank than the op. The logic below assumes a 1-1
+ // correspondence with the op's mask operands, so pad with leading unit
+ // dims.
+ SmallVector<int64_t> adjustedTargetShape(originalSize.size(), 1);
+ llvm::copy(*targetShape, adjustedTargetShape.end() - targetShape->size());
+
Value result = arith::ConstantOp::create(rewriter, loc, resultType,
rewriter.getZeroAttr(resultType));
VectorType targetVectorType =
- VectorType::get(*targetShape, rewriter.getI1Type());
- SmallVector<int64_t> strides(targetShape->size(), 1);
+ VectorType::get(adjustedTargetShape, rewriter.getI1Type());
+ SmallVector<int64_t> strides(adjustedTargetShape.size(), 1);
// In each dimension (d), each unrolled vector computes its mask size as:
// min(max(originalMaskOperands[d] - offset[d], 0), unrolledDimSize[d]).
for (SmallVector<int64_t> offsets :
- StaticTileOffsetRange(originalSize, *targetShape)) {
+ StaticTileOffsetRange(originalSize, adjustedTargetShape)) {
SmallVector<Value> unrolledOperands;
for (auto [i, originalMaskOperand] :
@@ -1110,8 +1117,8 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
Value adjustedMaskSize = rewriter.createOrFold<arith::SubIOp>(
loc, originalMaskOperand, offsetVal);
Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
- Value unrolledDimSize =
- arith::ConstantIndexOp::create(rewriter, loc, (*targetShape)[i]);
+ Value unrolledDimSize = arith::ConstantIndexOp::create(
+ rewriter, loc, adjustedTargetShape[i]);
Value nonNegative =
rewriter.createOrFold<arith::MaxSIOp>(loc, adjustedMaskSize, zero);
Value unrolledOperand = rewriter.createOrFold<arith::MinSIOp>(
@@ -1184,16 +1191,23 @@ struct UnrollConstantMaskPattern
*constantMaskOp.getIterationShapeToUnroll();
Location loc = constantMaskOp.getLoc();
+ // `getTargetShape` is only required to divide the trailing dimensions, so
+ // it may have a smaller rank than the op. The logic below assumes a 1-1
+ // correspondence with the op's mask dimensions, so pad with leading unit
+ // dims.
+ SmallVector<int64_t> adjustedTargetShape(originalSize.size(), 1);
+ llvm::copy(*targetShape, adjustedTargetShape.end() - targetShape->size());
+
Value result = arith::ConstantOp::create(rewriter, loc, resultType,
rewriter.getZeroAttr(resultType));
VectorType targetVectorType =
- VectorType::get(*targetShape, rewriter.getI1Type());
- SmallVector<int64_t> strides(targetShape->size(), 1);
+ VectorType::get(adjustedTargetShape, rewriter.getI1Type());
+ SmallVector<int64_t> strides(adjustedTargetShape.size(), 1);
// In each dimension (d), each unrolled vector computes its mask size as:
// min(max(originalMaskDim[d] - offset[d], 0), unrolledDimSize[d]).
for (const SmallVector<int64_t> &offsets :
- StaticTileOffsetRange(originalSize, *targetShape)) {
+ StaticTileOffsetRange(originalSize, adjustedTargetShape)) {
SmallVector<int64_t> unrolledMaskDims;
for (auto [i, originalMaskDim] :
@@ -1203,10 +1217,17 @@ struct UnrollConstantMaskPattern
int64_t adjustedMaskSize =
std::max(originalMaskDim - offsets[i], static_cast<int64_t>(0));
int64_t unrolledMaskDim =
- std::min(adjustedMaskSize, static_cast<int64_t>((*targetShape)[i]));
+ std::min(adjustedMaskSize, adjustedTargetShape[i]);
unrolledMaskDims.push_back(unrolledMaskDim);
}
+ // A tile that starts past the mask in some dimension gets a zero there
+ // while the remaining dimensions keep their extent. `constant_mask` is
+ // the conjunction of its dimensions, so such a tile is empty, and the op
+ // accepts only the all zeros spelling for an empty mask.
+ if (llvm::is_contained(unrolledMaskDims, 0))
+ unrolledMaskDims.assign(unrolledMaskDims.size(), 0);
+
auto unrolledMask = rewriter.createOrFold<vector::ConstantMaskOp>(
loc, targetVectorType, unrolledMaskDims);
result = rewriter.createOrFold<vector::InsertStridedSliceOp>(
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 61a74744e7a72..4d029f84c1501 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -621,6 +621,47 @@ func.func @vector_constant_mask() -> vector<16x16xi1> {
// -----
+func.func @vector_constant_mask_rank_mismatch() -> vector<2x8x8xi1> {
+ %0 = vector.constant_mask [1, 8, 8] : vector<2x8x8xi1>
+ return %0 : vector<2x8x8xi1>
+}
+
+// CHECK-LABEL: func @vector_constant_mask_rank_mismatch
+// CHECK-SAME: () -> vector<2x8x8xi1>
+// CHECK: %[[EMPTY:.*]] = arith.constant dense<false> : vector<1x8x8xi1>
+// CHECK: %[[ACC:.*]] = arith.constant dense<false> : vector<2x8x8xi1>
+// CHECK: %[[FULL:.*]] = arith.constant dense<true> : vector<1x8x8xi1>
+// CHECK: %[[INS0:.*]] = vector.insert_strided_slice %[[FULL]], %[[ACC]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: %[[INS1:.*]] = vector.insert_strided_slice %[[EMPTY]], %[[INS0]] offsets = [1, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: return %[[INS1]] : vector<2x8x8xi1>
+
+// -----
+
+func.func @vector_create_mask_rank_mismatch(%size: index) -> vector<2x8x8xi1> {
+ %c8 = arith.constant 8 : index
+ %0 = vector.create_mask %size, %c8, %c8 : vector<2x8x8xi1>
+ return %0 : vector<2x8x8xi1>
+}
+
+// CHECK-LABEL: func @vector_create_mask_rank_mismatch
+// CHECK-SAME: (%[[SIZE:.*]]: index) -> vector<2x8x8xi1>
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[ACC:.*]] = arith.constant dense<false> : vector<2x8x8xi1>
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[MAX0:.*]] = arith.maxsi %[[SIZE]], %[[C0]] : index
+// CHECK: %[[MIN0:.*]] = arith.minsi %[[MAX0]], %[[C1]] : index
+// CHECK: %[[MASK0:.*]] = vector.create_mask %[[MIN0]], %[[C8]], %[[C8]] : vector<1x8x8xi1>
+// CHECK: %[[INS0:.*]] = vector.insert_strided_slice %[[MASK0]], %[[ACC]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: %[[SUB1:.*]] = arith.subi %[[SIZE]], %[[C1]] : index
+// CHECK: %[[MAX1:.*]] = arith.maxsi %[[SUB1]], %[[C0]] : index
+// CHECK: %[[MIN1:.*]] = arith.minsi %[[MAX1]], %[[C1]] : index
+// CHECK: %[[MASK1:.*]] = vector.create_mask %[[MIN1]], %[[C8]], %[[C8]] : vector<1x8x8xi1>
+// CHECK: %[[INS1:.*]] = vector.insert_strided_slice %[[MASK1]], %[[INS0]] offsets = [1, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: return %[[INS1]] : vector<2x8x8xi1>
+
+// -----
+
func.func @shape_cast_1D(%v: vector<16xf32>) -> vector<2x2x4xf32> {
%0 = vector.shape_cast %v : vector<16xf32> to vector<2x2x4xf32>
return %0 : vector<2x2x4xf32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/234448
More information about the Mlir-commits
mailing list