[aarch64][globalisel] Refactor getRegBankBaseIdxOffset() to remove the power-of-2 assumption. NFC

[dsanders]

Summary

We don't exploit it yet though

Depends on D27976

[qcolombet]

LGTM.

[dsanders]

Rebased onto the updated D27976

[dsanders]

Refresh before commit

lib/Target/AArch64/AArch64GenRegisterBankInfo.def

 RegisterBank GPRRegBank(AArch64::GPRRegBankID, "GPR", 64, GPRCoverageData);
 RegisterBank FPRRegBank(AArch64::FPRRegBankID, "FPR", 512, FPRCoverageData);
 RegisterBank CCRRegBank(AArch64::CCRRegBankID, "CCR", 32, CCRCoverageData);
 } // end namespace AArch64
 
 RegisterBank *AArch64GenRegisterBankInfo::RegBanks[] = {
     &AArch64::GPRRegBank, &AArch64::FPRRegBank, &AArch64::CCRRegBank};
 
-namespace AArch64 {
-static unsigned getRegBankBaseIdxOffset(unsigned Size) {
-  assert(Size && "0-sized type!!");
-  // Make anything smaller than 32 gets 32
-  Size = ((Size + 31) / 32) * 32;
-  // 32 is 0, 64 is 1, 128 is 2, and so on.
-  return Log2_32(Size) - /*Log2_32(32)=*/ 5;
-}
-}
-
 RegisterBankInfo::PartialMapping AArch64GenRegisterBankInfo::PartMappings[]{
     /* StartIdx, Length, RegBank */
     // 0: GPR 32-bit value.
     {0, 32, AArch64::GPRRegBank},
     // 1: GPR 64-bit value.
     {0, 64, AArch64::GPRRegBank},
     // 2: FPR 32-bit value.
     {0, 32, AArch64::FPRRegBank},

 const RegisterBankInfo::ValueMapping *
 getValueMapping(AArch64GenRegisterBankInfo::PartialMappingIdx RBIdx,
                 unsigned Size) {
   assert(RBIdx != AArch64GenRegisterBankInfo::PartialMappingIdx::PMI_None &&
          "No mapping needed for that");
   unsigned ValMappingIdx =
       AArch64GenRegisterBankInfo::First3OpsIdx +
       (RBIdx - AArch64GenRegisterBankInfo::PartialMappingIdx::PMI_Min +
-       getRegBankBaseIdxOffset(Size)) *
+       AArch64GenRegisterBankInfo::getRegBankBaseIdxOffset(RBIdx, Size)) *
           AArch64GenRegisterBankInfo::ValueMappingIdx::DistanceBetweenRegBanks;
   assert(ValMappingIdx >= AArch64GenRegisterBankInfo::First3OpsIdx &&
          ValMappingIdx <= AArch64GenRegisterBankInfo::Last3OpsIdx &&
          "Mapping out of bound");
 
   return &AArch64GenRegisterBankInfo::ValMappings[ValMappingIdx];
 }
 
AArch64GenRegisterBankInfo::PartialMappingIdx DstRBIdx =
       DstIsGPR ? AArch64GenRegisterBankInfo::PMI_FirstGPR
                : AArch64GenRegisterBankInfo::PMI_FirstFPR;
   AArch64GenRegisterBankInfo::PartialMappingIdx SrcRBIdx =
       SrcIsGPR ? AArch64GenRegisterBankInfo::PMI_FirstGPR
                : AArch64GenRegisterBankInfo::PMI_FirstFPR;
   if (DstRBIdx == SrcRBIdx)
     return getValueMapping(DstRBIdx, Size);
   assert(Size <= 64 && "GPR cannot handle that size");
