[InstCombine] Improve TryToSinkInstruction with multiple uses

[anna]

Summary

This patch allows sinking an instruction which can have multiple uses in a
single user. We were previously over-restrictive by looking for exactly one use,
rather than one user.

Also added an API for retrieving a unique undroppable user.

[anna]

removed stray change in file

[lebedev.ri]

Please split the rename into a preparatory change

[anna]

llvm/include/llvm/IR/Value.h

456–457

This comment wasn't true even before the patch. We are traversing the entire use list here to find the "non-droppable use". I can land this separately.

[anna]

In D109700#2997739, @lebedev.ri wrote:

Please split the rename into a preparatory change

I don't follow. This is not a rename. It is an API change where we return a unique *User* instead of a use.

[mkazantsev]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3834

Can we at least allow it if they come from same block? It's typical for switches.

[anna]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3834

yes, will do.

[anna]

addressed review comment. Added test

[nikic]

The change looks reasonable to me.

llvm/lib/IR/Value.cpp

166–167

Might as well iterate over users() and avoid the three U.getUser() calls?

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3821–3822

Dropping dyn_ should be fine here -- user of instruction is always an instruction.

3830

nit: Missing space after =. Also, please add braces for the for, as the inner code has braces.

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

165

Is the cast here needed?

llvm/test/Transforms/InstCombine/sink_instruction.ll

82

Please precommit tests, including generation of all check lines in this file using update_test_checks.

[anna]

addressed review comments.

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

165

not needed.

[lebedev.ri]

In D109700#2997743, @anna wrote:

In D109700#2997739, @lebedev.ri wrote:

Please split the rename into a preparatory change

I don't follow. This is not a rename. It is an API change where we return a unique *User* instead of a use.

Please disregard. My actual question was intended to be more along the lines of:
"do we actually want to completely replace that API like this",
but i'm not really seeing a reason not to.

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3834

The "more than one use" comment isn't quite accurate anymore.

[anna]

In D109700#3002240, @lebedev.ri wrote:

In D109700#2997743, @anna wrote:

In D109700#2997739, @lebedev.ri wrote:

Please split the rename into a preparatory change

I don't follow. This is not a rename. It is an API change where we return a unique *User* instead of a use.

Please disregard. My actual question was intended to be more along the lines of:
"do we actually want to completely replace that API like this",
but i'm not really seeing a reason not to.

Ah. Yes, I went through the same process when deciding to replace the original API :)

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3834

Good point, will update comment before landing.

[hctim]

Early warning signs, looks like it broke the sanitizer fast bot: https://lab.llvm.org/buildbot/#/builders/5/builds/11759

Unclear whether this patch is fully responsible, or the bot's ccache is messed up.

[hctim]

Breakage on regular fastbot was probably occluded as the bot was already red, but it's broken there too: https://lab.llvm.org/buildbot/#/builders/109/builds/22459

[anna]

Thank you for the failure notifications. I have landed the fixes in tree:
https://reviews.llvm.org/rG3273430406c186f1f2af597adeedf21d4fa52b18
https://reviews.llvm.org/rGb6cb03e6b94d91a1757d41d4b0f139166ae15de0

If it still fails, I will revert the change and investigate tomorrow.

So, getSingleUndroppableUser is needed for AssumeBundleBuilder. Added the API back and also continued using it in AssumeBundleBuilder (the second commit above).

Sorry for the noise here. EoD is never a good time to land. Should have waited till ninja check passed.

[lebedev.ri]

So i guess my question was correct then :/
I would suggest to reduce this change to an introduction of the new method (while keeping the old one!), and porting a single user (instcombine pass) to use it.

[anna]

Narrower change compared to previous one: getUniqueUndroppableUser API introduced is only used within InstCombine.

[anna]

The change is narrower now and updated test cases in SLPVectorizer (no surprises in the updates).

[nikic]

Still LGTM, though I'd be interested in knowing what was wrong with the AssumeBundleBuilder case. It's not entirely obvious to me why the switch to getUniqueUndroppableUser() caused issues there.

