[llvm] [AMDGPU] Refine GCNHazardRecognizer hasHazard() (PR #138841)

[Carl Ritson via llvm-commits]

================
@@ -441,42 +441,96 @@ using IsExpiredFn = function_ref<bool(const MachineInstr &, int WaitStates)>;
 using GetNumWaitStatesFn = function_ref<unsigned int(const MachineInstr &)>;
 
 // Search for a hazard in a block and its predecessors.
+// StateT must implement getHashValue() and isEqual().
 template <typename StateT>
 static bool
-hasHazard(StateT State,
+hasHazard(StateT InitialState,
           function_ref<HazardFnResult(StateT &, const MachineInstr &)> IsHazard,
           function_ref<void(StateT &, const MachineInstr &)> UpdateState,
-          const MachineBasicBlock *MBB,
-          MachineBasicBlock::const_reverse_instr_iterator I,
-          DenseSet<const MachineBasicBlock *> &Visited) {
-  for (auto E = MBB->instr_rend(); I != E; ++I) {
-    // No need to look at parent BUNDLE instructions.
-    if (I->isBundle())
-      continue;
-
-    switch (IsHazard(State, *I)) {
-    case HazardFound:
-      return true;
-    case HazardExpired:
-      return false;
-    default:
-      // Continue search
-      break;
+          const MachineBasicBlock *InitialMBB,
+          MachineBasicBlock::const_reverse_instr_iterator InitialI) {
+  SmallVector<std::pair<const MachineBasicBlock *, unsigned>> Worklist;
+  SmallDenseSet<std::pair<const MachineBasicBlock *, unsigned>> Visited;
+  SmallVector<std::pair<unsigned, unsigned>, 1> Collisions;
+  SmallDenseMap<unsigned, unsigned> StateHash2Idx;
+  SmallVector<StateT> States;
+
+  // States contains a vector of unique state structures.
+  // StateT is hashed via getHashValue() and StateHash2Idx maps each hash
+  // to an index in the States vector.
+  // In the unlikely event of a hash collision the Collision vector provides
+  // additional hash to index associations which must be retrieved by a linear
+  // scan.
+
+  // Retrieve unique constant index for a StateT structure in the States vector.
+  auto ResolveStateIdx = [&](const StateT State) {
+    unsigned StateHash = State.getHashValue();
+    unsigned StateIdx;
+    if (!StateHash2Idx.contains(StateHash)) {
+      StateIdx = States.size();
+      States.push_back(State);
+      StateHash2Idx[StateHash] = StateIdx;
+    } else {
+      StateIdx = StateHash2Idx[StateHash];
+      if (LLVM_UNLIKELY(!StateT::isEqual(State, States[StateIdx]))) {
+        // Hash collision
+        auto Collision = llvm::find_if(Collisions, [&](auto &C) {
+          return C.first == StateHash &&
+                 StateT::isEqual(State, States[C.second]);
+        });
+        if (Collision != Collisions.end()) {
+          StateIdx = Collision->second;
+        } else {
+          StateIdx = States.size();
+          States.push_back(State);
+          Collisions.emplace_back(StateHash, StateIdx);
+        }
+      }
     }
+    return StateIdx;
+  };
 
-    if (I->isInlineAsm() || I->isMetaInstruction())
-      continue;
+  const MachineBasicBlock *MBB = InitialMBB;
+  StateT State = InitialState;
+  auto I = InitialI;
 
-    UpdateState(State, *I);
-  }
+  for (;;) {
+    bool Expired = false;
+    for (auto E = MBB->instr_rend(); I != E; ++I) {
+      // No need to look at parent BUNDLE instructions.
+      if (I->isBundle())
+        continue;
 
-  for (MachineBasicBlock *Pred : MBB->predecessors()) {
-    if (!Visited.insert(Pred).second)
-      continue;
+      auto Result = IsHazard(State, *I);
+      if (Result == HazardFound)
+        return true;
+      if (Result == HazardExpired) {
+        Expired = true;
+        break;
+      }
 
-    if (hasHazard(State, IsHazard, UpdateState, Pred, Pred->instr_rbegin(),
-                  Visited))
-      return true;
+      if (I->isInlineAsm() || I->isMetaInstruction())
+        continue;
+
+      UpdateState(State, *I);
+    }
+
+    if (!Expired) {
+      unsigned StateIdx = ResolveStateIdx(State);
+      for (MachineBasicBlock *Pred : MBB->predecessors()) {
+        if (!Visited.insert(std::pair(Pred, StateIdx)).second)
+          continue;
+        Worklist.emplace_back(Pred, StateIdx);
----------------
perlfu wrote:

Done

https://github.com/llvm/llvm-project/pull/138841