Fix CombineContractBroadcast folding reduction iterators.

[Benoit]

Summary

This enables me to drop a custom vectorization pattern that we had in IREE and just use upstream linalg vectorization instead:
https://github.com/iree-org/iree/pull/9647

[nicolasvasilache]

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

1136

seems this also works for parallel dims ?

1162

Please use bool isFunctionOfDim(unsigned position) const to make close to a 1-liner

1170

What corner case does this condition capture?

If a dim is truly unused then we cannot give it an extent anyway and the op seems like it would be invalid?

[ThomasRaoux]

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

1135

you can replace this by broadcast.getVectorType() to avoid the cast

[ThomasRaoux]

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

1141

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

1167

nit: remove single line {:

1170

The corner case is because vector.contract needs to have at least one reduction dim:
https://github.com/llvm/llvm-project/blob/2124b2f0e6083470cf27448e4217706ee646ab03/mlir/lib/Dialect/Vector/IR/VectorOps.cpp#L767
So even if the dim size is 1 we still want to keep one reduction (unless we want to relax the op but I think the semantic of it would not be ideal)

1183

nit: remove single line {:

[Benoit]

addressed review comments

[Benoit]

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

1136

Yes, that's the point! Parallel dims are fine here, so we ignore them in this code which is checking for a potential bad case with reduction dims.

1170

Thanks for the explanation Thomas, I added your testcase and added a rationale comment on the testcases.

[Benoit]

Addressed review comments.

[ThomasRaoux]

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

1136

You are missing on level of indirection here, the iterator type associated to the dimension is based on the affine map. is should be contractOp.getIteratorTypes().getValue()[map.getDimPosition(i)).
Maybe add or modifying the test to have a permutation in the affine map would help catching that kind of problem.

1170

nit: Might be worth a comment?

[Benoit]

more review comments

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

 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallBitVector.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;
 using namespace mlir::vector;
for (Value *operand : {&lhs, &rhs}) {
         }
         originalDims.push_back(
             rewriter.getAffineDimExpr(dim.index() + rankDiff));
       }
       // Contract doesn't support inner dimension broadcast. Once this is
       // relaxed we can remove this case.
       if (innerDimBroadcast)
         continue;
+
+      // It would be incorrect to fold a broadcast onto a reduction dimension
+      // of non-unit size.
+      bool nonUnitDimReductionBroadcast = false;
+      for (int64_t i = 0; i < rankDiff; ++i) {
+        if (broadcast.getVectorType().getDimSize(i) != 1 &&
+            isReductionIterator(contractOp.getIteratorTypes()
+                                    .getValue()[map.getDimPosition(i)])) {
+          nonUnitDimReductionBroadcast = true;
+          break;
+        }
+      }
+      if (nonUnitDimReductionBroadcast)
+        continue;
+
       AffineMap broadcastMap =
           AffineMap::get(broadcast.getVectorType().getRank(), 0, originalDims,
                          contractOp.getContext());
       map = broadcastMap.compose(map);
       *operand = broadcast.getSource();
       changed = true;
     }
+
     if (!changed)
       return failure();
+
+    // Determine which dims are usused, now that the maps have been composed
+    // with the broadcast maps.
+    unsigned numDims = maps[0].getNumDims();
+    llvm::SmallBitVector unusedDims(numDims, true);
+    for (const auto &m : maps) {
+      for (unsigned i = 0; i < numDims; ++i) {
+        if (m.isFunctionOfDim(i))
+          unusedDims.reset(i);
+      }
+    }
+    // Compress unused dims.
+    for (auto &m : maps)
+      m = compressDims(m, unusedDims);
+    // Compute the combined iterators.
+    SmallVector<Attribute, 4> iterators;
+    for (unsigned i = 0; i < numDims; ++i) {
+      if (!unusedDims.test(i))
+        iterators.push_back(contractOp.getIteratorTypes().getValue()[i]);
+    }
+    // Check that compressing unused dims isn't removing all reduction
+    // iterators. For example, if the vector.contract had only one reduction
+    // iterator and that was a unit-dimension created by a broadcast,
+    // then we should bail here, otherwise we would create a contract without
+    // a reduction iterator.
+    if (!llvm::any_of(iterators, isReductionIterator))
+      return failure();
+
     rewriter.replaceOpWithNewOp<vector::ContractionOp>(
         contractOp, lhs, rhs, contractOp.getAcc(),
-        rewriter.getAffineMapArrayAttr(maps), contractOp.getIteratorTypes());
+        rewriter.getAffineMapArrayAttr(maps), rewriter.getArrayAttr(iterators));
     return success();
   }
 };
 
 /// Reorders cast(broadcast) to broadcast(cast). This makes broadcast ops and
 /// contraction ops closer, which kicks in CombineContractBroadcast pattern when
 /// casting ops are around these operations.
 /// Ex:

mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir

func.func @contract_broadcast(
   %cst = arith.constant dense<0.000000e+00> : vector<8x32xf32>
   %0 = vector.broadcast %arg0 : vector<32x16xf32> to vector<8x32x16xf32>
   %1 = vector.contract {indexing_maps = [#map0, #map0, #map1],
     iterator_types = ["parallel", "parallel", "reduction"],
     kind = #vector.kind<add>} %0, %arg1, %cst : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
   return %1 : vector<8x32xf32>
 }
 
+// -----
+// Test that CombineContractBroadcast is able to combine a broadcast that
+// creates a unit dim that is consumed by a reduction iterator, dropping that
+// reduction iterator, as long as there is another reduction iterator left.
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+// CHECK-DAG: #[[$map0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: contract_broadcast_unit_dim_reduction
+//  CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<8x4xi32>, %[[ARG2:.+]]: vector<8x8xi32>)
+//  CHECK: vector.contract
+//  CHECK-SAME: indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]]
+//  CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+//  CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x4xi32>, vector<8x4xi32> into vector<8x8xi32>
+func.func @contract_broadcast_unit_dim_reduction(%arg0 : vector<8x4xi32>, %arg1 : vector<8x4xi32>, %arg2 : vector<8x8xi32>) -> vector<8x8xi32> {
+    %0 = vector.broadcast %arg0 : vector<8x4xi32> to vector<1x8x4xi32>
+    %1 = vector.broadcast %arg1 : vector<8x4xi32> to vector<1x8x4xi32>
+    %result = vector.contract {
+        indexing_maps = [#map0, #map1, #map2],
+        iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+        kind = #vector.kind<add>
+      } %0, %1, %arg2 : vector<1x8x4xi32>, vector<1x8x4xi32> into vector<8x8xi32>
+    return %result : vector<8x8xi32>
+}
+
+// -----
+// Test that CombineContractBroadcast will not combine a broadcast that creates
+// a non-unit dim that is consumed by a reduction iterator.
+// Moreover, the affine_map's are permuting the position of that reduction
+// iterator with the position of a parallel iterator to ensure that
+// the logic guarding that case does not mix up dimensions here.
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>
+
+// CHECK-DAG: #[[$map0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>
+// CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>
+// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2)>
+
+// CHECK-LABEL: contract_broadcast_non_unit_dim_reduction_with_permutation
+//  CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<8x4xi32>, %[[ARG2:.+]]: vector<8x8xi32>)
+//  CHECK: %[[BROADCAST0:.+]] = vector.broadcast %[[ARG0]] : vector<8x4xi32> to vector<2x8x4xi32>
+//  CHECK: %[[BROADCAST1:.+]] = vector.broadcast %[[ARG1]] : vector<8x4xi32> to vector<2x8x4xi32>
+//  CHECK: vector.contract
+//  CHECK-SAME: indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]]
+//  CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel", "reduction"]
+//  CHECK-SAME: %[[BROADCAST0]], %[[BROADCAST1]], %[[ARG2]] : vector<2x8x4xi32>, vector<2x8x4xi32> into vector<8x8xi32>
+func.func @contract_broadcast_non_unit_dim_reduction_with_permutation(%arg0 : vector<8x4xi32>, %arg1 : vector<8x4xi32>, %arg2 : vector<8x8xi32>) -> vector<8x8xi32> {
+    %0 = vector.broadcast %arg0 : vector<8x4xi32> to vector<2x8x4xi32>
+    %1 = vector.broadcast %arg1 : vector<8x4xi32> to vector<2x8x4xi32>
+    %result = vector.contract {
+        indexing_maps = [#map0, #map1, #map2],
+        iterator_types = ["parallel", "reduction", "parallel", "reduction"],
+        kind = #vector.kind<add>
+      } %0, %1, %arg2 : vector<2x8x4xi32>, vector<2x8x4xi32> into vector<8x8xi32>
+    return %result : vector<8x8xi32>
+}
+
+// -----
+
+// Test that CombineContractBroadcast is not combining this case, as that would
+// result in dropping this contract's only reduction iterator.
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+
+// CHECK-LABEL: contract_broadcast_unit_dim_reduction_as_only_reduction
+//  CHECK-SAME: (%[[ARG0:.+]]: vector<8xi32>, %[[ARG1:.+]]: vector<8xi32>, %[[ARG2:.+]]: vector<8x8xi32>)
+//  CHECK: %[[BROADCAST0:.+]] = vector.broadcast %[[ARG0]] : vector<8xi32> to vector<1x8xi32>
+//  CHECK: %[[BROADCAST1:.+]] = vector.broadcast %[[ARG1]] : vector<8xi32> to vector<1x8xi32>
+//  CHECK: vector.contract
+//  CHECK-SAME: indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]]
+//  CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel"]
+//  CHECK-SAME: %[[BROADCAST0]], %[[BROADCAST1]], %[[ARG2]] : vector<1x8xi32>, vector<1x8xi32> into vector<8x8xi32>
+func.func @contract_broadcast_unit_dim_reduction_as_only_reduction(%arg0 : vector<8xi32>, %arg1 : vector<8xi32>, %arg2 : vector<8x8xi32>) -> vector<8x8xi32> {
+    %0 = vector.broadcast %arg0 : vector<8xi32> to vector<1x8xi32>
+    %1 = vector.broadcast %arg1 : vector<8xi32> to vector<1x8xi32>
+    %result = vector.contract {
+        indexing_maps = [#map0, #map1, #map2],
+        iterator_types = ["reduction", "parallel", "parallel"],
+        kind = #vector.kind<add>
+      } %0, %1, %arg2 : vector<1x8xi32>, vector<1x8xi32> into vector<8x8xi32>
+    return %result : vector<8x8xi32>
+}
+
 //===----------------------------------------------------------------------===//
 // Reorder casting ops and vector ops. The casting ops have almost identical
 // pattern, so only arith.extsi op is tested.
 //===----------------------------------------------------------------------===//
 
 // -----
 
 func.func @broadcast_vector_extsi(%a : vector<4xi8>) -> vector<2x4xi32> {