[SimplifyCFG] performBranchToCommonDestFolding(): form block-closed SSA form before cloning instructions (PR51125)

[lebedev.ri]

Summary

LLVM IR SSA form is "implicit" in @pr51125. While is a valid LLVM IR,
and does not require any PHI nodes, that completely breaks the further logic
in CloneInstructionsIntoPredecessorBlockAndUpdateSSAUses()
that updates the live-out uses of the bonus instructions.

What i believe we need to do, is to first make the SSA form explicit,
by inserting tautological PHI nodes, and rewriting the offending uses.

$ /builddirs/llvm-project/build-Clang12/bin/opt -load /repositories/alive2/build-Clang-release/tv/tv.so -load-pass-plugin /repositories/alive2/build-Clang-release/tv/tv.so -tv -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -bonus-inst-threshold=10 -tv -o /dev/null /tmp/test.ll 

----------------------------------------
@global_pr51125 = global 4 bytes, align 4

define i32 @pr51125() {
%entry:
  br label %L

%L:
  %ld = load i32, * @global_pr51125, align 4
  %iszero = icmp eq i32 %ld, 0
  br i1 %iszero, label %exit, label %L2

%L2:
  store i32 4294967295, * @global_pr51125, align 4
  %cmp = icmp eq i32 %ld, 4294967295
  br i1 %cmp, label %L, label %exit

%exit:
  %r = phi i32 [ %ld, %L2 ], [ %ld, %L ]
  ret i32 %r
}
=>
@global_pr51125 = global 4 bytes, align 4

define i32 @pr51125() {
%entry:
  %ld.old = load i32, * @global_pr51125, align 4
  %iszero.old = icmp eq i32 %ld.old, 0
  br i1 %iszero.old, label %exit, label %L2

%L2:
  %ld2 = phi i32 [ %ld.old, %entry ], [ %ld, %L2 ]
  store i32 4294967295, * @global_pr51125, align 4
  %cmp = icmp ne i32 %ld2, 4294967295
  %ld = load i32, * @global_pr51125, align 4
  %iszero = icmp eq i32 %ld, 0
  %or.cond = select i1 %cmp, i1 1, i1 %iszero
  br i1 %or.cond, label %exit, label %L2

%exit:
  %ld1 = phi i32 [ poison, %L2 ], [ %ld.old, %entry ]
  %r = phi i32 [ %ld2, %L2 ], [ %ld.old, %entry ]
  ret i32 %r
}
Transformation seems to be correct!

Fixes https://bugs.llvm.org/show_bug.cgi?id=51125

[jdoerfert]

It seems odd that we need to "prepare the IR" so late. At first glance, only looking at pr51125, it seems we treat uses in PHIs as uses of the phi block and not the predecessor (edge). Is that the problem?

[lebedev.ri]

In D106317#2888819, @jdoerfert wrote:

It seems odd that we need to "prepare the IR" so late.

Well, i view this part of CloneInstructionsIntoPredecessorBlockAndUpdateSSAUses(), just done before the CFG changes.

At first glance, only looking at pr51125, it seems we treat uses in PHIs as uses of the phi block and not the predecessor (edge). Is that the problem?

"as uses *in* the phi block" you mean? I'm afraid i'm not really seeing it?

[jdoerfert]

In D106317#2888900, @lebedev.ri wrote:

In D106317#2888819, @jdoerfert wrote:

At first glance, only looking at pr51125, it seems we treat uses in PHIs as uses of the phi block and not the predecessor (edge). Is that the problem?

"as uses *in* the phi block" you mean? I'm afraid i'm not really seeing it?

