[Clang][AArch64] Include SME attributes in the name mangling of function types

clang/lib/AST/ItaniumMangle.cpp

@@ -574,6 +574,7 @@ class CXXNameMangler {
   static StringRef getCallingConvQualifierName(CallingConv CC);
   void mangleExtParameterInfo(FunctionProtoType::ExtParameterInfo info);
   void mangleExtFunctionInfo(const FunctionType *T);
+  void mangleSMEAttrs(unsigned SMEAttrs);
   void mangleBareFunctionType(const FunctionProtoType *T, bool MangleReturnType,
                               const FunctionDecl *FD = nullptr);
   void mangleNeonVectorType(const VectorType *T);
@@ -3532,6 +3533,72 @@ void CXXNameMangler::mangleExtFunctionInfo(const FunctionType *T) {
   // FIXME: noreturn
 }
 
+enum class AAPCSBitmaskSME : unsigned {
+  ArmStreamingBit = 1 << 0,
+  ArmStreamingCompatibleBit = 1 << 1,
+  ArmAgnosticSMEZAStateBit = 1 << 2,
+  ZA_Shift = 3,
+  ZT0_Shift = 6,
+  NoState = 0b000,
+  ArmIn = 0b001,
+  ArmOut = 0b010,
+  ArmInOut = 0b011,
+  ArmPreserves = 0b100,
+  LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue=*/ArmPreserves)
+};
+
+static AAPCSBitmaskSME encodeAAPCSZAState(unsigned SMEAttrs) {
+  switch (SMEAttrs) {
+  case FunctionType::ARM_None:
+    return AAPCSBitmaskSME::NoState;
+  case FunctionType::ARM_In:
+    return AAPCSBitmaskSME::ArmIn;
+  case FunctionType::ARM_Out:
+    return AAPCSBitmaskSME::ArmOut;
+  case FunctionType::ARM_InOut:
+    return AAPCSBitmaskSME::ArmInOut;
+  case FunctionType::ARM_Preserves:
+    return AAPCSBitmaskSME::ArmPreserves;
+  default:
+    llvm_unreachable("Unrecognised SME attribute");
+  }
+}
+
+// The mangling scheme for function types which have SME attributes is
+// implemented as a "pseudo" template:
+//
+//   '__SME_ATTRS<<normal_function_type>, <sme_state>>'
+//
+// Combining the function type with a bitmask representing the streaming and ZA
+// properties of the function's interface.
+//
+// Mangling of SME keywords is described in more detail in the AArch64 ACLE:
+// https://github.com/ARM-software/acle/blob/main/main/acle.md#c-mangling-of-sme-keywords
+//
+void CXXNameMangler::mangleSMEAttrs(unsigned SMEAttrs) {
+  if (!SMEAttrs)
+    return;
+
+  unsigned Bitmask = 0;
+  if (SMEAttrs & FunctionType::SME_PStateSMEnabledMask)
+    Bitmask |= static_cast<unsigned>(AAPCSBitmaskSME::ArmStreamingBit);
+  else if (SMEAttrs & FunctionType::SME_PStateSMCompatibleMask)
+    Bitmask |=
+        static_cast<unsigned>(AAPCSBitmaskSME::ArmStreamingCompatibleBit);
+
+  // TODO: Must represent __arm_agnostic("sme_za_state")
+
+  Bitmask |= static_cast<unsigned>(
+      encodeAAPCSZAState(FunctionType::getArmZAState(SMEAttrs))
+      << AAPCSBitmaskSME::ZA_Shift);
+
+  Bitmask |= static_cast<unsigned>(
+      encodeAAPCSZAState(FunctionType::getArmZT0State(SMEAttrs))
+      << AAPCSBitmaskSME::ZT0_Shift);
+
+  Out << "Lj" << Bitmask << "EE";
+}
+
 void
 CXXNameMangler::mangleExtParameterInfo(FunctionProtoType::ExtParameterInfo PI) {
   // Vendor-specific qualifiers are emitted in reverse alphabetical order.
@@ -3569,6 +3636,11 @@ CXXNameMangler::mangleExtParameterInfo(FunctionProtoType::ExtParameterInfo PI) {
 // <function-type> ::= [<CV-qualifiers>] F [Y]
 //                      <bare-function-type> [<ref-qualifier>] E
 void CXXNameMangler::mangleType(const FunctionProtoType *T) {
+  unsigned SMEAttrs = T->getAArch64SMEAttributes();
+
+  if (SMEAttrs)
+    Out << "11__SME_ATTRSI";
+
   mangleExtFunctionInfo(T);
 
   // Mangle CV-qualifiers, if present.  These are 'this' qualifiers,
@@ -3603,6 +3675,8 @@ void CXXNameMangler::mangleType(const FunctionProtoType *T) {
   mangleRefQualifier(T->getRefQualifier());
 
   Out << 'E';
+
+  mangleSMEAttrs(SMEAttrs);
 }
 
 void CXXNameMangler::mangleType(const FunctionNoProtoType *T) {

clang/test/CodeGenCXX/aarch64-mangle-sme-atts.cpp

@@ -0,0 +1,74 @@
+// RUN: %clang_cc1 -triple aarch64-none-linux-gnu -target-feature +sve -target-feature +sme -target-feature +sme2 %s -emit-llvm -o - | FileCheck %s
+
+typedef __attribute__((neon_vector_type(2))) int int32x2_t;
+
+//
+// Streaming-Mode Attributes
+//
+
+// CHECK: define dso_local void @_Z12fn_streamingP11__SME_ATTRSIFvvELj1EE
+void fn_streaming(void (*foo)() __arm_streaming) { foo(); }
+
+// CHECK: define dso_local void @_Z23fn_streaming_compatibleP11__SME_ATTRSIFivELj2EE(
+void fn_streaming_compatible(int (*foo)() __arm_streaming_compatible) { foo(); }
+
+//
+// ZA Attributes
+//
+
+// CHECK: define dso_local void @_Z15fn_za_preservedP11__SME_ATTRSIF11__Int32x2_tvELj32EE(
+__arm_new("za") void fn_za_preserved(int32x2_t (*foo)() __arm_preserves("za")) { foo(); }
+
+// CHECK: define dso_local void @_Z8fn_za_inP11__SME_ATTRSIFvu13__SVFloat64_tELj8EES_(
+__arm_new("za") void fn_za_in(void (*foo)(__SVFloat64_t) __arm_in("za"), __SVFloat64_t x) { foo(x); }
+
+// CHECK: define dso_local noundef i32 @_Z9fn_za_outP11__SME_ATTRSIFivELj16EE(
+__arm_new("za") int fn_za_out(int (*foo)() __arm_out("za")) { return foo(); }
+
+// CHECK: define dso_local void @_Z11fn_za_inoutP11__SME_ATTRSIFvvELj24EE(
+__arm_new("za") void fn_za_inout(void (*foo)() __arm_inout("za")) { foo(); }
+
+
+//
+// ZT0 Attributes
+//
+
+// CHECK: define dso_local void @_Z16fn_zt0_preservedP11__SME_ATTRSIFivELj256EE(
+__arm_new("zt0") void fn_zt0_preserved(int (*foo)() __arm_preserves("zt0")) { foo(); }
+
+// CHECK: define dso_local void @_Z9fn_zt0_inP11__SME_ATTRSIFivELj64EE(
+__arm_new("zt0") void fn_zt0_in(int (*foo)() __arm_in("zt0")) { foo(); }
+
+// CHECK: define dso_local void @_Z10fn_zt0_outP11__SME_ATTRSIFivELj128EE(
+__arm_new("zt0") void fn_zt0_out(int (*foo)() __arm_out("zt0")) { foo(); }
+
+// CHECK: define dso_local void @_Z12fn_zt0_inoutP11__SME_ATTRSIFivELj192EE(
+__arm_new("zt0") void fn_zt0_inout(int (*foo)() __arm_inout("zt0")) { foo(); }
+
+//
+// Streaming-mode, ZA & ZT0 Attributes
+//
+
+// CHECK: define dso_local void @_Z17fn_all_attr_typesP11__SME_ATTRSIFivELj282EE(
+__arm_new("za") __arm_new("zt0")
+void fn_all_attr_types(int (*foo)() __arm_streaming_compatible __arm_inout("za") __arm_preserves("zt0"))
+{ foo(); }
+
+//
+// No SME Attributes
+//
+
+// CHECK: define dso_local void @_Z12no_sme_attrsPFvvE(
+void no_sme_attrs(void (*foo)()) { foo(); }
+
+// CHECK: define dso_local void @_Z24locally_streaming_callerPFvvE(
+__arm_locally_streaming void locally_streaming_caller(void (*foo)()) { foo(); }
+
+// CHECK: define dso_local void @_Z16streaming_callerv(
+void streaming_caller() __arm_streaming {}
+
+// CHECK: define dso_local void @_Z16za_shared_callerv(
+void za_shared_caller() __arm_in("za") {}
+
+// CHECK: define dso_local void @_Z17zt0_shared_callerv(
+void zt0_shared_caller() __arm_out("zt0") {}