Implement target(branch-protection) attribute for AArch64

This patch implements `__attribute__((target("branch-protection=...")))`
in a manner, compatible with the analogous GCC feature:

https://gcc.gnu.org/onlinedocs/gcc-9.2.0/gcc/AArch64-Function-Attributes.html#AArch64-Function-Attributes

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

clang/include/clang/Basic/Attr.td

@@ -2138,6 +2138,7 @@ def Target : InheritableAttr {
     struct ParsedTargetAttr {
       std::vector<std::string> Features;
       StringRef Architecture;
+      StringRef BranchProtection;
       bool DuplicateArchitecture = false;
       bool operator ==(const ParsedTargetAttr &Other) const {
         return DuplicateArchitecture == Other.DuplicateArchitecture &&
@@ -2210,6 +2211,11 @@ def Target : InheritableAttr {
         if (Feature.startswith("fpmath=") || Feature.startswith("tune="))
           continue;
 
+        if (Feature.startswith("branch-protection=")) {
+          Ret.BranchProtection = Feature.split('=').second.trim();
+          continue;
+        }
+
         // While we're here iterating check for a different target cpu.
         if (Feature.startswith("arch=")) {
           if (!Ret.Architecture.empty())

clang/include/clang/Basic/AttrDocs.td

@@ -1806,6 +1806,10 @@ the target with or without a "-mno-" in front corresponding to the absence
 of the feature, as well as ``arch="CPU"`` which will change the default "CPU"
 for the function.
 
+For AArch64, the attribute also allows the "branch-protection=<args>" option,
+where the permissible arguments and their effect on code generation are the same
+as for the command-line option ``-mbranch-protection``.
+
 Example "subtarget features" from the x86 backend include: "mmx", "sse", "sse4.2",
 "avx", "xop" and largely correspond to the machine specific options handled by
 the front end.

clang/include/clang/Basic/CodeGenOptions.def

@@ -370,8 +370,8 @@ CODEGENOPT(ForceEmitVTables, 1, 0)
 /// Whether to emit an address-significance table into the object file.
 CODEGENOPT(Addrsig, 1, 0)
 
-ENUM_CODEGENOPT(SignReturnAddress, SignReturnAddressScope, 2, None)
-ENUM_CODEGENOPT(SignReturnAddressKey, SignReturnAddressKeyValue, 1, AKey)
+ENUM_CODEGENOPT(SignReturnAddress, SignReturnAddressScope, 2, SignReturnAddressScope::None)
+ENUM_CODEGENOPT(SignReturnAddressKey, SignReturnAddressKeyValue, 1, SignReturnAddressKeyValue::AKey)
 CODEGENOPT(BranchTargetEnforcement, 1, 0)
 
 /// Whether to emit unused static constants.

clang/include/clang/Basic/CodeGenOptions.h

@@ -109,13 +109,13 @@ public:
     Embed_Marker    // Embed a marker as a placeholder for bitcode.
   };
 
-  enum SignReturnAddressScope {
+  enum class SignReturnAddressScope {
     None,    // No signing for any function
     NonLeaf, // Sign the return address of functions that spill LR
     All      // Sign the return address of all functions
   };
 
-  enum SignReturnAddressKeyValue { AKey, BKey };
+  enum class SignReturnAddressKeyValue { AKey, BKey };
 
   enum class FramePointerKind {
     None,        // Omit all frame pointers.

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -2595,6 +2595,8 @@ def err_attribute_requires_positive_integer : Error<
   "integral compile time constant expression">;
 def err_attribute_requires_opencl_version : Error<
   "%0 attribute requires OpenCL version %1%select{| or above}2">;
+def err_invalid_branch_protection_spec : Error<
+  "invalid or misplaced branch protection specification '%0'">;
 def warn_unsupported_target_attribute
     : Warning<"%select{unsupported|duplicate}0%select{| architecture}1 '%2' in"
               " the 'target' attribute string; 'target' attribute ignored">,

clang/include/clang/Basic/TargetInfo.h

@@ -16,6 +16,7 @@
 
 #include "clang/Basic/AddressSpaces.h"
 #include "clang/Basic/LLVM.h"
+#include "clang/Basic/CodeGenOptions.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Basic/TargetCXXABI.h"
 #include "clang/Basic/TargetOptions.h"
@@ -1103,6 +1104,23 @@ public:
     return true;
   }
 
+  struct BranchProtectionInfo {
+    CodeGenOptions::SignReturnAddressScope SignReturnAddr =
+        CodeGenOptions::SignReturnAddressScope::None;
+    CodeGenOptions::SignReturnAddressKeyValue SignKey =
+        CodeGenOptions::SignReturnAddressKeyValue::AKey;
+    bool BranchTargetEnforcement = false;
+  };
+
+  /// Determine if this TargetInfo supports the given branch protection
+  /// specification
+  virtual bool validateBranchProtection(StringRef Spec,
+                                        BranchProtectionInfo &BPI,
+                                        StringRef &Err) const {
+    Err = "";
+    return false;
+  }
+
   /// Perform initialization based on the user configured
   /// set of features (e.g., +sse4).
   ///

clang/lib/Basic/Targets/AArch64.cpp

@@ -15,6 +15,8 @@
 #include "clang/Basic/TargetInfo.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/AArch64TargetParser.h"
 
 using namespace clang;
 using namespace clang::targets;
@@ -107,6 +109,28 @@ bool AArch64TargetInfo::setABI(const std::string &Name) {
   return true;
 }
 
+bool AArch64TargetInfo::validateBranchProtection(StringRef Spec,
+                                                 BranchProtectionInfo &BPI,
+                                                 StringRef &Err) const {
+  llvm::AArch64::ParsedBranchProtection PBP;
+  if (!llvm::AArch64::parseBranchProtection(Spec, PBP, Err))
+    return false;
+
+  BPI.SignReturnAddr =
+      llvm::StringSwitch<CodeGenOptions::SignReturnAddressScope>(PBP.Scope)
+          .Case("non-leaf", CodeGenOptions::SignReturnAddressScope::NonLeaf)
+          .Case("all", CodeGenOptions::SignReturnAddressScope::All)
+          .Default(CodeGenOptions::SignReturnAddressScope::None);
+
+  if (PBP.Key == "a_key")
+    BPI.SignKey = CodeGenOptions::SignReturnAddressKeyValue::AKey;
+  else
+    BPI.SignKey = CodeGenOptions::SignReturnAddressKeyValue::BKey;
+
+  BPI.BranchTargetEnforcement = PBP.BranchTargetEnforcement;
+  return true;
+}
+
 bool AArch64TargetInfo::isValidCPUName(StringRef Name) const {
   return Name == "generic" ||
          llvm::AArch64::parseCPUArch(Name) != llvm::AArch64::ArchKind::INVALID;

clang/lib/Basic/Targets/AArch64.h

@@ -49,6 +49,9 @@ public:
   StringRef getABI() const override;
   bool setABI(const std::string &Name) override;
 
+  bool validateBranchProtection(StringRef, BranchProtectionInfo &,
+                                StringRef &) const override;
+
   bool isValidCPUName(StringRef Name) const override;
   void fillValidCPUList(SmallVectorImpl<StringRef> &Values) const override;
   bool setCPU(const std::string &Name) override;

clang/lib/CodeGen/TargetInfo.cpp

@@ -5056,23 +5056,38 @@ public:
     const FunctionDecl *FD = dyn_cast_or_null<FunctionDecl>(D);
     if (!FD)
       return;
-    llvm::Function *Fn = cast<llvm::Function>(GV);
 
-    auto Kind = CGM.getCodeGenOpts().getSignReturnAddress();
-    if (Kind != CodeGenOptions::SignReturnAddressScope::None) {
+    CodeGenOptions::SignReturnAddressScope Scope = CGM.getCodeGenOpts().getSignReturnAddress();
+    CodeGenOptions::SignReturnAddressKeyValue Key = CGM.getCodeGenOpts().getSignReturnAddressKey();
+    bool BranchTargetEnforcement = CGM.getCodeGenOpts().BranchTargetEnforcement;
+    if (const auto *TA = FD->getAttr<TargetAttr>()) {
+      TargetAttr::ParsedTargetAttr Attr = TA->parse();
+      if (!Attr.BranchProtection.empty()) {
+        TargetInfo::BranchProtectionInfo BPI;
+        StringRef Error;
+        (void)CGM.getTarget().validateBranchProtection(Attr.BranchProtection,
+                                                       BPI, Error);
+        assert(Error.empty());
+        Scope = BPI.SignReturnAddr;
+        Key = BPI.SignKey;
+        BranchTargetEnforcement = BPI.BranchTargetEnforcement;
+      }
+    }
+
+    auto *Fn = cast<llvm::Function>(GV);
+    if (Scope != CodeGenOptions::SignReturnAddressScope::None) {
       Fn->addFnAttr("sign-return-address",
-                    Kind == CodeGenOptions::SignReturnAddressScope::All
+                    Scope == CodeGenOptions::SignReturnAddressScope::All
                         ? "all"
                         : "non-leaf");
 
-      auto Key = CGM.getCodeGenOpts().getSignReturnAddressKey();
       Fn->addFnAttr("sign-return-address-key",
                     Key == CodeGenOptions::SignReturnAddressKeyValue::AKey
                         ? "a_key"
                         : "b_key");
     }
 
-    if (CGM.getCodeGenOpts().BranchTargetEnforcement)
+    if (BranchTargetEnforcement)
       Fn->addFnAttr("branch-target-enforcement");
   }
 };

clang/lib/Driver/ToolChains/Clang.cpp

@@ -1529,56 +1529,6 @@ void Clang::RenderTargetOptions(const llvm::Triple &EffectiveTriple,
   }
 }
 
-// Parse -mbranch-protection=<protection>[+<protection>]* where
-//   <protection> ::= standard | none | [bti,pac-ret[+b-key,+leaf]*]
-// Returns a triple of (return address signing Scope, signing key, require
-// landing pads)
-static std::tuple<StringRef, StringRef, bool>
-ParseAArch64BranchProtection(const Driver &D, const ArgList &Args,
-                             const Arg *A) {
-  StringRef Scope = "none";
-  StringRef Key = "a_key";
-  bool IndirectBranches = false;
-
-  StringRef Value = A->getValue();
-  // This maps onto -mbranch-protection=<scope>+<key>
-
-  if (Value.equals("standard")) {
-    Scope = "non-leaf";
-    Key = "a_key";
-    IndirectBranches = true;
-
-  } else if (!Value.equals("none")) {
-    SmallVector<StringRef, 4> BranchProtection;
-    StringRef(A->getValue()).split(BranchProtection, '+');
-
-    auto Protection = BranchProtection.begin();
-    while (Protection != BranchProtection.end()) {
-      if (Protection->equals("bti"))
-        IndirectBranches = true;
-      else if (Protection->equals("pac-ret")) {
-        Scope = "non-leaf";
-        while (++Protection != BranchProtection.end()) {
-          // Inner loop as "leaf" and "b-key" options must only appear attached
-          // to pac-ret.
-          if (Protection->equals("leaf"))
-            Scope = "all";
-          else if (Protection->equals("b-key"))
-            Key = "b_key";
-          else
-            break;
-        }
-        Protection--;
-      } else
-        D.Diag(diag::err_invalid_branch_protection)
-            << *Protection << A->getAsString(Args);
-      Protection++;
-    }
-  }
-
-  return std::make_tuple(Scope, Key, IndirectBranches);
-}
-
 namespace {
 void RenderAArch64ABI(const llvm::Triple &Triple, const ArgList &Args,
                       ArgStringList &CmdArgs) {
@@ -1650,9 +1600,16 @@ void Clang::AddAArch64TargetArgs(const ArgList &Args,
             << Scope << A->getAsString(Args);
       Key = "a_key";
       IndirectBranches = false;
-    } else
-      std::tie(Scope, Key, IndirectBranches) =
-          ParseAArch64BranchProtection(D, Args, A);
+    } else {
+      StringRef Err;
+      llvm::AArch64::ParsedBranchProtection PBP;
+      if (!llvm::AArch64::parseBranchProtection(A->getValue(), PBP, Err))
+        D.Diag(diag::err_invalid_branch_protection)
+            << Err << A->getAsString(Args);
+      Scope = PBP.Scope;
+      Key = PBP.Key;
+      IndirectBranches = PBP.BranchTargetEnforcement;
+    }
 
     CmdArgs.push_back(
         Args.MakeArgString(Twine("-msign-return-address=") + Scope));

clang/lib/Sema/SemaDeclAttr.cpp

@@ -3041,6 +3041,19 @@ bool Sema::checkTargetAttr(SourceLocation LiteralLoc, StringRef AttrStr) {
              << Unsupported << None << CurFeature;
   }
 
+  TargetInfo::BranchProtectionInfo BPI;
+  StringRef Error;
+  if (!ParsedAttrs.BranchProtection.empty() &&
+      !Context.getTargetInfo().validateBranchProtection(
+          ParsedAttrs.BranchProtection, BPI, Error)) {
+    if (Error.empty())
+      return Diag(LiteralLoc, diag::warn_unsupported_target_attribute)
+             << Unsupported << None << "branch-protection";
+    else
+      return Diag(LiteralLoc, diag::err_invalid_branch_protection_spec)
+             << Error;
+  }
+
   return false;
 }
 

clang/test/CodeGen/aarch64-branch-protection-attr.c

@@ -0,0 +1,81 @@
+// REQUIRES: aarch64-registered-target
+// RUN: %clang_cc1 -triple aarch64-unknown-unknown-eabi -emit-llvm  -target-cpu generic -target-feature +v8.5a %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK --check-prefix=NO-OVERRIDE
+// RUN: %clang_cc1 -triple aarch64-unknown-unknown-eabi -emit-llvm  -target-cpu generic -target-feature +v8.5a %s -o - \
+// RUN:   -msign-return-address=non-leaf -msign-return-address-key=a_key -mbranch-target-enforce \
+// RUN:                               | FileCheck %s --check-prefix=CHECK --check-prefix=OVERRIDE
+
+void missing() {}
+// NO-OVERRIDE: define void @missing() #[[#NONE:]]
+// OVERRIDE: define void @missing() #[[#STD:]]
+
+__attribute__ ((target("branch-protection=none")))
+void none() {}
+// NO-OVERRIDE: define void @none() #[[#NONE]]
+// OVERRIDE: define void @none() #[[#NONE:]]
+
+  __attribute__ ((target("branch-protection=standard")))
+void std() {}
+// NO-OVERRIDE: define void @std() #[[#STD:]]
+// OVERRIDE: define void @std() #[[#STD]]
+
+__attribute__ ((target("branch-protection=bti")))
+void btionly() {}
+// NO-OVERRIDE: define void @btionly() #[[#BTI:]]
+// OVERRIDE: define void @btionly() #[[#BTI:]]
+
+__attribute__ ((target("branch-protection=pac-ret")))
+void paconly() {}
+// NO-OVERRIDE: define void @paconly() #[[#PAC:]]
+// OVERRIDE: define void @paconly() #[[#PAC:]]
+
+__attribute__ ((target("branch-protection=pac-ret+bti")))
+void pacbti0() {}
+// NO-OVERRIDE: define void @pacbti0() #[[#PACBTI:]]
+// OVERRIDE: define void @pacbti0() #[[#PACBTI:]]
+
+__attribute__ ((target("branch-protection=bti+pac-ret")))
+void pacbti1() {}
+// NO-OVERRIDE: define void @pacbti1() #[[#PACBTI]]
+// OVERRIDE: define void @pacbti1() #[[#PACBTI]]
+
+__attribute__ ((target("branch-protection=pac-ret+leaf")))
+void leaf() {}
+// NO-OVERRIDE: define void @leaf() #[[#PACLEAF:]]
+// OVERRIDE: define void @leaf() #[[#PACLEAF:]]
+
+__attribute__ ((target("branch-protection=pac-ret+b-key")))
+void bkey() {}
+// NO-OVERRIDE: define void @bkey() #[[#PACBKEY:]]
+// OVERRIDE: define void @bkey() #[[#PACBKEY:]]
+
+__attribute__ ((target("branch-protection=pac-ret+b-key+leaf")))
+void bkeyleaf0() {}
+// NO-OVERRIDE: define void @bkeyleaf0()  #[[#PACBKEYLEAF:]]
+// OVERRIDE: define void @bkeyleaf0()  #[[#PACBKEYLEAF:]]
+
+__attribute__ ((target("branch-protection=pac-ret+leaf+b-key")))
+void bkeyleaf1() {}
+// NO-OVERRIDE: define void @bkeyleaf1()  #[[#PACBKEYLEAF]]
+// OVERRIDE: define void @bkeyleaf1()  #[[#PACBKEYLEAF]]
+
+__attribute__ ((target("branch-protection=pac-ret+leaf+bti")))
+void btileaf() {}
+// NO-OVERRIDE: define void @btileaf() #[[#BTIPACLEAF:]]
+// OVERRIDE: define void @btileaf() #[[#BTIPACLEAF:]]
+
+// CHECK-DAG: attributes #[[#NONE]]
+
+// CHECK-DAG: attributes #[[#STD]] = { {{.*}} "branch-target-enforcement" {{.*}} "sign-return-address"="non-leaf" "sign-return-address-key"="a_key"
+
+// CHECK-DAG: attributes #[[#BTI]] = { {{.*}}"branch-target-enforcement"
+
+// CHECK-DAG: attributes #[[#PAC]] = { {{.*}} "sign-return-address"="non-leaf" "sign-return-address-key"="a_key"
+
+// CHECK-DAG: attributes #[[#PACLEAF]] = { {{.*}} "sign-return-address"="all" "sign-return-address-key"="a_key"
+
+// CHECK-DAG: attributes #[[#PACBKEY]] = { {{.*}} "sign-return-address"="non-leaf" "sign-return-address-key"="b_key"
+
+// CHECK-DAG: attributes #[[#PACBKEYLEAF]] = { {{.*}} "sign-return-address"="all" "sign-return-address-key"="b_key"
+
+// CHECK-DAG: attributes #[[#BTIPACLEAF]] = { {{.*}}"branch-target-enforcement" {{.*}} "sign-return-address"="all" "sign-return-address-key"="a_key"

clang/test/Sema/attr-target.c

@@ -17,5 +17,6 @@ int __attribute__((target("arch="))) turtle() { return 4; }
 int __attribute__((target("arch=hiss,arch=woof"))) pine_tree() { return 4; }
 //expected-warning@+1 {{duplicate 'arch=' in the 'target' attribute string; 'target' attribute ignored}}
 int __attribute__((target("arch=ivybridge,arch=haswell"))) oak_tree() { return 4; }
-
+//expected-warning@+1 {{unsupported 'branch-protection' in the 'target' attribute string; 'target' attribute ignored}}
+int __attribute__((target("branch-protection=none"))) birch_tree() { return 5; }
 

clang/test/Sema/branch-protection-attr-err.c

@@ -0,0 +1,22 @@
+// RUN: %clang_cc1 -triple aarch64 -verify -fsyntax-only %s
+
+__attribute__((target("branch-protection=foo"))) // expected-error {{invalid or misplaced branch protection specification 'foo'}}
+void badvalue0() {}
+
+__attribute__((target("branch-protection=+bti"))) // expected-error {{invalid or misplaced branch protection specification '<empty>'}}
+void badvalue1() {}
+
+__attribute__((target("branch-protection=bti+"))) // expected-error {{invalid or misplaced branch protection specification '<empty>'}}
+void badvalue2() {}
+
+__attribute__((target("branch-protection=pac-ret+bkey"))) // expected-error {{invalid or misplaced branch protection specification 'bkey'}}
+void badvalue3() {}
+
+__attribute__((target("branch-protection=bti+leaf"))) // expected-error {{invalid or misplaced branch protection specification 'leaf'}}
+void badoption0() {}
+
+__attribute__((target("branch-protection=bti+leaf+pac-ret"))) // expected-error {{invalid or misplaced branch protection specification 'leaf'}}
+void badorder0() {}
+
+__attribute__ ((target("branch-protection=pac-ret+bti+leaf")))  // expected-error {{invalid or misplaced branch protection specification 'leaf'}}
+void badorder1() {}

llvm/include/llvm/Support/AArch64TargetParser.h

@@ -123,6 +123,15 @@ void fillValidCPUArchList(SmallVectorImpl<StringRef> &Values);
 
 bool isX18ReservedByDefault(const Triple &TT);
 
+struct ParsedBranchProtection {
+  StringRef Scope;
+  StringRef Key;
+  bool BranchTargetEnforcement;
+};
+
+bool parseBranchProtection(StringRef Spec, ParsedBranchProtection &PBP,
+                           StringRef &Err);
+
 } // namespace AArch64
 } // namespace llvm
 

llvm/lib/Support/AArch64TargetParser.cpp

@@ -213,3 +213,51 @@ AArch64::ArchKind AArch64::parseCPUArch(StringRef CPU) {
   }
   return ArchKind::INVALID;
 }
+
+// Parse a branch protection specification, which has the form
+//   standard | none | [bti,pac-ret[+b-key,+leaf]*]
+// Returns true on success, with individual elements of the specification
+// returned in `PBP`. Returns false in error, with `Err` containing
+// an erroneous part of the spec.
+bool AArch64::parseBranchProtection(StringRef Spec, ParsedBranchProtection &PBP,
+                                    StringRef &Err) {
+  PBP = {"none", "a_key", false};
+  if (Spec == "none")
+    return true; // defaults are ok
+
+  if (Spec == "standard") {
+    PBP.Scope = "non-leaf";
+    PBP.BranchTargetEnforcement = true;
+    return true;
+  }
+
+  SmallVector<StringRef, 4> Opts;
+  Spec.split(Opts, "+");
+  for (int I = 0, E = Opts.size(); I != E; ++I) {
+    StringRef Opt = Opts[I].trim();
+    if (Opt == "bti") {
+      PBP.BranchTargetEnforcement = true;
+      continue;
+    }
+    if (Opt == "pac-ret") {
+      PBP.Scope = "non-leaf";
+      for (; I + 1 != E; ++I) {
+        StringRef PACOpt = Opts[I + 1].trim();
+        if (PACOpt == "leaf")
+          PBP.Scope = "all";
+        else if (PACOpt == "b-key")
+          PBP.Key = "b_key";
+        else
+          break;
+      }
+      continue;
+    }
+    if (Opt == "")
+      Err = "<empty>";
+    else
+      Err = Opt;
+    return false;
+  }
+
+  return true;
+}