Index: lib/Target/AMDGPU/SIMachineScheduler.cpp
===================================================================
--- lib/Target/AMDGPU/SIMachineScheduler.cpp
+++ lib/Target/AMDGPU/SIMachineScheduler.cpp
@@ -1657,8 +1657,8 @@
   SITII = static_cast<const SIInstrInfo*>(TII);
   SITRI = static_cast<const SIRegisterInfo*>(TRI);
 
-  VGPRSetID = SITRI->getVGPR32PressureSet();
-  SGPRSetID = SITRI->getSGPR32PressureSet();
+  VGPRSetID = SITRI->getVGPRPressureSet();
+  SGPRSetID = SITRI->getSGPRPressureSet();
 }
 
 SIScheduleDAGMI::~SIScheduleDAGMI() {
Index: lib/Target/AMDGPU/SIRegisterInfo.h
===================================================================
--- lib/Target/AMDGPU/SIRegisterInfo.h
+++ lib/Target/AMDGPU/SIRegisterInfo.h
@@ -25,8 +25,8 @@
 
 struct SIRegisterInfo final : public AMDGPURegisterInfo {
 private:
-  unsigned SGPR32SetID;
-  unsigned VGPR32SetID;
+  unsigned SGPRSetID;
+  unsigned VGPRSetID;
   BitVector SGPRPressureSets;
   BitVector VGPRPressureSets;
 
@@ -182,11 +182,24 @@
                               const TargetRegisterClass *RC,
                               const MachineFunction &MF) const;
 
-  unsigned getSGPR32PressureSet() const { return SGPR32SetID; };
-  unsigned getVGPR32PressureSet() const { return VGPR32SetID; };
+  unsigned getSGPRPressureSet() const { return SGPRSetID; };
+  unsigned getVGPRPressureSet() const { return VGPRSetID; };
 
   bool isVGPR(const MachineRegisterInfo &MRI, unsigned Reg) const;
 
+  bool isSGPRPressureSet(unsigned SetID) const {
+    return SGPRPressureSets.test(SetID) && !VGPRPressureSets.test(SetID);
+  }
+  bool isVGPRPressureSet(unsigned SetID) const {
+    return VGPRPressureSets.test(SetID) && !SGPRPressureSets.test(SetID);
+  }
+
+  std::pair<unsigned, unsigned>
+  getRegPressureSetLimits(const MachineFunction &MF) const;
+
+  std::pair<unsigned, unsigned>
+  getMaxRegPressure(ArrayRef<unsigned> SetPressure) const;
+
 private:
   void buildScratchLoadStore(MachineBasicBlock::iterator MI,
                              unsigned LoadStoreOp, const MachineOperand *SrcDst,
Index: lib/Target/AMDGPU/SIRegisterInfo.cpp
===================================================================
--- lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -95,19 +95,40 @@
                                    VGPRPressureSets(getNumRegPressureSets()) {
   unsigned NumRegPressureSets = getNumRegPressureSets();
 
-  SGPR32SetID = NumRegPressureSets;
-  VGPR32SetID = NumRegPressureSets;
-  for (unsigned i = 0; i < NumRegPressureSets; ++i) {
-    if (strncmp("SGPR_32", getRegPressureSetName(i), 7) == 0)
-      SGPR32SetID = i;
-    else if (strncmp("VGPR_32", getRegPressureSetName(i), 7) == 0)
-      VGPR32SetID = i;
+  SGPRSetID = NumRegPressureSets;
+  VGPRSetID = NumRegPressureSets;
+  std::map<unsigned, unsigned> PressureSetRegUnits;
 
+  for (unsigned i = 0; i < NumRegPressureSets; ++i) {
     classifyPressureSet(i, AMDGPU::SGPR0, SGPRPressureSets);
     classifyPressureSet(i, AMDGPU::VGPR0, VGPRPressureSets);
+    PressureSetRegUnits[i] = 0;
+  }
+
+  // Determine the number of reg units for each pressure set.
+  for (unsigned i = 0, e = getNumRegUnits(); i != e; ++i) {
+    const int *PSets = getRegUnitPressureSets(i);
+    for (unsigned j = 0; PSets[j] != -1; ++j) {
+      PressureSetRegUnits[PSets[j]]++;
+    }
+  }
+
+  unsigned VGPRMax = 0, SGPRMax = 0;
+  for (unsigned i = 0; i < NumRegPressureSets; ++i) {
+    //dbgs() << getRegPressureSetName(i) << " has " << PressureSetRegUnits[i] << " units.\n";
+    if (isVGPRPressureSet(i) && PressureSetRegUnits[i] > VGPRMax) {
+      VGPRSetID = i;
+      VGPRMax = PressureSetRegUnits[i];
+      continue;
+    }
+    if (isSGPRPressureSet(i) && PressureSetRegUnits[i] > SGPRMax) {
+      SGPRSetID = i;
+      SGPRMax = PressureSetRegUnits[i];
+    }
   }
-  assert(SGPR32SetID < NumRegPressureSets &&
-         VGPR32SetID < NumRegPressureSets);
+
+  assert(SGPRSetID < NumRegPressureSets &&
+         VGPRSetID < NumRegPressureSets);
 }
 
 void SIRegisterInfo::reserveRegisterTuples(BitVector &Reserved, unsigned Reg) const {
@@ -954,3 +975,52 @@
 
   return hasVGPRs(RC);
 }
+
+std::pair<unsigned, unsigned>
+SIRegisterInfo::getRegPressureSetLimits(const MachineFunction &MF) const {
+  std::pair<unsigned, unsigned> Limits = std::make_pair(0, 0);
+
+  for (unsigned i = 0, e = getNumRegPressureSets(); i < e; ++i) {
+    bool isSGPRSet = SGPRPressureSets.test(i);
+    bool isVGPRSet = VGPRPressureSets.test(i);
+
+    if (!isSGPRSet && !isVGPRSet)
+      continue;
+
+    if (isSGPRSet && isVGPRSet)
+      continue;
+
+    unsigned Limit = getRegPressureSetLimit(MF, i);
+    if (isSGPRSet) {
+      Limits.first = std::max(Limits.first, Limit);
+      continue;
+    }
+
+    Limits.second = std::max(Limits.second, Limit);
+  }
+  return Limits;
+}
+
+std::pair<unsigned, unsigned>
+SIRegisterInfo::getMaxRegPressure(ArrayRef<unsigned> SetPressure) const {
+  std::pair<unsigned, unsigned> MaxPressure = std::make_pair(0, 0);
+
+  for (unsigned i = 0; i < SetPressure.size(); i++) {
+    bool isSGPRSet = SGPRPressureSets.test(i);
+    bool isVGPRSet = VGPRPressureSets.test(i);
+
+    if (!isSGPRSet && !isVGPRSet)
+      continue;
+
+    // Ignore the sets that contain sgprs and vgprs.
+    if (isSGPRSet && isVGPRSet)
+      continue;
+
+    if (isSGPRSet) {
+      MaxPressure.first = std::max(MaxPressure.first, SetPressure[i]);
+      continue;
+    }
+    MaxPressure.second = std::max(MaxPressure.second, SetPressure[i]);
+  }
+  return MaxPressure;
+}