[InstCombine] fold gep-of-select-of-constants (PR45084)

[spatel]

Summary

As shown in PR45084:
https://bugs.llvm.org/show_bug.cgi?id=45084
...we failed to combine a gep with constant indexes with a pointer operand that is a select of constants.

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1816

Should the select be limited to one use? I'm not sure whether this transform is favorable if we need to keep both selects. In either case, this should be tested.

1822

It is better to construct these GEPs via IRBuilder, so the resulting constants get folded right away (GEPs are target-dependent).

Also, aren't we losing inbounds here?

1823

We can preserve branch weight metadata here.

[spatel]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1816

Unless there's evidence that the select-of-constants form is less favorable, we want to allow the transform because it eliminates a use of the existing select. I'll add a test.

1822

Can you describe the benefit of IRBuilder vs. ConstantExpr::get()? I'm not seeing any difference in the debug output for the examples that I have.

I'm happy to change it (but we should probably change a bunch of existing similar code too?).

I'm not clear on the inbounds rules here (it's a constant, so is it not required to be inbounds?). But yes, I'll fix this to inherit from the existing instruction.

1823

Yes, good catch - I'll add to the test.

[spatel]

Patch updated:

1. Propagate metadata from select (test adjusted).
2. Propagate 'inbounds' from gep (not sure how to test non-inbounds case).
3. Added test for >1 use of the select.

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1822

Can you describe the benefit of IRBuilder vs. ConstantExpr::get()? I'm not seeing any difference in the debug output for the examples that I have.

Sure. Take something like

%struct.f = type { i32, i32 }
  
@g0 = internal unnamed_addr constant %struct.f zeroinitializer, align 4

define i32* @PR45084(i1 %cond) {
  %sel = select i1 %cond, %struct.f* @g0, %struct.f* null
  %gep = getelementptr inbounds %struct.f, %struct.f* %sel, i64 0, i32 1
  ret i32* %gep
}

In the log we see:

IC: Old =   %gep = getelementptr inbounds %struct.f, %struct.f* %sel, i64 0, i32 1
    New =   <badref> = select i1 %cond, i32* getelementptr inbounds (%struct.f, %struct.f* @g0, i64 0, i32 1), i32* getelementptr inbounds (%struct.f, %struct.f* null, i64 0, i32 1)
...
IC: ConstFold operand of:   %gep = select i1 %cond, i32* getelementptr inbounds (%struct.f, %struct.f* @g0, i64 0, i32 1), i32* getelementptr inbounds (%struct.f, %struct.f* null, i64 0, i32 1)
    Old = i32* getelementptr inbounds (%struct.f, %struct.f* null, i64 0, i32 1)
    New = i32* inttoptr (i64 4 to i32*)

This is a target-depending fold, so it is not performed during ConstantExpr construction. It will be automatically performed if you go through IRBuilder, and save the need to do this on the next IC iteration.

I'm happy to change it (but we should probably change a bunch of existing similar code too?).

Existing code seems to get this mostly right, though I did fix one occurrence in 9b5de84e274fe9df971e5d32cc4ffe1ab972519d.

I'm not clear on the inbounds rules here (it's a constant, so is it not required to be inbounds?). But yes, I'll fix this to inherit from the existing instruction.

If that were the case, I would assume that the inbounds flag wouldn't exist on ConstantExpr GEP in the first place. I think it's fine to have non-inbounds GEP on constants. E.g. my example above is wrong, the GEP on null can't be inbounds. I don't know if any of the ConstantExpr actually optimizes on inbounds though, I didn't find anything from a quick search.

[spatel]

Patch updated:
Use IRBuilder to efficiently create the GEP constant expressions.

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1829

Surprised that IRBuilder has no overloads accepting bool...

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

static bool isMergedGEPInBounds(GEPOperator &GEP1, GEPOperator &GEP2) {
   // At least one GEP must be inbounds.
   if (!GEP1.isInBounds() && !GEP2.isInBounds())
     return false;
 
   return (GEP1.isInBounds() || GEP1.hasAllZeroIndices()) &&
          (GEP2.isInBounds() || GEP2.hasAllZeroIndices());
 }
 
+/// Thread a GEP operation with constant indices through the constant true/false
+/// arms of a select.
+static Instruction *foldSelectGEP(GetElementPtrInst &GEP,
+                                  InstCombiner::BuilderTy &Builder) {
+  if (!GEP.hasAllConstantIndices())
+    return nullptr;
+
+  Instruction *Sel;
+  Value *Cond;
+  Constant *TrueC, *FalseC;
+  if (!match(GEP.getPointerOperand(), m_Instruction(Sel)) ||
+      !match(Sel,
+             m_Select(m_Value(Cond), m_Constant(TrueC), m_Constant(FalseC))))
+    return nullptr;
+
+  // gep (select Cond, TrueC, FalseC), IndexC --> select Cond, TrueC', FalseC'
+  // Propagate 'inbounds' and metadata from existing instructions.
+  // Note: using IRBuilder to create the constants for efficiency.
+  SmallVector<Value *, 4> IndexC(GEP.idx_begin(), GEP.idx_end());
+  bool IsInBounds = GEP.isInBounds();
+  Value *NewTrueC = IsInBounds ? Builder.CreateInBoundsGEP(TrueC, IndexC)
+                               : Builder.CreateGEP(TrueC, IndexC);
+  Value *NewFalseC = IsInBounds ? Builder.CreateInBoundsGEP(FalseC, IndexC)
+                                : Builder.CreateGEP(FalseC, IndexC);
+  return SelectInst::Create(Cond, NewTrueC, NewFalseC, "", nullptr, Sel);
+}
+
 Instruction *InstCombiner::visitGetElementPtrInst(GetElementPtrInst &GEP) {
   SmallVector<Value*, 8> Ops(GEP.op_begin(), GEP.op_end());
   Type *GEPType = GEP.getType();
   Type *GEPEltType = GEP.getSourceElementType();
   if (Value *V = SimplifyGEPInst(GEPEltType, Ops, SQ.getWithInstruction(&GEP)))
     return replaceInstUsesWith(GEP, V);
 
   // For vector geps, use the generic demanded vector support.
if (auto *AI = dyn_cast<AllocaInst>(UnderlyingPtrOp)) {
           return GetElementPtrInst::CreateInBounds(
               GEP.getSourceElementType(), PtrOp, makeArrayRef(Ops).slice(1),
               GEP.getName());
         }
       }
     }
   }
 
+  if (Instruction *R = foldSelectGEP(GEP, Builder))
+    return R;
+
   return nullptr;
 }
 
 static bool isNeverEqualToUnescapedAlloc(Value *V, const TargetLibraryInfo *TLI,
                                          Instruction *AI) {
   if (isa<ConstantPointerNull>(V))
     return true;
   if (auto *LI = dyn_cast<LoadInst>(V))

llvm/test/Transforms/InstCombine/getelementptr.ll

 
 %struct.f = type { i32 }
 
 @g0 = internal unnamed_addr constant %struct.f zeroinitializer, align 4
 @g1 = internal unnamed_addr constant %struct.f { i32 -1 }, align 4
 
 define i32* @PR45084(i1 %cond) {
 ; CHECK-LABEL: @PR45084(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], %struct.f* @g0, %struct.f* @g1, !prof !0
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [[STRUCT_F:%.*]], %struct.f* [[SEL]], i64 0, i32 0
+; CHECK-NEXT:    [[GEP:%.*]] = select i1 [[COND:%.*]], i32* getelementptr inbounds (%struct.f, %struct.f* @g0, i64 0, i32 0), i32* getelementptr inbounds (%struct.f, %struct.f* @g1, i64 0, i32 0), !prof !0
 ; CHECK-NEXT:    ret i32* [[GEP]]
 ;
   %sel = select i1 %cond, %struct.f* @g0, %struct.f* @g1, !prof !0
   %gep = getelementptr inbounds %struct.f, %struct.f* %sel, i64 0, i32 0
   ret i32* %gep
 }
 
 define i32* @PR45084_extra_use(i1 %cond, %struct.f** %p) {
 ; CHECK-LABEL: @PR45084_extra_use(
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], %struct.f* @g0, %struct.f* @g1
 ; CHECK-NEXT:    store %struct.f* [[SEL]], %struct.f** [[P:%.*]], align 8
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [[STRUCT_F:%.*]], %struct.f* [[SEL]], i64 0, i32 0
+; CHECK-NEXT:    [[GEP:%.*]] = select i1 [[COND]], i32* getelementptr inbounds (%struct.f, %struct.f* @g0, i64 0, i32 0), i32* getelementptr inbounds (%struct.f, %struct.f* @g1, i64 0, i32 0)
 ; CHECK-NEXT:    ret i32* [[GEP]]
 ;
   %sel = select i1 %cond, %struct.f* @g0, %struct.f* @g1
   store %struct.f* %sel, %struct.f** %p
   %gep = getelementptr %struct.f, %struct.f* %sel, i64 0, i32 0
   ret i32* %gep
 }
 
 !0 = !{!"branch_weights", i32 2, i32 10}

llvm/test/Transforms/InstCombine/mem-gep-zidx.ll

;
   %arrayidx = getelementptr inbounds i64, i64* %p, i64 %idxprom
   store i64 %v, i64* %arrayidx
   call void @foo(i64* %p)
   ret void
 }
 
 define signext i32 @test3(i32 signext %x, i1 %y) #0 {
 ; CHECK-LABEL: @test3(
-; CHECK-NEXT:    [[P:%.*]] = select i1 [[Y:%.*]], [1 x i32]* @f.a, [1 x i32]* @f.b
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [1 x i32], [1 x i32]* [[P]], i64 0, i64 0
-; CHECK-NEXT:    [[R:%.*]] = load i32, i32* [[TMP1]], align 4
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[Y:%.*]], i32 12, i32 55
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %idxprom = sext i32 %x to i64
   %p = select i1 %y, [1 x i32]* @f.a, [1 x i32]* @f.b
   %arrayidx = getelementptr inbounds [1 x i32], [1 x i32]* %p, i64 0, i64 %idxprom
   %r = load i32, i32* %arrayidx, align 4
   ret i32 %r
 }