[InstCombine] Merging constant-indexed GEP of GEP in both directions

[huangjd]

Summary

Previously GEP of GEP with constant indices are only merged to the front (Second GEP is merged to the first one whenever possible). This patch expanded the optimization to also try merging to the direction under some constraints on GEP usage, if the first attempt failed.

[nikic]

I'm not a big fan of the "try to combine in both directions approach". It will still miss cases where neither direction works, but the GEPs can still be folded. Something like this should work for all constant GEPs: https://gist.github.com/nikic/8a1a82382468c88f19094d4b0a2f547c

[huangjd]

@nikic
I will add converting constant GEP to i8* in the next patch. The existing merging functionality can handle GEP with not all constant indices, such as

%1 = gep {i32, i32, i32}, ptr %p, i64 %a, i32 1
%2 = gep i32, ptr %1, i64 1

Such expression cannot be convert to GEP of i8 without creating a complicated index expression, which is basically moving instruction selection to the front. It is unreasonable that merge only allows one direction but not the other, given use-def constraints satisfied.

[nikic]

@huangjd The not all constant case continues to be handled. If you mean the case where we have ((p + C) + x), then this should first be converted to ((p + x) + C` by the previously discussed transform to hoist constant GEPs, at which point the constant GEP is in the right position of this transform again. There is no need to additionally perform a swap at this point.

[huangjd]

Added fallback if merging GEP of GEP with all constant indices fails on both directions, convert the expression into GEP of i8, using the sum of both GEP's offsets

[huangjd]

updated comment on a test case after adding feature

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Instruction *InstCombinerImpl::visitGEPOfGEP(GetElementPtrInst &GEP,
   // 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() &&
+  if (PtrTy->isOpaquePointerTy() &&
       (Src->hasOneUse() || Src->hasAllConstantIndices())) {
-    // Split Src into a variable part and a constant suffix.
-    gep_type_iterator GTI = gep_type_begin(*Src);
+    // Try to merge GEP (Back of BB) into Src (Front of BB) first, if that
+    // fails, swap the two GEP whenever valid and try again. This is valid as
+    // long as GEP's indices do not depend on Src.
+    GetElementPtrInst *Back = &GEP;
+    GetElementPtrInst *Front = cast<GetElementPtrInst>(Src);
+    for (int i = 0; i < 2; std::swap(Back, Front), i++) {
+      // Swap is not valid if GEP's indices depend on Src. For example
+      // a = gep x, 1
+      // b = gep a, func(a), 1
+      // Also if Src is used by other instructions than GEP, merging constant-
+      // indexed Src to non-constant GEP does not reduce instruction count.
+      if (Front == &GEP && !Src->hasOneUse())
+        break;
+
+      if (Back->hasAllConstantIndices()) {
+        // Split Front into a variable part and a constant suffix.
+        gep_type_iterator GTI = gep_type_begin(Front);
         Type *BaseType = GTI.getIndexedType();
         bool IsFirstType = true;
         unsigned NumVarIndices = 0;
-    for (auto Pair : enumerate(Src->indices())) {
+        for (auto Pair : enumerate(Front->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.
+        // Determine the offset for the constant suffix of Front.
         APInt Offset(DL.getIndexTypeSizeInBits(PtrTy), 0);
-    if (NumVarIndices != Src->getNumIndices()) {
-      // FIXME: getIndexedOffsetInType() does not handled scalable vectors.
+        if (NumVarIndices != Front->getNumIndices()) {
+          // FIXME: getIndexedOffsetInType() does not handle scalable vectors.
           if (isa<ScalableVectorType>(BaseType))
-        return nullptr;
+            continue;
 
           SmallVector<Value *> ConstantIndices;
           if (!IsFirstType)
             ConstantIndices.push_back(
                 Constant::getNullValue(Type::getInt32Ty(GEP.getContext())));
-      append_range(ConstantIndices, drop_begin(Src->indices(), NumVarIndices));
+          append_range(ConstantIndices,
+                       drop_begin(Front->indices(), NumVarIndices));
           Offset += DL.getIndexedOffsetInType(BaseType, ConstantIndices);
         }
 
-    // Add the offset for GEP (which is fully constant).
-    if (!GEP.accumulateConstantOffset(DL, Offset))
-      return nullptr;
+        // Add the offset for Back (which is fully constant).
+        if (!Back->accumulateConstantOffset(DL, Offset))
+          continue;
 
+        APInt OffsetOld = Offset;
         // Convert the total offset back into indices.
         SmallVector<APInt> ConstIndices =
             DL.getGEPIndicesForOffset(BaseType, Offset);
-    if (!Offset.isZero() || (!IsFirstType && !ConstIndices[0].isZero()))
-      return nullptr;
+        if (!Offset.isZero() || (!IsFirstType && !ConstIndices[0].isZero())) {
+          // If both GEP are constant-indexed, and cannot be merged in either
+          // way, convert them to a GEP of i8.
+          if (Front == &GEP && Front->hasAllConstantIndices()) {
+            return isMergedGEPInBounds(*Src, *cast<GEPOperator>(&GEP))
+                       ? GetElementPtrInst::CreateInBounds(
+                             Builder.getInt8Ty(), Src->getOperand(0),
+                             Builder.getInt(OffsetOld), GEP.getName())
+                       : GetElementPtrInst::Create(
+                             Builder.getInt8Ty(), Src->getOperand(0),
+                             Builder.getInt(OffsetOld), GEP.getName());
+          }
+          continue;
+        }
 
         SmallVector<Value *> Indices;
-    append_range(Indices, drop_end(Src->indices(),
-                                   Src->getNumIndices() - NumVarIndices));
+        append_range(Indices, drop_end(Front->indices(),
+                                       Front->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(),
+                   ? GetElementPtrInst::CreateInBounds(
+                         Front->getSourceElementType(), Src->getOperand(0),
+                         Indices, Back->getName())
+                   : GetElementPtrInst::Create(Front->getSourceElementType(),
                                                Src->getOperand(0), Indices,
-                                           GEP.getName());
+                                               Back->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.

llvm/test/Transforms/InstCombine/gep-merge-constant-indices.ll

;
   %2 = getelementptr inbounds i64, ptr %1, i64 1
   ret ptr %2
 }
 
 ; Second GEP is merged into the first.
 ; result = (i8*) p + 10
 define ptr @mergeReverse(ptr %p) {
 ; CHECK-LABEL: @mergeReverse(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i64, ptr [[P:%.*]], i64 1
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i8, ptr [[TMP1]], i64 2
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[P:%.*]], i64 10
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i64, ptr %p, i64 1
   %2 = getelementptr inbounds i8, ptr %1, i64 2
   ret ptr %2
 }
 
 ; Offsets of first and last GEP cancel out.
 ; result = p
;
   %1 = getelementptr inbounds [20 x i8], ptr %p, i64 1, i64 1
   %2 = getelementptr inbounds i64, ptr %1, i64 2
   ret ptr %2
 }
 
 ; result = (i8*) p + 14
 define ptr @array2(ptr %p, i64 %a) {
 ; CHECK-LABEL: @array2(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [7 x i32], ptr [[P:%.*]], i64 0, i64 3
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i8, ptr [[TMP1]], i64 2
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[P:%.*]], i64 14
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds [7 x i32], ptr %p, i64 0, i64 3
   %2 = getelementptr inbounds i8, ptr %1, i64 2
   ret ptr %2
 }
 
 ; result = (i8*) (([3 x i8]*) p + 6) + 2
 define ptr @array3(ptr %p) {
 ; CHECK-LABEL: @array3(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i64, ptr [[P:%.*]], i64 2
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [3 x i8], ptr [[TMP1]], i64 1, i64 1
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [3 x i8], ptr [[P:%.*]], i64 6, i64 2
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i64, ptr %p, i64 2
   %2 = getelementptr inbounds [3 x i8], ptr %1, i64 1, i64 1
   ret ptr %2
 }
 
 ; result = (struct.C*) p + 3
 define ptr @struct1(ptr %p) {
 ; CHECK-LABEL: @struct1(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i64, ptr [[P:%.*]], i64 3
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [[STRUCT_C:%.*]], ptr [[TMP1]], i64 1
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [[STRUCT_C:%.*]], ptr [[P:%.*]], i64 3
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i64, ptr %p, i64 3
   %2 = getelementptr inbounds %struct.C, ptr %1, i64 1
   ret ptr %2
 }
 
 ; result = &((struct.A*) p - 1).member1
 define ptr @struct2(ptr %p) {

 ;
   %1 = getelementptr inbounds i8, ptr %p, i64 -128
   %2 = getelementptr inbounds %struct.A, ptr %1, i64 0, i32 1
   ret ptr %2
 }
 ; result = ((struct.C*) p + 1).member2
 define ptr @struct4(ptr %p, i64 %a) {
 ; CHECK-LABEL: @struct4(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i64, ptr [[P:%.*]], i64 1
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [[STRUCT_C:%.*]], ptr [[TMP1]], i64 1
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [[STRUCT_C:%.*]], ptr [[P:%.*]], i64 1, i32 2
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i64, ptr %p, i64 1
   %2 = getelementptr inbounds %struct.C, ptr %1, i64 1
   ret ptr %2
 }
 
 ; result = (i8*) &((struct.B) p)[0].member2.member0 + 7
 define ptr @structStruct(ptr %p) {
;
   %1 = getelementptr inbounds %struct.B, ptr %p, i64 0, i32 1
   %2 = getelementptr inbounds i32, ptr %1, i64 1
   ret ptr %2
 }
 
 ; result = (i8*) &((struct.A*) &((struct.B*) p)[0].member2).member0 + 2
 define ptr @appendIndexReverse(ptr %p) {
 ; CHECK-LABEL: @appendIndexReverse(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i64, ptr [[P:%.*]], i64 1
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [[STRUCT_B:%.*]], ptr [[TMP1]], i64 0, i32 1
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [[STRUCT_B:%.*]], ptr [[P:%.*]], i64 0, i32 2, i32 0, i64 2
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i64, ptr %p, i64 1
   %2 = getelementptr inbounds %struct.B, ptr %1, i64 0, i32 1
   ret ptr %2
 }
 
 ; Constant-indexed GEP can be merged if the sum of offsets aliases a member's
 ; address of one of the GEP instructions.
 ; result = &((struct.C*) p + 2).member1
 define ptr @structMemberAliasing(ptr %p, i64 %a) {
 ; CHECK-LABEL: @structMemberAliasing(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i64, ptr [[P:%.*]], i64 1
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [[STRUCT_C:%.*]], ptr [[TMP1]], i64 1, i32 2
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [[STRUCT_C:%.*]], ptr [[P:%.*]], i64 2, i32 1
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i64, ptr %p, i64 1
   %2 = getelementptr inbounds %struct.C, ptr %1, i64 1, i32 2
   ret ptr %2
 }
 
-; Negative test. Offset of either GEP is not divisible by the other's size.
+; GEP cannot be merged either direction, but can be canonicalized as GEP of i8.
 ; Here i24 is 8-bit aligned.
 define ptr @notDivisible(ptr %p) {
 ; CHECK-LABEL: @notDivisible(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i24, ptr [[P:%.*]], i64 1
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i32, ptr [[TMP1]], i64 1
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[P:%.*]], i64 7
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i24, ptr %p, i64 1
   %2 = getelementptr inbounds i32, ptr %1, i64 1
   ret ptr %2
 }
 
 ; Two GEP instructions can be merged if one is constant-indexed and the other
 ; is a sequential type with a constant last index, and the constant offset is
 ; divisible by the sequential type size.
 ; result = (i32*) (([4 x i32]*) p + a) + 3
 define ptr @partialConstant1(ptr %p, i64 %a) {
 ; CHECK-LABEL: @partialConstant1(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, ptr [[P:%.*]], i64 1
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [4 x i32], ptr [[TMP1]], i64 [[A:%.*]], i64 2
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [4 x i32], ptr [[P:%.*]], i64 [[A:%.*]], i64 3
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i32, ptr %p, i64 1
   %2 = getelementptr inbounds [4 x i32], ptr %1, i64 %a, i64 2
   ret ptr %2
 }
 
 ; Negative test. Similar to above, but two GEP should not be merged if the
 ; constant offset is not divisible.

 }
 
 ; Two GEP instructions can be merged if one is constant-indexed and the other
 ; is an aggregate type with a constant last index, and the resulting pointer
 ; address by adding the constant offset aliases the address of another member.
 ; result = &((struct.C*) p + a).member2
 define ptr @partialConstantMemberAliasing1(ptr %p, i64 %a) {
 ; CHECK-LABEL: @partialConstantMemberAliasing1(
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, ptr [[P:%.*]], i64 1
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [[STRUCT_C:%.*]], ptr [[TMP1]], i64 [[A:%.*]], i32 1
-; CHECK-NEXT:    ret ptr [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [[STRUCT_C:%.*]], ptr [[P:%.*]], i64 [[A:%.*]], i32 2
+; CHECK-NEXT:    ret ptr [[TMP1]]
 ;
   %1 = getelementptr inbounds i32, ptr %p, i64 1
   %2 = getelementptr inbounds %struct.C, ptr %1, i64 %a, i32 1
   ret ptr %2
 }
 
 ; Negative test. Similar to above, but the new address does not alias the
 ; address of another member.

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

 ;
   %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:    [[A3:%.*]] = getelementptr i8, ptr [[A:%.*]], i64 14
 ; 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) {