[InstCombine] Make foldSelectOpOp able to handle two-operand getelementptr

[john.brawn]

Summary

Three (or more) operand getelementptrs could plausibly also be handled, but handling only two-operand fits in easily with the existing BinaryOperator handling.

[majnemer]

lib/Transforms/InstCombine/InstCombineSelect.cpp

346–348

Please consistently brace.

347–348

Don't else after a return: https://llvm.org/docs/CodingStandards.html#don-t-use-else-after-a-return

351–353

Is this clang-format'd?

354–355

Brace this too.

[john.brawn]

Adjusted braces and clang-formatted.

[majnemer]

lib/Transforms/InstCombine/InstCombineSelect.cpp

346–356

auto *BO

349

I think this dyn_cast can just be cast<>, then you can avoid the llvm_unreachable.

350

You dyn_cast but unconditionally deference. Should this be cast<> ?

[john.brawn]

Update based on comments:

• Use auto when RHS is a cast
• Use cast instead of dyn_cast once we know we're dealing with GetElementPtrInst.

[john.brawn]

lib/Transforms/InstCombine/InstCombineSelect.cpp

349

I could do that, but that would make GetElementPtrInst be in essence the default case which doesn't really make sense, especially if we end up supporting other instructions here as well.

[john.brawn]

Ping.

[spatel]

I added the tests with current output in rL322733. Can you rebase here so we just see the diffs (please use the script at utils/update_test_checks.py to regenerate the check lines)? I think we'll be good then.

[john.brawn]

Rebased patch.

[spatel]

LGTM.

lib/Transforms/InstCombine/InstCombineSelect.cpp

if (TI->getNumOperands() == 1 && TI->isCast()) {
     // Fold this by inserting a select from the input values.
     Value *NewSI =
         Builder.CreateSelect(SI.getCondition(), TI->getOperand(0),
                              FI->getOperand(0), SI.getName() + ".v", &SI);
     return CastInst::Create(Instruction::CastOps(TI->getOpcode()), NewSI,
                             TI->getType());
   }
 
-  // Only handle binary operators with 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.
-  BinaryOperator *BO = dyn_cast<BinaryOperator>(TI);
-  if (!BO || !TI->hasOneUse() || !FI->hasOneUse())
+  // Only handle binary operators (including two-operand getelementptr) with
+  // 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;
 
   // 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);
if (TI->getOperand(0) == FI->getOperand(0)) {
     return nullptr;
   }
 
   // If we reach here, they do have operations in common.
   Value *NewSI = Builder.CreateSelect(SI.getCondition(), OtherOpT, OtherOpF,
                                       SI.getName() + ".v", &SI);
   Value *Op0 = MatchIsOpZero ? MatchOp : NewSI;
   Value *Op1 = MatchIsOpZero ? NewSI : MatchOp;
+  if (BinaryOperator *BO = dyn_cast<BinaryOperator>(TI))
     return BinaryOperator::Create(BO->getOpcode(), Op0, Op1);
+  else if (GetElementPtrInst *TGEP = dyn_cast<GetElementPtrInst>(TI)) {
+    GetElementPtrInst *FGEP = dyn_cast<GetElementPtrInst>(FI);
+    Type *ElementType = TGEP->getResultElementType();
+    return TGEP->isInBounds() && FGEP->isInBounds()
+      ? GetElementPtrInst::CreateInBounds(ElementType, Op0, { Op1 })
+      : GetElementPtrInst::Create(ElementType, Op0, { Op1 });
+  } else
+    llvm_unreachable("Expected BinaryOperator or GEP");
+  return nullptr;
 }
 
 static bool isSelect01(const APInt &C1I, const APInt &C2I) {
   if (!C1I.isNullValue() && !C2I.isNullValue()) // One side must be zero.
     return false;
   return C1I.isOneValue() || C1I.isAllOnesValue() ||
          C2I.isOneValue() || C2I.isAllOnesValue();
 }

test/Transforms/InstCombine/select-gep.ll

+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+; CHECK-LABEL: @test1a
+; CHECK: [[CMP:%.*]] = icmp ugt i32* %p, %q
+; CHECK: [[SEL:%.*]] = select i1 [[CMP]], i32* %p, i32* %q
+; CHECK: [[GEP:%.*]] = getelementptr i32, i32* [[SEL]], i64 4
+; CHECK: ret i32* [[GEP]]
+define i32* @test1a(i32* %p, i32* %q) {
+  %gep1 = getelementptr i32, i32* %p, i64 4
+  %gep2 = getelementptr i32, i32* %q, i64 4
+  %cmp = icmp ugt i32* %p, %q
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}
+
+; CHECK-LABEL: @test1b
+; CHECK: [[CMP:%.*]] = icmp ugt i32* %p, %q
+; CHECK: [[SEL:%.*]] = select i1 [[CMP]], i32* %p, i32* %q
+; CHECK: [[GEP:%.*]] = getelementptr i32, i32* [[SEL]], i64 4
+; CHECK: ret i32* [[GEP]]
+define i32* @test1b(i32* %p, i32* %q) {
+  %gep1 = getelementptr inbounds i32, i32* %p, i64 4
+  %gep2 = getelementptr i32, i32* %q, i64 4
+  %cmp = icmp ugt i32* %p, %q
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}
+
+; CHECK-LABEL: @test1c
+; CHECK: [[CMP:%.*]] = icmp ugt i32* %p, %q
+; CHECK: [[SEL:%.*]] = select i1 [[CMP]], i32* %p, i32* %q
+; CHECK: [[GEP:%.*]] = getelementptr i32, i32* [[SEL]], i64 4
+; CHECK: ret i32* [[GEP]]
+define i32* @test1c(i32* %p, i32* %q) {
+  %gep1 = getelementptr i32, i32* %p, i64 4
+  %gep2 = getelementptr inbounds i32, i32* %q, i64 4
+  %cmp = icmp ugt i32* %p, %q
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}
+
+; CHECK-LABEL: @test1d
+; CHECK: [[CMP:%.*]] = icmp ugt i32* %p, %q
+; CHECK: [[SEL:%.*]] = select i1 [[CMP]], i32* %p, i32* %q
+; CHECK: [[GEP:%.*]] = getelementptr inbounds i32, i32* [[SEL]], i64 4
+; CHECK: ret i32* [[GEP]]
+define i32* @test1d(i32* %p, i32* %q) {
+  %gep1 = getelementptr inbounds i32, i32* %p, i64 4
+  %gep2 = getelementptr inbounds i32, i32* %q, i64 4
+  %cmp = icmp ugt i32* %p, %q
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}
+
+; CHECK-LABEL: @test2
+; CHECK: [[CMP:%.*]] = icmp ugt i64 %x, %y
+; CHECK: [[SEL:%.*]] = select i1 [[CMP]], i64 %x, i64 %y
+; CHECK: [[GEP:%.*]] = getelementptr inbounds i32, i32* %p, i64 [[SEL]]
+; CHECK: ret i32* [[GEP]]
+define i32* @test2(i32* %p, i64 %x, i64 %y) {
+  %gep1 = getelementptr inbounds i32, i32* %p, i64 %x
+  %gep2 = getelementptr inbounds i32, i32* %p, i64 %y
+  %cmp = icmp ugt i64 %x, %y
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}
+
+; Three (or more) operand GEPs are currently expected to not be optimised,
+; though they could be in principle.
+; CHECK-LABEL: @test3a
+; CHECK: %gep1 = getelementptr inbounds [4 x i32], [4 x i32]* %p, i64 2, i64 %x
+; CHECK: %gep2 = getelementptr inbounds [4 x i32], [4 x i32]* %p, i64 2, i64 %y
+; CHECK: %cmp = icmp ugt i64 %x, %y
+; CHECK: %select = select i1 %cmp, i32* %gep1, i32* %gep2
+; CHECK: ret i32* %select
+define i32* @test3a([4 x i32]* %p, i64 %x, i64 %y) {
+  %gep1 = getelementptr inbounds [4 x i32], [4 x i32]* %p, i64 2, i64 %x
+  %gep2 = getelementptr inbounds [4 x i32], [4 x i32]* %p, i64 2, i64 %y
+  %cmp = icmp ugt i64 %x, %y
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}
+
+; CHECK-LABEL: @test3b
+; CHECK: %gep1 = getelementptr inbounds [4 x i32], [4 x i32]* %p, i64 %x, i64 2
+; CHECK: %gep2 = getelementptr inbounds i32, i32* %q, i64 %x
+; CHECK: %cmp = icmp ugt i64 %x, %y
+; CHECK: %select = select i1 %cmp, i32* %gep1, i32* %gep2
+; CHECK: ret i32* %select
+define i32* @test3b([4 x i32]* %p, i32* %q, i64 %x, i64 %y) {
+  %gep1 = getelementptr inbounds [4 x i32], [4 x i32]* %p, i64 %x, i64 2
+  %gep2 = getelementptr inbounds i32, i32* %q, i64 %x
+  %cmp = icmp ugt i64 %x, %y
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}
+
+; CHECK-LABEL: @test3c
+; CHECK: %gep1 = getelementptr inbounds i32, i32* %p, i64 %x
+; CHECK: %gep2 = getelementptr inbounds [4 x i32], [4 x i32]* %q, i64 %x, i64 2
+; CHECK: %cmp = icmp ugt i64 %x, %y
+; CHECK: %select = select i1 %cmp, i32* %gep1, i32* %gep2
+; CHECK: ret i32* %select
+define i32* @test3c(i32* %p, [4 x i32]* %q, i64 %x, i64 %y) {
+  %gep1 = getelementptr inbounds i32, i32* %p, i64 %x
+  %gep2 = getelementptr inbounds [4 x i32], [4 x i32]* %q, i64 %x, i64 2
+  %cmp = icmp ugt i64 %x, %y
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}
+
+; Shouldn't be optimised as it would mean introducing an extra select
+; CHECK-LABEL: @test4
+; CHECK: %gep1 = getelementptr inbounds i32, i32* %p, i64 %x
+; CHECK: %gep2 = getelementptr inbounds i32, i32* %q, i64 %y
+; CHECK: %cmp = icmp ugt i64 %x, %y
+; CHECK: %select = select i1 %cmp, i32* %gep1, i32* %gep2
+; CHECK: ret i32* %select
+define i32* @test4(i32* %p, i32* %q, i64 %x, i64 %y) {
+  %gep1 = getelementptr inbounds i32, i32* %p, i64 %x
+  %gep2 = getelementptr inbounds i32, i32* %q, i64 %y
+  %cmp = icmp ugt i64 %x, %y
+  %select = select i1 %cmp, i32* %gep1, i32* %gep2
+  ret i32* %select
+}