[FastRA] Fix handling of bundled MIs

[pdhaliwal]

Summary

Fast register allocator skips bundled MIs, as the main assignment
loop uses MachineBasicBlock::iterator (= MachineInstrBundleIterator)
This was causing SIInsertWaitcnts to crash which expects all
instructions to have registers assigned.

This patch makes sure to set everything inside bundle to the same
assignments done on BUNDLE header.

[pdhaliwal]

Removed redundant header

[pdhaliwal]

Ping!

[qcolombet]

llvm/lib/CodeGen/RegAllocFast.cpp

1450

Could we fix allocateInstruction instead of having a somewhat parallel way of assigning registers?

[pdhaliwal]

llvm/lib/CodeGen/RegAllocFast.cpp

1450

If fixing the allocateInstruction means that for each assignment happening over bundle header, a corresponding assignment will happen for bundled instruction, then I am not sure if that is the good way. allocateInstruction has multiple assignment passes over operands. And for each operand's assignment, I would need to get uses and apply the same assignment. Is that what you are thinking of? Or Should I just move the if condition down below to the allocateInstruction?

[pdhaliwal]

llvm/lib/CodeGen/RegAllocFast.cpp

1450

Ping!

[pdhaliwal]

Ping!

[pdhaliwal]

Ping!

[pdhaliwal]

Ping!

[pdhaliwal]

Update AMDGPU barrier intrinsic

llvm/lib/CodeGen/RegAllocFast.cpp

struct LiveReg {
       }
     };
 
     using LiveRegMap = SparseSet<LiveReg>;
     /// This map contains entries for each virtual register that is currently
     /// available in a physical register.
     LiveRegMap LiveVirtRegs;
 
+    /// Stores assigned virtual registers present in the bundle MI.
+    DenseMap<Register, MCPhysReg> BundleVirtRegsMap;
+
     DenseMap<unsigned, SmallVector<MachineInstr *, 2>> LiveDbgValueMap;
     /// List of DBG_VALUE that we encountered without the vreg being assigned
     /// because they were placed after the last use of the vreg.
     DenseMap<unsigned, SmallVector<MachineInstr *, 1>> DanglingDbgValues;
 
     /// Has a bit set for every virtual register for which it was determined
     /// that it is alive across blocks.
     BitVector MayLiveAcrossBlocks;
private:
 
     void allocateBasicBlock(MachineBasicBlock &MBB);
 
     void addRegClassDefCounts(std::vector<unsigned> &RegClassDefCounts,
                               Register Reg) const;
 
     void allocateInstruction(MachineInstr &MI);
     void handleDebugValue(MachineInstr &MI);
+    void handleBundle(MachineInstr &MI);
+
     bool usePhysReg(MachineInstr &MI, MCPhysReg PhysReg);
     bool definePhysReg(MachineInstr &MI, MCPhysReg PhysReg);
     bool displacePhysReg(MachineInstr &MI, MCPhysReg PhysReg);
     void freePhysReg(MCPhysReg PhysReg);
 
     unsigned calcSpillCost(MCPhysReg PhysReg) const;
 
     LiveRegMap::iterator findLiveVirtReg(Register VirtReg) {
if (!MI.isImplicitDef()) {
                         << LRI->Reloaded << '\n');
       bool Kill = LRI->LastUse == nullptr;
       spill(SpillBefore, VirtReg, PhysReg, Kill, LRI->LiveOut);
       LRI->LastUse = nullptr;
     }
     LRI->LiveOut = false;
     LRI->Reloaded = false;
   }
+  if (MI.getOpcode() == TargetOpcode::BUNDLE) {
+    BundleVirtRegsMap[VirtReg] = PhysReg;
+  }
   markRegUsedInInstr(PhysReg);
   setPhysReg(MI, MO, PhysReg);
 }
 
 /// Allocates a register for a VirtReg use.
 void RegAllocFast::useVirtReg(MachineInstr &MI, unsigned OpNum,
                               Register VirtReg) {
   assert(VirtReg.isVirtual() && "Not a virtual register");
if (LRI->Error) {
       const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
       ArrayRef<MCPhysReg> AllocationOrder = RegClassInfo.getOrder(&RC);
       setPhysReg(MI, MO, *AllocationOrder.begin());
       return;
     }
   }
 
   LRI->LastUse = &MI;
+
+  if (MI.getOpcode() == TargetOpcode::BUNDLE) {
+    BundleVirtRegsMap[VirtReg] = LRI->PhysReg;
+  }
   markRegUsedInInstr(LRI->PhysReg);
   setPhysReg(MI, MO, LRI->PhysReg);
 }
 
 /// Changes operand OpNum in MI the refer the PhysReg, considering subregs. This
 /// may invalidate any operand pointers.  Return true if the operand kills its
 /// register.
 void RegAllocFast::setPhysReg(MachineInstr &MI, MachineOperand &MO,
void RegAllocFast::allocateInstruction(MachineInstr &MI) {
   // However we need to handle some corner cases:
   // - pre-assigned defs and uses need to be handled before the other def/use
   //   operands are processed to avoid the allocation heuristics clashing with
   //   the pre-assignment.
   // - The "free def operands" step has to come last instead of first for tied
   //   operands and early-clobbers.
 
   UsedInInstr.clear();
+  BundleVirtRegsMap.clear();
 
   // Scan for special cases; Apply pre-assigned register defs to state.
   bool HasPhysRegUse = false;
   bool HasRegMask = false;
   bool HasVRegDef = false;
   bool HasDef = false;
   bool HasEarlyClobber = false;
   bool NeedToAssignLiveThroughs = false;
if (LRI != LiveVirtRegs.end() && LRI->PhysReg) {
     DanglingDbgValues[Reg].push_back(&MI);
   }
 
   // If Reg hasn't been spilled, put this DBG_VALUE in LiveDbgValueMap so
   // that future spills of Reg will have DBG_VALUEs.
   LiveDbgValueMap[Reg].push_back(&MI);
 }
 
+void RegAllocFast::handleBundle(MachineInstr &MI) {
+  MachineBasicBlock::instr_iterator BundledMI = MI.getIterator();
+  ++BundledMI;
+  while (BundledMI->isBundledWithPred()) {
+    for (unsigned I = 0; I < BundledMI->getNumOperands(); ++I) {
+      MachineOperand &MO = BundledMI->getOperand(I);
+      if (!MO.isReg())
+        continue;
+
+      Register Reg = MO.getReg();
+      if (!Reg.isVirtual())
+        continue;
+
+      DenseMap<Register, MCPhysReg>::iterator DI;
+      DI = BundleVirtRegsMap.find(Reg);
+      assert(DI != BundleVirtRegsMap.end() && "Unassigned virtual register");
+
+      setPhysReg(MI, MO, DI->second);
+    }
+
+    ++BundledMI;
+  }
+}
+
 void RegAllocFast::allocateBasicBlock(MachineBasicBlock &MBB) {
   this->MBB = &MBB;
   LLVM_DEBUG(dbgs() << "\nAllocating " << MBB);
 
   RegUnitStates.assign(TRI->getNumRegUnits(), regFree);
   assert(LiveVirtRegs.empty() && "Mapping not cleared from last block?");
 
   for (MachineBasicBlock *Succ : MBB.successors()) {
for (MachineInstr &MI : reverse(MBB)) {
 
     // Special handling for debug values. Note that they are not allowed to
     // affect codegen of the other instructions in any way.
     if (MI.isDebugValue()) {
       handleDebugValue(MI);
       continue;
     }
 
     allocateInstruction(MI);
+
+    // Once BUNDLE header is assigned registers, same assignments need to be
+    // done for bundled MIs.
+    if (MI.getOpcode() == TargetOpcode::BUNDLE) {
+      handleBundle(MI);
+    }
   }
 
   LLVM_DEBUG(
     dbgs() << "Begin Regs:";
     dumpState()
   );
 
   // Spill all physical registers holding virtual registers now.

llvm/test/CodeGen/AMDGPU/fast-regalloc-bundles.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -mcpu=gfx902 -verify-machineinstrs -run-pass=regallocfast %s -o - | FileCheck -check-prefixes=GCN,XNACK,GCX9 %s
+
+---
+name: fast_regalloc_bundle_handling
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: vgpr_32 }
+  - { id: 1, class: vgpr_32 }
+  - { id: 2, class: vgpr_32 }
+body: |
+  bb.0:
+    ; GCN-LABEL: name: fast_regalloc_bundle_handling
+    ; GCN: renamable $vgpr0 = IMPLICIT_DEF
+    ; GCN: renamable $vgpr1 = IMPLICIT_DEF
+    ; GCN: renamable $vgpr0 = BUNDLE implicit killed renamable $vgpr0, implicit killed renamable $vgpr1, implicit $exec {
+    ; GCN:   renamable $vgpr0 = V_ADD_U32_e32 $vgpr0, $vgpr1, implicit $exec
+    ; GCN: }
+    ; GCN: S_ENDPGM 0, implicit killed renamable $vgpr0
+    %0 = IMPLICIT_DEF
+    %1 = IMPLICIT_DEF
+    %2 = BUNDLE implicit %0, implicit %1, implicit $exec {
+      %2 = V_ADD_U32_e32 %0, %1, implicit $exec
+    }
+    S_ENDPGM 0, implicit %2
+...

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ds.gws.barrier-fastregalloc.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=amdgcn-mesa-mesa3d -mcpu=gfx900 -O0 -stop-after=postrapseudos -o - -verify-machineinstrs < %s | FileCheck -enable-var-scope -check-prefix=MIR %s
+
+
+; MIR-LABEL: name: gws_barrier_offset0{{$}}
+; MIR: BUNDLE implicit{{( killed)?( renamable)?}} $vgpr0, implicit $m0, implicit $exec {
+; MIR-NEXT: DS_GWS_BARRIER renamable $vgpr0, 0, -1, implicit $m0, implicit $exec :: (load 4 from custom "GWSResource")
+; MIR-NEXT: S_WAITCNT 0
+; MIR-NEXT: }
+define amdgpu_kernel void @gws_barrier_offset0(i32 %val) #0 {
+  call void @llvm.amdgcn.ds.gws.barrier(i32 %val, i32 0)
+  ret void
+}
+
+
+declare void @llvm.amdgcn.ds.gws.barrier(i32, i32) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { convergent inaccessiblememonly nounwind }