[LoopSink] Allow sinking to PHI-use

[wenlei]

Summary

This change allows sinking defs from loop preheader with PHI-use into loop body. Loop sink can now see through PHI-use and select incoming blocks of value being used as candidate sink destination.

It makes loop sink more effective so more LICM can be undone if proven unprofitable with profile info. It addresses the motivating case in D87551, without resorting to profile guided LICM which breaks canonicalization.

[nikic]

Looks reasonable conceptually. Could be extended with critical edge splitting, but not sure how profitable that would be.

llvm/lib/Transforms/Scalar/LoopSink.cpp

182

Might it make sense to move the phi node handling before this check? It should be fine to sink an instruction that is used outside the loop into the corresponding exiting block (i.e. the predecessor of the lcssa phi node), and this may be profitable if the exiting block is cold.

194

You don't need this loop, use PN->getIncomingBlock(U) instead.

llvm/test/Transforms/LICM/loopsink-phi.ll

2

aa-pipeline unnecessary

10

Use update_test_checks.py.

12

Remove !dbg metadata and name instructions in test.

[wenlei]

llvm/lib/Transforms/Scalar/LoopSink.cpp

182

I was thinking about that too. But with some digging I found this was considered in the original loop sink patch (D22778) and then was decided to keep it strictly a loop sink pass, instead of general sinking. Perhaps makes sense to keep it that way.

Original comment/reply:

Isn't it still okay to try to sink in outside the loop if the user block is cold enough?

That would become general purpose sinking instead of loop-sinking. And we need to handle alias/invariant differently.

194

We need to find all incoming blocks of a value, but PN->getIncomingBlock(U) only return one (first) instance? The loop doesn't "break" when a value is found.

E.g. for value %326, we need to find the 4 incoming blocks for all 4 uses in the PHI.

%343         = phi i64 [ %326, %338 ], [ %291, %305 ], [ %326, %332 ], [ %326, %334 ], [ %326, %337 ]

[wenlei]

address comments

[nikic]

llvm/lib/Transforms/Scalar/LoopSink.cpp

194

The outer loop already iterates over all uses. The other phi operands will be visited by it later.

llvm/test/Transforms/LICM/loopsink-phi.ll

49

This !prof metadata probably not needed. In fact, the entire call looks unnecessary, can just be an arg.

143

Drop dead debug metadata.

[wenlei]

llvm/lib/Transforms/Scalar/LoopSink.cpp

194

Ah, I see what you meant. Changed - thanks.

[wenlei]

address comments

[nikic]

LGTM

llvm/test/Transforms/LICM/loopsink-phi.ll

49

You probably don't need this check and can just make .j.check.preheader the entry block?

99

Can drop this one from the module flags.

130

!31 is unused.

[wenlei]

Appreciate the quick review!

llvm/test/Transforms/LICM/loopsink-phi.ll

49

indeed, removed bb and ret block.

[wenlei]

simplify tests

[joanahalili]

Heads-up: we are noticing some clang crashes due to this commit. We are currently working on a reproducer and will post as soon as we get one.

[alexfh]

We're still trying to get a reproducer, but it's somewhat complicated given that the issue happens in ThinLTO builds. The problem manifests as:

Instruction does not dominate all uses!
  %1480 = getelementptr inbounds %"class.absl::AlphaNum", ptr %24, i64 0, i32 1
  %1478 = phi ptr [ %1476, %1474 ], [ %1480, %1472 ]
fatal error: error in backend: Broken module found, compilation aborted!

Please revert while we're preparing a test case!

[alexfh]

In D152772#4437840, @alexfh wrote:

We're still trying to get a reproducer, but it's somewhat complicated given that the issue happens in ThinLTO builds. The problem manifests as:

Instruction does not dominate all uses!
  %1480 = getelementptr inbounds %"class.absl::AlphaNum", ptr %24, i64 0, i32 1
  %1478 = phi ptr [ %1476, %1474 ], [ %1480, %1472 ]
fatal error: error in backend: Broken module found, compilation aborted!

Please revert while we're preparing a test case!

I reverted the commit in c96c85aba28089ec840a867e289782831e9637ee. We'll post a reduced test case as soon as we have one.

[aeubanks]

$ cat /tmp/a.ll
target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
target triple = "aarch64-grtev4-linux-gnu"

