[XRay][AArch64] Support -fxray-shared (llvm#114431)

This patch adds support for `-fxray-shared` on AArch64. This feature,
introduced in llvm#113548 for x86_64, enables the instrumentation of shared
libraries with XRay.

Changes:
- Adds AArch64 to the list of targets supporting `-fxray-shared`
- Introduces PIC versions of the AArch64 XRay trampolines 
- Adjusts relevant XRay tests

Author: sebastiankreutzer

clang/lib/Driver/XRayArgs.cpp

@@ -67,8 +67,12 @@ XRayArgs::XRayArgs(const ToolChain &TC, const ArgList &Args) {
                    false)) {
     XRayShared = true;
 
-    // DSO instrumentation is currently limited to x86_64
-    if (Triple.getArch() != llvm::Triple::x86_64) {
+    // Certain targets support DSO instrumentation
+    switch (Triple.getArch()) {
+    case llvm::Triple::aarch64:
+    case llvm::Triple::x86_64:
+      break;
+    default:
       D.Diag(diag::err_drv_unsupported_opt_for_target)
           << "-fxray-shared" << Triple.str();
     }

clang/test/Driver/XRay/xray-shared.cpp

@@ -1,15 +1,21 @@
+// Check supported targets
 // RUN: %clang -### --target=x86_64-unknown-linux-gnu -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s
+// RUN: %clang -### --target=aarch64-unknown-linux-gnu -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s
+
+// Check unsupported targets
+// RUN: not %clang -### --target=arm-unknown-linux-gnu -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-TARGET
+// RUN: not %clang -### --target=mips-unknown-linux-gnu -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-TARGET
+// RUN: not %clang -### --target=loongarch64-unknown-linux-gnu -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-TARGET
+// RUN: not %clang -### --target=hexagon-unknown-linux-gnu -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-TARGET
+// RUN: not %clang -### --target=powerpc64le-unknown-linux-gnu -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-TARGET
+
+// Check PIC requirement
 // RUN: %clang -### --target=x86_64-unknown-linux-gnu -fpic -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s
 // RUN: not %clang -### --target=x86_64-unknown-linux-gnu -fno-PIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-PIC
 // RUN: not %clang -### --target=x86_64-unknown-linux-gnu -fno-pic -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-PIC
-
 // On 64 bit darwin, PIC is always enabled
 // RUN: %clang -### --target=x86_64-apple-darwin -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s
 
-// Check unsupported targets
-// RUN: not %clang -### --target=aarch64-pc-freebsd -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-TARGET
-// RUN: not %clang -### --target=arm64-apple-macos -fPIC -fxray-instrument -fxray-shared -c %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR-TARGET
-
 // CHECK: "-cc1" {{.*}}"-fxray-instrument" {{.*}}"-fxray-shared"
 // ERR-TARGET:   error: unsupported option '-fxray-shared' for target
 // ERR-PIC:   error: option '-fxray-shared' cannot be specified without '-fPIC'

compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake

@@ -104,7 +104,7 @@ else()
 set(ALL_XRAY_SUPPORTED_ARCH ${X86_64} ${ARM32} ${ARM64} ${MIPS32} ${MIPS64}
 		powerpc64le ${HEXAGON} ${LOONGARCH64})
 endif()
-set(ALL_XRAY_DSO_SUPPORTED_ARCH ${X86_64})
+set(ALL_XRAY_DSO_SUPPORTED_ARCH ${X86_64} ${ARM64})
 set(ALL_SHADOWCALLSTACK_SUPPORTED_ARCH ${ARM64})
 
 if (UNIX)

compiler-rt/lib/xray/CMakeLists.txt

@@ -56,6 +56,10 @@ set(aarch64_SOURCES
   xray_trampoline_AArch64.S
   )
 
+set(aarch64_DSO_SOURCES
+  xray_trampoline_AArch64.S
+  )
+
 set(loongarch64_SOURCES
   xray_loongarch64.cpp
   xray_trampoline_loongarch64.S

compiler-rt/lib/xray/xray_trampoline_AArch64.S

@@ -26,6 +26,17 @@
   ldp x1, x2, [sp], #16
 .endm
 
+.macro LOAD_HANDLER_ADDR reg handler
+#if !defined(XRAY_PIC)
+    adrp \reg, ASM_SYMBOL(\handler)
+    ldr \reg, [\reg, :lo12:ASM_SYMBOL(\handler)]
+#else
+    adrp \reg, :got:ASM_SYMBOL(\handler)
+    ldr \reg, [\reg, :got_lo12:ASM_SYMBOL(\handler)]
+    ldr \reg, [\reg]
+#endif
+.endm
+
 .text
 .p2align 2
 .global ASM_SYMBOL(__xray_FunctionEntry)
@@ -42,8 +53,7 @@ ASM_SYMBOL(__xray_FunctionEntry):
   SAVE_REGISTERS
 
   // Load the handler function pointer.
-  adrp x2, ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)
-  ldr x2, [x2, #:lo12:ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)]
+  LOAD_HANDLER_ADDR x2, _ZN6__xray19XRayPatchedFunctionE
   cbz x2, 1f
   // Set w0 to the function ID (w17). Set x1 to XRayEntryType::ENTRY = 0.
   mov w0, w17
@@ -69,8 +79,7 @@ ASM_SYMBOL(__xray_FunctionExit):
   SAVE_REGISTERS
 
   // Load the handler function pointer into x2.
-  adrp x2, ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)
-  ldr x2, [x2, #:lo12:ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)]
+  LOAD_HANDLER_ADDR x2, _ZN6__xray19XRayPatchedFunctionE
   cbz x2, 1f
   // Set w0 to the function ID (w17). Set x1 to XRayEntryType::EXIT = 1.
   mov w0, w17
@@ -96,8 +105,7 @@ ASM_SYMBOL(__xray_FunctionTailExit):
   // Save the registers which may be modified by the handler function.
   SAVE_REGISTERS
   // Load the handler function pointer into x2.
-  adrp x2, ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)
-  ldr x2, [x2, #:lo12:ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)]
+  LOAD_HANDLER_ADDR x2, _ZN6__xray19XRayPatchedFunctionE
   cbz x2, 1f
   // Set w0 to the function ID (w17). Set x1 to XRayEntryType::TAIL = 2.
   mov w0, w17
@@ -118,13 +126,11 @@ ASM_SYMBOL(__xray_ArgLoggerEntry):
   // Push the registers which may be modified by the handler function.
   SAVE_REGISTERS
 
-  adrp x8, ASM_SYMBOL(_ZN6__xray13XRayArgLoggerE)
-  ldr x8, [x8, #:lo12:ASM_SYMBOL(_ZN6__xray13XRayArgLoggerE)]
+  LOAD_HANDLER_ADDR x8, _ZN6__xray13XRayArgLoggerE
   cbnz x8, 2f
 
   // Load the handler function pointer.
-  adrp x8, ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)
-  ldr x8, [x8, #:lo12:ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)]
+  LOAD_HANDLER_ADDR x8, _ZN6__xray19XRayPatchedFunctionE
   cbz x8, 1f
 
 2:
@@ -144,8 +150,7 @@ ASM_SIZE(__xray_ArgLoggerEntry)
 ASM_TYPE_FUNCTION(__xray_CustomEvent)
 ASM_SYMBOL(__xray_CustomEvent):
   SAVE_REGISTERS
-  adrp x8, ASM_SYMBOL(_ZN6__xray22XRayPatchedCustomEventE)
-  ldr x8, [x8, #:lo12:ASM_SYMBOL(_ZN6__xray22XRayPatchedCustomEventE)]
+  LOAD_HANDLER_ADDR x8, _ZN6__xray22XRayPatchedCustomEventE
   cbz x8, 1f
   blr x8
 1:
@@ -157,8 +162,7 @@ ASM_SIZE(__xray_CustomEvent)
 ASM_TYPE_FUNCTION(__xray_TypedEvent)
 ASM_SYMBOL(__xray_TypedEvent):
   SAVE_REGISTERS
-  adrp x8, ASM_SYMBOL(_ZN6__xray21XRayPatchedTypedEventE)
-  ldr x8, [x8, #:lo12:ASM_SYMBOL(_ZN6__xray21XRayPatchedTypedEventE)]
+  LOAD_HANDLER_ADDR x8, _ZN6__xray21XRayPatchedTypedEventE
   cbz x8, 1f
   blr x8
 1:

compiler-rt/test/xray/TestCases/Posix/basic-mode-dso.cpp

@@ -8,7 +8,8 @@
 // RUN: %llvm_xray account --format=csv --sort=funcid "`ls basic-mode-dso-* | head -1`" | FileCheck --check-prefix=ACCOUNT %s
 // RUN: rm basic-mode-dso-*
 
-// REQUIRES: target=x86_64{{.*}}
+// REQUIRES: target={{(aarch64|x86_64)-.*}}
+// REQUIRES: built-in-llvm-tree
 
 //--- main.cpp
 

compiler-rt/test/xray/TestCases/Posix/clang-xray-shared.cpp

@@ -6,7 +6,7 @@
 // RUN: %clangxx -fxray-instrument %s -shared -o %t.so
 // RUN: llvm-nm %t.so | FileCheck %s --check-prefix DISABLED
 //
-// REQUIRES: target=x86_64{{.*}}
+// REQUIRES: target={{(aarch64|x86_64)-.*}}
 
 [[clang::xray_always_instrument]] int always_instrumented() { return 42; }
 

compiler-rt/test/xray/TestCases/Posix/dlopen.cpp

@@ -7,7 +7,7 @@
 //
 // RUN: XRAY_OPTIONS="patch_premain=true" %run %t/main.o %t/testlib.so 2>&1 | FileCheck %s
 
-// REQUIRES: target=x86_64{{.*}}
+// REQUIRES: target={{(aarch64|x86_64)-.*}}
 
 //--- main.cpp
 

compiler-rt/test/xray/TestCases/Posix/dso-dep-chains.cpp

@@ -17,7 +17,7 @@
 //
 // RUN:  XRAY_OPTIONS="patch_premain=true" %run %t/main.o %t/testlibd.so %t/testlibe.so 2>&1 | FileCheck %s
 
-// REQUIRES: target=x86_64{{.*}}
+// REQUIRES: target={{(aarch64|x86_64)-.*}}
 
 //--- main.cpp
 

compiler-rt/test/xray/TestCases/Posix/patch-premain-dso.cpp

@@ -6,7 +6,7 @@
 
 // RUN: XRAY_OPTIONS="patch_premain=true,verbosity=1" %run %t/main.o 2>&1 | FileCheck %s
 
-// REQUIRES: target=x86_64{{.*}}
+// REQUIRES: target={{(aarch64|x86_64)-.*}}
 
 //--- main.cpp