yes, uses *in* the phi block. In pr51125 we have 4 uses of %ld and right now we treat one different than the remaining 3. The reason, as I imagine, is
that we do not consider the use of %ld in the phi %r to be inside of %L` but inside of %exit instead. Conceptually it's on the edge but it's certainly
closer to the end of %L than the beginning of %exit.

I was playing around with this a bit and while the follow code seems to work on that example it is by no means thought through very well:

     SSAUpdate.AddAvailableValue(PredBlock, NewBonusInst);
     for (Use &U : make_early_inc_range(BonusInst.uses())) {
       auto *UI = cast<Instruction>(U.getUser());
+      if (auto *PHI = dyn_cast<PHINode>(UI))
+        if (PHI->getIncomingBlock(U) == PredBlock) {
+          Value *V = SSAUpdate.GetValueInMiddleOfBlock(PredBlock);
+          if (!isa<UndefValue>(V)) {
+            U.set(V);
+            continue;
+          }
+        }
       if (UI->getParent() != PredBlock)
         SSAUpdate.RewriteUseAfterInsertions(U);
       else // Use is in the same block as, and comes before, NewBonusInst.

[lebedev.ri]

Streamline the code, NFCI.

[lebedev.ri]

With an upcoming branch (tomorrow!), does anyone believe that this is an incorrect fix?

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

1101–1115

Do we still need SSAUpdater here? Can this be just if (PN->getIncomingBlock(U) != BB) U.set(NewBonusInst); or so?

3044

This comment would really benefit from a description of what problem this is solving and how it does so. "makes miscompiles easy" doesn't really tell us anything and "sudo-LCSSA" form is somewhat confusing when not talking about loops.

3052

What is an "offending" use? Based on the implementation, it's a use that is potentially not dominated by BonusInst.

I'd personally invert the predicate and call it IsTriviallyDominatedUse or something.

3062

receives

[lebedev.ri]

Thank you for taking a look!
Addressing review notes.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

1101–1115

Hmm, yes, this is better.

3044

Reworded, but not sure how better this is.

[lebedev.ri]

Better naming, NFC.

[nikic]

This is rather weird, but LGTM.

llvm/test/Transforms/SimplifyCFG/fold-branch-to-common-dest.ll

169 ↗

(On Diff #365755)

Unrelated regeneration, precommit?

[lebedev.ri]

In D106317#2945577, @nikic wrote:

This is rather weird, but LGTM.

Thank you for the review!

[nikic]

This had a suspiciously large impact on code size (-4% on mafft with LTO): https://llvm-compile-time-tracker.com/compare.php?from=60dd0121c92e93a33cf39b5f7006923ac9e4f127&to=3d9beefc7d713ad8462d92427ccd17b9532ce904&stat=size-text Not necessarily a problem, but I didn't expect to see any significant impact from this patch.

[lebedev.ri]

Perhaps we simply can't do this here without proper domtree.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

NewBonusInst->dropUndefImplyingAttrsAndUnknownMetadata(
         LLVMContext::MD_annotation);
 
     PredBlock->getInstList().insert(PTI->getIterator(), NewBonusInst);
     NewBonusInst->takeName(&BonusInst);
     BonusInst.setName(NewBonusInst->getName() + ".old");
 
     // Update (liveout) uses of bonus instructions,
     // now that the bonus instruction has been cloned into predecessor.
-    SSAUpdater SSAUpdate;
-    SSAUpdate.Initialize(BonusInst.getType(),
-                         (NewBonusInst->getName() + ".merge").str());
-    SSAUpdate.AddAvailableValue(BB, &BonusInst);
-    SSAUpdate.AddAvailableValue(PredBlock, NewBonusInst);
+    // Note that we expect to be in a block-closed SSA form for this to work!
     for (Use &U : make_early_inc_range(BonusInst.uses())) {
       auto *UI = cast<Instruction>(U.getUser());
-      if (UI->getParent() != PredBlock)
-        SSAUpdate.RewriteUseAfterInsertions(U);
-      else // Use is in the same block as, and comes before, NewBonusInst.
-        SSAUpdate.RewriteUse(U);
+      auto *PN = dyn_cast<PHINode>(UI);
+      if (!PN) {
+        assert(UI->getParent() == BB && BonusInst.comesBefore(UI) &&
+               "If the user is not a PHI node, then it should be in the same "
+               "block as, and come after, the original bonus instruction.");
+        continue; // Keep using the original bonus instruction.
+      }
+      // Is this the block-closed SSA form PHI node?
+      if (PN->getIncomingBlock(U) == BB)
+        continue; // Great, keep using the original bonus instruction.
+      // The only other alternative is an "use" when coming from
+      // the predecessor block - here we should refer to the cloned bonus instr.
+      assert(PN->getIncomingBlock(U) == PredBlock &&
+             "Not in block-closed SSA form?");
+      U.set(NewBonusInst);
     }
   }
 }
 
 bool SimplifyCFGOpt::PerformValueComparisonIntoPredecessorFolding(
     Instruction *TI, Value *&CV, Instruction *PTI, IRBuilder<> &Builder) {
   BasicBlock *BB = TI->getParent();
   BasicBlock *Pred = PTI->getParent();
static bool performBranchToCommonDestFolding(BranchInst *BI, BranchInst *PBI,
   // Determine if the two branches share a common destination.
   Instruction::BinaryOps Opc;
   bool InvertPredCond;
   std::tie(Opc, InvertPredCond) =
       *shouldFoldCondBranchesToCommonDestination(BI, PBI, TTI);
 
   LLVM_DEBUG(dbgs() << "FOLDING BRANCH TO COMMON DEST:\n" << *PBI << *BB);
 
+  // We want to duplicate all the bonus instructions in this block,
+  // and rewrite their uses, but in some cases with self-loops,
+  // the naive use rewrite approach won't work (will result in miscompilations).
+  // To avoid this problem, let's form block-closed SSA form.
+  for (Instruction &BonusInst :
+       reverse(iterator_range<BasicBlock::iterator>(*BB))) {
+    auto IsBCSSAUse = [BB, &BonusInst](Use &U) {
+      auto *UI = cast<Instruction>(U.getUser());
+      if (auto *PN = dyn_cast<PHINode>(UI))
+        return PN->getIncomingBlock(U) == BB;
+      return UI->getParent() == BB && BonusInst.comesBefore(UI);
+    };
+
+    // Does this instruction require rewriting of uses?
+    if (all_of(BonusInst.uses(), IsBCSSAUse))
+      continue;
+
+    SSAUpdater SSAUpdate;
+    Type *Ty = BonusInst.getType();
+    SmallVector<PHINode *, 8> BCSSAPHIs;
+    SSAUpdate.Initialize(Ty, BonusInst.getName());
+
+    // Into each successor block of BB, insert a PHI node, that receives
+    // the BonusInst when coming from it's basic block, or poison otherwise.
+    for (BasicBlock *Succ : successors(BB)) {
+      // The block may have the same successor multiple times. Do it only once.
+      if (SSAUpdate.HasValueForBlock(Succ))
+        continue;
+      BCSSAPHIs.emplace_back(PHINode::Create(
+          Ty, 0, BonusInst.getName() + ".bcssa", &Succ->front()));
+      PHINode *PN = BCSSAPHIs.back();
+      for (BasicBlock *PredOfSucc : predecessors(Succ))
+        PN->addIncoming(PredOfSucc == BB ? (Value *)&BonusInst
+                                         : PoisonValue::get(Ty),
+                        PredOfSucc);
+      SSAUpdate.AddAvailableValue(Succ, PN);
+    }
+
+    // And rewrite all uses that break block-closed SSA form.
+    for (Use &U : make_early_inc_range(BonusInst.uses()))
+      if (!IsBCSSAUse(U))
+        SSAUpdate.RewriteUseAfterInsertions(U);
+
+    // We might not have ended up needing PHI's in all of the succ blocks,
+    // drop the ones that are certainly unused, but don't bother otherwise.
+    for (PHINode *PN : BCSSAPHIs)
+      if (PN->use_empty())
+        PN->eraseFromParent();
+  }
+
   IRBuilder<> Builder(PBI);
   // The builder is used to create instructions to eliminate the branch in BB.
   // If BB's terminator has !annotation metadata, add it to the new
   // instructions.
   Builder.CollectMetadataToCopy(BB->getTerminator(),
                                 {LLVMContext::MD_annotation});
 
   // If we need to invert the condition in the pred block to match, do so now.