forked from OSchip/llvm-project
[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:
parent
70638d997c
commit
9cfd8d7c6c
|
@ -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 =
|
||||
|
|
|
@ -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 %{{.+}}
|
||||
|
|
Loading…
Reference in New Issue