[anna]

In D109700#3004612, @nikic wrote:

Still LGTM, though I'd be interested in knowing what was wrong with the AssumeBundleBuilder case. It's not entirely obvious to me why the switch to getUniqueUndroppableUser() caused issues there.

I will check that separately and address. The lit testcase may just require an update. However, the unit-test which fails building (because of absence of the use API: getSingleUndroppableUse)
calls AbstractCallSite(const Use *U) constructor. That still requires the getSingleUndroppableUse to be around.

llvm/include/llvm/IR/Value.h

 //===- llvm/Value.h - Definition of the Value class -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
#endif
   /// Note that this is not the same as "has one use". If a value has one use,
   /// then there certainly is a single user. But if value has several uses,
   /// it is possible that all uses are in a single user, or not.
   ///
   /// This check is potentially costly, since it requires traversing,
   /// in the worst case, the whole use list of a value.
   bool hasOneUser() const;
 
-  /// Return true if there is exactly one use of this value that cannot be
-  /// dropped.
-  Use *getSingleUndroppableUse();
-  const Use *getSingleUndroppableUse() const {
-    return const_cast<Value *>(this)->getSingleUndroppableUse();
+  /// Return true if there is exactly one unique user of this value that cannot be
+  /// dropped (that user can have multiple uses of this value).
+  User *getUniqueUndroppableUser();
+  const User *getUniqueUndroppableUser() const {
+    return const_cast<Value *>(this)->getUniqueUndroppableUser();
   }
 
   /// Return true if there this value.
   ///
   /// This is specialized because it is a common request and does not require
   /// traversing the whole use list.
   bool hasNUndroppableUses(unsigned N) const;
 

llvm/lib/IR/Value.cpp

 //===-- Value.cpp - Implement the Value class -----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
bool Value::hasOneUser() const {
   if (use_empty())
     return false;
   if (hasOneUse())
     return true;
   return std::equal(++user_begin(), user_end(), user_begin());
 }
 
 static bool isUnDroppableUser(const User *U) { return !U->isDroppable(); }
 
-Use *Value::getSingleUndroppableUse() {
-  Use *Result = nullptr;
-  for (Use &U : uses()) {
-    if (!U.getUser()->isDroppable()) {
-      if (Result)
+User *Value::getUniqueUndroppableUser() {
+  User *Result = nullptr;
+  for (auto *U : users()) {
+    if (!U->isDroppable()) {
+      if (Result && Result != U)
         return nullptr;
-      Result = &U;
+      Result = U;
     }
   }
   return Result;
 }
 
 bool Value::hasNUndroppableUses(unsigned int N) const {
   return hasNItems(user_begin(), user_end(), N, isUnDroppableUser);
 }

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

 //===- InstructionCombining.cpp - Combine multiple instructions -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
   return nullptr;
 }
 
 /// Try to move the specified instruction from its current block into the
 /// beginning of DestBlock, which can only happen if it's safe to move the
 /// instruction past all of the instructions between it and the end of its
 /// block.
 static bool TryToSinkInstruction(Instruction *I, BasicBlock *DestBlock) {
-  assert(I->getSingleUndroppableUse() && "Invariants didn't hold!");
+  assert(I->getUniqueUndroppableUser() && "Invariants didn't hold!");
   BasicBlock *SrcBlock = I->getParent();
 
   // Cannot move control-flow-involving, volatile loads, vaarg, etc.
   if (isa<PHINode>(I) || I->isEHPad() || I->mayHaveSideEffects() ||
       I->isTerminator())
     return false;
 
   // Do not sink static or dynamic alloca instructions. Static allocas must
while (!Worklist.isEmpty()) {
 
     // See if we can trivially sink this instruction to its user if we can
     // prove that the successor is not executed more frequently than our block.
     // Return the UserBlock if successful.
     auto getOptionalSinkBlockForInst =
         [this](Instruction *I) -> Optional<BasicBlock *> {
       if (!EnableCodeSinking)
         return None;
-      Use *SingleUse = I->getSingleUndroppableUse();
-      if (!SingleUse)
+      auto *UserInst = cast_or_null<Instruction>(I->getUniqueUndroppableUser());
+      if (!UserInst)
         return None;
 
       BasicBlock *BB = I->getParent();
-      Instruction *UserInst = cast<Instruction>(SingleUse->getUser());
-      BasicBlock *UserParent;
+      BasicBlock *UserParent = nullptr;
 
-      // Get the block the use occurs in.
-      if (PHINode *PN = dyn_cast<PHINode>(UserInst))
-        UserParent = PN->getIncomingBlock(*SingleUse);
-      else
+      // Special handling for Phi nodes - get the block the use occurs in.
+      if (PHINode *PN = dyn_cast<PHINode>(UserInst)) {
+        for (unsigned i = 0; i < PN->getNumIncomingValues(); i++) {
+          if (PN->getIncomingValue(i) == I) {
+            // Bail out if we have more than one use in PN. We don't do any
+            // sophisticed analysis (i.e finding NearestCommonDominator of these
+            // use blocks).
+            if (UserParent && UserParent != PN->getIncomingBlock(i))
+              return None;
+            UserParent = PN->getIncomingBlock(i);
+          }
+        }
+        assert(UserParent && "expected to find user block!");
+      } else
         UserParent = UserInst->getParent();
 
       // Try sinking to another block. If that block is unreachable, then do
       // not bother. SimplifyCFG should handle it.
       if (UserParent == BB || !DT.isReachableFromEntry(UserParent))
         return None;
 
       auto *Term = UserParent->getTerminator();

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

 //===- AssumeBundleBuilder.cpp - tools to preserve informations -*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
if (auto *Arg = dyn_cast<Argument>(RK.WasOn)) {
            Arg->getAttribute(RK.AttrKind).getValueAsInt() >= RK.ArgValue))
         return false;
       return true;
     }
     if (auto *Inst = dyn_cast<Instruction>(RK.WasOn))
       if (wouldInstructionBeTriviallyDead(Inst)) {
         if (RK.WasOn->use_empty())
           return false;
-        Use *SingleUse = RK.WasOn->getSingleUndroppableUse();
-        if (SingleUse && SingleUse->getUser() == InstBeingModified)
+        auto *UniqueUser =
+            RK.WasOn->getUniqueUndroppableUser();
+        if (UniqueUser == InstBeingModified)
           return false;
       }
     return true;
   }
 
   void addKnowledge(RetainedKnowledge RK) {
     RK = canonicalizedKnowledge(RK, M->getDataLayout());
 

llvm/test/Transforms/InstCombine/icmp-mul-zext.ll

 }
 
 ; https://bugs.llvm.org/show_bug.cgi?id=33765
 
 @glob = external global i16
 
 define void @PR33765(i8 %beth) {
 ; CHECK-LABEL: @PR33765(
-; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[BETH:%.*]] to i32
 ; CHECK-NEXT:    br i1 false, label [[IF_THEN9:%.*]], label [[IF_THEN9]]
 ; CHECK:       if.then9:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[BETH:%.*]] to i32
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw nsw i32 [[CONV]], [[CONV]]
 ; CHECK-NEXT:    [[TINKY:%.*]] = load i16, i16* @glob, align 2
 ; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[MUL]] to i16
 ; CHECK-NEXT:    [[CONV14:%.*]] = and i16 [[TINKY]], [[TMP1]]
 ; CHECK-NEXT:    store i16 [[CONV14]], i16* @glob, align 2
 ; CHECK-NEXT:    ret void
 ;
   %conv = zext i8 %beth to i32

llvm/test/Transforms/InstCombine/sink_instruction.ll

 
 bb4:                                              ; preds = %bb2
   ret i32 %x_addr.0
 }
 
 declare i32 @bar()
 
 define i32 @test3(i32* nocapture readonly %P, i32 %i) {
 ; CHECK-LABEL: @test3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    switch i32 [[I:%.*]], label [[SW_EPILOG:%.*]] [
 ; CHECK-NEXT:    i32 5, label [[SW_BB:%.*]]
 ; CHECK-NEXT:    i32 2, label [[SW_BB]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw.bb:
 ; CHECK-NEXT:    [[IDXPROM:%.*]] = sext i32 [[I]] to i64
 ; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, i32* [[P:%.*]], i64 [[IDXPROM]]
sw.bb:                                            ; preds = %entry, %entry
   br label %sw.epilog
 
 sw.epilog:                                        ; preds = %entry, %sw.bb
   %sum.0 = phi i32 [ %add, %sw.bb ], [ 0, %entry ]
   ret i32 %sum.0
 }
 
 declare i32 @foo(i32, i32)
-; TODO: Two uses in a single user. We can still sink the instruction (tmp.9).
+; Two uses in a single user. We can still sink the instruction (tmp.9).
 define i32 @test4(i32 %A, i32 %B, i1 %C) {
 ; CHECK-LABEL: @test4(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP_9:%.*]] = add i32 [[B:%.*]], [[A:%.*]]
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[THEN:%.*]], label [[ENDIF:%.*]]
 ; CHECK:       then:
+; CHECK-NEXT:    [[TMP_9:%.*]] = add i32 [[B:%.*]], [[A:%.*]]
 ; CHECK-NEXT:    [[RES:%.*]] = call i32 @foo(i32 [[TMP_9]], i32 [[TMP_9]])
 ; CHECK-NEXT:    ret i32 [[RES]]
 ; CHECK:       endif:
 ; CHECK-NEXT:    [[TMP_2:%.*]] = sdiv i32 [[A]], [[B]]
 ; CHECK-NEXT:    ret i32 [[TMP_2]]
 ;
 entry:
   %tmp.2 = sdiv i32 %A, %B                ; <i32> [#uses=1]

 sw.bb:                                            ; preds = %entry, %entry
   br label %sw.epilog
 
 sw.epilog:                                        ; preds = %entry, %sw.bb
   %sum.0 = phi i32 [ %0, %sw.bb ], [ %add, %dispatchBB ], [ %0, %entry ]
   ret i32 %sum.0
 }
 
-; TODO: Multiple uses but from same BB. We can sink.
+; Multiple uses but from same BB. We can sink.
 define i32 @test6(i32* nocapture readonly %P, i32 %i, i1 %cond) {
 ; CHECK-LABEL: @test6(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[IDXPROM:%.*]] = sext i32 [[I:%.*]] to i64
-; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, i32* [[P:%.*]], i64 [[IDXPROM]]
-; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* [[ARRAYIDX]], align 4
 ; CHECK-NEXT:    [[ADD:%.*]] = shl nsw i32 [[I]], 1
 ; CHECK-NEXT:    br label [[DISPATCHBB:%.*]]
 ; CHECK:       dispatchBB:
+; CHECK-NEXT:    [[IDXPROM:%.*]] = sext i32 [[I:%.*]] to i64
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, i32* [[P:%.*]], i64 [[IDXPROM]]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* [[ARRAYIDX]], align 4
 ; CHECK-NEXT:    switch i32 [[I]], label [[SW_BB:%.*]] [
 ; CHECK-NEXT:    i32 5, label [[SW_EPILOG:%.*]]
 ; CHECK-NEXT:    i32 2, label [[SW_EPILOG]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw.bb:
 ; CHECK-NEXT:    br label [[SW_EPILOG]]
 ; CHECK:       sw.epilog:
 ; CHECK-NEXT:    [[SUM_0:%.*]] = phi i32 [ [[ADD]], [[SW_BB]] ], [ [[TMP0]], [[DISPATCHBB]] ], [ [[TMP0]], [[DISPATCHBB]] ]