Using DxilValueCache for unroll. (#2694)

include/dxc/HLSL/DxilValueCache.h → include/llvm/Analysis/DxilValueCache.h

@@ -8,12 +8,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Pass.h"
-#include "llvm/Transforms/Utils/ValueMapper.h"
+#include "llvm/IR/ValueMap.h"
 
 namespace llvm {
 
 class Module;
 class DominatorTree;
+class Constant;
 
 struct DxilValueCache : public ImmutablePass {
   static char ID;
@@ -68,6 +69,7 @@ public:
 
   void dump() const;
   Value *GetValue(Value *V, DominatorTree *DT=nullptr);
+  Constant *GetConstValue(Value *V, DominatorTree *DT = nullptr);
   void ResetUnknowns() { ValueMap.ResetUnknowns(); }
   bool IsAlwaysReachable(BasicBlock *BB, DominatorTree *DT=nullptr);
   bool IsUnreachable(BasicBlock *BB, DominatorTree *DT=nullptr);

lib/Analysis/CMakeLists.txt

@@ -28,6 +28,7 @@ add_llvm_library(LLVMAnalysis
   DxilConstantFolding.cpp
   DxilConstantFoldingExt.cpp
   DxilSimplify.cpp
+  DxilValueCache.cpp
   IVUsers.cpp
   InstCount.cpp
   InstructionSimplify.cpp

lib/HLSL/DxilValueCache.cpp → lib/Analysis/DxilValueCache.cpp

@@ -25,7 +25,7 @@
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/ADT/Statistic.h"
 
-#include "dxc/HLSL/DxilValueCache.h"
+#include "llvm/Analysis/DxilValueCache.h"
 
 #include <unordered_set>
 
@@ -375,6 +375,12 @@ Value *DxilValueCache::GetValue(Value *V, DominatorTree *DT) {
   return ProcessValue(V, DT);
 }
 
+Constant *DxilValueCache::GetConstValue(Value *V, DominatorTree *DT) {
+  if (Value *NewV = GetValue(V))
+    return dyn_cast<Constant>(NewV);
+  return nullptr;
+}
+
 bool DxilValueCache::IsAlwaysReachable(BasicBlock *BB, DominatorTree *DT) {
   ProcessValue(BB, DT);
   return IsAlwaysReachable_(BB);

lib/Analysis/ScalarEvolution.cpp

@@ -88,6 +88,8 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Analysis/DxilValueCache.h" // HLSL Change
+
 #include <algorithm>
 using namespace llvm;
 
@@ -125,6 +127,7 @@ INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DxilValueCache) // HLSL Change
 INITIALIZE_PASS_END(ScalarEvolution, "scalar-evolution",
                 "Scalar Evolution Analysis", false, true)
 char ScalarEvolution::ID = 0;
@@ -5160,6 +5163,14 @@ ScalarEvolution::ComputeExitLimitFromICmp(const Loop *L,
         if (!isa<SCEVCouldNotCompute>(Ret)) return Ret;
       }
 
+  // HLSL Change - begin
+  // Try to compute the value exhaustively *right now*. Before trying the more pessimistic
+  // partial evaluation.
+  const SCEV *AggresiveResult = ComputeExitCountExhaustively(L, ExitCond, !L->contains(TBB));
+  if (AggresiveResult != getCouldNotCompute())
+    return AggresiveResult;
+  // HLSL Change - end
+
   switch (Cond) {
   case ICmpInst::ICMP_NE: {                     // while (X != Y)
     // Convert to: while (X-Y != 0)
@@ -5198,7 +5209,8 @@ ScalarEvolution::ComputeExitLimitFromICmp(const Loop *L,
 #endif
     break;
   }
-  return ComputeExitCountExhaustively(L, ExitCond, !L->contains(TBB));
+  // return ComputeExitCountExhaustively(L, ExitCond, !L->contains(TBB)); // HLSL Change
+  return getCouldNotCompute(); // HLSL Change - We already tried the exhaustive approach earlier, so don't try again and just give up.
 }
 
 ScalarEvolution::ExitLimit
@@ -5356,6 +5368,7 @@ static bool canConstantEvolve(Instruction *I, const Loop *L) {
 /// recursing through each instruction operand until reaching a loop header phi.
 static PHINode *
 getConstantEvolvingPHIOperands(Instruction *UseInst, const Loop *L,
+                               DxilValueCache *DVC, // HLSL Change
                                DenseMap<Instruction *, PHINode *> &PHIMap) {
 
   // Otherwise, we can evaluate this instruction if all of its operands are
@@ -5366,6 +5379,10 @@ getConstantEvolvingPHIOperands(Instruction *UseInst, const Loop *L,
 
     if (isa<Constant>(*OpI)) continue;
 
+    // HLSL Change begin
+    if (DVC->GetConstValue(*OpI)) continue;
+    // HLSL Change end
+
     Instruction *OpInst = dyn_cast<Instruction>(*OpI);
     if (!OpInst || !canConstantEvolve(OpInst, L)) return nullptr;
 
@@ -5378,7 +5395,8 @@ getConstantEvolvingPHIOperands(Instruction *UseInst, const Loop *L,
     if (!P) {
       // Recurse and memoize the results, whether a phi is found or not.
       // This recursive call invalidates pointers into PHIMap.
-      P = getConstantEvolvingPHIOperands(OpInst, L, PHIMap);
+      //P = getConstantEvolvingPHIOperands(OpInst, L, PHIMap); // HLSL Change
+      P = getConstantEvolvingPHIOperands(OpInst, L, DVC, PHIMap); // HLSL Change - Pass DVC
       PHIMap[OpInst] = P;
     }
     if (!P)
@@ -5396,7 +5414,8 @@ getConstantEvolvingPHIOperands(Instruction *UseInst, const Loop *L,
 /// way, but the operands of an operation must either be constants or a value
 /// derived from a constant PHI.  If this expression does not fit with these
 /// constraints, return null.
-static PHINode *getConstantEvolvingPHI(Value *V, const Loop *L) {
+// static PHINode *getConstantEvolvingPHI(Value *V, const Loop *L) { // HLSL Change
+static PHINode *getConstantEvolvingPHI(Value *V, const Loop *L, DxilValueCache *DVC) { // HLSL Change
   Instruction *I = dyn_cast<Instruction>(V);
   if (!I || !canConstantEvolve(I, L)) return nullptr;
 
@@ -5406,7 +5425,8 @@ static PHINode *getConstantEvolvingPHI(Value *V, const Loop *L) {
 
   // Record non-constant instructions contained by the loop.
   DenseMap<Instruction *, PHINode *> PHIMap;
-  return getConstantEvolvingPHIOperands(I, L, PHIMap);
+  // return getConstantEvolvingPHIOperands(I, L, PHIMap); // HLSL Change
+  return getConstantEvolvingPHIOperands(I, L, DVC, PHIMap); // HLSL Change
 }
 
 /// EvaluateExpression - Given an expression that passes the
@@ -5561,7 +5581,8 @@ ScalarEvolution::getConstantEvolutionLoopExitValue(PHINode *PN,
 const SCEV *ScalarEvolution::ComputeExitCountExhaustively(const Loop *L,
                                                           Value *Cond,
                                                           bool ExitWhen) {
-  PHINode *PN = getConstantEvolvingPHI(Cond, L);
+  // PHINode *PN = getConstantEvolvingPHI(Cond, L); // HLSL Change
+  PHINode *PN = getConstantEvolvingPHI(Cond, L, &getAnalysis<DxilValueCache>()); // HLSL Change
   if (!PN) return getCouldNotCompute();
 
   // If the loop is canonicalized, the PHI will have exactly two entries.
@@ -5578,20 +5599,48 @@ const SCEV *ScalarEvolution::ComputeExitCountExhaustively(const Loop *L,
   PHINode *PHI = nullptr;
   for (BasicBlock::iterator I = Header->begin();
        (PHI = dyn_cast<PHINode>(I)); ++I) {
+
     Constant *StartCST =
       dyn_cast<Constant>(PHI->getIncomingValue(!SecondIsBackedge));
+
+    // HLSL Change begin
+    // If we don't have a constant, try getting a constant from the value cache.
+    if (!StartCST)
+      if (Constant *C = getAnalysis<DxilValueCache>().GetConstValue(PHI->getIncomingValue(!SecondIsBackedge)))
+        StartCST = C;
+    // HLSL Change end
+
     if (!StartCST) continue;
     CurrentIterVals[PHI] = StartCST;
   }
   if (!CurrentIterVals.count(PN))
     return getCouldNotCompute();
 
+  // HLSL Change begin
+  SmallVector<std::pair<Instruction *, Constant *>, 4> KnownInvariantOps;
+  if (Instruction *CondI = dyn_cast<Instruction>(Cond)) {
+    for (Use &U : CondI->operands()) {
+      if (Instruction *OpI = dyn_cast<Instruction>(U.get())) {
+        if (Value *V = getAnalysis<DxilValueCache>().GetValue(OpI)) {
+          if (Constant *C = dyn_cast<Constant>(V))
+            KnownInvariantOps.push_back({ OpI, C });
+        }
+      }
+    }
+  }
+  // HLSL Change end
+
   // Okay, we find a PHI node that defines the trip count of this loop.  Execute
   // the loop symbolically to determine when the condition gets a value of
   // "ExitWhen".
   unsigned MaxIterations = MaxBruteForceIterations;   // Limit analysis.
   const DataLayout &DL = F->getParent()->getDataLayout();
   for (unsigned IterationNum = 0; IterationNum != MaxIterations;++IterationNum){
+    // HLSL Change begin
+    for (std::pair<Instruction *, Constant *> &Pair : KnownInvariantOps)
+      CurrentIterVals[Pair.first] = Pair.second;
+    // HLSL Change end
+
     ConstantInt *CondVal = dyn_cast_or_null<ConstantInt>(
         EvaluateExpression(Cond, L, CurrentIterVals, DL, TLI));
 
@@ -8116,6 +8165,7 @@ void ScalarEvolution::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequiredTransitive<LoopInfoWrapperPass>();
   AU.addRequiredTransitive<DominatorTreeWrapperPass>();
   AU.addRequired<TargetLibraryInfoWrapperPass>();
+  AU.addRequired<DxilValueCache>(); // HLSL Change
 }
 
 bool ScalarEvolution::hasLoopInvariantBackedgeTakenCount(const Loop *L) {

lib/HLSL/CMakeLists.txt

@@ -32,7 +32,6 @@ add_llvm_library(LLVMHLSL
   DxilTranslateRawBuffer.cpp
   DxilExportMap.cpp
   DxilValidation.cpp
-  DxilValueCache.cpp
   DxcOptimizer.cpp
   HLDeadFunctionElimination.cpp
   HLExpandStoreIntrinsics.cpp

lib/HLSL/DxcOptimizer.cpp

@@ -20,7 +20,7 @@
 #include "dxc/HLSL/HLMatrixLowerPass.h"
 #include "dxc/HLSL/DxilGenerationPass.h"
 #include "dxc/HLSL/ComputeViewIdState.h"
-#include "dxc/HLSL/DxilValueCache.h"
+#include "llvm/Analysis/DxilValueCache.h"
 #include "dxc/DXIL/DxilUtil.h"
 #include "dxc/Support/dxcapi.impl.h"
 

lib/HLSL/DxilCondenseResources.cpp

@@ -21,7 +21,7 @@
 #include "dxc/DXIL/DxilUtil.h"
 #include "dxc/HLSL/HLMatrixType.h"
 #include "dxc/HLSL/HLModule.h"
-#include "dxc/HLSL/DxilValueCache.h"
+#include "llvm/Analysis/DxilValueCache.h"
 #include "dxc/DXIL/DxilMetadataHelper.h"
 
 #include "llvm/IR/Instructions.h"

lib/HLSL/DxilLegalizeSampleOffsetPass.cpp

@@ -10,7 +10,7 @@
 ///////////////////////////////////////////////////////////////////////////////
 
 #include "dxc/HLSL/DxilGenerationPass.h"
-#include "dxc/HLSL/DxilValueCache.h"
+#include "llvm/Analysis/DxilValueCache.h"
 #include "dxc/DXIL/DxilModule.h"
 #include "dxc/DXIL/DxilOperations.h"
 

lib/Transforms/IPO/PassManagerBuilder.cpp

@@ -29,7 +29,7 @@
 #include "dxc/HLSL/DxilGenerationPass.h" // HLSL Change
 #include "dxc/HLSL/HLMatrixLowerPass.h" // HLSL Change
 #include "dxc/HLSL/ComputeViewIdState.h" // HLSL Change
-#include "dxc/HLSL/DxilValueCache.h" // HLSL Change
+#include "llvm/Analysis/DxilValueCache.h" // HLSL Change
 
 using namespace llvm;
 

lib/Transforms/Scalar/DxilEliminateVector.cpp

@@ -21,7 +21,7 @@
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/DIBuilder.h"
 
-#include "dxc/HLSL/DxilValueCache.h"
+#include "llvm/Analysis/DxilValueCache.h"
 
 #include <vector>
 

lib/Transforms/Scalar/DxilLoopUnroll.cpp

@@ -77,9 +77,11 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/IR/LegacyPassManager.h"
+#include "llvm/Analysis/ScalarEvolutionExpressions.h"
 
 #include "dxc/DXIL/DxilUtil.h"
 #include "dxc/HLSL/HLModule.h"
+#include "llvm/Analysis/DxilValueCache.h"
 
 using namespace llvm;
 using namespace hlsl;
@@ -127,6 +129,7 @@ public:
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addPreserved<DominatorTreeWrapperPass>();
     AU.addRequired<ScalarEvolution>();
+    AU.addRequired<DxilValueCache>();
     AU.addRequiredID(LoopSimplifyID);
   }
 };
@@ -536,7 +539,7 @@ inline static int64_t GetGEPIndex(GEPOperator *GEP, unsigned idx) {
 // Otherwise it emits an error and fails compilation.
 //
 template<typename IteratorT>
-static bool BreakUpArrayAllocas(bool AllowOOBIndex, IteratorT ItBegin, IteratorT ItEnd, DominatorTree *DT, AssumptionCache *AC) { 
+static bool BreakUpArrayAllocas(bool AllowOOBIndex, IteratorT ItBegin, IteratorT ItEnd, DominatorTree *DT, AssumptionCache *AC, DxilValueCache *DVC) { 
   bool Success = true;
 
   SmallVector<AllocaInst *, 8> WorkList(ItBegin, ItEnd);
@@ -558,7 +561,19 @@ static bool BreakUpArrayAllocas(bool AllowOOBIndex, IteratorT ItBegin, IteratorT
 
     GEPs.clear(); // Re-use array
     for (User *U : AI->users()) {
+      // Try to set all GEP operands to constant
       if (GEPOperator *GEP = dyn_cast<GEPOperator>(U)) {
+        if (!GEP->hasAllConstantIndices() && isa<GetElementPtrInst>(GEP)) {
+          for (unsigned i = 0; i < GEP->getNumIndices(); i++) {
+            Value *IndexOp = GEP->getOperand(i + 1);
+            if (isa<Constant>(IndexOp))
+              continue;
+
+            if (Constant *C = DVC->GetConstValue(IndexOp))
+              GEP->setOperand(i + 1, C);
+          }
+        }
+
         if (!GEP->hasAllConstantIndices() || GEP->getNumIndices() < 2 ||
           GetGEPIndex(GEP, 0) != 0)
         {
@@ -689,6 +704,7 @@ bool DxilLoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   DebugLoc LoopLoc = L->getStartLoc(); // Debug location for the start of the loop.
   Function *F = L->getHeader()->getParent();
   ScalarEvolution *SE = &getAnalysis<ScalarEvolution>();
+  DxilValueCache *DVC = &getAnalysis<DxilValueCache>();
 
   bool HasExplicitLoopCount = false;
   int ExplicitUnrollCountSigned = 0;
@@ -720,15 +736,13 @@ bool DxilLoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   }
 
   unsigned TripCount = 0;
-  unsigned TripMultiple = 0;
-  bool HasTripCount = false;
+
   BasicBlock *ExitingBlock = L->getLoopLatch();
   if (!ExitingBlock || !L->isLoopExiting(ExitingBlock))
     ExitingBlock = L->getExitingBlock();
+
   if (ExitingBlock) {
     TripCount = SE->getSmallConstantTripCount(L, ExitingBlock);
-    TripMultiple = SE->getSmallConstantTripMultiple(L, ExitingBlock);
-    HasTripCount = TripMultiple != 1 || TripCount == 1;
   }
 
   // Analysis passes
@@ -853,7 +867,7 @@ bool DxilLoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   unsigned MaxAttempt = this->MaxIterationAttempt;
   // If we were able to figure out the definitive trip count,
   // just unroll that many times.
-  if (HasTripCount) {
+  if (TripCount != 0) {
     MaxAttempt = TripCount;
   }
   else if (HasExplicitLoopCount) {
@@ -966,7 +980,12 @@ bool DxilLoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
     // Check exit condition to see if we fully unrolled the loop
     if (BranchInst *BI = dyn_cast<BranchInst>(CurIteration.Latch->getTerminator())) {
       bool Cond = false;
-      if (GetConstantI1(BI->getCondition(), &Cond)) {
+
+      Value *ConstantCond = BI->getCondition();
+      if (Value *C = DVC->GetValue(ConstantCond))
+        ConstantCond = C;
+
+      if (GetConstantI1(ConstantCond, &Cond)) {
         if (BI->getSuccessor(Cond ? 1 : 0) == CurIteration.Header) {
           Succeeded = true;
           break;
@@ -976,7 +995,7 @@ bool DxilLoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
 
     // We've reached the N defined in [unroll(N)]
     if ((HasExplicitLoopCount && IterationI+1 >= ExplicitUnrollCount) ||
-      (HasTripCount && IterationI+1 >= TripCount))
+      (TripCount != 0 && IterationI+1 >= TripCount))
     {
       Succeeded = true;
       BranchInst *BI = cast<BranchInst>(CurIteration.Latch->getTerminator());
@@ -1074,7 +1093,7 @@ bool DxilLoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
 
     // Now that we potentially turned some GEP indices into constants,
     // try to clean up their allocas.
-    if (!BreakUpArrayAllocas(FxcCompatMode /* allow oob index */, ProblemAllocas.begin(), ProblemAllocas.end(), DT, AC)) {
+    if (!BreakUpArrayAllocas(FxcCompatMode /* allow oob index */, ProblemAllocas.begin(), ProblemAllocas.end(), DT, AC, DVC)) {
       FailLoopUnroll(false, F->getContext(), LoopLoc, "Could not unroll loop due to out of bound array access.");
     }
 
@@ -1245,5 +1264,17 @@ Pass *llvm::createDxilLoopUnrollPass(unsigned MaxIterationAttempt) {
   return new DxilLoopUnroll(MaxIterationAttempt);
 }
 
-INITIALIZE_PASS(DxilConditionalMem2Reg, "dxil-cond-mem2reg", "Dxil Conditional Mem2Reg", false, false)
-INITIALIZE_PASS(DxilLoopUnroll, "dxil-loop-unroll", "Dxil Unroll loops", false, false)
+INITIALIZE_PASS_BEGIN(DxilConditionalMem2Reg, "dxil-cond-mem2reg", "Dxil Conditional Mem2Reg", false, false)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
+INITIALIZE_PASS_END(DxilConditionalMem2Reg, "dxil-cond-mem2reg", "Dxil Conditional Mem2Reg", false, false)
+
+INITIALIZE_PASS_BEGIN(DxilLoopUnroll, "dxil-loop-unroll", "Dxil Unroll loops", false, false)
+INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(ScalarEvolution)
+INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
+INITIALIZE_PASS_DEPENDENCY(DxilValueCache)
+INITIALIZE_PASS_END(DxilLoopUnroll, "dxil-loop-unroll", "Dxil Unroll loops", false, false)
+

tools/clang/test/HLSLFileCheck/hlsl/control_flow/loops/loop7.hlsl

@@ -0,0 +1,50 @@
+// RUN: %dxc -E main -Od -T ps_6_0 %s | FileCheck %s
+
+// CHECK: @main
+
+struct S {
+  uint flags;
+};
+
+#define FLAG_0 0x4
+#define FLAG_1 0x20
+#define FLAG_2 0x100
+
+Texture2D t0 : register(t0);
+Texture2D t1 : register(t1);
+
+bool has_flag(uint flags, uint flag) {
+  return (flags & flag) != 0;
+}
+
+float foo(S s, bool add_flag) {
+
+  Texture2D textures[10] = {
+    t0, t1, t0, t1, t0,
+    t1, t0, t1, t0, t1,
+  };
+
+  if (add_flag) {
+    s.flags |= FLAG_2;
+  }
+
+  uint low = has_flag(s.flags, FLAG_0) ? 0 : 1; // Should be 0
+  uint high = has_flag(s.flags, FLAG_1) ? 10 : 11; // Should be 11
+
+  float ret = 0;
+  [unroll]
+  for (uint i = low; i*2 <= high; i++) { // Only 0, 1, 2
+    ret += textures[i].Load(i).x;
+  }
+
+  return ret;
+}
+
+[RootSignature("DescriptorTable(SRV(t0), SRV(t1))")]
+float main() : SV_Target {
+  S s = (S)0;
+  s.flags |= FLAG_0;
+  return foo(s, false);
+}
+
+

tools/clang/test/HLSLFileCheck/hlsl/control_flow/loops/loop8.hlsl

@@ -0,0 +1,52 @@
+// RUN: %dxc -E main -Od -T ps_6_0 %s | FileCheck %s
+
+// CHECK: @main
+
+struct S {
+  uint flags;
+};
+
+#define FLAG_0 0x4
+#define FLAG_1 0x20
+#define FLAG_2 0x100
+
+Texture2D t0 : register(t0);
+Texture2D t1 : register(t1);
+
+bool has_flag(uint flags, uint flag) {
+  return (flags & flag) != 0;
+}
+
+float foo(S s, bool add_flag) {
+
+  Texture2D textures[15] = {
+    t0, t1, t0, t1, t0,
+    t1, t0, t1, t0, t1,
+    t0, t1, t0, t1, t0,
+  };
+
+  if (add_flag) {
+    s.flags |= FLAG_2;
+  }
+
+  uint low = has_flag(s.flags, FLAG_0) ? 0 : 1; // Should be 0
+  uint high = has_flag(s.flags, FLAG_1) ? 10 : 11; // Should be 11
+
+  float ret = 0;
+  [unroll]
+  for (uint i = low; i*2 <= high; i++) { // Only 0, 1, 2
+    ret += textures[i].Load(i).x;
+    high += 1;
+  }
+
+  return ret;
+}
+
+[RootSignature("DescriptorTable(SRV(t0), SRV(t1))")]
+float main() : SV_Target {
+  S s = (S)0;
+  s.flags |= FLAG_0;
+  return foo(s, false);
+}
+
+