[llvm] [AArch64][SME] Elide private ZA setup when possible (PR #196090)
via llvm-commits
llvm-commits at lists.llvm.org
Wed May 6 08:06:26 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-aarch64
Author: Benjamin Maxwell (MacDue)
<details>
<summary>Changes</summary>
In private ZA functions without any instructions that require "active" ZA we can omit all ZA setup (and saves/restores). This is equivalent to removing the `__arm_new("za/zt0")` attribute when ZA state is unused.
---
Full diff: https://github.com/llvm/llvm-project/pull/196090.diff
3 Files Affected:
- (modified) llvm/lib/Target/AArch64/MachineSMEABIPass.cpp (+16-1)
- (modified) llvm/test/CodeGen/AArch64/sme-peephole-opts.ll (+2-34)
- (modified) llvm/test/CodeGen/AArch64/sme-zt0-state.ll (+11-36)
``````````diff
diff --git a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
index 4462af1ca306f..3aba11bd7de9f 100644
--- a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
+++ b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
@@ -490,7 +490,6 @@ FunctionInfo MachineSMEABI::collectNeededZAStates(SMEAttrs SMEFnAttrs) {
auto [NeededState, InsertPt] = getInstNeededZAState(*TRI, MI, SMEFnAttrs);
assert((InsertPt == MBBI || isCallStartOpcode(InsertPt->getOpcode())) &&
"Unexpected state change insertion point!");
- // TODO: Do something to avoid state changes where NZCV is live.
if (MBBI == FirstTerminatorInsertPt)
Block.PhysLiveRegsAtExit = PhysLiveRegs;
if (MBBI == FirstNonPhiInsertPt)
@@ -1162,6 +1161,19 @@ void MachineSMEABI::emitStateChange(EmitContext &Context,
}
}
+/// Returns true if private ZA setup and be elided. This occurs when there is
+/// no instruction within the function that requires ZA to be active.
+static bool canElidePrivateZASetup(const FunctionInfo &FnInfo) {
+ for (const BlockInfo &BlockInfo : FnInfo.Blocks) {
+ for (const InstInfo &InstInfo : BlockInfo.Insts) {
+ if (InstInfo.NeededState == ZAState::ACTIVE ||
+ InstInfo.NeededState == ZAState::ACTIVE_ZT0_SAVED)
+ return false;
+ }
+ }
+ return true;
+}
+
} // end anonymous namespace
INITIALIZE_PASS(MachineSMEABI, "aarch64-machine-sme-abi", "Machine SME ABI",
@@ -1193,6 +1205,9 @@ bool MachineSMEABI::runOnMachineFunction(MachineFunction &MF) {
FunctionInfo FnInfo = collectNeededZAStates(SMEFnAttrs);
+ if (SMEFnAttrs.hasPrivateZAInterface() && canElidePrivateZASetup(FnInfo))
+ return false;
+
SmallVector<ZAState> BundleStates = assignBundleZAStates(Bundles, FnInfo);
EmitContext Context;
diff --git a/llvm/test/CodeGen/AArch64/sme-peephole-opts.ll b/llvm/test/CodeGen/AArch64/sme-peephole-opts.ll
index bdfddad32ff3a..36539d94338a0 100644
--- a/llvm/test/CodeGen/AArch64/sme-peephole-opts.ll
+++ b/llvm/test/CodeGen/AArch64/sme-peephole-opts.ll
@@ -643,42 +643,10 @@ define void @test15(ptr %callee) nounwind "aarch64_za_state_agnostic" {
define void @test16(ptr %callee) nounwind "aarch64_pstate_sm_body" "aarch64_new_za" {
; CHECK-LABEL: test16:
; CHECK: // %bb.0:
-; CHECK-NEXT: stp d15, d14, [sp, #-96]! // 16-byte Folded Spill
-; CHECK-NEXT: stp d13, d12, [sp, #16] // 16-byte Folded Spill
-; CHECK-NEXT: stp d11, d10, [sp, #32] // 16-byte Folded Spill
-; CHECK-NEXT: stp d9, d8, [sp, #48] // 16-byte Folded Spill
-; CHECK-NEXT: stp x29, x30, [sp, #64] // 16-byte Folded Spill
-; CHECK-NEXT: add x29, sp, #64
-; CHECK-NEXT: str x19, [sp, #80] // 8-byte Spill
-; CHECK-NEXT: sub sp, sp, #16
-; CHECK-NEXT: rdsvl x8, #1
-; CHECK-NEXT: mov x9, sp
-; CHECK-NEXT: msub x9, x8, x8, x9
-; CHECK-NEXT: mov sp, x9
-; CHECK-NEXT: stp x9, x8, [x29, #-80]
-; CHECK-NEXT: mrs x8, TPIDR2_EL0
-; CHECK-NEXT: cbz x8, .LBB17_2
-; CHECK-NEXT: // %bb.1:
-; CHECK-NEXT: bl __arm_tpidr2_save
-; CHECK-NEXT: msr TPIDR2_EL0, xzr
-; CHECK-NEXT: zero {za}
-; CHECK-NEXT: .LBB17_2:
-; CHECK-NEXT: smstart za
-; CHECK-NEXT: smstart sm
-; CHECK-NEXT: sub x8, x29, #80
-; CHECK-NEXT: msr TPIDR2_EL0, x8
-; CHECK-NEXT: smstop sm
+; CHECK-NEXT: str x30, [sp, #-16]! // 8-byte Folded Spill
; CHECK-NEXT: bl callee
; CHECK-NEXT: bl callee
-; CHECK-NEXT: msr TPIDR2_EL0, xzr
-; CHECK-NEXT: smstop za
-; CHECK-NEXT: sub sp, x29, #64
-; CHECK-NEXT: ldp x29, x30, [sp, #64] // 16-byte Folded Reload
-; CHECK-NEXT: ldr x19, [sp, #80] // 8-byte Reload
-; CHECK-NEXT: ldp d9, d8, [sp, #48] // 16-byte Folded Reload
-; CHECK-NEXT: ldp d11, d10, [sp, #32] // 16-byte Folded Reload
-; CHECK-NEXT: ldp d13, d12, [sp, #16] // 16-byte Folded Reload
-; CHECK-NEXT: ldp d15, d14, [sp], #96 // 16-byte Folded Reload
+; CHECK-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
; CHECK-NEXT: ret
call void @callee()
call void @callee()
diff --git a/llvm/test/CodeGen/AArch64/sme-zt0-state.ll b/llvm/test/CodeGen/AArch64/sme-zt0-state.ll
index d3c3c111c205b..7ef3262e5811c 100644
--- a/llvm/test/CodeGen/AArch64/sme-zt0-state.ll
+++ b/llvm/test/CodeGen/AArch64/sme-zt0-state.ll
@@ -131,57 +131,32 @@ define void @zt0_in_caller_zt0_new_callee(ptr %callee) "aarch64_in_zt0" nounwind
; New-ZT0 Callee
-; Expect commit of lazy-save if ZA is dormant
-; Expect smstart ZA & clear ZT0
-; Expect spill & fill of ZT0 around call
-; Before return, expect smstop ZA
+; Expect ZA state setup to be elided (no instructions in this function require
+; ZA state).
define void @zt0_new_caller_zt0_new_callee(ptr %callee) "aarch64_new_zt0" nounwind {
; CHECK-LABEL: zt0_new_caller_zt0_new_callee:
; CHECK: // %bb.0:
-; CHECK-NEXT: sub sp, sp, #80
-; CHECK-NEXT: str x30, [sp, #64] // 8-byte Spill
-; CHECK-NEXT: mrs x8, TPIDR2_EL0
-; CHECK-NEXT: cbz x8, .LBB6_2
-; CHECK-NEXT: // %bb.1:
-; CHECK-NEXT: bl __arm_tpidr2_save
-; CHECK-NEXT: msr TPIDR2_EL0, xzr
-; CHECK-NEXT: zero { zt0 }
-; CHECK-NEXT: .LBB6_2:
-; CHECK-NEXT: smstart za
-; CHECK-NEXT: mov x8, sp
-; CHECK-NEXT: str zt0, [x8]
-; CHECK-NEXT: smstop za
+; CHECK-NEXT: str x30, [sp, #-16]! // 8-byte Folded Spill
; CHECK-NEXT: blr x0
-; CHECK-NEXT: ldr x30, [sp, #64] // 8-byte Reload
-; CHECK-NEXT: add sp, sp, #80
+; CHECK-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload
; CHECK-NEXT: ret
call void %callee() "aarch64_new_zt0";
ret void;
}
-; Expect commit of lazy-save if ZA is dormant
-; Expect smstart ZA & clear ZT0
-; No spill & fill of ZT0 around __arm_tpidr2_save
; Expect spill & fill of ZT0 around __arm_sme_state call
-; Before return, expect smstop ZA
-define i64 @zt0_new_caller_abi_routine_callee() "aarch64_new_zt0" nounwind {
+define i64 @zt0_new_caller_abi_routine_callee() "aarch64_inout_zt0" nounwind {
; CHECK-LABEL: zt0_new_caller_abi_routine_callee:
; CHECK: // %bb.0:
; CHECK-NEXT: sub sp, sp, #80
-; CHECK-NEXT: str x30, [sp, #64] // 8-byte Spill
-; CHECK-NEXT: mrs x8, TPIDR2_EL0
-; CHECK-NEXT: cbz x8, .LBB7_2
-; CHECK-NEXT: // %bb.1:
-; CHECK-NEXT: bl __arm_tpidr2_save
-; CHECK-NEXT: msr TPIDR2_EL0, xzr
-; CHECK-NEXT: zero { zt0 }
-; CHECK-NEXT: .LBB7_2:
-; CHECK-NEXT: smstart za
-; CHECK-NEXT: mov x8, sp
-; CHECK-NEXT: str zt0, [x8]
+; CHECK-NEXT: stp x30, x19, [sp, #64] // 16-byte Folded Spill
+; CHECK-NEXT: mov x19, sp
+; CHECK-NEXT: str zt0, [x19]
; CHECK-NEXT: smstop za
; CHECK-NEXT: bl __arm_sme_state
-; CHECK-NEXT: ldr x30, [sp, #64] // 8-byte Reload
+; CHECK-NEXT: smstart za
+; CHECK-NEXT: ldr zt0, [x19]
+; CHECK-NEXT: ldp x30, x19, [sp, #64] // 16-byte Folded Reload
; CHECK-NEXT: add sp, sp, #80
; CHECK-NEXT: ret
%res = call {i64, i64} @__arm_sme_state()
``````````
</details>
https://github.com/llvm/llvm-project/pull/196090
More information about the llvm-commits
mailing list