[AArch64][DAGCombine] Add pushSignExtendThroughExtractSubvector

[peterwaller-arm]

Summary

Eliminate a SIGN_EXTEND_INREG by reifying an anyext.

   (sign_extend_inreg (extract_subvector (anyext v) Idx))
=> (extract_subvector (sign_extend v) Idx)

Depends on D130699.

[peterwaller-arm]

Remove commented code.

[peterwaller-arm]

• Update patch for changes to base patch D130698.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

15033–15034

Is this necessary? What are you trying to protect against?

15048

SIGN_EXTEND_INREG has a "type" operand that defines the size of the data being extended. You'll need to ensure this matches the element type of data being any extended.

[peterwaller-arm]

• Stylistic changes to shrink the code a little.
• Add check for signext_inreg type.
• Update tests.

[RKSimon]

Any reason this couldn't be done generically in DAGCombiner?

[paulwalker-arm]

In D130700#3685939, @RKSimon wrote:

Any reason this couldn't be done generically in DAGCombiner?

No reason. We updated the zext variant (D130782) to be a general combine and will do likewise here.

[peterwaller-arm]

New patch at https://reviews.llvm.org/D131503

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static SDValue pushAndThroughExtractSubvector(SDNode *N, SelectionDAG &DAG) {
   // => (extract_subvector (iN_zeroext v))
   SDValue ZeroExtExtendee =
       DAG.getNode(ISD::ZERO_EXTEND, SDLoc(N), AnyExtVT, Extendee);
 
   return DAG.getNode(ISD::EXTRACT_SUBVECTOR, SDLoc(N), VT, ZeroExtExtendee,
                      LHS->getOperand(1));
 }
 
+static SDValue pushSignExtendThroughExtractSubvector(SDNode *N,
+                                                     SelectionDAG &DAG) {
+  EVT VT = N->getValueType(0);
+  SDValue LHS = N->getOperand(0);
+
+  if (LHS->getOpcode() != ISD::EXTRACT_SUBVECTOR || !LHS->hasOneUse())
+    return SDValue();
+  SDValue AnyExt = LHS->getOperand(0);
+  if (AnyExt->getOpcode() != ISD::ANY_EXTEND)
+    return SDValue();
+  EVT AnyExtVT = AnyExt->getValueType(0);
+  SDValue Extendee = AnyExt->getOperand(0);
+
+  EVT SrcVT = cast<VTSDNode>(N->getOperand(1))->getVT();
+  if (Extendee.getValueType().getScalarType() != SrcVT.getScalarType())
+    return SDValue();
+
+  //    (sign_extend_inreg (extract_subvector (anyext v) Idx))
+  // => (extract_subvector (sign_extend v) Idx)
+  SDValue SignExtExtendee =
+      DAG.getNode(ISD::SIGN_EXTEND, SDLoc(N), AnyExtVT, Extendee);
+
+  return DAG.getNode(ISD::EXTRACT_SUBVECTOR, SDLoc(N), VT, SignExtExtendee,
+                     LHS->getOperand(1));
+}
+
 static SDValue performANDCombine(SDNode *N,
                                  TargetLowering::DAGCombinerInfo &DCI) {
   SelectionDAG &DAG = DCI.DAG;
   SDValue LHS = N->getOperand(0);
   SDValue RHS = N->getOperand(1);
   EVT VT = N->getValueType(0);
 
   if (SDValue R = performANDORCSELCombine(N, DAG))
if (Opc == AArch64ISD::UUNPKHI || Opc == AArch64ISD::UUNPKLO) {
     EVT ExtVT = VT.getDoubleNumVectorElementsVT(*DAG.getContext());
 
     SDValue Ext = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, ExtOp.getValueType(),
                               ExtOp, DAG.getValueType(ExtVT));
 
     return DAG.getNode(SOpc, DL, N->getValueType(0), Ext);
   }
 
+  if (SDValue R = pushSignExtendThroughExtractSubvector(N, DAG))
+    return R;
+
   if (DCI.isBeforeLegalizeOps())
     return SDValue();
 
   if (!EnableCombineMGatherIntrinsics)
     return SDValue();
 
   // SVE load nodes (e.g. AArch64ISD::GLD1) are straightforward candidates
   // for DAG Combine with SIGN_EXTEND_INREG. Bail out for all other nodes.

llvm/test/CodeGen/AArch64/sve-fixed-length-masked-gather.ll

 
 define void @masked_gather_v2f16(<2 x half>* %a, <2 x half*>* %b) vscale_range(2,0) #0 {
 ; CHECK-LABEL: masked_gather_v2f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr s1, [x0]
 ; CHECK-NEXT:    movi d0, #0000000000000000
 ; CHECK-NEXT:    ptrue p0.d, vl4
 ; CHECK-NEXT:    fcmeq v1.4h, v1.4h, #0.0
-; CHECK-NEXT:    ushll v1.4s, v1.4h, #0
-; CHECK-NEXT:    shl v1.2s, v1.2s, #16
-; CHECK-NEXT:    sshr v1.2s, v1.2s, #16
+; CHECK-NEXT:    sshll v1.4s, v1.4h, #0
 ; CHECK-NEXT:    fmov w8, s1
 ; CHECK-NEXT:    mov w9, v1.s[1]
 ; CHECK-NEXT:    ldr q1, [x1]
 ; CHECK-NEXT:    mov v0.h[0], w8
 ; CHECK-NEXT:    mov v0.h[1], w9
 ; CHECK-NEXT:    shl v0.4h, v0.4h, #15
 ; CHECK-NEXT:    cmlt v0.4h, v0.4h, #0
 ; CHECK-NEXT:    sunpklo z0.s, z0.h

llvm/test/CodeGen/AArch64/sve-fixed-length-masked-loads.ll

 define <2 x half> @masked_load_v2f16(<2 x half>* %ap, <2 x half>* %bp) vscale_range(2,0) #0 {
 ; CHECK-LABEL: masked_load_v2f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr s1, [x0]
 ; CHECK-NEXT:    movi d0, #0000000000000000
 ; CHECK-NEXT:    ldr s2, [x1]
 ; CHECK-NEXT:    ptrue p0.h, vl4
 ; CHECK-NEXT:    fcmeq v1.4h, v1.4h, v2.4h
-; CHECK-NEXT:    ushll v1.4s, v1.4h, #0
-; CHECK-NEXT:    shl v1.2s, v1.2s, #16
-; CHECK-NEXT:    sshr v1.2s, v1.2s, #16
+; CHECK-NEXT:    sshll v1.4s, v1.4h, #0
 ; CHECK-NEXT:    fmov w8, s1
 ; CHECK-NEXT:    mov w9, v1.s[1]
 ; CHECK-NEXT:    mov v0.h[0], w8
 ; CHECK-NEXT:    mov v0.h[1], w9
-; CHECK-NEXT:    shl v0.4h, v0.4h, #15
-; CHECK-NEXT:    cmlt v0.4h, v0.4h, #0
 ; CHECK-NEXT:    cmpne p0.h, p0/z, z0.h, #0
 ; CHECK-NEXT:    ld1h { z0.h }, p0/z, [x0]
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $z0
 ; CHECK-NEXT:    ret
   %a = load <2 x half>, <2 x half>* %ap
   %b = load <2 x half>, <2 x half>* %bp
   %mask = fcmp oeq <2 x half> %a, %b
   %load = call <2 x half> @llvm.masked.load.v2f16(<2 x half>* %ap, i32 8, <2 x i1> %mask, <2 x half> zeroinitializer)

llvm/test/CodeGen/AArch64/sve-fixed-length-masked-scatter.ll

 define void @masked_scatter_v2f16(<2 x half>* %a, <2 x half*>* %b) vscale_range(2,0) #0 {
 ; CHECK-LABEL: masked_scatter_v2f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr s1, [x0]
 ; CHECK-NEXT:    movi d0, #0000000000000000
 ; CHECK-NEXT:    ptrue p0.d, vl4
 ; CHECK-NEXT:    fcmeq v2.4h, v1.4h, #0.0
 ; CHECK-NEXT:    uunpklo z1.s, z1.h
-; CHECK-NEXT:    ushll v2.4s, v2.4h, #0
-; CHECK-NEXT:    shl v2.2s, v2.2s, #16
-; CHECK-NEXT:    sshr v2.2s, v2.2s, #16
+; CHECK-NEXT:    sshll v2.4s, v2.4h, #0
 ; CHECK-NEXT:    fmov w8, s2
 ; CHECK-NEXT:    mov w9, v2.s[1]
 ; CHECK-NEXT:    ldr q2, [x1]
 ; CHECK-NEXT:    mov v0.h[0], w8
 ; CHECK-NEXT:    mov v0.h[1], w9
 ; CHECK-NEXT:    shl v0.4h, v0.4h, #15
 ; CHECK-NEXT:    cmlt v0.4h, v0.4h, #0
 ; CHECK-NEXT:    sunpklo z0.s, z0.h

