[AArch64] Lower INSERT_SUBVECTOR to LD1_POST

[asavonic]

Summary

Before this patch, (insert_subvector (scalar_to_vector) lane) and an
ADD used to generate LDR + INS + ADD (or an offset for a subsequent
LDR). Now this is folded into a single LD1_POST instruction.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
 
   setTargetDAGCombine(ISD::MUL);
 
   setTargetDAGCombine(ISD::SELECT);
   setTargetDAGCombine(ISD::VSELECT);
 
   setTargetDAGCombine(ISD::INTRINSIC_VOID);
   setTargetDAGCombine(ISD::INTRINSIC_W_CHAIN);
+  setTargetDAGCombine(ISD::INSERT_SUBVECTOR);
   setTargetDAGCombine(ISD::INSERT_VECTOR_ELT);
   setTargetDAGCombine(ISD::EXTRACT_VECTOR_ELT);
   setTargetDAGCombine(ISD::VECREDUCE_ADD);
   setTargetDAGCombine(ISD::STEP_VECTOR);
 
   setTargetDAGCombine(ISD::GlobalAddress);
 
   // In case of strict alignment, avoid an excessive number of byte wide stores.
if (DCI.isBeforeLegalizeOps())
     return SDValue();
 
   SelectionDAG &DAG = DCI.DAG;
   EVT VT = N->getValueType(0);
 
   if (VT.isScalableVector())
     return SDValue();
 
+  SDNode *LD = nullptr;
+  SDNode *LDUser = N;
+  // If N is an INSERT_SUBVECTOR, its index operand must be multiplied
+  // to get the index for LD1.
+  unsigned LaneScale = 1;
+
+  if (N->getOpcode() == ISD::INSERT_SUBVECTOR) {
+    // Handle (insert_subvector V (scalar_to_vector (load)) Idx)
+    SDValue ScalarToVec = N->getOperand(1);
+    if (ScalarToVec->getOpcode() != ISD::SCALAR_TO_VECTOR)
+      return SDValue();
+    LD = ScalarToVec->getOperand(0).getNode();
+    LDUser = ScalarToVec.getNode();
+
+    // Adjust the lane index so that (insert_subvector (... Addr) Idx)
+    // matches (LD1LANE Addr NewIdx).
+    EVT SubvecVT = ScalarToVec->getValueType(0);
+    LaneScale *= VT.getVectorNumElements() / SubvecVT.getVectorNumElements();
+    assert(LaneScale && "Invalid insert_subvector/scalar_to_vector sequence");
+  } else {
     unsigned LoadIdx = IsLaneOp ? 1 : 0;
-  SDNode *LD = N->getOperand(LoadIdx).getNode();
+    LD = N->getOperand(LoadIdx).getNode();
+  }
+
   // If it is not LOAD, can not do such combine.
   if (LD->getOpcode() != ISD::LOAD)
     return SDValue();
 
   // The vector lane must be a constant in the LD1LANE opcode.
   SDValue Lane;
   if (IsLaneOp) {
     Lane = N->getOperand(2);
     auto *LaneC = dyn_cast<ConstantSDNode>(Lane);
-    if (!LaneC || LaneC->getZExtValue() >= VT.getVectorNumElements())
+    if (!LaneC)
       return SDValue();
+
+    uint64_t LaneVal = LaneC->getZExtValue();
+    if (LaneVal >= VT.getVectorNumElements())
+      return SDValue();
+
+    if (LaneScale != 1) {
+      Lane = DAG.getVectorIdxConstant(LaneVal * LaneScale, SDLoc(LaneC));
+    }
   }
 
   LoadSDNode *LoadSDN = cast<LoadSDNode>(LD);
   EVT MemVT = LoadSDN->getMemoryVT();
+  EVT ResultVT = VT;
   // Check if memory operand is the same type as the vector element.
-  if (MemVT != VT.getVectorElementType())
+  if (MemVT != VT.getVectorElementType()) {
+    if (LoadSDN->getExtensionType() != ISD::EXTLOAD)
       return SDValue();
 
+    // For EXTLOAD we need to adjust load type and lane index to match
+    // the memory type.
+    LaneScale *= VT.getScalarSizeInBits() / MemVT.getScalarSizeInBits();
+    assert(LaneScale && "Invalid load anyext");
+    ResultVT =
+        EVT::getVectorVT(*DAG.getContext(), MemVT,
+                         VT.getFixedSizeInBits() / MemVT.getScalarSizeInBits());
+  }
+
   // Check if there are other uses. If so, do not combine as it will introduce
   // an extra load.
   for (SDNode::use_iterator UI = LD->use_begin(), UE = LD->use_end(); UI != UE;
        ++UI) {
     if (UI.getUse().getResNo() == 1) // Ignore uses of the chain result.
       continue;
-    if (*UI != N)
+    if (*UI != LDUser)
       return SDValue();
   }
 
   SDValue Addr = LD->getOperand(1);
   SDValue Vector = N->getOperand(0);
   // Search for a use of the address operand that is an increment.
   for (SDNode::use_iterator UI = Addr.getNode()->use_begin(), UE =
        Addr.getNode()->use_end(); UI != UE; ++UI) {
     SDNode *User = *UI;
     if (User->getOpcode() != ISD::ADD
         || UI.getUse().getResNo() != Addr.getResNo())
       continue;
 
     // If the increment is a constant, it must match the memory ref size.
     SDValue Inc = User->getOperand(User->getOperand(0) == Addr ? 1 : 0);
     if (ConstantSDNode *CInc = dyn_cast<ConstantSDNode>(Inc.getNode())) {
       uint32_t IncVal = CInc->getZExtValue();
-      unsigned NumBytes = VT.getScalarSizeInBits() / 8;
+      unsigned NumBytes = MemVT.getScalarSizeInBits() / 8;
       if (IncVal != NumBytes)
         continue;
       Inc = DAG.getRegister(AArch64::XZR, MVT::i64);
     }
 
     // To avoid cycle construction make sure that neither the load nor the add
     // are predecessors to each other or the Vector.
     SmallPtrSet<const SDNode *, 32> Visited;
for (SDNode::use_iterator UI = Addr.getNode()->use_begin(), UE =
     Ops.push_back(LD->getOperand(0));  // Chain
     if (IsLaneOp) {
       Ops.push_back(Vector);           // The vector to be inserted
       Ops.push_back(Lane);             // The lane to be inserted in the vector
     }
     Ops.push_back(Addr);
     Ops.push_back(Inc);
 
-    EVT Tys[3] = { VT, MVT::i64, MVT::Other };
+    EVT Tys[3] = {ResultVT, MVT::i64, MVT::Other};
     SDVTList SDTys = DAG.getVTList(Tys);
     unsigned NewOp = IsLaneOp ? AArch64ISD::LD1LANEpost : AArch64ISD::LD1DUPpost;
     SDValue UpdN = DAG.getMemIntrinsicNode(NewOp, SDLoc(N), SDTys, Ops,
                                            MemVT,
                                            LoadSDN->getMemOperand());
 
+    SDValue Result = (ResultVT == VT)
+                         ? SDValue(UpdN.getNode(), 0)
+                         : DAG.getNode(ISD::BITCAST, SDLoc(N), VT, UpdN);
+
     // Update the uses.
     SDValue NewResults[] = {
         SDValue(LD, 0),            // The result of load
         SDValue(UpdN.getNode(), 2) // Chain
     };
     DCI.CombineTo(LD, NewResults);
-    DCI.CombineTo(N, SDValue(UpdN.getNode(), 0));     // Dup/Inserted Result
+    DCI.CombineTo(N, Result);                         // Dup/Inserted Result
     DCI.CombineTo(User, SDValue(UpdN.getNode(), 1));  // Write back register
 
     break;
   }
   return SDValue();
 }
 
 /// Simplify ``Addr`` given that the top byte of it is ignored by HW during
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   case AArch64ISD::GLD1S_SXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1S_IMM_MERGE_ZERO:
     return performGLD1Combine(N, DAG);
   case AArch64ISD::VASHR:
   case AArch64ISD::VLSHR:
     return performVectorShiftCombine(N, *this, DCI);
   case ISD::INSERT_VECTOR_ELT:
     return performInsertVectorEltCombine(N, DCI);
