1 files changed, 44 insertions, 22 deletions

diff --git a/llvm/lib/Target/ARM/ARMCodeGenPrepare.cpp b/llvm/lib/Target/ARM/ARMCodeGenPrepare.cpp
index 69ecc337820..3151621ec10 100644
--- a/llvm/lib/Target/ARM/ARMCodeGenPrepare.cpp
+++ b/llvm/lib/Target/ARM/ARMCodeGenPrepare.cpp
@@ -181,6 +181,8 @@ static bool isSink(Value *V) {
     return UsesNarrowValue(Return->getReturnValue());
   if (auto *Trunc = dyn_cast<TruncInst>(V))
     return UsesNarrowValue(Trunc->getOperand(0));
+  if (auto *ICmp = dyn_cast<ICmpInst>(V))
+    return ICmp->isSigned();
 
   return isa<CallInst>(V);
 }
@@ -294,6 +296,11 @@ void IRPromoter::Mutate(Type *OrigTy,
   LLVM_DEBUG(dbgs() << "ARM CGP: Promoting use-def chains to from "
              << ARMCodeGenPrepare::TypeSize << " to 32-bits\n");
 
+  // Cache original types.
+  DenseMap<Value*, Type*> TruncTysMap;
+  for (auto *V : Visited)
+    TruncTysMap[V] = V->getType();
+
   auto ReplaceAllUsersOfWith = [&](Value *From, Value *To) {
     SmallVector<Instruction*, 4> Users;
     Instruction *InstTo = dyn_cast<Instruction>(To);
@@ -337,6 +344,7 @@ void IRPromoter::Mutate(Type *OrigTy,
     ReplaceAllUsersOfWith(I, Call);
     InstsToRemove.push_back(I);
     NewInsts.insert(Call);
+    TruncTysMap[Call] = OrigTy;
   };
 
   auto InsertZExt = [&](Value *V, Instruction *InsertPt) {
@@ -351,6 +359,7 @@ void IRPromoter::Mutate(Type *OrigTy,
       ZExt->moveAfter(InsertPt);
     ReplaceAllUsersOfWith(V, ZExt);
     NewInsts.insert(ZExt);
+    TruncTysMap[ZExt] = TruncTysMap[V];
   };
 
   // First, insert extending instructions between the leaves and their users.
@@ -409,6 +418,22 @@ void IRPromoter::Mutate(Type *OrigTy,
     InsertDSPIntrinsic(cast<Instruction>(V));
   }
 
+  auto InsertTrunc = [&](Value *V, Type *TruncTy) -> Instruction* {
+    if (TruncTy == ExtTy || !isa<Instruction>(V) ||
+        !isa<IntegerType>(V->getType()))
+      return nullptr;
+
+    if (!Promoted.count(V) && !NewInsts.count(V))
+      return nullptr;
+
+    LLVM_DEBUG(dbgs() << "ARM CGP: Creating " << *TruncTy << " Trunc for "
+               << *V << "\n");
+    Builder.SetInsertPoint(cast<Instruction>(V));
+    auto *Trunc = cast<Instruction>(Builder.CreateTrunc(V, TruncTy));
+    NewInsts.insert(Trunc);
+    return Trunc;
+  };
+
   LLVM_DEBUG(dbgs() << "ARM CGP: Fixing up the roots:\n");
   // Fix up any stores or returns that use the results of the promoted
   // chain.
@@ -423,28 +448,25 @@ void IRPromoter::Mutate(Type *OrigTy,
       TruncTy = F->getFunctionType()->getReturnType();
     }
 
-    for (unsigned i = 0; i < I->getNumOperands(); ++i) {
-      Value *V = I->getOperand(i);
-      if (!isa<IntegerType>(V->getType()))
-        continue;
-
-      if (Promoted.count(V) || NewInsts.count(V)) {
-        if (auto *Op = dyn_cast<Instruction>(V)) {
-
-          if (auto *Call = dyn_cast<CallInst>(I))
-            TruncTy = Call->getFunctionType()->getParamType(i);
+    // These will only have one operand to fix.
+    if (isa<StoreInst>(I) || isa<ReturnInst>(I) || isa<TruncInst>(I)) {
+      if (Instruction *Trunc = InsertTrunc(I->getOperand(0), TruncTy)) {
+        Trunc->moveBefore(I);
+        I->setOperand(0, Trunc);
+      }
+      continue;
+    }
 
-          if (TruncTy == ExtTy)
-            continue;
+    // Now handle calls and signed icmps.
+    for (unsigned i = 0; i < I->getNumOperands(); ++i) {
+      if (auto *Call = dyn_cast<CallInst>(I))
+        TruncTy = Call->getFunctionType()->getParamType(i);
+      else
+        TruncTy = TruncTysMap[I->getOperand(i)];
 
-          LLVM_DEBUG(dbgs() << "ARM CGP: Creating " << *TruncTy
-                     << " Trunc for " << *Op << "\n");
-          Builder.SetInsertPoint(Op);
-          auto *Trunc = cast<Instruction>(Builder.CreateTrunc(Op, TruncTy));
-          Trunc->moveBefore(I);
-          I->setOperand(i, Trunc);
-          NewInsts.insert(Trunc);
-        }
+      if (Instruction *Trunc = InsertTrunc(I->getOperand(i), TruncTy)) {
+        Trunc->moveBefore(I);
+        I->setOperand(i, Trunc);
       }
     }
   }
@@ -458,8 +480,8 @@ void IRPromoter::Mutate(Type *OrigTy,
 bool ARMCodeGenPrepare::isSupportedValue(Value *V) {
   LLVM_DEBUG(dbgs() << "ARM CGP: Is " << *V << " supported?\n");
 
-  if (auto *ICmp = dyn_cast<ICmpInst>(V))
-    return ICmp->isEquality() || !ICmp->isSigned();
+  if (isa<ICmpInst>(V))
+    return true;
 
   // Memory instructions
   if (isa<StoreInst>(V) || isa<GetElementPtrInst>(V))