[AArch64][SVE] Lower index_vector to step_vector

[junparser]

Summary

As discussed in D100107, this patch first convert index_vector to
step_vector, and convert step_vector back to index_vector after LegalizeDAG.

[junparser]

@paulwalker-arm @sdesmalen kindly ping~

[junparser]

rebased.

[junparser]

Kindly ping.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13658

You should be able to use Op2 directly and thus remove the need for the SPLAT_VECTOR and MUL.

13665

Is this a phased the rollout? As in I'm assuming that this will ultimately be expressed in terms of STEP_VECTOR as well.

15461–15462

This is already a requirement for the node and asserted with getNode() so doesn't need to be duplicated here.

[junparser]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13658

Sounds reasonable to me, update later.

13665

For non-constant op, we need lower index_vector into mul(splat(op2) step_vector()) + splat(base), mul may be lower to shl, even combine as mla with add which may confused the codegen. I'm not sure whether it is worth to deal this.

[junparser]

Remove assertion.

[junparser]

Address comments.

[paulwalker-arm]

Assuming you agree I think you can just delete a bit of code and then the patch is good to go.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13653–13663

With your previous DAGCombine patches in mind I figured I'd take this patch for a test drive and I believe there's no longer need for this specialisation because the generic block (i.e. the else case) lowers to the expected output (or at least running this patches version of sve-intrinsics-index.ll reports a pass).

15451–15453

FYI: I imagine we're only a few isel patterns away from not need this, but there's no need to hold up this patch for it.

[junparser]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13653–13663

Yes, previous DAGCombine patches can handle this. Then I'll remove the if part.

[junparser]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

15451–15453

Yep, We can remove index_vector node finally

[junparser]

Address comment.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13650–13652

Up to you but this change is not needed because Op1 and Op2 must be the same type.

[junparser]

@paulwalker-arm, Thanks for the review, I added D101593 which removes index_vector, FYI.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

