[mlir][linalg] BufferizeToAllocationOp: Add option to specify custom alloc op

[springerm]

Summary

Supported ops are "memref.alloc" and "memref.alloca".

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

let description = [{
     later. Which concrete op will be used for the memory copy is up to the
     bufferization framework. Alternatively, a custom memcpy op can be specified
     via `memcpy_op`. Currently supported are "memref.copy" and "linalg.copy".
     In that case, the source of each memcpy must not have a custom memory space.
     Furthermore, because the future buffer layout unknown for a given tensor,
     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.
+
     #### 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);
+                           $memcpy_op,
+                       DefaultValuedAttr<StrAttr, "\"memref.alloc\"">:
+                           $alloc_op);
   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)>

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

 /// Return vector::CombiningKind for the given op.
 std::optional<vector::CombiningKind> getCombinerOpKind(Operation *combinerOp);
 
 //===----------------------------------------------------------------------===//
 // Bufferization-related transforms.
 //===----------------------------------------------------------------------===//
 
 struct BufferizeToAllocationOptions {
-  enum class MemcpyOp { MemrefTensorStore = 0, MemrefCopy = 1, LinalgCopy = 2 };
+  enum class AllocOp { MemrefAlloc = 0, MemrefAlloca = 1 };
+  AllocOp allocOp = AllocOp::MemrefAlloc;
 
+  enum class MemcpyOp { MemrefTensorStore = 0, MemrefCopy = 1, LinalgCopy = 2 };
   MemcpyOp memcpyOp = MemcpyOp::MemrefTensorStore;
 };
 
 /// 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 ...
 ///

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

if (getMemcpyOp() == "memref.tensor_store") {
     options.memcpyOp =
         linalg::BufferizeToAllocationOptions::MemcpyOp::MemrefCopy;
   } else if (getMemcpyOp() == "linalg.copy") {
     options.memcpyOp =
         linalg::BufferizeToAllocationOptions::MemcpyOp::LinalgCopy;
   } else {
     llvm_unreachable("invalid memcpy op");
   }
+  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");
+  }
 
   // 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);
void transform::BufferizeToAllocationOp::getEffects(
   producesHandle(getNewOps(), effects);
   modifiesPayload(effects);
 }
 
 LogicalResult transform::BufferizeToAllocationOp::verify() {
   if (getMemcpyOp() != "memref.tensor_store" &&
       getMemcpyOp() != "memref.copy" && getMemcpyOp() != "linalg.copy")
     return emitOpError() << "unsupported memcpy op";
+  if (getAllocOp() != "memref.alloc" && getAllocOp() != "memref.alloca")
+    return emitOpError() << "unsupported alloc op";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure

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

for (int64_t i = 0; i < tensorType.getRank(); ++i) {
     if (tensorType.isDynamicDim(i))
       dynSizes.push_back(
           b.create<DimOp>(value.getLoc(), value,
                           b.create<arith::ConstantIndexOp>(value.getLoc(), i)));
   }
   return dynSizes;
 }
 
-static Value createAllocationForTensor(RewriterBase &rewriter, Location loc,
-                                       Value value,
-                                       Attribute memorySpace = {}) {
+static Value
+createAllocationForTensor(RewriterBase &rewriter, Location loc, Value value,
+                          const linalg::BufferizeToAllocationOptions &options,
   OpBuilder::InsertionGuard g(rewriter);
   auto tensorType = cast<RankedTensorType>(value.getType());
 
   // Create buffer allocation.
   auto memrefType =
       cast<MemRefType>(bufferization::getMemRefTypeWithStaticIdentityLayout(
           tensorType, memorySpace));
   SmallVector<Value> dynamicSizes = reifyOrComputeDynamicSizes(rewriter, value);
-  Value alloc = rewriter.create<memref::AllocOp>(loc, memrefType, dynamicSizes);
 
+  Value alloc;
+  if (options.allocOp ==
+      linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloc) {
+    alloc = rewriter.create<memref::AllocOp>(loc, memrefType, dynamicSizes);
     // Place deallocation at the end of the block.
     rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
     rewriter.create<memref::DeallocOp>(loc, alloc);
+  } else if (options.allocOp ==
+             linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloca) {
+    alloc = rewriter.create<memref::AllocaOp>(loc, memrefType, dynamicSizes);
+    // No dealloc is needed.
+  }
 
   return alloc;
 }
 
 Value linalg::bufferizeToAllocation(
     RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
     PadOp padOp, Attribute memorySpace, Operation *insertionPoint) {
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(insertionPoint ? insertionPoint : padOp);
   Location loc = padOp.getLoc();
 
   // Create buffer allocation.
-  Value alloc =
-      createAllocationForTensor(rewriter, loc, padOp.getResult(), memorySpace);
+  Value alloc = createAllocationForTensor(rewriter, loc, padOp.getResult(),
+                                          options, memorySpace);
   rewriter.setInsertionPoint(padOp);
 
   if (!padOp.hasZeroLowPad() || !padOp.hasZeroHighPad()) {
     // Create linalg.fill or linalg.generic. Not needed if there is no padding.
     Operation *fillOp =
         movePaddingToFillOrGenericOp(rewriter, loc, padOp, alloc);
     rewriter.setInsertionPointAfter(fillOp);
   }
Value linalg::bufferizeToAllocation(
   if (outOfPlaceOperands.size() != 1)
     return nullptr;
 
   // Allocate buffers.
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(insertionPoint ? insertionPoint : op);
   SmallVector<Value> allocs;
   for (OpOperand *operand : outOfPlaceOperands) {
-    Value alloc = createAllocationForTensor(rewriter, op->getLoc(),
-                                            operand->get(), memorySpace);
+    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));

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

^bb1(%arg1: !transform.any_op):
   test_print_number_of_associated_payload_ir_ops %tensor_store : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @tensor_pad_constant_with_custom_copy(
 //   CHECK-NOT:   memref.tensor_store
 //   CHECK-NOT:   memref.copy
+//       CHECK:   memref.alloca
 //       CHECK:   linalg.copy
 func.func @tensor_pad_constant_with_custom_copy(
     %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
-  %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 3, memcpy_op = "linalg.copy"}: !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 3, alloc_op = "memref.alloca", memcpy_op = "linalg.copy"}: !transform.any_op
 
   // Ensure that one linalg.fill was generated.
   %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op
   // expected-remark @below{{1}}
   test_print_number_of_associated_payload_ir_ops %fill_op : !transform.any_op
 
   // Ensure that one linalg.copy was generated.
   %linalg_copy = transform.select "linalg.copy" in %new : (!transform.any_op) -> !transform.any_op
   // expected-remark @below{{1}}
   test_print_number_of_associated_payload_ir_ops %linalg_copy : !transform.any_op
 
+  // Ensure that one memref.alloca was generated.
+  %alloca = transform.select "memref.alloca" in %new : (!transform.any_op) -> !transform.any_op
+  // expected-remark @below{{1}}
+  test_print_number_of_associated_payload_ir_ops %alloca : !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_pad_constant(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>