[mlir][linalg] Fix insertion point bug in D144022

This should have been part of D144022.

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
index 261ad97..1915c2a 100644 (file)
--- a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
@@ -264,6 +264,7 @@ Value linalg::bufferizeToAllocation(RewriterBase &rewriter, PadOp padOp,
   // Create buffer allocation.
   Value alloc =
       createAllocationForTensor(rewriter, loc, padOp.getResult(), memorySpace);
+  rewriter.setInsertionPointAfter(alloc.getDefiningOp());
 
   // Create linalg.fill or linalg.generic.
   Operation *fillOp = movePaddingToFillOrGenericOp(rewriter, loc, padOp, alloc);
@@ -344,7 +345,7 @@ Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Value value,
   if (auto bbArg = value.dyn_cast<BlockArgument>()) {
     rewriter.setInsertionPointToStart(bbArg.getOwner());
   } else {
-    rewriter.setInsertionPoint(value.getDefiningOp());
+    rewriter.setInsertionPointAfter(value.getDefiningOp());
   }
   Location loc = value.getLoc();
 
@@ -352,6 +353,7 @@ Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Value value,
   Value alloc = createAllocationForTensor(rewriter, loc, value, memorySpace);
 
   // Create memref.tensor_store.
+  rewriter.setInsertionPointAfter(alloc.getDefiningOp());
   rewriter.create<memref::TensorStoreOp>(loc, value, alloc);
 
   // Create bufferization.to_tensor with "restrict" and "writable". The returned

diff --git a/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir b/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir
index be277ba..0b3942f 100644 (file)
--- a/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir
@@ -110,3 +110,25 @@ transform.sequence failures(propagate) {
   // Make sure that One-Shot Bufferize can bufferize the rest.
   transform.bufferization.one_shot_bufferize %arg1
 }
+
+// -----
+
+// CHECK-LABEL: func @materialization_of_opresult(
+//       CHECK:   %[[t:.*]] = "dummy.some_op"
+//       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4>
+//       CHECK:   memref.tensor_store %[[t]], %[[alloc]]
+//       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]]
+//       CHECK:   return %[[r]]
+func.func @materialization_of_opresult(%idx: index) -> tensor<?x10xindex> {
+  %t = "dummy.some_op"() : () -> (tensor<?x10xindex>)
+  return %t : tensor<?x10xindex>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["dummy.some_op"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.get_result %0[0] : (!pdl.operation) -> !transform.any_value
+  %2 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
+}
+
+