[Mlir-commits] [mlir] [mlir][vector] Disallow pure reduction dims in vector.contract (PR #237709)
Andrzej Warzyński
llvmlistbot at llvm.org
Tue Nov 3 04:00:01 PST 2026
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/237709
>From 1cb1efbe04a3238307f70949ed4b97f7facdbad3 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 30 Oct 2026 18:33:26 +0000
Subject: [PATCH 1/2] [mlir][vector] Disallow pure reduction dims in
vector.contract
Require every reduction dimension in `vector.contract` to appear in both
LHS and RHS. This disallows one-sided ("pure") reduction dimensions and
aligns the restriction with `linalg.contract`.
For example, the following contraction has two reduction dimensions:
`r2` appears in both inputs, while `r1` appears only in LHS. This PR makes
such contractions fail verification:
```mlir
%res = vector.contract {
indexing_maps = [
affine_map<(r1, p, r2) -> (r1, r2)>,
affine_map<(r1, p, r2) -> (p, r2)>,
affine_map<(r1, p, r2) -> (p)>
],
iterator_types = ["reduction", "parallel", "reduction"],
kind = #vector.kind<add>
} %lhs, %rhs, %acc
: vector<1x2xi32>, vector<2x2xi32> into vector<2xi32>
```
Existing lowering support for pure reduction dimensions was limited to
unit-size cases handled by `ContractionOpLowering::lowerParallel`. This PR
removes that special case (note the function name).
The associated comment attributed these cases to
`CastAwayContractionLeadingOneDim`. However, its current implementation
only removes parallel iteration dimensions corresponding to leading unit
dimensions in the accumulator; it does not introduce pure reduction
dimensions.
Tests are updated to preserve their original edge cases without relying
on pure reduction dimensions. `@contract_one_sided_unit_reduction_dim`
in `vector-contract-to-dot-transforms.mlir` is replaced with a verifier
test checking that such contractions are rejected.
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 1 -
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 4 +--
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 26 ++++++++++++++
.../Vector/Transforms/LowerVectorContract.cpp | 16 ++-------
.../VectorToXeGPU/contract-to-xegpu.mlir | 18 ----------
mlir/test/Dialect/Vector/invalid.mlir | 35 +++++++++++++++----
.../vector-contract-to-dot-transforms.mlir | 28 ---------------
7 files changed, 58 insertions(+), 70 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index db422ceb086c37..6dd70f198c9329 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -109,7 +109,6 @@ def Vector_ContractionOp :
| "parallel" | yes | lhs only | free (M-like) |
| "parallel" | yes | rhs only | free (N-like) |
| "reduction" | no | both | contracting |
- | "reduction" | no | one of them | reduction on one operand |
The M-like and N-like terms follow
[`linalg.contract`](https://mlir.llvm.org/docs/Dialects/Linalg/#linalgcontract-linalgcontractop)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 3c9797cc6acb5c..9501be5d87b58a 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -4852,9 +4852,7 @@ verifyContractionDims(size_t iterationSpaceDims, ArrayRef<size_t> inOccurrences,
// for the output. In terms of einsum semantics such dims have a
// sensible meaning - namely an additional reduction per each such dim.
// By contrast, the ContractionOpInterface does not know about this
- // iter type - cf. inferContractionDims' supported dim kinds. Similarly,
- // while vector.contract's verifier accepts dims of this kind many of
- // its lowerings give up on encountering these dims.
+ // iter type - cf. inferContractionDims' supported dim kinds.
// TODO: Remove following once we have comprehensive support for input-only
// reduction dims, at both the linalg- and vector-dialect levels.
if (inOccCount == 1 && outOccCount != 1)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 7945ca3ea79634..edd6b3c4668bee 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1140,6 +1140,25 @@ static LogicalResult verifyOutputShape(
return success();
}
+// Find a pure reduction dims, i.e. a reduction dim that is only
+// present in only one of the inputs.
+std::optional<int64_t> getPureReductionDim(vector::ContractionOp *contract) {
+ auto maps = contract->getIndexingMapsArray();
+ AffineMap lhsMap = maps[0];
+ AffineMap rhsMap = maps[1];
+
+ for (auto [dim, attr] : llvm::enumerate(contract->getIteratorTypes())) {
+ auto iteratorType = cast<vector::IteratorTypeAttr>(attr).getValue();
+ if (iteratorType != vector::IteratorType::reduction)
+ continue;
+
+ if (!lhsMap.isFunctionOfDim(dim) || !rhsMap.isFunctionOfDim(dim))
+ return dim;
+ }
+
+ return std::nullopt;
+}
+
LogicalResult ContractionOp::verify() {
VectorType lhsType = getLhsType();
VectorType rhsType = getRhsType();
@@ -1180,6 +1199,13 @@ LogicalResult ContractionOp::verify() {
<< index << " to be a projected permutation of its inputs";
}
+ // TODO: Unify this with linalg.contract
+ auto pureReduceDim = getPureReductionDim(this);
+ if (pureReduceDim.has_value())
+ return emitOpError("iteration space dim at index ")
+ << pureReduceDim.value()
+ << " is neither a contracting dim nor of parallel iteration type";
+
auto contractingDimMap = getContractingDimMap();
auto batchDimMap = getBatchDimMap();
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index b3cd4fb1e8f9b0..2c2eaddb214498 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -1001,7 +1001,6 @@ FailureOr<Value> ContractionOpLowering::matchAndRewriteMaskableOp(
}
// Lower one parallel dimension.
-// Incidentally also tolerates unit-size (hence trivial) reduction dimensions.
// TODO: consider reusing existing contract unrolling
FailureOr<Value> ContractionOpLowering::lowerParallel(PatternRewriter &rewriter,
vector::ContractionOp op,
@@ -1042,19 +1041,10 @@ FailureOr<Value> ContractionOpLowering::lowerParallel(PatternRewriter &rewriter,
diag << "expected either lhsIndex=" << lhsIndex
<< " or rhsIndex=" << rhsIndex << " to be nonnegative";
});
- // value_or(-1) means that we tolerate a dimension not appearing
- // in the result map. That can't happen for actual parallel iterators, but
- // the caller ContractionOpLowering::matchAndRewrite is currently calling
- // lowerParallel also for the case of unit-size reduction dims appearing only
- // on one of LHS or RHS, not both. At the moment, such cases are created by
- // CastAwayContractionLeadingOneDim, so we need to either support that or
- // modify that pattern.
+
int64_t resIndex = getResultIndex(iMap[2], iterIndex).value_or(-1);
- if (resIndex == -1 && dimSize != 1)
- return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
- diag << "expected the dimension for iterIndex=" << iterIndex
- << " to either appear in the result map, or to be a unit dimension";
- });
+ assert(resIndex != -1 && "A dimension issing from the accumulator indexing "
+ "map is not a parallel dim!");
// Construct new iterator types and affine map array attribute.
std::array<AffineMap, 3> lowIndexingMaps = {
diff --git a/mlir/test/Conversion/VectorToXeGPU/contract-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/contract-to-xegpu.mlir
index 98b3133465b3ef..151df780813074 100644
--- a/mlir/test/Conversion/VectorToXeGPU/contract-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/contract-to-xegpu.mlir
@@ -165,24 +165,6 @@ func.func @negative_combining_kind(%lhs: vector<8x16xf16>, %rhs: vector<16x16xf1
// -----
-#map = affine_map<(d0, d1, d2) -> (d0, d2)>
-#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
-#map2 = affine_map<(d0, d1, d2) -> ()>
-func.func @negative_accumulator_shape(%lhs: vector<8x16xf16>, %rhs: vector<16x16xf16>,
- %acc: vector<f32>) -> vector<f32> {
- %3 = vector.contract
- {indexing_maps = [#map, #map1, #map2],
- iterator_types = ["reduction", "reduction", "reduction"],
- kind = #vector.kind<add>} %lhs, %rhs, %acc
- : vector<8x16xf16>, vector<16x16xf16> into vector<f32>
- return %3 : vector<f32>
-}
-
-// CHECK-LABEL: @negative_accumulator_shape(
-// CHECK: vector.contract
-
-// -----
-
// A batched matmul whose leading dimensions are batch dimensions shared across
// lhs, rhs, and acc lowers to a batched xegpu.dpas. Two batch dims exercise the
// upper bound (rank 4 = 2 batch + 2 core dims).
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index c44bc2e8f07f68..c22b6422ef6747 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1021,17 +1021,17 @@ func.func @contraction(%arg0: vector<2x1xf32>, %arg1: vector<1x3xf32>, %arg2: ve
// -----
-func.func @contract_with_dim_unused_by_lhs_and_rhs(%arg0 : vector<1x2xi32>, %arg1 : vector<2xi32>, %arg2 : vector<1xi32>) -> vector<1xi32> {
+func.func @contract_with_dim_unused_by_lhs_and_rhs(%arg0 : vector<1x2xi32>, %arg1 : vector<1x2xi32>, %arg2 : vector<1x3xi32>) -> vector<1x3xi32> {
// expected-error at +1 {{'vector.contract' op expected all dimensions to be either a LHS or a RHS dimension}}
%result = vector.contract {
indexing_maps = [
- affine_map<(d0, d1, d2) -> (d0, d2)>,
- affine_map<(d0, d1, d2) -> (d2)>,
- affine_map<(d0, d1, d2) -> (d1)>
+ affine_map<(d0, d1, d2) -> (d0, d1)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>,
+ affine_map<(d0, d1, d2) -> (d0, d2)>
],
- iterator_types = ["reduction", "parallel", "reduction"],
- kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<1x2xi32>, vector<2xi32> into vector<1xi32>
- return %result : vector<1xi32>
+ iterator_types = ["parallel", "reduction", "parallel"],
+ kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<1x2xi32>, vector<1x2xi32> into vector<1x3xi32>
+ return %result : vector<1x3xi32>
}
// -----
@@ -1065,6 +1065,27 @@ func.func @contraction_invalid_kind(%arg0: vector<4x3xf32>,
// -----
+#contraction_accesses = [
+ affine_map<(i, j, k) -> (i, k)>,
+ affine_map<(i, j, k) -> (j)>,
+ affine_map<(i, j, k) -> (i, j)>
+ ]
+#contraction_trait = {
+ kind = "add",
+ indexing_maps = #contraction_accesses,
+ iterator_types = ["parallel", "parallel", "reduction"]
+ }
+func.func @contraction_one_sided_reduction(%arg0: vector<4x3xf32>,
+ %arg1: vector<7xf32>,
+ %arg2: vector<4x7xf32>) {
+ // expected-error @+1 {{iteration space dim at index 2 is neither a contracting dim nor of parallel iteration type}}
+ %0 = vector.contract #contraction_trait %arg0, %arg1, %arg2
+ : vector<4x3xf32>, vector<7xf32> into vector<4x7xf32>
+ return
+}
+
+// -----
+
func.func @create_mask_0d_no_operands() {
%c1 = arith.constant 1 : index
// expected-error at +1 {{must specify exactly one operand for 0-D create_mask}}
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..df9e2d6659e539 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-dot-transforms.mlir
@@ -280,34 +280,6 @@ func.func @full_contract2(%arg0: vector<2x3xf32>,
return %0 : f32
}
-// CHECK-LABEL: @contract_one_sided_unit_reduction_dim
-// CHECK-SAME: (%[[A0:.+]]: vector<1x2xi32>, %[[A1:.+]]: vector<2x2xi32>, %[[A2:.+]]: vector<2xi32>)
-// CHECK-DAG: %[[C:.+]] = arith.constant dense<0> : vector<2xi32>
-// CHECK-DAG: %[[E00:.+]] = vector.extract %[[A0]][0] : vector<2xi32> from vector<1x2xi32>
-// CHECK-DAG: %[[E10:.+]] = vector.extract %[[A1]][0] : vector<2xi32> from vector<2x2xi32>
-// CHECK: %[[M0:.+]] = arith.muli %[[E10]], %[[E00]] : vector<2xi32>
-// CHECK: %[[R0:.+]] = vector.reduction <add>, %[[M0]] : vector<2xi32> into i32
-// CHECK: %[[I0:.+]] = vector.insert %[[R0]], %[[C]] [0] : i32 into vector<2xi32>
-// CHECK: %[[E11:.+]] = vector.extract %[[A1]][1] : vector<2xi32> from vector<2x2xi32>
-// CHECK: %[[M1:.+]] = arith.muli %[[E11]], %[[E00]] : vector<2xi32>
-// CHECK: %[[R1:.+]] = vector.reduction <add>, %[[M1]] : vector<2xi32> into i32
-// CHECK: %[[I1:.+]] = vector.insert %[[R1]], %[[I0]] [1] : i32 into vector<2xi32>
-// CHECK: %[[S:.+]] = arith.addi %[[I1]], %[[A2]] : vector<2xi32>
-// CHECK: return %[[S]] : vector<2xi32>
-
-func.func @contract_one_sided_unit_reduction_dim(%arg0 : vector<1x2xi32>, %arg1 : vector<2x2xi32>, %arg2 : vector<2xi32>) -> vector<2xi32> {
- %res = vector.contract {
- indexing_maps = [
- affine_map<(d0, d1, d2) -> (d0, d2)>,
- affine_map<(d0, d1, d2) -> (d1, d2)>,
- affine_map<(d0, d1, d2) -> (d1)>
- ],
- iterator_types = ["reduction", "parallel", "reduction"],
- kind = #vector.kind<add>
- } %arg0, %arg1, %arg2 : vector<1x2xi32>, vector<2x2xi32>, vector<2xi32> into vector<2xi32>
- return %res : vector<2xi32>
-}
-
// Verify that fastmath flags on vector.contract propagate to the lowered ops.
// CHECK-LABEL: func @extract_contract2_fmf
// CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,
>From a69d42bb0cbddd7706c569d38947acf8120ced6b Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 3 Nov 2026 11:56:58 +0000
Subject: [PATCH 2/2] Address PR comments
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 5 ++--
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 24 +++++++++----------
.../Vector/Transforms/LowerVectorContract.cpp | 2 +-
mlir/test/Dialect/Linalg/invalid.mlir | 2 +-
mlir/test/Dialect/Vector/invalid.mlir | 2 +-
5 files changed, 16 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 9501be5d87b58a..7f7032c273cf00 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -4856,9 +4856,8 @@ verifyContractionDims(size_t iterationSpaceDims, ArrayRef<size_t> inOccurrences,
// TODO: Remove following once we have comprehensive support for input-only
// reduction dims, at both the linalg- and vector-dialect levels.
if (inOccCount == 1 && outOccCount != 1)
- return emitError()
- << "iteration space dim at index " << dimIndex
- << " is neither a contracting dim nor of parallel iteration type";
+ return emitError() << "reduction dim at index " << dimIndex
+ << " must appear in both LHS and RHS";
}
if (!hasContractingDim)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index edd6b3c4668bee..539175460b084d 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1140,19 +1140,18 @@ static LogicalResult verifyOutputShape(
return success();
}
-// Find a pure reduction dims, i.e. a reduction dim that is only
-// present in only one of the inputs.
-std::optional<int64_t> getPureReductionDim(vector::ContractionOp *contract) {
+// Find a one-sided reduction dim, i.e. a reduction dim that is present in only
+// one of the inputs.
+static std::optional<int64_t>
+getOneSidedReductionDim(vector::ContractionOp *contract) {
auto maps = contract->getIndexingMapsArray();
AffineMap lhsMap = maps[0];
AffineMap rhsMap = maps[1];
for (auto [dim, attr] : llvm::enumerate(contract->getIteratorTypes())) {
auto iteratorType = cast<vector::IteratorTypeAttr>(attr).getValue();
- if (iteratorType != vector::IteratorType::reduction)
- continue;
-
- if (!lhsMap.isFunctionOfDim(dim) || !rhsMap.isFunctionOfDim(dim))
+ if (iteratorType == vector::IteratorType::reduction &&
+ (!lhsMap.isFunctionOfDim(dim) || !rhsMap.isFunctionOfDim(dim)))
return dim;
}
@@ -1199,12 +1198,11 @@ LogicalResult ContractionOp::verify() {
<< index << " to be a projected permutation of its inputs";
}
- // TODO: Unify this with linalg.contract
- auto pureReduceDim = getPureReductionDim(this);
- if (pureReduceDim.has_value())
- return emitOpError("iteration space dim at index ")
- << pureReduceDim.value()
- << " is neither a contracting dim nor of parallel iteration type";
+ // TODO: For better code re-use-, unify this with linalg.contract.
+ auto dimIdx = getOneSidedReductionDim(this);
+ if (dimIdx.has_value())
+ return emitError() << "reduction dim at index " << dimIdx.value()
+ << " must appear in both LHS and RHS";
auto contractingDimMap = getContractingDimMap();
auto batchDimMap = getBatchDimMap();
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index 2c2eaddb214498..6d344252603b34 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -1043,7 +1043,7 @@ FailureOr<Value> ContractionOpLowering::lowerParallel(PatternRewriter &rewriter,
});
int64_t resIndex = getResultIndex(iMap[2], iterIndex).value_or(-1);
- assert(resIndex != -1 && "A dimension issing from the accumulator indexing "
+ assert(resIndex != -1 && "A dimension missing from the accumulator indexing "
"map is not a parallel dim!");
// Construct new iterator types and affine map array attribute.
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 70b375ef24941b..581728e412d73f 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -745,7 +745,7 @@ func.func @unused_iteration_space_dim_contraction(
func.func @unused_iteration_space_dim_contraction(
%lhs: tensor<8x4x1xf32>, %rhs: tensor<1x64xf32>, %init: tensor<4x64xf32>) {
- // expected-error @+1 {{iteration space dim at index 3 is neither a contracting dim nor of parallel iteration type}}
+ // expected-error @+1 {{reduction dim at index 3 must appear in both LHS and RHS}}
linalg.contract
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>,
affine_map<(d0, d1, d2, d3) -> (d2, d1)>,
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index c22b6422ef6747..4dd2571722bc47 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1078,7 +1078,7 @@ func.func @contraction_invalid_kind(%arg0: vector<4x3xf32>,
func.func @contraction_one_sided_reduction(%arg0: vector<4x3xf32>,
%arg1: vector<7xf32>,
%arg2: vector<4x7xf32>) {
- // expected-error @+1 {{iteration space dim at index 2 is neither a contracting dim nor of parallel iteration type}}
+ // expected-error @+1 {{reduction dim at index 2 must appear in both LHS and RHS}}
%0 = vector.contract #contraction_trait %arg0, %arg1, %arg2
: vector<4x3xf32>, vector<7xf32> into vector<4x7xf32>
return
More information about the Mlir-commits
mailing list