+  case ISD::INSERT_SUBVECTOR:
+    return performPostLD1Combine(N, DCI, true);
   case ISD::EXTRACT_VECTOR_ELT:
     return performExtractVectorEltCombine(N, DAG);
   case ISD::VECREDUCE_ADD:
     return performVecReduceAddCombine(N, DCI.DAG, Subtarget);
   case ISD::INTRINSIC_VOID:
   case ISD::INTRINSIC_W_CHAIN:
     switch (cast<ConstantSDNode>(N->getOperand(1))->getZExtValue()) {
     case Intrinsic::aarch64_sve_prfb_gather_scalar_offset:

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

; CHECK-NEXT:    ret
   %i_0 = load i16, i16* %i16_ptr
   %v0 = insertelement <2 x i16> undef, i16 %i_0, i32 0
   ret <2 x i16> %v0
 }
 
 define <2 x i16> @test1(<2 x i16>* %v2i16_ptr) {
 ; CHECK-LABEL: test1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ld1 { v0.h }[0], [x0]
-; CHECK-NEXT:    add x8, x0, #2 // =2
-; CHECK-NEXT:    ld1 { v0.h }[2], [x8]
+; CHECK-NEXT:    ld1 { v0.h }[0], [x0], #2
+; CHECK-NEXT:    ld1 { v0.h }[2], [x0]
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %v2i16 = load <2 x i16>, <2 x i16>* %v2i16_ptr
   ret <2 x i16> %v2i16
 }
 
 define <2 x i16> @test2(i16* %i16_ptr, i64 %inc) {
 ; CHECK-LABEL: test2:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ld1 { v0.h }[0], [x0]
-; CHECK-NEXT:    add x8, x0, x1, lsl #1
-; CHECK-NEXT:    ld1 { v0.h }[2], [x8]
+; CHECK-NEXT:    lsl x8, x1, #1
+; CHECK-NEXT:    ld1 { v0.h }[0], [x0], x8
+; CHECK-NEXT:    ld1 { v0.h }[2], [x0]
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %i_0 = load i16, i16* %i16_ptr
   %i16_ptr_inc = getelementptr i16, i16* %i16_ptr, i64 %inc
   %i_1 = load i16, i16* %i16_ptr_inc
   %v0 = insertelement <2 x i16> undef, i16 %i_0, i32 0
   %v1 = insertelement <2 x i16> %v0, i16 %i_1, i32 1
   ret <2 x i16> %v1
 }
 
 define <2 x i8> @test3(<2 x i8>* %v2i8_ptr) {
 ; CHECK-LABEL: test3:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ld1 { v0.b }[0], [x0]
-; CHECK-NEXT:    add x8, x0, #1 // =1
-; CHECK-NEXT:    ld1 { v0.b }[4], [x8]
+; CHECK-NEXT:    ld1 { v0.b }[0], [x0], #1
+; CHECK-NEXT:    ld1 { v0.b }[4], [x0]
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %v2i8 = load <2 x i8>, <2 x i8>* %v2i8_ptr
   ret <2 x i8> %v2i8
 }
 
 define <4 x i8> @test4(<4 x i8>* %v4i8_ptr) {
 ; CHECK-LABEL: test4:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ld1 { v0.b }[0], [x0]
-; CHECK-NEXT:    add x8, x0, #1 // =1
+; CHECK-NEXT:    mov x8, x0
+; CHECK-NEXT:    ld1 { v0.b }[0], [x8], #1
 ; CHECK-NEXT:    ld1 { v0.b }[2], [x8]
 ; CHECK-NEXT:    add x8, x0, #2 // =2
 ; CHECK-NEXT:    ld1 { v0.b }[4], [x8]
 ; CHECK-NEXT:    add x8, x0, #3 // =3
 ; CHECK-NEXT:    ld1 { v0.b }[6], [x8]
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %v4i8 = load <4 x i8>, <4 x i8>* %v4i8_ptr
   ret <4 x i8> %v4i8
 }

