[AMDGPU] Split GCNRegBankReassign

[rampitec]

Summary

Allow pass to work separately with SGPR, VGPR registers or both.
This is NFC now but will be needed to split RA for separate
SGPR and VGPR passes.

[arsenm]

Can't this just skip virtual registers? Does it really need to know which mode its in?

[rampitec]

In D100063#2674983, @arsenm wrote:

Can't this just skip virtual registers? Does it really need to know which mode its in?

All registers are virtual here. It will skip a register if there is no assignment in VRM (and obviously skip physregs).
This is slower as it is though. I can add the check "Reg.isPhysical() || !VRM->isAssignedReg(Reg)" earlier to mitigate it.
Will that help?

[arsenm]

In D100063#2674987, @rampitec wrote:

In D100063#2674983, @arsenm wrote:

Can't this just skip virtual registers? Does it really need to know which mode its in?

All registers are virtual here. It will skip a register if there is no assignment in VRM (and obviously skip physregs).
This is slower as it is though. I can add the check "Reg.isPhysical() || !VRM->isAssignedReg(Reg)" earlier to mitigate it.
Will that help?

Why is that slower?

I also would prefer with the allocation split that the final rewrite of registers was done at once at the end. It does commit the intermediate state now, although I don't remember why I did it that way

[rampitec]

In D100063#2675023, @arsenm wrote:

In D100063#2674987, @rampitec wrote:

In D100063#2674983, @arsenm wrote:

Can't this just skip virtual registers? Does it really need to know which mode its in?

All registers are virtual here. It will skip a register if there is no assignment in VRM (and obviously skip physregs).
This is slower as it is though. I can add the check "Reg.isPhysical() || !VRM->isAssignedReg(Reg)" earlier to mitigate it.
Will that help?

Why is that slower?

Simply because I do that check late. It does not matter much until you split the RA. I can move the check earlier, that is not a big deal.

I also would prefer with the allocation split that the final rewrite of registers was done at once at the end. It does commit the intermediate state now, although I don't remember why I did it that way

If you will run rewriter once no changes are needed at all. Just run the pass right before the rewriter. Although currently you are running it twice. What are your plans about this?

[arsenm]

In D100063#2675053, @rampitec wrote:

If you will run rewriter once no changes are needed at all. Just run the pass right before the rewriter. Although currently you are running it twice. What are your plans about this?

I would need to look at it again. For now I guess assume the rewrite will run twice

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

400

Unrelated, but this is likely a bug. This should cover all meta instructions (most likely this will do the wrong thing for kill)

[rampitec]

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

400

Ack, will fix it separately.

llvm/lib/Target/AMDGPU/AMDGPU.h

 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/CodeGen.h"
 
 namespace llvm {
 
 class FunctionPass;
 class GCNTargetMachine;
 class ImmutablePass;
+class MachineFunctionPass;
 class ModulePass;
 class Pass;
 class Target;
 class TargetMachine;
 class TargetOptions;
 class PassRegistry;
 class Module;
 

 FunctionPass *createAMDGPULateCodeGenPreparePass();
 FunctionPass *createAMDGPUMachineCFGStructurizerPass();
 FunctionPass *createAMDGPUPropagateAttributesEarlyPass(const TargetMachine *);
 ModulePass *createAMDGPUPropagateAttributesLatePass(const TargetMachine *);
 FunctionPass *createAMDGPURewriteOutArgumentsPass();
 ModulePass *createAMDGPULowerModuleLDSPass();
 FunctionPass *createSIModeRegisterPass();
 
+namespace AMDGPU {
+enum RegBankReassignMode {
+  RM_VGPR = 1,
+  RM_SGPR = 2,
+  RM_BOTH = RM_VGPR | RM_SGPR
+};
+}
+MachineFunctionPass *
+createGCNRegBankReassignPass(AMDGPU::RegBankReassignMode Mode);
+
 struct AMDGPUSimplifyLibCallsPass : PassInfoMixin<AMDGPUSimplifyLibCallsPass> {
   AMDGPUSimplifyLibCallsPass(TargetMachine &TM) : TM(TM) {}
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 
 private:
   TargetMachine &TM;
 };
 

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

if (EnableDCEInRA)
     insertPass(&DetectDeadLanesID, &DeadMachineInstructionElimID);
 
   TargetPassConfig::addOptimizedRegAlloc();
 }
 
 bool GCNPassConfig::addPreRewrite() {
   if (EnableRegReassign) {
     addPass(&GCNNSAReassignID);
-    addPass(&GCNRegBankReassignID);
+    addPass(createGCNRegBankReassignPass(AMDGPU::RM_BOTH));
   }
   return true;
 }
 
 void GCNPassConfig::addPostRegAlloc() {
   addPass(&SIFixVGPRCopiesID);
   if (getOptLevel() > CodeGenOpt::None)
     addPass(&SIOptimizeExecMaskingID);

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/InitializePasses.h"
 
 using namespace llvm;
+using namespace AMDGPU;
 
 static cl::opt<unsigned> VerifyStallCycles("amdgpu-verify-regbanks-reassign",
   cl::desc("Verify stall cycles in the regbanks reassign pass"),
   cl::value_desc("0|1|2"),
   cl::init(0), cl::Hidden);
 
 // Threshold to keep compile time reasonable.
 static cl::opt<unsigned> VRegThresh("amdgpu-regbanks-reassign-threshold",
#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
     }
 #endif
   };
 
 public:
   static char ID;
 
 public:
-  GCNRegBankReassign() : MachineFunctionPass(ID) {
+  GCNRegBankReassign(RegBankReassignMode Mode = RM_BOTH)
+    : MachineFunctionPass(ID), Mode(Mode) {
     initializeGCNRegBankReassignPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   StringRef getPassName() const override { return "GCN RegBank Reassign"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
private:
   MachineLoopInfo *MLI;
 
   VirtRegMap *VRM;
 
   LiveRegMatrix *LRM;
 
   LiveIntervals *LIS;
 
+  RegBankReassignMode Mode;
+
   unsigned MaxNumVGPRs;
 
   unsigned MaxNumSGPRs;
 
   BitVector RegsUsed;
 
   SmallVector<OperandMask, 8> OperandMasks;
 

 }
 
 std::pair<unsigned, unsigned>
 GCNRegBankReassign::analyzeInst(const MachineInstr &MI, Register Reg,
                                 unsigned SubReg, int Bank) {
   unsigned StallCycles = 0;
   unsigned UsedBanks = 0;
 
   if (MI.isDebugValue())
     return std::make_pair(StallCycles, UsedBanks);
 
+  if (!(Mode & RM_SGPR) &&
+      MI.getDesc().TSFlags & (SIInstrFlags::SMRD | SIInstrFlags::SALU))
+    return std::make_pair(StallCycles, UsedBanks);
+
   RegsUsed.reset();
   OperandMasks.clear();
   for (const auto& Op : MI.explicit_uses()) {
     // Undef can be assigned to any register, so two vregs can be assigned
     // the same phys reg within the same instruction.
     if (!Op.isReg() || Op.isUndef())
       continue;
 
     const Register R = Op.getReg();
     const TargetRegisterClass *RC = TRI->getRegClassForReg(*MRI, R);
 
     // Do not compute stalls for AGPRs
     if (TRI->hasAGPRs(RC))
       continue;
+    if ((Mode != RM_BOTH) && !(Mode & (TRI->hasVGPRs(RC) ? RM_VGPR : RM_SGPR)))
+      continue;
 
     // Do not compute stalls if sub-register covers all banks
     if (Op.getSubReg()) {
       LaneBitmask LM = TRI->getSubRegIndexLaneMask(Op.getSubReg());
       if (TRI->hasVGPRs(RC)) {
         if (TRI->getNumCoveredRegs(LM) >= NUM_VGPR_BANKS)
           continue;
       } else {

 
 bool GCNRegBankReassign::runOnMachineFunction(MachineFunction &MF) {
   ST = &MF.getSubtarget<GCNSubtarget>();
   if (!ST->hasRegisterBanking() || skipFunction(MF.getFunction()))
     return false;
 
   MRI = &MF.getRegInfo();
 
-  LLVM_DEBUG(dbgs() << "=== RegBanks reassign analysis on function " << MF.getName()
-                    << "\nNumVirtRegs = " << MRI->getNumVirtRegs() << "\n\n");
+  LLVM_DEBUG(dbgs() << "=== RegBanks reassign analysis on function "
+                    << MF.getName() << '\n'
+                    << ((Mode & RM_VGPR) ? "VGPR " : "")
+                    << ((Mode & RM_SGPR) ? "SGPR " : "") << "mode\n"
+                    << "NumVirtRegs = " << MRI->getNumVirtRegs() << "\n\n");
 
   if (MRI->getNumVirtRegs() > VRegThresh) {
     LLVM_DEBUG(dbgs() << "NumVirtRegs > " << VRegThresh
                       << " threshold, skipping function.\n\n");
     return false;
   }
 
   TRI = ST->getRegisterInfo();
bool GCNRegBankReassign::runOnMachineFunction(MachineFunction &MF) {
 
   if (VerifyStallCycles == 1 && !verifyCycles(MF, StallCycles, CyclesSaved))
     report_fatal_error("RegBank reassign stall cycles verification failed.");
 
   RegsUsed.clear();
 
   return CyclesSaved > 0;
 }
+
+MachineFunctionPass *
+llvm::createGCNRegBankReassignPass(RegBankReassignMode Mode) {
+  return new GCNRegBankReassign(Mode);
+}