[CodeGenPrepare] Unmerging GEPs across indirect branches must respect types

[momo5502]

Summary

The optimization in CodeGenPrepare, where GEPs are unmerged across indirect branches must respect the types of both GEPs and their sizes when adjusting the indices.

The sample here shows the bug:

https://godbolt.org/z/fvYvvGGjE

The value %elementValuePtr addresses the second field of the %struct.Blub. It is therefore a GEP with index 1 and type i8.
The value %nextArrayElement addresses the next array element. It is therefore a GEP with index 1 and type %struct.Blub.

Both values point to completely different addresses, even if the indices are the same, due to the types being different.
However, after CodeGenPrepare has run, %nextArrayElement is a bitcast from %elementValuePtr, meaning both were treated as equal.

The cause for this is that the unmerging optimization does not take types into consideration.
It sees both GEPs have %currentArrayElement as source operand and therefore tries to rewrite %nextArrayElement in terms of %elementValuePtr.
It changes the index to the difference of the two GEPs. As both indices are 1, the difference is 0. As the indices are 0 the GEP is later replaced with a simple bitcast in CodeGenPrepare.

Before adjusting the indices, the types of the GEPs have to be aligned and the indices scaled accordingly for the optimization to be correct.
Due to the size of the struct being 16 and the %elementValuePtr pointing to offset 1, the correct index for the unmerged %nextArrayElement would be 15.

I adjusted the optimization and added a simple test case that asserts the correct index is preserved.
I assume this bug emerged from the opaque pointer change as GEPs like %elementValuePtr that access the struct field based of type i8 did not naturally occur before.

[momo5502]

Will be continued as a GitHub PR

llvm/lib/CodeGen/CodeGenPrepare.cpp

 // Return true if the GEP has two operands, the first operand is of a sequential
 // type, and the second operand is a constant.
 static bool GEPSequentialConstIndexed(GetElementPtrInst *GEP) {
   gep_type_iterator I = gep_type_begin(*GEP);
   return GEP->getNumOperands() == 2 && I.isSequential() &&
          isa<ConstantInt>(GEP->getOperand(1));
 }
 
+// Returns the size of the GEP type.
+// GEP must meet GEPSequentialConstIndexed requirements.
+static TypeSize GEPSequentialConstIndexedTypeSize(GetElementPtrInst &GEP) {
+  const DataLayout &DL = GEP.getModule()->getDataLayout();
+  gep_type_iterator GTI = gep_type_begin(GEP);
+  return DL.getTypeAllocSize(GTI.getIndexedType());
+}
+
+// Returns true if the Target can be addressed from Source
+// in case types mismatch.
+// Both GEPs must meet GEPSequentialConstIndexed requirements.
+static bool GEPIsAddressableFromSource(GetElementPtrInst &Target,
+                                       GetElementPtrInst &Source) {
+  if (Source.getSourceElementType() == Target.getSourceElementType()) {
+    return true;
+  }
+
+  TypeSize SourceTypeSize = GEPSequentialConstIndexedTypeSize(Source);
+  TypeSize TargetTypeSize = GEPSequentialConstIndexedTypeSize(Target);
+
+  return (TargetTypeSize % SourceTypeSize) == 0;
+}
+
+// Calculates the relative index required to address Target
+// from Source with respect to their element types.
+// Both GEPs must meet GEPSequentialConstIndexed requirements.
+static APInt GEPCalculateRelativeIndex(GetElementPtrInst &Target,
+                                       GetElementPtrInst &Source) {
+  uint64_t Scale = 1;
+
+  if (Source.getSourceElementType() != Target.getSourceElementType()) {
+    TypeSize SourceTypeSize = GEPSequentialConstIndexedTypeSize(Source);
+    TypeSize TargetTypeSize = GEPSequentialConstIndexedTypeSize(Target);
+
+    Scale = TargetTypeSize / SourceTypeSize;
+    assert(Scale * SourceTypeSize == TargetTypeSize &&
+           "Target GEP must be addressable from Source GEP");
+  }
+
+  ConstantInt *SourceIdx = cast<ConstantInt>(Source.getOperand(1));
+  ConstantInt *TargetIdx = cast<ConstantInt>(Target.getOperand(1));
+
+  return (TargetIdx->getValue() * Scale) - SourceIdx->getValue();
+}
+
 // Try unmerging GEPs to reduce liveness interference (register pressure) across
 // IndirectBr edges. Since IndirectBr edges tend to touch on many blocks,
 // reducing liveness interference across those edges benefits global register
 // allocation. Currently handles only certain cases.
 //
 // For example, unmerge %GEPI and %UGEPI as below.
 //
 // ---------- BEFORE ----------
