Order parallel indices before transposing the input in multireductions

[harsh]

Summary

The current code does not preserve the order of the parallel
dimensions when doing multi-reductions and thus we can end
up in scenarios where the result shape does not match the
desired shape after reduction.

This patch fixes that by ensuring that the parallel indices
are in order and then concatenates them to the reduction dimensions
so that the reduction dimensions are innermost.

[harsh]

Fixed lint changes

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3945–3947

Why don't we sort the indices so that the non reduction dimensions are in order in this transpose instead? This way we don't need a second transpose if they are not in the right order.

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3932–3936

I think you can simplify this logic and solve the problem by just pushing the non reduced dimensions first in order then push the reduction dimensions.

[asaadaldien]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3932–3933

Good catch the bug is here! the linear swap isn't correct because it shuffles parallel dims around. You can replace this by inserting parallel induces first then reductions, e.g something like the following:

[nicolasvasilache]

+1 to what @ThomasRaoux said.
It should also be possible to rewrite that logic in a much cleaner and simpler form like:

auto reductionDimsRange = op.reduction_dims.getAsValueRange<int64_t>();
SmallVector<int64_t, 4> reductionDims = llvm::to_vector<4>(llvm::map_range(reductionDimsRange.begin(), reductionDimsRange.end(), [](APInt a) { return a.getZExtValue(); }));
SmallDenseSet<int64_T> reductionDimsSet;
// fill it from reductionDims or from llvm::map_range, whichever makes most sense. 
SmallVector<int64_t> parallelDims;
for (int64_t idx = 0, e = ; idx != e; ++idx)
  if (!reductionDimsSet.contains(idx))
    parallelDims.push_back(idx);

if (parallelDims != llvm::seq<int64_t>(0, parallelDims.size())) { // This is the condition for "reductions must be innermost", see comment below.
  // concat parallelDims, reductionDims + add transpose. 
}

One thing to note is that the current lowering seems to always put reductions inside and use vector.reduction.
I think in most cases this will result in quite bad code because it will rely on horizontal reductions and add a ton of insert/extract.

We should have an option to emit the other way around too: put the reductionDims outside and just unroll pointwise operations on the parallel dims.
This is closer to the "outerproduct-like" semantics and has much better chances of being rewritten as fmas when possible; missing fma instructions is immediately a 2x performance penalty on many systems.

With the code structuring I propose above, the proper transposition amounts to just changing the order of concatenated dims.

[harsh]

Thanks @ThomasRaoux , @asaadaldien , @nicolasvasilache for the comments. @nicolasvasilache I have modified the patch as per your changes but kept the current patch to just handle moving reductions to the inner most dimensions. I will put up another patch to handle moving reductions to the outermost dimensions and based on the performance of inner vs outer, we can decide which path we want to take.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3932–3933

Thanks! Since I will be attempting moving the reduction dims to be outermost also, I have modified the patch as per Nicolas' comments.

3932–3936

Thanks you are right. I have fixed this as per Nicolas' comments.

3945–3947

Makes sense.

[ThomasRaoux]

In D104884#2845548, @harsh wrote:

Thanks @ThomasRaoux , @asaadaldien , @nicolasvasilache for the comments. @nicolasvasilache I have modified the patch as per your changes but kept the current patch to just handle moving reductions to the inner most dimensions. I will put up another patch to handle moving reductions to the outermost dimensions and based on the performance of inner vs outer, we can decide which path we want to take.

Did you upload the new patch?

[harsh]

Fixed as per comments

[harsh]

@ThomasRaoux - yes I just uploaded it.

[ThomasRaoux]

LGTM

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3920

nit: formatting looks off.

3925

nit: skip braces on single line if statements (https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements)

3933

nit: In general prefer early exit (https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code)

[harsh]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3933

Just to clarify, this means I should check if the inner most dims are reduction and if so, exit early here?

[harsh]

Address comments and linting errors

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3933

Yes, looks good.

[harsh]

Thanks and feel free to merge because I don't have merge priveleges.

[ThomasRaoux]

In D104884#2845835, @harsh wrote:

Thanks and feel free to merge because I don't have merge priveleges.

Done

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 #include "mlir/IR/Location.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "vector-to-vector"
 
 using namespace mlir;
struct InnerDimReductionConversion
   using OpRewritePattern<vector::MultiDimReductionOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::MultiDimReductionOp multiReductionOp,
                                 PatternRewriter &rewriter) const override {
     auto src = multiReductionOp.source();
     auto loc = multiReductionOp.getLoc();
     auto srcRank = multiReductionOp.getSourceVectorType().getRank();
 
-    auto reductionDims = llvm::to_vector<4>(
-        llvm::map_range(multiReductionOp.reduction_dims().cast<ArrayAttr>(),
-                        [](Attribute attr) -> int64_t {
-                          return attr.cast<IntegerAttr>().getInt();
-                        }));
-    llvm::sort(reductionDims);
-
-    int64_t reductionSize = multiReductionOp.reduction_dims().size();
-
-    // Fails if already inner most reduction.
-    bool innerMostReduction = true;
-    for (int i = 0; i < reductionSize; ++i) {
-      if (reductionDims[reductionSize - i - 1] != srcRank - i - 1) {
-        innerMostReduction = false;
-      }
+    // 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);
     }
-    if (innerMostReduction)
-      return failure();
 
-    // Permutes the indices so reduction dims are inner most dims.
-    SmallVector<int64_t> indices;
-    for (int i = 0; i < srcRank; ++i) {
-      indices.push_back(i);
-    }
-    int ir = reductionSize - 1;
-    int id = srcRank - 1;
-    while (ir >= 0) {
-      std::swap(indices[reductionDims[ir--]], indices[id--]);
-    }
+    // Add transpose only if inner-most dimensions are not reductions
+    if (parallelDims ==
+        llvm::to_vector<4>(llvm::seq<int64_t>(0, parallelDims.size())))
+      return failure();
 
-    // Sets inner most dims as reduction.
+    SmallVector<int64_t, 4> indices;
+    indices.append(parallelDims.begin(), parallelDims.end());
+    indices.append(reductionDims.begin(), reductionDims.end());
+    auto transposeOp = rewriter.create<vector::TransposeOp>(loc, src, indices);
     SmallVector<bool> reductionMask(srcRank, false);
     for (int i = 0; i < reductionSize; ++i) {
       reductionMask[srcRank - i - 1] = true;
     }
-    auto transposeOp = rewriter.create<vector::TransposeOp>(loc, src, indices);
     rewriter.replaceOpWithNewOp<vector::MultiDimReductionOp>(
         multiReductionOp, transposeOp.result(), reductionMask,
         multiReductionOp.kind());
     return success();
   }
 };
 
 // Reduces the rank of vector.mult_reduction nd -> 2d given all reduction
 // dimensions are inner most.
 struct ReduceMultiDimReductionRank

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

 func @vector_multi_reduction_transposed(%arg0: vector<2x3x4x5xf32>) -> vector<2x5xf32> {
     %0 = vector.multi_reduction #vector.kind<add>, %arg0 [1, 2] : vector<2x3x4x5xf32> to vector<2x5xf32>
     return %0 : vector<2x5xf32>
 }
 
 // CHECK-LABEL: func @vector_multi_reduction_transposed
 //  CHECK-SAME:    %[[INPUT:.+]]: vector<2x3x4x5xf32>
 //       CHECK:     %[[TRANSPOSED_INPUT:.+]] = vector.transpose %[[INPUT]], [0, 3, 1, 2] : vector<2x3x4x5xf32> to vector<2x5x3x4xf32>
-//       CHEKC:     vector.shape_cast %[[TRANSPOSED_INPUT]] : vector<2x5x3x4xf32> to vector<10x12xf32>
+//       CHECK:     vector.shape_cast %[[TRANSPOSED_INPUT]] : vector<2x5x3x4xf32> to vector<10x12xf32>
 //       CHECK:     %[[RESULT:.+]] = vector.shape_cast %{{.*}} : vector<10xf32> to vector<2x5xf32>
 //       CHECK:       return %[[RESULT]]     
+
+func @vector_multi_reduction_ordering(%arg0: vector<3x2x4xf32>) -> vector<2x4xf32> {
+    %0 = vector.multi_reduction #vector.kind<mul>, %arg0 [0] : vector<3x2x4xf32> to vector<2x4xf32>
+    return %0 : vector<2x4xf32>
+}
+// CHECK-LABEL: func @vector_multi_reduction_ordering
+//  CHECK-SAME:   %[[INPUT:.+]]: vector<3x2x4xf32>
+//       CHECK:       %[[RESULT_VEC_0:.+]] = constant dense<{{.*}}> : vector<8xf32>
+//       CHECK:       %[[C0:.+]] = constant 0 : i32
+//       CHECK:       %[[C1:.+]] = constant 1 : i32
+//       CHECK:       %[[C2:.+]] = constant 2 : i32
+//       CHECK:       %[[C3:.+]] = constant 3 : i32
+//       CHECK:       %[[C4:.+]] = constant 4 : i32
+//       CHECK:       %[[C5:.+]] = constant 5 : i32
+//       CHECK:       %[[C6:.+]] = constant 6 : i32
+//       CHECK:       %[[C7:.+]] = constant 7 : i32
+//       CHECK:       %[[TRANSPOSED_INPUT:.+]] = vector.transpose %[[INPUT]], [1, 2, 0] : vector<3x2x4xf32> to vector<2x4x3xf32>
+//       CHECK:       %[[V0:.+]] = vector.extract %[[TRANSPOSED_INPUT]][0, 0]
+//       CHECK:       %[[RV0:.+]] = vector.reduction "mul", %[[V0]] : vector<3xf32> into f32
+//       CHECK:       %[[RESULT_VEC_1:.+]] = vector.insertelement %[[RV0:.+]], %[[RESULT_VEC_0]][%[[C0]] : i32] : vector<8xf32>
+//       CHECK:       %[[V1:.+]] = vector.extract %[[TRANSPOSED_INPUT]][0, 1]
+//       CHECK:       %[[RV1:.+]] = vector.reduction "mul", %[[V1]] : vector<3xf32> into f32
+//       CHECK:       %[[RESULT_VEC_2:.+]] = vector.insertelement %[[RV1:.+]], %[[RESULT_VEC_1]][%[[C1]] : i32] : vector<8xf32>
+//       CHECK:       %[[V2:.+]] = vector.extract %[[TRANSPOSED_INPUT]][0, 2]
+//       CHECK:       %[[RV2:.+]] = vector.reduction "mul", %[[V2]] : vector<3xf32> into f32
+//       CHECK:       %[[RESULT_VEC_3:.+]] = vector.insertelement %[[RV2:.+]], %[[RESULT_VEC_2]][%[[C2]] : i32] : vector<8xf32>
+//       CHECK:       %[[V3:.+]] = vector.extract %[[TRANSPOSED_INPUT]][0, 3]
+//       CHECK:       %[[RV3:.+]] = vector.reduction "mul", %[[V3]] : vector<3xf32> into f32
+//       CHECK:       %[[RESULT_VEC_4:.+]] = vector.insertelement %[[RV3:.+]], %[[RESULT_VEC_3]][%[[C3]] : i32] : vector<8xf32>
+//       CHECK:       %[[V4:.+]] = vector.extract %[[TRANSPOSED_INPUT]][1, 0]
+//       CHECK:       %[[RV4:.+]] = vector.reduction "mul", %[[V4]] : vector<3xf32> into f32
+//       CHECK:       %[[RESULT_VEC_5:.+]] = vector.insertelement %[[RV4:.+]], %[[RESULT_VEC_4]][%[[C4]] : i32] : vector<8xf32>
+//       CHECK:       %[[V5:.+]] = vector.extract %[[TRANSPOSED_INPUT]][1, 1]
+//       CHECK:       %[[RV5:.+]] = vector.reduction "mul", %[[V5]] : vector<3xf32> into f32
+//       CHECK:       %[[RESULT_VEC_6:.+]] = vector.insertelement %[[RV5:.+]], %[[RESULT_VEC_5]][%[[C5]] : i32] : vector<8xf32>
+//       CHECK:       %[[V6:.+]] = vector.extract %[[TRANSPOSED_INPUT]][1, 2]
+//       CHECK:       %[[RV6:.+]] = vector.reduction "mul", %[[V6]] : vector<3xf32> into f32
+//       CHECK:       %[[RESULT_VEC_7:.+]] = vector.insertelement %[[RV6:.+]], %[[RESULT_VEC_6]][%[[C6]] : i32] : vector<8xf32>
+//       CHECK:       %[[V7:.+]] = vector.extract %[[TRANSPOSED_INPUT]][1, 3]
+//       CHECK:       %[[RV7:.+]] = vector.reduction "mul", %[[V7]] : vector<3xf32> into f32
+//       CHECK:       %[[RESULT_VEC:.+]] = vector.insertelement %[[RV7:.+]], %[[RESULT_VEC_7]][%[[C7]] : i32] : vector<8xf32>
+//       CHECK:       %[[RESHAPED_VEC:.+]] = vector.shape_cast %[[RESULT_VEC]] : vector<8xf32> to vector<2x4xf32>
+//       CHECK:       return %[[RESHAPED_VEC]]