[Mlir-commits] [mlir] Fix linalg masked contiguous extract (PR #206634)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 29 19:50:52 PDT 2026
https://github.com/LouisLu060211 updated https://github.com/llvm/llvm-project/pull/206634
>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/3] 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/3] [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
+ }
+}
>From 5fd20e47a0e529caf15798f8fa7b8b73ec751dbc Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Tue, 30 Jun 2026 10:50:35 +0800
Subject: [PATCH 3/3] [mlir][linalg] Fix masking for rank-reducing contiguous
extracts
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 20 ++++++-----------
mlir/test/Transforms/remove-dead-values.mlir | 23 --------------------
2 files changed, 7 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index ae46630c949de..f0a210a2ededb 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -520,20 +520,14 @@ 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) -> Value {
- return createPoisonedValue(b, value);
+ return llvm::map_to_vector(values, [&](Value value) {
+ if (value.use_empty())
+ return Value();
+ return ub::PoisonOp::create(b, value.getLoc(), value.getType()).getResult();
});
}
@@ -695,9 +689,9 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
if (o.replaceWithPoison) {
rewriter.setInsertionPoint(o.op);
for (auto deadIdx : o.nonLive.set_bits()) {
- Value poisoned = createPoisonedValue(rewriter, o.op->getOperand(deadIdx));
- if (poisoned)
- o.op->setOperand(deadIdx, poisoned);
+ o.op->setOperand(
+ deadIdx, createPoisonedValues(rewriter, o.op->getOperand(deadIdx))
+ .front());
}
} 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 bbed2dc2e2feb..64088ce15cd48 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -868,26 +868,3 @@ 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>
- }
-}
More information about the Mlir-commits
mailing list