Add support for GEP of GEP instruction combining when types do not exactly match

[simeon]

Summary

InstCombine can combine a GEP of a GEP into a single one if the type of the source pointer matches the source element type of the outer pointer. However, in certain cases the source pointer may have had one or more of its indices folded (as CreateGEP calls FoldGEP) and thus have a different resulting type than the source element type, leading to a missed optimization. This patch attempts to handle this case by "unfolding" the source GEP by incrementally appending zero indices until its type matches.

[simeon]

Gentle ping

[nikic]

Can you please explain what the larger context/motivation for this is?

Longer term at least, I think we would be better off going in the opposite direction and making sure that each GEP has at most one variable index. For example, if you have two GEPs with variable indices and one of them is loop invariant, LICM will be able to move it out of the loop (possibly after reassociating the two GEPs). This is not possible if both indices are in the same GEP.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

for (const APInt &Idx : drop_begin(ConstIndices, !IsFirstType)) {
       IsInBounds &= Idx.isNonNegative() == ConstIndices[0].isNonNegative();
     }
 
     return replaceInstUsesWith(
         GEP, Builder.CreateGEP(Src->getSourceElementType(), Src->getOperand(0),
                                Indices, "", IsInBounds));
   }
 
-  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();
 
   // Can we combine the two pointer arithmetics offsets?
-  if (EndsWithSequential) {
+  if (Src->getResultElementType() == GEP.getSourceElementType() && 
+      EndsWithSequential) {
     // Replace: gep (gep %P, long B), long A, ...
     // With:    T = long A+B; gep %P, T, ...
     Value *SO1 = Src->getOperand(Src->getNumOperands()-1);
     Value *GO1 = GEP.getOperand(1);
 
     // If they aren't the same type, then the input hasn't been processed
     // by the loop above yet (which canonicalizes sequential index types to
     // intptr_t).  Just avoid transforming this until the input has been
if (Src->getNumOperands() == 2) {
       GEP.setIsInBounds(isMergedGEPInBounds(*Src, *cast<GEPOperator>(&GEP)));
       replaceOperand(GEP, 0, Src->getOperand(0));
       replaceOperand(GEP, 1, Sum);
       return &GEP;
     }
     Indices.append(Src->op_begin()+1, Src->op_end()-1);
     Indices.push_back(Sum);
     Indices.append(GEP.op_begin()+2, GEP.op_end());
-  } else if (isa<Constant>(*GEP.idx_begin()) &&
+  } else if (Src->getResultElementType() == GEP.getSourceElementType() &&
+             isa<Constant>(*GEP.idx_begin()) &&
              cast<Constant>(*GEP.idx_begin())->isNullValue() &&
              Src->getNumOperands() != 1) {
     // Otherwise we can do the fold if the first index of the GEP is a zero
     Indices.append(Src->op_begin()+1, Src->op_end());
     Indices.append(GEP.idx_begin()+1, GEP.idx_end());
+  } else { 
+    Indices.append(Src->op_begin()+1, Src->op_end());
+
+    // The indices of GEP may have been folded; in such case try to get
+    // the types to match by appending zeros to the indices of Src
+    Type *SrcType = Src->getResultElementType();
+    while (SrcType && SrcType->isAggregateType() && 
+           SrcType != GEP.getSourceElementType()) {
+        Indices.push_back(Builder.getInt32(0));
+        SrcType = GetElementPtrInst::getIndexedType(
+          Src->getSourceElementType(), Indices);
+    }
+    if (SrcType == Src->getResultElementType() || 
+        SrcType != GEP.getSourceElementType())
+      return nullptr;
+
+    Indices.pop_back();
+    Indices.append(GEP.idx_begin(), GEP.idx_end());
   }
 
   if (!Indices.empty())
     return replaceInstUsesWith(
         GEP, Builder.CreateGEP(
                  Src->getSourceElementType(), Src->getOperand(0), Indices, "",
                  isMergedGEPInBounds(*Src, *cast<GEPOperator>(&GEP))));
 

llvm/test/Transforms/InstCombine/gepgep-after-fold.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32-S128"
+
+%s1 = type [3 x [3 x i32]]
+%s2 = type [3 x [3 x [3 x i32]]]
+
+define void @test_unfold_one(ptr %x, i32 %i, i32 %j) {
+; CHECK-LABEL: @test_unfold_one(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [3 x [3 x i32]], ptr [[X:%.*]], i32 [[I:%.*]], i32 2, i32 [[J:%.*]]
+; CHECK-NEXT:    store i32 1, ptr [[GEP]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %src = getelementptr inbounds %s1, ptr %x, i32 %i, i32 2
+  %gep = getelementptr inbounds i32, ptr %src, i32 %j
+  store i32 1, ptr %gep
+  ret void
+}
+
+define void @test_unfold_one2(ptr %x, i32 %i, i32 %j) {
+; CHECK-LABEL: @test_unfold_one2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [3 x [3 x [3 x i32]]], ptr [[X:%.*]], i32 [[I:%.*]], i32 2, i32 [[J:%.*]]
+; CHECK-NEXT:    store i32 1, ptr [[GEP]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %src = getelementptr inbounds %s2, ptr %x, i32 %i, i32 2
+  %gep = getelementptr inbounds [3 x i32], ptr %src, i32 %j
+  store i32 1, ptr %gep
+  ret void
+}
+
+define void @test_unfold_two(ptr %x, i32 %i, i32 %j) {
+; CHECK-LABEL: @test_unfold_two(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [3 x [3 x [3 x i32]]], ptr [[X:%.*]], i32 [[I:%.*]], i32 2, i32 0, i32 [[J:%.*]]
+; CHECK-NEXT:    store i32 1, ptr [[GEP]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %src = getelementptr inbounds %s2, ptr %x, i32 %i, i32 2
+  %gep = getelementptr inbounds i32, ptr %src, i32 %j
+  store i32 1, ptr %gep
+  ret void
+}
+
+define void @test_negative(ptr %x, i32 %i, i32 %j) {
+; CHECK-LABEL: @test_negative(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SRC:%.*]] = getelementptr inbounds [3 x [3 x i32]], ptr [[X:%.*]], i32 [[I:%.*]], i32 2
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [3 x i32], ptr [[SRC]], i32 [[J:%.*]]
+; CHECK-NEXT:    store i32 1, ptr [[GEP]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %src = getelementptr inbounds %s1, ptr %x, i32 %i, i32 2
+  ; should be %gep = getelementptr inbounds [3 x i32], ptr %src, i32 0, i32 %j
+  %gep = getelementptr inbounds [3 x i32], ptr %src, i32 %j
+  store i32 1, ptr %gep
+  ret void
+}