[mlir][linalg] BufferizeToAllocationOp: Add option to materialize buffers for operands

[springerm]

Summary

Add an option that does not bufferize the targeted op itself, but just materializes a buffer for the destination operands. This is useful for partial bufferization of complex ops such as scf.forall, which need special handling (and an analysis if the region).

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

let description = [{
     a fully dynamic layout is assumed for best compatibility. Users should use
     "memref.tensor_store" when possible.
 
     "memref.alloc" is used for new buffer allocations. The buffer is deallocated
     at the end of the block. Custom allocation ops can be specified via
     `alloc_op`. Currently supported are "memref.alloc" and "memref.alloca". In
     case of a "memref.alloca", the buffer is not deallocated.
 
+    If `bufferize_destination_only` is set, only the destination operands of the
+    op are bufferized to a new memory allocation, but not the op itself.
+
     #### Return modes
 
     This operation consumes the `target` handle and produces the
     `allocated_buffer` and `new_ops` handles. It always succeeds.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        OptionalAttr<AnyAttr>:$memory_space,
                        DefaultValuedAttr<StrAttr, "\"memref.tensor_store\"">:
                            $memcpy_op,
                        DefaultValuedAttr<StrAttr, "\"memref.alloc\"">:
-                           $alloc_op);
+                           $alloc_op,
+                       UnitAttr:$bufferize_destination_only);
   let hasVerifier = 1;
   let results = (outs Transform_AnyValue:$allocated_buffer,
                       Transform_AnyOpType:$new_ops);
   let assemblyFormat = "$target attr-dict `:` type($target)";
-  
   let builders = [
     OpBuilder<(ins "Value":$target, "Attribute":$memorySpace)>,
     OpBuilder<(ins "Value":$target, "int64_t":$memorySpace)>
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp

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

 //===----------------------------------------------------------------------===//
 
 struct BufferizeToAllocationOptions {
   enum class AllocOp { MemrefAlloc = 0, MemrefAlloca = 1 };
   AllocOp allocOp = AllocOp::MemrefAlloc;
 
   enum class MemcpyOp { MemrefTensorStore = 0, MemrefCopy = 1, LinalgCopy = 2 };
   MemcpyOp memcpyOp = MemcpyOp::MemrefTensorStore;
+
+  /// If set to "true", only the destination tensor operands are bufferized to
+  /// a new allocation (and wrapped in "bufferization.to_tensor"), but not the
+  /// targeted op itself.
+  bool bufferizeDestinationOnly = false;
 };
 
 /// Materialize a buffer allocation for the given tensor.pad op and lower the
 /// op to linalg.fill/linalg.generic + memref.tensor_store. E.g.:
 ///
 /// %0 = tensor.pad low[%l] high[%h] %t ...
 ///
 /// is lowered to:

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

if (getAllocOp() == "memref.alloc") {
     options.allocOp =
         linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloc;
   } else if (getAllocOp() == "memref.alloca") {
     options.allocOp =
         linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloca;
   } else {
     llvm_unreachable("invalid alloc op");
   }
+  options.bufferizeDestinationOnly = getBufferizeDestinationOnly();
 
   // Bufferize ops.
   Attribute memorySpace =
       getMemorySpace().has_value() ? getMemorySpace().value() : Attribute();
   SmallVector<Value> allocatedBuffers;
   for (Operation *op : state.getPayloadOps(getTarget())) {
     Value buffer =
         linalg::bufferizeToAllocation(rewriter, options, op, memorySpace);

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

createAllocationForTensor(RewriterBase &rewriter, Location loc, Value value,
   }
 
   return alloc;
 }
 
 Value linalg::bufferizeToAllocation(
     RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
     PadOp padOp, Attribute memorySpace, Operation *insertionPoint) {
+  // tensor.pad does not have a destination operand.
+  assert(!options.bufferizeDestinationOnly && "invalid options");
+
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(insertionPoint ? insertionPoint : padOp);
   Location loc = padOp.getLoc();
 
   // Create buffer allocation.
   Value alloc = createAllocationForTensor(rewriter, loc, padOp.getResult(),
                                           options, memorySpace);
   rewriter.setInsertionPoint(padOp);
Value linalg::bufferizeToAllocation(
   assert(isa<vector::YieldOp>(yieldOp) && "expected yield op terminator");
 
   // Bufferize maskable op. By default, place the buffer allocation right before
   // the mask op.
   Value alloc = bufferizeToAllocation(
       rewriter, options, maskOp.getMaskableOp(), memorySpace,
       /*insertionPoint=*/insertionPoint ? insertionPoint : maskOp);
 
+  if (options.bufferizeDestinationOnly)
+    return alloc;
+
   // Bufferize terminator.
   rewriter.setInsertionPoint(yieldOp);
   if (failed(cast<bufferization::BufferizableOpInterface>(yieldOp).bufferize(
           rewriter, bufferizationOptions)))
     return nullptr;
 
   // Erase dead to_tensor ops inside of the mask op. This is necessary because
   // there only be one op (apart from the terminator) inside the mask op.
Value linalg::bufferizeToAllocation(
 
   // Only bufferizable ops are supported.
   auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
   if (!bufferizableOp)
     return nullptr;
   BufferizationOptions bufferizationOptions;
   AnalysisState state(bufferizationOptions);
 
+#ifndef NDEBUG
+  // Ops with nested tensor ops are not supported yet. At the moment, this
+  // function just bufferizes the given op itself, but not its body.
+  op->walk([&](Operation *nestedOp) {
+    if (op == nestedOp)
+      return;
+    if (llvm::any_of(nestedOp->getOperands(),
+                     [](Value v) { return v.getType().isa<TensorType>(); }))
+      llvm_unreachable("ops with nested tensor ops are not supported yet");
+    if (llvm::any_of(nestedOp->getResults(),
+                     [](Value v) { return v.getType().isa<TensorType>(); }))
+      llvm_unreachable("ops with nested tensor ops are not supported yet");
+  });
+#endif // NDEBUG
+
   // 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;
Value alloc = createAllocationForTensor(
         rewriter, op->getLoc(), operand->get(), options, memorySpace);
     allocs.push_back(alloc);
     if (!state.findDefinitions(operand->get()).empty()) {
       // Initialize buffer with a copy of the operand data. Not needed if the
       // tensor is uninitialized.
       createMemcpy(rewriter, op->getLoc(), operand->get(), alloc, options);
     }
     rewriter.updateRootInPlace(op, [&]() {
-      operand->set(rewriter.create<ToTensorOp>(op->getLoc(), alloc));
+      auto toTensorOp = rewriter.create<ToTensorOp>(op->getLoc(), alloc);
+      operand->set(toTensorOp);
+      if (options.bufferizeDestinationOnly) {
+        rewriter.updateRootInPlace(toTensorOp, [&]() {
+          toTensorOp.setRestrict(true);
+          toTensorOp.setWritable(true);
+        });
+      }
     });
   }
 
+  if (options.bufferizeDestinationOnly)
+    return allocs.front();
+
   // Bufferize the op.
   rewriter.setInsertionPoint(op);
   if (failed(bufferizableOp.bufferize(rewriter, bufferizationOptions)))
     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) {

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

func.func @vector_mask(%t: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vector<16xi1>) -> tensor<?xf32> {
   return %r : tensor<?xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["vector.mask"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
 }
+
+// -----
+
+// CHECK-LABEL: func @tensor_insert_destination(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
+//       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4>
+//       CHECK:   memref.tensor_store %[[t]], %[[alloc]]
+//       CHECK:   %[[t2:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
+//       CHECK:   %[[inserted:.*]] = tensor.insert %{{.*}} into %[[t2]]
+//       CHECK:   memref.dealloc %[[alloc]]
+//       CHECK:   return %[[inserted]]
+func.func @tensor_insert_destination(%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, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only} : !transform.any_op
+}