[AMDGPU] WQM/WWM: Fix marking of partial definitions

[critson]

Summary

Track lanes when processing definitions for marking WQM/WWM.
If all lanes have been defined then marking can stop.
This prevents marking unnecessary instructions as WQM/WWM.

In particular this fixes a bug where values passing through
V_SET_INACTIVE would me marked as requiring WWM.

[foad]

Typo "me" in description.

llvm/test/CodeGen/AMDGPU/wqm.mir

214

Remove most of this? Usually you only need name, tracksRegLiveness and body.

[critson]

• Tidy test.

[arsenm]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

265

You lost this assert (but we should convert it to the pass required properties set anyway)

349

Should this skip all implicit operands?

[critson]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

265

Yes, while I think I deleted the assertion by accident, I believe it is technically wrong.
I will add clearedProperties to the pass in D94746 which is where !IsSSA is definitely required.

349

No we must process implicit uses as well and mark them as WWM/WQM with $vcc being the primary use case, e.g. $vcc input of WWM VOP2 v_addc_co must also be in computed in WWM.
At the moment $exec is the only register I am confident that can be ignored as it explicitly managed by this pass; however, it is possible other register (e.g. $m0 and $mode) might also be valid here, but I'd prefer to leave that as exercise for later.

[piotr]

Looks good to me (with two nits inline).

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

311

DefLanes is a bit too similar to DefinedLanes. Can you rename one of them please?

317

HasDef |= Overlap.any();

[critson]

• Rebase
• Address reviewer comments

[critson]

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

311

Changed to OpLanes

llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp

private:
   SmallVector<MachineInstr *, 4> LowerToCopyInstrs;
 
   void printInfo();
 
   void markInstruction(MachineInstr &MI, char Flag,
                        std::vector<WorkItem> &Worklist);
   void markDefs(const MachineInstr &UseMI, LiveRange &LR, Register Reg,
                 unsigned SubReg, char Flag, std::vector<WorkItem> &Worklist);
+  void markOperand(const MachineInstr &MI, const MachineOperand &Op, char Flag,
+                   std::vector<WorkItem> &Worklist);
   void markInstructionUses(const MachineInstr &MI, char Flag,
                            std::vector<WorkItem> &Worklist);
   char scanInstructions(MachineFunction &MF, std::vector<WorkItem> &Worklist);
   void propagateInstruction(MachineInstr &MI, std::vector<WorkItem> &Worklist);
   void propagateBlock(MachineBasicBlock &MBB, std::vector<WorkItem> &Worklist);
   char analyzeFunction(MachineFunction &MF);
 
   MachineBasicBlock::iterator saveSCC(MachineBasicBlock &MBB,

 #endif
 
 void SIWholeQuadMode::markInstruction(MachineInstr &MI, char Flag,
                                       std::vector<WorkItem> &Worklist) {
   InstrInfo &II = Instructions[&MI];
 
   assert(!(Flag & StateExact) && Flag != 0);
 
-  LLVM_DEBUG(dbgs() << "markInstruction " << PrintState(Flag) << ": " << MI);
-
   // Remove any disabled states from the flag. The user that required it gets
   // an undefined value in the helper lanes. For example, this can happen if
   // the result of an atomic is used by instruction that requires WQM, where
   // ignoring the request for WQM is correct as per the relevant specs.
   Flag &= ~II.Disabled;
 
   // Ignore if the flag is already encompassed by the existing needs, or we
   // just disabled everything.
   if ((II.Needs & Flag) == Flag)
     return;
 
+  LLVM_DEBUG(dbgs() << "markInstruction " << PrintState(Flag) << ": " << MI);
   II.Needs |= Flag;
   Worklist.push_back(&MI);
 }
 
 /// 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) {
-  assert(!MRI->isSSA());
-
-  LLVM_DEBUG(dbgs() << "markDefs " << PrintState(Flag) << ": " << UseMI);
-
   LiveQueryResult UseLRQ = LR.Query(LIS->getInstructionIndex(UseMI));
   if (!UseLRQ.valueIn())
     return;
 
+  // Note: this code assumes that lane masks on AMDGPU completely
+  // cover registers.
+  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());
   do {
     const VNInfo *Value = ToProcess.pop_back_val();
     Visited.insert(Value);
 
     if (Value->isPHIDef()) {
       // Need to mark all defs used in the PHI node
       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) {
         if (const VNInfo *VN = LR.getVNInfoBefore(LIS->getMBBEndIdx(*PI))) {
           if (!Visited.count(VN))
             ToProcess.push_back(VN);
         }
       }
     } else {
       MachineInstr *MI = LIS->getInstructionFromIndex(Value->def);
       assert(MI && "Def has no defining instruction");
-      markInstruction(*MI, Flag, Worklist);
 
+      if (Reg.isVirtual()) {
         // Iterate over all operands to find relevant definitions
+        bool HasDef = false;
         for (const MachineOperand &Op : MI->operands()) {
-        if (!(Op.isReg() && Op.getReg() == Reg))
+          if (!(Op.isReg() && Op.isDef() && Op.getReg() == Reg))
             continue;
 
-        // Does this def cover whole register?
-        bool DefinesFullReg =
-            Op.isUndef() || !Op.getSubReg() || Op.getSubReg() == SubReg;
-        if (!DefinesFullReg) {
-          // Partial definition; need to follow and mark input value
+          // Compute lanes defined and overlap with use
+          LaneBitmask DefLanes =
+              Op.isUndef() ? LaneBitmask::getAll()
+                           : TRI->getSubRegIndexLaneMask(Op.getSubReg());
+          LaneBitmask Overlap = (UseLanes & DefLanes);
+
+          // Record if this instruction defined any of use
+          HasDef = HasDef | Overlap.any();
+
+          // Check if all lanes of use have been defined
+          DefinedLanes |= DefLanes;
+          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);
             }
           }
         }