private:
   SDValue LowerFLT_ROUNDS_(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSET_ROUNDING(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINSERT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerEXTRACT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSCALAR_TO_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerBUILD_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVECTOR_SHUFFLE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSPLAT_VECTOR(SDValue Op, SelectionDAG &DAG) const;
-  SDValue LowerSTEP_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDUPQLane(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerToPredicatedOp(SDValue Op, SelectionDAG &DAG, unsigned NewOp,
                               bool OverrideNEON = false) const;
   SDValue LowerToScalableOp(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerEXTRACT_SUBVECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINSERT_SUBVECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDIV(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerMUL(SDValue Op, SelectionDAG &DAG) const;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
   setTargetDAGCombine(ISD::SELECT);
   setTargetDAGCombine(ISD::VSELECT);
 
   setTargetDAGCombine(ISD::INTRINSIC_VOID);
   setTargetDAGCombine(ISD::INTRINSIC_W_CHAIN);
   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.
   MaxStoresPerMemsetOptSize = 8;
   MaxStoresPerMemset = Subtarget->requiresStrictAlign()
                        ? MaxStoresPerMemsetOptSize : 32;
 
for (auto VT : {MVT::nxv16i8, MVT::nxv8i16, MVT::nxv4i32, MVT::nxv2i64}) {
       setOperationAction(ISD::VECREDUCE_ADD, VT, Custom);
       setOperationAction(ISD::VECREDUCE_AND, VT, Custom);
       setOperationAction(ISD::VECREDUCE_OR, VT, Custom);
       setOperationAction(ISD::VECREDUCE_XOR, VT, Custom);
       setOperationAction(ISD::VECREDUCE_UMIN, VT, Custom);
       setOperationAction(ISD::VECREDUCE_UMAX, VT, Custom);
       setOperationAction(ISD::VECREDUCE_SMIN, VT, Custom);
       setOperationAction(ISD::VECREDUCE_SMAX, VT, Custom);
-      setOperationAction(ISD::STEP_VECTOR, VT, Custom);
 
       setOperationAction(ISD::UMUL_LOHI, VT, Expand);
       setOperationAction(ISD::SMUL_LOHI, VT, Expand);
       setOperationAction(ISD::SELECT_CC, VT, Expand);
     }
 
     // Illegal unpacked integer vector types.
     for (auto VT : {MVT::nxv8i8, MVT::nxv4i16, MVT::nxv2i32}) {
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::EXTRACT_VECTOR_ELT:
     return LowerEXTRACT_VECTOR_ELT(Op, DAG);
   case ISD::BUILD_VECTOR:
     return LowerBUILD_VECTOR(Op, DAG);
   case ISD::VECTOR_SHUFFLE:
     return LowerVECTOR_SHUFFLE(Op, DAG);
   case ISD::SPLAT_VECTOR:
     return LowerSPLAT_VECTOR(Op, DAG);
-  case ISD::STEP_VECTOR:
-    return LowerSTEP_VECTOR(Op, DAG);
   case ISD::EXTRACT_SUBVECTOR:
     return LowerEXTRACT_SUBVECTOR(Op, DAG);
   case ISD::INSERT_SUBVECTOR:
     return LowerINSERT_SUBVECTOR(Op, DAG);
   case ISD::SDIV:
   case ISD::UDIV:
     return LowerDIV(Op, DAG);
   case ISD::SMIN:
if (NumElts == 4) {
 
     if (Cost <= 4)
       return GeneratePerfectShuffle(PFEntry, V1, V2, DAG, dl);
   }
 
   return GenerateTBL(Op, ShuffleMask, DAG);
 }
 
-SDValue AArch64TargetLowering::LowerSTEP_VECTOR(SDValue Op,
-                                                SelectionDAG &DAG) const {
-  SDLoc dl(Op);
-  EVT VT = Op.getValueType();
-  assert(VT.isScalableVector() &&
-         "Only expect scalable vectors for STEP_VECTOR");
-  assert(VT.getScalarType() != MVT::i1 &&
-         "Vectors of i1 types not supported for STEP_VECTOR");
-
-  SDValue StepVal = Op.getOperand(0);
-  SDValue Zero = DAG.getConstant(0, dl, StepVal.getValueType());
-  return DAG.getNode(AArch64ISD::INDEX_VECTOR, dl, VT, Zero, StepVal);
-}
-
 SDValue AArch64TargetLowering::LowerSPLAT_VECTOR(SDValue Op,
                                                  SelectionDAG &DAG) const {
   SDLoc dl(Op);
   EVT VT = Op.getValueType();
   EVT ElemVT = VT.getScalarType();
   SDValue SplatVal = Op.getOperand(0);
 
   if (useSVEForFixedLengthVectorVT(VT))
SDValue AArch64TargetLowering::LowerDUPQLane(SDValue Op,
   // The ACLE says this must produce the same result as:
   //   svtbl(data, svadd_x(svptrue_b64(),
   //                       svand_x(svptrue_b64(), svindex_u64(0, 1), 1),
   //                       index * 2))
   SDValue One = DAG.getConstant(1, DL, MVT::i64);
   SDValue SplatOne = DAG.getNode(ISD::SPLAT_VECTOR, DL, MVT::nxv2i64, One);
 
   // create the vector 0,1,0,1,...
-  SDValue Zero = DAG.getConstant(0, DL, MVT::i64);
-  SDValue SV = DAG.getNode(AArch64ISD::INDEX_VECTOR,
-                           DL, MVT::nxv2i64, Zero, One);
+  SDValue SV = DAG.getNode(ISD::STEP_VECTOR, DL, MVT::nxv2i64, One);
   SV = DAG.getNode(ISD::AND, DL, MVT::nxv2i64, SV, SplatOne);
 
   // create the vector idx64,idx64+1,idx64,idx64+1,...
   SDValue Idx64 = DAG.getNode(ISD::ADD, DL, MVT::i64, Idx128, Idx128);
   SDValue SplatIdx64 = DAG.getNode(ISD::SPLAT_VECTOR, DL, MVT::nxv2i64, Idx64);
   SDValue ShuffleMask = DAG.getNode(ISD::ADD, DL, MVT::nxv2i64, SV, SplatIdx64);
 
   // create the vector Val[idx64],Val[idx64+1],Val[idx64],Val[idx64+1],...
return DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, N->getValueType(0),
                                  N->getOperand(1)),
                      DAG.getConstant(0, dl, MVT::i64));
 }
 
 static SDValue LowerSVEIntrinsicIndex(SDNode *N, SelectionDAG &DAG) {
   SDLoc DL(N);
   SDValue Op1 = N->getOperand(1);
   SDValue Op2 = N->getOperand(2);
-  EVT ScalarTy = Op1.getValueType();
-
-  if ((ScalarTy == MVT::i8) || (ScalarTy == MVT::i16)) {
-    Op1 = DAG.getNode(ISD::ANY_EXTEND, DL, MVT::i32, Op1);
-    Op2 = DAG.getNode(ISD::ANY_EXTEND, DL, MVT::i32, Op2);
-  }
+  EVT ScalarTy = Op2.getValueType();
+  if ((ScalarTy == MVT::i8) || (ScalarTy == MVT::i16))
+    ScalarTy = MVT::i32;
 
-  return DAG.getNode(AArch64ISD::INDEX_VECTOR, DL, N->getValueType(0),
-                     Op1, Op2);
+  // Lower index_vector(base, step) to mul(step step_vector(1)) + splat(base).
+  SDValue One = DAG.getConstant(1, DL, ScalarTy);
+  SDValue StepVector =
+      DAG.getNode(ISD::STEP_VECTOR, DL, N->getValueType(0), One);
+  SDValue Step = DAG.getNode(ISD::SPLAT_VECTOR, DL, N->getValueType(0), Op2);
+  SDValue Mul = DAG.getNode(ISD::MUL, DL, N->getValueType(0), StepVector, Step);
+  SDValue Base = DAG.getNode(ISD::SPLAT_VECTOR, DL, N->getValueType(0), Op1);
+  return DAG.getNode(ISD::ADD, DL, N->getValueType(0), Mul, Base);
 }
 
 static SDValue LowerSVEIntrinsicDUP(SDNode *N, SelectionDAG &DAG) {
   SDLoc dl(N);
   SDValue Scalar = N->getOperand(3);
   EVT ScalarTy = Scalar.getValueType();
 
   if ((ScalarTy == MVT::i8) || (ScalarTy == MVT::i16))
     Scalar = DAG.getNode(ISD::ANY_EXTEND, dl, MVT::i32, Scalar);
 
   SDValue Passthru = N->getOperand(1);
   SDValue Pred = N->getOperand(2);
if (!T->isSized() ||
     return SDValue();
 
   SDLoc DL(GN);
   SDValue Result = DAG.getGlobalAddress(GV, DL, MVT::i64, Offset);
   return DAG.getNode(ISD::SUB, DL, MVT::i64, Result,
                      DAG.getConstant(MinOffset, DL, MVT::i64));
 }
 
+static SDValue performStepVectorCombine(SDNode *N,
+                                        TargetLowering::DAGCombinerInfo &DCI,
+                                        SelectionDAG &DAG) {
+  if (!DCI.isAfterLegalizeDAG())
+    return SDValue();
+
+  SDLoc DL(N);
+  EVT VT = N->getValueType(0);
+  SDValue StepVal = N->getOperand(0);
+  SDValue Zero = DAG.getConstant(0, DL, StepVal.getValueType());
+  return DAG.getNode(AArch64ISD::INDEX_VECTOR, DL, VT, Zero, StepVal);
+}
+
 // Turns the vector of indices into a vector of byte offstes by scaling Offset
 // by (BitWidth / 8).
 static SDValue getScaledOffsetForBitWidth(SelectionDAG &DAG, SDValue Offset,
                                           SDLoc DL, unsigned BitWidth) {
   assert(Offset.getValueType().isScalableVector() &&
          "This method is only for scalable vectors of offsets");
 
   SDValue Shift = DAG.getConstant(Log2_32(BitWidth / 8), DL, MVT::i64);
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   case AArch64ISD::GLD1S_IMM_MERGE_ZERO:
     return performGLD1Combine(N, DAG);
   case ISD::INSERT_VECTOR_ELT:
     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::STEP_VECTOR:
+    return performStepVectorCombine(N, DCI, DAG);
   case ISD::INTRINSIC_VOID:
   case ISD::INTRINSIC_W_CHAIN:
     switch (cast<ConstantSDNode>(N->getOperand(1))->getZExtValue()) {
     case Intrinsic::aarch64_sve_prfb_gather_scalar_offset:
       return combineSVEPrefetchVecBaseImmOff(N, DAG, 1 /*=ScalarSizeInBytes*/);
     case Intrinsic::aarch64_sve_prfh_gather_scalar_offset:
       return combineSVEPrefetchVecBaseImmOff(N, DAG, 2 /*=ScalarSizeInBytes*/);
     case Intrinsic::aarch64_sve_prfw_gather_scalar_offset:

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

def : Pat<(nxv2f64 (AArch64dup (f64 fpimm:$val))),
 
   defm SQINCP_ZP   : sve_int_count_v<0b00000, "sqincp", int_aarch64_sve_sqincp>;
   defm UQINCP_ZP   : sve_int_count_v<0b00100, "uqincp", int_aarch64_sve_uqincp>;
   defm SQDECP_ZP   : sve_int_count_v<0b01000, "sqdecp", int_aarch64_sve_sqdecp>;
   defm UQDECP_ZP   : sve_int_count_v<0b01100, "uqdecp", int_aarch64_sve_uqdecp>;
   defm INCP_ZP     : sve_int_count_v<0b10000, "incp">;
   defm DECP_ZP     : sve_int_count_v<0b10100, "decp">;
 
-  defm INDEX_RR : sve_int_index_rr<"index", index_vector, index_vector_oneuse>;
-  defm INDEX_IR : sve_int_index_ir<"index", index_vector, index_vector_oneuse>;
+  defm INDEX_RR : sve_int_index_rr<"index", index_vector, index_vector_oneuse, AArch64mul_p_oneuse>;
+  defm INDEX_IR : sve_int_index_ir<"index", index_vector, index_vector_oneuse, AArch64mul_p, AArch64mul_p_oneuse>;
   defm INDEX_RI : sve_int_index_ri<"index", index_vector, index_vector_oneuse>;
   defm INDEX_II : sve_int_index_ii<"index", index_vector, index_vector_oneuse>;
 
   // Unpredicated shifts
   defm ASR_ZZI : sve_int_bin_cons_shift_imm_right<0b00, "asr", AArch64asr_p>;
   defm LSR_ZZI : sve_int_bin_cons_shift_imm_right<0b01, "lsr", AArch64lsr_p>;
   defm LSL_ZZI : sve_int_bin_cons_shift_imm_left< 0b11, "lsl", AArch64lsl_p>;
 

llvm/lib/Target/AArch64/SVEInstrFormats.td

: I<(outs zprty:$Zd), (ins imm_ty:$imm5, srcRegType:$Rm),
   let Inst{23-22} = sz8_64;
   let Inst{21}    = 0b1;
   let Inst{20-16} = Rm;
   let Inst{15-10} = 0b010010;
   let Inst{9-5}   = imm5;
   let Inst{4-0}   = Zd;
 }
 
-multiclass sve_int_index_ir<string asm, SDPatternOperator op, SDPatternOperator oneuseop> {
+multiclass sve_int_index_ir<string asm, SDPatternOperator op, SDPatternOperator oneuseop, SDPatternOperator mulop, SDPatternOperator muloneuseop> {
   def _B : sve_int_index_ir<0b00, asm, ZPR8, GPR32, simm5_8b>;
   def _H : sve_int_index_ir<0b01, asm, ZPR16, GPR32, simm5_16b>;
   def _S : sve_int_index_ir<0b10, asm, ZPR32, GPR32, simm5_32b>;
   def _D : sve_int_index_ir<0b11, asm, ZPR64, GPR64, simm5_64b>;
 
   def : Pat<(nxv16i8 (op simm5_8b:$imm5, GPR32:$Rm)),
             (!cast<Instruction>(NAME # "_B") simm5_8b:$imm5, GPR32:$Rm)>;
   def : Pat<(nxv8i16 (op simm5_16b:$imm5, GPR32:$Rm)),
             (!cast<Instruction>(NAME # "_H") simm5_16b:$imm5, GPR32:$Rm)>;
   def : Pat<(nxv4i32 (op simm5_32b:$imm5, GPR32:$Rm)),
             (!cast<Instruction>(NAME # "_S") simm5_32b:$imm5, GPR32:$Rm)>;
   def : Pat<(nxv2i64 (op simm5_64b:$imm5, GPR64:$Rm)),
             (!cast<Instruction>(NAME # "_D") simm5_64b:$imm5, GPR64:$Rm)>;
 
   // add(index_vector(zero, step), dup(X)) -> index_vector(X, step).
   def : Pat<(add (nxv16i8 (oneuseop (i32 0), GPR32:$Rm)), (nxv16i8 (AArch64dup(simm5_8b:$imm5)))),
             (!cast<Instruction>(NAME # "_B") simm5_8b:$imm5, GPR32:$Rm)>;
   def : Pat<(add (nxv8i16 (oneuseop (i32 0), GPR32:$Rm)), (nxv8i16 (AArch64dup(simm5_16b:$imm5)))),
             (!cast<Instruction>(NAME # "_H") simm5_16b:$imm5, GPR32:$Rm)>;
   def : Pat<(add (nxv4i32 (oneuseop (i32 0), GPR32:$Rm)), (nxv4i32 (AArch64dup(simm5_32b:$imm5)))),
             (!cast<Instruction>(NAME # "_S") simm5_32b:$imm5, GPR32:$Rm)>;
   def : Pat<(add (nxv2i64 (oneuseop (i64 0), GPR64:$Rm)), (nxv2i64 (AArch64dup(simm5_64b:$imm5)))),
             (!cast<Instruction>(NAME # "_D") simm5_64b:$imm5, GPR64:$Rm)>;
 
+  // mul(index_vector(0, 1), dup(Y)) -> index_vector(0, Y).
+  def : Pat<(mulop (nxv16i1 (AArch64ptrue 31)), (nxv16i8 (oneuseop (i32 0), (i32 1))), (nxv16i8 (AArch64dup(i32 GPR32:$Rm)))),
+            (!cast<Instruction>(NAME # "_B") (i32 0), GPR32:$Rm)>;
+  def : Pat<(mulop (nxv8i1 (AArch64ptrue 31)), (nxv8i16 (oneuseop (i32 0), (i32 1))), (nxv8i16 (AArch64dup(i32 GPR32:$Rm)))),
+            (!cast<Instruction>(NAME # "_H") (i32 0), GPR32:$Rm)>;
+  def : Pat<(mulop (nxv4i1 (AArch64ptrue 31)), (nxv4i32 (oneuseop (i32 0), (i32 1))), (nxv4i32 (AArch64dup(i32 GPR32:$Rm)))),
+            (!cast<Instruction>(NAME # "_S") (i32 0), GPR32:$Rm)>;
+  def : Pat<(mulop (nxv2i1 (AArch64ptrue 31)), (nxv2i64 (oneuseop (i64 0), (i64 1))), (nxv2i64 (AArch64dup(i64 GPR64:$Rm)))),
+            (!cast<Instruction>(NAME # "_D") (i64 0), GPR64:$Rm)>;
+
+  // add(mul(index_vector(0, 1), dup(Y), dup(X)) -> index_vector(X, Y).
+  def : Pat<(add (muloneuseop (nxv16i1 (AArch64ptrue 31)), (nxv16i8 (oneuseop (i32 0), (i32 1))), (nxv16i8 (AArch64dup(i32 GPR32:$Rm)))), (nxv16i8 (AArch64dup(simm5_8b:$imm5)))),
+            (!cast<Instruction>(NAME # "_B") simm5_8b:$imm5, GPR32:$Rm)>;
+  def : Pat<(add (muloneuseop (nxv8i1 (AArch64ptrue 31)), (nxv8i16 (oneuseop (i32 0), (i32 1))), (nxv8i16 (AArch64dup(i32 GPR32:$Rm)))), (nxv8i16 (AArch64dup(simm5_16b:$imm5)))),
+            (!cast<Instruction>(NAME # "_H") simm5_16b:$imm5, GPR32:$Rm)>;
+  def : Pat<(add (muloneuseop (nxv4i1 (AArch64ptrue 31)), (nxv4i32 (oneuseop (i32 0), (i32 1))), (nxv4i32 (AArch64dup(i32 GPR32:$Rm)))), (nxv4i32 (AArch64dup(simm5_32b:$imm5)))),
+            (!cast<Instruction>(NAME # "_S") simm5_32b:$imm5, GPR32:$Rm)>;
+  def : Pat<(add (muloneuseop (nxv2i1 (AArch64ptrue 31)), (nxv2i64 (oneuseop (i64 0), (i64 1))), (nxv2i64 (AArch64dup(i64 GPR64:$Rm)))), (nxv2i64 (AArch64dup(simm5_64b:$imm5)))),
+            (!cast<Instruction>(NAME # "_D") simm5_64b:$imm5, GPR64:$Rm)>;
 }
 
 class sve_int_index_ri<bits<2> sz8_64, string asm, ZPRRegOp zprty,
                        RegisterClass srcRegType, Operand imm_ty>
 : I<(outs zprty:$Zd), (ins srcRegType:$Rn, imm_ty:$imm5),
   asm, "\t$Zd, $Rn, $imm5",
   "", []>, Sched<[]> {
   bits<5> Rn;
: I<(outs zprty:$Zd), (ins srcRegType:$Rn, srcRegType:$Rm),
   let Inst{23-22} = sz8_64;
   let Inst{21}    = 0b1;
   let Inst{20-16} = Rm;
   let Inst{15-10} = 0b010011;
   let Inst{9-5}   = Rn;
   let Inst{4-0}   = Zd;
 }
 
-multiclass sve_int_index_rr<string asm, SDPatternOperator op, SDPatternOperator oneuseop> {
+multiclass sve_int_index_rr<string asm, SDPatternOperator op, SDPatternOperator oneuseop, SDPatternOperator mulop> {
   def _B : sve_int_index_rr<0b00, asm, ZPR8, GPR32>;
   def _H : sve_int_index_rr<0b01, asm, ZPR16, GPR32>;
   def _S : sve_int_index_rr<0b10, asm, ZPR32, GPR32>;
   def _D : sve_int_index_rr<0b11, asm, ZPR64, GPR64>;
 
   def : SVE_2_Op_Pat<nxv16i8, op, i32, i32, !cast<Instruction>(NAME # _B)>;
   def : SVE_2_Op_Pat<nxv8i16, op, i32, i32, !cast<Instruction>(NAME # _H)>;
   def : SVE_2_Op_Pat<nxv4i32, op, i32, i32, !cast<Instruction>(NAME # _S)>;
   def : SVE_2_Op_Pat<nxv2i64, op, i64, i64, !cast<Instruction>(NAME # _D)>;
 
   // add(index_vector(zero, step), dup(X)) -> index_vector(X, step).
   def : Pat<(add (nxv16i8 (oneuseop (i32 0), GPR32:$Rm)), (nxv16i8 (AArch64dup(i32 GPR32:$Rn)))),
             (!cast<Instruction>(NAME # "_B") GPR32:$Rn, GPR32:$Rm)>;
   def : Pat<(add (nxv8i16 (oneuseop (i32 0), GPR32:$Rm)), (nxv8i16 (AArch64dup(i32 GPR32:$Rn)))),
             (!cast<Instruction>(NAME # "_H") GPR32:$Rn, GPR32:$Rm)>;
   def : Pat<(add (nxv4i32 (oneuseop (i32 0), GPR32:$Rm)), (nxv4i32 (AArch64dup(i32 GPR32:$Rn)))),
             (!cast<Instruction>(NAME # "_S") GPR32:$Rn, GPR32:$Rm)>;
   def : Pat<(add (nxv2i64 (oneuseop (i64 0), GPR64:$Rm)), (nxv2i64 (AArch64dup(i64 GPR64:$Rn)))),
             (!cast<Instruction>(NAME # "_D") GPR64:$Rn, GPR64:$Rm)>;
+
+  // add(mul(index_vector(0, 1), dup(Y), dup(X)) -> index_vector(X, Y).
+  def : Pat<(add (mulop (nxv16i1 (AArch64ptrue 31)), (nxv16i8 (oneuseop (i32 0), (i32 1))), (nxv16i8 (AArch64dup(i32 GPR32:$Rm)))), (nxv16i8 (AArch64dup(i32 GPR32:$Rn)))),
+            (!cast<Instruction>(NAME # "_B") GPR32:$Rn, GPR32:$Rm)>;
+  def : Pat<(add (mulop (nxv8i1 (AArch64ptrue 31)), (nxv8i16 (oneuseop (i32 0), (i32 1))), (nxv8i16 (AArch64dup(i32 GPR32:$Rm)))),(nxv8i16 (AArch64dup(i32 GPR32:$Rn)))),
+            (!cast<Instruction>(NAME # "_H") GPR32:$Rn, GPR32:$Rm)>;
+  def : Pat<(add (mulop (nxv4i1 (AArch64ptrue 31)), (nxv4i32 (oneuseop (i32 0), (i32 1))), (nxv4i32 (AArch64dup(i32 GPR32:$Rm)))),(nxv4i32 (AArch64dup(i32 GPR32:$Rn)))),
+            (!cast<Instruction>(NAME # "_S") GPR32:$Rn, GPR32:$Rm)>;
+  def : Pat<(add (mulop (nxv2i1 (AArch64ptrue 31)), (nxv2i64 (oneuseop (i64 0), (i64 1))), (nxv2i64 (AArch64dup(i64 GPR64:$Rm)))),(nxv2i64 (AArch64dup(i64 GPR64:$Rn)))),
+            (!cast<Instruction>(NAME # "_D") GPR64:$Rn, GPR64:$Rm)>;
 }
 
 //===----------------------------------------------------------------------===//
 // SVE Bitwise Shift - Predicated Group
 //===----------------------------------------------------------------------===//
 
 class sve_int_bin_pred_shift_imm<bits<4> tsz8_64, bits<4> opc, string asm,
                                  ZPRRegOp zprty, Operand immtype>

llvm/test/CodeGen/AArch64/sve-intrinsics-index.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
 
 ;
 ; INDEX (IMMEDIATES)
 ;
 
 define <vscale x 16 x i8> @index_ii_i8() {
 ; CHECK-LABEL: index_ii_i8:
-; CHECK: index z0.b, #-16, #15
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.b, #-16, #15
+; CHECK-NEXT:    ret
   %out = call <vscale x 16 x i8> @llvm.aarch64.sve.index.nxv16i8(i8 -16, i8 15)
   ret <vscale x 16 x i8> %out
 }
 
 define <vscale x 8 x i16> @index_ii_i16() {
 ; CHECK-LABEL: index_ii_i16:
-; CHECK: index z0.h, #15, #-16
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.h, #15, #-16
+; CHECK-NEXT:    ret
   %out = call <vscale x 8 x i16> @llvm.aarch64.sve.index.nxv8i16(i16 15, i16 -16)
   ret <vscale x 8 x i16> %out
 }
 
 define <vscale x 4 x i32> @index_ii_i32() {
 ; CHECK-LABEL: index_ii_i32:
-; CHECK: index z0.s, #-16, #15
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.s, #-16, #15
+; CHECK-NEXT:    ret
   %out = call <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32 -16, i32 15)
   ret <vscale x 4 x i32> %out
 }
 
 define <vscale x 2 x i64> @index_ii_i64() {
 ; CHECK-LABEL: index_ii_i64:
-; CHECK: index z0.d, #15, #-16
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.d, #15, #-16
+; CHECK-NEXT:    ret
   %out = call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 15, i64 -16)
   ret <vscale x 2 x i64> %out
 }
 
 define <vscale x 2 x i64> @index_ii_range() {
 ; CHECK-LABEL: index_ii_range:
-; CHECK: mov w8, #16
-; CHECK-NEXT: mov x9, #-17
-; CHECK-NEXT: index z0.d, x9, x8
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #16
+; CHECK-NEXT:    mov x9, #-17
+; CHECK-NEXT:    index z0.d, x9, x8
+; CHECK-NEXT:    ret
   %out = call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 -17, i64 16)
   ret <vscale x 2 x i64> %out
 }
 
+define <vscale x 8 x i16> @index_ii_range_combine(i16 %a) {
+; CHECK-LABEL: index_ii_range_combine:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.h, #2, #8
+; CHECK-NEXT:    ret
+  %val = insertelement <vscale x 8 x i16> poison, i16 2, i32 0
+  %val1 = shufflevector <vscale x 8 x i16> %val, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+  %val2 = call <vscale x 8 x i16> @llvm.aarch64.sve.index.nxv8i16(i16 0, i16 2)
+  %val3 = shl <vscale x 8 x i16> %val2, %val1
+  %out = add <vscale x 8 x i16> %val3, %val1
+  ret <vscale x 8 x i16> %out
+}
+
 ;
 ; INDEX (IMMEDIATE, SCALAR)
 ;
 
 define <vscale x 16 x i8> @index_ir_i8(i8 %a) {
 ; CHECK-LABEL: index_ir_i8:
-; CHECK: index z0.b, #15, w0
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.b, #15, w0
+; CHECK-NEXT:    ret
   %out = call <vscale x 16 x i8> @llvm.aarch64.sve.index.nxv16i8(i8 15, i8 %a)
   ret <vscale x 16 x i8> %out
 }
 
 define <vscale x 8 x i16> @index_ir_i16(i16 %a) {
 ; CHECK-LABEL: index_ir_i16:
-; CHECK: index z0.h, #-16, w0
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.h, #-16, w0
+; CHECK-NEXT:    ret
   %out = call <vscale x 8 x i16> @llvm.aarch64.sve.index.nxv8i16(i16 -16, i16 %a)
   ret <vscale x 8 x i16> %out
 }
 
 define <vscale x 4 x i32> @index_ir_i32(i32 %a) {
 ; CHECK-LABEL: index_ir_i32:
-; CHECK: index z0.s, #15, w0
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.s, #15, w0
+; CHECK-NEXT:    ret
   %out = call <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32 15, i32 %a)
   ret <vscale x 4 x i32> %out
 }
 
 define <vscale x 2 x i64> @index_ir_i64(i64 %a) {
 ; CHECK-LABEL: index_ir_i64:
-; CHECK: index z0.d, #-16, x0
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.d, #-16, x0
+; CHECK-NEXT:    ret
   %out = call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 -16, i64 %a)
   ret <vscale x 2 x i64> %out
 }
 
 define <vscale x 4 x i32> @index_ir_range(i32 %a) {
 ; CHECK-LABEL: index_ir_range:
-; CHECK: mov w8, #-17
-; CHECK: index z0.s, w8, w0
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #-17
+; CHECK-NEXT:    index z0.s, w8, w0
+; CHECK-NEXT:    ret
   %out = call <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32 -17, i32 %a)
   ret <vscale x 4 x i32> %out
 }
 
+define <vscale x 4 x i32> @index_ir_range_combine(i32 %a) {
+; CHECK-LABEL: index_ir_range_combine:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.s, #0, w0
+; CHECK-NEXT:    ret
+  %val = insertelement <vscale x 4 x i32> poison, i32 2, i32 0
+  %val1 = shufflevector <vscale x 4 x i32> %val, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+  %tmp = call <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32 2, i32 1)
+  %tmp1 = sub <vscale x 4 x i32> %tmp, %val1
+  %val2 = insertelement <vscale x 4 x i32> poison, i32 %a, i32 0
+  %val3 = shufflevector <vscale x 4 x i32> %val2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+  %out = mul <vscale x 4 x i32> %tmp1, %val3
+  ret <vscale x 4 x i32> %out
+}
+
 ;
 ; INDEX (SCALAR, IMMEDIATE)
 ;
 
 define <vscale x 16 x i8> @index_ri_i8(i8 %a) {
 ; CHECK-LABEL: index_ri_i8:
-; CHECK: index z0.b, w0, #-16
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.b, w0, #-16
+; CHECK-NEXT:    ret
   %out = call <vscale x 16 x i8> @llvm.aarch64.sve.index.nxv16i8(i8 %a, i8 -16)
   ret <vscale x 16 x i8> %out
 }
 
 define <vscale x 8 x i16> @index_ri_i16(i16 %a) {
 ; CHECK-LABEL: index_ri_i16:
-; CHECK: index z0.h, w0, #15
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.h, w0, #15
+; CHECK-NEXT:    ret
   %out = call <vscale x 8 x i16> @llvm.aarch64.sve.index.nxv8i16(i16 %a, i16 15)
   ret <vscale x 8 x i16> %out
 }
 
 define <vscale x 4 x i32> @index_ri_i32(i32 %a) {
 ; CHECK-LABEL: index_ri_i32:
-; CHECK: index z0.s, w0, #-16
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.s, w0, #-16
+; CHECK-NEXT:    ret
   %out = call <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32 %a, i32 -16)
   ret <vscale x 4 x i32> %out
 }
 
 define <vscale x 2 x i64> @index_ri_i64(i64 %a) {
 ; CHECK-LABEL: index_ri_i64:
-; CHECK: index z0.d, x0, #15
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.d, x0, #15
+; CHECK-NEXT:    ret
   %out = call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 %a, i64 15)
   ret <vscale x 2 x i64> %out
 }
 
 define <vscale x 8 x i16> @index_ri_range(i16 %a) {
 ; CHECK-LABEL: index_ri_range:
-; CHECK: mov w8, #16
-; CHECK: index z0.h, w0, w8
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #16
+; CHECK-NEXT:    index z0.h, w0, w8
+; CHECK-NEXT:    ret
   %out = call <vscale x 8 x i16> @llvm.aarch64.sve.index.nxv8i16(i16 %a, i16 16)
   ret <vscale x 8 x i16> %out
 }
 
 ;
 ; INDEX (SCALARS)
 ;
 
 define <vscale x 16 x i8> @index_rr_i8(i8 %a, i8 %b) {
 ; CHECK-LABEL: index_rr_i8:
-; CHECK: index z0.b, w0, w1
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.b, w0, w1
+; CHECK-NEXT:    ret
   %out = call <vscale x 16 x i8> @llvm.aarch64.sve.index.nxv16i8(i8 %a, i8 %b)
   ret <vscale x 16 x i8> %out
 }
 
 define <vscale x 8 x i16> @index_rr_i16(i16 %a, i16 %b) {
 ; CHECK-LABEL: index_rr_i16:
-; CHECK: index z0.h, w0, w1
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.h, w0, w1
+; CHECK-NEXT:    ret
   %out = call <vscale x 8 x i16> @llvm.aarch64.sve.index.nxv8i16(i16 %a, i16 %b)
   ret <vscale x 8 x i16> %out
 }
 
 define <vscale x 4 x i32> @index_rr_i32(i32 %a, i32 %b) {
 ; CHECK-LABEL: index_rr_i32:
-; CHECK: index z0.s, w0, w1
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.s, w0, w1
+; CHECK-NEXT:    ret
   %out = call <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32 %a, i32 %b)
   ret <vscale x 4 x i32> %out
 }
 
 define <vscale x 2 x i64> @index_rr_i64(i64 %a, i64 %b) {
 ; CHECK-LABEL: index_rr_i64:
-; CHECK: index z0.d, x0, x1
-; CHECK-NEXT: ret
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.d, x0, x1
+; CHECK-NEXT:    ret
   %out = call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 %a, i64 %b)
   ret <vscale x 2 x i64> %out
 }
 
+define <vscale x 4 x i32> @index_rr_i32_combine(i32 %a, i32 %b) {
+; CHECK-LABEL: index_rr_i32_combine:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    index z0.s, w0, w1
+; CHECK-NEXT:    ret
+  %val = insertelement <vscale x 4 x i32> poison, i32 %a, i32 0
+  %val1 = shufflevector <vscale x 4 x i32> %val, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+  %val2 = insertelement <vscale x 4 x i32> poison, i32 %b, i32 0
+  %val3 = shufflevector <vscale x 4 x i32> %val2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+  %tmp = call <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32 0, i32 1)
+  %tmp1 = mul <vscale x 4 x i32> %tmp, %val3
+  %out = add <vscale x 4 x i32> %tmp1, %val1
+  ret <vscale x 4 x i32> %out
+}
+
+define <vscale x 4 x i32> @index_rr_i32_not_combine(i32 %a, i32 %b) {
+; CHECK-LABEL: index_rr_i32_not_combine:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.s, w0
+; CHECK-NEXT:    mov z1.s, w1
+; CHECK-NEXT:    index z2.s, #0, #1
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    mla z0.s, p0/m, z2.s, z1.s
+; CHECK-NEXT:    add z0.s, z0.s, z2.s
+; CHECK-NEXT:    ret
+  %val = insertelement <vscale x 4 x i32> poison, i32 %a, i32 0
+  %val1 = shufflevector <vscale x 4 x i32> %val, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+  %val2 = insertelement <vscale x 4 x i32> poison, i32 %b, i32 0
+  %val3 = shufflevector <vscale x 4 x i32> %val2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+  %tmp = call <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32 0, i32 1)
+  %tmp1 = mul <vscale x 4 x i32> %tmp, %val3
+  %tmp2 = add <vscale x 4 x i32> %tmp1, %val1
+  %out = add <vscale x 4 x i32> %tmp2, %tmp
+  ret <vscale x 4 x i32> %out
+}
+
 declare <vscale x 16 x i8> @llvm.aarch64.sve.index.nxv16i8(i8, i8)
 declare <vscale x 8 x i16> @llvm.aarch64.sve.index.nxv8i16(i16, i16)
 declare <vscale x 4 x i32> @llvm.aarch64.sve.index.nxv4i32(i32, i32)
 declare <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64, i64)