[AMDGPU] Avoid unnecessary graph visits during WQM marking

[critson]

Summary

Avoid revisiting nodes with the same set of defined lanes by
using a unified visited set which integrates lanes into the key.
This retains the intent of the original code by still revisiting
a subgraph if a different set of lanes is defined and hence
marking might progress differently.

Note: default size of the visited set has been confirmed to
cover >99% of invocations in large array of test shaders.

[foad]

Looks reasonable.

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

350

Remove = VisitKey?

[critson]

• Address reviewer comments.

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

const LaneBitmask UseLanes =
                                 : LaneBitmask::getNone());
 
   // Perform a depth-first iteration of the LiveRange graph marking defs.
   // Stop processing of a given branch when all use lanes have been defined.
   // The first definition stops processing for a physical register.
   struct PhiEntry {
     const VNInfo *Phi;
     unsigned PredIdx;
-    unsigned VisitIdx;
     LaneBitmask DefinedLanes;
 
-    PhiEntry(const VNInfo *Phi, unsigned PredIdx, unsigned VisitIdx,
-             LaneBitmask DefinedLanes)
-        : Phi(Phi), PredIdx(PredIdx), VisitIdx(VisitIdx),
-          DefinedLanes(DefinedLanes) {}
+    PhiEntry(const VNInfo *Phi, unsigned PredIdx, LaneBitmask DefinedLanes)
+        : Phi(Phi), PredIdx(PredIdx), DefinedLanes(DefinedLanes) {}
   };
-  SmallSetVector<const VNInfo *, 4> Visited;
+  using VisitKey = std::pair<const VNInfo *, LaneBitmask>;
   SmallVector<PhiEntry, 2> PhiStack;
+  SmallSet<VisitKey, 4> Visited;
   LaneBitmask DefinedLanes;
-  unsigned NextPredIdx; // Only used for processing phi nodes
+  unsigned NextPredIdx = 0; // Only used for processing phi nodes
   do {
     const VNInfo *NextValue = nullptr;
+    const VisitKey Key(Value, DefinedLanes);
 
-    if (!Visited.count(Value)) {
-      Visited.insert(Value);
+    if (!Visited.count(Key)) {
+      Visited.insert(Key);
       // On first visit to a phi then start processing first predecessor
       NextPredIdx = 0;
     }
 
     if (Value->isPHIDef()) {
       // Each predecessor node in the phi must be processed as a subgraph
       const MachineBasicBlock *MBB = LIS->getMBBFromIndex(Value->def);
       assert(MBB && "Phi-def has no defining MBB");
 
       // Find next predecessor to process
       unsigned Idx = NextPredIdx;
       auto PI = MBB->pred_begin() + Idx;
       auto PE = MBB->pred_end();
       for (; PI != PE && !NextValue; ++PI, ++Idx) {
         if (const VNInfo *VN = LR.getVNInfoBefore(LIS->getMBBEndIdx(*PI))) {
-          if (!Visited.count(VN))
+          if (!Visited.count(VisitKey(VN, DefinedLanes)))
             NextValue = VN;
         }
       }
 
       // If there are more predecessors to process; add phi to stack
       if (PI != PE)
-        PhiStack.emplace_back(Value, Idx, Visited.size(), DefinedLanes);
+        PhiStack.emplace_back(Value, Idx, DefinedLanes);
     } else {
       MachineInstr *MI = LIS->getInstructionFromIndex(Value->def);
       assert(MI && "Def has no defining instruction");
 
       if (Reg.isVirtual()) {
         // Iterate over all operands to find relevant definitions
         bool HasDef = false;
         for (const MachineOperand &Op : MI->operands()) {
if (Value->isPHIDef()) {
           DefinedLanes |= OpLanes;
         }
 
         // Check if all lanes of use have been defined
         if ((DefinedLanes & UseLanes) != UseLanes) {
           // Definition not complete; need to process input value
           LiveQueryResult LRQ = LR.Query(LIS->getInstructionIndex(*MI));
           if (const VNInfo *VN = LRQ.valueIn()) {
-            if (!Visited.count(VN))
+            if (!Visited.count(VisitKey(VN, DefinedLanes)))
               NextValue = VN;
           }
         }
 
         // Only mark the instruction if it defines some part of the use
         if (HasDef)
           markInstruction(*MI, Flag, Worklist);
       } else {
         // For physical registers simply mark the defining instruction
         markInstruction(*MI, Flag, Worklist);
       }
     }
 
     if (!NextValue && !PhiStack.empty()) {
       // Reach end of chain; revert to processing last phi
       PhiEntry &Entry = PhiStack.back();
       NextValue = Entry.Phi;
       NextPredIdx = Entry.PredIdx;
       DefinedLanes = Entry.DefinedLanes;
-      // Rewind visited set to correct state
-      while (Visited.size() > Entry.VisitIdx)
-        Visited.pop_back();
       PhiStack.pop_back();
     }
 
     Value = NextValue;
   } while (Value);
 }
 
 void SIWholeQuadMode::markOperand(const MachineInstr &MI,