[DAGCombiner] Transform (zext (select c, load1, load2)) -> (select c, zextload1, zextload2)

[Carrot]

Summary

If extload is legal, following transform

(zext (select c, load1, load2)) -> (select c, zextload1, zextload2)

can save one ext instruction.

[xbolva00]

Can you precommit testcase, please?

[Carrot]

Rebase the code.

[RKSimon]

llvm/test/CodeGen/X86/select-ext.ll

73–146

Worth adding sextload test case? What about vectors?

[Carrot]

llvm/test/CodeGen/X86/select-ext.ll

73–146

Never thought about vector cases. Will look into it.

[Carrot]

Add more test cases.

llvm/test/CodeGen/X86/select-ext.ll

73–146

The current patch can also handle vectors. In the new test case, sext_vector is compiled to following without this patch

sext_vector:                            # @sext_vector
        .cfi_startproc
# %bb.0:
        movq    (%rdi), %xmm1                   # xmm1 = mem[0],zero
        movq    8(%rdi), %xmm0                  # xmm0 = mem[0],zero
        testl   %esi, %esi
        jne     .LBB3_2
# %bb.1:
        movdqa  %xmm1, %xmm0
.LBB3_2:
        pmovsxdq        %xmm0, %xmm0
        retq

With this patch I got

sext_vector:                            # @sext_vector
        .cfi_startproc
# %bb.0:
        pmovsxdq        (%rdi), %xmm1
        pmovsxdq        8(%rdi), %xmm0
        testl   %esi, %esi
        jne     .LBB3_2
# %bb.1:
        movdqa  %xmm1, %xmm0
.LBB3_2:
        retq

The sign extension instruction before return is removed.

[RKSimon]

@Carrot I've added the new tests to trunk along with some vselect test cases - please can you rebase?

[Carrot]

Rebase the test case.

[RKSimon]

llvm/test/CodeGen/X86/select-ext.ll

106

Any chance we can support ISD::VSELECT as well?

[Carrot]

Add support for ISD::VSELECT.

[xbolva00]

llvm/test/CodeGen/X86/select-ext.ll

57

Remove todo

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10057

We also need to ensure that the loads aren't already a extload - or are either the same type of extension (or just ISD::EXTLOAD which I guess we can replace). This will need more tests (inc negative tests).

[Carrot]

Check conflicting extload.

[Carrot]

ping

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10097

I'm not sure if this correctly handles the case where Opcode is ANY_EXTEND and the existing loads are SEXT/ZEXT - in which case we need to ensure the final (legal) loads are fully SEXT/ZEXT loads - even though we don't care about the new extended bits we do still need the 'middle' extended bits to be correct.

Does that make sense? I'm not certain if this is handled - I may have gotten this wrong and haven't gotten a test case for this.

[Carrot]

ANY_EXTEND of SEXTLOAD/ZEXTLOAD should not be allowed.
In theory if both loads have same SEXTLOAD/ZEXTLOAD, we can change ExpOpcode to corresponding SIGN_EXTEND/ZERO_EXTEND, but it will make the code more complex, and this case should be extremely rare.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10097

You are right, ANY_EXTEND of SEXTLOAD/ZEXTLOAD was not correctly handled.

[RKSimon]

Does anyone else have any comments?

[RKSimon]

LGTM - cheers

llvm/test/CodeGen/X86/select-ext.ll

91

Please pre-commit these additional tests against trunk then update with this patch so the commit shows the diff/nodiff.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitSETCCCARRY(SDNode *N) {
 
   // If Carry is false, fold to a regular SETCC.
   if (isNullConstant(Carry))
     return DAG.getNode(ISD::SETCC, SDLoc(N), N->getVTList(), LHS, RHS, Cond);
 
   return SDValue();
 }
 
+/// Check if N satisfies:
+///   N is used once.
+///   N is a Load.
+///   The load is compatible with ExtOpcode. It means
+///     If load has explicit zero/sign extension, ExpOpcode must have the same
+///     extension.
+///     Otherwise returns true.
+static bool isCompatibleLoad(SDValue N, unsigned ExtOpcode) {
+  if (!N.hasOneUse())
+    return false;
+
+  if (!isa<LoadSDNode>(N))
+    return false;
+
+  LoadSDNode *Load = cast<LoadSDNode>(N);
+  ISD::LoadExtType LoadExt = Load->getExtensionType();
+  if (LoadExt == ISD::NON_EXTLOAD || LoadExt == ISD::EXTLOAD)
+    return true;
+
+  // Now LoadExt is either SEXTLOAD or ZEXTLOAD, ExtOpcode must have the same
+  // extension.
+  if ((LoadExt == ISD::SEXTLOAD && ExtOpcode != ISD::SIGN_EXTEND) ||
+      (LoadExt == ISD::ZEXTLOAD && ExtOpcode != ISD::ZERO_EXTEND))
+    return false;
+
+  return true;
+}
+
+/// Fold
+///   (sext (select c, load x, load y)) -> (select c, sextload x, sextload y)
+///   (zext (select c, load x, load y)) -> (select c, zextload x, zextload y)
+///   (aext (select c, load x, load y)) -> (select c, extload x, extload y)
+/// This function is called by the DAGCombiner when visiting sext/zext/aext
+/// dag nodes (see for example method DAGCombiner::visitSIGN_EXTEND).
+static SDValue tryToFoldExtendSelectLoad(SDNode *N, const TargetLowering &TLI,
+                                         SelectionDAG &DAG) {
+  unsigned Opcode = N->getOpcode();
+  SDValue N0 = N->getOperand(0);
+  EVT VT = N->getValueType(0);
+  SDLoc DL(N);
+
+  assert((Opcode == ISD::SIGN_EXTEND || Opcode == ISD::ZERO_EXTEND ||
+          Opcode == ISD::ANY_EXTEND) &&
+         "Expected EXTEND dag node in input!");
+
+  if (!(N0->getOpcode() == ISD::SELECT || N0->getOpcode() == ISD::VSELECT) ||
+      !N0.hasOneUse())
+    return SDValue();
+
+  SDValue Op1 = N0->getOperand(1);
+  SDValue Op2 = N0->getOperand(2);
+  if (!isCompatibleLoad(Op1, Opcode) || !isCompatibleLoad(Op2, Opcode))
+    return SDValue();
+
+  auto ExtLoadOpcode = ISD::EXTLOAD;
+  if (Opcode == ISD::SIGN_EXTEND)
+    ExtLoadOpcode = ISD::SEXTLOAD;
+  else if (Opcode == ISD::ZERO_EXTEND)
+    ExtLoadOpcode = ISD::ZEXTLOAD;
+
+  LoadSDNode *Load1 = cast<LoadSDNode>(Op1);
+  LoadSDNode *Load2 = cast<LoadSDNode>(Op2);
+  if (!TLI.isLoadExtLegal(ExtLoadOpcode, VT, Load1->getMemoryVT()) ||
+      !TLI.isLoadExtLegal(ExtLoadOpcode, VT, Load2->getMemoryVT()))
+    return SDValue();
+
+  SDValue Ext1 = DAG.getNode(Opcode, DL, VT, Op1);
+  SDValue Ext2 = DAG.getNode(Opcode, DL, VT, Op2);
+  return DAG.getSelect(DL, VT, N0->getOperand(0), Ext1, Ext2);
+}
+
 /// Try to fold a sext/zext/aext dag node into a ConstantSDNode or
 /// a build_vector of constants.
 /// This function is called by the DAGCombiner when visiting sext/zext/aext
 /// dag nodes (see for example method DAGCombiner::visitSIGN_EXTEND).
 /// Vector extends are not folded if operations are legal; this is to
 /// avoid introducing illegal build_vector dag nodes.
 static SDValue tryToFoldExtendOfConstant(SDNode *N, const TargetLowering &TLI,
                                          SelectionDAG &DAG, bool LegalTypes) {
if (SDValue NewXor = visitXOR(N0.getNode())) {
         return DAG.getNode(ISD::SIGN_EXTEND, DL, VT, NewXor);
       }
     }
 
     SDValue Zext = DAG.getNode(ISD::ZERO_EXTEND, DL, VT, N0.getOperand(0));
     return DAG.getNode(ISD::ADD, DL, VT, Zext, DAG.getAllOnesConstant(DL, VT));
   }
 
+  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG))
+    return Res;
+
   return SDValue();
 }
 
 // isTruncateOf - If N is a truncate of some other value, return true, record
 // the value being truncated in Op and which of Op's bits are zero/one in Known.
 // This function computes KnownBits to avoid a duplicated call to
 // computeKnownBits in the caller.
 static bool isTruncateOf(SelectionDAG &DAG, SDValue N, SDValue &Op,
SDValue DAGCombiner::visitZERO_EXTEND(SDNode *N) {
   }
 
   if (SDValue NewVSel = matchVSelectOpSizesWithSetCC(N))
     return NewVSel;
 
   if (SDValue NewCtPop = widenCtPop(N, DAG))
     return NewCtPop;
 
+  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG))
+    return Res;
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitANY_EXTEND(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   if (SDValue Res = tryToFoldExtendOfConstant(N, TLI, DAG, LegalTypes))
if (SDValue SCC = SimplifySelectCC(
             DAG.getConstant(0, DL, VT),
             cast<CondCodeSDNode>(N0.getOperand(2))->get(), true))
       return SCC;
   }
 
   if (SDValue NewCtPop = widenCtPop(N, DAG))
     return NewCtPop;
 
+  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG))
+    return Res;
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitAssertExt(SDNode *N) {
   unsigned Opcode = N->getOpcode();
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   EVT AssertVT = cast<VTSDNode>(N1)->getVT();

llvm/test/CodeGen/X86/select-ext.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+sse4.1 | FileCheck %s
 
-; TODO: (zext(select c, load1, load2)) -> (select c, zextload1, zextload2)
+; (zext(select c, load1, load2)) -> (select c, zextload1, zextload2)
 define i64 @zext_scalar(i8* %p, i1 zeroext %c) {
 ; CHECK-LABEL: zext_scalar:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movzbl (%rdi), %eax
-; CHECK-NEXT:    movzbl 1(%rdi), %ecx
+; CHECK-NEXT:    movzbl (%rdi), %ecx
+; CHECK-NEXT:    movzbl 1(%rdi), %eax
 ; CHECK-NEXT:    testl %esi, %esi
-; CHECK-NEXT:    cmovel %eax, %ecx
-; CHECK-NEXT:    movzbl %cl, %eax
+; CHECK-NEXT:    cmoveq %rcx, %rax
 ; CHECK-NEXT:    retq
   %ld1 = load volatile i8, i8* %p
   %arrayidx1 = getelementptr inbounds i8, i8* %p, i64 1
   %ld2 = load volatile i8, i8* %arrayidx1
   %cond.v = select i1 %c, i8 %ld2, i8 %ld1
   %cond = zext i8 %cond.v to i64
   ret i64 %cond
 }
 
 define i64 @zext_scalar2(i8* %p, i16* %q, i1 zeroext %c) {
 ; CHECK-LABEL: zext_scalar2:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movzbl (%rdi), %eax
-; CHECK-NEXT:    testl %edx, %edx
-; CHECK-NEXT:    je .LBB1_2
-; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    movzbl (%rdi), %ecx
 ; CHECK-NEXT:    movzwl (%rsi), %eax
-; CHECK-NEXT:  .LBB1_2:
-; CHECK-NEXT:    movzwl %ax, %eax
+; CHECK-NEXT:    testl %edx, %edx
+; CHECK-NEXT:    cmoveq %rcx, %rax
 ; CHECK-NEXT:    retq
   %ld1 = load volatile i8, i8* %p
   %ext_ld1 = zext i8 %ld1 to i16
   %ld2 = load volatile i16, i16* %q
   %cond.v = select i1 %c, i16 %ld2, i16 %ext_ld1
   %cond = zext i16 %cond.v to i64
   ret i64 %cond
 }
; CHECK-NEXT:    retq
   %ld1 = load volatile i8, i8* %p
   %ext_ld1 = sext i8 %ld1 to i16
   %ld2 = load volatile i16, i16* %q
   %cond.v = select i1 %c, i16 %ld2, i16 %ext_ld1
   %cond = zext i16 %cond.v to i64
   ret i64 %cond
 }
 
-; TODO: (sext(select c, load1, load2)) -> (select c, sextload1, sextload2)
+; (sext(select c, load1, load2)) -> (select c, sextload1, sextload2)
 define i64 @sext_scalar(i8* %p, i1 zeroext %c) {
 ; CHECK-LABEL: sext_scalar:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movzbl (%rdi), %eax
-; CHECK-NEXT:    movzbl 1(%rdi), %ecx
+; CHECK-NEXT:    movsbq (%rdi), %rcx
+; CHECK-NEXT:    movsbq 1(%rdi), %rax
 ; CHECK-NEXT:    testl %esi, %esi
-; CHECK-NEXT:    cmovel %eax, %ecx
-; CHECK-NEXT:    movsbq %cl, %rax
+; CHECK-NEXT:    cmoveq %rcx, %rax
 ; CHECK-NEXT:    retq
   %ld1 = load volatile i8, i8* %p
   %arrayidx1 = getelementptr inbounds i8, i8* %p, i64 1
   %ld2 = load volatile i8, i8* %arrayidx1
   %cond.v = select i1 %c, i8 %ld2, i8 %ld1
   %cond = sext i8 %cond.v to i64
   ret i64 %cond
 }
 
 ; Same as zext_scalar, but operate on vectors.
 define <2 x i64> @zext_vector_i1(<2 x i32>* %p, i1 zeroext %c) {
 ; CHECK-LABEL: zext_vector_i1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movq {{.*#+}} xmm1 = mem[0],zero
-; CHECK-NEXT:    movq {{.*#+}} xmm0 = mem[0],zero
+; CHECK-NEXT:    pmovzxdq {{.*#+}} xmm1 = mem[0],zero,mem[1],zero
+; CHECK-NEXT:    pmovzxdq {{.*#+}} xmm0 = mem[0],zero,mem[1],zero
 ; CHECK-NEXT:    testl %esi, %esi
 ; CHECK-NEXT:    jne .LBB4_2
 ; CHECK-NEXT:  # %bb.1:
 ; CHECK-NEXT:    movdqa %xmm1, %xmm0
 ; CHECK-NEXT:  .LBB4_2:
-; CHECK-NEXT:    pmovzxdq {{.*#+}} xmm0 = xmm0[0],zero,xmm0[1],zero
 ; CHECK-NEXT:    retq
   %ld1 = load volatile <2 x i32>, <2 x i32>* %p
   %arrayidx1 = getelementptr inbounds <2 x i32>, <2 x i32>* %p, i64 1
   %ld2 = load volatile <2 x i32>, <2 x i32>* %arrayidx1
   %cond.v = select i1 %c, <2 x i32> %ld2, <2 x i32> %ld1
   %cond = zext <2 x i32> %cond.v to <2 x i64>
   ret <2 x i64> %cond
 }
 
 define <2 x i64> @zext_vector_v2i1(<2 x i32>* %p, <2 x i1> %c) {
 ; CHECK-LABEL: zext_vector_v2i1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    pshufd {{.*#+}} xmm0 = xmm0[0,2,2,3]
-; CHECK-NEXT:    pslld $31, %xmm0
-; CHECK-NEXT:    movsd {{.*#+}} xmm1 = mem[0],zero
-; CHECK-NEXT:    movsd {{.*#+}} xmm2 = mem[0],zero
-; CHECK-NEXT:    blendvps %xmm0, %xmm2, %xmm1
-; CHECK-NEXT:    pmovzxdq {{.*#+}} xmm0 = xmm1[0],zero,xmm1[1],zero
+; CHECK-NEXT:    psllq $63, %xmm0
+; CHECK-NEXT:    pmovzxdq {{.*#+}} xmm1 = mem[0],zero,mem[1],zero
+; CHECK-NEXT:    pmovzxdq {{.*#+}} xmm2 = mem[0],zero,mem[1],zero
+; CHECK-NEXT:    blendvpd %xmm0, %xmm2, %xmm1
+; CHECK-NEXT:    movapd %xmm1, %xmm0
 ; CHECK-NEXT:    retq
   %ld1 = load volatile <2 x i32>, <2 x i32>* %p
   %arrayidx1 = getelementptr inbounds <2 x i32>, <2 x i32>* %p, i64 1
   %ld2 = load volatile <2 x i32>, <2 x i32>* %arrayidx1
   %cond.v = select <2 x i1> %c, <2 x i32> %ld2, <2 x i32> %ld1
   %cond = zext <2 x i32> %cond.v to <2 x i64>
   ret <2 x i64> %cond
 }
 
 ; Same as sext_scalar, but operate on vectors.
 define <2 x i64> @sext_vector_i1(<2 x i32>* %p, i1 zeroext %c) {
 ; CHECK-LABEL: sext_vector_i1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movq {{.*#+}} xmm1 = mem[0],zero
-; CHECK-NEXT:    movq {{.*#+}} xmm0 = mem[0],zero
+; CHECK-NEXT:    pmovsxdq (%rdi), %xmm1
+; CHECK-NEXT:    pmovsxdq 8(%rdi), %xmm0
 ; CHECK-NEXT:    testl %esi, %esi
 ; CHECK-NEXT:    jne .LBB6_2
 ; CHECK-NEXT:  # %bb.1:
 ; CHECK-NEXT:    movdqa %xmm1, %xmm0
 ; CHECK-NEXT:  .LBB6_2:
-; CHECK-NEXT:    pmovsxdq %xmm0, %xmm0
 ; CHECK-NEXT:    retq
   %ld1 = load volatile <2 x i32>, <2 x i32>* %p
   %arrayidx1 = getelementptr inbounds <2 x i32>, <2 x i32>* %p, i64 1
   %ld2 = load volatile <2 x i32>, <2 x i32>* %arrayidx1
   %cond.v = select i1 %c, <2 x i32> %ld2, <2 x i32> %ld1
   %cond = sext <2 x i32> %cond.v to <2 x i64>
   ret <2 x i64> %cond
 }
 
 define <2 x i64> @sext_vector_v2i1(<2 x i32>* %p, <2 x i1> %c) {
 ; CHECK-LABEL: sext_vector_v2i1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    pshufd {{.*#+}} xmm0 = xmm0[0,2,2,3]
-; CHECK-NEXT:    pslld $31, %xmm0
-; CHECK-NEXT:    movsd {{.*#+}} xmm1 = mem[0],zero
-; CHECK-NEXT:    movsd {{.*#+}} xmm2 = mem[0],zero
-; CHECK-NEXT:    blendvps %xmm0, %xmm2, %xmm1
-; CHECK-NEXT:    pmovsxdq %xmm1, %xmm0
+; CHECK-NEXT:    psllq $63, %xmm0
+; CHECK-NEXT:    pmovsxdq (%rdi), %xmm1
+; CHECK-NEXT:    pmovsxdq 8(%rdi), %xmm2
+; CHECK-NEXT:    blendvpd %xmm0, %xmm2, %xmm1
+; CHECK-NEXT:    movapd %xmm1, %xmm0
 ; CHECK-NEXT:    retq
   %ld1 = load volatile <2 x i32>, <2 x i32>* %p
   %arrayidx1 = getelementptr inbounds <2 x i32>, <2 x i32>* %p, i64 1
   %ld2 = load volatile <2 x i32>, <2 x i32>* %arrayidx1
   %cond.v = select <2 x i1> %c, <2 x i32> %ld2, <2 x i32> %ld1
   %cond = sext <2 x i32> %cond.v to <2 x i64>
   ret <2 x i64> %cond
 }