[mlir][linalg] Fix bufferize pattern to allow unknown operations in body of generic

[herhut]

Summary

The original version of the bufferization pattern for linalg.generic would manually clone operations within the region to the bufferized clone of the operation. This triggers legality requirements on those operations in the conversion infra. Instead, this now uses the rewriter to inline the region instead, avoiding those legality requirements.

[herhut]

Note that this is already tested by existing tests that use the math dialect.

[pifon2a]

Thank you, Stephan!

[herhut]

also support linalg ops with hidden regions like dot.

[herhut]

clang tidy fix

mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp

if (auto genericOp = mlir::dyn_cast<GenericOp>(*linalgOp)) {
     auto newGenericOp = rewriter.create<GenericOp>(
         genericOp.getLoc(),
         /*resultTensorTypes=*/llvm::None,
         /*inputs=*/inputs,
         /*outputs=*/outputs, genericOp.indexing_maps(),
         genericOp.iterator_types(), genericOp.docAttr(),
         genericOp.library_callAttr());
 
-    // Create a new block in the region of the new Generic Op.
-    Block *oldBlock = genericOp.getBody();
-    Region &newRegion = newGenericOp.region();
-    Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
-                                           oldBlock->getArgumentTypes());
-
-    // Clone the body of the old block to the new block.
-    BlockAndValueMapping mapping;
-    mapping.map(oldBlock->getArguments(), newBlock->getArguments());
-
-    OpBuilder::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPointToEnd(newBlock);
-    for (auto &op : oldBlock->getOperations()) {
-      Operation *clonedOp = rewriter.clone(op, mapping);
-      mapping.map(op.getResults(), clonedOp->getResults());
-    }
+    rewriter.inlineRegionBefore(genericOp.getRegion(), newGenericOp.region(),
+                                newGenericOp.region().begin());
     return newGenericOp;
   }
+  assert(linalgOp->getNumRegions() == 0);
   SmallVector<Value, 8> newOperands = inputs;
   newOperands.append(outputs.begin(), outputs.end());
   return linalgOp.clone(rewriter, linalgOp.getLoc(),
                         /*resultTypes=*/ArrayRef<Type>{}, newOperands);
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization patterns.

 /// work on buffers.
 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     BufferizeTypeConverter typeConverter;
 
     // Mark all Standard operations legal.
-    target.addLegalDialect<AffineDialect, math::MathDialect,
-                           memref::MemRefDialect, StandardOpsDialect,
-                           tensor::TensorDialect>();
+    target.addLegalDialect<AffineDialect, memref::MemRefDialect,
+                           StandardOpsDialect, tensor::TensorDialect>();
     target.addIllegalOp<InitTensorOp, tensor::ExtractSliceOp,
                         tensor::InsertSliceOp, PadTensorOp>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return typeConverter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);