Index: lib/Target/AArch64/AArch64GenRegisterBankInfo.def
===================================================================
--- lib/Target/AArch64/AArch64GenRegisterBankInfo.def
+++ lib/Target/AArch64/AArch64GenRegisterBankInfo.def
@@ -149,16 +149,6 @@
 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.
@@ -242,7 +232,7 @@
   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 &&
@@ -268,11 +258,12 @@
   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 +
+      (DstRBIdx - AArch64GenRegisterBankInfo::PMI_Min +
+       AArch64GenRegisterBankInfo::getRegBankBaseIdxOffset(DstRBIdx, Size)) *
+          AArch64GenRegisterBankInfo::ValueMappingIdx::
+              DistanceBetweenCrossRegCpy;
   assert(ValMappingIdx >= AArch64GenRegisterBankInfo::FirstCrossRegCpyIdx &&
          ValMappingIdx <= AArch64GenRegisterBankInfo::LastCrossRegCpyIdx &&
          "Mapping out of bound");
Index: lib/Target/AArch64/AArch64RegisterBankInfo.h
===================================================================
--- lib/Target/AArch64/AArch64RegisterBankInfo.h
+++ lib/Target/AArch64/AArch64RegisterBankInfo.h
@@ -92,6 +92,29 @@
     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.
Index: lib/Target/AArch64/AArch64RegisterBankInfo.cpp
===================================================================
--- lib/Target/AArch64/AArch64RegisterBankInfo.cpp
+++ lib/Target/AArch64/AArch64RegisterBankInfo.cpp
@@ -381,6 +381,8 @@
   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.
@@ -392,17 +394,17 @@
   // 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) &&
+        "Operand has incompatible 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};
 }