[AArch64] Match shuffle/cast as zext.

[fhahn]

Summary

Extend tryToWidenSetCCOperands logic to consider implicit zexts via
the shuffle/cast pattern introduced in D120571.

Depends on D120571.

[fhahn]

Adjust pattern after changes to D120571.

[t.p.northover]

This one looks OK, I think.

[t.p.northover]

Reopening, as we think endianness stuff needs checking.

[fhahn]

This is no longer relevant.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static SDValue performCSELCombine(SDNode *N,
                                   SelectionDAG &DAG) {
   // CSEL x, x, cc -> x
   if (N->getOperand(0) == N->getOperand(1))
     return N->getOperand(0);
 
   return performCONDCombine(N, DCI, DAG, 2, 3);
 }
 
+// If there a node zero extending V to v16i32 either through a explicit
+// ZERO_EXTEND node or implicitly through a shuffle and cast, return the node of
+// the extension result. Otherwise return nullptr.
+static SDNode *getZExtOrEquivalentShuffleCast(SDValue V, SelectionDAG &DAG) {
+  if (auto *ZExt =
+          DAG.getNodeIfExists(ISD::ZERO_EXTEND, DAG.getVTList(MVT::v16i32), V))
+    return ZExt;
+
+  // Check if there's a sequence of nodes implicitly zero extending V. That is
+  //  v = CONCAT_VECTORS V, undef, undef, undef
+  //  mask = BUILD_VECTOR <0, 64, 64, 64, 1, 64, 64, 64, ....>
+  //  v2 = VECTOR_SHUFFLE mask, v, <0, undef, undef, .....>
+  //  res = BIT_CAST v2 to v16i32
+  auto Und = DAG.getUNDEF(MVT::v16i8);
+  auto *Concat = DAG.getNodeIfExists(
+      ISD::CONCAT_VECTORS, DAG.getVTList(MVT::v64i8), {V, Und, Und, Und});
+  if (!Concat)
+    return nullptr;
+
+  SmallVector<int, 64> NewMask(64, 64);
+  for (unsigned i = 0; i < 16; i++)
+    NewMask[i * 4] = i;
+
+  SDValue Undef = DAG.getUNDEF(MVT::i8);
+  SmallVector<SDValue> Constants(64, Undef);
+  Constants[0] = DAG.getConstant(0, SDLoc(), MVT::i8);
+  SDValue ZeroSplat =
+      DAG.getNode(ISD::BUILD_VECTOR, SDLoc(), MVT::v64i8, Constants);
+  auto Shuffle = DAG.getVectorShuffle(MVT::v64i8, SDLoc(), SDValue(Concat, 0),
+                                      ZeroSplat, NewMask);
+  return DAG.getNodeIfExists(ISD::BITCAST, DAG.getVTList(MVT::v16i32), Shuffle);
+
+}
+
 // Try to re-use an already extended operand of a v16i8 SetCC feeding a
 // extended select. Doing so avoids requiring another full extension of the
 // SET_CC result when lowering the select.
 static SDValue tryToWidenSetCCOperands(SDNode *Op, SelectionDAG &DAG) {
   if (Op->getOperand(0).getValueType() != MVT::v16i8)
     return SDValue();
 
   if (any_of(Op->uses(), [](const SDNode *N) {
static SDValue tryToWidenSetCCOperands(SDNode *Op, SelectionDAG &DAG) {
   SDLoc DL(Op);
   SDValue Op0ExtV;
   SDValue Op1ExtV;
   ISD::CondCode CC = cast<CondCodeSDNode>(Op->getOperand(2))->get();
   // Check if the first operand of the SET_CC is already extended. If it is,
   // split the SET_CC and re-use the extended version of the operand.
   SDNode *Op0SExt = DAG.getNodeIfExists(
       ISD::SIGN_EXTEND, DAG.getVTList(MVT::v16i32), Op->getOperand(0));
-  SDNode *Op0ZExt = DAG.getNodeIfExists(
-      ISD::ZERO_EXTEND, DAG.getVTList(MVT::v16i32), Op->getOperand(0));
+  SDNode *Op0ZExt = getZExtOrEquivalentShuffleCast(Op->getOperand(0), DAG);
   if (Op0SExt && isSignedIntSetCC(CC)) {
     Op0ExtV = SDValue(Op0SExt, 0);
     Op1ExtV = DAG.getNode(ISD::SIGN_EXTEND, DL, MVT::v16i32, Op->getOperand(1));
   } else if (Op0ZExt) {
     Op0ExtV = SDValue(Op0ZExt, 0);
     // Try to convert signed to unsigned predicate, so zext can be used: %a >s
     // -1 -> zext(%a) <u 128.
     if (CC == ISD::SETGT && SplatVal.isAllOnes()) {

llvm/test/CodeGen/AArch64/vselect-ext.ll

 ; CHECK-NEXT:  Lloh1:
 ; CHECK-NEXT:    adrp x10, lCPI7_1@PAGE
 ; CHECK-NEXT:  Lloh2:
 ; CHECK-NEXT:    adrp x11, lCPI7_2@PAGE
 ; CHECK-NEXT:  Lloh3:
 ; CHECK-NEXT:    adrp x12, lCPI7_3@PAGE
 ; CHECK-NEXT:    movi.2d v2, #0000000000000000
 ; CHECK-NEXT:    mov x8, xzr
-; CHECK-NEXT:    movi.2d v4, #0xffffffffffffffff
+; CHECK-NEXT:    movi.2d v4, #0x0000ff000000ff
 ; CHECK-NEXT:  Lloh4:
 ; CHECK-NEXT:    ldr q0, [x9, lCPI7_0@PAGEOFF]
 ; CHECK-NEXT:  Lloh5:
 ; CHECK-NEXT:    ldr q3, [x10, lCPI7_1@PAGEOFF]
 ; CHECK-NEXT:  Lloh6:
 ; CHECK-NEXT:    ldr q5, [x11, lCPI7_2@PAGEOFF]
 ; CHECK-NEXT:  Lloh7:
 ; CHECK-NEXT:    ldr q6, [x12, lCPI7_3@PAGEOFF]
 ; CHECK-NEXT:  LBB7_1: ; %loop
 ; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    ldr q1, [x0, x8]
 ; CHECK-NEXT:    add x8, x8, #16
 ; CHECK-NEXT:    cmp x8, #128
-; CHECK-NEXT:    cmgt.16b v7, v1, v4
-; CHECK-NEXT:    tbl.16b v16, { v1, v2 }, v0
-; CHECK-NEXT:    tbl.16b v17, { v1, v2 }, v3
-; CHECK-NEXT:    sshll2.8h v20, v7, #0
-; CHECK-NEXT:    tbl.16b v18, { v1, v2 }, v5
-; CHECK-NEXT:    sshll2.4s v21, v20, #0
-; CHECK-NEXT:    sshll.4s v20, v20, #0
-; CHECK-NEXT:    tbl.16b v19, { v1, v2 }, v6
-; CHECK-NEXT:    sshll.8h v7, v7, #0
-; CHECK-NEXT:    and.16b v16, v16, v21
-; CHECK-NEXT:    and.16b v17, v17, v20
-; CHECK-NEXT:    stp q17, q16, [x1, #32]
-; CHECK-NEXT:    sshll2.4s v16, v7, #0
-; CHECK-NEXT:    sshll.4s v7, v7, #0
-; CHECK-NEXT:    and.16b v16, v18, v16
-; CHECK-NEXT:    and.16b v7, v19, v7
-; CHECK-NEXT:    stp q7, q16, [x1], #64
+; CHECK-NEXT:    tbl.16b v7, { v1, v2 }, v0
+; CHECK-NEXT:    tbl.16b v16, { v1, v2 }, v3
+; CHECK-NEXT:    tbl.16b v17, { v1, v2 }, v5
+; CHECK-NEXT:    tbl.16b v18, { v1, v2 }, v6
+; CHECK-NEXT:    cmhi.4s v19, v4, v7
+; CHECK-NEXT:    cmhi.4s v20, v4, v16
+; CHECK-NEXT:    cmhi.4s v21, v4, v17
+; CHECK-NEXT:    cmhi.4s v22, v4, v18
+; CHECK-NEXT:    and.16b v7, v7, v19
+; CHECK-NEXT:    and.16b v16, v16, v20
+; CHECK-NEXT:    stp q16, q7, [x1, #32]
+; CHECK-NEXT:    and.16b v7, v17, v21
+; CHECK-NEXT:    and.16b v16, v18, v22
+; CHECK-NEXT:    stp q16, q7, [x1], #64
 ; CHECK-NEXT:    b.ne LBB7_1
 ; CHECK-NEXT:  ; %bb.2: ; %exit
 ; CHECK-NEXT:    ret
 ; CHECK-NEXT:    .loh AdrpLdr Lloh3, Lloh7
 ; CHECK-NEXT:    .loh AdrpLdr Lloh2, Lloh6
 ; CHECK-NEXT:    .loh AdrpLdr Lloh1, Lloh5
 ; CHECK-NEXT:    .loh AdrpLdr Lloh0, Lloh4
 entry: