[Mlir-commits] [mlir] [mlir][Vector] Reject scalable outer dim in multi_reduction unrolling (PR #224174)
Cheng-Chieh Liu
llvmlistbot at llvm.org
Tue Nov 3 07:53:52 PST 2026
https://github.com/jay87406 updated https://github.com/llvm/llvm-project/pull/224174
>From cff5c6387deafa346d73dc89b193ef7ccbd4b60a Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Wed, 16 Sep 2026 04:49:31 -0500
Subject: [PATCH 1/3] [mlir][Vector] Reject scalable outer dim in
multi_reduction unrolling
TwoDimMultiReductionToReduction and TwoDimMultiReductionToElementWise
lower a rank-2 vector.multi_reduction into a sequence of per-lane
vector.extract/vector.reduction (or arith) ops, using
getSourceVectorType().getShape()[0] as the loop trip count. For a
scalable outer dimension like vector<[2]x4xf32>, that's only the
static minimum, not the actual vscale-scaled extent, so lanes past
the fixed trip count were silently dropped (or excluded from the
reduction) whenever vscale > 1 -- no error, just a wrong result.
Reject a scalable outer dimension in both patterns, and add negative
lit coverage for each lowering strategy.
---
.../Transforms/LowerVectorMultiReduction.cpp | 8 ++++++++
.../vector-multi-reduction-pass-lowering.mlir | 18 ++++++++++++++++++
2 files changed, 26 insertions(+)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
index 76599822fbfe40..02aa8df70a7d26 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
@@ -354,6 +354,10 @@ struct TwoDimMultiReductionToElementWise
if (multiReductionOp.isReducedDim(1) || !multiReductionOp.isReducedDim(0))
return failure();
+ if (multiReductionOp.getSourceVectorType().getScalableDims()[0])
+ return rewriter.notifyMatchFailure(
+ multiReductionOp, "cannot scalarize a scalable outer dimension");
+
Value mask = maskingOp ? maskingOp.getMask() : Value();
auto loc = multiReductionOp.getLoc();
@@ -411,6 +415,10 @@ struct TwoDimMultiReductionToReduction
if (multiReductionOp.isReducedDim(0) || !multiReductionOp.isReducedDim(1))
return failure();
+ if (multiReductionOp.getSourceVectorType().getScalableDims()[0])
+ return rewriter.notifyMatchFailure(
+ multiReductionOp, "cannot scalarize a scalable outer dimension");
+
Value mask = maskingOp ? maskingOp.getMask() : nullptr;
auto loc = multiReductionOp.getLoc();
diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
index ddbc5c7bdb2c08..453684535a1a64 100644
--- a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
@@ -61,3 +61,21 @@ func.func @vector_multi_reduction_masked(%arg0: vector<2x4xf32>, %acc: vector<2x
// INNER-PARALLEL: %[[INNERMASK:.+]] = vector.extract %[[TPMASK]][0] : vector<2xi1> from vector<4x2xi1>
// INNER-PARALLEL: %[[REDUCED:.+]] = arith.mulf %[[INNERVEC]], %[[ACC]] : vector<2xf32>
// INNER-PARALLEL: arith.select %[[INNERMASK]], %[[REDUCED]], %[[ACC]] : vector<2xi1>, vector<2xf32>
+
+// -----
+
+// INNER-REDUCTION-LABEL: @negative_multi_reduction_scalable_outer
+// INNER-REDUCTION: vector.multi_reduction <mul>, %{{.*}}, %{{.*}} [1] : vector<[2]x4xf32> to vector<[2]xf32>
+func.func @negative_multi_reduction_scalable_outer(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>) -> vector<[2]xf32> {
+ %0 = vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32>
+ return %0 : vector<[2]xf32>
+}
+
+// -----
+
+// INNER-PARALLEL-LABEL: @negative_multi_reduction_scalable_outer_reduce
+// INNER-PARALLEL: vector.multi_reduction <mul>, %{{.*}}, %{{.*}} [0] : vector<[2]x4xf32> to vector<4xf32>
+func.func @negative_multi_reduction_scalable_outer_reduce(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+ %0 = vector.multi_reduction <mul>, %arg0, %acc [0] : vector<[2]x4xf32> to vector<4xf32>
+ return %0 : vector<4xf32>
+}
>From 85eb4531a16b86826f409044f3e699041f16a4f8 Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Sat, 19 Sep 2026 04:22:23 -0500
Subject: [PATCH 2/3] Check the other lowering strategy succeeds, add masked
negative tests
Per review: show that rejecting the scalable outer dim in one
lowering strategy doesn't drop functionality, since the other
strategy still lowers the same input correctly. Also add masked
variants of both negative tests.
---
.../vector-multi-reduction-pass-lowering.mlir | 34 ++++++++++++++++---
1 file changed, 30 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
index 453684535a1a64..0496e9fa750cec 100644
--- a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
@@ -64,8 +64,10 @@ func.func @vector_multi_reduction_masked(%arg0: vector<2x4xf32>, %acc: vector<2x
// -----
-// INNER-REDUCTION-LABEL: @negative_multi_reduction_scalable_outer
-// INNER-REDUCTION: vector.multi_reduction <mul>, %{{.*}}, %{{.*}} [1] : vector<[2]x4xf32> to vector<[2]xf32>
+// ALL-LABEL: @negative_multi_reduction_scalable_outer
+// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>
+// INNER-REDUCTION: vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32>
+// INNER-PARALLEL: vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
func.func @negative_multi_reduction_scalable_outer(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>) -> vector<[2]xf32> {
%0 = vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32>
return %0 : vector<[2]xf32>
@@ -73,9 +75,33 @@ func.func @negative_multi_reduction_scalable_outer(%arg0: vector<[2]x4xf32>, %ac
// -----
-// INNER-PARALLEL-LABEL: @negative_multi_reduction_scalable_outer_reduce
-// INNER-PARALLEL: vector.multi_reduction <mul>, %{{.*}}, %{{.*}} [0] : vector<[2]x4xf32> to vector<4xf32>
+// ALL-LABEL: @negative_multi_reduction_scalable_outer_reduce
+// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<4xf32>
+// INNER-REDUCTION: vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
+// INNER-PARALLEL: vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [0] : vector<[2]x4xf32> to vector<4xf32>
func.func @negative_multi_reduction_scalable_outer_reduce(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
%0 = vector.multi_reduction <mul>, %arg0, %acc [0] : vector<[2]x4xf32> to vector<4xf32>
return %0 : vector<4xf32>
}
+
+// -----
+
+// ALL-LABEL: @negative_multi_reduction_scalable_outer_masked
+// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>, %[[MASK:.+]]: vector<[2]x4xi1>
+// INNER-REDUCTION: vector.mask %[[MASK]] { vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32> } : vector<[2]x4xi1> -> vector<[2]xf32>
+// INNER-PARALLEL: vector.transpose %[[MASK]], [1, 0] : vector<[2]x4xi1> to vector<4x[2]xi1>
+func.func @negative_multi_reduction_scalable_outer_masked(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>, %mask: vector<[2]x4xi1>) -> vector<[2]xf32> {
+ %0 = vector.mask %mask { vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32> } : vector<[2]x4xi1> -> vector<[2]xf32>
+ return %0 : vector<[2]xf32>
+}
+
+// -----
+
+// ALL-LABEL: @negative_multi_reduction_scalable_outer_reduce_masked
+// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<4xf32>, %[[MASK:.+]]: vector<[2]x4xi1>
+// INNER-REDUCTION: vector.transpose %[[MASK]], [1, 0] : vector<[2]x4xi1> to vector<4x[2]xi1>
+// INNER-PARALLEL: vector.mask %[[MASK]] { vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [0] : vector<[2]x4xf32> to vector<4xf32> } : vector<[2]x4xi1> -> vector<4xf32>
+func.func @negative_multi_reduction_scalable_outer_reduce_masked(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>, %mask: vector<[2]x4xi1>) -> vector<4xf32> {
+ %0 = vector.mask %mask { vector.multi_reduction <mul>, %arg0, %acc [0] : vector<[2]x4xf32> to vector<4xf32> } : vector<[2]x4xi1> -> vector<4xf32>
+ return %0 : vector<4xf32>
+}
>From 805f48c5e6b9d0765e7ea3b433eb9d038750edfd Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Tue, 3 Nov 2026 15:53:29 +0000
Subject: [PATCH 3/3] Document the scalable outer dim bail, check full lowering
in tests
Per review: add a comment and a TODO above both bails explaining
why a scalable outer dim is rejected, and mention it in the docs for
populateVectorMultiReductionUnrollingPatterns. In the tests, check
the unrolled sequence through to the returned value for the strategy
that lowers each input, instead of stopping at the transpose. Drop
the negative_ prefix, since each test only bails for one strategy,
and name the tests after which dim is scalable.
---
.../Vector/Transforms/LoweringPatterns.h | 2 +
.../Transforms/LowerVectorMultiReduction.cpp | 8 ++
.../vector-multi-reduction-pass-lowering.mlir | 102 ++++++++++++++----
3 files changed, 92 insertions(+), 20 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
index 556979ad3e9be0..648ab84066d6b9 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
@@ -96,6 +96,8 @@ void populateVectorMultiReductionFlatteningPatterns(
/// dimension, unroll the outer dimension to obtain a sequence of extract +
/// vector.reduction + insert. This can further lower to horizontal reduction
/// ops.
+///
+/// Both 2-D patterns bail out if the outer (unrolled) dimension is scalable.
void populateVectorMultiReductionUnrollingPatterns(
RewritePatternSet &patterns, VectorMultiReductionLowering options,
PatternBenefit benefit = 1);
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
index d444584a1e20a4..23167aafdf50db 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
@@ -356,6 +356,10 @@ struct TwoDimMultiReductionToElementWise
if (multiReductionOp.isReducedDim(1) || !multiReductionOp.isReducedDim(0))
return failure();
+ // The outer dim is unrolled using its static size as the trip count. For
+ // a scalable dim that is only the minimum size (vscale = 1), so lanes
+ // would be silently dropped when vscale > 1.
+ // TODO: Support scalable outer dims.
if (multiReductionOp.getSourceVectorType().getScalableDims()[0])
return rewriter.notifyMatchFailure(
multiReductionOp, "cannot scalarize a scalable outer dimension");
@@ -418,6 +422,10 @@ struct TwoDimMultiReductionToReduction
if (multiReductionOp.isReducedDim(0) || !multiReductionOp.isReducedDim(1))
return failure();
+ // The outer dim is unrolled using its static size as the trip count. For
+ // a scalable dim that is only the minimum size (vscale = 1), so lanes
+ // would be silently dropped when vscale > 1.
+ // TODO: Support scalable outer dims.
if (multiReductionOp.getSourceVectorType().getScalableDims()[0])
return rewriter.notifyMatchFailure(
multiReductionOp, "cannot scalarize a scalable outer dimension");
diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
index a0b9f336039ce6..8c9c7693668fd1 100644
--- a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
@@ -64,47 +64,109 @@ func.func @vector_multi_reduction_masked(%arg0: vector<2x4xf32>, %acc: vector<2x
// -----
-// ALL-LABEL: @negative_multi_reduction_scalable_outer
-// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>
-// INNER-REDUCTION: vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32>
-// INNER-PARALLEL: vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
-func.func @negative_multi_reduction_scalable_outer(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>) -> vector<[2]xf32> {
+// Reduces the fixed inner dim, the scalable outer dim is parallel.
+// * inner-reduction: TwoDimMultiReductionToReduction would have to unroll the
+// scalable outer dim, so it bails and the op is left as is.
+// * inner-parallel: the op is transposed to vector<4x[2]xf32>, and the fixed
+// outer dim is then unrolled by TwoDimMultiReductionToElementWise.
+func.func @vector_multi_reduction_scalable_parallel_dim(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>) -> vector<[2]xf32> {
%0 = vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32>
return %0 : vector<[2]xf32>
}
+// ALL-LABEL: func @vector_multi_reduction_scalable_parallel_dim
+// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>
+// INNER-REDUCTION: %[[RES:.+]] = vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32>
+// INNER-REDUCTION: return %[[RES]] : vector<[2]xf32>
+
+// INNER-PARALLEL: %[[TRANSPOSED:.+]] = vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
+// INNER-PARALLEL: %[[V0:.+]] = vector.extract %[[TRANSPOSED]][0] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-PARALLEL: %[[RV0:.+]] = arith.mulf %[[V0]], %[[ACC]] : vector<[2]xf32>
+// INNER-PARALLEL: %[[V1:.+]] = vector.extract %[[TRANSPOSED]][1] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-PARALLEL: %[[RV01:.+]] = arith.mulf %[[V1]], %[[RV0]] : vector<[2]xf32>
+// INNER-PARALLEL: %[[V2:.+]] = vector.extract %[[TRANSPOSED]][2] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-PARALLEL: %[[RV012:.+]] = arith.mulf %[[V2]], %[[RV01]] : vector<[2]xf32>
+// INNER-PARALLEL: %[[V3:.+]] = vector.extract %[[TRANSPOSED]][3] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-PARALLEL: %[[RESULT_VEC:.+]] = arith.mulf %[[V3]], %[[RV012]] : vector<[2]xf32>
+// INNER-PARALLEL: return %[[RESULT_VEC]] : vector<[2]xf32>
// -----
-// ALL-LABEL: @negative_multi_reduction_scalable_outer_reduce
-// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<4xf32>
-// INNER-REDUCTION: vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
-// INNER-PARALLEL: vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [0] : vector<[2]x4xf32> to vector<4xf32>
-func.func @negative_multi_reduction_scalable_outer_reduce(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+// Reduces the scalable outer dim, the fixed inner dim is parallel.
+// * inner-reduction: the op is transposed to vector<4x[2]xf32>, and the fixed
+// outer dim is then unrolled by TwoDimMultiReductionToReduction.
+// * inner-parallel: TwoDimMultiReductionToElementWise would have to unroll the
+// scalable outer dim, so it bails and the op is left as is.
+func.func @vector_multi_reduction_scalable_reduction_dim(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
%0 = vector.multi_reduction <mul>, %arg0, %acc [0] : vector<[2]x4xf32> to vector<4xf32>
return %0 : vector<4xf32>
}
+// ALL-LABEL: func @vector_multi_reduction_scalable_reduction_dim
+// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<4xf32>
+// INNER-REDUCTION: %[[TRANSPOSED:.+]] = vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
+// INNER-REDUCTION: %[[V0:.+]] = vector.extract %[[TRANSPOSED]][0] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-REDUCTION: %[[ACC0:.+]] = vector.extract %[[ACC]][0] : f32 from vector<4xf32>
+// INNER-REDUCTION: %[[RV0:.+]] = vector.reduction <mul>, %[[V0]], %[[ACC0]] : vector<[2]xf32> into f32
+// INNER-REDUCTION: %[[RESULT_VEC_0:.+]] = vector.insert %[[RV0]], %{{.+}} [0] : f32 into vector<4xf32>
+// INNER-REDUCTION: %[[V1:.+]] = vector.extract %[[TRANSPOSED]][1] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-REDUCTION: %[[ACC1:.+]] = vector.extract %[[ACC]][1] : f32 from vector<4xf32>
+// INNER-REDUCTION: %[[RV1:.+]] = vector.reduction <mul>, %[[V1]], %[[ACC1]] : vector<[2]xf32> into f32
+// INNER-REDUCTION: %[[RESULT_VEC_1:.+]] = vector.insert %[[RV1]], %[[RESULT_VEC_0]] [1] : f32 into vector<4xf32>
+// INNER-REDUCTION: %[[V2:.+]] = vector.extract %[[TRANSPOSED]][2] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-REDUCTION: %[[ACC2:.+]] = vector.extract %[[ACC]][2] : f32 from vector<4xf32>
+// INNER-REDUCTION: %[[RV2:.+]] = vector.reduction <mul>, %[[V2]], %[[ACC2]] : vector<[2]xf32> into f32
+// INNER-REDUCTION: %[[RESULT_VEC_2:.+]] = vector.insert %[[RV2]], %[[RESULT_VEC_1]] [2] : f32 into vector<4xf32>
+// INNER-REDUCTION: %[[V3:.+]] = vector.extract %[[TRANSPOSED]][3] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-REDUCTION: %[[ACC3:.+]] = vector.extract %[[ACC]][3] : f32 from vector<4xf32>
+// INNER-REDUCTION: %[[RV3:.+]] = vector.reduction <mul>, %[[V3]], %[[ACC3]] : vector<[2]xf32> into f32
+// INNER-REDUCTION: %[[RESULT_VEC:.+]] = vector.insert %[[RV3]], %[[RESULT_VEC_2]] [3] : f32 into vector<4xf32>
+// INNER-REDUCTION: return %[[RESULT_VEC]] : vector<4xf32>
+
+// INNER-PARALLEL: %[[RES:.+]] = vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [0] : vector<[2]x4xf32> to vector<4xf32>
+// INNER-PARALLEL: return %[[RES]] : vector<4xf32>
// -----
-// ALL-LABEL: @negative_multi_reduction_scalable_outer_masked
-// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>, %[[MASK:.+]]: vector<[2]x4xi1>
-// INNER-REDUCTION: vector.mask %[[MASK]] { vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32> } : vector<[2]x4xi1> -> vector<[2]xf32>
-// INNER-PARALLEL: vector.transpose %[[MASK]], [1, 0] : vector<[2]x4xi1> to vector<4x[2]xi1>
-func.func @negative_multi_reduction_scalable_outer_masked(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>, %mask: vector<[2]x4xi1>) -> vector<[2]xf32> {
+// Masked version of @vector_multi_reduction_scalable_parallel_dim.
+func.func @vector_multi_reduction_scalable_parallel_dim_masked(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>, %mask: vector<[2]x4xi1>) -> vector<[2]xf32> {
%0 = vector.mask %mask { vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32> } : vector<[2]x4xi1> -> vector<[2]xf32>
return %0 : vector<[2]xf32>
}
+// ALL-LABEL: func @vector_multi_reduction_scalable_parallel_dim_masked
+// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>, %[[MASK:.+]]: vector<[2]x4xi1>
+// INNER-REDUCTION: %[[RES:.+]] = vector.mask %[[MASK]] { vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32> } : vector<[2]x4xi1> -> vector<[2]xf32>
+// INNER-REDUCTION: return %[[RES]] : vector<[2]xf32>
+
+// INNER-PARALLEL: %[[TPMASK:.+]] = vector.transpose %[[MASK]], [1, 0] : vector<[2]x4xi1> to vector<4x[2]xi1>
+// INNER-PARALLEL: %[[TPINPUT:.+]] = vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
+// INNER-PARALLEL: %[[INNERVEC:.+]] = vector.extract %[[TPINPUT]][0] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-PARALLEL: %[[INNERMASK:.+]] = vector.extract %[[TPMASK]][0] : vector<[2]xi1> from vector<4x[2]xi1>
+// INNER-PARALLEL: %[[REDUCED:.+]] = arith.mulf %[[INNERVEC]], %[[ACC]] : vector<[2]xf32>
+// INNER-PARALLEL: arith.select %[[INNERMASK]], %[[REDUCED]], %[[ACC]] : vector<[2]xi1>, vector<[2]xf32>
+// INNER-PARALLEL-COUNT-2: arith.select %{{.+}}, %{{.+}}, %{{.+}} : vector<[2]xi1>, vector<[2]xf32>
+// INNER-PARALLEL: %[[RES:.+]] = arith.select %{{.+}}, %{{.+}}, %{{.+}} : vector<[2]xi1>, vector<[2]xf32>
+// INNER-PARALLEL: return %[[RES]] : vector<[2]xf32>
// -----
-// ALL-LABEL: @negative_multi_reduction_scalable_outer_reduce_masked
-// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<4xf32>, %[[MASK:.+]]: vector<[2]x4xi1>
-// INNER-REDUCTION: vector.transpose %[[MASK]], [1, 0] : vector<[2]x4xi1> to vector<4x[2]xi1>
-// INNER-PARALLEL: vector.mask %[[MASK]] { vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [0] : vector<[2]x4xf32> to vector<4xf32> } : vector<[2]x4xi1> -> vector<4xf32>
-func.func @negative_multi_reduction_scalable_outer_reduce_masked(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>, %mask: vector<[2]x4xi1>) -> vector<4xf32> {
+// Masked version of @vector_multi_reduction_scalable_reduction_dim.
+func.func @vector_multi_reduction_scalable_reduction_dim_masked(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>, %mask: vector<[2]x4xi1>) -> vector<4xf32> {
%0 = vector.mask %mask { vector.multi_reduction <mul>, %arg0, %acc [0] : vector<[2]x4xf32> to vector<4xf32> } : vector<[2]x4xi1> -> vector<4xf32>
return %0 : vector<4xf32>
}
+// ALL-LABEL: func @vector_multi_reduction_scalable_reduction_dim_masked
+// ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<4xf32>, %[[MASK:.+]]: vector<[2]x4xi1>
+// INNER-REDUCTION: %[[TPMASK:.+]] = vector.transpose %[[MASK]], [1, 0] : vector<[2]x4xi1> to vector<4x[2]xi1>
+// INNER-REDUCTION: %[[TPINPUT:.+]] = vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
+// INNER-REDUCTION: %[[INNERVEC:.+]] = vector.extract %[[TPINPUT]][0] : vector<[2]xf32> from vector<4x[2]xf32>
+// INNER-REDUCTION: %[[INNERACC:.+]] = vector.extract %[[ACC]][0] : f32 from vector<4xf32>
+// INNER-REDUCTION: %[[INNERMASK:.+]] = vector.extract %[[TPMASK]][0] : vector<[2]xi1> from vector<4x[2]xi1>
+// INNER-REDUCTION: vector.mask %[[INNERMASK]] { vector.reduction <mul>, %[[INNERVEC]], %[[INNERACC]] : vector<[2]xf32> into f32 } : vector<[2]xi1> -> f32
+// INNER-REDUCTION-COUNT-3: vector.mask %{{.+}} { vector.reduction <mul>, %{{.+}}, %{{.+}} : vector<[2]xf32> into f32 } : vector<[2]xi1> -> f32
+// INNER-REDUCTION: %[[RES:.+]] = vector.insert %{{.+}}, %{{.+}} [3] : f32 into vector<4xf32>
+// INNER-REDUCTION: return %[[RES]] : vector<4xf32>
+
+// INNER-PARALLEL: %[[RES:.+]] = vector.mask %[[MASK]] { vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [0] : vector<[2]x4xf32> to vector<4xf32> } : vector<[2]x4xi1> -> vector<4xf32>
+// INNER-PARALLEL: return %[[RES]] : vector<4xf32>
// -----
More information about the Mlir-commits
mailing list