[compiler-rt][ctx_instr] Add ctx_profile component (#89304)

Add the component structure for contextual instrumented PGO and the bump allocator + test.

(Tracking Issue: #89287, RFC referenced there)

Author: mtrofin

compiler-rt/CMakeLists.txt

@@ -50,6 +50,8 @@ option(COMPILER_RT_BUILD_LIBFUZZER "Build libFuzzer" ON)
 mark_as_advanced(COMPILER_RT_BUILD_LIBFUZZER)
 option(COMPILER_RT_BUILD_PROFILE "Build profile runtime" ON)
 mark_as_advanced(COMPILER_RT_BUILD_PROFILE)
+option(COMPILER_RT_BUILD_CTX_PROFILE "Build ctx profile runtime" ON)
+mark_as_advanced(COMPILER_RT_BUILD_CTX_PROFILE)
 option(COMPILER_RT_BUILD_MEMPROF "Build memory profiling runtime" ON)
 mark_as_advanced(COMPILER_RT_BUILD_MEMPROF)
 option(COMPILER_RT_BUILD_XRAY_NO_PREINIT "Build xray with no preinit patching" OFF)

compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake

@@ -66,6 +66,7 @@ set(ALL_MEMPROF_SUPPORTED_ARCH ${X86_64})
 set(ALL_PROFILE_SUPPORTED_ARCH ${X86} ${X86_64} ${ARM32} ${ARM64} ${PPC32} ${PPC64}
     ${MIPS32} ${MIPS64} ${S390X} ${SPARC} ${SPARCV9} ${HEXAGON}
     ${RISCV32} ${RISCV64} ${LOONGARCH64})
+set(ALL_CTX_PROFILE_SUPPORTED_ARCH ${X86_64})
 set(ALL_TSAN_SUPPORTED_ARCH ${X86_64} ${MIPS64} ${ARM64} ${PPC64} ${S390X}
     ${LOONGARCH64} ${RISCV64})
 set(ALL_UBSAN_SUPPORTED_ARCH ${X86} ${X86_64} ${ARM32} ${ARM64} ${RISCV64}

compiler-rt/cmake/config-ix.cmake

@@ -632,6 +632,9 @@ if(APPLE)
   list_intersect(PROFILE_SUPPORTED_ARCH
     ALL_PROFILE_SUPPORTED_ARCH
     SANITIZER_COMMON_SUPPORTED_ARCH)
+  list_intersect(CTX_PROFILE_SUPPORTED_ARCH
+    ALL_CTX_PROFILE_SUPPORTED_ARCH
+    SANITIZER_COMMON_SUPPORTED_ARCH)
   list_intersect(TSAN_SUPPORTED_ARCH
     ALL_TSAN_SUPPORTED_ARCH
     SANITIZER_COMMON_SUPPORTED_ARCH)
@@ -678,6 +681,7 @@ else()
   filter_available_targets(HWASAN_SUPPORTED_ARCH ${ALL_HWASAN_SUPPORTED_ARCH})
   filter_available_targets(MEMPROF_SUPPORTED_ARCH ${ALL_MEMPROF_SUPPORTED_ARCH})
   filter_available_targets(PROFILE_SUPPORTED_ARCH ${ALL_PROFILE_SUPPORTED_ARCH})
+  filter_available_targets(CTX_PROFILE_SUPPORTED_ARCH ${ALL_CTX_PROFILE_SUPPORTED_ARCH})
   filter_available_targets(TSAN_SUPPORTED_ARCH ${ALL_TSAN_SUPPORTED_ARCH})
   filter_available_targets(UBSAN_SUPPORTED_ARCH ${ALL_UBSAN_SUPPORTED_ARCH})
   filter_available_targets(SAFESTACK_SUPPORTED_ARCH
@@ -803,6 +807,13 @@ else()
   set(COMPILER_RT_HAS_PROFILE FALSE)
 endif()
 
+if (COMPILER_RT_HAS_SANITIZER_COMMON AND CTX_PROFILE_SUPPORTED_ARCH AND
+    OS_NAME MATCHES "Linux")
+  set(COMPILER_RT_HAS_CTX_PROFILE TRUE)
+else()
+  set(COMPILER_RT_HAS_CTX_PROFILE FALSE)
+endif()
+
 if (COMPILER_RT_HAS_SANITIZER_COMMON AND TSAN_SUPPORTED_ARCH)
   if (OS_NAME MATCHES "Linux|Darwin|FreeBSD|NetBSD")
     set(COMPILER_RT_HAS_TSAN TRUE)

compiler-rt/lib/CMakeLists.txt

@@ -51,6 +51,10 @@ if(COMPILER_RT_BUILD_PROFILE AND COMPILER_RT_HAS_PROFILE)
   compiler_rt_build_runtime(profile)
 endif()
 
+if(COMPILER_RT_BUILD_CTX_PROFILE AND COMPILER_RT_HAS_CTX_PROFILE)
+  compiler_rt_build_runtime(ctx_profile)
+endif()
+
 if(COMPILER_RT_BUILD_XRAY)
   compiler_rt_build_runtime(xray)
 endif()

compiler-rt/lib/ctx_profile/CMakeLists.txt

@@ -0,0 +1,28 @@
+add_compiler_rt_component(ctx_profile)
+
+set(CTX_PROFILE_SOURCES
+  CtxInstrProfiling.cpp
+  )
+
+set(CTX_PROFILE_HEADERS
+  CtxInstrProfiling.h
+  )
+
+include_directories(..)
+include_directories(../../include)
+
+# We don't use the C++ Standard Library here, so avoid including it by mistake.
+append_list_if(COMPILER_RT_HAS_NOSTDINCXX_FLAG -nostdinc++ EXTRA_FLAGS)
+
+add_compiler_rt_runtime(clang_rt.ctx_profile
+  STATIC
+  ARCHS ${CTX_PROFILE_SUPPORTED_ARCH}
+  OBJECT_LIBS RTSanitizerCommon RTSanitizerCommonLibc
+  CFLAGS ${EXTRA_FLAGS}
+  SOURCES ${CTX_PROFILE_SOURCES}
+  ADDITIONAL_HEADERS ${CTX_PROFILE_HEADERS}
+  PARENT_TARGET ctx_profile)
+
+if(COMPILER_RT_INCLUDE_TESTS)
+  add_subdirectory(tests)
+endif()

compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp

@@ -0,0 +1,40 @@
+//===- CtxInstrProfiling.cpp - contextual instrumented PGO ----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "CtxInstrProfiling.h"
+#include "sanitizer_common/sanitizer_allocator_internal.h"
+#include "sanitizer_common/sanitizer_common.h"
+#include "sanitizer_common/sanitizer_dense_map.h"
+#include "sanitizer_common/sanitizer_mutex.h"
+#include "sanitizer_common/sanitizer_placement_new.h"
+#include "sanitizer_common/sanitizer_thread_safety.h"
+
+#include <assert.h>
+
+using namespace __ctx_profile;
+
+// FIXME(mtrofin): use malloc / mmap instead of sanitizer common APIs to reduce
+// the dependency on the latter.
+Arena *Arena::allocateNewArena(size_t Size, Arena *Prev) {
+  assert(!Prev || Prev->Next == nullptr);
+  Arena *NewArena =
+      new (__sanitizer::InternalAlloc(Size + sizeof(Arena))) Arena(Size);
+  if (Prev)
+    Prev->Next = NewArena;
+  return NewArena;
+}
+
+void Arena::freeArenaList(Arena *&A) {
+  assert(A);
+  for (auto *I = A; I != nullptr;) {
+    auto *Current = I;
+    I = I->Next;
+    __sanitizer::InternalFree(Current);
+  }
+  A = nullptr;
+}

compiler-rt/lib/ctx_profile/CtxInstrProfiling.h

@@ -0,0 +1,55 @@
+/*===- CtxInstrProfiling.h- Contextual instrumentation-based PGO  ---------===*\
+|*
+|* 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
+|*
+\*===----------------------------------------------------------------------===*/
+
+#ifndef CTX_PROFILE_CTXINSTRPROFILING_H_
+#define CTX_PROFILE_CTXINSTRPROFILING_H_
+
+#include <sanitizer/common_interface_defs.h>
+
+namespace __ctx_profile {
+
+/// Arena (bump allocator) forming a linked list. Intentionally not thread safe.
+/// Allocation and de-allocation happen using sanitizer APIs. We make that
+/// explicit.
+class Arena final {
+public:
+  // When allocating a new Arena, optionally specify an existing one to append
+  // to, assumed to be the last in the Arena list. We only need to support
+  // appending to the arena list.
+  static Arena *allocateNewArena(size_t Size, Arena *Prev = nullptr);
+  static void freeArenaList(Arena *&A);
+
+  uint64_t size() const { return Size; }
+
+  // Allocate S bytes or return nullptr if we don't have that many available.
+  char *tryBumpAllocate(size_t S) {
+    if (Pos + S > Size)
+      return nullptr;
+    Pos += S;
+    return start() + (Pos - S);
+  }
+
+  Arena *next() const { return Next; }
+
+  // the beginning of allocatable memory.
+  const char *start() const { return const_cast<Arena *>(this)->start(); }
+  const char *pos() const { return start() + Pos; }
+
+private:
+  explicit Arena(uint32_t Size) : Size(Size) {}
+  ~Arena() = delete;
+
+  char *start() { return reinterpret_cast<char *>(&this[1]); }
+
+  Arena *Next = nullptr;
+  uint64_t Pos = 0;
+  const uint64_t Size;
+};
+
+} // namespace __ctx_profile
+#endif // CTX_PROFILE_CTXINSTRPROFILING_H_

compiler-rt/lib/ctx_profile/tests/CMakeLists.txt

@@ -0,0 +1,70 @@
+include(CheckCXXCompilerFlag)
+include(CompilerRTCompile)
+include(CompilerRTLink)
+
+set(CTX_PROFILE_UNITTEST_CFLAGS
+  ${COMPILER_RT_UNITTEST_CFLAGS}
+  ${COMPILER_RT_GTEST_CFLAGS}
+  ${COMPILER_RT_GMOCK_CFLAGS}
+  ${SANITIZER_TEST_CXX_CFLAGS}
+  -I${COMPILER_RT_SOURCE_DIR}/lib/
+  -DSANITIZER_COMMON_NO_REDEFINE_BUILTINS
+  -O2
+  -g
+  -fno-rtti
+  -Wno-pedantic
+  -fno-omit-frame-pointer)
+
+# Suppress warnings for gmock variadic macros for clang and gcc respectively.
+append_list_if(SUPPORTS_GNU_ZERO_VARIADIC_MACRO_ARGUMENTS_FLAG -Wno-gnu-zero-variadic-macro-arguments CTX_PROFILE_UNITTEST_CFLAGS)
+append_list_if(COMPILER_RT_HAS_WVARIADIC_MACROS_FLAG -Wno-variadic-macros CTX_PROFILE_UNITTEST_CFLAGS)
+
+file(GLOB PROFILE_HEADERS ../*.h)
+
+set(CTX_PROFILE_SOURCES
+  ../CtxInstrProfiling.cpp)
+
+set(CTX_PROFILE_UNITTESTS
+  CtxInstrProfilingTest.cpp
+  driver.cpp)
+
+include_directories(../../../include)
+
+set(CTX_PROFILE_UNIT_TEST_HEADERS
+  ${CTX_PROFILE_HEADERS})
+
+set(CTX_PROFILE_UNITTEST_LINK_FLAGS
+  ${COMPILER_RT_UNITTEST_LINK_FLAGS})
+
+list(APPEND CTX_PROFILE_UNITTEST_LINK_FLAGS -pthread)
+
+set(CTX_PROFILE_UNITTEST_LINK_LIBRARIES
+  ${COMPILER_RT_UNWINDER_LINK_LIBS}
+  ${SANITIZER_TEST_CXX_LIBRARIES})
+list(APPEND CTX_PROFILE_UNITTEST_LINK_LIBRARIES "dl")
+
+if(COMPILER_RT_DEFAULT_TARGET_ARCH IN_LIST CTX_PROFILE_SUPPORTED_ARCH)
+  # Profile unit tests are only run on the host machine.
+  set(arch ${COMPILER_RT_DEFAULT_TARGET_ARCH})
+
+  add_executable(CtxProfileUnitTests 
+    ${CTX_PROFILE_UNITTESTS}
+    ${COMPILER_RT_GTEST_SOURCE}
+    ${COMPILER_RT_GMOCK_SOURCE}
+    ${CTX_PROFILE_SOURCES}
+    $<TARGET_OBJECTS:RTSanitizerCommon.${arch}>
+    $<TARGET_OBJECTS:RTSanitizerCommonCoverage.${arch}>
+    $<TARGET_OBJECTS:RTSanitizerCommonLibc.${arch}>
+    $<TARGET_OBJECTS:RTSanitizerCommonSymbolizer.${arch}>
+    $<TARGET_OBJECTS:RTSanitizerCommonSymbolizerInternal.${arch}>)
+  set_target_compile_flags(CtxProfileUnitTests ${CTX_PROFILE_UNITTEST_CFLAGS})
+  set_target_link_flags(CtxProfileUnitTests ${CTX_PROFILE_UNITTEST_LINK_FLAGS})
+  target_link_libraries(CtxProfileUnitTests ${CTX_PROFILE_UNITTEST_LINK_LIBRARIES})
+
+  if (TARGET cxx-headers OR HAVE_LIBCXX)
+    add_dependencies(CtxProfileUnitTests cxx-headers)
+  endif()
+
+  set_target_properties(CtxProfileUnitTests PROPERTIES
+    RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+endif()

compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp

@@ -0,0 +1,22 @@
+#include "../CtxInstrProfiling.h"
+#include "gtest/gtest.h"
+
+using namespace __ctx_profile;
+
+TEST(ArenaTest, Basic) {
+  Arena *A = Arena::allocateNewArena(1024);
+  EXPECT_EQ(A->size(), 1024U);
+  EXPECT_EQ(A->next(), nullptr);
+
+  auto *M1 = A->tryBumpAllocate(1020);
+  EXPECT_NE(M1, nullptr);
+  auto *M2 = A->tryBumpAllocate(4);
+  EXPECT_NE(M2, nullptr);
+  EXPECT_EQ(M1 + 1020, M2);
+  EXPECT_EQ(A->tryBumpAllocate(1), nullptr);
+  Arena *A2 = Arena::allocateNewArena(2024, A);
+  EXPECT_EQ(A->next(), A2);
+  EXPECT_EQ(A2->next(), nullptr);
+  Arena::freeArenaList(A);
+  EXPECT_EQ(A, nullptr);
+}

compiler-rt/lib/ctx_profile/tests/driver.cpp

@@ -0,0 +1,14 @@
+//===-- driver.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "gtest/gtest.h"
+
+int main(int argc, char **argv) {
+  testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}