[AMDGPU] Cleanup assumptions about generated subregs

[rampitec]

Summary

We are using countPopulation on a LaneBitmask to determine
a number of registers it covers. This is the assumption which
does not necessarily need to be true. It is not changed but
factored into a single call SIRegisterInfo::getNumCoveredRegs().

Some other places are cleaned up with respect to assumptions
about subreg indexes values and tablegen behavior.

[arsenm]

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

303

Typo Channeles

[rampitec]

Fixed typo.
Found 2 more places: SIAddIMGInit and computeIndirectRegAndOffset.

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

for (const auto& Op : MI.explicit_uses()) {
 
     Register R = Op.getReg();
     if (TRI->hasAGPRs(TRI->getRegClassForReg(*MRI, R)))
       continue;
 
     unsigned ShiftedBank = Bank;
 
     if (Bank != -1 && R == Reg && Op.getSubReg()) {
-      unsigned LM = TRI->getSubRegIndexLaneMask(Op.getSubReg()).getAsInteger();
-      if (!(LM & 1) && (Bank < NUM_VGPR_BANKS)) {
+      unsigned Offset = TRI->getChannelFromSubReg(Op.getSubReg());
+      LaneBitmask LM = TRI->getSubRegIndexLaneMask(Op.getSubReg());
+      if (Offset && Bank < NUM_VGPR_BANKS) {
         // If a register spans all banks we cannot shift it to avoid conflict.
-        if (countPopulation(LM) >= NUM_VGPR_BANKS)
+        if (TRI->getNumCoveredRegs(LM) >= NUM_VGPR_BANKS)
           continue;
-        ShiftedBank = (Bank + countTrailingZeros(LM)) % NUM_VGPR_BANKS;
-      } else if (!(LM & 3) && (Bank >= SGPR_BANK_OFFSET)) {
+        ShiftedBank = (Bank + Offset) % NUM_VGPR_BANKS;
+      } else if (Offset > 1 && Bank >= SGPR_BANK_OFFSET) {
         // If a register spans all banks we cannot shift it to avoid conflict.
-        if (countPopulation(LM) / 2 >= NUM_SGPR_BANKS)
+        if (TRI->getNumCoveredRegs(LM) / 2 >= NUM_SGPR_BANKS)
           continue;
-        ShiftedBank = SGPR_BANK_OFFSET + (Bank - SGPR_BANK_OFFSET +
-                                          (countTrailingZeros(LM) >> 1)) %
-                                             NUM_SGPR_BANKS;
+        ShiftedBank = SGPR_BANK_OFFSET +
+          (Bank - SGPR_BANK_OFFSET + (Offset >> 1)) % NUM_SGPR_BANKS;
       }
     }
 
     uint32_t Mask = getRegBankMask(R, Op.getSubReg(),
                                    (Reg == R) ? ShiftedBank : -1);
     StallCycles += countPopulation(UsedBanks & Mask);
     UsedBanks |= Mask;
     OperandMasks.push_back(OperandMask(Op.getReg(), Op.getSubReg(), Mask));
unsigned GCNRegBankReassign::getFreeBanks(unsigned Reg,
                                           unsigned SubReg,
                                           unsigned Mask,
                                           unsigned UsedBanks) const {
   if (!isReassignable(Reg))
     return 0;
 
   unsigned FreeBanks = getFreeBanks(Mask, UsedBanks);
 
-  unsigned LM = TRI->getSubRegIndexLaneMask(SubReg).getAsInteger();
-  if (!(LM & 1) && (Mask & VGPR_BANK_MASK)) {
-    unsigned Shift = countTrailingZeros(LM);
+  unsigned Offset = TRI->getChannelFromSubReg(SubReg);
+  if (Offset && (Mask & VGPR_BANK_MASK)) {
+    unsigned Shift = Offset;
     if (Shift >= NUM_VGPR_BANKS)
       return 0;
     unsigned VB = FreeBanks & VGPR_BANK_MASK;
     FreeBanks = ((VB >> Shift) | (VB << (NUM_VGPR_BANKS - Shift))) &
                 VGPR_BANK_MASK;
-  } else if (!(LM & 3) && (Mask & SGPR_BANK_MASK)) {
-    unsigned Shift = countTrailingZeros(LM) >> 1;
+  } else if (Offset > 1 && (Mask & SGPR_BANK_MASK)) {
+    unsigned Shift = Offset >> 1;
     if (Shift >= NUM_SGPR_BANKS)
       return 0;
     unsigned SB = FreeBanks >> SGPR_BANK_OFFSET;
     FreeBanks = ((SB >> Shift) | (SB << (NUM_SGPR_BANKS - Shift))) &
                 SGPR_BANK_SHIFTED_MASK;
     FreeBanks <<= SGPR_BANK_OFFSET;
   }
 

llvm/lib/Target/AMDGPU/GCNRegPressure.cpp

#endif
 
   case SGPR_TUPLE:
   case VGPR_TUPLE:
   case AGPR_TUPLE:
     assert(NewMask < MaxMask || NewMask == MaxMask);
     assert(PrevMask < NewMask);
 
     Value[Kind == SGPR_TUPLE ? SGPR32 : Kind == AGPR_TUPLE ? AGPR32 : VGPR32] +=
-      Sign * (~PrevMask & NewMask).getNumLanes();
+      Sign * SIRegisterInfo::getNumCoveredRegs(~PrevMask & NewMask);
 
     if (PrevMask.none()) {
       assert(NewMask.any());
       Value[Kind] += Sign * MRI.getPressureSets(Reg).getWeight();
     }
     break;
 
   default: llvm_unreachable("Unknown register kind");
static LaneBitmask getUsedRegMask(const MachineOperand &MO,
                                   const MachineRegisterInfo &MRI,
                                   const LiveIntervals &LIS) {
   assert(MO.isUse() && MO.isReg() && Register::isVirtualRegister(MO.getReg()));
 
   if (auto SubReg = MO.getSubReg())
     return MRI.getTargetRegisterInfo()->getSubRegIndexLaneMask(SubReg);
 
   auto MaxMask = MRI.getMaxLaneMaskForVReg(MO.getReg());
-  if (MaxMask == LaneBitmask::getLane(0)) // cannot have subregs
+  if (SIRegisterInfo::getNumCoveredRegs(MaxMask) > 1) // cannot have subregs
     return MaxMask;
 
   // For a tentative schedule LIS isn't updated yet but livemask should remain
   // the same on any schedule. Subreg defs can be reordered but they all must
   // dominate uses anyway.
   auto SI = LIS.getInstructionIndex(*MO.getParent()).getBaseIndex();
   return getLiveLaneMask(MO.getReg(), SI, LIS, MRI);
 }

llvm/lib/Target/AMDGPU/SIAddIMGInit.cpp

for (I = MBB.begin(); I != MBB.end(); I = Next) {
           assert(MO_Dmask && "Expected dmask operand in instruction");
 
           unsigned dmask = MO_Dmask->getImm();
           // Determine the number of active lanes taking into account the
           // Gather4 special case
           unsigned ActiveLanes =
               TII->isGather4(Opcode) ? 4 : countPopulation(dmask);
 
-          // Subreg indices are counted from 1
-          // When D16 then we want next whole VGPR after write data.
-          static_assert(AMDGPU::sub0 == 1 && AMDGPU::sub4 == 5, "Subreg indices different from expected");
-
           bool Packed = !ST.hasUnpackedD16VMem();
 
           unsigned InitIdx =
               D16Val && Packed ? ((ActiveLanes + 1) >> 1) + 1 : ActiveLanes + 1;
 
           // Abandon attempt if the dst size isn't large enough
           // - this is in fact an error but this is picked up elsewhere and
           // reported correctly.
           uint32_t DstSize =
               RI->getRegSizeInBits(*TII->getOpRegClass(MI, DstIdx)) / 32;
           if (DstSize < InitIdx)
             continue;
 
           // Create a register for the intialization value.
           Register PrevDst =
               MRI.createVirtualRegister(TII->getOpRegClass(MI, DstIdx));
           unsigned NewDst = 0; // Final initialized value will be in here
 
           // If PRTStrictNull feature is enabled (the default) then initialize
           // all the result registers to 0, otherwise just the error indication
           // register (VGPRn+1)
           unsigned SizeLeft = ST.usePRTStrictNull() ? InitIdx : 1;
-          unsigned CurrIdx = ST.usePRTStrictNull() ? 1 : InitIdx;
+          unsigned CurrIdx = ST.usePRTStrictNull() ? 0 : (InitIdx - 1);
 
           if (DstSize == 1) {
             // In this case we can just initialize the result directly
             BuildMI(MBB, MI, DL, TII->get(AMDGPU::V_MOV_B32_e32), PrevDst)
                 .addImm(0);
             NewDst = PrevDst;
           } else {
             BuildMI(MBB, MI, DL, TII->get(AMDGPU::IMPLICIT_DEF), PrevDst);
             for (; SizeLeft; SizeLeft--, CurrIdx++) {
               NewDst =
                   MRI.createVirtualRegister(TII->getOpRegClass(MI, DstIdx));
               // Initialize dword
               Register SubReg =
                   MRI.createVirtualRegister(&AMDGPU::VGPR_32RegClass);
               BuildMI(MBB, MI, DL, TII->get(AMDGPU::V_MOV_B32_e32), SubReg)
                   .addImm(0);
               // Insert into the super-reg
               BuildMI(MBB, I, DL, TII->get(TargetOpcode::INSERT_SUBREG), NewDst)
                   .addReg(PrevDst)
                   .addReg(SubReg)
-                  .addImm(CurrIdx);
+                  .addImm(AMDGPURegisterInfo::getSubRegFromChannel(CurrIdx));
 
               PrevDst = NewDst;
             }
           }
 
           // Add as an implicit operand
           MachineInstrBuilder(MF, MI).addReg(NewDst, RegState::Implicit);
 

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

computeIndirectRegAndOffset(const SIRegisterInfo &TRI,
                             int Offset) {
   int NumElts = TRI.getRegSizeInBits(*SuperRC) / 32;
 
   // Skip out of bounds offsets, or else we would end up using an undefined
   // register.
   if (Offset >= NumElts || Offset < 0)
     return std::make_pair(AMDGPU::sub0, Offset);
 
-  return std::make_pair(AMDGPU::sub0 + Offset, 0);
+  return std::make_pair(AMDGPURegisterInfo::getSubRegFromChannel(Offset), 0);
 }
 
 // Return true if the index is an SGPR and was set.
 static bool setM0ToIndexFromSGPR(const SIInstrInfo *TII,
                                  MachineRegisterInfo &MRI,
                                  MachineInstr &MI,
                                  int Offset,
                                  bool UseGPRIdxMode,

llvm/lib/Target/AMDGPU/SIInstrInfo.h

public:
   unsigned getOpSize(const MachineInstr &MI, unsigned OpNo) const {
     const MachineOperand &MO = MI.getOperand(OpNo);
     if (MO.isReg()) {
       if (unsigned SubReg = MO.getSubReg()) {
         assert(RI.getRegSizeInBits(*RI.getSubClassWithSubReg(
                                    MI.getParent()->getParent()->getRegInfo().
                                      getRegClass(MO.getReg()), SubReg)) >= 32 &&
                "Sub-dword subregs are not supported");
-        return RI.getSubRegIndexLaneMask(SubReg).getNumLanes() * 4;
+        return RI.getNumChannelsFromSubReg(SubReg) * 4;
       }
     }
     return RI.getRegSizeInBits(*getOpRegClass(MI, OpNo)) / 8;
   }
 
   /// Legalize the \p OpIndex operand of this instruction by inserting
   /// a MOV.  For example:
   /// ADD_I32_e32 VGPR0, 15

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

if (RI.isSGPRClass(RegClass)) {
     }
   }
 
   ArrayRef<int16_t> SubIndices = RI.getRegSplitParts(RegClass, EltSize);
   for (unsigned Idx = 0; Idx < SubIndices.size(); ++Idx) {
     int64_t IdxValue = Idx == 0 ? Value : 0;
 
     MachineInstrBuilder Builder = BuildMI(MBB, MI, DL,
-      get(Opcode), RI.getSubReg(DestReg, Idx));
+      get(Opcode), RI.getSubReg(DestReg, SubIndices[Idx]));
     Builder.addImm(IdxValue);
   }
 }
 
 const TargetRegisterClass *
 SIInstrInfo::getPreferredSelectRegClass(unsigned Size) const {
   return &AMDGPU::VGPR_32RegClass;
 }

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

public:
   MachineInstr *findReachingDef(unsigned Reg, unsigned SubReg,
                                 MachineInstr &Use,
                                 MachineRegisterInfo &MRI,
                                 LiveIntervals *LIS) const;
 
   const uint32_t *getAllVGPRRegMask() const;
   const uint32_t *getAllAllocatableSRegMask() const;
 
+  // \returns number of 32 bit registers covered by a \p LM
+  static unsigned getNumCoveredRegs(LaneBitmask LM) {
+    return LM.getNumLanes();
+  }
+
+  // \returns a DWORD offset of a \p SubReg
+  unsigned getChannelFromSubReg(unsigned SubReg) const {
+    return SubReg ? alignTo(getSubRegIdxOffset(SubReg), 32) / 32 : 0;
+  }
+
+  // \returns a DWORD size of a \p SubReg
+  unsigned getNumChannelsFromSubReg(unsigned SubReg) const {
+    return getNumCoveredRegs(getSubRegIndexLaneMask(SubReg));
+  }
+
 private:
   void buildSpillLoadStore(MachineBasicBlock::iterator MI,
                            unsigned LoadStoreOp,
                            int Index,
                            unsigned ValueReg,
                            bool ValueIsKill,
                            unsigned ScratchRsrcReg,
                            unsigned ScratchOffsetReg,
                            int64_t InstrOffset,
                            MachineMemOperand *MMO,
                            RegScavenger *RS) const;
 };
 
 } // End namespace llvm
 
 #endif

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

 }
 
 const TargetRegisterClass *SIRegisterInfo::getSubRegClass(
                          const TargetRegisterClass *RC, unsigned SubIdx) const {
   if (SubIdx == AMDGPU::NoSubRegister)
     return RC;
 
   // We can assume that each lane corresponds to one 32-bit register.
-  unsigned Count = getSubRegIndexLaneMask(SubIdx).getNumLanes();
+  unsigned Count = getNumChannelsFromSubReg(SubIdx);
   if (isSGPRClass(RC)) {
     switch (Count) {
     case 1:
       return &AMDGPU::SGPR_32RegClass;
     case 2:
       return &AMDGPU::SReg_64RegClass;
     case 3:
       return &AMDGPU::SReg_96RegClass;

llvm/lib/Target/AMDGPU/SIShrinkInstructions.cpp

 
 static TargetInstrInfo::RegSubRegPair
 getSubRegForIndex(unsigned Reg, unsigned Sub, unsigned I,
                   const SIRegisterInfo &TRI, const MachineRegisterInfo &MRI) {
   if (TRI.getRegSizeInBits(Reg, MRI) != 32) {
     if (Register::isPhysicalRegister(Reg)) {
       Reg = TRI.getSubReg(Reg, TRI.getSubRegFromChannel(I));
     } else {
-      LaneBitmask LM = TRI.getSubRegIndexLaneMask(Sub);
-      Sub = TRI.getSubRegFromChannel(I + countTrailingZeros(LM.getAsInteger()));
+      Sub = TRI.getSubRegFromChannel(I + TRI.getChannelFromSubReg(Sub));
     }
   }
   return TargetInstrInfo::RegSubRegPair(Reg, Sub);
 }
 
 // Match:
 // mov t, x
 // mov x, y

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

 }
 
 bool isGCN3Encoding(const MCSubtargetInfo &STI) {
   return STI.getFeatureBits()[AMDGPU::FeatureGCN3Encoding];
 }
 
 bool isSGPR(unsigned Reg, const MCRegisterInfo* TRI) {
   const MCRegisterClass SGPRClass = TRI->getRegClass(AMDGPU::SReg_32RegClassID);
-  const unsigned FirstSubReg = TRI->getSubReg(Reg, 1);
+  const unsigned FirstSubReg = TRI->getSubReg(Reg, AMDGPU::sub0);
   return SGPRClass.contains(FirstSubReg != 0 ? FirstSubReg : Reg) ||
     Reg == AMDGPU::SCC;
 }
 
 bool isRegIntersect(unsigned Reg0, unsigned Reg1, const MCRegisterInfo* TRI) {
   for (MCRegAliasIterator R(Reg0, TRI, true); R.isValid(); ++R) {
     if (*R == Reg1) return true;
   }