diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp index 42d1d9a437bb2..70b4552190a4e 100644 --- a/llvm/lib/Transforms/Scalar/SROA.cpp +++ b/llvm/lib/Transforms/Scalar/SROA.cpp @@ -4070,18 +4070,27 @@ class AggLoadStoreRewriter : public InstVisitor { // => select cond, gep(ptr1, idx), gep(ptr2, idx) // and gep ptr, (select cond, idx1, idx2) // => select cond, gep(ptr, idx1), gep(ptr, idx2) + // We also allow for i1 zext indices, which are equivalent to selects. bool unfoldGEPSelect(GetElementPtrInst &GEPI) { // Check whether the GEP has exactly one select operand and all indices // will become constant after the transform. - SelectInst *Sel = dyn_cast(GEPI.getPointerOperand()); + Instruction *Sel = dyn_cast(GEPI.getPointerOperand()); for (Value *Op : GEPI.indices()) { if (auto *SI = dyn_cast(Op)) { if (Sel) return false; Sel = SI; - if (!isa(Sel->getTrueValue()) || - !isa(Sel->getFalseValue())) + if (!isa(SI->getTrueValue()) || + !isa(SI->getFalseValue())) + return false; + continue; + } + if (auto *ZI = dyn_cast(Op)) { + if (Sel) + return false; + Sel = ZI; + if (!ZI->getSrcTy()->isIntegerTy(1)) return false; continue; } @@ -4107,8 +4116,16 @@ class AggLoadStoreRewriter : public InstVisitor { return NewOps; }; - Value *True = Sel->getTrueValue(); - Value *False = Sel->getFalseValue(); + Value *Cond, *True, *False; + if (auto *SI = dyn_cast(Sel)) { + Cond = SI->getCondition(); + True = SI->getTrueValue(); + False = SI->getFalseValue(); + } else { + Cond = Sel->getOperand(0); + True = ConstantInt::get(Sel->getType(), 1); + False = ConstantInt::get(Sel->getType(), 0); + } SmallVector TrueOps = GetNewOps(True); SmallVector FalseOps = GetNewOps(False); @@ -4123,8 +4140,8 @@ class AggLoadStoreRewriter : public InstVisitor { IRB.CreateGEP(Ty, FalseOps[0], ArrayRef(FalseOps).drop_front(), False->getName() + ".sroa.gep", NW); - Value *NSel = IRB.CreateSelect(Sel->getCondition(), NTrue, NFalse, - Sel->getName() + ".sroa.sel"); + Value *NSel = + IRB.CreateSelect(Cond, NTrue, NFalse, Sel->getName() + ".sroa.sel"); Visited.erase(&GEPI); GEPI.replaceAllUsesWith(NSel); GEPI.eraseFromParent(); diff --git a/llvm/test/Transforms/SROA/select-gep.ll b/llvm/test/Transforms/SROA/select-gep.ll index 1342a2ca4ea2b..b48b0f77aa991 100644 --- a/llvm/test/Transforms/SROA/select-gep.ll +++ b/llvm/test/Transforms/SROA/select-gep.ll @@ -201,6 +201,47 @@ define i32 @test_select_idx_mem2reg(i1 %c) { ret i32 %res } +; Test gep with a select-like zext index unfolding on an alloca that is +; splittable and promotable. +define i64 @test_select_like_zext_idx_mem2reg(i1 %c) { +; CHECK-LABEL: @test_select_like_zext_idx_mem2reg( +; CHECK-NEXT: [[IDX:%.*]] = zext i1 [[C:%.*]] to i64 +; CHECK-NEXT: [[RES:%.*]] = select i1 [[C]], i64 2, i64 1 +; CHECK-NEXT: ret i64 [[RES]] +; + %alloca = alloca [2 x i64], align 8 + store i64 1, ptr %alloca + %gep1 = getelementptr inbounds i64, ptr %alloca, i64 1 + store i64 2, ptr %gep1 + %idx = zext i1 %c to i64 + %gep2 = getelementptr inbounds i64, ptr %alloca, i64 %idx + %res = load i64, ptr %gep2 + ret i64 %res +} + +; Test gep with a zext index that is not equivalent to a select. No unfolding +; or promotion should take place. +define i64 @test_zext_unlike_select_idx_mem2reg(i8 %c) { +; CHECK-LABEL: @test_zext_unlike_select_idx_mem2reg( +; CHECK-NEXT: [[ALLOCA:%.*]] = alloca [2 x i64], align 8 +; CHECK-NEXT: store i64 1, ptr [[ALLOCA]], align 4 +; CHECK-NEXT: [[GEP1:%.*]] = getelementptr inbounds i64, ptr [[ALLOCA]], i64 1 +; CHECK-NEXT: store i64 2, ptr [[GEP1]], align 4 +; CHECK-NEXT: [[IDX:%.*]] = zext i8 [[C:%.*]] to i64 +; CHECK-NEXT: [[GEP2:%.*]] = getelementptr inbounds i64, ptr [[ALLOCA]], i64 [[IDX]] +; CHECK-NEXT: [[RES:%.*]] = load i64, ptr [[GEP2]], align 4 +; CHECK-NEXT: ret i64 [[RES]] +; + %alloca = alloca [2 x i64], align 8 + store i64 1, ptr %alloca + %gep1 = getelementptr inbounds i64, ptr %alloca, i64 1 + store i64 2, ptr %gep1 + %idx = zext i8 %c to i64 + %gep2 = getelementptr inbounds i64, ptr %alloca, i64 %idx + %res = load i64, ptr %gep2 + ret i64 %res +} + ; Test gep of index select unfolding on an alloca that escaped, and as such ; is not splittable or promotable. ; FIXME: Ideally, no transform would take place in this case.