1 files changed, 22 insertions, 4 deletions

diff --git a/llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp b/llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp
index 137ac5466f4e..085796dee08b 100644
--- a/llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp
@@ -42,6 +42,7 @@
 #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"),
@@ -135,7 +136,8 @@ public:
   static char ID;
 
 public:
-  GCNRegBankReassign() : MachineFunctionPass(ID) {
+  GCNRegBankReassign(RegBankReassignMode Mode = RM_BOTH)
+    : MachineFunctionPass(ID), Mode(Mode) {
     initializeGCNRegBankReassignPass(*PassRegistry::getPassRegistry());
   }
 
@@ -167,6 +169,8 @@ private:
 
   LiveIntervals *LIS;
 
+  RegBankReassignMode Mode;
+
   unsigned MaxNumVGPRs;
 
   unsigned MaxNumSGPRs;
@@ -393,7 +397,11 @@ GCNRegBankReassign::analyzeInst(const MachineInstr &MI, Register Reg,
   unsigned StallCycles = 0;
   unsigned UsedBanks = 0;
 
-  if (MI.isDebugValue())
+  if (MI.isMetaInstruction())
+    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();
@@ -410,6 +418,8 @@ GCNRegBankReassign::analyzeInst(const MachineInstr &MI, Register Reg,
     // 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()) {
@@ -813,8 +823,11 @@ bool GCNRegBankReassign::runOnMachineFunction(MachineFunction &MF) {
 
   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
@@ -880,3 +893,8 @@ bool GCNRegBankReassign::runOnMachineFunction(MachineFunction &MF) {
 
   return CyclesSaved > 0;
 }
+
+MachineFunctionPass *
+llvm::createGCNRegBankReassignPass(RegBankReassignMode Mode) {
+  return new GCNRegBankReassign(Mode);
+}