[HotColdSplitting] Ensure PHIs have unique incoming values

[vsk]

Summary

Do not add a block to the outlining region if doing so would mean that
a PHI outside of the region would have incompatible incoming values from
the code replacer block.

Test case by Jun Bum Lim!

Fixes llvm.org/PR39564, rdar://45718012.

Depends on https://reviews.llvm.org/D53887.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

410

Can we outline this function to have better readability?

420

When can we have a case where PHI does not have incoming value from immediate predecessor?

[vsk]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

420

When the successor block has no uses of values defined in the predecessor. A special case of this is when the successor block just contains a terminator:

foo:
  %sink = ...
  br label %bar
bar:
  br label %baz

[vsk]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

420

Sorry, that example is incorrect because there isn't a phi there at all. How about:

foo:
  %sink = ...
  br label %bar
bar:
  %phi = phi [ 0, %p1 ], [ 1, %p2 ]
  br i1 undef, label %p1, label %p2
p1:
  br i1 undef, label %bar, %p1
p2:
  br i1 undef, label %bar, %p2

I think foo dominates bar here, but bar has a phi without an incoming value from foo.

[vsk]

• Address review feedback from Aditya.

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

420

No, my example is wrong again. A phi has to have an incoming value from each of its parent's predecessor blocks.

[vsk]

• Do not add duplicate blocks to the outlining region. The test for this is 'forward-dfs-reaches-marked-block.ll'.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

420

:)

[vsk]

Friendly ping.

[vsk]

Actually, when I defined “IsCold” as “has more than one predecessor” to stress-test the pass, I found another instance of this assert firing while building a stage2 clang.

[vsk]

• Use an iterative approach to remove blocks from the outlining region, if they would cause PHIs outside of the region to become invalid (i.e. have more than one distinct incoming value from the codeRepl block).

This successfully builds a stage2 clang compiler, which passes the test suite without regressions. I had to disable stack coloring due to llvm.org/PR39671. To stress the splitting pass, I replaced the definition of "unlikelyExecuted" with "BB->hasNPredecessorsOrMore(2)", and rejected outlining regions which would mark the whole function cold.

At this point, I'm confident in https://reviews.llvm.org/D53887 and this patch.

[vsk]

While this patch is correct, it might make sense to abandon it in favor of https://reviews.llvm.org/D55018, which solves the underlying issue more comprehensively.

[vsk]

I'll abandon this in favor of D55018 and fold some of the tests from this patch into D53887.

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

static unsigned getEntryPointScore(BasicBlock &BB, unsigned Score) {
     return isViableEntryPoint(BB) ? Score : 0;
   }
 
   /// These scores should be lower than the score for predecessor blocks,
   /// because regions starting at predecessor blocks are typically larger.
   static constexpr unsigned ScoreForSuccBlock = 1;
   static constexpr unsigned ScoreForSinkBlock = 1;
 
