[Mlir-commits] [mlir] [mlir][Vector] Add a dedicated Transform Dialect op for `UnrollContractPattern` (PR #230543)

[Federico Bruzzone]

================
@@ -168,6 +168,38 @@ void transform::ApplyUnrollFromElementsPatternsOp::populatePatterns(
   vector::populateVectorFromElementsUnrollPatterns(patterns);
 }
 
+void transform::ApplyUnrollContractPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  int64_t m = getM();
+  int64_t n = getN();
+  int64_t k = getK();
+
+  auto nativeShapeFn =
+      [m, n, k](Operation *op) -> std::optional<SmallVector<int64_t>> {
+    return SmallVector<int64_t>({m, n, k});
+  };
+
+  vector::populateVectorContractUnrollPatterns(
+      patterns, vector::UnrollVectorOptions()
+                    .setNativeShapeFn(nativeShapeFn)
+                    .setFilterConstraint([](Operation *op) {
+                      auto contract = dyn_cast<vector::ContractionOp>(op);
+                      if (!contract)
+                        return failure();
+
+                      if (contract.getIteratorTypesArray() !=
+                          SmallVector<vector::IteratorType>{
+                              vector::IteratorType::parallel,
+                              vector::IteratorType::parallel,
+                              vector::IteratorType::reduction})
+                        return failure();
----------------
FedericoBruzzone wrote:

If we want to avoid a new stack allocation for `SmallVector`, but a little bit less readable:
```suggestion
auto iterators = contract.getIteratorTypesArray();
if (iterators.size() != 3 ||
    iterators[0] != vector::IteratorType::parallel ||
    iterators[1] != vector::IteratorType::parallel ||
    iterators[2] != vector::IteratorType::reduction)
  return failure();
```

_Don't merge it. it's not formatted correctly_

https://github.com/llvm/llvm-project/pull/230543