add worklist sorting to support nested specializations

include/proteus/JitInterface.h

@@ -1,3 +1,3 @@
 //===-- jit.h -- user interface to Proteus JIT library --===//
 //
 // Part of the Proteus Project, under the Apache License v2.0 with LLVM
@@ -65,9 +65,9 @@
 
 #if defined(__CUDACC__) || defined(__HIP__)
 template <typename T>
-__attribute__((noinline)) __device__ std::enable_if_t<
-    std::is_trivially_copyable_v<std::remove_pointer_t<T>>, void>
-jit_object(T *V, size_t Size = sizeof(T)) noexcept;
+__attribute__((noinline)) __device__
+    std::enable_if_t<std::is_trivially_copyable_v<std::remove_pointer_t<T>>,
+                     void> jit_object(T *V, size_t Size = sizeof(T)) noexcept;
 #endif
 
 template <typename T>
@@ -82,8 +82,7 @@
 __attribute__((noinline)) __device__ std::enable_if_t<
     !std::is_pointer_v<T> &&
         std::is_trivially_copyable_v<std::remove_reference_t<T>>,
-    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 {

src/include/proteus/impl/CompilerInterfaceDevice.h

@@ -32,4 +32,9 @@
 __proteus_launch_kernel(void *Kernel, dim3 GridDim, dim3 BlockDim,
                         void **KernelArgs, uint64_t ShmemSize, void *Stream);
 
+extern "C" proteus::DeviceTraits<JitDeviceImplT>::DeviceError_t
+__proteus_launch_kernel_by_name(const char *KernelLookupKey, dim3 GridDim,
+                                dim3 BlockDim, void **KernelArgs,
+                                uint64_t ShmemSize, void *Stream);
+
 #endif

src/include/proteus/impl/JitEngineDevice.h

@@ -477,6 +477,7 @@ template <typename ImplT> class JitEngineDevice : public JitEngine {
                          const char *ModuleId);
   void finalizeRegistration();
   void registerFunction(void *Handle, void *Kernel, char *KernelName,
+                        const char *KernelLookupKey,
                         ArrayRef<RuntimeConstantInfo *> RCInfoArray);
   void registerLambdaCallsiteLocation(void *Kernel, uint64_t LambdaID,
                                       uint32_t CallsiteIndex,
@@ -499,6 +500,15 @@ template <typename ImplT> class JitEngineDevice : public JitEngine {
     return JITKernelInfoMap[Func];
   }
 
+  std::optional<std::reference_wrapper<JITKernelInfo>>
+  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 <typename ImplT> 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 <typename ImplT> class JitEngineDevice : public JitEngine {
   std::string DeviceArch;
 
   DenseMap<const void *, JITKernelInfo> JITKernelInfoMap;
+  std::unordered_map<std::string, const void *> KernelLookupKeyToKernel;
   DenseMap<const void *, LambdaCallsiteLocationMap>
       PendingLambdaCallsiteLocationInfo;
   std::unique_ptr<CompilerAsync> AsyncCompiler;
@@ -739,10 +750,12 @@ template <typename ImplT> void JitEngineDevice<ImplT>::finalizeRegistration() {
 
 template <typename ImplT>
 void JitEngineDevice<ImplT>::registerFunction(
-    void *Handle, void *Kernel, char *KernelName,
+    void *Handle, void *Kernel, char *KernelName, const char *KernelLookupKey,
     ArrayRef<RuntimeConstantInfo *> 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

src/include/proteus/impl/JitEngineInfoRegistry.h

@@ -37,6 +37,7 @@ struct RegisterFunctionInfo {
   void *Handle;
   void *Kernel;
   char *KernelName;
+  const char *KernelLookupKey;
   ArrayRef<RuntimeConstantInfo *> RCInfoArray;
 };
 
@@ -75,9 +76,11 @@ class JitEngineInfoRegistry {
   }
 
   void registerFunction(void *Handle, void *Kernel, char *KernelName,
+                        const char *KernelLookupKey,
                         ArrayRef<RuntimeConstantInfo *> 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,

src/pass/ProteusPass.cpp

@@ -58,6 +58,7 @@
 #include <llvm/IR/Function.h>
 #include <llvm/IR/GlobalValue.h>
 #include <llvm/IR/GlobalVariable.h>
+#include <llvm/IR/InstIterator.h>
 #include <llvm/IR/InstrTypes.h>
 #include <llvm/IR/Instruction.h>
 #include <llvm/IR/Instructions.h>
@@ -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<Function *, 8>
+  findNestedJitFunctions(Function &F,
+                         const SmallPtrSetImpl<Function *> &JitFunctions) {
+    SmallPtrSet<Function *, 8> Nested;
+    SmallPtrSet<Function *, 16> Visited;
+    SmallVector<Function *, 16> 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<CallBase>(&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<JitWorkListEntry *, 16> sortJitWorkListInnermostFirst(
+      const SmallVectorImpl<JitWorkListEntry *> &JitWorkList) {
+    SmallPtrSet<Function *, 16> JitFunctions;
+    DenseMap<Function *, JitWorkListEntry *> FnToEntry;
+    for (auto *JFI : JitWorkList) {
+      JitFunctions.insert(JFI->first);
+      FnToEntry[JFI->first] = JFI;
+    }
+
+    DenseMap<Function *, SmallPtrSet<Function *, 8>> Nested;
+    for (auto *JFI : JitWorkList)
+      Nested[JFI->first] = findNestedJitFunctions(*JFI->first, JitFunctions);
+
+    SmallVector<JitWorkListEntry *, 16> Ordered;
+    SmallPtrSet<Function *, 16> Done;
+    SmallPtrSet<Function *, 16> OnStack;
+
+    std::function<void(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<Function *, JitFunctionInfo> &JITInfo) {
     Function *JITFn = JITInfo.first;
@@ -786,9 +864,17 @@ class ProteusPassImpl {
           if (isCoverageGlobal(*GV))
             return true;
 
-          if (const GlobalVariable *G = dyn_cast<GlobalVariable>(GV))
+          if (const GlobalVariable *G = dyn_cast<GlobalVariable>(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,21 +1603,32 @@ 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);
     CallBase *CallOrInvoke = nullptr;
 
     SmallVector<Value *> Args = {LaunchKernelCB->arg_begin(),
                                  LaunchKernelCB->arg_end()};
+    if (KernelStub)
+      Args[0] = Builder.CreateGlobalString(getKernelLookupKey(M, *KernelStub),
+                                           ".proteus.kernel.lookup");
 
     if (isa<CallInst>(LaunchKernelCB)) {
       CallOrInvoke = Builder.CreateCall(JitLaunchKernelFn, Args);
@@ -1551,7 +1648,8 @@ class ProteusPassImpl {
     LaunchKernelCB->eraseFromParent();
   }
 
-  void emitJitLaunchKernelCall(Module &M) {
+  void emitJitLaunchKernelCall(
+      Module &M, const DenseMap<Value *, GlobalVariable *> &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<Function>(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 =