SROA: Process bitcast (select ptr1, ptr2)

[cdevadas]

Summary

After D79171, there is a type reversal in the insertelement
operation and its inputs resulting in a new IR pattern that during
SROA prevents speculateSelectInstLoads and PromoteMemToReg
for certain input scenarios. This patch will restore them.

[cdevadas]

The lit test select-bitcast.ll is an extract from a potential application that caused a bad performance due to the missing alloca promotion.
The intervening bitcast instruction in the load-bitcast-select sequence turned out to be a hurdle to perform alloca promotion during SROA.

[lebedev.ri]

Can you explain why AggLoadStoreRewriter currently fails due to that bitcast?
Which check in particular prevents the transformation?

While i agree it's obviously an improvement,
i'm wondering if there is some simpler solution.

[cdevadas]

In D106545#2896286, @lebedev.ri wrote:

Can you explain why AggLoadStoreRewriter currently fails due to that bitcast?
Which check in particular prevents the transformation?

While i agree it's obviously an improvement,
i'm wondering if there is some simpler solution.

I don't think currently there is anything that exists in the AggLoadStoreRewriter for bitcast (select cond V1 V2) pattern.
As I mentioned earlier, D79171 introduced this new pattern. Earlier, it was GEP in place of bitcast and the gep (select) folding already exists in the Rewriter.

[lebedev.ri]

I took a look, and like i suspected, this isn't the fix i'd envision.
I believe isSafeSelectToSpeculate() should, instead of only looking at direct uses of the SI,
add them all into worklist, iterate through it, and if any item is a bitcast, recurse into it,
else if it's a load then do what it currently does, and abort otherwise.

[cdevadas]

In D106545#2897048, @lebedev.ri wrote:

I took a look, and like i suspected, this isn't the fix i'd envision.
I believe isSafeSelectToSpeculate() should, instead of only looking at direct uses of the SI,
add them all into worklist, iterate through it, and if any item is a bitcast, recurse into it,
else if it's a load then do what it currently does, and abort otherwise.

Posted D106667 to handle it during speculateSelectInstLoads. This review can be abandoned later.

[cdevadas]

Handled it with D106667.

llvm/lib/Transforms/Scalar/SROA.cpp

bool visitStoreInst(StoreInst &SI) {
     StoreOpSplitter Splitter(&SI, *U, V->getType(), AATags,
                              getAdjustedAlignment(&SI, 0), DL);
     Splitter.emitSplitOps(V->getType(), V, V->getName() + ".fca");
     Visited.erase(&SI);
     SI.eraseFromParent();
     return true;
   }
 
+  // Fold bitcast (select cond, ptr1, ptr2) => select cond, bitcast(ptr1),
+  // bitcast(ptr2)
+  bool foldBitCastSelect(BitCastInst &BC) {
+    SelectInst *Sel = cast<SelectInst>(BC.getOperand(0));
+
+    LLVM_DEBUG(dbgs() << "  Rewriting bitcast(select) -> select(bitcast):"
+                      << "\n    original: " << *Sel << "\n              "
+                      << BC);
+
+    IRBuilderTy Builder(&BC);
+
+    Type *Ty = BC.getType();
+    Value *True = Sel->getTrueValue();
+    Value *NTrue =
+        Builder.CreateBitCast(True, Ty, True->getName() + ".sroa.cast");
+
+    Value *False = Sel->getFalseValue();
+    Value *NFalse =
+        Builder.CreateBitCast(False, Ty, False->getName() + ".sroa.cast");
+
+    Value *NSel = Builder.CreateSelect(Sel->getCondition(), NTrue, NFalse,
+                                       Sel->getName() + ".sroa.sel");
+    Visited.erase(&BC);
+    BC.replaceAllUsesWith(NSel);
+    BC.eraseFromParent();
+    Instruction *NSelI = cast<Instruction>(NSel);
+    Visited.insert(NSelI);
+    enqueueUsers(*NSelI);
+
+    LLVM_DEBUG(dbgs() << "\n          to: " << *NTrue << "\n              "
+                      << *NFalse << "\n              " << *NSel << '\n');
+
+    return true;
+  }
+
   bool visitBitCastInst(BitCastInst &BC) {
+    if (isa<SelectInst>(BC.getOperand(0)) && foldBitCastSelect(BC))
+      return true;
+
     enqueueUsers(BC);
     return false;
   }
 
   bool visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
     enqueueUsers(ASC);
     return false;
   }

llvm/test/Transforms/SROA/phi-and-select.ll

entry:
 
   %cond = icmp sle i32 %v0, %v1
   %select = select i1 %cond, i32* %a1, i32* %a0
 
   %result = load i32, i32* %select
   ret i32 %result
 }
 
-; If bitcast isn't considered a safe phi/select use, the alloca
-; remains as an array.
-; FIXME: Why isn't this identical to test2?
 define float @test2_bitcast() {
 ; CHECK-LABEL: @test2_bitcast(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[A_SROA_0:%.*]] = alloca i32, align 4
-; CHECK-NEXT:    [[A_SROA_3:%.*]] = alloca i32, align 4
-; CHECK-NEXT:    store i32 0, i32* [[A_SROA_0]], align 4
-; CHECK-NEXT:    store i32 1, i32* [[A_SROA_3]], align 4
-; CHECK-NEXT:    [[A_SROA_0_0_A_SROA_0_0_V0:%.*]] = load i32, i32* [[A_SROA_0]], align 4
-; CHECK-NEXT:    [[A_SROA_3_0_A_SROA_3_4_V1:%.*]] = load i32, i32* [[A_SROA_3]], align 4
-; CHECK-NEXT:    [[COND:%.*]] = icmp sle i32 [[A_SROA_0_0_A_SROA_0_0_V0]], [[A_SROA_3_0_A_SROA_3_4_V1]]
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[COND]], i32* [[A_SROA_3]], i32* [[A_SROA_0]]
-; CHECK-NEXT:    [[SELECT_BC:%.*]] = bitcast i32* [[SELECT]] to float*
-; CHECK-NEXT:    [[RESULT:%.*]] = load float, float* [[SELECT_BC]], align 4
-; CHECK-NEXT:    ret float [[RESULT]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp sle i32 0, 1
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32 1 to float
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i32 0 to float
+; CHECK-NEXT:    [[RESULT_SROA_SPECULATED:%.*]] = select i1 [[COND]], float [[TMP0]], float [[TMP1]]
+; CHECK-NEXT:    ret float [[RESULT_SROA_SPECULATED]]
 ;
 entry:
   %a = alloca [2 x i32]
   %a0 = getelementptr [2 x i32], [2 x i32]* %a, i64 0, i32 0
   %a1 = getelementptr [2 x i32], [2 x i32]* %a, i64 0, i32 1
   store i32 0, i32* %a0
   store i32 1, i32* %a1
   %v0 = load i32, i32* %a0

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

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -sroa < %s | FileCheck %s
+
+%st.half = type { half }
+
+define <2 x i16> @test_bitcast_select_arg_alloca(i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test_bitcast_select_arg_alloca(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast half 0xHFFFF to i16
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast half 0xH0000 to i16
+; CHECK-NEXT:    [[LD1_SROA_SPECULATED:%.*]] = select i1 [[COND1:%.*]], i16 [[TMP1]], i16 [[TMP2]]
+; CHECK-NEXT:    [[V1:%.*]] = insertelement <2 x i16> undef, i16 [[LD1_SROA_SPECULATED]], i32 0
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast half 0xHFFFF to i16
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast half 0xH0000 to i16
+; CHECK-NEXT:    [[LD2_SROA_SPECULATED:%.*]] = select i1 [[COND2:%.*]], i16 [[TMP3]], i16 [[TMP4]]
+; CHECK-NEXT:    [[V2:%.*]] = insertelement <2 x i16> [[V1]], i16 [[LD2_SROA_SPECULATED]], i32 1
+; CHECK-NEXT:    ret <2 x i16> [[V2]]
+;
+  %true = alloca half, align 2
+  %false = alloca half, align 2
+  store half 0xHFFFF, half* %true, align 2
+  store half 0xH0000, half* %false, align 2
+  %false.cast = bitcast half* %false to %st.half*
+  %true.cast = bitcast half* %true to %st.half*
+  %sel1 = select i1 %cond1, %st.half* %true.cast, %st.half* %false.cast
+  %cast1 = bitcast %st.half* %sel1 to i16*
+  %ld1 = load i16, i16* %cast1, align 2
+  %v1 = insertelement <2 x i16> undef, i16 %ld1, i32 0
+  %sel2 = select i1 %cond2, %st.half* %true.cast, %st.half* %false.cast
+  %cast2 = bitcast %st.half* %sel2 to i16*
+  %ld2 = load i16, i16* %cast2, align 2
+  %v2 = insertelement <2 x i16> %v1, i16 %ld2, i32 1
+  ret <2 x i16> %v2
+}