[AMDGPU] Fix shortfalls in WQM marking

[critson]

Summary

When tracking defined lanes through phi nodes in the live range
graph each branch of the phi must be handled independently.
Also rewrite the marking algorithm to reduce unnecessary
operations.

Previously a shared set of defined lanes was used which caused
marking to stop prematurely. This was observable in existing lit
tests, but test patterns did not cover this detail.

[foad]

Just some nits inline.

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

367–368

I think you could use a SetVector here instead of a separate set and vector.

377

"predecessors"

425

PhiStack.back()

[critson]

• Switch to SmallSetVector
• Address other review comments

[piotr]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

349

I think it would be clearer if you mentioned that NextPredIdx is only needed for the phi case, as opposed to DefinedLanes for instance. Can you please add a comment or change its name to include "phi" somehow (whatever seems right for you)?

368–370

Instead of starting with Idx = 0 and skipping, can you start with Idx = NextPredIdx? That would remove the need for the "continue".

[critson]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

349

Yep, I can rename this variable.

368–370

What is happening here is we are skipping through the list of predecessor blocks to the correct number. I do not think there is a way to access a specific predecessor by index?

[critson]

• Add comment to NextPredIdx

[piotr]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

368–370

auto PI = MBB->pred_begin() + NextPredIdx; ?

[critson]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

368–370

Thanks!
I didn't know you could do that with iterators.
(You can see the limits of my C++ knowledge.)

[critson]

• Address comments

[piotr]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

365

To keep the original behaviour, you also need to adjust Idx here and assign it to NextPredIdx, because Idx is used after the for loop in line 377, right?

[critson]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

365

Yes, I noticed that immediately after pushing this diff.

[critson]

• Fix mistake in previous diff

[piotr]

LGTM

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

 
 /// Mark all relevant definitions of register \p Reg in usage \p UseMI.
 void SIWholeQuadMode::markDefs(const MachineInstr &UseMI, LiveRange &LR,
                                Register Reg, unsigned SubReg, char Flag,
                                std::vector<WorkItem> &Worklist) {
   LLVM_DEBUG(dbgs() << "markDefs " << PrintState(Flag) << ": " << UseMI);
 
   LiveQueryResult UseLRQ = LR.Query(LIS->getInstructionIndex(UseMI));
-  if (!UseLRQ.valueIn())
+  const VNInfo *Value = UseLRQ.valueIn();
+  if (!Value)
     return;
 
   // Note: this code assumes that lane masks on AMDGPU completely
   // cover registers.
+  const LaneBitmask UseLanes =
+      SubReg ? TRI->getSubRegIndexLaneMask(SubReg)
+             : (Reg.isVirtual() ? MRI->getMaxLaneMaskForVReg(Reg)
+                                : 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;
-  LaneBitmask UseLanes;
-  if (SubReg) {
-    UseLanes = TRI->getSubRegIndexLaneMask(SubReg);
-  } else if (Reg.isVirtual()) {
-    UseLanes = MRI->getMaxLaneMaskForVReg(Reg);
-  }
 
-  SmallPtrSet<const VNInfo *, 4> Visited;
-  SmallVector<const VNInfo *, 4> ToProcess;
-  ToProcess.push_back(UseLRQ.valueIn());
+    PhiEntry(const VNInfo *Phi, unsigned PredIdx, unsigned VisitIdx,
+             LaneBitmask DefinedLanes)
+        : Phi(Phi), PredIdx(PredIdx), VisitIdx(VisitIdx),
+          DefinedLanes(DefinedLanes) {}
+  };
+  SmallSetVector<const VNInfo *, 4> Visited;
+  SmallVector<PhiEntry, 2> PhiStack;
+  LaneBitmask DefinedLanes;
+  unsigned NextPredIdx; // Only used for processing phi nodes
   do {
-    const VNInfo *Value = ToProcess.pop_back_val();
+    const VNInfo *NextValue = nullptr;
+
+    if (!Visited.count(Value)) {
       Visited.insert(Value);
+      // On first visit to a phi then start processing first predecessor
+      NextPredIdx = 0;
+    }
 
     if (Value->isPHIDef()) {
-      // Need to mark all defs used in the PHI node
+      // 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");
-      for (MachineBasicBlock::const_pred_iterator PI = MBB->pred_begin(),
-                                                  PE = MBB->pred_end();
-           PI != PE; ++PI) {
+
+      // Find next predecessor to process
+      unsigned Idx = 0;
+      auto PI = MBB->pred_begin();
+      auto PE = MBB->pred_end();
+      for (; PI != PE && !NextValue; ++PI, ++Idx) {
+        if (Idx < NextPredIdx)
+          continue;
         if (const VNInfo *VN = LR.getVNInfoBefore(LIS->getMBBEndIdx(*PI))) {
           if (!Visited.count(VN))
-            ToProcess.push_back(VN);
+            NextValue = VN;
         }
       }
+
+      // If there are more predecessors to process; add phi to stack
+      if (PI != PE)
+        PhiStack.emplace_back(Value, Idx, Visited.size(), 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 (!(Op.isReg() && Op.isDef() && Op.getReg() == Reg))
             continue;
 
           // Compute lanes defined and overlap with use
           LaneBitmask OpLanes =
               Op.isUndef() ? LaneBitmask::getAll()
                            : TRI->getSubRegIndexLaneMask(Op.getSubReg());
           LaneBitmask Overlap = (UseLanes & OpLanes);
 
           // Record if this instruction defined any of use
           HasDef |= Overlap.any();
 
-          // Check if all lanes of use have been defined
+          // Mark any lanes defined
           DefinedLanes |= OpLanes;
-          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))
-                ToProcess.push_back(VN);
-            }
+        }
+
+        // Check if all lanes of use have been defined
+              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);
       }
     }
-  } while (!ToProcess.empty());
 
-  assert(!Reg.isVirtual() || ((DefinedLanes & UseLanes) == UseLanes));
+    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,
                                   const MachineOperand &Op, char Flag,
                                   std::vector<WorkItem> &Worklist) {
   assert(Op.isReg());
   Register Reg = Op.getReg();
 

llvm/test/CodeGen/AMDGPU/wqm.ll

main_body:
   %s = fadd float %a, %b
   ret float %s
 }
 
 ; CHECK-LABEL: {{^}}test_loop_vcc:
 ; CHECK-NEXT: ; %entry
 ; CHECK-NEXT: s_mov_b64 [[LIVE:s\[[0-9]+:[0-9]+\]]], exec
 ; CHECK: s_wqm_b64 exec, exec
+; CHECK: v_mov
+; CHECK: v_mov
+; CHECK: v_mov
+; CHECK: v_mov
 ; CHECK: s_and_b64 exec, exec, [[LIVE]]
 ; CHECK: image_store
 ; CHECK: s_wqm_b64 exec, exec
 ; CHECK-DAG: v_mov_b32_e32 [[CTR:v[0-9]+]], 0
 ; CHECK-DAG: s_mov_b32 [[SEVEN:s[0-9]+]], 0x40e00000
 
 ; CHECK: [[LOOPHDR:BB[0-9]+_[0-9]+]]: ; %body
 ; CHECK: v_add_f32_e32 [[CTR]], 2.0, [[CTR]]

llvm/test/CodeGen/AMDGPU/wqm.mir

bb.0:
     %2:vreg_64 = BUFFER_LOAD_DWORDX2_OFFEN %1:vgpr_32, %0:sgpr_128, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec
     %2.sub0:vreg_64 = V_SET_INACTIVE_B32 %2.sub0:vreg_64, 0, implicit $exec, implicit-def $scc
     %2.sub1:vreg_64 = V_SET_INACTIVE_B32 %2.sub1:vreg_64, 0, implicit $exec, implicit-def $scc
     %3:vreg_64 = nnan nsz arcp contract reassoc nofpexcept V_MAX_F64_e64 0, %2:vreg_64, 0, %2:vreg_64, 0, 0, implicit $mode, implicit $exec
     $vgpr0 = STRICT_WWM %3.sub0:vreg_64, implicit $exec
     $vgpr1 = STRICT_WWM %3.sub1:vreg_64, implicit $exec
     SI_RETURN_TO_EPILOG $vgpr0, $vgpr1
 ...
+
+---
+# Check that WQM marking occurs correctly through phi nodes in live range graph.
+# If not then initial V_MOV will not be in WQM.
+#
+#CHECK-LABEL: name: test_wqm_lr_phi
+#CHECK: COPY $exec
+#CHECK-NEXT: S_WQM
+#CHECK-NEXT: V_MOV_B32_e32 -10
+#CHECK-NEXT: V_MOV_B32_e32 0
+name:            test_wqm_lr_phi
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    undef %0.sub0:vreg_64 = V_MOV_B32_e32 -10, implicit $exec
+    %0.sub1:vreg_64 = V_MOV_B32_e32 0, implicit $exec
+    %1:sreg_64 = S_GETPC_B64
+    %2:sgpr_256 = S_LOAD_DWORDX8_IMM %1:sreg_64, 32, 0, 0
+
+   bb.1:
+    $vcc = V_CMP_LT_U32_e64 4, 4, implicit $exec
+    S_CBRANCH_VCCNZ %bb.3, implicit $vcc
+    S_BRANCH %bb.2
+
+   bb.2:
+    %0.sub0:vreg_64 = V_ADD_U32_e32 1, %0.sub1, implicit $exec
+    S_BRANCH %bb.3
+
+   bb.3:
+    %0.sub1:vreg_64 = V_ADD_U32_e32 1, %0.sub1, implicit $exec
+    S_BRANCH %bb.4
+
+   bb.4:
+    %3:sgpr_128 = IMPLICIT_DEF
+    %4:vreg_128 = IMAGE_SAMPLE_V4_V2 %0:vreg_64, %2:sgpr_256, %3:sgpr_128, 15, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec :: (dereferenceable load 16 from custom "ImageResource")
+    $vgpr0 = COPY %4.sub0:vreg_128
+    $vgpr1 = COPY %4.sub1:vreg_128
+    SI_RETURN_TO_EPILOG $vgpr0, $vgpr1
+...