[llvm] 79685b5 - [JITLink][aarch64] Fix dependence tracking for Pointer64Authenticated edges.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2025-01-28T14:57:49+11:00
New Revision: 79685b59b78804f99e46f12ac888edbc0a6d1472

URL: https://github.com/llvm/llvm-project/commit/79685b59b78804f99e46f12ac888edbc0a6d1472
DIFF: https://github.com/llvm/llvm-project/commit/79685b59b78804f99e46f12ac888edbc0a6d1472.diff

LOG: [JITLink][aarch64] Fix dependence tracking for Pointer64Authenticated edges.

Transform Pointer64Authenticated edges into KeepAlive edges, rather than
removing them, in order to preserve symbol dependence information.

The lowerPointer64AuthEdgesToSigningFunction pass is responsible for
transforming Pointer64Authenticated edges to a signing function that can be run
in the executing process to initialize global PAC pointers. Removing the edges
entirely in this pass results in loss of dependence tracking, which can in turn
cause ORC to report PAC pointers as ready before the pointers targets have
completed materialization (resulting in a use-before-finalize error, often
manifesting as access to uninitialized / unprotected memory).

This commit addresses the issue by leaving the edges in the graph and simply
changing their kind to KeepAlive, a no-op for fixup purposes but followed for
dependence tracking purposes.

Added: 
    llvm/test/ExecutionEngine/JITLink/AArch64/MachO_ptrauth_dependencies.s

Modified: 
    llvm/lib/ExecutionEngine/JITLink/aarch64.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/ExecutionEngine/JITLink/aarch64.cpp b/llvm/lib/ExecutionEngine/JITLink/aarch64.cpp
index d101332dbc7e9f..c97cedca81e2b1 100644
--- a/llvm/lib/ExecutionEngine/JITLink/aarch64.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/aarch64.cpp
@@ -317,8 +317,7 @@ Error lowerPointer64AuthEdgesToSigningFunction(LinkGraph &G) {
   };
 
   for (auto *B : G.blocks()) {
-    for (auto EI = B->edges().begin(); EI != B->edges().end();) {
-      auto &E = *EI;
+    for (auto &E : B->edges()) {
       if (E.getKind() == aarch64::Pointer64Authenticated) {
         uint64_t EncodedInfo = E.getAddend();
         int32_t RealAddend = (uint32_t)(EncodedInfo & 0xffffffff);
@@ -358,10 +357,9 @@ Error lowerPointer64AuthEdgesToSigningFunction(LinkGraph &G) {
         // Store signed pointer.
         cantFail(writeStoreRegSeq(AppendInstr, Reg2, Reg1));
 
-        // Remove this edge.
-        EI = B->removeEdge(EI);
-      } else
-        ++EI;
+        // Replace edge with a keep-alive to preserve dependence info.
+        E.setKind(Edge::KeepAlive);
+      }
     }
   }
 

diff  --git a/llvm/test/ExecutionEngine/JITLink/AArch64/MachO_ptrauth_dependencies.s b/llvm/test/ExecutionEngine/JITLink/AArch64/MachO_ptrauth_dependencies.s
new file mode 100644
index 00000000000000..01618d78e094d4
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/AArch64/MachO_ptrauth_dependencies.s
@@ -0,0 +1,19 @@
+# RUN: llvm-mc -triple=arm64e-apple-macosx -filetype=obj -o %t.o %s
+# RUN: llvm-jitlink -num-threads=0 -debug-only=orc -noexec \
+# RUN:              -abs _foo=0x1 %t.o 2>&1 \
+# RUN:              | FileCheck %s
+#
+# Ensure that we don't lose dependence tracking information when ptrauth edges
+# are lowered: _main should still depend on _foo.
+
+# CHECK:    Symbols: { _main }, Dependencies: { (main, { _foo }) }
+
+	.section	__TEXT,__text,regular,pure_instructions
+
+	.section	__DATA,__data
+	.globl	_main
+	.p2align	3, 0x0
+_main:
+	.quad	_foo at AUTH(ia,0)
+
+.subsections_via_symbols