[AArch64][SME2] Improve register allocation of multi-vector SME intrinsics

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

@@ -67,6 +67,10 @@ class AArch64ExpandPseudo : public MachineFunctionPass {
                             TargetRegisterClass ContiguousClass,
                             TargetRegisterClass StridedClass,
                             unsigned ContiguousOpc, unsigned StridedOpc);
+  bool expandFormTuplePseudo(MachineBasicBlock &MBB,
+                             MachineBasicBlock::iterator MBBI,
+                             MachineBasicBlock::iterator &NextMBBI,
+                             unsigned Size);
   bool expandMOVImm(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
                     unsigned BitSize);
 
@@ -1142,6 +1146,32 @@ bool AArch64ExpandPseudo::expandMultiVecPseudo(
   return true;
 }
 
+bool AArch64ExpandPseudo::expandFormTuplePseudo(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
+    MachineBasicBlock::iterator &NextMBBI, unsigned Size) {
+  assert(Size == 2 || Size == 4 && "Invalid Tuple Size");
+  MachineInstr &MI = *MBBI;
+  Register ReturnTuple = MI.getOperand(0).getReg();
+
+  const TargetRegisterInfo *TRI =
+      MBB.getParent()->getSubtarget().getRegisterInfo();
+  for (unsigned i = 0; i < Size; i++) {
+    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)
+          .addReg(FormTupleOpReg)
+          .addReg(FormTupleOpReg);
+  }
+
+  MI.eraseFromParent();
+  return true;
+}
+
 /// If MBBI references a pseudo instruction that should be expanded here,
 /// do the expansion and return true.  Otherwise return false.
 bool AArch64ExpandPseudo::expandMI(MachineBasicBlock &MBB,
@@ -1724,6 +1754,10 @@ bool AArch64ExpandPseudo::expandMI(MachineBasicBlock &MBB,
      return expandMultiVecPseudo(
          MBB, MBBI, AArch64::ZPR4RegClass, AArch64::ZPR4StridedRegClass,
          AArch64::LDNT1D_4Z, AArch64::LDNT1D_4Z_STRIDED);
+   case AArch64::FORM_STRIDED_TUPLE_X2_PSEUDO:
+     return expandFormTuplePseudo(MBB, MBBI, NextMBBI, 2);
+   case AArch64::FORM_STRIDED_TUPLE_X4_PSEUDO:
+     return expandFormTuplePseudo(MBB, MBBI, NextMBBI, 4);
   }
   return false;
 }

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -7639,6 +7639,11 @@ static unsigned getIntrinsicID(const SDNode *N) {
       return IID;
     return Intrinsic::not_intrinsic;
   }
+  case ISD::INTRINSIC_W_CHAIN: {
+    unsigned IID = N->getConstantOperandVal(1);
+    if (IID < Intrinsic::num_intrinsics)
+      return IID;
+  }
   }
 }
 
@@ -8641,6 +8646,55 @@ static bool checkZExtBool(SDValue Arg, const SelectionDAG &DAG) {
   return ZExtBool;
 }
 
+bool shouldUseFormStridedPseudo(MachineInstr &MI) {
+  MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
+  bool UseFormStrided = false;
+  unsigned NumOperands =
+      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
+
+  MCRegister SubReg = MCRegister::NoRegister;
+  for (unsigned I = 1; I < MI.getNumOperands(); ++I) {
+    MachineOperand &MO = MI.getOperand(I);
+    assert(MO.isReg() && "Unexpected operand to FORM_STRIDED_TUPLE");
+
+    MachineOperand *Def = MRI.getOneDef(MO.getReg());
+    if (!Def || !Def->isReg() || !Def->getParent()->isCopy()) {
+      UseFormStrided = false;
+      break;
+    }
+
+    MachineOperand CpyOp = Def->getParent()->getOperand(1);
+    MachineOperand *Ld = MRI.getOneDef(CpyOp.getReg());
+    unsigned OpSubReg = CpyOp.getSubReg();
+    if (SubReg == MCRegister::NoRegister)
+      SubReg = OpSubReg;
+    if (!Ld || !Ld->isReg() || OpSubReg != SubReg) {
+      UseFormStrided = false;
+      break;
+    }
+
+    const TargetRegisterClass *RegClass =
+        NumOperands == 2 ? &AArch64::ZPR2StridedOrContiguousRegClass
+                         : &AArch64::ZPR4StridedOrContiguousRegClass;
+
+    if (MRI.getRegClass(Ld->getReg()) == RegClass)
+      UseFormStrided = true;
+  }
+
+  return UseFormStrided;
+}
+
 void AArch64TargetLowering::AdjustInstrPostInstrSelection(MachineInstr &MI,
                                                           SDNode *Node) const {
   // Live-in physreg copies that are glued to SMSTART are applied as
@@ -8666,6 +8720,29 @@ void AArch64TargetLowering::AdjustInstrPostInstrSelection(MachineInstr &MI,
     }
   }
 
