[ObjC][ARC] Prevent moving objc_retain calls past objc_release calls that release the retained object

[ahatanak]

Summary

This patch fixes what looks like a longstanding bug in ARC optimizer where it reverses the order of objc_retain calls and objc_release calls that retain and release the same object.

The code in ARC optimizer that is responsible for code motion takes the following steps:

1. Traverse the CFG bottom-up and determine how far up objc_release calls can be moved. Determine the insertion points for the objc_release calls, but don't actually move them.
2. Traverse the CFG top-down and determine how far down objc_retain calls can be moved. Determine the insertion points for the objc_retain calls, but don't actually move them.
3. Try to move the objc_retain and objc_release calls if they can't be removed.

The problem is that the insertion points for the objc_retain calls are determined in step 2 without taking into consideration the insertion points for objc_release calls determined in step 1, so the order of an objc_retain call and an objc_release call can be reversed, which is incorrect, even though each step is correct in isolation.

To fix this bug, this patch teaches the top-down traversal step to take into consideration the insertion points for objc_release calls determined in the bottom-up traversal step. Code motion for an objc_retain call is disabled if there is a possibility that it can be moved past an objc_release call that releases the retained object.

rdar://79292791

[ahatanak]

Simplify test case.

[ahatanak]

llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp

1513

This isn't quite correct as it's not impossible for two different retains with different RC identity roots to have the same insertion point.

[ahatanak]

Handle the case where two different retains with different RC identity roots have the same insertion point and add a test case for that.

[fhahn]

If we move a release past a retain, does this mean the retain/release pair is redundant? If that's the case, would it be easier to detect this case, during or after moving and then remove the redundant pair?

llvm/test/Transforms/ObjCARC/code-motion.ll

43

For this test, would be it be correct to move the retain/release calls to %if.else?

if.else:
  call void @alterRefCount()
  call i8* @llvm.objc.retain(i8* %[[OBJ]])
  call void @use(i8* %obj)
  call void @llvm.objc.release(i8* %[[OBJ]])
  br label %join

51

Given that the function is quite small and the key here is to check precisely *if/where* the calls are moved, it would probably be good to create check lines for the whole function?

You could also pre-commit the tests, so the impact of the change is more obvious in the diff itself.

[ahatanak]

Rebase after precommitting tests.

[ahatanak]

The pair should be redundant if the release can be moved to the retain's insertion point. But I don't think removing the pair would be easier than just disabling code motion as doing so would require extra bookkeeping to keep track of the retain/release pairs that can be removed together. I think it's harder to get it right when the control flow isn't as simple as the tests in code-motion.ll.

llvm/test/Transforms/ObjCARC/code-motion.ll

43

It would be correct to move the pair to %if.else, but the retain cannot be moved past the call to @alterRefCount since that could cause obj to be deallocated before the call to @use.

[ahatanak]

Fix test cases.

[fhahn]

LGTM, thanks! It might be good to also add a test case where the insert points are not in the same basic block, ie there’s some control flow between them.

I ’m not super familiar with the code, but @ahatanak was kind enough to walk me through the details offline. The approach looks like a conservative approach to address this miscompile and less risky than trying to undo the bad movements.

I still think undoing the bad movements might be an interesting direction in the future, because it would improve codegen for the examples in the tests.

llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp

1501–1536

Might be helpful to document how it is filled?

1517

Identity root*s*?

[ahatanak]

Rebase after b931c2a714b93327c602ee54e867f812bded4262. Address review comments.

[ahatanak]

I agree that undoing/preventing the bad move or eliminating the redundant retain/release pair is something we should look into.

llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp

void CheckForCFGHazards(const BasicBlock *BB,
                           DenseMap<const BasicBlock *, BBState> &BBStates,
                           BBState &MyStates) const;
   bool VisitInstructionBottomUp(Instruction *Inst, BasicBlock *BB,
                                 BlotMapVector<Value *, RRInfo> &Retains,
                                 BBState &MyStates);
   bool VisitBottomUp(BasicBlock *BB,
                      DenseMap<const BasicBlock *, BBState> &BBStates,
                      BlotMapVector<Value *, RRInfo> &Retains);
