[Mlir-commits] [mlir] Fix linalg masked contiguous extract (PR #206634)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 29 19:33:20 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core
Author: LouisLu060211
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/206634.diff
4 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (+9-3)
- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+13-7)
- (modified) mlir/test/Dialect/Linalg/transform-op-vectorize.mlir (+35)
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+23)
``````````diff
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/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/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
+ }
+}
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>
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/206634
More information about the Mlir-commits
mailing list