[Mlir-commits] [mlir] 4affd0c - [mlir] fix a memory leak in NestedPattern

[Alex Zinenko]

Author: Alex Zinenko
Date: 2021-03-12T18:52:14+01:00
New Revision: 4affd0c40eccbce921f185e1e09f84f8194cd4cf

URL: https://github.com/llvm/llvm-project/commit/4affd0c40eccbce921f185e1e09f84f8194cd4cf
DIFF: https://github.com/llvm/llvm-project/commit/4affd0c40eccbce921f185e1e09f84f8194cd4cf.diff

LOG: [mlir] fix a memory leak in NestedPattern

NestedPattern uses a BumpPtrAllocator to store child (nested) pattern
objects to decrease the overhead of dynamic allocation. This assumes all
allocations happen inside the allocator that will be freed as a whole.
However, NestedPattern contains `std::function` as a member, which
allocates internally using `new`, unaware of the BumpPtrAllocator. Since
NestedPattern only holds pointers to the nested patterns allocated in
the BumpPtrAllocator, it never calls their destructors, so the
destructor of the `std::function`s they contain are never called either,
leaking the allocated memory.

Make NestedPattern explicitly call destructors of nested patterns. This
additionally requires to actually copy the nested patterns in
copy-construction and copy-assignment instead of just sharing the
pointer to the arena-allocated list of children to avoid double-free. An
alternative solution would be to add reference counting to the list of
arena-allocated list of children.

Reviewed By: nicolasvasilache

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

Added: 
    

Modified: 
    mlir/include/mlir/Analysis/NestedMatcher.h
    mlir/lib/Analysis/NestedMatcher.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Analysis/NestedMatcher.h b/mlir/include/mlir/Analysis/NestedMatcher.h
index f140a434f857..1587a4ea1313 100644
--- a/mlir/include/mlir/Analysis/NestedMatcher.h
+++ b/mlir/include/mlir/Analysis/NestedMatcher.h
@@ -93,8 +93,15 @@ class NestedPattern {
 public:
   NestedPattern(ArrayRef<NestedPattern> nested,
                 FilterFunctionType filter = defaultFilterFunction);
-  NestedPattern(const NestedPattern &) = default;
-  NestedPattern &operator=(const NestedPattern &) = default;
+  NestedPattern(const NestedPattern &other);
+  NestedPattern &operator=(const NestedPattern &other);
+
+  ~NestedPattern() {
+    // Call destructors manually, ArrayRef is non-owning so it wouldn't call
+    // them, but we should free the memory allocated by std::function outside of
+    // the arena allocator.
+    freeNested();
+  }
 
   /// Returns all the top-level matches in `func`.
   void match(FuncOp func, SmallVectorImpl<NestedMatch> *matches) {
@@ -114,6 +121,13 @@ class NestedPattern {
   friend class NestedMatch;
   friend struct State;
 
+  /// Copies the list of nested patterns to the arena allocator associated with
+  /// this pattern.
+  void copyNestedToThis(ArrayRef<NestedPattern> nested);
+
+  /// Calls destructors on nested patterns.
+  void freeNested();
+
   /// Underlying global bump allocator managed by a NestedPatternContext.
   static llvm::BumpPtrAllocator *&allocator();
 

diff  --git a/mlir/lib/Analysis/NestedMatcher.cpp b/mlir/lib/Analysis/NestedMatcher.cpp
index 7e15ea1094c9..bd5ac0b39ec4 100644
--- a/mlir/lib/Analysis/NestedMatcher.cpp
+++ b/mlir/lib/Analysis/NestedMatcher.cpp
@@ -39,14 +39,37 @@ llvm::BumpPtrAllocator *&NestedPattern::allocator() {
   return allocator;
 }
 
+void NestedPattern::copyNestedToThis(ArrayRef<NestedPattern> nested) {
+  if (nested.empty())
+    return;
+
+  auto *newNested = allocator()->Allocate<NestedPattern>(nested.size());
+  std::uninitialized_copy(nested.begin(), nested.end(), newNested);
+  nestedPatterns = ArrayRef<NestedPattern>(newNested, nested.size());
+}
+
+void NestedPattern::freeNested() {
+  for (const auto &p : nestedPatterns)
+    p.~NestedPattern();
+}
+
 NestedPattern::NestedPattern(ArrayRef<NestedPattern> nested,
                              FilterFunctionType filter)
     : nestedPatterns(), filter(filter), skip(nullptr) {
-  if (!nested.empty()) {
-    auto *newNested = allocator()->Allocate<NestedPattern>(nested.size());
-    std::uninitialized_copy(nested.begin(), nested.end(), newNested);
-    nestedPatterns = ArrayRef<NestedPattern>(newNested, nested.size());
-  }
+  copyNestedToThis(nested);
+}
+
+NestedPattern::NestedPattern(const NestedPattern &other)
+    : nestedPatterns(), filter(other.filter), skip(other.skip) {
+  copyNestedToThis(other.nestedPatterns);
+}
+
+NestedPattern &NestedPattern::operator=(const NestedPattern &other) {
+  freeNested();
+  filter = other.filter;
+  skip = other.skip;
+  copyNestedToThis(other.nestedPatterns);
+  return *this;
 }
 
 unsigned NestedPattern::getDepth() const {