[mlir] fix a memory leak in NestedPattern

[ftynse]

Summary

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::functions 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.

[dcaballe]

Thanks, Alex!

[mehdi_amini]

Thanks!

Do we relies on the owning behavior of std::function for users of NestedPattern or could we switch it to function_ref?

[ftynse]

In D98485#2623491, @mehdi_amini wrote:

Thanks!

Do we relies on the owning behavior of std::function for users of NestedPattern or could we switch it to function_ref?

It does. function_ref was the first thing I tried.

[rriddle]

mlir/include/mlir/Analysis/NestedMatcher.h

132

Is this always allocating the same type? You could switch to SpecificBumpPtrAllocator, which does call the destructors (IIRC).

mlir/include/mlir/Analysis/NestedMatcher.h

 /// a plain walk over operations to match flat patterns but the current
 /// implementation is competitive nonetheless.
 using FilterFunctionType = std::function<bool(Operation &)>;
 inline bool defaultFilterFunction(Operation &) { return true; }
 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) {
     func.walk([&](Operation *op) { matchOne(op, matches); });
   }
 
   /// Returns all the top-level matches in `op`.
   void match(Operation *op, SmallVectorImpl<NestedMatch> *matches) {
     op->walk([&](Operation *child) { matchOne(child, matches); });
   }
 
   /// Returns the depth of the pattern.
   unsigned getDepth() const;
 
 private:
   friend class NestedPatternContext;
   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();
 
   /// Matches this pattern against a single `op` and fills matches with the
   /// result.
   void matchOne(Operation *op, SmallVectorImpl<NestedMatch> *matches);
 
   /// Nested patterns to be matched.
   ArrayRef<NestedPattern> nestedPatterns;
 

mlir/lib/Analysis/NestedMatcher.cpp

NestedMatch NestedMatch::build(Operation *operation,
   return *result;
 }
 
 llvm::BumpPtrAllocator *&NestedPattern::allocator() {
   thread_local llvm::BumpPtrAllocator *allocator = nullptr;
   return allocator;
 }
 
-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());
-  }
+void NestedPattern::copyNestedToThis(ArrayRef<NestedPattern> nested) {
+  if (nested.empty())
+    return;
+
+
+void NestedPattern::freeNested() {
+  for (const auto &p : nestedPatterns)
+    p.~NestedPattern();
+}
+
+NestedPattern::NestedPattern(ArrayRef<NestedPattern> nested,
+                             FilterFunctionType filter)
+    : nestedPatterns(), filter(filter), skip(nullptr) {
+  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 {
   if (nestedPatterns.empty()) {
     return 1;
   }
   unsigned depth = 0;
   for (auto &c : nestedPatterns) {