[llvm] [AArch64][SME2] Add ZT0 attributes to SMEAttrs (PR #77607)

Kerry McLaughlin via llvm-commits llvm-commits at lists.llvm.org
Fri Jan 12 09:07:54 PST 2024


https://github.com/kmclaughlin-arm updated https://github.com/llvm/llvm-project/pull/77607

>From 08ca20b751f6c03282647589bf45ec408553816d Mon Sep 17 00:00:00 2001
From: Kerry McLaughlin <kerry.mclaughlin at arm.com>
Date: Tue, 9 Jan 2024 11:09:29 +0000
Subject: [PATCH 1/4] [AArch64][SME2] Add ZT0 attributes to SMEAttrs

This patch extends SMEAttrs to interpret the following new attributes,
which are mutually exclusive and apply to SME2 only:
  - aarch64_sme_zt0_in (ZT0_In)
  - aarch64_sme_zt0_out (ZT0_Out)
  - aarch64_sme_zt0_inout (ZT0_InOut)
  - aarch64_sme_zt0_new (ZT0_New)
  - aarch64_sme_zt0_preserved (ZT0_Preserved)

ZT0_In, ZT0_Out, ZT0_InOut & ZT0_Preserved are all considered to share
ZT0. These attributes will be required by later patches to determine
if ZT0 should be preserved around function calls, or cleared on entry
to the function.
---
 .../AArch64/Utils/AArch64SMEAttributes.cpp    |  18 +++
 .../AArch64/Utils/AArch64SMEAttributes.h      |  32 ++++--
 .../Target/AArch64/SMEAttributesTest.cpp      | 103 ++++++++++++++++++
 3 files changed, 145 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
index 0082b4017986c6..82b07419a244a0 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
@@ -18,8 +18,11 @@ void SMEAttrs::set(unsigned M, bool Enable) {
   else
     Bitmask &= ~M;
 
+  // Streaming Mode Attrs
   assert(!(hasStreamingInterface() && hasStreamingCompatibleInterface()) &&
          "SM_Enabled and SM_Compatible are mutually exclusive");
+
+  // ZA Attrs
   assert(!(hasNewZABody() && hasSharedZAInterface()) &&
          "ZA_New and ZA_Shared are mutually exclusive");
   assert(!(hasNewZABody() && preservesZA()) &&
@@ -28,6 +31,11 @@ void SMEAttrs::set(unsigned M, bool Enable) {
          "ZA_New and ZA_NoLazySave are mutually exclusive");
   assert(!(hasSharedZAInterface() && (Bitmask & ZA_NoLazySave)) &&
          "ZA_Shared and ZA_NoLazySave are mutually exclusive");
+
+  // ZT0 Attrs
+  assert((!sharesZT0() || (hasNewZT0Body() ^ isZT0In() ^ isZT0InOut() ^
+                           isZT0Out() ^ preservesZT0())) &&
+         "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
 }
 
 SMEAttrs::SMEAttrs(const CallBase &CB) {
@@ -60,6 +68,16 @@ SMEAttrs::SMEAttrs(const AttributeList &Attrs) {
     Bitmask |= ZA_New;
   if (Attrs.hasFnAttr("aarch64_pstate_za_preserved"))
     Bitmask |= ZA_Preserved;
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_in"))
+    Bitmask |= ZT0_In;
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_out"))
+    Bitmask |= ZT0_Out;
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_inout"))
+    Bitmask |= ZT0_InOut;
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_preserved"))
+    Bitmask |= ZT0_Preserved;
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_new"))
+    Bitmask |= ZT0_New;
 }
 
 std::optional<bool>
diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
index e766b778b54102..88d67ad4b51f53 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
@@ -29,14 +29,19 @@ class SMEAttrs {
   // Enum with bitmasks for each individual SME feature.
   enum Mask {
     Normal = 0,
-    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
-    ZA_NoLazySave = 1 << 6, // Used for SME ABI routines to avoid lazy saves
-    All = ZA_Preserved - 1
+    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
+    ZA_NoLazySave = 1 << 6,  // Used for SME ABI routines to avoid lazy saves
+    ZT0_New = 1 << 7,        // aarch64_sme_zt0_new
+    ZT0_In = 1 << 8,         // aarch64_sme_zt0_in
+    ZT0_Out = 1 << 9,        // aarch64_sme_zt0_out
+    ZT0_InOut = 1 << 10,     // aarch64_sme_zt0_inout
+    ZT0_Preserved = 1 << 11, // aarch64_sme_zt0_preserved
+    All = ZT0_Preserved - 1
   };
 
   SMEAttrs(unsigned Mask = Normal) : Bitmask(0) { set(Mask); }
@@ -86,6 +91,17 @@ class SMEAttrs {
     return hasZAState() && Callee.hasPrivateZAInterface() &&
            !(Callee.Bitmask & ZA_NoLazySave);
   }
+
+  // Interfaces to query ZT0 State
+  bool hasNewZT0Body() const { return Bitmask & ZT0_New; }
+  bool isZT0In() const { return Bitmask & ZT0_In; }
+  bool isZT0Out() const { return Bitmask & ZT0_Out; }
+  bool isZT0InOut() const { return Bitmask & ZT0_InOut; }
+  bool preservesZT0() const { return Bitmask & ZT0_Preserved; }
+  bool sharesZT0() const {
+    return Bitmask & (ZT0_In | ZT0_Out | ZT0_InOut | ZT0_Preserved);
+  }
+  bool hasZT0State() const { return hasNewZT0Body() || sharesZT0(); }
 };
 
 } // namespace llvm
diff --git a/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp b/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
index 7780c71bbc00e0..9c88ec53589a73 100644
--- a/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
+++ b/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
@@ -50,6 +50,37 @@ TEST(SMEAttributes, Constructors) {
                       ->getFunction("foo"))
                   .preservesZA());
 
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_in\"")
+                      ->getFunction("foo"))
+                  .isZT0In());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_out\"")
+                      ->getFunction("foo"))
+                  .isZT0Out());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_inout\"")
+                      ->getFunction("foo"))
+                  .isZT0InOut());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_in\"")
+                      ->getFunction("foo"))
+                  .sharesZT0());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_out\"")
+                      ->getFunction("foo"))
+                  .sharesZT0());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_inout\"")
+                      ->getFunction("foo"))
+                  .sharesZT0());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_preserved\"")
+                      ->getFunction("foo"))
+                  .sharesZT0());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_preserved\"")
+                      ->getFunction("foo"))
+                  .preservesZT0());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_new\"")
+                      ->getFunction("foo"))
+                  .hasNewZT0Body());
+
   // Invalid combinations.
   EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled | SA::SM_Compatible),
                      "SM_Enabled and SM_Compatible are mutually exclusive");
@@ -58,6 +89,29 @@ TEST(SMEAttributes, Constructors) {
   EXPECT_DEBUG_DEATH(SA(SA::ZA_New | SA::ZA_Preserved),
                      "ZA_New and ZA_Preserved are mutually exclusive");
 
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_In),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_Out),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_InOut),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_Preserved),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_In | SA::ZT0_Out),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_In | SA::ZT0_InOut),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_Out | SA::ZT0_InOut),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_In),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_Out),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_InOut),
+                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+
   // Test that the set() methods equally check validity.
   EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled).set(SA::SM_Compatible),
                      "SM_Enabled and SM_Compatible are mutually exclusive");
@@ -95,6 +149,55 @@ TEST(SMEAttributes, Basics) {
   ASSERT_FALSE(SA(SA::Normal).hasNewZABody());
   ASSERT_FALSE(SA(SA::Normal).hasZAState());
   ASSERT_FALSE(SA(SA::Normal).preservesZA());
+
+  // Test ZT0 State interfaces
+  ASSERT_TRUE(SA(SA::ZT0_In).isZT0In());
+  ASSERT_FALSE(SA(SA::ZT0_In).isZT0Out());
+  ASSERT_FALSE(SA(SA::ZT0_In).isZT0InOut());
+  ASSERT_FALSE(SA(SA::ZT0_In).preservesZT0());
+  ASSERT_FALSE(SA(SA::ZT0_In).hasNewZT0Body());
+  ASSERT_TRUE(SA(SA::ZT0_In).sharesZT0());
+  ASSERT_TRUE(SA(SA::ZT0_In).hasZT0State());
+
+  ASSERT_TRUE(SA(SA::ZT0_Out).isZT0Out());
+  ASSERT_FALSE(SA(SA::ZT0_Out).isZT0In());
+  ASSERT_FALSE(SA(SA::ZT0_Out).isZT0InOut());
+  ASSERT_FALSE(SA(SA::ZT0_Out).preservesZT0());
+  ASSERT_FALSE(SA(SA::ZT0_Out).hasNewZT0Body());
+  ASSERT_TRUE(SA(SA::ZT0_Out).sharesZT0());
+  ASSERT_TRUE(SA(SA::ZT0_Out).hasZT0State());
+
+  ASSERT_TRUE(SA(SA::ZT0_InOut).isZT0InOut());
+  ASSERT_FALSE(SA(SA::ZT0_InOut).isZT0In());
+  ASSERT_FALSE(SA(SA::ZT0_InOut).isZT0Out());
+  ASSERT_FALSE(SA(SA::ZT0_InOut).preservesZT0());
+  ASSERT_FALSE(SA(SA::ZT0_InOut).hasNewZT0Body());
+  ASSERT_TRUE(SA(SA::ZT0_InOut).sharesZT0());
+  ASSERT_TRUE(SA(SA::ZT0_InOut).hasZT0State());
+
+  ASSERT_TRUE(SA(SA::ZT0_Preserved).preservesZT0());
+  ASSERT_FALSE(SA(SA::ZT0_Preserved).isZT0In());
+  ASSERT_FALSE(SA(SA::ZT0_Preserved).isZT0Out());
+  ASSERT_FALSE(SA(SA::ZT0_Preserved).isZT0InOut());
+  ASSERT_FALSE(SA(SA::ZT0_Preserved).hasNewZT0Body());
+  ASSERT_TRUE(SA(SA::ZT0_Preserved).sharesZT0());
+  ASSERT_TRUE(SA(SA::ZT0_Preserved).hasZT0State());
+
+  ASSERT_TRUE(SA(SA::ZT0_New).hasNewZT0Body());
+  ASSERT_FALSE(SA(SA::ZT0_New).isZT0In());
+  ASSERT_FALSE(SA(SA::ZT0_New).isZT0Out());
+  ASSERT_FALSE(SA(SA::ZT0_New).isZT0InOut());
+  ASSERT_FALSE(SA(SA::ZT0_New).preservesZT0());
+  ASSERT_FALSE(SA(SA::ZT0_New).sharesZT0());
+  ASSERT_TRUE(SA(SA::ZT0_New).hasZT0State());
+
+  ASSERT_FALSE(SA(SA::Normal).isZT0In());
+  ASSERT_FALSE(SA(SA::Normal).isZT0Out());
+  ASSERT_FALSE(SA(SA::Normal).isZT0InOut());
+  ASSERT_FALSE(SA(SA::Normal).preservesZT0());
+  ASSERT_FALSE(SA(SA::Normal).hasNewZT0Body());
+  ASSERT_FALSE(SA(SA::Normal).sharesZT0());
+  ASSERT_FALSE(SA(SA::Normal).hasZT0State());
 }
 
 TEST(SMEAttributes, Transitions) {

>From da7076c59250b6b96a79793e2bf26c855540d058 Mon Sep 17 00:00:00 2001
From: Kerry McLaughlin <kerry.mclaughlin at arm.com>
Date: Wed, 10 Jan 2024 14:44:38 +0000
Subject: [PATCH 2/4] - Added a sharedZA() function which returns true if
 aarch64_pstate_za_shared is set - Changed hasSharedZAInterface() to return
 true if either sharedZA() or sharedZT0() is true - Changed wording of the
 assert checking for mutually exclusive ZT0 attributes

---
 .../AArch64/Utils/AArch64SMEAttributes.cpp    |  7 +--
 .../AArch64/Utils/AArch64SMEAttributes.h      |  7 ++-
 .../Target/AArch64/SMEAttributesTest.cpp      | 52 +++++++++++++++----
 3 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
index 82b07419a244a0..6061ec607ce000 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
@@ -23,19 +23,20 @@ void SMEAttrs::set(unsigned M, bool Enable) {
          "SM_Enabled and SM_Compatible are mutually exclusive");
 
   // ZA Attrs
-  assert(!(hasNewZABody() && hasSharedZAInterface()) &&
+  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 & ZA_NoLazySave)) &&
          "ZA_New and ZA_NoLazySave are mutually exclusive");
-  assert(!(hasSharedZAInterface() && (Bitmask & ZA_NoLazySave)) &&
+  assert(!(sharesZA() && (Bitmask & ZA_NoLazySave)) &&
          "ZA_Shared and ZA_NoLazySave are mutually exclusive");
 
   // ZT0 Attrs
   assert((!sharesZT0() || (hasNewZT0Body() ^ isZT0In() ^ isZT0InOut() ^
                            isZT0Out() ^ preservesZT0())) &&
-         "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+         "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved are all "
+         "mutually exclusive");
 }
 
 SMEAttrs::SMEAttrs(const CallBase &CB) {
diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
index 88d67ad4b51f53..9d92fbc29eab02 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
@@ -81,12 +81,11 @@ class SMEAttrs {
 
   // Interfaces to query PSTATE.ZA
   bool hasNewZABody() const { return Bitmask & ZA_New; }
-  bool hasSharedZAInterface() const { return Bitmask & ZA_Shared; }
+  bool sharesZA() const { return Bitmask & ZA_Shared; }
+  bool hasSharedZAInterface() const { return sharesZA() || sharesZT0(); }
   bool hasPrivateZAInterface() const { return !hasSharedZAInterface(); }
   bool preservesZA() const { return Bitmask & ZA_Preserved; }
-  bool hasZAState() const {
-    return hasNewZABody() || hasSharedZAInterface();
-  }
+  bool hasZAState() const { return hasNewZABody() || sharesZA(); }
   bool requiresLazySave(const SMEAttrs &Callee) const {
     return hasZAState() && Callee.hasPrivateZAInterface() &&
            !(Callee.Bitmask & ZA_NoLazySave);
diff --git a/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp b/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
index 9c88ec53589a73..c21359d0c4d195 100644
--- a/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
+++ b/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
@@ -38,6 +38,10 @@ TEST(SMEAttributes, Constructors) {
               ->getFunction("foo"))
           .hasStreamingCompatibleInterface());
 
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_pstate_za_shared\"")
+                      ->getFunction("foo"))
+                  .sharesZA());
+
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_pstate_za_shared\"")
                       ->getFunction("foo"))
                   .hasSharedZAInterface());
@@ -73,6 +77,19 @@ TEST(SMEAttributes, Constructors) {
                       ->getFunction("foo"))
                   .sharesZT0());
 
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_in\"")
+                      ->getFunction("foo"))
+                  .hasSharedZAInterface());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_out\"")
+                      ->getFunction("foo"))
+                  .hasSharedZAInterface());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_inout\"")
+                      ->getFunction("foo"))
+                  .hasSharedZAInterface());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_preserved\"")
+                      ->getFunction("foo"))
+                  .hasSharedZAInterface());
+
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_preserved\"")
                       ->getFunction("foo"))
                   .preservesZT0());
@@ -80,6 +97,9 @@ TEST(SMEAttributes, Constructors) {
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_new\"")
                       ->getFunction("foo"))
                   .hasNewZT0Body());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_new\"")
+                      ->getFunction("foo"))
+                  .hasPrivateZAInterface());
 
   // Invalid combinations.
   EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled | SA::SM_Compatible),
@@ -90,27 +110,37 @@ TEST(SMEAttributes, Constructors) {
                      "ZA_New and ZA_Preserved are mutually exclusive");
 
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_In),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_Out),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_InOut),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_Preserved),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
 
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_In | SA::ZT0_Out),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_In | SA::ZT0_InOut),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_Out | SA::ZT0_InOut),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
 
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_In),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_Out),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_InOut),
-                     "ZT0_New,In,Out,InOut,Preserved are mutually exclusive");
+                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
+                     "are all mutually exclusive");
 
   // Test that the set() methods equally check validity.
   EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled).set(SA::SM_Compatible),
@@ -135,7 +165,7 @@ TEST(SMEAttributes, Basics) {
 
   // Test PSTATE.ZA interfaces.
   ASSERT_FALSE(SA(SA::ZA_Shared).hasPrivateZAInterface());
-  ASSERT_TRUE(SA(SA::ZA_Shared).hasSharedZAInterface());
+  ASSERT_TRUE(SA(SA::ZA_Shared).sharesZA());
   ASSERT_TRUE(SA(SA::ZA_Shared).hasZAState());
   ASSERT_FALSE(SA(SA::ZA_Shared).preservesZA());
   ASSERT_TRUE(SA(SA::ZA_Shared | SA::ZA_Preserved).preservesZA());

>From f39fbb8d5c44abfb55fbb126390948adfefa798a Mon Sep 17 00:00:00 2001
From: Kerry McLaughlin <kerry.mclaughlin at arm.com>
Date: Wed, 10 Jan 2024 16:03:51 +0000
Subject: [PATCH 3/4] - Removed duplicate tests from SMEAttributesTest

---
 .../Target/AArch64/SMEAttributesTest.cpp      | 64 +++++++------------
 1 file changed, 23 insertions(+), 41 deletions(-)

diff --git a/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp b/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
index c21359d0c4d195..264b0fdaeb56c5 100644
--- a/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
+++ b/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
@@ -63,43 +63,12 @@ TEST(SMEAttributes, Constructors) {
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_inout\"")
                       ->getFunction("foo"))
                   .isZT0InOut());
