[TargetLowering] Fix propagation of undefs in zero extension ops (PR40091)

[RKSimon]

Summary

As described on PR40091, we have several places where zext (and zext_vector_inreg) fold an undef input into an undef output. For zero extensions this incorrect as the output should guarantee to least have the new upper bits set to zero.

SimplifyDemandedVectorElts is the worst offender (and its the most likely to cause new undefs to appear) but DAGCombiner's tryToFoldExtendOfConstant has a similar issue.

Thanks to @dmgreen for catching this.

[andreadb]

LGTM

lib/CodeGen/SelectionDAG/TargetLowering.cpp

1852–1866 ↗

(On Diff #178860)

You can probably try to merge this case with the previous one, since they have a lot of code in common.
That being said, I leave this decision to you. If you think it would make it less readable then it is fine to keep it like this.

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

static SDValue tryToFoldExtendOfConstant(SDNode *N, const TargetLowering &TLI,
 
   // We can fold this node into a build_vector.
   unsigned VTBits = SVT.getSizeInBits();
   unsigned EVTBits = N0->getValueType(0).getScalarSizeInBits();
   SmallVector<SDValue, 8> Elts;
   unsigned NumElts = VT.getVectorNumElements();
   SDLoc DL(N);
 
-  for (unsigned i=0; i != NumElts; ++i) {
-    SDValue Op = N0->getOperand(i);
-    if (Op->isUndef()) {
-      Elts.push_back(DAG.getUNDEF(SVT));
+  // For zero-extensions, UNDEF elements still guarantee to have the upper
+  // bits set to zero.
+  bool IsZext =
+      Opcode == ISD::ZERO_EXTEND || Opcode == ISD::ZERO_EXTEND_VECTOR_INREG;
+
+  for (unsigned i = 0; i != NumElts; ++i) {
+    SDValue Op = N0.getOperand(i);
+    if (Op.isUndef()) {
+      Elts.push_back(IsZext ? DAG.getConstant(0, DL, SVT) : DAG.getUNDEF(SVT));
       continue;
     }
 
     SDLoc DL(Op);
     // Get the constant value and if needed trunc it to the size of the type.
     // Nodes like build_vector might have constants wider than the scalar type.
     APInt C = cast<ConstantSDNode>(Op)->getAPIntValue().zextOrTrunc(EVTBits);
     if (Opcode == ISD::SIGN_EXTEND || Opcode == ISD::SIGN_EXTEND_VECTOR_INREG)

llvm/trunk/lib/CodeGen/SelectionDAG/TargetLowering.cpp

case ISD::ZERO_EXTEND_VECTOR_INREG: {
     SDValue Src = Op.getOperand(0);
     unsigned NumSrcElts = Src.getValueType().getVectorNumElements();
     APInt DemandedSrcElts = DemandedElts.zextOrSelf(NumSrcElts);
     if (SimplifyDemandedVectorElts(Src, DemandedSrcElts, SrcUndef,
                                    SrcZero, TLO, Depth + 1))
       return true;
     KnownZero = SrcZero.zextOrTrunc(NumElts);
     KnownUndef = SrcUndef.zextOrTrunc(NumElts);
+
+    if (Op.getOpcode() == ISD::ZERO_EXTEND_VECTOR_INREG) {
+      // zext(undef) upper bits are guaranteed to be zero.
+      if (DemandedElts.isSubsetOf(KnownUndef))
+        return TLO.CombineTo(Op, TLO.DAG.getConstant(0, SDLoc(Op), VT));
+      KnownUndef.clearAllBits();
+    }
     break;
   }
   case ISD::OR:
   case ISD::XOR:
   case ISD::ADD:
   case ISD::SUB:
   case ISD::FADD:
   case ISD::FSUB:
case ISD::AND: {
     break;
   }
   case ISD::TRUNCATE:
   case ISD::SIGN_EXTEND:
   case ISD::ZERO_EXTEND:
     if (SimplifyDemandedVectorElts(Op.getOperand(0), DemandedElts, KnownUndef,
                                    KnownZero, TLO, Depth + 1))
       return true;
+
+    if (Op.getOpcode() == ISD::ZERO_EXTEND) {
+      // zext(undef) upper bits are guaranteed to be zero.
+      if (DemandedElts.isSubsetOf(KnownUndef))
+        return TLO.CombineTo(Op, TLO.DAG.getConstant(0, SDLoc(Op), VT));
+      KnownUndef.clearAllBits();
+    }
     break;
   default: {
     if (Op.getOpcode() >= ISD::BUILTIN_OP_END) {
       if (SimplifyDemandedVectorEltsForTargetNode(Op, DemandedElts, KnownUndef,
                                                   KnownZero, TLO, Depth))
         return true;
     } else {
       KnownBits Known;

llvm/trunk/test/CodeGen/AArch64/pr40091.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=aarch64-arm-none-eabi | FileCheck %s
 
 define i64 @test(i64 %aa) {
 ; CHECK-LABEL: test:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    movi v0.2d, #0xffffffffffffffff
+; CHECK-NEXT:    movi v0.8b, #137
 ; CHECK-NEXT:    fmov x0, d0
 ; CHECK-NEXT:    ret
 entry:
   %a = bitcast i64 %aa to  <1 x i64>
   %k = icmp sgt <1 x i64> %a, zeroinitializer
   %l = zext <1 x i1> %k to <1 x i64>
   %o = and <1 x i64> %l, %a
   %p = xor <1 x i64> %l, <i64 -1>
   %q = and <1 x i64> %p, <i64 81985529216486895>
   %r = or <1 x i64> %q, %o
   %s = bitcast <1 x i64> %r to <8 x i8>
   %t = shufflevector <8 x i8> %s, <8 x i8> %s, <8 x i32> <i32 3, i32 3, i32 3, i32 3, i32 3, i32 3, i32 3, i32 3>
   %u = bitcast <8 x i8> %t to i64
   ret i64 %u
 }

llvm/trunk/test/CodeGen/X86/fold-vector-sext-zext.ll

; X64-NEXT:    retq
   %4 = insertelement <4 x i8> %3, i8 -3, i32 3
   %5 = zext <4 x i8> %4 to <4 x i64>
   ret <4 x i64> %5
 }
 
 define <4 x i16> @test_zext_4i8_4i16_undef() {
 ; X32-LABEL: test_zext_4i8_4i16_undef:
 ; X32:       # %bb.0:
-; X32-NEXT:    vmovaps {{.*#+}} xmm0 = <u,255,u,253>
+; X32-NEXT:    vmovaps {{.*#+}} xmm0 = [0,255,0,253]
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: test_zext_4i8_4i16_undef:
 ; X64:       # %bb.0:
-; X64-NEXT:    vmovaps {{.*#+}} xmm0 = <u,255,u,253>
+; X64-NEXT:    vmovaps {{.*#+}} xmm0 = [0,255,0,253]
 ; X64-NEXT:    retq
   %1 = insertelement <4 x i8> undef, i8 undef, i32 0
   %2 = insertelement <4 x i8> %1, i8 -1, i32 1
   %3 = insertelement <4 x i8> %2, i8 undef, i32 2
   %4 = insertelement <4 x i8> %3, i8 -3, i32 3
   %5 = zext <4 x i8> %4 to <4 x i16>
   ret <4 x i16> %5
 }
 
 define <4 x i32> @test_zext_4i8_4i32_undef() {
 ; X32-LABEL: test_zext_4i8_4i32_undef:
 ; X32:       # %bb.0:
-; X32-NEXT:    vmovaps {{.*#+}} xmm0 = <0,u,2,u>
+; X32-NEXT:    vmovaps {{.*#+}} xmm0 = [0,0,2,0]
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: test_zext_4i8_4i32_undef:
 ; X64:       # %bb.0:
-; X64-NEXT:    vmovaps {{.*#+}} xmm0 = <0,u,2,u>
+; X64-NEXT:    vmovaps {{.*#+}} xmm0 = [0,0,2,0]
 ; X64-NEXT:    retq
   %1 = insertelement <4 x i8> undef, i8 0, i32 0
   %2 = insertelement <4 x i8> %1, i8 undef, i32 1
   %3 = insertelement <4 x i8> %2, i8 2, i32 2
   %4 = insertelement <4 x i8> %3, i8 undef, i32 3
   %5 = zext <4 x i8> %4 to <4 x i32>
   ret <4 x i32> %5
 }
 
 define <4 x i64> @test_zext_4i8_4i64_undef() {
 ; X32-LABEL: test_zext_4i8_4i64_undef:
 ; X32:       # %bb.0:
-; X32-NEXT:    vmovaps {{.*#+}} ymm0 = <u,u,255,0,2,0,u,u>
+; X32-NEXT:    vmovaps {{.*#+}} ymm0 = [0,0,255,0,2,0,0,0]
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: test_zext_4i8_4i64_undef:
 ; X64:       # %bb.0:
-; X64-NEXT:    vmovaps {{.*#+}} ymm0 = <u,255,2,u>
+; X64-NEXT:    vmovaps {{.*#+}} ymm0 = [0,255,2,0]
 ; X64-NEXT:    retq
   %1 = insertelement <4 x i8> undef, i8 undef, i32 0
   %2 = insertelement <4 x i8> %1, i8 -1, i32 1
   %3 = insertelement <4 x i8> %2, i8 2, i32 2
   %4 = insertelement <4 x i8> %3, i8 undef, i32 3
   %5 = zext <4 x i8> %4 to <4 x i64>
   ret <4 x i64> %5
 }
; X64-NEXT:    retq
   %8 = insertelement <8 x i8> %7, i8 -7, i32 7
   %9 = zext <8 x i8> %8 to <8 x i32>
   ret <8 x i32> %9
 }
 
 define <8 x i16> @test_zext_8i8_8i16_undef() {
 ; X32-LABEL: test_zext_8i8_8i16_undef:
 ; X32:       # %bb.0:
-; X32-NEXT:    vmovaps {{.*#+}} xmm0 = <u,255,u,253,u,251,u,249>
+; X32-NEXT:    vmovaps {{.*#+}} xmm0 = [0,255,0,253,0,251,0,249]
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: test_zext_8i8_8i16_undef:
 ; X64:       # %bb.0:
-; X64-NEXT:    vmovaps {{.*#+}} xmm0 = <u,255,u,253,u,251,u,249>
+; X64-NEXT:    vmovaps {{.*#+}} xmm0 = [0,255,0,253,0,251,0,249]
 ; X64-NEXT:    retq
   %1 = insertelement <8 x i8> undef, i8 undef, i32 0
   %2 = insertelement <8 x i8> %1, i8 -1, i32 1
   %3 = insertelement <8 x i8> %2, i8 undef, i32 2
   %4 = insertelement <8 x i8> %3, i8 -3, i32 3
   %5 = insertelement <8 x i8> %4, i8 undef, i32 4
   %6 = insertelement <8 x i8> %5, i8 -5, i32 5
   %7 = insertelement <8 x i8> %6, i8 undef, i32 6
   %8 = insertelement <8 x i8> %7, i8 -7, i32 7
   %9 = zext <8 x i8> %8 to <8 x i16>
   ret <8 x i16> %9
 }
 
 define <8 x i32> @test_zext_8i8_8i32_undef() {
 ; X32-LABEL: test_zext_8i8_8i32_undef:
 ; X32:       # %bb.0:
-; X32-NEXT:    vmovaps {{.*#+}} ymm0 = <0,u,2,253,4,u,6,u>
+; X32-NEXT:    vmovaps {{.*#+}} ymm0 = [0,0,2,253,4,0,6,0]
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: test_zext_8i8_8i32_undef:
 ; X64:       # %bb.0:
-; X64-NEXT:    vmovaps {{.*#+}} ymm0 = <0,u,2,253,4,u,6,u>
+; X64-NEXT:    vmovaps {{.*#+}} ymm0 = [0,0,2,253,4,0,6,0]
 ; X64-NEXT:    retq
   %1 = insertelement <8 x i8> undef, i8 0, i32 0
   %2 = insertelement <8 x i8> %1, i8 undef, i32 1
   %3 = insertelement <8 x i8> %2, i8 2, i32 2
   %4 = insertelement <8 x i8> %3, i8 -3, i32 3
   %5 = insertelement <8 x i8> %4, i8 4, i32 4
   %6 = insertelement <8 x i8> %5, i8 undef, i32 5
   %7 = insertelement <8 x i8> %6, i8 6, i32 6
   %8 = insertelement <8 x i8> %7, i8 undef, i32 7
   %9 = zext <8 x i8> %8 to <8 x i32>
   ret <8 x i32> %9
 }