[llvm][AArch64] Autoupgrade function attributes from Module attributes.

[DanielKristofKiss]

Refactored version of (#80640).
Run the attribute copy only during IRMove.

[llvmbot]

@llvm/pr-subscribers-llvm-ir

@llvm/pr-subscribers-lto

Author: Dani (DanielKristofKiss)

Changes

Refactored version of (#80640).
Run the attribute copy only during IRMove.

---

Full diff: https://github.com/llvm/llvm-project/pull/82763.diff

6 Files Affected:

• (modified) llvm/include/llvm/IR/AutoUpgrade.h (+3)
• (modified) llvm/lib/IR/AutoUpgrade.cpp (+66)
• (modified) llvm/lib/Linker/IRMover.cpp (+5)
• (modified) llvm/test/LTO/AArch64/link-branch-target-enforcement.ll (+1)
• (added) llvm/test/LTO/AArch64/link-sign-return-address.ll (+43)
• (modified) llvm/test/Linker/link-arm-and-thumb.ll (+4-3)

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 b90bbe71ac1896..3b784b06397445 100644
--- a/llvm/lib/IR/AutoUpgrade.cpp
+++ b/llvm/lib/IR/AutoUpgrade.cpp
@@ -5169,6 +5169,72 @@ 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->getZExtValue();
+}
+
+// Copy an attribute from module to the function if exists.
+// First value of the pair is used when the module attribute is not zero
+// the second otherwise.
+static void
+CopyModuleAttributeToFunction(Function &F, StringRef FnAttrName,
+                              StringRef ModAttrName,
+                              std::pair<StringRef, StringRef> Values) {
+  if (F.hasFnAttribute(FnAttrName))
+    return;
+  F.addFnAttr(FnAttrName, isModuleAttributeSet(*F.getParent(), ModAttrName)
+                              ? Values.first
+                              : Values.second);
+}
+
+// Copy a boolean attribute from module to the function if exists.
+// Module attribute treated false if zero otherwise true.
+static void CopyModuleAttributeToFunction(Function &F, StringRef AttrName) {
+  CopyModuleAttributeToFunction(
+      F, AttrName, AttrName,
+      std::make_pair<StringRef, StringRef>("true", "false"));
+}
+
+// Copy an attribute from module to the function if exists.
+// First value of the pair is used when the module attribute is not zero
+// the second otherwise.
+static void
+CopyModuleAttributeToFunction(Function &F, StringRef AttrName,
+                              std::pair<StringRef, StringRef> Values) {
+  CopyModuleAttributeToFunction(F, AttrName, AttrName, Values);
+}
+
+void llvm::CopyModuleAttrToFunctions(Module &M) {
+  Triple T(M.getTargetTriple());
+  if (!T.isThumb() && !T.isARM() && !T.isAArch64())
+    return;
+
+  for (Function &F : M.getFunctionList()) {
+    if (F.isDeclaration())
+      continue;
+
+    if (!F.hasFnAttribute("sign-return-address")) {
+      StringRef SignType = "none";
+      if (isModuleAttributeSet(M, "sign-return-address"))
+        SignType = "non-leaf";
+
+      if (isModuleAttributeSet(M, "sign-return-address-all"))
+        SignType = "all";
+
+      F.addFnAttr("sign-return-address", SignType);
+    }
+    CopyModuleAttributeToFunction(F, "branch-target-enforcement");
+    CopyModuleAttributeToFunction(F, "branch-protection-pauth-lr");
+    CopyModuleAttributeToFunction(F, "guarded-control-stack");
+    CopyModuleAttributeToFunction(
+        F, "sign-return-address-key",
+        std::make_pair<StringRef, StringRef>("b_key", "a_key"));
+  }
+}
+
 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 37d21119447b9c..026d466ed8539e 100644
--- a/llvm/lib/Linker/IRMover.cpp
+++ b/llvm/lib/Linker/IRMover.cpp
@@ -1606,6 +1606,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.
+  CopyModuleAttributesToFunction(*SrcM);
+
   std::reverse(Worklist.begin(), Worklist.end());
   while (!Worklist.empty()) {
     GlobalValue *GV = Worklist.back();
diff --git a/llvm/test/LTO/AArch64/link-branch-target-enforcement.ll b/llvm/test/LTO/AArch64/link-branch-target-enforcement.ll
index ccf8cf67ede6dc..74d9c86881d52c 100644
--- a/llvm/test/LTO/AArch64/link-branch-target-enforcement.ll
+++ b/llvm/test/LTO/AArch64/link-branch-target-enforcement.ll
@@ -32,6 +32,7 @@ entry:
 ; CHECK-DUMP: <main>:
 ; CHECK-DUMP:      bl      0x8 <main+0x8>
 ; CHECK-DUMP: <foo>:
+; CHECK-DUMP:     paciasp
 
 ; `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..c25857ceed7b40
--- /dev/null
+++ b/llvm/test/LTO/AArch64/link-sign-return-address.ll
@@ -0,0 +1,43 @@
+; Testcase to check that module with different branch-target-enforcement 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:     paciasp
+; CHECK-DUMP:     mov     w0, #0x2a
+; CHECK-DUMP:     autiasp
+; 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..37bd8c37f8b5e5 100644
--- a/llvm/test/Linker/link-arm-and-thumb.ll
+++ b/llvm/test/Linker/link-arm-and-thumb.ll
@@ -13,11 +13,12 @@ entry:
   ret i32 %add
 }
 
-; CHECK: define i32 @main() {
+; CHECK: define i32 @main() [[MAIN_ATTRS:#[0-9]+]]
 ; 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 [[MAIN_ATTRS]] = { {{.*}} }
+; 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:

[aeubanks]

we seem to be seeing large performance regressions from this patch on aarch64, any idea what could be going on? we'll also look into why this patch might be regressing performance

[DanielKristofKiss]

aeubanks it adds attributes to function that was not there before but in certain cases needed.
I have a series of patches to change how clang and llvm handles these attributes so once those lands I plan to tune this behaviour as well. #83154.

[partaror]

The original pull-request description states:

sign-return-address and similar module attributes should be propagated to
the function level before got merged because module flags may contradict and
this information is not recoverable.

In this patch, why are we adding attributes to the function which are neither present in the function and nor the module? This patch adds function attribute with value false, if an attribute is neither present in the function and nor the module. This behavior does not conform to the PR description -- attributes which were never present are being propagated. This behavior seems incorrect because default value for the missing attributes might be true in some cases.

[nickdesaulniers]

Hi @DanielKristofKiss , I think we should back this out again, based on the test case provided in #84494 (comment). (cc @pranavk @dhoekwater).

I think we need to revisit the design of these function attributes. In particular function attributes with binary values are problematic. (as in abstractly "foo"="true" or concretely "guarded-control-stack"="true"). For inlining, we'd like to avoid inlining when caller and callee differ in function attribute value. (What color is your function comes to mind).

But, by having the "binary" function attribute value take the form of "foo"="true", you haven't actually created a "binary" function attribute but actually a "ternary" function attribute, because now "foo"="false" and the complete lack of the foo function attribute are distinct in a way that's not clear if they mean the same thing or not. (Using "what color is your function" terminology, what we wanted was red or blue functions. What we got is red, blue, or green functions). To truly have binary function values, the attribute should simply take the form of "foo" (no ="true" or ="false"). The presence of the attribute on the function is implicitly "true" and the lack of attribute is implicitly "false".

Otherwise, if it's intentional that the function attribute have 3 possible values, we need to revisit the inlining behavior. Rather than do so trying to fix forward this, I think for the purposes of having fewer backports for the clang-18 branch, it would be better to:

1. back this out.
2. commit the test case from @pranavk linked above.
3. rework the design of the function attributes to be boolean rather than ternary
4. reland this with updated design for the function attributes.

WDYT?

[DanielKristofKiss]

Hi @nickdesaulniers,

Ok, let's back this out.

These patches changes the flags to "binary" fronted #82819 #83277 and backed #83154. ( with Chromium we got the same binary before/after). In the original design of the attribute the "ternary" behaviour is just the consequence of the dependency on the module attribute.

#84804 updates the module value, so the upgrader could detect the "old" module flags and easily do the right import.

[DanielKristofKiss]

So with llvm-lto all modules moved into an empty temp module but with lld.ld a the mover is initialised with the real module which never goes thru the update ( since this upgrade is removed from the bitcode reader update path )

IRMover::IRMover( 
...
    for (const auto *MD : StructTypes.getVisitedMetadata()) {
     SharedMDs[MD].reset(const_cast<MDNode *>(MD));
   }
+
+  CopyModuleAttrToFunctions(M);
 }

after this the reproducer looks good.