- Renamed CpySrc -> CopySrc

- Moved switch in shouldUseFormStridedPseudo outside of loop
- Added a description of the pseudo nodes to SMEInstrFormats.td

Author: kmclaughlin-arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -8656,6 +8656,18 @@ static bool checkZExtBool(SDValue Arg, const SelectionDAG &DAG) {
 bool shouldUseFormStridedPseudo(MachineInstr &MI) {
   MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
 
+  const TargetRegisterClass *RegClass = nullptr;
+  switch (MI.getOpcode()) {
+  case AArch64::FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO:
+    RegClass = &AArch64::ZPR2StridedOrContiguousRegClass;
+    break;
+  case AArch64::FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO:
+    RegClass = &AArch64::ZPR4StridedOrContiguousRegClass;
+    break;
+  default:
+    llvm_unreachable("Unexpected opcode.");
+  }
+
   MCRegister SubReg = MCRegister::NoRegister;
   for (unsigned I = 1; I < MI.getNumOperands(); ++I) {
     MachineOperand &MO = MI.getOperand(I);
@@ -8665,26 +8677,15 @@ bool shouldUseFormStridedPseudo(MachineInstr &MI) {
     if (!Def || !Def->getParent()->isCopy())
       return false;
 
-    const MachineOperand &CpySrc = Def->getParent()->getOperand(1);
-    MachineOperand *CopySrcOp = MRI.getOneDef(CpySrc.getReg());
-    unsigned OpSubReg = CpySrc.getSubReg();
+    const MachineOperand &CopySrc = Def->getParent()->getOperand(1);
+    unsigned OpSubReg = CopySrc.getSubReg();
     if (SubReg == MCRegister::NoRegister)
       SubReg = OpSubReg;
+
+    MachineOperand *CopySrcOp = MRI.getOneDef(CopySrc.getReg());
     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.");
-    }
-
     if (MRI.getRegClass(CopySrcOp->getReg()) != RegClass)
       return false;
   }
@@ -8721,22 +8722,23 @@ void AArch64TargetLowering::AdjustInstrPostInstrSelection(MachineInstr &MI,
       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};
-
-      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]);
-      }
+    if (shouldUseFormStridedPseudo(MI))
+      return;
 
-      MI.eraseFromParent();
+    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;
   }
 

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

@@ -1130,7 +1130,8 @@ bool AArch64RegisterInfo::getRegAllocationHints(
         MI.getOpcode() != AArch64::FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO)
       continue;
 
-    switch (MI.getOperand(1).getSubReg()) {
+    unsigned FirstOpSubReg = MI.getOperand(1).getSubReg();
+    switch (FirstOpSubReg) {
     case AArch64::zsub0:
     case AArch64::zsub1:
     case AArch64::zsub2:
@@ -1146,7 +1147,7 @@ bool AArch64RegisterInfo::getRegAllocationHints(
       continue;
 
     MCRegister TupleStartReg =
-        getSubReg(VRM->getPhys(FirstOpVirtReg), MI.getOperand(1).getSubReg());
+        getSubReg(VRM->getPhys(FirstOpVirtReg), FirstOpSubReg);
     for (unsigned I = 0; I < Order.size(); ++I)
       if (MCRegister R = getSubReg(Order[I], AArch64::zsub0))
         if (R == TupleStartReg)

llvm/lib/Target/AArch64/SMEInstrFormats.td

@@ -34,6 +34,16 @@ def tileslicerange0s4 : ComplexPattern<i32, 2, "SelectSMETileSlice<0,  4>", []>;
 
 def am_sme_indexed_b4 :ComplexPattern<iPTR, 2, "SelectAddrModeIndexedSVE<0,15>", [], [SDNPWantRoot]>;
 
+// The FORM_TRANSPOSED_REG_TUPLE pseudos defined below are intended to
+// improve register allocation for intrinsics which use strided and contiguous
+// multi-vector registers, avoiding unnecessary copies.
+// If the operands of the pseudo are copies where the source register is in
+// the StridedOrContiguous class, the pseudo is used to provide a hint to the
+// register allocator suggesting a contigious multi-vector register which
+// matches the subregister sequence used by the operands.
+// If the operands do not match this pattern, the pseudos are expanded
+// to a REG_SEQUENCE using the post-isel hook.
+
 def FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO :
   Pseudo<(outs ZPR2Mul2:$tup),
          (ins ZPR:$zn0, ZPR:$zn1), []>, Sched<[]>{