Tweaks

Change-Id: I51fbdc4b701a096d1f1638a4abab5c105ab37c5b

Author: MacDue

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

@@ -1032,8 +1032,7 @@ MachineBasicBlock *AArch64ExpandPseudo::expandCommitOrRestoreZASave(
   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)
+  for (unsigned I = 2; I < MI.getNumOperands(); ++I)
     MIB.add(MI.getOperand(I));
 
   if (IsRestoreZA) {
@@ -1045,8 +1044,9 @@ MachineBasicBlock *AArch64ExpandPseudo::expandCommitOrRestoreZASave(
     BuildMI(*SMBB, SMBB->end(), DL, TII->get(AArch64::MSR))
         .addImm(AArch64SysReg::TPIDR2_EL0)
         .addReg(AArch64::XZR);
-    bool ZeroZA = MI.definesRegister(AArch64::ZAB0, TRI);
+    bool ZeroZA = MI.getOperand(1).getImm() != 0;
     if (ZeroZA) {
+      assert(MI.definesRegister(AArch64::ZAB0, TRI) && "should define ZA!");
       BuildMI(*SMBB, SMBB->end(), DL, TII->get(AArch64::ZERO_M))
           .addImm(ZERO_ALL_ZA_MASK)
           .addDef(AArch64::ZAB0, RegState::ImplicitDefine);

llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td

@@ -101,7 +101,7 @@ let hasSideEffects = 1 in {
 
 def CommitZASavePseudo
   : Pseudo<(outs),
-           (ins GPR64:$tpidr2_el0, i64imm:$commit_routine, variable_ops), []>,
+           (ins GPR64:$tpidr2_el0, i1imm:$zero_za, i64imm:$commit_routine, variable_ops), []>,
            Sched<[]>;
 
 def AArch64_inout_za_use

llvm/lib/Target/AArch64/MachineSMEABIPass.cpp

@@ -547,13 +547,16 @@ void MachineSMEABI::emitNewZAPrologue(MachineBasicBlock &MBB,
       .addReg(TPIDR2EL0, RegState::Define)
       .addImm(AArch64SysReg::TPIDR2_EL0);
   // If TPIDR2_EL0 is non-zero, commit the lazy save.
+  // NOTE: Functions that only use ZT0 don't need to zero ZA.
+  bool ZeroZA =
+      MF->getInfo<AArch64FunctionInfo>()->getSMEFnAttrs().hasZAState();
   auto CommitZASave =
       BuildMI(MBB, MBBI, DL, TII->get(AArch64::CommitZASavePseudo))
           .addReg(TPIDR2EL0)
+          .addImm(ZeroZA ? 1 : 0)
           .addExternalSymbol(TLI->getLibcallName(RTLIB::SMEABI_TPIDR2_SAVE))
           .addRegMask(TRI->SMEABISupportRoutinesCallPreservedMaskFromX0());
-  // NOTE: Functions that only use ZT0 don't need to zero ZA.
-  if (MF->getInfo<AArch64FunctionInfo>()->getSMEFnAttrs().hasZAState())
+  if (ZeroZA)
     CommitZASave.addDef(AArch64::ZAB0, RegState::ImplicitDefine);
   // Enable ZA (as ZA could have previously been in the OFF state).
   BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSRpstatesvcrImm1))

llvm/test/CodeGen/AArch64/aarch64-sme-za-call-lowering.ll

@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
 ; RUN: llc -mtriple=aarch64 -aarch64-streaming-hazard-size=0 -mattr=+sme,+sve -aarch64-new-sme-abi -stop-before=aarch64-machine-sme-abi -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK-COMMON,CHECK-BEFORE-SMEABI
-; RUN: llc -mtriple=aarch64 -aarch64-streaming-hazard-size=0 -mattr=+sme,+sve -aarch64-new-sme-abi -stop-after=aarch64-machine-sme-abi -verify-machineinstrs < %s | FileCheck %s --check-prefix=CHECK-COMMON,CHECK-AFTER-SMEABI
+; RUN: llc -mtriple=aarch64 -aarch64-streaming-hazard-size=0 -mattr=+sme,+sve -aarch64-new-sme-abi -stop-after=aarch64-machine-sme-abi -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK-COMMON,CHECK-AFTER-SMEABI
 
 declare void @private_za_callee()
 declare void @shared_za_callee() "aarch64_inout_za"

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

@@ -37,7 +37,6 @@ body:             |
 ...
 ---
 
-# X0 = TPIDR2 block pointer
 # X8 = TPIDR2_EL0
 name:            commit_za_save
 alignment:       4
@@ -47,25 +46,23 @@ body:             |
     ; 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: liveins: $x8
     ; 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: BL &__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
+    CommitZASavePseudo $x8, 0, &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0
 
     RET_ReallyLR
 
@@ -97,7 +94,7 @@ body:             |
     ; 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
+    CommitZASavePseudo $x8, 1, &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zab0
 
     RET_ReallyLR