1 files changed, 21 insertions, 13 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index 44be33eca4a..5713b7b7f9a 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -17,6 +17,7 @@
 #include "AMDGPUSubtarget.h"
 #include "AMDGPUTargetMachine.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/DivergenceAnalysis.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/ValueTracking.h"
@@ -58,6 +59,7 @@ static cl::opt<bool> WidenLoads(
 class AMDGPUCodeGenPrepare : public FunctionPass,
                              public InstVisitor<AMDGPUCodeGenPrepare, bool> {
   const GCNSubtarget *ST = nullptr;
+  AssumptionCache *AC = nullptr;
   DivergenceAnalysis *DA = nullptr;
   Module *Mod = nullptr;
   bool HasUnsafeFPMath = false;
@@ -134,11 +136,12 @@ class AMDGPUCodeGenPrepare : public FunctionPass,
   bool promoteUniformBitreverseToI32(IntrinsicInst &I) const;
 
   /// Expands 24 bit div or rem.
-  Value* expandDivRem24(IRBuilder<> &Builder, Value *Num, Value *Den,
+  Value* expandDivRem24(IRBuilder<> &Builder, BinaryOperator &I,
+                        Value *Num, Value *Den,
                         bool IsDiv, bool IsSigned) const;
 
   /// Expands 32 bit div or rem.
-  Value* expandDivRem32(IRBuilder<> &Builder, Instruction::BinaryOps Opc,
+  Value* expandDivRem32(IRBuilder<> &Builder, BinaryOperator &I,
                         Value *Num, Value *Den) const;
 
   /// Widen a scalar load.
@@ -173,6 +176,7 @@ public:
   StringRef getPassName() const override { return "AMDGPU IR optimizations"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<DivergenceAnalysis>();
     AU.setPreservesAll();
  }
@@ -500,16 +504,17 @@ static Value* getMulHu(IRBuilder<> &Builder, Value *LHS, Value *RHS) {
 // The fractional part of a float is enough to accurately represent up to
 // a 24-bit signed integer.
 Value* AMDGPUCodeGenPrepare::expandDivRem24(IRBuilder<> &Builder,
+                                            BinaryOperator &I,
                                             Value *Num, Value *Den,
                                             bool IsDiv, bool IsSigned) const {
   assert(Num->getType()->isIntegerTy(32));
 
   const DataLayout &DL = Mod->getDataLayout();
-  unsigned LHSSignBits = ComputeNumSignBits(Num, DL);
+  unsigned LHSSignBits = ComputeNumSignBits(Num, DL, 0, AC, &I);
   if (LHSSignBits < 9)
     return nullptr;
 
-  unsigned RHSSignBits = ComputeNumSignBits(Den, DL);
+  unsigned RHSSignBits = ComputeNumSignBits(Den, DL, 0, AC, &I);
   if (RHSSignBits < 9)
     return nullptr;
 
@@ -603,8 +608,9 @@ Value* AMDGPUCodeGenPrepare::expandDivRem24(IRBuilder<> &Builder,
 }
 
 Value* AMDGPUCodeGenPrepare::expandDivRem32(IRBuilder<> &Builder,
-                                            Instruction::BinaryOps Opc,
+                                            BinaryOperator &I,
                                             Value *Num, Value *Den) const {
+  Instruction::BinaryOps Opc = I.getOpcode();
   assert(Opc == Instruction::URem || Opc == Instruction::UDiv ||
          Opc == Instruction::SRem || Opc == Instruction::SDiv);
 
@@ -632,7 +638,7 @@ Value* AMDGPUCodeGenPrepare::expandDivRem32(IRBuilder<> &Builder,
     }
   }
 
-  if (Value *Res = expandDivRem24(Builder, Num, Den, IsDiv, IsSigned)) {
+  if (Value *Res = expandDivRem24(Builder, I, Num, Den, IsDiv, IsSigned)) {
     Res = Builder.CreateTrunc(Res, Ty);
     return Res;
   }
@@ -767,16 +773,16 @@ bool AMDGPUCodeGenPrepare::visitBinaryOperator(BinaryOperator &I) {
     if (VectorType *VT = dyn_cast<VectorType>(Ty)) {
       NewDiv = UndefValue::get(VT);
 
-      for (unsigned I = 0, E = VT->getNumElements(); I != E; ++I) {
-        Value *NumEltI = Builder.CreateExtractElement(Num, I);
-        Value *DenEltI = Builder.CreateExtractElement(Den, I);
-        Value *NewElt = expandDivRem32(Builder, Opc, NumEltI, DenEltI);
+      for (unsigned N = 0, E = VT->getNumElements(); N != E; ++N) {
+        Value *NumEltN = Builder.CreateExtractElement(Num, N);
+        Value *DenEltN = Builder.CreateExtractElement(Den, N);
+        Value *NewElt = expandDivRem32(Builder, I, NumEltN, DenEltN);
         if (!NewElt)
-          NewElt = Builder.CreateBinOp(Opc, NumEltI, DenEltI);
-        NewDiv = Builder.CreateInsertElement(NewDiv, NewElt, I);
+          NewElt = Builder.CreateBinOp(Opc, NumEltN, DenEltN);
+        NewDiv = Builder.CreateInsertElement(NewDiv, NewElt, N);
       }
     } else {
-      NewDiv = expandDivRem32(Builder, Opc, Num, Den);
+      NewDiv = expandDivRem32(Builder, I, Num, Den);
     }
 
     if (NewDiv) {
@@ -891,6 +897,7 @@ bool AMDGPUCodeGenPrepare::runOnFunction(Function &F) {
 
   const AMDGPUTargetMachine &TM = TPC->getTM<AMDGPUTargetMachine>();
   ST = &TM.getSubtarget<GCNSubtarget>(F);
+  AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   DA = &getAnalysis<DivergenceAnalysis>();
   HasUnsafeFPMath = hasUnsafeFPMath(F);
   AMDGPUASI = TM.getAMDGPUAS();
@@ -910,6 +917,7 @@ bool AMDGPUCodeGenPrepare::runOnFunction(Function &F) {
 
 INITIALIZE_PASS_BEGIN(AMDGPUCodeGenPrepare, DEBUG_TYPE,
                       "AMDGPU IR optimizations", false, false)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DivergenceAnalysis)
 INITIALIZE_PASS_END(AMDGPUCodeGenPrepare, DEBUG_TYPE, "AMDGPU IR optimizations",
                     false, false)