[InstCombine] Simplify chain of GEP with constant indices

[huangjd]

Summary

Simplify chain of GEP in the form of GEP (GEP ... (GEP (GEP p C1) X1) ... Xn) C2 where C1 and C2 are constants by splitting GEP p C1 and breaking its user depedency, allowing visitGEPofGEP to reassociate the cloned inner GEP to outer GEP and perform constant index merging.

Example code that is optimized

struct S {
    char padding[123];
    int x;
};

int* P(S* s) {
    return (int*)((char*)&s->x + s->x - offsetof(S, x));
}

originally generates

P(S*):                                # @P(S*)
        movslq  124(%rdi), %rax
        addq    %rdi, %rax
        addq    $124, %rax
        addq    $-124, %rax
        retq

Now generates

P(S*):                                # @P(S*)
	movslq	124(%rdi), %rax
	addq	%rdi, %rax
	retq

[davidxl]

This is essentially forward propagation (with multi-use) and is a generally a good way to reduce dependence height. For instance it can be used to enable reassociation and LICM:

for (k = .... ;k++) {

for (i = ....  ; i++) {

      b = (a+i);
      call(b);
      y = b + k + x;
    }

}

>

for (k = .... ;k++) {

for (i = ....  ; i++) {

      b = (a+i);
      call(b);
      y  = (a + i) + k + x;
    }

}

>

for (k = .... ;k++) {

for (i = ....  ; i++) {

      b = (a+i);
      call(b);
      y  = (a + k) + i + x;
    }

}

>

for (k = .... ;k++) {

tmp = a + k;
for (i = ....  ; i++) {
      b = (a+i);
      call(b);
      y  = tmp + k + x;
    }

}

[davidxl]

y  = tmp + k + x;

should be

y  = tmp + i + x;

}

}

[spatel]

Please pre-commit the baseline tests as an NFC patch, so we'll just show the test diffs here.

[huangjd]

See D137664 for baseline

[huangjd]

Merged after adding baseline

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *visitSelectInst(SelectInst &SI);
   Instruction *visitCallInst(CallInst &CI);
   Instruction *visitInvokeInst(InvokeInst &II);
   Instruction *visitCallBrInst(CallBrInst &CBI);
 
   Instruction *SliceUpIllegalIntegerPHI(PHINode &PN);
   Instruction *visitPHINode(PHINode &PN);
   Instruction *visitGetElementPtrInst(GetElementPtrInst &GEP);
+  Value *splitConstantGEP(Value *Src);
   Instruction *visitGEPOfGEP(GetElementPtrInst &GEP, GEPOperator *Src);
   Instruction *visitGEPOfBitcast(BitCastInst *BCI, GetElementPtrInst &GEP);
   Instruction *visitAllocaInst(AllocaInst &AI);
   Instruction *visitAllocSite(Instruction &FI);
   Instruction *visitFree(CallInst &FI, Value *FreedOp);
   Instruction *visitLoadInst(LoadInst &LI);
   Instruction *visitStoreInst(StoreInst &SI);
   Instruction *visitAtomicRMWInst(AtomicRMWInst &SI);

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

static Instruction *foldSelectGEP(GetElementPtrInst &GEP,
   SmallVector<Value *, 4> IndexC(GEP.indices());
   bool IsInBounds = GEP.isInBounds();
   Type *Ty = GEP.getSourceElementType();
   Value *NewTrueC = Builder.CreateGEP(Ty, TrueC, IndexC, "", IsInBounds);
   Value *NewFalseC = Builder.CreateGEP(Ty, FalseC, IndexC, "", IsInBounds);
   return SelectInst::Create(Cond, NewTrueC, NewFalseC, "", nullptr, Sel);
 }
 
+// For a chain of GEP in the form GEP (GEP ... (GEP (GEP p C1) X1) ... Xn) C2
+// where C1 and C2 are constants, and GEP X1 ... Xn each has no other uses,
+// we can split GEP p C1 from the use-def chain so that it can be reassociated
+// and merged with the outermost GEP by visitGEPofGEP.
+Value *InstCombinerImpl::splitConstantGEP(Value *Src) {
+  if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(Src)) {
+    APInt Offset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
+    if (GEP->hasAllConstantIndices() &&
+        GEP->accumulateConstantOffset(DL, Offset))
+      return Builder.CreateGEP(GEP->getSourceElementType(),
+                               GEP->getPointerOperand(),
+                               SmallVector<Value *>(GEP->indices()),
+                               GEP->getName(), GEP->isInBounds());
+    if (Src->hasOneUse())
+      if (Value *NewV = splitConstantGEP(GEP->getPointerOperand()))
+        return Builder.CreateGEP(GEP->getSourceElementType(), NewV,
+                                 SmallVector<Value *>(GEP->indices()),
+                                 GEP->getName(), GEP->isInBounds());
+  }
+  return nullptr;
+}
+
 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 (!shouldMergeGEPs(*cast<GEPOperator>(&GEP), *Src))
     return nullptr;
 