llvm/test/CodeGen/AArch64/sadd_sat_vec.ll

; CHECK-NEXT:    ret
   %z = call <4 x i8> @llvm.sadd.sat.v4i8(<4 x i8> %x, <4 x i8> %y)
   store <4 x i8> %z, <4 x i8>* %pz
   ret void
 }
 
 define void @v2i8(<2 x i8>* %px, <2 x i8>* %py, <2 x i8>* %pz) nounwind {
 ; CHECK-LABEL: v2i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ld1 { v0.b }[0], [x1]
-; CHECK-NEXT:    ld1 { v1.b }[0], [x0]
-; CHECK-NEXT:    add x8, x0, #1 // =1
-; CHECK-NEXT:    add x9, x1, #1 // =1
-; CHECK-NEXT:    ld1 { v0.b }[4], [x9]
-; CHECK-NEXT:    ld1 { v1.b }[4], [x8]
+; CHECK-NEXT:    ld1 { v0.b }[0], [x0], #1
+; CHECK-NEXT:    ld1 { v0.b }[4], [x0]
+; CHECK-NEXT:    ld1 { v1.b }[0], [x1], #1
 ; CHECK-NEXT:    shl v0.2s, v0.2s, #24
+; CHECK-NEXT:    ld1 { v1.b }[4], [x1]
 ; CHECK-NEXT:    shl v1.2s, v1.2s, #24
-; CHECK-NEXT:    sqadd v0.2s, v1.2s, v0.2s
+; CHECK-NEXT:    sqadd v0.2s, v0.2s, v1.2s
 ; CHECK-NEXT:    ushr v0.2s, v0.2s, #24
 ; CHECK-NEXT:    mov w8, v0.s[1]
 ; CHECK-NEXT:    fmov w9, s0
 ; CHECK-NEXT:    strb w9, [x2]
 ; CHECK-NEXT:    strb w8, [x2, #1]
 ; CHECK-NEXT:    ret
   %x = load <2 x i8>, <2 x i8>* %px
   %y = load <2 x i8>, <2 x i8>* %py
; CHECK-NEXT:    ret
   %z = call <4 x i16> @llvm.sadd.sat.v4i16(<4 x i16> %x, <4 x i16> %y)
   store <4 x i16> %z, <4 x i16>* %pz
   ret void
 }
 
 define void @v2i16(<2 x i16>* %px, <2 x i16>* %py, <2 x i16>* %pz) nounwind {
 ; CHECK-LABEL: v2i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ld1 { v0.h }[0], [x1]
-; CHECK-NEXT:    ld1 { v1.h }[0], [x0]
-; CHECK-NEXT:    add x8, x0, #2 // =2
-; CHECK-NEXT:    add x9, x1, #2 // =2
-; CHECK-NEXT:    ld1 { v0.h }[2], [x9]
-; CHECK-NEXT:    ld1 { v1.h }[2], [x8]
+; CHECK-NEXT:    ld1 { v0.h }[0], [x0], #2
+; CHECK-NEXT:    ld1 { v0.h }[2], [x0]
+; CHECK-NEXT:    ld1 { v1.h }[0], [x1], #2
 ; CHECK-NEXT:    shl v0.2s, v0.2s, #16
+; CHECK-NEXT:    ld1 { v1.h }[2], [x1]
 ; CHECK-NEXT:    shl v1.2s, v1.2s, #16
-; CHECK-NEXT:    sqadd v0.2s, v1.2s, v0.2s
+; CHECK-NEXT:    sqadd v0.2s, v0.2s, v1.2s
 ; CHECK-NEXT:    ushr v0.2s, v0.2s, #16
 ; CHECK-NEXT:    mov w8, v0.s[1]
 ; CHECK-NEXT:    fmov w9, s0
 ; CHECK-NEXT:    strh w9, [x2]
 ; CHECK-NEXT:    strh w8, [x2, #2]
 ; CHECK-NEXT:    ret
   %x = load <2 x i16>, <2 x i16>* %px
   %y = load <2 x i16>, <2 x i16>* %py

llvm/test/CodeGen/AArch64/ssub_sat_vec.ll

; CHECK-NEXT:    ret
   %z = call <4 x i8> @llvm.ssub.sat.v4i8(<4 x i8> %x, <4 x i8> %y)
   store <4 x i8> %z, <4 x i8>* %pz
   ret void
 }
 
 define void @v2i8(<2 x i8>* %px, <2 x i8>* %py, <2 x i8>* %pz) nounwind {
 ; CHECK-LABEL: v2i8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ld1 { v0.b }[0], [x1]
-; CHECK-NEXT:    ld1 { v1.b }[0], [x0]
-; CHECK-NEXT:    add x8, x0, #1 // =1
-; CHECK-NEXT:    add x9, x1, #1 // =1
-; CHECK-NEXT:    ld1 { v0.b }[4], [x9]
-; CHECK-NEXT:    ld1 { v1.b }[4], [x8]
+; CHECK-NEXT:    ld1 { v0.b }[0], [x0], #1
+; CHECK-NEXT:    ld1 { v0.b }[4], [x0]
+; CHECK-NEXT:    ld1 { v1.b }[0], [x1], #1
 ; CHECK-NEXT:    shl v0.2s, v0.2s, #24
+; CHECK-NEXT:    ld1 { v1.b }[4], [x1]
 ; CHECK-NEXT:    shl v1.2s, v1.2s, #24
-; CHECK-NEXT:    sqsub v0.2s, v1.2s, v0.2s
+; CHECK-NEXT:    sqsub v0.2s, v0.2s, v1.2s
 ; CHECK-NEXT:    ushr v0.2s, v0.2s, #24
 ; CHECK-NEXT:    mov w8, v0.s[1]
 ; CHECK-NEXT:    fmov w9, s0
 ; CHECK-NEXT:    strb w9, [x2]
 ; CHECK-NEXT:    strb w8, [x2, #1]
 ; CHECK-NEXT:    ret
   %x = load <2 x i8>, <2 x i8>* %px
   %y = load <2 x i8>, <2 x i8>* %py
; CHECK-NEXT:    ret
   %z = call <4 x i16> @llvm.ssub.sat.v4i16(<4 x i16> %x, <4 x i16> %y)
   store <4 x i16> %z, <4 x i16>* %pz
   ret void
 }
 
 define void @v2i16(<2 x i16>* %px, <2 x i16>* %py, <2 x i16>* %pz) nounwind {
 ; CHECK-LABEL: v2i16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ld1 { v0.h }[0], [x1]
-; CHECK-NEXT:    ld1 { v1.h }[0], [x0]
-; CHECK-NEXT:    add x8, x0, #2 // =2
-; CHECK-NEXT:    add x9, x1, #2 // =2
-; CHECK-NEXT:    ld1 { v0.h }[2], [x9]
-; CHECK-NEXT:    ld1 { v1.h }[2], [x8]
+; CHECK-NEXT:    ld1 { v0.h }[0], [x0], #2
+; CHECK-NEXT:    ld1 { v0.h }[2], [x0]
+; CHECK-NEXT:    ld1 { v1.h }[0], [x1], #2
 ; CHECK-NEXT:    shl v0.2s, v0.2s, #16
+; CHECK-NEXT:    ld1 { v1.h }[2], [x1]
 ; CHECK-NEXT:    shl v1.2s, v1.2s, #16
-; CHECK-NEXT:    sqsub v0.2s, v1.2s, v0.2s
+; CHECK-NEXT:    sqsub v0.2s, v0.2s, v1.2s
 ; CHECK-NEXT:    ushr v0.2s, v0.2s, #16
 ; CHECK-NEXT:    mov w8, v0.s[1]
 ; CHECK-NEXT:    fmov w9, s0
 ; CHECK-NEXT:    strh w9, [x2]
 ; CHECK-NEXT:    strh w8, [x2, #2]
 ; CHECK-NEXT:    ret
   %x = load <2 x i16>, <2 x i16>* %px
   %y = load <2 x i16>, <2 x i16>* %py