[XRay] Add initial support for loongarch64

Only support patching FunctionEntry/FunctionExit/FunctionTailExit for now.

Reviewed By: MaskRay, xen0n
Co-Authored-By: zhanglimin <zhanglimin@loongson.cn>

Differential Revision: https://reviews.llvm.org/D140727

8 files changed, 296 insertions, 4 deletions

diff --git a/compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake b/compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake
index a26ec96..70847d2 100644
--- a/compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake
+++ b/compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake
@@ -80,7 +80,7 @@ if(APPLE)
 set(ALL_XRAY_SUPPORTED_ARCH ${X86_64} ${ARM64})
 else()
 set(ALL_XRAY_SUPPORTED_ARCH ${X86_64} ${ARM32} ${ARM64} ${MIPS32} ${MIPS64}
-		powerpc64le ${HEXAGON})
+		powerpc64le ${HEXAGON} ${LOONGARCH64})
 endif()
 set(ALL_SHADOWCALLSTACK_SUPPORTED_ARCH ${ARM64})
 
diff --git a/compiler-rt/lib/xray/CMakeLists.txt b/compiler-rt/lib/xray/CMakeLists.txt
index e94d39a..cf7b506 100644
--- a/compiler-rt/lib/xray/CMakeLists.txt
+++ b/compiler-rt/lib/xray/CMakeLists.txt
@@ -47,6 +47,11 @@ set(aarch64_SOURCES
   xray_trampoline_AArch64.S
   )
 
+set(loongarch64_SOURCES
+  xray_loongarch64.cpp
+  xray_trampoline_loongarch64.S
+  )
+
 set(mips_SOURCES
   xray_mips.cpp
   xray_trampoline_mips.S
@@ -130,6 +135,7 @@ set(XRAY_ALL_SOURCE_FILES
   ${arm_SOURCES}
   ${armhf_SOURCES}
   ${hexagon_SOURCES}
+  ${loongarch64_SOURCES}
   ${mips_SOURCES}
   ${mipsel_SOURCES}
   ${mips64_SOURCES}
diff --git a/compiler-rt/lib/xray/xray_interface.cpp b/compiler-rt/lib/xray/xray_interface.cpp
index 095fc62..5839043 100644
--- a/compiler-rt/lib/xray/xray_interface.cpp
+++ b/compiler-rt/lib/xray/xray_interface.cpp
@@ -46,6 +46,8 @@ static const int16_t cSledLength = 12;
 static const int16_t cSledLength = 32;
 #elif defined(__arm__)
 static const int16_t cSledLength = 28;
+#elif SANITIZER_LOONGARCH64
+static const int16_t cSledLength = 48;
 #elif SANITIZER_MIPS32
 static const int16_t cSledLength = 48;
 #elif SANITIZER_MIPS64
diff --git a/compiler-rt/lib/xray/xray_loongarch64.cpp b/compiler-rt/lib/xray/xray_loongarch64.cpp
new file mode 100644
index 0000000..b839adb
--- /dev/null
+++ b/compiler-rt/lib/xray/xray_loongarch64.cpp
@@ -0,0 +1,160 @@
+//===-------- xray_loongarch64.cpp ------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of XRay, a dynamic runtime instrumentation system.
+//
+// Implementation of loongarch-specific routines.
+//
+//===----------------------------------------------------------------------===//
+#include "sanitizer_common/sanitizer_common.h"
+#include "xray_defs.h"
+#include "xray_interface_internal.h"
+#include <atomic>
+
+namespace __xray {
+
+enum RegNum : uint32_t {
+  RN_RA = 1,
+  RN_SP = 3,
+  RN_T0 = 12,
+  RN_T1 = 13,
+};
+
+// Encode instructions in the 2RIx format, where the primary formats here
+// are 2RI12-type and 2RI16-type.
+static inline uint32_t
+encodeInstruction2RIx(uint32_t Opcode, uint32_t Rd, uint32_t Rj,
+                      uint32_t Imm) XRAY_NEVER_INSTRUMENT {
+  return Opcode | (Imm << 10) | (Rj << 5) | Rd;
+}
+
+// Encode instructions in 1RI20 format, e.g. lu12i.w/lu32i.d.
+static inline uint32_t
+encodeInstruction1RI20(uint32_t Opcode, uint32_t Rd,
+                       uint32_t Imm) XRAY_NEVER_INSTRUMENT {
+  return Opcode | (Imm << 5) | Rd;
+}
+
+static inline bool patchSled(const bool Enable, const uint32_t FuncId,
+                             const XRaySledEntry &Sled,
+                             void (*TracingHook)()) XRAY_NEVER_INSTRUMENT {
+  // When |Enable| == true,
+  // We replace the following compile-time stub (sled):
+  //
+  // .Lxray_sled_beginN:
+  //	B .Lxray_sled_endN
+  //	11 NOPs (44 bytes)
+  // .Lxray_sled_endN:
+  //
+  // With the following runtime patch:
+  //
+  // xray_sled_n:
+  //   addi.d  sp, sp, -16                       ; create the stack frame
+  //   st.d    ra, sp, 8                         ; save the return address
+  //   lu12i.w t0, %abs_hi20(__xray_FunctionEntry/Exit)
+  //   ori     t0, t0, %abs_lo12(__xray_FunctionEntry/Exit)
+  //   lu32i.d t0, %abs64_lo20(__xray_FunctionEntry/Exit)
+  //   lu52i.d t0, t0, %abs64_hi12(__xray_FunctionEntry/Exit)
+  //   lu12i.w t1, %abs_hi20(function_id)
+  //   ori     t1, t1, %abs_lo12(function_id)    ; pass the function id
+  //   jirl    ra, t0, 0                         ; call the tracing hook
+  //   ld.d    ra, sp, 8                         ; restore the return address
+  //   addi.d  sp, sp, 16                        ; de-allocate the stack frame
+  //
+  // Replacement of the first 4-byte instruction should be the last and atomic
+  // operation, so that the user code which reaches the sled concurrently
+  // either jumps over the whole sled, or executes the whole sled when the
+  // latter is ready.
+  //
+  // When |Enable|==false, we set the first instruction in the sled back to
+  //   B #48
+
+  uint32_t *Address = reinterpret_cast<uint32_t *>(Sled.address());
+  if (Enable) {
+    uint32_t LoTracingHookAddr = reinterpret_cast<int64_t>(TracingHook) & 0xfff;
+    uint32_t HiTracingHookAddr =
+        (reinterpret_cast<int64_t>(TracingHook) >> 12) & 0xfffff;
+    uint32_t HigherTracingHookAddr =
+        (reinterpret_cast<int64_t>(TracingHook) >> 32) & 0xfffff;
+    uint32_t HighestTracingHookAddr =
+        (reinterpret_cast<int64_t>(TracingHook) >> 52) & 0xfff;
+    uint32_t LoFunctionID = FuncId & 0xfff;
+    uint32_t HiFunctionID = (FuncId >> 12) & 0xfffff;
+    Address[1] = encodeInstruction2RIx(0x29c00000, RegNum::RN_RA, RegNum::RN_SP,
+                                       0x8); // st.d ra, sp, 8
+    Address[2] = encodeInstruction1RI20(
+        0x14000000, RegNum::RN_T0,
+        HiTracingHookAddr); // lu12i.w t0, HiTracingHookAddr
+    Address[3] = encodeInstruction2RIx(
+        0x03800000, RegNum::RN_T0, RegNum::RN_T0,
+        LoTracingHookAddr); // ori t0, t0, LoTracingHookAddr
+    Address[4] = encodeInstruction1RI20(
+        0x16000000, RegNum::RN_T0,
+        HigherTracingHookAddr); // lu32i.d t0, HigherTracingHookAddr
+    Address[5] = encodeInstruction2RIx(
+        0x03000000, RegNum::RN_T0, RegNum::RN_T0,
+        HighestTracingHookAddr); // lu52i.d t0, t0, HighestTracingHookAddr
+    Address[6] =
+        encodeInstruction1RI20(0x14000000, RegNum::RN_T1,
+                               HiFunctionID); // lu12i.w t1, HiFunctionID
+    Address[7] =
+        encodeInstruction2RIx(0x03800000, RegNum::RN_T1, RegNum::RN_T1,
+                              LoFunctionID); // ori t1, t1, LoFunctionID
+    Address[8] = encodeInstruction2RIx(0x4c000000, RegNum::RN_RA, RegNum::RN_T0,
+                                       0); // jirl ra, t0, 0
+    Address[9] = encodeInstruction2RIx(0x28c00000, RegNum::RN_RA, RegNum::RN_SP,
+                                       0x8); // ld.d ra, sp, 8
+    Address[10] = encodeInstruction2RIx(
+        0x02c00000, RegNum::RN_SP, RegNum::RN_SP, 0x10); // addi.d sp, sp, 16
+    uint32_t CreateStackSpace = encodeInstruction2RIx(
+        0x02c00000, RegNum::RN_SP, RegNum::RN_SP, 0xff0); // addi.d sp, sp, -16
+    std::atomic_store_explicit(
+        reinterpret_cast<std::atomic<uint32_t> *>(Address), CreateStackSpace,
+        std::memory_order_release);
+  } else {
+    std::atomic_store_explicit(
+        reinterpret_cast<std::atomic<uint32_t> *>(Address),
+        uint32_t(0x50003000), std::memory_order_release); // b #48
+  }
+  return true;
+}
+
+bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
+                        const XRaySledEntry &Sled,
+                        void (*Trampoline)()) XRAY_NEVER_INSTRUMENT {
+  return patchSled(Enable, FuncId, Sled, Trampoline);
+}
+
+bool patchFunctionExit(const bool Enable, const uint32_t FuncId,
+                       const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  return patchSled(Enable, FuncId, Sled, __xray_FunctionExit);
+}
+
+bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
+                           const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  // TODO: In the future we'd need to distinguish between non-tail exits and
+  // tail exits for better information preservation.
+  return patchSled(Enable, FuncId, Sled, __xray_FunctionExit);
+}
+
+bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
+                      const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  // FIXME: Implement in loongarch?
+  return false;
+}
+
+bool patchTypedEvent(const bool Enable, const uint32_t FuncId,
+                     const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  // FIXME: Implement in loongarch?
+  return false;
+}
+} // namespace __xray
+
+extern "C" void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
+  // TODO: This will have to be implemented in the trampoline assembly file.
+}
diff --git a/compiler-rt/lib/xray/xray_trampoline_loongarch64.S b/compiler-rt/lib/xray/xray_trampoline_loongarch64.S
new file mode 100644
index 0000000..fcbefcc
--- /dev/null
+++ b/compiler-rt/lib/xray/xray_trampoline_loongarch64.S
@@ -0,0 +1,124 @@
+//===-- xray_trampoline_loongarch64.s ---------------------------*- ASM -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of XRay, a dynamic runtime instrumentation system.
+//
+// This implements the loongarch-specific assembler for the trampolines.
+//
+//===----------------------------------------------------------------------===//
+
+#include "../sanitizer_common/sanitizer_asm.h"
+
+#define FROM_0_TO_7 0,1,2,3,4,5,6,7
+#define FROM_7_TO_0 7,6,5,4,3,2,1,0
+
+.macro SAVE_ARG_REGISTERS
+  .irp i,FROM_7_TO_0
+    st.d  $a\i, $sp, (8 * 8 + 8 * \i)
+  .endr
+  .irp i,FROM_7_TO_0
+    fst.d $f\i, $sp, (8 * \i)
+  .endr
+.endm
+
+.macro RESTORE_ARG_REGISTERS
+  .irp i,FROM_0_TO_7
+    fld.d $f\i, $sp, (8 * \i)
+  .endr
+  .irp i,FROM_0_TO_7
+    ld.d  $a\i, $sp, (8 * 8 + 8 * \i)
+  .endr
+.endm
+
+.macro SAVE_RET_REGISTERS
+  st.d    $a1, $sp, 24
+  st.d    $a0, $sp, 16
+  fst.d   $f1, $sp, 8
+  fst.d   $f0, $sp, 0
+.endm
+
+.macro RESTORE_RET_REGISTERS
+  fld.d   $f0, $sp, 0
+  fld.d   $f1, $sp, 8
+  ld.d    $a0, $sp, 16
+  ld.d    $a1, $sp, 24
+.endm
+
+  .text
+  .file "xray_trampoline_loongarch64.S"
+  .globl ASM_SYMBOL(__xray_FunctionEntry)
+  ASM_HIDDEN(__xray_FunctionEntry)
+  .p2align 2
+  ASM_TYPE_FUNCTION(__xray_FunctionEntry)
+ASM_SYMBOL(__xray_FunctionEntry):
+  .cfi_startproc
+  // Save argument registers before doing any actual work.
+  .cfi_def_cfa_offset 136
+  addi.d  $sp, $sp, -136
+  st.d    $ra, $sp, 128
+  .cfi_offset 1, -8
+  SAVE_ARG_REGISTERS
+
+  la.got  $t2, ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)
+  ld.d    $t2, $t2, 0
+
+  beqz    $t2, FunctionEntry_restore
+
+  // a1=0 means that we are tracing an entry event.
+  move    $a1, $zero
+  // Function ID is in t1 (the first parameter).
+  move    $a0, $t1
+  jirl    $ra, $t2, 0
+
+FunctionEntry_restore:
+  // Restore argument registers.
+  RESTORE_ARG_REGISTERS
+  ld.d    $ra, $sp, 128
+  addi.d  $sp, $sp, 136
+  ret
+FunctionEntry_end:
+  ASM_SIZE(__xray_FunctionEntry)
+  .cfi_endproc
+
+  .text
+  .globl ASM_SYMBOL(__xray_FunctionExit)
+  ASM_HIDDEN(__xray_FunctionExit)
+  .p2align 2
+  ASM_TYPE_FUNCTION(__xray_FunctionExit)
+ASM_SYMBOL(__xray_FunctionExit):
+  .cfi_startproc
+  // Save return registers before doing any actual work.
+  .cfi_def_cfa_offset 48
+  addi.d  $sp, $sp, -48
+  st.d    $ra, $sp, 40
+  .cfi_offset 1, -8
+  st.d    $fp, $sp, 32
+  SAVE_RET_REGISTERS
+
+  la.got  $t2, ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)
+  ld.d    $t2, $t2, 0
+
+  beqz    $t2, FunctionExit_restore
+
+  // a1=1 means that we are tracing an exit event.
+  li.w    $a1, 1
+  // Function ID is in t1 (the first parameter).
+  move    $a0, $t1
+  jirl    $ra, $t2, 0
+
+FunctionExit_restore:
+  // Restore return registers.
+  RESTORE_RET_REGISTERS
+  ld.d    $fp, $sp, 32
+  ld.d    $ra, $sp, 40
+  addi.d  $sp, $sp, 48
+  ret
+
+FunctionExit_end:
+  ASM_SIZE(__xray_FunctionExit)
+  .cfi_endproc
diff --git a/compiler-rt/lib/xray/xray_tsc.h b/compiler-rt/lib/xray/xray_tsc.h
index 58347dc..e1cafe1 100644
--- a/compiler-rt/lib/xray/xray_tsc.h
+++ b/compiler-rt/lib/xray/xray_tsc.h
@@ -43,7 +43,7 @@ inline uint64_t getTSCFrequency() XRAY_NEVER_INSTRUMENT {
 #elif defined(__powerpc64__)
 #include "xray_powerpc64.inc"
 #elif defined(__arm__) || defined(__aarch64__) || defined(__mips__) ||         \
-    defined(__hexagon__)
+    defined(__hexagon__) || defined(__loongarch_lp64)
 // Emulated TSC.
 // There is no instruction like RDTSCP in user mode on ARM. ARM's CP15 does
 //   not have a constant frequency like TSC on x86(_64), it may go faster
diff --git a/compiler-rt/test/xray/TestCases/Posix/c-test.cpp b/compiler-rt/test/xray/TestCases/Posix/c-test.cpp
index 9a8a16d..6427566 100644
--- a/compiler-rt/test/xray/TestCases/Posix/c-test.cpp
+++ b/compiler-rt/test/xray/TestCases/Posix/c-test.cpp
@@ -4,7 +4,7 @@
 // RUN:     2>&1 | FileCheck %s
 // RUN: rm -f xray-log.c-test.*
 //
-// REQUIRES: target={{(aarch64|x86_64)-.*}}
+// REQUIRES: target={{(aarch64|loongarch64|x86_64)-.*}}
 // REQUIRES: built-in-llvm-tree
 __attribute__((xray_always_instrument)) void always() {}
 
diff --git a/compiler-rt/test/xray/TestCases/Posix/fdr-thread-order.cpp b/compiler-rt/test/xray/TestCases/Posix/fdr-thread-order.cpp
index 1dc79db..85284fc 100644
--- a/compiler-rt/test/xray/TestCases/Posix/fdr-thread-order.cpp
+++ b/compiler-rt/test/xray/TestCases/Posix/fdr-thread-order.cpp
@@ -8,7 +8,7 @@
 // RUN: %llvm_xray convert --symbolize --output-format=yaml -instr_map=%t.exe %t/* | \
 // RUN:   FileCheck %s --check-prefix TRACE
 
-// REQUIRES: target={{(aarch64|x86_64)-.*}}
+// REQUIRES: target={{(aarch64|loongarch64|x86_64)-.*}}
 // REQUIRES: built-in-llvm-tree
 
 #include "xray/xray_log_interface.h"