Added generic function to get subreg idx from offset and size

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

diff --git a/llvm/include/llvm/CodeGen/TargetRegisterInfo.h b/llvm/include/llvm/CodeGen/TargetRegisterInfo.h
index 2e97629..11f63c6 100644
--- a/llvm/include/llvm/CodeGen/TargetRegisterInfo.h
+++ b/llvm/include/llvm/CodeGen/TargetRegisterInfo.h
@@ -240,6 +240,10 @@
   /// access sub-registers at different offsets), return -1.
   unsigned getSubRegIdxOffset(unsigned Idx) const;
 
+  /// Find a SubReg index for the given bit size and bit offset.
+  /// Returns 0(no-op sub-register) if no matching index exists.
+  unsigned getSubRegIdxFromOffsetSize(unsigned Offset, unsigned Size) const;
+
   /// Return a bitmask representing the parts of a register that are covered by
   /// SubIdx \see LaneBitmask.
   ///
@@ -249,9 +253,6 @@
     return SubRegIndexLaneMasks[SubIdx];
   }
 
-  /// Try to find a matching subreg from the given lanemask.
-  unsigned getSubRegIdxFromLaneMask(LaneBitmask LaneMask) const;
-
   /// Try to find one or more subregister indexes to cover \p LaneMask.
   ///
   /// If this is possible, returns true and appends the best matching set of

llvm/lib/CodeGen/TargetRegisterInfo.cpp

diff --git a/llvm/lib/CodeGen/TargetRegisterInfo.cpp b/llvm/lib/CodeGen/TargetRegisterInfo.cpp
index 7c522f0..2a49e7e 100644
--- a/llvm/lib/CodeGen/TargetRegisterInfo.cpp
+++ b/llvm/lib/CodeGen/TargetRegisterInfo.cpp
@@ -495,16 +495,6 @@
   return getRegSizeInBits(*RC);
 }
 
-unsigned
-TargetRegisterInfo::getSubRegIdxFromLaneMask(LaneBitmask LaneMask) const {
-  for (unsigned Idx = 1, E = getNumSubRegIndices(); Idx < E; ++Idx) {
-    if (getSubRegIndexLaneMask(Idx) == LaneMask)
-      return Idx;
-  }
-
-  return 0 /*NoSubRegister*/;
-}
-
 bool TargetRegisterInfo::getCoveringSubRegIndexes(
     const TargetRegisterClass *RC, LaneBitmask LaneMask,
     SmallVectorImpl<unsigned> &NeededIndexes) const {
@@ -613,6 +603,18 @@
   return SubRegIdxRanges[HwMode * getNumSubRegIndices() + Idx].Offset;
 }
 
+unsigned TargetRegisterInfo::getSubRegIdxFromOffsetSize(unsigned Offset,
+                                                        unsigned Size) const {
+  unsigned NumIdx = getNumSubRegIndices();
+  unsigned Base = HwMode * NumIdx;
+  for (unsigned Idx = 1; Idx < NumIdx; Idx++) {
+    if (SubRegIdxRanges[Base + Idx].Offset == Offset &&
+        SubRegIdxRanges[Base + Idx].Size == Size)
+      return Idx;
+  }
+  return 0;
+}
+
 Register
 TargetRegisterInfo::lookThruCopyLike(Register SrcReg,
                                      const MachineRegisterInfo *MRI) const {

llvm/lib/Target/AMDGPU/AMDGPURewriteAGPRCopyMFMA.cpp

diff --git a/llvm/lib/Target/AMDGPU/AMDGPURewriteAGPRCopyMFMA.cpp b/llvm/lib/Target/AMDGPU/AMDGPURewriteAGPRCopyMFMA.cpp
index 0c6d23e..9f66d09 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURewriteAGPRCopyMFMA.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURewriteAGPRCopyMFMA.cpp
@@ -454,27 +454,18 @@
 AMDGPURewriteAGPRCopyMFMAImpl::getSubRegFromReload(MachineInstr &MI,
                                                    Register Reg) const {
   unsigned NumRegs = TRI.getRegSizeInBits(*MRI.getRegClass(Reg)) / 32;
-  unsigned SubReg = AMDGPU::NoSubRegister;
+  unsigned NumSpilledRegs = TII.getNumSubRegsForSpillOp(MI);
   // SubReg accesses for the tuple registers are of interest here.
+  // Skip if the entire tuple is reloaded.
   // Note: We don't support 16-bit subreg reloads. If that assumption is
   // changed in the future, this function should be revised.
-  if (NumRegs == 1)
-    return SubReg;
+  if (NumRegs == 1 || NumRegs == NumSpilledRegs)
+    return AMDGPU::NoSubRegister;
 
-  unsigned NumSpilledRegs = TII.getNumSubRegsForSpillOp(MI);
-  // Skip if the entire tuple is reloaded.
-  if (NumRegs == NumSpilledRegs)
-    return SubReg;
-
-  // Construct the covering lanes for the reloaded portion.
-  unsigned SubRegIdx =
-      TII.getNamedOperand(MI, AMDGPU::OpName::offset)->getImm() / 4;
-  // Subreg lane masks are maintained in terms of regunits and each 32-bit
-  // register consists of two regunits.
-  uint64_t Lanes = (1ULL << NumSpilledRegs * 2) - 1;
-  LaneBitmask CoveringLanes = LaneBitmask(Lanes << SubRegIdx * 2);
-  SubReg = TRI.getSubRegIdxFromLaneMask(CoveringLanes);
-  return SubReg;
+  unsigned StackSlotBitOffset =
+      TII.getNamedOperand(MI, AMDGPU::OpName::offset)->getImm() * 8;
+  return TRI.getSubRegIdxFromOffsetSize(StackSlotBitOffset,
+                                        NumSpilledRegs * 32);
 }
 
 void AMDGPURewriteAGPRCopyMFMAImpl::replaceSpillWithCopyToVReg(