[AArch64][SME] Use PNR Reg classes for predicate constraint (llvm#67606)

This patch fixes an error where ASM with constraints cannot select SME
instructions which use the top eight predicate-as-counter registers.

Author: MDevereau

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -10052,18 +10052,21 @@ static PredicateConstraint parsePredicateConstraint(StringRef Constraint) {
 
 static const TargetRegisterClass *
 getPredicateRegisterClass(PredicateConstraint Constraint, EVT VT) {
-  if (!VT.isScalableVector() || VT.getVectorElementType() != MVT::i1)
+  if (VT != MVT::aarch64svcount &&
+      (!VT.isScalableVector() || VT.getVectorElementType() != MVT::i1))
     return nullptr;
 
   switch (Constraint) {
   default:
     return nullptr;
   case PredicateConstraint::Uph:
-    return &AArch64::PPR_p8to15RegClass;
+    return VT == MVT::aarch64svcount ? &AArch64::PNR_p8to15RegClass
+                                     : &AArch64::PPR_p8to15RegClass;
   case PredicateConstraint::Upl:
-    return &AArch64::PPR_3bRegClass;
+    return VT == MVT::aarch64svcount ? nullptr : &AArch64::PPR_3bRegClass;
   case PredicateConstraint::Upa:
-    return &AArch64::PPRRegClass;
+    return VT == MVT::aarch64svcount ? &AArch64::PNRRegClass
+                                     : &AArch64::PPRRegClass;
   }
 }
 

llvm/test/CodeGen/AArch64/aarch64-sme2-asm.ll

@@ -0,0 +1,29 @@
+; RUN: llc < %s -mtriple aarch64-none-linux-gnu -mattr=+sme2 -stop-after=finalize-isel | FileCheck %s
+
+define dso_local void @UphPNR(target("aarch64.svcount") %predcnt) {
+entry:
+; CHECK:  %0:ppr = COPY $p0
+; CHECK:  STR_PXI %0, %stack.0.predcnt.addr, 0 :: (store unknown-size into %ir.predcnt.addr, align 2)
+; CHECK:  %1:pnr_p8to15 = COPY %0
+; CHECK:  INLINEASM &"ld1w {z0.s,z1.s,z2.s,z3.s}, $0/z, [x10]", 1 /* sideeffect attdialect */, 393225 /* reguse:PNR_p8to15 */, %1
+; CHECK:  RET_ReallyLR
+  %predcnt.addr = alloca target("aarch64.svcount"), align 2
+  store target("aarch64.svcount") %predcnt, ptr %predcnt.addr, align 2
+  %0 = load target("aarch64.svcount"), ptr %predcnt.addr, align 2
+  call void asm sideeffect "ld1w {z0.s,z1.s,z2.s,z3.s}, $0/z, [x10]", "@3Uph"(target("aarch64.svcount") %0)
+  ret void
+}
+
+define dso_local void @UpaPNR(target("aarch64.svcount") %predcnt) {
+entry:
+; CHECK:  %0:ppr = COPY $p0
+; CHECK:  STR_PXI %0, %stack.0.predcnt.addr, 0 :: (store unknown-size into %ir.predcnt.addr, align 2)
+; CHECK:  %1:pnr = COPY %0
+; CHECK:  INLINEASM &"ld1w {z0.s,z1.s,z2.s,z3.s}, $0/z, [x10]", 1 /* sideeffect attdialect */, 262153 /* reguse:PNR */, %1
+; CHECK:  RET_ReallyLR
+  %predcnt.addr = alloca target("aarch64.svcount"), align 2
+  store target("aarch64.svcount") %predcnt, ptr %predcnt.addr, align 2
+  %0 = load target("aarch64.svcount"), ptr %predcnt.addr, align 2
+  call void asm sideeffect "ld1w {z0.s,z1.s,z2.s,z3.s}, $0/z, [x10]", "@3Upa"(target("aarch64.svcount") %0)
+  ret void
+}