[clang] [lld] [llvm] [ARM][AArch64] Change module flags values. (PR #84804)

Daniel Kiss via cfe-commits cfe-commits at lists.llvm.org
Tue Mar 19 04:18:28 PDT 2024


https://github.com/DanielKristofKiss updated https://github.com/llvm/llvm-project/pull/84804

>From e0596b2e216c041cfeb63faa8cf6d31a2601934f Mon Sep 17 00:00:00 2001
From: Daniel Kiss <daniel.kiss at arm.com>
Date: Fri, 8 Mar 2024 15:06:28 +0100
Subject: [PATCH 1/2] [LLVM] Autoupgrade function attributes from Module
 attributes.

Refactoring #82763 to cache module attributes.
---
 lld/test/ELF/lto/aarch64_inline.ll            | 51 +++++++++++++++++++
 llvm/include/llvm/IR/AutoUpgrade.h            |  3 ++
 llvm/lib/IR/AutoUpgrade.cpp                   | 47 +++++++++++++++++
 llvm/lib/Linker/IRMover.cpp                   | 10 ++++
 .../AArch64/link-branch-target-enforcement.ll |  3 ++
 .../LTO/AArch64/link-sign-return-address.ll   | 43 ++++++++++++++++
 llvm/test/Linker/link-arm-and-thumb.ll        |  6 +--
 7 files changed, 160 insertions(+), 3 deletions(-)
 create mode 100644 lld/test/ELF/lto/aarch64_inline.ll
 create mode 100644 llvm/test/LTO/AArch64/link-sign-return-address.ll

diff --git a/lld/test/ELF/lto/aarch64_inline.ll b/lld/test/ELF/lto/aarch64_inline.ll
new file mode 100644
index 00000000000000..781c283bc56a3e
--- /dev/null
+++ b/lld/test/ELF/lto/aarch64_inline.ll
@@ -0,0 +1,51 @@
+; REQUIRES: aarch64
+;; Test verifies inlining happens cross module when module flags are upgraded
+;; by the thin-lto linker/IRMover.
+;; Regression test for #82763
+
+; RUN: split-file %s %t
+; RUN: opt -thinlto-bc -thinlto-split-lto-unit -unified-lto %t/foo.s -o %t/foo.o
+; RUN: opt -thinlto-bc -thinlto-split-lto-unit -unified-lto %t/main.s -o %t/main.o
+; RUN: ld.lld -O2 --lto=thin --entry=main %t/main.o %t/foo.o -o %t/exe
+; RUN: llvm-objdump -d %t/exe | FileCheck %s
+
+
+; CHECK-LABEL:  <main>:
+; CHECK-NEXT:         pacibsp
+; CHECK-NEXT:         mov     w0, #0x22
+; CHECK-NEXT:         autibsp
+; CHECK-NEXT:         ret
+
+;--- foo.s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+define dso_local noundef i32 @foo() local_unnamed_addr #0 {
+entry:
+  ret i32 34
+}
+
+attributes #0 = { mustprogress nofree norecurse nosync nounwind willreturn memory(none) }
+!llvm.module.flags = !{!0, !1, !2, !3 }
+!0 = !{i32 8, !"branch-target-enforcement", i32 1}
+!1 = !{i32 8, !"sign-return-address", i32 1}
+!2 = !{i32 8, !"sign-return-address-all", i32 1}
+!3 = !{i32 8, !"sign-return-address-with-bkey", i32 1}
+
+;--- main.s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+declare i32 @foo();
+
+define i32 @main() {
+entry:
+  %1 = call i32 @foo()
+  ret i32 %1
+}
+
+!llvm.module.flags = !{!0, !1, !2, !3 }
+!0 = !{i32 8, !"branch-target-enforcement", i32 1}
+!1 = !{i32 8, !"sign-return-address", i32 1}
+!2 = !{i32 8, !"sign-return-address-all", i32 1}
+!3 = !{i32 8, !"sign-return-address-with-bkey", i32 1}
diff --git a/llvm/include/llvm/IR/AutoUpgrade.h b/llvm/include/llvm/IR/AutoUpgrade.h
index 152f781ffa9b30..1ef32bcb121bec 100644
--- a/llvm/include/llvm/IR/AutoUpgrade.h
+++ b/llvm/include/llvm/IR/AutoUpgrade.h
@@ -88,6 +88,9 @@ namespace llvm {
   /// info. Return true if module is modified.
   bool UpgradeDebugInfo(Module &M);
 
+  /// Copies module attributes to the functions in the module.
+  void CopyModuleAttrToFunctions(Module &M);
+
   /// Check whether a string looks like an old loop attachment tag.
   inline bool mayBeOldLoopAttachmentTag(StringRef Name) {
     return Name.starts_with("llvm.vectorizer.");
diff --git a/llvm/lib/IR/AutoUpgrade.cpp b/llvm/lib/IR/AutoUpgrade.cpp
index be0abb4b71dae2..ed7699c113584b 100644
--- a/llvm/lib/IR/AutoUpgrade.cpp
+++ b/llvm/lib/IR/AutoUpgrade.cpp
@@ -5178,6 +5178,53 @@ void llvm::UpgradeFunctionAttributes(Function &F) {
     Arg.removeAttrs(AttributeFuncs::typeIncompatible(Arg.getType()));
 }
 
+// Check if the module attribute is present and not zero.
+static bool isModuleAttributeSet(Module &M, const StringRef &ModAttr) {
+  const auto *Attr =
+      mdconst::extract_or_null<ConstantInt>(M.getModuleFlag(ModAttr));
+  return Attr && !Attr->isZero();
+}
+
+// Check if the function attribute is not present and set it.
+static void SetFunctionAttrIfNotSet(Function &F, StringRef FnAttrName,
+                                    StringRef Value) {
+  if (!F.hasFnAttribute(FnAttrName))
+    F.addFnAttr(FnAttrName, Value);
+}
+
+void llvm::CopyModuleAttrToFunctions(Module &M) {
+  Triple T(M.getTargetTriple());
+  if (!T.isThumb() && !T.isARM() && !T.isAArch64())
+    return;
+
+  StringRef SignTypeValue = "none";
+  if (isModuleAttributeSet(M, "sign-return-address-all"))
+    SignTypeValue = "all";
+  else if (isModuleAttributeSet(M, "sign-return-address"))
+    SignTypeValue = "non-leaf";
+
+  StringRef BTEValue =
+      isModuleAttributeSet(M, "branch-target-enforcement") ? "true" : "false";
+  StringRef BPPLValue =
+      isModuleAttributeSet(M, "branch-protection-pauth-lr") ? "true" : "false";
+  StringRef GCSValue =
+      isModuleAttributeSet(M, "guarded-control-stack") ? "true" : "false";
+  StringRef SignKeyValue =
+      isModuleAttributeSet(M, "sign-return-address-with-bkey") ? "b_key"
+                                                               : "a_key";
+
+  for (Function &F : M.getFunctionList()) {
+    if (F.isDeclaration())
+      continue;
+
+    SetFunctionAttrIfNotSet(F, "sign-return-address", SignTypeValue);
+    SetFunctionAttrIfNotSet(F, "branch-target-enforcement", BTEValue);
+    SetFunctionAttrIfNotSet(F, "branch-protection-pauth-lr", BPPLValue);
+    SetFunctionAttrIfNotSet(F, "guarded-control-stack", GCSValue);
+    SetFunctionAttrIfNotSet(F, "sign-return-address-key", SignKeyValue);
+  }
+}
+
 static bool isOldLoopArgument(Metadata *MD) {
   auto *T = dyn_cast_or_null<MDTuple>(MD);
   if (!T)
diff --git a/llvm/lib/Linker/IRMover.cpp b/llvm/lib/Linker/IRMover.cpp
index a7e6db82e5c23c..b019cd201313df 100644
--- a/llvm/lib/Linker/IRMover.cpp
+++ b/llvm/lib/Linker/IRMover.cpp
@@ -1623,6 +1623,11 @@ Error IRLinker::run() {
   // Loop over all of the linked values to compute type mappings.
   computeTypeMapping();
 
+  // Convert module level attributes to function level attributes because
+  // after merging modules the attributes might change and would have different
+  // effect on the functions as the original module would have.
+  CopyModuleAttrToFunctions(*SrcM);
+
   std::reverse(Worklist.begin(), Worklist.end());
   while (!Worklist.empty()) {
     GlobalValue *GV = Worklist.back();
@@ -1787,6 +1792,11 @@ IRMover::IRMover(Module &M) : Composite(M) {
   for (const auto *MD : StructTypes.getVisitedMetadata()) {
     SharedMDs[MD].reset(const_cast<MDNode *>(MD));
   }
+
+  // Convert module level attributes to function level attributes because
+  // after merging modules the attributes might change and would have different
+  // effect on the functions as the original module would have.
+  CopyModuleAttrToFunctions(M);
 }
 
 Error IRMover::move(std::unique_ptr<Module> Src,
diff --git a/llvm/test/LTO/AArch64/link-branch-target-enforcement.ll b/llvm/test/LTO/AArch64/link-branch-target-enforcement.ll
index ccf8cf67ede6dc..8313d812e189a4 100644
--- a/llvm/test/LTO/AArch64/link-branch-target-enforcement.ll
+++ b/llvm/test/LTO/AArch64/link-branch-target-enforcement.ll
@@ -30,8 +30,11 @@ entry:
 
 ; CHECK-NOT: linking module flags 'branch-target-enforcement': IDs have conflicting values in
 ; CHECK-DUMP: <main>:
+; CHECK-DUMP:      paciasp
+; CHECK-DUMP:      str
 ; CHECK-DUMP:      bl      0x8 <main+0x8>
 ; CHECK-DUMP: <foo>:
+; CHECK-DUMP:     pacibsp
 
 ; `main` doesn't support BTI while `foo` does, so in the binary
 ; we should see only PAC which is supported by both.
diff --git a/llvm/test/LTO/AArch64/link-sign-return-address.ll b/llvm/test/LTO/AArch64/link-sign-return-address.ll
new file mode 100644
index 00000000000000..271b98700d95c6
--- /dev/null
+++ b/llvm/test/LTO/AArch64/link-sign-return-address.ll
@@ -0,0 +1,43 @@
+; Testcase to check that module with different sign return address can
+; be mixed.
+;
+; RUN: llvm-as %s -o %t1.bc
+; RUN: llvm-as %p/Inputs/foo.ll -o %t2.bc
+; RUN: llvm-lto -exported-symbol main \
+; RUN:          -exported-symbol foo \
+; RUN:          -filetype=obj \
+; RUN:           %t2.bc %t1.bc \
+; RUN:           -o %t1.exe 2>&1
+; RUN: llvm-objdump -d %t1.exe | FileCheck --check-prefix=CHECK-DUMP %s
+; RUN: llvm-readelf -n %t1.exe | FileCheck --allow-empty --check-prefix=CHECK-PROP %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+declare i32 @foo();
+
+define i32 @main() {
+entry:
+  %add = call i32 @foo()
+  ret i32 %add
+}
+
+!llvm.module.flags = !{!0, !1, !2, !3 }
+!0 = !{i32 8, !"branch-target-enforcement", i32 0}
+!1 = !{i32 8, !"sign-return-address", i32 0}
+!2 = !{i32 8, !"sign-return-address-all", i32 0}
+!3 = !{i32 8, !"sign-return-address-with-bkey", i32 0}
+
+; CHECK-DUMP: <foo>:
+; CHECK-DUMP:     pacibsp
+; CHECK-DUMP:     mov     w0, #0x2a
+; CHECK-DUMP:     autibsp
+; CHECK-DUMP:     ret
+; CHECK-DUMP: <main>:
+; CHECK-DUMP-NOT:  paciasp
+; CHECK-DUMP:      str     x30,
+; CHECK-DUMP:      bl      0x14 <main+0x4>
+
+; `main` doesn't support PAC sign-return-address while `foo` does, so in the binary
+; we should not see anything.
+; CHECK-PROP-NOT:   Properties: aarch64 feature: PAC
\ No newline at end of file
diff --git a/llvm/test/Linker/link-arm-and-thumb.ll b/llvm/test/Linker/link-arm-and-thumb.ll
index a90f2128e4430a..b5984bf5579477 100644
--- a/llvm/test/Linker/link-arm-and-thumb.ll
+++ b/llvm/test/Linker/link-arm-and-thumb.ll
@@ -13,11 +13,11 @@ entry:
   ret i32 %add
 }
 
-; CHECK: define i32 @main() {
+; CHECK: define i32 @main()
 ; CHECK: define i32 @foo(i32 %a, i32 %b) [[ARM_ATTRS:#[0-9]+]]
 ; CHECK: define i32 @bar(i32 %a, i32 %b) [[THUMB_ATTRS:#[0-9]+]]
 
-; CHECK: attributes [[ARM_ATTRS]] = { "target-features"="-thumb-mode" }
-; CHECK: attributes [[THUMB_ATTRS]] = { "target-features"="+thumb-mode" }
+; CHECK: attributes [[ARM_ATTRS]] = {{{.*}}"target-features"="-thumb-mode" }
+; CHECK: attributes [[THUMB_ATTRS]] = {{{.*}}"target-features"="+thumb-mode" }
 
 ; STDERR-NOT: warning: Linking two modules of different target triples:

>From e3c65cc8965b3fde4163f4536f6729547e70078e Mon Sep 17 00:00:00 2001
From: Daniel Kiss <daniel.kiss at arm.com>
Date: Fri, 8 Mar 2024 15:22:30 +0100
Subject: [PATCH 2/2] [ARM][AArch64] Change module flags values.

Module flag used to indicate the feature to be propagated to the
function. As now the frontend emits all attributes accoringly let's
help the automerger to only do work when old and new bitcodes are
merged.
---
 clang/lib/CodeGen/CodeGenModule.cpp           | 19 ++++---
 .../CodeGen/aarch64-sign-return-address.c     | 12 ++---
 .../CodeGen/arm-branch-protection-attr-2.c    |  8 +--
 .../arm-ignore-branch-protection-option.c     |  2 +-
 lld/test/ELF/lto/aarch64_inline.ll            | 26 ++++++++--
 llvm/lib/IR/AutoUpgrade.cpp                   | 49 ++++++++++---------
 llvm/test/LTO/AArch64/Inputs/bar.ll           | 22 +++++++++
 .../LTO/AArch64/link-sign-return-address.ll   | 17 ++++++-
 8 files changed, 111 insertions(+), 44 deletions(-)
 create mode 100644 llvm/test/LTO/AArch64/Inputs/bar.ll

diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 8ceecff28cbc63..c0070b071e7a8e 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -1171,22 +1171,29 @@ void CodeGenModule::Release() {
                               "tag-stack-memory-buildattr", 1);
 
   if (T.isARM() || T.isThumb() || T.isAArch64()) {
+    // Previously 1 is used and meant for the backed to derive the function
+    // attribute form it. 2 now means function attributes already set for all
+    // functions in this module, so no need to propagate those from the module
+    // flag. Value is only used in case of LTO module merge because the backend
+    // will see all required function attribute set already. Value is used
+    // before modules got merged. Any posive value means the feature is active
+    // and required binary markings need to be emit accordingly.
     if (LangOpts.BranchTargetEnforcement)
       getModule().addModuleFlag(llvm::Module::Min, "branch-target-enforcement",
-                                1);
+                                2);
     if (LangOpts.BranchProtectionPAuthLR)
       getModule().addModuleFlag(llvm::Module::Min, "branch-protection-pauth-lr",
-                                1);
+                                2);
     if (LangOpts.GuardedControlStack)
-      getModule().addModuleFlag(llvm::Module::Min, "guarded-control-stack", 1);
+      getModule().addModuleFlag(llvm::Module::Min, "guarded-control-stack", 2);
     if (LangOpts.hasSignReturnAddress())
-      getModule().addModuleFlag(llvm::Module::Min, "sign-return-address", 1);
+      getModule().addModuleFlag(llvm::Module::Min, "sign-return-address", 2);
     if (LangOpts.isSignReturnAddressScopeAll())
       getModule().addModuleFlag(llvm::Module::Min, "sign-return-address-all",
-                                1);
+                                2);
     if (!LangOpts.isSignReturnAddressWithAKey())
       getModule().addModuleFlag(llvm::Module::Min,
-                                "sign-return-address-with-bkey", 1);
+                                "sign-return-address-with-bkey", 2);
   }
 
   if (CodeGenOpts.StackClashProtector)
diff --git a/clang/test/CodeGen/aarch64-sign-return-address.c b/clang/test/CodeGen/aarch64-sign-return-address.c
index 8bc54b1a56c38c..35c56889e07071 100644
--- a/clang/test/CodeGen/aarch64-sign-return-address.c
+++ b/clang/test/CodeGen/aarch64-sign-return-address.c
@@ -22,17 +22,17 @@
 // NONE-NOT:  !"branch-target-enforcement"
 // ALL-NOT:   !"branch-target-enforcement"
 // PART-NOT:  !"branch-target-enforcement"
-// BTE:       !{i32 8, !"branch-target-enforcement", i32 1}
+// BTE:       !{i32 8, !"branch-target-enforcement", i32 2}
 // B-KEY-NOT: !"branch-target-enforcement"
 
 // NONE-NOT:  !"sign-return-address"
-// ALL:   !{i32 8, !"sign-return-address", i32 1}
-// PART:  !{i32 8, !"sign-return-address", i32 1}
+// ALL:   !{i32 8, !"sign-return-address", i32 2}
+// PART:  !{i32 8, !"sign-return-address", i32 2}
 // BTE-NOT:   !"sign-return-address"
-// B-KEY: !{i32 8, !"sign-return-address", i32 1}
+// B-KEY: !{i32 8, !"sign-return-address", i32 2}
 
 // NONE-NOT:  !"sign-return-address-all"
-// ALL:   !{i32 8, !"sign-return-address-all", i32 1}
+// ALL:   !{i32 8, !"sign-return-address-all", i32 2}
 // PART-NOT:  !"sign-return-address-all"
 // BTE-NOT:   !"sign-return-address-all"
 // B-KEY-NOT: !"sign-return-address-all"
@@ -41,6 +41,6 @@
 // ALL-NOT:   !"sign-return-address-with-bkey"
 // PART-NOT:  !"sign-return-address-with-bkey"
 // BTE-NOT:   !"sign-return-address-with-bkey"
-// B-KEY: !{i32 8, !"sign-return-address-with-bkey", i32 1}
+// B-KEY: !{i32 8, !"sign-return-address-with-bkey", i32 2}
 
 void foo() {}
diff --git a/clang/test/CodeGen/arm-branch-protection-attr-2.c b/clang/test/CodeGen/arm-branch-protection-attr-2.c
index 1f3c00873043e8..741c0026c4d05c 100644
--- a/clang/test/CodeGen/arm-branch-protection-attr-2.c
+++ b/clang/test/CodeGen/arm-branch-protection-attr-2.c
@@ -18,16 +18,16 @@
 // NONE-NOT: !"branch-target-enforcement"
 // PART-NOT: !"branch-target-enforcement"
 // ALL-NOT:  !"branch-target-enforcement"
-// BTE:      !{i32 8, !"branch-target-enforcement", i32 1}
+// BTE:      !{i32 8, !"branch-target-enforcement", i32 2}
 
 // NONE-NOT: !"sign-return-address"
-// PART:     !{i32 8, !"sign-return-address", i32 1}
-// ALL:      !{i32 8, !"sign-return-address", i32 1}
+// PART:     !{i32 8, !"sign-return-address", i32 2}
+// ALL:      !{i32 8, !"sign-return-address", i32 2}
 // BTE-NOT:  !"sign-return-address"
 
 // NONE-NOT: !"sign-return-address-all", i32 0}
 // PART-NOT: !"sign-return-address-all", i32 0}
-// ALL:      !{i32 8, !"sign-return-address-all", i32 1}
+// ALL:      !{i32 8, !"sign-return-address-all", i32 2}
 // BTE-NOT:  !"sign-return-address-all", i32 0}
 
 void foo() {}
diff --git a/clang/test/Frontend/arm-ignore-branch-protection-option.c b/clang/test/Frontend/arm-ignore-branch-protection-option.c
index 99a2accef3ae2f..45bdb37f5ed1a3 100644
--- a/clang/test/Frontend/arm-ignore-branch-protection-option.c
+++ b/clang/test/Frontend/arm-ignore-branch-protection-option.c
@@ -15,4 +15,4 @@ __attribute__((target("arch=cortex-m0"))) void f() {}
 // CHECK-NOT:  attributes { {{.*}} "branch-target-enforcement"
 
 /// Check that there are branch protection module attributes despite the warning.
-// CHECK: !{i32 8, !"branch-target-enforcement", i32 1}
+// CHECK: !{i32 8, !"branch-target-enforcement", i32 2}
diff --git a/lld/test/ELF/lto/aarch64_inline.ll b/lld/test/ELF/lto/aarch64_inline.ll
index 781c283bc56a3e..c8b6d7a81ac424 100644
--- a/lld/test/ELF/lto/aarch64_inline.ll
+++ b/lld/test/ELF/lto/aarch64_inline.ll
@@ -5,14 +5,15 @@
 
 ; RUN: split-file %s %t
 ; RUN: opt -thinlto-bc -thinlto-split-lto-unit -unified-lto %t/foo.s -o %t/foo.o
+; RUN: opt -thinlto-bc -thinlto-split-lto-unit -unified-lto %t/bar.s -o %t/bar.o
 ; RUN: opt -thinlto-bc -thinlto-split-lto-unit -unified-lto %t/main.s -o %t/main.o
-; RUN: ld.lld -O2 --lto=thin --entry=main %t/main.o %t/foo.o -o %t/exe
+; RUN: ld.lld -O2 --lto=thin --entry=main %t/main.o %t/foo.o %t/bar.o -o %t/exe
 ; RUN: llvm-objdump -d %t/exe | FileCheck %s
 
 
 ; CHECK-LABEL:  <main>:
 ; CHECK-NEXT:         pacibsp
-; CHECK-NEXT:         mov     w0, #0x22
+; CHECK-NEXT:         mov     w0, #0x23
 ; CHECK-NEXT:         autibsp
 ; CHECK-NEXT:         ret
 
@@ -32,16 +33,35 @@ attributes #0 = { mustprogress nofree norecurse nosync nounwind willreturn memor
 !2 = !{i32 8, !"sign-return-address-all", i32 1}
 !3 = !{i32 8, !"sign-return-address-with-bkey", i32 1}
 
+;--- bar.s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+define dso_local noundef i32 @bar() local_unnamed_addr #0 {
+entry:
+  ret i32 1
+}
+
+attributes #0 = { mustprogress nofree norecurse nosync nounwind willreturn memory(none) "branch-target-enforcement"="true" "sign-return-address"="all" "sign-return-address-key"="b_key" }
+!llvm.module.flags = !{!0, !1, !2, !3 }
+!0 = !{i32 8, !"branch-target-enforcement", i32 2}
+!1 = !{i32 8, !"sign-return-address", i32 2}
+!2 = !{i32 8, !"sign-return-address-all", i32 2}
+!3 = !{i32 8, !"sign-return-address-with-bkey", i32 2}
+
 ;--- main.s
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-unknown-linux-gnu"
 
 declare i32 @foo();
+declare i32 @bar();
 
 define i32 @main() {
 entry:
   %1 = call i32 @foo()
-  ret i32 %1
+  %2 = call i32 @bar()
+  %3 = add i32 %1, %2
+  ret i32 %3
 }
 
 !llvm.module.flags = !{!0, !1, !2, !3 }
diff --git a/llvm/lib/IR/AutoUpgrade.cpp b/llvm/lib/IR/AutoUpgrade.cpp
index ed7699c113584b..3c97c60ec044d2 100644
--- a/llvm/lib/IR/AutoUpgrade.cpp
+++ b/llvm/lib/IR/AutoUpgrade.cpp
@@ -5178,11 +5178,11 @@ void llvm::UpgradeFunctionAttributes(Function &F) {
     Arg.removeAttrs(AttributeFuncs::typeIncompatible(Arg.getType()));
 }
 
-// Check if the module attribute is present and not zero.
-static bool isModuleAttributeSet(Module &M, const StringRef &ModAttr) {
+// Check if the module attribute is present and set to one.
+static bool isModuleAttributeOne(Module &M, const StringRef &ModAttr) {
   const auto *Attr =
       mdconst::extract_or_null<ConstantInt>(M.getModuleFlag(ModAttr));
-  return Attr && !Attr->isZero();
+  return Attr && Attr->isOne();
 }
 
 // Check if the function attribute is not present and set it.
@@ -5197,31 +5197,36 @@ void llvm::CopyModuleAttrToFunctions(Module &M) {
   if (!T.isThumb() && !T.isARM() && !T.isAArch64())
     return;
 
-  StringRef SignTypeValue = "none";
-  if (isModuleAttributeSet(M, "sign-return-address-all"))
+  bool BTE = isModuleAttributeOne(M, "branch-target-enforcement");
+  bool BPPLR = isModuleAttributeOne(M, "branch-protection-pauth-lr");
+  bool GCS = isModuleAttributeOne(M, "guarded-control-stack");
+  bool SRA = isModuleAttributeOne(M, "sign-return-address");
+
+  if (!BTE && !BPPLR && !GCS && !SRA)
+    return;
+
+  StringRef SignTypeValue = "non-leaf";
+  if (SRA && isModuleAttributeOne(M, "sign-return-address-all"))
     SignTypeValue = "all";
-  else if (isModuleAttributeSet(M, "sign-return-address"))
-    SignTypeValue = "non-leaf";
-
-  StringRef BTEValue =
-      isModuleAttributeSet(M, "branch-target-enforcement") ? "true" : "false";
-  StringRef BPPLValue =
-      isModuleAttributeSet(M, "branch-protection-pauth-lr") ? "true" : "false";
-  StringRef GCSValue =
-      isModuleAttributeSet(M, "guarded-control-stack") ? "true" : "false";
-  StringRef SignKeyValue =
-      isModuleAttributeSet(M, "sign-return-address-with-bkey") ? "b_key"
-                                                               : "a_key";
+
+  StringRef SignKeyValue = "a_key";
+  if (SRA && isModuleAttributeOne(M, "sign-return-address-with-bkey"))
+    SignKeyValue = "b_key";
 
   for (Function &F : M.getFunctionList()) {
     if (F.isDeclaration())
       continue;
 
-    SetFunctionAttrIfNotSet(F, "sign-return-address", SignTypeValue);
-    SetFunctionAttrIfNotSet(F, "branch-target-enforcement", BTEValue);
-    SetFunctionAttrIfNotSet(F, "branch-protection-pauth-lr", BPPLValue);
-    SetFunctionAttrIfNotSet(F, "guarded-control-stack", GCSValue);
-    SetFunctionAttrIfNotSet(F, "sign-return-address-key", SignKeyValue);
+    if (SRA) {
+      SetFunctionAttrIfNotSet(F, "sign-return-address", SignTypeValue);
+      SetFunctionAttrIfNotSet(F, "sign-return-address-key", SignKeyValue);
+    }
+    if (BTE)
+      SetFunctionAttrIfNotSet(F, "branch-target-enforcement", "true");
+    if (BPPLR)
+      SetFunctionAttrIfNotSet(F, "branch-protection-pauth-lr", "true");
+    if (GCS)
+      SetFunctionAttrIfNotSet(F, "guarded-control-stack", "true");
   }
 }
 
diff --git a/llvm/test/LTO/AArch64/Inputs/bar.ll b/llvm/test/LTO/AArch64/Inputs/bar.ll
new file mode 100644
index 00000000000000..1f0712660ebdc5
--- /dev/null
+++ b/llvm/test/LTO/AArch64/Inputs/bar.ll
@@ -0,0 +1,22 @@
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+define dso_local void @bar() #0 {
+entry:
+  ret void
+}
+
+define dso_local void @baz() #1 {
+entry:
+  ret void
+}
+
+attributes #0 = { noinline nounwind optnone uwtable }
+attributes #1 = { noinline nounwind optnone uwtable "branch-target-enforcement"="true" }
+
+!llvm.module.flags = !{!0, !1, !2, !3}
+
+!0 = !{i32 8, !"branch-target-enforcement", i32 2}
+!1 = !{i32 8, !"sign-return-address", i32 2}
+!2 = !{i32 8, !"sign-return-address-all", i32 2}
+!3 = !{i32 8, !"sign-return-address-with-bkey", i32 2}
diff --git a/llvm/test/LTO/AArch64/link-sign-return-address.ll b/llvm/test/LTO/AArch64/link-sign-return-address.ll
index 271b98700d95c6..d5ddb7e3155a11 100644
--- a/llvm/test/LTO/AArch64/link-sign-return-address.ll
+++ b/llvm/test/LTO/AArch64/link-sign-return-address.ll
@@ -3,10 +3,11 @@
 ;
 ; RUN: llvm-as %s -o %t1.bc
 ; RUN: llvm-as %p/Inputs/foo.ll -o %t2.bc
+; RUN: llvm-as %p/Inputs/bar.ll -o %t3.bc
 ; RUN: llvm-lto -exported-symbol main \
 ; RUN:          -exported-symbol foo \
 ; RUN:          -filetype=obj \
-; RUN:           %t2.bc %t1.bc \
+; RUN:           %t3.bc %t2.bc %t1.bc \
 ; RUN:           -o %t1.exe 2>&1
 ; RUN: llvm-objdump -d %t1.exe | FileCheck --check-prefix=CHECK-DUMP %s
 ; RUN: llvm-readelf -n %t1.exe | FileCheck --allow-empty --check-prefix=CHECK-PROP %s
@@ -15,10 +16,14 @@ target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-unknown-linux-gnu"
 
 declare i32 @foo();
+declare void @baz();
+declare void @bar();
 
 define i32 @main() {
 entry:
   %add = call i32 @foo()
+  call void @bar()
+  call void @baz()
   ret i32 %add
 }
 
@@ -28,6 +33,12 @@ entry:
 !2 = !{i32 8, !"sign-return-address-all", i32 0}
 !3 = !{i32 8, !"sign-return-address-with-bkey", i32 0}
 
+
+; CHECK-DUMP: <bar>:
+; CHECK-DUMP:     ret
+; CHECK-DUMP: <baz>:
+; CHECK-DUMP:     bti     c
+; CHECK-DUMP:     ret
 ; CHECK-DUMP: <foo>:
 ; CHECK-DUMP:     pacibsp
 ; CHECK-DUMP:     mov     w0, #0x2a
@@ -36,7 +47,9 @@ entry:
 ; CHECK-DUMP: <main>:
 ; CHECK-DUMP-NOT:  paciasp
 ; CHECK-DUMP:      str     x30,
-; CHECK-DUMP:      bl      0x14 <main+0x4>
+; CHECK-DUMP:      bl      0x20 <main+0x4>
+; CHECK-DUMP:      bl      0x0 <bar>
+; CHECK-DUMP:      bl      0x4 <baz>
 
 ; `main` doesn't support PAC sign-return-address while `foo` does, so in the binary
 ; we should not see anything.



More information about the cfe-commits mailing list