[Mlir-commits] [mlir] [mlir][Vector] Add a dedicated Transform Dialect op for `UnrollContractPattern` (PR #230543)
Andrzej Warzyński
llvmlistbot at llvm.org
Fri Oct 9 08:00:44 PDT 2026
https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/230543
Adds a dedicated Transform Dialect op to run `UnrollContractPattern`.
The new op is currently restricted to matmul-like `vector.contract` ops
and has limitations similar to those of:
`transform.apply_patterns.gpu.unroll_vectors_subgroup_mma`
We could consider merging the two ops in the future.
While the new `transform.apply_patterns.vector.unroll_contract` op is
quite constrained, it suffices for my use case, so I have refrained from
generalising it. It takes three unroll sizes, corresponding to the M, N,
and K dimensions in GEMM/matmul terminology.
There are also a couple of minor NFC changes. For example,
vector-multi-reduction-unrolling.mlir
is renamed to
vector-multi-reduction-unroll.mlir
to align with the naming convention used by other tests for unrolling
patterns.
>From 4f35749f57f1861aa4cf73f0a9c5341b0c5d9cb8 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 9 Oct 2026 13:22:58 +0000
Subject: [PATCH] [mlir][Vector] Add a dedicated Transform Dialect op for
`UnrollContractPattern`
Adds a dedicated Transform Dialect op to run `UnrollContractPattern`.
The new op is currently restricted to matmul-like `vector.contract` ops
and has limitations similar to those of:
`transform.apply_patterns.gpu.unroll_vectors_subgroup_mma`
We could consider merging the two ops in the future.
While the new `transform.apply_patterns.vector.unroll_contract` op is
quite constrained, it suffices for my use case, so I have refrained from
generalising it. It takes three unroll sizes, corresponding to the M, N,
and K dimensions in GEMM/matmul terminology.
There are also a couple of minor NFC changes. For example,
vector-multi-reduction-unrolling.mlir
is renamed to
vector-multi-reduction-unroll.mlir
to align with the naming convention used by other tests for unrolling
patterns.
---
.../Vector/TransformOps/VectorTransformOps.td | 22 +++++++++
.../Vector/Transforms/VectorRewritePatterns.h | 6 +++
.../TransformOps/VectorTransformOps.cpp | 32 +++++++++++++
.../Vector/Transforms/VectorUnroll.cpp | 22 ++++++++-
.../Dialect/Vector/td/unroll-contract.mlir | 13 +++++
.../Vector/vector-contract-unroll.mlir | 48 +++++++++++++++++++
...lir => vector-multi-reduction-unroll.mlir} | 0
7 files changed, 141 insertions(+), 2 deletions(-)
create mode 100644 mlir/test/Dialect/Vector/td/unroll-contract.mlir
create mode 100644 mlir/test/Dialect/Vector/vector-contract-unroll.mlir
rename mlir/test/Dialect/Vector/{vector-multi-reduction-unrolling.mlir => vector-multi-reduction-unroll.mlir} (100%)
diff --git a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
index 3bd53bb69c576d..1b965937d6e6d2 100644
--- a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
@@ -14,6 +14,7 @@ include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
include "mlir/Dialect/Vector/Transforms/VectorTransformsBase.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Transform/IR/TransformTypes.td"
def ApplyVectorToLLVMConversionPatternsOp : Op<Transform_Dialect,
"apply_conversion_patterns.vector.vector_to_llvm",
@@ -324,6 +325,27 @@ def ApplyUnrollFromElementsPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplyUnrollContractPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.vector.unroll_contract",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Apply unroll patterns for `vector.contract` using the given shapes as the
+ target shape. If no `vector.contract` Op is present in the target block, no
+ unrolling takes place.
+
+ This op assumes that the target `vector.contract` Op models a matrix
+ multiplication. Generalising this Op to other cases is left as a TODO.
+ }];
+
+ let arguments = (ins I64Attr:$m,
+ I64Attr:$n,
+ I64Attr:$k);
+
+ let assemblyFormat = [{
+ `[` $m `,` $n `,` $k `]` attr-dict
+ }];
+}
+
def ApplyUnrollToElementsPatternsOp : Op<Transform_Dialect,
"apply_patterns.vector.unroll_to_elements",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
index 6363279f6826ba..9a67d40136c6dd 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
@@ -328,6 +328,12 @@ SmallVector<Value> sliceTransferIndices(ArrayRef<int64_t> elementOffsets,
AffineMap permutationMap, Location loc,
OpBuilder &builder);
+/// Unrolls 2 or more dimensional `vector.to_elements` ops by unrolling the
+/// outermost dimension of the operand.
+void populateVectorContractUnrollPatterns(RewritePatternSet &patterns,
+ const UnrollVectorOptions &options,
+ PatternBenefit benefit = 1);
+
/// Unrolls 2 or more dimensional `vector.to_elements` ops by unrolling the
/// outermost dimension of the operand.
void populateVectorToElementsUnrollPatterns(RewritePatternSet &patterns,
diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
index c4c930738d2c92..0d9e44cb7b9e5f 100644
--- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
@@ -168,6 +168,38 @@ void transform::ApplyUnrollFromElementsPatternsOp::populatePatterns(
vector::populateVectorFromElementsUnrollPatterns(patterns);
}
+void transform::ApplyUnrollContractPatternsOp::populatePatterns(
+ RewritePatternSet &patterns) {
+ int64_t m = getM();
+ int64_t n = getN();
+ int64_t k = getK();
+
+ auto nativeShapeFn =
+ [m, n, k](Operation *op) -> std::optional<SmallVector<int64_t>> {
+ return SmallVector<int64_t>({m, n, k});
+ };
+
+ vector::populateVectorContractUnrollPatterns(
+ patterns, vector::UnrollVectorOptions()
+ .setNativeShapeFn(nativeShapeFn)
+ .setFilterConstraint([](Operation *op) {
+ auto contract = dyn_cast<vector::ContractionOp>(op);
+ if (!contract)
+ return failure();
+
+ if (contract.getIteratorTypesArray() !=
+ SmallVector<vector::IteratorType>{
+ vector::IteratorType::parallel,
+ vector::IteratorType::parallel,
+ vector::IteratorType::reduction})
+ return failure();
+ auto type = dyn_cast<VectorType>(contract.getType());
+ if (!type || type.getRank() != 2)
+ return failure();
+ return success();
+ }));
+}
+
void transform::ApplyUnrollToElementsPatternsOp::populatePatterns(
RewritePatternSet &patterns) {
vector::populateVectorToElementsUnrollPatterns(patterns);
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index d58aedcc9d0234..abeeecbc5522e6 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -90,10 +90,17 @@ static Operation *cloneOpWithOperandsAndTypes(OpBuilder &builder, Location loc,
static std::optional<SmallVector<int64_t>>
getTargetShape(const vector::UnrollVectorOptions &options, Operation *op) {
LDBG() << "Get unroll shape for op " << op->getName().getStringRef();
- if (options.filterConstraint && failed(options.filterConstraint(op))) {
+
+ if (!options.filterConstraint) {
LDBG() << "--no filter constraint -> BAIL";
return std::nullopt;
}
+
+ if (failed(options.filterConstraint(op))) {
+ LDBG() << "--filtering failed -> BAIL";
+ return std::nullopt;
+ }
+
assert(options.nativeShape &&
"vector unrolling expects the native shape or native"
"shape call back function to be set");
@@ -284,8 +291,12 @@ struct UnrollContractionPattern
LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
PatternRewriter &rewriter) const override {
auto targetShape = getTargetShape(options, contractOp);
- if (!targetShape)
+ if (!targetShape) {
+ // targetShape = SmallVector<int64_t>({1, 8, 4});
return failure();
+ }
+ contractOp.getResultType();
+ // return failure();
auto dstVecType = cast<VectorType>(contractOp.getResultType());
SmallVector<int64_t> originalSize = *contractOp.getShapeForUnroll();
@@ -1728,6 +1739,13 @@ void mlir::vector::populateVectorUnrollPatterns(
benefit);
}
+void mlir::vector::populateVectorContractUnrollPatterns(
+ RewritePatternSet &patterns, const UnrollVectorOptions &options,
+ PatternBenefit benefit) {
+ patterns.add<UnrollContractionPattern>(patterns.getContext(), options,
+ benefit);
+}
+
void mlir::vector::populateVectorToElementsUnrollPatterns(
RewritePatternSet &patterns, PatternBenefit benefit) {
patterns.add<UnrollToElements>(patterns.getContext(), UnrollVectorOptions(),
diff --git a/mlir/test/Dialect/Vector/td/unroll-contract.mlir b/mlir/test/Dialect/Vector/td/unroll-contract.mlir
new file mode 100644
index 00000000000000..c88a272d960639
--- /dev/null
+++ b/mlir/test/Dialect/Vector/td/unroll-contract.mlir
@@ -0,0 +1,13 @@
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @unroll_contract(%module_op: !transform.any_op {transform.readonly}) {
+
+ %func = transform.structured.match ops{["func.func"]} in %module_op : (!transform.any_op) -> !transform.any_op
+ %contract = transform.structured.match ops{["vector.contract"]} in %func : (!transform.any_op) -> !transform.any_op
+
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.unroll_contract [1, 8, 4]
+ } : !transform.any_op
+
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Vector/vector-contract-unroll.mlir b/mlir/test/Dialect/Vector/vector-contract-unroll.mlir
new file mode 100644
index 00000000000000..1f95532868fd86
--- /dev/null
+++ b/mlir/test/Dialect/Vector/vector-contract-unroll.mlir
@@ -0,0 +1,48 @@
+// RUN: mlir-opt %s -transform-preload-library='transform-library-paths=%p/td/unroll-contract.mlir' \
+// RUN: -transform-interpreter=entry-point=unroll_contract --split-input-file | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// Test UnrollContractionPattern
+//===----------------------------------------------------------------------===//
+// CHECK-LABEL: @matmul
+func.func @matmul(%mask: vector<8x8x4xi1>, %lhs: vector<8x4xf32>, %rhs: vector<4x8xf32>, %acc: vector<8x8xf32>) -> vector<8x8xf32> {
+
+// CHECK-REPEAT-8: vector.contract {{.*}} : vector<1x4xf32>, vector<4x8xf32> into vector<1x8xf32>
+ %res = vector.contract {
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>, affine_map<(m, n, k) -> (k, n)>, affine_map<(m, n, k) -> (m, n)>],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %lhs, %rhs, %acc : vector<8x4xf32>, vector<4x8xf32> into vector<8x8xf32>
+
+ return %res : vector<8x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @negative_batched_matmul
+func.func @negative_batched_matmul(%mask: vector<1x8x8x4xi1>, %lhs: vector<1x8x4xf32>, %rhs: vector<1x4x8xf32>, %acc: vector<1x8x8xf32>) -> vector<1x8x8xf32> {
+
+// CHECK: vector.contract {{.*}} : vector<1x8x4xf32>, vector<1x4x8xf32> into vector<1x8x8xf32>
+ %res = vector.contract {
+ indexing_maps = [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)>],
+ iterator_types = ["parallel", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %lhs, %rhs, %acc : vector<1x8x4xf32>, vector<1x4x8xf32> into vector<1x8x8xf32>
+
+ return %res : vector<1x8x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @negative_matmul_transpose_a
+func.func @negative_matmul_transpose_a(%mask: vector<8x8x4xi1>, %lhs: vector<4x8xf32>, %rhs: vector<4x8xf32>, %acc: vector<8x8xf32>) -> vector<8x8xf32> {
+
+// CHECK: vector.contract {{.*}} : vector<4x8xf32>, vector<4x8xf32> into vector<8x8xf32>
+ %res = vector.contract {
+ indexing_maps = [affine_map<(k, m, n) -> (k, m)>, affine_map<(k, m, n) -> (k, n)>, affine_map<(k, m, n) -> (m, n)>],
+ iterator_types = ["reduction", "parallel", "parallel"],
+ kind = #vector.kind<add>}
+ %lhs, %rhs, %acc : vector<4x8xf32>, vector<4x8xf32> into vector<8x8xf32>
+
+ return %res : vector<8x8xf32>
+}
diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-unrolling.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-unroll.mlir
similarity index 100%
rename from mlir/test/Dialect/Vector/vector-multi-reduction-unrolling.mlir
rename to mlir/test/Dialect/Vector/vector-multi-reduction-unroll.mlir
More information about the Mlir-commits
mailing list