Fix KnownSafety optimization bugs in ARCSequenceOpts

While doing bottom up dataflow, if we encounter an
unmatched retain instruction, that can pair with a 'KnownSafe'
already visited release instruction, we turn off KnownSafety if the two
RCIdentities mayAlias.
This is done in BottomUpRefCountState::checkAndResetKnownSafety.
In order to determine if a retain is umatched, we look at
IncToDecStateMap. If a retain was matched during bottom up dataflow, it
is always found in IncToDecStateMap with value of the matched release's
BottomUpRefCountState.

Similarly, during top down dataflow, if we encounter an unmatched
release instruction, that can pair with a 'KnownSafe' already
visited retain instruction, we turn off KnownSafety if the two RCIdentities
mayAlias.
This is done in TopDownRefCountState::checkAndResetKnownSafety.
In order to determine if a release is umatched, we look at
DecToIncStateMap. If a release was matched during top down dataflow, it
is always found in DecToIncStateMap with value of the matched retain's
TopDownRefCountState.

For ARCLoopOpts, during bottom up and top down traversal of a region with
a nested loop, we find if the retain/release in the loop summary was
matched or not by looking at the persistent RefCountInstToMatched map.
This map is populated when processing the nested loop region from the
IncToDecStateMap/DecToStateMap which gets thrown away after the loop
region is processed.

This fixes the bugs in both ARCSequenceOpts without loop
support and with loop support.

Author: meg-gupta

lib/SILOptimizer/ARC/ARCRegionState.cpp

