[Mlir-commits] [mlir] 42fd68b - [mlir][linalg][bufferize] LinalgOp: Move existing region to new op

[Matthias Springer]

Author: Matthias Springer
Date: 2022-01-07T07:00:24+09:00
New Revision: 42fd68b34457f4d07c052cca6f490054ec2e215c

URL: https://github.com/llvm/llvm-project/commit/42fd68b34457f4d07c052cca6f490054ec2e215c
DIFF: https://github.com/llvm/llvm-project/commit/42fd68b34457f4d07c052cca6f490054ec2e215c.diff

LOG: [mlir][linalg][bufferize] LinalgOp: Move existing region to new op

This has two advantages.

1. It is more efficient. No need to clone the entire region.
2. Recreating ops (via cloning) invalidates analysis results. Previously, an OpResult could have bufferized out-of-place, even though the analysis requested an in-place bufferization. That is because BufferizationState keeps track of OpResults for storing bufferization analysis results (and cloned ops have new OpResults).

Differential Revision: https://reviews.llvm.org/D116453

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp
index a3cb3c36065b..546cde439669 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp
@@ -66,9 +66,14 @@ static LogicalResult bufferizeLinalgOp(RewriterBase &rewriter, LinalgOp op,
 
   // Set insertion point now that potential alloc/dealloc are introduced.
   rewriter.setInsertionPoint(op);
-  // Clone the op, but use the new operands. Since the new op does not have any
-  // tensor results, it does not return anything.
-  op.clone(rewriter, op.getLoc(), /*resultTypes=*/TypeRange{}, newOperands);
+  // Clone the op, but use the new operands. Move the existing block into the
+  // new op. Since the new op does not have any tensor results, it does not
+  // return anything.
+  assert(op->getNumRegions() == 1 && "expected that op has 1 region");
+  auto newOp = cast<LinalgOp>(op.cloneWithoutRegions(
+      rewriter, op.getLoc(), /*resultTypes=*/TypeRange{}, newOperands));
+  rewriter.inlineRegionBefore(op->getRegion(0), newOp->getRegion(0),
+                              newOp->getRegion(0).begin());
 
   // Replace the results of the old op with the new output buffers.
   replaceOpWithBufferizedValues(rewriter, op, newOutputBuffers);