[Mlir-commits] [mlir] [mlir][vector] support unrolling vector.contract with scalar result (PR #232203)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 20 03:04:35 PDT 2026
https://github.com/StarOne01 updated https://github.com/llvm/llvm-project/pull/232203
>From d39e80a2063885a0b755d8d81456c0df5a5587f6 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Thu, 15 Oct 2026 12:39:41 +0530
Subject: [PATCH 1/3] [mlir][vector] add a check to UnrollContractionPattern to
bail out if we have a scalar result and add contract to scalar test case
---
.../Dialect/Vector/Transforms/VectorUnroll.cpp | 3 +++
.../Dialect/Vector/vector-unroll-options.mlir | 17 +++++++++++++++++
2 files changed, 20 insertions(+)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 19abe44081eb43..21ea96e0b0b818 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -290,6 +290,9 @@ struct UnrollContractionPattern
LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
PatternRewriter &rewriter) const override {
+ if (!isa<VectorType>(contractOp.getResultType()))
+ return rewriter.notifyMatchFailure(contractOp,
+ "scalar result not supported");
auto targetShape = getTargetShape(options, contractOp);
if (!targetShape)
return failure();
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 61a74744e7a72e..22fcb483ea1e57 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -935,3 +935,20 @@ func.func @masked_vector_multi_reduction(%v : vector<4x6xf32>, %acc: vector<4xf3
// CHECK-LABEL: func @masked_vector_multi_reduction
// CHECK-NEXT: %[[R:.*]] = vector.mask %{{.*}} { vector.multi_reduction <add>, {{.*}} : vector<4x6xf32> to vector<4xf32> }
// CHECK-NEXT: return %[[R]]
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+#map0 = affine_map<(d0) -> ()>
+
+func.func @vector_contract_to_scalar(%a: vector<4xf32>, %b: vector<4xf32>, %acc: f32) -> f32 {
+ %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4xf32>, vector<4xf32> into f32
+ return %r : f32
+}
+
+// CHECK-LABEL: func @vector_contract_to_scalar(
+// CHECK-SAME: %[[A:.*]]: vector<4xf32>, %[[B:.*]]: vector<4xf32>, %[[ACC:.*]]: f32) -> f32 {
+// CHECK-NOT: vector.extract_strided_slice
+// CHECK: %[[R:.*]] = vector.contract
+// CHECK-SAME: %[[A]], %[[B]], %[[ACC]] : vector<4xf32>, vector<4xf32> into f32
+// CHECK: return %[[R]] : f32
>From e4fd0afc90af3edc237a605e0deb0dd63a49a2f7 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sun, 18 Oct 2026 08:59:04 +0530
Subject: [PATCH 2/3] [mlir][vector] enable full unrolling when we have a
scalar result and add tests for 1dxf32, 2dxf32, 1dxi32, 1dxf16
---
.../Vector/Transforms/VectorUnroll.cpp | 25 ++++----
.../Dialect/Vector/vector-unroll-options.mlir | 64 ++++++++++++++++---
2 files changed, 69 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 21ea96e0b0b818..1608c365531284 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -152,7 +152,6 @@ getUnrollOrder(unsigned numLoops, Operation *op,
}
namespace {
-
struct UnrollTransferReadPattern
: public OpRewritePattern<vector::TransferReadOp> {
UnrollTransferReadPattern(MLIRContext *context,
@@ -290,13 +289,11 @@ struct UnrollContractionPattern
LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
PatternRewriter &rewriter) const override {
- if (!isa<VectorType>(contractOp.getResultType()))
- return rewriter.notifyMatchFailure(contractOp,
- "scalar result not supported");
+ bool scalarResult = !isa<VectorType>(contractOp.getResultType());
auto targetShape = getTargetShape(options, contractOp);
if (!targetShape)
return failure();
- auto dstVecType = cast<VectorType>(contractOp.getResultType());
+ auto dstVecType = dyn_cast<VectorType>(contractOp.getResultType());
SmallVector<int64_t> originalSize = *contractOp.getIterationShapeToUnroll();
Location loc = contractOp.getLoc();
@@ -346,21 +343,28 @@ struct UnrollContractionPattern
auto *accIt = accCache.find(accOffets);
if (accIt != accCache.end())
slicesOperands[2] = accIt->second;
+ else if (scalarResult)
+ slicesOperands[2] = contractOp.getAcc();
else
extractOperand(2, contractOp.getAcc(), accPermutationMap, accOffets);
-
- SmallVector<int64_t> dstShape =
- applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(*targetShape));
- auto targetType = VectorType::get(dstShape, dstVecType.getElementType());
+ Type targetType = contractOp.getResultType();
+ if (!scalarResult) {
+ SmallVector<int64_t> dstShape =
+ applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(*targetShape));
+ targetType = VectorType::get(dstShape, dstVecType.getElementType());
+ }
Operation *newOp = cloneOpWithOperandsAndTypes(
rewriter, loc, contractOp, slicesOperands, targetType);
-
SmallVector<int64_t> dstOffets =
applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(offsets));
// Save the accumulated value untill all the loops are unrolled since
// reduction loop keep updating the accumulator.
accCache[dstOffets] = newOp->getResult(0);
}
+ if (scalarResult) {
+ rewriter.replaceOp(contractOp, accCache.begin()->second);
+ return success();
+ }
// Assemble back the accumulator into a single vector.
Value result = arith::ConstantOp::create(rewriter, loc, dstVecType,
rewriter.getZeroAttr(dstVecType));
@@ -1731,7 +1735,6 @@ struct UnrollDeinterleavePattern
private:
vector::UnrollVectorOptions options;
};
-
} // namespace
void mlir::vector::populateVectorUnrollPatterns(
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 22fcb483ea1e57..4129942e5bd3e8 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -940,15 +940,61 @@ func.func @masked_vector_multi_reduction(%v : vector<4x6xf32>, %acc: vector<4xf3
#map = affine_map<(d0) -> (d0)>
#map0 = affine_map<(d0) -> ()>
-
-func.func @vector_contract_to_scalar(%a: vector<4xf32>, %b: vector<4xf32>, %acc: f32) -> f32 {
- %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4xf32>, vector<4xf32> into f32
- return %r : f32
+module {
+ func.func @vector_contract_to_scalar_f32_from_1d(%a: vector<4xf32>, %b: vector<4xf32>, %acc: f32) -> f32 {
+ %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4xf32>, vector<4xf32> into f32
+ return %r : f32
+ }
}
-// CHECK-LABEL: func @vector_contract_to_scalar(
+// CHECK-LABEL: func @vector_contract_to_scalar_f32_from_1d(
// CHECK-SAME: %[[A:.*]]: vector<4xf32>, %[[B:.*]]: vector<4xf32>, %[[ACC:.*]]: f32) -> f32 {
-// CHECK-NOT: vector.extract_strided_slice
-// CHECK: %[[R:.*]] = vector.contract
-// CHECK-SAME: %[[A]], %[[B]], %[[ACC]] : vector<4xf32>, vector<4xf32> into f32
-// CHECK: return %[[R]] : f32
+// CHECK: %[[A0:.*]] = vector.extract_strided_slice %[[A]] offsets = [0], sizes = [2], strides = [1] : vector<4xf32> to vector<2xf32>
+// CHECK: %[[B0:.*]] = vector.extract_strided_slice %[[B]] offsets = [0], sizes = [2], strides = [1] : vector<4xf32> to vector<2xf32>
+// CHECK: %[[R0:.*]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ACC]] : vector<2xf32>, vector<2xf32> into f32
+// CHECK: %[[A1:.*]] = vector.extract_strided_slice %[[A]] offsets = [2], sizes = [2], strides = [1] : vector<4xf32> to vector<2xf32>
+// CHECK: %[[B1:.*]] = vector.extract_strided_slice %[[B]] offsets = [2], sizes = [2], strides = [1] : vector<4xf32> to vector<2xf32>
+// CHECK: %[[R1:.*]] = vector.contract {{.*}} %[[A1]], %[[B1]], %[[R0]] : vector<2xf32>, vector<2xf32> into f32
+// CHECK: return %[[R1]] : f32
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+#map0 = affine_map<(d0) -> ()>
+module {
+ func.func @vector_contract_to_scalar_f16(%a: vector<8xf16>, %b: vector<8xf16>, %acc: f16) -> f16 {
+ %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<8xf16>, vector<8xf16> into f16
+ return %r : f16
+ }
+}
+
+// CHECK-LABEL: func @vector_contract_to_scalar_f16(
+// CHECK-SAME: %[[A:.*]]: vector<8xf16>, %[[B:.*]]: vector<8xf16>, %[[ACC:.*]]: f16) -> f16 {
+// CHECK: %[[A0:.*]] = vector.extract_strided_slice %[[A]] offsets = [0], sizes = [4], strides = [1] : vector<8xf16> to vector<4xf16>
+// CHECK: %[[B0:.*]] = vector.extract_strided_slice %[[B]] offsets = [0], sizes = [4], strides = [1] : vector<8xf16> to vector<4xf16>
+// CHECK: %[[R0:.*]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ACC]] : vector<4xf16>, vector<4xf16> into f16
+// CHECK: %[[A1:.*]] = vector.extract_strided_slice %[[A]] offsets = [4], sizes = [4], strides = [1] : vector<8xf16> to vector<4xf16>
+// CHECK: %[[B1:.*]] = vector.extract_strided_slice %[[B]] offsets = [4], sizes = [4], strides = [1] : vector<8xf16> to vector<4xf16>
+// CHECK: %[[R1:.*]] = vector.contract {{.*}} %[[A1]], %[[B1]], %[[R0]] : vector<4xf16>, vector<4xf16> into f16
+// CHECK: return %[[R1]] : f16
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map0 = affine_map<(d0,d1) -> ()>
+module {
+ func.func @vector_contract_to_scalar_f32_from_2d(%a: vector<4x4xf32>, %b: vector<4x4xf32>, %acc: f32) -> f32 {
+ %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction", "reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4x4xf32>, vector<4x4xf32> into f32
+ return %r : f32
+ }
+}
+
+// CHECK-LABEL: func @vector_contract_to_scalar_f32_from_2d(
+// CHECK-SAME: %[[A:.*]]: vector<4x4xf32>, %[[B:.*]]: vector<4x4xf32>, %[[ACC:.*]]: f32) -> f32 {
+// CHECK: %[[A0:.*]] = vector.extract_strided_slice %[[A]] offsets = [0, 0], sizes = [4, 2], strides = [1, 1] : vector<4x4xf32> to vector<4x2xf32>
+// CHECK: %[[B0:.*]] = vector.extract_strided_slice %[[B]] offsets = [0, 0], sizes = [4, 2], strides = [1, 1] : vector<4x4xf32> to vector<4x2xf32>
+// CHECK: %[[R0:.*]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ACC]] : vector<4x2xf32>, vector<4x2xf32> into f32
+// CHECK: %[[A1:.*]] = vector.extract_strided_slice %[[A]] offsets = [0, 2], sizes = [4, 2], strides = [1, 1] : vector<4x4xf32> to vector<4x2xf32>
+// CHECK: %[[B1:.*]] = vector.extract_strided_slice %[[B]] offsets = [0, 2], sizes = [4, 2], strides = [1, 1] : vector<4x4xf32> to vector<4x2xf32>
+// CHECK: %[[R1:.*]] = vector.contract {{.*}} %[[A1]], %[[B1]], %[[R0]] : vector<4x2xf32>, vector<4x2xf32> into f32
+// CHECK: return %[[R1]] : f32
>From 20c83b316c05d279e553c4285e77007ad41be2fe Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Mon, 19 Oct 2026 13:22:42 +0530
Subject: [PATCH 3/3] [mlir][vector] refactor `UnrollContractionPattern` to
improve scalar result handling and update test cases for indexing maps as per
the review
---
.../Vector/Transforms/VectorUnroll.cpp | 8 +++--
.../Dialect/Vector/vector-unroll-options.mlir | 33 +++++++++----------
2 files changed, 20 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 1608c365531284..4f431a47a1945c 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -289,11 +289,12 @@ struct UnrollContractionPattern
LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
PatternRewriter &rewriter) const override {
- bool scalarResult = !isa<VectorType>(contractOp.getResultType());
+ Type targetType = contractOp.getResultType();
+ bool scalarResult = !isa<VectorType>(targetType);
auto targetShape = getTargetShape(options, contractOp);
if (!targetShape)
return failure();
- auto dstVecType = dyn_cast<VectorType>(contractOp.getResultType());
+ auto dstVecType = dyn_cast<VectorType>(targetType);
SmallVector<int64_t> originalSize = *contractOp.getIterationShapeToUnroll();
Location loc = contractOp.getLoc();
@@ -347,7 +348,6 @@ struct UnrollContractionPattern
slicesOperands[2] = contractOp.getAcc();
else
extractOperand(2, contractOp.getAcc(), accPermutationMap, accOffets);
- Type targetType = contractOp.getResultType();
if (!scalarResult) {
SmallVector<int64_t> dstShape =
applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(*targetShape));
@@ -362,6 +362,8 @@ struct UnrollContractionPattern
accCache[dstOffets] = newOp->getResult(0);
}
if (scalarResult) {
+ // With a scalar result, the last tile's value replaces the op
+ // and we return since there is nothing to reassemble.
rewriter.replaceOp(contractOp, accCache.begin()->second);
return success();
}
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 4129942e5bd3e8..231c4390c378a5 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -939,12 +939,11 @@ func.func @masked_vector_multi_reduction(%v : vector<4x6xf32>, %acc: vector<4xf3
// -----
#map = affine_map<(d0) -> (d0)>
-#map0 = affine_map<(d0) -> ()>
-module {
- func.func @vector_contract_to_scalar_f32_from_1d(%a: vector<4xf32>, %b: vector<4xf32>, %acc: f32) -> f32 {
- %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4xf32>, vector<4xf32> into f32
- return %r : f32
- }
+#acc = affine_map<(d0) -> ()>
+
+func.func @vector_contract_to_scalar_f32_from_1d(%a: vector<4xf32>, %b: vector<4xf32>, %acc: f32) -> f32 {
+ %r = vector.contract {indexing_maps = [#map, #map, #acc], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4xf32>, vector<4xf32> into f32
+ return %r : f32
}
// CHECK-LABEL: func @vector_contract_to_scalar_f32_from_1d(
@@ -960,12 +959,11 @@ module {
// -----
#map = affine_map<(d0) -> (d0)>
-#map0 = affine_map<(d0) -> ()>
-module {
- func.func @vector_contract_to_scalar_f16(%a: vector<8xf16>, %b: vector<8xf16>, %acc: f16) -> f16 {
- %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<8xf16>, vector<8xf16> into f16
- return %r : f16
- }
+#acc = affine_map<(d0) -> ()>
+
+func.func @vector_contract_to_scalar_f16(%a: vector<8xf16>, %b: vector<8xf16>, %acc: f16) -> f16 {
+ %r = vector.contract {indexing_maps = [#map, #map, #acc], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<8xf16>, vector<8xf16> into f16
+ return %r : f16
}
// CHECK-LABEL: func @vector_contract_to_scalar_f16(
@@ -981,12 +979,11 @@ module {
// -----
#map = affine_map<(d0, d1) -> (d0, d1)>
-#map0 = affine_map<(d0,d1) -> ()>
-module {
- func.func @vector_contract_to_scalar_f32_from_2d(%a: vector<4x4xf32>, %b: vector<4x4xf32>, %acc: f32) -> f32 {
- %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction", "reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4x4xf32>, vector<4x4xf32> into f32
- return %r : f32
- }
+#acc = affine_map<(d0,d1) -> ()>
+
+func.func @vector_contract_to_scalar_f32_from_2d(%a: vector<4x4xf32>, %b: vector<4x4xf32>, %acc: f32) -> f32 {
+ %r = vector.contract {indexing_maps = [#map, #map, #acc], iterator_types = ["reduction", "reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4x4xf32>, vector<4x4xf32> into f32
+ return %r : f32
}
// CHECK-LABEL: func @vector_contract_to_scalar_f32_from_2d(
More information about the Mlir-commits
mailing list