+  using PHIFrontierMap = SmallDenseMap<const PHINode *, Value *>;
+
   OutliningRegion(const OutliningRegion &) = delete;
   OutliningRegion &operator=(const OutliningRegion &) = delete;
 
 public:
   OutliningRegion() = default;
   OutliningRegion(OutliningRegion &&) = default;
   OutliningRegion &operator=(OutliningRegion &&) = default;
 
   static OutliningRegion create(BasicBlock &SinkBB, const DominatorTree &DT,
                                 const PostDomTree &PDT) {
     OutliningRegion ColdRegion;
 
+    // Map PHIs outside of the outlining region to a unique incoming value
+    // originating from the outlining region. Don't update the PHI frontier for
+    // sink-predecessor blocks, because their successors are post-dominated by
+    // the sink block.
+    PHIFrontierMap PHIFrontier;
+
+    SmallPtrSet<BasicBlock *, 4> RegionBlocks;
+
+    auto addBlockToRegion = [&](BasicBlock *BB, unsigned Score) {
+      RegionBlocks.insert(BB);
+      ColdRegion.Blocks.emplace_back(BB, Score);
+      assert(RegionBlocks.size() == ColdRegion.Blocks.size() && "Duplicate BB");
+    };
+
     // The ancestor farthest-away from SinkBB, and also post-dominated by it.
     unsigned SinkScore = getEntryPointScore(SinkBB, ScoreForSinkBlock);
     ColdRegion.SuggestedEntryPoint = (SinkScore > 0) ? &SinkBB : nullptr;
     unsigned BestScore = SinkScore;
 
     // Visit SinkBB's ancestors using inverse DFS.
     auto PredIt = ++idf_begin(&SinkBB);
     auto PredEnd = idf_end(&SinkBB);
while (PredIt != PredEnd) {
       // The path length is always >= 2, ensuring that predecessor blocks are
       // considered as entry points before the sink block.
       unsigned PredScore = getEntryPointScore(PredBB, PredIt.getPathLength());
       if (PredScore > BestScore) {
         ColdRegion.SuggestedEntryPoint = &PredBB;
         BestScore = PredScore;
       }
 
-      ColdRegion.Blocks.emplace_back(&PredBB, PredScore);
+      addBlockToRegion(&PredBB, PredScore);
       ++PredIt;
     }
 
     // Add SinkBB to the cold region. It's considered as an entry point before
     // any sink-successor blocks.
-    ColdRegion.Blocks.emplace_back(&SinkBB, SinkScore);
+    addBlockToRegion(&SinkBB, SinkScore);
+    updatePHIFrontier(SinkBB, RegionBlocks, PHIFrontier);
 
     // Find all successors of SinkBB dominated by SinkBB using DFS.
     auto SuccIt = ++df_begin(&SinkBB);
     auto SuccEnd = df_end(&SinkBB);
     while (SuccIt != SuccEnd) {
       BasicBlock &SuccBB = **SuccIt;
       bool SinkDom = DT.dominates(&SinkBB, &SuccBB);
 
+      // Don't allow the backwards & forwards DFSes to mark the same block.
+      bool DuplicateBlock = RegionBlocks.count(&SuccBB);
+
       // If SinkBB does not dominate a successor, do not mark the successor (or
       // any of its successors) cold.
-      if (!SinkDom || !mayExtractBlock(SuccBB)) {
+      if (DuplicateBlock || !SinkDom || !mayExtractBlock(SuccBB) ||
+          !updatePHIFrontier(SuccBB, RegionBlocks, PHIFrontier)) {
         SuccIt.skipChildren();
         continue;
       }
 
       unsigned SuccScore = getEntryPointScore(SuccBB, ScoreForSuccBlock);
       if (SuccScore > BestScore) {
         ColdRegion.SuggestedEntryPoint = &SuccBB;
         BestScore = SuccScore;
       }
 
-      ColdRegion.Blocks.emplace_back(&SuccBB, SuccScore);
+      addBlockToRegion(&SuccBB, SuccScore);
       ++SuccIt;
     }
 
     return ColdRegion;
   }
 
   /// Whether this region has nothing to extract.
   bool empty() const { return !SuggestedEntryPoint; }
BlockSequence takeSingleEntrySubRegion(DominatorTree &DT) {
     });
     Blocks.erase(RegionStartIt, RegionEndIt);
 
     // Update the suggested entry point.
     SuggestedEntryPoint = NextEntryPoint;
 
     return SubRegion;
   }
