[AArch64][SME] Add utility class for handling SME attributes.

[sdesmalen]

Summary

This patch adds a utility class that will be used in subsequent patches
for parsing the function/callsite attributes and determining whether
changes to PSTATE.SM are needed, or whether a lazy-save mechanism is
required.

It also implements some of the restrictions on the SME attributes
in the IR Verifier pass.

More details about the SME attributes and design can be found
in D131562.

[tschuett]

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

37

Could this help you?
https://llvm.org/doxygen/BitmaskEnum_8h.html

[sdesmalen]

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

37

I don't immediately see a use-case for this utility, given the limited use of fiddling with the masks here. Is there any use-case that stood out to you where this may help?

[david-arm]

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

21

nit: Is it worth just adding a few comments above the class describing what it does?

[sdesmalen]

Added descripton to SMEAttrs class.

[david-arm]

LGTM! Thanks for adding the comments.

[aemerson]

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

61

is hasStreamingInterfaceAndBody() useful as a method too?

73

Again here.

77

Just return the return value?

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

44–50

Would these not be more natural to create from constructors?

67–68

This needs documenting, especially BodyOverridesInterface.

llvm/unittests/Target/AArch64/CMakeLists.txt

11

Why?

[sdesmalen]

Changed the SMEAttrs::getFrom*() functions to constructors.
Also addressed other review comments.

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

77

<facepalm>

llvm/unittests/Target/AArch64/CMakeLists.txt

11

It is needed for parseAssemblyString

[aemerson]

LGTM.

llvm/lib/IR/Verifier.cpp

if (Attrs.hasFnAttr(Attribute::OptimizeNone)) {
 
     Check(!Attrs.hasFnAttr(Attribute::OptimizeForSize),
           "Attributes 'optsize and optnone' are incompatible!", V);
 
     Check(!Attrs.hasFnAttr(Attribute::MinSize),
           "Attributes 'minsize and optnone' are incompatible!", V);
   }
 
+  if (Attrs.hasFnAttr("aarch64_pstate_sm_enabled")) {
+    Check(!Attrs.hasFnAttr("aarch64_pstate_sm_compatible"),
+           "Attributes 'aarch64_pstate_sm_enabled and "
+           "aarch64_pstate_sm_compatible' are incompatible!",
+           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);
+  }
+
   if (Attrs.hasFnAttr(Attribute::JumpTable)) {
     const GlobalValue *GV = cast<GlobalValue>(V);
     Check(GV->hasGlobalUnnamedAddr(),
           "Attribute 'jumptable' requires 'unnamed_addr'", V);
   }
 
   if (Attrs.hasFnAttr(Attribute::AllocSize)) {
     std::pair<unsigned, Optional<unsigned>> Args =

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

+//===-- AArch64SMEAttributes.h - Helper for interpreting SME attributes -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/Optional.h"
+#include "llvm/IR/Function.h"
+
+#ifndef LLVM_LIB_TARGET_AARCH64_UTILS_AARCH64SMEATTRIBUTES_H
+#define LLVM_LIB_TARGET_AARCH64_UTILS_AARCH64SMEATTRIBUTES_H
+namespace llvm {
+
+class Function;
+class CallBase;
+class AttributeList;
+
+/// SMEAttrs is a utility class to parse the SME ACLE attributes on functions.
+/// It helps determine a function's requirements for PSTATE.ZA and PSTATE.SM. It
+/// has interfaces to query whether a streaming mode change or lazy-save
+/// mechanism is required when going from one function to another (e.g. through
+/// a call).
+class SMEAttrs {
+  unsigned Bitmask;
+
+public:
+  // 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_locally
+    ZA_Shared = 1 << 3,     // aarch64_pstate_sm_shared
+    ZA_New = 1 << 4,        // aarch64_pstate_sm_new
+    ZA_Preserved = 1 << 5,  // aarch64_pstate_sm_preserved
+    All = ZA_Preserved - 1
+  };
+
+  SMEAttrs(unsigned Mask = Normal) : Bitmask(0) { set(Mask); }
+  SMEAttrs(const Function &F) : SMEAttrs(F.getAttributes()) {}
+  SMEAttrs(const CallBase &CB);
+  SMEAttrs(const AttributeList &L);
+
+  void set(unsigned M, bool Enable = true);
+
+  // Interfaces to query PSTATE.SM
+  bool hasStreamingBody() const { return Bitmask & SM_Body; }
+  bool hasStreamingInterface() const { return Bitmask & SM_Enabled; }
+  bool hasStreamingInterfaceOrBody() const {
+    return hasStreamingBody() || hasStreamingInterface();
+  }
+  bool hasStreamingCompatibleInterface() const {
+    return Bitmask & SM_Compatible;
+  }
+  bool hasNonStreamingInterface() const {
+    return !hasStreamingInterface() && !hasStreamingCompatibleInterface();
+  }
+  bool hasNonStreamingInterfaceAndBody() const {
+    return hasNonStreamingInterface() && !hasStreamingBody();
+  }
+
+  /// \return true if a call from Caller -> Callee requires a change in
+  /// streaming mode.
+  /// If \p BodyOverridesInterface is true and Callee has a streaming body,
+  /// then requiresSMChange considers a call to Callee as having a Streaming
+  /// interface. This can be useful when considering e.g. inlining, where we
+  /// explicitly want the body to overrule the interface (because after inlining
+  /// the interface is no longer relevant).
+  Optional<bool> requiresSMChange(const SMEAttrs &Callee,
+                                  bool BodyOverridesInterface = false) const;
+
+  // Interfaces to query PSTATE.ZA
+  bool hasNewZAInterface() const { return Bitmask & ZA_New; }
+  bool hasSharedZAInterface() const { return Bitmask & ZA_Shared; }
+  bool hasPrivateZAInterface() const { return !hasSharedZAInterface(); }
+  bool preservesZA() const { return Bitmask & ZA_Preserved; }
+  bool hasZAState() const {
+    return hasNewZAInterface() || hasSharedZAInterface();
+  }
+  bool requiresLazySave(const SMEAttrs &Callee) const {
+    return hasZAState() && Callee.hasPrivateZAInterface() &&
+           !Callee.preservesZA();
+  }
+};
+
+} // namespace llvm
+
+#endif // LLVM_LIB_TARGET_AARCH64_UTILS_AARCH64SMEATTRIBUTES_H

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

+//===-- AArch64SMEAttributes.cpp - Helper for interpreting SME attributes -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "AArch64SMEAttributes.h"
+#include "llvm/ADT/None.h"
+#include "llvm/IR/InstrTypes.h"
+#include <cassert>
+
+using namespace llvm;
+
+void SMEAttrs::set(unsigned M, bool Enable) {
+  if (Enable)
+    Bitmask |= M;
+  else
+    Bitmask &= ~M;
+
+  assert(!(hasStreamingInterface() && hasStreamingCompatibleInterface()) &&
+         "SM_Enabled and SM_Compatible are mutually exclusive");
+  assert(!(hasNewZAInterface() && hasSharedZAInterface()) &&
+         "ZA_New and ZA_Shared are mutually exclusive");
+  assert(!(hasNewZAInterface() && preservesZA()) &&
+         "ZA_New and ZA_Preserved are mutually exclusive");
+}
+
+SMEAttrs::SMEAttrs(const CallBase &CB) {
+  *this = SMEAttrs(CB.getAttributes());
+  if (auto *F = CB.getCalledFunction())
+    set(SMEAttrs(*F).Bitmask);
+}
+
+SMEAttrs::SMEAttrs(const AttributeList &Attrs) {
+  Bitmask = 0;
+  if (Attrs.hasFnAttr("aarch64_pstate_sm_enabled"))
+    Bitmask |= SM_Enabled;
+  if (Attrs.hasFnAttr("aarch64_pstate_sm_compatible"))
+    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;
+}
+
+Optional<bool> SMEAttrs::requiresSMChange(const SMEAttrs &Callee,
+                                          bool BodyOverridesInterface) const {
+  // If the transition is not through a call (e.g. when considering inlining)
+  // and Callee has a streaming body, then we can ignore the interface of
+  // Callee.
+  if (BodyOverridesInterface && Callee.hasStreamingBody()) {
+    return hasStreamingInterfaceOrBody() ? None : Optional<bool>(true);
+  }
+
+  if (Callee.hasStreamingCompatibleInterface())
+    return None;
+
+  // Both non-streaming
+  if (hasNonStreamingInterfaceAndBody() && Callee.hasNonStreamingInterface())
+    return None;
+
+  // Both streaming
+  if (hasStreamingInterfaceOrBody() && Callee.hasStreamingInterface())
+    return None;
+
+  return Callee.hasStreamingInterface();
+}

llvm/lib/Target/AArch64/Utils/CMakeLists.txt

 add_llvm_component_library(LLVMAArch64Utils
   AArch64BaseInfo.cpp
+  AArch64SMEAttributes.cpp
 
   LINK_COMPONENTS
   Support
+  Core
 
   ADD_TO_COMPONENT
   AArch64
   )

llvm/test/Verifier/sme-attributes.ll

+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+declare void @sm_attrs() "aarch64_pstate_sm_enabled" "aarch64_pstate_sm_compatible";
+; CHECK: Attributes 'aarch64_pstate_sm_enabled and aarch64_pstate_sm_compatible' are incompatible!
+
+declare void @za_preserved() "aarch64_pstate_za_new" "aarch64_pstate_za_preserved";
+; CHECK: Attributes 'aarch64_pstate_za_new and aarch64_pstate_za_preserved' are incompatible!
+
+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!

llvm/unittests/Target/AArch64/CMakeLists.txt

 include_directories(
   ${CMAKE_SOURCE_DIR}/lib/Target/AArch64
   ${CMAKE_BINARY_DIR}/lib/Target/AArch64
   )
 
 set(LLVM_LINK_COMPONENTS
   AArch64CodeGen
+  AArch64Utils
   AArch64Desc
   AArch64Info
+  AsmParser
   CodeGen
   Core
   GlobalISel
   MC
   MIRParser
   SelectionDAG
   Support
   Target
   )
 
 add_llvm_target_unittest(AArch64Tests
   InstSizes.cpp
   DecomposeStackOffsetTest.cpp
   MatrixRegisterAliasing.cpp
+  SMEAttributesTest.cpp
   )
 
 set_property(TARGET AArch64Tests PROPERTY FOLDER "Tests/UnitTests/TargetTests")

llvm/unittests/Target/AArch64/SMEAttributesTest.cpp

+#include "Utils/AArch64SMEAttributes.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/SourceMgr.h"
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using SA = SMEAttrs;
+
+std::unique_ptr<Module> parseIR(const char *IR) {
+  static LLVMContext C;
+  SMDiagnostic Err;
+  return parseAssemblyString(IR, Err, C);
+}
+
+TEST(SMEAttributes, Constructors) {
+  LLVMContext Context;
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo()")->getFunction("foo"))
+                  .hasNonStreamingInterfaceAndBody());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_pstate_sm_body\"")
+                      ->getFunction("foo"))
+                  .hasNonStreamingInterface());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_pstate_sm_enabled\"")
+                      ->getFunction("foo"))
+                  .hasStreamingInterface());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_pstate_sm_body\"")
+                      ->getFunction("foo"))
+                  .hasStreamingBody());
+
+  ASSERT_TRUE(
+      SA(*parseIR("declare void @foo() \"aarch64_pstate_sm_compatible\"")
+              ->getFunction("foo"))
+          .hasStreamingCompatibleInterface());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_pstate_za_shared\"")
+                      ->getFunction("foo"))
+                  .hasSharedZAInterface());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_pstate_za_new\"")
+                      ->getFunction("foo"))
+                  .hasNewZAInterface());
+
+  ASSERT_TRUE(SA(*parseIR("declare void @foo() \"aarch64_pstate_za_preserved\"")
+                      ->getFunction("foo"))
+                  .preservesZA());
+
+  // Invalid combinations.
+  EXPECT_DEBUG_DEATH(SA(SA::SM_Enabled | SA::SM_Compatible),
+                     "SM_Enabled and SM_Compatible are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZA_New | SA::ZA_Shared),
+                     "ZA_New and ZA_Shared are mutually exclusive");
+  EXPECT_DEBUG_DEATH(SA(SA::ZA_New | SA::ZA_Preserved),
+                     "ZA_New and ZA_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");
+  EXPECT_DEBUG_DEATH(SA(SA::SM_Compatible).set(SA::SM_Enabled),
+                     "SM_Enabled and SM_Compatible are mutually exclusive");
+}
+
+TEST(SMEAttributes, Basics) {
+  // Test PSTATE.SM interfaces.
+  ASSERT_TRUE(SA(SA::Normal).hasNonStreamingInterfaceAndBody());
+  ASSERT_TRUE(SA(SA::SM_Enabled).hasStreamingInterface());
+  ASSERT_TRUE(SA(SA::SM_Body).hasStreamingBody());
+  ASSERT_TRUE(SA(SA::SM_Body).hasNonStreamingInterface());
+  ASSERT_FALSE(SA(SA::SM_Body).hasNonStreamingInterfaceAndBody());
+  ASSERT_FALSE(SA(SA::SM_Body).hasStreamingInterface());
+  ASSERT_TRUE(SA(SA::SM_Compatible).hasStreamingCompatibleInterface());
+  ASSERT_TRUE(
+      SA(SA::SM_Compatible | SA::SM_Body).hasStreamingCompatibleInterface());
+  ASSERT_TRUE(SA(SA::SM_Compatible | SA::SM_Body).hasStreamingBody());
+  ASSERT_FALSE(SA(SA::SM_Compatible | SA::SM_Body).hasNonStreamingInterface());
+
+  // Test PSTATE.ZA interfaces.
+  ASSERT_FALSE(SA(SA::ZA_Shared).hasPrivateZAInterface());
+  ASSERT_TRUE(SA(SA::ZA_Shared).hasSharedZAInterface());
+  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_TRUE(SA(SA::ZA_New).hasNewZAInterface());
+  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).hasNewZAInterface());
+  ASSERT_FALSE(SA(SA::Normal).hasZAState());
+  ASSERT_FALSE(SA(SA::Normal).preservesZA());
+}
+
+TEST(SMEAttributes, Transitions) {
+  // Normal -> Normal
+  ASSERT_FALSE(SA(SA::Normal).requiresSMChange(SA(SA::Normal)));
+  // Normal -> Normal + LocallyStreaming
+  ASSERT_FALSE(SA(SA::Normal).requiresSMChange(SA(SA::Normal | SA::SM_Body)));
+  ASSERT_EQ(*SA(SA::Normal)
+                 .requiresSMChange(SA(SA::Normal | SA::SM_Body),
+                                   /*BodyOverridesInterface=*/true),
+            true);
+
+  // Normal -> Streaming
+  ASSERT_EQ(*SA(SA::Normal).requiresSMChange(SA(SA::SM_Enabled)), true);
+  // Normal -> Streaming + LocallyStreaming
+  ASSERT_EQ(*SA(SA::Normal).requiresSMChange(SA(SA::SM_Enabled | SA::SM_Body)),
+            true);
+  ASSERT_EQ(*SA(SA::Normal)
+                 .requiresSMChange(SA(SA::SM_Enabled | SA::SM_Body),
+                                   /*BodyOverridesInterface=*/true),
+            true);
+
+  // Normal -> Streaming-compatible
+  ASSERT_FALSE(SA(SA::Normal).requiresSMChange(SA(SA::SM_Compatible)));
+  // Normal -> Streaming-compatible + LocallyStreaming
+  ASSERT_FALSE(
+      SA(SA::Normal).requiresSMChange(SA(SA::SM_Compatible | SA::SM_Body)));
+  ASSERT_EQ(*SA(SA::Normal)
+                 .requiresSMChange(SA(SA::SM_Compatible | SA::SM_Body),
+                                   /*BodyOverridesInterface=*/true),
+            true);
+
+  // Streaming -> Normal
+  ASSERT_EQ(*SA(SA::SM_Enabled).requiresSMChange(SA(SA::Normal)), false);
+  // Streaming -> Normal + LocallyStreaming
+  ASSERT_EQ(*SA(SA::SM_Enabled).requiresSMChange(SA(SA::Normal | SA::SM_Body)),
+            false);
+  ASSERT_FALSE(SA(SA::SM_Enabled)
+                   .requiresSMChange(SA(SA::Normal | SA::SM_Body),
+                                     /*BodyOverridesInterface=*/true));
+
+  // Streaming -> Streaming
+  ASSERT_FALSE(SA(SA::SM_Enabled).requiresSMChange(SA(SA::SM_Enabled)));
+  // Streaming -> Streaming + LocallyStreaming
+  ASSERT_FALSE(
+      SA(SA::SM_Enabled).requiresSMChange(SA(SA::SM_Enabled | SA::SM_Body)));
+  ASSERT_FALSE(SA(SA::SM_Enabled)
+                   .requiresSMChange(SA(SA::SM_Enabled | SA::SM_Body),
+                                     /*BodyOverridesInterface=*/true));
+
+  // Streaming -> Streaming-compatible
+  ASSERT_FALSE(SA(SA::SM_Enabled).requiresSMChange(SA(SA::SM_Compatible)));
+  // Streaming -> Streaming-compatible + LocallyStreaming
+  ASSERT_FALSE(
+      SA(SA::SM_Enabled).requiresSMChange(SA(SA::SM_Compatible | SA::SM_Body)));
+  ASSERT_FALSE(SA(SA::SM_Enabled)
+                   .requiresSMChange(SA(SA::SM_Compatible | SA::SM_Body),
+                                     /*BodyOverridesInterface=*/true));
+
+  // Streaming-compatible -> Normal
+  ASSERT_EQ(*SA(SA::SM_Compatible).requiresSMChange(SA(SA::Normal)), false);
+  ASSERT_EQ(
+      *SA(SA::SM_Compatible).requiresSMChange(SA(SA::Normal | SA::SM_Body)),
+      false);
+  ASSERT_EQ(*SA(SA::SM_Compatible)
+                 .requiresSMChange(SA(SA::Normal | SA::SM_Body),
+                                   /*BodyOverridesInterface=*/true),
+            true);
+
+  // Streaming-compatible -> Streaming
+  ASSERT_EQ(*SA(SA::SM_Compatible).requiresSMChange(SA(SA::SM_Enabled)), true);
+  // Streaming-compatible -> Streaming + LocallyStreaming
+  ASSERT_EQ(
+      *SA(SA::SM_Compatible).requiresSMChange(SA(SA::SM_Enabled | SA::SM_Body)),
+      true);
+  ASSERT_EQ(*SA(SA::SM_Compatible)
+                 .requiresSMChange(SA(SA::SM_Enabled | SA::SM_Body),
+                                   /*BodyOverridesInterface=*/true),
+            true);
+
+  // Streaming-compatible -> Streaming-compatible
+  ASSERT_FALSE(SA(SA::SM_Compatible).requiresSMChange(SA(SA::SM_Compatible)));
+  // Streaming-compatible -> Streaming-compatible + LocallyStreaming
+  ASSERT_FALSE(SA(SA::SM_Compatible)
+                   .requiresSMChange(SA(SA::SM_Compatible | SA::SM_Body)));
+  ASSERT_EQ(*SA(SA::SM_Compatible)
+                 .requiresSMChange(SA(SA::SM_Compatible | SA::SM_Body),
+                                   /*BodyOverridesInterface=*/true),
+            true);
+}