[mlir][linalg] BufferizeToAllocationOp: Bufferize ops, not values

[springerm]

Summary

The bufferize_to_allocation transform op now operates on payload ops, not payload values. Only ops can be bufferized, not values.

Also remove the replacement result from the transform op.

Depends On: D153874

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 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:
-    ```
-    %0 = "some_op"() : () -> (tensor<10xf32>)
-    %1 = memref.alloc() : memref<10xf32>
-    memref.tensor_store %0, %1 : memref<10xf32>
-    %2 = bufferization.to_tensor %1 restrict writable : memref<10xf32>
-    "some_use"(%2) : (tensor<10xf32>) -> ()
-    ```
-
-    This transform has optimized lowerings for certain targets that are results
-    of non-DPS ops. For such targets, not only a buffer allocation is emitted
-    but also the defining op is bufferized. This is to avoid a second
-    allocation for the missing destination of the non-DPS op (when subsequently
-    running a bufferization pass/transform). Currently supported ops with
-    optimized lowerings:
-    - tensor.pad
+    This transform bufferizes the targeted operation and materializes the
+    result in a new allocation. It replaces all original uses of the target
+    result with the newly allocated buffer, 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.
+
+    Only `tensor.pad` targets are supported at the moment.
+    - `tensor.pad` is lowered to an allocation, followed by a `linalg.fill` and
+      and a buffer copy (all on memrefs).
 
     An optional memory space attribute can be specified for the materialized
     buffer allocation.
 
     #### Return modes
 
-    This operation consumes the `target` handle and produces the `replacement`
-    and `allocated_buffer` handles. It always succeeds.
+    This operation consumes the `target` handle and produces the
+    `allocated_buffer` handle. It always succeeds.
   }];
 
-  let arguments = (ins Transform_AnyValue:$target,
+  let arguments = (ins TransformHandleTypeInterface:$target,
                        OptionalAttr<AnyAttr>:$memory_space);
-  let results = (outs Transform_AnyValue:$allocated_buffer,
-                      Transform_AnyValue:$replacement);
-  let assemblyFormat = "$target attr-dict";
+  let results = (outs Transform_AnyValue:$allocated_buffer);
+  let assemblyFormat = "$target attr-dict `:` type($target)";
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 def DecomposeOp : Op<Transform_Dialect, "structured.decompose",
     [FunctionalStyleTransformOpTrait, 

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

 ///
 /// %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, this function returns the
-/// newly allocated buffer. Furthermore, the result of the
-/// bufferization.to_tensor op is optionally returned via `replacement`.
+/// newly allocated buffer.
 Value bufferizeToAllocation(RewriterBase &rewriter, tensor::PadOp padOp,
-                            Attribute memorySpace = {},
-                            Value *replacement = nullptr);
+                            Attribute memorySpace = {});
 
-/// Materialize a buffer allocation for the given tensor value. E.g.:
+/// Bufferize the given op with tensor semantics and materialize the result in
+/// a newly allocated buffer. E.g.:
 ///
-/// %alloc = memref.alloc
-/// memref.tensor_store %value, %alloc
-/// %0 = bufferization.to_tensor %alloc restrict writable
+/// Only tensor.pad is supported at the moment.
 ///
-/// 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 = {},
-                            Value *replacement = nullptr);
+/// This function returns the newly allocated buffer.
+Value bufferizeToAllocation(RewriterBase &rewriter, Operation *op,
+                            Attribute memorySpace = {});
 
 /// Try to eliminate tensor::EmptyOps inside `op` that are anchored on a
 /// LinalgOp. This transforms looks for LinalgOps that have an unused output
 /// operand and an input operand that is rooted in a tensor::EmptyOp. The
 /// tensor::EmptyOp uses are replaced with the output operand and the two
 /// operands of the LinalgOp are swapped.
 ///
 /// Example:

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();
-  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);
+  for (Operation *op : state.getPayloadOps(getTarget())) {
+    Value buffer = linalg::bufferizeToAllocation(rewriter, op, memorySpace);
+    if (!buffer) {
+      DiagnosedSilenceableFailure diag = emitSilenceableError()
+                                         << "failed to bufferize operation";
+      diag.attachNote(op->getLoc()) << "target payload op";
+      return diag;
+    }
     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(getReplacement(), effects);
   producesHandle(getAllocatedBuffer(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 

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, Value *replacement) {
+                                    Attribute memorySpace) {
   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);
-
-  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 =
mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
   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, Value *replacement) {
+Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Operation *op,
+                                    Attribute memorySpace) {
   // Call specialized overload for certain ops.
-  if (auto padOp = value.getDefiningOp<PadOp>())
-    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());
-  } else {
-    rewriter.setInsertionPointAfter(value.getDefiningOp());
-  }
-  Location loc = value.getLoc();
+  if (auto padOp = dyn_cast<tensor::PadOp>(op))
+    return bufferizeToAllocation(rewriter, padOp, memorySpace);
 
-  // Create buffer allocation.
-  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
-  // 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); });
-  }
-
-  if (replacement)
-    *replacement = toTensorOp;
-  return alloc;
+  // TODO: Support other ops.
+  return nullptr;
 }
 
 namespace {
 
 template <typename OpTy>
 LogicalResult rewriteOpInDestinationPassingStyle(OpTy op,
                                                  PatternRewriter &rewriter) {
   return linalg::rewriteInDestinationPassingStyle(rewriter, op);

mlir/test/Dialect/Linalg/pad-to-specific-memory-space.mlir

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 1]
   } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
-  %pad_result = transform.get_result %pad[0] : (!transform.any_op) -> !transform.any_value
-  %buffer, %replacement = transform.structured.bufferize_to_allocation %pad_result {memory_space = 3}
+  %buffer = transform.structured.bufferize_to_allocation %pad {memory_space = 3} : !transform.any_op
   %2 = transform.bufferization.one_shot_bufferize %arg1 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 
 }

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

^bb0(%arg0: index, %arg1: index):
     tensor.yield %c : 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, %3 = transform.structured.bufferize_to_allocation %1
+  %2 = transform.structured.bufferize_to_allocation %0 : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @tensor_pad_constant(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[src:.*]] = bufferization.to_memref %[[t]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc
 //       CHECK:   %[[subview:.*]] = memref.subview %[[alloc]]
 //       CHECK:   memref.copy %[[src]], %[[subview]]
 //       CHECK:   bufferization.to_tensor %[[alloc]] restrict writable
 func.func @tensor_pad_constant(%t: tensor<?x10xindex>, %l2: index, %h1: index,
                                %h2: index) -> tensor<?x?xindex> {
   %0 = tensor.pad %t low[5, %l2] high[%h1, %h2] {
   ^bb0(%arg0: index, %arg1: index):
     %c = arith.constant 50 : index
     tensor.yield %c : 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, %3 = transform.structured.bufferize_to_allocation %1
+  %2 = transform.structured.bufferize_to_allocation %0 : !transform.any_op
   // Make sure that One-Shot Bufferize can bufferize the rest.
   %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, %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, %3 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
-  // Make sure that One-Shot Bufferize can bufferize the rest.
-  %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, %3 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
-}
-
-