[Mlir-commits] [mlir] 9cfd8d7 - [mlir][Vector] Avoid infinite loop in InnerOuterDimReductionConversion.

[llvmlistbot at llvm.org]

Author: MaheshRavishankar
Date: 2021-12-09T09:30:05-08:00
New Revision: 9cfd8d7c6ca46c6ba51142ec21fd6543afcba08a

URL: https://github.com/llvm/llvm-project/commit/9cfd8d7c6ca46c6ba51142ec21fd6543afcba08a
DIFF: https://github.com/llvm/llvm-project/commit/9cfd8d7c6ca46c6ba51142ec21fd6543afcba08a.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp
    mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp b/mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp
index 0f1c4011b231c..9e4edda64c894 100644
--- a/mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp
@@ -41,6 +41,9 @@ class InnerOuterDimReductionConversion
     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 f94a0b6e1a960..1562fedc99d83 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<maxf>, %arg0 [0] : vector<2xf32> to f32
+  return %0 : f32
+}
+// CHECK-LABEL: func @vector_reduction_1D
+//       CHECK:   return %{{.+}}