diff --git a/.gitignore b/.gitignore index 3eded22e2..b99f287ae 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,5 @@ miniconda3* install-* docs/doxygen* site +**.proteus* +**.proteus-logs* diff --git a/docs/user/config.md b/docs/user/config.md index 29f59420c..f07afd2df 100644 --- a/docs/user/config.md +++ b/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"` | diff --git a/include/proteus/CompilerInterfaceTypes.h b/include/proteus/CompilerInterfaceTypes.h index 5a3908691..7e2b1fe4a 100644 --- a/include/proteus/CompilerInterfaceTypes.h +++ b/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}; diff --git a/src/include/proteus/impl/AutoReadOnlyCaptures.h b/src/include/proteus/impl/AutoReadOnlyCaptures.h new file mode 100644 index 000000000..8c4ace2fe --- /dev/null +++ b/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 +#include + +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 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 &Explicit, + llvm::ArrayRef Auto) { + llvm::SmallSet 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(Ptr); + } else if (RCType == RuntimeConstantType::INT8) { + RC.Type = RuntimeConstantType::INT8; + RC.Value.Int8Val = *static_cast(Ptr); + } else if (RCType == RuntimeConstantType::INT32) { + RC.Type = RuntimeConstantType::INT32; + RC.Value.Int32Val = *static_cast(Ptr); + } else if (RCType == RuntimeConstantType::INT64) { + RC.Type = RuntimeConstantType::INT64; + RC.Value.Int64Val = *static_cast(Ptr); + } else if (RCType == RuntimeConstantType::FLOAT) { + RC.Type = RuntimeConstantType::FLOAT; + RC.Value.FloatVal = *static_cast(Ptr); + } else if (RCType == RuntimeConstantType::DOUBLE) { + RC.Type = RuntimeConstantType::DOUBLE; + RC.Value.DoubleVal = *static_cast(Ptr); + } + + return RC; +} + +inline llvm::SmallVector +parseAutoReadOnlyCapturesMetadata(Function &F) { + llvm::SmallVector Captures; + MDNode *Root = F.getMetadata("proteus.auto_readonly_captures"); + if (!Root) + return Captures; + + auto ParseI32 = [](Metadata *M) -> std::optional { + auto *CAM = dyn_cast(M); + if (!CAM) + return std::nullopt; + auto *CI = dyn_cast(CAM->getValue()); + if (!CI) + return std::nullopt; + return static_cast(CI->getSExtValue()); + }; + + for (unsigned I = 0; I < Root->getNumOperands(); ++I) { + auto *EntryNode = dyn_cast_or_null(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(*RCTypeInt); + if (!isSupportedAutoReadOnlyRCType(RCType)) + continue; + + Captures.push_back( + AutoReadOnlyCaptureMetadataEntry{*SlotIndex, *ByteOffset, RCType}); + } + + return Captures; +} + +inline llvm::SmallVector +extractAutoDetectedCapturesFromMetadata( + const void *LambdaClosure, + const llvm::SmallVector + &DetectedCaptures) { + llvm::SmallVector Result; + if (!LambdaClosure || DetectedCaptures.empty()) + return Result; + + const char *ClosureBytes = static_cast(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(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 << ""; + break; + } + + OS << "\n"; + return S; +} + +inline llvm::SmallVector +resolveLambdaSpecializationValues(ArrayRef ExplicitValues, + Function *LambdaFn, const void *LambdaClosure, + bool EnableAutoReadOnlyCaptures, + bool TraceSpecializations) { + llvm::SmallVector 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 diff --git a/src/include/proteus/impl/CompilationTask.h b/src/include/proteus/impl/CompilationTask.h index 2d0659b41..c89c5d523 100644 --- a/src/include/proteus/impl/CompilationTask.h +++ b/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 @@ -25,7 +26,7 @@ class CompilationTask { dim3 BlockDim; dim3 GridDim; SmallVector RCVec; - SmallVector> LambdaCalleeInfo; + SmallVector LambdaSpecializations; std::unordered_map VarNameToGlobalInfo; SmallPtrSet 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 &RCVec, - const SmallVector> &LambdaCalleeInfo, + const SmallVector + &LambdaSpecializations, const std::unordered_map &VarNameToGlobalInfo, const SmallPtrSet &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); diff --git a/src/include/proteus/impl/Config.h b/src/include/proteus/impl/Config.h index e209d174c..38f65adba 100644 --- a/src/include/proteus/impl/Config.h +++ b/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) diff --git a/src/include/proteus/impl/CoreLLVMDevice.h b/src/include/proteus/impl/CoreLLVMDevice.h index 56126388d..086f3e10e 100644 --- a/src/include/proteus/impl/CoreLLVMDevice.h +++ b/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 RCArray, - const SmallVector> 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 RCArray, + ArrayRef 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 &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, diff --git a/src/include/proteus/impl/Hashing.h b/src/include/proteus/impl/Hashing.h index 4bef4fb62..8821acead 100644 --- a/src/include/proteus/impl/Hashing.h +++ b/src/include/proteus/impl/Hashing.h @@ -3,6 +3,7 @@ #include "proteus/CompilerInterfaceTypes.h" #include "proteus/TimeTracing.h" +#include "proteus/impl/LambdaSpecializationInfo.h" #include "proteus/impl/RuntimeConstantTypeHelpers.h" #include @@ -134,6 +135,28 @@ inline HashT hashValue(ArrayRef Arr) { return HashValue; } +inline HashT +hashValue(const ResolvedLambdaSpecializationInfo &LambdaSpecialization) { + HashT HashValue = hashValue(LambdaSpecialization.CalleeName); + return stable_hash_combine( + HashValue.getValue(), + hashValue(ArrayRef{LambdaSpecialization.Values}) + .getValue()); +} + +inline HashT +hashValue(ArrayRef LambdaSpecializations) { + if (LambdaSpecializations.empty()) + return 0; + + HashT HashValue = hashValue(LambdaSpecializations.front()); + for (size_t I = 1; I < LambdaSpecializations.size(); ++I) + HashValue = stable_hash_combine( + HashValue.getValue(), hashValue(LambdaSpecializations[I]).getValue()); + + return HashValue; +} + inline HashT hashCombine(HashT A, HashT B) { return stable_hash_combine(A.getValue(), B.getValue()); } diff --git a/src/include/proteus/impl/JitEngineDevice.h b/src/include/proteus/impl/JitEngineDevice.h index bf0726ad9..14b558f80 100644 --- a/src/include/proteus/impl/JitEngineDevice.h +++ b/src/include/proteus/impl/JitEngineDevice.h @@ -14,6 +14,7 @@ #include "proteus/CompilerInterfaceTypes.h" #include "proteus/Init.h" #include "proteus/TimeTracing.h" +#include "proteus/impl/AutoReadOnlyCaptures.h" #include "proteus/impl/Caching/MemoryCache.h" #include "proteus/impl/Caching/ObjectCacheChain.h" #include "proteus/impl/Cloning.h" @@ -25,9 +26,12 @@ #include "proteus/impl/Hashing.h" #include "proteus/impl/JitEngine.h" #include "proteus/impl/JitEngineInfoRegistry.h" +#include "proteus/impl/LambdaRegistry.h" +#include "proteus/impl/LambdaSpecializationInfo.h" #include "proteus/impl/Utils.h" #include +#include #include #include #include @@ -42,6 +46,7 @@ #include #include #include +#include #include #include #include @@ -229,8 +234,7 @@ class JITKernelInfo { std::optional> Bitcode; std::optional> BinInfo; std::optional StaticHash; - std::optional>> - LambdaCalleeInfo; + std::optional> CachedLambdaCalleeInfo; public: JITKernelInfo(void *Kernel, BinaryInfo &BinInfo, char const *Name, @@ -238,7 +242,7 @@ class JITKernelInfo { : Kernel(Kernel), Ctx(std::make_unique()), Name(Name), RCInfoArray(RCInfoArray), ExtractedModule(std::nullopt), Bitcode{std::nullopt}, BinInfo(BinInfo), - LambdaCalleeInfo(std::nullopt) {} + CachedLambdaCalleeInfo(std::nullopt) {} JITKernelInfo() = default; void *getKernel() const { @@ -268,11 +272,11 @@ class JITKernelInfo { StaticHash = hashCombine(StaticHash.value(), ModuleHash); } - bool hasLambdaCalleeInfo() { return LambdaCalleeInfo.has_value(); } - const auto &getLambdaCalleeInfo() { return LambdaCalleeInfo.value(); } - void setLambdaCalleeInfo( - SmallVector> &&LambdaInfo) { - LambdaCalleeInfo = std::move(LambdaInfo); + bool hasLambdaCalleeInfo() { return CachedLambdaCalleeInfo.has_value(); } + const auto &getLambdaCalleeInfo() { return CachedLambdaCalleeInfo.value(); } + void + setLambdaCalleeInfo(SmallVector &&LambdaInfo) { + CachedLambdaCalleeInfo = std::move(LambdaInfo); } }; @@ -414,9 +418,297 @@ template class JitEngineDevice : public JitEngine { return KernelInfo.getBitcode(); } - void getLambdaJitValues(JITKernelInfo &KernelInfo, - SmallVector &LambdaJitValuesVec) { - TIMESCOPE(JitEngineDevice, getLambdaJitValues); + std::optional mergeKernelArgIndex(std::optional Current, + std::optional Candidate, + bool &HasAmbiguity) const { + if (!Candidate) + return Current; + if (!Current) + return Candidate; + if (*Current != *Candidate) + HasAmbiguity = true; + return Current; + } + + std::optional matchKernelArgIndexByType(Function &KernelFn, + Type *ClosureTy) const { + if (!ClosureTy) + return std::nullopt; + + std::optional Result; + bool HasAmbiguity = false; + for (Argument &Arg : KernelFn.args()) { + Type *KernelArgTy = Arg.getParamByRefType(); + if (!KernelArgTy) + KernelArgTy = Arg.getParamByValType(); + if (KernelArgTy != ClosureTy) + continue; + + Result = mergeKernelArgIndex(Result, Arg.getArgNo(), HasAmbiguity); + if (HasAmbiguity) + return std::nullopt; + } + + return Result; + } + + std::optional + findClosureStorageType(Value *V, SmallPtrSetImpl &Seen) const { + if (!V) + return std::nullopt; + if (!Seen.insert(V).second) + return std::nullopt; + + if (auto *AI = dyn_cast(V)) + return AI->getAllocatedType(); + + if (auto *Arg = dyn_cast(V)) { + if (Type *Ty = Arg->getParamByRefType()) + return Ty; + if (Type *Ty = Arg->getParamByValType()) + return Ty; + return std::nullopt; + } + + if (auto *BC = dyn_cast(V)) + return findClosureStorageType(BC->getOperand(0), Seen); + + if (auto *ASC = dyn_cast(V)) + return findClosureStorageType(ASC->getOperand(0), Seen); + + if (auto *LI = dyn_cast(V)) + return findClosureStorageType(LI->getPointerOperand(), Seen); + + if (auto *GEP = dyn_cast(V)) + return findClosureStorageType(GEP->getPointerOperand(), Seen); + + if (auto *PN = dyn_cast(V)) { + std::optional Result; + for (Value *Incoming : PN->incoming_values()) { + auto Candidate = findClosureStorageType(Incoming, Seen); + if (!Candidate) + continue; + if (!Result) { + Result = Candidate; + continue; + } + if (*Result != *Candidate) + return std::nullopt; + } + return Result; + } + + if (auto *SI = dyn_cast(V)) { + auto TrueTy = findClosureStorageType(SI->getTrueValue(), Seen); + auto FalseTy = findClosureStorageType(SI->getFalseValue(), Seen); + if (!TrueTy) + return FalseTy; + if (!FalseTy) + return TrueTy; + if (*TrueTy != *FalseTy) + return std::nullopt; + return TrueTy; + } + + return std::nullopt; + } + + std::optional + traceStoredKernelArgIndex(Function &KernelFn, Value *Ptr, + SmallPtrSetImpl &SeenPointers) const { + if (!SeenPointers.insert(Ptr).second) + return std::nullopt; + + std::optional Result; + bool HasAmbiguity = false; + auto MergeCandidate = [&](Value *Candidate) { + if (!Candidate) + return; + + SmallPtrSet CandidateSeenValues; + SmallPtrSet CandidateSeenFunctions; + auto CandidateIndex = + traceKernelArgIndex(KernelFn, Candidate->stripPointerCasts(), + CandidateSeenValues, CandidateSeenFunctions); + Result = mergeKernelArgIndex(Result, CandidateIndex, HasAmbiguity); + }; + + for (User *User : Ptr->users()) { + if (auto *MemCpy = dyn_cast(User)) { + if (MemCpy->getDest() == Ptr) + MergeCandidate(MemCpy->getSource()); + continue; + } + + if (auto *SI = dyn_cast(User)) { + if (SI->getPointerOperand() == Ptr) + MergeCandidate(SI->getValueOperand()); + continue; + } + + if (isa(User) || isa(User) || + isa(User) || isa(User) || + isa(User)) { + auto Candidate = traceStoredKernelArgIndex(KernelFn, cast(User), + SeenPointers); + Result = mergeKernelArgIndex(Result, Candidate, HasAmbiguity); + } + + if (HasAmbiguity) + return std::nullopt; + } + + return Result; + } + + std::optional + traceKernelArgIndex(Function &KernelFn, Value *V, + SmallPtrSetImpl &SeenValues, + SmallPtrSetImpl &SeenFunctions) const { + if (!V) + return std::nullopt; + + if (!SeenValues.insert(V).second) + return std::nullopt; + + if (auto *BC = dyn_cast(V)) + return traceKernelArgIndex(KernelFn, BC->getOperand(0), SeenValues, + SeenFunctions); + + if (auto *ASC = dyn_cast(V)) + return traceKernelArgIndex(KernelFn, ASC->getOperand(0), SeenValues, + SeenFunctions); + + if (auto *LI = dyn_cast(V)) + return traceKernelArgIndex(KernelFn, LI->getPointerOperand(), SeenValues, + SeenFunctions); + + if (auto *Arg = dyn_cast(V)) { + if (Arg->getParent() == &KernelFn) + return Arg->getArgNo(); + + Function *ParentFn = Arg->getParent(); + if (!SeenFunctions.insert(ParentFn).second) + return std::nullopt; + + std::optional Result; + bool HasAmbiguity = false; + for (User *User : ParentFn->users()) { + auto *CB = dyn_cast(User); + if (!CB) + continue; + if (CB->getCalledOperand()->stripPointerCasts() != ParentFn) + continue; + if (Arg->getArgNo() >= CB->arg_size()) + continue; + + auto Candidate = + traceKernelArgIndex(KernelFn, CB->getArgOperand(Arg->getArgNo()), + SeenValues, SeenFunctions); + Result = mergeKernelArgIndex(Result, Candidate, HasAmbiguity); + if (HasAmbiguity) + return std::nullopt; + } + + return Result; + } + + if (auto *AI = dyn_cast(V)) { + SmallPtrSet SeenPointers; + return traceStoredKernelArgIndex(KernelFn, AI, SeenPointers); + } + + if (auto *GEP = dyn_cast(V)) + return traceKernelArgIndex(KernelFn, GEP->getPointerOperand(), SeenValues, + SeenFunctions); + + if (auto *PN = dyn_cast(V)) { + std::optional Result; + bool HasAmbiguity = false; + for (Value *Incoming : PN->incoming_values()) { + auto Candidate = + traceKernelArgIndex(KernelFn, Incoming, SeenValues, SeenFunctions); + Result = mergeKernelArgIndex(Result, Candidate, HasAmbiguity); + if (HasAmbiguity) + return std::nullopt; + } + return Result; + } + + if (auto *SI = dyn_cast(V)) { + auto TrueCandidate = traceKernelArgIndex(KernelFn, SI->getTrueValue(), + SeenValues, SeenFunctions); + bool HasAmbiguity = false; + auto Result = + mergeKernelArgIndex(std::nullopt, TrueCandidate, HasAmbiguity); + auto FalseCandidate = traceKernelArgIndex(KernelFn, SI->getFalseValue(), + SeenValues, SeenFunctions); + Result = mergeKernelArgIndex(Result, FalseCandidate, HasAmbiguity); + if (HasAmbiguity) + return std::nullopt; + return Result; + } + + return std::nullopt; + } + + SmallVector + discoverLambdaCalleeInfo(JITKernelInfo &KernelInfo, Module &KernelModule) { + Function *KernelFn = KernelModule.getFunction(KernelInfo.getName()); + if (!KernelFn) + reportFatalError("Expected non-null kernel function"); + + SmallVector LambdaInfo; + for (auto &F : KernelModule.getFunctionList()) { + PROTEUS_DBG(Logger::logs("proteus") + << " Trying F " << demangle(F.getName().str()) << "\n "); + auto OptionalMapIt = + LambdaRegistry::instance().matchJitVariableMap(F.getName()); + if (!OptionalMapIt) + continue; + + std::optional KernelArgIndex; + bool HasAmbiguity = false; + for (User *User : F.users()) { + auto *CB = dyn_cast(User); + if (!CB) + continue; + if (CB->getCalledOperand()->stripPointerCasts() != &F) + continue; + if (CB->arg_empty()) + continue; + + SmallPtrSet SeenValues; + SmallPtrSet SeenFunctions; + auto Candidate = traceKernelArgIndex(*KernelFn, CB->getArgOperand(0), + SeenValues, SeenFunctions); + if (!Candidate) { + SmallPtrSet SeenTypes; + auto ClosureTy = + findClosureStorageType(CB->getArgOperand(0), SeenTypes); + Candidate = + matchKernelArgIndexByType(*KernelFn, ClosureTy.value_or(nullptr)); + } + KernelArgIndex = + mergeKernelArgIndex(KernelArgIndex, Candidate, HasAmbiguity); + if (HasAmbiguity) + break; + } + + LambdaCalleeInfo Info{ + F.getName().str(), OptionalMapIt.value()->first.str(), + KernelArgIndex ? static_cast(*KernelArgIndex) : -1}; + LambdaInfo.push_back(std::move(Info)); + } + + return LambdaInfo; + } + + void resolveLambdaSpecializations( + JITKernelInfo &KernelInfo, + SmallVector &LambdaSpecializations, + void **KernelArgs) { + TIMESCOPE(JitEngineDevice, resolveLambdaSpecializations); LambdaRegistry &LR = LambdaRegistry::instance(); if (LR.empty()) { KernelInfo.setLambdaCalleeInfo({}); @@ -430,25 +722,29 @@ template class JitEngineDevice : public JitEngine { << "Caller trigger " << KernelInfo.getName() << " -> " << demangle(KernelInfo.getName()) << "\n"); - SmallVector> LambdaCalleeInfo; - for (auto &F : KernelModule.getFunctionList()) { - PROTEUS_DBG(Logger::logs("proteus") - << " Trying F " << demangle(F.getName().str()) << "\n "); - auto OptionalMapIt = - LambdaRegistry::instance().matchJitVariableMap(F.getName()); - if (OptionalMapIt) - LambdaCalleeInfo.emplace_back(F.getName(), - OptionalMapIt.value()->first); - } - - KernelInfo.setLambdaCalleeInfo(std::move(LambdaCalleeInfo)); + KernelInfo.setLambdaCalleeInfo( + discoverLambdaCalleeInfo(KernelInfo, KernelModule)); } - for (auto &[FnName, LambdaType] : KernelInfo.getLambdaCalleeInfo()) { - const SmallVector &Values = - LR.getJitVariables(LambdaType); - LambdaJitValuesVec.insert(LambdaJitValuesVec.end(), Values.begin(), - Values.end()); + Module &KernelModule = getModule(KernelInfo); + for (const auto &Info : KernelInfo.getLambdaCalleeInfo()) { + // Get explicit jit_variable captures + const SmallVector &ExplicitValues = + LR.getJitVariables(Info.LambdaType); + + Function *LambdaFn = KernelModule.getFunction(Info.CalleeName); + const void *LambdaClosure = + (KernelArgs && Info.KernelArgIndex >= 0) + ? KernelArgs[static_cast(Info.KernelArgIndex)] + : nullptr; + SmallVector MergedValues = + resolveLambdaSpecializationValues( + ExplicitValues, LambdaFn, LambdaClosure, + Config::get().ProteusAutoReadOnlyCaptures, + Config::get().traceSpecializations()); + + LambdaSpecializations.push_back(ResolvedLambdaSpecializationInfo{ + Info.CalleeName, std::move(MergedValues)}); } } @@ -576,13 +872,15 @@ JitEngineDevice::compileAndRun( SmallVector RCVec = getRuntimeConstantValues(KernelArgs, KernelInfo.getRCInfoArray()); - SmallVector LambdaJitValuesVec; - getLambdaJitValues(KernelInfo, LambdaJitValuesVec); + SmallVector LambdaSpecializations; + resolveLambdaSpecializations(KernelInfo, LambdaSpecializations, KernelArgs); // Determine the hash based on dimension specialization. If we do not // specialize IR based on grid dimensions, avoid hashing on those to // eliminate repeated compilation overhead. - HashT HashValue = hash(getStaticHash(KernelInfo), RCVec, LambdaJitValuesVec, - BlockDim.x, BlockDim.y, BlockDim.z); + HashT HashValue = + hash(getStaticHash(KernelInfo), RCVec, + ArrayRef{LambdaSpecializations}, + BlockDim.x, BlockDim.y, BlockDim.z); if (Config::get().getCGConfig().specializeDims() || Config::get().getCGConfig().specializeDimsRange()) HashValue = hash(HashValue, GridDim.x, GridDim.y, GridDim.z); @@ -631,7 +929,7 @@ JitEngineDevice::compileAndRun( AsyncCompiler->compile(CompilationTask{ KernelBitcode, HashValue, KernelInfo.getName(), Suffix, BlockDim, - GridDim, RCVec, KernelInfo.getLambdaCalleeInfo(), + GridDim, RCVec, LambdaSpecializations, BinInfo.getVarNameToGlobalInfo(), GlobalLinkedBinaries, DeviceArch, /*CodeGenConfig */ Config::get().getCGConfig(KernelInfo.getName()), /*DumpIR*/ Config::get().ProteusDumpLLVMIR, @@ -651,8 +949,8 @@ JitEngineDevice::compileAndRun( // Process through synchronous compilation. ObjBuf = CompilerSync::instance().compile(CompilationTask{ KernelBitcode, HashValue, KernelInfo.getName(), Suffix, BlockDim, - GridDim, RCVec, KernelInfo.getLambdaCalleeInfo(), - BinInfo.getVarNameToGlobalInfo(), GlobalLinkedBinaries, DeviceArch, + GridDim, RCVec, LambdaSpecializations, BinInfo.getVarNameToGlobalInfo(), + GlobalLinkedBinaries, DeviceArch, /*CodeGenConfig */ Config::get().getCGConfig(KernelInfo.getName()), /*DumpIR*/ Config::get().ProteusDumpLLVMIR, /*RelinkGlobalsByCopy*/ Config::get().ProteusRelinkGlobalsByCopy}); diff --git a/src/include/proteus/impl/JitEngineHost.h b/src/include/proteus/impl/JitEngineHost.h index 17f489f7a..07e60b166 100644 --- a/src/include/proteus/impl/JitEngineHost.h +++ b/src/include/proteus/impl/JitEngineHost.h @@ -46,7 +46,8 @@ class JitEngineHost : public JitEngine { ~JitEngineHost(); void specializeIR(Module &M, StringRef FnName, StringRef Suffix, - ArrayRef RCArray); + ArrayRef RCArray, + const SmallVector &LambdaJitValuesVec); void *compileAndLink(StringRef FnName, char *IR, int IRSize, void **Args, ArrayRef RCInfoArray); diff --git a/src/include/proteus/impl/LambdaRegistry.h b/src/include/proteus/impl/LambdaRegistry.h index 2f8389f14..649fbe51e 100644 --- a/src/include/proteus/impl/LambdaRegistry.h +++ b/src/include/proteus/impl/LambdaRegistry.h @@ -86,13 +86,14 @@ class LambdaRegistry { const StringRef LambdaTypeRef{LambdaType}; PROTEUS_DBG(Logger::logs("proteus") << "=> RegisterLambda " << LambdaTypeRef << "\n"); - // Copy PendingJitVariables if there were changed, otherwise the runtime - // values for the lambda definition have not changed. + // Register new lambda types even when there are no explicit captures, + // but do not overwrite an existing capture set with an empty update. PROTEUS_DBG(dump()); - if (!PendingJitVariableMap[LambdaTypeRef].empty()) { - JitVariableMap[LambdaTypeRef] = PendingJitVariableMap[LambdaTypeRef]; - PendingJitVariableMap[LambdaTypeRef].clear(); - } + auto &PendingValues = PendingJitVariableMap[LambdaTypeRef]; + if (!JitVariableMap.contains(LambdaTypeRef) || !PendingValues.empty()) + JitVariableMap[LambdaTypeRef] = PendingValues; + if (!PendingValues.empty()) + PendingValues.clear(); } const SmallVector &getJitVariables(StringRef LambdaTypeRef) { diff --git a/src/include/proteus/impl/LambdaSpecializationInfo.h b/src/include/proteus/impl/LambdaSpecializationInfo.h new file mode 100644 index 000000000..4c8afde38 --- /dev/null +++ b/src/include/proteus/impl/LambdaSpecializationInfo.h @@ -0,0 +1,26 @@ +#ifndef PROTEUS_LAMBDA_SPECIALIZATION_INFO_H +#define PROTEUS_LAMBDA_SPECIALIZATION_INFO_H + +#include "proteus/CompilerInterfaceTypes.h" + +#include + +#include +#include + +namespace proteus { + +struct LambdaCalleeInfo { + std::string CalleeName; + std::string LambdaType; + int32_t KernelArgIndex = -1; +}; + +struct ResolvedLambdaSpecializationInfo { + std::string CalleeName; + llvm::SmallVector Values; +}; + +} // namespace proteus + +#endif diff --git a/src/include/proteus/impl/TransformLambdaSpecialization.h b/src/include/proteus/impl/TransformLambdaSpecialization.h index d9bcc97cb..9ad4ef1f5 100644 --- a/src/include/proteus/impl/TransformLambdaSpecialization.h +++ b/src/include/proteus/impl/TransformLambdaSpecialization.h @@ -15,13 +15,18 @@ #include "proteus/impl/Debug.h" #include "proteus/impl/Utils.h" +#include +#include #include -#include +#include #include #include #include #include +#include +#include + namespace proteus { using namespace llvm; @@ -84,41 +89,75 @@ class TransformLambdaSpecialization { return S; }; - static void handleLoad(Module &M, User *User, - const SmallVector &RCVec) { - auto *Arg = findArgByPos(RCVec, 0); + static void replaceLoad(Module &M, LoadInst *LI, int32_t ByteOffset, + const SmallVector &RCVec) { + auto *Arg = findArgByOffset(RCVec, ByteOffset); + if (!Arg && ByteOffset == 0) + Arg = findArgByPos(RCVec, 0); if (!Arg) return; - Constant *C = getConstant(M.getContext(), User->getType(), *Arg); - User->replaceAllUsesWith(C); + Constant *C = getConstant(M.getContext(), LI->getType(), *Arg); + LI->replaceAllUsesWith(C); PROTEUS_DBG(Logger::logs("proteus") << traceOut(Arg->Pos, C)); if (Config::get().traceSpecializations()) Logger::trace(traceOut(Arg->Pos, C)); } - static void handleGEP(Module &M, GetElementPtrInst *GEP, User *User, - const SmallVector &RCVec) { - auto *GEPSlot = GEP->getOperand(User->getNumOperands() - 1); - ConstantInt *CI = dyn_cast(GEPSlot); - int Slot = CI->getZExtValue(); - Type *SrcTy = GEP->getSourceElementType(); + static std::optional getGEPByteOffset(const DataLayout &DL, + GetElementPtrInst *GEP) { + APInt Offset(DL.getPointerTypeSizeInBits(GEP->getType()), 0, true); + if (!GEP->accumulateConstantOffset(DL, Offset)) + return std::nullopt; - auto *Arg = SrcTy->isStructTy() ? findArgByPos(RCVec, Slot) - : findArgByOffset(RCVec, Slot); - if (!Arg) + int64_t Offset64 = Offset.getSExtValue(); + if (Offset64 < std::numeric_limits::min() || + Offset64 > std::numeric_limits::max()) + return std::nullopt; + + return static_cast(Offset64); + } + + static void visitPointerUsers(Module &M, Value *Ptr, int32_t ByteOffset, + const SmallVector &RCVec, + SmallPtrSetImpl &Seen) { + if (!Seen.insert(Ptr).second) return; - for (auto *GEPUser : GEP->users()) { - auto *LI = dyn_cast(GEPUser); - if (!LI) - reportFatalError("Expected load instruction"); - Type *LoadType = LI->getType(); - Constant *C = getConstant(M.getContext(), LoadType, *Arg); - LI->replaceAllUsesWith(C); - PROTEUS_DBG(Logger::logs("proteus") << traceOut(Arg->Pos, C)); - if (Config::get().traceSpecializations()) - Logger::trace(traceOut(Arg->Pos, C)); + const DataLayout &DL = M.getDataLayout(); + for (User *User : Ptr->users()) { + if (auto *LI = dyn_cast(User)) { + if (LI->getPointerOperand() != Ptr) + continue; + replaceLoad(M, LI, ByteOffset, RCVec); + continue; + } + + if (auto *GEP = dyn_cast(User)) { + if (GEP->getPointerOperand() != Ptr) + continue; + + auto LocalOffset = getGEPByteOffset(DL, GEP); + if (!LocalOffset) + continue; + + visitPointerUsers(M, GEP, ByteOffset + *LocalOffset, RCVec, Seen); + continue; + } + + if (isa(User) || isa(User) || + isa(User) || isa(User)) { + bool UsesPtr = false; + for (Value *Operand : User->operands()) { + if (Operand == Ptr) { + UsesPtr = true; + break; + } + } + if (!UsesPtr) + continue; + visitPointerUsers(M, cast(User), ByteOffset, RCVec, Seen); + } } } @@ -139,14 +178,8 @@ class TransformLambdaSpecialization { } } - PROTEUS_DBG(Logger::logs("proteus") << "\t users" << "\n"); - for (User *User : LambdaClass->users()) { - PROTEUS_DBG(Logger::logs("proteus") << *User << "\n"); - if (isa(User)) - handleLoad(M, User, RCVec); - else if (auto *GEP = dyn_cast(User)) - handleGEP(M, GEP, User, RCVec); - } + SmallPtrSet Seen; + visitPointerUsers(M, LambdaClass, /*ByteOffset=*/0, RCVec, Seen); } }; diff --git a/src/pass/AutoReadOnlyCapturesAnalysis.cpp b/src/pass/AutoReadOnlyCapturesAnalysis.cpp new file mode 100644 index 000000000..9f184090f --- /dev/null +++ b/src/pass/AutoReadOnlyCapturesAnalysis.cpp @@ -0,0 +1,529 @@ +#include "AutoReadOnlyCapturesAnalysis.h" + +#include "Helpers.h" +#include "proteus/impl/RuntimeConstantTypeHelpers.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include + +using namespace llvm; + +namespace proteus { +namespace { + +constexpr char AutoReadOnlyCapturesMetadataName[] = + "proteus.auto_readonly_captures"; + +struct SlotState { + int32_t ByteOffset; + RuntimeConstantType RCType = RuntimeConstantType::NONE; + bool IsReadOnly = true; +}; + +enum class PointerUseEffect { + Ignore, + ReadOnly, + BenignTransform, + WriteOrEscape, +}; + +std::optional classifySupportedScalar(Type *Ty) { + if (!(Ty->isIntegerTy(1) || Ty->isIntegerTy(8) || Ty->isIntegerTy(32) || + Ty->isIntegerTy(64) || Ty->isFloatTy() || Ty->isDoubleTy())) { + return std::nullopt; + } + + RuntimeConstantType RCType = convertTypeToRuntimeConstantType(Ty); + if (!isSupportedAutoReadOnlyRCType(RCType)) + return std::nullopt; + return RCType; +} + +std::optional getGEPByteOffset(Function &F, GetElementPtrInst *GEP) { + const DataLayout &DL = F.getParent()->getDataLayout(); + APInt Offset(DL.getPointerTypeSizeInBits(GEP->getType()), 0, true); + if (!GEP->accumulateConstantOffset(DL, Offset)) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Cannot compute constant GEP offset: " << *GEP + << "\n"); + return std::nullopt; + } + + int64_t ByteOffset64 = Offset.getSExtValue(); + if (ByteOffset64 < 0 || ByteOffset64 > std::numeric_limits::max()) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] GEP offset out of range: " << ByteOffset64 + << " for GEP: " << *GEP << "\n"); + return std::nullopt; + } + + return static_cast(ByteOffset64); +} + +std::optional getTopLevelSlotIndex(Argument *ClosureArg, + Value *BasePtr, + GetElementPtrInst *GEP) { + if (BasePtr != ClosureArg || GEP->getNumIndices() != 2) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] GEP not top-level closure access: " << *GEP + << "\n"); + return std::nullopt; + } + + auto *SourceStruct = dyn_cast(GEP->getSourceElementType()); + if (!SourceStruct) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] GEP source not a struct: " << *GEP << "\n"); + return std::nullopt; + } + + auto *SlotIdxConst = dyn_cast(GEP->getOperand(2)); + if (!SlotIdxConst) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] GEP slot index not constant: " << *GEP << "\n"); + return std::nullopt; + } + + int64_t SlotIdx64 = SlotIdxConst->getSExtValue(); + if (SlotIdx64 < 0) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] GEP slot index negative: " << SlotIdx64 + << " in: " << *GEP << "\n"); + return std::nullopt; + } + + int32_t SlotIndex = static_cast(SlotIdx64); + if (static_cast(SlotIndex) != SlotIdx64) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] GEP slot index overflow: " << SlotIdx64 + << " in: " << *GEP << "\n"); + return std::nullopt; + } + + if (static_cast(SlotIndex) >= SourceStruct->getNumElements()) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] GEP slot index out of bounds: " << SlotIndex + << " >= " << SourceStruct->getNumElements() << " in: " << *GEP + << "\n"); + return std::nullopt; + } + + if (!classifySupportedScalar(GEP->getResultElementType())) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Unsupported GEP result type: " + << *GEP->getResultElementType() << " in: " << *GEP << "\n"); + return std::nullopt; + } + + return SlotIndex; +} + +PointerUseEffect classifyPointerUse(User *U, Value *Ptr) { + if (auto *LI = dyn_cast(U)) { + if (LI->getPointerOperand() == Ptr) + return PointerUseEffect::ReadOnly; + return PointerUseEffect::WriteOrEscape; + } + + if (auto *SI = dyn_cast(U)) { + if (SI->getPointerOperand() == Ptr) + return PointerUseEffect::WriteOrEscape; + if (SI->getValueOperand() == Ptr) + return PointerUseEffect::WriteOrEscape; + return PointerUseEffect::Ignore; + } + + if (isa(U)) + return PointerUseEffect::WriteOrEscape; + + if (isa(U) || isa(U) || + isa(U) || isa(U) || isa(U)) + return PointerUseEffect::BenignTransform; + + return PointerUseEffect::WriteOrEscape; +} + +class ClosureEscapeTracker final : public CaptureTracker { + bool Captured = false; + +public: + bool pointerEscapes() const { return Captured; } + + void tooManyUses() override { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Closure has too many uses for tracking\n"); + Captured = true; + } + + bool shouldExplore(const Use *U) override { + // Field accesses are analyzed per-slot below. Escapes through a field GEP + // should disqualify only that slot, not the whole closure object. + if (auto *GEP = dyn_cast(U->getUser()); + GEP && GEP->getPointerOperand() == U->get()) { + return false; + } + + return true; + } + + bool captured(const Use *U) override { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Closure escapes through: " << *U->getUser() + << "\n"); + Captured = true; + return true; + } +}; + +bool pointerEscapes(Value *RootPtr) { + assert(RootPtr->getType()->isPointerTy() && "Expected pointer root"); + + ClosureEscapeTracker Tracker; + PointerMayBeCaptured(RootPtr, &Tracker); + return Tracker.pointerEscapes(); +} + +std::optional +parseDirectLoadCapture(LoadInst *LI, int32_t SlotIndex, int32_t ByteOffset) { + auto RCType = classifySupportedScalar(LI->getType()); + if (!RCType) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Unsupported load type: " << *LI->getType() + << " in: " << *LI << "\n"); + return std::nullopt; + } + + return AutoReadOnlyCaptureMetadataEntry{SlotIndex, ByteOffset, *RCType}; +} + +void updateReadOnlySlot(DenseMap &Slots, + const AutoReadOnlyCaptureMetadataEntry &Entry) { + auto It = Slots.find(Entry.SlotIndex); + if (It == Slots.end()) { + SlotState NewState; + NewState.ByteOffset = Entry.ByteOffset; + NewState.RCType = Entry.RCType; + Slots.insert({Entry.SlotIndex, NewState}); + return; + } + + SlotState &State = It->second; + if (!State.IsReadOnly) + return; + + if (State.ByteOffset != Entry.ByteOffset) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Slot " << Entry.SlotIndex + << " disqualified: inconsistent byte offsets (" << State.ByteOffset + << " vs " << Entry.ByteOffset << ")\n"); + State.IsReadOnly = false; + return; + } + + if (State.RCType == RuntimeConstantType::NONE) { + State.RCType = Entry.RCType; + return; + } + + if (State.RCType != Entry.RCType) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Slot " << Entry.SlotIndex + << " disqualified: inconsistent types (" << toString(State.RCType) + << " vs " << toString(Entry.RCType) << ")\n"); + State.IsReadOnly = false; + } +} + +void markSlotNonReadOnly(DenseMap &Slots, int32_t SlotIdx, + int32_t ByteOffset) { + auto It = Slots.find(SlotIdx); + if (It == Slots.end()) { + SlotState NewState; + NewState.ByteOffset = ByteOffset; + NewState.IsReadOnly = false; + Slots.insert({SlotIdx, NewState}); + return; + } + + It->second.IsReadOnly = false; +} + +void analyzePointerUsersForSlot(Function &F, Argument *ClosureArg, + Value *RootPtr, int32_t SlotIndex, + int32_t ByteOffset, + DenseMap &Slots) { + auto Existing = Slots.find(SlotIndex); + if (Existing != Slots.end() && !Existing->second.IsReadOnly) + return; + + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Analyzing slot " << SlotIndex << " at offset " + << ByteOffset << "\n"); + + SmallVector WorkList{RootPtr}; + SmallPtrSet Seen; + + while (!WorkList.empty()) { + Value *V = WorkList.pop_back_val(); + if (!Seen.insert(V).second) + continue; + + for (User *U : V->users()) { + if (auto *LI = dyn_cast(U)) { + if (LI->getPointerOperand() != V) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Slot " << SlotIndex + << " disqualified: pointer used as value in load: " << *LI + << "\n"); + markSlotNonReadOnly(Slots, SlotIndex, ByteOffset); + break; + } + + auto Parsed = parseDirectLoadCapture(LI, SlotIndex, ByteOffset); + if (!Parsed) + continue; + + updateReadOnlySlot(Slots, *Parsed); + continue; + } + + if (auto *SI = dyn_cast(U)) { + if (SI->getPointerOperand() == V || SI->getValueOperand() == V) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Slot " << SlotIndex + << " disqualified: written by store: " << *SI << "\n"); + markSlotNonReadOnly(Slots, SlotIndex, ByteOffset); + break; + } + continue; + } + + PointerUseEffect Effect = classifyPointerUse(U, V); + if (Effect == PointerUseEffect::Ignore || + Effect == PointerUseEffect::ReadOnly) + continue; + + if (Effect == PointerUseEffect::BenignTransform) { + if (auto *GEP = dyn_cast(U)) { + auto LocalOffset = getGEPByteOffset(F, GEP); + if (!LocalOffset) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Slot " << SlotIndex + << " disqualified: non-constant nested GEP: " << *GEP + << "\n"); + markSlotNonReadOnly(Slots, SlotIndex, ByteOffset); + break; + } + + int32_t NestedByteOffset = ByteOffset + *LocalOffset; + auto NestedSlot = getTopLevelSlotIndex(ClosureArg, V, GEP); + analyzePointerUsersForSlot(F, ClosureArg, GEP, + NestedSlot.value_or(NestedByteOffset), + NestedByteOffset, Slots); + continue; + } + WorkList.push_back(cast(U)); + continue; + } + + if (Effect == PointerUseEffect::WriteOrEscape) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Slot " << SlotIndex + << " disqualified: write/escape use: " << *U << "\n"); + } + markSlotNonReadOnly(Slots, SlotIndex, ByteOffset); + break; + } + + auto It = Slots.find(SlotIndex); + if (It != Slots.end() && !It->second.IsReadOnly) + break; + } +} + +SmallVector +collectReadOnlyCaptures(const DenseMap &Slots) { + SmallVector Captures; + + for (const auto &Entry : Slots) { + int32_t SlotIndex = Entry.first; + const SlotState &State = Entry.second; + if (!State.IsReadOnly) + continue; + if (State.RCType == RuntimeConstantType::NONE) + continue; + + Captures.push_back(AutoReadOnlyCaptureMetadataEntry{ + SlotIndex, State.ByteOffset, State.RCType}); + } + + llvm::sort(Captures, [](const AutoReadOnlyCaptureMetadataEntry &L, + const AutoReadOnlyCaptureMetadataEntry &R) { + return L.SlotIndex < R.SlotIndex; + }); + + return Captures; +} + +} // namespace + +SmallVector +analyzeAutoReadOnlyCaptures(Function &F) { + SmallVector Captures; + + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Analyzing function: " << F.getName() << "\n"); + + if (F.arg_empty()) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Skipping " << F.getName() << ": no arguments\n"); + return Captures; + } + + Argument *ClosureArg = &*F.arg_begin(); + if (!ClosureArg->getType()->isPointerTy()) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Skipping " << F.getName() + << ": first argument not a pointer (type: " << *ClosureArg->getType() + << ")\n"); + return Captures; + } + + if (pointerEscapes(ClosureArg)) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Function " << F.getName() + << ": closure pointer escapes\n"); + return Captures; + } + + DenseMap Slots; + + SmallVector WorkList{ClosureArg}; + SmallPtrSet Seen; + + while (!WorkList.empty()) { + Value *V = WorkList.pop_back_val(); + if (!Seen.insert(V).second) + continue; + + for (User *U : V->users()) { + if (auto *LI = dyn_cast(U)) { + if (LI->getPointerOperand() != V) + continue; + + auto Parsed = + parseDirectLoadCapture(LI, /*SlotIndex=*/0, /*ByteOffset=*/0); + if (!Parsed) + continue; + + updateReadOnlySlot(Slots, *Parsed); + continue; + } + + if (auto *SI = dyn_cast(U)) { + if (SI->getPointerOperand() != V) + continue; + + // A direct store through the closure pointer writes the first slot. + markSlotNonReadOnly(Slots, /*SlotIdx=*/0, /*ByteOffset=*/0); + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Slot 0 disqualified: direct store to closure: " + << *SI << "\n"); + continue; + } + + if (auto *GEP = dyn_cast(U)) { + auto LocalOffset = getGEPByteOffset(F, GEP); + if (!LocalOffset) { + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Skipping GEP with non-constant offset: " + << *GEP << "\n"); + continue; + } + + auto TopLevelSlot = getTopLevelSlotIndex(ClosureArg, V, GEP); + int32_t ByteOffset = *LocalOffset; + analyzePointerUsersForSlot(F, ClosureArg, GEP, + TopLevelSlot.value_or(ByteOffset), + ByteOffset, Slots); + continue; + } + + if (isa(U) || isa(U) || isa(U) || + isa(U)) { + WorkList.push_back(cast(U)); + continue; + } + } + } + + auto Result = collectReadOnlyCaptures(Slots); + + DEBUG(Logger::logs("proteus-pass") + << "[AutoReadOnly] Function " << F.getName() << ": " << Slots.size() + << " slots analyzed, " << Result.size() << " qualified as readonly\n"); + + if (!Result.empty()) { + DEBUG(Logger::logs("proteus-pass") << "[AutoReadOnly] Qualified slots:"); + for (const auto &Cap : Result) { + DEBUG(Logger::logs("proteus-pass") + << " [slot " << Cap.SlotIndex << " @ offset " << Cap.ByteOffset + << ", type " << toString(Cap.RCType) << "]"); + } + DEBUG(Logger::logs("proteus-pass") << "\n"); + } + + return Result; +} + +void emitAutoReadOnlyCapturesMetadata( + Function &F, ArrayRef Captures) { + if (Captures.empty()) { + F.setMetadata(AutoReadOnlyCapturesMetadataName, nullptr); + return; + } + + LLVMContext &Ctx = F.getContext(); + Type *I32Ty = Type::getInt32Ty(Ctx); + + SmallVector Entries; + Entries.reserve(Captures.size()); + for (const auto &Capture : Captures) { + Entries.push_back(MDNode::get( + Ctx, + {ConstantAsMetadata::get(ConstantInt::get(I32Ty, Capture.SlotIndex)), + ConstantAsMetadata::get(ConstantInt::get(I32Ty, Capture.ByteOffset)), + ConstantAsMetadata::get( + ConstantInt::get(I32Ty, static_cast(Capture.RCType)))})); + } + + F.setMetadata(AutoReadOnlyCapturesMetadataName, MDNode::get(Ctx, Entries)); +} + +void annotateAutoReadOnlyCaptures(Module &M) { + for (Function &F : M) { + if (F.isDeclaration()) + continue; + + auto Captures = analyzeAutoReadOnlyCaptures(F); + emitAutoReadOnlyCapturesMetadata(F, Captures); + } +} + +} // namespace proteus diff --git a/src/pass/AutoReadOnlyCapturesAnalysis.h b/src/pass/AutoReadOnlyCapturesAnalysis.h new file mode 100644 index 000000000..e788a5476 --- /dev/null +++ b/src/pass/AutoReadOnlyCapturesAnalysis.h @@ -0,0 +1,29 @@ +#ifndef PROTEUS_PASS_AUTO_READONLY_CAPTURES_ANALYSIS_H +#define PROTEUS_PASS_AUTO_READONLY_CAPTURES_ANALYSIS_H + +#include "proteus/impl/AutoReadOnlyCaptures.h" + +#include +#include + +#include + +namespace llvm { +class Function; +class Module; +} // namespace llvm + +namespace proteus { + +llvm::SmallVector +analyzeAutoReadOnlyCaptures(llvm::Function &F); + +void emitAutoReadOnlyCapturesMetadata( + llvm::Function &F, + llvm::ArrayRef Captures); + +void annotateAutoReadOnlyCaptures(llvm::Module &M); + +} // namespace proteus + +#endif diff --git a/src/pass/CMakeLists.txt b/src/pass/CMakeLists.txt index 58963c81b..452d67042 100644 --- a/src/pass/CMakeLists.txt +++ b/src/pass/CMakeLists.txt @@ -26,6 +26,7 @@ endif() # dependencies for building tests. set(PROTEUS_PASS_SOURCES ${PROJECT_SOURCE_DIR}/src/pass/AnnotationHandler.cpp + ${PROJECT_SOURCE_DIR}/src/pass/AutoReadOnlyCapturesAnalysis.cpp ${PROJECT_SOURCE_DIR}/src/pass/ProteusPass.cpp ${PROJECT_SOURCE_DIR}/src/runtime/Error.cpp ) diff --git a/src/pass/ProteusPass.cpp b/src/pass/ProteusPass.cpp index fe58fb9a1..79eeb8695 100644 --- a/src/pass/ProteusPass.cpp +++ b/src/pass/ProteusPass.cpp @@ -26,6 +26,7 @@ //===----------------------------------------------------------------------===// #include "AnnotationHandler.h" +#include "AutoReadOnlyCapturesAnalysis.h" #include "Helpers.h" #include "proteus/CompilerInterfaceTypes.h" @@ -281,6 +282,8 @@ class ProteusPassImpl { // runtime constants. emitJitFunctionArgMetadata(*JitMod, JFI, *JitF); + annotateAutoReadOnlyCaptures(*JitMod); + if (verifyModule(*JitMod, &errs())) reportFatalError("Broken JIT module found, compilation aborted!"); @@ -313,6 +316,7 @@ class ProteusPassImpl { bool HasSourceFileID) { SmallVector Bitcode; raw_svector_ostream OS(Bitcode); + annotateAutoReadOnlyCaptures(EmbedM); WriteBitcodeToFile(EmbedM, OS); HashT HashValue = hash(StringRef{Bitcode.data(), Bitcode.size()}); @@ -795,7 +799,16 @@ class ProteusPassImpl { ArrayType *ArgPtrsTy = ArrayType::get(Types.PtrTy, StubFn->arg_size()); Value *ArgPtrs = nullptr; - if (NumRuntimeConstants > 0) { + + // Check if this is a lambda function (contains ::operator() in demangled + // name) + std::string DemangledName = llvm::demangle(StubFn->getName().str()); + bool IsLambda = DemangledName.find("::operator()") != std::string::npos; + + // For lambdas, we always need to pass Args to enable auto-readonly capture + // detection, even when there are no explicit jit_variable() captures. + // For non-lambdas, only create Args when there are runtime constants. + if (NumRuntimeConstants > 0 || IsLambda) { ArgPtrs = Builder.CreateAlloca(ArgPtrsTy); // Create an alloca for each argument to store a pointer to the argument, // mimicking how arguments are passed for GPU kernels. This is done so diff --git a/src/runtime/JitEngineHost.cpp b/src/runtime/JitEngineHost.cpp index a1e95b18f..eef341764 100644 --- a/src/runtime/JitEngineHost.cpp +++ b/src/runtime/JitEngineHost.cpp @@ -36,6 +36,8 @@ #include +#include "proteus/impl/AutoReadOnlyCaptures.h" + using namespace proteus; using namespace llvm; using namespace llvm::orc; @@ -133,8 +135,10 @@ JitEngineHost::~JitEngineHost() { CacheChain->printStats(); } -void JitEngineHost::specializeIR(Module &M, StringRef FnName, StringRef Suffix, - ArrayRef RCArray) { +void JitEngineHost::specializeIR( + Module &M, StringRef FnName, StringRef Suffix, + ArrayRef RCArray, + const SmallVector &LambdaJitValuesVec) { TIMESCOPE(JitEngineHost, specializeIR); Function *F = M.getFunction(FnName); assert(F && "Expected non-null function!"); @@ -146,22 +150,13 @@ void JitEngineHost::specializeIR(Module &M, StringRef FnName, StringRef Suffix, PROTEUS_DBG(Logger::logs("proteus") << "Metadata jit for F " << F->getName() << " = " << *Node << "\n"); - // Replace argument uses with runtime constants. - SmallVector ArgPos; - for (unsigned int I = 0; I < Node->getNumOperands(); ++I) { - ConstantAsMetadata *CAM = cast(Node->getOperand(I)); - ConstantInt *ConstInt = cast(CAM->getValue()); - int ArgNo = ConstInt->getZExtValue(); - ArgPos.push_back(ArgNo); - } - TransformArgumentSpecialization::transform(M, *F, RCArray); if (!LambdaRegistry::instance().empty()) { if (auto OptionalMapIt = LambdaRegistry::instance().matchJitVariableMap(F->getName())) { - auto &RCVec = OptionalMapIt.value()->getSecond(); - TransformLambdaSpecialization::transform(M, *F, RCVec); + // Use the merged lambda values from getLambdaJitValues + TransformLambdaSpecialization::transform(M, *F, LambdaJitValuesVec); } } @@ -175,7 +170,7 @@ void JitEngineHost::specializeIR(Module &M, StringRef FnName, StringRef Suffix, } } -void getLambdaJitValues(StringRef FnName, +void getLambdaJitValues(Module &M, StringRef FnName, void **Args, SmallVector &LambdaJitValuesVec) { TIMESCOPE("proteus::getLambdaJitValues"); LambdaRegistry &LR = LambdaRegistry::instance(); @@ -186,14 +181,25 @@ void getLambdaJitValues(StringRef FnName, << "Caller trigger " << FnName << " -> " << demangle(FnName.str()) << "\n"); - SmallVector LambdaCalleeInfo; PROTEUS_DBG(Logger::logs("proteus") << " Trying F " << demangle(FnName.str()) << "\n "); auto OptionalMapIt = LR.matchJitVariableMap(FnName); if (!OptionalMapIt) return; - LambdaJitValuesVec = OptionalMapIt.value()->getSecond(); + // Get the explicit jit_variable captures + const SmallVector &ExplicitValues = + OptionalMapIt.value()->getSecond(); + + Function *LambdaFn = M.getFunction(FnName); + // For host JIT, the lambda closure is passed as the first argument and + // Args[0] contains a pointer-to-pointer to the closure due to the ABI. + const void *LambdaClosure = + (Args && Args[0]) ? *static_cast(Args[0]) : nullptr; + LambdaJitValuesVec = resolveLambdaSpecializationValues( + ExplicitValues, LambdaFn, LambdaClosure, + Config::get().ProteusAutoReadOnlyCaptures, + Config::get().traceSpecializations()); } void * @@ -217,7 +223,7 @@ JitEngineHost::compileAndLink(StringRef FnName, char *IR, int IRSize, SmallVector RCVec = getRuntimeConstantValues(Args, RCInfoArray); SmallVector LambdaJitValuesVec; - getLambdaJitValues(FnName, LambdaJitValuesVec); + getLambdaJitValues(*M, FnName, Args, LambdaJitValuesVec); HashT HashValue = hash(StrIR, FnName, RCVec, LambdaJitValuesVec); if (Config::get().ProteusDebugOutput) { @@ -247,7 +253,7 @@ JitEngineHost::compileAndLink(StringRef FnName, char *IR, int IRSize, } else { PROTEUS_DBG(Logger::logfile(HashValue.toString() + ".input.ll", *M)); // Specialize the module using runtime values. - specializeIR(*M, FnName, Suffix, RCVec); + specializeIR(*M, FnName, Suffix, RCVec, LambdaJitValuesVec); PROTEUS_DBG(Logger::logfile(HashValue.toString() + ".specialized.ll", *M)); // Compile the object. auto ObjectModule = compileOnly(*M); diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 3162c964e..290c8354e 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -26,7 +26,9 @@ function(CREATE_CPU_TEST exe check_source) $ ) - add_test(NAME ${exe} COMMAND ${LIT} -vv -D FILECHECK=${FILECHECK} ${check_source}) + # Provide an explicit %ext substitution so LIT RUN lines like + # `%build/.%ext` resolve to the built executable. + add_test(NAME ${exe} COMMAND ${LIT} -vv -D EXT=${CMAKE_EXECUTABLE_SUFFIX} -D FILECHECK=${FILECHECK} -D EXE=${exe} ${check_source}) set_tests_properties(${exe} PROPERTIES LABELS "cpu") endfunction() @@ -49,7 +51,11 @@ exec_root = tempfile.mkdtemp(prefix='lit.tmp.', dir='${CMAKE_CURRENT_BINARY_DIR} config.test_exec_root = exec_root atexit.register(lambda: shutil.rmtree(exec_root, ignore_errors=False)) +ext = lit_config.params['EXT'] +exe = lit_config.params['EXE'] FILECHECK = lit_config.params['FILECHECK'] +config.substitutions.append(('%ext', ext)) +config.substitutions.append(('%exe', exe)) config.substitutions.append(('%FILECHECK', FILECHECK)) config.substitutions.append(('%target_arch', platform.machine())) config.substitutions.append(('%build', '${CMAKE_CURRENT_BINARY_DIR}')) @@ -66,7 +72,13 @@ CREATE_CPU_TEST(lambda lambda.cpp) CREATE_CPU_TEST(lambda_def lambda_def.cpp) CREATE_CPU_TEST(lambda_multiple lambda_multiple.cpp) CREATE_CPU_TEST(lambda_multiple_api lambda_multiple_api.cpp) +CREATE_CPU_TEST(lambda_written_captures lambda_written_captures.cpp) +CREATE_CPU_TEST(lambda_auto_readonly lambda_auto_readonly.cpp) +CREATE_CPU_TEST(lambda_nested_captures lambda_nested_captures.cpp) +CREATE_CPU_TEST(lambda_mixed_captures lambda_mixed_captures.cpp) +CREATE_CPU_TEST(lambda_field_pointer_escape lambda_field_pointer_escape.cpp) CREATE_CPU_TEST(lambda_spec_test lambda_spec_test.cpp) +CREATE_CPU_TEST(lambda_pointer_captures lambda_pointer_captures.cpp) CREATE_CPU_TEST(types_jit_array types_jit_array.cpp) CREATE_CPU_TEST(dynamic_jit_array dynamic_jit_array.cpp) CREATE_CPU_TEST(jit_struct jit_struct.cpp) diff --git a/tests/cpu/lambda_auto_readonly.cpp b/tests/cpu/lambda_auto_readonly.cpp new file mode 100644 index 000000000..0c69ceeb3 --- /dev/null +++ b/tests/cpu/lambda_auto_readonly.cpp @@ -0,0 +1,42 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/%exe lambda_auto_readonly 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +int main() { + int A = 42; + double B = 3.14; + float C = 2.5f; + bool D = true; + + double X[4] = {0.0, 0.0, 0.0, 0.0}; + + auto lambda = [=, &X]() __attribute__((annotate("jit"))) { + X[0] = A; + X[1] = B; + X[2] = C; + X[3] = D ? 1.0 : 0.0; + }; + + proteus::register_lambda(lambda); + lambda(); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + std::cout << "x[2] = " << X[2] << "\n"; + std::cout << "x[3] = " << X[3] << "\n"; + + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with float 2.5 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i8 1 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 +// CHECK: x[2] = 2.5 +// CHECK: x[3] = 1 diff --git a/tests/cpu/lambda_field_pointer_escape.cpp b/tests/cpu/lambda_field_pointer_escape.cpp new file mode 100644 index 000000000..1c500b242 --- /dev/null +++ b/tests/cpu/lambda_field_pointer_escape.cpp @@ -0,0 +1,38 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/%exe lambda_field_pointer_escape 2>&1 | %FILECHECK %s +// clang-format on + +#include +#include + +#include "proteus/JitInterface.h" + +__attribute__((noinline)) void observe(const int *Ptr) { + if (!Ptr) + std::abort(); +} + +int main() { + int A = 42; + double B = 3.14; + double X[2] = {0.0, 0.0}; + + auto lambda = [&X, A, B]() __attribute__((annotate("jit"))) { + observe(&A); + X[0] = A; + X[1] = B; + }; + + proteus::register_lambda(lambda); + lambda(); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK-NOT: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 diff --git a/tests/cpu/lambda_mixed_captures.cpp b/tests/cpu/lambda_mixed_captures.cpp new file mode 100644 index 000000000..ed39046d6 --- /dev/null +++ b/tests/cpu/lambda_mixed_captures.cpp @@ -0,0 +1,38 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/%exe lambda_mixed_captures 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +int main() { + int A = 10; + int B = 20; + double C = 3.14; + double D = 2.71; + + double X[2] = {0.0, 0.0}; + + auto lambda = + [=, &X, A = proteus::jit_variable(A), C = proteus::jit_variable(C)]() + __attribute__((annotate("jit"))) { + X[0] = A + B; + X[1] = C + D; + }; + + proteus::register_lambda(lambda); + lambda(); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + return 0; +} + +// CHECK-DAG: [LambdaSpec] Replacing slot {{[0-9]+}} with i32 10 +// CHECK-DAG: [LambdaSpec] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 20 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 2.71 +// CHECK: x[0] = 30 +// CHECK: x[1] = 5.85 diff --git a/tests/cpu/lambda_nested_captures.cpp b/tests/cpu/lambda_nested_captures.cpp new file mode 100644 index 000000000..01e61b944 --- /dev/null +++ b/tests/cpu/lambda_nested_captures.cpp @@ -0,0 +1,35 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/%exe lambda_nested_captures 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +struct Payload { + int A; + double B; +}; + +int main() { + Payload P{42, 3.14}; + double X[2] = {0.0, 0.0}; + + auto lambda = [=, &X]() __attribute__((annotate("jit"))) { + X[0] = P.A; + X[1] = P.B; + }; + + proteus::register_lambda(lambda); + lambda(); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 diff --git a/tests/cpu/lambda_pointer_captures.cpp b/tests/cpu/lambda_pointer_captures.cpp new file mode 100644 index 000000000..f94a706c7 --- /dev/null +++ b/tests/cpu/lambda_pointer_captures.cpp @@ -0,0 +1,40 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/%exe lambda_pointer_captures 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +int main() { + int Scalar = 42; + int *Ptr = &Scalar; + double Value = 3.14; + + double X[2] = {0.0, 0.0}; + + auto lambda = [=, &X]() __attribute__((annotate("jit"))) { + X[0] = Scalar; + X[1] = Value; + (void)Ptr; + }; + + int *PtrOnly = &Scalar; + auto lambda2 = [=]() __attribute__((annotate("jit"))) { (void)PtrOnly; }; + + proteus::register_lambda(lambda); + proteus::register_lambda(lambda2); + + lambda(); + lambda2(); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 diff --git a/tests/cpu/lambda_written_captures.cpp b/tests/cpu/lambda_written_captures.cpp new file mode 100644 index 000000000..a40f41ac1 --- /dev/null +++ b/tests/cpu/lambda_written_captures.cpp @@ -0,0 +1,35 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/%exe lambda_written_captures 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +int main() { + int A = 10; // Read-only - should be auto-detected + int B = 20; // Written in lambda - should NOT be auto-detected + int C = 30; // Read-only - should be auto-detected + + double X[2] = {0.0, 0.0}; + + auto lambda = [=, &X]() __attribute__((annotate("jit"))) mutable { + B = B + 1; + X[0] = A + B; + X[1] = C; + }; + + proteus::register_lambda(lambda); + lambda(); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 10 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 30 +// CHECK-NOT: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 20 +// CHECK: x[0] = 31 +// CHECK: x[1] = 30 diff --git a/tests/gpu/CMakeLists.txt b/tests/gpu/CMakeLists.txt index b150e4d94..b1e01407a 100644 --- a/tests/gpu/CMakeLists.txt +++ b/tests/gpu/CMakeLists.txt @@ -264,6 +264,13 @@ CREATE_GPU_TEST(shared_array shared_array.cpp) CREATE_GPU_TEST(enable_disable enable_disable.cpp) CREATE_GPU_TEST(lambda_multiple lambda_multiple.cpp) CREATE_GPU_TEST(lambda_def lambda_def.cpp) +CREATE_GPU_TEST(lambda_written_captures lambda_written_captures.cpp) +CREATE_GPU_TEST(lambda_auto_readonly lambda_auto_readonly.cpp) +CREATE_GPU_TEST(lambda_auto_readonly_second_arg lambda_auto_readonly_second_arg.cpp) +CREATE_GPU_TEST(lambda_nested_captures lambda_nested_captures.cpp) +CREATE_GPU_TEST(lambda_mixed_captures lambda_mixed_captures.cpp) +CREATE_GPU_TEST(lambda_field_pointer_escape lambda_field_pointer_escape.cpp) +CREATE_GPU_TEST(lambda_pointer_captures lambda_pointer_captures.cpp) CREATE_GPU_TEST(lambda_host_device lambda_host_device.cpp) CREATE_GPU_TEST(lambda_spec_test lambda_spec_test.cpp) CREATE_GPU_TEST(builtin_globals builtin_globals.cpp) diff --git a/tests/gpu/lambda_auto_readonly.cpp b/tests/gpu/lambda_auto_readonly.cpp new file mode 100644 index 000000000..1f5452aac --- /dev/null +++ b/tests/gpu/lambda_auto_readonly.cpp @@ -0,0 +1,54 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/lambda_auto_readonly.%ext 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +#include "gpu_common.h" + +template +__global__ __attribute__((annotate("jit"))) void kernel(T LB) { + std::size_t I = blockIdx.x + threadIdx.x; + if (I == 0) + LB(); +} + +int main() { + int A = 42; + double B = 3.14; + float C = 2.5f; + bool D = true; + + double *X; + gpuErrCheck(gpuMallocManaged(&X, sizeof(double) * 4)); + + auto lambda = [=] __device__ __attribute__((annotate("jit"))) () { + X[0] = A; + X[1] = B; + X[2] = C; + X[3] = D ? 1.0 : 0.0; + }; + + proteus::register_lambda(lambda); + kernel<<<1, 1>>>(lambda); + gpuErrCheck(gpuDeviceSynchronize()); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + std::cout << "x[2] = " << X[2] << "\n"; + std::cout << "x[3] = " << X[3] << "\n"; + + gpuErrCheck(gpuFree(X)); + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with float 2.5 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i8 1 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 +// CHECK: x[2] = 2.5 +// CHECK: x[3] = 1 diff --git a/tests/gpu/lambda_auto_readonly_second_arg.cpp b/tests/gpu/lambda_auto_readonly_second_arg.cpp new file mode 100644 index 000000000..1b173c0ab --- /dev/null +++ b/tests/gpu/lambda_auto_readonly_second_arg.cpp @@ -0,0 +1,79 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/lambda_auto_readonly_second_arg.%ext 2>&1 | %FILECHECK %s +// clang-format on + +#include +#include + +#include "proteus/JitInterface.h" + +#include "gpu_common.h" + +struct PrefixArgs { + double *PadPtr; + int A; + double B; + float C; + bool D; +}; + +static_assert(offsetof(PrefixArgs, A) == 8); +static_assert(offsetof(PrefixArgs, B) == 16); +static_assert(offsetof(PrefixArgs, C) == 24); +static_assert(offsetof(PrefixArgs, D) == 28); + +template +__global__ __attribute__((annotate("jit"))) void kernel(PrefixArgs Prefix, + T LB) { + std::size_t I = blockIdx.x + threadIdx.x; + if (I == 0) + LB(); + if (Prefix.PadPtr == nullptr) + return; +} + +int main() { + int A = 42; + double B = 3.14; + float C = 2.5f; + bool D = true; + + PrefixArgs Prefix{ + /*PadPtr=*/nullptr, + /*A=*/7, + /*B=*/1.25, + /*C=*/9.5f, + /*D=*/false, + }; + + double *X; + gpuErrCheck(gpuMallocManaged(&X, sizeof(double) * 4)); + + auto lambda = [=] __device__ __attribute__((annotate("jit"))) () { + X[0] = A; + X[1] = B; + X[2] = C; + X[3] = D ? 1.0 : 0.0; + }; + + proteus::register_lambda(lambda); + kernel<<<1, 1>>>(Prefix, lambda); + gpuErrCheck(gpuDeviceSynchronize()); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + std::cout << "x[2] = " << X[2] << "\n"; + std::cout << "x[3] = " << X[3] << "\n"; + + gpuErrCheck(gpuFree(X)); + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with float 2.5 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i8 1 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 +// CHECK: x[2] = 2.5 +// CHECK: x[3] = 1 diff --git a/tests/gpu/lambda_field_pointer_escape.cpp b/tests/gpu/lambda_field_pointer_escape.cpp new file mode 100644 index 000000000..d691913d3 --- /dev/null +++ b/tests/gpu/lambda_field_pointer_escape.cpp @@ -0,0 +1,50 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/lambda_field_pointer_escape.%ext 2>&1 | %FILECHECK %s +// clang-format on + +#include +#include + +#include "gpu_common.h" +#include "proteus/JitInterface.h" + +__device__ __attribute__((noinline)) void observe(const int *Ptr) { + if (Ptr == nullptr) + printf("null\n"); +} + +template +__global__ __attribute__((annotate("jit"))) void kernel(T LB) { + std::size_t I = blockIdx.x + threadIdx.x; + if (I == 0) + LB(); +} + +int main() { + int A = 42; + double B = 3.14; + + double *X; + gpuErrCheck(gpuMallocManaged(&X, sizeof(double) * 2)); + + auto lambda = [X, A, B] __device__ __attribute__((annotate("jit"))) { + observe(&A); + X[0] = A; + X[1] = B; + }; + + proteus::register_lambda(lambda); + kernel<<<1, 1>>>(lambda); + gpuErrCheck(gpuDeviceSynchronize()); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + gpuErrCheck(gpuFree(X)); + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK-NOT: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 diff --git a/tests/gpu/lambda_mixed_captures.cpp b/tests/gpu/lambda_mixed_captures.cpp new file mode 100644 index 000000000..f675172a1 --- /dev/null +++ b/tests/gpu/lambda_mixed_captures.cpp @@ -0,0 +1,50 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/lambda_mixed_captures.%ext 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +#include "gpu_common.h" + +template +__global__ __attribute__((annotate("jit"))) void kernel(T LB) { + std::size_t I = blockIdx.x + threadIdx.x; + if (I == 0) + LB(); +} + +int main() { + int A = 10; + int B = 20; + double C = 3.14; + double D = 2.71; + + double *X; + gpuErrCheck(gpuMallocManaged(&X, sizeof(double) * 2)); + + auto lambda = + [=, A = proteus::jit_variable(A), C = proteus::jit_variable(C)] __device__ + __attribute__((annotate("jit"))) () { + X[0] = A + B; + X[1] = C + D; + }; + + proteus::register_lambda(lambda); + kernel<<<1, 1>>>(lambda); + gpuErrCheck(gpuDeviceSynchronize()); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + gpuErrCheck(gpuFree(X)); + return 0; +} + +// CHECK-DAG: [LambdaSpec] Replacing slot {{[0-9]+}} with i32 10 +// CHECK-DAG: [LambdaSpec] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 20 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 2.71 +// CHECK: x[0] = 30 +// CHECK: x[1] = 5.85 diff --git a/tests/gpu/lambda_nested_captures.cpp b/tests/gpu/lambda_nested_captures.cpp new file mode 100644 index 000000000..489e97a28 --- /dev/null +++ b/tests/gpu/lambda_nested_captures.cpp @@ -0,0 +1,48 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/lambda_nested_captures.%ext 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +#include "gpu_common.h" + +struct Payload { + int A; + double B; +}; + +template +__global__ __attribute__((annotate("jit"))) void kernel(T LB) { + std::size_t I = blockIdx.x + threadIdx.x; + if (I == 0) + LB(); +} + +int main() { + Payload P{42, 3.14}; + + double *X; + gpuErrCheck(gpuMallocManaged(&X, sizeof(double) * 2)); + + auto lambda = [=] __device__ __attribute__((annotate("jit"))) () { + X[0] = P.A; + X[1] = P.B; + }; + + proteus::register_lambda(lambda); + kernel<<<1, 1>>>(lambda); + gpuErrCheck(gpuDeviceSynchronize()); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + gpuErrCheck(gpuFree(X)); + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 diff --git a/tests/gpu/lambda_pointer_captures.cpp b/tests/gpu/lambda_pointer_captures.cpp new file mode 100644 index 000000000..940c87c02 --- /dev/null +++ b/tests/gpu/lambda_pointer_captures.cpp @@ -0,0 +1,53 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT=specialization %build/lambda_pointer_captures.%ext 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "proteus/JitInterface.h" + +#include "gpu_common.h" + +template +__global__ __attribute__((annotate("jit"))) void kernel(T LB) { + std::size_t I = blockIdx.x + threadIdx.x; + if (I == 0) + LB(); +} + +int main() { + int Scalar = 42; + int *Ptr = &Scalar; + double Value = 3.14; + + double *X; + gpuErrCheck(gpuMallocManaged(&X, sizeof(double) * 2)); + + auto lambda = [=] __device__ __attribute__((annotate("jit"))) () { + X[0] = Scalar; + X[1] = Value; + (void)Ptr; + }; + + int *PtrOnly = &Scalar; + auto lambda2 = [=] __device__ + __attribute__((annotate("jit"))) () { (void)PtrOnly; }; + + proteus::register_lambda(lambda); + proteus::register_lambda(lambda2); + + kernel<<<1, 1>>>(lambda); + kernel<<<1, 1>>>(lambda2); + gpuErrCheck(gpuDeviceSynchronize()); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + gpuErrCheck(gpuFree(X)); + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 42 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with double 3.14 +// CHECK: x[0] = 42 +// CHECK: x[1] = 3.14 diff --git a/tests/gpu/lambda_written_captures.cpp b/tests/gpu/lambda_written_captures.cpp new file mode 100644 index 000000000..11c7f3422 --- /dev/null +++ b/tests/gpu/lambda_written_captures.cpp @@ -0,0 +1,46 @@ +// clang-format off +// RUN: PROTEUS_AUTO_READONLY_CAPTURES=1 PROTEUS_TRACE_OUTPUT="specialization" %build/lambda_written_captures.%ext 2>&1 | %FILECHECK %s +// clang-format on + +#include + +#include "gpu_common.h" +#include "proteus/JitInterface.h" + +template +__global__ __attribute__((annotate("jit"))) void kernel(T LB) { + std::size_t I = blockIdx.x + threadIdx.x; + if (I == 0) + LB(); +} + +int main() { + int A = 10; // Read-only - should be auto-detected + int B = 20; // Written in lambda - should NOT be auto-detected + int C = 30; // Read-only - should be auto-detected + + double *X; + gpuErrCheck(gpuMallocManaged(&X, sizeof(double) * 2)); + + auto lambda = [=] __device__ __attribute__((annotate("jit"))) () mutable { + B = B + 1; + X[0] = A + B; + X[1] = C; + }; + + proteus::register_lambda(lambda); + kernel<<<1, 1>>>(lambda); + gpuErrCheck(gpuDeviceSynchronize()); + + std::cout << "x[0] = " << X[0] << "\n"; + std::cout << "x[1] = " << X[1] << "\n"; + + gpuErrCheck(gpuFree(X)); + return 0; +} + +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 10 +// CHECK-DAG: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 30 +// CHECK-NOT: [LambdaSpec][Auto] Replacing slot {{[0-9]+}} with i32 20 +// CHECK: x[0] = 31 +// CHECK: x[1] = 30