[Arm][AArch64][Clang] Respect function's branch protection attributes.

clang/lib/CodeGen/CGCall.cpp

@@ -2010,7 +2010,7 @@ static void getTrivialDefaultFunctionAttributes(
   }
 
   TargetInfo::BranchProtectionInfo BPI(LangOpts);
-  TargetCodeGenInfo::setBranchProtectionFnAttributes(BPI, FuncAttrs);
+  TargetCodeGenInfo::initBranchProtectionFnAttributes(BPI, FuncAttrs);
 }
 
 /// Merges `target-features` from \TargetOpts and \F, and sets the result in

clang/lib/CodeGen/TargetInfo.cpp

@@ -209,13 +209,37 @@ llvm::Value *TargetCodeGenInfo::createEnqueuedBlockKernel(
 
 void TargetCodeGenInfo::setBranchProtectionFnAttributes(
     const TargetInfo::BranchProtectionInfo &BPI, llvm::Function &F) {
-  llvm::AttrBuilder FuncAttrs(F.getContext());
-  setBranchProtectionFnAttributes(BPI, FuncAttrs);
-  F.addFnAttrs(FuncAttrs);
+  // Called on already created and initialized function where attributes already
+  // set from command line attributes but some might need to be removed as the
+  // actual BPI is different.
+  if (BPI.SignReturnAddr != LangOptions::SignReturnAddressScopeKind::None) {
+    F.addFnAttr("sign-return-address", BPI.getSignReturnAddrStr());
+    F.addFnAttr("sign-return-address-key", BPI.getSignKeyStr());
+  } else {
+    if (F.hasFnAttribute("sign-return-address"))
+      F.removeFnAttr("sign-return-address");
+    if (F.hasFnAttribute("sign-return-address-key"))
+      F.removeFnAttr("sign-return-address-key");
+  }
+
+  auto AddRemoveAttributeAsSet = [&](bool Set, const StringRef &ModAttr) {
+    if (Set)
+      F.addFnAttr(ModAttr);
+    else if (F.hasFnAttribute(ModAttr))
+      F.removeFnAttr(ModAttr);
+  };
+
+  AddRemoveAttributeAsSet(BPI.BranchTargetEnforcement,
+                          "branch-target-enforcement");
+  AddRemoveAttributeAsSet(BPI.BranchProtectionPAuthLR,
+                          "branch-protection-pauth-lr");
+  AddRemoveAttributeAsSet(BPI.GuardedControlStack, "guarded-control-stack");
 }
 
-void TargetCodeGenInfo::setBranchProtectionFnAttributes(
+void TargetCodeGenInfo::initBranchProtectionFnAttributes(
     const TargetInfo::BranchProtectionInfo &BPI, llvm::AttrBuilder &FuncAttrs) {
+  // Only used for initializing attributes in the AttrBuilder, which will not
+  // contain any of these attributes so no need to remove anything.
   if (BPI.SignReturnAddr != LangOptions::SignReturnAddressScopeKind::None) {
     FuncAttrs.addAttribute("sign-return-address", BPI.getSignReturnAddrStr());
     FuncAttrs.addAttribute("sign-return-address-key", BPI.getSignKeyStr());

clang/lib/CodeGen/TargetInfo.h

@@ -423,13 +423,16 @@ class TargetCodeGenInfo {
     return nullptr;
   }
 
+  // Set the Branch Protection Attributes of the Function accordingly to the
+  // BPI. Remove attributes that contradict with current BPI.
   static void
   setBranchProtectionFnAttributes(const TargetInfo::BranchProtectionInfo &BPI,
                                   llvm::Function &F);
 
+  // Add the Branch Protection Attributes of the FuncAttrs.
   static void
-  setBranchProtectionFnAttributes(const TargetInfo::BranchProtectionInfo &BPI,
-                                  llvm::AttrBuilder &FuncAttrs);
+  initBranchProtectionFnAttributes(const TargetInfo::BranchProtectionInfo &BPI,
+                                   llvm::AttrBuilder &FuncAttrs);
 
 protected:
   static std::string qualifyWindowsLibrary(StringRef Lib);

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

@@ -1,6 +1,18 @@
 // REQUIRES: aarch64-registered-target
 // RUN: %clang_cc1 -triple aarch64 -emit-llvm  -target-cpu generic -target-feature +v8.5a %s -o - \
 // RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple aarch64 -emit-llvm  -target-cpu generic -target-feature +v8.5a -mbranch-target-enforce %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple aarch64 -emit-llvm  -target-cpu generic -target-feature +v8.5a -mguarded-control-stack %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple aarch64 -emit-llvm  -target-cpu generic -target-feature +v8.5a -msign-return-address=non-leaf -msign-return-address-key=a_key %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple aarch64 -emit-llvm  -target-cpu generic -target-feature +v8.5a -msign-return-address=all -msign-return-address-key=b_key %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple aarch64 -emit-llvm  -target-cpu generic -target-feature +v8.5a -mbranch-protection-pauth-lr -msign-return-address=all -msign-return-address-key=a_key %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple aarch64 -emit-llvm  -target-cpu generic -target-feature +v8.5a -mguarded-control-stack -mbranch-target-enforce -mbranch-protection-pauth-lr -msign-return-address=all -msign-return-address-key=a_key %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
 
 __attribute__ ((target("branch-protection=none")))
 void none() {}
@@ -83,7 +95,6 @@ void gcs() {}
 
 // CHECK-DAG: attributes #[[#BTIPACLEAF]] = { {{.*}} "branch-target-enforcement" {{.*}}"sign-return-address"="all" "sign-return-address-key"="a_key"
 
-
 // CHECK-DAG: attributes #[[#PAUTHLR]] = { {{.*}} "branch-protection-pauth-lr" {{.*}}"sign-return-address"="non-leaf" "sign-return-address-key"="a_key"
 
 // CHECK-DAG: attributes #[[#PAUTHLR_BKEY]] = { {{.*}} "branch-protection-pauth-lr" {{.*}}"sign-return-address"="non-leaf" "sign-return-address-key"="b_key"

clang/test/CodeGen/arm-branch-protection-attr-1.c

@@ -1,6 +1,12 @@
 // REQUIRES: arm-registered-target
 // RUN: %clang_cc1 -triple thumbv7m-unknown-unknown-eabi -emit-llvm %s -o - \
 // RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple thumbv7m-unknown-unknown-eabi -mbranch-target-enforce -emit-llvm %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple thumbv7m-unknown-unknown-eabi -msign-return-address=all -emit-llvm %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
+// RUN: %clang_cc1 -triple thumbv7m-unknown-unknown-eabi -mbranch-target-enforce -msign-return-address=all -emit-llvm %s -o - \
+// RUN:                               | FileCheck %s --check-prefix=CHECK
 
 __attribute__((target("branch-protection=none"))) void none() {}
 // CHECK: define{{.*}} void @none() #[[#NONE:]]