[llvm] [AArch64][SME] Emit unimplemented error for non-lazy-save of ZA. (PR #227608)

[Sander de Smalen via llvm-commits]

https://github.com/sdesmalen-arm created https://github.com/llvm/llvm-project/pull/227608

A non-lazy-save of ZA is required when the caller has both ZA and ZT0 state and calls a function with a shared-ZA interface that only shares ZT0.

This case is already rejected by Clang but LLVM compiles it without generating a save of ZA.

>From ac6348af273eb4e6e8b457f751b62c9f869c370c Mon Sep 17 00:00:00 2001
From: Sander de Smalen <sander.desmalen at arm.com>
Date: Wed, 30 Sep 2026 08:45:15 +0000
Subject: [PATCH] [AArch64][SME] Emit unimplemented error for non-lazy-save of
 ZA.

A non-lazy-save of ZA is required when the caller has both ZA and
ZT0 state and calls a function with a shared-ZA interface that only
shares ZT0.

This case is already rejected by Clang but LLVM compiles it without
generating a save of ZA.
---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp   |  3 +++
 llvm/lib/Target/AArch64/AArch64SMEAttributes.h    |  8 +++++++-
 llvm/test/CodeGen/AArch64/sme-non-lazy-save-za.ll | 14 ++++++++++++++
 3 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/sme-non-lazy-save-za.ll

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 26ca67ed132e5..9e5aac507e8dd 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -10588,6 +10588,9 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
   // Determine whether we need any streaming mode changes.
   SMECallAttrs CallAttrs =
       getSMECallAttrs(MF.getFunction(), getRuntimeLibcallsInfo(), CLI);
+  if (CallAttrs.requiresNonLazySaveZA())
+    reportFatalUsageError(
+        "Calls that require saving ZA non-lazily is not yet implemented");
 
   std::optional<unsigned> ZAMarkerNode = getZAMarkerForCall(CallAttrs);
 
diff --git a/llvm/lib/Target/AArch64/AArch64SMEAttributes.h b/llvm/lib/Target/AArch64/AArch64SMEAttributes.h
index a340866335793..5a1e3d60df70a 100644
--- a/llvm/lib/Target/AArch64/AArch64SMEAttributes.h
+++ b/llvm/lib/Target/AArch64/AArch64SMEAttributes.h
@@ -189,6 +189,11 @@ class SMECallAttrs {
            !callee().isSMEABIRoutine();
   }
 
+  bool requiresNonLazySaveZA() const {
+    return caller().hasZAState() && callee().hasSharedZAInterface() &&
+           !callee().sharesZA();
+  }
+
   bool requiresPreservingZT0() const {
     return caller().hasZT0State() && !callee().sharesZT0() &&
            !callee().hasAgnosticZAInterface();
@@ -200,7 +205,8 @@ class SMECallAttrs {
   }
 
   bool requiresZASave() const {
-    return requiresLazySave() || requiresPreservingAllZAState();
+    return requiresLazySave() || requiresNonLazySaveZA() ||
+           requiresPreservingAllZAState();
   }
 };
 
diff --git a/llvm/test/CodeGen/AArch64/sme-non-lazy-save-za.ll b/llvm/test/CodeGen/AArch64/sme-non-lazy-save-za.ll
new file mode 100644
index 0000000000000..7a8827dfedd66
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/sme-non-lazy-save-za.ll
@@ -0,0 +1,14 @@
+; RUN: not llc < %s 2>&1 | FileCheck %s
+; RUN: not llc -fast-isel=true < %s 2>&1 | FileCheck %s
+; RUN: not llc -global-isel=true -global-isel-abort=2 < %s 2>&1 | FileCheck %s
+
+; CHECK: LLVM ERROR: Calls that require saving ZA non-lazily is not yet implemented
+
+target triple = "aarch64"
+
+define void @foo(ptr %f) #0 {
+  call void %f() "aarch64_inout_zt0"
+  ret void
+}
+
+attributes #0 = { "aarch64_inout_zt0" "aarch64_inout_za" "target-feature"="+sme,+sve" }