From 9cfd8d7c6ca46c6ba51142ec21fd6543afcba08a Mon Sep 17 00:00:00 2001 From: MaheshRavishankar Date: Thu, 9 Dec 2021 09:18:00 -0800 Subject: [PATCH] [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 --- .../Vector/VectorMultiDimReductionTransforms.cpp | 3 +++ .../Vector/vector-multi-reduction-outer-lowering.mlir | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp b/mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp index 0f1c4011b231..9e4edda64c89 100644 --- a/mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp +++ b/mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp @@ -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 = diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir index f94a0b6e1a96..1562fedc99d8 100644 --- a/mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir +++ b/mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir @@ -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, %arg0 [0] : vector<2xf32> to f32 + return %0 : f32 +} +// CHECK-LABEL: func @vector_reduction_1D +// CHECK: return %{{.+}}