if (auto *PN = dyn_cast<PHINode>(PtrOp)) {
     replaceOperand(GEP, 0, NewGEP);
     PtrOp = NewGEP;
   }
 
   if (auto *Src = dyn_cast<GEPOperator>(PtrOp))
     if (Instruction *I = visitGEPOfGEP(GEP, Src))
       return I;
 
+  // This check must be after visitGEPOfGEP, since GEP (GEP p C1) C2 can be
+  // merged by it. If GEP has zero offset, do not transform because it is a zero
+  // cost bitcast.
+  APInt Offset(DL.getIndexTypeSizeInBits(GEP.getType()), 0);
+  if (GEP.hasAllConstantIndices() && !GEP.hasAllZeroIndices() &&
+      GEP.accumulateConstantOffset(DL, Offset) && !Offset.isZero())
+    if (Value *V = splitConstantGEP(GEP.getPointerOperand())) {
+      GetElementPtrInst *NewGEP = GetElementPtrInst::Create(
+          GEP.getSourceElementType(), V, SmallVector<Value *>(GEP.indices()));
+      NewGEP->setIsInBounds(GEP.isInBounds());
+      return NewGEP;
+    }
+
   // Skip if GEP source element type is scalable. The type alloc size is unknown
   // at compile-time.
   if (GEP.getNumIndices() == 1 && !IsGEPSrcEleScalable) {
     unsigned AS = GEP.getPointerAddressSpace();
     if (GEP.getOperand(1)->getType()->getScalarSizeInBits() ==
         DL.getIndexSizeInBits(AS)) {
       uint64_t TyAllocSize = DL.getTypeAllocSize(GEPEltType).getFixedSize();
 

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

 
 ; If the first and last in a chain of GEP are constant indexed, they can be
 ; merged, even if the first GEP has multiple uses.
 ; result = (i8*) ((i32*) p + a + b) + 22
 define ptr @gepChain1(ptr %p, i64 %a, i64 %b) {
 ; CHECK-LABEL: @gepChain1(
 ; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i64, ptr [[P:%.*]], i64 2
 ; CHECK-NEXT:    call void @use(ptr nonnull [[TMP1]])
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i32, ptr [[TMP1]], i64 [[A:%.*]]
-; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds i32, ptr [[TMP2]], i64 [[B:%.*]]
-; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr inbounds i16, ptr [[TMP3]], i64 3
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr i32, ptr [[P]], i64 [[A:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr i32, ptr [[TMP2]], i64 [[B:%.*]]
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[TMP3]], i64 22
 ; CHECK-NEXT:    ret ptr [[TMP4]]
 ;
   %1 = getelementptr inbounds i64, ptr %p, i64 2
   call void @use(ptr %1)
   %2 = getelementptr inbounds i32, ptr %1, i64 %a
   %3 = getelementptr inbounds i32, ptr %2, i64 %b
   %4 = getelementptr inbounds i16, ptr %3, i64 3
   ret ptr %4
 }
 
 ; Constant indexed GEP cancel out.
 ; result = (i32*) p + a + b
 define ptr @gepChain2(ptr %p, i64 %a, i64 %b) {
 ; CHECK-LABEL: @gepChain2(
 ; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i64, ptr [[P:%.*]], i64 2
 ; CHECK-NEXT:    call void @use(ptr nonnull [[TMP1]])
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i32, ptr [[TMP1]], i64 [[A:%.*]]
-; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds i32, ptr [[TMP2]], i64 [[B:%.*]]
-; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr inbounds i16, ptr [[TMP3]], i64 -8
-; CHECK-NEXT:    ret ptr [[TMP4]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr i32, ptr [[P]], i64 [[A:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr i32, ptr [[TMP2]], i64 [[B:%.*]]
+; CHECK-NEXT:    ret ptr [[TMP3]]
 ;
   %1 = getelementptr inbounds i64, ptr %p, i64 2
   call void @use(ptr %1)
   %2 = getelementptr inbounds i32, ptr %1, i64 %a
   %3 = getelementptr inbounds i32, ptr %2, i64 %b
   %4 = getelementptr inbounds i16, ptr %3, i64 -8
   ret ptr %4
 }