From b75051923a4cfc0f090b521f8fe4a7892d3f7f41 Mon Sep 17 00:00:00 2001 From: Julian Andrej Date: Mon, 24 Aug 2026 16:05:41 -0700 Subject: [PATCH 1/7] add worklist sorting to support nested specializations for cpu --- src/pass/ProteusPass.cpp | 102 +++++++++++++++++++++++++++++++++--- tests/cpu/CMakeLists.txt | 2 + tests/cpu/annot_nested.cpp | 43 +++++++++++++++ tests/cpu/lambda_nested.cpp | 44 ++++++++++++++++ 4 files changed, 183 insertions(+), 8 deletions(-) create mode 100644 tests/cpu/annot_nested.cpp create mode 100644 tests/cpu/lambda_nested.cpp diff --git a/src/pass/ProteusPass.cpp b/src/pass/ProteusPass.cpp index 529c1cb3..d1039205 100644 --- a/src/pass/ProteusPass.cpp +++ b/src/pass/ProteusPass.cpp @@ -58,6 +58,7 @@ #include #include #include +#include #include #include #include @@ -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; }); 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..4de20e56 --- /dev/null +++ b/tests/cpu/lambda_nested.cpp @@ -0,0 +1,44 @@ +// 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 From 6b591bc7443f64b5ffd562b438d3d19e7d391c0d Mon Sep 17 00:00:00 2001 From: Julian Andrej Date: Tue, 25 Aug 2026 08:43:11 -0700 Subject: [PATCH 2/7] gpu pass --- .../proteus/impl/CompilerInterfaceDevice.h | 5 +++ src/include/proteus/impl/CoreDeviceCUDA.h | 5 +++ src/include/proteus/impl/JitEngineDevice.h | 12 +++++ src/pass/ProteusPass.cpp | 32 +++++++++---- src/runtime/CompilerInterfaceDevice.cpp | 17 +++++++ src/runtime/JitEngineHost.cpp | 16 +++++++ src/runtime/ProteusCUDARuntimeBuiltins.cpp | 12 +++++ tests/gpu/CMakeLists.txt | 2 + tests/gpu/annot_nested.cpp | 45 +++++++++++++++++++ 9 files changed, 138 insertions(+), 8 deletions(-) create mode 100644 tests/gpu/annot_nested.cpp diff --git a/src/include/proteus/impl/CompilerInterfaceDevice.h b/src/include/proteus/impl/CompilerInterfaceDevice.h index 88613475..2f10aaa1 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 *KernelName, dim3 GridDim, + dim3 BlockDim, void **KernelArgs, + uint64_t ShmemSize, void *Stream); + #endif diff --git a/src/include/proteus/impl/CoreDeviceCUDA.h b/src/include/proteus/impl/CoreDeviceCUDA.h index fe9eff92..a0b19f15 100644 --- a/src/include/proteus/impl/CoreDeviceCUDA.h +++ b/src/include/proteus/impl/CoreDeviceCUDA.h @@ -20,6 +20,11 @@ inline cudaError_t (*__proteus_cudaGetSymbolAddress_ptr)( inline cudaError_t (*__proteus_cudaLaunchKernel_ptr)(const void *, dim3, dim3, void **, size_t, cudaStream_t) = nullptr; +inline unsigned (*__proteus_cudaPushCallConfiguration_ptr)(dim3, dim3, size_t, + void *) = nullptr; +inline cudaError_t (*__proteus_cudaPopCallConfiguration_ptr)(dim3 *, dim3 *, + size_t *, + void *) = nullptr; } // NOLINTEND(readability-identifier-naming) diff --git a/src/include/proteus/impl/JitEngineDevice.h b/src/include/proteus/impl/JitEngineDevice.h index f3ccbf63..9ec39d66 100644 --- a/src/include/proteus/impl/JitEngineDevice.h +++ b/src/include/proteus/impl/JitEngineDevice.h @@ -499,6 +499,15 @@ template class JitEngineDevice : public JitEngine { return JITKernelInfoMap[Func]; } + std::optional> + getJITKernelInfo(StringRef FuncName) { + auto It = KernelNameToKernel.find(FuncName.str()); + if (It == KernelNameToKernel.end()) + return std::nullopt; + + return getJITKernelInfo(It->second); + } + HashT getStaticHash(JITKernelInfo &KernelInfo) { if (KernelInfo.hasStaticHash()) return KernelInfo.getStaticHash(); @@ -571,6 +580,7 @@ template class JitEngineDevice : public JitEngine { std::string DeviceArch; DenseMap JITKernelInfoMap; + std::unordered_map KernelNameToKernel; DenseMap PendingLambdaCallsiteLocationInfo; std::unique_ptr AsyncCompiler; @@ -755,6 +765,8 @@ void JitEngineDevice::registerFunction( return; } + KernelNameToKernel.try_emplace(KernelName, Kernel); + if (!HandleToBinaryInfo.count(Handle)) reportFatalError("Expected Handle in map"); BinaryInfo &BinInfo = HandleToBinaryInfo[Handle]; diff --git a/src/pass/ProteusPass.cpp b/src/pass/ProteusPass.cpp index d1039205..70316514 100644 --- a/src/pass/ProteusPass.cpp +++ b/src/pass/ProteusPass.cpp @@ -206,7 +206,7 @@ class ProteusPassImpl { if (hasDeviceLaunchKernelCalls(M)) { instrumentLambdaLaunchCallsites(M, StubToKernelMap); - emitJitLaunchKernelCall(M); + emitJitLaunchKernelCall(M, StubToKernelMap); } instrumentRegisterFunction(M); @@ -1585,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"); @@ -1603,14 +1603,18 @@ 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); + void replaceWithJitLaunchKernel(Module &M, CallBase *LaunchKernelCB, + GlobalVariable *KernelName) { + FunctionCallee JitLaunchKernelFn = + getJitLaunchKernelFn(M, KernelName != nullptr); // Insert before the launch kernel call instruction. IRBuilder<> Builder(LaunchKernelCB); @@ -1618,6 +1622,8 @@ class ProteusPassImpl { SmallVector Args = {LaunchKernelCB->arg_begin(), LaunchKernelCB->arg_end()}; + if (KernelName) + Args[0] = KernelName; if (isa(LaunchKernelCB)) { CallOrInvoke = Builder.CreateCall(JitLaunchKernelFn, Args); @@ -1637,7 +1643,8 @@ class ProteusPassImpl { LaunchKernelCB->eraseFromParent(); } - void emitJitLaunchKernelCall(Module &M) { + void emitJitLaunchKernelCall( + Module &M, const DenseMap &StubToKernelMap) { Function *LaunchKernelFn = nullptr; if (!LaunchFunctionName) { reportFatalError( @@ -1666,8 +1673,17 @@ class ProteusPassImpl { ToBeReplaced.push_back(CB); } - for (CallBase *CB : ToBeReplaced) - replaceWithJitLaunchKernel(M, CB); + for (CallBase *CB : ToBeReplaced) { + GlobalVariable *KernelName = 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)) + KernelName = It->second; + + replaceWithJitLaunchKernel(M, CB, KernelName); + } } FunctionCallee getJitRegisterFatBinaryFn(Module &M) { diff --git a/src/runtime/CompilerInterfaceDevice.cpp b/src/runtime/CompilerInterfaceDevice.cpp index 5b124583..9483d736 100644 --- a/src/runtime/CompilerInterfaceDevice.cpp +++ b/src/runtime/CompilerInterfaceDevice.cpp @@ -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 *KernelName, 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{KernelName}); + if (!OptionalKernelInfo) + reportFatalError("Missing registered GPU kernel " + Twine(KernelName)); + + 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/JitEngineHost.cpp b/src/runtime/JitEngineHost.cpp index a633b0ce..cd3b852c 100644 --- a/src/runtime/JitEngineHost.cpp +++ b/src/runtime/JitEngineHost.cpp @@ -82,6 +82,18 @@ void JitEngineHost::addStaticLibrarySymbols() { __proteus_cudaLaunchKernel_ptr)}, JITSymbolFlags::Exported); } + if (__proteus_cudaPushCallConfiguration_ptr) { + SymbolMap[LLJITPtr->mangleAndIntern("__cudaPushCallConfiguration")] = + orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( + __proteus_cudaPushCallConfiguration_ptr)}, + JITSymbolFlags::Exported); + } + if (__proteus_cudaPopCallConfiguration_ptr) { + SymbolMap[LLJITPtr->mangleAndIntern("__cudaPopCallConfiguration")] = + orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( + __proteus_cudaPopCallConfiguration_ptr)}, + JITSymbolFlags::Exported); + } #endif @@ -91,6 +103,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..30ee4172 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. @@ -26,6 +32,10 @@ extern cudaError_t (*__proteus_cudaGetSymbolAddress_ptr)(void **, const void *); extern cudaError_t (*__proteus_cudaLaunchKernel_ptr)(const void *, dim3, dim3, void **, size_t, cudaStream_t); +extern unsigned (*__proteus_cudaPushCallConfiguration_ptr)(dim3, dim3, size_t, + void *); +extern cudaError_t (*__proteus_cudaPopCallConfiguration_ptr)(dim3 *, dim3 *, + size_t *, void *); } // Initialization function to set the function pointers for the CUDA runtime @@ -34,6 +44,8 @@ extern cudaError_t (*__proteus_cudaLaunchKernel_ptr)(const void *, dim3, dim3, extern "C" void __proteus_cudart_builtins_init() { __proteus_cudaGetSymbolAddress_ptr = checkCudaGetSymbolAddress; __proteus_cudaLaunchKernel_ptr = cudaLaunchKernel; + __proteus_cudaPushCallConfiguration_ptr = __cudaPushCallConfiguration; + __proteus_cudaPopCallConfiguration_ptr = __cudaPopCallConfiguration; } // NOLINTEND(readability-identifier-naming) 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 From fcb0f349508ec09aea74a608f17de0afd3537154 Mon Sep 17 00:00:00 2001 From: Julian Andrej Date: Tue, 25 Aug 2026 09:38:32 -0700 Subject: [PATCH 3/7] hip --- src/include/proteus/impl/HIPRuntimeAPI.h | 4 ++++ src/runtime/HIPRuntimeAPI.cpp | 14 ++++++++++++++ src/runtime/JitEngineHost.cpp | 14 ++++++++++++++ 3 files changed, 32 insertions(+) diff --git a/src/include/proteus/impl/HIPRuntimeAPI.h b/src/include/proteus/impl/HIPRuntimeAPI.h index 9ee12a41..eb01fe96 100644 --- a/src/include/proteus/impl/HIPRuntimeAPI.h +++ b/src/include/proteus/impl/HIPRuntimeAPI.h @@ -25,6 +25,10 @@ hipError_t moduleLaunchKernel(hipFunction_t Function, unsigned int GridDimX, hipError_t launchKernel(const void *FunctionAddress, dim3 NumBlocks, dim3 DimBlocks, void **Args, size_t SharedMemBytes, hipStream_t Stream); +hipError_t pushCallConfiguration(dim3 GridDim, dim3 BlockDim, size_t SharedMem, + hipStream_t Stream); +hipError_t popCallConfiguration(dim3 *GridDim, dim3 *BlockDim, + size_t *SharedMem, hipStream_t *Stream); hipError_t funcSetAttribute(const void *Function, hipFuncAttribute Attribute, int Value); diff --git a/src/runtime/HIPRuntimeAPI.cpp b/src/runtime/HIPRuntimeAPI.cpp index 82a26411..5d7a5d1e 100644 --- a/src/runtime/HIPRuntimeAPI.cpp +++ b/src/runtime/HIPRuntimeAPI.cpp @@ -175,6 +175,20 @@ hipError_t launchKernel(const void *FunctionAddress, dim3 NumBlocks, Stream); } +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); +} + hipError_t funcSetAttribute(const void *Function, hipFuncAttribute Attribute, int Value) { using Fn = hipError_t (*)(const void *, hipFuncAttribute, int); diff --git a/src/runtime/JitEngineHost.cpp b/src/runtime/JitEngineHost.cpp index cd3b852c..edcd2c1e 100644 --- a/src/runtime/JitEngineHost.cpp +++ b/src/runtime/JitEngineHost.cpp @@ -17,6 +17,9 @@ #include "proteus/impl/LambdaRegistry.h" #include "proteus/impl/TransformArgumentSpecialization.h" #include "proteus/impl/TransformLambdaSpecialization.h" +#if PROTEUS_ENABLE_HIP +#include "proteus/impl/HIPRuntimeAPI.h" +#endif #include #if PROTEUS_ENABLE_HIP || PROTEUS_ENABLE_CUDA #include "proteus/impl/CompilerInterfaceDevice.h" @@ -97,6 +100,17 @@ void JitEngineHost::addStaticLibrarySymbols() { #endif +#if PROTEUS_ENABLE_HIP + SymbolMap[LLJITPtr->mangleAndIntern("__hipPushCallConfiguration")] = + orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( + &proteus::hipdyn::pushCallConfiguration)}, + JITSymbolFlags::Exported); + SymbolMap[LLJITPtr->mangleAndIntern("__hipPopCallConfiguration")] = + orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( + &proteus::hipdyn::popCallConfiguration)}, + JITSymbolFlags::Exported); +#endif + #if PROTEUS_ENABLE_CUDA || PROTEUS_ENABLE_HIP // Add __proteus_launch_kernel as a static symbol. SymbolMap[LLJITPtr->mangleAndIntern("__proteus_launch_kernel")] = From 0291f4b4cdcaf0631598039d4ccc07526fcce187 Mon Sep 17 00:00:00 2001 From: Julian Andrej Date: Tue, 25 Aug 2026 10:35:47 -0700 Subject: [PATCH 4/7] linkage --- src/include/proteus/impl/CoreDeviceCUDA.h | 5 --- src/include/proteus/impl/HIPRuntimeAPI.h | 4 -- src/runtime/HIPRuntimeAPI.cpp | 37 ++++++++++------ src/runtime/JitEngineHost.cpp | 50 +++++++++++----------- src/runtime/ProteusCUDARuntimeBuiltins.cpp | 15 ++++--- 5 files changed, 57 insertions(+), 54 deletions(-) diff --git a/src/include/proteus/impl/CoreDeviceCUDA.h b/src/include/proteus/impl/CoreDeviceCUDA.h index a0b19f15..fe9eff92 100644 --- a/src/include/proteus/impl/CoreDeviceCUDA.h +++ b/src/include/proteus/impl/CoreDeviceCUDA.h @@ -20,11 +20,6 @@ inline cudaError_t (*__proteus_cudaGetSymbolAddress_ptr)( inline cudaError_t (*__proteus_cudaLaunchKernel_ptr)(const void *, dim3, dim3, void **, size_t, cudaStream_t) = nullptr; -inline unsigned (*__proteus_cudaPushCallConfiguration_ptr)(dim3, dim3, size_t, - void *) = nullptr; -inline cudaError_t (*__proteus_cudaPopCallConfiguration_ptr)(dim3 *, dim3 *, - size_t *, - void *) = nullptr; } // NOLINTEND(readability-identifier-naming) diff --git a/src/include/proteus/impl/HIPRuntimeAPI.h b/src/include/proteus/impl/HIPRuntimeAPI.h index eb01fe96..9ee12a41 100644 --- a/src/include/proteus/impl/HIPRuntimeAPI.h +++ b/src/include/proteus/impl/HIPRuntimeAPI.h @@ -25,10 +25,6 @@ hipError_t moduleLaunchKernel(hipFunction_t Function, unsigned int GridDimX, hipError_t launchKernel(const void *FunctionAddress, dim3 NumBlocks, dim3 DimBlocks, void **Args, size_t SharedMemBytes, hipStream_t Stream); -hipError_t pushCallConfiguration(dim3 GridDim, dim3 BlockDim, size_t SharedMem, - hipStream_t Stream); -hipError_t popCallConfiguration(dim3 *GridDim, dim3 *BlockDim, - size_t *SharedMem, hipStream_t *Stream); hipError_t funcSetAttribute(const void *Function, hipFuncAttribute Attribute, int Value); diff --git a/src/runtime/HIPRuntimeAPI.cpp b/src/runtime/HIPRuntimeAPI.cpp index 5d7a5d1e..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) { @@ -175,20 +198,6 @@ hipError_t launchKernel(const void *FunctionAddress, dim3 NumBlocks, Stream); } -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); -} - hipError_t funcSetAttribute(const void *Function, hipFuncAttribute Attribute, int Value) { using Fn = hipError_t (*)(const void *, hipFuncAttribute, int); diff --git a/src/runtime/JitEngineHost.cpp b/src/runtime/JitEngineHost.cpp index edcd2c1e..3bf6ebb4 100644 --- a/src/runtime/JitEngineHost.cpp +++ b/src/runtime/JitEngineHost.cpp @@ -17,9 +17,6 @@ #include "proteus/impl/LambdaRegistry.h" #include "proteus/impl/TransformArgumentSpecialization.h" #include "proteus/impl/TransformLambdaSpecialization.h" -#if PROTEUS_ENABLE_HIP -#include "proteus/impl/HIPRuntimeAPI.h" -#endif #include #if PROTEUS_ENABLE_HIP || PROTEUS_ENABLE_CUDA #include "proteus/impl/CompilerInterfaceDevice.h" @@ -45,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 @@ -85,30 +93,22 @@ void JitEngineHost::addStaticLibrarySymbols() { __proteus_cudaLaunchKernel_ptr)}, JITSymbolFlags::Exported); } - if (__proteus_cudaPushCallConfiguration_ptr) { - SymbolMap[LLJITPtr->mangleAndIntern("__cudaPushCallConfiguration")] = - orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( - __proteus_cudaPushCallConfiguration_ptr)}, - JITSymbolFlags::Exported); - } - if (__proteus_cudaPopCallConfiguration_ptr) { - SymbolMap[LLJITPtr->mangleAndIntern("__cudaPopCallConfiguration")] = - orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( - __proteus_cudaPopCallConfiguration_ptr)}, - JITSymbolFlags::Exported); - } - #endif -#if PROTEUS_ENABLE_HIP - SymbolMap[LLJITPtr->mangleAndIntern("__hipPushCallConfiguration")] = - orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( - &proteus::hipdyn::pushCallConfiguration)}, - JITSymbolFlags::Exported); - SymbolMap[LLJITPtr->mangleAndIntern("__hipPopCallConfiguration")] = - orc::ExecutorSymbolDef(orc::ExecutorAddr{reinterpret_cast( - &proteus::hipdyn::popCallConfiguration)}, - JITSymbolFlags::Exported); +#if PROTEUS_ENABLE_CUDA || PROTEUS_ENABLE_HIP + if (__proteus_get_device_launch_config_symbols) { + 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 diff --git a/src/runtime/ProteusCUDARuntimeBuiltins.cpp b/src/runtime/ProteusCUDARuntimeBuiltins.cpp index 30ee4172..e0e617cf 100644 --- a/src/runtime/ProteusCUDARuntimeBuiltins.cpp +++ b/src/runtime/ProteusCUDARuntimeBuiltins.cpp @@ -32,10 +32,15 @@ extern cudaError_t (*__proteus_cudaGetSymbolAddress_ptr)(void **, const void *); extern cudaError_t (*__proteus_cudaLaunchKernel_ptr)(const void *, dim3, dim3, void **, size_t, cudaStream_t); -extern unsigned (*__proteus_cudaPushCallConfiguration_ptr)(dim3, dim3, size_t, - void *); -extern cudaError_t (*__proteus_cudaPopCallConfiguration_ptr)(dim3 *, dim3 *, - size_t *, void *); +} + +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 @@ -44,8 +49,6 @@ extern cudaError_t (*__proteus_cudaPopCallConfiguration_ptr)(dim3 *, dim3 *, extern "C" void __proteus_cudart_builtins_init() { __proteus_cudaGetSymbolAddress_ptr = checkCudaGetSymbolAddress; __proteus_cudaLaunchKernel_ptr = cudaLaunchKernel; - __proteus_cudaPushCallConfiguration_ptr = __cudaPushCallConfiguration; - __proteus_cudaPopCallConfiguration_ptr = __cudaPopCallConfiguration; } // NOLINTEND(readability-identifier-naming) From 5371ab7b5fc75d81bff61c039349c3ec44348fbe Mon Sep 17 00:00:00 2001 From: Julian Andrej Date: Tue, 25 Aug 2026 10:44:47 -0700 Subject: [PATCH 5/7] format --- include/proteus/JitInterface.h | 9 ++++----- tests/cpu/lambda_nested.cpp | 5 ++--- 2 files changed, 6 insertions(+), 8 deletions(-) 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/tests/cpu/lambda_nested.cpp b/tests/cpu/lambda_nested.cpp index 4de20e56..d6d88f21 100644 --- a/tests/cpu/lambda_nested.cpp +++ b/tests/cpu/lambda_nested.cpp @@ -15,9 +15,8 @@ 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); - }); + run([=, W = proteus::jit_variable(W)]() + __attribute__((annotate("jit"))) { printf("V %d W %d\n", V, W); }); }); } From 85e7cf60024d46e19637f62d22117d09deb7c90f Mon Sep 17 00:00:00 2001 From: Julian Andrej Date: Tue, 25 Aug 2026 11:13:06 -0700 Subject: [PATCH 6/7] make kernels translation unit agnostic --- .../proteus/impl/CompilerInterfaceDevice.h | 2 +- src/include/proteus/impl/JitEngineDevice.h | 17 +++++----- .../proteus/impl/JitEngineInfoRegistry.h | 5 ++- src/pass/ProteusPass.cpp | 33 ++++++++++++------- src/runtime/CompilerInterfaceDevice.cpp | 16 ++++----- 5 files changed, 43 insertions(+), 30 deletions(-) diff --git a/src/include/proteus/impl/CompilerInterfaceDevice.h b/src/include/proteus/impl/CompilerInterfaceDevice.h index 2f10aaa1..4141f4f7 100644 --- a/src/include/proteus/impl/CompilerInterfaceDevice.h +++ b/src/include/proteus/impl/CompilerInterfaceDevice.h @@ -33,7 +33,7 @@ __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 *KernelName, dim3 GridDim, +__proteus_launch_kernel_by_name(const char *KernelLookupKey, dim3 GridDim, dim3 BlockDim, void **KernelArgs, uint64_t ShmemSize, void *Stream); diff --git a/src/include/proteus/impl/JitEngineDevice.h b/src/include/proteus/impl/JitEngineDevice.h index 9ec39d66..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, @@ -500,9 +501,9 @@ template class JitEngineDevice : public JitEngine { } std::optional> - getJITKernelInfo(StringRef FuncName) { - auto It = KernelNameToKernel.find(FuncName.str()); - if (It == KernelNameToKernel.end()) + getJITKernelInfo(StringRef KernelLookupKey) { + auto It = KernelLookupKeyToKernel.find(KernelLookupKey.str()); + if (It == KernelLookupKeyToKernel.end()) return std::nullopt; return getJITKernelInfo(It->second); @@ -545,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); @@ -580,7 +581,7 @@ template class JitEngineDevice : public JitEngine { std::string DeviceArch; DenseMap JITKernelInfoMap; - std::unordered_map KernelNameToKernel; + std::unordered_map KernelLookupKeyToKernel; DenseMap PendingLambdaCallsiteLocationInfo; std::unique_ptr AsyncCompiler; @@ -749,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 @@ -765,8 +768,6 @@ void JitEngineDevice::registerFunction( return; } - KernelNameToKernel.try_emplace(KernelName, Kernel); - if (!HandleToBinaryInfo.count(Handle)) reportFatalError("Expected Handle in map"); BinaryInfo &BinInfo = HandleToBinaryInfo[Handle]; 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 70316514..edf6f52a 100644 --- a/src/pass/ProteusPass.cpp +++ b/src/pass/ProteusPass.cpp @@ -1611,10 +1611,14 @@ class ProteusPassImpl { return JitLaunchKernelFn; } + std::string getKernelLookupKey(Module &M, const Function &KernelStub) { + return getUniqueFileID(M) + ":" + KernelStub.getName().str(); + } + void replaceWithJitLaunchKernel(Module &M, CallBase *LaunchKernelCB, - GlobalVariable *KernelName) { + Function *KernelStub) { FunctionCallee JitLaunchKernelFn = - getJitLaunchKernelFn(M, KernelName != nullptr); + getJitLaunchKernelFn(M, KernelStub != nullptr); // Insert before the launch kernel call instruction. IRBuilder<> Builder(LaunchKernelCB); @@ -1622,8 +1626,9 @@ class ProteusPassImpl { SmallVector Args = {LaunchKernelCB->arg_begin(), LaunchKernelCB->arg_end()}; - if (KernelName) - Args[0] = KernelName; + if (KernelStub) + Args[0] = Builder.CreateGlobalString(getKernelLookupKey(M, *KernelStub), + ".proteus.kernel.lookup"); if (isa(LaunchKernelCB)) { CallOrInvoke = Builder.CreateCall(JitLaunchKernelFn, Args); @@ -1674,15 +1679,15 @@ class ProteusPassImpl { } for (CallBase *CB : ToBeReplaced) { - GlobalVariable *KernelName = nullptr; + 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)) - KernelName = It->second; + KernelStub = StubFn; - replaceWithJitLaunchKernel(M, CB, KernelName); + replaceWithJitLaunchKernel(M, CB, KernelStub); } } @@ -1849,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); @@ -1925,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 9483d736..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 @@ -147,14 +147,14 @@ __proteus_launch_kernel(void *Kernel, dim3 GridDim, dim3 BlockDim, } extern "C" proteus::DeviceTraits::DeviceError_t -__proteus_launch_kernel_by_name(const char *KernelName, dim3 GridDim, +__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{KernelName}); + auto OptionalKernelInfo = Jit.getJITKernelInfo(StringRef{KernelLookupKey}); if (!OptionalKernelInfo) - reportFatalError("Missing registered GPU kernel " + Twine(KernelName)); + reportFatalError("Missing registered GPU kernel " + Twine(KernelLookupKey)); void *Kernel = OptionalKernelInfo->get().getKernel(); auto &LR = LambdaRegistry::instance(); From 54e25c9e6396fcb5469e1ef3c1f0732c96fff5e1 Mon Sep 17 00:00:00 2001 From: Julian Andrej Date: Tue, 25 Aug 2026 11:48:16 -0700 Subject: [PATCH 7/7] don't expose symbol to hip --- src/runtime/JitEngineHost.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/runtime/JitEngineHost.cpp b/src/runtime/JitEngineHost.cpp index 3bf6ebb4..0fb78e69 100644 --- a/src/runtime/JitEngineHost.cpp +++ b/src/runtime/JitEngineHost.cpp @@ -95,8 +95,12 @@ void JitEngineHost::addStaticLibrarySymbols() { } #endif -#if PROTEUS_ENABLE_CUDA || PROTEUS_ENABLE_HIP +#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;