-  bool VisitInstructionTopDown(Instruction *Inst,
-                               DenseMap<Value *, RRInfo> &Releases,
-                               BBState &MyStates);
-  bool VisitTopDown(BasicBlock *BB,
-                    DenseMap<const BasicBlock *, BBState> &BBStates,
-                    DenseMap<Value *, RRInfo> &Releases);
+  bool VisitInstructionTopDown(
+      Instruction *Inst, DenseMap<Value *, RRInfo> &Releases, BBState &MyStates,
+      const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+          &ReleaseInsertPtToRCIdentityRoots);
+  bool VisitTopDown(
+      BasicBlock *BB, DenseMap<const BasicBlock *, BBState> &BBStates,
+      const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+          &ReleaseInsertPtToRCIdentityRoots);
   bool Visit(Function &F, DenseMap<const BasicBlock *, BBState> &BBStates,
              BlotMapVector<Value *, RRInfo> &Retains,
              DenseMap<Value *, RRInfo> &Releases);
 
   void MoveCalls(Value *Arg, RRInfo &RetainsToMove, RRInfo &ReleasesToMove,
                  BlotMapVector<Value *, RRInfo> &Retains,
                  DenseMap<Value *, RRInfo> &Releases,
                  SmallVectorImpl<Instruction *> &DeadInsts, Module *M);
if (InvokeInst *II = dyn_cast<InvokeInst>(&Pred->back()))
       NestingDetected |= VisitInstructionBottomUp(II, BB, Retains, MyStates);
   }
 
   LLVM_DEBUG(dbgs() << "\nFinal State:\n" << BBStates[BB] << "\n");
 
   return NestingDetected;
 }
 