-  unsigned ValMappingIdx = AArch64GenRegisterBankInfo::FirstCrossRegCpyIdx +
-                           (DstRBIdx - AArch64GenRegisterBankInfo::PMI_Min +
-                            getRegBankBaseIdxOffset(Size)) *
-                               AArch64GenRegisterBankInfo::ValueMappingIdx::
-                                   DistanceBetweenCrossRegCpy;
+  unsigned ValMappingIdx =
+      AArch64GenRegisterBankInfo::FirstCrossRegCpyIdx +
+       AArch64GenRegisterBankInfo::getRegBankBaseIdxOffset(DstRBIdx, Size)) *
   assert(ValMappingIdx >= AArch64GenRegisterBankInfo::FirstCrossRegCpyIdx &&
          ValMappingIdx <= AArch64GenRegisterBankInfo::LastCrossRegCpyIdx &&
          "Mapping out of bound");
   return &AArch64GenRegisterBankInfo::ValMappings[ValMappingIdx];
 }
 } // End AArch64 namespace.
 
 bool AArch64GenRegisterBankInfo::checkPartialMap(unsigned Idx,

lib/Target/AArch64/AArch64RegisterBankInfo.h

for (const auto &Current : Order) {
       }
       if (Previous + 1 != Current)
         return false;
       Previous = Current;
     }
     return true;
   }
 
+  static unsigned getRegBankBaseIdxOffset(unsigned RBIdx, unsigned Size) {
+    if (RBIdx == PMI_FirstGPR) {
+      if (Size <= 32)
+        return 0;
+      if (Size <= 64)
+        return 1;
+      llvm_unreachable("Unexpected size");
+    }
+    if (RBIdx == PMI_FirstFPR) {
+      if (Size <= 32)
+        return 0;
+      if (Size <= 64)
+        return 1;
+      if (Size <= 128)
+        return 2;
+      if (Size <= 256)
+        return 3;
+      if (Size <= 512)
+        return 4;
+      llvm_unreachable("Unexpected size");
+    }
+    llvm_unreachable("Unexpected bank");
+}
 };
 
 /// This class provides the information for the target register banks.
 class AArch64RegisterBankInfo final : public AArch64GenRegisterBankInfo {
   /// See RegisterBankInfo::applyMapping.
   void applyMappingImpl(const OperandsMapper &OpdMapper) const override;
 
   /// Get an instruction mapping where all the operands map to

lib/Target/AArch64/AArch64RegisterBankInfo.cpp

AArch64RegisterBankInfo::getSameKindOfOperandsMapping(const MachineInstr &MI) {
   unsigned NumOperands = MI.getNumOperands();
   assert(NumOperands <= 3 &&
          "This code is for instructions with 3 or less operands");
 
   LLT Ty = MRI.getType(MI.getOperand(0).getReg());
   unsigned Size = Ty.getSizeInBits();
   bool IsFPR = Ty.isVector() || isPreISelGenericFloatingPointOpcode(Opc);
 
+  PartialMappingIdx RBIdx = IsFPR ? PMI_FirstFPR : PMI_FirstGPR;
+
 #ifndef NDEBUG
   // Make sure all the operands are using similar size and type.
   // Should probably be checked by the machine verifier.
   // This code won't catch cases where the number of lanes is
   // different between the operands.
   // If we want to go to that level of details, it is probably
   // best to check that the types are the same, period.
   // Currently, we just check that the register banks are the same
   // for each types.
   for (unsigned Idx = 1; Idx != NumOperands; ++Idx) {
     LLT OpTy = MRI.getType(MI.getOperand(Idx).getReg());
-    assert(AArch64::getRegBankBaseIdxOffset(OpTy.getSizeInBits()) ==
-               AArch64::getRegBankBaseIdxOffset(Size) &&
-           "Operand has incompatible size");
+    assert(
+        AArch64GenRegisterBankInfo::getRegBankBaseIdxOffset(
+            RBIdx, OpTy.getSizeInBits()) ==
+            AArch64GenRegisterBankInfo::getRegBankBaseIdxOffset(RBIdx, Size) &&
     bool OpIsFPR = OpTy.isVector() || isPreISelGenericFloatingPointOpcode(Opc);
     (void)OpIsFPR;
     assert(IsFPR == OpIsFPR && "Operand has incompatible type");
   }
 #endif // End NDEBUG.
 
-  PartialMappingIdx RBIdx = IsFPR ? PMI_FirstFPR : PMI_FirstGPR;
-
   return InstructionMapping{DefaultMappingID, 1,
                             AArch64::getValueMapping(RBIdx, Size), NumOperands};
 }
 
 RegisterBankInfo::InstructionMapping
 AArch64RegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
   const unsigned Opc = MI.getOpcode();
   const MachineFunction &MF = *MI.getParent()->getParent();