[LLVM] Autoupgrade function attributes from Module attributes.

DanielKristofKiss · DanielKristofKiss · commit 8ecca581e395 · 2024-03-19T09:35:47.000+01:00
Refactoring #82763 to cache module attributes.
diff --git a/lld/test/ELF/lto/aarch64_inline.ll 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
@@ -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
@@ -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->isOne();
+}
+
+// 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
@@ -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
@@ -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
@@ -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
diff --git 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:
