[Mlir-commits] [mlir] [mlir][vector] Support fixed-size masks in `eliminateVectorMasks` (PR #221595)
Dhairyashil R G
llvmlistbot at llvm.org
Sun Sep 6 11:25:03 PDT 2026
https://github.com/dhairyashilRG created https://github.com/llvm/llvm-project/pull/221595
`eliminateVectorMasks` returned immediately unless a `VscaleRange` was supplied. The TODO above that early return argued the fixed-size case was unlikely to be useful, because "for fixed-size code dimensions are all static so masks tend to fold away".
That is only true when the mask operand is constant. When it comes from an induction variable, which is what tiling produces, the operand is not constant and no fold removes the mask, even though value-bounds analysis can prove it all-true. For example, `-canonicalize` leaves this alone:
```mlir
scf.for %i = %c0 to %c1024 step %c4 {
%rem = arith.subi %c1024, %i : index
%mask = vector.create_mask %rem : vector<4xi1>
%v = vector.transfer_read %t[%i], %f0, %mask : tensor<1024xf32>, vector<4xf32>
...
}
```
`%i <= 1020`, so `%rem >= 4` and the mask is all-true.
The algorithm was never scalable-specific; only the bound query was. This makes `vscaleRange` genuinely optional: without one, the lower bound is computed with `ValueBoundsConstraintSet::computeConstantBound` and compared against the mask dimension. Scalable dimensions still need a range and bail without one, so a fixed-size caller cannot prove a dimension whose runtime size is `vscale` times the size in the type. The scalable path is unchanged, and the existing test still passes with its RUN line untouched.
### Behaviour change for downstream callers
An omitted `vscaleRange` used to make the call a no-op. It now runs the fixed-size analysis.
No in-tree caller relies on the old meaning, but IREE does call this unconditionally in `Codegen/Common/GenericVectorization.cpp`, passing `getDefaultVscaleRange()`, which returns `std::nullopt` for every target other than AArch64+SVE. For those targets this patch turns a no-op into live mask elimination. That is the intended benefit, but it is worth knowing before the next integrate. The added cost is bounded: constant and `vscale`-multiple operands short-circuit before any analysis, so only `create_mask` ops with unknown operands pay for a bound query.
### Tests
New `eliminate-masks-fixed-size.mlir`, run through a `fixed-size` option added to the test pass: two cases that are eliminated (the tiled loop above, and a dynamic operand bounded by `arith.maxsi`) and three that are not (an unbounded operand, a lower bound smaller than the mask dimension, and a scalable dimension with no vscale range). Each negative was checked to fail for the stated reason, by making the smallest change that flips it.
Context: https://discourse.llvm.org/t/rfc-should-vector-transfer-read-write-keep-in-bounds-measurements-on-the-masking-alternative/91649
Assisted-by: Claude
>From 13240544aa0abfdb946ee55f3d545a35b8058969 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Sun, 6 Sep 2026 23:16:56 +0530
Subject: [PATCH] [mlir][vector] Support fixed-size masks in
eliminateVectorMasks
`eliminateVectorMasks` returned immediately unless a `VscaleRange` was
supplied, under a TODO arguing the fixed-size case was unlikely to be
useful because "for fixed-size code dimensions are all static so masks
tend to fold away".
That holds only when the mask operand is constant. When it is derived
from an induction variable, which is the shape tiling produces, the
operand is not constant and no fold removes the mask, even though
value-bounds analysis can prove it all-true. For a `vector<4xi1>` mask of
`1024 - %i` inside `scf.for %i = 0 to 1024 step 4`, canonicalization
leaves the `create_mask` alone.
The algorithm itself was never scalable-specific; only the bound query
was. Make `vscaleRange` genuinely optional: without one, compute an
ordinary constant lower bound with `ValueBoundsConstraintSet` and compare
it against the mask dimension. Scalable dimensions still require a range
and bail without one, so a fixed-size caller cannot prove a dimension
whose runtime size is `vscale` times the size in the type. The scalable
path is unchanged.
Note this changes what an omitted `vscaleRange` means. It previously made
the call a no-op; it now runs the fixed-size analysis. No in-tree caller
relies on the old meaning, but downstream ones may: IREE calls this
unconditionally in `GenericVectorization.cpp`, passing a range that is
`std::nullopt` for every target other than AArch64+SVE.
Assisted-by: Claude
---
.../Vector/Transforms/VectorTransforms.h | 9 +-
.../Transforms/VectorMaskElimination.cpp | 38 +++++---
.../Vector/eliminate-masks-fixed-size.mlir | 97 +++++++++++++++++++
.../Dialect/Vector/TestVectorTransforms.cpp | 10 +-
4 files changed, 135 insertions(+), 19 deletions(-)
create mode 100644 mlir/test/Dialect/Vector/eliminate-masks-fixed-size.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
index e815e026305fa..e88058fe1ff70 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
@@ -124,11 +124,10 @@ struct VscaleRange {
/// Attempts to eliminate redundant vector masks by replacing them with all-true
/// constants at the top of the function (which results in the masks folding
-/// away). Note: Currently, this only runs for vector.create_mask ops and
-/// requires `vscaleRange`. If `vscaleRange` is not provided this transform does
-/// nothing. This is because these redundant masks are much more likely for
-/// scalable code which requires memref/tensor dynamic sizes, whereas fixed-size
-/// code has static sizes, so simpler folds remove the masks.
+/// away). Note: Currently, this only runs for vector.create_mask ops.
+///
+/// `vscaleRange` is only needed to prove that a scalable mask dimension is
+/// all-true; without it, only fixed-size dimensions are considered.
void eliminateVectorMasks(IRRewriter &rewriter, FunctionOpInterface function,
std::optional<VscaleRange> vscaleRange = {});
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp
index 6f75ce7a04511..f8f0ac8e33183 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp
@@ -15,11 +15,14 @@ using namespace mlir;
using namespace mlir::vector;
namespace {
-/// Attempts to resolve a (scalable) CreateMaskOp to an all-true constant mask.
-/// All-true masks can then be eliminated by simple folds.
-LogicalResult resolveAllTrueCreateMaskOp(IRRewriter &rewriter,
- vector::CreateMaskOp createMaskOp,
- VscaleRange vscaleRange) {
+/// Attempts to resolve a CreateMaskOp to an all-true constant mask. All-true
+/// masks can then be eliminated by simple folds. `vscaleRange` is required to
+/// reason about scalable dimensions; without it only fixed-size dimensions can
+/// be proven all-true.
+LogicalResult
+resolveAllTrueCreateMaskOp(IRRewriter &rewriter,
+ vector::CreateMaskOp createMaskOp,
+ std::optional<VscaleRange> vscaleRange) {
auto maskType = createMaskOp.getVectorType();
auto maskTypeDimScalableFlags = maskType.getScalableDims();
auto maskTypeDimSizes = maskType.getShape();
@@ -54,9 +57,25 @@ LogicalResult resolveAllTrueCreateMaskOp(IRRewriter &rewriter,
for (auto [i, dimSize] : unknownDims) {
// Compute the lower bound for the unknown dimension (i.e. the smallest
// value it could be).
+ if (!vscaleRange) {
+ // A constant bound cannot prove a scalable dim, whose runtime size is
+ // `vscale` times the size in the type. Checked first to skip the query.
+ if (maskTypeDimScalableFlags[i])
+ return failure();
+ FailureOr<int64_t> dimLowerBound =
+ ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::LB, dimSize);
+ if (failed(dimLowerBound))
+ return failure();
+ // If LB < the mask dim size then this dim is not all-true.
+ if (*dimLowerBound < maskTypeDimSizes[i])
+ return failure();
+ continue;
+ }
+
FailureOr<ConstantOrScalableBound> dimLowerBound =
vector::ScalableValueBoundsConstraintSet::computeScalableBound(
- dimSize, {}, vscaleRange.vscaleMin, vscaleRange.vscaleMax,
+ dimSize, {}, vscaleRange->vscaleMin, vscaleRange->vscaleMax,
presburger::BoundType::LB);
if (failed(dimLowerBound))
return failure();
@@ -94,11 +113,6 @@ namespace mlir::vector {
void eliminateVectorMasks(IRRewriter &rewriter, FunctionOpInterface function,
std::optional<VscaleRange> vscaleRange) {
- // TODO: Support fixed-size case. This is less likely to be useful as for
- // fixed-size code dimensions are all static so masks tend to fold away.
- if (!vscaleRange)
- return;
-
// Early exit for functions without a body.
if (function.isExternal())
return;
@@ -114,7 +128,7 @@ void eliminateVectorMasks(IRRewriter &rewriter, FunctionOpInterface function,
rewriter.setInsertionPointToStart(&function.front());
for (auto mask : worklist)
- (void)resolveAllTrueCreateMaskOp(rewriter, mask, *vscaleRange);
+ (void)resolveAllTrueCreateMaskOp(rewriter, mask, vscaleRange);
}
} // namespace mlir::vector
diff --git a/mlir/test/Dialect/Vector/eliminate-masks-fixed-size.mlir b/mlir/test/Dialect/Vector/eliminate-masks-fixed-size.mlir
new file mode 100644
index 0000000000000..192285ccc512a
--- /dev/null
+++ b/mlir/test/Dialect/Vector/eliminate-masks-fixed-size.mlir
@@ -0,0 +1,97 @@
+// RUN: mlir-opt %s -split-input-file -test-eliminate-vector-masks=fixed-size | FileCheck %s
+
+// The mask bound comes from the induction variable, so it is not constant and
+// no fold removes it. Value-bounds analysis proves `%i <= 1020`, hence
+// `1024 - %i >= 4`, so the mask is all-true.
+
+// CHECK-LABEL: @eliminate_mask_bounded_by_induction_variable
+// CHECK: %[[ALL_TRUE:.*]] = vector.constant_mask [4] : vector<4xi1>
+// CHECK: vector.transfer_read {{.*}}, %[[ALL_TRUE]]
+func.func @eliminate_mask_bounded_by_induction_variable(%t: tensor<1024xf32>) -> f32 {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c1024 = arith.constant 1024 : index
+ %f0 = arith.constant 0.0 : f32
+ %r = scf.for %i = %c0 to %c1024 step %c4 iter_args(%acc = %f0) -> f32 {
+ %rem = arith.subi %c1024, %i : index
+ %mask = vector.create_mask %rem : vector<4xi1>
+ %v = vector.transfer_read %t[%i], %f0, %mask : tensor<1024xf32>, vector<4xf32>
+ %s = vector.reduction <add>, %v : vector<4xf32> into f32
+ %n = arith.addf %acc, %s : f32
+ scf.yield %n : f32
+ }
+ return %r : f32
+}
+
+// -----
+
+// `max(%n, 4) >= 4` for any `%n`, so the mask is all-true even though the
+// operand is dynamic.
+
+// CHECK-LABEL: @eliminate_mask_with_dynamic_operand
+// CHECK: %[[ALL_TRUE:.*]] = vector.constant_mask [4] : vector<4xi1>
+// CHECK: vector.transfer_read {{.*}}, %[[ALL_TRUE]]
+func.func @eliminate_mask_with_dynamic_operand(%t: tensor<?xf32>, %n: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %f0 = arith.constant 0.0 : f32
+ %m = arith.maxsi %n, %c4 : index
+ %mask = vector.create_mask %m : vector<4xi1>
+ %v = vector.transfer_read %t[%c0], %f0, %mask : tensor<?xf32>, vector<4xf32>
+ return %v : vector<4xf32>
+}
+
+// -----
+
+// Nothing bounds `%n` from below, so the mask cannot be proven all-true.
+
+// CHECK-LABEL: @negative_unbounded_operand
+// CHECK-NOT: vector.constant_mask
+// CHECK: vector.create_mask
+func.func @negative_unbounded_operand(%t: tensor<?xf32>, %n: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %f0 = arith.constant 0.0 : f32
+ %mask = vector.create_mask %n : vector<4xi1>
+ %v = vector.transfer_read %t[%c0], %f0, %mask : tensor<?xf32>, vector<4xf32>
+ return %v : vector<4xf32>
+}
+
+// -----
+
+// The lower bound of `max(%n, 2)` is 2, which is less than the mask dimension,
+// so the mask is not all-true. Unlike the unbounded case above, a bound *is*
+// computed here; it is the comparison against the mask size that rejects it.
+
+// CHECK-LABEL: @negative_lower_bound_too_small
+// CHECK-NOT: vector.constant_mask
+// CHECK: vector.create_mask
+func.func @negative_lower_bound_too_small(%t: tensor<?xf32>, %n: index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %f0 = arith.constant 0.0 : f32
+ %m = arith.maxsi %n, %c2 : index
+ %mask = vector.create_mask %m : vector<4xi1>
+ %v = vector.transfer_read %t[%c0], %f0, %mask : tensor<?xf32>, vector<4xf32>
+ return %v : vector<4xf32>
+}
+
+// -----
+
+// `max(%n, 4 * vscale) >= 4 * vscale`, which is exactly the runtime size of
+// `vector<[4]xi1>`, so with a vscale range this mask is provably all-true.
+// Without one a scalable dimension cannot be proven, so it must be left alone.
+
+// CHECK-LABEL: @negative_scalable_dim_without_vscale_range
+// CHECK-NOT: vector.constant_mask
+// CHECK: vector.create_mask
+func.func @negative_scalable_dim_without_vscale_range(%t: tensor<?xf32>, %n: index) -> vector<[4]xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %f0 = arith.constant 0.0 : f32
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %vscale, %c4 : index
+ %m = arith.maxsi %n, %c4_vscale : index
+ %mask = vector.create_mask %m : vector<[4]xi1>
+ %v = vector.transfer_read %t[%c0], %f0, %mask : tensor<?xf32>, vector<[4]xf32>
+ return %v : vector<[4]xf32>
+}
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index c0f8a85206cdf..773cb5e29c1a9 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -1055,6 +1055,10 @@ struct TestEliminateVectorMasks
Option<unsigned> vscaleMax{
*this, "vscale-max", llvm::cl::desc("Maximum possible value of vscale."),
llvm::cl::init(16)};
+ Option<bool> fixedSize{
+ *this, "fixed-size",
+ llvm::cl::desc("Run without a vscale range, as fixed-size code would."),
+ llvm::cl::init(false)};
StringRef getArgument() const final { return "test-eliminate-vector-masks"; }
StringRef getDescription() const final {
@@ -1062,8 +1066,10 @@ struct TestEliminateVectorMasks
}
void runOnOperation() override {
IRRewriter rewriter(&getContext());
- eliminateVectorMasks(rewriter, getOperation(),
- VscaleRange{vscaleMin, vscaleMax});
+ std::optional<VscaleRange> vscaleRange;
+ if (!fixedSize)
+ vscaleRange = VscaleRange{vscaleMin, vscaleMax};
+ eliminateVectorMasks(rewriter, getOperation(), vscaleRange);
}
};
More information about the Mlir-commits
mailing list