[InstCombine] Skip merging non-free GEP

[junbuml]

Summary

Currently, InstCombine aggressively merge GEPs because it can help backends make
better decision with complex addressing modes. However, in some cases, merging
GEPs can cause more instructions. This change will skip merging GEPs if a
non-free GEP is used in multiple basic blocks.

For the test-case (unit_skip_gep_merge.ll), this patch change code generation (-O3) for AArch64 :
from :

foo:
        orr     w8, wzr, #0x18
        madd    x8, x2, x8, x0
        ldr     x8, [x8, #8]
        str     x8, [x4]
        tbz     w1, #0, .LBB0_2
// %bb.1:
        orr     w8, wzr, #0x18
        madd    x8, x2, x8, x0
        ldr     x8, [x8, #16]
        str     x8, [x3]
.LBB0_2:
        ret

to :

foo:
        orr     w8, wzr, #0x18
        madd    x8, x2, x8, x0
        ldr     x9, [x8, #8]
        str     x9, [x4]
        tbz     w1, #0, .LBB0_2
// %bb.1:
        ldr     x8, [x8, #16]
        str     x8, [x3]
.LBB0_2:
        ret

[efriedma]

Do you have perf/size numbers?

I'm sort of confused that MachineCSE doesn't trigger on your testcase, but that isn't really relevant here, I guess.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1779

For array GEPs, we have a check to try to make sure we aren't making a GEP more expensive (see EndsWithSequential etc.) Maybe instead of adding a different kind of check, it would make more sense to extend the existing check to struct GEPs?

[junbuml]

Do you have perf/size numbers?

I didn't see any degradation in size/perf in my spec2000. Just observed minor performance improvement in some internal benchmarks.

I'm sort of confused that MachineCSE doesn't trigger on your testcase, but that isn't really relevant here, I guess.

I also think MachineCSE should be able to catch it in the machine level, but somehow it's not. I will check this too. However, I believe it's better not to distribute the same calculations in early phase unless we have clear benefits with the duplicated calculations.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1779

Did you mean to move this whole check (line 1680~1687) into EndsWithSequential?

[efriedma]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1779

No.

We currently have a check in the "if (EndsWithSequential)" which restricts GEP folding. That doesn't apply to your testcase because the last index is a struct index. Instead, it falls into the "else if (isa<Constant>(*GEP.idx_begin())" case, which doesn't have a profitability check. My suggestion is to adapt the EndsWithSequential profitability check to also apply to struct GEPs.

[junbuml]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1779

The last index (%v) in the test testcase is still a pointer of struct, which is sequential. This testcase actually get into "if (EndsWithSequential)" and the Src (%arrayidx8) is merged with GEP (%tmp1 and %tmp2). Looks like there no such check we have in this patch in "if (EndsWithSequential)".

[efriedma]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1779

Oh, sorry, you're right, I managed to confuse myself based on what we should be doing, rather than what we're actually doing at the moment.

I don't like adding heuristics based on the uses of the GEP; that's likely to lead to inconsistent results, and possibly O(N^2) compile-time.

What I think we should be doing instead is trying to maintain the invariant that each GEP either has all constant offsets, or has exactly one non-zero offset. Roughly, the rule allows GEPs which will lower to one addition instruction. It composes well, other optimizations can already deal with it (IIRC), and it's a straightforward extension of what we're already doing in the "if (EndsWithSequential)" block.

[junbuml]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1779

What I think we should be doing instead is trying to maintain the invariant that each GEP either has all constant offsets, or has exactly one non-zero offset.

I'm afraid if I fully catch what you meant here. I guess you meant that overall we need to be less aggressive in merging GEPs in instcombine, and in instcombine, we try to keep GEPs have either all constant offset or exactly one non-constant offset? If merging GEPs can break this, we do not merge.

It composes well, other optimizations can already deal with it (IIRC)

I guess CodeGenprepare and ISel. Please let me know other passes which handle something like this.

[efriedma]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1779

If merging GEPs can break this, we do not merge.

Yes, that's what I meant.

[junbuml]

Tried to handle Eli's comment. Please take a look and let me know any comment.

[junbuml]

With the last modification I made based on Eli's comment, I didn't see any significant changes in size / performance in my spec2000 test on AArch64. Observed minor performance improvement in some internal benchmarks.

[junbuml]

Kindly ping.

[junbuml]

Kindly ping one more time.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

static bool shouldMergeGEPs(GEPOperator &GEP, GEPOperator &Src) {
   // Src. If Src is not a trivial GEP too, don't combine
   // the indices.
   if (GEP.hasAllZeroIndices() && !Src.hasAllZeroIndices() &&
       !Src.hasOneUse())
     return false;
   return true;
 }
 
+static bool couldBeExpensiveGEP(ArrayRef<const Value *> Indices) {
+  // We try to maintain that GEPs have either all constant offset or exactly
+  // one non-constant offset. If merging GEPs can break this, we do not merge.
+  bool FoundNonConst = false;
+  bool FoundNonZeroConst = false;
+  for (unsigned Idx = 0, Size = Indices.size(); Idx != Size; ++Idx) {
+    if (auto *C = dyn_cast<Constant>(Indices[Idx])) {
+      if (!C->isNullValue())
+        FoundNonZeroConst = true;
+    } else {
+      if (FoundNonConst)
+        return true;
+      FoundNonConst = true;
+    }
+    if (FoundNonConst && FoundNonZeroConst)
+      return true;
+  }
+  return false;
+}
+
 /// Return a value X such that Val = X * Scale, or null if none.
 /// If the multiplication is known not to overflow, then NoSignedWrap is set.
 Value *InstCombiner::Descale(Value *Val, APInt Scale, bool &NoSignedWrap) {
   assert(isa<IntegerType>(Val->getType()) && "Can only descale integers!");
   assert(cast<IntegerType>(Val->getType())->getBitWidth() ==
          Scale.getBitWidth() && "Scale not compatible with value!");
 
   // If Val is zero or Scale is one then Val = Val * Scale.
if (auto *Src = dyn_cast<GEPOperator>(PtrOp)) {
 
     // 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) {
       // 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
       // normalized.
       if (SO1->getType() != GO1->getType())
         return nullptr;
 
       Value *Sum =
           SimplifyAddInst(GO1, SO1, false, false, SQ.getWithInstruction(&GEP));
       // Only do the combine when we are sure the cost after the
       // merge is never more than that before the merge.
       if (Sum == nullptr)
         return nullptr;
 
       // Update the GEP in place if possible.
       if (Src->getNumOperands() == 2) {
+        if (!Src->hasOneUse()) {
+          Indices.push_back(Sum);
+          Indices.append(GEP.op_begin() + 2, GEP.op_end());
+          if (couldBeExpensiveGEP(Indices))
+            return nullptr;
+        }
         GEP.setOperand(0, Src->getOperand(0));
         GEP.setOperand(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()) &&
                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());
     }
 
-    if (!Indices.empty())
+    if (!Indices.empty()) {
+      if (!Src->hasOneUse() && couldBeExpensiveGEP(Indices))
+        return nullptr;
+
       return GEP.isInBounds() && Src->isInBounds()
                  ? GetElementPtrInst::CreateInBounds(
                        Src->getSourceElementType(), Src->getOperand(0), Indices,
                        GEP.getName())
                  : GetElementPtrInst::Create(Src->getSourceElementType(),
                                              Src->getOperand(0), Indices,
                                              GEP.getName());
     }
+  }
 
   if (GEP.getNumIndices() == 1) {
     unsigned AS = GEP.getPointerAddressSpace();
     if (GEP.getOperand(1)->getType()->getScalarSizeInBits() ==
         DL.getIndexSizeInBits(AS)) {
       uint64_t TyAllocSize = DL.getTypeAllocSize(GEPEltType);
 
       bool Matched = false;

llvm/test/Transforms/InstCombine/gepphigep.ll

bb:
   %tmp20 = getelementptr inbounds %struct2, %struct2* %tmp1, i64 %tmp19
   %tmp21 = getelementptr inbounds %struct2, %struct2* %tmp20, i64 0, i32 0
   store i32 0, i32* %tmp21, align 4
   %tmp24 = getelementptr inbounds %struct2, %struct2* %tmp10, i64 0, i32 1
   %tmp25 = load i32, i32* %tmp24, align 4
   ret i32 %tmp25
 
 ; CHECK-LABEL: @test2(
-; CHECK: getelementptr inbounds %struct2, %struct2* %tmp1, i64 %tmp9, i32 0
+; CHECK: getelementptr inbounds %struct2, %struct2* %tmp10, i64 0, i32 0
 ; CHECK: getelementptr inbounds %struct2, %struct2* %tmp1, i64 %tmp19, i32 0
-; CHECK: getelementptr inbounds %struct2, %struct2* %tmp1, i64 %tmp9, i32 1
+; CHECK: getelementptr inbounds %struct2, %struct2* %tmp10, i64 0, i32 1
 }
 
 ; Check that instcombine doesn't insert GEPs before landingpad.
 
 define i32 @test3(%struct3* %dm, i1 %tmp4, i64 %tmp9, i64 %tmp19, i64 %tmp20, i64 %tmp21) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
 bb:
   %tmp = getelementptr inbounds %struct3, %struct3* %dm, i64 0
   br i1 %tmp4, label %bb1, label %bb2

llvm/test/Transforms/InstCombine/unit_skip_gep_merge.ll

+; RUN: opt -instcombine -S < %s |  FileCheck %s
+
+%ST = type {i8*, i8*, i8*}
+declare void @bar(i8*)
+
+; CHECK-LABEL: @foo
+; CHECK-LABEL: entry
+; CHECK: %tmp1 = getelementptr inbounds %ST, %ST* %arrayidx8
+; CHECK-LABEL: BB0
+; CHECK: %tmp2 = getelementptr inbounds %ST, %ST* %arrayidx8
+define void @foo(%ST* %B, i1 %c, i64 %v, i8** %S, i8** %S2) {
+entry:
+  %arrayidx8 = getelementptr inbounds %ST, %ST* %B, i64 %v
+  %tmp1 = getelementptr inbounds %ST, %ST* %arrayidx8, i64 0, i32 1
+  %r = load i8*, i8** %tmp1
+  store i8* %r, i8** %S2
+  br i1 %c, label %BB0, label %BB2
+
+BB0:
+  %tmp2 = getelementptr inbounds %ST, %ST* %arrayidx8, i64 0, i32 2
+  %l = load i8*, i8** %tmp2
+  store i8* %l, i8** %S
+  br label %BB2
+
+BB2:
+  ret void
+}
+