[Mlir-commits] [mlir] [MLIR][Linalg] Add static shape masking option in vectorization (PR #214812)
Artem Kroviakov
llvmlistbot at llvm.org
Wed Aug 12 08:20:25 PDT 2026
https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/214812
>From a870478baf1d39aa851972c0b8f1d94b78b0d455 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Wed, 12 Aug 2026 15:20:02 +0000
Subject: [PATCH] [MLIR][Linalg] Add static shape masking option in
vectorization
---
.../Linalg/TransformOps/LinalgTransformOps.td | 18 +-
.../Dialect/Linalg/Transforms/Transforms.h | 7 +-
.../TransformOps/LinalgTransformOps.cpp | 55 ++-
.../Linalg/Transforms/Vectorization.cpp | 102 ++++-
.../Linalg/vectorization/mask-bounds.mlir | 402 ++++++++++++++++++
5 files changed, 567 insertions(+), 17 deletions(-)
create mode 100644 mlir/test/Dialect/Linalg/vectorization/mask-bounds.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 09a6c4c4d652f..69686667c1f22 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -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,6 +2623,7 @@ 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)^)?
}];
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..b065a4c286731
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/vectorization/mask-bounds.mlir
@@ -0,0 +1,402 @@
+// 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.
+///----------------------------------------------------------------------------------------
+
+/// A statically shaped reduction, masked on the reduction dim (1) with a
+/// dynamic bound. Without the bound no mask would be created at all, since the
+/// operand shape matches the vector shape exactly - see
+/// @unmasked_static_reduction below.
+
+// CHECK-LABEL: func.func @masked_static_reduction(
+// CHECK-SAME: %[[SRC:.*]]: tensor<128x64xf32>,
+// CHECK-SAME: %[[ACC:.*]]: tensor<128xf32>,
+// CHECK-SAME: %[[UB:.*]]: index
+// CHECK: %[[BOUND:.*]] = affine.min
+// CHECK-DAG: %[[C128:.*]] = arith.constant 128 : index
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[C128]], %[[BOUND]] : vector<128x64xi1>
+// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] {
+// CHECK-SAME: vector.transfer_read %[[SRC]]
+// CHECK-SAME: } : vector<128x64xi1> -> vector<128x64xf32>
+// CHECK: vector.mask %[[MASK]] {
+// CHECK-SAME: vector.multi_reduction <maximumf>, %[[READ]]
+// CHECK-SAME: } : vector<128x64xi1> -> vector<128xf32>
+func.func @masked_static_reduction(%src: tensor<128x64xf32>,
+ %acc: tensor<128xf32>,
+ %ub: index) -> tensor<128xf32> {
+ // The bound is <= 64, so masking is meaningful (and provably in range).
+ %bound = affine.min affine_map<()[s0] -> (s0, 64)>()[%ub]
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%src : tensor<128x64xf32>) outs(%acc : tensor<128xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.maximumf %in, %out : f32
+ linalg.yield %1 : f32
+ } -> tensor<128xf32>
+ return %0 : tensor<128xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %min = transform.structured.match ops{["affine.min"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %bound = transform.get_result %min[0]
+ : (!transform.any_op) -> !transform.any_value
+ transform.structured.vectorize %generic vector_sizes [128, 64]
+ mask_bounds [1] (%bound : !transform.any_value) : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+/// Baseline: without a mask bound, a fully static reduction is vectorized
+/// unmasked. This is the behavior a mask bound opts out of.
+
+// CHECK-LABEL: func.func @unmasked_static_reduction(
+// CHECK-NOT: vector.create_mask
+// CHECK-NOT: vector.mask
+// CHECK: vector.multi_reduction <maximumf>
+func.func @unmasked_static_reduction(%src: tensor<128x64xf32>,
+ %acc: tensor<128xf32>) -> tensor<128xf32> {
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%src : tensor<128x64xf32>) outs(%acc : tensor<128xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.maximumf %in, %out : f32
+ linalg.yield %1 : f32
+ } -> tensor<128xf32>
+ return %0 : tensor<128xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %generic vector_sizes [128, 64] : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+/// A bound on the parallel dim (0) instead of the reduction dim. The bound
+/// lands in the leading position of the mask, and the accumulator read/write -
+/// which only maps dim 0 - is masked too.
+
+// CHECK-LABEL: func.func @masked_parallel_dim(
+// CHECK-SAME: %[[SRC:.*]]: tensor<128x64xf32>,
+// CHECK-SAME: %[[ACC:.*]]: tensor<128xf32>,
+// CHECK-SAME: %[[UB:.*]]: index
+// CHECK: %[[BOUND:.*]] = affine.min
+// CHECK-DAG: %[[C64:.*]] = arith.constant 64 : index
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[BOUND]], %[[C64]] : vector<128x64xi1>
+// CHECK: %[[ACC_MASK:.*]] = vector.create_mask %[[BOUND]] : vector<128xi1>
+// CHECK: vector.mask %[[MASK]] {
+// CHECK-SAME: vector.multi_reduction <add>
+func.func @masked_parallel_dim(%src: tensor<128x64xf32>,
+ %acc: tensor<128xf32>,
+ %ub: index) -> tensor<128xf32> {
+ %bound = affine.min affine_map<()[s0] -> (s0, 128)>()[%ub]
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%src : tensor<128x64xf32>) outs(%acc : tensor<128xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.addf %in, %out : f32
+ linalg.yield %1 : f32
+ } -> tensor<128xf32>
+ return %0 : tensor<128xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %min = transform.structured.match ops{["affine.min"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %bound = transform.get_result %min[0]
+ : (!transform.any_op) -> !transform.any_value
+ transform.structured.vectorize %generic vector_sizes [128, 64]
+ mask_bounds [0] (%bound : !transform.any_value) : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+/// Bounds on several dims at once, listed out of order to check that
+/// `mask_bound_dims` indexes the iteration space rather than the operand order.
+
+// CHECK-LABEL: func.func @masked_multiple_dims(
+// CHECK: %[[B0:.*]] = affine.min
+// CHECK: %[[B1:.*]] = affine.min
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[B0]], %[[B1]] : vector<128x64xi1>
+// CHECK: vector.mask %[[MASK]] { vector.multi_reduction <add>
+func.func @masked_multiple_dims(%src: tensor<128x64xf32>,
+ %acc: tensor<128xf32>,
+ %ub0: index, %ub1: index) -> tensor<128xf32> {
+ %b0 = affine.min affine_map<()[s0] -> (s0, 128)>()[%ub0]
+ %b1 = affine.min affine_map<()[s0] -> (s0, 64)>()[%ub1]
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%src : tensor<128x64xf32>) outs(%acc : tensor<128xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.addf %in, %out : f32
+ linalg.yield %1 : f32
+ } -> tensor<128xf32>
+ return %0 : tensor<128xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %mins = transform.structured.match ops{["affine.min"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %min0, %min1 = transform.split_handle %mins
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %b0 = transform.get_result %min0[0]
+ : (!transform.any_op) -> !transform.any_value
+ %b1 = transform.get_result %min1[0]
+ : (!transform.any_op) -> !transform.any_value
+ transform.structured.vectorize %generic vector_sizes [128, 64]
+ mask_bounds [1, 0] (%b1, %b0 : !transform.any_value, !transform.any_value)
+ : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+/// An elementwise op (no reduction) with a bound: the bound reaches the
+/// transfer_read and transfer_write, not just a reduction.
+
+// CHECK-LABEL: func.func @masked_elementwise(
+// CHECK: %[[BOUND:.*]] = affine.min
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[BOUND]] : vector<128xi1>
+// CHECK: vector.mask %[[MASK]] { vector.transfer_read
+// CHECK: vector.mask %[[MASK]] { vector.transfer_write
+func.func @masked_elementwise(%src: tensor<128xf32>, %ub: index) -> tensor<128xf32> {
+ %bound = affine.min affine_map<()[s0] -> (s0, 128)>()[%ub]
+ %empty = tensor.empty() : tensor<128xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]
+ } ins(%src : tensor<128xf32>) outs(%empty : tensor<128xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.negf %in : f32
+ linalg.yield %1 : f32
+ } -> tensor<128xf32>
+ return %0 : tensor<128xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %min = transform.structured.match ops{["affine.min"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %bound = transform.get_result %min[0]
+ : (!transform.any_op) -> !transform.any_value
+ transform.structured.vectorize %generic vector_sizes [128]
+ mask_bounds [0] (%bound : !transform.any_value) : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+/// A bound on a dim that is *already* dynamic. The bound wins over the
+/// `tensor.dim`-derived size that would otherwise be inferred.
+
+// CHECK-LABEL: func.func @masked_bound_overrides_dynamic_dim(
+// CHECK: %[[BOUND:.*]] = affine.min
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[BOUND]] : vector<128xi1>
+// CHECK-NOT: vector.create_mask
+// CHECK: vector.mask %[[MASK]]
+func.func @masked_bound_overrides_dynamic_dim(%src: tensor<?xf32>,
+ %ub: index) -> tensor<?xf32> {
+ %c0 = arith.constant 0 : index
+ %bound = affine.min affine_map<()[s0] -> (s0, 128)>()[%ub]
+ %dim = tensor.dim %src, %c0 : tensor<?xf32>
+ %empty = tensor.empty(%dim) : tensor<?xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]
+ } ins(%src : tensor<?xf32>) outs(%empty : tensor<?xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.negf %in : f32
+ linalg.yield %1 : f32
+ } -> tensor<?xf32>
+ return %0 : tensor<?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %min = transform.structured.match ops{["affine.min"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %bound = transform.get_result %min[0]
+ : (!transform.any_op) -> !transform.any_value
+ transform.structured.vectorize %generic vector_sizes [128]
+ mask_bounds [0] (%bound : !transform.any_value) : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+/// A bound takes precedence over `assume_dynamic_dims_match_vec_sizes`, which
+/// would otherwise suppress masking entirely. Dim 1 carries an explicit bound
+/// and so is still masked. The precedence is per masking map, not per op: the
+/// accumulator only projects onto dim 0, which has no bound, so it keeps the
+/// unmasked in-bounds form the assumption asks for.
+
+// CHECK-LABEL: func.func @masked_bound_beats_assume_dynamic_dims(
+// CHECK: %[[BOUND:.*]] = affine.min
+// CHECK: %[[DIM:.*]] = tensor.dim
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[DIM]], %[[BOUND]] : vector<128x64xi1>
+// CHECK: vector.mask %[[MASK]] { vector.transfer_read
+// The accumulator read is left unmasked and in-bounds, not wrapped in a mask.
+// CHECK: vector.transfer_read {{.*}}in_bounds = [true]{{.*}} tensor<?xf32>
+// CHECK-NOT: vector.create_mask
+// CHECK: vector.mask %[[MASK]] { vector.multi_reduction <maximumf>
+func.func @masked_bound_beats_assume_dynamic_dims(%src: tensor<?x64xf32>,
+ %acc: tensor<?xf32>,
+ %ub: index) -> tensor<?xf32> {
+ %bound = affine.min affine_map<()[s0] -> (s0, 64)>()[%ub]
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%src : tensor<?x64xf32>) outs(%acc : tensor<?xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.maximumf %in, %out : f32
+ linalg.yield %1 : f32
+ } -> tensor<?xf32>
+ return %0 : tensor<?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %min = transform.structured.match ops{["affine.min"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %bound = transform.get_result %min[0]
+ : (!transform.any_op) -> !transform.any_value
+ transform.structured.vectorize %generic vector_sizes [128, 64]
+ mask_bounds [1] (%bound : !transform.any_value)
+ {assume_dynamic_dims_match_vec_sizes} : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+/// A fixed-shape op (`linalg.matmul`) whose result
+/// feeds two reductions with *different* neutral values. `maximumf` needs -inf,
+/// `minimumf` needs +inf. Padding the operands up to the instruction's fixed shape
+/// satisfies the shape requirement, but no single padding value satisfies both
+/// neutral values, so the reductions have to be masked to the number of valid
+/// columns instead.
+
+// CHECK-LABEL: func.func @masked_reductions_after_matmul(
+/// The valid extent, captured before padding.
+// CHECK: %[[DIM:.*]] = tensor.dim
+// CHECK: %[[BOUND:.*]] = affine.min
+// CHECK: %[[PADDED:.*]] = tensor.pad
+/// The matmul is statically shaped, so it is vectorized without any masking.
+// CHECK: vector.transfer_read %[[PADDED]]
+// CHECK: %[[MUL:.*]] = arith.mulf
+// CHECK: %[[MM:.*]] = vector.multi_reduction <add>, %[[MUL]]
+// CHECK: %[[MM_RES:.*]] = vector.transfer_write %[[MM]]
+// CHECK: %[[MASK_MAX:.*]] = vector.create_mask %{{.*}}, %[[BOUND]] : vector<128x64xi1>
+// CHECK: %[[READ_MAX:.*]] = vector.mask %[[MASK_MAX]] {
+// CHECK-SAME: vector.transfer_read %[[MM_RES]]
+// CHECK-SAME: } : vector<128x64xi1> -> vector<128x64xf32>
+// CHECK: vector.mask %[[MASK_MAX]] {
+// CHECK-SAME: vector.multi_reduction <maximumf>, %[[READ_MAX]]
+// CHECK-SAME: } : vector<128x64xi1> -> vector<128xf32>
+// CHECK: %[[MASK_MIN:.*]] = vector.create_mask %{{.*}}, %[[BOUND]] : vector<128x64xi1>
+// CHECK: %[[READ_MIN:.*]] = vector.mask %[[MASK_MIN]] {
+// CHECK-SAME: vector.transfer_read %[[MM_RES]]
+// CHECK-SAME: } : vector<128x64xi1> -> vector<128x64xf32>
+// CHECK: vector.mask %[[MASK_MIN]] {
+// CHECK-SAME: vector.multi_reduction <minimumf>, %[[READ_MIN]]
+// CHECK-SAME: } : vector<128x64xi1> -> vector<128xf32>
+func.func @masked_reductions_after_matmul(%A: tensor<128x32xf32>,
+ %B_dyn: tensor<32x?xf32>,
+ %C: tensor<128x64xf32>,
+ %max_init: tensor<128xf32>,
+ %sum_init: tensor<128xf32>)
+ -> (tensor<128xf32>, tensor<128xf32>) {
+ %c1 = arith.constant 1 : index
+ %pad_val = arith.constant 0.0 : f32
+ %dim = tensor.dim %B_dyn, %c1 : tensor<32x?xf32>
+ %extent = affine.min affine_map<()[s0] -> (s0, 64)>()[%dim]
+ %high = affine.apply affine_map<()[s0] -> (64 - s0)>()[%extent]
+ %B = tensor.pad %B_dyn low[0, 0] high[0, %high] {
+ ^bb0(%i: index, %j: index):
+ tensor.yield %pad_val : f32
+ } : tensor<32x?xf32> to tensor<32x64xf32>
+ %mm = linalg.matmul ins(%A, %B : tensor<128x32xf32>, tensor<32x64xf32>)
+ outs(%C : tensor<128x64xf32>) -> tensor<128x64xf32>
+ %max = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%mm : tensor<128x64xf32>) outs(%max_init : tensor<128xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %m = arith.maximumf %in, %out : f32
+ linalg.yield %m : f32
+ } -> tensor<128xf32>
+ %sum = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%mm : tensor<128x64xf32>) outs(%sum_init : tensor<128xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %s = arith.minimumf %in, %out : f32
+ linalg.yield %s : f32
+ } -> tensor<128xf32>
+ return %max, %sum : tensor<128xf32>, tensor<128xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %reductions = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %min = transform.structured.match ops{["affine.min"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %extent = transform.get_result %min[0]
+ : (!transform.any_op) -> !transform.any_value
+ // The fixed-shape op needs no bound: its shapes are static and complete.
+ transform.structured.vectorize %matmul vector_sizes [128, 64, 32]
+ : !transform.any_op
+ // The reductions must only cover the valid columns.
+ transform.structured.vectorize %reductions vector_sizes [128, 64]
+ mask_bounds [1] (%extent : !transform.any_value) : !transform.any_op
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list