[DAG] FoldConstantArithmetic - add initial support for undef elements in bitcasted binop constant folding

[RKSimon]

Summary

FoldConstantArithmetic can fold constant vectors hidden behind bitcasts (e.g. vXi64 -> v2Xi32 on 32-bit platforms), but currently bails if either vector contains undef elements. These undefs can often occur due to SimplifyDemandedBits/VectorElts calls recognising that the upper bits are often unnecessary (e.g. funnel-shift/rotate implicit-modulo and AND masks).

This patch adds a basic 'FoldValueWithUndef' handler that will attempt to constant fold if one or both of the ops are undef - so far this just handles the AND and MUL cases where we always fold to zero.

The RISCV codegen increase is interesting - it looks like the BUILD_VECTOR lowering was loading a constant pool entry but now (with all elements defined constant) it materializes it instead?

[tschuett]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5602

Do you need to hide this behind a lambda or could you pull it into a function because it is getting more powerful over time?

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5602

No, it could be exposed like FoldValue() is - that is currently a static helper above FoldConstantArithmetic. We could move it into TLI later if it'd help target nodes.

[xiangzhangllvm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5601

Maybe better to add a small case here, some like
"old nodes" --> "new nodes"

5607

Do we need to make sure the C1 C2 has no other uses ?

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5607

I don't think we have any one use limits on constant folding?

[RKSimon]

Pull out FoldValueWithUndef lambda to match FoldValue helper

[RKSimon]

llvm/test/CodeGen/RISCV/srem-seteq-illegal-types.ll

667

@craig.topper I'm assuming this amount of code isn't an improvement just because we avoided a load?

[xiangzhangllvm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5607

Yes, sorry my mistake at first.
take following case for example:

A = and(x, undef)
C = A op B

If we directly let "A = 0" (I think A=undef before), seems it may break the undef propagation in other use. for example C when "op" can do undef propagation (but not constant zero propagation).

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5607

A = and(x, undef)
C = A op B

A can't (or shouldn't) propagate undef: https://alive2.llvm.org/ce/z/0EB60y - we almost certainly have cases of that in code, which need addressing

[xiangzhangllvm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5607

Yes, thanks for your cases! it really not propagate.

But sorry for I am still some puzzle here:
Do you know why we can't (or shouldn't) propagate undef ?

To me, I can't see the real difference between

A = undef

and

A = and(x, undef)

[efriedma]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5607

and(0, undef) is not undef; it's zero. By contrast, and(0, poison) is poison. This is the fundamental difference between undef and poison. See https://llvm.org/docs/LangRef.html#undefined-values .

[xiangzhangllvm]

LGTM

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5607

Very clear now, thanks very much for your explain!

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::AVGCEILU: {
     APInt C1Ext = C1.zext(FullWidth);
     APInt C2Ext = C2.zext(FullWidth);
     return (C1Ext + C2Ext + 1).extractBits(C1.getBitWidth(), 1);
   }
   }
   return llvm::None;
 }
 
+// Handle constant folding with UNDEF.
+// TODO: Handle more cases.
+static llvm::Optional<APInt> FoldValueWithUndef(unsigned Opcode,
+                                                const APInt &C1, bool IsUndef1,
+                                                const APInt &C2,
+                                                bool IsUndef2) {
+  if (!(IsUndef1 || IsUndef2))
+    return FoldValue(Opcode, C1, C2);
+
+  // Fold and(x, undef) -> 0
+  // Fold mul(x, undef) -> 0
+  if (Opcode == ISD::AND || Opcode == ISD::MUL)
+    return APInt::getZero(C1.getBitWidth());
+
+  return llvm::None;
+};
+
 SDValue SelectionDAG::FoldSymbolOffset(unsigned Opcode, EVT VT,
                                        const GlobalAddressSDNode *GA,
                                        const SDNode *N2) {
   if (GA->getOpcode() != ISD::GlobalAddress)
     return SDValue();
   if (!TLI->isOffsetFoldingLegal(GA))
     return SDValue();
   auto *C2 = dyn_cast<ConstantSDNode>(N2);
SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
 
   // This is for vector folding only from here on.
   if (!VT.isVector())
     return SDValue();
 
   ElementCount NumElts = VT.getVectorElementCount();
 
   // See if we can fold through bitcasted integer ops.
-  // TODO: Can we handle undef elements?
   if (NumOps == 2 && VT.isFixedLengthVector() && VT.isInteger() &&
       Ops[0].getValueType() == VT && Ops[1].getValueType() == VT &&
       Ops[0].getOpcode() == ISD::BITCAST &&
       Ops[1].getOpcode() == ISD::BITCAST) {
     SDValue N1 = peekThroughBitcasts(Ops[0]);
     SDValue N2 = peekThroughBitcasts(Ops[1]);
     auto *BV1 = dyn_cast<BuildVectorSDNode>(N1);
     auto *BV2 = dyn_cast<BuildVectorSDNode>(N2);
     EVT BVVT = N1.getValueType();
     if (BV1 && BV2 && BVVT.isInteger() && BVVT == N2.getValueType()) {
       bool IsLE = getDataLayout().isLittleEndian();
       unsigned EltBits = VT.getScalarSizeInBits();
       SmallVector<APInt> RawBits1, RawBits2;
       BitVector UndefElts1, UndefElts2;
       if (BV1->getConstantRawBits(IsLE, EltBits, RawBits1, UndefElts1) &&
-          BV2->getConstantRawBits(IsLE, EltBits, RawBits2, UndefElts2) &&
-          UndefElts1.none() && UndefElts2.none()) {
+          BV2->getConstantRawBits(IsLE, EltBits, RawBits2, UndefElts2)) {
         SmallVector<APInt> RawBits;
         for (unsigned I = 0, E = NumElts.getFixedValue(); I != E; ++I) {
-          Optional<APInt> Fold = FoldValue(Opcode, RawBits1[I], RawBits2[I]);
+          Optional<APInt> Fold = FoldValueWithUndef(
+              Opcode, RawBits1[I], UndefElts1[I], RawBits2[I], UndefElts2[I]);
           if (!Fold)
             break;
           RawBits.push_back(*Fold);
         }
         if (RawBits.size() == NumElts.getFixedValue()) {
           // We have constant folded, but we need to cast this again back to
           // the original (possibly legalized) type.
           SmallVector<APInt> DstBits;
           BitVector DstUndefs;
           BuildVectorSDNode::recastRawBits(IsLE, BVVT.getScalarSizeInBits(),
                                            DstBits, RawBits, DstUndefs,

llvm/test/CodeGen/ARM/srem-seteq-illegal-types.ll

 ; ARM7-NEXT:  .LCPI3_1:
 ; ARM7-NEXT:    .long 3 @ 0x3
 ; ARM7-NEXT:    .long 0 @ 0x0
 ; ARM7-NEXT:    .long 4294967293 @ 0xfffffffd
 ; ARM7-NEXT:    .long 1 @ 0x1
 ; ARM7-NEXT:  .LCPI3_2:
 ; ARM7-NEXT:    .long 3 @ 0x3
 ; ARM7-NEXT:    .long 0 @ 0x0
-; ARM7-NEXT:    .zero 4
+; ARM7-NEXT:    .long 0 @ 0x0
 ; ARM7-NEXT:    .long 0 @ 0x0
 ;
 ; ARM8-LABEL: test_srem_vec:
 ; ARM8:       @ %bb.0:
 ; ARM8-NEXT:    push {r4, r5, r6, r7, r11, lr}
 ; ARM8-NEXT:    vpush {d8, d9}
 ; ARM8-NEXT:    mov r6, r0
 ; ARM8-NEXT:    and r0, r3, #1

 ; ARM8-NEXT:  .LCPI3_1:
 ; ARM8-NEXT:    .long 3 @ 0x3
 ; ARM8-NEXT:    .long 0 @ 0x0
 ; ARM8-NEXT:    .long 4294967293 @ 0xfffffffd
 ; ARM8-NEXT:    .long 1 @ 0x1
 ; ARM8-NEXT:  .LCPI3_2:
 ; ARM8-NEXT:    .long 3 @ 0x3
 ; ARM8-NEXT:    .long 0 @ 0x0
-; ARM8-NEXT:    .zero 4
+; ARM8-NEXT:    .long 0 @ 0x0
 ; ARM8-NEXT:    .long 0 @ 0x0
 ;
 ; NEON7-LABEL: test_srem_vec:
 ; NEON7:       @ %bb.0:
 ; NEON7-NEXT:    push {r4, r5, r6, r7, r11, lr}
 ; NEON7-NEXT:    vpush {d8, d9}
 ; NEON7-NEXT:    mov r6, r0
 ; NEON7-NEXT:    and r0, r3, #1

 ; NEON7-NEXT:  .LCPI3_1:
 ; NEON7-NEXT:    .long 3 @ 0x3
 ; NEON7-NEXT:    .long 0 @ 0x0
 ; NEON7-NEXT:    .long 4294967293 @ 0xfffffffd
 ; NEON7-NEXT:    .long 1 @ 0x1
 ; NEON7-NEXT:  .LCPI3_2:
 ; NEON7-NEXT:    .long 3 @ 0x3
 ; NEON7-NEXT:    .long 0 @ 0x0
-; NEON7-NEXT:    .zero 4
+; NEON7-NEXT:    .long 0 @ 0x0
 ; NEON7-NEXT:    .long 0 @ 0x0
 ;
 ; NEON8-LABEL: test_srem_vec:
 ; NEON8:       @ %bb.0:
 ; NEON8-NEXT:    push {r4, r5, r6, r7, r11, lr}
 ; NEON8-NEXT:    vpush {d8, d9}
 ; NEON8-NEXT:    mov r6, r0
 ; NEON8-NEXT:    and r0, r3, #1

 ; NEON8-NEXT:  .LCPI3_1:
 ; NEON8-NEXT:    .long 3 @ 0x3
 ; NEON8-NEXT:    .long 0 @ 0x0
 ; NEON8-NEXT:    .long 4294967293 @ 0xfffffffd
 ; NEON8-NEXT:    .long 1 @ 0x1
 ; NEON8-NEXT:  .LCPI3_2:
 ; NEON8-NEXT:    .long 3 @ 0x3
 ; NEON8-NEXT:    .long 0 @ 0x0
-; NEON8-NEXT:    .zero 4
+; NEON8-NEXT:    .long 0 @ 0x0
 ; NEON8-NEXT:    .long 0 @ 0x0
   %srem = srem <3 x i33> %X, <i33 9, i33 9, i33 -9>
   %cmp = icmp ne <3 x i33> %srem, <i33 3, i33 -3, i33 3>
   ret <3 x i1> %cmp
 }

llvm/test/CodeGen/RISCV/srem-seteq-illegal-types.ll

 ; RV32MV-NEXT:    li a3, 0
 ; RV32MV-NEXT:    call __moddi3@plt
 ; RV32MV-NEXT:    sw a1, 12(sp)
 ; RV32MV-NEXT:    sw a0, 8(sp)
 ; RV32MV-NEXT:    li a0, 85
 ; RV32MV-NEXT:    vsetivli zero, 1, e8, mf8, ta, mu
 ; RV32MV-NEXT:    vmv.s.x v0, a0
 ; RV32MV-NEXT:    vsetivli zero, 8, e32, m2, ta, mu
-; RV32MV-NEXT:    vmv.v.i v8, 1
 ; RV32MV-NEXT:    mv a0, sp
-; RV32MV-NEXT:    vle32.v v10, (a0)
-; RV32MV-NEXT:    lui a0, %hi(.LCPI3_0)
-; RV32MV-NEXT:    addi a0, a0, %lo(.LCPI3_0)
-; RV32MV-NEXT:    vle32.v v12, (a0)
-; RV32MV-NEXT:    vmerge.vim v8, v8, -1, v0
-; RV32MV-NEXT:    vand.vv v8, v10, v8
+; RV32MV-NEXT:    vle32.v v8, (a0)
+; RV32MV-NEXT:    vmv.v.i v10, 1
+; RV32MV-NEXT:    vmerge.vim v10, v10, -1, v0
+; RV32MV-NEXT:    vand.vv v8, v8, v10
+; RV32MV-NEXT:    li a0, 2
+; RV32MV-NEXT:    vmv.s.x v10, a0
+; RV32MV-NEXT:    li a0, 1
+; RV32MV-NEXT:    vmv.s.x v12, a0
+; RV32MV-NEXT:    vmv.v.i v14, 0
+; RV32MV-NEXT:    vsetivli zero, 3, e32, m2, tu, mu
+; RV32MV-NEXT:    vslideup.vi v14, v12, 2
+; RV32MV-NEXT:    vsetivli zero, 5, e32, m2, tu, mu
+; RV32MV-NEXT:    vslideup.vi v14, v10, 4
 ; RV32MV-NEXT:    vsetivli zero, 4, e64, m2, ta, mu
-; RV32MV-NEXT:    vmsne.vv v0, v8, v12
+; RV32MV-NEXT:    vmsne.vv v0, v8, v14
 ; RV32MV-NEXT:    vmv.v.i v8, 0
 ; RV32MV-NEXT:    vmerge.vim v8, v8, -1, v0
 ; RV32MV-NEXT:    vsetivli zero, 1, e32, m2, ta, mu
 ; RV32MV-NEXT:    vse32.v v8, (s2)
 ; RV32MV-NEXT:    vslidedown.vi v10, v8, 1
 ; RV32MV-NEXT:    vmv.x.s a0, v10
 ; RV32MV-NEXT:    vslidedown.vi v10, v8, 2
 ; RV32MV-NEXT:    vmv.x.s a1, v10

llvm/test/CodeGen/Thumb2/srem-seteq-illegal-types.ll

 ; CHECK-NEXT:  .LCPI3_1:
 ; CHECK-NEXT:    .long 3 @ 0x3
 ; CHECK-NEXT:    .long 0 @ 0x0
 ; CHECK-NEXT:    .long 4294967293 @ 0xfffffffd
 ; CHECK-NEXT:    .long 1 @ 0x1
 ; CHECK-NEXT:  .LCPI3_2:
 ; CHECK-NEXT:    .long 3 @ 0x3
 ; CHECK-NEXT:    .long 0 @ 0x0
-; CHECK-NEXT:    .zero 4
+; CHECK-NEXT:    .long 0 @ 0x0
 ; CHECK-NEXT:    .long 0 @ 0x0
   %srem = srem <3 x i33> %X, <i33 9, i33 9, i33 -9>
   %cmp = icmp ne <3 x i33> %srem, <i33 3, i33 -3, i33 3>
   ret <3 x i1> %cmp
 }

llvm/test/CodeGen/X86/fshl-splat-undef.ll

 ; We should make sure not "merging" <i32 12, i32 undef> to <i64 undef>
 ; (We can not convert t2 to {i64 undef, i64 undef, ...})
 ; That is not equal with the origin result)
 ;
 define void @test_fshl(<8 x i64> %lo, <8 x i64> %hi, <8 x i64>* %arr) {
 ; CHECK-LABEL: test_fshl:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; CHECK-NEXT:    vpbroadcastd {{.*#+}} zmm2 = [12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12]
-; CHECK-NEXT:    vpandq {{\.?LCPI[0-9]+_[0-9]+}}, %zmm2, %zmm2
-; CHECK-NEXT:    vpsllvq %zmm2, %zmm1, %zmm1
+; CHECK-NEXT:    vpsllq $12, %zmm1, %zmm1
 ; CHECK-NEXT:    vpsrlq $52, %zmm0, %zmm0
 ; CHECK-NEXT:    vpternlogq $168, {{\.?LCPI[0-9]+_[0-9]+}}, %zmm1, %zmm0
 ; CHECK-NEXT:    vmovdqa64 %zmm0, (%eax)
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retl
 entry:
   %fshl = call <8 x i64> @llvm.fshl.v8i64(<8 x i64> %hi, <8 x i64> %lo, <8 x i64> <i64 12, i64 12, i64 12, i64 12, i64 12, i64 12, i64 12, i64 12>)
   %res = shufflevector <8 x i64> %fshl, <8 x i64> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 10, i32 11, i32 12, i32 13, i32 6, i32 7>
   store <8 x i64> %res, <8 x i64>* %arr, align 64
   ret void
 }
 
 
 ; Function Attrs: nocallback nofree nosync nounwind readnone speculatable willreturn
 declare <8 x i64> @llvm.fshl.v8i64(<8 x i64>, <8 x i64>, <8 x i64>)