[llvm] [AArch64][SME] Fix incorrect "attributes at callsite do not match" assert (PR #164991)

[Benjamin Maxwell via llvm-commits]

https://github.com/MacDue created https://github.com/llvm/llvm-project/pull/164991

Clang always duplicates SME attributes to each callsite, which means removing "ZA_State_Agnostic" from CalledFn before the assert resulted in the assertion failing for IR emitted by clang.

I've updated the existing test to match the form emitted by clang (which previously hit the assert).



>From d62226ce32ad207dbc1589053861c1f5d053e67f Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 24 Oct 2025 14:43:22 +0000
Subject: [PATCH] [AArch64][SME] Fix incorrect "attributes at callsite do not
 match" assert

Clang always duplicates SME attributes to each callsite, which means
removing "ZA_State_Agnostic" from CalledFn before the assert resulted
in the assertion failing for IR emitted by clang.

I've updated the existing test to match the form emitted by clang (which
previously hit the assert).

Change-Id: Ia1934bd140428db0e69853f345df9c62fc603a8d
---
 .../Target/AArch64/Utils/AArch64SMEAttributes.cpp    | 12 ++++++------
 llvm/test/CodeGen/AArch64/sme-za-exceptions.ll       |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
index 92e09eed92a6b..085c858820568 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
@@ -131,15 +131,15 @@ SMECallAttrs::SMECallAttrs(const CallBase &CB,
   if (auto *CalledFunction = CB.getCalledFunction())
     CalledFn = SMEAttrs(*CalledFunction, RTLCI);
 
-  // An `invoke` of an agnostic ZA function may not return normally (it may
-  // resume in an exception block). In this case, it acts like a private ZA
-  // callee and may require a ZA save to be set up before it is called.
-  if (isa<InvokeInst>(CB))
-    CalledFn.set(SMEAttrs::ZA_State_Agnostic, /*Enable=*/false);
-
   // FIXME: We probably should not allow SME attributes on direct calls but
   // clang duplicates streaming mode attributes at each callsite.
   assert((IsIndirect ||
           ((Callsite.withoutPerCallsiteFlags() | CalledFn) == CalledFn)) &&
          "SME attributes at callsite do not match declaration");
+
+  // An `invoke` of an agnostic ZA function may not return normally (it may
+  // resume in an exception block). In this case, it acts like a private ZA
+  // callee and may require a ZA save to be set up before it is called.
+  if (isa<InvokeInst>(CB))
+    CalledFn.set(SMEAttrs::ZA_State_Agnostic, /*Enable=*/false);
 }
diff --git a/llvm/test/CodeGen/AArch64/sme-za-exceptions.ll b/llvm/test/CodeGen/AArch64/sme-za-exceptions.ll
index b8d6c8824759a..3f35cb576694e 100644
--- a/llvm/test/CodeGen/AArch64/sme-za-exceptions.ll
+++ b/llvm/test/CodeGen/AArch64/sme-za-exceptions.ll
@@ -829,7 +829,7 @@ define void @try_catch_agnostic_za_invoke() "aarch64_za_state_agnostic" personal
 ; CHECK-SDAG-NEXT:    bl __arm_sme_restore
 ; CHECK-SDAG-NEXT:    b .LBB5_1
 entry:
-  invoke void @agnostic_za_call()
+  invoke void @agnostic_za_call() "aarch64_za_state_agnostic"
           to label %exit unwind label %catch
 
 catch: