[Mlir-commits] [mlir] 27a1a85 - buffer-deallocation: consider aliases introduced by arith.select.

[Johannes Reifferscheid]

Author: Johannes Reifferscheid
Date: 2022-08-23T14:37:02+02:00
New Revision: 27a1a85662e68e96eeba0732f92ac61b2bbde9c3

URL: https://github.com/llvm/llvm-project/commit/27a1a85662e68e96eeba0732f92ac61b2bbde9c3
DIFF: https://github.com/llvm/llvm-project/commit/27a1a85662e68e96eeba0732f92ac61b2bbde9c3.diff

LOG: buffer-deallocation: consider aliases introduced by arith.select.

Currently, buffer deallocation considers arith.select to be
non-aliasing, which results in deallocs being inserted incorrectly. Since
arith.select doesn't implement any useful interfaces, this change just handles
it explicitly. Eventually this should probably be fixed properly, if this pass
is going to be used long term.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D132460

Added: 
    

Modified: 
    mlir/lib/Analysis/BufferViewFlowAnalysis.cpp
    mlir/lib/Analysis/CMakeLists.txt
    mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
    utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Analysis/BufferViewFlowAnalysis.cpp b/mlir/lib/Analysis/BufferViewFlowAnalysis.cpp
index 5b2b31db29498..80f538dfb3bcb 100644
--- a/mlir/lib/Analysis/BufferViewFlowAnalysis.cpp
+++ b/mlir/lib/Analysis/BufferViewFlowAnalysis.cpp
@@ -8,6 +8,7 @@
 
 #include "mlir/Analysis/BufferViewFlowAnalysis.h"
 
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/SetOperations.h"
@@ -51,9 +52,9 @@ void BufferViewFlowAnalysis::remove(const SmallPtrSetImpl<Value> &aliasValues) {
 /// successor regions and branch-like return operations from nested regions.
 void BufferViewFlowAnalysis::build(Operation *op) {
   // Registers all dependencies of the given values.
-  auto registerDependencies = [&](auto values, auto dependencies) {
-    for (auto entry : llvm::zip(values, dependencies))
-      this->dependencies[std::get<0>(entry)].insert(std::get<1>(entry));
+  auto registerDependencies = [&](ValueRange values, ValueRange dependencies) {
+    for (auto [value, dep] : llvm::zip(values, dependencies))
+      this->dependencies[value].insert(dep);
   };
 
   // Add additional dependencies created by view changes to the alias list.
@@ -119,4 +120,10 @@ void BufferViewFlowAnalysis::build(Operation *op) {
       }
     }
   });
+
+  // TODO: This should be an interface.
+  op->walk([&](arith::SelectOp selectOp) {
+    registerDependencies({selectOp.getOperand(1)}, {selectOp.getResult()});
+    registerDependencies({selectOp.getOperand(2)}, {selectOp.getResult()});
+  });
 }

diff  --git a/mlir/lib/Analysis/CMakeLists.txt b/mlir/lib/Analysis/CMakeLists.txt
index 4ead4c7127905..701584c53f362 100644
--- a/mlir/lib/Analysis/CMakeLists.txt
+++ b/mlir/lib/Analysis/CMakeLists.txt
@@ -40,6 +40,7 @@ add_mlir_library(MLIRAnalysis
   mlir-headers
 
   LINK_LIBS PUBLIC
+  MLIRArithmeticDialect
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
index 240cc2a60681f..61493d964093c 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
@@ -1298,3 +1298,19 @@ func.func @while_three_arg(%arg0: index) {
 // CHECK-NEXT: return
   return
 }
+
+// -----
+
+func.func @select_aliases(%arg0: index, %arg1: memref<?xi8>, %arg2: i1) {
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: arith.select
+  // CHECK: test.copy
+  // CHECK: memref.dealloc
+  // CHECK: memref.dealloc
+  %0 = memref.alloc(%arg0) : memref<?xi8>
+  %1 = memref.alloc(%arg0) : memref<?xi8>
+  %2 = arith.select %arg2, %0, %1 : memref<?xi8>
+  test.copy(%2, %arg1) : (memref<?xi8>, memref<?xi8>)
+  return
+}

diff  --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 7d45b2ef5cc4c..cefa2a93afc1f 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -5879,6 +5879,7 @@ cc_library(
     ),
     includes = ["include"],
     deps = [
+        ":ArithmeticDialect",
         ":CallOpInterfaces",
         ":ControlFlowInterfaces",
         ":DataLayoutInterfaces",