From 57e714bcc813403f3003f522fcc4390d47a80c6e Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Wed, 12 Jan 2022 18:40:08 +0900
Subject: [PATCH] [mlir][linalg][bufferize] Add pass options for
 `createDeallocs`

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.

Differential Revision: https://reviews.llvm.org/D117096
---
 mlir/include/mlir/Dialect/Linalg/Passes.td         |  3 +++
 .../Transforms/ComprehensiveBufferizePass.cpp      |  1 +
 .../comprehensive-function-bufferize-compat.mlir   | 31 ++++++++++++++++++++++
 .../Dialect/Linalg/TestComprehensiveBufferize.cpp  |  5 ++++
 4 files changed, 40 insertions(+)
 create mode 100644 mlir/test/Dialect/Linalg/comprehensive-function-bufferize-compat.mlir

diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index b9ac216..b2befd0 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -58,6 +58,9 @@ def LinalgComprehensiveModuleBufferize :
             /*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()";
 }
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp
index 2da7df6..e92f852 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp
@@ -92,6 +92,7 @@ void LinalgComprehensiveModuleBufferize::runOnOperation() {
   options->analysisFuzzerSeed = analysisFuzzerSeed;
   options->testAnalysisOnly = testAnalysisOnly;
   options->printConflicts = printConflicts;
+  options->createDeallocs = createDeallocs;
 
   // Enable InitTensorOp elimination.
   if (initTensorElimination) {
diff --git a/mlir/test/Dialect/Linalg/comprehensive-function-bufferize-compat.mlir b/mlir/test/Dialect/Linalg/comprehensive-function-bufferize-compat.mlir
new file mode 100644
index 0000000..6a07a96
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/comprehensive-function-bufferize-compat.mlir
@@ -0,0 +1,31 @@
+// 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
+}
diff --git a/mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp b/mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp
index 506b99e..844b21c 100644
--- a/mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp
+++ b/mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp
@@ -92,6 +92,10 @@ struct TestComprehensiveFunctionBufferize
       *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
 
@@ -105,6 +109,7 @@ void TestComprehensiveFunctionBufferize::runOnFunction() {
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;
   options->analysisFuzzerSeed = analysisFuzzerSeed;
+  options->createDeallocs = createDeallocs;
 
   if (dialectFilter.hasValue()) {
     options->dialectFilter.emplace();
-- 
2.7.4