[clang] [llvm] [mlir] [AArch64] Replace LLVM IR function attributes for PSTATE.ZA. (PR #79166)

via cfe-commits cfe-commits at lists.llvm.org
Tue Jan 23 08:44:41 PST 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Sander de Smalen (sdesmalen-arm)

<details>
<summary>Changes</summary>

Since https://github.com/ARM-software/acle/pull/276 the ACLE defines attributes to better describe the use of a given SME state.

Previously the attributes merely described the possibility of it being 'shared' or 'preserved', whereas the new attributes have more semantics and also describe how the data flows through the program.

For ZT0 we already had to add new LLVM IR attributes:
* arm_new_zt0
* arm_in_zt0
* arm_out_zt0
* arm_inout_zt0
* arm_preserves_zt0

We have now done the same for ZA, such that we add:
* arm_new_za       (previously `aarch64_pstate_za_new`)
* arm_in_za        (more specific variation of `aarch64_pstate_za_shared`)
* arm_out_za       (more specific variation of `aarch64_pstate_za_shared`)
* arm_inout_za     (more specific variation of `aarch64_pstate_za_shared`)
* arm_preserves_za (previously `aarch64_pstate_za_shared, aarch64_pstate_za_preserved`)

This explicitly removes 'pstate' from the name, because with SME2 and the new ACLE attributes there is a difference between "sharing ZA" (sharing the ZA matrix register with the caller) and "sharing PSTATE.ZA" (sharing either the ZA or ZT0 register, both part of PSTATE.ZA with the caller).

---

Patch is 60.59 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/79166.diff


28 Files Affected:

- (modified) clang/lib/CodeGen/CGBuiltin.cpp (+2-4) 
- (modified) clang/lib/CodeGen/CGCall.cpp (+8-8) 
- (modified) clang/lib/CodeGen/CodeGenModule.cpp (+1-1) 
- (modified) clang/test/CodeGen/aarch64-sme-intrinsics/aarch64-sme-attrs.cpp (+9-9) 
- (modified) clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_zero.c (+2-2) 
- (modified) clang/test/Modules/aarch64-sme-keywords.cppm (+5-5) 
- (modified) llvm/docs/AArch64SME.rst (+21-16) 
- (modified) llvm/lib/IR/Verifier.cpp (+8-11) 
- (modified) llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp (+1-1) 
- (modified) llvm/lib/Target/AArch64/SMEABIPass.cpp (+4-5) 
- (modified) llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp (+19-13) 
- (modified) llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h (+24-9) 
- (modified) llvm/test/CodeGen/AArch64/sme-disable-gisel-fisel.ll (+5-5) 
- (modified) llvm/test/CodeGen/AArch64/sme-lazy-save-call-remarks.ll (+3-3) 
- (modified) llvm/test/CodeGen/AArch64/sme-lazy-save-call.ll (+4-4) 
- (modified) llvm/test/CodeGen/AArch64/sme-new-za-function.ll (+4-4) 
- (modified) llvm/test/CodeGen/AArch64/sme-shared-za-interface.ll (+2-2) 
- (modified) llvm/test/CodeGen/AArch64/sme-zt0-state.ll (+11-11) 
- (modified) llvm/test/Transforms/Inline/AArch64/sme-pstateza-attrs.ll (+11-11) 
- (modified) llvm/test/Verifier/sme-attributes.ll (+28-4) 
- (modified) llvm/unittests/Target/AArch64/SMEAttributesTest.cpp (+77-39) 
- (modified) mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td (+25-10) 
- (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td (+3-1) 
- (modified) mlir/lib/Target/LLVMIR/ModuleImport.cpp (+13-8) 
- (modified) mlir/lib/Target/LLVMIR/ModuleTranslation.cpp (+8-4) 
- (modified) mlir/test/Dialect/ArmSME/enable-arm-za.mlir (+12-4) 
- (modified) mlir/test/Target/LLVMIR/Import/function-attributes.ll (+19-7) 
- (modified) mlir/test/Target/LLVMIR/llvmir.mlir (+19-5) 


``````````diff
diff --git a/clang/lib/CodeGen/CGBuiltin.cpp b/clang/lib/CodeGen/CGBuiltin.cpp
index 7ef764b8e1ac80..ce90a36f5ac37f 100644
--- a/clang/lib/CodeGen/CGBuiltin.cpp
+++ b/clang/lib/CodeGen/CGBuiltin.cpp
@@ -10676,10 +10676,8 @@ Value *CodeGenFunction::EmitAArch64BuiltinExpr(unsigned BuiltinID,
         llvm::FunctionType::get(StructType::get(CGM.Int64Ty, CGM.Int64Ty), {},
                                 false),
         "__arm_sme_state"));
-    auto Attrs =
-        AttributeList()
-            .addFnAttribute(getLLVMContext(), "aarch64_pstate_sm_compatible")
-            .addFnAttribute(getLLVMContext(), "aarch64_pstate_za_preserved");
+    auto Attrs = AttributeList().addFnAttribute(getLLVMContext(),
+                                                "aarch64_pstate_sm_compatible");
     CI->setAttributes(Attrs);
     CI->setCallingConv(
         llvm::CallingConv::
diff --git a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp
index 28c211aa631e4d..657666c9bda4ec 100644
--- a/clang/lib/CodeGen/CGCall.cpp
+++ b/clang/lib/CodeGen/CGCall.cpp
@@ -1774,14 +1774,14 @@ static void AddAttributesFromFunctionProtoType(ASTContext &Ctx,
     FuncAttrs.addAttribute("aarch64_pstate_sm_compatible");
 
   // ZA
-  if (FunctionType::getArmZAState(SMEBits) == FunctionType::ARM_Out ||
-      FunctionType::getArmZAState(SMEBits) == FunctionType::ARM_InOut)
-    FuncAttrs.addAttribute("aarch64_pstate_za_shared");
-  if (FunctionType::getArmZAState(SMEBits) == FunctionType::ARM_Preserves ||
-      FunctionType::getArmZAState(SMEBits) == FunctionType::ARM_In) {
-    FuncAttrs.addAttribute("aarch64_pstate_za_shared");
-    FuncAttrs.addAttribute("aarch64_pstate_za_preserved");
-  }
+  if (FunctionType::getArmZAState(SMEBits) == FunctionType::ARM_Preserves)
+    FuncAttrs.addAttribute("aarch64_preserves_za");
+  if (FunctionType::getArmZAState(SMEBits) == FunctionType::ARM_In)
+    FuncAttrs.addAttribute("aarch64_in_za");
+  if (FunctionType::getArmZAState(SMEBits) == FunctionType::ARM_Out)
+    FuncAttrs.addAttribute("aarch64_out_za");
+  if (FunctionType::getArmZAState(SMEBits) == FunctionType::ARM_InOut)
+    FuncAttrs.addAttribute("aarch64_inout_za");
 
   // ZT0
   if (FunctionType::getArmZT0State(SMEBits) == FunctionType::ARM_Preserves)
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 1280bcd36de943..83160ae8b3ec4a 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -2414,7 +2414,7 @@ void CodeGenModule::SetLLVMFunctionAttributesForDefinition(const Decl *D,
 
   if (auto *Attr = D->getAttr<ArmNewAttr>()) {
     if (Attr->isNewZA())
-      B.addAttribute("aarch64_pstate_za_new");
+      B.addAttribute("aarch64_new_za");
     if (Attr->isNewZT0())
       B.addAttribute("aarch64_new_zt0");
   }
diff --git a/clang/test/CodeGen/aarch64-sme-intrinsics/aarch64-sme-attrs.cpp b/clang/test/CodeGen/aarch64-sme-intrinsics/aarch64-sme-attrs.cpp
index f69703a8a7d89b..fdd2de11365dd4 100644
--- a/clang/test/CodeGen/aarch64-sme-intrinsics/aarch64-sme-attrs.cpp
+++ b/clang/test/CodeGen/aarch64-sme-intrinsics/aarch64-sme-attrs.cpp
@@ -284,20 +284,20 @@ int test_variadic_template() __arm_inout("za") {
 // CHECK: attributes #[[SM_COMPATIBLE]] = { mustprogress noinline nounwind "aarch64_pstate_sm_compatible" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
 // CHECK: attributes #[[SM_COMPATIBLE_DECL]] = { "aarch64_pstate_sm_compatible" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
 // CHECK: attributes #[[SM_BODY]] = { mustprogress noinline nounwind "aarch64_pstate_sm_body" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
-// CHECK: attributes #[[ZA_SHARED]] = { mustprogress noinline nounwind "aarch64_pstate_za_shared" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
-// CHECK: attributes #[[ZA_SHARED_DECL]] = { "aarch64_pstate_za_shared" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
-// CHECK: attributes #[[ZA_PRESERVED]] = { mustprogress noinline nounwind "aarch64_pstate_za_preserved" "aarch64_pstate_za_shared" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
-// CHECK: attributes #[[ZA_PRESERVED_DECL]] = { "aarch64_pstate_za_preserved" "aarch64_pstate_za_shared" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
-// CHECK: attributes #[[ZA_NEW]] = { mustprogress noinline nounwind "aarch64_pstate_za_new" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
+// CHECK: attributes #[[ZA_SHARED]] = { mustprogress noinline nounwind "aarch64_inout_za" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
+// CHECK: attributes #[[ZA_SHARED_DECL]] = { "aarch64_inout_za" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
+// CHECK: attributes #[[ZA_PRESERVED]] = { mustprogress noinline nounwind "aarch64_preserves_za" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
+// CHECK: attributes #[[ZA_PRESERVED_DECL]] = { "aarch64_preserves_za" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
+// CHECK: attributes #[[ZA_NEW]] = { mustprogress noinline nounwind "aarch64_new_za" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
 // CHECK: attributes #[[NORMAL_DEF]] = { mustprogress noinline nounwind "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+bf16,+sme" }
 // CHECK: attributes #[[SM_ENABLED_CALL]] = { "aarch64_pstate_sm_enabled" }
 // CHECK: attributes #[[SM_COMPATIBLE_CALL]] = { "aarch64_pstate_sm_compatible" }
 // CHECK: attributes #[[SM_BODY_CALL]] = { "aarch64_pstate_sm_body" }
-// CHECK: attributes #[[ZA_SHARED_CALL]] = { "aarch64_pstate_za_shared" }
-// CHECK: attributes #[[ZA_PRESERVED_CALL]] = { "aarch64_pstate_za_preserved" "aarch64_pstate_za_shared" }
+// CHECK: attributes #[[ZA_SHARED_CALL]] = { "aarch64_inout_za" }
+// CHECK: attributes #[[ZA_PRESERVED_CALL]] = { "aarch64_preserves_za" }
 // CHECK: attributes #[[NOUNWIND_CALL]] = { nounwind }
 // CHECK: attributes #[[NOUNWIND_SM_ENABLED_CALL]] = { nounwind "aarch64_pstate_sm_enabled" }
 // CHECK: attributes #[[NOUNWIND_SM_COMPATIBLE_CALL]] = { nounwind "aarch64_pstate_sm_compatible" }
-// CHECK: attributes #[[NOUNWIND_ZA_SHARED_CALL]] = { nounwind "aarch64_pstate_za_shared" }
-// CHECK: attributes #[[NOUNWIND_ZA_PRESERVED_CALL]] = { nounwind "aarch64_pstate_za_preserved" "aarch64_pstate_za_shared" }
+// CHECK: attributes #[[NOUNWIND_ZA_SHARED_CALL]] = { nounwind "aarch64_inout_za" }
+// CHECK: attributes #[[NOUNWIND_ZA_PRESERVED_CALL]] = { nounwind "aarch64_preserves_za" }
 
diff --git a/clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_zero.c b/clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_zero.c
index 7f56941108828a..9963c0e48b8e7b 100644
--- a/clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_zero.c
+++ b/clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_zero.c
@@ -55,13 +55,13 @@ void test_svzero_mask_za_2(void) __arm_inout("za") {
 }
 
 // CHECK-C-LABEL: define dso_local void @test_svzero_za(
-// CHECK-C-SAME: ) local_unnamed_addr #[[ATTR0]] {
+// CHECK-C-SAME: ) local_unnamed_addr #[[ATTR2:[0-9]+]] {
 // CHECK-C-NEXT:  entry:
 // CHECK-C-NEXT:    tail call void @llvm.aarch64.sme.zero(i32 255)
 // CHECK-C-NEXT:    ret void
 //
 // CHECK-CXX-LABEL: define dso_local void @_Z14test_svzero_zav(
-// CHECK-CXX-SAME: ) local_unnamed_addr #[[ATTR0]] {
+// CHECK-CXX-SAME: ) local_unnamed_addr #[[ATTR2:[0-9]+]] {
 // CHECK-CXX-NEXT:  entry:
 // CHECK-CXX-NEXT:    tail call void @llvm.aarch64.sme.zero(i32 255)
 // CHECK-CXX-NEXT:    ret void
diff --git a/clang/test/Modules/aarch64-sme-keywords.cppm b/clang/test/Modules/aarch64-sme-keywords.cppm
index df4dd32b16cff7..759701a633cebe 100644
--- a/clang/test/Modules/aarch64-sme-keywords.cppm
+++ b/clang/test/Modules/aarch64-sme-keywords.cppm
@@ -43,14 +43,14 @@ import A;
 //
 // CHECK:declare void @_ZW1A22f_streaming_compatiblev() #[[STREAMING_COMPATIBLE_DECL:[0-9]+]]
 //
-// CHECK-DAG: attributes #[[SHARED_ZA_DEF]] = {{{.*}} "aarch64_pstate_za_shared" {{.*}}}
-// CHECK-DAG: attributes #[[SHARED_ZA_DECL]] = {{{.*}} "aarch64_pstate_za_shared" {{.*}}}
-// CHECK-DAG: attributes #[[PRESERVES_ZA_DECL]] = {{{.*}} "aarch64_pstate_za_preserved" {{.*}}}
+// CHECK-DAG: attributes #[[SHARED_ZA_DEF]] = {{{.*}} "aarch64_inout_za" {{.*}}}
+// CHECK-DAG: attributes #[[SHARED_ZA_DECL]] = {{{.*}} "aarch64_inout_za" {{.*}}}
+// CHECK-DAG: attributes #[[PRESERVES_ZA_DECL]] = {{{.*}} "aarch64_preserves_za" {{.*}}}
 // CHECK-DAG: attributes #[[NORMAL_DEF]] = {{{.*}}}
 // CHECK-DAG: attributes #[[STREAMING_DECL]] = {{{.*}} "aarch64_pstate_sm_enabled" {{.*}}}
 // CHECK-DAG: attributes #[[STREAMING_COMPATIBLE_DECL]] = {{{.*}} "aarch64_pstate_sm_compatible" {{.*}}}
-// CHECK-DAG: attributes #[[SHARED_ZA_USE]] = { "aarch64_pstate_za_shared" }
-// CHECK-DAG: attributes #[[PRESERVES_ZA_USE]] = { "aarch64_pstate_za_preserved" "aarch64_pstate_za_shared" }
+// CHECK-DAG: attributes #[[SHARED_ZA_USE]] = { "aarch64_inout_za" }
+// CHECK-DAG: attributes #[[PRESERVES_ZA_USE]] = { "aarch64_preserves_za" }
 // CHECK-DAG: attributes #[[STREAMING_USE]] = { "aarch64_pstate_sm_enabled" }
 // CHECK-DAG: attributes #[[STREAMING_COMPATIBLE_USE]] = { "aarch64_pstate_sm_compatible" }
 
diff --git a/llvm/docs/AArch64SME.rst b/llvm/docs/AArch64SME.rst
index 63573bf91eacb5..4b6fa5e10f880a 100644
--- a/llvm/docs/AArch64SME.rst
+++ b/llvm/docs/AArch64SME.rst
@@ -22,26 +22,32 @@ Below we describe the LLVM IR attributes and their relation to the C/C++
 level ACLE attributes:
 
 ``aarch64_pstate_sm_enabled``
-    is used for functions with ``__attribute__((arm_streaming))``
+    is used for functions with ``__arm_streaming``
 
 ``aarch64_pstate_sm_compatible``
-    is used for functions with ``__attribute__((arm_streaming_compatible))``
+    is used for functions with ``__arm_streaming_compatible``
 
 ``aarch64_pstate_sm_body``
-  is used for functions with ``__attribute__((arm_locally_streaming))`` and is
+  is used for functions with ``__arm_locally_streaming`` and is
   only valid on function definitions (not declarations)
 
-``aarch64_pstate_za_new``
-  is used for functions with ``__attribute__((arm_new_za))``
+``aarch64_pstate_new_za``
+  is used for functions with ``__arm_new("za")``
 
-``aarch64_pstate_za_shared``
-  is used for functions with ``__attribute__((arm_shared_za))``
+``aarch64_in_za``
+  is used for functions with ``__arm_in("za")``
 
-``aarch64_pstate_za_preserved``
-  is used for functions with ``__attribute__((arm_preserves_za))``
+``aarch64_out_za``
+  is used for functions with ``__arm_out("za")``
+
+``aarch64_inout_za``
+  is used for functions with ``__arm_inout("za")``
+
+``aarch64_preserves_za``
+  is used for functions with ``__arm_preserves("za")``
 
 ``aarch64_expanded_pstate_za``
-  is used for functions with ``__attribute__((arm_new_za))``
+  is used for functions with ``__arm_new_za``
 
 Clang must ensure that the above attributes are added both to the
 function's declaration/definition as well as to their call-sites. This is
@@ -89,11 +95,10 @@ Restrictions on attributes
 * It is not allowed for a function to be decorated with both
   ``aarch64_pstate_sm_compatible`` and ``aarch64_pstate_sm_enabled``.
 
-* It is not allowed for a function to be decorated with both
-  ``aarch64_pstate_za_new`` and ``aarch64_pstate_za_preserved``.
-
-* It is not allowed for a function to be decorated with both
-  ``aarch64_pstate_za_new`` and ``aarch64_pstate_za_shared``.
+* It is not allowed for a function to be decorated with more than one of the
+  following attributes:
+  ``aarch64_new_za``, ``aarch64_in_za``, ``aarch64_out_za``, ``aarch64_inout_za``,
+  ``aarch64_preserves_za``.
 
 These restrictions also apply in the higher level SME ACLE, which means we can
 emit diagnostics in Clang to signal users about incorrect behaviour.
@@ -426,7 +431,7 @@ to toggle PSTATE.ZA using intrinsics. This also makes it simpler to setup a
 lazy-save mechanism for calls to private-ZA functions (i.e. functions that may
 either directly or indirectly clobber ZA state).
 
-For the purpose of handling functions marked with ``aarch64_pstate_za_new``,
+For the purpose of handling functions marked with ``aarch64_new_za``,
 we have introduced a new LLVM IR pass (SMEABIPass) that is run just before
 SelectionDAG. Any such functions dealt with by this pass are marked with
 ``aarch64_expanded_pstate_za``.
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 91cf91fbc788bd..53d923ced8d2e3 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -2155,17 +2155,14 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
            V);
   }
 
-  if (Attrs.hasFnAttr("aarch64_pstate_za_new")) {
-    Check(!Attrs.hasFnAttr("aarch64_pstate_za_preserved"),
-           "Attributes 'aarch64_pstate_za_new and aarch64_pstate_za_preserved' "
-           "are incompatible!",
-           V);
-
-    Check(!Attrs.hasFnAttr("aarch64_pstate_za_shared"),
-           "Attributes 'aarch64_pstate_za_new and aarch64_pstate_za_shared' "
-           "are incompatible!",
-           V);
-  }
+  Check(
+      (Attrs.hasFnAttr("aarch64_new_za") + Attrs.hasFnAttr("aarch64_in_za") +
+       Attrs.hasFnAttr("aarch64_inout_za") +
+       Attrs.hasFnAttr("aarch64_out_za") +
+       Attrs.hasFnAttr("aarch64_preserves_za")) <= 1,
+      "Attributes 'aarch64_new_za', 'aarch64_in_za', 'aarch64_out_za', "
+      "'aarch64_inout_za' and 'aarch64_preserves_za' are mutually exclusive",
+      V);
 
   Check(
       (Attrs.hasFnAttr("aarch64_new_zt0") + Attrs.hasFnAttr("aarch64_in_zt0") +
diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index d611338fc268f9..fa65c8103697b1 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
@@ -235,7 +235,7 @@ bool AArch64TTIImpl::areInlineCompatible(const Function *Caller,
                                          const Function *Callee) const {
   SMEAttrs CallerAttrs(*Caller);
   SMEAttrs CalleeAttrs(*Callee);
-  if (CalleeAttrs.hasNewZABody())
+  if (CalleeAttrs.isNewZA())
     return false;
 
   if (CallerAttrs.requiresLazySave(CalleeAttrs) ||
diff --git a/llvm/lib/Target/AArch64/SMEABIPass.cpp b/llvm/lib/Target/AArch64/SMEABIPass.cpp
index 0247488ce93f1d..bee5d63c5a749b 100644
--- a/llvm/lib/Target/AArch64/SMEABIPass.cpp
+++ b/llvm/lib/Target/AArch64/SMEABIPass.cpp
@@ -62,8 +62,7 @@ void emitTPIDR2Save(Module *M, IRBuilder<> &Builder) {
       FunctionType::get(Builder.getVoidTy(), {}, /*IsVarArgs=*/false);
   auto Attrs =
       AttributeList()
-          .addFnAttribute(M->getContext(), "aarch64_pstate_sm_compatible")
-          .addFnAttribute(M->getContext(), "aarch64_pstate_za_preserved");
+          .addFnAttribute(M->getContext(), "aarch64_pstate_sm_compatible");
   FunctionCallee Callee =
       M->getOrInsertFunction("__arm_tpidr2_save", TPIDR2SaveTy, Attrs);
   CallInst *Call = Builder.CreateCall(Callee);
@@ -78,7 +77,7 @@ void emitTPIDR2Save(Module *M, IRBuilder<> &Builder) {
 }
 
 /// This function generates code at the beginning and end of a function marked
-/// with either `aarch64_pstate_za_new` or `aarch64_new_zt0`.
+/// with either `aarch64_new_za` or `aarch64_new_zt0`.
 /// At the beginning of the function, the following code is generated:
 ///  - Commit lazy-save if active   [Private-ZA Interface*]
 ///  - Enable PSTATE.ZA             [Private-ZA Interface]
@@ -133,7 +132,7 @@ bool SMEABI::updateNewStateFunctions(Module *M, Function *F,
     Builder.CreateCall(EnableZAIntr->getFunctionType(), EnableZAIntr);
   }
 
-  if (FnAttrs.hasNewZABody()) {
+  if (FnAttrs.isNewZA()) {
     Function *ZeroIntr =
         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_zero);
     Builder.CreateCall(ZeroIntr->getFunctionType(), ZeroIntr,
@@ -174,7 +173,7 @@ bool SMEABI::runOnFunction(Function &F) {
 
   bool Changed = false;
   SMEAttrs FnAttrs(F);
-  if (FnAttrs.hasNewZABody() || FnAttrs.isNewZT0())
+  if (FnAttrs.isNewZA() || FnAttrs.isNewZT0())
     Changed |= updateNewStateFunctions(M, &F, Builder, FnAttrs);
 
   return Changed;
diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
index 3ee54e5df0a13d..ac07274a36c231 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
@@ -23,13 +23,15 @@ void SMEAttrs::set(unsigned M, bool Enable) {
          "SM_Enabled and SM_Compatible are mutually exclusive");
 
   // ZA Attrs
-  assert(!(hasNewZABody() && sharesZA()) &&
-         "ZA_New and ZA_Shared are mutually exclusive");
-  assert(!(hasNewZABody() && preservesZA()) &&
-         "ZA_New and ZA_Preserved are mutually exclusive");
-  assert(!(hasNewZABody() && (Bitmask & SME_ABI_Routine)) &&
+  assert(!(isNewZA() && (Bitmask & SME_ABI_Routine)) &&
          "ZA_New and SME_ABI_Routine are mutually exclusive");
 
+  assert(
+      (!sharesZA() ||
+       (isNewZA() ^ isInZA() ^ isInOutZA() ^ isOutZA() ^ isPreservesZA())) &&
+      "Attributes 'aarch64_new_za', 'aarch64_in_za', 'aarch64_out_za', "
+      "'aarch64_inout_za' and 'aarch64_preserves_za' are mutually exclusive");
+
   // ZT0 Attrs
   assert(
       (!sharesZT0() || (isNewZT0() ^ isInZT0() ^ isInOutZT0() ^ isOutZT0() ^
@@ -49,8 +51,8 @@ SMEAttrs::SMEAttrs(StringRef FuncName) : Bitmask(0) {
   if (FuncName == "__arm_tpidr2_save" || FuncName == "__arm_sme_state")
     Bitmask |= (SMEAttrs::SM_Compatible | SMEAttrs::SME_ABI_Routine);
   if (FuncName == "__arm_tpidr2_restore")
-    Bitmask |= (SMEAttrs::SM_Compatible | SMEAttrs::ZA_Shared |
-                SMEAttrs::SME_ABI_Routine);
+    Bitmask |= SMEAttrs::SM_Compatible | encodeZAState(StateValue::In) |
+                SMEAttrs::SME_ABI_Routine;
 }
 
 SMEAttrs::SMEAttrs(const AttributeList &Attrs) {
@@ -61,12 +63,16 @@ SMEAttrs::SMEAttrs(const AttributeList &Attrs) {
     Bitmask |= SM_Compatible;
   if (Attrs.hasFnAttr("aarch64_pstate_sm_body"))
     Bitmask |= SM_Body;
-  if (Attrs.hasFnAttr("aarch64_pstate_za_shared"))
-    Bitmask |= ZA_Shared;
-  if (Attrs.hasFnAttr("aarch64_pstate_za_new"))
-    Bitmask |= ZA_New;
-  if (Attrs.hasFnAttr("aarch64_pstate_za_preserved"))
-    Bitmask |= ZA_Preserved;
+  if (Attrs.hasFnAttr("aarch64_in_za"))
+    Bitmask |= encodeZAState(StateValue::In);
+  if (Attrs.hasFnAttr("aarch64_out_za"))
+    Bitmask |= encodeZAState(StateValue::Out);
+  if (Attrs.hasFnAttr("aarch64_inout_za"))
+    Bitmask |= encodeZAState(StateValue::InOut);
+  if (Attrs.hasFnAttr("aarch64_preserves_za"))
+    Bitmask |= encodeZAState(StateValue::Preserved);
+  if (Attrs.hasFnAttr("aarch64_new_za"))
+    Bitmask |= encodeZAState(StateValue::New);
   if (Attrs.hasFnAttr("aarch64_in_zt0"))
     Bitmask |= encodeZT0State(StateValue::In);
   if (Attrs.hasFnAttr("aarch64_out_zt0"))
diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
index 27b7075a0944ff..4c7c1c9b079538 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
@@ -41,10 +41,9 @@ class SMEAttrs {
     SM_Enabled = 1 << 0,      // aarch64_pstate_sm_enabled
     SM_Compatible = 1 << 1,   // aarch64_pstate_sm_compatible
     SM_Body = 1 << 2,         // aarch64_pstate_sm_body
-    ZA_Shared = 1 << 3,       // aarch64_pstate_sm_shared
-    ZA_New = 1 << 4,          // aarch64_pstate_sm_new
-    ZA_Preserved = 1 << 5,    // aarch64_pstate_sm_preserved
-    SME_ABI_Routine = 1 << 6, // Used for SME ABI routines to avoid lazy saves
+    SME_ABI_Routine = 1 << 3, // Used for SME ABI routines to avoid lazy saves
+    ZA_Shift = 4,
+    ZA_Mask = 0b111 << ZA_Shift,
     ZT0_Shift = 7,
     ZT0_Mask = 0b111 << ZT0_Shift
   };
@@ -77,13 +76,29 @@ class SMEAttrs {
   /// streaming mode.
   bool requiresSMChange(const SMEAttrs &Callee) const;
 
-  // Interfac...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/79166


More information about the cfe-commits mailing list