[ctxprof][nfc] Share the definition of FunctionData between compiler-rt and llvm (#132136)

Mechanism to keep the compiler-rt and llvm view of `FunctionData` in sync. Since CtxInstrContextNode.h is exactly the same on both sides (there's an existing test, `compiler-rt/test/ctx_profile/TestCases/check-same-ctx-node.test`, checking that), we capture the structure in a macro that is then generated as `struct` fields on the compiler-rt side, and as `Type` objects on the llvm side. The macro needs to be told how to render a few kinds of fields. 

If we add more fields to FunctionData that can be described by the current known types of fields, then the llvm side would automatically be updated. If we need to add more kinds of fields, which we do by adding parameters to the macro, the llvm side (if not updated) would trigger a compilation error.

4 files changed, 56 insertions, 14 deletions

diff --git a/compiler-rt/lib/ctx_profile/CtxInstrContextNode.h b/compiler-rt/lib/ctx_profile/CtxInstrContextNode.h
index a176662..a42bf9e 100644
--- a/compiler-rt/lib/ctx_profile/CtxInstrContextNode.h
+++ b/compiler-rt/lib/ctx_profile/CtxInstrContextNode.h
@@ -113,6 +113,24 @@ public:
   uint64_t entrycount() const { return counters()[0]; }
 };
 
+/// The internal structure of FunctionData. This makes sure that changes to
+/// the fields of FunctionData either get automatically captured on the llvm
+/// side, or force a manual corresponding update.
+///
+/// The macro arguments (see CtxInstrProfiling.h for example):
+///
+/// PTRDECL is a macro taking 2 parameters: a type and the name of the field.
+/// The field is a pointer of that type;
+///
+/// VOLATILE_PTRDECL is the same as above, but for volatile pointers;
+///
+/// MUTEXDECL takes one parameter, the name of a field that is a mutex.
+#define CTXPROF_FUNCTION_DATA(PTRDECL, VOLATILE_PTRDECL, MUTEXDECL)            \
+  PTRDECL(FunctionData, Next)                                                  \
+  VOLATILE_PTRDECL(ContextRoot, CtxRoot)                                       \
+  VOLATILE_PTRDECL(ContextNode, FlatCtx)                                       \
+  MUTEXDECL(Mutex)
+
 /// Abstraction for the parameter passed to `__llvm_ctx_profile_fetch`.
 /// `startContextSection` is called before any context roots are sent for
 /// writing. Then one or more `writeContextual` calls are made; finally,
diff --git a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.h b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.h
index 1e15180..6326bea 100644
--- a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.h
+++ b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.h
@@ -142,17 +142,19 @@ struct ContextRoot {
 // The current design trades off a bit of overhead at the first time a function
 // is encountered *for flat profiling* for avoiding size penalties.
 struct FunctionData {
+#define _PTRDECL(T, N) T *N = nullptr;
+#define _VOLATILE_PTRDECL(T, N) T *volatile N = nullptr;
+#define _MUTEXDECL(N) ::__sanitizer::SpinMutex N;
+  CTXPROF_FUNCTION_DATA(_PTRDECL, _VOLATILE_PTRDECL, _MUTEXDECL)
+#undef _PTRDECL
+#undef _VOLATILE_PTRDECL
+#undef _MUTEXDECL
+
   // Constructor for test only - since this is expected to be
   // initialized by the compiler.
-  FunctionData() { Mutex.Init(); }
-
-  FunctionData *Next = nullptr;
-  ContextRoot *volatile CtxRoot = nullptr;
-  ContextNode *volatile FlatCtx = nullptr;
-
+  FunctionData() = default;
   ContextRoot *getOrAllocateContextRoot();
 
-  ::__sanitizer::StaticSpinMutex Mutex;
   // If (unlikely) StaticSpinMutex internals change, we need to modify the LLVM
   // instrumentation lowering side because it is responsible for allocating and
   // zero-initializing ContextRoots.
diff --git a/llvm/include/llvm/ProfileData/CtxInstrContextNode.h b/llvm/include/llvm/ProfileData/CtxInstrContextNode.h
index a176662..a42bf9e 100644
--- a/llvm/include/llvm/ProfileData/CtxInstrContextNode.h
+++ b/llvm/include/llvm/ProfileData/CtxInstrContextNode.h
@@ -113,6 +113,24 @@ public:
   uint64_t entrycount() const { return counters()[0]; }
 };
 
+/// The internal structure of FunctionData. This makes sure that changes to
+/// the fields of FunctionData either get automatically captured on the llvm
+/// side, or force a manual corresponding update.
+///
+/// The macro arguments (see CtxInstrProfiling.h for example):
+///
+/// PTRDECL is a macro taking 2 parameters: a type and the name of the field.
+/// The field is a pointer of that type;
+///
+/// VOLATILE_PTRDECL is the same as above, but for volatile pointers;
+///
+/// MUTEXDECL takes one parameter, the name of a field that is a mutex.
+#define CTXPROF_FUNCTION_DATA(PTRDECL, VOLATILE_PTRDECL, MUTEXDECL)            \
+  PTRDECL(FunctionData, Next)                                                  \
+  VOLATILE_PTRDECL(ContextRoot, CtxRoot)                                       \
+  VOLATILE_PTRDECL(ContextNode, FlatCtx)                                       \
+  MUTEXDECL(Mutex)
+
 /// Abstraction for the parameter passed to `__llvm_ctx_profile_fetch`.
 /// `startContextSection` is called before any context roots are sent for
 /// writing. Then one or more `writeContextual` calls are made; finally,
diff --git a/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp b/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp
index 1b66e7b..58748a1 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp
@@ -18,6 +18,7 @@
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/ProfileData/CtxInstrContextNode.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/Support/CommandLine.h"
 #include <utility>
@@ -113,13 +114,16 @@ CtxInstrumentationLowerer::CtxInstrumentationLowerer(Module &M,
   auto *I32Ty = Type::getInt32Ty(M.getContext());
   auto *I64Ty = Type::getInt64Ty(M.getContext());
 
-  FunctionDataTy =
-      StructType::get(M.getContext(), {
-                                          PointerTy,          /*Next*/
-                                          PointerTy,          /*CtxRoot*/
-                                          PointerTy,          /*FlatCtx*/
-                                          SanitizerMutexType, /*Mutex*/
-                                      });
+#define _PTRDECL(_, __) PointerTy,
+#define _VOLATILE_PTRDECL(_, __) PointerTy,
+#define _MUTEXDECL(_) SanitizerMutexType,
+
+  FunctionDataTy = StructType::get(
+      M.getContext(),
+      {CTXPROF_FUNCTION_DATA(_PTRDECL, _VOLATILE_PTRDECL, _MUTEXDECL)});
+#undef _PTRDECL
+#undef _VOLATILE_PTRDECL
+#undef _MUTEXDECL
 
   // The Context header.
   ContextNodeTy = StructType::get(M.getContext(), {