llvm/test/CodeGen/AArch64/sve-fixed-length-masked-stores.ll

 define void @masked_store_v2f16(<2 x half>* %ap, <2 x half>* %bp) vscale_range(2,0) #0 {
 ; CHECK-LABEL: masked_store_v2f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr s1, [x0]
 ; CHECK-NEXT:    movi d0, #0000000000000000
 ; CHECK-NEXT:    ldr s2, [x1]
 ; CHECK-NEXT:    ptrue p0.h, vl4
 ; CHECK-NEXT:    fcmeq v2.4h, v1.4h, v2.4h
-; CHECK-NEXT:    ushll v2.4s, v2.4h, #0
-; CHECK-NEXT:    shl v2.2s, v2.2s, #16
-; CHECK-NEXT:    sshr v2.2s, v2.2s, #16
+; CHECK-NEXT:    sshll v2.4s, v2.4h, #0
 ; CHECK-NEXT:    fmov w8, s2
 ; CHECK-NEXT:    mov w9, v2.s[1]
 ; CHECK-NEXT:    mov v0.h[0], w8
 ; CHECK-NEXT:    mov v0.h[1], w9
 ; CHECK-NEXT:    shl v0.4h, v0.4h, #15
 ; CHECK-NEXT:    cmlt v0.4h, v0.4h, #0
 ; CHECK-NEXT:    cmpne p0.h, p0/z, z0.h, #0
 ; CHECK-NEXT:    st1h { z1.h }, p0, [x1]

llvm/test/CodeGen/AArch64/vector-fcvt.ll

 ; CHECK-NEXT:    ret
   %1 = sitofp <16 x i8> %a to <16 x double>
   ret <16 x double> %1
 }
 
 define <8 x double> @sitofp_i16_double(<8 x i16> %a) {
 ; CHECK-LABEL: sitofp_i16_double:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ushll2 v1.4s, v0.8h, #0
-; CHECK-NEXT:    ushll v0.4s, v0.4h, #0
-; CHECK-NEXT:    shl v2.2s, v1.2s, #16
-; CHECK-NEXT:    shl v3.2s, v0.2s, #16
-; CHECK-NEXT:    ext v1.16b, v1.16b, v1.16b, #8
-; CHECK-NEXT:    sshr v2.2s, v2.2s, #16
-; CHECK-NEXT:    ext v0.16b, v0.16b, v0.16b, #8
-; CHECK-NEXT:    sshr v3.2s, v3.2s, #16
-; CHECK-NEXT:    sshll v2.2d, v2.2s, #0
-; CHECK-NEXT:    shl v1.2s, v1.2s, #16
-; CHECK-NEXT:    shl v0.2s, v0.2s, #16
-; CHECK-NEXT:    scvtf v2.2d, v2.2d
-; CHECK-NEXT:    sshr v1.2s, v1.2s, #16
-; CHECK-NEXT:    sshr v0.2s, v0.2s, #16
-; CHECK-NEXT:    sshll v3.2d, v3.2s, #0
+; CHECK-NEXT:    sshll2 v1.4s, v0.8h, #0
+; CHECK-NEXT:    sshll v0.4s, v0.4h, #0
+; CHECK-NEXT:    sshll2 v2.2d, v1.4s, #0
+; CHECK-NEXT:    sshll2 v3.2d, v0.4s, #0
 ; CHECK-NEXT:    sshll v4.2d, v1.2s, #0
-; CHECK-NEXT:    sshll v1.2d, v0.2s, #0
-; CHECK-NEXT:    scvtf v0.2d, v3.2d
-; CHECK-NEXT:    scvtf v1.2d, v1.2d
-; CHECK-NEXT:    scvtf v3.2d, v4.2d
+; CHECK-NEXT:    sshll v0.2d, v0.2s, #0
+; CHECK-NEXT:    scvtf v1.2d, v3.2d
+; CHECK-NEXT:    scvtf v0.2d, v0.2d
+; CHECK-NEXT:    scvtf v3.2d, v2.2d
+; CHECK-NEXT:    scvtf v2.2d, v4.2d
 ; CHECK-NEXT:    ret
   %1 = sitofp <8 x i16> %a to <8 x double>
   ret <8 x double> %1
 }
 
 define <8 x double> @sitofp_i32_double(<8 x i32> %a) {
 ; CHECK-LABEL: sitofp_i32_double:
 ; CHECK:       // %bb.0: