Test + comments

Change-Id: I68a4ae06ecabe87b02ec351b30de28ffa45ca812

Author: MacDue

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

@@ -1031,7 +1031,13 @@ MachineBasicBlock *AArch64ExpandPseudo::expandCommitOrRestoreZASave(
   // Replace the pseudo with a call (BL).
   MachineInstrBuilder MIB =
       BuildMI(*SMBB, SMBB->end(), DL, TII->get(AArch64::BL));
+  // Copy operands (mainly the regmask) from the pseudo.
+  unsigned FirstBLOperand = IsRestoreZA ? 2 : 1;
+  for (unsigned I = FirstBLOperand; I < MI.getNumOperands(); ++I)
+    MIB.add(MI.getOperand(I));
+
   if (IsRestoreZA) {
+    // Mark the TPIDR2 block pointer (X0) as an implicit use.
     MIB.addReg(MI.getOperand(1).getReg(), RegState::Implicit);
   } else /*CommitZA*/ {
     auto *TRI = MBB.getParent()->getSubtarget().getRegisterInfo();
@@ -1046,11 +1052,8 @@ MachineBasicBlock *AArch64ExpandPseudo::expandCommitOrRestoreZASave(
           .addDef(AArch64::ZAB0, RegState::ImplicitDefine);
     }
   }
-  unsigned FirstBLOperand = IsRestoreZA ? 2 : 1;
-  for (unsigned I = FirstBLOperand; I < MI.getNumOperands(); ++I)
-    MIB.add(MI.getOperand(I));
-  BuildMI(SMBB, DL, TII->get(AArch64::B)).addMBB(EndBB);
 
+  BuildMI(SMBB, DL, TII->get(AArch64::B)).addMBB(EndBB);
   MI.eraseFromParent();
   return EndBB;
 }

llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir

@@ -0,0 +1,104 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+# RUN: llc -run-pass=aarch64-expand-pseudo -mtriple=aarch64-unknown-linux-gnu -o - %s | FileCheck %s
+
+---
+
+# X0 = TPIDR2 block pointer
+# X8 = TPIDR2_EL0
+name:            restore_za_save
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: restore_za_save
+    ; CHECK: successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: $x0 = IMPLICIT_DEF
+    ; CHECK-NEXT: $x8 = MRS 56965, implicit-def $nzcv
+    ; CHECK-NEXT: CBZX $x8, %bb.1
+    ; CHECK-NEXT: B %bb.2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: .1:
+    ; CHECK-NEXT: successors: %bb.2(0x80000000)
+    ; CHECK-NEXT: liveins: $x8, $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: BL &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $lr, implicit $sp, implicit $x0
+    ; CHECK-NEXT: B %bb.2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: .2:
+    ; CHECK-NEXT: RET undef $lr
+    $x0 = IMPLICIT_DEF
+    $x8 = MRS 56965, implicit-def $nzcv
+
+    RestoreZAPseudo $x8, $x0, &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0
+
+    RET_ReallyLR
+
+...
+---
+
+# X0 = TPIDR2 block pointer
+# X8 = TPIDR2_EL0
+name:            commit_za_save
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: commit_za_save
+    ; CHECK: successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: $x0 = IMPLICIT_DEF
+    ; CHECK-NEXT: $x8 = MRS 56965, implicit-def $nzcv
+    ; CHECK-NEXT: CBNZX $x8, %bb.1
+    ; CHECK-NEXT: B %bb.2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: .1:
+    ; CHECK-NEXT: successors: %bb.2(0x80000000)
+    ; CHECK-NEXT: liveins: $x8, $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: BL $x0, &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $lr, implicit $sp
+    ; CHECK-NEXT: MSR 56965, $xzr
+    ; CHECK-NEXT: B %bb.2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: .2:
+    ; CHECK-NEXT: RET undef $lr
+    $x0 = IMPLICIT_DEF
+    $x8 = MRS 56965, implicit-def $nzcv
+
+    CommitZASavePseudo $x8, $x0, &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0
+
+    RET_ReallyLR
+
+...
+---
+# X8 = TPIDR2_EL0
+name:            commit_za_save_zero_za
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: commit_za_save_zero_za
+    ; CHECK: successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: $x8 = MRS 56965, implicit-def $nzcv
+    ; CHECK-NEXT: CBNZX $x8, %bb.1
+    ; CHECK-NEXT: B %bb.2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: .1:
+    ; CHECK-NEXT: successors: %bb.2(0x80000000)
+    ; CHECK-NEXT: liveins: $x8
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: BL &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $lr, implicit $sp, implicit-def $zab0
+    ; CHECK-NEXT: MSR 56965, $xzr
+    ; CHECK-NEXT: ZERO_M 255, implicit-def $zab0
+    ; CHECK-NEXT: B %bb.2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: .2:
+    ; CHECK-NEXT: RET undef $lr
+    $x8 = MRS 56965, implicit-def $nzcv
+
+    CommitZASavePseudo $x8, &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zab0
+
+    RET_ReallyLR
+
+...