+        // 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));
 }
 
-/// Mark all instructions defining the uses in \p MI with \p Flag.
-void SIWholeQuadMode::markInstructionUses(const MachineInstr &MI, char Flag,
-                                          std::vector<WorkItem> &Worklist) {
+void SIWholeQuadMode::markOperand(const MachineInstr &MI,
+                                  const MachineOperand &Op, char Flag,
+  assert(Op.isReg());
+  Register Reg = Op.getReg();
 
-  LLVM_DEBUG(dbgs() << "markInstructionUses " << PrintState(Flag) << ": "
-                    << MI);
-
-  for (const MachineOperand &Use : MI.uses()) {
-    if (!Use.isReg() || !Use.isUse())
-      continue;
-
-    Register Reg = Use.getReg();
+  // Ignore some hardware registers
+  switch (Reg) {
+  case AMDGPU::EXEC:
+  case AMDGPU::EXEC_LO:
+    return;
+  default:
+    break;
+  }
 
+  LLVM_DEBUG(dbgs() << "markOperand " << PrintState(Flag) << ": " << Op
+                    << " for " << MI);
+  if (Reg.isVirtual()) {
+    LiveRange &LR = LIS->getInterval(Reg);
+    markDefs(MI, LR, Reg, Op.getSubReg(), Flag, Worklist);
+  } else {
     // Handle physical registers that we need to track; this is mostly relevant
     // for VCC, which can appear as the (implicit) input of a uniform branch,
     // e.g. when a loop counter is stored in a VGPR.
-    if (!Reg.isVirtual()) {
-      if (Reg == AMDGPU::EXEC || Reg == AMDGPU::EXEC_LO)
-        continue;
-
-      for (MCRegUnitIterator RegUnit(Reg.asMCReg(), TRI); RegUnit.isValid();
-           ++RegUnit) {
-        LiveRange &LR = LIS->getRegUnit(*RegUnit);
-        const VNInfo *Value = LR.Query(LIS->getInstructionIndex(MI)).valueIn();
-        if (!Value)
-          continue;
 
-        if (MRI->isSSA()) {
-          // Since we're in machine SSA, we do not need to track physical
-          // registers across basic blocks.
-          if (Value->isPHIDef())
-            continue;
-          markInstruction(*LIS->getInstructionFromIndex(Value->def), Flag,
-                          Worklist);
-        } else {
-          markDefs(MI, LR, *RegUnit, AMDGPU::NoSubRegister, Flag, Worklist);
-        }
-      }
-
-      continue;
-    }
 
-    if (MRI->isSSA()) {
-      for (MachineInstr &DefMI : MRI->def_instructions(Use.getReg()))
-        markInstruction(DefMI, Flag, Worklist);
-    } else {
-      LiveRange &LR = LIS->getInterval(Reg);
-      markDefs(MI, LR, Reg, Use.getSubReg(), Flag, Worklist);
-    }
+/// Mark all instructions defining the uses in \p MI with \p Flag.
+void SIWholeQuadMode::markInstructionUses(const MachineInstr &MI, char Flag,
+                                          std::vector<WorkItem> &Worklist) {
+
+  LLVM_DEBUG(dbgs() << "markInstructionUses " << PrintState(Flag) << ": "
+                    << MI);
+
+  for (const MachineOperand &Use : MI.uses()) {
+    if (!Use.isReg() || !Use.isUse())
+      continue;
+    markOperand(MI, Use, Flag, Worklist);
   }
 }
 
 // Scan instructions to determine which ones require an Exact execmask and
 // which ones seed WQM requirements.
 char SIWholeQuadMode::scanInstructions(MachineFunction &MF,
                                        std::vector<WorkItem> &Worklist) {
   char GlobalFlags = 0;
for (auto II = MBB.begin(), IE = MBB.end(); II != IE; ++II) {
       } else if (Opcode == AMDGPU::V_SET_INACTIVE_B32 ||
                  Opcode == AMDGPU::V_SET_INACTIVE_B64) {
         III.Disabled = StateWWM;
         MachineOperand &Inactive = MI.getOperand(2);
         if (Inactive.isReg()) {
           if (Inactive.isUndef()) {
             LowerToCopyInstrs.push_back(&MI);
           } else {
-            Register Reg = Inactive.getReg();
-            if (Reg.isVirtual()) {
-              for (MachineInstr &DefMI : MRI->def_instructions(Reg))
-                markInstruction(DefMI, StateWWM, Worklist);
-            }
+            markOperand(MI, Inactive, StateWWM, Worklist);
           }
         }
         SetInactiveInstrs.push_back(&MI);
         continue;
       } else if (TII->isDisableWQM(MI)) {
         BBI.Needs |= StateExact;
         if (!(BBI.InNeeds & StateExact)) {
           BBI.InNeeds |= StateExact;

llvm/test/CodeGen/AMDGPU/wqm.mir

body:             |
 
   bb.2:
     $vgpr0 = COPY %4.sub0:vreg_128
     $vgpr1 = COPY %4.sub1:vreg_128
     $vgpr2 = COPY %9.sub0:vreg_128
     $vgpr3 = COPY %9.sub1:vreg_128
     SI_RETURN_TO_EPILOG $vgpr0, $vgpr1, $vgpr2, $vgpr3
 ...
+
+---
+# Check that unnecessary instruction do not get marked for WWM
+#
+#CHECK-NOT: ENTER_WWM
+#CHECK: BUFFER_LOAD_DWORDX2
+#CHECK-NOT: ENTER_WWM
+#CHECK: V_SET_INACTIVE_B32
+#CHECK: V_SET_INACTIVE_B32
+#CHECK: ENTER_WWM
+#CHECK: V_MAX
+name:            test_wwm_set_inactive_propagation
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $sgpr0_sgpr1_sgpr2_sgpr3, $vgpr0
+    %0:sgpr_128 = COPY $sgpr0_sgpr1_sgpr2_sgpr3
+    %1:vgpr_32 = COPY $vgpr0
+    %2:vreg_64 = BUFFER_LOAD_DWORDX2_OFFEN %1:vgpr_32, %0:sgpr_128, 0, 0, 0, 0, 0, 0, 0, implicit $exec
+    %2.sub0:vreg_64 = V_SET_INACTIVE_B32 %2.sub0:vreg_64, 0, implicit $exec
+    %2.sub1:vreg_64 = V_SET_INACTIVE_B32 %2.sub1:vreg_64, 0, implicit $exec
+    %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 = WWM %3.sub0:vreg_64, implicit $exec
+    $vgpr1 = WWM %3.sub1:vreg_64, implicit $exec
+    SI_RETURN_TO_EPILOG $vgpr0, $vgpr1
+...