[AArch64][SME] Fix incorrect "attributes at callsite do not match" assert

[MacDue]

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).

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Benjamin Maxwell (MacDue)

Changes

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).

---

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

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp (+6-6)
• (modified) llvm/test/CodeGen/AArch64/sme-za-exceptions.ll (+1-1)

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: