[llvm] 603ac57 - [AArch64][SME] Add support for zeroing ZT0 to CommitZASavePseudo (#166360)

via llvm-commits llvm-commits at lists.llvm.org
Tue Nov 18 01:26:25 PST 2025


Author: Benjamin Maxwell
Date: 2025-11-18T09:26:21Z
New Revision: 603ac57ef9768dc557d223f61c5e3c5cb5e50a12

URL: https://github.com/llvm/llvm-project/commit/603ac57ef9768dc557d223f61c5e3c5cb5e50a12
DIFF: https://github.com/llvm/llvm-project/commit/603ac57ef9768dc557d223f61c5e3c5cb5e50a12.diff

LOG: [AArch64][SME] Add support for zeroing ZT0 to CommitZASavePseudo (#166360)

This will be used to support ZT0 in the MachineSMEABIPass.

Added: 
    

Modified: 
    llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
    llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td
    llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
    llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir

Removed: 
    


################################################################################
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 5bb70ee11b06d..737169253ddb3 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 24d30c731b945..2afbec92392f0 100644
--- a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
+++ b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
@@ -842,6 +842,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
 


        


More information about the llvm-commits mailing list