- Removed the FORM_STRIDED_TUPLE_X# nodes, leaving only the pseudo nodes.

- Changed the tablegen patterns used by the dot intrinsics to always output
  the FORM_STRIDED_TUPLE_X#_PSEUDO nodes.
- Check that the operands to the pseudo are copies from a StridedOrContiguous
  register class in AdjustInstrPostInstrSelection, falling back on creating
  a REG_SEQUENCE node if not.

Author: kmclaughlin-arm

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

@@ -1159,6 +1159,8 @@ bool AArch64ExpandPseudo::expandFormTuplePseudo(
     Register FormTupleOpReg = MI.getOperand(i + 1).getReg();
     Register ReturnTupleSubReg =
         TRI->getSubReg(ReturnTuple, AArch64::zsub0 + i);
+    // Add copies to ensure the subregisters remain in the correct order
+    // for any contigious operation they are used by.
     if (FormTupleOpReg != ReturnTupleSubReg)
       BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::ORR_ZZZ))
           .addReg(ReturnTupleSubReg, RegState::Define)

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

@@ -505,7 +505,7 @@ class AArch64DAGToDAGISel : public SelectionDAGISel {
   bool SelectAllActivePredicate(SDValue N);
   bool SelectAnyPredicate(SDValue N);
 
-  void SelectFormTuplePseudo(SDNode *N, unsigned Size);
+  void SelectFormTuplePseudo(SDNode *N);
 };
 
 class AArch64DAGToDAGISelLegacy : public SelectionDAGISelLegacy {
@@ -7183,14 +7183,6 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
     }
     break;
   }
-  case AArch64ISD::FORM_STRIDED_TUPLE_X2: {
-    SelectFormTuplePseudo(Node, 2);
-    return;
-  }
-  case AArch64ISD::FORM_STRIDED_TUPLE_X4: {
-    SelectFormTuplePseudo(Node, 4);
-    return;
-  }
   }
 
   // Select the default instruction
@@ -7448,20 +7440,3 @@ bool AArch64DAGToDAGISel::SelectSMETileSlice(SDValue N, unsigned MaxSize,
   Offset = CurDAG->getTargetConstant(0, SDLoc(N), MVT::i64);
   return true;
 }
-
-void AArch64DAGToDAGISel::SelectFormTuplePseudo(SDNode *Node, unsigned Size) {
-  assert((Size == 2 || Size == 4) && "Invalid Tuple size");
-  EVT VT = Node->getValueType(0);
-  SmallVector<SDValue> Ops;
-  for (unsigned I = 0; I < Size; I++)
-    Ops.push_back(Node->getOperand(I));
-  SDLoc DL(Node);
-  unsigned Opc = Size == 2 ? AArch64::FORM_STRIDED_TUPLE_X2_PSEUDO
-                           : AArch64::FORM_STRIDED_TUPLE_X4_PSEUDO;
-  SDNode *Tuple = CurDAG->getMachineNode(Opc, DL, MVT::Untyped, Ops);
-  SDValue SuperReg = SDValue(Tuple, 0);
-  for (unsigned I = 0; I < Size; ++I)
-    ReplaceUses(SDValue(Node, I), CurDAG->getTargetExtractSubreg(
-                                      AArch64::zsub0 + I, DL, VT, SuperReg));
-  CurDAG->RemoveDeadNode(Node);
-}

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -5711,46 +5711,6 @@ SDValue AArch64TargetLowering::getRuntimePStateSM(SelectionDAG &DAG,
                      Mask);
 }
 
-static unsigned getIntrinsicID(const SDNode *N);
-
-SDValue TryLowerMultiVecSMEDotIntrinsic(SDValue Op, SelectionDAG &DAG,
-                                        unsigned Size) {
-  assert((Size == 2 || Size == 4) && "Invalid Tuple Size");
-  auto IsStridedLoad = [Size](SDValue Op) -> bool {
-    unsigned Intrinsic = getIntrinsicID(Op.getNode());
-    if (Size == 2)
-      return Intrinsic == Intrinsic::aarch64_sve_ld1_pn_x2;
-    else
-      return Intrinsic == Intrinsic::aarch64_sve_ld1_pn_x4;
-  };
-
-  SmallVector<SDValue> Ops;
-  unsigned LastLoadIdx = Size == 2 ? 5 : 7;
-  unsigned LoadResNo = Op.getOperand(3).getResNo();
-  for (unsigned I = 3; I < LastLoadIdx; I++) {
-    if (!IsStridedLoad(Op->getOperand(I)) ||
-        Op.getOperand(I).getResNo() != LoadResNo)
-      return SDValue();
-    Ops.push_back(Op->getOperand(I));
-  }
-
-  EVT VT = Op->getOperand(3).getValueType();
-  SDVTList VTList =
-      Size == 2 ? DAG.getVTList(VT, VT) : DAG.getVTList(VT, VT, VT, VT);
-  unsigned Opc = Size == 2 ? AArch64ISD::FORM_STRIDED_TUPLE_X2
-                           : AArch64ISD::FORM_STRIDED_TUPLE_X4;
-  SDLoc DL(Op);
-  SDValue Pseudo = DAG.getNode(Opc, DL, VTList, Ops);
-
-  SmallVector<SDValue> DotOps = {Op.getOperand(0), Op->getOperand(1),
-                                 Op->getOperand(2)};
-  for (unsigned I = 0; I < Size; I++)
-    DotOps.push_back(Pseudo.getValue(I));
-  DotOps.push_back(Op->getOperand(DotOps.size()));
-  DotOps.push_back(Op->getOperand(DotOps.size()));
-  return DAG.getNode(Op->getOpcode(), DL, MVT::Other, DotOps);
-}
-
 // Lower an SME LDR/STR ZA intrinsic
 // Case 1: If the vector number (vecnum) is an immediate in range, it gets
 // folded into the instruction