+
+private:
+  /// Update the PHI frontier by recording \p BB's outgoing values. If \p BB
+  /// can be added to the outlining region, return true.
+  static bool
+  updatePHIFrontier(const BasicBlock &BB,
+                    const SmallPtrSetImpl<BasicBlock *> &RegionBlocks,
+                    PHIFrontierMap &PHIFrontier) {
+    // Make sure that each successor of BB in the frontier has a unique
+    // incoming value originating from the outlining region.
+    for (const BasicBlock *SuccBB : successors(&BB)) {
+      // Ignore self-edges and successors which aren't in the frontier.
+      if (SuccBB == &BB || RegionBlocks.count(SuccBB))
+        continue;
+
+      for (const PHINode &SuccPhi : SuccBB->phis()) {
+        // Record the incoming value from BB. Flag it if it isn't unique.
+        Value *&RecordedIncomingVal = PHIFrontier[&SuccPhi];
+        Value *ActualIncomingVal = SuccPhi.getIncomingValueForBlock(&BB);
+        if (RecordedIncomingVal)
+          if (ActualIncomingVal != RecordedIncomingVal)
+            return false;
+
+        RecordedIncomingVal = ActualIncomingVal;
+      }
+    }
+    return true;
+  }
 };
 
 bool HotColdSplitting::outlineColdRegions(Function &F, ProfileSummaryInfo &PSI,
                                           BlockFrequencyInfo *BFI,
                                           TargetTransformInfo &TTI,
                                           DominatorTree &DT, PostDomTree &PDT,
                                           OptimizationRemarkEmitter &ORE) {
   bool Changed = false;

llvm/test/Transforms/HotColdSplit/forward-dfs-reaches-marked-block.ll

+; RUN: opt -hotcoldsplit -S < %s | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.14.0"
+
+; CHECK-LABEL: define {{.*}}@fun
+; CHECK: call {{.*}}@fun.cold.1(
+define void @fun() {
+entry:
+  br i1 undef, label %if.then, label %if.else
+
+if.then:
+  ; This will be marked by the inverse DFS on sink-predecesors.
+  br label %sink
+
+sink:
+  call void @sink()
+  
+  ; Do not allow the forward-DFS on sink-successors to mark the block again.
+  br i1 undef, label %if.then, label %if.then.exit
+
+if.then.exit:
+  ret void
+
+if.else:
+  ret void
+}
+
+declare void @sink() cold

llvm/test/Transforms/HotColdSplit/outline-while-loop.ll

while.end:                                        ; preds = %while.end.loopexit, %while.cond.preheader
   tail call void (...) @sink()
   ret void
 
 if.end:                                           ; preds = %entry
   tail call void @sideeffect(i32 1)
   ret void
 }
 
+; This is the same as @foo, but the while loop comes after the sink block.
+; CHECK-LABEL: define {{.*}}@while_loop_after_sink(
+; CHECK: br i1 {{.*}}, label %if.end, label %codeRepl
+; CHECK-LABEL: codeRepl:
+; CHECK-NEXT: call void @while_loop_after_sink.cold.1
+; CHECK-LABEL: if.end:
+; CHECK: call void @sideeffect(i32 1)
+define void @while_loop_after_sink(i32 %cond) {
+entry:
+  %tobool = icmp eq i32 %cond, 0
+  br i1 %tobool, label %if.end, label %sink
+
+sink:
+  tail call void (...) @sink()
+  br label %while.cond.preheader
+
+while.cond.preheader:
+  %cmp3 = icmp sgt i32 %cond, 10
+  br i1 %cmp3, label %while.body.preheader, label %while.end
+
+while.body.preheader:                             ; preds = %while.cond.preheader
+  br label %while.body
+
+while.body:                                       ; preds = %while.body.preheader, %while.body
+  %cond.addr.04 = phi i32 [ %dec, %while.body ], [ %cond, %while.body.preheader ]
+  %dec = add nsw i32 %cond.addr.04, -1
+  tail call void @sideeffect(i32 0) #3
+  %cmp = icmp sgt i32 %dec, 10
+  br i1 %cmp, label %while.body, label %while.end.loopexit
+
+while.end.loopexit:                               ; preds = %while.body
+  br label %while.end
+
+while.end:                                        ; preds = %while.end.loopexit, %while.cond.preheader
+  ret void
+
+if.end:                                           ; preds = %entry
+  tail call void @sideeffect(i32 1)
+  ret void
+}
+
 ; CHECK-LABEL: define {{.*}}@foo.cold.1
 ; CHECK: phi i32
 ; CHECK-NEXT: add nsw i32
 ; CHECK-NEXT: call {{.*}}@sideeffect
 ; CHECK-NEXT: icmp
 ; CHECK-NEXT: br
 
+; CHECK-LABEL: define {{.*}}@while_loop_after_sink.cold.1
+; CHECK: call {{.*}}@sink
+; CHECK: phi i32
+; CHECK-NEXT: add nsw i32
+; CHECK-NEXT: call {{.*}}@sideeffect
+; CHECK-NEXT: icmp
+; CHECK-NEXT: br
+
 declare void @sideeffect(i32)
 
 declare void @sink(...) cold

llvm/test/Transforms/HotColdSplit/phi-with-distinct-outlined-values.ll

+; RUN: opt -S -hotcoldsplit < %s | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.14.0"
+
+; CHECK-LABEL: define {{.*}}@foo(
+; CHECK: phi i32 [ 0, %entry ], [ 1, %codeRepl ], [ 3, %coldbb2 ]
+
+; CHECK-LABEL: define {{.*}}@foo.cold.1(
+; CHECK: call {{.*}}@sink
+
+define void @foo(i32 %cond) {
+entry:
+  %tobool = icmp eq i32 %cond, 0
+  br i1 %tobool, label %if.end, label %coldbb
+
+coldbb:
+  call void @sink()
+  call void @sideeffect()
+  call void @sideeffect()
+  br i1 undef, label %if.end, label %coldbb2
+
+coldbb2:
+  br label %if.end
+
+if.end:
+  %p = phi i32 [0, %entry], [1, %coldbb], [3, %coldbb2]
+  ret void
+}
+
+declare void @sink() cold
+
+declare void @sideeffect()

llvm/test/Transforms/HotColdSplit/succ-block-with-self-edge.ll

+; RUN: opt -S -hotcoldsplit < %s | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.14.0"
+
+; CHECK-LABEL: define {{.*}}@exit_block_with_same_incoming_vals
+; CHECK: call {{.*}}@exit_block_with_same_incoming_vals.cold.1(
+; CHECK-NOT: br i1 undef
+; CHECK: phi i32 [ 0, %entry ], [ 1, %codeRepl ]
+define void @exit_block_with_same_incoming_vals(i32 %cond) {
+entry:
+  %tobool = icmp eq i32 %cond, 0
+  br i1 %tobool, label %if.end, label %coldbb
+
+coldbb:
+  call void @sink()
+  call void @sideeffect()
+  call void @sideeffect()
+  br i1 undef, label %if.end, label %coldbb2
+
+coldbb2:
+  %p2 = phi i32 [0, %coldbb], [1, %coldbb2]
+  br i1 undef, label %if.end, label %coldbb2
+
+if.end:
+  %p = phi i32 [0, %entry], [1, %coldbb], [1, %coldbb2]
+  ret void
+}
+
+; CHECK-LABEL: define {{.*}}@exit_block_with_distinct_incoming_vals
+; CHECK: call {{.*}}@exit_block_with_distinct_incoming_vals.cold.1(
+; CHECK: br i1 undef
+; CHECK: phi i32 [ 0, %entry ], [ 1, %codeRepl ], [ 2, %coldbb2 ]
+define void @exit_block_with_distinct_incoming_vals(i32 %cond) {
+entry:
+  %tobool = icmp eq i32 %cond, 0
+  br i1 %tobool, label %if.end, label %coldbb
+
+coldbb:
+  call void @sink()
+  call void @sideeffect()
+  call void @sideeffect()
+  br i1 undef, label %if.end, label %coldbb2
+
+coldbb2:
+  %p2 = phi i32 [0, %coldbb], [1, %coldbb2]
+  br i1 undef, label %if.end, label %coldbb2
+
+if.end:
+  %p = phi i32 [0, %entry], [1, %coldbb], [2, %coldbb2]
+  ret void
+}
+
+declare void @sink() cold
+
+declare void @sideeffect()