[Mlir-commits] [mlir] [MLIR][Linalg] Add static shape masking option in vectorization (PR #214812)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 7 11:56:12 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Artem Kroviakov (akroviakov)
<details>
<summary>Changes</summary>
This PR adds the counter option to `assume_dynamic_dims_match_vec_sizes` : that flag
lets a caller suppress masking a dynamic dim, this PR lets a caller request
masking a static one. When both apply, the explicit bound is preferred.
The PR adds an optional `inputMaskBounds` to `linalg::vectorize`, one entry per
iteration space dimension. A non-null entry forces masking of that dimension and
supplies the mask upper bound, regardless of what the static shape says.
---
Patch is 33.39 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/214812.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td (+26-11)
- (modified) mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h (+6-1)
- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+49-6)
- (modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (+94-8)
- (added) mlir/test/Dialect/Linalg/vectorization/mask-bounds.mlir (+311)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 09a6c4c4d652f..c9ceefd627ecb 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -110,7 +110,7 @@ def ApplyPadVectorizationPatternsOp : Op<Transform_Dialect,
These patterns rewrite tensor.pad Ops using vector.transfer_read and
vector.transfer_write operations. This is done either by:
1. Folding tensor.pad with an existing vector.transfer_read /
- vector.transfer_write Op (generated prior to running these patterns).
+ vector.transfer_write Op (generated prior to running these patterns).
2. Rewriting it (when matched together with q tensor.insert_slice
consumer Op) as a vector.transfer_read + vector.transfer_write pair.
@@ -594,7 +594,7 @@ def GeneralizeOp : Op<Transform_Dialect, "structured.generalize",
let arguments = (ins TransformHandleTypeInterface:$target);
let results = (outs TransformHandleTypeInterface:$transformed);
let assemblyFormat = [{
- $target attr-dict `:`
+ $target attr-dict `:`
custom<SemiFunctionType>(type($target), type($transformed), "false")
}];
@@ -634,7 +634,7 @@ def SpecializeOp : Op<Transform_Dialect, "structured.specialize",
DefaultValuedAttr<BoolAttr, "false">:$emit_category);
let results = (outs TransformHandleTypeInterface:$transformed);
let assemblyFormat = [{
- $target attr-dict `:`
+ $target attr-dict `:`
custom<SemiFunctionType>(type($target), type($transformed), "false")
}];
@@ -1265,7 +1265,7 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
TransformHandleTypeInterface:$copy);
let assemblyFormat = [{
- $target
+ $target
(`pad_to_multiple_of` custom<DynamicIndexList>($pad_to_multiple_of, $static_pad_to_multiple_of)^)?
(`use_prescribed_tensor_shapes` $use_prescribed_tensor_shapes^)?
attr-dict
@@ -1325,12 +1325,12 @@ def PadTilingInterfaceOp : Op<Transform_Dialect, "structured.pad_tiling_interfac
across the op semantics and, unlike for simple elementwise ops, may not be
trivially deducible or specifiable on operands only (e.g. convolutions).
Currently, only a limited set of projected permutation maps are supported.
-
+
The specification of `padding_sizes` follows that of `tile_sizes` during
tiling: the value "0" on a particular iterator encode "no padding". Like in
the case of tiling, an automatic completion by 0 to the operation rank
occurs.
-
+
This transformation returns a handle to the padded operation and to the
padding operation ("tensor.pad").
@@ -1341,8 +1341,8 @@ def PadTilingInterfaceOp : Op<Transform_Dialect, "structured.pad_tiling_interfac
This operation ignores non-IndexingMapOpInterface ops and drops them in the
return. In the future, this operation will support all TilingInterfaceOps
- for which the contract between iteration domain and operands can be
- reified.
+ for which the contract between iteration domain and operands can be
+ reified.
This operation may produce a definite failure if the padding fails for any
reason.
@@ -2556,7 +2556,8 @@ def VectorizeChildrenAndApplyPatternsOp :
}
def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
- [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+ [AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
TransformOpInterface, ReportTrackingListenerFailuresOpTrait]> {
let description = [{
Vectorize the target ops, which must be Linalg ops.
@@ -2581,6 +2582,16 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
Note: The input vector sizes must be bigger than or equal to their
counterpart iteration space sizes.
+ Optionally, `mask_bounds` forces masking of specific iteration space
+ dimensions even when those dimensions are statically sized.
+ `mask_bound_dims` lists the dimension each bound applies to:
+
+ ```mlir
+ // Mask iteration space dim 1 with %bound, despite a static shape.
+ transform.structured.vectorize %target vector_sizes [128, 64]
+ mask_bounds [1] (%bound : !transform.any_value) : !transform.any_op
+ ```
+
Typically this operator should be applied to linalg operations that have
already been tiled to the appropriate sizes.
@@ -2598,7 +2609,9 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
OptionalAttr<UnitAttr>:$vectorize_nd_extract,
OptionalAttr<UnitAttr>:$assume_dynamic_dims_match_vec_sizes,
OptionalAttr<UnitAttr>:$create_named_contraction,
- DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:$scalable_sizes);
+ DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:$scalable_sizes,
+ Variadic<TransformValueHandleTypeInterface>:$mask_bounds,
+ DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$mask_bound_dims);
let results = (outs);
@@ -2610,8 +2623,9 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
$vector_sizes,
$static_vector_sizes,
$scalable_sizes))
+ (`mask_bounds` $mask_bound_dims `(` $mask_bounds^ `:` type($mask_bounds) `)`)?
attr-dict
- `:` type($target)(`,`type($vector_sizes)^)?
+ `:` type($target)(`,`type($vector_sizes)^)?
}];
let hasVerifier = 1;
@@ -2627,6 +2641,7 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
}];
}
+
//===----------------------------------------------------------------------===//
// HoistRedundantVectorTransfersOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 130f7f6fd2792..2bf7edd809d7b 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1009,13 +1009,18 @@ struct VectorizationResult {
/// shapes.
/// Optionally, `createNamedContraction` can force compatible contractions to be
/// vectorized directly to vector.contract operation.
+/// `inputMaskBounds`, if provided, must match the rank of the iteration space.
+/// A non-null entry forces masking of the corresponding iteration space
+/// dimension using that value as the mask upper bound, even when the dimension
+/// is statically sized.
FailureOr<VectorizationResult>
vectorize(RewriterBase &rewriter, Operation *op,
ArrayRef<int64_t> inputVectorSizes = {},
ArrayRef<bool> inputScalableVecDims = {},
bool vectorizeNDExtract = false, bool flatten1DDepthwiseConv = false,
bool assumeDynamicDimsMatchVecSizes = false,
- bool createNamedContraction = false);
+ bool createNamedContraction = false,
+ ArrayRef<Value> inputMaskBounds = {});
/// Emit a suitable vector form for a Copy op with fully static shape.
LogicalResult vectorizeCopy(RewriterBase &builder, memref::CopyOp copyOp);
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 8c5bd50bd6f8b..98136089f410e 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4187,18 +4187,44 @@ DiagnosedSilenceableFailure transform::VectorizeOp::apply(
if (!status.succeeded())
return status;
+ SmallVector<Value> boundForDim;
+ for (auto [dim, handle] :
+ llvm::zip_equal(getMaskBoundDims(), getMaskBounds())) {
+ auto payloadValues = state.getPayloadValues(handle);
+ if (!llvm::hasSingleElement(payloadValues)) {
+ return emitSilenceableFailure(getLoc())
+ << "expected exactly one payload value for each mask bound";
+ }
+ if (static_cast<size_t>(dim) >= boundForDim.size())
+ boundForDim.resize(dim + 1);
+ boundForDim[dim] = *payloadValues.begin();
+ }
+
// TODO: Check that the correct number of vectorSizes was provided.
for (Operation *target : targets) {
if (!linalg::hasVectorizationImpl(target)) {
return mlir::emitSilenceableFailure(target->getLoc())
<< "Unsupported Op, cannot vectorize";
}
- FailureOr<VectorizationResult> vectorResults =
- linalg::vectorize(rewriter, target, vectorSizes, getScalableSizes(),
- getVectorizeNdExtract().value_or(false),
- /*flatten1DDepthwiseConv=*/false,
- getAssumeDynamicDimsMatchVecSizes().value_or(false),
- getCreateNamedContraction().value_or(false));
+ SmallVector<Value> maskBounds(boundForDim);
+ if (!maskBounds.empty()) {
+ if (auto linalgTarget = dyn_cast<linalg::LinalgOp>(target)) {
+ int64_t numLoops = linalgTarget.getNumLoops();
+ if (static_cast<int64_t>(maskBounds.size()) > numLoops) {
+ return mlir::emitSilenceableFailure(target->getLoc())
+ << "mask bound dim " << maskBounds.size() - 1
+ << " is out of range for an iteration space of rank "
+ << numLoops;
+ }
+ maskBounds.resize(numLoops);
+ }
+ }
+ FailureOr<VectorizationResult> vectorResults = linalg::vectorize(
+ rewriter, target, vectorSizes, getScalableSizes(),
+ getVectorizeNdExtract().value_or(false),
+ /*flatten1DDepthwiseConv=*/false,
+ getAssumeDynamicDimsMatchVecSizes().value_or(false),
+ getCreateNamedContraction().value_or(false), maskBounds);
if (failed(vectorResults)) {
return mlir::emitSilenceableFailure(target->getLoc())
<< "Attempted to vectorize, but failed";
@@ -4213,6 +4239,7 @@ void transform::VectorizeOp::getEffects(
SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
consumesHandle(getTargetMutable(), effects);
onlyReadsHandle(getVectorSizesMutable(), effects);
+ onlyReadsHandle(getMaskBoundsMutable(), effects);
modifiesPayload(effects);
}
@@ -4226,6 +4253,22 @@ LogicalResult transform::VectorizeOp::verify() {
return emitOpError("expected same number of vector sizes (")
<< getStaticVectorSizes().size() << ") and scalable sizes ("
<< getScalableSizes().size() << ")";
+ if (getMaskBounds().size() != getMaskBoundDims().size())
+ return emitOpError("expected same number of mask bounds (")
+ << getMaskBounds().size() << ") and mask bound dims ("
+ << getMaskBoundDims().size() << ")";
+ llvm::SmallDenseSet<int64_t> seenDims;
+ for (int64_t dim : getMaskBoundDims()) {
+ if (dim < 0)
+ return emitOpError("mask bound dim must be non-negative, got ") << dim;
+ if (!getStaticVectorSizes().empty() &&
+ dim >= static_cast<int64_t>(getStaticVectorSizes().size()))
+ return emitOpError("mask bound dim ")
+ << dim << " is out of range for " << getStaticVectorSizes().size()
+ << " vector sizes";
+ if (!seenDims.insert(dim).second)
+ return emitOpError("duplicate mask bound for dim ") << dim;
+ }
return success();
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 21ca3108efcd6..963a1ba87f1a5 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -33,6 +33,7 @@
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/STLExtras.h"
@@ -225,7 +226,8 @@ struct VectorizationState {
LogicalResult initState(RewriterBase &rewriter, LinalgOp linalgOp,
ArrayRef<int64_t> inputVectorSizes,
ArrayRef<bool> inputScalableVecDims,
- bool assumeDynamicDimsMatchVecSizes = false);
+ bool assumeDynamicDimsMatchVecSizes = false,
+ ArrayRef<Value> inputMaskBounds = {});
/// Returns the canonical vector shape used to vectorize the iteration space.
ArrayRef<int64_t> getCanonicalVecShape() const { return canonicalVecShape; }
@@ -342,6 +344,18 @@ struct VectorizationState {
/// shapes. Use this flag with care and only for cases where you are
/// confident the assumption holds.
bool assumeDynamicDimsMatchVecSizes = false;
+
+ /// Caller-provided mask upper bounds, one per iteration space dimension.
+ /// A non-null entry forces masking of that dimension with the
+ /// given bound. The opposite of `assumeDynamicDimsMatchVecSizes`:
+ /// it allows masking static dims.
+ SmallVector<Value> maskBounds;
+
+ /// Returns the caller-provided mask bound for `vecDim`, or a null Value if
+ /// none was provided.
+ Value getMaskBound(unsigned vecDim) const {
+ return vecDim < maskBounds.size() ? maskBounds[vecDim] : Value();
+ }
};
LogicalResult
@@ -349,6 +363,11 @@ VectorizationState::precomputeIterSpaceValueSizes(RewriterBase &rewriter,
LinalgOp linalgOp) {
// TODO: Support 0-d vectors.
for (int vecDim = 0, end = canonicalVecShape.size(); vecDim < end; ++vecDim) {
+ if (Value bound = getMaskBound(vecDim)) {
+ iterSpaceValueSizes.push_back(bound);
+ continue;
+ }
+
if (ShapedType::isStatic(iterSpaceStaticSizes[vecDim])) {
// Create constant index op for static dimensions.
iterSpaceValueSizes.push_back(arith::ConstantIndexOp::create(
@@ -383,8 +402,10 @@ LogicalResult VectorizationState::initState(RewriterBase &rewriter,
LinalgOp linalgOp,
ArrayRef<int64_t> inputVectorSizes,
ArrayRef<bool> inputScalableVecDims,
- bool assumeDimsMatchVec) {
+ bool assumeDimsMatchVec,
+ ArrayRef<Value> inputMaskBounds) {
assumeDynamicDimsMatchVecSizes = assumeDimsMatchVec;
+ maskBounds.assign(inputMaskBounds.begin(), inputMaskBounds.end());
// Initialize the insertion point.
rewriter.setInsertionPoint(linalgOp);
@@ -465,8 +486,20 @@ Value VectorizationState::getOrCreateMaskFor(
// operation.
// TODO: Improve this check. Only projected permutation indexing maps are
// supported.
+ //
+ // A mask bound explicitly requests masking. Re-route dims with a mask bound
+ // to dynamic dim handling. Track which bounded dims to not skip masking.
+ SmallVector<int64_t> effectiveStaticSizes(iterSpaceStaticSizes);
+ SmallVector<int64_t> isBoundedDim(iterSpaceStaticSizes.size(), 0);
+ for (auto [vecDim, size] : llvm::enumerate(effectiveStaticSizes)) {
+ if (getMaskBound(vecDim)) {
+ size = ShapedType::kDynamic;
+ isBoundedDim[vecDim] = 1;
+ }
+ }
+
SmallVector<int64_t> permutedStaticSizes =
- applyPermutationMap<int64_t>(maskingMap, iterSpaceStaticSizes);
+ applyPermutationMap<int64_t>(maskingMap, ArrayRef(effectiveStaticSizes));
auto maskType = getCanonicalVecType(rewriter.getI1Type(), maskingMap);
auto maskShape = maskType.getShape();
@@ -478,7 +511,11 @@ Value VectorizationState::getOrCreateMaskFor(
return Value();
}
- if (assumeDynamicDimsMatchVecSizes) {
+ // A bound on a not accessed says nothing about it.
+ bool hasMaskBound = llvm::is_contained(
+ applyPermutationMap<int64_t>(maskingMap, ArrayRef(isBoundedDim)), 1);
+
+ if (assumeDynamicDimsMatchVecSizes && !hasMaskBound) {
// While for _dynamic_ dim sizes we can _assume_ that the corresponding
// vector sizes match, we still need to check the _static_ dim sizes. Only
// then we can be 100% sure that masking is not required.
@@ -2488,6 +2525,49 @@ vectorizeScalableVectorPrecondition(Operation *op,
isa<linalg::BatchMmt4DOp>(op) || hasReductionIterator(linalgOp));
}
+/// Verify bounds: one per iteration space dimension, index-typed
+/// and not wider than their vector dim. Accept non-static boundaries.
+static LogicalResult
+vectorizeMaskBoundsPrecondition(Operation *op, ArrayRef<int64_t> vectorSizes,
+ ArrayRef<Value> inputMaskBounds) {
+ if (inputMaskBounds.empty())
+ return success();
+
+ auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+ if (!linalgOp) {
+ LDBG() << "Mask bounds are only supported for LinalgOp";
+ return failure();
+ }
+
+ if (inputMaskBounds.size() != linalgOp.getNumLoops()) {
+ LDBG() << "Mask bounds size (" << inputMaskBounds.size()
+ << ") does not match the iteration space rank ("
+ << linalgOp.getNumLoops() << ")";
+ return failure();
+ }
+
+ for (auto [dim, bound] : llvm::enumerate(inputMaskBounds)) {
+ if (!bound)
+ continue;
+ if (!bound.getType().isIndex()) {
+ LDBG() << "Mask bound for dim " << dim << " is not index-typed";
+ return failure();
+ }
+ if (dim >= vectorSizes.size())
+ continue;
+ FailureOr<int64_t> boundUb = ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::UB, bound, /*stopCondition=*/nullptr,
+ ValueBoundsOptions{/*closedUB=*/true});
+ if (succeeded(boundUb) && *boundUb > vectorSizes[dim]) {
+ LDBG() << "Mask bound for dim " << dim << " has upper bound " << *boundUb
+ << ", which exceeds the vector size " << vectorSizes[dim];
+ return failure();
+ }
+ }
+
+ return success();
+}
+
LogicalResult mlir::linalg::vectorizeOpPrecondition(
Operation *op, ArrayRef<int64_t> inputVectorSizes,
ArrayRef<bool> inputScalableVecDims, bool vectorizeNDExtract,
@@ -2545,7 +2625,7 @@ FailureOr<VectorizationResult> mlir::linalg::vectorize(
RewriterBase &rewriter, Operation *op, ArrayRef<int64_t> inputVectorSizes,
ArrayRef<bool> inputScalableVecDims, bool vectorizeNDExtract,
bool flatten1DDepthwiseConv, bool assumeDynamicDimsMatchVecSizes,
- bool createNamedContraction) {
+ bool createNamedContraction, ArrayRef<Value> inputMaskBounds) {
LDBG() << "Attempting to vectorize: " << *op;
LDBG() << "Input vector sizes: " << llvm::interleaved(inputVectorSizes);
LDBG() << "Input scalable vector dims: "
@@ -2558,12 +2638,18 @@ FailureOr<VectorizationResult> mlir::linalg::vectorize(
return failure();
}
+ if (failed(vectorizeMaskBoundsPrecondition(op, inputVectorSizes,
+ inputMaskBounds))) {
+ LDBG() << "Mask bounds pre-conditions failed";
+ return failure();
+ }
+
// Initialize vectorization state.
VectorizationState state(rewriter);
if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {
- if (failed(state.initState(rewriter, linalgOp, inputVectorSizes,
- inputScalableVecDims,
- assumeDynamicDimsMatchVecSizes))) {
+ if (failed(state.initState(
+ rewriter, linalgOp, inputVectorSizes, inputScalableVecDims,
+ assumeDynamicDimsMatchVecSizes, inputMaskBounds))) {
LDBG() << "Vectorization state couldn't be initialized";
return failure();
}
diff --git a/mlir/test/Dialect/Linalg/vectorization/mask-bounds.mlir b/mlir/test/Dialect/Linalg/vectorization/mask-bounds.mlir
new file mode 100644
index 0000000000000..9d73cb4396b94
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/vectorization/mask-bounds.mlir
@@ -0,0 +1,311 @@
+// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s
+
+///----------------------------------------------------------------------------------------
+/// Tests for `linalg::vectorize` with caller-provided mask bounds.
+///
+/// A mask bound forces masking of an iteration space dimension even when that
+/// dimension is statically sized. This is for operands that were padded to a
+/// static shape but whose trailing elements must not contribute to the result.
+/// It is the dual of `assume_dynamic_dims_match_vec_sizes`,
+/// which suppresses masking of a dynamic dim.
+///-----------...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/214812
More information about the Mlir-commits
mailing list