Process gep (select ptr1, ptr2) in SROA

[arsenm]

llvm/lib/Transforms/Scalar/SROA.cpp

3471

Needs a newline after?

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

15

Store something for non-undef test output?

17

This increases the instruction count if the load won't be eliminated. Can you add a test with a volatile load or something else to prevent it?

[rampitec]

Addressed review comments.

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

17

It may increase instruction count, but if it will remain after SROA there is a pattern in InstCombine to revert it back. That is actually the reason why D79145 was over complicated and I have switched to processing it right in SROA.

[rampitec]

PSDB passed.

[rampitec]

Added one more test for select gep (select gep, gep), gep (select gep, gep).

[rampitec]

Ping.

[arsenm]

llvm/lib/Transforms/Scalar/SROA.cpp

3468

Early exit if not select and reduce indentation

[rampitec]

llvm/lib/Transforms/Scalar/SROA.cpp

3468

I will have to return this if anyway after D79218. There is one if for Select and one for PHI.

[rampitec]

llvm/lib/Transforms/Scalar/SROA.cpp

3468

I will extract it into a separate function, here and in another patch.

[rampitec]

Extracted folding into a separate function.

[arsenm]

llvm/lib/Transforms/Scalar/SROA.cpp

3481

The inbounds vs. noninbounds paths aren't stressed in the test

[rampitec]

Added a non-inbound GEP test.

[arsenm]

llvm/lib/Transforms/Scalar/SROA.cpp

3477

Probably should check there are only constant indexes?

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

7

Should have a test with non-constant GEP indices. This probably doesn't help if they aren't constants

[rampitec]

Added check for constant indexes and test.

[rampitec]

Ping! @chandlerc do you mind to take a look? Thanks in advance!

[arsenm]

lgtm

llvm/lib/Transforms/Scalar/SROA.cpp

private:
   }
 
   bool visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
     enqueueUsers(ASC);
     return false;
   }
 
   bool visitGetElementPtrInst(GetElementPtrInst &GEPI) {
+    // Fold gep (select cond, ptr1, ptr2) => select cond, gep(ptr1), gep(ptr2)
+    if (SelectInst *Sel = dyn_cast<SelectInst>(GEPI.getPointerOperand())) {
+      LLVM_DEBUG(dbgs() << "  Rewriting gep(select) -> select(gep):"
+                        << "\n    original: " << *Sel
+                        << "\n              " << GEPI);
+      IRBuilderTy Builder(&GEPI);
+      SmallVector<Value *, 4> Index(GEPI.idx_begin(), GEPI.idx_end());
+      bool IsInBounds = GEPI.isInBounds();
+
+      Value *True = Sel->getTrueValue();
+      Value *NTrue =
+          IsInBounds
+              ? Builder.CreateInBoundsGEP(True, Index,
+                                          True->getName() + ".sroa.gep")
+              : Builder.CreateGEP(True, Index, True->getName() + ".sroa.gep");
+
+      Value *False = Sel->getFalseValue();
+
+      Value *NFalse =
+          IsInBounds
+              ? Builder.CreateInBoundsGEP(False, Index,
+                                          False->getName() + ".sroa.gep")
+              : Builder.CreateGEP(False, Index, False->getName() + ".sroa.gep");
+
+      Value *NSel = Builder.CreateSelect(Sel->getCondition(), NTrue, NFalse,
+                                         Sel->getName() + ".sroa.sel");
+      GEPI.replaceAllUsesWith(NSel);
+      GEPI.eraseFromParent();
+
+      LLVM_DEBUG(dbgs() << "\n          to: " << *NTrue
+                        << "\n              " << *NFalse
+                        << "\n              " << *NSel << '\n');
+
+      if (isa<Instruction>(NTrue))
+        visit(cast<Instruction>(NTrue));
+      if (isa<Instruction>(NFalse))
+        visit(cast<Instruction>(NFalse));
+      return true;
+    }
+
     enqueueUsers(GEPI);
     return false;
   }
 
   bool visitPHINode(PHINode &PN) {
     enqueueUsers(PN);
     return false;
   }

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

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -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_SROA_SPECULATED:%.*]] = select i1 [[COND:%.*]], i32 undef, i32 undef
+; CHECK-NEXT:    ret i32 [[LOAD_SROA_SPECULATED]]
+;
+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
+}
+
+define i32 @test_sroa_select_gep_undef(i1 %cond) {
+; CHECK-LABEL: @test_sroa_select_gep_undef(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[A_SROA_0:%.*]] = alloca i32
+; CHECK-NEXT:    [[SELECT_SROA_SEL:%.*]] = select i1 [[COND:%.*]], i32* [[A_SROA_0]], i32* undef
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, i32* [[SELECT_SROA_SEL]], align 4
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+bb:
+  %a = alloca %pair, align 4
+  %select = select i1 %cond, %pair* %a, %pair* undef
+  %gep = getelementptr inbounds %pair, %pair* %select, i32 0, i32 1
+  %load = load i32, i32* %gep, align 4
+  ret i32 %load
+}