[llvm] 2efea51 - [AArch64] Fix spilling/filling of virtual registers in PNR regclass. (#70679)

[via llvm-commits]

Author: Sander de Smalen
Date: 2023-11-01T10:57:12Z
New Revision: 2efea512c25bb6f6178e0ff881a96e99b2aecee4

URL: https://github.com/llvm/llvm-project/commit/2efea512c25bb6f6178e0ff881a96e99b2aecee4
DIFF: https://github.com/llvm/llvm-project/commit/2efea512c25bb6f6178e0ff881a96e99b2aecee4.diff

LOG: [AArch64] Fix spilling/filling of virtual registers in PNR regclass. (#70679)

We made the assumption that the registers were always physical
registers, which doesn't have to be true.

Added: 
    

Modified: 
    llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
    llvm/test/CodeGen/AArch64/spillfill-sve.mir

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index a9b0cfbbd33ea6a..4680af7d121c757 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -4773,7 +4773,10 @@ void AArch64InstrInfo::storeRegToStackSlot(MachineBasicBlock &MBB,
     } else if (AArch64::PNRRegClass.hasSubClassEq(RC)) {
       assert((Subtarget.hasSVE2p1() || Subtarget.hasSME2()) &&
              "Unexpected register store without SVE2p1 or SME2");
-      SrcReg = (SrcReg - AArch64::PN0) + AArch64::P0;
+      if (SrcReg.isVirtual())
+        MF.getRegInfo().constrainRegClass(SrcReg, &AArch64::PPRRegClass);
+      else
+        SrcReg = (SrcReg - AArch64::PN0) + AArch64::P0;
       Opc = AArch64::STR_PXI;
       StackID = TargetStackID::ScalableVector;
     }
@@ -4946,7 +4949,10 @@ void AArch64InstrInfo::loadRegFromStackSlot(MachineBasicBlock &MBB,
       assert((Subtarget.hasSVE2p1() || Subtarget.hasSME2()) &&
              "Unexpected register load without SVE2p1 or SME2");
       PNRReg = DestReg;
-      DestReg = (DestReg - AArch64::PN0) + AArch64::P0;
+      if (DestReg.isVirtual())
+        MF.getRegInfo().constrainRegClass(DestReg, &AArch64::PPRRegClass);
+      else
+        DestReg = (DestReg - AArch64::PN0) + AArch64::P0;
       Opc = AArch64::LDR_PXI;
       StackID = TargetStackID::ScalableVector;
     }

diff  --git a/llvm/test/CodeGen/AArch64/spillfill-sve.mir b/llvm/test/CodeGen/AArch64/spillfill-sve.mir
index dfd41633c46c3ec..de93cbfe1aa19f5 100644
--- a/llvm/test/CodeGen/AArch64/spillfill-sve.mir
+++ b/llvm/test/CodeGen/AArch64/spillfill-sve.mir
@@ -8,6 +8,7 @@
 
   define aarch64_sve_vector_pcs void @spills_fills_stack_id_ppr() #0 { entry: unreachable }
   define aarch64_sve_vector_pcs void @spills_fills_stack_id_pnr() #1 { entry: unreachable }
+  define aarch64_sve_vector_pcs void @spills_fills_stack_id_virtreg_pnr() #1 { entry: unreachable }
   define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr() #0 { entry: unreachable }
   define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr2() #0 { entry: unreachable }
   define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr2strided() #0 { entry: unreachable }
@@ -106,6 +107,46 @@ body:             |
     RET_ReallyLR
 ...
 ---
+name: spills_fills_stack_id_virtreg_pnr
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: pnr_p8to15 }
+stack:
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: spills_fills_stack_id_virtreg_pnr
+    ; CHECK: stack:
+    ; CHECK:      - { id: 0, name: '', type: spill-slot, offset: 0, size: 2, alignment: 2
+    ; CHECK-NEXT:     stack-id: scalable-vector, callee-saved-register: ''
+
+    ; EXPAND-LABEL: name: spills_fills_stack_id_virtreg_pnr
+    ; EXPAND: renamable $pn8 = WHILEGE_CXX_B
+    ; EXPAND: STR_PXI killed renamable $pn8, $sp, 7
+    ; EXPAND: $p0 = LDR_PXI $sp, 7, implicit-def $pn0
+
+    %0:pnr_p8to15 = WHILEGE_CXX_B undef $x0, undef $x0, 0, implicit-def dead $nzcv
+
+    $pn0 = IMPLICIT_DEF
+    $pn1 = IMPLICIT_DEF
+    $pn2 = IMPLICIT_DEF
+    $pn3 = IMPLICIT_DEF
+    $pn4 = IMPLICIT_DEF
+    $pn5 = IMPLICIT_DEF
+    $pn6 = IMPLICIT_DEF
+    $pn7 = IMPLICIT_DEF
+    $pn8 = IMPLICIT_DEF
+    $pn9 = IMPLICIT_DEF
+    $pn10 = IMPLICIT_DEF
+    $pn11 = IMPLICIT_DEF
+    $pn12 = IMPLICIT_DEF
+    $pn13 = IMPLICIT_DEF
+    $pn14 = IMPLICIT_DEF
+    $pn15 = IMPLICIT_DEF
+
+    $pn0 = COPY %0
+    RET_ReallyLR
+...
+---
 name: spills_fills_stack_id_zpr
 tracksRegLiveness: true
 registers: