diff --git a/include/proteus/JitInterface.h b/include/proteus/JitInterface.h index 9c23b5eb..059f148f 100644 --- a/include/proteus/JitInterface.h +++ b/include/proteus/JitInterface.h @@ -65,9 +65,9 @@ jit_object(T *V, size_t Size = sizeof(std::remove_pointer_t)) noexcept; #if defined(__CUDACC__) || defined(__HIP__) template -__attribute__((noinline)) __device__ std::enable_if_t< - std::is_trivially_copyable_v>, void> -jit_object(T *V, size_t Size = sizeof(T)) noexcept; +__attribute__((noinline)) __device__ + std::enable_if_t>, + void> jit_object(T *V, size_t Size = sizeof(T)) noexcept; #endif template @@ -82,8 +82,7 @@ template __attribute__((noinline)) __device__ std::enable_if_t< !std::is_pointer_v && std::is_trivially_copyable_v>, - void> -jit_object(T &V, size_t Size = sizeof(T)) noexcept; + void> jit_object(T &V, size_t Size = sizeof(T)) noexcept; #endif namespace detail { diff --git a/src/include/proteus/impl/CompilerInterfaceDevice.h b/src/include/proteus/impl/CompilerInterfaceDevice.h index 88613475..4141f4f7 100644 --- a/src/include/proteus/impl/CompilerInterfaceDevice.h +++ b/src/include/proteus/impl/CompilerInterfaceDevice.h @@ -32,4 +32,9 @@ extern "C" proteus::DeviceTraits::DeviceError_t __proteus_launch_kernel(void *Kernel, dim3 GridDim, dim3 BlockDim, void **KernelArgs, uint64_t ShmemSize, void *Stream); +extern "C" proteus::DeviceTraits::DeviceError_t +__proteus_launch_kernel_by_name(const char *KernelLookupKey, dim3 GridDim, + dim3 BlockDim, void **KernelArgs, + uint64_t ShmemSize, void *Stream); + #endif diff --git a/src/include/proteus/impl/JitEngineDevice.h b/src/include/proteus/impl/JitEngineDevice.h index f3ccbf63..0ac73232 100644 --- a/src/include/proteus/impl/JitEngineDevice.h +++ b/src/include/proteus/impl/JitEngineDevice.h @@ -477,6 +477,7 @@ template class JitEngineDevice : public JitEngine { const char *ModuleId); void finalizeRegistration(); void registerFunction(void *Handle, void *Kernel, char *KernelName, + const char *KernelLookupKey, ArrayRef RCInfoArray); void registerLambdaCallsiteLocation(void *Kernel, uint64_t LambdaID, uint32_t CallsiteIndex, @@ -499,6 +500,15 @@ template class JitEngineDevice : public JitEngine { return JITKernelInfoMap[Func]; } + std::optional> + getJITKernelInfo(StringRef KernelLookupKey) { + auto It = KernelLookupKeyToKernel.find(KernelLookupKey.str()); + if (It == KernelLookupKeyToKernel.end()) + return std::nullopt; + + return getJITKernelInfo(It->second); + } + HashT getStaticHash(JITKernelInfo &KernelInfo) { if (KernelInfo.hasStaticHash()) return KernelInfo.getStaticHash(); @@ -536,7 +546,7 @@ template class JitEngineDevice : public JitEngine { for (auto &Func : FatbinInfo.Functions) registerFunction(Handle, Func.Kernel, Func.KernelName, - Func.RCInfoArray); + Func.KernelLookupKey, Func.RCInfoArray); for (auto &Var : FatbinInfo.Vars) registerVar(Var.Handle, Var.VarName, Var.HostAddr, Var.VarSize); @@ -571,6 +581,7 @@ template class JitEngineDevice : public JitEngine { std::string DeviceArch; DenseMap JITKernelInfoMap; + std::unordered_map KernelLookupKeyToKernel; DenseMap PendingLambdaCallsiteLocationInfo; std::unique_ptr AsyncCompiler; @@ -739,10 +750,12 @@ template void JitEngineDevice::finalizeRegistration() { template void JitEngineDevice::registerFunction( - void *Handle, void *Kernel, char *KernelName, + void *Handle, void *Kernel, char *KernelName, const char *KernelLookupKey, ArrayRef RCInfoArray) { PROTEUS_DBG(Logger::logs("proteus") << "Register function " << Kernel << " To Handle " << Handle << "\n"); + KernelLookupKeyToKernel.try_emplace(KernelLookupKey, Kernel); + // NOTE: HIP RDC might call multiple times the registerFunction for the same // kernel, which has weak linkage, when it comes from different translation // units. Either the first or the second call can prevail and should be diff --git a/src/include/proteus/impl/JitEngineInfoRegistry.h b/src/include/proteus/impl/JitEngineInfoRegistry.h index 3d87e32a..f70ed8a1 100644 --- a/src/include/proteus/impl/JitEngineInfoRegistry.h +++ b/src/include/proteus/impl/JitEngineInfoRegistry.h @@ -37,6 +37,7 @@ struct RegisterFunctionInfo { void *Handle; void *Kernel; char *KernelName; + const char *KernelLookupKey; ArrayRef RCInfoArray; }; @@ -75,9 +76,11 @@ class JitEngineInfoRegistry { } void registerFunction(void *Handle, void *Kernel, char *KernelName, + const char *KernelLookupKey, ArrayRef RCInfoArray) { auto &FatbinInfo = FatbinaryMap.at(Handle); - FatbinInfo.Functions.push_back({Handle, Kernel, KernelName, RCInfoArray}); + FatbinInfo.Functions.push_back( + {Handle, Kernel, KernelName, KernelLookupKey, RCInfoArray}); } void registerVar(void *Handle, const void *HostAddr, const char *VarName, diff --git a/src/pass/ProteusPass.cpp b/src/pass/ProteusPass.cpp index 529c1cb3..edf6f52a 100644 --- a/src/pass/ProteusPass.cpp +++ b/src/pass/ProteusPass.cpp @@ -58,6 +58,7 @@ #include #include #include +#include #include #include #include @@ -205,7 +206,7 @@ class ProteusPassImpl { if (hasDeviceLaunchKernelCalls(M)) { instrumentLambdaLaunchCallsites(M, StubToKernelMap); - emitJitLaunchKernelCall(M); + emitJitLaunchKernelCall(M, StubToKernelMap); } instrumentRegisterFunction(M); @@ -228,15 +229,17 @@ class ProteusPassImpl { JitWorkList.push_back(&JFI); } - // IMPORTANT: Build all per-function JIT modules before rewriting any of the - // original functions into stubs. Otherwise, later JIT module extraction can - // accidentally clone the already-rewritten stub (and its mutable globals), - // producing invalid JIT IR (e.g. external globals with InternalLinkage). - // See unit test lambda_def_register_once, which tests this ordering. - for (auto *JFI : JitWorkList) + // IMPORTANT: A JIT function is rewritten into its dispatch stub before the + // module of any JIT function that contains it is extracted, so the + // enclosing module carries the nested dispatch and the inner region + // specializes on its own runtime constants. Ordering innermost-first is + // what makes that happen; the stub's mutable bookkeeping globals are + // cloned by definition (see emitJitModuleHost) so the cloned IR stays + // valid. See unit tests lambda_def_register_once and lambda_nested. + for (auto *JFI : sortJitWorkListInnermostFirst(JitWorkList)) { emitJitModuleHost(M, *JFI); - for (auto *JFI : JitWorkList) emitJitEntryCall(M, *JFI); + } DEBUG(Logger::logs("proteus-pass") << "=== Post Original Host Module\n" @@ -776,6 +779,81 @@ class ProteusPassImpl { } } + using JitWorkListEntry = decltype(JitFunctionInfoMap)::value_type; + + // JIT functions reachable from F's body, looking through ordinary calls but + // stopping at another JIT function -- those are the regions nested in F. + static SmallPtrSet + findNestedJitFunctions(Function &F, + const SmallPtrSetImpl &JitFunctions) { + SmallPtrSet Nested; + SmallPtrSet Visited; + SmallVector Worklist{&F}; + + while (!Worklist.empty()) { + Function *Current = Worklist.pop_back_val(); + if (!Visited.insert(Current).second) + continue; + + for (Instruction &I : instructions(*Current)) { + auto *CB = dyn_cast(&I); + if (!CB) + continue; + + Function *Callee = CB->getCalledFunction(); + if (!Callee || Callee->isDeclaration() || Callee == &F) + continue; + + if (JitFunctions.contains(Callee)) { + Nested.insert(Callee); + continue; + } + + Worklist.push_back(Callee); + } + } + + return Nested; + } + + // Post-order over the nesting relation, so an inner JIT function is always + // processed before the ones containing it. Recursion through JIT functions + // has no innermost region, so a cycle keeps its original relative order. + SmallVector sortJitWorkListInnermostFirst( + const SmallVectorImpl &JitWorkList) { + SmallPtrSet JitFunctions; + DenseMap FnToEntry; + for (auto *JFI : JitWorkList) { + JitFunctions.insert(JFI->first); + FnToEntry[JFI->first] = JFI; + } + + DenseMap> Nested; + for (auto *JFI : JitWorkList) + Nested[JFI->first] = findNestedJitFunctions(*JFI->first, JitFunctions); + + SmallVector Ordered; + SmallPtrSet Done; + SmallPtrSet OnStack; + + std::function Visit = [&](Function *F) { + if (Done.contains(F) || !OnStack.insert(F).second) + return; + + for (Function *Inner : Nested[F]) + Visit(Inner); + + OnStack.erase(F); + if (Done.insert(F).second) + Ordered.push_back(FnToEntry[F]); + }; + + for (auto *JFI : JitWorkList) + Visit(JFI->first); + + return Ordered; + } + void emitJitModuleHost(Module &M, std::pair &JITInfo) { Function *JITFn = JITInfo.first; @@ -786,9 +864,17 @@ class ProteusPassImpl { if (isCoverageGlobal(*GV)) return true; - if (const GlobalVariable *G = dyn_cast(GV)) + if (const GlobalVariable *G = dyn_cast(GV)) { + // Bookkeeping globals of a nested dispatch stub are per-callsite + // state, so the enclosing JIT module gets its own definitions. They + // are mutable and internal, so cloning them as declarations would + // produce invalid IR. + if (G->getName().starts_with(".proteus.")) + return true; + if (!G->isConstant()) return false; + } return true; }); @@ -1499,7 +1585,7 @@ class ProteusPassImpl { return true; } - FunctionCallee getJitLaunchKernelFn(Module &M) { + FunctionCallee getJitLaunchKernelFn(Module &M, bool LookupByName) { FunctionType *JitLaunchKernelFnTy = nullptr; assert(LaunchFunctionName && "Expected valid launch function name"); @@ -1517,14 +1603,22 @@ class ProteusPassImpl { "PROTEUS_ENABLE_CUDA|PROTEUS_ENABLE_HIP compilation flags " "for ProteusPass"); + StringRef EntryName = LookupByName ? "__proteus_launch_kernel_by_name" + : "__proteus_launch_kernel"; FunctionCallee JitLaunchKernelFn = - M.getOrInsertFunction("__proteus_launch_kernel", JitLaunchKernelFnTy); + M.getOrInsertFunction(EntryName, JitLaunchKernelFnTy); return JitLaunchKernelFn; } - void replaceWithJitLaunchKernel(Module &M, CallBase *LaunchKernelCB) { - FunctionCallee JitLaunchKernelFn = getJitLaunchKernelFn(M); + std::string getKernelLookupKey(Module &M, const Function &KernelStub) { + return getUniqueFileID(M) + ":" + KernelStub.getName().str(); + } + + void replaceWithJitLaunchKernel(Module &M, CallBase *LaunchKernelCB, + Function *KernelStub) { + FunctionCallee JitLaunchKernelFn = + getJitLaunchKernelFn(M, KernelStub != nullptr); // Insert before the launch kernel call instruction. IRBuilder<> Builder(LaunchKernelCB); @@ -1532,6 +1626,9 @@ class ProteusPassImpl { SmallVector Args = {LaunchKernelCB->arg_begin(), LaunchKernelCB->arg_end()}; + if (KernelStub) + Args[0] = Builder.CreateGlobalString(getKernelLookupKey(M, *KernelStub), + ".proteus.kernel.lookup"); if (isa(LaunchKernelCB)) { CallOrInvoke = Builder.CreateCall(JitLaunchKernelFn, Args); @@ -1551,7 +1648,8 @@ class ProteusPassImpl { LaunchKernelCB->eraseFromParent(); } - void emitJitLaunchKernelCall(Module &M) { + void emitJitLaunchKernelCall( + Module &M, const DenseMap &StubToKernelMap) { Function *LaunchKernelFn = nullptr; if (!LaunchFunctionName) { reportFatalError( @@ -1580,8 +1678,17 @@ class ProteusPassImpl { ToBeReplaced.push_back(CB); } - for (CallBase *CB : ToBeReplaced) - replaceWithJitLaunchKernel(M, CB); + for (CallBase *CB : ToBeReplaced) { + Function *KernelStub = nullptr; + Value *Stub = getStubGV(CB->getArgOperand(0)); + auto *StubFn = dyn_cast_or_null(Stub); + auto It = StubToKernelMap.find(Stub); + if (StubFn && It != StubToKernelMap.end() && + JitFunctionInfoMap.contains(StubFn)) + KernelStub = StubFn; + + replaceWithJitLaunchKernel(M, CB, KernelStub); + } } FunctionCallee getJitRegisterFatBinaryFn(Module &M) { @@ -1747,12 +1854,14 @@ class ProteusPassImpl { // __proteus_register_function(void *Handle, // void *Kernel, // char const *KernelName, + // char const *KernelLookupKey, // RuntimeConstantInfo **RCInfoArrayPtr, // int32_t NumRCs) - FunctionType *JitRegisterFunctionFnTy = FunctionType::get( - Types.VoidTy, - {Types.PtrTy, Types.PtrTy, Types.PtrTy, Types.PtrTy, Types.Int32Ty}, - /* isVarArg=*/false); + FunctionType *JitRegisterFunctionFnTy = + FunctionType::get(Types.VoidTy, + {Types.PtrTy, Types.PtrTy, Types.PtrTy, Types.PtrTy, + Types.PtrTy, Types.Int32Ty}, + /* isVarArg=*/false); FunctionCallee JitRegisterKernelFn = M.getOrInsertFunction( "__proteus_register_function", JitRegisterFunctionFnTy); @@ -1823,11 +1932,13 @@ class ProteusPassImpl { ConstantInt::get(Builder.getInt32Ty(), NumRuntimeConstants); FunctionCallee JitRegisterFunction = getJitRegisterFunctionFn(M); + auto *KernelLookupKey = Builder.CreateGlobalString( + getKernelLookupKey(M, *FunctionToRegister), ".proteus.kernel.lookup"); Builder.CreateCall(JitRegisterFunction, {RegisterCB->getArgOperand(0), RegisterCB->getArgOperand(1), - RegisterCB->getArgOperand(2), + RegisterCB->getArgOperand(2), KernelLookupKey, RuntimeConstantInfoPtrArray, NumRCsValue}); auto HelperIt = diff --git a/src/runtime/CompilerInterfaceDevice.cpp b/src/runtime/CompilerInterfaceDevice.cpp index 5b124583..75a684a9 100644 --- a/src/runtime/CompilerInterfaceDevice.cpp +++ b/src/runtime/CompilerInterfaceDevice.cpp @@ -58,14 +58,14 @@ __proteus_register_linked_binary(void *FatbinWrapper, const char *ModuleId) { JitEngineInfo.registerLinkedBinary(FatbinWrapper, ModuleId); } -extern "C" __attribute((used)) void -__proteus_register_function(void *Handle, void *Kernel, char *KernelName, - RuntimeConstantInfo **RCInfoArrayPtr, - int32_t NumRCs) { +extern "C" __attribute((used)) void __proteus_register_function( + void *Handle, void *Kernel, char *KernelName, const char *KernelLookupKey, + RuntimeConstantInfo **RCInfoArrayPtr, int32_t NumRCs) { ArrayRef RCInfoArray{RCInfoArrayPtr, static_cast(NumRCs)}; auto &JitEngineInfo = JitEngineInfoRegistry::instance(); - JitEngineInfo.registerFunction(Handle, Kernel, KernelName, RCInfoArray); + JitEngineInfo.registerFunction(Handle, Kernel, KernelName, KernelLookupKey, + RCInfoArray); } extern "C" __attribute__((used)) void @@ -146,6 +146,23 @@ __proteus_launch_kernel(void *Kernel, dim3 GridDim, dim3 BlockDim, ShmemSize, Stream); } +extern "C" proteus::DeviceTraits::DeviceError_t +__proteus_launch_kernel_by_name(const char *KernelLookupKey, dim3 GridDim, + dim3 BlockDim, void **KernelArgs, + uint64_t ShmemSize, void *Stream) { + TIMESCOPE("__proteus_launch_kernel_by_name"); + auto &Jit = JitDeviceImplT::instance(); + auto OptionalKernelInfo = Jit.getJITKernelInfo(StringRef{KernelLookupKey}); + if (!OptionalKernelInfo) + reportFatalError("Missing registered GPU kernel " + Twine(KernelLookupKey)); + + void *Kernel = OptionalKernelInfo->get().getKernel(); + auto &LR = LambdaRegistry::instance(); + LR.invokeRegisterLambdaConstants(Kernel, KernelArgs); + return __proteus_launch_kernel_internal(Kernel, GridDim, BlockDim, KernelArgs, + ShmemSize, Stream); +} + extern "C" void __proteus_enable_device() { auto &Jit = JitDeviceImplT::instance(); Jit.enable(); diff --git a/src/runtime/HIPRuntimeAPI.cpp b/src/runtime/HIPRuntimeAPI.cpp index 82a26411..6dfdad31 100644 --- a/src/runtime/HIPRuntimeAPI.cpp +++ b/src/runtime/HIPRuntimeAPI.cpp @@ -104,8 +104,31 @@ template Fn resolveHIPRTCSymbol(const char *Name) { return resolveSymbol(getHIPRTCHandle(), Name, "libhiprtc"); } +hipError_t pushCallConfiguration(dim3 GridDim, dim3 BlockDim, size_t SharedMem, + hipStream_t Stream) { + using Fn = decltype(&::__hipPushCallConfiguration); + static Fn Func = resolveHIPRuntimeSymbol("__hipPushCallConfiguration"); + return Func(GridDim, BlockDim, SharedMem, Stream); +} + +hipError_t popCallConfiguration(dim3 *GridDim, dim3 *BlockDim, + size_t *SharedMem, hipStream_t *Stream) { + using Fn = decltype(&::__hipPopCallConfiguration); + static Fn Func = resolveHIPRuntimeSymbol("__hipPopCallConfiguration"); + return Func(GridDim, BlockDim, SharedMem, Stream); +} + } // namespace +extern "C" void __proteus_get_device_launch_config_symbols( + const char **PushName, uintptr_t *PushAddress, const char **PopName, + uintptr_t *PopAddress) { + *PushName = "__hipPushCallConfiguration"; + *PushAddress = reinterpret_cast(&pushCallConfiguration); + *PopName = "__hipPopCallConfiguration"; + *PopAddress = reinterpret_cast(&popCallConfiguration); +} + namespace proteus::hipdyn { const char *getErrorString(hipError_t Error) { diff --git a/src/runtime/JitEngineHost.cpp b/src/runtime/JitEngineHost.cpp index a633b0ce..0fb78e69 100644 --- a/src/runtime/JitEngineHost.cpp +++ b/src/runtime/JitEngineHost.cpp @@ -42,6 +42,17 @@ using namespace proteus; using namespace llvm; using namespace llvm::orc; +#if PROTEUS_ENABLE_CUDA +extern "C" LLVM_ATTRIBUTE_WEAK void +__proteus_get_device_launch_config_symbols(const char **, uintptr_t *, + const char **, uintptr_t *); +#elif PROTEUS_ENABLE_HIP +extern "C" void __proteus_get_device_launch_config_symbols(const char **, + uintptr_t *, + const char **, + uintptr_t *); +#endif + namespace { DenseMap @@ -82,7 +93,26 @@ void JitEngineHost::addStaticLibrarySymbols() { __proteus_cudaLaunchKernel_ptr)}, JITSymbolFlags::Exported); } +#endif +#if PROTEUS_ENABLE_CUDA + if (__proteus_get_device_launch_config_symbols) { +#elif PROTEUS_ENABLE_HIP + { +#endif +#if PROTEUS_ENABLE_CUDA || PROTEUS_ENABLE_HIP + const char *PushName = nullptr; + uintptr_t PushCallConfiguration = 0; + const char *PopName = nullptr; + uintptr_t PopCallConfiguration = 0; + __proteus_get_device_launch_config_symbols( + &PushName, &PushCallConfiguration, &PopName, &PopCallConfiguration); + + SymbolMap[LLJITPtr->mangleAndIntern(PushName)] = orc::ExecutorSymbolDef( + orc::ExecutorAddr{PushCallConfiguration}, JITSymbolFlags::Exported); + SymbolMap[LLJITPtr->mangleAndIntern(PopName)] = orc::ExecutorSymbolDef( + orc::ExecutorAddr{PopCallConfiguration}, JITSymbolFlags::Exported); + } #endif #if PROTEUS_ENABLE_CUDA || PROTEUS_ENABLE_HIP @@ -91,6 +121,10 @@ void JitEngineHost::addStaticLibrarySymbols() { orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( __proteus_launch_kernel)}, JITSymbolFlags::Exported); + SymbolMap[LLJITPtr->mangleAndIntern("__proteus_launch_kernel_by_name")] = + orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( + __proteus_launch_kernel_by_name)}, + JITSymbolFlags::Exported); #endif // Register the symbol in the main JIT dynamic library. diff --git a/src/runtime/ProteusCUDARuntimeBuiltins.cpp b/src/runtime/ProteusCUDARuntimeBuiltins.cpp index c9eacf6b..e0e617cf 100644 --- a/src/runtime/ProteusCUDARuntimeBuiltins.cpp +++ b/src/runtime/ProteusCUDARuntimeBuiltins.cpp @@ -6,6 +6,12 @@ // NOLINTBEGIN(readability-identifier-naming) +extern "C" unsigned __cudaPushCallConfiguration(dim3 GridDim, dim3 BlockDim, + size_t SharedMem, void *Stream); +extern "C" cudaError_t __cudaPopCallConfiguration(dim3 *GridDim, dim3 *BlockDim, + size_t *SharedMem, + void *Stream); + // Resolve at runtime CUDA runtime symbols to avoid a dependency on the CUDA // runtime library for the proteus runtime library, and allow users to link with // either the static or dynamic CUDA runtime library. @@ -28,6 +34,15 @@ extern cudaError_t (*__proteus_cudaLaunchKernel_ptr)(const void *, dim3, dim3, cudaStream_t); } +extern "C" void __proteus_get_device_launch_config_symbols( + const char **PushName, uintptr_t *PushAddress, const char **PopName, + uintptr_t *PopAddress) { + *PushName = "__cudaPushCallConfiguration"; + *PushAddress = reinterpret_cast(&__cudaPushCallConfiguration); + *PopName = "__cudaPopCallConfiguration"; + *PopAddress = reinterpret_cast(&__cudaPopCallConfiguration); +} + // Initialization function to set the function pointers for the CUDA runtime // symbols used by the Proteus runtime library. This function is called from the // ProteusPass when detecting a CUDA module. diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 7f729d0a..c58f245c 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -66,6 +66,8 @@ CREATE_CPU_TEST(types_api types_api.cpp) CREATE_CPU_TEST(lambda lambda.cpp) CREATE_CPU_TEST(lambda_ptrtoint_capture lambda_ptrtoint_capture.cpp) CREATE_CPU_TEST(lambda_def lambda_def.cpp) +CREATE_CPU_TEST(lambda_nested lambda_nested.cpp) +CREATE_CPU_TEST(annot_nested annot_nested.cpp) CREATE_CPU_TEST(lambda_def_register_once lambda_def_register_once.cpp) CREATE_CPU_TEST(lambda_factory lambda_factory.cpp) CREATE_CPU_TEST(lambda_multiple lambda_multiple.cpp) diff --git a/tests/cpu/annot_nested.cpp b/tests/cpu/annot_nested.cpp new file mode 100644 index 00000000..7900a6af --- /dev/null +++ b/tests/cpu/annot_nested.cpp @@ -0,0 +1,43 @@ +// clang-format off +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_TRACE_OUTPUT="specialization;kernel-trace" %build/annot_nested | %FILECHECK %s +// RUN: rm -rf "%t.$$.proteus" +// Same as lambda_nested, for the annotated function interface: an annotated +// function called from a JIT'd function specializes on its own argument. +// clang-format on + +#include + +#include + +int Result; + +__attribute__((annotate("jit", 1))) void innerFn(int W) { Result += W; } + +__attribute__((annotate("jit", 1))) void outerFn(int V, int W) { + Result = V * 100; + innerFn(W); +} + +int main() { + for (int W : {10, 20, 30}) { + outerFn(1, W); + printf("Result %d\n", Result); + } + + return 0; +} + +// clang-format off +// CHECK: [ArgSpec] Replaced Function _Z7outerFnii ArgNo 0 with value i32 1 +// CHECK: [ArgSpec] Replaced Function _Z7innerFni ArgNo 0 with value i32 10 +// CHECK: Result 110 +// CHECK: [ArgSpec] Replaced Function _Z7innerFni ArgNo 0 with value i32 20 +// CHECK: Result 120 +// CHECK: [ArgSpec] Replaced Function _Z7innerFni ArgNo 0 with value i32 30 +// CHECK: Result 130 +// CHECK: === Kernel Trace (rank 0) === +// CHECK-DAG: outerFn(int, int) rank=0 specializations=1 launches=3 +// CHECK-DAG: innerFn(int) rank=0 specializations=3 launches=3 +// CHECK: === End Kernel Trace === +// clang-format on diff --git a/tests/cpu/lambda_nested.cpp b/tests/cpu/lambda_nested.cpp new file mode 100644 index 00000000..d6d88f21 --- /dev/null +++ b/tests/cpu/lambda_nested.cpp @@ -0,0 +1,43 @@ +// clang-format off +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_TRACE_OUTPUT="specialization;kernel-trace" %build/lambda_nested | %FILECHECK %s +// RUN: rm -rf "%t.$$.proteus" +// A lambda registered inside a JIT'd lambda body specializes on its own +// runtime constants: the outer region is compiled once for V, and the inner one +// is recompiled for each W. +// clang-format on + +#include + +#include + +template void run(F &&Func) { proteus::register_lambda(Func)(); } + +void nested(int V, int W) { + run([=, V = proteus::jit_variable(V)]() __attribute__((annotate("jit"))) { + run([=, W = proteus::jit_variable(W)]() + __attribute__((annotate("jit"))) { printf("V %d W %d\n", V, W); }); + }); +} + +int main() { + nested(1, 10); + nested(1, 20); + nested(1, 30); + + return 0; +} + +// clang-format off +// CHECK: [LambdaSpec] Replacing slot 0 with i32 1 +// CHECK: [LambdaSpec] Replacing slot 0 with i32 10 +// CHECK: V 1 W 10 +// CHECK: [LambdaSpec] Replacing slot 0 with i32 20 +// CHECK: V 1 W 20 +// CHECK: [LambdaSpec] Replacing slot 0 with i32 30 +// CHECK: V 1 W 30 +// CHECK: === Kernel Trace (rank 0) === +// CHECK-DAG: nested(int, int)::$_0::operator()() const rank=0 specializations=1 launches=3 +// CHECK-DAG: nested(int, int)::$_0::operator()() const::{{.*}}operator()() const rank=0 specializations=3 launches=3 +// CHECK: === End Kernel Trace === +// clang-format on diff --git a/tests/gpu/CMakeLists.txt b/tests/gpu/CMakeLists.txt index 3b8816c5..1750382a 100644 --- a/tests/gpu/CMakeLists.txt +++ b/tests/gpu/CMakeLists.txt @@ -244,6 +244,7 @@ CREATE_GPU_TEST(daxpy_api daxpy_api.cpp) CREATE_GPU_TEST(kernel_host_jit kernel_host_jit.cpp) CREATE_GPU_TEST(kernel_host_device_jit kernel_host_device_jit.cpp) CREATE_GPU_TEST(kernel_host_device_jit_api kernel_host_device_jit_api.cpp) +CREATE_GPU_TEST(annot_nested annot_nested.cpp) CREATE_GPU_TEST(types types.cpp) CREATE_GPU_TEST(types_api types_api.cpp) CREATE_GPU_TEST(kernel_unused_gvar kernel_unused_gvar.cpp kernel_unused_gvar_def.cpp) @@ -321,6 +322,7 @@ CREATE_GPU_TEST_RDC(daxpy_api daxpy_api.cpp) CREATE_GPU_TEST_RDC(kernel_host_jit kernel_host_jit.cpp) CREATE_GPU_TEST_RDC(kernel_host_device_jit kernel_host_device_jit.cpp) CREATE_GPU_TEST_RDC(kernel_host_device_jit_api kernel_host_device_jit_api.cpp) +CREATE_GPU_TEST_RDC(annot_nested annot_nested.cpp) CREATE_GPU_TEST_RDC(types types.cpp) CREATE_GPU_TEST_RDC(types_api types_api.cpp) CREATE_GPU_TEST_RDC(kernel_calls_func kernel_calls_func.cpp device_func.cpp) diff --git a/tests/gpu/annot_nested.cpp b/tests/gpu/annot_nested.cpp new file mode 100644 index 00000000..60562495 --- /dev/null +++ b/tests/gpu/annot_nested.cpp @@ -0,0 +1,45 @@ +// clang-format off +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_TRACE_OUTPUT="specialization;kernel-trace" %build/annot_nested.%ext | %FILECHECK %s +// RUN: rm -rf "%t.$$.proteus" +// A directly launched JIT kernel inside a JIT host function specializes on its +// own argument while the enclosing host function reuses its specialization. +// clang-format on + +#include + +#include "gpu_common.h" +#include + +__global__ __attribute__((annotate("jit", 1))) void innerFn(int W) { + printf("Inner %d\n", W); +} + +__attribute__((annotate("jit", 1))) void outerFn(int V, int W) { + printf("Outer %d\n", V); + innerFn<<<1, 1>>>(W); +} + +int main() { + for (int W : {10, 20, 30}) { + outerFn(1, W); + gpuErrCheck(gpuDeviceSynchronize()); + } + + return 0; +} + +// clang-format off +// CHECK: [ArgSpec] Replaced Function _Z7outerFnii ArgNo 0 with value i32 1 +// CHECK: Outer 1 +// CHECK: [ArgSpec] Replaced Function _Z7innerFni ArgNo 0 with value i32 10 +// CHECK: Inner 10 +// CHECK: Outer 1 +// CHECK: [ArgSpec] Replaced Function _Z7innerFni ArgNo 0 with value i32 20 +// CHECK: Inner 20 +// CHECK: Outer 1 +// CHECK: [ArgSpec] Replaced Function _Z7innerFni ArgNo 0 with value i32 30 +// CHECK: Inner 30 +// CHECK-DAG: [proteus][JitEngineHost] outerFn(int, int) rank=0 specializations=1 launches=3 +// CHECK-DAG: [proteus][JitEngineDevice] innerFn(int) rank=0 specializations=3 launches=3 +// clang-format on