[AArch64][SME] Use PNR Reg classes for predicate constraint

[MDevereau]

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

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/67606.diff

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+7-4)
• (added) llvm/test/CodeGen/AArch64/aarch64-sme2-asm.ll (+29)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index f887cf50662dc4c..12cb9d66725ad1d 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/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;
   }
 }
 
diff --git a/llvm/test/CodeGen/AArch64/aarch64-sme2-asm.ll b/llvm/test/CodeGen/AArch64/aarch64-sme2-asm.ll
new file mode 100644
index 000000000000000..98e7ad740681e68
--- /dev/null
+++ b/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
+}
\ No newline at end of file

[sdesmalen-arm]

We should probably add a PNR_3bRegClass as well, even if this isn't used by any of the instructions, but at least it will allow users to limit the chosen register to pn0-pn7. Could you add it? (I'm happy for this to be done in a separate patch if this leads to those flaky tests needing to be updated)

[MDevereau]

I'll do this as a follow up patch for the reasons you mentioned.

[MDevereau]

Created follow up PR here: #67785