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

[via llvm-commits]

Author: Matthew Devereau
Date: 2023-09-29T10:33:25+01:00
New Revision: 0d328e3875285e8b0163b536aef747e0deba5261

URL: https://github.com/llvm/llvm-project/commit/0d328e3875285e8b0163b536aef747e0deba5261
DIFF: https://github.com/llvm/llvm-project/commit/0d328e3875285e8b0163b536aef747e0deba5261.diff

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

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

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

Modified: 
    llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 24dfdfcf1b8ba01..662306c0dadf0d4 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