[Mlir-commits] [mlir] experimental vector contract (PR #194536)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri May 8 11:47:33 PDT 2026
https://github.com/efric updated https://github.com/llvm/llvm-project/pull/194536
>From 549ecfe1f0709e96c392cfc5dfd5c50b043a2f36 Mon Sep 17 00:00:00 2001
From: Eric Feng <Eric.Feng at amd.com>
Date: Mon, 27 Apr 2026 21:35:05 -0700
Subject: [PATCH 1/5] experimental expose vector contract lowerings to have
multiple options
Signed-off-by: Eric Feng <Eric.Feng at amd.com>
---
.../Vector/Transforms/LoweringPatterns.h | 27 ++
.../Vector/Transforms/LowerVectorContract.cpp | 280 +++++++++++++-----
.../vector-contract-composable-lowering.mlir | 91 ++++++
.../Dialect/Vector/TestVectorTransforms.cpp | 64 ++++
4 files changed, 381 insertions(+), 81 deletions(-)
create mode 100644 mlir/test/Dialect/Vector/vector-contract-composable-lowering.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
index aa75eff409ef9..4999899849cb9 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
@@ -12,6 +12,8 @@
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
+#include <functional>
+
namespace mlir {
class RewritePatternSet;
@@ -47,6 +49,31 @@ namespace vector {
/// [ContractionOpToOuterProductOpLowering]
/// Progressively lower a `vector.contract` with row-major matmul semantics to
/// linearized `vector.extract` + `vector.outerproduct` + `vector.insert`.
+using VectorContractLoweringFilter =
+ std::function<LogicalResult(ContractionOp)>;
+
+LogicalResult acceptAllVectorContractLoweringFilter(ContractionOp op);
+
+void populateVectorContractToDotPatterns(
+ RewritePatternSet &patterns,
+ VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1);
+
+void populateVectorContractToOuterProductPatterns(
+ RewritePatternSet &patterns,
+ VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1);
+
+void populateVectorContractToParallelArithPatterns(
+ RewritePatternSet &patterns,
+ VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1);
+
+void populateVectorContractGenericLoweringPatterns(
+ RewritePatternSet &patterns,
+ VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1);
+
void populateVectorContractLoweringPatterns(
RewritePatternSet &patterns,
VectorContractLowering vectorContractLoweringOption,
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index eaf7bb8109514..2adf71c5da324 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -201,6 +201,38 @@ static Value createMul(Location loc, Value x, Value y, bool isInt,
return arith::MulFOp::create(rewriter, loc, x, y, fmf);
}
+static LogicalResult
+checkSameOperandAndAccumulatorElementType(vector::ContractionOp op,
+ PatternRewriter &rewriter) {
+ if (op.getLhsType().getElementType() ==
+ getElementTypeOrSelf(op.getAccType()) &&
+ op.getRhsType().getElementType() == getElementTypeOrSelf(op.getAccType()))
+ return success();
+
+ return rewriter.notifyMatchFailure(
+ op, "mixed-mode contract lowering is not supported");
+}
+
+static LogicalResult checkAddKind(vector::ContractionOp op,
+ PatternRewriter &rewriter) {
+ if (op.getKind() == vector::CombiningKind::ADD)
+ return success();
+
+ return rewriter.notifyMatchFailure(
+ op, "contractions other than 'add' not supported");
+}
+
+static bool isContractArithOpSupported(vector::CombiningKind kind, bool isInt) {
+ using vector::CombiningKind;
+ if (isInt)
+ return kind != CombiningKind::MINNUMF && kind != CombiningKind::MAXNUMF &&
+ kind != CombiningKind::MINIMUMF && kind != CombiningKind::MAXIMUMF;
+ return kind != CombiningKind::AND && kind != CombiningKind::MINUI &&
+ kind != CombiningKind::MINSI && kind != CombiningKind::MAXUI &&
+ kind != CombiningKind::MAXSI && kind != CombiningKind::OR &&
+ kind != CombiningKind::XOR;
+}
+
namespace {
/// Progressive lowering of a `vector.contract %a, %b, %c` with row-major matmul
@@ -223,29 +255,31 @@ class ContractionOpToOuterProductOpLowering
public:
using MaskableOpRewritePattern::MaskableOpRewritePattern;
- using FilterConstraintType =
- std::function<LogicalResult(vector::ContractionOp op)>;
-
- static LogicalResult defaultFilter(vector::ContractionOp op) {
- return success();
- }
-
ContractionOpToOuterProductOpLowering(
vector::VectorContractLowering vectorContractLowering,
MLIRContext *context, PatternBenefit benefit = 1,
- FilterConstraintType constraint = defaultFilter)
+ VectorContractLoweringFilter constraint =
+ acceptAllVectorContractLoweringFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
vectorContractLowering(vectorContractLowering),
filter(std::move(constraint)) {}
+ ContractionOpToOuterProductOpLowering(
+ MLIRContext *context,
+ VectorContractLoweringFilter constraint =
+ acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1)
+ : MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
+ filter(std::move(constraint)) {}
+
FailureOr<Value>
matchAndRewriteMaskableOp(vector::ContractionOp op, MaskingOpInterface maskOp,
PatternRewriter &rewriter) const override;
private:
/// Options to control the vector patterns.
- vector::VectorContractLowering vectorContractLowering;
- FilterConstraintType filter;
+ std::optional<vector::VectorContractLowering> vectorContractLowering;
+ VectorContractLoweringFilter filter;
};
/// Progressive lowering of a `vector.contract %a, %b, %c` with row-major matmul
@@ -271,19 +305,21 @@ class ContractionOpToDotLowering
public:
using MaskableOpRewritePattern::MaskableOpRewritePattern;
- using FilterConstraintType =
- std::function<LogicalResult(vector::ContractionOp op)>;
-
- static LogicalResult defaultFilter(vector::ContractionOp op) {
- return success();
- }
-
ContractionOpToDotLowering(
vector::VectorContractLowering vectorContractLowering,
MLIRContext *context, PatternBenefit benefit = 1,
- const FilterConstraintType &constraint = defaultFilter)
+ VectorContractLoweringFilter constraint =
+ acceptAllVectorContractLoweringFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
- vectorContractLowering(vectorContractLowering), filter(defaultFilter) {}
+ vectorContractLowering(vectorContractLowering),
+ filter(std::move(constraint)) {}
+
+ ContractionOpToDotLowering(MLIRContext *context,
+ VectorContractLoweringFilter constraint =
+ acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1)
+ : MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
+ filter(std::move(constraint)) {}
FailureOr<Value>
matchAndRewriteMaskableOp(vector::ContractionOp op, MaskingOpInterface maskOp,
@@ -291,8 +327,8 @@ class ContractionOpToDotLowering
private:
/// Options to control the vector patterns.
- vector::VectorContractLowering vectorContractLowering;
- FilterConstraintType filter;
+ std::optional<vector::VectorContractLowering> vectorContractLowering;
+ VectorContractLoweringFilter filter;
};
/// Progressive lowering of ContractionOp.
@@ -309,33 +345,34 @@ class ContractionOpToDotLowering
///
/// This only kicks in when either VectorTransformsOptions is set
/// to Dot or when other contraction patterns fail.
-class ContractionOpLowering
+class ContractionOpGenericLowering
: public MaskableOpRewritePattern<vector::ContractionOp> {
public:
using MaskableOpRewritePattern::MaskableOpRewritePattern;
- using FilterConstraintType =
- std::function<LogicalResult(vector::ContractionOp op)>;
- static LogicalResult defaultFilter(vector::ContractionOp op) {
- return success();
- }
-
- ContractionOpLowering(
- vector::VectorContractLowering vectorContractLoweringOption,
- MLIRContext *context, PatternBenefit benefit = 1,
- FilterConstraintType constraint = defaultFilter)
+ ContractionOpGenericLowering(MLIRContext *context,
+ VectorContractLoweringFilter constraint =
+ acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
- vectorContractLoweringOption(vectorContractLoweringOption),
filter(std::move(constraint)) {}
FailureOr<Value>
matchAndRewriteMaskableOp(vector::ContractionOp op, MaskingOpInterface maskOp,
PatternRewriter &rewriter) const override;
+protected:
+ LogicalResult
+ matchSupportedGenericContraction(vector::ContractionOp op,
+ PatternRewriter &rewriter) const;
+
+ FailureOr<Value> lowerGenericContraction(PatternRewriter &rewriter,
+ vector::ContractionOp op,
+ MaskingOpInterface maskOp) const;
+
private:
- /// Options to control the vector patterns.
- vector::VectorContractLowering vectorContractLoweringOption;
- FilterConstraintType filter;
+ VectorContractLoweringFilter filter;
+
// Lower one parallel dimension.
FailureOr<Value> lowerParallel(PatternRewriter &rewriter,
vector::ContractionOp op, int64_t lhsIndex,
@@ -345,6 +382,25 @@ class ContractionOpLowering
vector::ContractionOp op, Value mask) const;
};
+class ContractionOpLowering : public ContractionOpGenericLowering {
+public:
+ ContractionOpLowering(
+ vector::VectorContractLowering vectorContractLoweringOption,
+ MLIRContext *context, PatternBenefit benefit = 1,
+ VectorContractLoweringFilter constraint =
+ acceptAllVectorContractLoweringFilter)
+ : ContractionOpGenericLowering(context, std::move(constraint), benefit),
+ vectorContractLoweringOption(vectorContractLoweringOption) {}
+
+ FailureOr<Value>
+ matchAndRewriteMaskableOp(vector::ContractionOp op, MaskingOpInterface maskOp,
+ PatternRewriter &rewriter) const override;
+
+private:
+ /// Options to control the vector patterns.
+ vector::VectorContractLowering vectorContractLoweringOption;
+};
+
/// Generate a vector implementation for matmat, matvec and tmatvec.
/// This unrolls outer-products along the reduction dimension.
struct UnrolledOuterProductGenerator
@@ -592,7 +648,8 @@ FailureOr<Value>
ContractionOpToOuterProductOpLowering::matchAndRewriteMaskableOp(
vector::ContractionOp op, MaskingOpInterface maskOp,
PatternRewriter &rewriter) const {
- if (vectorContractLowering != vector::VectorContractLowering::OuterProduct)
+ if (vectorContractLowering &&
+ *vectorContractLowering != vector::VectorContractLowering::OuterProduct)
return failure();
if (failed(filter(op)))
@@ -622,9 +679,19 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
if (failed(filter(op)))
return failure();
- if (vectorContractLowering != vector::VectorContractLowering::Dot)
+ if (vectorContractLowering &&
+ *vectorContractLowering != vector::VectorContractLowering::Dot)
+ return failure();
+
+ if (failed(checkSameOperandAndAccumulatorElementType(op, rewriter)) ||
+ failed(checkAddKind(op, rewriter)))
return failure();
+ VectorType dstType = dyn_cast<VectorType>(op.getResultType());
+ if (!dstType || dstType.getRank() < 1 || dstType.getRank() > 2)
+ return rewriter.notifyMatchFailure(
+ op, "expected result type of rank 1 or 2 for dot lowering");
+
auto iteratorTypes = op.getIteratorTypes().getValue();
static constexpr std::array<int64_t, 2> perm = {1, 0};
Location loc = op.getLoc();
@@ -641,7 +708,7 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
// In the following we wish to make the reduction dimension innermost so we
// can load vectors and just fmul + reduce into a scalar.
//
- if (isParallelIterator(iteratorTypes[0]) &&
+ if (iteratorTypes.size() == 3 && isParallelIterator(iteratorTypes[0]) &&
isParallelIterator(iteratorTypes[1]) &&
isReductionIterator(iteratorTypes[2])) {
//
@@ -674,7 +741,8 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
} else {
return failure();
}
- } else if (isParallelIterator(iteratorTypes[0]) &&
+ } else if (iteratorTypes.size() == 2 &&
+ isParallelIterator(iteratorTypes[0]) &&
isReductionIterator(iteratorTypes[1])) {
//
// One outer parallel, one inner reduction (matvec flavor)
@@ -695,10 +763,6 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
return failure();
}
- VectorType dstType = cast<VectorType>(op.getResultType());
- assert(dstType.getRank() >= 1 && dstType.getRank() <= 2 &&
- "Expected dst type of rank 1 or 2");
-
unsigned rank = dstType.getRank();
unsigned dstRows = dstType.getShape()[0];
unsigned dstColumns = rank == 1 ? 1 : dstType.getShape()[1];
@@ -744,17 +808,21 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
struct ContractOpToElementwise
: public MaskableOpRewritePattern<vector::ContractionOp> {
using MaskableOpRewritePattern::MaskableOpRewritePattern;
- using FilterConstraintType =
- std::function<LogicalResult(vector::ContractionOp op)>;
- static LogicalResult defaultFilter(vector::ContractionOp op) {
- return success();
- }
- ContractOpToElementwise(
- vector::VectorContractLowering vectorContractLowering,
- MLIRContext *context, PatternBenefit benefit = 1,
- const FilterConstraintType &constraint = defaultFilter)
+
+ ContractOpToElementwise(vector::VectorContractLowering vectorContractLowering,
+ MLIRContext *context, PatternBenefit benefit = 1,
+ VectorContractLoweringFilter constraint =
+ acceptAllVectorContractLoweringFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
- vectorContractLowering(vectorContractLowering), filter(defaultFilter) {}
+ vectorContractLowering(vectorContractLowering),
+ filter(std::move(constraint)) {}
+
+ ContractOpToElementwise(MLIRContext *context,
+ VectorContractLoweringFilter constraint =
+ acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1)
+ : MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
+ filter(std::move(constraint)) {}
FailureOr<Value>
matchAndRewriteMaskableOp(vector::ContractionOp contractOp,
@@ -767,7 +835,12 @@ struct ContractOpToElementwise
if (failed(filter(contractOp)))
return failure();
- if (vectorContractLowering != vector::VectorContractLowering::ParallelArith)
+ if (vectorContractLowering &&
+ *vectorContractLowering !=
+ vector::VectorContractLowering::ParallelArith)
+ return failure();
+
+ if (failed(checkSameOperandAndAccumulatorElementType(contractOp, rewriter)))
return failure();
ArrayRef<int64_t> lhsShape = contractOp.getLhsType().getShape();
@@ -825,6 +898,10 @@ struct ContractOpToElementwise
rhsTranspose.push_back(rhsDims.size() - 1);
}
}
+ bool isInt = contractOp.getLhsType().getElementType().isIntOrIndex();
+ if (!isContractArithOpSupported(contractOp.getKind(), isInt))
+ return failure();
+
Value newLhs = contractOp.getLhs();
Value newRhs = contractOp.getRhs();
Location loc = contractOp.getLoc();
@@ -840,7 +917,6 @@ struct ContractOpToElementwise
VectorType::get(rhsDims, contractOp.getRhsType().getElementType());
newRhs = vector::BroadcastOp::create(rewriter, loc, expandedType, newRhs);
}
- bool isInt = contractOp.getLhsType().getElementType().isIntOrIndex();
newLhs = vector::TransposeOp::create(rewriter, loc, newLhs, lhsTranspose);
newRhs = vector::TransposeOp::create(rewriter, loc, newRhs, rhsTranspose);
SmallVector<int64_t> lhsOffsets(lhsReductionDims.size(), 0);
@@ -851,16 +927,14 @@ struct ContractOpToElementwise
createContractArithOp(loc, newLhs, newRhs, contractOp.getAcc(),
contractOp.getKind(), rewriter, isInt,
/*mask=*/Value(), contractOp.getFastmathAttr());
- if (result)
- return *result;
-
- return failure();
+ assert(result && "kind and type support should have been checked");
+ return *result;
}
private:
/// Options to control the vector patterns.
- vector::VectorContractLowering vectorContractLowering;
- FilterConstraintType filter;
+ std::optional<vector::VectorContractLowering> vectorContractLowering;
+ VectorContractLoweringFilter filter;
};
/// Progressive lowering of ContractionOp.
@@ -880,24 +954,31 @@ struct ContractOpToElementwise
// TODO: break down into transpose/reshape/cast ops
// when they become available to avoid code dup
// TODO: investigate lowering order impact on performance
-FailureOr<Value> ContractionOpLowering::matchAndRewriteMaskableOp(
- vector::ContractionOp op, MaskingOpInterface maskOp,
- PatternRewriter &rewriter) const {
+LogicalResult ContractionOpGenericLowering::matchSupportedGenericContraction(
+ vector::ContractionOp op, PatternRewriter &rewriter) const {
if (failed(filter(op)))
return failure();
- // TODO: support mixed mode contract lowering.
- if (op.getLhsType().getElementType() !=
- getElementTypeOrSelf(op.getAccType()) ||
- op.getRhsType().getElementType() != getElementTypeOrSelf(op.getAccType()))
+ if (failed(checkSameOperandAndAccumulatorElementType(op, rewriter)))
return failure();
- // TODO: the code below assumes the default contraction, make sure it supports
- // other kinds before enabling this lowering.
- if (op.getKind() != vector::CombiningKind::ADD) {
- return rewriter.notifyMatchFailure(
- op, "contractions other than 'add' not supported");
- }
+ return checkAddKind(op, rewriter);
+}
+
+FailureOr<Value> ContractionOpGenericLowering::matchAndRewriteMaskableOp(
+ vector::ContractionOp op, MaskingOpInterface maskOp,
+ PatternRewriter &rewriter) const {
+ if (failed(matchSupportedGenericContraction(op, rewriter)))
+ return failure();
+
+ return lowerGenericContraction(rewriter, op, maskOp);
+}
+
+FailureOr<Value> ContractionOpLowering::matchAndRewriteMaskableOp(
+ vector::ContractionOp op, MaskingOpInterface maskOp,
+ PatternRewriter &rewriter) const {
+ if (failed(matchSupportedGenericContraction(op, rewriter)))
+ return failure();
// TODO: implement benefits, cost models.
MLIRContext *ctx = op.getContext();
@@ -920,6 +1001,12 @@ FailureOr<Value> ContractionOpLowering::matchAndRewriteMaskableOp(
if (!failed(newVal4))
return newVal4;
+ return lowerGenericContraction(rewriter, op, maskOp);
+}
+
+FailureOr<Value> ContractionOpGenericLowering::lowerGenericContraction(
+ PatternRewriter &rewriter, vector::ContractionOp op,
+ MaskingOpInterface maskOp) const {
// Vector mask setup.
Value mask;
@@ -982,11 +1069,9 @@ 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,
- int64_t lhsIndex,
- int64_t rhsIndex,
- Value mask) const {
+FailureOr<Value> ContractionOpGenericLowering::lowerParallel(
+ PatternRewriter &rewriter, vector::ContractionOp op, int64_t lhsIndex,
+ int64_t rhsIndex, Value mask) const {
VectorType lhsType = op.getLhsType();
VectorType rhsType = op.getRhsType();
VectorType resType = cast<VectorType>(op.getResultType());
@@ -1069,7 +1154,7 @@ FailureOr<Value> ContractionOpLowering::lowerParallel(PatternRewriter &rewriter,
}
// Lower one reduction dimension.
-FailureOr<Value> ContractionOpLowering::lowerReduction(
+FailureOr<Value> ContractionOpGenericLowering::lowerReduction(
PatternRewriter &rewriter, vector::ContractionOp op, Value mask) const {
auto loc = op.getLoc();
VectorType lhsType = op.getLhsType();
@@ -1229,6 +1314,39 @@ class OuterProductOpLowering : public OpRewritePattern<vector::OuterProductOp> {
} // namespace
+LogicalResult
+mlir::vector::acceptAllVectorContractLoweringFilter(ContractionOp) {
+ return success();
+}
+
+void mlir::vector::populateVectorContractToDotPatterns(
+ RewritePatternSet &patterns, VectorContractLoweringFilter filter,
+ PatternBenefit benefit) {
+ patterns.add<ContractionOpToDotLowering>(patterns.getContext(),
+ std::move(filter), benefit);
+}
+
+void mlir::vector::populateVectorContractToOuterProductPatterns(
+ RewritePatternSet &patterns, VectorContractLoweringFilter filter,
+ PatternBenefit benefit) {
+ patterns.add<ContractionOpToOuterProductOpLowering>(
+ patterns.getContext(), std::move(filter), benefit);
+}
+
+void mlir::vector::populateVectorContractToParallelArithPatterns(
+ RewritePatternSet &patterns, VectorContractLoweringFilter filter,
+ PatternBenefit benefit) {
+ patterns.add<ContractOpToElementwise>(patterns.getContext(),
+ std::move(filter), benefit);
+}
+
+void mlir::vector::populateVectorContractGenericLoweringPatterns(
+ RewritePatternSet &patterns, VectorContractLoweringFilter filter,
+ PatternBenefit benefit) {
+ patterns.add<ContractionOpGenericLowering>(patterns.getContext(),
+ std::move(filter), benefit);
+}
+
void mlir::vector::populateVectorContractLoweringPatterns(
RewritePatternSet &patterns,
VectorContractLowering vectorContractLoweringOption, PatternBenefit benefit,
diff --git a/mlir/test/Dialect/Vector/vector-contract-composable-lowering.mlir b/mlir/test/Dialect/Vector/vector-contract-composable-lowering.mlir
new file mode 100644
index 0000000000000..1c303a0e082f4
--- /dev/null
+++ b/mlir/test/Dialect/Vector/vector-contract-composable-lowering.mlir
@@ -0,0 +1,91 @@
+// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=dot-outerproduct" --split-input-file | FileCheck %s --check-prefix=DOT
+// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=generic" --split-input-file | FileCheck %s --check-prefix=GENERIC
+// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=parallel-arith-reject" --split-input-file | FileCheck %s --check-prefix=PARALLEL
+
+#matmat_accesses = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (k, n)>,
+ affine_map<(m, n, k) -> (m, n)>
+]
+#matmat_trait = {
+ indexing_maps = #matmat_accesses,
+ iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+// DOT-LABEL: func @dot_accept
+// DOT-NOT: vector.outerproduct
+// DOT: vector.reduction <add>
+func.func @dot_accept(%A: vector<2x4xf32>,
+ %B: vector<4x3xf32>,
+ %C: vector<2x3xf32>) -> vector<2x3xf32> {
+ %0 = vector.contract #matmat_trait %A, %B, %C
+ : vector<2x4xf32>, vector<4x3xf32> into vector<2x3xf32>
+ return %0 : vector<2x3xf32>
+}
+
+// DOT-LABEL: func @dot_reject_to_outerproduct
+// DOT: vector.outerproduct
+func.func @dot_reject_to_outerproduct(%A: vector<2x4xf32>,
+ %B: vector<4x3xf32>,
+ %C: vector<2x3xf32>)
+ -> vector<2x3xf32> {
+ %0 = vector.contract #matmat_trait %A, %B, %C
+ : vector<2x4xf32>, vector<4x3xf32> into vector<2x3xf32>
+ return %0 : vector<2x3xf32>
+}
+
+// -----
+
+#dotp_accesses = [
+ affine_map<(i) -> (i)>,
+ affine_map<(i) -> (i)>,
+ affine_map<(i) -> ()>
+]
+#dotp_add_trait = {
+ indexing_maps = #dotp_accesses,
+ iterator_types = ["reduction"]
+}
+#dotp_mul_trait = {
+ indexing_maps = #dotp_accesses,
+ iterator_types = ["reduction"],
+ kind = #vector.kind<mul>
+}
+
+// GENERIC-LABEL: func @generic_add
+// GENERIC: arith.mulf
+// GENERIC: vector.reduction <add>
+func.func @generic_add(%A: vector<4xf32>, %B: vector<4xf32>,
+ %C: f32) -> f32 {
+ %0 = vector.contract #dotp_add_trait %A, %B, %C
+ : vector<4xf32>, vector<4xf32> into f32
+ return %0 : f32
+}
+
+// GENERIC-LABEL: func @generic_non_add
+// GENERIC: vector.contract
+// GENERIC-SAME: kind = #vector.kind<mul>
+func.func @generic_non_add(%A: vector<4xf32>, %B: vector<4xf32>,
+ %C: f32) -> f32 {
+ %0 = vector.contract #dotp_mul_trait %A, %B, %C
+ : vector<4xf32>, vector<4xf32> into f32
+ return %0 : f32
+}
+
+// -----
+
+// PARALLEL-LABEL: func @parallel_arith_filter_reject
+// PARALLEL: vector.contract
+func.func @parallel_arith_filter_reject(
+ %A: vector<1x1x4xf32>, %B: vector<1x1x4xf32>,
+ %C: vector<4xf32>) -> vector<4xf32> {
+ %0 = vector.contract {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d1, d2, d0)>,
+ affine_map<(d0, d1, d2) -> (d1, d2, d0)>,
+ affine_map<(d0, d1, d2) -> (d0)>
+ ],
+ iterator_types = ["parallel", "reduction", "reduction"],
+ kind = #vector.kind<add>
+ } %A, %B, %C : vector<1x1x4xf32>, vector<1x1x4xf32> into vector<4xf32>
+ return %0 : vector<4xf32>
+}
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index ff3520a286cc8..7f674a5c8c10c 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include <optional>
+#include <string>
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
@@ -138,6 +139,67 @@ struct TestVectorContractionPrepareForMMTLowering
}
};
+static bool parentFunctionNameContains(vector::ContractionOp op,
+ StringRef substring) {
+ if (auto funcOp = op->getParentOfType<func::FuncOp>())
+ return funcOp.getName().contains(substring);
+ return false;
+}
+
+struct TestVectorContractLoweringComposition final
+ : public PassWrapper<TestVectorContractLoweringComposition,
+ OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+ TestVectorContractLoweringComposition)
+
+ TestVectorContractLoweringComposition() = default;
+ TestVectorContractLoweringComposition(
+ const TestVectorContractLoweringComposition &pass)
+ : PassWrapper(pass) {}
+
+ StringRef getArgument() const final {
+ return "test-vector-contract-lowering-composition";
+ }
+
+ StringRef getDescription() const final {
+ return "Test composable vector.contract lowering pattern population.";
+ }
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<arith::ArithDialect, vector::VectorDialect>();
+ }
+
+ void runOnOperation() override {
+ RewritePatternSet patterns(&getContext());
+ if (mode == "dot-outerproduct") {
+ populateVectorContractToDotPatterns(
+ patterns,
+ [](vector::ContractionOp op) {
+ return success(!parentFunctionNameContains(op, "dot_reject"));
+ },
+ PatternBenefit(2));
+ populateVectorContractToOuterProductPatterns(
+ patterns, acceptAllVectorContractLoweringFilter, PatternBenefit(1));
+ } else if (mode == "generic") {
+ populateVectorContractGenericLoweringPatterns(patterns);
+ } else if (mode == "parallel-arith-reject") {
+ populateVectorContractToParallelArithPatterns(
+ patterns, [](vector::ContractionOp) { return failure(); });
+ } else {
+ getOperation().emitError()
+ << "unknown contract lowering test mode: " << mode;
+ return signalPassFailure();
+ }
+
+ (void)applyPatternsGreedily(getOperation(), std::move(patterns));
+ }
+
+ Option<std::string> mode{
+ *this, "mode",
+ llvm::cl::desc("Contract lowering composition mode to test"),
+ llvm::cl::init("dot-outerproduct")};
+};
+
struct TestVectorUnrollingPatterns
: public PassWrapper<TestVectorUnrollingPatterns,
OperationPass<func::FuncOp>> {
@@ -1053,6 +1115,8 @@ void registerTestVectorLowerings() {
PassRegistration<TestVectorContractionPrepareForMMTLowering>();
+ PassRegistration<TestVectorContractLoweringComposition>();
+
PassRegistration<TestVectorUnrollingPatterns>();
PassRegistration<TestVectorTransferUnrollingPatterns>();
>From 055d49c55d4dccdf9f35f16425528199884ab65d Mon Sep 17 00:00:00 2001
From: Eric Feng <Eric.Feng at amd.com>
Date: Thu, 7 May 2026 16:29:49 -0700
Subject: [PATCH 2/5] add more tests and parallelarith
Signed-off-by: Eric Feng <Eric.Feng at amd.com>
---
.../Vector/Transforms/LoweringPatterns.h | 5 ++
.../vector-contract-composable-lowering.mlir | 47 +++++++++++++++----
.../Dialect/Vector/TestVectorTransforms.cpp | 8 +++-
3 files changed, 48 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
index 4999899849cb9..689fabb1413dc 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
@@ -49,6 +49,11 @@ namespace vector {
/// [ContractionOpToOuterProductOpLowering]
/// Progressively lower a `vector.contract` with row-major matmul semantics to
/// linearized `vector.extract` + `vector.outerproduct` + `vector.insert`.
+
+/// A `VectorContractLoweringFilter` lets clients compose multiple lowering
+/// strategies by benefit. Returning failure means this strategy silently
+/// declines the op without consuming it or diagnosing invalid IR; lower-benefit
+/// strategies may still match the same op.
using VectorContractLoweringFilter =
std::function<LogicalResult(ContractionOp)>;
diff --git a/mlir/test/Dialect/Vector/vector-contract-composable-lowering.mlir b/mlir/test/Dialect/Vector/vector-contract-composable-lowering.mlir
index 1c303a0e082f4..2f755afb92951 100644
--- a/mlir/test/Dialect/Vector/vector-contract-composable-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-composable-lowering.mlir
@@ -1,6 +1,7 @@
-// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=dot-outerproduct" --split-input-file | FileCheck %s --check-prefix=DOT
+// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=composed" --split-input-file | FileCheck %s --check-prefix=COMPOSED
// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=generic" --split-input-file | FileCheck %s --check-prefix=GENERIC
-// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=parallel-arith-reject" --split-input-file | FileCheck %s --check-prefix=PARALLEL
+// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=parallel-arith" --split-input-file | FileCheck %s --check-prefix=PARALLEL_ACCEPT
+// RUN: mlir-opt %s --test-vector-contract-lowering-composition="mode=parallel-arith-reject" --split-input-file | FileCheck %s --check-prefix=PARALLEL_REJECT
#matmat_accesses = [
affine_map<(m, n, k) -> (m, k)>,
@@ -12,9 +13,9 @@
iterator_types = ["parallel", "parallel", "reduction"]
}
-// DOT-LABEL: func @dot_accept
-// DOT-NOT: vector.outerproduct
-// DOT: vector.reduction <add>
+// COMPOSED-LABEL: func @dot_accept
+// COMPOSED-NOT: vector.outerproduct
+// COMPOSED: vector.reduction <add>
func.func @dot_accept(%A: vector<2x4xf32>,
%B: vector<4x3xf32>,
%C: vector<2x3xf32>) -> vector<2x3xf32> {
@@ -23,8 +24,8 @@ func.func @dot_accept(%A: vector<2x4xf32>,
return %0 : vector<2x3xf32>
}
-// DOT-LABEL: func @dot_reject_to_outerproduct
-// DOT: vector.outerproduct
+// COMPOSED-LABEL: func @dot_reject_to_outerproduct
+// COMPOSED: vector.outerproduct
func.func @dot_reject_to_outerproduct(%A: vector<2x4xf32>,
%B: vector<4x3xf32>,
%C: vector<2x3xf32>)
@@ -34,6 +35,29 @@ func.func @dot_reject_to_outerproduct(%A: vector<2x4xf32>,
return %0 : vector<2x3xf32>
}
+#batch_matmul_accesses = [
+ affine_map<(b, m, n, k) -> (b, m, k)>,
+ affine_map<(b, m, n, k) -> (b, k, n)>,
+ affine_map<(b, m, n, k) -> (b, m, n)>
+]
+#batch_matmul_trait = {
+ indexing_maps = #batch_matmul_accesses,
+ iterator_types = ["parallel", "parallel", "parallel", "reduction"]
+}
+
+// COMPOSED-LABEL: func @dot_structural_failure_to_generic
+// COMPOSED-NOT: vector.outerproduct
+// COMPOSED: vector.extract
+// COMPOSED: vector.reduction <add>
+func.func @dot_structural_failure_to_generic(%A: vector<2x2x4xf32>,
+ %B: vector<2x4x3xf32>,
+ %C: vector<2x2x3xf32>)
+ -> vector<2x2x3xf32> {
+ %0 = vector.contract #batch_matmul_trait %A, %B, %C
+ : vector<2x2x4xf32>, vector<2x4x3xf32> into vector<2x2x3xf32>
+ return %0 : vector<2x2x3xf32>
+}
+
// -----
#dotp_accesses = [
@@ -73,9 +97,12 @@ func.func @generic_non_add(%A: vector<4xf32>, %B: vector<4xf32>,
// -----
-// PARALLEL-LABEL: func @parallel_arith_filter_reject
-// PARALLEL: vector.contract
-func.func @parallel_arith_filter_reject(
+// PARALLEL_ACCEPT-LABEL: func @parallel_arith
+// PARALLEL_ACCEPT-NOT: vector.contract
+// PARALLEL_ACCEPT: vector.fma
+// PARALLEL_REJECT-LABEL: func @parallel_arith
+// PARALLEL_REJECT: vector.contract
+func.func @parallel_arith(
%A: vector<1x1x4xf32>, %B: vector<1x1x4xf32>,
%C: vector<4xf32>) -> vector<4xf32> {
%0 = vector.contract {
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index 7f674a5c8c10c..95af6dff8eaeb 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -171,17 +171,21 @@ struct TestVectorContractLoweringComposition final
void runOnOperation() override {
RewritePatternSet patterns(&getContext());
- if (mode == "dot-outerproduct") {
+ if (mode == "composed") {
populateVectorContractToDotPatterns(
patterns,
[](vector::ContractionOp op) {
return success(!parentFunctionNameContains(op, "dot_reject"));
},
- PatternBenefit(2));
+ PatternBenefit(3));
populateVectorContractToOuterProductPatterns(
+ patterns, acceptAllVectorContractLoweringFilter, PatternBenefit(2));
+ populateVectorContractGenericLoweringPatterns(
patterns, acceptAllVectorContractLoweringFilter, PatternBenefit(1));
} else if (mode == "generic") {
populateVectorContractGenericLoweringPatterns(patterns);
+ } else if (mode == "parallel-arith") {
+ populateVectorContractToParallelArithPatterns(patterns);
} else if (mode == "parallel-arith-reject") {
populateVectorContractToParallelArithPatterns(
patterns, [](vector::ContractionOp) { return failure(); });
>From 6627acce4f3a2a345dadf9f4e1e669dd266863e4 Mon Sep 17 00:00:00 2001
From: Eric Feng <Eric.Feng at amd.com>
Date: Thu, 7 May 2026 16:38:13 -0700
Subject: [PATCH 3/5] Keep contract lowering filter patch focused
---
.../Vector/Transforms/LowerVectorContract.cpp | 70 ++++++++-----------
1 file changed, 28 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index 2adf71c5da324..8507dad5296ca 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -201,38 +201,6 @@ static Value createMul(Location loc, Value x, Value y, bool isInt,
return arith::MulFOp::create(rewriter, loc, x, y, fmf);
}
-static LogicalResult
-checkSameOperandAndAccumulatorElementType(vector::ContractionOp op,
- PatternRewriter &rewriter) {
- if (op.getLhsType().getElementType() ==
- getElementTypeOrSelf(op.getAccType()) &&
- op.getRhsType().getElementType() == getElementTypeOrSelf(op.getAccType()))
- return success();
-
- return rewriter.notifyMatchFailure(
- op, "mixed-mode contract lowering is not supported");
-}
-
-static LogicalResult checkAddKind(vector::ContractionOp op,
- PatternRewriter &rewriter) {
- if (op.getKind() == vector::CombiningKind::ADD)
- return success();
-
- return rewriter.notifyMatchFailure(
- op, "contractions other than 'add' not supported");
-}
-
-static bool isContractArithOpSupported(vector::CombiningKind kind, bool isInt) {
- using vector::CombiningKind;
- if (isInt)
- return kind != CombiningKind::MINNUMF && kind != CombiningKind::MAXNUMF &&
- kind != CombiningKind::MINIMUMF && kind != CombiningKind::MAXIMUMF;
- return kind != CombiningKind::AND && kind != CombiningKind::MINUI &&
- kind != CombiningKind::MINSI && kind != CombiningKind::MAXUI &&
- kind != CombiningKind::MAXSI && kind != CombiningKind::OR &&
- kind != CombiningKind::XOR;
-}
-
namespace {
/// Progressive lowering of a `vector.contract %a, %b, %c` with row-major matmul
@@ -683,8 +651,15 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
*vectorContractLowering != vector::VectorContractLowering::Dot)
return failure();
- if (failed(checkSameOperandAndAccumulatorElementType(op, rewriter)) ||
- failed(checkAddKind(op, rewriter)))
+ // TODO: support mixed mode contract lowering.
+ if (op.getLhsType().getElementType() !=
+ getElementTypeOrSelf(op.getAccType()) ||
+ op.getRhsType().getElementType() != getElementTypeOrSelf(op.getAccType()))
+ return failure();
+
+ // TODO: the code below assumes the default contraction, make sure it supports
+ // other kinds before enabling this lowering.
+ if (op.getKind() != vector::CombiningKind::ADD)
return failure();
VectorType dstType = dyn_cast<VectorType>(op.getResultType());
@@ -840,7 +815,11 @@ struct ContractOpToElementwise
vector::VectorContractLowering::ParallelArith)
return failure();
- if (failed(checkSameOperandAndAccumulatorElementType(contractOp, rewriter)))
+ // TODO: support mixed mode contract lowering.
+ if (contractOp.getLhsType().getElementType() !=
+ getElementTypeOrSelf(contractOp.getAccType()) ||
+ contractOp.getRhsType().getElementType() !=
+ getElementTypeOrSelf(contractOp.getAccType()))
return failure();
ArrayRef<int64_t> lhsShape = contractOp.getLhsType().getShape();
@@ -899,9 +878,6 @@ struct ContractOpToElementwise
}
}
bool isInt = contractOp.getLhsType().getElementType().isIntOrIndex();
- if (!isContractArithOpSupported(contractOp.getKind(), isInt))
- return failure();
-
Value newLhs = contractOp.getLhs();
Value newRhs = contractOp.getRhs();
Location loc = contractOp.getLoc();
@@ -927,8 +903,10 @@ struct ContractOpToElementwise
createContractArithOp(loc, newLhs, newRhs, contractOp.getAcc(),
contractOp.getKind(), rewriter, isInt,
/*mask=*/Value(), contractOp.getFastmathAttr());
- assert(result && "kind and type support should have been checked");
- return *result;
+ if (result)
+ return *result;
+
+ return failure();
}
private:
@@ -959,10 +937,18 @@ LogicalResult ContractionOpGenericLowering::matchSupportedGenericContraction(
if (failed(filter(op)))
return failure();
- if (failed(checkSameOperandAndAccumulatorElementType(op, rewriter)))
+ // TODO: support mixed mode contract lowering.
+ if (op.getLhsType().getElementType() !=
+ getElementTypeOrSelf(op.getAccType()) ||
+ op.getRhsType().getElementType() != getElementTypeOrSelf(op.getAccType()))
return failure();
- return checkAddKind(op, rewriter);
+ // TODO: the code below assumes the default contraction, make sure it supports
+ // other kinds before enabling this lowering.
+ if (op.getKind() != vector::CombiningKind::ADD)
+ return rewriter.notifyMatchFailure(
+ op, "contractions other than 'add' not supported");
+ return success();
}
FailureOr<Value> ContractionOpGenericLowering::matchAndRewriteMaskableOp(
>From 01120edb106f62185b359566a8bc165a9988f2c5 Mon Sep 17 00:00:00 2001
From: Eric Feng <Eric.Feng at amd.com>
Date: Fri, 8 May 2026 11:33:16 -0700
Subject: [PATCH 4/5] Reuse filter constraint naming for contract lowering
policy
---
.../Vector/Transforms/LoweringPatterns.h | 19 ++-
.../Vector/Transforms/LowerVectorContract.cpp | 126 ++++++------------
2 files changed, 49 insertions(+), 96 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
index 689fabb1413dc..187cd8b7e9b39 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
@@ -50,33 +50,32 @@ namespace vector {
/// Progressively lower a `vector.contract` with row-major matmul semantics to
/// linearized `vector.extract` + `vector.outerproduct` + `vector.insert`.
-/// A `VectorContractLoweringFilter` lets clients compose multiple lowering
-/// strategies by benefit. Returning failure means this strategy silently
-/// declines the op without consuming it or diagnosing invalid IR; lower-benefit
-/// strategies may still match the same op.
-using VectorContractLoweringFilter =
- std::function<LogicalResult(ContractionOp)>;
+/// A `FilterConstraintType` lets clients compose multiple lowering strategies
+/// by benefit. Returning failure means this strategy silently declines the op
+/// without consuming it or diagnosing invalid IR; lower-benefit strategies may
+/// still match the same op.
+using FilterConstraintType = std::function<LogicalResult(ContractionOp op)>;
LogicalResult acceptAllVectorContractLoweringFilter(ContractionOp op);
void populateVectorContractToDotPatterns(
RewritePatternSet &patterns,
- VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType filter = acceptAllVectorContractLoweringFilter,
PatternBenefit benefit = 1);
void populateVectorContractToOuterProductPatterns(
RewritePatternSet &patterns,
- VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType filter = acceptAllVectorContractLoweringFilter,
PatternBenefit benefit = 1);
void populateVectorContractToParallelArithPatterns(
RewritePatternSet &patterns,
- VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType filter = acceptAllVectorContractLoweringFilter,
PatternBenefit benefit = 1);
void populateVectorContractGenericLoweringPatterns(
RewritePatternSet &patterns,
- VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType filter = acceptAllVectorContractLoweringFilter,
PatternBenefit benefit = 1);
void populateVectorContractLoweringPatterns(
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index 8507dad5296ca..1debecf1b2e28 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -226,28 +226,19 @@ class ContractionOpToOuterProductOpLowering
ContractionOpToOuterProductOpLowering(
vector::VectorContractLowering vectorContractLowering,
MLIRContext *context, PatternBenefit benefit = 1,
- VectorContractLoweringFilter constraint =
- acceptAllVectorContractLoweringFilter)
+ FilterConstraintType constraint = acceptAllVectorContractLoweringFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
vectorContractLowering(vectorContractLowering),
filter(std::move(constraint)) {}
- ContractionOpToOuterProductOpLowering(
- MLIRContext *context,
- VectorContractLoweringFilter constraint =
- acceptAllVectorContractLoweringFilter,
- PatternBenefit benefit = 1)
- : MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
- filter(std::move(constraint)) {}
-
FailureOr<Value>
matchAndRewriteMaskableOp(vector::ContractionOp op, MaskingOpInterface maskOp,
PatternRewriter &rewriter) const override;
private:
/// Options to control the vector patterns.
- std::optional<vector::VectorContractLowering> vectorContractLowering;
- VectorContractLoweringFilter filter;
+ vector::VectorContractLowering vectorContractLowering;
+ FilterConstraintType filter;
};
/// Progressive lowering of a `vector.contract %a, %b, %c` with row-major matmul
@@ -276,27 +267,19 @@ class ContractionOpToDotLowering
ContractionOpToDotLowering(
vector::VectorContractLowering vectorContractLowering,
MLIRContext *context, PatternBenefit benefit = 1,
- VectorContractLoweringFilter constraint =
- acceptAllVectorContractLoweringFilter)
+ FilterConstraintType constraint = acceptAllVectorContractLoweringFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
vectorContractLowering(vectorContractLowering),
filter(std::move(constraint)) {}
- ContractionOpToDotLowering(MLIRContext *context,
- VectorContractLoweringFilter constraint =
- acceptAllVectorContractLoweringFilter,
- PatternBenefit benefit = 1)
- : MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
- filter(std::move(constraint)) {}
-
FailureOr<Value>
matchAndRewriteMaskableOp(vector::ContractionOp op, MaskingOpInterface maskOp,
PatternRewriter &rewriter) const override;
private:
/// Options to control the vector patterns.
- std::optional<vector::VectorContractLowering> vectorContractLowering;
- VectorContractLoweringFilter filter;
+ vector::VectorContractLowering vectorContractLowering;
+ FilterConstraintType filter;
};
/// Progressive lowering of ContractionOp.
@@ -318,10 +301,10 @@ class ContractionOpGenericLowering
public:
using MaskableOpRewritePattern::MaskableOpRewritePattern;
- ContractionOpGenericLowering(MLIRContext *context,
- VectorContractLoweringFilter constraint =
- acceptAllVectorContractLoweringFilter,
- PatternBenefit benefit = 1)
+ ContractionOpGenericLowering(
+ MLIRContext *context,
+ FilterConstraintType constraint = acceptAllVectorContractLoweringFilter,
+ PatternBenefit benefit = 1)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
filter(std::move(constraint)) {}
@@ -339,7 +322,7 @@ class ContractionOpGenericLowering
MaskingOpInterface maskOp) const;
private:
- VectorContractLoweringFilter filter;
+ FilterConstraintType filter;
// Lower one parallel dimension.
FailureOr<Value> lowerParallel(PatternRewriter &rewriter,
@@ -355,8 +338,7 @@ class ContractionOpLowering : public ContractionOpGenericLowering {
ContractionOpLowering(
vector::VectorContractLowering vectorContractLoweringOption,
MLIRContext *context, PatternBenefit benefit = 1,
- VectorContractLoweringFilter constraint =
- acceptAllVectorContractLoweringFilter)
+ FilterConstraintType constraint = acceptAllVectorContractLoweringFilter)
: ContractionOpGenericLowering(context, std::move(constraint), benefit),
vectorContractLoweringOption(vectorContractLoweringOption) {}
@@ -616,8 +598,7 @@ FailureOr<Value>
ContractionOpToOuterProductOpLowering::matchAndRewriteMaskableOp(
vector::ContractionOp op, MaskingOpInterface maskOp,
PatternRewriter &rewriter) const {
- if (vectorContractLowering &&
- *vectorContractLowering != vector::VectorContractLowering::OuterProduct)
+ if (vectorContractLowering != vector::VectorContractLowering::OuterProduct)
return failure();
if (failed(filter(op)))
@@ -647,26 +628,9 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
if (failed(filter(op)))
return failure();
- if (vectorContractLowering &&
- *vectorContractLowering != vector::VectorContractLowering::Dot)
- return failure();
-
- // TODO: support mixed mode contract lowering.
- if (op.getLhsType().getElementType() !=
- getElementTypeOrSelf(op.getAccType()) ||
- op.getRhsType().getElementType() != getElementTypeOrSelf(op.getAccType()))
- return failure();
-
- // TODO: the code below assumes the default contraction, make sure it supports
- // other kinds before enabling this lowering.
- if (op.getKind() != vector::CombiningKind::ADD)
+ if (vectorContractLowering != vector::VectorContractLowering::Dot)
return failure();
- VectorType dstType = dyn_cast<VectorType>(op.getResultType());
- if (!dstType || dstType.getRank() < 1 || dstType.getRank() > 2)
- return rewriter.notifyMatchFailure(
- op, "expected result type of rank 1 or 2 for dot lowering");
-
auto iteratorTypes = op.getIteratorTypes().getValue();
static constexpr std::array<int64_t, 2> perm = {1, 0};
Location loc = op.getLoc();
@@ -683,7 +647,7 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
// In the following we wish to make the reduction dimension innermost so we
// can load vectors and just fmul + reduce into a scalar.
//
- if (iteratorTypes.size() == 3 && isParallelIterator(iteratorTypes[0]) &&
+ if (isParallelIterator(iteratorTypes[0]) &&
isParallelIterator(iteratorTypes[1]) &&
isReductionIterator(iteratorTypes[2])) {
//
@@ -716,8 +680,7 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
} else {
return failure();
}
- } else if (iteratorTypes.size() == 2 &&
- isParallelIterator(iteratorTypes[0]) &&
+ } else if (isParallelIterator(iteratorTypes[0]) &&
isReductionIterator(iteratorTypes[1])) {
//
// One outer parallel, one inner reduction (matvec flavor)
@@ -738,6 +701,10 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
return failure();
}
+ VectorType dstType = cast<VectorType>(op.getResultType());
+ assert(dstType.getRank() >= 1 && dstType.getRank() <= 2 &&
+ "Expected dst type of rank 1 or 2");
+
unsigned rank = dstType.getRank();
unsigned dstRows = dstType.getShape()[0];
unsigned dstColumns = rank == 1 ? 1 : dstType.getShape()[1];
@@ -784,21 +751,14 @@ struct ContractOpToElementwise
: public MaskableOpRewritePattern<vector::ContractionOp> {
using MaskableOpRewritePattern::MaskableOpRewritePattern;
- ContractOpToElementwise(vector::VectorContractLowering vectorContractLowering,
- MLIRContext *context, PatternBenefit benefit = 1,
- VectorContractLoweringFilter constraint =
- acceptAllVectorContractLoweringFilter)
+ ContractOpToElementwise(
+ vector::VectorContractLowering vectorContractLowering,
+ MLIRContext *context, PatternBenefit benefit = 1,
+ FilterConstraintType constraint = acceptAllVectorContractLoweringFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
vectorContractLowering(vectorContractLowering),
filter(std::move(constraint)) {}
- ContractOpToElementwise(MLIRContext *context,
- VectorContractLoweringFilter constraint =
- acceptAllVectorContractLoweringFilter,
- PatternBenefit benefit = 1)
- : MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
- filter(std::move(constraint)) {}
-
FailureOr<Value>
matchAndRewriteMaskableOp(vector::ContractionOp contractOp,
MaskingOpInterface maskOp,
@@ -810,16 +770,7 @@ struct ContractOpToElementwise
if (failed(filter(contractOp)))
return failure();
- if (vectorContractLowering &&
- *vectorContractLowering !=
- vector::VectorContractLowering::ParallelArith)
- return failure();
-
- // TODO: support mixed mode contract lowering.
- if (contractOp.getLhsType().getElementType() !=
- getElementTypeOrSelf(contractOp.getAccType()) ||
- contractOp.getRhsType().getElementType() !=
- getElementTypeOrSelf(contractOp.getAccType()))
+ if (vectorContractLowering != vector::VectorContractLowering::ParallelArith)
return failure();
ArrayRef<int64_t> lhsShape = contractOp.getLhsType().getShape();
@@ -877,7 +828,6 @@ struct ContractOpToElementwise
rhsTranspose.push_back(rhsDims.size() - 1);
}
}
- bool isInt = contractOp.getLhsType().getElementType().isIntOrIndex();
Value newLhs = contractOp.getLhs();
Value newRhs = contractOp.getRhs();
Location loc = contractOp.getLoc();
@@ -893,6 +843,7 @@ struct ContractOpToElementwise
VectorType::get(rhsDims, contractOp.getRhsType().getElementType());
newRhs = vector::BroadcastOp::create(rewriter, loc, expandedType, newRhs);
}
+ bool isInt = contractOp.getLhsType().getElementType().isIntOrIndex();
newLhs = vector::TransposeOp::create(rewriter, loc, newLhs, lhsTranspose);
newRhs = vector::TransposeOp::create(rewriter, loc, newRhs, rhsTranspose);
SmallVector<int64_t> lhsOffsets(lhsReductionDims.size(), 0);
@@ -911,8 +862,8 @@ struct ContractOpToElementwise
private:
/// Options to control the vector patterns.
- std::optional<vector::VectorContractLowering> vectorContractLowering;
- VectorContractLoweringFilter filter;
+ vector::VectorContractLowering vectorContractLowering;
+ FilterConstraintType filter;
};
/// Progressive lowering of ContractionOp.
@@ -1306,28 +1257,31 @@ mlir::vector::acceptAllVectorContractLoweringFilter(ContractionOp) {
}
void mlir::vector::populateVectorContractToDotPatterns(
- RewritePatternSet &patterns, VectorContractLoweringFilter filter,
+ RewritePatternSet &patterns, FilterConstraintType filter,
PatternBenefit benefit) {
- patterns.add<ContractionOpToDotLowering>(patterns.getContext(),
- std::move(filter), benefit);
+ patterns.add<ContractionOpToDotLowering>(vector::VectorContractLowering::Dot,
+ patterns.getContext(), benefit,
+ std::move(filter));
}
void mlir::vector::populateVectorContractToOuterProductPatterns(
- RewritePatternSet &patterns, VectorContractLoweringFilter filter,
+ RewritePatternSet &patterns, FilterConstraintType filter,
PatternBenefit benefit) {
patterns.add<ContractionOpToOuterProductOpLowering>(
- patterns.getContext(), std::move(filter), benefit);
+ vector::VectorContractLowering::OuterProduct, patterns.getContext(),
+ benefit, std::move(filter));
}
void mlir::vector::populateVectorContractToParallelArithPatterns(
- RewritePatternSet &patterns, VectorContractLoweringFilter filter,
+ RewritePatternSet &patterns, FilterConstraintType filter,
PatternBenefit benefit) {
- patterns.add<ContractOpToElementwise>(patterns.getContext(),
- std::move(filter), benefit);
+ patterns.add<ContractOpToElementwise>(
+ vector::VectorContractLowering::ParallelArith, patterns.getContext(),
+ benefit, std::move(filter));
}
void mlir::vector::populateVectorContractGenericLoweringPatterns(
- RewritePatternSet &patterns, VectorContractLoweringFilter filter,
+ RewritePatternSet &patterns, FilterConstraintType filter,
PatternBenefit benefit) {
patterns.add<ContractionOpGenericLowering>(patterns.getContext(),
std::move(filter), benefit);
>From 0826ada10bb68850765e47d89f09964fc83f29f6 Mon Sep 17 00:00:00 2001
From: Eric Feng <Eric.Feng at amd.com>
Date: Fri, 8 May 2026 11:43:14 -0700
Subject: [PATCH 5/5] Rename vector contract default filter
---
.../Dialect/Vector/Transforms/LoweringPatterns.h | 10 +++++-----
.../Vector/Transforms/LowerVectorContract.cpp | 12 ++++++------
.../test/lib/Dialect/Vector/TestVectorTransforms.cpp | 4 ++--
3 files changed, 13 insertions(+), 13 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
index 187cd8b7e9b39..1c98b364d7e0f 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
@@ -56,26 +56,26 @@ namespace vector {
/// still match the same op.
using FilterConstraintType = std::function<LogicalResult(ContractionOp op)>;
-LogicalResult acceptAllVectorContractLoweringFilter(ContractionOp op);
+LogicalResult defaultFilter(ContractionOp op);
void populateVectorContractToDotPatterns(
RewritePatternSet &patterns,
- FilterConstraintType filter = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType filter = defaultFilter,
PatternBenefit benefit = 1);
void populateVectorContractToOuterProductPatterns(
RewritePatternSet &patterns,
- FilterConstraintType filter = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType filter = defaultFilter,
PatternBenefit benefit = 1);
void populateVectorContractToParallelArithPatterns(
RewritePatternSet &patterns,
- FilterConstraintType filter = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType filter = defaultFilter,
PatternBenefit benefit = 1);
void populateVectorContractGenericLoweringPatterns(
RewritePatternSet &patterns,
- FilterConstraintType filter = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType filter = defaultFilter,
PatternBenefit benefit = 1);
void populateVectorContractLoweringPatterns(
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index 1debecf1b2e28..9ae75fc5423f8 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -226,7 +226,7 @@ class ContractionOpToOuterProductOpLowering
ContractionOpToOuterProductOpLowering(
vector::VectorContractLowering vectorContractLowering,
MLIRContext *context, PatternBenefit benefit = 1,
- FilterConstraintType constraint = acceptAllVectorContractLoweringFilter)
+ FilterConstraintType constraint = defaultFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
vectorContractLowering(vectorContractLowering),
filter(std::move(constraint)) {}
@@ -267,7 +267,7 @@ class ContractionOpToDotLowering
ContractionOpToDotLowering(
vector::VectorContractLowering vectorContractLowering,
MLIRContext *context, PatternBenefit benefit = 1,
- FilterConstraintType constraint = acceptAllVectorContractLoweringFilter)
+ FilterConstraintType constraint = defaultFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
vectorContractLowering(vectorContractLowering),
filter(std::move(constraint)) {}
@@ -303,7 +303,7 @@ class ContractionOpGenericLowering
ContractionOpGenericLowering(
MLIRContext *context,
- FilterConstraintType constraint = acceptAllVectorContractLoweringFilter,
+ FilterConstraintType constraint = defaultFilter,
PatternBenefit benefit = 1)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
filter(std::move(constraint)) {}
@@ -338,7 +338,7 @@ class ContractionOpLowering : public ContractionOpGenericLowering {
ContractionOpLowering(
vector::VectorContractLowering vectorContractLoweringOption,
MLIRContext *context, PatternBenefit benefit = 1,
- FilterConstraintType constraint = acceptAllVectorContractLoweringFilter)
+ FilterConstraintType constraint = defaultFilter)
: ContractionOpGenericLowering(context, std::move(constraint), benefit),
vectorContractLoweringOption(vectorContractLoweringOption) {}
@@ -754,7 +754,7 @@ struct ContractOpToElementwise
ContractOpToElementwise(
vector::VectorContractLowering vectorContractLowering,
MLIRContext *context, PatternBenefit benefit = 1,
- FilterConstraintType constraint = acceptAllVectorContractLoweringFilter)
+ FilterConstraintType constraint = defaultFilter)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit),
vectorContractLowering(vectorContractLowering),
filter(std::move(constraint)) {}
@@ -1252,7 +1252,7 @@ class OuterProductOpLowering : public OpRewritePattern<vector::OuterProductOp> {
} // namespace
LogicalResult
-mlir::vector::acceptAllVectorContractLoweringFilter(ContractionOp) {
+mlir::vector::defaultFilter(ContractionOp) {
return success();
}
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index 95af6dff8eaeb..13f13b7fba60a 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -179,9 +179,9 @@ struct TestVectorContractLoweringComposition final
},
PatternBenefit(3));
populateVectorContractToOuterProductPatterns(
- patterns, acceptAllVectorContractLoweringFilter, PatternBenefit(2));
+ patterns, defaultFilter, PatternBenefit(2));
populateVectorContractGenericLoweringPatterns(
- patterns, acceptAllVectorContractLoweringFilter, PatternBenefit(1));
+ patterns, defaultFilter, PatternBenefit(1));
} else if (mode == "generic") {
populateVectorContractGenericLoweringPatterns(patterns);
} else if (mode == "parallel-arith") {
More information about the Mlir-commits
mailing list