Automatically mark read-only lambda captures as constant

.gitignore

@@ -7,3 +7,5 @@ miniconda3*
 install-*
 docs/doxygen*
 site
+**.proteus*
+**.proteus-logs*

docs/user/config.md

@@ -18,6 +18,7 @@ through environment variables.
 | `PROTEUS_KERNEL_CLONE` | `"link-clone-prune"`, `"link-clone-light"`, `"cross-clone"` (default: `"cross-clone"`) | Cloning method for JIT module creation |
 | `PROTEUS_ASYNC_COMPILATION` | `0` or `1` (default: `0`) | Enable asynchronous compilation |
 | `PROTEUS_ASYNC_THREADS` | Integer `>= 1` (default: `1`) | Number of threads used for asynchronous compilation |
+| `PROTEUS_AUTO_READONLY_CAPTURES` | `0` or `1` (default: `1`) | Enable automatic detection of read-only lambda captures for JIT specialization. When enabled, scalar captures (`int`, `float`, `double`, `bool`) that are read-only within the lambda body are automatically specialized without requiring explicit `jit_variable()` annotation. |
 | `PROTEUS_ASYNC_TEST_BLOCKING` | `0` or `1` (default: `0`) | Make asynchronous compilation blocking for testing |
 | `PROTEUS_ENABLE_TIMERS` | `0` or `1` (default: `0`) | Enable timer-based profiling output in JIT operations |
 | `PROTEUS_TRACE_OUTPUT` | Semicolon-separated tokens: `specialization`, `ir-dump`, `kernel-trace`, `cache-stats` (default: empty) | Enable trace output. `specialization` prints specialization info, `ir-dump` dumps LLVM IR post-optimization, `kernel-trace` prints an end-of-run per-kernel summary with specialization and launch counts, and `cache-stats` prints object-cache hit/access statistics. Example: `"specialization;kernel-trace"` |

include/proteus/CompilerInterfaceTypes.h

