[mlir][IR] Fix memory leak in DominanceInfo

[mehdi_amini]

Do you know why this wasn't caught until your other change?

[springerm]

In D154185#4464028, @mehdi_amini wrote:

Do you know why this wasn't caught until your other change?

No idea. My change is NFC, or at least I thought so. I was debugging this for quite some time and that still doesn’t make any sense to me….

[springerm]

In D154185#4464079, @springerm wrote:

In D154185#4464028, @mehdi_amini wrote:

Do you know why this wasn't caught until your other change?

No idea. My change is NFC, or at least I thought so. I was debugging this for quite some time and that still doesn’t make any sense to me….

Another thing that’s surprising is that the open source build bot detected the leak for this test case:

Testing: 0.. 10.. 20.. 30.. 40.. 50.. 60.. 70.. 80.. 90
FAIL: MLIR :: Dialect/LLVM/transform-e2e.mlir (70727 of 75254)

• TEST 'MLIR :: Dialect/LLVM/transform-e2e.mlir' FAILED ****

But the Google internal build passed even with asan turned on. But another internal test case was reporting the leak. This is ver mysterious…

[mehdi_amini]

Google internal disable ASAN leak testing for some obscure reasons, it was supposed to be fixed a while ago so I didn't bother adding it to MLIR TAP ASAN at the time (I think there is an env var or an option to get it to work...)

[mehdi_amini]

mlir/lib/IR/Dominance.cpp

34

Since it is deleted here, I suspect we don't have a single upstream test which calls "invalidate"?

[springerm]

mlir/lib/IR/Dominance.cpp

34

Good point I’ll try to at one later. And I’ll try to use a unique_ptr instead of C++ new/delete

[springerm]

In D154185#4464028, @mehdi_amini wrote:

Do you know why this wasn't caught until your other change?

I found the problem: As part of my revision, I added this code to CSE:

/// Invalidate dominance info if the IR was changed.
if (!opsToErase.empty())
  domInfo->invalidate();

In the original code, the dominance info was never invalidated, even though the IR changed (so the invalidate function was never run). That's probably fine in most cases because the only thing that's cached here is a mapping of Region * to DomTree *. However, this could lead to problems in cases such as:

1. CSE erases an op with a region. CSE preserves the dominance info analysis.
2. Another pass is running. A new op with a region is created and the region happens to be allocated at the same memory location (pointer is reused). This pass also preserves the dominance info analysis. (E.g., that is safe to do if only new ops are added.)
3. CSE is run again. The same dominance info is retrieved from the pass, but now the new region is mapped to an incorrect (outdated) DomTree*.

[springerm]

Actually, there is also this comment:

// We currently don't remove region operations, so mark dominance as
// preserved.

So invalidation is actually not needed.

But it may be better to have code in place that invalidates all erased regions (even if there are none today). Only the erased regions, not a blanket "invalidate everything", as I did in my change. Then we can remove this comment and we run a lower risk of breaking something, should we decide to CSE ops with regions in the future.

[mehdi_amini]

Do you see a way to test this?

[springerm]

In D154185#4464848, @mehdi_amini wrote:

Do you see a way to test this?

That's tricky. I could call invalidate from a unit test. But it's unclear what I'm actually testing. Basically invalidate just removes an item from a cache. There's not much to test there. (It would just be a regression test for the leak.)

In D154185#4464846, @springerm wrote:

But it may be better to have code in place that invalidates all erased regions (even if there are none today). Only the erased regions, not a blanket "invalidate everything", as I did in my change. Then we can remove this comment and we run a lower risk of breaking something, should we decide to CSE ops with regions in the future.

This is actually also not a good idea, because for consistency I would also have to invalidate the regions in PostDominanceInfo (because it is also marked as preserved by the pass). But CSE does not even use that. It would be inefficient to retrieve the PostDominanceInfo analysis (because it will probably be computed on the spot), only to invalidate a few regions.

And I’ll try to use a unique_ptr instead of C++ new/delete

DominanceInfo uses llvm::PointerIntPair to store the DomInfo* pointer and hasSSADominance more efficiently:

mutable DenseMap<Region *, llvm::PointerIntPair<DomTree *, 1, bool>> dominanceInfos;

This prevents us from using unique_ptr here. I could change it, but we would loose some efficiency. Maybe this is overly optimized, not sure what's more important here...

mlir/include/mlir/IR/Dominance.h

public:
   DominanceInfoBase &operator=(DominanceInfoBase &&) = default;
   ~DominanceInfoBase();
 
   DominanceInfoBase(const DominanceInfoBase &) = delete;
   DominanceInfoBase &operator=(const DominanceInfoBase &) = delete;
 
   /// Invalidate dominance info. This can be used by clients that make major
   /// changes to the CFG and don't have a good way to update it.
-  void invalidate() { dominanceInfos.clear(); }
-  void invalidate(Region *region) { dominanceInfos.erase(region); }
+  void invalidate();
+  void invalidate(Region *region);
 
   /// Finds the nearest common dominator block for the two given blocks a
   /// and b. If no common dominator can be found, this function will return
   /// nullptr.
   Block *findNearestCommonDominator(Block *a, Block *b) const;
 
   /// Get the root dominance node of the given region. Note that this operation
   /// is only defined for multi-block regions!

mlir/lib/IR/Dominance.cpp

 
 //===----------------------------------------------------------------------===//
 // DominanceInfoBase
 //===----------------------------------------------------------------------===//
 
 template <bool IsPostDom>
 DominanceInfoBase<IsPostDom>::~DominanceInfoBase() {
   for (auto entry : dominanceInfos)
     delete entry.second.getPointer();
 }
 
+template <bool IsPostDom>
+void DominanceInfoBase<IsPostDom>::invalidate() {
+  for (auto entry : dominanceInfos)
+    delete entry.second.getPointer();
+  dominanceInfos.clear();
+}
+
+template <bool IsPostDom>
+void DominanceInfoBase<IsPostDom>::invalidate(Region *region) {
+  auto it = dominanceInfos.find(region);
+  if (it != dominanceInfos.end()) {
+    delete it->second.getPointer();
+    dominanceInfos.erase(it);
+  }
+}
+
 /// Return the dom tree and "hasSSADominance" bit for the given region.  The
 /// DomTree will be null for single-block regions.  This lazily constructs the
 /// DomTree on demand when needsDomTree=true.
 template <bool IsPostDom>
 auto DominanceInfoBase<IsPostDom>::getDominanceInfo(Region *region,
                                                     bool needsDomTree) const
     -> llvm::PointerIntPair<DomTree *, 1, bool> {
   // Check to see if we already have this information.