[AArch64] Combine Trunc(DUP) -> DUP

[dmgreen]

Summary

This adds a simple fold of TRUNCATE(AArch64ISD::DUP) -> AArch64ISD::DUP, which can help generate more optimal UMULL sequences, and seems useful in general.

[samtebbs]

Nice one, LGTM.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
   setTargetDAGCombine(ISD::SETCC);
 
   setTargetDAGCombine(ISD::INTRINSIC_WO_CHAIN);
 
   setTargetDAGCombine({ISD::ANY_EXTEND, ISD::ZERO_EXTEND, ISD::SIGN_EXTEND,
                        ISD::VECTOR_SPLICE, ISD::SIGN_EXTEND_INREG,
                        ISD::CONCAT_VECTORS, ISD::EXTRACT_SUBVECTOR,
                        ISD::INSERT_SUBVECTOR, ISD::STORE, ISD::BUILD_VECTOR});
+  setTargetDAGCombine(ISD::TRUNCATE);
   setTargetDAGCombine(ISD::LOAD);
 
   setTargetDAGCombine(ISD::MSTORE);
 
   setTargetDAGCombine(ISD::MUL);
 
   setTargetDAGCombine({ISD::SELECT, ISD::VSELECT});
 
if (Elt0->getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
     SDValue Ext = DAG.getNode(ISD::ANY_EXTEND, DL, ExtVT, VecToExtend);
     return DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MVT::v2i32, Ext,
                        SubvectorIdx);
   }
 
   return SDValue();
 }
 
