[Mlir-commits] [mlir] [mlir][vector] add ApplyLowerMultiReductionTransformationPatternsOp tests (PR #180977)
Erick Ochoa Lopez
llvmlistbot at llvm.org
Wed Feb 11 11:49:01 PST 2026
https://github.com/amd-eochoalo updated https://github.com/llvm/llvm-project/pull/180977
>From 67ec909e00f17a2e641ff3588c66edeb3abc038a Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 11 Feb 2026 10:30:50 -0500
Subject: [PATCH 1/4] [mlir][vector] Add
apply_patterns.vector.lower_multi_reduction_transformation
Creates transform op for testing finer grained patterns.
---
.../Vector/TransformOps/VectorTransformOps.td | 23 +++++++++++++++++++
.../TransformOps/VectorTransformOps.cpp | 8 +++++++
2 files changed, 31 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
index c9668fe30e648..dfe54835ac89f 100644
--- a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
@@ -243,6 +243,29 @@ def ApplyLowerMultiReductionPatternsOp : Op<Transform_Dialect,
}];
}
+def ApplyLowerMultiReductionTransformationPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.vector.lower_multi_reduction_transformation",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Indicates that vector multi_reduction-like operations should be
+ transformed such that all reduction dimensions become innermost or
+ outermost, and 1-D reductions are lifted to 2-D.
+
+ This populates the patterns from
+ `populateVectorMultiReductionTransformationPatterns`, i.e.:
+ * `InnerOuterDimReductionConversion`
+ * `OneDimMultiReductionToTwoDim`
+ }];
+
+ let arguments = (ins DefaultValuedAttr<VectorMultiReductionLoweringAttr,
+ "vector::VectorMultiReductionLowering::InnerParallel">:$lowering_strategy
+ );
+
+ let assemblyFormat = [{
+ (`lowering_strategy` `=` $lowering_strategy^)? attr-dict
+ }];
+}
+
def ApplyLowerOuterProductPatternsOp : Op<Transform_Dialect,
"apply_patterns.vector.lower_outerproduct",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
index 5495f8abdb335..3e9af181195c3 100644
--- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
@@ -138,6 +138,14 @@ void transform::ApplyLowerMultiReductionPatternsOp::populatePatterns(
patterns, vectorTransformOptions.vectorMultiReductionLowering);
}
+void transform::ApplyLowerMultiReductionTransformationPatternsOp::
+ populatePatterns(RewritePatternSet &patterns) {
+ vector::VectorTransformsOptions vectorTransformOptions;
+ vectorTransformOptions.setVectorMultiReductionLowering(getLoweringStrategy());
+ vector::populateVectorMultiReductionTransformationPatterns(
+ patterns, vectorTransformOptions.vectorMultiReductionLowering);
+}
+
void transform::ApplyLowerOuterProductPatternsOp::populatePatterns(
RewritePatternSet &patterns) {
populateVectorOuterProductLoweringPatterns(patterns);
>From 93d41427c799013ca63cff74bbf1cd485ff8e28b Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 11 Feb 2026 10:51:05 -0500
Subject: [PATCH 2/4] [mlir][vector] Add test NFC
---
...vector-multi-reduction-transformation.mlir | 124 ++++++++++++++++++
1 file changed, 124 insertions(+)
create mode 100644 mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir
diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir
new file mode 100644
index 0000000000000..24009821a38b1
--- /dev/null
+++ b/mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir
@@ -0,0 +1,124 @@
+// RUN: mlir-opt %s --transform-interpreter='entry-point=innerreduction' | FileCheck %s --check-prefix=INNERREDUCTION
+// RUN: mlir-opt %s --transform-interpreter='entry-point=innerparallel' | FileCheck %s --check-prefix=INNERPARALLEL
+
+// INNERREDUCTION-LABEL: func @transpose_reduction_dims_innerreduction
+// INNERREDUCTION-SAME: %[[INPUT:.+]]: vector<3x2x4xf32>
+// INNERREDUCTION-SAME: %[[ACC:.+]]: vector<2x4xf32>
+func.func @transpose_reduction_dims_innerreduction(%arg0: vector<3x2x4xf32>, %acc: vector<2x4xf32>) -> vector<2x4xf32> {
+ // INNERREDUCTION: %[[TRANSPOSED:.+]] = vector.transpose %[[INPUT]], [1, 2, 0]
+ // INNERREDUCTION: %[[RESULT:.+]] = vector.multi_reduction <mul>, %[[TRANSPOSED]], %[[ACC]] [2]
+ %0 = vector.multi_reduction <mul>, %arg0, %acc [0] : vector<3x2x4xf32> to vector<2x4xf32>
+ // INNERREDUCTION: return %[[RESULT]]
+ return %0 : vector<2x4xf32>
+}
+
+// INNERPARALLEL-LABEL: func @transpose_reduction_dims_innerparallel
+// INNERPARALLEL-SAME: %[[INPUT:.+]]: vector<3x2x4xf32>
+// INNERPARALLEL-SAME: %[[ACC:.+]]: vector<3x2xf32>
+func.func @transpose_reduction_dims_innerparallel(%arg0: vector<3x2x4xf32>, %acc: vector<3x2xf32>) -> vector<3x2xf32> {
+ // INNERPARALLEL: %[[TRANSPOSED:.+]] = vector.transpose %[[INPUT]], [2, 0, 1]
+ // INNERPARALLEL: vector.multi_reduction <mul>, %[[TRANSPOSED]], %[[ACC]] [0]
+ %0 = vector.multi_reduction <mul>, %arg0, %acc [2] : vector<3x2x4xf32> to vector<3x2xf32>
+ return %0 : vector<3x2xf32>
+}
+
+// INNERREDUCTION-LABEL: func @transpose_multi_reduction_dims
+// INNERREDUCTION-SAME: %[[INPUT:.+]]: vector<2x3x4x5xf32>
+// INNERREDUCTION-SAME: %[[ACC:.+]]: vector<2x5xf32>
+func.func @transpose_multi_reduction_dims(%arg0: vector<2x3x4x5xf32>, %acc: vector<2x5xf32>) -> vector<2x5xf32> {
+ // INNERREDUCTION: %[[TRANSPOSED:.+]] = vector.transpose %[[INPUT]], [0, 3, 1, 2]
+ // INNERREDUCTION: %[[RESULT:.+]] = vector.multi_reduction <add>, %[[TRANSPOSED]], %[[ACC]] [2, 3]
+ %0 = vector.multi_reduction <add>, %arg0, %acc [1, 2] : vector<2x3x4x5xf32> to vector<2x5xf32>
+ // INNERREDUCTION: return %[[RESULT]]
+ return %0 : vector<2x5xf32>
+}
+
+// INNERREDUCTION-LABEL: func @transpose_parallel_middle
+// INNERREDUCTION-SAME: %[[INPUT:.+]]: vector<3x4x5xf32>
+// INNERREDUCTION-SAME: %[[ACC:.+]]: vector<4xf32>
+// INNERPARALLEL-LABEL: func @transpose_parallel_middle
+// INNERPARALLEL-SAME: %[[INPUT:.+]]: vector<3x4x5xf32>
+// INNERPARALLEL-SAME: %[[ACC:.+]]: vector<4xf32>
+func.func @transpose_parallel_middle(%arg0: vector<3x4x5xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+ // INNERREDUCTION: %[[TRANSPOSED:.+]] = vector.transpose %[[INPUT]], [1, 0, 2]
+ // INNERREDUCTION: %[[RESULT:.+]] = vector.multi_reduction <add>, %[[TRANSPOSED]], %[[ACC]] [1, 2]
+ // INNERPARALLEL: %[[TRANSPOSED:.+]] = vector.transpose %[[INPUT]], [0, 2, 1]
+ // INNERPARALLEL: vector.multi_reduction <add>, %[[TRANSPOSED]], %[[ACC]] [0, 1]
+ %0 = vector.multi_reduction <add>, %arg0, %acc [0, 2] : vector<3x4x5xf32> to vector<4xf32>
+ // INNERREDUCTION: return %[[RESULT]]
+ return %0 : vector<4xf32>
+}
+
+// INNERREDUCTION-LABEL: func @one_dim_to_two_dim_innerreduction
+// INNERREDUCTION-SAME: %[[INPUT:.+]]: vector<8xf32>
+// INNERREDUCTION-SAME: %[[ACC:.+]]: f32
+func.func @one_dim_to_two_dim_innerreduction(%arg0: vector<8xf32>, %acc: f32) -> f32 {
+ // INNERREDUCTION: %[[CAST:.+]] = vector.shape_cast %[[INPUT]] : vector<8xf32> to vector<1x8xf32>
+ // INNERREDUCTION: %[[BROADCAST:.+]] = vector.broadcast %[[ACC]] : f32 to vector<1xf32>
+ // INNERREDUCTION: %[[RESULT:.+]] = vector.multi_reduction <add>, %[[CAST]], %[[BROADCAST]] [1]
+ %0 = vector.multi_reduction <add>, %arg0, %acc [0] : vector<8xf32> to f32
+ // INNERREDUCTION: %[[EXTRACT:.+]] = vector.extract %[[RESULT]][0]
+ // INNERREDUCTION: return %[[EXTRACT]]
+ return %0 : f32
+}
+
+// INNERPARALLEL-LABEL: func @one_dim_to_two_dim_innerparallel
+// INNERPARALLEL-SAME: %[[INPUT:.+]]: vector<2xf32>
+// INNERPARALLEL-SAME: %[[ACC:.+]]: f32
+func.func @one_dim_to_two_dim_innerparallel(%arg0: vector<2xf32>, %acc: f32) -> f32 {
+ // INNERPARALLEL: %[[CAST:.+]] = vector.shape_cast %[[INPUT]] : vector<2xf32> to vector<1x2xf32>
+ // INNERPARALLEL: %[[BROADCAST:.+]] = vector.broadcast %[[ACC]] : f32 to vector<1xf32>
+ // INNERPARALLEL: %[[TRANSPOSED:.+]] = vector.transpose %[[CAST]], [1, 0]
+ // INNERPARALLEL: vector.multi_reduction <maxnumf>, %[[TRANSPOSED]], %[[BROADCAST]] [0]
+ %0 = vector.multi_reduction <maxnumf>, %arg0, %acc [0] : vector<2xf32> to f32
+ return %0 : f32
+}
+
+// INNERREDUCTION-LABEL: func @one_dim_to_two_dim_scalable
+// INNERREDUCTION-SAME: %[[INPUT:.+]]: vector<[4]xf32>
+// INNERREDUCTION-SAME: %[[ACC:.+]]: f32
+func.func @one_dim_to_two_dim_scalable(%arg0: vector<[4]xf32>, %acc: f32) -> f32 {
+ // INNERREDUCTION: %[[CAST:.+]] = vector.shape_cast %[[INPUT]] : vector<[4]xf32> to vector<1x[4]xf32>
+ // INNERREDUCTION: %[[BROADCAST:.+]] = vector.broadcast %[[ACC]] : f32 to vector<1xf32>
+ // INNERREDUCTION: %[[RESULT:.+]] = vector.multi_reduction <add>, %[[CAST]], %[[BROADCAST]] [1]
+ %0 = vector.multi_reduction <add>, %arg0, %acc [0] : vector<[4]xf32> to f32
+ // INNERREDUCTION: %[[EXTRACT:.+]] = vector.extract %[[RESULT]][0]
+ // INNERREDUCTION: return %[[EXTRACT]]
+ return %0 : f32
+}
+
+// INNERREDUCTION-LABEL: func @one_dim_to_two_dim_masked
+// INNERREDUCTION-SAME: %[[INPUT:.+]]: vector<8xf32>
+// INNERREDUCTION-SAME: %[[ACC:.+]]: f32
+// INNERREDUCTION-SAME: %[[MASK:.+]]: vector<8xi1>
+func.func @one_dim_to_two_dim_masked(%arg0: vector<8xf32>, %acc: f32, %mask: vector<8xi1>) -> f32 {
+ // INNERREDUCTION: %[[CAST:.+]] = vector.shape_cast %[[INPUT]] : vector<8xf32> to vector<1x8xf32>
+ // INNERREDUCTION: %[[BROADCAST_ACC:.+]] = vector.broadcast %[[ACC]] : f32 to vector<1xf32>
+ // INNERREDUCTION: %[[BROADCAST_MASK:.+]] = vector.broadcast %[[MASK]] : vector<8xi1> to vector<1x8xi1>
+ // INNERREDUCTION: %[[RESULT:.+]] = vector.mask %[[BROADCAST_MASK]] {
+ // INNERREDUCTION: vector.multi_reduction <add>, %[[CAST]], %[[BROADCAST_ACC]] [1]
+ %0 = vector.mask %mask {
+ vector.multi_reduction <add>, %arg0, %acc [0] : vector<8xf32> to f32
+ } : vector<8xi1> -> f32
+ // INNERREDUCTION: %[[EXTRACT:.+]] = vector.extract %[[RESULT]][0]
+ // INNERREDUCTION: return %[[EXTRACT]]
+ return %0 : f32
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @innerreduction(%root : !transform.any_op {transform.readonly}) {
+ %func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">
+ transform.apply_patterns to %func_op {
+ transform.apply_patterns.vector.lower_multi_reduction_transformation lowering_strategy = "innerreduction"
+ } : !transform.op<"func.func">
+ transform.yield
+ }
+
+ transform.named_sequence @innerparallel(%root : !transform.any_op {transform.readonly}) {
+ %func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">
+ transform.apply_patterns to %func_op {
+ transform.apply_patterns.vector.lower_multi_reduction_transformation lowering_strategy = "innerparallel"
+ } : !transform.op<"func.func">
+ transform.yield
+ }
+}
>From f3e3a3c927d1370023170ca85841c9ddbe084cf6 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 11 Feb 2026 10:58:53 -0500
Subject: [PATCH 3/4] [mlir][vector] Add test for python bindings.
Adds tests for ApplyLowerMultiReductionTransformationPatternsOp's python
bindings.
---
mlir/test/python/dialects/transform_vector_ext.py | 8 ++++++++
1 file changed, 8 insertions(+)
diff --git a/mlir/test/python/dialects/transform_vector_ext.py b/mlir/test/python/dialects/transform_vector_ext.py
index 2bcb2a2ac5812..6166007673c2e 100644
--- a/mlir/test/python/dialects/transform_vector_ext.py
+++ b/mlir/test/python/dialects/transform_vector_ext.py
@@ -100,6 +100,14 @@ def enum_configurable_patterns():
lowering_strategy=vector.VectorMultiReductionLowering.InnerReduction
)
+ # CHECK: transform.apply_patterns.vector.lower_multi_reduction_transformation
+ vector.ApplyLowerMultiReductionTransformationPatternsOp()
+ # CHECK: transform.apply_patterns.vector.lower_multi_reduction_transformation
+ # CHECK-SAME: lowering_strategy = innerreduction
+ vector.ApplyLowerMultiReductionTransformationPatternsOp(
+ lowering_strategy=vector.VectorMultiReductionLowering.InnerReduction
+ )
+
# CHECK: transform.apply_patterns.vector.lower_transpose
vector.ApplyLowerTransposePatternsOp()
# CHECK: transform.apply_patterns.vector.lower_transpose
>From 2e954ee260370058e7c562218f63fe5d6b8f5249 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 11 Feb 2026 14:48:39 -0500
Subject: [PATCH 4/4] Rename
---
.../Dialect/Vector/TransformOps/VectorTransformOps.td | 4 ++--
.../Dialect/Vector/TransformOps/VectorTransformOps.cpp | 2 +-
.../Vector/vector-multi-reduction-transformation.mlir | 4 ++--
mlir/test/python/dialects/transform_vector_ext.py | 8 ++++----
4 files changed, 9 insertions(+), 9 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
index dfe54835ac89f..0cc8058d99b99 100644
--- a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
@@ -243,8 +243,8 @@ def ApplyLowerMultiReductionPatternsOp : Op<Transform_Dialect,
}];
}
-def ApplyLowerMultiReductionTransformationPatternsOp : Op<Transform_Dialect,
- "apply_patterns.vector.lower_multi_reduction_transformation",
+def ApplyReorderMultiReductionDimsPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.vector.reorder_multi_reduction_dims",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
let description = [{
Indicates that vector multi_reduction-like operations should be
diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
index 3e9af181195c3..8b63a1c60c56b 100644
--- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
@@ -138,7 +138,7 @@ void transform::ApplyLowerMultiReductionPatternsOp::populatePatterns(
patterns, vectorTransformOptions.vectorMultiReductionLowering);
}
-void transform::ApplyLowerMultiReductionTransformationPatternsOp::
+void transform::ApplyReorderMultiReductionDimsPatternsOp::
populatePatterns(RewritePatternSet &patterns) {
vector::VectorTransformsOptions vectorTransformOptions;
vectorTransformOptions.setVectorMultiReductionLowering(getLoweringStrategy());
diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir
index 24009821a38b1..f849f2b88a166 100644
--- a/mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir
+++ b/mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir
@@ -109,7 +109,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @innerreduction(%root : !transform.any_op {transform.readonly}) {
%func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op {
- transform.apply_patterns.vector.lower_multi_reduction_transformation lowering_strategy = "innerreduction"
+ transform.apply_patterns.vector.reorder_multi_reduction_dims lowering_strategy = "innerreduction"
} : !transform.op<"func.func">
transform.yield
}
@@ -117,7 +117,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @innerparallel(%root : !transform.any_op {transform.readonly}) {
%func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op {
- transform.apply_patterns.vector.lower_multi_reduction_transformation lowering_strategy = "innerparallel"
+ transform.apply_patterns.vector.reorder_multi_reduction_dims lowering_strategy = "innerparallel"
} : !transform.op<"func.func">
transform.yield
}
diff --git a/mlir/test/python/dialects/transform_vector_ext.py b/mlir/test/python/dialects/transform_vector_ext.py
index 6166007673c2e..04ff316e76ec7 100644
--- a/mlir/test/python/dialects/transform_vector_ext.py
+++ b/mlir/test/python/dialects/transform_vector_ext.py
@@ -100,11 +100,11 @@ def enum_configurable_patterns():
lowering_strategy=vector.VectorMultiReductionLowering.InnerReduction
)
- # CHECK: transform.apply_patterns.vector.lower_multi_reduction_transformation
- vector.ApplyLowerMultiReductionTransformationPatternsOp()
- # CHECK: transform.apply_patterns.vector.lower_multi_reduction_transformation
+ # CHECK: transform.apply_patterns.vector.reorder_multi_reduction_dims
+ vector.ApplyReorderMultiReductionDimsPatternsOp()
+ # CHECK: transform.apply_patterns.vector.reorder_multi_reduction_dims
# CHECK-SAME: lowering_strategy = innerreduction
- vector.ApplyLowerMultiReductionTransformationPatternsOp(
+ vector.ApplyReorderMultiReductionDimsPatternsOp(
lowering_strategy=vector.VectorMultiReductionLowering.InnerReduction
)
More information about the Mlir-commits
mailing list