[Mlir-commits] [mlir] [mlir][transform] Add linalg.fold_reshape_ops_by_expansion transform op (PR #233732)

[Tuomas Kärnä]

================
@@ -276,6 +276,18 @@ void transform::ApplyEraseUnusedOperandsAndResultsPatternsOp::populatePatterns(
   linalg::populateEraseUnusedOperandsAndResultsPatterns(patterns);
 }
 
+void transform::ApplyFoldReshapeOpsByExpansionPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  linalg::ControlFusionFn controlFn =
+      [foldSingleUseOnly = getFoldSingleUseOnly()](OpOperand *fusedOperand) {
+        if (!foldSingleUseOnly)
+          return true;
+        Operation *producer = fusedOperand->get().getDefiningOp();
+        return producer && producer->hasOneUse();
+      };
+  linalg::populateFoldReshapeOpsByExpansionPatterns(patterns, controlFn);
----------------
tkarna wrote:

`populateFoldReshapeOpsByExpansionPatterns` takes a generic lambda control function to select applicable ops. For now, only `foldSingleUseOnly` option is implemented. Setting it to `true` corresponds to the default control function in `linalg-fuse-elementwise-ops` pass. In the long run we could think of a way to support more generic user-defined control functions.

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