feature: redesign Proteus's lambda JIT compilation interface and pipeline

include/proteus/CompilerInterfaceTypes.h

@@ -78,6 +78,10 @@ struct RuntimeConstant {
   ArrayInfo ArrInfo{0, RuntimeConstantType::NONE, nullptr};
   ObjectInfo ObjInfo{0, false, nullptr};
 
+  // We need a trivial constructor to use RuntimeConstant in DenseMap keys
+  explicit RuntimeConstant()
+      : Type(RuntimeConstantType::NONE), Pos(-1), Offset(-1) {}
+
   explicit RuntimeConstant(RuntimeConstantType Type, int32_t Pos)
       : Type(Type), Pos(Pos) {
     std::memset(&Value, 0, sizeof(RuntimeConstantValue));

include/proteus/JitInterface.h

@@ -17,17 +17,28 @@
 #include "proteus/Init.h"
 
 #include <cassert>
+#include <cstdint>
 #include <cstring>
 #include <type_traits>
 #include <utility>
 
-extern "C" void __proteus_register_variable(proteus::RuntimeConstant RC,
-                                            const char *AssociatedLambda);
-extern "C" void __proteus_register_lambda(const char *Symbol);
+extern "C" __attribute__((used)) void
+__proteus_register_lambda_runtime_constant(int32_t Type, int32_t Pos,
+                                           int32_t Offset, const void *ValuePtr,
+                                           uint64_t functor_id);
+
 extern "C" void __proteus_take_address(void const *) noexcept;
+extern "C" __attribute__((used)) void
+__proteus_finalize_register(uint64_t Tag) noexcept;
 
 namespace proteus {
 
+#if defined(__CUDACC__) || defined(__HIP__)
+#define PROTEUS_HOST_DEVICE __host__ __device__
+#else
+#define PROTEUS_HOST_DEVICE
+#endif
+
 template <typename T> __attribute__((noinline)) void jit_arg(T V) noexcept;
 #if defined(__CUDACC__) || defined(__HIP__)
 template <typename T>
@@ -75,51 +86,77 @@ __attribute__((noinline)) __device__ std::enable_if_t<
 jit_object(T &V, size_t Size = sizeof(T)) noexcept;
 #endif
 
-template <typename T> inline static RuntimeConstantType convertCTypeToRCType() {
-  if constexpr (std::is_same_v<T, bool>) {
-    return RuntimeConstantType::BOOL;
-  } else if constexpr (std::is_integral_v<T> && sizeof(T) == sizeof(int8_t)) {
-    return RuntimeConstantType::INT8;
-  } else if constexpr (std::is_integral_v<T> && sizeof(T) == sizeof(int32_t)) {
-    return RuntimeConstantType::INT32;
-  } else if constexpr (std::is_integral_v<T> && sizeof(T) == sizeof(int64_t)) {
-    return RuntimeConstantType::INT64;
-  } else if constexpr (std::is_same_v<T, float>) {
-    return RuntimeConstantType::FLOAT;
-  } else if constexpr (std::is_same_v<T, double>) {
-    return RuntimeConstantType::DOUBLE;
-  } else if constexpr (std::is_same_v<T, long double>) {
-    return RuntimeConstantType::LONG_DOUBLE;
-  } else if constexpr (std::is_pointer_v<T>) {
-    return RuntimeConstantType::PTR;
-  } else {
-    return RuntimeConstantType::NONE;
+namespace detail {
+// todo: use LLVM hashing?
+constexpr std::uint64_t fnv1a64(const char *s) {
+  std::uint64_t h = 14695981039346656037ull;
+  for (; *s; ++s) {
+    h ^= (unsigned char)(*s);
+    h *= 1099511628211ull;
   }
+  return h;
 }
 
-template <typename T>
-static __attribute__((noinline)) T
-jit_variable(T V, int Pos = -1, int Offset = -1,
-             const char *AssociatedLambda = "") noexcept {
-  RuntimeConstant RC{convertCTypeToRCType<T>(), Pos, Offset};
-  std::memcpy(static_cast<void *>(&RC), &V, sizeof(T));
-  __proteus_register_variable(RC, AssociatedLambda);
+// todo: make a test with lambda factory in a header, two separate cpp files
+template <class Lambda> constexpr std::uint64_t functor_id() {
+  return fnv1a64(__PRETTY_FUNCTION__); // includes Lambda + Ctr in the text
+}
 
-  return V;
+template <uint64_t FunctorID, typename Lambda> struct LambdaFunctorWrapper {
+  using LambdaType = Lambda;
+  static constexpr std::uint64_t functor_id = FunctorID;
+  LambdaType lambda;
+
+  template <typename... Args>
+  PROTEUS_HOST_DEVICE __attribute__((annotate("jit")))
+  __attribute__((annotate("proteus.wrapper_call", functor_id))) decltype(auto)
+  operator()(Args &&...args) const noexcept {
+    return lambda(std::forward<Args>(args)...);
+  }
+};
+
+template <typename... T> struct is_lambda_functor_wrapper : std::false_type {};
+
+template <uint64_t FunctorID, typename Lambda>
+struct is_lambda_functor_wrapper<LambdaFunctorWrapper<FunctorID, Lambda>>
+    : std::true_type {};
+
+template <std::uint64_t FunctorId, typename L>
+PROTEUS_HOST_DEVICE inline auto tag_functor(L &&lambda) {
+  return LambdaFunctorWrapper<FunctorId, std::decay_t<L>>{
+      std::forward<L>(lambda)};
 }
 
-template <typename T>
-static __attribute__((noinline)) T &&
-register_lambda(T &&t, const char *Symbol = "") noexcept {
-  assert(Symbol && "Expected non-null Symbol");
-  __proteus_register_lambda(Symbol);
+template <uint64_t ID, typename T>
+[[nodiscard]] static __attribute__((noinline))
+__attribute__((annotate("proteus.register_call_impl", ID))) auto
+__register_lambda_impl(T &&t) noexcept {
+  static_assert(!detail::is_lambda_functor_wrapper<std::decay_t<T>>::value);
   // Force LLVM to generate an AllocaInst of the underlying Clang--generated
   // anonymous class for T.  We remove this after recording the demangled
   // lambda name.
   using LambdaType = std::decay_t<T>;
   LambdaType local = t;
   __proteus_take_address(&local);
-  return std::forward<T>(t);
+  auto result = tag_functor<ID>(std::forward<T>(t));
+  return result;
+}
+
+} // namespace detail
+
+template <class L>
+[[nodiscard]] inline auto __attribute__((annotate("proteus.register_call")))
+register_lambda(L &&lambda) noexcept {
+  auto registered_lambda = ::proteus::detail::__register_lambda_impl<
+      ::proteus::detail::functor_id<std::decay_t<L>>()>(
+      std::forward<L>(lambda));
+  __proteus_take_address(&registered_lambda);
+  return registered_lambda;
+}
+
+template <typename T>
+static __attribute__((noinline)) T jit_variable(T V) noexcept {
+  return V;
 }
 
 #if defined(__CUDACC__) || defined(__HIP__)

src/include/proteus/impl/CompilationTask.h

@@ -7,6 +7,7 @@
 #include "proteus/impl/CoreLLVMDevice.h"
 #include "proteus/impl/Debug.h"
 #include "proteus/impl/Hashing.h"
+#include "proteus/impl/LambdaCallsite.h"
 #include "proteus/impl/Utils.h"
 
 #include <llvm/Bitcode/BitcodeReader.h>
@@ -25,7 +26,8 @@ class CompilationTask {
   dim3 BlockDim;
   dim3 GridDim;
   SmallVector<RuntimeConstant> RCVec;
-  SmallVector<std::pair<std::string, StringRef>> LambdaCalleeInfo;
+  SmallVector<uint64_t> LambdaCalleeInfo;
+  LambdaCallsiteRuntimeConstantsMap LambdaCallsiteRuntimeConstants;
   std::unordered_map<std::string, GlobalVarInfo> VarNameToGlobalInfo;
   SmallPtrSet<void *, 8> GlobalLinkedBinaries;
   std::string DeviceArch;
@@ -71,14 +73,16 @@ class CompilationTask {
       MemoryBufferRef Bitcode, HashT HashValue, const std::string &KernelName,
       std::string &Suffix, dim3 BlockDim, dim3 GridDim,
       const SmallVector<RuntimeConstant> &RCVec,
-      const SmallVector<std::pair<std::string, StringRef>> &LambdaCalleeInfo,
+      const SmallVector<uint64_t> &LambdaCalleeInfo,
+      const LambdaCallsiteRuntimeConstantsMap &LambdaCallsiteRuntimeConstants,
       const std::unordered_map<std::string, GlobalVarInfo> &VarNameToGlobalInfo,
       const SmallPtrSet<void *, 8> &GlobalLinkedBinaries,
       const std::string &DeviceArch, const CodeGenerationConfig &CGConfig,
       bool DumpIR, bool RelinkGlobalsByCopy)
       : Bitcode(Bitcode), HashValue(HashValue), KernelName(KernelName),
         Suffix(Suffix), BlockDim(BlockDim), GridDim(GridDim), RCVec(RCVec),
         LambdaCalleeInfo(LambdaCalleeInfo),
+        LambdaCallsiteRuntimeConstants(LambdaCallsiteRuntimeConstants),
         VarNameToGlobalInfo(VarNameToGlobalInfo),
         GlobalLinkedBinaries(GlobalLinkedBinaries), DeviceArch(DeviceArch),
         CGOption(CGConfig.codeGenOption()), DumpIR(DumpIR),
@@ -143,9 +147,9 @@ class CompilationTask {
     PROTEUS_DBG(Logger::logfile(HashValue.toString() + ".input.ll", *M));
 
     proteus::specializeIR(*M, KernelName, Suffix, BlockDim, GridDim, RCVec,
-                          LambdaCalleeInfo, SpecializeArgs, SpecializeDims,
-                          SpecializeDimsRange, SpecializeLaunchBounds,
-                          MinBlocksPerSM);
+                          LambdaCalleeInfo, LambdaCallsiteRuntimeConstants,
+                          SpecializeArgs, SpecializeDims, SpecializeDimsRange,
+                          SpecializeLaunchBounds, MinBlocksPerSM);
 
     PROTEUS_DBG(Logger::logfile(HashValue.toString() + ".specialized.ll", *M));
 

src/include/proteus/impl/CoreLLVM.h

@@ -11,8 +11,11 @@ static_assert(__cplusplus >= 201703L,
 #include "proteus/impl/JITPassPluginRegistry.h"
 #include "proteus/impl/Logger.h"
 
+#include <llvm/Analysis/ValueTracking.h>
 #include <llvm/CodeGen/CommandFlags.h>
+#include <llvm/IR/Constants.h>
 #include <llvm/IR/DebugInfo.h>
+#include <llvm/IR/Metadata.h>
 #include <llvm/IR/Module.h>
 #include <llvm/IR/PassManager.h>
 #include <llvm/Linker/Linker.h>
@@ -387,6 +390,48 @@ inline void runCleanupPassPipeline(Module &M) {
   StripDebugInfo(M);
 }
 
+inline void findFunctionsWithU64Metadata(
+    llvm::Module &M, llvm::StringRef Key,
+    llvm::SmallVectorImpl<std::pair<llvm::Function *, std::uint64_t>> &Out) {
+  llvm::SmallDenseMap<llvm::Function *, std::uint64_t, 32> Seen;
+  for (llvm::Function &F : M) {
+    llvm::MDNode *Node = F.getMetadata(Key);
+    if (!Node || Node->getNumOperands() < 1)
+      continue;
+    auto *CAM = llvm::dyn_cast<llvm::ConstantAsMetadata>(Node->getOperand(0));
+    auto *CI =
+        CAM ? llvm::dyn_cast<llvm::ConstantInt>(CAM->getValue()) : nullptr;
+    if (!CI)
+      continue;
+    std::uint64_t Id = CI->getZExtValue();
+    if (Seen.try_emplace(&F, Id).second)
+      Out.emplace_back(&F, Id);
+  }
+}
+
+inline void
+findFunctionsWithU64Metadata(llvm::Module &M, llvm::StringRef Key,
+                             llvm::SmallVectorImpl<std::uint64_t> &Out) {
+  llvm::SmallDenseMap<llvm::Function *, std::uint64_t, 32> Seen;
+  for (llvm::Function &F : M) {
+    llvm::MDNode *Node = F.getMetadata(Key);
+    if (!Node || Node->getNumOperands() < 1)
+      continue;
+    auto *CAM = llvm::dyn_cast<llvm::ConstantAsMetadata>(Node->getOperand(0));
+    auto *CI =
+        CAM ? llvm::dyn_cast<llvm::ConstantInt>(CAM->getValue()) : nullptr;
+    if (!CI)
+      continue;
+    std::uint64_t Id = CI->getZExtValue();
+    if (Seen.try_emplace(&F, Id).second)
+      Out.emplace_back(Id);
+  }
+}
+
+inline bool hasU64Metadata(Function *F, StringRef Key) {
+  return F->getMetadata(Key) != nullptr;
+}
+
 inline void pruneIR(Module &M, bool UnsetExternallyInitialized = true) {
   // Remove llvm.global.annotations now that we have read them.
   if (auto *GlobalAnnotations = M.getGlobalVariable("llvm.global.annotations"))
@@ -435,6 +480,20 @@ inline void internalize(Module &M, StringRef PreserveFunctionName) {
   });
 }
 
+inline std::optional<uint64_t> getFunctionU64Metadata(Function &F,
+                                                      StringRef Key) {
+  MDNode *Node = F.getMetadata(Key);
+  if (!Node || Node->getNumOperands() < 1)
+    return std::nullopt;
+
+  auto *CAM = dyn_cast<ConstantAsMetadata>(Node->getOperand(0));
+  auto *CI = CAM ? dyn_cast<ConstantInt>(CAM->getValue()) : nullptr;
+  if (!CI)
+    return std::nullopt;
+
+  return CI->getZExtValue();
+}
+
 } // namespace proteus
 
 #endif

src/include/proteus/impl/CoreLLVMDevice.h

@@ -286,7 +286,8 @@ inline void relinkGlobalsObject(
 inline void specializeIR(
     Module &M, StringRef FnName, StringRef Suffix, dim3 &BlockDim,
     dim3 &GridDim, ArrayRef<RuntimeConstant> RCArray,
-    const SmallVector<std::pair<std::string, StringRef>> LambdaCalleeInfo,
+    const SmallVector<uint64_t> LambdaCalleeInfo,
+    const LambdaCallsiteRuntimeConstantsMap &LambdaCallsiteRuntimeConstants,
     bool SpecializeArgs, bool SpecializeDims, bool SpecializeDimsRange,
     bool SpecializeLaunchBounds, int MinBlocksPerSM) {
   TIMESCOPE("proteus::specializeIR");
@@ -298,13 +299,11 @@ inline void specializeIR(
   if (SpecializeArgs)
     TransformArgumentSpecialization::transform(M, *F, RCArray);
 
-  auto &LR = LambdaRegistry::instance();
-  for (auto &[FnName, LambdaType] : LambdaCalleeInfo) {
-    const SmallVector<RuntimeConstant> &RCVec = LR.getJitVariables(LambdaType);
-    Function *F = M.getFunction(FnName);
-    if (!F)
-      reportFatalError("Expected non-null Function");
-    TransformLambdaSpecialization::transform(M, *F, RCVec);
+  // We add a per-function cache for memory ssa so that we don't have to
+  // duplicate results. TransformLambdaSpecialization LambdaTransformer;
+  for (auto &ID : LambdaCalleeInfo) {
+    TransformLambdaSpecialization::transformDeviceKernel(
+        M, ID, LambdaCallsiteRuntimeConstants);
   }
 
   // Run the shared array transform after any value specialization (arguments,