@@ -5940,22 +5900,6 @@ SDValue AArch64TargetLowering::LowerINTRINSIC_VOID(SDValue Op,
         Op->getOperand(0), // Chain
         DAG.getTargetConstant((int32_t)(AArch64SVCR::SVCRZA), DL, MVT::i32),
         DAG.getConstant(AArch64SME::Always, DL, MVT::i64));
-  case Intrinsic::aarch64_sme_uvdot_lane_za32_vg1x4:
-  case Intrinsic::aarch64_sme_suvdot_lane_za32_vg1x4:
-  case Intrinsic::aarch64_sme_usvdot_lane_za32_vg1x4:
-  case Intrinsic::aarch64_sme_svdot_lane_za32_vg1x4:
-  case Intrinsic::aarch64_sme_usdot_lane_za32_vg1x4:
-  case Intrinsic::aarch64_sme_udot_lane_za32_vg1x4:
-  case Intrinsic::aarch64_sme_sudot_lane_za32_vg1x4:
-  case Intrinsic::aarch64_sme_sdot_lane_za32_vg1x4:
-    return TryLowerMultiVecSMEDotIntrinsic(Op, DAG, 4);
-  case Intrinsic::aarch64_sme_uvdot_lane_za32_vg1x2:
-  case Intrinsic::aarch64_sme_sdot_lane_za32_vg1x2:
-  case Intrinsic::aarch64_sme_svdot_lane_za32_vg1x2:
-  case Intrinsic::aarch64_sme_usdot_lane_za32_vg1x2:
-  case Intrinsic::aarch64_sme_sudot_lane_za32_vg1x2:
-  case Intrinsic::aarch64_sme_udot_lane_za32_vg1x2:
-    return TryLowerMultiVecSMEDotIntrinsic(Op, DAG, 2);
   }
 }
 
@@ -8729,6 +8673,77 @@ void AArch64TargetLowering::AdjustInstrPostInstrSelection(MachineInstr &MI,
     }
   }
 
+  if (MI.getOpcode() == AArch64::FORM_STRIDED_TUPLE_X2_PSEUDO ||
+      MI.getOpcode() == AArch64::FORM_STRIDED_TUPLE_X4_PSEUDO) {
+    MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
+    bool UseFormStrided = false;
+    unsigned Size =
+        MI.getOpcode() == AArch64::FORM_STRIDED_TUPLE_X2_PSEUDO ? 2 : 4;
+
+    // The FORM_STRIDED_TUPLE pseudo should only be used if the input operands
+    // are copy nodes where the source register is in a StridedOrContiguous
+    // class. For example:
+    //   %3:zpr2stridedorcontiguous = LD1B_2Z_IMM_PSEUDO ..
+    //   %4:zpr = COPY %3.zsub1:zpr2stridedorcontiguous
+    //   %5:zpr = COPY %3.zsub0:zpr2stridedorcontiguous
+    //   %6:zpr2stridedorcontiguous = LD1B_2Z_PSEUDO ..
+    //   %7:zpr = COPY %6.zsub1:zpr2stridedorcontiguous
+    //   %8:zpr = COPY %6.zsub0:zpr2stridedorcontiguous
+    //   %9:zpr2mul2 = FORM_STRIDED_TUPLE_X2_PSEUDO %5:zpr, %8:zpr
+
+    SmallVector<unsigned, 4> OpSubRegs;
+    for (unsigned I = 1; I < MI.getNumOperands(); ++I) {
+      MachineOperand &MO = MI.getOperand(I);
+      if (!MO.isReg())
+        continue;
+
+      MachineOperand *Def = MRI.getOneDef(MO.getReg());
+      if (!Def || !Def->isReg() || !Def->getParent()->isCopy())
+        continue;
+
+      MachineInstr *Cpy = Def->getParent();
+      MachineOperand CpyOp = Cpy->getOperand(1);
+      if (!CpyOp.isReg())
+        continue;
+
+      MachineOperand *Ld = MRI.getOneDef(CpyOp.getReg());
+      OpSubRegs.push_back(CpyOp.getSubReg());
+      if (!Ld || !Ld->isReg())
+        continue;
+
+      const TargetRegisterClass *RegClass =
+          Size == 2 ? &AArch64::ZPR2StridedOrContiguousRegClass
+                    : &AArch64::ZPR4StridedOrContiguousRegClass;
+
+      if (MRI.getRegClass(Ld->getReg()) == RegClass)
+        UseFormStrided = true;
+    }
+
+    // Ensure the operands all use the same subreg index.
+    if (!std::equal(OpSubRegs.begin(), OpSubRegs.end(), OpSubRegs.begin()))
+      UseFormStrided = false;
+
+    // If input values to the FORM_STRIDED_TUPLE pseudo aren't copies from a
+    // StridedOrContiguous class, fall back on REG_SEQUENCE node.
+    if (!UseFormStrided) {
+      static const unsigned SubRegs[] = {AArch64::zsub0, AArch64::zsub1,
+                                         AArch64::zsub2, AArch64::zsub3};
+
+      const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+      MachineInstrBuilder MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+                                        TII->get(TargetOpcode::REG_SEQUENCE),
+                                        MI.getOperand(0).getReg());
+
+      for (unsigned I = 1; I < MI.getNumOperands(); ++I) {
+        MIB.add(MI.getOperand(I));
+        MIB.addImm(SubRegs[I - 1]);
+      }
+
+      MI.eraseFromParent();
+    }
+    return;
+  }
+
   // Add an implicit use of 'VG' for ADDXri/SUBXri, which are instructions that
   // have nothing to do with VG, were it not that they are used to materialise a
   // frame-address. If they contain a frame-index to a scalable vector, this

llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td

@@ -29,7 +29,7 @@ def AArch64_save_zt : SDNode<"AArch64ISD::SAVE_ZT", SDTypeProfile<0, 2,
                              [SDTCisInt<0>, SDTCisPtrTy<1>]>,
                              [SDNPHasChain, SDNPSideEffect, SDNPMayStore]>;
 
-def SDT_FORM_STRIDED_TUPLE_X2 : SDTypeProfile<4, 4,
+def SDT_FORM_STRIDED_TUPLE_X2 : SDTypeProfile<2, 2,
                              [SDTCisVec<0>, SDTCisSameAs<0, 1>,
                               SDTCisSameAs<0, 2>, SDTCisSameAs<0, 3>]>;
 

llvm/lib/Target/AArch64/SMEInstrFormats.td

@@ -38,12 +38,14 @@ def FORM_STRIDED_TUPLE_X2_PSEUDO :
   Pseudo<(outs ZPR2Mul2:$tup),
          (ins ZPR:$zn0, ZPR:$zn1), []>, Sched<[]>{
   let hasSideEffects = 0;
+  let hasPostISelHook = 1;
 }
 
 def FORM_STRIDED_TUPLE_X4_PSEUDO :
   Pseudo<(outs ZPR4Mul4:$tup),
          (ins ZPR:$zn0, ZPR:$zn1, ZPR:$zn2, ZPR:$zn3), []>, Sched<[]>{
   let hasSideEffects = 0;
+  let hasPostISelHook = 1;
 }
 
 def SDTZALoadStore : SDTypeProfile<0, 3, [SDTCisInt<0>, SDTCisPtrTy<1>, SDTCisInt<2>]>;
@@ -186,6 +188,12 @@ class SME2_ZA_TwoOp_VG2_Multi_Index_Pat<string name, SDPatternOperator intrinsic
           (!cast<Instruction>(name # _PSEUDO) $base, $offset,
                                               (REG_SEQUENCE ZPR2Mul2, vt:$Zn1, zsub0, vt:$Zn2, zsub1), zpr_ty:$Zm, imm_ty:$i)>;
 
+class SME2_ZA_TwoOp_VG2_Multi_Index_FormStrided_Pat<string name, SDPatternOperator intrinsic, Operand index_ty, ZPRRegOp zpr_ty, ValueType vt,
+                                        Operand imm_ty, ComplexPattern tileslice>
+    : Pat<(intrinsic (i32 (tileslice MatrixIndexGPR32Op8_11:$base, index_ty:$offset)), vt:$Zn1, vt:$Zn2, vt:$Zm, (i32 imm_ty:$i)),
+          (!cast<Instruction>(name # _PSEUDO) $base, $offset,
+                                              (FORM_STRIDED_TUPLE_X2_PSEUDO vt:$Zn1,vt:$Zn2), zpr_ty:$Zm, imm_ty:$i)>;
+
 class SME2_ZA_TwoOp_VG4_Multi_Index_Pat<string name, SDPatternOperator intrinsic, Operand index_ty, ZPRRegOp zpr_ty, ValueType vt,
                                         Operand imm_ty, ComplexPattern tileslice>
     : Pat<(intrinsic (i32 (tileslice MatrixIndexGPR32Op8_11:$base, index_ty:$offset)),
@@ -194,6 +202,14 @@ class SME2_ZA_TwoOp_VG4_Multi_Index_Pat<string name, SDPatternOperator intrinsic
                                               (REG_SEQUENCE ZPR4Mul4, vt:$Zn1, zsub0, vt:$Zn2, zsub1, vt:$Zn3, zsub2, vt:$Zn4, zsub3),
                                               zpr_ty:$Zm, imm_ty:$i)>;
 
+class SME2_ZA_TwoOp_VG4_Multi_Index_FormStrided_Pat<string name, SDPatternOperator intrinsic, Operand index_ty, ZPRRegOp zpr_ty, ValueType vt,
+                                        Operand imm_ty, ComplexPattern tileslice>
+    : Pat<(intrinsic (i32 (tileslice MatrixIndexGPR32Op8_11:$base, index_ty:$offset)),
+                     vt:$Zn1, vt:$Zn2, vt:$Zn3, vt:$Zn4, vt:$Zm, (i32 imm_ty:$i)),
+          (!cast<Instruction>(name # _PSEUDO) $base, $offset,
+                                              (FORM_STRIDED_TUPLE_X4_PSEUDO vt:$Zn1, vt:$Zn2, vt:$Zn3, vt:$Zn4),
+                                              zpr_ty:$Zm, imm_ty:$i)>;
+
 class SME2_Sat_Shift_VG2_Pat<string name, SDPatternOperator intrinsic, ValueType out_vt, ValueType in_vt, Operand imm_ty>
     : Pat<(out_vt (intrinsic in_vt:$Zn1, in_vt:$Zn2, (i32 imm_ty:$i))),
                   (!cast<Instruction>(name) (REG_SEQUENCE ZPR2Mul2, in_vt:$Zn1, zsub0, in_vt:$Zn2, zsub1), imm_ty:$i)>;
@@ -2635,7 +2651,7 @@ multiclass sme2_multi_vec_array_vg2_index_32b<string mnemonic, bits<2> sz, bits<
   }
   def _PSEUDO : sme2_za_array_2op_multi_index_pseudo<NAME, sme_elm_idx0_7, multi_vector_ty, vector_ty, VectorIndexS32b_timm, SMEMatrixArray>;
 
-  def : SME2_ZA_TwoOp_VG2_Multi_Index_Pat<NAME, intrinsic, sme_elm_idx0_7, vector_ty, vt, VectorIndexS32b_timm, tileslice16>;
+  def : SME2_ZA_TwoOp_VG2_Multi_Index_FormStrided_Pat<NAME, intrinsic, sme_elm_idx0_7, vector_ty, vt, VectorIndexS32b_timm, tileslice16>;
 
   def : InstAlias<mnemonic # "\t$ZAda[$Rv, $imm3], $Zn, $Zm$i",
         (!cast<Instruction>(NAME) MatrixOp32:$ZAda,  MatrixIndexGPR32Op8_11:$Rv, sme_elm_idx0_7:$imm3,
@@ -2778,7 +2794,7 @@ multiclass sme2_multi_vec_array_vg4_index_32b<string mnemonic, bits<4> op,
 
   def _PSEUDO : sme2_za_array_2op_multi_index_pseudo<NAME, sme_elm_idx0_7, multi_vector_ty, vector_ty, VectorIndexS32b_timm, SMEMatrixArray>;
 
-  def : SME2_ZA_TwoOp_VG4_Multi_Index_Pat<NAME, intrinsic, sme_elm_idx0_7, vector_ty, vt, VectorIndexS32b_timm, tileslice16>;
+  def : SME2_ZA_TwoOp_VG4_Multi_Index_FormStrided_Pat<NAME, intrinsic, sme_elm_idx0_7, vector_ty, vt, VectorIndexS32b_timm, tileslice16>;
 
   def : InstAlias<mnemonic # "\t$ZAda[$Rv, $imm3], $Zn, $Zm$i",
         (!cast<Instruction>(NAME) MatrixOp32:$ZAda,  MatrixIndexGPR32Op8_11:$Rv, sme_elm_idx0_7:$imm3,