[InstCombine] Improve TryToSinkInstruction with multiple uses

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, the API for retrieving undroppable user has been updated accordingly since
in both usecases (Attributor and InstCombine), we seem to care about the user,
rather than the use.

Reviewed-By: nikic

Differential Revision: https://reviews.llvm.org/D109700

Author: annamthomas

llvm/include/llvm/IR/Value.h

@@ -452,11 +452,11 @@ class Value {
   /// 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.

llvm/lib/IR/Value.cpp

@@ -164,13 +164,13 @@ bool Value::hasOneUser() const {
 
 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;

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -3659,7 +3659,7 @@ Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
 /// 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.
@@ -3818,18 +3818,27 @@ bool InstCombinerImpl::run() {
         [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;
-
-      // Get the block the use occurs in.
-      if (PHINode *PN = dyn_cast<PHINode>(UserInst))
-        UserParent = PN->getIncomingBlock(*SingleUse);
-      else
+      BasicBlock *UserParent = nullptr;
+
+      // 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 uses in different blocks. We don't do any
+            // sophisticated 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

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

@@ -161,8 +161,9 @@ struct AssumeBuilderState {
       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;

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

@@ -56,9 +56,9 @@ lor.end:
 
 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

llvm/test/Transforms/InstCombine/sink_instruction.ll

@@ -114,13 +114,13 @@ sw.epilog:                                        ; preds = %entry, %sw.bb
 }
 
 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:
@@ -175,16 +175,16 @@ sw.epilog:                                        ; preds = %entry, %sw.bb
   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]]