Index: llvm/lib/CodeGen/TypePromotion.cpp
===================================================================
--- llvm/lib/CodeGen/TypePromotion.cpp
+++ llvm/lib/CodeGen/TypePromotion.cpp
@@ -174,6 +174,11 @@
   // wrapping.
   bool isLegalToPromote(Value *V);
   bool TryToPromote(Value *V, unsigned PromotedWidth);
+  bool TryToPromoteICmp(ICmpInst *ICmpInst, const TargetLowering *TLI,
+                        const DataLayout &DL);
+  bool TryToPromotePHI(Instruction *Ext, PHINode *Phi,
+                       const TargetLowering *TLI, const DataLayout &DL,
+                       const TargetTransformInfo &TTI);
 
 public:
   static char ID;
@@ -919,6 +924,146 @@
   return true;
 }
 
+bool TypePromotion::TryToPromoteICmp(ICmpInst *ICmp, const TargetLowering *TLI,
+                                     const DataLayout &DL) {
+
+  LLVM_DEBUG(dbgs() << "IR Promotion: Searching from: " << *ICmp << "\n");
+
+  for (auto &Op : ICmp->operands()) {
+    if (auto *I = dyn_cast<Instruction>(Op)) {
+      EVT SrcVT = TLI->getValueType(DL, I->getType());
+      if (SrcVT.isSimple() && TLI->isTypeLegal(SrcVT.getSimpleVT()))
+        return false;
+
+      if (TLI->getTypeAction(*Ctx, SrcVT) != TargetLowering::TypePromoteInteger)
+        return false;
+      EVT PromotedVT = TLI->getTypeToTransformTo(*Ctx, SrcVT);
+      if (RegisterBitWidth < PromotedVT.getFixedSizeInBits()) {
+        LLVM_DEBUG(dbgs() << "IR Promotion: Couldn't find target register "
+                          << "for promoted type\n");
+        return false;
+      }
+
+      return TryToPromote(I, PromotedVT.getFixedSizeInBits());
+    }
+  }
+  return false;
+}
+
+bool TypePromotion::TryToPromotePHI(Instruction *ExtI, PHINode *Phi,
+                                    const TargetLowering *TLI,
+                                    const DataLayout &DL,
+                                    const TargetTransformInfo &TTI) {
+
+  EVT ExtVT = TLI->getValueType(DL, ExtI->getType());
+  if (!ExtVT.isVector() && RegisterBitWidth < ExtVT.getFixedSizeInBits()) {
+    LLVM_DEBUG(dbgs() << "IR Promotion: extension type does not fit into "
+                      << "single register\n");
+    return false;
+  }
+
+  // If the ExtI value is likely to get promoted then make sure to try to promote
+  // the PHI to this type.
+  if (TLI->getTypeAction(*Ctx, ExtVT) == TargetLowering::TypePromoteInteger) {
+    EVT PromotedVT = TLI->getTypeToTransformTo(*Ctx, ExtVT);
+    TypeSize = PromotedVT.getFixedSizeInBits();
+    if (TypeSize <= RegisterBitWidth && shouldPromote(ExtI)) {
+      ExtVT = PromotedVT;
+    }
+  }
+
+  // Check all incoming values for the existing PHI, we need to make sure we
+  // can extend them without occurrign extra cost.
+  for (Value *IncValue : Phi->incoming_values()) {
+    if (isa<Instruction>(IncValue)) {
+      auto *LoadI = cast<Instruction>(IncValue);
+      bool masked = false;
+      if (isa<IntrinsicInst>(LoadI)) {
+        Intrinsic::ID Id = cast<IntrinsicInst>(LoadI)->getIntrinsicID();
+        if (Id != Intrinsic::masked_load) {
+          LLVM_DEBUG(dbgs() << "IR Promotion: Incoming value to PHI is not a "
+                            << "load\n");
+          return false;
+        }
+        masked = true;
+      } else if (LoadI->getOpcode() != Instruction::Load) {
+        LLVM_DEBUG(dbgs() << "IR Promotion: Incoming value to PHI is not a "
+                          << "load\n");
+        return false;
+      }
+
+      // Check whether the target can merge the Load and Extend instructions.
+      InstructionCost Cost = TTI.getCastInstrCost(
+          ExtI->getOpcode(), LoadI->getType(), ExtI->getType(),
+          masked ? TTI::CastContextHint::Masked : TTI::CastContextHint::Normal,
+          TTI::TCK_SizeAndLatency, ExtI);
+      if (!Cost.isValid() || Cost != 0) {
+        LLVM_DEBUG(dbgs() << "IR Promotion: Will not promote PHI as Load "
+                          << "and Extend instructions can not be merged.\n");
+        return false;
+      }
+
+      // Make sure the other uses of the Load can also be promoted.
+      if (!LoadI->hasOneUse()) {
+        for (Use &U : LoadI->uses()) {
+          if (U.getUser() == Phi)
+            continue;
+          if (isa<Instruction>(U.getUser())) {
+            auto *UseI = cast<Instruction>(U.getUser());
+
+            if (UseI->getOpcode() == ExtI->getOpcode())
+              continue;
+            if (isa<ICmpInst>(U.getUser())) {
+              auto CmpI = cast<ICmpInst>(U.getUser());
+              if (isa<ConstantInt>(CmpI->getOperand(0)) ||
+                  isa<ConstantInt>(CmpI->getOperand(1)))
+                continue;
+            }
+
+            LLVM_DEBUG(dbgs() << "IR Promotion: PHI's incoming load has "
+                              << "invalidating use\n");
+
+            return false;
+          }
+        }
+      }
+    } else if (isa<Constant>(IncValue))
+      continue;
+    else {
+      return false;
+    }
+  }
+
+  IRBuilder<> Builder{*Ctx};
+  Builder.SetInsertPoint(Phi);
+  bool Sign = ExtI->getOpcode() == Instruction::SExt;
+  IntegerType *ExtType = IntegerType::get(*Ctx, ExtVT.getFixedSizeInBits());
+  llvm::PHINode *NewPhi = Builder.CreatePHI(ExtType, Phi->getNumOperands());
+  for (BasicBlock *BB : Phi->blocks()) {
+    Value *IncValue = Phi->getIncomingValueForBlock(BB);
+    if (isa<Instruction>(IncValue)) {
+      auto *IncI = cast<Instruction>(IncValue);
+      Builder.SetInsertPoint(IncI->getParent(), ++IncI->getIterator());
+      Value *NewExt = Sign ? Builder.CreateSExt(IncI, ExtType)
+                           : Builder.CreateZExt(IncI, ExtType);
+      NewPhi->addIncoming(NewExt, BB);
+    } else if (isa<Constant>(IncValue)) {
+      Constant *Cst = cast<Constant>(IncValue);
+      Constant *NewCst = Sign ? ConstantExpr::getSExt(Cst, ExtType)
+                              : ConstantExpr::getZExt(Cst, ExtType);
+      NewPhi->addIncoming(NewCst, BB);
+    }
+  }
+
+  // If we can't promote the ExtI then don't make the replacement.
+  if (ExtI->getType() != ExtType) {
+    if (!TryToPromote(ExtI, ExtVT.getFixedSizeInBits()))
+      return false;
+  }
+  ExtI->replaceAllUsesWith(NewPhi);
+  return true;
+}
+
 bool TypePromotion::runOnFunction(Function &F) {
   if (skipFunction(F) || DisablePromotion)
     return false;
@@ -949,37 +1094,21 @@
       if (AllVisited.count(&I))
         continue;
 
-      if (!isa<ICmpInst>(&I))
-        continue;
-
-      auto *ICmp = cast<ICmpInst>(&I);
-      // Skip signed or pointer compares
-      if (ICmp->isSigned() ||
-          !isa<IntegerType>(ICmp->getOperand(0)->getType()))
-        continue;
-
-      LLVM_DEBUG(dbgs() << "IR Promotion: Searching from: " << *ICmp << "\n");
-
-      for (auto &Op : ICmp->operands()) {
-        if (auto *I = dyn_cast<Instruction>(Op)) {
-          EVT SrcVT = TLI->getValueType(DL, I->getType());
-          if (SrcVT.isSimple() && TLI->isTypeLegal(SrcVT.getSimpleVT()))
-            break;
-
-          if (TLI->getTypeAction(ICmp->getContext(), SrcVT) !=
-              TargetLowering::TypePromoteInteger)
-            break;
-          EVT PromotedVT = TLI->getTypeToTransformTo(ICmp->getContext(), SrcVT);
-          if (RegisterBitWidth < PromotedVT.getFixedSizeInBits()) {
-            LLVM_DEBUG(dbgs() << "IR Promotion: Couldn't find target register "
-                       << "for promoted type\n");
-            break;
-          }
-
-          MadeChange |= TryToPromote(I, PromotedVT.getFixedSizeInBits());
-          break;
-        }
-      }
+      if (isa<ICmpInst>(&I)) {
+        auto *ICmp = cast<ICmpInst>(&I);
+
+        // Skip signed or pointer compares
+        if (ICmp->isSigned() || !isa<IntegerType>(ICmp->getOperand(0)->getType()))
+          continue;
+
+        MadeChange |= TryToPromoteICmp(ICmp, TLI, DL);
+      } else if (isa<Instruction>(&I) &&
+                 (I.getOpcode() == Instruction::ZExt ||
+                  I.getOpcode() == Instruction::SExt) &&
+                 isa<PHINode>(I.getOperand(0)) && I.getOperand(0)->hasOneUse())
+        MadeChange |=
+            TryToPromotePHI(cast<Instruction>(&I),
+                            cast<PHINode>(I.getOperand(0)), TLI, DL, TTI);
     }
     LLVM_DEBUG(if (verifyFunction(F, &dbgs())) {
                 dbgs() << F;