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

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.

llvm/lib/IR/Verifier.cpp

@@ -2161,6 +2161,15 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
            V);
   }
 
+  Check(
+      (Attrs.hasFnAttr("aarch64_new_zt0") + Attrs.hasFnAttr("aarch64_in_zt0") +
+       Attrs.hasFnAttr("aarch64_inout_zt0") +
+       Attrs.hasFnAttr("aarch64_out_zt0") +
+       Attrs.hasFnAttr("aarch64_preserves_zt0")) <= 1,
+      "Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', "
+      "'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive",
+      V);
+
   if (Attrs.hasFnAttr(Attribute::JumpTable)) {
     const GlobalValue *GV = cast<GlobalValue>(V);
     Check(GV->hasGlobalUnnamedAddr(),

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp

@@ -18,16 +18,26 @@ void SMEAttrs::set(unsigned M, bool Enable) {
   else
     Bitmask &= ~M;
 
+  // Streaming Mode Attrs
   assert(!(hasStreamingInterface() && hasStreamingCompatibleInterface()) &&
          "SM_Enabled and SM_Compatible are mutually exclusive");
-  assert(!(hasNewZABody() && hasSharedZAInterface()) &&
+
+  // 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 & 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() || (isNewZT0() ^ isInZT0() ^ isInOutZT0() ^ isOutZT0() ^
+                        isPreservesZT0())) &&
+      "Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', "
+      "'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive");
 }
 
 SMEAttrs::SMEAttrs(const CallBase &CB) {
@@ -60,6 +70,16 @@ SMEAttrs::SMEAttrs(const AttributeList &Attrs) {
     Bitmask |= ZA_New;
   if (Attrs.hasFnAttr("aarch64_pstate_za_preserved"))
     Bitmask |= ZA_Preserved;
+  if (Attrs.hasFnAttr("aarch64_in_zt0"))
+    Bitmask |= encodeZT0State(StateValue::In);
+  if (Attrs.hasFnAttr("aarch64_out_zt0"))
+    Bitmask |= encodeZT0State(StateValue::Out);
+  if (Attrs.hasFnAttr("aarch64_inout_zt0"))
+    Bitmask |= encodeZT0State(StateValue::InOut);
+  if (Attrs.hasFnAttr("aarch64_preserves_zt0"))
+    Bitmask |= encodeZT0State(StateValue::Preserved);
+  if (Attrs.hasFnAttr("aarch64_new_zt0"))
+    Bitmask |= encodeZT0State(StateValue::New);
 }
 
 std::optional<bool>

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h

@@ -26,6 +26,15 @@ class SMEAttrs {
   unsigned Bitmask;
 
 public:
+  enum class StateValue {
+    None = 0,
+    In = 1,        // aarch64_in_zt0
+    Out = 2,       // aarch64_out_zt0
+    InOut = 3,     // aarch64_inout_zt0
+    Preserved = 4, // aarch64_preserves_zt0
+    New = 5        // aarch64_new_zt0
+  };
+
   // Enum with bitmasks for each individual SME feature.
   enum Mask {
     Normal = 0,
@@ -36,7 +45,8 @@ public:
     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
+    ZT0_Shift = 7,
+    ZT0_Mask = 0b111 << ZT0_Shift
   };
 
   SMEAttrs(unsigned Mask = Normal) : Bitmask(0) { set(Mask); }
@@ -76,16 +86,39 @@ public:
 
   // 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);
   }
+
+  // Interfaces to query ZT0 State
+  static StateValue decodeZT0State(unsigned Bitmask) {
+    return static_cast<StateValue>((Bitmask & ZT0_Mask) >> ZT0_Shift);
+  }
+  static unsigned encodeZT0State(StateValue S) {
+    return static_cast<unsigned>(S) << ZT0_Shift;
+  }
+
+  bool isNewZT0() const { return decodeZT0State(Bitmask) == StateValue::New; }
+  bool isInZT0() const { return decodeZT0State(Bitmask) == StateValue::In; }
+  bool isOutZT0() const { return decodeZT0State(Bitmask) == StateValue::Out; }
+  bool isInOutZT0() const {
+    return decodeZT0State(Bitmask) == StateValue::InOut;
+  }
+  bool isPreservesZT0() const {
+    return decodeZT0State(Bitmask) == StateValue::Preserved;
+  }
+  bool sharesZT0() const {
+    StateValue State = decodeZT0State(Bitmask);
+    return State == StateValue::In || State == StateValue::Out ||
+           State == StateValue::InOut || State == StateValue::Preserved;
+  }
+  bool hasZT0State() const { return isNewZT0() || sharesZT0(); }
 };
 
 } // namespace llvm

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_new_zt0" "aarch64_preserves_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_new_in() "aarch64_new_zt0" "aarch64_in_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_new_inout() "aarch64_new_zt0" "aarch64_inout_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_new_out() "aarch64_new_zt0" "aarch64_out_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_preserved_in() "aarch64_preserves_zt0" "aarch64_in_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_preserved_inout() "aarch64_preserves_zt0" "aarch64_inout_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_preserved_out() "aarch64_preserves_zt0" "aarch64_out_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_in_inout() "aarch64_in_zt0" "aarch64_inout_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_in_out() "aarch64_in_zt0" "aarch64_out_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive
+
+declare void @zt0_inout_out() "aarch64_inout_zt0" "aarch64_out_zt0";
+; CHECK: Attributes 'aarch64_new_zt0', 'aarch64_in_zt0', 'aarch64_out_zt0', 'aarch64_inout_zt0' and 'aarch64_preserves_zt0' are mutually exclusive

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());
@@ -50,6 +54,22 @@ TEST(SMEAttributes, Constructors) {
                       ->getFunction("foo"))
                   .preservesZA());
 
+  ASSERT_TRUE(
+      SA(*parseIR("declare void @foo() \"aarch64_in_zt0\"")->getFunction("foo"))
+          .isInZT0());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_out_zt0\"")
+                      ->getFunction("foo"))
+                  .isOutZT0());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_inout_zt0\"")
+                      ->getFunction("foo"))
+                  .isInOutZT0());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_preserves_zt0\"")
+                      ->getFunction("foo"))
+                  .isPreservesZT0());
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_new_zt0\"")
+                      ->getFunction("foo"))
+                  .isNewZT0());
+
   // Invalid combinations.
   EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled | SA::SM_Compatible),
                      "SM_Enabled and SM_Compatible are mutually exclusive");
@@ -82,19 +102,90 @@ 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_FALSE(SA(SA::ZA_Shared).sharesZT0());
+  ASSERT_FALSE(SA(SA::ZA_Shared).hasZT0State());
 
   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_FALSE(SA(SA::ZA_New).sharesZT0());
+  ASSERT_FALSE(SA(SA::ZA_New).hasZT0State());
 
   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());
+
+  // Test ZT0 State interfaces
+  SA ZT0_In = SA(SA::encodeZT0State(SA::StateValue::In));
+  ASSERT_TRUE(ZT0_In.isInZT0());
+  ASSERT_FALSE(ZT0_In.isOutZT0());
+  ASSERT_FALSE(ZT0_In.isInOutZT0());
+  ASSERT_FALSE(ZT0_In.isPreservesZT0());
+  ASSERT_FALSE(ZT0_In.isNewZT0());
+  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::encodeZT0State(SA::StateValue::Out));
+  ASSERT_TRUE(ZT0_Out.isOutZT0());
+  ASSERT_FALSE(ZT0_Out.isInZT0());
+  ASSERT_FALSE(ZT0_Out.isInOutZT0());
+  ASSERT_FALSE(ZT0_Out.isPreservesZT0());
+  ASSERT_FALSE(ZT0_Out.isNewZT0());
+  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::encodeZT0State(SA::StateValue::InOut));
+  ASSERT_TRUE(ZT0_InOut.isInOutZT0());
+  ASSERT_FALSE(ZT0_InOut.isInZT0());
+  ASSERT_FALSE(ZT0_InOut.isOutZT0());
+  ASSERT_FALSE(ZT0_InOut.isPreservesZT0());
+  ASSERT_FALSE(ZT0_InOut.isNewZT0());
+  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::encodeZT0State(SA::StateValue::Preserved));
+  ASSERT_TRUE(ZT0_Preserved.isPreservesZT0());
+  ASSERT_FALSE(ZT0_Preserved.isInZT0());
+  ASSERT_FALSE(ZT0_Preserved.isOutZT0());
+  ASSERT_FALSE(ZT0_Preserved.isInOutZT0());
+  ASSERT_FALSE(ZT0_Preserved.isNewZT0());
+  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::encodeZT0State(SA::StateValue::New));
+  ASSERT_TRUE(ZT0_New.isNewZT0());
+  ASSERT_FALSE(ZT0_New.isInZT0());
+  ASSERT_FALSE(ZT0_New.isOutZT0());
+  ASSERT_FALSE(ZT0_New.isInOutZT0());
+  ASSERT_FALSE(ZT0_New.isPreservesZT0());
+  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).isPreservesZT0());
+  ASSERT_FALSE(SA(SA::Normal).isNewZT0());
+  ASSERT_FALSE(SA(SA::Normal).sharesZT0());
+  ASSERT_FALSE(SA(SA::Normal).hasZT0State());
 }
 
 TEST(SMEAttributes, Transitions) {