[Mlir-commits] [mlir] [mlir][vector] refactor vector.multi_reduction lowering tests (PR #180977)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Feb 11 09:31:18 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Erick Ochoa Lopez (amd-eochoalo)
<details>
<summary>Changes</summary>
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
Assisted-By: claude-4.5-sonnet
---
Full diff: https://github.com/llvm/llvm-project/pull/180977.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td (+23)
- (modified) mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp (+8)
- (added) mlir/test/Dialect/Vector/vector-multi-reduction-transformation.mlir (+124)
- (modified) mlir/test/python/dialects/transform_vector_ext.py (+8)
``````````diff
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);
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
+ }
+}
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/180977
More information about the Mlir-commits
mailing list