[InstCombine] Combine opaque pointer GEPs with mismatching element types

[nikic]

Summary

Currently, two GEPs will only be combined if the result element type of one is the same as the source element type of the other. However, this means we may miss folding opportunities where the second GEP could be rewritten using a different element type. This is especially relevant for opaque pointers, where constant GEPs often use i8 element type.

Address this by converting GEP indices to offsets, adding them, and then converting them back to indices. The first (inner) GEP is allowed to have variable indices as well, in which case only the constant suffix is converted into an offset.

I believe this should address the regression reported in https://reviews.llvm.org/D123300#3467615.

[nikic]

llvm/test/Transforms/InstCombine/opaque-ptr.ll

215

Worth noting that we still miss cases like these, where we would have to change the element type of the first GEP. I think we'll want to handle this by canonicalizing constant GEPs to i8 in InstCombine. I think this patch should ensure that doing so will not cause regressions when such i8 GEPs are combined with non-constant GEPs.

Type-based GEPs are such a mess :(

[nikic]

Don't require one-use for two constant GEPs, so the case from https://reviews.llvm.org/D123300#3467615 is actually fixed.

[aeubanks]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2067

do we want to skip the transform after this if any of this fails?

2070

I'm surprised we don't have some sort of first_n

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2067

I think it's fine to skip it, because this one should always be more powerful (for the "all constant GEP" subset it handles). At least that was the intention, maybe I'm overlooking some edge case.

2070

Yeah, I went looking for that one as well. There is take_front on indexed_accessor_range, but there doesn't seem to be anything generic currently.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

static Instruction *foldSelectGEP(GetElementPtrInst &GEP,
   return SelectInst::Create(Cond, NewTrueC, NewFalseC, "", nullptr, Sel);
 }
 
 Instruction *InstCombinerImpl::visitGEPOfGEP(GetElementPtrInst &GEP,
                                              GEPOperator *Src) {
   // Combine Indices - If the source pointer to this getelementptr instruction
   // is a getelementptr instruction with matching element type, combine the
   // indices of the two getelementptr instructions into a single instruction.
-  if (Src->getResultElementType() != GEP.getSourceElementType())
-    return nullptr;
-
   if (!shouldMergeGEPs(*cast<GEPOperator>(&GEP), *Src))
     return nullptr;
 
-  if (Src->getNumOperands() == 2 && GEP.getNumOperands() == 2 &&
+  if (Src->getResultElementType() == GEP.getSourceElementType() &&
+      Src->getNumOperands() == 2 && GEP.getNumOperands() == 2 &&
       Src->hasOneUse()) {
     Value *GO1 = GEP.getOperand(1);
     Value *SO1 = Src->getOperand(1);
 
     if (LI) {
       // Try to reassociate loop invariant GEP chains to enable LICM.
       if (Loop *L = LI->getLoopFor(GEP.getParent())) {
         // Reassociate the two GEPs if SO1 is variant in the loop and GO1 is
Instruction *InstCombinerImpl::visitGEPOfGEP(GetElementPtrInst &GEP,
 
   // Note that if our source is a gep chain itself then we wait for that
   // chain to be resolved before we perform this transformation.  This
   // avoids us creating a TON of code in some cases.
   if (auto *SrcGEP = dyn_cast<GEPOperator>(Src->getOperand(0)))
     if (SrcGEP->getNumOperands() == 2 && shouldMergeGEPs(*Src, *SrcGEP))
       return nullptr;   // Wait until our source is folded to completion.
 
+  // For constant GEPs, use a more general offset-based folding approach.
+  // Only do this for opaque pointers, as the result element type may change.
+  Type *PtrTy = Src->getType()->getScalarType();
+  if (PtrTy->isOpaquePointerTy() && GEP.hasAllConstantIndices() &&
+      (Src->hasOneUse() || Src->hasAllConstantIndices())) {
+    // Split Src into a variable part and a constant suffix.
+    gep_type_iterator GTI = gep_type_begin(*Src);
+    Type *BaseType = GTI.getIndexedType();
+    bool IsFirstType = true;
+    unsigned NumVarIndices = 0;
+    for (auto Pair : enumerate(Src->indices())) {
+      if (!isa<ConstantInt>(Pair.value())) {
+        BaseType = GTI.getIndexedType();
+        IsFirstType = false;
+        NumVarIndices = Pair.index() + 1;
+      }
+      ++GTI;
+    }
+
+    // Determine the offset for the constant suffix of Src.
+    APInt Offset(DL.getIndexTypeSizeInBits(PtrTy), 0);
+    if (NumVarIndices != Src->getNumIndices()) {
+      // FIXME: getIndexedOffsetInType() does not handled scalable vectors.
+      if (isa<ScalableVectorType>(BaseType))
+        return nullptr;
+
+      SmallVector<Value *> ConstantIndices;
+      if (!IsFirstType)
+        ConstantIndices.push_back(
+            Constant::getNullValue(Type::getInt32Ty(GEP.getContext())));
+      append_range(ConstantIndices, drop_begin(Src->indices(), NumVarIndices));
+      Offset += DL.getIndexedOffsetInType(BaseType, ConstantIndices);
+    }
+
+    // Add the offset for GEP (which is fully constant).
+    if (!GEP.accumulateConstantOffset(DL, Offset))
+      return nullptr;
+
+    // Convert the total offset back into indices.
+    SmallVector<APInt> ConstIndices =
+        DL.getGEPIndicesForOffset(BaseType, Offset);
+    if (!Offset.isZero() || (!IsFirstType && !ConstIndices[0].isZero()))
+      return nullptr;
+
+    SmallVector<Value *> Indices;
+    append_range(Indices, drop_end(Src->indices(),
+                                   Src->getNumIndices() - NumVarIndices));
+    for (const APInt &Idx : drop_begin(ConstIndices, !IsFirstType))
+      Indices.push_back(ConstantInt::get(GEP.getContext(), Idx));
+
+    return isMergedGEPInBounds(*Src, *cast<GEPOperator>(&GEP))
+               ? GetElementPtrInst::CreateInBounds(Src->getSourceElementType(),
+                                                   Src->getOperand(0), Indices,
+                                                   GEP.getName())
+               : GetElementPtrInst::Create(Src->getSourceElementType(),
+                                           Src->getOperand(0), Indices,
+                                           GEP.getName());
+  }
+
+  if (Src->getResultElementType() != GEP.getSourceElementType())
+    return nullptr;
+
   SmallVector<Value*, 8> Indices;
 
   // Find out whether the last index in the source GEP is a sequential idx.
   bool EndsWithSequential = false;
   for (gep_type_iterator I = gep_type_begin(*Src), E = gep_type_end(*Src);
        I != E; ++I)
     EndsWithSequential = I.isSequential();
 

llvm/test/Transforms/InstCombine/opaque-ptr.ll

 ;
   %a2 = getelementptr { i32, { i32, i32 } }, ptr %a, i32 0, i32 1
   %a3 = getelementptr { i32, i32 }, ptr %a2, i32 0, i32 1
   ret ptr %a3
 }
 
 define ptr @geps_combinable_different_elem_type1(ptr %a) {
 ; CHECK-LABEL: @geps_combinable_different_elem_type1(
-; CHECK-NEXT:    [[A2:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 0, i32 1
-; CHECK-NEXT:    [[A3:%.*]] = getelementptr { i32, i32 }, ptr [[A2]], i64 0, i32 1
+; CHECK-NEXT:    [[A3:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 1
 ; CHECK-NEXT:    ret ptr [[A3]]
 ;
   %a2 = getelementptr { i32, i32 }, ptr %a, i32 0, i32 1
   %a3 = getelementptr { i32, i32 }, ptr %a2, i32 0, i32 1
   ret ptr %a3
 }
 
 define ptr @geps_combinable_different_elem_type2(ptr %a) {
 ; CHECK-LABEL: @geps_combinable_different_elem_type2(
-; CHECK-NEXT:    [[A2:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 0, i32 1
-; CHECK-NEXT:    [[A3:%.*]] = getelementptr i8, ptr [[A2]], i64 4
+; CHECK-NEXT:    [[A3:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 1
 ; CHECK-NEXT:    ret ptr [[A3]]
 ;
   %a2 = getelementptr { i32, i32 }, ptr %a, i32 0, i32 1
   %a3 = getelementptr i8, ptr %a2, i64 4
   ret ptr %a3
 }
 
 define ptr @geps_combinable_different_elem_type3(ptr %a) {
 ; CHECK-LABEL: @geps_combinable_different_elem_type3(
-; CHECK-NEXT:    [[A2:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 0, i32 1
-; CHECK-NEXT:    [[A3:%.*]] = getelementptr i8, ptr [[A2]], i64 8
+; CHECK-NEXT:    [[A3:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 1, i32 1
 ; CHECK-NEXT:    ret ptr [[A3]]
 ;
   %a2 = getelementptr { i32, i32 }, ptr %a, i32 0, i32 1
   %a3 = getelementptr i8, ptr %a2, i64 8
   ret ptr %a3
 }
 
 define ptr @geps_combinable_different_elem_type4(ptr %a) {
 ; CHECK-LABEL: @geps_combinable_different_elem_type4(
 ; CHECK-NEXT:    [[A2:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 0, i32 1
 ; CHECK-NEXT:    [[A3:%.*]] = getelementptr i8, ptr [[A2]], i64 10
 ; CHECK-NEXT:    ret ptr [[A3]]
 ;
   %a2 = getelementptr { i32, i32 }, ptr %a, i32 0, i32 1
   %a3 = getelementptr i8, ptr %a2, i64 10
   ret ptr %a3
 }
 
 define ptr @geps_combinable_different_elem_type5(ptr %a) {
 ; CHECK-LABEL: @geps_combinable_different_elem_type5(
-; CHECK-NEXT:    [[A2:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 0, i32 1
-; CHECK-NEXT:    [[A3:%.*]] = getelementptr i8, ptr [[A2]], i64 -4
-; CHECK-NEXT:    ret ptr [[A3]]
+; CHECK-NEXT:    ret ptr [[A:%.*]]
 ;
   %a2 = getelementptr { i32, i32 }, ptr %a, i32 0, i32 1
   %a3 = getelementptr i8, ptr %a2, i64 -4
   ret ptr %a3
 }
 
 define ptr @geps_combinable_different_elem_type6(ptr %a, i64 %idx) {
 ; CHECK-LABEL: @geps_combinable_different_elem_type6(
-; CHECK-NEXT:    [[A2:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 [[IDX:%.*]]
-; CHECK-NEXT:    [[A3:%.*]] = getelementptr i8, ptr [[A2]], i64 4
+; CHECK-NEXT:    [[A3:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 [[IDX:%.*]], i32 1
 ; CHECK-NEXT:    ret ptr [[A3]]
 ;
   %a2 = getelementptr { i32, i32 }, ptr %a, i64 %idx
   %a3 = getelementptr i8, ptr %a2, i64 4
   ret ptr %a3
 }
 
 define ptr @geps_combinable_different_elem_type7(ptr %a, i64 %idx) {

 ;
   %a2 = getelementptr inbounds { { i32, i32 } }, ptr %a, i64 %idx, i32 0, i32 1
   %a3 = getelementptr inbounds i8, ptr %a2, i32 4
   ret ptr %a3
 }
 
 define ptr @geps_combinable_different_elem_type9(ptr %a, i64 %idx) {
 ; CHECK-LABEL: @geps_combinable_different_elem_type9(
-; CHECK-NEXT:    [[A2:%.*]] = getelementptr inbounds { { i32, i32 } }, ptr [[A:%.*]], i64 [[IDX:%.*]], i32 0, i32 1
-; CHECK-NEXT:    [[A3:%.*]] = getelementptr inbounds i8, ptr [[A2]], i64 -4
+; CHECK-NEXT:    [[A3:%.*]] = getelementptr inbounds { { i32, i32 } }, ptr [[A:%.*]], i64 [[IDX:%.*]]
 ; CHECK-NEXT:    ret ptr [[A3]]
 ;
   %a2 = getelementptr inbounds { { i32, i32 } }, ptr %a, i64 %idx, i32 0, i32 1
   %a3 = getelementptr inbounds i8, ptr %a2, i32 -4
   ret ptr %a3
 }
 
 declare void @use(ptr)
 
 define ptr @geps_combinable_different_elem_type_extra_use1(ptr %a) {
 ; CHECK-LABEL: @geps_combinable_different_elem_type_extra_use1(
 ; CHECK-NEXT:    [[A2:%.*]] = getelementptr { i32, i32 }, ptr [[A:%.*]], i64 0, i32 1
 ; CHECK-NEXT:    call void @use(ptr [[A2]])
-; CHECK-NEXT:    [[A3:%.*]] = getelementptr i8, ptr [[A2]], i64 4
+; CHECK-NEXT:    [[A3:%.*]] = getelementptr { i32, i32 }, ptr [[A]], i64 1
 ; CHECK-NEXT:    ret ptr [[A3]]
 ;
   %a2 = getelementptr { i32, i32 }, ptr %a, i32 0, i32 1
   call void @use(ptr %a2)
   %a3 = getelementptr i8, ptr %a2, i64 4
   ret ptr %a3
 }