[DAGCombine] Combine signext_inreg of extract-extend

[peterwaller-arm]

Summary

The outer signext_inreg is redundant in the following:

Fold (signext_inreg (extract_subvector (zext|anyext|sext iN_value to _) _) from iN)
     -> (extract_subvector (signext iN_value to iM))

Tests are precommitted and clone those by analogy from the AND case in
the same file. Add a negative test to check extension width is handled
correctly.

This patch supersedes D130700.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13226

(style) Create this node inside the if() below - otherwise you end up with dead nodes that could possibly affect later transforms

[peterwaller-arm]

Thanks for the review comments, Simon!

• Avoid creating dead nodes.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13228

Does this need an operation legality check for ISD::SIGN_EXTEND ? (!LegalOperations || TLI.isOperationLegal(ISD::SIGN_EXTEND, InnerExtVT))?

[peterwaller-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13228

The combine only replaces an extant any-extend with an equivalent sign-extend. Is it an unreasonable assumption that if a target supports an any-extend it won't support the equivalent sign or zero extend? I can add the condition if you think this is suspect.

[david-arm]

The patch looks good to me @peterwaller-arm! Before I accept I'm just waiting a bit longer about the legality question.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13228

I guess it's not just replacing an ANY_EXTEND with SIGN_EXTEND - it also replaces ZERO_EXTEND with SIGN_EXTEND too, although the transformation looks valid. I guess I have the same question as @peterwaller-arm - would a target have a legal SIGN_EXTEND, but illegal ZERO_ or ANY_EXTEND?

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13228

Its unlikely I agree - are the tests affected if we add the legality check?

[peterwaller-arm]

Ihanks @RKSimon. The tests are unaffected. I've opted to follow your suggestion.

• Add legality check.
• Also add legality check for analagous zero-extending case in visitAND cc @david-arm.

[peterwaller-arm]

• Fix copy-paste error for the legality check for the ZERO_EXTEND case.

[david-arm]

LGTM!

[RKSimon]

LGTM - cheers!

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

6420

Please commit this separately referencing the phab patch where it was added (D130782?)

13226

Extendee.getValueType()

[peterwaller-arm]

• Rebase onto 898699831b54, which fixes the zero-extending case in visitAND.
  • A backport has been requested for this one in https://github.com/llvm/llvm-project/issues/57089.
• Use .getValueType().

[peterwaller-arm]

I will submit first thing next week unless there are further comments.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

6420

Done, and a backport requested in https://github.com/llvm/llvm-project/issues/57089 since this case went into LLVM-15-rc3.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitAND(SDNode *N) {
   if (N0.getOpcode() == ISD::EXTRACT_SUBVECTOR && N0.hasOneUse() && N1C &&
       ISD::isExtOpcode(N0.getOperand(0).getOpcode())) {
     SDValue Ext = N0.getOperand(0);
     EVT ExtVT = Ext->getValueType(0);
     SDValue Extendee = Ext->getOperand(0);
 
     unsigned ScalarWidth = Extendee.getValueType().getScalarSizeInBits();
     if (N1C->getAPIntValue().isMask(ScalarWidth) &&
         (!LegalOperations || TLI.isOperationLegal(ISD::ZERO_EXTEND, ExtVT))) {
       //    (and (extract_subvector (zext|anyext|sext v) _) iN_mask)
       // => (extract_subvector (iN_zeroext v))
       SDValue ZeroExtExtendee =
           DAG.getNode(ISD::ZERO_EXTEND, SDLoc(N), ExtVT, Extendee);
 
       return DAG.getNode(ISD::EXTRACT_SUBVECTOR, SDLoc(N), VT, ZeroExtExtendee,
                          N0.getOperand(1));
     }
SDValue DAGCombiner::visitSIGN_EXTEND_INREG(SDNode *N) {
 
   // Form (sext_inreg (bswap >> 16)) or (sext_inreg (rotl (bswap) 16))
   if (ExtVTBits <= 16 && N0.getOpcode() == ISD::OR) {
     if (SDValue BSwap = MatchBSwapHWordLow(N0.getNode(), N0.getOperand(0),
                                            N0.getOperand(1), false))
       return DAG.getNode(ISD::SIGN_EXTEND_INREG, SDLoc(N), VT, BSwap, N1);
   }
 
+  // Fold (iM_signext_inreg
+  //        (extract_subvector (zext|anyext|sext iN_v to _) _)
+  //        from iN)
+  //      -> (extract_subvector (signext iN_v to iM))
+  if (N0.getOpcode() == ISD::EXTRACT_SUBVECTOR && N0.hasOneUse() &&
+      ISD::isExtOpcode(N0.getOperand(0).getOpcode())) {
+    SDValue InnerExt = N0.getOperand(0);
+    EVT InnerExtVT = InnerExt->getValueType(0);
+    SDValue Extendee = InnerExt->getOperand(0);
+
+    if (ExtVTBits == Extendee.getValueType().getScalarSizeInBits() &&
+        (!LegalOperations ||
+         TLI.isOperationLegal(ISD::SIGN_EXTEND, InnerExtVT))) {
+      SDValue SignExtExtendee =
+          DAG.getNode(ISD::SIGN_EXTEND, SDLoc(N), InnerExtVT, Extendee);
+      return DAG.getNode(ISD::EXTRACT_SUBVECTOR, SDLoc(N), VT, SignExtExtendee,
+                         N0.getOperand(1));
+    }
+  }
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitEXTEND_VECTOR_INREG(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   // {s/z}ext_vector_inreg(undef) = 0 because the top bits must be the same.

llvm/test/CodeGen/AArch64/extract-subvec-combine.ll

; CHECK-NEXT:    ret
   %extract = call <4 x i16> @llvm.vector.extract.v4i16.v8i16(<8 x i16> %sext, i64 4)
   %and = and <4 x i16> %extract, <i16 255, i16 255, i16 255, i16 255>
   ret <4 x i16> %and
 }
 
 define <2 x i32> @sext_extract_zext_idx0(<4 x i16> %vec) nounwind {
 ; CHECK-LABEL: sext_extract_zext_idx0:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ushll v0.4s, v0.4h, #0
-; CHECK-NEXT:    shl v0.2s, v0.2s, #16
-; CHECK-NEXT:    sshr v0.2s, v0.2s, #16
+; CHECK-NEXT:    sshll v0.4s, v0.4h, #0
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %zext = zext <4 x i16> %vec to <4 x i32>
   %extract = call <2 x i32> @llvm.vector.extract.v2i32.v4i32(<4 x i32> %zext, i64 0)
   %sext_inreg_step0 = shl <2 x i32> %extract, <i32 16, i32 16>
   %sext_inreg = ashr <2 x i32> %sext_inreg_step0, <i32 16, i32 16>
   ret <2 x i32> %sext_inreg
 }
 
+; Negative test, combine should not fire if sign extension is for a different width.
+define <2 x i32> @sext_extract_zext_idx0_negtest(<4 x i16> %vec) nounwind {
+; CHECK-LABEL: sext_extract_zext_idx0_negtest:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ushll v0.4s, v0.4h, #0
+; CHECK-NEXT:    shl v0.2s, v0.2s, #17
+; CHECK-NEXT:    sshr v0.2s, v0.2s, #17
+; CHECK-NEXT:    ret
+  %zext = zext <4 x i16> %vec to <4 x i32>
+  %extract = call <2 x i32> @llvm.vector.extract.v2i32.v4i32(<4 x i32> %zext, i64 0)
+  %sext_inreg_step0 = shl <2 x i32> %extract, <i32 17, i32 17>
+  %sext_inreg = ashr <2 x i32> %sext_inreg_step0, <i32 17, i32 17>
+  ret <2 x i32> %sext_inreg
+}
+
 define <4 x i16> @sext_extract_sext_idx0(<8 x i8> %vec) nounwind {
 ; CHECK-LABEL: sext_extract_sext_idx0:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    sshll v0.8h, v0.8b, #0
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %sext = sext <8 x i8> %vec to <8 x i16>
   %extract = call <4 x i16> @llvm.vector.extract.v4i16.v8i16(<8 x i16> %sext, i64 0)
   %sext_inreg_step0 = shl <4 x i16> %extract, <i16 8, i16 8, i16 8, i16 8>
   %sext_inreg = ashr <4 x i16> %sext_inreg_step0, <i16 8, i16 8, i16 8, i16 8>
   ret <4 x i16> %sext_inreg
 }
 
 define <2 x i32> @sext_extract_zext_idx2(<4 x i16> %vec) nounwind {
 ; CHECK-LABEL: sext_extract_zext_idx2:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ushll v0.4s, v0.4h, #0
+; CHECK-NEXT:    sshll v0.4s, v0.4h, #0
 ; CHECK-NEXT:    ext v0.16b, v0.16b, v0.16b, #8
-; CHECK-NEXT:    shl v0.2s, v0.2s, #16
-; CHECK-NEXT:    sshr v0.2s, v0.2s, #16
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %zext = zext <4 x i16> %vec to <4 x i32>
   %extract = call <2 x i32> @llvm.vector.extract.v2i32.v4i32(<4 x i32> %zext, i64 2)
   %sext_inreg_step0 = shl <2 x i32> %extract, <i32 16, i32 16>
   %sext_inreg = ashr <2 x i32> %sext_inreg_step0, <i32 16, i32 16>
   ret <2 x i32> %sext_inreg
 }
 

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: