[mlir][linalg] BufferizeToAllocationOp: Return handle to buffer

[springerm]

Summary

Add an additional result handle to the op. This new handle is mapped to the newly allocated buffer.

Depends On: D153853

[springerm]

rebase

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

 def BufferizeToAllocationOp : Op<Transform_Dialect,
     "structured.bufferize_to_allocation",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     This transform materializes an allocation for the targeted tensor value. It
     replaces all original uses of the target with the newly allocated buffer,
-    wrapped in a `bufferization.to_tensor` op. It returns a handle to the result
-    of the `to_tensor` op.
+    wrapped in a `bufferization.to_tensor` op. It returns a handle to the newly
+    allocated buffer. Furthermore, it returns a handle to the result of the
+    `to_tensor` op.
 
     Example:
     ```
     %0 = "some_op"() : () -> (tensor<10xf32>)
     "some_use"(%0) : (tensor<10xf32>) -> ()
     ```
 
     Is rewritten to:
let description = [{
     optimized lowerings:
     - tensor.pad
 
     An optional memory space attribute can be specified for the materialized
     buffer allocation.
 
     #### Return modes
 
-    This operation consumes the `target` handle and produces the `transformed`
-    handle. It always succeeds.
+    This operation consumes the `target` handle and produces the `replacement`
+    and `allocated_buffer` handles. It always succeeds.
   }];
 
   let arguments = (ins Transform_AnyValue:$target,
                        OptionalAttr<AnyAttr>:$memory_space);
-  let results = (outs Transform_AnyValue:$transformed);
+  let results = (outs Transform_AnyValue:$allocated_buffer,
+                      Transform_AnyValue:$replacement);
   let assemblyFormat = "$target attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 def DecomposeOp : Op<Transform_Dialect, "structured.decompose",

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// is lowered to:
 ///
 /// %alloc = memref.alloc
 /// linalg.fill ... outs(%alloc)
 /// %subview = memref.subview %alloc [%l] [...] [1]
 /// memref.tensor_store %t, %subview
 /// %0 = bufferization.to_tensor %alloc restrict writable
 ///
-/// In addition to rewriting the IR as shown above, the result of the
-/// bufferization.to_tensor op is returned.
+/// In addition to rewriting the IR as shown above, this function returns the
+/// newly allocated buffer. Furthermore, the result of the
+/// bufferization.to_tensor op is optionally returned via `replacement`.
 Value bufferizeToAllocation(RewriterBase &rewriter, tensor::PadOp padOp,
-                            Attribute memorySpace = {});
+                            Attribute memorySpace = {},
+                            Value *replacement = nullptr);
 
 /// Materialize a buffer allocation for the given tensor value. E.g.:
 ///
 /// %alloc = memref.alloc
 /// memref.tensor_store %value, %alloc
 /// %0 = bufferization.to_tensor %alloc restrict writable
 ///
 /// In case `value` is a tensor.pad result, the corresponding overload is used
 /// internally to produce a better bufferization.
+///
+/// In addition to rewriting the IR as shown above, this function returns the
+/// newly allocated buffer. Furthermore, the result of the
+/// bufferization.to_tensor op is optionally returned via `replacement`.
 Value bufferizeToAllocation(RewriterBase &rewriter, Value value,
-                            Attribute memorySpace = {});
+                            Attribute memorySpace = {},
+                            Value *replacement = nullptr);
 
 /// Fuse two `linalg.generic` operations that have a producer-consumer
 /// relationship captured through `fusedOperand`. The method expects
 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.
 struct ElementwiseOpFusionResult {
   Operation *fusedOp;
   llvm::DenseMap<Value, Value> replacements;
 };

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 // BufferizeToAllocationOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::BufferizeToAllocationOp::apply(
     transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   Attribute memorySpace =
       getMemorySpace().has_value() ? getMemorySpace().value() : Attribute();
-  auto transformed = llvm::to_vector(
-      llvm::map_range(state.getPayloadValues(getTarget()), [&](Value v) {
-        return linalg::bufferizeToAllocation(rewriter, v, memorySpace);
-      }));
-  results.setValues(cast<OpResult>(getTransformed()), transformed);
+  SmallVector<Value> replacements;
+  SmallVector<Value> allocatedBuffers;
+  for (Value value : state.getPayloadValues(getTarget())) {
+    Value replacement;
+    Value buffer = linalg::bufferizeToAllocation(rewriter, value, memorySpace,
+                                                 &replacement);
+    replacements.push_back(replacement);
+    allocatedBuffers.push_back(buffer);
+  }
+  results.setValues(cast<OpResult>(getReplacement()), replacements);
+  results.setValues(cast<OpResult>(getAllocatedBuffer()), allocatedBuffers);
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::BufferizeToAllocationOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
-  producesHandle(getTransformed(), effects);
+  producesHandle(getReplacement(), effects);
+  producesHandle(getAllocatedBuffer(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure

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

static Value createAllocationForTensor(RewriterBase &rewriter, Location loc,
   // Place deallocation at the end of the block.
   rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
   rewriter.create<memref::DeallocOp>(loc, alloc);
 
   return alloc;
 }
 
 Value linalg::bufferizeToAllocation(RewriterBase &rewriter, PadOp padOp,
-                                    Attribute memorySpace) {
+                                    Attribute memorySpace, Value *replacement) {
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(padOp);
   Location loc = padOp.getLoc();
 
   // Create buffer allocation.
   Value alloc =
       createAllocationForTensor(rewriter, loc, padOp.getResult(), memorySpace);
   rewriter.setInsertionPointAfter(alloc.getDefiningOp());
Value subview = rewriter.create<memref::SubViewOp>(
       loc, alloc, /*offsets=*/padOp.getMixedLowPad(), sizes, strides);
   rewriter.create<memref::TensorStoreOp>(loc, padOp.getSource(), subview);
 
   // Create bufferization.to_tensor with "restrict" and "writable". The returned
   // tensor is a new buffer allocation, so it does not alias with any buffer.
   Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
       loc, alloc, /*restrict=*/true, /*writable=*/true);
   rewriter.replaceOp(padOp, toTensorOp);
-  return toTensorOp;
+
+  if (replacement)
+    *replacement = toTensorOp;
+  return alloc;
 }
 
 /// Lower tensor.from_elements to a sequence of chained tensor.insert.
 FailureOr<Operation *> mlir::linalg::rewriteInDestinationPassingStyle(
     RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp) {
   Location loc = fromElementsOp.getLoc();
   RankedTensorType tensorType =
       cast<RankedTensorType>(fromElementsOp.getType());
SmallVector<OpFoldResult> sliceStrides(resultType.getRank(),
                                          rewriter.getIndexAttr(1));
   auto insertSliceOp = rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
       padOp, padOp.getSource(), fillOp->getResult(0),
       /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
   return insertSliceOp.getOperation();
 }
 
 Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Value value,
-                                    Attribute memorySpace) {
+                                    Attribute memorySpace, Value *replacement) {
   // Call specialized overload for certain ops.
   if (auto padOp = value.getDefiningOp<PadOp>())
-    return bufferizeToAllocation(rewriter, padOp, memorySpace);
+    return bufferizeToAllocation(rewriter, padOp, memorySpace, replacement);
 
   // Collect all uses.
   SmallVector<OpOperand *> uses = llvm::to_vector(
       llvm::map_range(value.getUses(), [](OpOperand &use) { return &use; }));
 
   OpBuilder::InsertionGuard g(rewriter);
   if (auto bbArg = dyn_cast<BlockArgument>(value)) {
     rewriter.setInsertionPointToStart(bbArg.getOwner());
Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Value value,
   // tensor is a new buffer allocation, so it does not alias with any buffer.
   Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
       loc, alloc, /*restrict=*/true, /*writable=*/true);
   for (OpOperand *use : uses) {
     rewriter.updateRootInPlace(use->getOwner(),
                                [&]() { use->set(toTensorOp); });
   }
 
-  return toTensorOp;
+  if (replacement)
+    *replacement = toTensorOp;
+  return alloc;
 }
 
 namespace {
 
 template <typename OpTy>
 LogicalResult rewriteOpInDestinationPassingStyle(OpTy op,
                                                  PatternRewriter &rewriter) {
   return linalg::rewriteInDestinationPassingStyle(rewriter, op);

mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir

func.func @tensor_pad_constant(%t: tensor<?x10xindex>, %l2: index, %h1: index,
   } : tensor<?x10xindex> to tensor<?x?xindex>
   return %0 : tensor<?x?xindex>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.get_result %0[0] : (!transform.any_op) -> !transform.any_value
-  %2 = transform.structured.bufferize_to_allocation %1
+  %2, %3 = transform.structured.bufferize_to_allocation %1
 }
 
 // -----
 
 // CHECK-LABEL: func @tensor_pad_constant(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[src:.*]] = bufferization.to_memref %[[t]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc
func.func @tensor_pad_constant(%t: tensor<?x10xindex>, %l2: index, %h1: index,
   } : tensor<?x10xindex> to tensor<?x?xindex>
   return %0 : tensor<?x?xindex>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.get_result %0[0] : (!transform.any_op) -> !transform.any_value
-  %2 = transform.structured.bufferize_to_allocation %1
+  %2, %3 = transform.structured.bufferize_to_allocation %1
   // Make sure that One-Shot Bufferize can bufferize the rest.
-  %3 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
+  %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @materialization_of_bbarg(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[c0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[dim:.*]] = tensor.dim %[[t]], %[[c0]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%[[dim]]) : memref<?x10xindex, 4>
 //       CHECK:   memref.tensor_store %[[t]], %[[alloc]]
 //       CHECK:   %[[alloc_t:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
 //       CHECK:   %[[r:.*]] = tensor.extract %[[alloc_t]]
 //       CHECK:   memref.dealloc %[[alloc]]
 //       CHECK:   return %[[r]]
 func.func @materialization_of_bbarg(%t: tensor<?x10xindex>, %idx: index) -> index {
   %r = tensor.extract %t[%idx, %idx] : tensor<?x10xindex>
   return %r : index
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = test_produce_value_handle_to_argument_of_parent_block %0, 0 : (!transform.any_op) -> !transform.any_value
-  %2 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
+  %2, %3 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
 }
 
 // -----
 
 // CHECK-LABEL: func @materialization_of_bbarg(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[m:.*]] = bufferization.to_memref %[[t]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4>
 //       CHECK:   memref.copy %[[m]], %[[alloc]]
 //       CHECK:   %[[r:.*]] = memref.load %[[alloc]]
 //       CHECK:   return %[[r]]
 func.func @materialization_of_bbarg(%t: tensor<?x10xindex>, %idx: index) -> index {
   %r = tensor.extract %t[%idx, %idx] : tensor<?x10xindex>
   return %r : index
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = test_produce_value_handle_to_argument_of_parent_block %0, 0 : (!transform.any_op) -> !transform.any_value
-  %2 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
+  %2, %3 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
   // Make sure that One-Shot Bufferize can bufferize the rest.
-  %3 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
+  %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
 // 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: !transform.any_op):
   %0 = transform.structured.match ops{["dummy.some_op"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.get_result %0[0] : (!transform.any_op) -> !transform.any_value
-  %2 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
+  %2, %3 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
 }