@@ -73,7 +73,7 @@ struct RuntimeConstant {
   RuntimeConstantValue Value;
   RuntimeConstantType Type;
   int32_t Pos;
-  int32_t Offset;
+  int32_t Offset = 0;
 
   ArrayInfo ArrInfo{0, RuntimeConstantType::NONE, nullptr};
   ObjectInfo ObjInfo{0, false, nullptr};

src/include/proteus/impl/AutoReadOnlyCaptures.h

@@ -0,0 +1,239 @@
+//===-- AutoReadOnlyCaptures.h -- Auto read-only capture metadata utils --===//
+//
+// Part of the Proteus 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 PROTEUS_IMPL_AUTOREADONLYCAPTURES_H
+#define PROTEUS_IMPL_AUTOREADONLYCAPTURES_H
+
+#include "proteus/CompilerInterfaceTypes.h"
+#include "proteus/impl/Logger.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Support/Format.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cstdint>
+#include <optional>
+
+namespace proteus {
+
+using namespace llvm;
+
+struct AutoReadOnlyCaptureMetadataEntry {
+  int32_t SlotIndex;
+  int32_t ByteOffset;
+  RuntimeConstantType RCType;
+};
+
+inline bool isSupportedAutoReadOnlyRCType(RuntimeConstantType RCType) {
+  switch (RCType) {
+  case RuntimeConstantType::BOOL:
+  case RuntimeConstantType::INT8:
+  case RuntimeConstantType::INT32:
+  case RuntimeConstantType::INT64:
+  case RuntimeConstantType::FLOAT:
+  case RuntimeConstantType::DOUBLE:
+    return true;
+  default:
+    return false;
+  }
+}
+
+inline bool containsCaptureForSlot(ArrayRef<RuntimeConstant> Captures,
+                                   int32_t SlotIndex) {
+  for (const auto &RC : Captures) {
+    if (RC.Pos == SlotIndex)
+      return true;
+  }
+
+  return false;
+}
+
+/// Merge auto-detected captures with explicit captures (explicit takes
+/// precedence).
+inline void mergeCaptures(llvm::SmallVectorImpl<RuntimeConstant> &Explicit,
+                          llvm::ArrayRef<RuntimeConstant> Auto) {
+  llvm::SmallSet<int32_t, 16> ExplicitSlots;
+  for (const auto &RC : Explicit)
+    ExplicitSlots.insert(RC.Pos);
+
+  for (const auto &RC : Auto) {
+    if (!ExplicitSlots.contains(RC.Pos))
+      Explicit.push_back(RC);
+  }
+}
+
+inline RuntimeConstant
+readValueFromMemory(const void *Ptr,
+                    const AutoReadOnlyCaptureMetadataEntry &Capture) {
+  RuntimeConstant RC(RuntimeConstantType::NONE, Capture.SlotIndex,
+                     Capture.ByteOffset);
+  RuntimeConstantType RCType = Capture.RCType;
+
+  if (RCType == RuntimeConstantType::BOOL) {
+    RC.Type = RuntimeConstantType::BOOL;
+    RC.Value.BoolVal = *static_cast<const bool *>(Ptr);
+  } else if (RCType == RuntimeConstantType::INT8) {
+    RC.Type = RuntimeConstantType::INT8;
+    RC.Value.Int8Val = *static_cast<const int8_t *>(Ptr);
+  } else if (RCType == RuntimeConstantType::INT32) {
+    RC.Type = RuntimeConstantType::INT32;
+    RC.Value.Int32Val = *static_cast<const int32_t *>(Ptr);
+  } else if (RCType == RuntimeConstantType::INT64) {
+    RC.Type = RuntimeConstantType::INT64;
+    RC.Value.Int64Val = *static_cast<const int64_t *>(Ptr);
+  } else if (RCType == RuntimeConstantType::FLOAT) {
+    RC.Type = RuntimeConstantType::FLOAT;
+    RC.Value.FloatVal = *static_cast<const float *>(Ptr);
+  } else if (RCType == RuntimeConstantType::DOUBLE) {
+    RC.Type = RuntimeConstantType::DOUBLE;
+    RC.Value.DoubleVal = *static_cast<const double *>(Ptr);
+  }
+
+  return RC;
+}
+
+inline llvm::SmallVector<AutoReadOnlyCaptureMetadataEntry>
+parseAutoReadOnlyCapturesMetadata(Function &F) {
+  llvm::SmallVector<AutoReadOnlyCaptureMetadataEntry> Captures;
+  MDNode *Root = F.getMetadata("proteus.auto_readonly_captures");
+  if (!Root)
+    return Captures;
+
+  auto ParseI32 = [](Metadata *M) -> std::optional<int32_t> {
+    auto *CAM = dyn_cast<ConstantAsMetadata>(M);
+    if (!CAM)
+      return std::nullopt;
+    auto *CI = dyn_cast<ConstantInt>(CAM->getValue());
+    if (!CI)
+      return std::nullopt;
+    return static_cast<int32_t>(CI->getSExtValue());
+  };
+
+  for (unsigned I = 0; I < Root->getNumOperands(); ++I) {
+    auto *EntryNode = dyn_cast_or_null<MDNode>(Root->getOperand(I));
+    if (!EntryNode || EntryNode->getNumOperands() != 3)
+      continue;
+
+    auto SlotIndex = ParseI32(EntryNode->getOperand(0));
+    auto ByteOffset = ParseI32(EntryNode->getOperand(1));
+    auto RCTypeInt = ParseI32(EntryNode->getOperand(2));
+    if (!SlotIndex || !ByteOffset || !RCTypeInt)
+      continue;
+    if (*SlotIndex < 0 || *ByteOffset < 0)
+      continue;
+
+    RuntimeConstantType RCType = static_cast<RuntimeConstantType>(*RCTypeInt);
+    if (!isSupportedAutoReadOnlyRCType(RCType))
+      continue;
+
+    Captures.push_back(
+        AutoReadOnlyCaptureMetadataEntry{*SlotIndex, *ByteOffset, RCType});
+  }
+
+  return Captures;
+}
+
+inline llvm::SmallVector<RuntimeConstant>
+extractAutoDetectedCapturesFromMetadata(
+    const void *LambdaClosure,
+    const llvm::SmallVector<AutoReadOnlyCaptureMetadataEntry>
+        &DetectedCaptures) {
+  llvm::SmallVector<RuntimeConstant> Result;
+  if (!LambdaClosure || DetectedCaptures.empty())
+    return Result;
+
+  const char *ClosureBytes = static_cast<const char *>(LambdaClosure);
+  for (const auto &Cap : DetectedCaptures) {
+    RuntimeConstant RC =
+        readValueFromMemory(ClosureBytes + Cap.ByteOffset, Cap);
+    if (RC.Type == RuntimeConstantType::NONE)
+      continue;
+    Result.push_back(RC);
+  }
+
+  return Result;
+}
+
+inline SmallString<128> traceOutAuto(int Slot, const RuntimeConstant &RC) {
+  SmallString<128> S;
+  raw_svector_ostream OS(S);
+  OS << "[LambdaSpec][Auto] Replacing slot " << Slot << " with ";
+
+  switch (RC.Type) {
+  case RuntimeConstantType::BOOL:
+    OS << "i1 " << (RC.Value.BoolVal ? "1" : "0");
+    break;
+  case RuntimeConstantType::INT8:
+    OS << "i8 " << static_cast<int>(RC.Value.Int8Val);
+    break;
+  case RuntimeConstantType::INT32:
+    OS << "i32 " << RC.Value.Int32Val;
+    break;
+  case RuntimeConstantType::INT64:
+    OS << "i64 " << RC.Value.Int64Val;
+    break;
+  case RuntimeConstantType::FLOAT:
+    OS << "float " << format("%g", RC.Value.FloatVal);
+    break;
+  case RuntimeConstantType::DOUBLE:
+    OS << "double " << format("%g", RC.Value.DoubleVal);
+    break;
+  default:
+    OS << "<unsupported type>";
+    break;
+  }
+
+  OS << "\n";
+  return S;
+}
+
+inline llvm::SmallVector<RuntimeConstant>
+resolveLambdaSpecializationValues(ArrayRef<RuntimeConstant> ExplicitValues,
+                                  Function *LambdaFn, const void *LambdaClosure,
+                                  bool EnableAutoReadOnlyCaptures,
+                                  bool TraceSpecializations) {
+  llvm::SmallVector<RuntimeConstant> MergedValues(ExplicitValues.begin(),
+                                                  ExplicitValues.end());
+
+  if (!EnableAutoReadOnlyCaptures || !LambdaFn || !LambdaClosure)
+    return MergedValues;
+
+  auto DetectedCaptures = parseAutoReadOnlyCapturesMetadata(*LambdaFn);
+  if (DetectedCaptures.empty())
+    return MergedValues;
+
+  auto AutoCaptures =
+      extractAutoDetectedCapturesFromMetadata(LambdaClosure, DetectedCaptures);
+  if (AutoCaptures.empty())
+    return MergedValues;
+
+  mergeCaptures(MergedValues, AutoCaptures);
+
+  if (!TraceSpecializations)
+    return MergedValues;
+
+  for (const auto &RC : AutoCaptures) {
+    if (!containsCaptureForSlot(ExplicitValues, RC.Pos))
+      Logger::trace(traceOutAuto(RC.Pos, RC));
+  }
+
+  return MergedValues;
+}
+
+} // namespace proteus
+
+#endif

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/LambdaSpecializationInfo.h"
 #include "proteus/impl/Utils.h"
 
 #include <llvm/Bitcode/BitcodeReader.h>
@@ -25,7 +26,7 @@ class CompilationTask {
   dim3 BlockDim;
   dim3 GridDim;
   SmallVector<RuntimeConstant> RCVec;
-  SmallVector<std::pair<std::string, StringRef>> LambdaCalleeInfo;
+  SmallVector<ResolvedLambdaSpecializationInfo> LambdaSpecializations;
   std::unordered_map<std::string, GlobalVarInfo> VarNameToGlobalInfo;
   SmallPtrSet<void *, 8> GlobalLinkedBinaries;
   std::string DeviceArch;
@@ -80,14 +81,15 @@ 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<ResolvedLambdaSpecializationInfo>
+          &LambdaSpecializations,
       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),
+        LambdaSpecializations(LambdaSpecializations),
         VarNameToGlobalInfo(VarNameToGlobalInfo),
         GlobalLinkedBinaries(GlobalLinkedBinaries), DeviceArch(DeviceArch),
         CGOption(CGConfig.codeGenOption()), DumpIR(DumpIR),
@@ -154,7 +156,7 @@ class CompilationTask {
     PROTEUS_DBG(Logger::logfile(HashValue.toString() + ".input.ll", *M));
 
     proteus::specializeIR(*M, KernelName, Suffix, BlockDim, GridDim, RCVec,
-                          LambdaCalleeInfo, SpecializeArgs, SpecializeDims,
+                          LambdaSpecializations, SpecializeArgs, SpecializeDims,
                           SpecializeDimsRange, SpecializeLaunchBounds,
                           MinBlocksPerSM);
 

src/include/proteus/impl/Config.h

@@ -352,6 +352,7 @@ class Config {
   std::string ProteusObjectCacheChain;
   bool ProteusEnableTimeTrace;
   std::string ProteusTimeTraceFile;
+  bool ProteusAutoReadOnlyCaptures;
   int ProteusTimeTraceGrainUs;
   int ProteusCommThreadPollMs;
 
@@ -392,6 +393,8 @@ class Config {
 
     OS << "PROTEUS_USE_STORED_CACHE " << ProteusUseStoredCache << "\n";
     OS << "PROTEUS_CACHE_DIR " << Config::get().ProteusCacheDir << "\n";
+    OS << "PROTEUS_AUTO_READONLY_CAPTURES " << ProteusAutoReadOnlyCaptures
+       << "\n";
 
     OS << PrintOut("Default", GlobalCodeGenConfig) << "\n";
     for (auto &KV : TunedConfigs) {
@@ -427,6 +430,8 @@ class Config {
         getEnvOrDefaultBool("PROTEUS_ENABLE_TIME_TRACE", false);
     ProteusTimeTraceFile =
         getEnvOrDefaultString("PROTEUS_TIME_TRACE_FILE").value_or("");
+    ProteusAutoReadOnlyCaptures =
+        getEnvOrDefaultBool("PROTEUS_AUTO_READONLY_CAPTURES", true);
     ProteusTimeTraceGrainUs =
         getEnvOrDefaultInt("PROTEUS_TIME_TRACE_GRAIN", 500);
     if (ProteusTimeTraceGrainUs <= 0)

src/include/proteus/impl/CoreLLVMDevice.h

@@ -13,7 +13,7 @@
 
 #include "proteus/impl/CoreDevice.h"
 #include "proteus/impl/GlobalVarInfo.h"
-#include "proteus/impl/LambdaRegistry.h"
+#include "proteus/impl/LambdaSpecializationInfo.h"
 #include "proteus/impl/TransformArgumentSpecialization.h"
 #include "proteus/impl/TransformLambdaSpecialization.h"
 #include "proteus/impl/TransformSharedArray.h"
@@ -283,12 +283,12 @@ 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,
-    bool SpecializeArgs, bool SpecializeDims, bool SpecializeDimsRange,
-    bool SpecializeLaunchBounds, int MinBlocksPerSM) {
+inline void
+specializeIR(Module &M, StringRef FnName, StringRef Suffix, dim3 &BlockDim,
+             dim3 &GridDim, ArrayRef<RuntimeConstant> RCArray,
+             ArrayRef<ResolvedLambdaSpecializationInfo> LambdaSpecializations,
+             bool SpecializeArgs, bool SpecializeDims, bool SpecializeDimsRange,
+             bool SpecializeLaunchBounds, int MinBlocksPerSM) {
   TIMESCOPE("proteus::specializeIR");
   Timer T(Config::get().ProteusEnableTimers);
   Function *F = M.getFunction(FnName);
@@ -298,13 +298,12 @@ 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)
+  for (const auto &LambdaSpecialization : LambdaSpecializations) {
+    Function *LambdaFn = M.getFunction(LambdaSpecialization.CalleeName);
+    if (!LambdaFn)
       reportFatalError("Expected non-null Function");
-    TransformLambdaSpecialization::transform(M, *F, RCVec);
+    TransformLambdaSpecialization::transform(M, *LambdaFn,
+                                             LambdaSpecialization.Values);
   }
 
   // Run the shared array transform after any value specialization (arguments,