+static SDValue performTruncateCombine(SDNode *N,
+                                      SelectionDAG &DAG) {
+  EVT VT = N->getValueType(0);
+  SDValue N0 = N->getOperand(0);
+  if (VT.isFixedLengthVector() && VT.is64BitVector() && N0.hasOneUse() &&
+      N0.getOpcode() == AArch64ISD::DUP) {
+    SDValue Op = N0.getOperand(0);
+    if (VT.getScalarType() == MVT::i32 &&
+        N0.getOperand(0).getValueType().getScalarType() == MVT::i64)
+      Op = DAG.getNode(ISD::TRUNCATE, SDLoc(N), MVT::i32, Op);
+    return DAG.getNode(N0.getOpcode(), SDLoc(N), VT, Op);
+  }
+
+  return SDValue();
+}
+
 // Check an node is an extend or shift operand
 static bool isExtendOrShiftOperand(SDValue N) {
   unsigned Opcode = N.getOpcode();
   if (Opcode == ISD::SIGN_EXTEND || Opcode == ISD::SIGN_EXTEND_INREG ||
       Opcode == ISD::ZERO_EXTEND || Opcode == ISD::ANY_EXTEND) {
     EVT SrcVT;
     if (Opcode == ISD::SIGN_EXTEND_INREG)
       SrcVT = cast<VTSDNode>(N.getOperand(1))->getVT();
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   default:
     LLVM_DEBUG(dbgs() << "Custom combining: skipping\n");
     break;
   case ISD::ADD:
   case ISD::SUB:
     return performAddSubCombine(N, DCI, DAG);
   case ISD::BUILD_VECTOR:
     return performBuildVectorCombine(N, DCI, DAG);
+  case ISD::TRUNCATE:
+    return performTruncateCombine(N, DAG);
   case AArch64ISD::ANDS:
     return performFlagSettingCombine(N, DCI, ISD::AND);
   case AArch64ISD::ADC:
     if (auto R = foldOverflowCheck(N, DAG, /* IsAdd */ true))
       return R;
     return foldADCToCINC(N, DAG);
   case AArch64ISD::SBC:
     return foldOverflowCheck(N, DAG, /* IsAdd */ false);

llvm/test/CodeGen/AArch64/aarch64-dup-ext.ll

entry:
     %broadcast.splat = shufflevector <2 x i16> %broadcast.splatinsert, <2 x i16> undef, <2 x i32> zeroinitializer
     %out = mul nsw <2 x i16> %broadcast.splat, %ext.b
     ret <2 x i16> %out
 }
 
 define <2 x i64> @dupzext_v2i16_v2i64(i16 %src, <2 x i16> %b) {
 ; CHECK-LABEL: dupzext_v2i16_v2i64:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    // kill: def $w0 killed $w0 def $x0
-; CHECK-NEXT:    and x8, x0, #0xffff
+; CHECK-NEXT:    and w8, w0, #0xffff
 ; CHECK-NEXT:    movi d1, #0x00ffff0000ffff
-; CHECK-NEXT:    dup v2.2d, x8
+; CHECK-NEXT:    dup v2.2s, w8
 ; CHECK-NEXT:    and v0.8b, v0.8b, v1.8b
-; CHECK-NEXT:    xtn v2.2s, v2.2d
 ; CHECK-NEXT:    umull v0.2d, v2.2s, v0.2s
 ; CHECK-NEXT:    ret
 entry:
     %in = zext i16 %src to i64
     %ext.b = zext <2 x i16> %b to <2 x i64>
     %broadcast.splatinsert = insertelement <2 x i64> undef, i64 %in, i64 0
     %broadcast.splat = shufflevector <2 x i64> %broadcast.splatinsert, <2 x i64> undef, <2 x i32> zeroinitializer
     %out = mul nuw <2 x i64> %broadcast.splat, %ext.b
; CHECK-NEXT:    ret
     ret void
 }
 
 define <8 x i16> @typei1_v8i1_v8i16(i1 %src, <8 x i1> %b) {
 ; CHECK-LABEL: typei1_v8i1_v8i16:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    and w8, w0, #0x1
 ; CHECK-NEXT:    movi v1.8b, #1
-; CHECK-NEXT:    dup v2.8h, w8
+; CHECK-NEXT:    dup v2.8b, w8
 ; CHECK-NEXT:    and v0.8b, v0.8b, v1.8b
-; CHECK-NEXT:    xtn v2.8b, v2.8h
 ; CHECK-NEXT:    umull v0.8h, v2.8b, v0.8b
 ; CHECK-NEXT:    ret
 entry:
     %in = zext i1 %src to i16
     %ext.b = zext <8 x i1> %b to <8 x i16>
     %broadcast.splatinsert = insertelement <8 x i16> undef, i16 %in, i16 0
     %broadcast.splat = shufflevector <8 x i16> %broadcast.splatinsert, <8 x i16> undef, <8 x i32> zeroinitializer
     %out = mul nsw <8 x i16> %broadcast.splat, %ext.b

llvm/test/CodeGen/AArch64/aarch64-smull.ll

entry:
   %out = mul nsw <8 x i32> %in1, %in2
   ret <8 x i32> %out
 }
 
 define <8 x i32> @umull_and_v8i32_dup(<8 x i16> %src1, i32 %src2) {
 ; CHECK-LABEL: umull_and_v8i32_dup:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    and w8, w0, #0xff
-; CHECK-NEXT:    ext v1.16b, v0.16b, v0.16b, #8
-; CHECK-NEXT:    dup v2.4s, w8
-; CHECK-NEXT:    xtn v2.4h, v2.4s
+; CHECK-NEXT:    dup v2.8h, w8
+; CHECK-NEXT:    umull2 v1.4s, v0.8h, v2.8h
 ; CHECK-NEXT:    umull v0.4s, v0.4h, v2.4h
-; CHECK-NEXT:    umull v1.4s, v1.4h, v2.4h
 ; CHECK-NEXT:    ret
 entry:
   %in1 = zext <8 x i16> %src1 to <8 x i32>
   %in2 = and i32 %src2, 255
   %broadcast.splatinsert = insertelement <8 x i32> undef, i32 %in2, i64 0
   %broadcast.splat = shufflevector <8 x i32> %broadcast.splatinsert, <8 x i32> undef, <8 x i32> zeroinitializer
   %out = mul nsw <8 x i32> %in1, %broadcast.splat
   ret <8 x i32> %out
entry:
   %in2 = and <4 x i64> %src2, <i64 255, i64 255, i64 255, i64 255>
   %out = mul nsw <4 x i64> %in1, %in2
   ret <4 x i64> %out
 }
 
 define <4 x i64> @umull_and_v4i64_dup(<4 x i32> %src1, i64 %src2) {
 ; CHECK-LABEL: umull_and_v4i64_dup:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    and x8, x0, #0xff
-; CHECK-NEXT:    ext v1.16b, v0.16b, v0.16b, #8
-; CHECK-NEXT:    dup v2.2d, x8
-; CHECK-NEXT:    xtn v2.2s, v2.2d
+; CHECK-NEXT:    and w8, w0, #0xff
+; CHECK-NEXT:    dup v2.4s, w8
+; CHECK-NEXT:    umull2 v1.2d, v0.4s, v2.4s
 ; CHECK-NEXT:    umull v0.2d, v0.2s, v2.2s
-; CHECK-NEXT:    umull v1.2d, v1.2s, v2.2s
 ; CHECK-NEXT:    ret
 entry:
   %in1 = zext <4 x i32> %src1 to <4 x i64>
   %in2 = and i64 %src2, 255
   %broadcast.splatinsert = insertelement <4 x i64> undef, i64 %in2, i64 0
   %broadcast.splat = shufflevector <4 x i64> %broadcast.splatinsert, <4 x i64> undef, <4 x i32> zeroinitializer
   %out = mul nsw <4 x i64> %in1, %broadcast.splat
   ret <4 x i64> %out
 }

llvm/test/CodeGen/AArch64/dag-combine-trunc-build-vec.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64 < %s | FileCheck %s
 
 ; This test is reduced fom https://github.com/android/ndk/issues/1207 for an
 ; issue with the following DAGCobmine operation:
 ;   truncate(build_vector(x,y)) -> build_vector(truncate(x),truncate(y))
 ; The combine should avoid creating illegal types if types have already been
 ; legalized.
 
 define void @no_combine(i32 %p) local_unnamed_addr {
 ; CHECK-LABEL: no_combine:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.4h, #4
-; CHECK-NEXT:    dup v1.4s, w0
-; CHECK-NEXT:    xtn v1.4h, v1.4s
+; CHECK-NEXT:    dup v1.4h, w0
 ; CHECK-NEXT:    mov v1.d[1], v0.d[0]
 ; CHECK-NEXT:    uzp1 v0.16b, v1.16b, v1.16b
 ; CHECK-NEXT:    str q0, [x8]
 ; CHECK-NEXT:    ret
 
 ; The two shufflevector operations are needed to force the DAGCombine to happen
 ; after type legalization and before operation legalization.  Removing either
 ; makes the combine to happen before type legalization and the issue no longer