-
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_in\"")
-                      ->getFunction("foo"))
-                  .sharesZT0());
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_out\"")
-                      ->getFunction("foo"))
-                  .sharesZT0());
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_inout\"")
-                      ->getFunction("foo"))
-                  .sharesZT0());
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_preserved\"")
-                      ->getFunction("foo"))
-                  .sharesZT0());
-
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_in\"")
-                      ->getFunction("foo"))
-                  .hasSharedZAInterface());
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_out\"")
-                      ->getFunction("foo"))
-                  .hasSharedZAInterface());
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_inout\"")
-                      ->getFunction("foo"))
-                  .hasSharedZAInterface());
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_preserved\"")
-                      ->getFunction("foo"))
-                  .hasSharedZAInterface());
-
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_preserved\"")
                       ->getFunction("foo"))
                   .preservesZT0());
-
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_new\"")
                       ->getFunction("foo"))
                   .hasNewZT0Body());
-  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_new\"")
-                      ->getFunction("foo"))
-                  .hasPrivateZAInterface());
 
   // Invalid combinations.
   EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled | SA::SM_Compatible),
@@ -111,36 +80,36 @@ TEST(SMEAttributes, Constructors) {
 
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_In),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_Out),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_InOut),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_Preserved),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
 
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_In | SA::ZT0_Out),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_In | SA::ZT0_InOut),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_Out | SA::ZT0_InOut),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
 
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_In),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_Out),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
   EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_InOut),
                      "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all mutually exclusive");
+                     "are all \" \"mutually exclusive");
 
   // Test that the set() methods equally check validity.
   EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled).set(SA::SM_Compatible),
@@ -165,17 +134,20 @@ TEST(SMEAttributes, Basics) {
 
   // Test PSTATE.ZA interfaces.
   ASSERT_FALSE(SA(SA::ZA_Shared).hasPrivateZAInterface());
+  ASSERT_TRUE(SA(SA::ZA_Shared).hasSharedZAInterface());
   ASSERT_TRUE(SA(SA::ZA_Shared).sharesZA());
   ASSERT_TRUE(SA(SA::ZA_Shared).hasZAState());
   ASSERT_FALSE(SA(SA::ZA_Shared).preservesZA());
   ASSERT_TRUE(SA(SA::ZA_Shared | SA::ZA_Preserved).preservesZA());
 
   ASSERT_TRUE(SA(SA::ZA_New).hasPrivateZAInterface());
+  ASSERT_FALSE(SA(SA::ZA_New).hasSharedZAInterface());
   ASSERT_TRUE(SA(SA::ZA_New).hasNewZABody());
   ASSERT_TRUE(SA(SA::ZA_New).hasZAState());
   ASSERT_FALSE(SA(SA::ZA_New).preservesZA());
 
   ASSERT_TRUE(SA(SA::Normal).hasPrivateZAInterface());
+  ASSERT_FALSE(SA(SA::Normal).hasSharedZAInterface());
   ASSERT_FALSE(SA(SA::Normal).hasNewZABody());
   ASSERT_FALSE(SA(SA::Normal).hasZAState());
   ASSERT_FALSE(SA(SA::Normal).preservesZA());
@@ -188,6 +160,8 @@ TEST(SMEAttributes, Basics) {
   ASSERT_FALSE(SA(SA::ZT0_In).hasNewZT0Body());
   ASSERT_TRUE(SA(SA::ZT0_In).sharesZT0());
   ASSERT_TRUE(SA(SA::ZT0_In).hasZT0State());
+  ASSERT_TRUE(SA(SA::ZT0_In).hasSharedZAInterface());
+  ASSERT_FALSE(SA(SA::ZT0_In).hasPrivateZAInterface());
 
   ASSERT_TRUE(SA(SA::ZT0_Out).isZT0Out());
   ASSERT_FALSE(SA(SA::ZT0_Out).isZT0In());
@@ -196,6 +170,8 @@ TEST(SMEAttributes, Basics) {
   ASSERT_FALSE(SA(SA::ZT0_Out).hasNewZT0Body());
   ASSERT_TRUE(SA(SA::ZT0_Out).sharesZT0());
   ASSERT_TRUE(SA(SA::ZT0_Out).hasZT0State());
+  ASSERT_TRUE(SA(SA::ZT0_Out).hasSharedZAInterface());
+  ASSERT_FALSE(SA(SA::ZT0_Out).hasPrivateZAInterface());
 
   ASSERT_TRUE(SA(SA::ZT0_InOut).isZT0InOut());
   ASSERT_FALSE(SA(SA::ZT0_InOut).isZT0In());
@@ -204,6 +180,8 @@ TEST(SMEAttributes, Basics) {
   ASSERT_FALSE(SA(SA::ZT0_InOut).hasNewZT0Body());
   ASSERT_TRUE(SA(SA::ZT0_InOut).sharesZT0());
   ASSERT_TRUE(SA(SA::ZT0_InOut).hasZT0State());
+  ASSERT_TRUE(SA(SA::ZT0_InOut).hasSharedZAInterface());
+  ASSERT_FALSE(SA(SA::ZT0_InOut).hasPrivateZAInterface());
 
   ASSERT_TRUE(SA(SA::ZT0_Preserved).preservesZT0());
   ASSERT_FALSE(SA(SA::ZT0_Preserved).isZT0In());
@@ -212,6 +190,8 @@ TEST(SMEAttributes, Basics) {
   ASSERT_FALSE(SA(SA::ZT0_Preserved).hasNewZT0Body());
   ASSERT_TRUE(SA(SA::ZT0_Preserved).sharesZT0());
   ASSERT_TRUE(SA(SA::ZT0_Preserved).hasZT0State());
+  ASSERT_TRUE(SA(SA::ZT0_Preserved).hasSharedZAInterface());
+  ASSERT_FALSE(SA(SA::ZT0_Preserved).hasPrivateZAInterface());
 
   ASSERT_TRUE(SA(SA::ZT0_New).hasNewZT0Body());
   ASSERT_FALSE(SA(SA::ZT0_New).isZT0In());
@@ -220,6 +200,8 @@ TEST(SMEAttributes, Basics) {
   ASSERT_FALSE(SA(SA::ZT0_New).preservesZT0());
   ASSERT_FALSE(SA(SA::ZT0_New).sharesZT0());
   ASSERT_TRUE(SA(SA::ZT0_New).hasZT0State());
+  ASSERT_FALSE(SA(SA::ZT0_New).hasSharedZAInterface());
+  ASSERT_TRUE(SA(SA::ZT0_New).hasPrivateZAInterface());
 
   ASSERT_FALSE(SA(SA::Normal).isZT0In());
   ASSERT_FALSE(SA(SA::Normal).isZT0Out());

>From 7c70ad16c90dd7f8ab0799906966b80d4ea71652 Mon Sep 17 00:00:00 2001
From: Kerry McLaughlin <kerry.mclaughlin at arm.com>
Date: Fri, 12 Jan 2024 15:41:36 +0000
Subject: [PATCH 4/4] - Added an enum (StateValue) to AArch64SMEAttributes for
 representing ZT0 attributes - Added restrictions on the ZT0 attributes in the
 IR Verifier pass

---
 llvm/lib/IR/Verifier.cpp                      |  58 +++++++
 .../AArch64/Utils/AArch64SMEAttributes.cpp    |  14 +-
 .../AArch64/Utils/AArch64SMEAttributes.h      |  52 +++---
 llvm/test/Verifier/sme-attributes.ll          |  30 ++++
 .../Target/AArch64/SMEAttributesTest.cpp      | 155 ++++++++----------
 5 files changed, 194 insertions(+), 115 deletions(-)

diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index aeaca21a99cc5e..701ca07b5082b1 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -2160,6 +2160,64 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
            V);
   }
 
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_new")) {
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_preserved"),
+          "Attributes 'aarch64_sme_zt0_new' and 'aarch64_sme_zt0_preserved' "
+          "are incompatible!",
+          V);
+
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_in"),
+          "Attributes 'aarch64_sme_zt0_new' and 'aarch64_sme_zt0_in' "
+          "are incompatible!",
+          V);
+
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_inout"),
+          "Attributes 'aarch64_sme_zt0_new' and 'aarch64_sme_zt0_inout' "
+          "are incompatible!",
+          V);
+
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_out"),
+          "Attributes 'aarch64_sme_zt0_new' and 'aarch64_sme_zt0_out' "
+          "are incompatible!",
+          V);
+  }
+
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_preserved")) {
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_in"),
+          "Attributes 'aarch64_sme_zt0_preserved' and 'aarch64_sme_zt0_in' "
+          "are incompatible!",
+          V);
+
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_inout"),
+          "Attributes 'aarch64_sme_zt0_preserved' and 'aarch64_sme_zt0_inout' "
+          "are incompatible!",
+          V);
+
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_out"),
+          "Attributes 'aarch64_sme_zt0_preserved' and 'aarch64_sme_zt0_out' "
+          "are incompatible!",
+          V);
+  }
+
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_in")) {
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_inout"),
+          "Attributes 'aarch64_sme_zt0_in' and 'aarch64_sme_zt0_inout' "
+          "are incompatible!",
+          V);
+
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_out"),
+          "Attributes 'aarch64_sme_zt0_in' and 'aarch64_sme_zt0_out' "
+          "are incompatible!",
+          V);
+  }
+
+  if (Attrs.hasFnAttr("aarch64_sme_zt0_inout")) {
+    Check(!Attrs.hasFnAttr("aarch64_sme_zt0_out"),
+          "Attributes 'aarch64_sme_zt0_inout' and 'aarch64_sme_zt0_out' "
+          "are incompatible!",
+          V);
+  }
+
   if (Attrs.hasFnAttr(Attribute::JumpTable)) {
     const GlobalValue *GV = cast<GlobalValue>(V);
     Check(GV->hasGlobalUnnamedAddr(),
diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
index 6061ec607ce000..0a1a5b662c9879 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp
@@ -33,8 +33,8 @@ void SMEAttrs::set(unsigned M, bool Enable) {
          "ZA_Shared and ZA_NoLazySave are mutually exclusive");
 
   // ZT0 Attrs
-  assert((!sharesZT0() || (hasNewZT0Body() ^ isZT0In() ^ isZT0InOut() ^
-                           isZT0Out() ^ preservesZT0())) &&
+  assert((!sharesZT0() || (hasNewZT0Body() ^ isInZT0() ^ isInOutZT0() ^
+                           isOutZT0() ^ preservesZT0())) &&
          "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved are all "
          "mutually exclusive");
 }
@@ -70,15 +70,15 @@ SMEAttrs::SMEAttrs(const AttributeList &Attrs) {
   if (Attrs.hasFnAttr("aarch64_pstate_za_preserved"))
     Bitmask |= ZA_Preserved;
   if (Attrs.hasFnAttr("aarch64_sme_zt0_in"))
-    Bitmask |= ZT0_In;
+    setZT0State(StateValue::In);
   if (Attrs.hasFnAttr("aarch64_sme_zt0_out"))
-    Bitmask |= ZT0_Out;
+    setZT0State(StateValue::Out);
   if (Attrs.hasFnAttr("aarch64_sme_zt0_inout"))
-    Bitmask |= ZT0_InOut;
+    setZT0State(StateValue::InOut);
   if (Attrs.hasFnAttr("aarch64_sme_zt0_preserved"))
-    Bitmask |= ZT0_Preserved;
+    setZT0State(StateValue::Preserved);
   if (Attrs.hasFnAttr("aarch64_sme_zt0_new"))
-    Bitmask |= ZT0_New;
+    setZT0State(StateValue::New);
 }
 
 std::optional<bool>
diff --git a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
index 9d92fbc29eab02..366d7bf9d83b3c 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
+++ b/llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h
@@ -26,22 +26,27 @@ class SMEAttrs {
   unsigned Bitmask;
 
 public:
+  enum class StateValue {
+    None = 0,
+    In = 1,        // aarch64_sme_zt0_in
+    Out = 2,       // aarch64_sme_zt0_out
+    InOut = 3,     // aarch64_sme_zt0_inout
+    Preserved = 4, // aarch64_sme_zt0_preserved
+    New = 5        // aarch64_sme_zt0_new
+  };
+
   // Enum with bitmasks for each individual SME feature.
   enum Mask {
     Normal = 0,
-    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
-    ZA_NoLazySave = 1 << 6,  // Used for SME ABI routines to avoid lazy saves
-    ZT0_New = 1 << 7,        // aarch64_sme_zt0_new
-    ZT0_In = 1 << 8,         // aarch64_sme_zt0_in
-    ZT0_Out = 1 << 9,        // aarch64_sme_zt0_out
-    ZT0_InOut = 1 << 10,     // aarch64_sme_zt0_inout
-    ZT0_Preserved = 1 << 11, // aarch64_sme_zt0_preserved
-    All = ZT0_Preserved - 1
+    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
+    ZA_NoLazySave = 1 << 6, // Used for SME ABI routines to avoid lazy saves
+    ZT0_Shift = 7,
+    ZT0_Mask = 0b111 << ZT0_Shift
   };
 
   SMEAttrs(unsigned Mask = Normal) : Bitmask(0) { set(Mask); }
@@ -92,13 +97,22 @@ class SMEAttrs {
   }
 
   // Interfaces to query ZT0 State
-  bool hasNewZT0Body() const { return Bitmask & ZT0_New; }
-  bool isZT0In() const { return Bitmask & ZT0_In; }
-  bool isZT0Out() const { return Bitmask & ZT0_Out; }
-  bool isZT0InOut() const { return Bitmask & ZT0_InOut; }
-  bool preservesZT0() const { return Bitmask & ZT0_Preserved; }
+  StateValue getZT0State() const {
+    return static_cast<StateValue>((Bitmask & ZT0_Mask) >> ZT0_Shift);
+  }
+  void setZT0State(StateValue S) {
+    Bitmask |= (static_cast<unsigned>(S) << ZT0_Shift);
+  }
+
+  bool hasNewZT0Body() const { return getZT0State() == StateValue::New; }
+  bool isInZT0() const { return getZT0State() == StateValue::In; }
+  bool isOutZT0() const { return getZT0State() == StateValue::Out; }
+  bool isInOutZT0() const { return getZT0State() == StateValue::InOut; }
+  bool preservesZT0() const { return getZT0State() == StateValue::Preserved; }
   bool sharesZT0() const {
-    return Bitmask & (ZT0_In | ZT0_Out | ZT0_InOut | ZT0_Preserved);
+    StateValue State = getZT0State();
+    return State == StateValue::In || State == StateValue::Out ||
+           State == StateValue::InOut || State == StateValue::Preserved;
   }
   bool hasZT0State() const { return hasNewZT0Body() || sharesZT0(); }
 };
diff --git a/llvm/test/Verifier/sme-attributes.ll b/llvm/test/Verifier/sme-attributes.ll
index 7f788cfd09f072..1f5e321d7c927a 100644
--- a/llvm/test/Verifier/sme-attributes.ll
+++ b/llvm/test/Verifier/sme-attributes.ll
@@ -8,3 +8,33 @@ declare void @za_preserved() "aarch64_pstate_za_new" "aarch64_pstate_za_preserve
 
 declare void @za_shared() "aarch64_pstate_za_new" "aarch64_pstate_za_shared";
 ; CHECK: Attributes 'aarch64_pstate_za_new and aarch64_pstate_za_shared' are incompatible!
+
+declare void @zt0_new_preserved() "aarch64_sme_zt0_new" "aarch64_sme_zt0_preserved";
+; CHECK: Attributes 'aarch64_sme_zt0_new' and 'aarch64_sme_zt0_preserved' are incompatible!
+
+declare void @zt0_new_in() "aarch64_sme_zt0_new" "aarch64_sme_zt0_in";
+; CHECK: Attributes 'aarch64_sme_zt0_new' and 'aarch64_sme_zt0_in' are incompatible!
+
+declare void @zt0_new_inout() "aarch64_sme_zt0_new" "aarch64_sme_zt0_inout";
+; CHECK: Attributes 'aarch64_sme_zt0_new' and 'aarch64_sme_zt0_inout' are incompatible!
+
+declare void @zt0_new_out() "aarch64_sme_zt0_new" "aarch64_sme_zt0_out";
+; CHECK: Attributes 'aarch64_sme_zt0_new' and 'aarch64_sme_zt0_out' are incompatible
+
+declare void @zt0_preserved_in() "aarch64_sme_zt0_preserved" "aarch64_sme_zt0_in";
+; CHECK: Attributes 'aarch64_sme_zt0_preserved' and 'aarch64_sme_zt0_in' are incompatible
+
+declare void @zt0_preserved_inout() "aarch64_sme_zt0_preserved" "aarch64_sme_zt0_inout";
+; CHECK: Attributes 'aarch64_sme_zt0_preserved' and 'aarch64_sme_zt0_inout' are incompatible
+
+declare void @zt0_preserved_out() "aarch64_sme_zt0_preserved" "aarch64_sme_zt0_out";
+; CHECK: Attributes 'aarch64_sme_zt0_preserved' and 'aarch64_sme_zt0_out' are incompatible
+
+declare void @zt0_in_inout() "aarch64_sme_zt0_in" "aarch64_sme_zt0_inout";
+; CHECK: Attributes 'aarch64_sme_zt0_in' and 'aarch64_sme_zt0_inout' are incompatible
+
+declare void @zt0_in_out() "aarch64_sme_zt0_in" "aarch64_sme_zt0_out";
+; CHECK: Attributes 'aarch64_sme_zt0_in' and 'aarch64_sme_zt0_out' are incompatible
+
+declare void @zt0_inout_out() "aarch64_sme_zt0_inout" "aarch64_sme_zt0_out";
+; CHECK: Attributes 'aarch64_sme_zt0_inout' and 'aarch64_sme_zt0_out' are incompatible
diff --git a/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp b/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
index 264b0fdaeb56c5..75bf201b44d8bc 100644
--- a/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
+++ b/llvm/unittests/Target/AArch64/SMEAttributesTest.cpp
@@ -56,13 +56,13 @@ TEST(SMEAttributes, Constructors) {
 
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_in\"")
                       ->getFunction("foo"))
-                  .isZT0In());
+                  .isInZT0());
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_out\"")
                       ->getFunction("foo"))
