[mlir][linalg][bufferize] Add pass options for `createDeallocs`

[springerm]

Summary

This change makes it possible to use a different buffer deallocation strategy. E.g., -buffer-deallocation can be used, which also works for allocations that are not in destination-passing style.

mlir/include/mlir/Dialect/Linalg/Passes.td

Option<"useAlloca", "use-alloca", "bool",
            "Use stack allocations for memrefs (for testing purposes only)">,
     Option<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
            /*default=*/"0",
            "Analyze ops in random order with a given seed (fuzzer)">,
     Option<"initTensorElimination", "init-tensor-elimination", "bool",
             /*default=*/"false",
            "(Experimental) Try to eliminate init_tensor operations that are "
            "anchored at an insert_slice op">,
+    Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
+           "Specify if buffers should be deallocated. For compatibility with "
+           "core bufferization passes.">,
   ];
   let constructor = "mlir::createLinalgComprehensiveModuleBufferizePass()";
 }
 
 def LinalgFoldUnitExtentDims : Pass<"linalg-fold-unit-extent-dims", ""> {
   let summary = "Remove unit-extent dimension in Linalg ops on tensors";
   let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
   let options = [

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

options->allocationFns->memCpyFn = [](OpBuilder &b, Location loc, Value from,
     b.create<linalg::CopyOp>(loc, from, to);
   };
 
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->analysisFuzzerSeed = analysisFuzzerSeed;
   options->testAnalysisOnly = testAnalysisOnly;
   options->printConflicts = printConflicts;
+  options->createDeallocs = createDeallocs;
 
   // Enable InitTensorOp elimination.
   if (initTensorElimination) {
     options->addPostAnalysisStep<
         linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   }
 
   if (!allowReturnMemref)

mlir/test/Dialect/Linalg/comprehensive-function-bufferize-compat.mlir

+// RUN: mlir-opt %s \
+// RUN:     -test-comprehensive-function-bufferize="allow-return-memref allow-unknown-ops create-deallocs=0" \
+// RUN:     -split-input-file | \
+// RUN: FileCheck %s --check-prefix=CHECK-NODEALLOC
+
+// RUN: mlir-opt %s \
+// RUN:     -test-comprehensive-function-bufferize="allow-return-memref allow-unknown-ops create-deallocs=0" \
+// RUN:     -buffer-deallocation | \
+// RUN: FileCheck %s --check-prefix=CHECK-BUFFERDEALLOC
+
+// CHECK-NODEALLOC-LABEL: func @out_of_place_bufferization
+// CHECK-BUFFERDEALLOC-LABEL: func @out_of_place_bufferization
+func @out_of_place_bufferization(%t1 : tensor<?xf32>) -> (f32, f32) {
+  //     CHECK-NODEALLOC: memref.alloc
+  //     CHECK-NODEALLOC: memref.copy
+  // CHECK-NODEALLOC-NOT: memref.dealloc
+
+  //     CHECK-BUFFERDEALLOC: %[[alloc:.*]] = memref.alloc
+  //     CHECK-BUFFERDEALLOC: memref.copy
+  //     CHECK-BUFFERDEALLOC: memref.dealloc %[[alloc]]
+
+  %cst = arith.constant 0.0 : f32
+  %idx = arith.constant 5 : index
+
+  // This bufferizes out-of-place. An allocation + copy will be inserted.
+  %0 = tensor.insert %cst into %t1[%idx] : tensor<?xf32>
+
+  %1 = tensor.extract %t1[%idx] : tensor<?xf32>
+  %2 = tensor.extract %0[%idx] : tensor<?xf32>
+  return %1, %2 : f32, f32
+}

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

struct TestComprehensiveFunctionBufferize
   Option<unsigned> analysisFuzzerSeed{
       *this, "analysis-fuzzer-seed",
       llvm::cl::desc("Analyze ops in random order with a given seed (fuzzer)"),
       llvm::cl::init(0)};
   ListOption<std::string> dialectFilter{
       *this, "dialect-filter",
       llvm::cl::desc("Bufferize only ops from the specified dialects"),
       llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
+  Option<bool> createDeallocs{
+      *this, "create-deallocs",
+      llvm::cl::desc("Specify if buffers should be deallocated"),
+      llvm::cl::init(true)};
 };
 } // namespace
 
 void TestComprehensiveFunctionBufferize::runOnFunction() {
   auto options = std::make_unique<BufferizationOptions>();
 
   if (!allowReturnMemref)
     options->addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;
   options->analysisFuzzerSeed = analysisFuzzerSeed;
+  options->createDeallocs = createDeallocs;
 
   if (dialectFilter.hasValue()) {
     options->dialectFilter.emplace();
     for (const std::string &dialectNamespace : dialectFilter)
       options->dialectFilter->insert(dialectNamespace);
   }
 
   Operation *op = getFunction().getOperation();