+  if (MI.getOpcode() == AArch64::FORM_STRIDED_TUPLE_X2_PSEUDO ||
+      MI.getOpcode() == AArch64::FORM_STRIDED_TUPLE_X4_PSEUDO) {
+    // If input values to the FORM_STRIDED_TUPLE pseudo aren't copies from a
+    // StridedOrContiguous class, fall back on REG_SEQUENCE node.
+    if (!shouldUseFormStridedPseudo(MI)) {
+      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/AArch64RegisterInfo.cpp

@@ -1107,6 +1107,69 @@ unsigned AArch64RegisterInfo::getRegPressureLimit(const TargetRegisterClass *RC,
   }
 }
 
+// FORM_STRIDED_TUPLE nodes are created to improve register allocation where
+// a consecutive multi-vector tuple is constructed from the same indices of
+// multiple strided loads. This may still result in unnecessary copies between
+// the loads and the tuple. Here we try to return a hint to assign the
+// contiguous ZPRMulReg starting at the same register as the first operand of
+// the pseudo, which should be a subregister of the first strided load.
+//
+// For example, if the first strided load has been assigned $z16_z20_z24_z28
+// and the operands of the pseudo are each accessing subregister zsub2, we
+// should look through through Order to find a contiguous register which
+// begins with $z24 (i.e. $z24_z25_z26_z27).
+//
+bool AArch64RegisterInfo::getRegAllocationHints(
+    Register VirtReg, ArrayRef<MCPhysReg> Order,
+    SmallVectorImpl<MCPhysReg> &Hints, const MachineFunction &MF,
+    const VirtRegMap *VRM, const LiveRegMatrix *Matrix) const {
+  const AArch64Subtarget &STI = MF.getSubtarget<AArch64Subtarget>();
+  const TargetRegisterInfo *TRI = STI.getRegisterInfo();
+  const MachineRegisterInfo &MRI = MF.getRegInfo();
+  bool DefaultHints =
+      TargetRegisterInfo::getRegAllocationHints(VirtReg, Order, Hints, MF, VRM);
+
+  unsigned RegID = MRI.getRegClass(VirtReg)->getID();
+  if (RegID != AArch64::ZPR2Mul2RegClassID &&
+      RegID != AArch64::ZPR4Mul4RegClassID)
+    return DefaultHints;
+
+  for (MachineInstr &MI : MRI.def_instructions(VirtReg)) {
+    if (MI.getOpcode() != AArch64::FORM_STRIDED_TUPLE_X2_PSEUDO &&
+        MI.getOpcode() != AArch64::FORM_STRIDED_TUPLE_X4_PSEUDO)
+      continue;
+
+    // Look up the physical register mapped to the first load of the pseudo.
+    Register FirstLoadVirtReg = MI.getOperand(1).getReg();
+    if (!VRM->hasPhys(FirstLoadVirtReg))
+      continue;
+
+    int64_t SubRegIdx = -1;
+    MCRegister FirstLoadPhysReg = VRM->getPhys(FirstLoadVirtReg);
+
+    // The subreg number is used to access the correct unit of the
+    // strided register found in the map above.
+    SubRegIdx = MI.getOperand(1).getSubReg() - AArch64::zsub0;
+    if (SubRegIdx < 0 || SubRegIdx > 3)
+      continue;
+
+    SmallVector<Register, 4> RegUnits;
+    for (MCRegUnit Unit : TRI->regunits(FirstLoadPhysReg))
+      RegUnits.push_back(Unit);
+
+    // Find the contiguous ZPRMul register which starts with the
+    // same register unit as the strided register and add to Hints.
+    Register StartReg = RegUnits[SubRegIdx];
+    for (unsigned I = 0; I < Order.size(); ++I) {
+      Register Reg = *TRI->regunits(Order[I]).begin();
+      if (Reg == StartReg)
+        Hints.push_back(Order[I]);
+    }
+  }
+
+  return DefaultHints;
+}
+
 unsigned AArch64RegisterInfo::getLocalAddressRegister(
   const MachineFunction &MF) const {
   const auto &MFI = MF.getFrameInfo();

llvm/lib/Target/AArch64/AArch64RegisterInfo.h

@@ -134,6 +134,11 @@ class AArch64RegisterInfo final : public AArch64GenRegisterInfo {
   unsigned getRegPressureLimit(const TargetRegisterClass *RC,
                                MachineFunction &MF) const override;
 
+  bool getRegAllocationHints(Register VirtReg, ArrayRef<MCPhysReg> Order,
+                             SmallVectorImpl<MCPhysReg> &Hints,
+                             const MachineFunction &MF, const VirtRegMap *VRM,
+                             const LiveRegMatrix *Matrix) const override;
+
   unsigned getLocalAddressRegister(const MachineFunction &MF) const;
   bool regNeedsCFI(unsigned Reg, unsigned &RegToUseForCFI) const;
 

llvm/lib/Target/AArch64/SMEInstrFormats.td

@@ -34,6 +34,20 @@ def tileslicerange0s4 : ComplexPattern<i32, 2, "SelectSMETileSlice<0,  4>", []>;
 
 def am_sme_indexed_b4 :ComplexPattern<iPTR, 2, "SelectAddrModeIndexedSVE<0,15>", [], [SDNPWantRoot]>;
 
+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>]>;
 def AArch64SMELdr : SDNode<"AArch64ISD::SME_ZA_LDR", SDTZALoadStore,
                              [SDNPHasChain, SDNPSideEffect, SDNPMayLoad]>;
@@ -172,14 +186,14 @@ class SME2_ZA_TwoOp_VG2_Multi_Index_Pat<string name, SDPatternOperator intrinsic
                                         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,
-                                              (REG_SEQUENCE ZPR2Mul2, vt:$Zn1, zsub0, vt:$Zn2, zsub1), zpr_ty:$Zm, imm_ty:$i)>;
+                                              (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)),
                      vt:$Zn1, vt:$Zn2, vt:$Zn3, vt:$Zn4, vt:$Zm, (i32 imm_ty:$i)),
           (!cast<Instruction>(name # _PSEUDO) $base, $offset,
-                                              (REG_SEQUENCE ZPR4Mul4, vt:$Zn1, zsub0, vt:$Zn2, zsub1, vt:$Zn3, zsub2, vt:$Zn4, zsub3),
+                                              (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>