Allow gep (select) combine with non-constant operands

[rampitec]

Summary

This transformation then allows SROA to eliminate allocas
whose address is being selected.

[arsenm]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1729

demorgan this

1731–1732

Calling GetUnderlyingObject so often seems potentially expensive

[rampitec]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1731–1732

I can just disable this combine unconditionally. In many cases gep is a no-op. However, I in some situations gep really results in some arithmetic.

I.e. it is a tradeoff between code quality and compile time. Also note, GetUnderlyingObject has a threshold, which will limit compile time impact on one hand, but also limit optimization on another.

[rampitec]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1729

JBTW, demorgan comment or logic?

[rampitec]

Fixed comment and restated condition. Not sure it looks better though.

[rampitec]

In fact I have created two more changes: D79217 and D79218 to process select and phi with gep operands right in SROA. It is more code but looks cleaner:

1. It will only touch patterns already considered by SROA.
2. It will give InstCombiner a chance to convert it back if SROA did not succeed.
3. There is zero chance it will loop forever.
4. It does not call expensive GetUnerlyingObject.

If those two look better I would prefer modifications in SROA and drop this.

[spatel]

In D79145#2014188, @rampitec wrote:

In fact I have created two more changes: D79217 and D79218 to process select and phi with gep operands right in SROA. It is more code but looks cleaner:

1. It will only touch patterns already considered by SROA.
2. It will give InstCombiner a chance to convert it back if SROA did not succeed.
3. There is zero chance it will loop forever.
4. It does not call expensive GetUnerlyingObject.

If those two look better I would prefer modifications in SROA and drop this.

If we can deal with this in SROA directly, that would be better. We generally do not want to increase instruction count in InstCombine (as the 2nd test is showing here).

[rampitec]

In D79145#2014834, @spatel wrote:

In D79145#2014188, @rampitec wrote:

In fact I have created two more changes: D79217 and D79218 to process select and phi with gep operands right in SROA. It is more code but looks cleaner:

1. It will only touch patterns already considered by SROA.
2. It will give InstCombiner a chance to convert it back if SROA did not succeed.
3. There is zero chance it will loop forever.
4. It does not call expensive GetUnerlyingObject.

If those two look better I would prefer modifications in SROA and drop this.

If we can deal with this in SROA directly, that would be better. We generally do not want to increase instruction count in InstCombine (as the 2nd test is showing here).

Right, I hope these two SROA patches can land instead of this.

[arsenm]

Can this be abandoned now?

[rampitec]

In D79145#2058689, @arsenm wrote:

Can this be abandoned now?

D79217 only partially solves the problem, there is also D79218 needed. I am going to abandon this anyway.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

Instruction *InstCombiner::foldSelectOpOp(SelectInst &SI, Instruction *TI,
   // one-use here. As with the cast case above, it may be possible to relax the
   // one-use constraint, but that needs be examined carefully since it may not
   // reduce the total number of instructions.
   if (TI->getNumOperands() != 2 || FI->getNumOperands() != 2 ||
       (!isa<BinaryOperator>(TI) && !isa<GetElementPtrInst>(TI)) ||
       !TI->hasOneUse() || !FI->hasOneUse())
     return nullptr;
 
+  // Do not fold GEP leading to an alloca, foldSelectGEP() will revert it.
+  if (isa<GetElementPtrInst>(TI) &&
+      (isa<AllocaInst>(GetUnderlyingObject(TI, DL)) ||
+       isa<AllocaInst>(GetUnderlyingObject(FI, DL))))
+    return nullptr;
+
   // Figure out if the operations have any operands in common.
   Value *MatchOp, *OtherOpT, *OtherOpF;
   bool MatchIsOpZero;
   if (TI->getOperand(0) == FI->getOperand(0)) {
     MatchOp  = TI->getOperand(0);
     OtherOpT = TI->getOperand(1);
     OtherOpF = FI->getOperand(1);
     MatchIsOpZero = true;

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

static bool isMergedGEPInBounds(GEPOperator &GEP1, GEPOperator &GEP2) {
 
   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) {
+                                  InstCombiner::BuilderTy &Builder,
+                                  const DataLayout &DL) {
   if (!GEP.hasAllConstantIndices())
     return nullptr;
 
   Instruction *Sel;
   Value *Cond;
-  Constant *TrueC, *FalseC;
+  Value *True, *False;
   if (!match(GEP.getPointerOperand(), m_Instruction(Sel)) ||
       !match(Sel,
-             m_Select(m_Value(Cond), m_Constant(TrueC), m_Constant(FalseC))))
+             m_Select(m_Value(Cond), m_Value(True), m_Value(False))))
+    return nullptr;
+
+  // Allow transformation if both sides are constant or if either refers
+  // to alloca to help SROA.
+  if ((!isa<Constant>(True) || !isa<Constant>(False)) &&
+      !isa<AllocaInst>(GetUnderlyingObject(True, DL)) &&
+      !isa<AllocaInst>(GetUnderlyingObject(False, DL)))
     return nullptr;
 
-  // gep (select Cond, TrueC, FalseC), IndexC --> select Cond, TrueC', FalseC'
+  // gep (select Cond, True, False), IndexC --> select Cond, True', False'
   // 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);
+  Value *NewTrue = IsInBounds ? Builder.CreateInBoundsGEP(True, IndexC)
+                              : Builder.CreateGEP(True, IndexC);
+  Value *NewFalse = IsInBounds ? Builder.CreateInBoundsGEP(False, IndexC)
+                               : Builder.CreateGEP(False, IndexC);
+  return SelectInst::Create(Cond, NewTrue, NewFalse, "", 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();
   bool IsGEPSrcEleScalable = isa<ScalableVectorType>(GEPEltType);
   if (Value *V = SimplifyGEPInst(GEPEltType, Ops, SQ.getWithInstruction(&GEP)))
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))
+  if (Instruction *R = foldSelectGEP(GEP, Builder, DL))
     return R;
 
   return nullptr;
 }
 
 static bool isNeverEqualToUnescapedAlloc(Value *V, const TargetLibraryInfo *TLI,
                                          Instruction *AI) {
   if (isa<ConstantPointerNull>(V))

llvm/test/Transforms/InstCombine/getelementptr.ll

 ; 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
 }
 
+define i32* @test_gep_select_arg(i1 %cond, %struct.f* %a, %struct.f* %b) {
+; CHECK-LABEL: @test_gep_select_arg(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], %struct.f* [[A:%.*]], %struct.f* [[B:%.*]], !prof !0
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [[STRUCT_F:%.*]], %struct.f* [[SEL]], i64 0, i32 0
+; CHECK-NEXT:    ret i32* [[GEP]]
+;
+  %sel = select i1 %cond, %struct.f* %a, %struct.f* %b, !prof !0
+  %gep = getelementptr inbounds %struct.f, %struct.f* %sel, i64 0, i32 0
+  ret i32* %gep
+}
+
+define i32* @test_gep_select_alloca(i1 %cond) {
+; CHECK-LABEL: @test_gep_select_alloca(
+; CHECK-NEXT:    [[A:%.*]] = alloca [[STRUCT_F:%.*]], align 8
+; CHECK-NEXT:    [[B:%.*]] = alloca [[STRUCT_F]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [[STRUCT_F]], %struct.f* [[A]], i64 0, i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [[STRUCT_F]], %struct.f* [[B]], i64 0, i32 0
+; CHECK-NEXT:    [[GEP:%.*]] = select i1 [[COND:%.*]], i32* [[TMP1]], i32* [[TMP2]], !prof !0
+; CHECK-NEXT:    ret i32* [[GEP]]
+;
+  %a = alloca %struct.f
+  %b = alloca %struct.f
+  %sel = select i1 %cond, %struct.f* %a, %struct.f* %b, !prof !0
+  %gep = getelementptr inbounds %struct.f, %struct.f* %sel, i64 0, i32 0
+  ret i32* %gep
+}
+
 !0 = !{!"branch_weights", i32 2, i32 10}

llvm/test/Transforms/SROA/select-gep.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -instcombine -sroa < %s | FileCheck %s
+
+%pair = type { i32, i32 }
+
+define i32 @test_sroa_select_gep(i1 %cond) {
+; CHECK-LABEL: @test_sroa_select_gep(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[LOAD:%.*]] = select i1 [[COND:%.*]], i32 undef, i32 undef
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+bb:
+  %a = alloca %pair, align 4
+  %b = alloca %pair, align 4
+  %select = select i1 %cond, %pair* %a, %pair* %b
+  %gep = getelementptr inbounds %pair, %pair* %select, i32 0, i32 1
+  %load = load i32, i32* %gep, align 4
+  ret i32 %load
+}