- Renamed LdOps to UseOps

- Rewrote the case where no registers are already assigned to avoid creating
  the Regs vector.
- Added more comments with an example

Author: kmclaughlin-arm

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

@@ -1123,11 +1123,18 @@ bool AArch64RegisterInfo::getRegAllocationHints(
           Use.getOpcode() != AArch64::FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO)
         continue;
 
-      unsigned LdOps = Use.getNumOperands() - 1;
-      const TargetRegisterClass *StridedRC =
-          RegID == AArch64::ZPR2StridedOrContiguousRegClassID
-              ? &AArch64::ZPR2StridedRegClass
-              : &AArch64::ZPR4StridedRegClass;
+      unsigned UseOps = Use.getNumOperands() - 1;
+      const TargetRegisterClass *StridedRC;
+      switch (RegID) {
+      case AArch64::ZPR2StridedOrContiguousRegClassID:
+        StridedRC = &AArch64::ZPR2StridedRegClass;
+        break;
+      case AArch64::ZPR4StridedOrContiguousRegClassID:
+        StridedRC = &AArch64::ZPR4StridedRegClass;
+        break;
+      default:
+        llvm_unreachable("Unexpected RegID");
+      }
 
       SmallVector<MCPhysReg, 4> StridedOrder;
       for (MCPhysReg Reg : Order)
@@ -1147,32 +1154,67 @@ bool AArch64RegisterInfo::getRegAllocationHints(
             return VRM->hasPhys(Op.getReg());
           });
 
+      // Example:
+      //
+      // When trying to find a suitable register allocation for VirtReg %v2 in:
+      //
+      //  %v0:zpr2stridedorcontiguous = ld1 p0/z, [...]
+      //  %v1:zpr2stridedorcontiguous = ld1 p0/z, [...]
+      //  %v2:zpr2stridedorcontiguous = ld1 p0/z, [...]
+      //  %v3:zpr2stridedorcontiguous = ld1 p0/z, [...]
+      //  %v4:zpr4mul4 = FORM_TRANSPOSED_X4 %v0:0, %v1:0, %v2:0, %v3:0
+      //
+      // One such suitable allocation would be:
+      //
+      //  { z0, z8 }  = ld1 p0/z, [...]
+      //  { z1, z9 }  = ld1 p0/z, [...]
+      //  { z2, z10 } = ld1 p0/z, [...]
+      //  { z3, z11 } = ld1 p0/z, [...]
+      //  { z0, z1, z2, z3 } =
+      //     FORM_TRANSPOSED_X4 {z0, z8}:0, {z1, z9}:0, {z2, z10}:0, {z3, z11}:0
+      //
+      // Below we distinguish two cases when trying to find a register:
+      // * None of the registers used by FORM_TRANSPOSED_X4 have been assigned
+      //   yet. In this case the code muse ensure that there are at least UseOps
+      //   free consecutive registers. If IsMulZPR is true, then the first of
+      //   registers must also be a multiple of UseOps, e.g. { z0, z1, z2, z3 }
+      //   is valid but { z1, z2, z3, z5 } is not.
+      // * One or more of the registers used by FORM_TRANSPOSED_X4 is already
+      //   assigned a physical register, which means only checking that a
+      //   consectutive range of free tuple registers exists which includes
+      //   the assigned register.
+      //   e.g. in the example above, if { z0, z8 } is already allocated for
+      //   %v0, we just need to ensure that { z1, z9 }, { z2, z10 } and
+      //   { z3, z11 } are also free. If so, we add { z2, z10 }.
+
       if (AssignedRegOp == Use.operands_end()) {
         // There are no registers already assigned to any of the pseudo
         // operands. Look for a valid starting register for the group.
         for (unsigned I = 0; I < StridedOrder.size(); ++I) {
           MCPhysReg Reg = StridedOrder[I];
           SmallVector<MCPhysReg> Regs;
-          unsigned FirstStridedReg = Reg - OpIdx + 1;
 
           // If the FORM_TRANSPOSE nodes use the ZPRMul classes, the starting
           // register of the first load should be a multiple of 2 or 4.
-          unsigned FirstSubReg = getSubReg(FirstStridedReg, AArch64::zsub0);
-          if (IsMulZPR && (FirstSubReg - AArch64::Z0) % LdOps != 0)
+          unsigned SubRegIdx = Use.getOperand(OpIdx).getSubReg();
+          if (IsMulZPR && (getSubReg(Reg, SubRegIdx) - AArch64::Z0) % UseOps !=
+                              ((unsigned)OpIdx - 1))
             continue;
 
-          for (unsigned Op = 0; Op < LdOps; ++Op) {
-            if (!is_contained(StridedOrder, FirstStridedReg + Op) ||
-                getSubReg(FirstStridedReg + Op, AArch64::zsub0) !=
-                    FirstSubReg + Op)
-              break;
-            Regs.push_back(FirstStridedReg + Op);
-          }
-
-          if (Regs.size() == LdOps && all_of(Regs, [&](MCPhysReg R) {
-                return !Matrix->isPhysRegUsed(R);
-              }))
-            Hints.push_back(FirstStridedReg + OpIdx - 1);
+          // In the example above, if VirtReg is the third operand of the
+          // tuple (%v2) and Reg == Z2_Z10, then we need to make sure that
+          // Z0_Z8, Z1_Z9 and Z3_Z11 are also available.
+          auto IsFreeConsecutiveReg = [&](unsigned UseOp) {
+            unsigned R = Reg - (OpIdx - 1) + UseOp;
+            return StridedRC->contains(R) &&
+                   (UseOp == 0 ||
+                    ((getSubReg(R, AArch64::zsub0) - AArch64::Z0) ==
+                     (getSubReg(R - 1, AArch64::zsub0) - AArch64::Z0) + 1)) &&
+                   !Matrix->isPhysRegUsed(R);
+          };
+          if (all_of(iota_range<unsigned>(0U, UseOps, /*Inclusive=*/false),
+                     IsFreeConsecutiveReg))
+            Hints.push_back(Reg);
         }
       } else {
         // At least one operand already has a physical register assigned.