[SeparateConstOffsetFromGEP] Update splitGEP to handle case where including base offset results in an offset that's too large (#177653)

Currently, separate-const-offset-from-gep tries to combine both the
offsets on the base address and the offsets on the current GEP itself
when it tries to separate constant offsets. This results in the pass
failing to separate the offset in cases where the base address has a
large offset that would cause the total offset to be larger than what
the back-end can represent for the respective addressing mode. However,
in many cases we can still benefit from extracting the offset of the GEP
without including the offset from the base-address-GEP when the
base-address is used in multiple different places to help reduce
register pressure and recalculating the base address.

This PR addresses the above by re-trying without accumulating the offset
from the base-address GEP included if isLegalAddressingMode returns
false the first time.

1 files changed, 20 insertions, 7 deletions

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index e01b5797e662..c298daff3010 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -979,12 +979,11 @@ bool SeparateConstOffsetFromGEP::splitGEP(GetElementPtrInst *GEP) {
   // offset from each is accumulated.
   Value *NewBase;
   const APInt *BaseOffset;
-  const bool ExtractBase =
-      match(GEP->getPointerOperand(),
-            m_PtrAdd(m_Value(NewBase), m_APInt(BaseOffset)));
+  bool ExtractBase = match(GEP->getPointerOperand(),
+                           m_PtrAdd(m_Value(NewBase), m_APInt(BaseOffset)));
 
   unsigned IdxWidth = DL->getIndexTypeSizeInBits(GEP->getType());
-  const APInt BaseByteOffset =
+  APInt BaseByteOffset =
       ExtractBase ? BaseOffset->sextOrTrunc(IdxWidth) : APInt(IdxWidth, 0);
 
   // The backend can already nicely handle the case where all indices are
@@ -995,8 +994,8 @@ bool SeparateConstOffsetFromGEP::splitGEP(GetElementPtrInst *GEP) {
   bool Changed = canonicalizeArrayIndicesToIndexSize(GEP);
 
   bool NeedsExtraction;
-  APInt AccumulativeByteOffset =
-      BaseByteOffset + accumulateByteOffset(GEP, NeedsExtraction);
+  APInt NonBaseByteOffset = accumulateByteOffset(GEP, NeedsExtraction);
+  APInt AccumulativeByteOffset = BaseByteOffset + NonBaseByteOffset;
 
   TargetTransformInfo &TTI = GetTTI(*GEP->getFunction());
 
@@ -1018,7 +1017,21 @@ bool SeparateConstOffsetFromGEP::splitGEP(GetElementPtrInst *GEP) {
             GEP->getResultElementType(),
             /*BaseGV=*/nullptr, AccumulativeByteOffset.getSExtValue(),
             /*HasBaseReg=*/true, /*Scale=*/0, AddrSpace)) {
-      return Changed;
+      // If the addressing mode was not legal and the base byte offset was not
+      // 0, it could be a case where the total offset became too large for
+      // the addressing mode. Try again without extracting the base offset.
+      if (!ExtractBase)
+        return Changed;
+      ExtractBase = false;
+      BaseByteOffset = APInt(IdxWidth, 0);
+      AccumulativeByteOffset = NonBaseByteOffset;
+      if (!TTI.isLegalAddressingMode(
+              GEP->getResultElementType(),
+              /*BaseGV=*/nullptr, AccumulativeByteOffset.getSExtValue(),
+              /*HasBaseReg=*/true, /*Scale=*/0, AddrSpace))
+        return Changed;
+      // We can proceed with just extracting the other (non-base) offsets.
+      NeedsExtraction = true;
     }
   }