[llvm] [llvm][ARM][AArch64] Add attributes to synthetic functions. (PR #83153)

[via llvm-commits]

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

>From d1f6a1f6851cd11178ca2474a5ae28c3026b94cc Mon Sep 17 00:00:00 2001
From: Daniel Kiss <daniel.kiss at arm.com>
Date: Tue, 27 Feb 2024 10:38:04 +0100
Subject: [PATCH 1/2] Add attributes to synthetic functions.

Module flags represent the original intention.
---
 llvm/lib/IR/Function.cpp | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/llvm/lib/IR/Function.cpp b/llvm/lib/IR/Function.cpp
index 056e4f31981a72..5e217cc5680186 100644
--- a/llvm/lib/IR/Function.cpp
+++ b/llvm/lib/IR/Function.cpp
@@ -65,6 +65,8 @@
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ModRef.h"
+#include "llvm/TargetParser/ARMTargetParser.h"
+#include "llvm/TargetParser/Triple.h"
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
@@ -378,6 +380,40 @@ Function *Function::createWithDefaultAttr(FunctionType *Ty,
   }
   if (M->getModuleFlag("function_return_thunk_extern"))
     B.addAttribute(Attribute::FnRetThunkExtern);
+
+  Triple T(M->getTargetTriple());
+
+  auto Profile = llvm::ARM::parseArchProfile(T.getArchName());
+  if ((T.isArmT32() && Profile == llvm::ARM::ProfileKind::M) || T.isAArch64()) {
+    // Check if the module attribute is present and not zero.
+    auto isModuleAttributeSet = [&](const StringRef &ModAttr) -> bool {
+      const auto *Attr =
+          mdconst::extract_or_null<ConstantInt>(M->getModuleFlag(ModAttr));
+      return Attr && Attr->getZExtValue();
+    };
+
+    auto AddAttributeIfSet = [&](const StringRef &ModAttr) {
+      if (isModuleAttributeSet(ModAttr))
+        B.addAttribute(ModAttr, "true");
+    };
+
+    StringRef SignType = "none";
+    if (isModuleAttributeSet("sign-return-address"))
+      SignType = "non-leaf";
+    if (isModuleAttributeSet("sign-return-address-all"))
+      SignType = "all";
+    if (SignType != "none") {
+      B.addAttribute("sign-return-address", SignType);
+      if (T.isAArch64())
+        B.addAttribute("sign-return-address-key",
+                       isModuleAttributeSet("sign-return-address-with-bkey")
+                           ? "b_key"
+                           : "a_key");
+    }
+    AddAttributeIfSet("branch-target-enforcement");
+    AddAttributeIfSet("branch-protection-pauth-lr");
+    AddAttributeIfSet("guarded-control-stack");
+  }
   F->addFnAttrs(B);
   return F;
 }

>From 2a46897de89163ac9d0f1aa3b56ec59be8722319 Mon Sep 17 00:00:00 2001
From: Daniel Kiss <daniel.kiss at arm.com>
Date: Wed, 28 Feb 2024 15:17:46 +0100
Subject: [PATCH 2/2] fixup! [llvm][ARM][AArch64] Add attributes attributes to
 emitted functions.

---
 .../AddressSanitizer/module-flags-aarch64.ll  | 31 +++++++++++++++++++
 .../AArch64/module-flags-aarch64.ll           | 31 +++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 llvm/test/Instrumentation/AddressSanitizer/module-flags-aarch64.ll
 create mode 100644 llvm/test/Instrumentation/MemorySanitizer/AArch64/module-flags-aarch64.ll

diff --git a/llvm/test/Instrumentation/AddressSanitizer/module-flags-aarch64.ll b/llvm/test/Instrumentation/AddressSanitizer/module-flags-aarch64.ll
new file mode 100644
index 00000000000000..17934857c8b17b
--- /dev/null
+++ b/llvm/test/Instrumentation/AddressSanitizer/module-flags-aarch64.ll
@@ -0,0 +1,31 @@
+; RUN: opt < %s -passes=asan -S | FileCheck %s
+; REQUIRES: aarch64-registered-target
+
+target triple = "aarch64-unknown-linux-gnu"
+
+ at g = dso_local global i32 0, align 4
+
+define i32 @test_load() sanitize_address {
+entry:
+  %tmp = load i32, ptr @g, align 4
+  ret i32 %tmp
+}
+
+!llvm.module.flags = !{!0, !1}
+
+;; Due to -fasynchronous-unwind-tables.
+!0 = !{i32 7, !"uwtable", i32 2}
+
+;; Due to -fno-omit-frame-pointer.
+!1 = !{i32 7, !"frame-pointer", i32 2}
+
+!llvm.module.flags = !{!2, !3, !4}
+
+!2 = !{i32 8, !"branch-target-enforcement", i32 1}
+!3 = !{i32 8, !"sign-return-address", i32 1}
+!4 = !{i32 8, !"sign-return-address-all", i32 0}
+
+;; Set the uwtable attribute on ctor/dtor.
+; CHECK: define internal void @asan.module_ctor() #[[#ATTR:]]
+; CHECK: define internal void @asan.module_dtor() #[[#ATTR]]
+; CHECK: attributes #[[#ATTR]] = { nounwind uwtable "branch-target-enforcement"="true" "frame-pointer"="all" "sign-return-address"="non-leaf" "sign-return-address-key"="a_key" }
diff --git a/llvm/test/Instrumentation/MemorySanitizer/AArch64/module-flags-aarch64.ll b/llvm/test/Instrumentation/MemorySanitizer/AArch64/module-flags-aarch64.ll
new file mode 100644
index 00000000000000..da725484a6eefd
--- /dev/null
+++ b/llvm/test/Instrumentation/MemorySanitizer/AArch64/module-flags-aarch64.ll
@@ -0,0 +1,31 @@
+; RUN: opt < %s -passes=asan -S | FileCheck %s
+; REQUIRES: aarch64-registered-target
+
+target triple = "aarch64-unknown-linux-gnu"
+
+ at g = dso_local global i32 0, align 4
+
+define i32 @test_load() sanitize_memory {
+entry:
+  %tmp = load i32, ptr @g, align 4
+  ret i32 %tmp
+}
+
+!llvm.module.flags = !{!0, !1}
+
+;; Due to -fasynchronous-unwind-tables.
+!0 = !{i32 7, !"uwtable", i32 2}
+
+;; Due to -fno-omit-frame-pointer.
+!1 = !{i32 7, !"frame-pointer", i32 2}
+
+!llvm.module.flags = !{!2, !3, !4}
+
+!2 = !{i32 8, !"branch-target-enforcement", i32 1}
+!3 = !{i32 8, !"sign-return-address", i32 1}
+!4 = !{i32 8, !"sign-return-address-all", i32 0}
+
+;; Set the uwtable attribute on ctor/dtor.
+; CHECK: define internal void @asan.module_ctor() #[[#ATTR:]]
+; CHECK: define internal void @asan.module_dtor() #[[#ATTR]]
+; CHECK: attributes #[[#ATTR]] = { nounwind uwtable "branch-target-enforcement"="true" "frame-pointer"="all" "sign-return-address"="non-leaf" "sign-return-address-key"="a_key" }