Implement dynamic alloca-related memory operation cost reduction for loop unrolling

Implement dynamic alloca-related memory operation cost reduction for loop unrolling

Author: lioujheyu

IGC/Compiler/GenTTI.cpp

@@ -26,6 +26,8 @@ SPDX-License-Identifier: MIT
 using namespace llvm;
 using namespace IGC;
 
+#define DEBUG_TYPE "GENtti"
+
 namespace llvm {
 
 bool GenIntrinsicsTTIImpl::isLoweredToCall(const Function *F) {
@@ -81,48 +83,78 @@ unsigned countTotalInstructions(const Function *F, bool CheckSendMsg = true) {
 
 unsigned GenIntrinsicsTTIImpl::getFlatAddressSpace() { return ADDRESS_SPACE_PRIVATE; }
 
-/// Returns true if load instruction source address calculation
-/// depends only on base address, constants, and loop induction variables
-bool canReplaceWithRegisters(const LoadInst *LI, const Loop *L, ScalarEvolution &SE) {
-  auto Pointer = LI->getPointerOperand();
-  auto Base = LI->getPointerOperand()->stripInBoundsOffsets();
+bool GenIntrinsicsTTIImpl::isGEPLoopConstDerived(GetElementPtrInst *GEP, const Loop *L, ScalarEvolution &SE) {
+  if (!GEP)
+    return false;
+
+  const SCEV *SGEP = SE.getSCEV(GEP);
+
+  if (auto *AR = dyn_cast<SCEVAddRecExpr>(SGEP)) {
+    if (AR->getLoop() == L)
+      return true;
+  }
 
-  // Start with load source address
-  SmallVector<const Value *, 16> WorkList = {Pointer};
+  // Don't let pointer base interfere the traversal. This is due to some frontends
+  // generate GEP without inbound.
+  const SCEV *SGEPMinusPointerBase = SE.removePointerBase(SGEP);
 
-  // Traverse the source address calculation dependency tree
-  while (!WorkList.empty()) {
-    auto V = WorkList.pop_back_val();
+  struct CheckConstDerived {
+    bool TraversalDone = false;
+    bool AddRecFound = false;
+    bool isConstDerived = true;
 
-    if (V == Base || isa<Constant>(V)) {
-      // Do nothing if we meet base address or some constant
-    } else if (isa<CallBase>(V)) {
-      // Stop at calls
+    const Loop *L = nullptr;
+    const SCEV *S = nullptr;
+
+    CheckConstDerived(const Loop *L) : L(L) {}
+
+    bool setNotConstDerived() {
+      TraversalDone = true;
+      isConstDerived = false;
       return false;
-    } else if (auto U = dyn_cast<User>(V)) {
-      // In case of Instuction/Operator append
-      // all the operands to the work list,
-      // skip PHI nodes to prevent infinite while-loop
-      for (unsigned i = 0; i < U->getNumOperands(); ++i) {
-        auto O = U->getOperand(i);
-        if (auto P = dyn_cast<PHINode>(O)) {
-          if (!L->isAuxiliaryInductionVariable(*P, SE)) {
-            // Stop at non-auxilary IV
-            return false;
-          }
-        } else
-          WorkList.push_back(O);
+    }
+
+    bool follow(const SCEV *S) {
+      switch (S->getSCEVType()) {
+      case scConstant:
+      case scPtrToInt:
+      case scTruncate:
+      case scZeroExtend:
+      case scSignExtend:
+      case scAddExpr:
+      case scMulExpr:
+      case scUMaxExpr:
+      case scSMaxExpr:
+      case scUMinExpr:
+      case scSMinExpr:
+      case scSequentialUMinExpr:
+      case scUDivExpr:
+        return true;
+
+      case scAddRecExpr: {
+        const auto *ARLoop = cast<SCEVAddRecExpr>(S)->getLoop();
+        if (L && (ARLoop == L || ARLoop->contains(L))) {
+          AddRecFound = true;
+          return false; // Don't traverse into it
+        }
+
+        return setNotConstDerived();
       }
-    } else {
-      // Stop if we meet something apart from
-      // base address, constant value, IV
-      return false;
+
+      case scUnknown:
+      case scCouldNotCompute:
+        return setNotConstDerived();
+      }
+      llvm_unreachable("Unknown SCEV kind!");
     }
-  }
 
-  // If nothing was found above, consider load instruction source
-  // being a candidate to be replaced by registers
-  return true;
+    bool isDone() { return TraversalDone; }
+  };
+
+  CheckConstDerived CCD(L);
+  SCEVTraversal<CheckConstDerived> ST(CCD);
+  ST.visitAll(SGEPMinusPointerBase);
+  return (CCD.isConstDerived && CCD.AddRecFound);
 }
 
 void GenIntrinsicsTTIImpl::getUnrollingPreferences(Loop *L, ScalarEvolution &SE, TTI::UnrollingPreferences &UP,
@@ -230,10 +262,9 @@ void GenIntrinsicsTTIImpl::getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
     UP.Partial = true;
   } else // for high registry pressure shaders, limit the unrolling to small loops and only fully unroll
   {
-    if (IGC_GET_FLAG_VALUE(SetLoopUnrollThresholdForHighRegPressure) != 0)
-      UP.Threshold = IGC_GET_FLAG_VALUE(SetLoopUnrollThresholdForHighRegPressure);
-    else
-      UP.Threshold = 200;
+    UP.Threshold = IGC_GET_FLAG_VALUE(SetLoopUnrollThresholdForHighRegPressure);
+    // This is similiar to LLVM OptForSize scenario in LoopUnrollPass
+    UP.MaxPercentThresholdBoost = IGC_GET_FLAG_VALUE(SetLoopUnrollMaxPercentThresholdBoostForHighRegPressure);
   }
 
   unsigned MaxTripCount = SE.getSmallConstantMaxTripCount(L);
@@ -243,40 +274,90 @@ void GenIntrinsicsTTIImpl::getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
     UP.Force = true;
   }
 
-  const unsigned MaxTripCountToUseUpperBoundForLoopWithLoads = 16;
-  if (MaxTripCount && MaxTripCount <= MaxTripCountToUseUpperBoundForLoopWithLoads) {
-    // Check if loop contains LoadInst from an array
-    // that can potentially be replaced by registers
-
-    // Group all load instructions by base address
-    // of the source posinter
-    DenseMap<Value *, SmallSet<LoadInst *, 4>> LoadInstructions;
+  // For all the load/store who (having a GEP to),
+  // 1. Accessing a fixed size Alloca
+  // 2. Having an loop-iteration-inducted-only index
+  // For exmaple,
+  //
+  // bb:
+  //   %ALLOC = alloca [32 x float], align 4
+  // Loop1:
+  //   %i8 = phi i32 [ 0, %bb ], [ %i23, %Loop1 ]
+  //   %i19 = getelementptr [32 x float], ptr %ALLOC, i64 0, i64 %i8
+  //   %i23 = add i32 %i8, 1
+  //   %i14 = fmul ...
+  //   store float %i14, ptr %i19, align 4
+  //   %i24 = icmp eq i32 %i23, 32
+  //   br i1 %i24, label %..., label %Loop1
+  // ...
+  // Loop5:
+  //   %i93 = phi i32 [ %i115, %Loop5 ], [ 0, %... ]
+  //   %i99 = getelementptr [32 x float], ptr %ALLOC, i64 0, i64 %i93
+  //   %i103 = load float, ptr %i99, align 4
+  //   %i107 = fmul float %i103, 0x3F699999A0000000
+  //   %i115 = add i32 %i93, 1
+  //   %i116 = icmp eq i32 %i115, 32
+  //   br i1 %i116, label %bb117, label %Loop5
+  //
+  // Fully unrolling both loops leads SROA pass eliminate the entire access chain of the alloca. This is one of most
+  // impacted yet super common pattern across all application types. In many cases, especially when the only values that
+  // stored into alloca are compiler-detectable constant, these loops need to be unroll regardless how high the register
+  // pressure is.
+  //
+  // TODO: Having an analysis pass to link alloca with loops globally so that they are either unrolled together or not.
+  //       It can potentially do some global cost estimations.
+  // TODO: Having compilation retry enables loop unrolling for this case and determines if unrolling actually helps
+  //       reduce register pressure.
+  const unsigned UnrollMaxCountForAlloca = 64; // May need to be higher for OpenCL
+  bool AllocaFound = false;
+  if (MaxTripCount && MaxTripCount <= UnrollMaxCountForAlloca &&
+      IGC_IS_FLAG_ENABLED(EnablePromoteLoopUnrollwithAlloca)) {
+    unsigned int ThresholdBoost = 0;
     for (auto BB : L->blocks()) {
       for (auto &I : *BB) {
-        if (auto LI = dyn_cast<LoadInst>(&I)) {
-          auto Base = LI->getPointerOperand()->stripInBoundsOffsets();
-          if (isa<AllocaInst>(Base)) {
-            auto LIIterator = LoadInstructions.find(Base);
-            if (LIIterator == LoadInstructions.end())
-              LIIterator = LoadInstructions.insert(std::make_pair(Base, SmallSet<LoadInst *, 4>())).first;
-            LIIterator->second.insert(LI);
-          }
-        }
-      }
-    }
+        AllocaInst *AI = nullptr;
+        GetElementPtrInst *GEP = nullptr;
+
+        if (auto *LI = dyn_cast<LoadInst>(&I))
+          AI = dyn_cast<AllocaInst>(LI->getPointerOperand());
+        else if ((GEP = dyn_cast<GetElementPtrInst>(&I))) {
+          // Test if the GEP index is a function of the loop induction variable.
+          if (!isGEPLoopConstDerived(GEP, L, SE))
+            continue;
 
-    // Find at least one base address, such that all loads
-    // from it can be replaced by registers
-    for (const auto &LIIterator : LoadInstructions) {
-      bool Found = true;
-      for (const auto &LI : LIIterator.second)
-        Found &= canReplaceWithRegisters(LI, L, SE);
-      if (Found) {
-        UP.UpperBound = true;
-        UP.Force = true;
-        break;
+          auto *SBase = dyn_cast<SCEVUnknown>(SE.getPointerBase(SE.getSCEV(GEP)));
+          AI = dyn_cast<AllocaInst>(SBase->getValue());
+        } else
+          continue;
+
+        if (!AI)
+          continue;
+
+        Type *Ty = AI->getAllocatedType();
+        unsigned AllocaSize = Ty->isSized() ? DL.getTypeAllocSize(Ty) : 0;
+        if (AllocaSize > 1024 || AllocaSize == 0)
+          continue;
+
+        ThresholdBoost += AllocaSize;
+        if (GEP)
+          isGEPLoopInduction[GEP] = true;
+        AllocaFound = true;
       }
     }
+    if (AllocaFound) {
+      // LLVM default only to 10, boost to UnrollMaxCountForAlloca
+      UP.MaxIterationsCountToAnalyze = UnrollMaxCountForAlloca;
+      UP.Threshold += ThresholdBoost;
+      UP.Runtime = true;
+      UP.UpperBound = true;
+      UP.Force = true;
+
+      LLVM_DEBUG(dbgs() << "Increasing L:" << L->getName() << " threshold to " << UP.Threshold
+                        << " due to Alloca accessed by:");
+      for (const auto &pair : isGEPLoopInduction)
+        LLVM_DEBUG(dbgs() << " " << pair.first->getName());
+      LLVM_DEBUG(dbgs() << " \n");
+    }
   }
 
   unsigned sendMessage = 0;
@@ -439,7 +520,7 @@ void GenIntrinsicsTTIImpl::getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
                         (instCount + sendMessage * 4 > unrollLimitInstCount);
 
   // if the loop doesn't have sample, skip the unrolling parameter change
-  if (!sendMessage) {
+  if (!sendMessage && !AllocaFound) {
     // if the estimated unrolled instruction count is larger than the unrolling threshold, limit unrolling.
     if (limitUnrolling) {
       UP.Count = MIN(unrollLimitInstCount / (instCount + sendMessage * 4), 4);
@@ -453,7 +534,7 @@ void GenIntrinsicsTTIImpl::getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
 
   // if the TripCount is known, and the estimated unrolled count exceed LoopUnrollThreshold, set the unrolling count to
   // 4
-  if (limitUnrolling) {
+  if (limitUnrolling && !AllocaFound) {
     UP.Count = MIN(TripCount, 4);
     UP.MaxCount = UP.Count;
   }
@@ -469,11 +550,13 @@ void GenIntrinsicsTTIImpl::getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
     }
   }
 
+  /*** TESTING for removal
   UP.Runtime = true;
   UP.Count = 4;
   UP.MaxCount = UP.Count;
   // The following is only available and required from LLVM 3.7+.
   UP.AllowExpensiveTripCount = true;
+  ***/
 
   if (MDNode *LoopID = L->getLoopID()) {
     const llvm::StringRef maxIterMetadataNames = "spv.loop.iterations.max";
@@ -559,6 +642,19 @@ llvm::InstructionCost GenIntrinsicsTTIImpl::internalCalculateCost(const User *U,
     }
   }
 
+  if (IGC_IS_FLAG_ENABLED(EnablePromoteLoopUnrollwithAlloca)) {
+    const GetElementPtrInst *GEP = nullptr;
+    if (Operator::getOpcode(U) == Instruction::Load)
+      GEP = dyn_cast<GetElementPtrInst>(cast<LoadInst>(U)->getPointerOperand());
+    if (Operator::getOpcode(U) == Instruction::Store)
+      GEP = dyn_cast<GetElementPtrInst>(cast<StoreInst>(U)->getPointerOperand());
+
+    if (GEP) {
+      if (isGEPLoopInduction.find(GEP) != isGEPLoopInduction.end())
+        return TTI::TCC_Free;
+    }
+  }
+
   const Function *F = dyn_cast<Function>(U);
   if (F != nullptr) {
     IGC::CodeGenContext *CGC = this->ctx;

IGC/Compiler/GenTTI.h

@@ -29,6 +29,8 @@ class GenIntrinsicsTTIImpl : public IGCLLVM::TTIImplCRTPBase<GenIntrinsicsTTIImp
 public:
   GenIntrinsicsTTIImpl(IGC::CodeGenContext *pCtx) : BaseT(pCtx->getModule()->getDataLayout()), ctx(pCtx) {}
 
+  DenseMap<Value *, bool> isGEPLoopInduction;
+
   bool shouldBuildLookupTables();
 
   bool isLoweredToCall(const Function *F);
@@ -40,6 +42,8 @@ class GenIntrinsicsTTIImpl : public IGCLLVM::TTIImplCRTPBase<GenIntrinsicsTTIImp
   // from variables to memory operations.
   unsigned getFlatAddressSpace();
 
+  bool isGEPLoopConstDerived(GetElementPtrInst *GEP, const Loop *L, ScalarEvolution &SE);
+
   void getUnrollingPreferences(Loop *L, ScalarEvolution &SE, TTI::UnrollingPreferences &UP,
                                OptimizationRemarkEmitter *ORE);
 

IGC/common/igc_flags.h

@@ -432,8 +432,19 @@ DECLARE_IGC_REGKEY(
 DECLARE_IGC_REGKEY(DWORD, SetLoopUnrollThreshold, 0,
                    "Set the loop unroll threshold. Value 0 will use the default threshold.", false)
 DECLARE_IGC_REGKEY(
-    DWORD, SetLoopUnrollThresholdForHighRegPressure, 0,
-    "Set the loop unroll threshold for shaders with high reg pressure. Value 0 will use the default threshold.", false)
+    DWORD, SetLoopUnrollThresholdForHighRegPressure, 200,
+    "Set the loop unroll threshold for shaders with high reg pressure.", false)
+DECLARE_IGC_REGKEY(DWORD, SetLoopUnrollMaxPercentThresholdBoostForHighRegPressure, 100,
+                   "Set the loop unroll max allowed threshold boost in percentage for shaders with high reg pressure. "
+                   "The LLVM internal value is 400. Setting 100 here because the cost model is currently based on "
+                   "latency instead of code size where the latter is what we need.",
+                   false)
+DECLARE_IGC_REGKEY(
+    bool, EnablePromoteLoopUnrollwithAlloca, true,
+    "Loop cost estimation assumes Load/Store who accesses Alloca with index deductible to loop count having 0 cost. "
+    "Disable this flag makes them always cost something as well as disables dynamic threshold increase based on the "
+    "size of alloca and number of GEP to the alloca in the loop, leading to the loop less likely to be unrolled.",
+    false)
 DECLARE_IGC_REGKEY(DWORD, SetRegisterPressureThresholdForLoopUnroll, 96,
                    "Set the register pressure threshold for limiting the loop unroll to smaller loops", false)
 DECLARE_IGC_REGKEY(DWORD, SetBranchSwapThreshold, 400, "Set the branch swaping threshold.", false)