[Mlir-commits] [mlir] [mlir][linalg][elementwise] Fold broadcast into new elementwise (PR #167626)

[Javed Absar]

================
@@ -29,7 +29,25 @@ using namespace mlir::linalg;
 #define DEBUG_TYPE "linalg-fold-into-elementwise"
 
 namespace {
-struct FoldTransposePattern : public OpRewritePattern<ElementwiseOp> {
----------------
javedabsar1 wrote:

```
template <typename ProducerOpTy>
struct ElementwiseOpFolder {
 // Helper function to fold broadcast etc into elementwise op.
 // Producer in this context is `broadcast op` etc, consumer is elwise operand.
  static bool fold(OpOperand *elwiseOperand, AffineMap elwiseMap,
                   SmallVector<Value> &newIns,
                   SmallVector<AffineMap> &newMaps) {
    auto producerOp = elwiseOperand->get().getDefiningOp<ProducerOpTy>();
    if (!producerOp || !elwiseMap.isProjectedPermutation())
      return false;
    newIns.push_back(producerOp.getInput());
    // push in the new composed affine map
    newMaps.push_back(
        producerOp.getMatchingIndexingMap(producerOp.getDpsInputOperand(0))
            .compose(elwiseMap));
    return true;
  }
};
```

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