[AArch64][SME] Add support for zeroing ZT0 to CommitZASavePseudo

[MacDue]

This will be used to support ZT0 in the MachineSMEABIPass.

[MacDue]

This is a stacked PR. See other PRs below:

1. 👉 [AArch64][SME] Add support for zeroing ZT0 to CommitZASavePseudo #166360
2. [AArch64][SME] Handle zeroing ZA and ZT0 in functions with ZT0 state #166361
3. [AArch64][SME] Support saving/restoring ZT0 in the MachineSMEABIPass #166362

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Benjamin Maxwell (MacDue)

Changes

This will be used to support ZT0 in the MachineSMEABIPass.

---

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

4 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp (+9-4)
• (modified) llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td (+2-1)
• (modified) llvm/lib/Target/AArch64/MachineSMEABIPass.cpp (+1)
• (modified) llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir (+67-2)

diff --git a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
index f63981b87c1c1..34d74d04c4419 100644
--- a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
@@ -1063,6 +1063,7 @@ AArch64ExpandPseudo::expandCommitZASave(MachineBasicBlock &MBB,
                                         MachineBasicBlock::iterator MBBI) {
   MachineInstr &MI = *MBBI;
   DebugLoc DL = MI.getDebugLoc();
+  [[maybe_unused]] auto *RI = MBB.getParent()->getSubtarget().getRegisterInfo();
 
   // Compare TPIDR2_EL0 against 0. Commit ZA if TPIDR2_EL0 is non-zero.
   MachineInstrBuilder Branch =
@@ -1073,21 +1074,25 @@ AArch64ExpandPseudo::expandCommitZASave(MachineBasicBlock &MBB,
   MachineInstrBuilder MIB =
       BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::BL));
   // Copy operands (mainly the regmask) from the pseudo.
-  for (unsigned I = 2; I < MI.getNumOperands(); ++I)
+  for (unsigned I = 3; I < MI.getNumOperands(); ++I)
     MIB.add(MI.getOperand(I));
   // Clear TPIDR2_EL0.
   BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::MSR))
       .addImm(AArch64SysReg::TPIDR2_EL0)
       .addReg(AArch64::XZR);
   bool ZeroZA = MI.getOperand(1).getImm() != 0;
+  bool ZeroZT0 = MI.getOperand(2).getImm() != 0;
   if (ZeroZA) {
-    [[maybe_unused]] auto *TRI =
-        MBB.getParent()->getSubtarget().getRegisterInfo();
-    assert(MI.definesRegister(AArch64::ZAB0, TRI) && "should define ZA!");
+    assert(MI.definesRegister(AArch64::ZAB0, RI) && "should define ZA!");
     BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::ZERO_M))
         .addImm(ZERO_ALL_ZA_MASK)
         .addDef(AArch64::ZAB0, RegState::ImplicitDefine);
   }
+  if (ZeroZT0) {
+    assert(MI.definesRegister(AArch64::ZT0, RI) && "should define ZT0!");
+    BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::ZERO_T))
+        .addDef(AArch64::ZT0);
+  }
 
   MI.eraseFromParent();
   return &EndBB;
diff --git a/llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td b/llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td
index 752b185832c30..8f8f211c5fceb 100644
--- a/llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td
@@ -108,7 +108,8 @@ def SMEStateAllocPseudo : Pseudo<(outs), (ins), []>, Sched<[]>;
 
 def CommitZASavePseudo
   : Pseudo<(outs),
-           (ins GPR64:$tpidr2_el0, i1imm:$zero_za, i64imm:$commit_routine, variable_ops), []>,
+           (ins GPR64:$tpidr2_el0, i1imm:$zero_za, i1imm:$zero_zt0,
+                i64imm:$commit_routine, variable_ops), []>,
            Sched<[]>;
 
 def AArch64_inout_za_use
diff --git a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
index 7cb500394cec2..8f9aae944ad6d 100644
--- a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
+++ b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
@@ -832,6 +832,7 @@ void MachineSMEABI::emitNewZAPrologue(MachineBasicBlock &MBB,
       BuildMI(MBB, MBBI, DL, TII->get(AArch64::CommitZASavePseudo))
           .addReg(TPIDR2EL0)
           .addImm(ZeroZA ? 1 : 0)
+          .addImm(/*ZeroZT0=*/false)
           .addExternalSymbol(TLI->getLibcallName(RTLIB::SMEABI_TPIDR2_SAVE))
           .addRegMask(TRI->SMEABISupportRoutinesCallPreservedMaskFromX0());
   if (ZeroZA)
diff --git a/llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir b/llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir
index 6ca9b9b6cb200..9b745d56c4b7f 100644
--- a/llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir
+++ b/llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir
@@ -62,7 +62,7 @@ body:             |
     ; CHECK-NEXT: RET undef $lr
     $x8 = MRS 56965, implicit-def $nzcv
 
-    CommitZASavePseudo $x8, 0, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0
+    CommitZASavePseudo $x8, 0, 0, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0
 
     RET_ReallyLR
 
@@ -94,7 +94,72 @@ body:             |
     ; CHECK-NEXT: RET undef $lr
     $x8 = MRS 56965, implicit-def $nzcv
 
-    CommitZASavePseudo $x8, 1, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zab0
+    CommitZASavePseudo $x8, 1, 0, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zab0
+
+    RET_ReallyLR
+
+...
+---
+# X8 = TPIDR2_EL0
+name:            commit_za_save_zero_zt0
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: commit_za_save_zero_zt0
+    ; 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_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $lr, implicit $sp, implicit-def $zt0
+    ; CHECK-NEXT: MSR 56965, $xzr
+    ; CHECK-NEXT: $zt0 = ZERO_T
+    ; CHECK-NEXT: B %bb.2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: .2:
+    ; CHECK-NEXT: RET undef $lr
+    $x8 = MRS 56965, implicit-def $nzcv
+
+    CommitZASavePseudo $x8, 0, 1, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zt0
+
+    RET_ReallyLR
+
+...
+---
+# X8 = TPIDR2_EL0
+name:            commit_za_save_zero_everything
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: commit_za_save_zero_everything
+    ; 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_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $lr, implicit $sp, implicit-def $zab0, implicit-def $zt0
+    ; CHECK-NEXT: MSR 56965, $xzr
+    ; CHECK-NEXT: ZERO_M 255, implicit-def $zab0
+    ; CHECK-NEXT: $zt0 = ZERO_T
+    ; CHECK-NEXT: B %bb.2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: .2:
+    ; CHECK-NEXT: RET undef $lr
+    $x8 = MRS 56965, implicit-def $nzcv
+
+    CommitZASavePseudo $x8, 1, 1, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zab0, implicit-def $zt0
 
     RET_ReallyLR