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

[mravishankar]

Summary

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.

[ThomasRaoux]

mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir

165

nit: You could make it a negative test and check that op is unchanged:
// CHECK: vector.multi_reduction #vector.kind<maxf>, %{{.*}} [0] : vector<2xf32> to f32

[mravishankar]

mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir

165

That was the hope, but the pass adds many different patterns and the op is changed (in ways that are not great). So this at least tests there is no infinite loop.

mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp

explicit InnerOuterDimReductionConversion(
         useInnerDimsForReduction(
             options == vector::VectorMultiReductionLowering::InnerReduction) {}
 
   LogicalResult matchAndRewrite(vector::MultiDimReductionOp multiReductionOp,
                                 PatternRewriter &rewriter) const override {
     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 =
         multiReductionOp.reduction_dims().getAsValueRange<IntegerAttr>();
     auto reductionDims = llvm::to_vector<4>(llvm::map_range(
         reductionDimsRange, [](APInt a) { return a.getZExtValue(); }));
     llvm::SmallDenseSet<int64_t> reductionDimsSet(reductionDims.begin(),
                                                   reductionDims.end());

mlir/test/Dialect/Vector/vector-multi-reduction-outer-lowering.mlir

 //       CHECK:   %[[V17:.+]] = vector.extract %[[RESHAPED]][17] : vector<20x6xi32>
 //       CHECK:   %[[R16:.+]] = arith.addi %[[V17]], %[[R15]] : vector<6xi32>
 //       CHECK:   %[[V18:.+]] = vector.extract %[[RESHAPED]][18] : vector<20x6xi32>
 //       CHECK:   %[[R17:.+]] = arith.addi %[[V18]], %[[R16]] : vector<6xi32>
 //       CHECK:   %[[V19:.+]] = vector.extract %[[RESHAPED]][19] : vector<20x6xi32>
 //       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 %{{.+}}