[Mlir-commits] [mlir] [mlir][transform] Implement `FlattenElementwiseLinalgOp` transform op (PR #81431)

[llvmlistbot at llvm.org]

================
@@ -1479,40 +1475,40 @@ Operation *createCollapsedOp(LinalgType op,
       resultTypes.push_back(newOutput.getType());
   }
 
-  if (isa<linalg::CopyOp>(op)) {
-    return rewriter.create<linalg::CopyOp>(loc, inputOperands[0],
-                                           outputOperands[0]);
-  }
+  Operation *collapsedOp = clone(
+      rewriter, op, resultTypes,
+      llvm::to_vector(llvm::concat<Value>(inputOperands, outputOperands)));
 
-  // Get the iterator types for the operand.
-  SmallVector<utils::IteratorType> iteratorTypes =
-      getCollapsedOpIteratorTypes(op.getIteratorTypesArray(), collapsingInfo);
+  if (op->hasAttr("indexing_maps")) {
----------------
MaheshRavishankar wrote:

This seems iffy to me. THis is looking into internal implementation details of the LinalgOp representation in C++ and is hard to maintain. Let me give you an example. I have no idea what the indexing maps or iterator types are stored as. I just use the utility methods on ops to get this information. Side-stepping that can  introduce silent bugs if something changes in the implementation of the operation. 

I understand you are trying to generalize these things.. One way to do that would be to define a method

```
template <typename OpTy>
Operation * cloneWithIndexingMapsIteratorTypesAndOperands(RewriterBase &rewriter, OpTy origOp, TypeRange resultTypes, ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorType, ValueRange inputOperands, ValueRange outputOperands) {
  return nullptr;
}

template <>
Operation *cloneWithIndexingMapsIteratorTypesAndOperands<LinalgOp>(RewriterBase &rewriter, LinalgOp origOp, TypeRange resultTypes, ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorType, ValueRange inputOperands, ValueRange outputOperands) {
 return clone(rewriter, origOp, resultTypes, inputOperands, outputOperands);
}

template <>
Operation *cloneWithIndexingMapsIteratorTypesAndOperands<GenericOp>(RewriterBase &rewriter, GenericOp origOp, TypeRange resultTypes, ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorType, ValueRange inputOperands, ValueRange outputOperands) {
SmallVector<utils::IteratorType> iteratorTypes =
      getCollapsedOpIteratorTypes(op.getIteratorTypesArray(), collapsingInfo);

  // Get the indexing maps.
  auto indexingMaps =
      llvm::map_to_vector(op.getIndexingMapsArray(), [&](AffineMap map) {
        return getCollapsedOpIndexingMap(map, collapsingInfo);
      });

  Operation *collapsedOp = rewriter.create<linalg::GenericOp>(
      loc, resultTypes, inputOperands, outputOperands, indexingMaps,
      iteratorTypes, [](OpBuilder &builder, Location loc, ValueRange args) {});
  Block *origOpBlock = &op->getRegion(0).front();
  Block *collapsedOpBlock = &collapsedOp->getRegion(0).front();
  rewriter.mergeBlocks(origOpBlock, collapsedOpBlock,
                       collapsedOpBlock->getArguments());
}
```

Then you can 

```
if (auto genericOp = dyn_cast<GenericOp>(op)) {
  cloneWithIndexingMapsIteratorTypesAndOperands(rewriter, genericOp, ....)
} else {
  cloneWithIndexingMapsIteratorTypesAndOperands(rewriter, cast<LinalgOp>(op), ...)
}
```

I think that gives you what you want and doesnt leak internal implementation details of the op.

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