%struct.blam = type { %struct.blam.0, [32 x i8] }
%struct.blam.0 = type { ptr, i64 }

define internal void @wibble() !prof !0 {
bb:
  %getelementptr = getelementptr %struct.blam, ptr null, i64 0, i32 1
  br label %bb1

bb1:                                              ; preds = %bb7, %bb3, %bb
  br i1 false, label %bb2, label %bb3

bb2:                                              ; preds = %bb1
  ret void

bb3:                                              ; preds = %bb4, %bb1
  switch i32 0, label %bb5 [
    i32 1, label %bb4
    i32 0, label %bb1
  ], !prof !1

bb4:                                              ; preds = %bb3
  br i1 false, label %bb3, label %bb6

bb5:                                              ; preds = %bb3
  br i1 false, label %bb6, label %bb7

bb6:                                              ; preds = %bb5, %bb4
  br label %bb7

bb7:                                              ; preds = %bb6, %bb5
  %phi = phi ptr [ null, %bb6 ], [ %getelementptr, %bb5 ]
  store ptr %getelementptr, ptr null, align 8
  br label %bb1
}

!0 = !{!"function_entry_count", i64 1}
!1 = !{!"branch_weights", i32 1, i32 188894, i32 287400}
$ opt -p loop-sink /tmp/a.ll
Instruction does not dominate all uses!                                                                                                                                                                                                                                                   
  %getelementptr1 = getelementptr %struct.blam, ptr null, i64 0, i32 1                                                                                                                                                                                                                    
  %phi = phi ptr [ null, %bb6 ], [ %getelementptr1, %bb5 ]                                                                                                                                                                                                                                

[wenlei]

Oops, sorry for the breakage. Thanks for the repro, I will take a look.

[wenlei]

llvm/lib/Transforms/Scalar/LoopSink.cpp

257

When PHI-use is allowed, we also need a tweak here, otherwise the reported ICE could happen.

return cast<Instruction>(U.getUser())->getParent() == N;

-->

return cast<Instruction>(U.getUser())->getParent() == N &&
             !isa<PHINode>(cast<Instruction>(U.getUser()));

This is because the code assumed inserted clone is before all uses in the block, which is not true for PHI-use. The PHI-use is being taken care of by defs in its incoming blocks so we don't need to update PHI-use.

I will reland with this fix included.

llvm/lib/Transforms/Scalar/LoopSink.cpp

 static bool sinkInstruction(
     Loop &L, Instruction &I, const SmallVectorImpl<BasicBlock *> &ColdLoopBBs,
     const SmallDenseMap<BasicBlock *, int, 16> &LoopBlockNumber, LoopInfo &LI,
     DominatorTree &DT, BlockFrequencyInfo &BFI, MemorySSAUpdater *MSSAU) {
   // Compute the set of blocks in loop L which contain a use of I.
   SmallPtrSet<BasicBlock *, 2> BBs;
   for (auto &U : I.uses()) {
     Instruction *UI = cast<Instruction>(U.getUser());
-    // We cannot sink I to PHI-uses.
-    if (isa<PHINode>(UI))
-      return false;
+
     // We cannot sink I if it has uses outside of the loop.
     if (!L.contains(LI.getLoopFor(UI->getParent())))
       return false;
+
+    if (!isa<PHINode>(UI)) {
       BBs.insert(UI->getParent());
+      continue;
+    }
+
+    // We cannot sink I to PHI-uses, try to look through PHI to find the incoming
+    // block of the value being used.
+    PHINode *PN = dyn_cast<PHINode>(UI);
+    BasicBlock *PhiBB = PN->getIncomingBlock(U);
+
+    // If value's incoming block is from loop preheader directly, there's no
+    // place to sink to, bailout.
+    if (L.getLoopPreheader() == PhiBB)
+      return false;
+
+    BBs.insert(PhiBB);
   }
 
   // findBBsToSinkInto is O(BBs.size() * ColdLoopBBs.size()). We cap the max
   // BBs.size() to avoid expensive computation.
   // FIXME: Handle code size growth for min_size and opt_size.
   if (BBs.size() > MaxNumberOfUseBBsForSinking)
     return false;
 
if (MSSAU && MSSAU->getMemorySSA()->getMemoryAccess(&I)) {
           auto *MemUse = cast<MemoryUse>(NewMemAcc);
           MSSAU->insertUse(MemUse, /*RenameUses=*/true);
         }
       }
     }
 
     // Replaces uses of I with IC in N
     I.replaceUsesWithIf(IC, [N](Use &U) {
       return cast<Instruction>(U.getUser())->getParent() == N;
     });
     // Replaces uses of I with IC in blocks dominated by N
     replaceDominatedUsesWith(&I, IC, DT, N);
     LLVM_DEBUG(dbgs() << "Sinking a clone of " << I << " To: " << N->getName()
                       << '\n');
     NumLoopSunkCloned++;
   }
   LLVM_DEBUG(dbgs() << "Sinking " << I << " To: " << MoveBB->getName() << '\n');

llvm/test/Transforms/LICM/loopsink-phi.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt -S -verify-memoryssa -passes=loop-sink < %s | FileCheck %s
+; Make sure that unprofitable loop ICM can be undone by loop sink, and loop sink can handle
+; sinking through PHI use.
+
+
+define dso_local i32 @_Z3fooii(i32 %arg, i32 %arg1, i32 %arg2) local_unnamed_addr #0 !prof !29 {
+; CHECK-LABEL: define dso_local i32 @_Z3fooii
+; CHECK-SAME: (i32 [[ARG:%.*]], i32 [[ARG1:%.*]], i32 [[ARG2:%.*]]) local_unnamed_addr #[[ATTR0:[0-9]+]] !prof [[PROF29:![0-9]+]] {
+; CHECK-NEXT:  .l.check.preheader:
+; CHECK-NEXT:    br label [[DOTL_CHECK:%.*]]
+; CHECK:       .l.ret.loopexit:
+; CHECK-NEXT:    [[DOTLCSSA:%.*]] = phi i32 [ [[I10:%.*]], [[DOTL_ITERATE:%.*]] ]
+; CHECK-NEXT:    ret i32 [[DOTLCSSA]]
+; CHECK:       .l.check:
+; CHECK-NEXT:    [[I4:%.*]] = phi i32 [ 0, [[DOTL_CHECK_PREHEADER:%.*]] ], [ [[I11:%.*]], [[DOTL_ITERATE]] ]
+; CHECK-NEXT:    [[I5:%.*]] = phi i32 [ [[ARG]], [[DOTL_CHECK_PREHEADER]] ], [ [[I10]], [[DOTL_ITERATE]] ]
+; CHECK-NEXT:    [[I6:%.*]] = icmp eq i32 [[I4]], [[ARG1]]
+; CHECK-NEXT:    br i1 [[I6]], label [[DOTL_COLD:%.*]], label [[DOTL_ITERATE]], !prof [[PROF30:![0-9]+]]
+; CHECK:       .l.cold:
+; CHECK-NEXT:    [[FLAG:%.*]] = icmp eq i32 [[ARG1]], 5
+; CHECK-NEXT:    br i1 [[FLAG]], label [[DOTL_COLD1:%.*]], label [[DOTL_COLD2:%.*]]
+; CHECK:       .l.cold1:
+; CHECK-NEXT:    [[TMP1:%.*]] = mul nsw i32 [[ARG2]], [[ARG2]]
+; CHECK-NEXT:    br label [[DOTL_COLD3:%.*]]
+; CHECK:       .l.cold2:
+; CHECK-NEXT:    [[TMP2:%.*]] = add nsw i32 [[ARG2]], [[ARG2]]
+; CHECK-NEXT:    br label [[DOTL_COLD3]]
+; CHECK:       .l.cold3:
+; CHECK-NEXT:    [[I7:%.*]] = phi i32 [ [[TMP1]], [[DOTL_COLD1]] ], [ [[TMP2]], [[DOTL_COLD2]] ]
+; CHECK-NEXT:    [[I8:%.*]] = tail call i32 @_Z3bari(i32 [[I5]])
+; CHECK-NEXT:    [[I9:%.*]] = add nsw i32 [[I8]], [[I7]]
+; CHECK-NEXT:    br label [[DOTL_ITERATE]]
+; CHECK:       .l.iterate:
+; CHECK-NEXT:    [[I10]] = phi i32 [ [[I9]], [[DOTL_COLD3]] ], [ [[I5]], [[DOTL_CHECK]] ]
+; CHECK-NEXT:    [[I11]] = add nuw nsw i32 [[I4]], 1
+; CHECK-NEXT:    [[I12:%.*]] = icmp eq i32 [[I11]], [[I10]]
+; CHECK-NEXT:    br i1 [[I12]], label [[DOTL_RET_LOOPEXIT:%.*]], label [[DOTL_CHECK]]
+;
+
+.l.check.preheader:
+  %flag = icmp eq i32 %arg1, 5
+  %tmp2 = add nsw i32 %arg2, %arg2
+  %tmp1 = mul nsw i32 %arg2, %arg2
+  br label %.l.check
+
+.l.ret.loopexit:                                  ; preds = %.l.iterate
+  %.lcssa = phi i32 [ %i10, %.l.iterate ]
+  ret i32 %.lcssa
+
+.l.check:                                         ; preds = %.l.iterate, %.l.check.preheader
+  %i4 = phi i32 [ 0, %.l.check.preheader ], [ %i11, %.l.iterate ]
+  %i5 = phi i32 [ %arg, %.l.check.preheader ], [ %i10, %.l.iterate ]
+  %i6 = icmp eq i32 %i4, %arg1
+  br i1 %i6, label %.l.cold, label %.l.iterate, !prof !30
+
+.l.cold:                                          ; preds = %.l.check
+  br i1 %flag, label %.l.cold1, label %.l.cold2
+
+.l.cold1:                                         ; preds = %.l.cold
+  br label %.l.cold3
+
+.l.cold2:                                         ; preds = %.l.cold
+  br label %.l.cold3
+
+.l.cold3:                                         ; preds = %.l.cold2, %.l.cold1
+  %i7 = phi i32 [ %tmp1, %.l.cold1 ], [ %tmp2, %.l.cold2 ]
+  %i8 = tail call i32 @_Z3bari(i32 %i5)
+  %i9 = add nsw i32 %i8, %i7
+  br label %.l.iterate
+
+.l.iterate:                                       ; preds = %.l.cold3, %.l.check
+  %i10 = phi i32 [ %i9, %.l.cold3 ], [ %i5, %.l.check ]
+  %i11 = add nuw nsw i32 %i4, 1
+  %i12 = icmp eq i32 %i11, %i10
+  br i1 %i12, label %.l.ret.loopexit, label %.l.check
+}
+
+declare dso_local i32 @_Z3bari(i32) local_unnamed_addr
+
+attributes #0 = { "use-sample-profile" }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9, !10, !11}
+!2 = !{!"ProfileFormat", !"SampleProfile"}
+!3 = !{!"TotalCount", i64 403}
+!4 = !{!"MaxCount", i64 200}
+!5 = !{!"MaxInternalCount", i64 0}
+!6 = !{!"MaxFunctionCount", i64 1}
+!7 = !{!"NumCounts", i64 6}
+!8 = !{!"NumFunctions", i64 1}
+!9 = !{!"IsPartialProfile", i64 0}
+!10 = !{!"PartialProfileRatio", double 0.000000e+00}
+!11 = !{!"DetailedSummary", !12}
+!12 = !{!13, !14, !15, !16, !17, !18, !19, !20, !21, !22, !23, !24, !25, !26, !27, !28}
+!13 = !{i32 10000, i64 200, i32 2}
+!14 = !{i32 100000, i64 200, i32 2}
+!15 = !{i32 200000, i64 200, i32 2}
+!16 = !{i32 300000, i64 200, i32 2}
+!17 = !{i32 400000, i64 200, i32 2}
+!18 = !{i32 500000, i64 200, i32 2}
+!19 = !{i32 600000, i64 200, i32 2}
+!20 = !{i32 700000, i64 200, i32 2}
+!21 = !{i32 800000, i64 200, i32 2}
+!22 = !{i32 900000, i64 200, i32 2}
+!23 = !{i32 950000, i64 200, i32 2}
+!24 = !{i32 990000, i64 200, i32 2}
+!25 = !{i32 999000, i64 1, i32 5}
+!26 = !{i32 999900, i64 1, i32 5}
+!27 = !{i32 999990, i64 1, i32 5}
+!28 = !{i32 999999, i64 1, i32 5}
+!29 = !{!"function_entry_count", i64 2}
+!30 = !{!"branch_weights", i32 1, i32 201}