[Mlir-commits] [mlir] [mlir][vector] refactor vector.multi_reduction lowering tests (PR #180977)
Erick Ochoa Lopez
llvmlistbot at llvm.org
Wed Feb 11 09:26:09 PST 2026
https://github.com/amd-eochoalo created https://github.com/llvm/llvm-project/pull/180977
With the new finer grained populate methods introduced in 8dde3051504cb9ae42e654bbce39001f3946beea (#180750), there was a discussion about refactoring tests such that only one of the patterns applies at a time. This commit starts this process by adding the structure for one of these populate methods. The goal is for the populate methods to have their own file (each showing inner and outer reduction); deprecating populateVectorMultiReductionLoweringPatterns and ApplyLowerOuterProductPatternsOp; and removing the test file for mlir/test/Dialect/Vector/vector-multi-reduction-lowering.mlir
Co-authored-by: Claude <noreply at anthropic.com>
>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/3] [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/3] [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/3] [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
More information about the Mlir-commits
mailing list