-bool
-ObjCARCOpt::VisitInstructionTopDown(Instruction *Inst,
-                                    DenseMap<Value *, RRInfo> &Releases,
-                                    BBState &MyStates) {
+// Fill ReleaseInsertPtToRCIdentityRoots, which is a map from insertion points
+// to the set of RC identity roots that would be released by the release calls
+// moved to the insertion points.
+static void collectReleaseInsertPts(
+    const BlotMapVector<Value *, RRInfo> &Retains,
+    DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+        &ReleaseInsertPtToRCIdentityRoots) {
+  for (auto &P : Retains) {
+    // Retains is a map from an objc_retain call to a RRInfo of the RC identity
+    // root of the call. Get the RC identity root of the objc_retain call.
+    Instruction *Retain = cast<Instruction>(P.first);
+    Value *Root = GetRCIdentityRoot(Retain->getOperand(0));
+    // Collect all the insertion points of the objc_release calls that release
+    // the RC identity root of the objc_retain call.
+    for (const Instruction *InsertPt : P.second.ReverseInsertPts)
+      ReleaseInsertPtToRCIdentityRoots[InsertPt].insert(Root);
+  }
+}
+
+// Get the RC identity roots from an insertion point of an objc_release call.
+// Return nullptr if the passed instruction isn't an insertion point.
+static const SmallPtrSet<const Value *, 2> *
+getRCIdentityRootsFromReleaseInsertPt(
+    const Instruction *InsertPt,
+    const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+        &ReleaseInsertPtToRCIdentityRoots) {
+  auto I = ReleaseInsertPtToRCIdentityRoots.find(InsertPt);
+  if (I == ReleaseInsertPtToRCIdentityRoots.end())
+    return nullptr;
+  return &I->second;
+}
+
+bool ObjCARCOpt::VisitInstructionTopDown(
+    Instruction *Inst, DenseMap<Value *, RRInfo> &Releases, BBState &MyStates,
+    const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+        &ReleaseInsertPtToRCIdentityRoots) {
   bool NestingDetected = false;
   ARCInstKind Class = GetARCInstKind(Inst);
   const Value *Arg = nullptr;
 
+  // Make sure a call to objc_retain isn't moved past insertion points of calls
+  // to objc_release.
+  if (const SmallPtrSet<const Value *, 2> *Roots =
+          getRCIdentityRootsFromReleaseInsertPt(
+              Inst, ReleaseInsertPtToRCIdentityRoots))
+    for (auto *Root : *Roots) {
+      TopDownPtrState &S = MyStates.getPtrTopDownState(Root);
+      // Disable code motion if the current position is S_Retain to prevent
+      // moving the objc_retain call past objc_release calls. If it's
+      // S_CanRelease or larger, it's not necessary to disable code motion as
+      // the insertion points that prevent the objc_retain call from moving down
+      // should have been set already.
+      if (S.GetSeq() == S_Retain)
+        S.SetCFGHazardAfflicted(true);
+    }
+
   LLVM_DEBUG(dbgs() << "        Class: " << Class << "\n");
 
   switch (Class) {
   case ARCInstKind::RetainBlock:
     // In OptimizeIndividualCalls, we have strength reduced all optimizable
     // objc_retainBlocks to objc_retains. Thus at this point any
     // objc_retainBlocks that we see are not optimizable. We need to break since
     // a retain can be a potential use.
if (S.HandlePotentialAlterRefCount(Inst, Ptr, PA, Class, *BundledInsts))
       continue;
 
     S.HandlePotentialUse(Inst, Ptr, PA, Class);
   }
 
   return NestingDetected;
 }
 
-bool
-ObjCARCOpt::VisitTopDown(BasicBlock *BB,
-                         DenseMap<const BasicBlock *, BBState> &BBStates,
-                         DenseMap<Value *, RRInfo> &Releases) {
+bool ObjCARCOpt::VisitTopDown(
+    BasicBlock *BB, DenseMap<const BasicBlock *, BBState> &BBStates,
+    DenseMap<Value *, RRInfo> &Releases,
+    const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+        &ReleaseInsertPtToRCIdentityRoots) {
   LLVM_DEBUG(dbgs() << "\n== ObjCARCOpt::VisitTopDown ==\n");
   bool NestingDetected = false;
   BBState &MyStates = BBStates[BB];
 
   // Merge the states from each predecessor to compute the initial state
   // for the current block.
   BBState::edge_iterator PI(MyStates.pred_begin()),
                          PE(MyStates.pred_end());
bool ObjCARCOpt::VisitTopDown(
   LLVM_DEBUG(dbgs() << "Before:\n"
                     << BBStates[BB] << "\n"
                     << "Performing Dataflow:\n");
 
   // Visit all the instructions, top-down.
   for (Instruction &Inst : *BB) {
     LLVM_DEBUG(dbgs() << "    Visiting " << Inst << "\n");
 
-    NestingDetected |= VisitInstructionTopDown(&Inst, Releases, MyStates);
+    NestingDetected |= VisitInstructionTopDown(
+        &Inst, Releases, MyStates, ReleaseInsertPtToRCIdentityRoots);
 
     // Bail out if the number of pointers being tracked becomes too large so
     // that this pass can complete in a reasonable amount of time.
     if (MyStates.top_down_ptr_list_size() > MaxPtrStates) {
       DisableRetainReleasePairing = true;
       return false;
     }
   }
bool ObjCARCOpt::Visit(Function &F,
   // Use reverse-postorder on the reverse CFG for bottom-up.
   bool BottomUpNestingDetected = false;
   for (BasicBlock *BB : llvm::reverse(ReverseCFGPostOrder)) {
     BottomUpNestingDetected |= VisitBottomUp(BB, BBStates, Retains);
     if (DisableRetainReleasePairing)
       return false;
   }
 
+  DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+      ReleaseInsertPtToRCIdentityRoots;
+  collectReleaseInsertPts(Retains, ReleaseInsertPtToRCIdentityRoots);
+
   // Use reverse-postorder for top-down.
   bool TopDownNestingDetected = false;
   for (BasicBlock *BB : llvm::reverse(PostOrder)) {
-    TopDownNestingDetected |= VisitTopDown(BB, BBStates, Releases);
+    TopDownNestingDetected |=
+        VisitTopDown(BB, BBStates, Releases, ReleaseInsertPtToRCIdentityRoots);
     if (DisableRetainReleasePairing)
       return false;
   }
 
   return TopDownNestingDetected && BottomUpNestingDetected;
 }
 
 /// Move the calls in RetainsToMove and ReleasesToMove.

llvm/test/Transforms/ObjCARC/code-motion.ll

define void @test2() {
   %v0 = load i8*, i8** @g0, align 8
   %v1 = tail call i8* @llvm.objc.retain(i8* %v0)
   tail call void @use(i8* %v0)
   tail call void @alterRefCount()
   tail call void @llvm.objc.release(i8* %v0)
   ret void
 }
 
-; ARC optimizer shouldn't reverse the order of retains and releases in
-; if.then in @test3 and @test4.
+; Check that code motion is disabled in @test3 and @test4.
+; Previously, ARC optimizer would move the release past the retain.
+
+; if.then:
+;   call void @readOnlyFunc(i8* %obj, i8* null)
+;   call void @llvm.objc.release(i8* %obj) #1, !clang.imprecise_release !2
+;   %1 = add i32 1, 2
+;   %2 = tail call i8* @llvm.objc.retain(i8* %obj)
+;
+; Ideally, the retain/release pairs in BB if.then should be removed.
 
 define void @test3(i8* %obj, i1 %cond) {
 ; CHECK-LABEL: @test3(
+; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ:%.*]])
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ:%.*]], i8* null)
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]]) {{.*}}, !clang.imprecise_release !2
+; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ]], i8* null)
 ; CHECK-NEXT:    [[TMP1:%.*]] = add i32 1, 2
-; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    br label [[JOIN:%.*]]
 ; CHECK:       if.else:
-; CHECK-NEXT:    [[TMP3:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    call void @use(i8* [[OBJ]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
+; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    ret void
 ;
   %v0 = call i8* @llvm.objc.retain(i8* %obj)
   br i1 %cond, label %if.then, label %if.else
 
 if.then:
   call void @readOnlyFunc(i8* %obj, i8* null) #0
   add i32 1, 2
   call void @alterRefCount()
   br label %join
 
 if.else:
   call void @alterRefCount()
   call void @use(i8* %obj)
   br label %join
 
 join:
   call void @llvm.objc.release(i8* %obj), !clang.imprecise_release !9
   ret void
 }
 
 define void @test4(i8* %obj0, i8* %obj1, i1 %cond) {
 ; CHECK-LABEL: @test4(
+; CHECK-NEXT:    [[TMP3:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ0:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ1:%.*]])
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ0:%.*]], i8* [[OBJ1:%.*]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ1]]) {{.*}}, !clang.imprecise_release !2
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ0]]) {{.*}}, !clang.imprecise_release !2
+; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ0]], i8* [[OBJ1]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = add i32 1, 2
-; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ1]])
-; CHECK-NEXT:    [[TMP3:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ0]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    br label [[JOIN:%.*]]
 ; CHECK:       if.else:
-; CHECK-NEXT:    [[TMP4:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ1]])
-; CHECK-NEXT:    [[TMP5:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ0]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    call void @use(i8* [[OBJ0]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ0]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    call void @use(i8* [[OBJ1]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ1]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
+; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ0]]) {{.*}}, !clang.imprecise_release !2
+; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ1]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    ret void
 ;
   %v0 = call i8* @llvm.objc.retain(i8* %obj0)
   %v1 = call i8* @llvm.objc.retain(i8* %obj1)
   br i1 %cond, label %if.then, label %if.else
 
 if.then:
   call void @readOnlyFunc(i8* %obj0, i8* %obj1) #0
join:
   ret void
 }
 
 ; In this test, insertion points for the retain and release calls that could be
 ; eliminated are in different blocks (bb1 and if.then).
 
 define void @test5(i8* %obj, i1 %cond0, i1 %cond1) {
 ; CHECK-LABEL: @test5(
+; CHECK-NEXT:    [[V0:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ:%.*]])
 ; CHECK-NEXT:    br i1 [[COND0:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ:%.*]], i8* null)
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]])
+; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ]], i8* null)
 ; CHECK-NEXT:    br i1 [[COND1:%.*]], label [[IF_THEN2:%.*]], label [[IF_ELSE2:%.*]]
 ; CHECK:       if.then2:
 ; CHECK-NEXT:    br label [[BB1:%.*]]
 ; CHECK:       if.else2:
 ; CHECK-NEXT:    br label [[BB1]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    [[TMP1:%.*]] = add i32 1, 2
-; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    br label [[JOIN:%.*]]
 ; CHECK:       if.else:
-; CHECK-NEXT:    [[TMP3:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    call void @use(i8* [[OBJ]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]])
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
+; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]])
 ; CHECK-NEXT:    ret void
 ;
   %v0 = call i8* @llvm.objc.retain(i8* %obj)
   br i1 %cond0, label %if.then, label %if.else
 
 if.then:
   call void @readOnlyFunc(i8* %obj, i8* null) #0
   br i1 %cond1, label %if.then2, label %if.else2