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

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.

Reviewed By: david-arm, aemerson

Differential Revision: https://reviews.llvm.org/D131570

llvm/lib/IR/Verifier.cpp

@@ -2058,6 +2058,25 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
           "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(),

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

@@ -0,0 +1,73 @@
+//===-- 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/AArch64SMEAttributes.h

@@ -0,0 +1,90 @@
+//===-- 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/CMakeLists.txt

@@ -1,8 +1,10 @@
 add_llvm_component_library(LLVMAArch64Utils
   AArch64BaseInfo.cpp
+  AArch64SMEAttributes.cpp
 
   LINK_COMPONENTS
   Support
+  Core
 
   ADD_TO_COMPONENT
   AArch64

llvm/test/Verifier/sme-attributes.ll

@@ -0,0 +1,10 @@
+; 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

@@ -5,8 +5,10 @@ include_directories(
 
 set(LLVM_LINK_COMPONENTS
   AArch64CodeGen
+  AArch64Utils
   AArch64Desc
   AArch64Info
+  AsmParser
   CodeGen
   Core
   GlobalISel
@@ -21,6 +23,7 @@ add_llvm_target_unittest(AArch64Tests
   InstSizes.cpp
   DecomposeStackOffsetTest.cpp
   MatrixRegisterAliasing.cpp
+  SMEAttributesTest.cpp
   )
 
 set_property(TARGET AArch64Tests PROPERTY FOLDER "Tests/UnitTests/TargetTests")