Refactor optimization handling in new inline raytracing

Refactor optimization handling in new inline raytracing to defer modifying the function until we are done with all liveness objects.
This way, we don't invalidate liveness objects and avoid costly recalculations.

Author: jaladreips

IGC/AdaptorCommon/LivenessUtils/AllocationLivenessAnalyzer.h

@@ -30,6 +30,9 @@ class AllocationLivenessAnalyzer : public llvm::FunctionPass {
     struct Edge {
       llvm::BasicBlock *from;
       llvm::BasicBlock *to;
+
+      bool operator==(const Edge &other) const { return from == other.from && to == other.to; }
+      bool operator!=(const Edge &other) const { return !(other == *this); }
     };
 
     llvm::SmallVector<Edge> lifetimeEndEdges;

IGC/AdaptorCommon/RayTracing/NewTraceRayInlineLoweringPass.cpp

@@ -14,6 +14,8 @@ SPDX-License-Identifier: MIT
 
 #include "common/LLVMWarningsPush.hpp"
 #include <llvm/ADT/STLExtras.h>
+#include <llvm/ADT/DenseMapInfo.h>
+#include <llvm/ADT/Hashing.h>
 #include <llvm/Analysis/LoopInfo.h>
 #include <llvm/IR/Dominators.h>
 #include <llvm/IR/InstIterator.h>
@@ -25,6 +27,29 @@ SPDX-License-Identifier: MIT
 using namespace IGC;
 using namespace llvm;
 
+namespace llvm {
+template<>
+struct DenseMapInfo<IGC::AllocationLivenessAnalyzer::LivenessData::Edge> {
+  using Edge = IGC::AllocationLivenessAnalyzer::LivenessData::Edge;
+
+  static inline Edge getEmptyKey() {
+    return Edge{ DenseMapInfo<BasicBlock *>::getEmptyKey(), DenseMapInfo<BasicBlock *>::getEmptyKey() };
+  }
+
+  static inline Edge getTombstoneKey() {
+    return Edge{ DenseMapInfo<BasicBlock *>::getTombstoneKey(), DenseMapInfo<BasicBlock *>::getTombstoneKey() };
+  }
+
+  static unsigned getHashValue(const Edge &E) {
+    return (unsigned)hash_combine(E.from, E.to);
+  }
+
+  static bool isEqual(const Edge &LHS, const Edge &RHS) {
+    return LHS == RHS;
+  }
+};
+} // namespace llvm
+
 void InlineRaytracing::getAdditionalAnalysisUsage(AnalysisUsage &AU) const { AU.addRequired<CodeGenContextWrapper>(); }
 
 bool InlineRaytracing::LowerAllocations(Function &F) {
@@ -771,8 +796,7 @@ void InlineRaytracing::StopAndStartRayquery(RTBuilder &IRB, Instruction *I, Valu
   }
 }
 
-void InlineRaytracing::HandleOptimizationsAndSpills(llvm::Function &F, LivenessDataMap &livenessDataMap,
-                                                    DominatorTree &DT, LoopInfo &LI) {
+void InlineRaytracing::HandleOptimizationsAndSpills(llvm::Function &F, LivenessDataMap &livenessDataMap) {
   RTBuilder IRB(&*F.getEntryBlock().begin(), *m_pCGCtx);
 
   SmallVector<Instruction *> continuationInstructions;
@@ -800,108 +824,122 @@ void InlineRaytracing::HandleOptimizationsAndSpills(llvm::Function &F, LivenessD
       m_pCGCtx->platform.enableRayQueryThrottling(m_pCGCtx->getModuleMetaData()->compOpt.EnableDynamicRQManagement) &&
       m_numSlotsUsed == 1;
 
+  MapVector<Instruction *, SmallVector<std::function<void(RTBuilder &)>>> instructionClosures;
+  MapVector<LivenessData::Edge, SmallVector<std::function<void(RTBuilder &)>>> edgeClosures;
+
   for (auto &entry : livenessDataMap) {
-    bool cfgChanged = false;
 
     auto *rqObject = entry.first;
     auto *LD = &entry.second;
 
     // process the allocation acquire point
     // handle rayquery check
-    if (doRQCheckRelease) {
-      // check before the allocation is acquired
-      IRB.SetInsertPoint(LD->lifetimeStart);
-      IRB.CreateRayQueryCheckIntrinsic();
-    }
+    instructionClosures[LD->lifetimeStart].push_back([this, doRQCheckRelease](RTBuilder &IRB) {
+
+      if (doRQCheckRelease)
+        IRB.CreateRayQueryCheckIntrinsic();
+    });
 
     // process the allocation release points
     for (auto *I : LD->lifetimeEndInstructions) {
-      IRB.SetInsertPoint(isa<ReturnInst>(I) ? I : I->getNextNode());
 
-      auto *stackPtr = getStackPtr(IRB, IRB.Insert(rqObject->clone()));
+      instructionClosures[isa<ReturnInst>(I) ? I : I->getNextNode()].push_back(
+          [this, rqObject, doRQCheckRelease](RTBuilder &IRB) {
 
-      // handle cache control
-      InsertCacheControl(IRB, stackPtr);
+            auto *stackPtr = getStackPtr(IRB, IRB.Insert(rqObject->clone()));
 
-      // handle rayquery release
-      if (doRQCheckRelease)
-        IRB.CreateRayQueryReleaseIntrinsic();
+            // handle cache control
+            InsertCacheControl(IRB, stackPtr);
 
-      IGC_ASSERT(DT.dominates(LD->lifetimeStart, I));
+            // handle rayquery release
+            if (doRQCheckRelease)
+              IRB.CreateRayQueryReleaseIntrinsic();
+          });
     }
 
-    for (const auto &[from, to] : LD->lifetimeEndEdges) {
-      auto *succ = to;
-      // to avoid multiple executions of rayquery release instructions,
-      // we need to ensure that the "to" block has a single predecessor
-      if (!to->getSinglePredecessor()) {
-        succ = SplitEdge(from, succ);
-
-        // we invalidated other the liveness data for other instructions
-        cfgChanged = true;
-      }
-
-      IRB.SetInsertPoint(succ->getFirstNonPHI());
+    for (const auto &edge : LD->lifetimeEndEdges) {
 
-      auto *stackPtr = getStackPtr(IRB, IRB.Insert(rqObject->clone()));
+      edgeClosures[edge].push_back([this, rqObject, doRQCheckRelease](RTBuilder &IRB) {
 
-      // handle cache control
-      InsertCacheControl(IRB, stackPtr);
+        auto *stackPtr = getStackPtr(IRB, IRB.Insert(rqObject->clone()));
 
-      // handle rayquery release
-      if (doRQCheckRelease)
-        IRB.CreateRayQueryReleaseIntrinsic();
+        // handle cache control
+        InsertCacheControl(IRB, stackPtr);
 
-      IGC_ASSERT(DT.dominates(LD->lifetimeStart, succ));
+        // handle rayquery release
+        if (doRQCheckRelease)
+          IRB.CreateRayQueryReleaseIntrinsic();
+      });
     }
 
     // handle continuation instructions
     for (auto *I : continuationInstructions) {
+
       if (!LD->ContainsInstruction(*I))
         continue;
 
-      IRB.SetInsertPoint(I);
-        StopAndStartRayquery(IRB, I, IRB.Insert(rqObject->clone()), true, doRQCheckRelease);
+      instructionClosures[I].push_back([this, rqObject, doRQCheckRelease, I](RTBuilder &IRB) {
+
+          StopAndStartRayquery(IRB, I, IRB.Insert(rqObject->clone()), true, doRQCheckRelease);
+      });
     }
 
     // handle indirect calls
     for (auto *I : indirectCallInstructions) {
+
       if (!LD->ContainsInstruction(*I))
         continue;
 
-      IRB.SetInsertPoint(I);
-      StopAndStartRayquery(IRB, I, IRB.Insert(rqObject->clone()), true, doRQCheckRelease);
+      instructionClosures[I].push_back([this, rqObject, doRQCheckRelease, I](RTBuilder &IRB) {
+
+        StopAndStartRayquery(IRB, I, IRB.Insert(rqObject->clone()), true, doRQCheckRelease);
+      });
     }
 
     // handle hidden control flow instructions
     for (auto *I : hiddenCFInstructions) {
+
       if (!LD->ContainsInstruction(*I))
         continue;
 
-      IRB.SetInsertPoint(I);
-      StopAndStartRayquery(IRB, I, IRB.Insert(rqObject->clone()), true, doRQCheckRelease);
+      instructionClosures[I].push_back([this, rqObject, doRQCheckRelease, I](RTBuilder &IRB) {
+
+        StopAndStartRayquery(IRB, I, IRB.Insert(rqObject->clone()), true, doRQCheckRelease);
+      });
     }
 
     // handle barriers
     for (auto *I : barrierInstructions) {
+
       if (!LD->ContainsInstruction(*I))
         continue;
 
-      IRB.SetInsertPoint(I);
-      StopAndStartRayquery(IRB, I, IRB.Insert(rqObject->clone()), false, doRQCheckRelease);
+      instructionClosures[I].push_back([this, rqObject, doRQCheckRelease, I](RTBuilder &IRB) {
+
+        StopAndStartRayquery(IRB, I, IRB.Insert(rqObject->clone()), false, doRQCheckRelease);
+      });
     }
+  }
+
+  for (const auto [I, closures] : instructionClosures) {
 
-    if (cfgChanged) {
-      auto nextentry = livenessDataMap.find(rqObject);
+    IRB.SetInsertPoint(I);
+    for (const auto c : closures)
+      c(IRB);
+  }
 
-      // TODO: can we incrementally update LoopInfo and DomTree?
-      DT.recalculate(F);
+  for (const auto [edge, closures] : edgeClosures) {
 
-      getAnalysis<LoopInfoWrapperPass>().releaseMemory();
-      getAnalysis<LoopInfoWrapperPass>().runOnFunction(F);
-      while (++nextentry != livenessDataMap.end())
-        nextentry->second = ProcessInstruction(nextentry->first, DT, getAnalysis<LoopInfoWrapperPass>().getLoopInfo());
-    }
+    auto *succ = edge.to;
+    // to avoid multiple executions of rayquery release instructions,
+    // we need to ensure that the "to" block has a single predecessor
+    if (!edge.to->getSinglePredecessor())
+      succ = SplitEdge(edge.from, succ);
+
+    IRB.SetInsertPoint(succ->getFirstNonPHI());
+
+    for (const auto c : closures)
+      c(IRB);
   }
 }
 
@@ -971,7 +1009,7 @@ bool InlineRaytracing::runOnFunction(Function &F) {
 
   auto livenessData = AnalyzeLiveness(F, DT, LI);
   AssignSlots(F, livenessData);
-  HandleOptimizationsAndSpills(F, livenessData, DT, LI);
+  HandleOptimizationsAndSpills(F, livenessData);
   LowerSlotAssignments(F);
   LowerStackPtrs(F);
 

IGC/AdaptorCommon/RayTracing/NewTraceRayInlineLoweringPass.h

@@ -44,8 +44,7 @@ class InlineRaytracing : public AllocationLivenessAnalyzer {
   bool LowerAllocations(llvm::Function &F);
   LivenessDataMap AnalyzeLiveness(llvm::Function &F, llvm::DominatorTree &DT, llvm::LoopInfo &LI);
   void AssignSlots(llvm::Function &F, const LivenessDataMap &livenessDataMap);
-  void HandleOptimizationsAndSpills(llvm::Function &F, LivenessDataMap &livenessDataMap, llvm::DominatorTree &DT,
-                                    llvm::LoopInfo &LI);
+  void HandleOptimizationsAndSpills(llvm::Function &F, LivenessDataMap &livenessDataMap);
   void LowerSlotAssignments(llvm::Function &F);
   void LowerStackPtrs(llvm::Function &F);