@@ -204,6 +204,12 @@ bool ARCRegionState::processBlockBottomUp(
     // that the instruction "visits".
     SILValue Op = Result.RCIdentity;
 
+    std::function<bool(SILInstruction *)> checkIfRefCountInstIsMatched =
+        [&IncToDecStateMap](SILInstruction *Inst) {
+          assert(isa<StrongRetainInst>(Inst) || isa<RetainValueInst>(Inst));
+          return IncToDecStateMap.find(Inst) != IncToDecStateMap.end();
+        };
+
     // For all other (reference counted value, ref count state) we are
     // tracking...
     for (auto &OtherState : getBottomupStates()) {
@@ -217,6 +223,8 @@ bool ARCRegionState::processBlockBottomUp(
         continue;
 
       OtherState->second.updateForSameLoopInst(I, AA);
+      OtherState->second.checkAndResetKnownSafety(
+          I, OtherState->first, checkIfRefCountInstIsMatched, RCIA, AA);
     }
   }
 
@@ -245,8 +253,9 @@ static bool hasEarlyExits(
 
 bool ARCRegionState::processLoopBottomUp(
     const LoopRegion *R, AliasAnalysis *AA, LoopRegionFunctionInfo *LRFI,
+    RCIdentityFunctionInfo *RCIA,
     llvm::DenseMap<const LoopRegion *, ARCRegionState *> &RegionStateInfo,
-    ImmutablePointerSetFactory<SILInstruction> &SetFactory) {
+    llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts) {
   ARCRegionState *State = RegionStateInfo[R];
 
   // If we find that we have non-leaking early exits, clear state
@@ -256,14 +265,23 @@ bool ARCRegionState::processLoopBottomUp(
     return false;
   }
 
+  std::function<bool(SILInstruction *)> checkIfRefCountInstIsMatched =
+      [&UnmatchedRefCountInsts](SILInstruction *Inst) {
+        assert(isa<StrongRetainInst>(Inst) || isa<RetainValueInst>(Inst));
+        return UnmatchedRefCountInsts.find(Inst) == UnmatchedRefCountInsts.end();
+      };
+
   // For each state that we are currently tracking, apply our summarized
   // instructions to it.
   for (auto &OtherState : getBottomupStates()) {
     if (!OtherState.hasValue())
       continue;
 
-    for (auto *I : State->getSummarizedInterestingInsts())
+    for (auto *I : State->getSummarizedInterestingInsts()) {
       OtherState->second.updateForDifferentLoopInst(I, AA);
+      OtherState->second.checkAndResetKnownSafety(
+          I, OtherState->first, checkIfRefCountInstIsMatched, RCIA, AA);
+    }
   }
 
   return false;
@@ -273,6 +291,7 @@ bool ARCRegionState::processBottomUp(
     AliasAnalysis *AA, RCIdentityFunctionInfo *RCIA,
     EpilogueARCFunctionInfo *EAFI, LoopRegionFunctionInfo *LRFI,
     bool FreezeOwnedArgEpilogueReleases,
+    llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts,
     BlotMapVector<SILInstruction *, BottomUpRefCountState> &IncToDecStateMap,
     llvm::DenseMap<const LoopRegion *, ARCRegionState *> &RegionStateInfo,
     ImmutablePointerSetFactory<SILInstruction> &SetFactory) {
@@ -281,7 +300,8 @@ bool ARCRegionState::processBottomUp(
   // We only process basic blocks for now. This ensures that we always propagate
   // the empty set from loops.
   if (!R->isBlock())
-    return processLoopBottomUp(R, AA, LRFI, RegionStateInfo, SetFactory);
+    return processLoopBottomUp(R, AA, LRFI, RCIA, RegionStateInfo,
+                               UnmatchedRefCountInsts);
 
   return processBlockBottomUp(R, AA, RCIA, EAFI, LRFI, FreezeOwnedArgEpilogueReleases,
                               IncToDecStateMap, SetFactory);
@@ -337,6 +357,12 @@ bool ARCRegionState::processBlockTopDown(
     // that the instruction "visits".
     SILValue Op = Result.RCIdentity;
 
+    std::function<bool(SILInstruction *)> checkIfRefCountInstIsMatched =
+        [&DecToIncStateMap](SILInstruction *Inst) {
+          assert(isa<StrongReleaseInst>(Inst) || isa<ReleaseValueInst>(Inst));
+          return DecToIncStateMap.find(Inst) != DecToIncStateMap.end();
+        };
+
     // For all other [(SILValue, TopDownState)] we are tracking...
     for (auto &OtherState : getTopDownStates()) {
       // If the other state's value is blotted, skip it.
@@ -350,6 +376,8 @@ bool ARCRegionState::processBlockTopDown(
         continue;
 
       OtherState->second.updateForSameLoopInst(I, AA);
+      OtherState->second.checkAndResetKnownSafety(
+          I, OtherState->first, checkIfRefCountInstIsMatched, RCIA, AA);
     }
   }
 
@@ -358,8 +386,8 @@ bool ARCRegionState::processBlockTopDown(
 
 bool ARCRegionState::processLoopTopDown(
     const LoopRegion *R, ARCRegionState *State, AliasAnalysis *AA,
-    LoopRegionFunctionInfo *LRFI,
-    ImmutablePointerSetFactory<SILInstruction> &SetFactory) {
+    LoopRegionFunctionInfo *LRFI, RCIdentityFunctionInfo *RCIA,
+    llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts) {
 
   assert(R->isLoop() && "We assume we are processing a loop");
 
@@ -375,14 +403,23 @@ bool ARCRegionState::processLoopTopDown(
   assert(PredRegion->isBlock() && "Expected the predecessor region to be a "
                                   "block");
 
+  std::function<bool(SILInstruction *)> checkIfRefCountInstIsMatched =
+      [&UnmatchedRefCountInsts](SILInstruction *Inst) {
+        assert(isa<StrongReleaseInst>(Inst) || isa<ReleaseValueInst>(Inst));
+        return UnmatchedRefCountInsts.find(Inst) == UnmatchedRefCountInsts.end();
+      };
+
   // For each state that we are currently tracking, apply our summarized
   // instructions to it.
   for (auto &OtherState : getTopDownStates()) {
     if (!OtherState.hasValue())
       continue;
 
-    for (auto *I : State->getSummarizedInterestingInsts())
+    for (auto *I : State->getSummarizedInterestingInsts()) {
       OtherState->second.updateForDifferentLoopInst(I, AA);
+      OtherState->second.checkAndResetKnownSafety(
+          I, OtherState->first, checkIfRefCountInstIsMatched, RCIA, AA);
+    }
   }
 
   return false;
@@ -391,6 +428,7 @@ bool ARCRegionState::processLoopTopDown(
 bool ARCRegionState::processTopDown(
     AliasAnalysis *AA, RCIdentityFunctionInfo *RCIA,
     LoopRegionFunctionInfo *LRFI,
+    llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts,
     BlotMapVector<SILInstruction *, TopDownRefCountState> &DecToIncStateMap,
     llvm::DenseMap<const LoopRegion *, ARCRegionState *> &RegionStateInfo,
     ImmutablePointerSetFactory<SILInstruction> &SetFactory) {
@@ -399,7 +437,8 @@ bool ARCRegionState::processTopDown(
   // We only process basic blocks for now. This ensures that we always propagate
   // the empty set from loops.
   if (!R->isBlock())
-    return processLoopTopDown(R, RegionStateInfo[R], AA, LRFI, SetFactory);
+    return processLoopTopDown(R, RegionStateInfo[R], AA, LRFI, RCIA,
+                              UnmatchedRefCountInsts);
 
   return processBlockTopDown(*R->getBlock(), AA, RCIA, DecToIncStateMap,
                              SetFactory);

lib/SILOptimizer/ARC/ARCRegionState.h

@@ -188,6 +188,7 @@ class ARCRegionState {
   bool processTopDown(
       AliasAnalysis *AA, RCIdentityFunctionInfo *RCIA,
       LoopRegionFunctionInfo *LRFI,
+      llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts,
       BlotMapVector<SILInstruction *, TopDownRefCountState> &DecToIncStateMap,
       llvm::DenseMap<const LoopRegion *, ARCRegionState *> &LoopRegionState,
       ImmutablePointerSetFactory<SILInstruction> &SetFactory);
@@ -200,6 +201,7 @@ class ARCRegionState {
       AliasAnalysis *AA, RCIdentityFunctionInfo *RCIA,
       EpilogueARCFunctionInfo *EAFI, LoopRegionFunctionInfo *LRFI,
       bool FreezeOwnedArgEpilogueReleases,
+      llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts,
       BlotMapVector<SILInstruction *, BottomUpRefCountState> &IncToDecStateMap,
       llvm::DenseMap<const LoopRegion *, ARCRegionState *> &RegionStateInfo,
       ImmutablePointerSetFactory<SILInstruction> &SetFactory);
@@ -227,17 +229,18 @@ class ARCRegionState {
       ImmutablePointerSetFactory<SILInstruction> &SetFactory);
   bool processLoopBottomUp(
       const LoopRegion *R, AliasAnalysis *AA, LoopRegionFunctionInfo *LRFI,
+      RCIdentityFunctionInfo *RCIA,
       llvm::DenseMap<const LoopRegion *, ARCRegionState *> &RegionStateInfo,
-      ImmutablePointerSetFactory<SILInstruction> &SetFactory);
+      llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts);
 
   bool processBlockTopDown(
       SILBasicBlock &BB, AliasAnalysis *AA, RCIdentityFunctionInfo *RCIA,
       BlotMapVector<SILInstruction *, TopDownRefCountState> &DecToIncStateMap,
       ImmutablePointerSetFactory<SILInstruction> &SetFactory);
-  bool
-  processLoopTopDown(const LoopRegion *R, ARCRegionState *State,
-                     AliasAnalysis *AA, LoopRegionFunctionInfo *LRFI,
-                     ImmutablePointerSetFactory<SILInstruction> &SetFactory);
+  bool processLoopTopDown(
+      const LoopRegion *R, ARCRegionState *State, AliasAnalysis *AA,
+      LoopRegionFunctionInfo *LRFI, RCIdentityFunctionInfo *RCIA,
+      llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts);
 
   void summarizeLoop(
       const LoopRegion *R, LoopRegionFunctionInfo *LRFI,

lib/SILOptimizer/ARC/ARCSequenceOpts.cpp

@@ -156,6 +156,7 @@ bool LoopARCPairingContext::processRegion(const LoopRegion *Region,
     }
 
     MadeChange |= MatchedPair;
+    Evaluator.saveMatchingInfo(Region);
     Evaluator.clearLoopState(Region);
     Context.DecToIncStateMap.clear();
     Context.IncToDecStateMap.clear();

lib/SILOptimizer/ARC/GlobalARCSequenceDataflow.cpp

@@ -76,6 +76,12 @@ bool ARCSequenceDataflowEvaluator::processBBTopDown(ARCBBState &BBState) {
     }
   }
 
+  std::function<bool(SILInstruction *)> checkIfRefCountInstIsMatched =
+      [&DecToIncStateMap = DecToIncStateMap](SILInstruction *Inst) {
+        assert(isa<StrongReleaseInst>(Inst) || isa<ReleaseValueInst>(Inst));
+        return DecToIncStateMap.find(Inst) != DecToIncStateMap.end();
+      };
+
   // For each instruction I in BB...
   for (auto &I : BB) {
 
@@ -94,7 +100,7 @@ bool ARCSequenceDataflowEvaluator::processBBTopDown(ARCBBState &BBState) {
 
     // This SILValue may be null if we were unable to find a specific RCIdentity
     // that the instruction "visits".
-    SILValue Op = Result.RCIdentity;
+    SILValue CurrentRC = Result.RCIdentity;
 
     // For all other [(SILValue, TopDownState)] we are tracking...
     for (auto &OtherState : BBState.getTopDownStates()) {
@@ -105,10 +111,12 @@ bool ARCSequenceDataflowEvaluator::processBBTopDown(ARCBBState &BBState) {
       // If we visited an increment or decrement successfully (and thus Op is
       // set), if this is the state for this operand, skip it. We already
       // processed it.
-      if (Op && OtherState->first == Op)
+      if (CurrentRC && OtherState->first == CurrentRC)
         continue;
 
       OtherState->second.updateForSameLoopInst(&I, AA);
+      OtherState->second.checkAndResetKnownSafety(
+          &I, OtherState->first, checkIfRefCountInstIsMatched, RCIA, AA);
     }
   }
 
@@ -221,6 +229,12 @@ bool ARCSequenceDataflowEvaluator::processBBBottomUp(
       RCIA, EAFI, BBState, FreezeOwnedArgEpilogueReleases, IncToDecStateMap,
       SetFactory);
 
+  std::function<bool(SILInstruction *)> checkIfRefCountInstIsMatched =
+      [&IncToDecStateMap = IncToDecStateMap](SILInstruction *Inst) {
+        assert(isa<StrongRetainInst>(Inst) || isa<RetainValueInst>(Inst));
+        return IncToDecStateMap.find(Inst) != IncToDecStateMap.end();
+      };
+
   auto II = BB.rbegin();
   if (!isARCSignificantTerminator(&cast<TermInst>(*II))) {
     II++;
@@ -246,7 +260,7 @@ bool ARCSequenceDataflowEvaluator::processBBBottomUp(
 
     // This SILValue may be null if we were unable to find a specific RCIdentity
     // that the instruction "visits".
-    SILValue Op = Result.RCIdentity;
+    SILValue CurrentRC = Result.RCIdentity;
 
     // For all other (reference counted value, ref count state) we are
     // tracking...
@@ -257,10 +271,12 @@ bool ARCSequenceDataflowEvaluator::processBBBottomUp(
 
       // If this is the state associated with the instruction that we are
       // currently visiting, bail.
-      if (Op && OtherState->first == Op)
+      if (CurrentRC && OtherState->first == CurrentRC)
         continue;
 
       OtherState->second.updateForSameLoopInst(&I, AA);
+      OtherState->second.checkAndResetKnownSafety(
+          &I, OtherState->first, checkIfRefCountInstIsMatched, RCIA, AA);
     }
   }
 

lib/SILOptimizer/ARC/GlobalLoopARCSequenceDataflow.cpp

@@ -110,7 +110,8 @@ bool LoopARCSequenceDataflowEvaluator::processLoopTopDown(const LoopRegion *R) {
 
     // Then perform the dataflow.
     NestingDetected |= SubregionData.processTopDown(
-        AA, RCFI, LRFI, DecToIncStateMap, RegionStateInfo, SetFactory);
+        AA, RCFI, LRFI, UnmatchedRefCountInsts, DecToIncStateMap,
+        RegionStateInfo, SetFactory);
   }
 
   return NestingDetected;
@@ -214,8 +215,9 @@ bool LoopARCSequenceDataflowEvaluator::processLoopBottomUp(
 
     // Then perform the region optimization.
     NestingDetected |= SubregionData.processBottomUp(
-        AA, RCFI, EAFI, LRFI, FreezeOwnedArgEpilogueReleases, IncToDecStateMap,
-        RegionStateInfo, SetFactory);
+        AA, RCFI, EAFI, LRFI, FreezeOwnedArgEpilogueReleases,
+        UnmatchedRefCountInsts, IncToDecStateMap, RegionStateInfo,
+        SetFactory);
   }
 
   return NestingDetected;
@@ -284,8 +286,7 @@ void LoopARCSequenceDataflowEvaluator::dumpDataflowResults() {
   }
 }
 
-void LoopARCSequenceDataflowEvaluator::summarizeLoop(
-    const LoopRegion *R) {
+void LoopARCSequenceDataflowEvaluator::summarizeLoop(const LoopRegion *R) {
   RegionStateInfo[R]->summarize(LRFI, RegionStateInfo);
 }
 
@@ -314,3 +315,34 @@ void LoopARCSequenceDataflowEvaluator::removeInterestingInst(
   auto *Region = LRFI->getRegion(I->getParent());
   RegionStateInfo[Region]->removeInterestingInst(I);
 }
+
+// Compute if a RefCountInst was unmatched and populate in the persistent
+// UnmatchedRefCountInsts.
+// This can be done by looking up the RefCountInst in IncToDecStateMap or
+// DecToIncStateMap. If the StrongIncrement was matched to a StrongDecrement,
+// it will be found in IncToDecStateMap. If the StrongDecrement was matched to
+// a StrongIncrement, it will be found in DecToIncStateMap.
+void LoopARCSequenceDataflowEvaluator::saveMatchingInfo(const LoopRegion *R) {
+  if (R->isFunction())
+    return;
+  for (unsigned SubregionID : R->getSubregions()) {
+    auto *Subregion = LRFI->getRegion(SubregionID);
+    if (!Subregion->isBlock())
+      continue;
+    auto *RegionState = RegionStateInfo[Subregion];
+    for (auto Inst : RegionState->getSummarizedInterestingInsts()) {
+      if (isa<StrongRetainInst>(Inst) || isa<RetainValueInst>(Inst)) {
+        // unmatched if not found in IncToDecStateMap
+        if (IncToDecStateMap.find(Inst) == IncToDecStateMap.end())
+          UnmatchedRefCountInsts.insert(Inst);
+        continue;
+      }
+      if (isa<StrongReleaseInst>(Inst) || isa<ReleaseValueInst>(Inst)) {
+        // unmatched if not found in DecToIncStateMap
+        if (DecToIncStateMap.find(Inst) == DecToIncStateMap.end())
+          UnmatchedRefCountInsts.insert(Inst);
+        continue;
+      }
+    }
+  }
+}

lib/SILOptimizer/ARC/GlobalLoopARCSequenceDataflow.h

@@ -72,6 +72,9 @@ class LoopARCSequenceDataflowEvaluator {
   /// Stashed information for each region.
   llvm::DenseMap<const LoopRegion *, ARCRegionState *> RegionStateInfo;
 
+  /// Set of unmatched RefCountInsts
+  llvm::DenseSet<SILInstruction *> UnmatchedRefCountInsts;
+
 public:
   LoopARCSequenceDataflowEvaluator(
       SILFunction &F, AliasAnalysis *AA, LoopRegionFunctionInfo *LRFI,
@@ -108,6 +111,10 @@ class LoopARCSequenceDataflowEvaluator {
   /// Remove \p I from the interesting instruction list of its parent block.
   void removeInterestingInst(SILInstruction *I);
 
+  /// Compute if a RefCountInst was unmatched and populate the persistent
+  /// UnmatchedRefCountInsts set.
+  void saveMatchingInfo(const LoopRegion *R);
+
   /// Clear the folding node set of the set factory we have stored internally.
   void clearSetFactory() {
     SetFactory.clear();