for (User *Usr : GEPIOp->users()) {
     if (!isa<GetElementPtrInst>(Usr))
       return false;
     auto *UGEPI = cast<GetElementPtrInst>(Usr);
     // Check if UGEPI is a simple gep with a single constant index and GEPIOp is
     // the pointer operand to it. If so, record it in the vector. If not, give
     // up.
     if (!GEPSequentialConstIndexed(UGEPI))
       return false;
+    // Check if GEP Types match or if types are compatible
+    if (!GEPIsAddressableFromSource(*UGEPI, *GEPI))
+      return false;
     if (UGEPI->getOperand(0) != GEPIOp)
       return false;
     if (GEPIIdx->getType() !=
         cast<ConstantInt>(UGEPI->getOperand(1))->getType())
       return false;
     ConstantInt *UGEPIIdx = cast<ConstantInt>(UGEPI->getOperand(1));
     if (TTI->getIntImmCost(UGEPIIdx->getValue(), UGEPIIdx->getType(),
                            TargetTransformInfo::TCK_SizeAndLatency) >
         TargetTransformInfo::TCC_Basic)
       return false;
     UGEPIs.push_back(UGEPI);
   }
   if (UGEPIs.size() == 0)
     return false;
   // Check the materializing cost of (Uidx-Idx).
   for (GetElementPtrInst *UGEPI : UGEPIs) {
-    ConstantInt *UGEPIIdx = cast<ConstantInt>(UGEPI->getOperand(1));
-    APInt NewIdx = UGEPIIdx->getValue() - GEPIIdx->getValue();
+    APInt NewIdx = GEPCalculateRelativeIndex(*UGEPI, *GEPI);
     InstructionCost ImmCost = TTI->getIntImmCost(
         NewIdx, GEPIIdx->getType(), TargetTransformInfo::TCK_SizeAndLatency);
     if (ImmCost > TargetTransformInfo::TCC_Basic)
       return false;
   }
   // Now unmerge between GEPI and UGEPIs.
   for (GetElementPtrInst *UGEPI : UGEPIs) {
+    APInt NewIdx = GEPCalculateRelativeIndex(*UGEPI, *GEPI);
+    Constant *NewUGEPIIdx = ConstantInt::get(GEPIIdx->getType(), NewIdx);
     UGEPI->setOperand(0, GEPI);
-    ConstantInt *UGEPIIdx = cast<ConstantInt>(UGEPI->getOperand(1));
-    Constant *NewUGEPIIdx = ConstantInt::get(
-        GEPIIdx->getType(), UGEPIIdx->getValue() - GEPIIdx->getValue());
     UGEPI->setOperand(1, NewUGEPIIdx);
+    UGEPI->setSourceElementType(GEPI->getSourceElementType());
+    UGEPI->setResultElementType(GEPI->getResultElementType());
     // If GEPI is not inbounds but UGEPI is inbounds, change UGEPI to not
     // inbounds to avoid UB.
     if (!GEPI->isInBounds()) {
       UGEPI->setIsInBounds(false);
     }
   }
   // After unmerging, verify that GEPIOp is actually only used in SrcBlock (not
   // alive on IndirectBr edges).

llvm/test/CodeGen/Generic/indirect-br-gep-unmerge.ll

+; RUN: llc %s -stop-after=codegenprepare -o - | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%struct.Blub = type { i8, i8, ptr }
+
+@indirectBrPtr = external hidden global ptr
+
+define dso_local noundef ptr @testFunc(ptr noundef readonly %array, i1 %skip) {
+entry:
+  br i1 %skip, label %loopHeader, label %endBlock
+
+loopHeader:                                                ; preds = %2, %1
+  %currentArrayElement = phi ptr [ %array, %entry ], [ %nextArrayElement, %loopFooter ]
+  %elementValuePtr = getelementptr inbounds i8, ptr %currentArrayElement, i64 1
+  %elementValue = load i8, ptr %elementValuePtr, align 1
+  indirectbr ptr @indirectBrPtr, [label %loopFooter, label %endBlock]
+
+loopFooter:
+  %isGoodValue = icmp eq i8 %elementValue, 0
+  ; CHECK: %nextArrayElement = getelementptr inbounds i8, ptr %elementValuePtr, i64 15
+  %nextArrayElement = getelementptr inbounds %struct.Blub, ptr %currentArrayElement, i64 1
+  br i1 %isGoodValue, label %loopHeader, label %endBlock
+
+endBlock:                                                ; preds = %2
+  %retVal = phi ptr [ %array, %entry ], [ %elementValuePtr, %loopFooter ], [ %elementValuePtr, %loopHeader ]
+  ret ptr %retVal
+}