[mlir][Vector] Further fix to avoid infinite loop in InnerOuterDimReductionConversion

[hanchung]

Summary

If all the dims are reduction dims, it is already in inner-most/outer-most
reduction form.

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());
     int64_t reductionSize = reductionDims.size();
     SmallVector<int64_t, 4> parallelDims;
     for (int64_t i = 0; i < srcRank; ++i)
       if (!reductionDimsSet.contains(i))
         parallelDims.push_back(i);
 
     // Add transpose only if inner-most/outer-most dimensions are not parallel
+    // and there are parallel dims.
+    if (parallelDims.empty())
+      return failure();
     if (useInnerDimsForReduction &&
         (parallelDims ==
          llvm::to_vector<4>(llvm::seq<int64_t>(0, parallelDims.size()))))
       return failure();
 
     if (!useInnerDimsForReduction &&
         (parallelDims !=
          llvm::to_vector<4>(llvm::seq<int64_t>(0, parallelDims.size()))))

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

 // `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 %{{.+}}
+
+func @vector_multi_reduction_to_scalar(%arg0: vector<2x3xf32>) -> f32 {
+  %0 = vector.multi_reduction <add>, %arg0 [0, 1] : vector<2x3xf32> to f32
+  return %0 : f32
+}
+// CHECK-LABEL: func @vector_multi_reduction_to_scalar
+//       CHECK:   return %{{.+}}