[mlir][Vector] Avoid infinite loop in InnerOuterDimReductionConversion.

This patterns tries to convert an inner (outer) dim reduction to an
outer (inner) dim reduction. Doing this on a 1D or 0D vector results
in an infinite loop since the converted op is same as the original
operation. Just returning failure when source rank <= 1 fixes the
issue.

Differential Revision: https://reviews.llvm.org/D115426
This commit is contained in:
MaheshRavishankar 2021-12-09 09:18:00 -08:00
parent 70638d997c
commit 9cfd8d7c6c
2 changed files with 13 additions and 0 deletions

View File

@ -41,6 +41,9 @@ public:
auto src = multiReductionOp.source();
auto loc = multiReductionOp.getLoc();
auto srcRank = multiReductionOp.getSourceVectorType().getRank();
// If the rank is less than or equal to 1, there is nothing to do.
if (srcRank <= 1)
return failure();
// Separate reduction and parallel dims
auto reductionDimsRange =

View File

@ -153,3 +153,13 @@ func @vector_reduction_outer(%arg0: vector<2x3x4x5xi32>) -> vector<2x3xi32> {
// CHECK: %[[R18:.+]] = arith.addi %[[V19]], %[[R17]] : vector<6xi32>
// CHECK: %[[RESULT_VEC:.+]] = vector.shape_cast %[[R18]] : vector<6xi32> to vector<2x3xi32>
// CHECK: return %[[RESULT_VEC]] : vector<2x3xi32>
// This test is mainly to catch a bug that running
// `InnerOuterDimReductionConversion` on this function results in an
// infinite loop. So just check that some value is returned.
func @vector_reduction_1D(%arg0 : vector<2xf32>) -> f32 {
%0 = vector.multi_reduction #vector.kind<maxf>, %arg0 [0] : vector<2xf32> to f32
return %0 : f32
}
// CHECK-LABEL: func @vector_reduction_1D
// CHECK: return %{{.+}}