-                  .isZT0Out());
+                  .isOutZT0());
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_inout\"")
                       ->getFunction("foo"))
-                  .isZT0InOut());
+                  .isInOutZT0());
   ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_sme_zt0_preserved\"")
                       ->getFunction("foo"))
                   .preservesZT0());
@@ -78,39 +78,6 @@ TEST(SMEAttributes, Constructors) {
   EXPECT_DEBUG_DEATH(SA(SA::ZA_New | SA::ZA_Preserved),
                      "ZA_New and ZA_Preserved are mutually exclusive");
 
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_In),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_Out),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_InOut),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_New | SA::ZT0_Preserved),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_In | SA::ZT0_Out),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_In | SA::ZT0_InOut),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_Out | SA::ZT0_InOut),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_In),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_Out),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-  EXPECT_DEBUG_DEATH(SA(SA::ZT0_Preserved | SA::ZT0_InOut),
-                     "ZT0_New, ZT0_In, ZT0_Out, ZT0_InOut and ZT0_Preserved "
-                     "are all \" \"mutually exclusive");
-
   // Test that the set() methods equally check validity.
   EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled).set(SA::SM_Compatible),
                      "SM_Enabled and SM_Compatible are mutually exclusive");
@@ -153,59 +120,69 @@ TEST(SMEAttributes, Basics) {
   ASSERT_FALSE(SA(SA::Normal).preservesZA());
 
   // Test ZT0 State interfaces
-  ASSERT_TRUE(SA(SA::ZT0_In).isZT0In());
-  ASSERT_FALSE(SA(SA::ZT0_In).isZT0Out());
-  ASSERT_FALSE(SA(SA::ZT0_In).isZT0InOut());
-  ASSERT_FALSE(SA(SA::ZT0_In).preservesZT0());
-  ASSERT_FALSE(SA(SA::ZT0_In).hasNewZT0Body());
-  ASSERT_TRUE(SA(SA::ZT0_In).sharesZT0());
-  ASSERT_TRUE(SA(SA::ZT0_In).hasZT0State());
-  ASSERT_TRUE(SA(SA::ZT0_In).hasSharedZAInterface());
-  ASSERT_FALSE(SA(SA::ZT0_In).hasPrivateZAInterface());
-
-  ASSERT_TRUE(SA(SA::ZT0_Out).isZT0Out());
-  ASSERT_FALSE(SA(SA::ZT0_Out).isZT0In());
-  ASSERT_FALSE(SA(SA::ZT0_Out).isZT0InOut());
-  ASSERT_FALSE(SA(SA::ZT0_Out).preservesZT0());
-  ASSERT_FALSE(SA(SA::ZT0_Out).hasNewZT0Body());
-  ASSERT_TRUE(SA(SA::ZT0_Out).sharesZT0());
-  ASSERT_TRUE(SA(SA::ZT0_Out).hasZT0State());
-  ASSERT_TRUE(SA(SA::ZT0_Out).hasSharedZAInterface());
-  ASSERT_FALSE(SA(SA::ZT0_Out).hasPrivateZAInterface());
-
-  ASSERT_TRUE(SA(SA::ZT0_InOut).isZT0InOut());
-  ASSERT_FALSE(SA(SA::ZT0_InOut).isZT0In());
-  ASSERT_FALSE(SA(SA::ZT0_InOut).isZT0Out());
-  ASSERT_FALSE(SA(SA::ZT0_InOut).preservesZT0());
-  ASSERT_FALSE(SA(SA::ZT0_InOut).hasNewZT0Body());
-  ASSERT_TRUE(SA(SA::ZT0_InOut).sharesZT0());
-  ASSERT_TRUE(SA(SA::ZT0_InOut).hasZT0State());
-  ASSERT_TRUE(SA(SA::ZT0_InOut).hasSharedZAInterface());
-  ASSERT_FALSE(SA(SA::ZT0_InOut).hasPrivateZAInterface());
-
-  ASSERT_TRUE(SA(SA::ZT0_Preserved).preservesZT0());
-  ASSERT_FALSE(SA(SA::ZT0_Preserved).isZT0In());
-  ASSERT_FALSE(SA(SA::ZT0_Preserved).isZT0Out());
-  ASSERT_FALSE(SA(SA::ZT0_Preserved).isZT0InOut());
-  ASSERT_FALSE(SA(SA::ZT0_Preserved).hasNewZT0Body());
-  ASSERT_TRUE(SA(SA::ZT0_Preserved).sharesZT0());
-  ASSERT_TRUE(SA(SA::ZT0_Preserved).hasZT0State());
-  ASSERT_TRUE(SA(SA::ZT0_Preserved).hasSharedZAInterface());
-  ASSERT_FALSE(SA(SA::ZT0_Preserved).hasPrivateZAInterface());
-
-  ASSERT_TRUE(SA(SA::ZT0_New).hasNewZT0Body());
-  ASSERT_FALSE(SA(SA::ZT0_New).isZT0In());
-  ASSERT_FALSE(SA(SA::ZT0_New).isZT0Out());
-  ASSERT_FALSE(SA(SA::ZT0_New).isZT0InOut());
-  ASSERT_FALSE(SA(SA::ZT0_New).preservesZT0());
-  ASSERT_FALSE(SA(SA::ZT0_New).sharesZT0());
-  ASSERT_TRUE(SA(SA::ZT0_New).hasZT0State());
-  ASSERT_FALSE(SA(SA::ZT0_New).hasSharedZAInterface());
-  ASSERT_TRUE(SA(SA::ZT0_New).hasPrivateZAInterface());
-
-  ASSERT_FALSE(SA(SA::Normal).isZT0In());
-  ASSERT_FALSE(SA(SA::Normal).isZT0Out());
-  ASSERT_FALSE(SA(SA::Normal).isZT0InOut());
+  SA ZT0_In = SA(SA::Normal);
+  ZT0_In.setZT0State(SA::StateValue::In);
+  ASSERT_TRUE(ZT0_In.isInZT0());
+  ASSERT_FALSE(ZT0_In.isOutZT0());
+  ASSERT_FALSE(ZT0_In.isInOutZT0());
+  ASSERT_FALSE(ZT0_In.preservesZT0());
+  ASSERT_FALSE(ZT0_In.hasNewZT0Body());
+  ASSERT_TRUE(ZT0_In.sharesZT0());
+  ASSERT_TRUE(ZT0_In.hasZT0State());
+  ASSERT_TRUE(ZT0_In.hasSharedZAInterface());
+  ASSERT_FALSE(ZT0_In.hasPrivateZAInterface());
+
+  SA ZT0_Out = SA(SA::Normal);
+  ZT0_Out.setZT0State(SA::StateValue::Out);
+  ASSERT_TRUE(ZT0_Out.isOutZT0());
+  ASSERT_FALSE(ZT0_Out.isInZT0());
+  ASSERT_FALSE(ZT0_Out.isInOutZT0());
+  ASSERT_FALSE(ZT0_Out.preservesZT0());
+  ASSERT_FALSE(ZT0_Out.hasNewZT0Body());
+  ASSERT_TRUE(ZT0_Out.sharesZT0());
+  ASSERT_TRUE(ZT0_Out.hasZT0State());
+  ASSERT_TRUE(ZT0_Out.hasSharedZAInterface());
+  ASSERT_FALSE(ZT0_Out.hasPrivateZAInterface());
+
+  SA ZT0_InOut = SA(SA::Normal);
+  ZT0_InOut.setZT0State(SA::StateValue::InOut);
+  ASSERT_TRUE(ZT0_InOut.isInOutZT0());
+  ASSERT_FALSE(ZT0_InOut.isInZT0());
+  ASSERT_FALSE(ZT0_InOut.isOutZT0());
+  ASSERT_FALSE(ZT0_InOut.preservesZT0());
+  ASSERT_FALSE(ZT0_InOut.hasNewZT0Body());
+  ASSERT_TRUE(ZT0_InOut.sharesZT0());
+  ASSERT_TRUE(ZT0_InOut.hasZT0State());
+  ASSERT_TRUE(ZT0_InOut.hasSharedZAInterface());
+  ASSERT_FALSE(ZT0_InOut.hasPrivateZAInterface());
+
+  SA ZT0_Preserved = SA(SA::Normal);
+  ZT0_Preserved.setZT0State(SA::StateValue::Preserved);
+  ASSERT_TRUE(ZT0_Preserved.preservesZT0());
+  ASSERT_FALSE(ZT0_Preserved.isInZT0());
+  ASSERT_FALSE(ZT0_Preserved.isOutZT0());
+  ASSERT_FALSE(ZT0_Preserved.isInOutZT0());
+  ASSERT_FALSE(ZT0_Preserved.hasNewZT0Body());
+  ASSERT_TRUE(ZT0_Preserved.sharesZT0());
+  ASSERT_TRUE(ZT0_Preserved.hasZT0State());
+  ASSERT_TRUE(ZT0_Preserved.hasSharedZAInterface());
+  ASSERT_FALSE(ZT0_Preserved.hasPrivateZAInterface());
+
+  SA ZT0_New = SA(SA::Normal);
+  ZT0_New.setZT0State(SA::StateValue::New);
+  ASSERT_TRUE(ZT0_New.hasNewZT0Body());
+  ASSERT_FALSE(ZT0_New.isInZT0());
+  ASSERT_FALSE(ZT0_New.isOutZT0());
+  ASSERT_FALSE(ZT0_New.isInOutZT0());
+  ASSERT_FALSE(ZT0_New.preservesZT0());
+  ASSERT_FALSE(ZT0_New.sharesZT0());
+  ASSERT_TRUE(ZT0_New.hasZT0State());
+  ASSERT_FALSE(ZT0_New.hasSharedZAInterface());
+  ASSERT_TRUE(ZT0_New.hasPrivateZAInterface());
+
+  ASSERT_FALSE(SA(SA::Normal).isInZT0());
+  ASSERT_FALSE(SA(SA::Normal).isOutZT0());
+  ASSERT_FALSE(SA(SA::Normal).isInOutZT0());
   ASSERT_FALSE(SA(SA::Normal).preservesZT0());
   ASSERT_FALSE(SA(SA::Normal).hasNewZT0Body());
   ASSERT_FALSE(SA(SA::Normal).sharesZT0());



More information about the llvm-commits mailing list