[mlir][vector] Allow transposing multi_reduction when the parallel dim is in the middle

[bkramer]

Summary

The check here was stricter than necessary, we just want all the
parallel dims to go to one side and all reduction dims to go to the
other.

[vmurali]

The original code is wrong when the case is outer-reductions. Instead, it should check that

!useInnerDimsForReduction &&
(parallelDims ==
 llvm::to_vector<4>(llvm::seq<int64_t>(reductionDims.size(), parallelDims.size() + reductionDims.size()))))

The code that you submitted wouldn't work if you switch the RUN command to use inner-reductions strategy for your test example. (There are other cases also where it wouldn't work.)

// RUN: mlir-opt %s -test-vector-multi-reduction-lowering-patterns | FileCheck %s

whereas the change I wrote should work.

Basically, you want to ensure that the transpose happens whenever the reduction dimensions are either not contiguous or not in the innermost/outermost dimensions depending on the strategy.

[bkramer]

• Address review comments
• Add a test to make sure the inner reduction strategy still works

mlir/lib/Dialect/Vector/Transforms/VectorMultiDimReductionTransforms.cpp

LogicalResult matchAndRewrite(vector::MultiDimReductionOp multiReductionOp,
     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()))))
+        (parallelDims == llvm::to_vector<4>(llvm::seq<int64_t>(
+                             reductionDims.size(),
+                             parallelDims.size() + reductionDims.size()))))
       return failure();
 
     SmallVector<int64_t, 4> indices;
     if (useInnerDimsForReduction) {
       indices.append(parallelDims.begin(), parallelDims.end());
       indices.append(reductionDims.begin(), reductionDims.end());
     } else {
       indices.append(reductionDims.begin(), reductionDims.end());

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

 // CHECK:           %[[VAL_153:.*]] = vector.extract %[[VAL_139]][0, 2] : vector<8x16x4xi1>
 // CHECK:           %[[VAL_154:.*]] = vector.mask %[[VAL_153]] { vector.reduction <add>
 // CHECK:           %[[VAL_155:.*]] = vector.insertelement %[[VAL_154]]
 
 // CHECK:           %[[VAL_158:.*]] = vector.extract %[[VAL_139]][0, 3] : vector<8x16x4xi1>
 // CHECK:           %[[VAL_159:.*]] = vector.mask %[[VAL_158]] { vector.reduction <add>
 // CHECK:           %[[VAL_160:.*]] = vector.insertelement %[[VAL_159]]
 
+// -----
+
+func.func @vector_multi_reduction_parallel_middle(%arg0: vector<3x4x5xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+    %0 = vector.multi_reduction <add>, %arg0, %acc [0, 2] : vector<3x4x5xf32> to vector<4xf32>
+    return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: func @vector_multi_reduction_parallel_middle
+//  CHECK-SAME:   %[[INPUT:.+]]: vector<3x4x5xf32>, %[[ACC:.+]]: vector<4xf32>
+//       CHECK: vector.transpose %[[INPUT]], [1, 0, 2] : vector<3x4x5xf32> to vector<4x3x5xf32>

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

 //       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>
 
+func.func @vector_multi_reduction_parallel_middle(%arg0: vector<3x4x5xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+    %0 = vector.multi_reduction <add>, %arg0, %acc [0, 2] : vector<3x4x5xf32> to vector<4xf32>
+    return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: func @vector_multi_reduction_parallel_middle
+//  CHECK-SAME:   %[[INPUT:.+]]: vector<3x4x5xf32>, %[[ACC:.+]]: vector<4xf32>
+//       CHECK: vector.transpose %[[INPUT]], [0, 2, 1] : vector<3x4x5xf32> to vector<3x5x4xf32>
+
 // 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.func @vector_reduction_1D(%arg0 : vector<2xf32>, %acc: f32) -> f32 {
   %0 = vector.multi_reduction #vector.kind<maxf>, %arg0, %acc [0] : vector<2xf32> to f32
   return %0 : f32
 }
 // CHECK-LABEL: func @vector_reduction_1D
 //       CHECK:   return %{{.+}}
 
 func.func @vector_multi_reduction_to_scalar(%arg0: vector<2x3xf32>, %acc: f32) -> f32 {
   %0 = vector.multi_reduction <add>, %arg0, %acc [0, 1] : vector<2x3xf32> to f32
   return %0 : f32
 }
 // CHECK-LABEL: func @vector_multi_reduction_to_scalar
 //       CHECK:   return %{{.+}}