[Mlir-commits] [mlir] [mlir][vector] Allow 0-D vector acc/res in `vector.contract` (PR #234638)
Federico Bruzzone
llvmlistbot at llvm.org
Thu Oct 22 23:25:57 PDT 2026
https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/234638
>From 8f0a6234ae4fa6aefb695fbd71e2bc7674b05d2d Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Thu, 22 Oct 2026 16:02:57 +0200
Subject: [PATCH 1/2] [mlir][vector] Allow 0-D vector acc/res in
`vector.contract`
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 3 +-
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 11 ++--
.../Vector/Transforms/LowerVectorContract.cpp | 14 +++--
.../vectorization/contraction-interface.mlir | 54 +++++++++++++++++++
mlir/test/Dialect/Vector/invalid.mlir | 10 ++++
mlir/test/Dialect/Vector/ops.mlir | 9 ++++
.../vector-contract-to-dot-transforms.mlir | 30 +++++++++++
7 files changed, 124 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 437d0f7c4ee029..9608e135829980 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -72,7 +72,8 @@ def Vector_ContractionOp :
intermediate result to the accumulator argument of rank K, and returns a
vector result of rank K (where K = num_lhs_free_dims + num_rhs_free_dims +
num_batch_dims (see dimension type descriptions below)). For K = 0 (no
- free or batch dimensions), the accumulator and output are a scalar.
+ free or batch dimensions), the accumulator and output are either both a
+ scalar or both a 0-D vector.
If operands and the result have types of different bitwidths, operands are
promoted to have the same bitwidth as the result before performing the
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 490202f2dc9778..058bd7f91a0162 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1081,9 +1081,14 @@ static LogicalResult verifyOutputShape(
// Verify 'expectedResultDims'.
if (expectedResultDims.empty()) {
- // No batch or free dimension implies a scalar result.
- if (llvm::isa<VectorType>(resType) || llvm::isa<VectorType>(accType))
- return op.emitOpError("invalid accumulator/result vector shape");
+ // No batch or free dimension implies a scalar or 0-D vector result. The
+ // accumulator and the result must agree.
+ if (resType != accType)
+ return op.emitOpError("expected accumulator and result to have the same "
+ "type");
+ if (auto resVectorType = llvm::dyn_cast<VectorType>(resType))
+ if (resVectorType.getRank() != 0)
+ return op.emitOpError("invalid accumulator/result vector shape");
} else {
// At least one batch or free dimension implies a vector result.
auto resVectorType = llvm::dyn_cast<VectorType>(resType);
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index b3cd4fb1e8f9b0..464fb6e6154863 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -1095,10 +1095,11 @@ FailureOr<Value> ContractionOpLowering::lowerReduction(
VectorType lhsType = op.getLhsType();
VectorType rhsType = op.getRhsType();
Type resType = op.getResultType();
- if (isa<VectorType>(resType))
+ auto resVectorType = dyn_cast<VectorType>(resType);
+ if (resVectorType && resVectorType.getRank() != 0)
return rewriter.notifyMatchFailure(op,
"did not expect a VectorType result");
- bool isInt = isa<IntegerType>(resType);
+ bool isInt = isa<IntegerType>(getElementTypeOrSelf(resType));
// Use iterator index 0.
int64_t iterIndex = 0;
SmallVector<AffineMap> iMap = op.getIndexingMapsArray();
@@ -1130,12 +1131,19 @@ FailureOr<Value> ContractionOpLowering::lowerReduction(
auto kind = vector::CombiningKind::ADD;
Value acc = op.getAcc();
+ // vector.reduction takes a scalar accumulator.
+ if (resVectorType)
+ acc = vector::ExtractOp::create(rewriter, loc, acc, ArrayRef<int64_t>{});
Operation *reductionOp =
acc ? vector::ReductionOp::create(rewriter, loc, kind, m, acc,
op.getFastmath())
: vector::ReductionOp::create(rewriter, loc, kind, m,
op.getFastmath());
- return maskOperation(rewriter, reductionOp, mask)->getResult(0);
+ Value reduced = maskOperation(rewriter, reductionOp, mask)->getResult(0);
+ if (resVectorType)
+ reduced =
+ vector::BroadcastOp::create(rewriter, loc, resVectorType, reduced);
+ return reduced;
}
// Construct new iterator types and affine map array attribute.
std::array<AffineMap, 3> lowIndexingMaps = {
diff --git a/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir b/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
index 198bb4a5eae7a2..e7c62d110f5070 100644
--- a/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
@@ -39,6 +39,60 @@ module attributes {transform.with_named_sequence} {
// -----
+func.func @dot(%A: tensor<8xf32>, %B: tensor<8xf32>,
+ %C: tensor<f32>) -> tensor<f32> {
+ %0 = linalg.dot
+ ins(%A, %B : tensor<8xf32>, tensor<8xf32>)
+ outs(%C: tensor<f32>) -> tensor<f32>
+ return %0 : tensor<f32>
+}
+
+// CHECK-LABEL: func.func @dot(
+// CHECK-SAME: %[[A:.*]]: tensor<8xf32>, %[[B:.*]]: tensor<8xf32>,
+// CHECK-SAME: %[[C:.*]]: tensor<f32>)
+// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %[[A]]{{.*}}: tensor<8xf32>, vector<8xf32>
+// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %[[B]]{{.*}}: tensor<8xf32>, vector<8xf32>
+// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %[[C]]{{.*}}: tensor<f32>, vector<f32>
+// CHECK: %[[CONTRACT:.*]] = vector.contract
+// CHECK-SAME: kind = #vector.kind<add>
+// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]] : vector<8xf32>, vector<8xf32> into vector<f32>
+// CHECK: vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<f32>, tensor<f32>
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.dot"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %0 create_named_contraction : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @dot_dynamic(%A: tensor<?xf32>, %B: tensor<?xf32>,
+ %C: tensor<f32>) -> tensor<f32> {
+ %0 = linalg.dot
+ ins(%A, %B : tensor<?xf32>, tensor<?xf32>)
+ outs(%C: tensor<f32>) -> tensor<f32>
+ return %0 : tensor<f32>
+}
+
+// CHECK-LABEL: func.func @dot_dynamic(
+// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %{{.*}}[]{{.*}}: tensor<f32>, vector<f32>
+// CHECK: %[[MASK:.*]] = vector.create_mask %{{.*}} : vector<8xi1>
+// CHECK: %[[CONTRACT:.*]] = vector.mask %[[MASK]] { vector.contract
+// CHECK-SAME: %[[LOAD_C]] : vector<8xf32>, vector<8xf32> into vector<f32> } : vector<8xi1> -> vector<f32>
+// CHECK: vector.transfer_write %[[CONTRACT]], %{{.*}}[] : vector<f32>, tensor<f32>
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.dot"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %0 vector_sizes [8] create_named_contraction : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
func.func @matmul_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
%C: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.matmul
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index c44bc2e8f07f68..64004f51d49a07 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -961,6 +961,16 @@ func.func @contraction(%arg0: vector<7x8x16x15xf32>, %arg1: vector<8x16x7x5xf32>
// -----
+func.func @contraction_scalar_acc_0d_result(%arg0: vector<10xf32>, %arg1: vector<10xf32>, %arg2: f32) {
+ // expected-error at +1 {{expected accumulator and result to have the same type}}
+ %0 = "vector.contract"(%arg0, %arg1, %arg2) <{
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+ iterator_types = [#vector.iterator_type<reduction>],
+ kind = #vector.kind<add>}> : (vector<10xf32>, vector<10xf32>, f32) -> vector<f32>
+}
+
+// -----
+
#contraction_accesses = [
affine_map<(b0, f0, f1, c0, c1) -> (c0, b0, c1, f0)>,
affine_map<(b0, f0, f1, c0, c1) -> (b0, c1, c0, f1)>,
diff --git a/mlir/test/Dialect/Vector/ops.mlir b/mlir/test/Dialect/Vector/ops.mlir
index bd123bd0a6bd83..016c7a344c4a90 100644
--- a/mlir/test/Dialect/Vector/ops.mlir
+++ b/mlir/test/Dialect/Vector/ops.mlir
@@ -340,6 +340,15 @@ func.func @extract_strided_slice_scalable(%arg0: vector<4x[8]x16xf32>) -> vector
indexing_maps = #contraction_to_scalar_accesses,
iterator_types = ["reduction"]
}
+// CHECK-LABEL: @contraction_to_0d_vector
+func.func @contraction_to_0d_vector(%arg0: vector<10xf32>, %arg1: vector<10xf32>, %arg2: vector<f32>) -> vector<f32> {
+ // CHECK: %[[X:.*]] = vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["reduction"], kind = #vector.kind<add>} %{{.*}}, %{{.*}}, %{{.*}} : vector<10xf32>, vector<10xf32> into vector<f32>
+ %0 = vector.contract #contraction_to_scalar_trait %arg0, %arg1, %arg2
+ : vector<10xf32>, vector<10xf32> into vector<f32>
+ // CHECK: return %[[X]] : vector<f32>
+ return %0 : vector<f32>
+}
+
// CHECK-LABEL: @contraction_to_scalar
func.func @contraction_to_scalar(%arg0: vector<10xf32>, %arg1: vector<10xf32>) -> f32 {
// CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f32
diff --git a/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir b/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
index d00fe588f2b5bf..98642c9be9de67 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
@@ -36,6 +36,36 @@ func.func @masked_extract_contract1(%arg0: vector<4xf32>, %arg1: vector<4xf32>,
return %0 : f32
}
+// CHECK-LABEL: func @extract_contract1_0d_acc
+// CHECK-SAME: %[[A:.*0]]: vector<4xf32>,
+// CHECK-SAME: %[[B:.*1]]: vector<4xf32>,
+// CHECK-SAME: %[[C:.*2]]: vector<f32>
+// CHECK: %[[F:.*]] = arith.mulf %[[A]], %[[B]] : vector<4xf32>
+// CHECK: %[[ACC:.*]] = vector.extract %[[C]][] : f32 from vector<f32>
+// CHECK: %[[R:.*]] = vector.reduction <add>, %[[F]], %[[ACC]] : vector<4xf32> into f32
+// CHECK: %[[B0:.*]] = vector.broadcast %[[R]] : f32 to vector<f32>
+// CHECK: return %[[B0]] : vector<f32>
+
+func.func @extract_contract1_0d_acc(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: vector<f32>) -> vector<f32> {
+ %0 = vector.contract #dotp_trait %arg0, %arg1, %arg2
+ : vector<4xf32>, vector<4xf32> into vector<f32>
+ return %0 : vector<f32>
+}
+
+// CHECK-LABEL: func @masked_extract_contract1_0d_acc
+// CHECK-SAME: %[[A:.*0]]: vector<4xf32>, %[[B:.*1]]: vector<4xf32>, %[[C:.*2]]: vector<f32>
+// CHECK-SAME: %[[M:.*]]: vector<4xi1>
+// CHECK: %[[F:.*]] = arith.mulf %[[A]], %[[B]] : vector<4xf32>
+// CHECK: %[[ACC:.*]] = vector.extract %[[C]][] : f32 from vector<f32>
+// CHECK: %[[R:.*]] = vector.mask %[[M]] { vector.reduction <add>, %[[F]], %[[ACC]] : vector<4xf32> into f32 } : vector<4xi1> -> f32
+// CHECK: %[[B0:.*]] = vector.broadcast %[[R]] : f32 to vector<f32>
+// CHECK: return %[[B0]] : vector<f32>
+
+func.func @masked_extract_contract1_0d_acc(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: vector<f32>, %mask: vector<4xi1>) -> vector<f32> {
+ %0 = vector.mask %mask { vector.contract #dotp_trait %arg0, %arg1, %arg2 : vector<4xf32>, vector<4xf32> into vector<f32> } : vector<4xi1> -> vector<f32>
+ return %0 : vector<f32>
+}
+
// CHECK-LABEL: func @extract_contract1_int
// CHECK-SAME: %[[A:.*0]]: vector<4xi32>,
// CHECK-SAME: %[[B:.*1]]: vector<4xi32>,
>From c16d70b082373f88538d9d1a711730dd037e12c2 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Fri, 23 Oct 2026 08:25:03 +0200
Subject: [PATCH 2/2] Address comments from Prashanth
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td | 4 ++--
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 6 +++---
mlir/test/Dialect/Vector/invalid.mlir | 6 +++---
mlir/test/Dialect/Vector/ops.mlir | 4 ++--
.../Dialect/Vector/vector-contract-to-dot-transforms.mlir | 8 ++++----
5 files changed, 14 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 9608e135829980..c200da4112b081 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -72,8 +72,8 @@ def Vector_ContractionOp :
intermediate result to the accumulator argument of rank K, and returns a
vector result of rank K (where K = num_lhs_free_dims + num_rhs_free_dims +
num_batch_dims (see dimension type descriptions below)). For K = 0 (no
- free or batch dimensions), the accumulator and output are either both a
- scalar or both a 0-D vector.
+ free or batch dimensions), the accumulator and output are either a scalar
+ or a 0-D vector.
If operands and the result have types of different bitwidths, operands are
promoted to have the same bitwidth as the result before performing the
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 058bd7f91a0162..ca185b8baba41d 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1084,11 +1084,11 @@ static LogicalResult verifyOutputShape(
// No batch or free dimension implies a scalar or 0-D vector result. The
// accumulator and the result must agree.
if (resType != accType)
- return op.emitOpError("expected accumulator and result to have the same "
- "type");
+ return op.emitOpError("expected same accumulator and result type");
if (auto resVectorType = llvm::dyn_cast<VectorType>(resType))
if (resVectorType.getRank() != 0)
- return op.emitOpError("invalid accumulator/result vector shape");
+ return op.emitOpError("expected scalar or 0-D vector result when "
+ "there are no batch or free dimensions");
} else {
// At least one batch or free dimension implies a vector result.
auto resVectorType = llvm::dyn_cast<VectorType>(resType);
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 64004f51d49a07..1992f0408e01f3 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -961,9 +961,9 @@ func.func @contraction(%arg0: vector<7x8x16x15xf32>, %arg1: vector<8x16x7x5xf32>
// -----
-func.func @contraction_scalar_acc_0d_result(%arg0: vector<10xf32>, %arg1: vector<10xf32>, %arg2: f32) {
- // expected-error at +1 {{expected accumulator and result to have the same type}}
- %0 = "vector.contract"(%arg0, %arg1, %arg2) <{
+func.func @contraction_scalar_acc_0d_result(%arg0: vector<10xf32>, %arg1: vector<10xf32>, %acc: f32) {
+ // expected-error at +1 {{expected same accumulator and result type}}
+ %0 = "vector.contract"(%arg0, %arg1, %acc) <{
indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
iterator_types = [#vector.iterator_type<reduction>],
kind = #vector.kind<add>}> : (vector<10xf32>, vector<10xf32>, f32) -> vector<f32>
diff --git a/mlir/test/Dialect/Vector/ops.mlir b/mlir/test/Dialect/Vector/ops.mlir
index 016c7a344c4a90..c9d038dfc9f802 100644
--- a/mlir/test/Dialect/Vector/ops.mlir
+++ b/mlir/test/Dialect/Vector/ops.mlir
@@ -341,9 +341,9 @@ func.func @extract_strided_slice_scalable(%arg0: vector<4x[8]x16xf32>) -> vector
iterator_types = ["reduction"]
}
// CHECK-LABEL: @contraction_to_0d_vector
-func.func @contraction_to_0d_vector(%arg0: vector<10xf32>, %arg1: vector<10xf32>, %arg2: vector<f32>) -> vector<f32> {
+func.func @contraction_to_0d_vector(%arg0: vector<10xf32>, %arg1: vector<10xf32>, %acc: vector<f32>) -> vector<f32> {
// CHECK: %[[X:.*]] = vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["reduction"], kind = #vector.kind<add>} %{{.*}}, %{{.*}}, %{{.*}} : vector<10xf32>, vector<10xf32> into vector<f32>
- %0 = vector.contract #contraction_to_scalar_trait %arg0, %arg1, %arg2
+ %0 = vector.contract #contraction_to_scalar_trait %arg0, %arg1, %acc
: vector<10xf32>, vector<10xf32> into vector<f32>
// CHECK: return %[[X]] : vector<f32>
return %0 : vector<f32>
diff --git a/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir b/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
index 98642c9be9de67..87ab74493ea9ad 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
@@ -46,8 +46,8 @@ func.func @masked_extract_contract1(%arg0: vector<4xf32>, %arg1: vector<4xf32>,
// CHECK: %[[B0:.*]] = vector.broadcast %[[R]] : f32 to vector<f32>
// CHECK: return %[[B0]] : vector<f32>
-func.func @extract_contract1_0d_acc(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: vector<f32>) -> vector<f32> {
- %0 = vector.contract #dotp_trait %arg0, %arg1, %arg2
+func.func @extract_contract1_0d_acc(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %acc: vector<f32>) -> vector<f32> {
+ %0 = vector.contract #dotp_trait %arg0, %arg1, %acc
: vector<4xf32>, vector<4xf32> into vector<f32>
return %0 : vector<f32>
}
@@ -61,8 +61,8 @@ func.func @extract_contract1_0d_acc(%arg0: vector<4xf32>, %arg1: vector<4xf32>,
// CHECK: %[[B0:.*]] = vector.broadcast %[[R]] : f32 to vector<f32>
// CHECK: return %[[B0]] : vector<f32>
-func.func @masked_extract_contract1_0d_acc(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: vector<f32>, %mask: vector<4xi1>) -> vector<f32> {
- %0 = vector.mask %mask { vector.contract #dotp_trait %arg0, %arg1, %arg2 : vector<4xf32>, vector<4xf32> into vector<f32> } : vector<4xi1> -> vector<f32>
+func.func @masked_extract_contract1_0d_acc(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %acc: vector<f32>, %mask: vector<4xi1>) -> vector<f32> {
+ %0 = vector.mask %mask { vector.contract #dotp_trait %arg0, %arg1, %acc : vector<4xf32>, vector<4xf32> into vector<f32> } : vector<4xi1> -> vector<f32>
return %0 : vector<f32>
}
More information about the Mlir-commits
mailing list