Reapply "[XRay][AArch64] Support -fxray-shared (#114431)" (#115300)

This patch implements support for `-fxray-shared` on AArch64 and fixes a
remaining issue in the previous PR #114431.

A bug in the XRay `CMakeLists.txt` caused the XRay assembly sources to
be built for every architecture in `XRAY_DSO_SUPPORTED_ARCH` on Apple.
This led to the compiler trying to compile AArch64 assembly for X86
targets and vice versa.
This is addressed here by ensuring that assembly sources are only built
for the matching architecture (see fixup commit).

**Original PR description:**
This patch adds support for `-fxray-shared` on AArch64. This feature,
introduced in #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

clang/lib/Driver/XRayArgs.cpp

@@ -68,8 +68,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
@@ -241,7 +245,7 @@ if (APPLE)
     if (${arch} IN_LIST XRAY_DSO_SUPPORTED_ARCH)
       add_compiler_rt_object_libraries(RTXrayDSO_${arch}
         OS ${XRAY_SUPPORTED_OS}
-        ARCHS ${XRAY_DSO_SUPPORTED_ARCH}
+        ARCHS ${arch}
         SOURCES ${${arch}_DSO_SOURCES}
         ADDITIONAL_HEADERS ${XRAY_IMPL_HEADERS}
         CFLAGS ${XRAY_CFLAGS}
@@ -403,7 +407,6 @@ else() # not Apple
       PARENT_TARGET xray)
 
     if (${arch} IN_LIST XRAY_DSO_SUPPORTED_ARCH)
-      # TODO: Only implemented for X86 at the moment
       add_compiler_rt_object_libraries(RTXrayDSO
         ARCHS ${arch}
         SOURCES ${XRAY_DSO_SOURCES} ${${arch}_DSO_SOURCES} 

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
 

compiler-rt/test/xray/TestCases/Posix/patching-unpatching-dso.cpp

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