[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

19–76

Worth adding sextload test case? What about vectors?

[Carrot]

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

19–76

Never thought about vector cases. Will look into it.

[Carrot]

Add more test cases.

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

19–76

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

19–76

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

[Carrot]

Add support for ISD::VSELECT.

[xbolva00]

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

2–3

Remove todo

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9977

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

10017

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

10017

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

36

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();
 }
 
+/// 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.hasOneUse())
+    return SDValue();
+
+  SDValue Op1 = N0->getOperand(1);
+  SDValue Op2 = N0->getOperand(2);
+  if (!isa<LoadSDNode>(Op1) || !isa<LoadSDNode>(Op2) || !Op1.hasOneUse()
+      || !Op2.hasOneUse())
+    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);
+  if (!TLI.isLoadExtLegal(ExtLoadOpcode, VT, Load1->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) {
   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 || Opcode == ISD::SIGN_EXTEND_VECTOR_INREG ||
          Opcode == ISD::ZERO_EXTEND_VECTOR_INREG)
          && "Expected EXTEND dag node in input!");
 
   // fold (sext c1) -> c1
   // fold (zext c1) -> c1
   // fold (aext c1) -> c1
   if (isa<ConstantSDNode>(N0))
     return DAG.getNode(Opcode, DL, VT, N0);
 
   // fold (sext (select cond, c1, c2)) -> (select cond, sext c1, sext c2)
   // fold (zext (select cond, c1, c2)) -> (select cond, zext c1, zext c2)
   // fold (aext (select cond, c1, c2)) -> (select cond, sext c1, sext c2)
   if (N0->getOpcode() == ISD::SELECT) {
     SDValue Op1 = N0->getOperand(1);
     SDValue Op2 = N0->getOperand(2);
     if (isa<ConstantSDNode>(Op1) && isa<ConstantSDNode>(Op2) &&
         (Opcode != ISD::ZERO_EXTEND || !TLI.isZExtFree(N0.getValueType(), VT))) {
       // For any_extend, choose sign extension of the constants to allow a
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

-; RUN: llc -mtriple=x86_64-unknown-unknown < %s | FileCheck %s
+; RUN: llc -mtriple=x86_64-unknown-unknown -mattr=+sse4.1 < %s | FileCheck %s
 
-; TODO: (zext(select c, load1, load2)) -> (select c, zextload1, zextload2)
-
-; CHECK-LABEL: foo
-; CHECK:        movzbl  (%rdi), %eax
-; CHECK-NEXT:   movzbl  1(%rdi), %ecx
+; (zext(select c, load1, load2)) -> (select c, zextload1, zextload2)
+; CHECK-LABEL: zext_scalar
+; CHECK:        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
 
-define i64 @foo(i8* %p, i1 zeroext %c) {
+define i64 @zext_scalar(i8* %p, i1 zeroext %c) {
   %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
 }
 
+
+; (sext(select c, load1, load2)) -> (select c, sextload1, sextload2)
+; CHECK-LABEL: sext_scalar
+; CHECK:        movsbq  (%rdi), %rcx
+; CHECK-NEXT:   movsbq  1(%rdi), %rax
+; CHECK-NEXT:   testl   %esi, %esi
+; CHECK-NEXT:   cmoveq  %rcx, %rax
+; CHECK-NEXT:   retq
+
+define i64 @sext_scalar(i8* %p, i1 zeroext %c) {
+  %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.
+; CHECK-LABEL: zext_vector
+; CHECK:        pmovzxdq (%rdi), %xmm1
+; CHECK-NEXT:   pmovzxdq 8(%rdi), %xmm0
+; CHECK-NEXT:   testl    %esi, %esi
+; CHECK-NEXT:   jne     .LBB2_2
+; CHECK:        movdqa  %xmm1, %xmm0
+; CHECK-NEXT:  .LBB2_2:
+; CHECK-NEXT:   retq
+
+define <2 x i64> @zext_vector(<2 x i32>* %p, i1 zeroext %c) {
+  %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
+}
+
+
+; Same as sext_scalar, but operate on vectors.
+; CHECK-LABEL: sext_vector
+; CHECK:        pmovsxdq (%rdi), %xmm1
+; CHECK-NEXT:   pmovsxdq 8(%rdi), %xmm0
+; CHECK-NEXT:   testl    %esi, %esi
+; CHECK-NEXT:   jne     .LBB3_2
+; CHECK:        movdqa  %xmm1, %xmm0
+; CHECK-NEXT:  .LBB3_2:
+; CHECK-NEXT:   retq
+define <2 x i64> @sext_vector(<2 x i32>* %p, i1 zeroext %c) {
+  %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
+}
+