[mlir][linalg] BufferizeToAllocationOp: Bufferize non-allocating ops

[springerm]

Summary

Until now, only tensor.pad ops could be bufferized to an allocation. This revision adds support for all bufferizable ops that do not already bufferize to an allocation. (Those still need special handling.)

Depends On: D153970

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

def BufferizeToAllocationOp : Op<Transform_Dialect,
   let description = [{
     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.
+    Only bufferizable ops are that bufferize to a memory write or have an
+    aliasing OpOperand (and do not themselves bufferize to an allocation) are
+    supported. They are bufferized using their BufferizableOpInterface
+    implementation. E.g.:
+
+    ```
+    %0 = tensor.insert %f into %dest[%pos] : tensor<10xf32>
+    ```
+
+    Is bufferized to:
+
+    ```
+    %alloc = memref.alloc() : memref<10xf32>
+    memref.tensor_store %dest, %alloc : memref<10xf32>
+    memref.store %f, %alloc[%pos] : memref<10xf32>
+    %0 = bufferization.to_tensor %alloc restrict writable : memref<10xf32>
+    ```
+
+    Selected ops that bufferize to an allocation are also supported:
     - `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
 

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

 /// %0 = bufferization.to_tensor %alloc restrict writable
 ///
 /// In addition to rewriting the IR as shown above, this function returns the
 /// newly allocated buffer.
 Value bufferizeToAllocation(RewriterBase &rewriter, tensor::PadOp padOp,
                             Attribute memorySpace = {});
 
 /// Bufferize the given op with tensor semantics and materialize the result in
-/// a newly allocated buffer. E.g.:
+/// a newly allocated buffer.
 ///
-/// Only tensor.pad is supported at the moment.
+/// Only bufferizable ops that bufferize to a memory write or have an
+/// aliasing OpOperand (and do not themselves bufferize to an allocation) are
+/// supported. They are bufferized using their BufferizableOpInterface
+/// implementation.
 ///
-/// This function returns the newly allocated buffer.
+/// Selected ops that bufferize to an allocation are also supported:
+/// - tensor.pad
 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.

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

WalkResult walkResult = op->walk([&](BufferizableOpInterface op) {
     // Skip ops that are not in the filter.
     if (!options.isOpAllowed(op.getOperation()))
       return WalkResult::advance();
 
     // Input IR may not contain any ToTensorOps without the "restrict"
     // attribute. Such tensors may alias any other tensor, which is currently
     // not handled in the analysis.
     if (auto toTensorOp = dyn_cast<ToTensorOp>(op.getOperation())) {
-      if (!toTensorOp.getRestrict()) {
+      if (!toTensorOp.getRestrict() && !toTensorOp->getUses().empty()) {
         op->emitError("to_tensor ops without `restrict` are not supported by "
                       "One-Shot Analysis");
         return WalkResult::interrupt();
       }
     }
 
     for (OpOperand &opOperand : op->getOpOperands()) {
       if (isa<TensorType>(opOperand.get().getType())) {

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

mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
   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, Operation *op,
                                     Attribute memorySpace) {
+  using namespace bufferization;
+
   // Call specialized overload for certain ops.
   if (auto padOp = dyn_cast<tensor::PadOp>(op))
     return bufferizeToAllocation(rewriter, padOp, memorySpace);
 
-  // TODO: Support other ops.
+  // Only bufferizable ops are supported.
+  auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
+  if (!bufferizableOp)
+    return nullptr;
+  BufferizationOptions options;
+  AnalysisState state(options);
+
+  // Gather tensor results.
+  SmallVector<OpResult> tensorResults;
+  for (OpResult result : op->getResults()) {
+    if (!result.getType().isa<TensorType>())
+      continue;
+    // Unranked tensors are not supported
+    if (!isa<RankedTensorType>(result.getType()))
+      return nullptr;
+    // Ops that bufferize to an allocation are not supported.
+    if (bufferizableOp.bufferizesToAllocation(result))
+      return nullptr;
+    tensorResults.push_back(result);
+  }
+
+  // Gather all operands that should bufferize to a new allocation. I.e.,
+  // bufferize out-of-place.
+  SmallVector<OpOperand *> outOfPlaceOperands, resultUses;
+  auto addOutOfPlaceOperand = [&](OpOperand *operand) {
+    if (llvm::find(outOfPlaceOperands, operand) == outOfPlaceOperands.end())
+      outOfPlaceOperands.push_back(operand);
+  };
+  for (OpResult result : tensorResults) {
+    AliasingOpOperandList aliasingOperands =
+        state.getAliasingOpOperands(result);
+    for (const AliasingOpOperand &operand : aliasingOperands) {
+      addOutOfPlaceOperand(operand.opOperand);
+      for (OpOperand &resultUse : result.getUses())
+        resultUses.push_back(&resultUse);
+    }
+  }
+  for (OpOperand &operand : op->getOpOperands()) {
+    if (!state.bufferizesToMemoryWrite(operand))
+      continue;
+    if (!isa<RankedTensorType>(operand.get().getType()))
+      return nullptr;
+    addOutOfPlaceOperand(&operand);
+  }
+  // TODO: Support multiple buffers.
+  if (outOfPlaceOperands.size() != 1)
+    return nullptr;
+
+  // Allocate buffers.
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+  SmallVector<Value> allocs;
+  for (OpOperand *operand : outOfPlaceOperands) {
+    Value alloc = createAllocationForTensor(rewriter, op->getLoc(),
+                                            operand->get(), memorySpace);
+    allocs.push_back(alloc);
+    // Initialize buffer with a copy of the operand data.
+    // TODO: Do not copy uninitialized tensors such as tensor.empty.
+    rewriter.create<memref::TensorStoreOp>(op->getLoc(), operand->get(), alloc);
+    rewriter.updateRootInPlace(op, [&]() {
+      operand->set(rewriter.create<ToTensorOp>(op->getLoc(), alloc));
+    });
+  }
+
+  // Bufferize the op.
+  if (failed(bufferizableOp.bufferize(rewriter, options)))
     return nullptr;
+
+  // Set "restrict" attribute, indicating that no other tensor aliases with
+  // this tensor. That is because we just allocated a new buffer for the tensor.
+  for (OpOperand *resultUse : resultUses) {
+    auto toTensorOp = resultUse->get().getDefiningOp<ToTensorOp>();
+    assert(toTensorOp && "expected to_tensor op");
+    rewriter.updateRootInPlace(toTensorOp, [&]() {
+      toTensorOp.setRestrict(true);
+      toTensorOp.setWritable(true);
+    });
+  }
+  return allocs.front();
 }
 
 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

-// RUN: mlir-opt -split-input-file \
+// RUN: mlir-opt -split-input-file -verify-diagnostics \
 // RUN:   -test-transform-dialect-interpreter -canonicalize \
 // RUN:   -allow-unregistered-dialect -split-input-file %s | FileCheck %s
 
 // CHECK:       #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>
 // CHECK:       #[[$map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 10)>
 // CHECK-LABEL: func @tensor_pad_constant(
 //  CHECK-SAME:   %[[t:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index
 //   CHECK-DAG:   %[[c0:.*]] = arith.constant 0 : index

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %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 @tensor_insert(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
+//       CHECK:   %[[m:.*]] = bufferization.to_memref %[[t]]
+//       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4>
+//       CHECK:   memref.copy %[[m]], %[[alloc]]
+//       CHECK:   memref.store %{{.*}}, %[[alloc]]
+//       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
+//       CHECK:   memref.dealloc %[[alloc]]
+//       CHECK:   return %[[r]]
+func.func @tensor_insert(%t: tensor<?x10xindex>, %idx: index, %v: index) -> tensor<?x10xindex> {
+  %r = tensor.insert %v into %t[%idx, %idx] : tensor<?x10xindex>
+  return %r : tensor<?x10xindex>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !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
+}
+
+// -----
+
+func.func @tensor_extract(%t: tensor<?x10xindex>, %idx: index) -> index {
+  // expected-note @below{{target payload op}}
+  %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
+  // expected-error @below{{failed to bufferize operation}}
+  %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
+}