- shouldUseFormStridedPseudo sets RegClass using a switch statement

kmclaughlin-arm · kmclaughlin-arm · commit 645e30b34899 · 2024-12-10T16:51:30.000Z
- Removed INTRINSIC_W_CHAIN from getIntrinsicID
- Renamed FORM_STRIDED_TUPLE_X#_PSEUDO -&gt; FORM_TRANSPOSED_REG_TUPLE_X#_PSEUDO
- Add switch statement back into getRegAllocationHints
- Use getSubReg in getRegAllocationHints and remove index into RegUnits
diff --git a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
@@ -1754,9 +1754,9 @@ 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:
+   case AArch64::FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO:
      return expandFormTuplePseudo(MBB, MBBI, NextMBBI, 2);
-   case AArch64::FORM_STRIDED_TUPLE_X4_PSEUDO:
+   case AArch64::FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO:
      return expandFormTuplePseudo(MBB, MBBI, NextMBBI, 4);
   }
   return false;
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -7639,11 +7639,6 @@ 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;
-  }
   }
 }
 
@@ -8646,53 +8641,55 @@ static bool checkZExtBool(SDValue Arg, const SelectionDAG &DAG) {
   return ZExtBool;
 }
 
+// The FORM_TRANSPOSED_REG_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_TRANSPOSED_REG_TUPLE_X2_PSEUDO %5:zpr, %8:zpr
+//
 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");
+    assert(MO.isReg() && "Unexpected operand to FORM_TRANSPOSED_REG_TUPLE");
 
     MachineOperand *Def = MRI.getOneDef(MO.getReg());
-    if (!Def || !Def->isReg() || !Def->getParent()->isCopy()) {
-      UseFormStrided = false;
-      break;
-    }
+    if (!Def || !Def->getParent()->isCopy())
+      return false;
 
-    MachineOperand CpyOp = Def->getParent()->getOperand(1);
-    MachineOperand *Ld = MRI.getOneDef(CpyOp.getReg());
-    unsigned OpSubReg = CpyOp.getSubReg();
+    const MachineOperand &CpySrc = Def->getParent()->getOperand(1);
+    MachineOperand *CopySrcOp = MRI.getOneDef(CpySrc.getReg());
+    unsigned OpSubReg = CpySrc.getSubReg();
     if (SubReg == MCRegister::NoRegister)
       SubReg = OpSubReg;
-    if (!Ld || !Ld->isReg() || OpSubReg != SubReg) {
-      UseFormStrided = false;
+    if (!CopySrcOp || !CopySrcOp->isReg() || OpSubReg != SubReg)
+      return false;
+
+    const TargetRegisterClass *RegClass = nullptr;
+    switch (MI.getNumOperands() - 1) {
+    case 2:
+      RegClass = &AArch64::ZPR2StridedOrContiguousRegClass;
+      break;
+    case 4:
+      RegClass = &AArch64::ZPR4StridedOrContiguousRegClass;
       break;
+    default:
+      llvm_unreachable("Unexpected number of operands to pseudo.");
     }
 
-    const TargetRegisterClass *RegClass =
-        NumOperands == 2 ? &AArch64::ZPR2StridedOrContiguousRegClass
-                         : &AArch64::ZPR4StridedOrContiguousRegClass;
-
-    if (MRI.getRegClass(Ld->getReg()) == RegClass)
-      UseFormStrided = true;
+    if (MRI.getRegClass(CopySrcOp->getReg()) != RegClass)
+      return false;
   }
 
-  return UseFormStrided;
+  return true;
 }
 
 void AArch64TargetLowering::AdjustInstrPostInstrSelection(MachineInstr &MI,
@@ -8720,10 +8717,10 @@ 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 (MI.getOpcode() == AArch64::FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO ||
+      MI.getOpcode() == AArch64::FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO) {
+    // If input values to the FORM_TRANSPOSED_REG_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};
diff --git a/llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp b/llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp
@@ -1107,10 +1107,10 @@ 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
+// FORM_TRANSPOSED_REG_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.
 //
@@ -1123,51 +1123,38 @@ 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)
+    if (MI.getOpcode() != AArch64::FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO &&
+        MI.getOpcode() != AArch64::FORM_TRANSPOSED_REG_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))
+    switch (MI.getOperand(1).getSubReg()) {
+    case AArch64::zsub0:
+    case AArch64::zsub1:
+    case AArch64::zsub2:
+    case AArch64::zsub3:
+      break;
+    default:
       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)
+    // Look up the physical register mapped to the first operand of the pseudo.
+    Register FirstOpVirtReg = MI.getOperand(1).getReg();
+    if (!VRM->hasPhys(FirstOpVirtReg))
       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]);
-    }
+    MCRegister TupleStartReg =
+        getSubReg(VRM->getPhys(FirstOpVirtReg), MI.getOperand(1).getSubReg());
+    for (unsigned I = 0; I < Order.size(); ++I)
+      if (MCRegister R = getSubReg(Order[I], AArch64::zsub0))
+        if (R == TupleStartReg)
+          Hints.push_back(Order[I]);
   }
 
-  return DefaultHints;
+  return TargetRegisterInfo::getRegAllocationHints(VirtReg, Order, Hints, MF,
+                                                   VRM);
 }
 
 unsigned AArch64RegisterInfo::getLocalAddressRegister(
diff --git a/llvm/lib/Target/AArch64/SMEInstrFormats.td b/llvm/lib/Target/AArch64/SMEInstrFormats.td
@@ -34,14 +34,14 @@ 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 :
+def FORM_TRANSPOSED_REG_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 :
+def FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO :
   Pseudo<(outs ZPR4Mul4:$tup),
          (ins ZPR:$zn0, ZPR:$zn1, ZPR:$zn2, ZPR:$zn3), []>, Sched<[]>{
   let hasSideEffects = 0;
@@ -186,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,
-                                              (FORM_STRIDED_TUPLE_X2_PSEUDO vt:$Zn1,vt:$Zn2), zpr_ty:$Zm, imm_ty:$i)>;
+                                              (FORM_TRANSPOSED_REG_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,
-                                              (FORM_STRIDED_TUPLE_X4_PSEUDO vt:$Zn1, vt:$Zn2, vt:$Zn3, vt:$Zn4),
+                                              (FORM_TRANSPOSED_REG_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>
