[AMDGPU] Enable whole wave register copy

[cdevadas]

Summary

So far, we haven't exposed the allocation of whole-wave
registers to regalloc. We hand-picked them for various
whole wave mode operations. With a future patch, we
want the allocator to efficiently allocate them rather
than using the custom pre-allocation pass.

Any liverange split of virtual registers involved in
whole-wave operations require the resulting COPY
introduced with the split to be performed for all
lanes. It isn't implemented in the compiler yet.

This patch would identify all such copies and
manipulate the exec mask around them to enable all
lanes without affecting the value of exec mask
elsewhere.

[yassingh]

• Rebase
• Incorporated the downstream code

[yassingh]

fix file comment

[yassingh]

Rebase and merge D150390 ([AMDGPU] Introduce and use the new PRED_COPY opcode) into this.

[yassingh]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

1940–1942 ↗

(On Diff #528812)

I wanted to keep this in SILowerPredicatedCopy pass but it wasn't working there, it needed virtregRewriter to be run before it.

[yassingh]

• Revert back to PRED_COPY to COPY lowering mechanism
• Rebase

[yassingh]

• Review comments (use getWaveMaskRegClass(), getExec() functions)

[yassingh]

In D143762#4438232, @yassingh wrote:

• Review comments (use getWaveMaskRegClass(), getExec() functions)

Accidentally updated this one instead of D143759. Not reverting as the extra diff will go away once it is rebased.

[arsenm]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

8004–8007 ↗

(On Diff #533295)

Don't need to query TRI, it's already the RI member. You can also just do RI.isSGPRReg.

Also, you're not checking the flags? The point was to check the flags and switch to the WWM copy. You don't care about SGPR or VGPR, you care about WWM VGPR or not

[yassingh]

• Check vreg flag in getLiveRangeSplitOpcode

[arsenm]

llvm/lib/CodeGen/SplitKit.cpp

567 ↗

(On Diff #533516)

You're missing the new opcode here. buildSingleSubRegCopy should gain a new MCInstrDesc & parameter for the opcode to use

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

2419 ↗

(On Diff #533516)

Regular copy shouldn't reach here?

8004 ↗

(On Diff #533516)

SrcReg.isVirtual should be implied

llvm/lib/Target/AMDGPU/SILowerPredicatedCopies.cpp

140 ↗

(On Diff #533516)

Can you early return false if the WWM reg set is empty?

143 ↗

(On Diff #533516)

continue and reduce indentation

144 ↗

(On Diff #533516)

Should add a TODO to try to reduce the saveexec/restore exec pairs for adjacent WWM ops

159 ↗

(On Diff #533516)

Do these need to gain an implicit exec use?

[yassingh]

• review comments

[yassingh]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

2419 ↗

(On Diff #533516)

Yes it won't, will remove.

8004 ↗

(On Diff #533516)

Added an assert.

llvm/lib/Target/AMDGPU/SILowerPredicatedCopies.cpp

140 ↗

(On Diff #533516)

Done, had to add a new helper function.

159 ↗

(On Diff #533516)

Do you mean we should add the implicit exec? In that case, SIFixVGPRCopies will take care?

[arsenm]

llvm/lib/Target/AMDGPU/SILowerPredicatedCopies.cpp

159 ↗

(On Diff #533516)

This isn't something that can be taken care of later. SIFixVGPRCopies is a horribly broken hack, the less we depend on it the better

159 ↗

(On Diff #533516)

Actually these split copies might have been the original problem which caused it to be added. Maybe we have a way to drop it now?

[arsenm]

llvm/lib/Target/AMDGPU/SILowerPredicatedCopies.cpp

159 ↗

(On Diff #533516)

That was the reasoning given in D28874 when it was added. How about as a next step we make sure all the VGPR splits end up with exec reads

[yassingh]

llvm/lib/Target/AMDGPU/SILowerPredicatedCopies.cpp

159 ↗

(On Diff #533516)

Will look into it, don't have sufficient expertise to answer this right now. Adding @AMDGPU.

159 ↗

(On Diff #533516)

I can try and see what comes. I can see 2 ways to go about it, add the exec operand here itself or, in tablegen definition of PRED_COPY. What do you suggest?

[arsenm]

llvm/lib/Target/AMDGPU/SILowerPredicatedCopies.cpp

159 ↗

(On Diff #533516)

Add to the instruction definition. Though we're still stuck with COPY=PRED_COPY 'sometimes'

[yassingh]

Rebase

[arsenm]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

8004 ↗

(On Diff #534588)

I guess the name is wrong now. It's technically an unpredicated copy. Rename to WWM_COPY or COPY_WWM?

llvm/lib/Target/AMDGPU/SILowerPredicatedCopies.cpp

82 ↗

(On Diff #534588)

SIInstrInfo

92 ↗

(On Diff #534588)

Either this check is unnecessary or should have happened earlier to avoid crashing on the checkFlag

[cdevadas]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

8004 ↗

(On Diff #534588)

Make sense. I prefer WWM_COPY.

llvm/lib/Target/AMDGPU/SILowerPredicatedCopies.cpp

82 ↗

(On Diff #534588)

This function is no longer needed. Now the PRED_COPY is inserted only for wwm-regs after we fine-tuned AMDGPU's getLiveRangeSplitOpcode implementation.

151 ↗

(On Diff #534588)

I guess this whole check now can be entirely avoided. PRED_COPY is inserted only for wwm-regs. If we see any such copy, insert the EXEC mask manipulation unconditionally.

[yassingh]

• Rename PRED_COPY to WWM_COPY
• Review comments delete redundant isWWMCopy() check

[cdevadas]

Remove the test llvm/test/CodeGen/AMDGPU/skip-subreg-copy-from-iswwmcopy-check.mir. It is no longer needed as you have dropped isWWMCopy function.

llvm/include/llvm/CodeGen/TargetInstrInfo.h

1974 ↗

(On Diff #534849)

Did you rebase this patch after the prototype change in D150388?
This diff shouldn't be here.

llvm/lib/CodeGen/SplitKit.cpp

539 ↗

(On Diff #534849)

Ditto

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

2419 ↗

(On Diff #534849)

Still see the regular COPY here.

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

93

Remove the VRM check here. We no longer include this pass in the O0 pipeline.

125

I think this assertion is redundant. We are inserting WWM_COPY only during the live range split.

[yassingh]

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

93

This pass is being invoked in O0 too currently(added while exploring an independent lowering mechanism for PRED_COPY). Will remove that too.

[arsenm]

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

129

avoid DebugLoc copy, use const ref

[arsenm]

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

144

In a follow up could release the reserved register

[yassingh]

Review commnets and remove pass from O0.

[yassingh]

Reupload

[yassingh]

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

144

I don't understand, since sgpr allocation is already done do we have to free the reserved sgprs? Also, can you refer me to some other reserved registers to see how do I release it.

[yassingh]

rebase

[yassingh]

re-upload

[arsenm]

LGTM with nit

llvm/lib/Target/AMDGPU/SIInstructions.td

3303 ↗

(On Diff #535430)

This should not be AMDGPUGenericInstruction, it's not a generic instruction. Should subclass from AMDGPUInst, and move to the other non-generic pseudos. AMDGPUGenericInstruction is for the G_AMDGPU_ globalisel pre-isel pseudos

[arsenm]

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

120

Typos Club adjancent

135

Missing newline

llvm/test/CodeGen/AMDGPU/llc-pipeline.ll

368 ↗

(On Diff #535430)

I think I lost track of what was going on with -O0, did this lose the -O0 run?

[arsenm]

Also not sure how this has no test changes

[arsenm]

llvm/lib/Target/AMDGPU/SIInstructions.td

3310 ↗

(On Diff #535430)

Should be marked convergent

[yassingh]

Review comments

[yassingh]

In D143762#4456567, @arsenm wrote:

Also not sure how this has no test changes

This patch was originally submitted with 0 tests. My understating is we won't see any effects of these changes until WWM copies are inserted in code by the next patch(Spill SGPRs to virtual VGPRs).
cc @cdevadas

llvm/test/CodeGen/AMDGPU/llc-pipeline.ll

368 ↗

(On Diff #535430)

It was introduced when we moved from PRED_COPY "simplification" to proper lowering mechanism. Was removed again after the approach changed back to lowering WWM_COPY to COPY first. I have added the O0 invocation back but it's not doing anything.

[cdevadas]

In D143762#4471029, @yassingh wrote:

In D143762#4456567, @arsenm wrote:

Also not sure how this has no test changes

This patch was originally submitted with 0 tests. My understating is we won't see any effects of these changes until WWM copies are inserted in code by the next patch(Spill SGPRs to virtual VGPRs).
cc @cdevadas

Yes. This is only a pre-patch for D124196 where the actual test was included, llvm/test/CodeGen/AMDGPU/whole-wave-register-copy.ll

[arsenm]

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

132

Back to back debug printing prints same instruction, just remove the second one?

[yassingh]

Removed extra debug statement

[cdevadas]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

1356 ↗

(On Diff #537120)

I'm still not convinced why this is needed in the -O0 flow?
By now, the VGPR allocation is done in the -O0 flow, and we no longer have any virtual registers. This pass act on virtual registers to see if wwm copies needed exec manipulation.

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

90

If this pass is enabled for for -O0 as well, put back the !VRM check to ensure an early return from here for -O0 flow.
And may be a comment too?

[yassingh]

Update instruction class of WWM_COPY

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

1356 ↗

(On Diff #537120)

It's conceptually needed and it's an implementation detail of current regalloc fast that these aren't introduced. Plus I think in general we should have other WWM copies for general WWM support in the future

[cdevadas]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

1356 ↗

(On Diff #537120)

ok

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

111

Also, do an early return if MRI.getNumVirtRegs() is zero. This would avoid iterating the whole function when there are no virtual registers at all (would take care of the -O0 path when physRegs are already assigned).

[yassingh]

Rebase before merge

[yassingh]

Patch relanded with 4d42e8b5d1fa87e49768d100dd1bc53515391e89

llvm/lib/Target/AMDGPU/SILowerWWMCopies.cpp

 //===-- SILowerWWMCopies.cpp - Lower Copies after regalloc ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file
-/// Lowering the WWM_COPY instructions for various register classes. 
+/// Lowering the WWM_COPY instructions for various register classes.
 /// AMDGPU target generates WWM_COPY instruction to differentiate WWM
 /// copy from COPY. This pass generates the necessary exec mask manipulation
 /// instructions to replicate 'Whole Wave Mode' and lowers WWM_COPY back to
 /// COPY.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
public:
   static char ID;
 
   SILowerWWMCopies() : MachineFunctionPass(ID) {
     initializeSILowerWWMCopiesPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
-  StringRef getPassName() const override {
-    return "SI Lower WWM Copies";
-  }
+  StringRef getPassName() const override { return "SI Lower WWM Copies"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
 
 private:
   bool isSCCLiveAtMI(const MachineInstr &MI);
   void addToWWMSpills(MachineFunction &MF, Register Reg);
 
   LiveIntervals *LIS;
   SlotIndexes *Indexes;
   VirtRegMap *VRM;
   const SIRegisterInfo *TRI;
   const MachineRegisterInfo *MRI;
   SIMachineFunctionInfo *MFI;
 };
 
 } // End anonymous namespace.
 
-INITIALIZE_PASS_BEGIN(SILowerWWMCopies, DEBUG_TYPE,
-                      "SI Lower WWM Copies", false, false)
+INITIALIZE_PASS_BEGIN(SILowerWWMCopies, DEBUG_TYPE, "SI Lower WWM Copies",
+                      false, false)
 INITIALIZE_PASS_DEPENDENCY(LiveIntervals)
 INITIALIZE_PASS_DEPENDENCY(VirtRegMap)
-INITIALIZE_PASS_END(SILowerWWMCopies, DEBUG_TYPE,
-                    "SI Lower WWM Copies", false, false)
+INITIALIZE_PASS_END(SILowerWWMCopies, DEBUG_TYPE, "SI Lower WWM Copies", false,
+                    false)
 
 char SILowerWWMCopies::ID = 0;
 
 char &llvm::SILowerWWMCopiesID = SILowerWWMCopies::ID;
 
 bool SILowerWWMCopies::isSCCLiveAtMI(const MachineInstr &MI) {
   // We can't determine the liveness info if LIS isn't available. Early return
   // in that case and always assume SCC is live.
   if (!LIS)
     return true;
 
   LiveRange &LR =
       LIS->getRegUnit(*MCRegUnitIterator(MCRegister::from(AMDGPU::SCC), TRI));
   SlotIndex Idx = LIS->getInstructionIndex(MI);
   return LR.liveAt(Idx);
 }
 
 // If \p Reg is assigned with a physical VGPR, add the latter into wwm-spills
 // for preserving its entire lanes at function prolog/epilog.
-void SILowerWWMCopies::addToWWMSpills(MachineFunction &MF,
-                                             Register Reg) {
+void SILowerWWMCopies::addToWWMSpills(MachineFunction &MF, Register Reg) {
   if (Reg.isPhysical())
     return;
 
   Register PhysReg = VRM->getPhys(Reg);
   assert(PhysReg != VirtRegMap::NO_PHYS_REG &&
          "should have allocated a physical register");
 
   MFI->allocateWWMSpill(MF, PhysReg);
 }
 
 bool SILowerWWMCopies::runOnMachineFunction(MachineFunction &MF) {
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIInstrInfo *TII = ST.getInstrInfo();
 
   MFI = MF.getInfo<SIMachineFunctionInfo>();
   LIS = getAnalysisIfAvailable<LiveIntervals>();
   Indexes = getAnalysisIfAvailable<SlotIndexes>();
   VRM = getAnalysisIfAvailable<VirtRegMap>();
   TRI = ST.getRegisterInfo();
   MRI = &MF.getRegInfo();
 
   if (!MFI->hasVRegFlags())
     return false;
 
   bool Changed = false;
   for (MachineBasicBlock &MBB : MF) {
     for (MachineInstr &MI : MBB) {
       if (MI.getOpcode() != AMDGPU::WWM_COPY)
         continue;
 
       // TODO: Club adjancent WWM ops between same exec save/restore
       assert(TII->isVGPRCopy(MI));
 
       // For WWM vector copies, manipulate the exec mask around the copy
       // instruction.
       const DebugLoc &DL = MI.getDebugLoc();
       MachineBasicBlock::iterator InsertPt = MI.getIterator();
       Register RegForExecCopy = MFI->getSGPRForEXECCopy();
       TII->insertScratchExecCopy(MF, MBB, InsertPt, DL, RegForExecCopy,
                                  isSCCLiveAtMI(MI), Indexes);
       TII->restoreExec(MF, MBB, ++InsertPt, DL, RegForExecCopy, Indexes);
       addToWWMSpills(MF, MI.getOperand(0).getReg());
       LLVM_DEBUG(dbgs() << "WWM copy manipulation for " << MI);
 
       // Lower WWM_COPY back to COPY
       LLVM_DEBUG(dbgs() << "Updated " << MI << " to use COPY opcode");
       MI.setDesc(TII->get(AMDGPU::COPY));
       Changed |= true;
     }
   }
 
   return Changed;
 }