[Mlir-commits] [mlir] Fix linalg masked contiguous extract (PR #206634)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 29 19:32:33 PDT 2026
https://github.com/LouisLu060211 created https://github.com/llvm/llvm-project/pull/206634
Fix masked vectorization of contiguous tensor.extract when the source tensor rank is smaller than the surrounding loop rank.
The contiguous-load tensor.extract path builds a rank-reducing vector.transfer_read by reading only the trailing source dimensions and broadcasting leading vector dimensions. Previously this path returned the transfer_read unmasked and relied on the generic masking path, which builds a mask over the full linalg iteration space.
That full-rank mask is invalid when the source rank is smaller than the loop rank. For example, a 1-D source read inside a 2-D loop nest can produce a vector.transfer_read with inferred mask type vector<4xi1>, while the generic path wraps it with vector<1x4xi1>.
Mask the contiguous transfer_read in place using a minor-identity masking map over the dimensions actually read. This matches the transfer_read permutation map and preserves the existing full-rank behavior, where the minor-identity map is equivalent to the full identity mask.
Add a regression test covering a masked contiguous tensor.extract from a 1-D source inside a 2-D linalg.generic vectorized with vectorize_nd_extract. The test checks that the transfer_read is wrapped with a rank-reduced vector<4xi1> mask.
Fixes #206209.
Test:
- llvm-lit -a -vv mlir/test/Dialect/Linalg/transform-op-vectorize.mlir
Assisted-by: Codex
>From 8159ee2b4e0188e999056a09dd075da0d79ac568 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Wed, 24 Jun 2026 15:09:36 +0800
Subject: [PATCH 1/2] Fix crash in vectorizeOpPrecondition when vector sizes
array is too short
[MLIR][Vector] Add regression test for bug #204100 (mixed static/dynamic vector sizes)
Remove accidentally created file vectorize-dynamic-mixed-sizes.mlir
Remove mlir_venv/ from .gitignore (directory already deleted)
Fix comment about regression test location in Vectorization.cpp
Correct the path in the comment from .../Linalg/transform/... to .../Linalg/...
(the test actually lives directly under Linalg). Also update the
related vector-to-llvm test for consistency.
remove venv
Remove lldb/python_api from .gitignore
Fix emitSilenceableFailure to use target->getLoc()
Fixing comments
Fix a.mlir crash
[mlir][RemoveDeadValues] Fix crash replacing dead operand with poison
updated
updated minor changes
minnor issues
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 20 +++++++++++------
mlir/test/Transforms/remove-dead-values.mlir | 23 ++++++++++++++++++++
2 files changed, 36 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index f0a210a2ededb..ae46630c949de 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -520,14 +520,20 @@ static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,
}
}
+/// Create ub.poison ops for the given values. If a value has no uses, return
+/// an "empty" value.
+static Value createPoisonedValue(OpBuilder &b, Value value) {
+ if (!value || value.use_empty())
+ return Value();
+ return ub::PoisonOp::create(b, value.getLoc(), value.getType()).getResult();
+}
+
/// Create ub.poison ops for the given values. If a value has no uses, return
/// an "empty" value.
static SmallVector<Value> createPoisonedValues(OpBuilder &b,
ValueRange values) {
- return llvm::map_to_vector(values, [&](Value value) {
- if (value.use_empty())
- return Value();
- return ub::PoisonOp::create(b, value.getLoc(), value.getType()).getResult();
+ return llvm::map_to_vector(values, [&](Value value) -> Value {
+ return createPoisonedValue(b, value);
});
}
@@ -689,9 +695,9 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
if (o.replaceWithPoison) {
rewriter.setInsertionPoint(o.op);
for (auto deadIdx : o.nonLive.set_bits()) {
- o.op->setOperand(
- deadIdx, createPoisonedValues(rewriter, o.op->getOperand(deadIdx))
- .front());
+ Value poisoned = createPoisonedValue(rewriter, o.op->getOperand(deadIdx));
+ if (poisoned)
+ o.op->setOperand(deadIdx, poisoned);
}
} else {
o.op->eraseOperands(o.nonLive);
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 64088ce15cd48..bbed2dc2e2feb 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -868,3 +868,26 @@ module @func_with_non_call_users {
}
spirv.EntryPoint "GLCompute" @callee
}
+
+// -----
+
+// CHECK: pdl_interp.func private @matcher()
+// CHECK-LABEL: func.func private @callee()
+// CHECK: return
+module {
+ pdl_interp.func private @matcher(%arg0: !llvm.ptr) {
+ pdl_interp.finalize
+ }
+ module @rewriters {
+ }
+ func.func private @callee(%arg0: memref<f32>) -> memref<f32> {
+ %false = arith.constant false
+ %0 = scf.if %false -> (memref<f32>) {
+ scf.yield %arg0 : memref<f32>
+ } else {
+ %1 = bufferization.clone %arg0 : memref<f32> to memref<f32>
+ scf.yield %1 : memref<f32>
+ }
+ return %0 : memref<f32>
+ }
+}
>From 03f89a2465231c78de5a4da03fb4647346827849 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Tue, 30 Jun 2026 10:27:01 +0800
Subject: [PATCH 2/2] [mlir][linalg] Fix masking for rank-reducing contiguous
extracts
---
.../Linalg/Transforms/Vectorization.cpp | 12 +++++--
.../Linalg/transform-op-vectorize.mlir | 35 +++++++++++++++++++
2 files changed, 44 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index b57e66a1c3580..035409a436c19 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1278,9 +1278,10 @@ vectorizeTensorExtract(RewriterBase &rewriter, VectorizationState &state,
}
// 2b. Handle contiguous access.
+ int64_t numReadDims = std::min(dstRank, srcRank);
auto permutationMap = AffineMap::getMinorIdentityMap(
- srcRank, std::min(dstRank, srcRank), rewriter.getContext());
-
+ srcRank, numReadDims, rewriter.getContext());
+
int32_t rankDiff = dstRank - srcRank;
// When dstRank > srcRank, broadcast the source tensor to the unitary leading
// dims so that the ranks match. This is done by extending the map with 0s.
@@ -1299,9 +1300,14 @@ vectorizeTensorExtract(RewriterBase &rewriter, VectorizationState &state,
rewriter, loc, resultType, extractOp.getTensor(), transferReadIdxs,
/*padding=*/std::nullopt, permutationMap, inBounds);
+ auto maskingMap = AffineMap::getMinorIdentityMap(
+ linalgOp.getNumLoops(), numReadDims, rewriter.getContext());
+ Operation *readOrMaskedReadOp =
+ state.maskOperation(rewriter, transferReadOp, linalgOp, maskingMap);
+
LDBG() << "Vectorised as contiguous load: " << extractOp;
return VectorizationHookResult{VectorizationHookStatus::NewOp,
- transferReadOp};
+ readOrMaskedReadOp};
}
/// Emit reduction operations if the shapes of the value to reduce is different
diff --git a/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir b/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir
index 0d59dbba8940d..e6facdb37bfb4 100644
--- a/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir
@@ -190,3 +190,38 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// CHECK-DAG: #[[$MAP:.*]] = affine_map<(d0) -> (0, d0)>
+// CHECK-LABEL: func.func @masked_contiguous_extract_rank_reducing_mask
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[C3]] : vector<4xi1>
+// CHECK: vector.mask %[[MASK]] { vector.transfer_read {{.*}} permutation_map = #[[$MAP]]} : tensor<16xf32>, vector<1x4xf32> } : vector<4xi1> -> vector<1x4xf32>
+func.func @masked_contiguous_extract_rank_reducing_mask(
+ %src: tensor<16xf32>,
+ %output : tensor<1x3xf32>,
+ %idx: index) -> tensor<1x3xf32> {
+ %1 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]
+ } outs(%output : tensor<1x3xf32>) {
+ ^bb0(%out: f32):
+ %2 = linalg.index 1 : index
+ %3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%2, %idx)
+ %extracted = tensor.extract %src[%3] : tensor<16xf32>
+ linalg.yield %extracted : f32
+ } -> tensor<1x3xf32>
+ return %1 : tensor<1x3xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %0 vector_sizes [1, 4] {vectorize_nd_extract}
+ : !transform.any_op
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list