feature: redesign Proteus's lambda JIT compilation interface and pipeline by johnbowen42 · Pull Request #467 · Olympus-HPC/proteus

johnbowen42 · 2026-04-14T21:12:04Z

This feature restructures Proteus's lambda support so that a device module containing multiple calls of a lambda with different runtime constants can be correctly JIT compiled. This feature makes Proteus's support of JIT compiled lambdas closer to supporting all possible C++ lambda syntax, instead of the narrow syntax often used by parallel programming abstraction layers. Concretely

Add a functor wrapper to Proteus lambdas so that lambda callsites can be transparently tracked
Add LambdaPass which runs at the start of the pipeline and (1) handles jit_variable callsite analysis on less-simplified IR (before InstCombine can cause indirection at callsite (2) instruments lambda runtime constant registration for host lambda callsites and (3) generates manifests of lambda constants registered to device lambda callsite. Also LambdaPass removes global annotations on lambda functor operator callsites and attaches them as metadata nodes
Change ProteusPass so that (1) it consumes the LambdaPass manifests and then (2) uses them to build lambda registration functions, which are executed right before kernel launch
Change the runtime to look for lambda callsites in a device module containing metadata nodes, use that metadata as a key for lookup, and specialize each callsite based on the result lookup.

ggeorgakoudis

This is great and will solve our lambda issues.

I made a first pass and left comments on things we need to restructure.

ggeorgakoudis · 2026-04-15T18:47:52Z

I think installing and using a separate pass is quirky. We already register ProteusPass in multiple places (early-simplifications and LTO), so we can register it also at pipeline-start and just do the lambda analysis

Could you expand on this? How else would be run an early LambdaAnalysis pass and Proteuspass later?

Register ProteusPass at pipeline-start and pass an argument when you construct it to indicate the insert point. Check how we do it for early-simplifications and LTO. Apparently state is not true-false anymore, if there are 3 possible insert points.

johnbowen42 · 2026-04-17T22:18:30Z

all tests are passing for me locally on Tioga with rocm 6.4.2

johnbowen42 · 2026-07-02T19:02:33Z

This analysis is now compatible with RAJA launch, kernel, and forall LLVM IR shapes
llnl/RAJA#2052

codecov · 2026-07-08T19:55:56Z

Codecov Report

❌ Patch coverage is 76.95606% with 430 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/pass/KernelArgVisitor.h	55.46%	159 Missing ⚠️
src/pass/LambdaAnalysis.cpp	83.47%	77 Missing ⚠️
src/pass/KernelArgPtrUseVisitor.h	69.84%	60 Missing ⚠️
src/pass/ProteusPass.cpp	88.80%	42 Missing ⚠️
src/include/proteus/impl/JitEngineDevice.h	47.91%	25 Missing ⚠️
src/pass/AnnotationHandler.cpp	66.10%	20 Missing ⚠️
...clude/proteus/impl/TransformLambdaSpecialization.h	86.71%	17 Missing ⚠️
src/include/proteus/impl/Hashing.h	65.21%	16 Missing ⚠️
src/include/proteus/impl/LambdaCallsite.h	82.75%	5 Missing ⚠️
src/runtime/JitEngineHost.cpp	84.37%	5 Missing ⚠️
... and 2 more

📢 Thoughts on this report? Let us know!

remove jit var analysis from later pass tmp update tests check point: all lambda tests GPU passing add wrapper functors Lambda redeisgn: add functors. This commit builds but the runtime needs work delete plugin get tests working Refactor tests and track lambdas by global annotations Remove changes to type traits Get rid of storing things in llvm.global.annotations. Fix lambda CPU test All lambda unit tests passing. Fixed cloning bug in ProteusPass delete xfail compile test remove script change Fix all unit tests

…need to rework host lambda specialization

… expecting jit_variable schema for every lambda

…alue and insert_value analysis

…rt through ProteusPass

ZwFink

This is great work. I have mostly minor things.
Because this PR is so big and has been so long in the works, I don't want to hold it up for any of these things.
I am concerned that we would perpetually run around in bike-shedding circles, all the while this important feature is held up. Anything you choose not to resolve should be opened as an issue for follow-up.

ZwFink · 2026-06-17T18:11:22Z

+    // We need collect any kernel host stubs to pass to parse annotations, used
+    // in forced annotations.
+    const auto StubToKernelMap = getKernelHostStubs(M);
+


Copy-paste artifact from ProteusPass? This is never used here. Can be removed along with getKernelHostStubs.

ZwFink · 2026-06-17T18:11:32Z

+  DenseMap<Value *, GlobalVariable *> getKernelHostStubs(Module &M) {
+    DenseMap<Value *, GlobalVariable *> StubToKernelMap;


ZwFink · 2026-06-17T18:11:45Z

+  Value *getStubGV([[maybe_unused]] Value *Operand) {
+    // NOTE: when called by isDeviceKernelHostStub, Operand may not be a global


ZwFink · 2026-06-17T18:11:55Z

+  bool hasDeviceLaunchKernelCalls(Module &M) {
+    Function *LaunchKernelFn = nullptr;


ZwFink · 2026-06-17T18:12:08Z

+
+    DenseMap<StructType *, SmallVector<LambdaJitVariableInfo, 16>>
+        LambdaStorageTypeToJitIndices;
+    DenseMap<Type *, GlobalVariable *> LambdaTypeToGlobalName;


ZwFink · 2026-06-17T18:41:19Z

+    DenseMap<StructType *, SmallVector<LambdaJitVariableInfo, 16>>
+        LambdaStorageTypeToJitIndices;
+    DenseMap<Type *, GlobalVariable *> LambdaTypeToGlobalName;
+    DenseSet<StructType *> RegisteredLambdaStorageClasses;
+    DenseMap<StructType *, StructType *> FunctorTypeToLambdaTypeMap;
+    DenseMap<uint64_t, StructType *> FunctorIDToLambdaTypeMap;
+    DenseMap<uint64_t, StructType *> FunctorIDToFunctorTypeMap;


These are all declared here but used only for host compilation?

ZwFink · 2026-07-06T17:44:11Z

+    KernelFunction = FirstAnalysis->KernelFunction;
+    AnalysisSuccess = true;
+    AnalysisFailed = false;


Don't you also need to set KernelArg, Offset, and ChangedRC?

ZwFink · 2026-07-06T17:50:40Z

+    for (User *U : F->users()) {
+      auto *CB = dyn_cast<CallBase>(U);
+      if (!CB)
+        continue;
+      DEBUG(Logger::logs("proteus-pass")
+            << "Analysis crossed interprocedural boundary at "
+            << *CB->getArgOperand(ArgNum) << "\n");
+      WorkList.push_back({CB->getArgOperand(ArgNum), CB});


Does this need to be context sensitive?

ZwFink · 2026-07-10T17:37:45Z

+  std::optional<RuntimeConstantType> getRCTypeForLLVMType(Type *Ty) {
+    if (Ty->isIntegerTy(1))
+      return RuntimeConstantType::BOOL;
+    if (Ty->isIntegerTy(8))
+      return RuntimeConstantType::INT8;
+    if (Ty->isIntegerTy(32))
+      return RuntimeConstantType::INT32;
+    if (Ty->isIntegerTy(64))
+      return RuntimeConstantType::INT64;
+    if (Ty->isFloatTy())
+      return RuntimeConstantType::FLOAT;
+    if (Ty->isDoubleTy())
+      return RuntimeConstantType::DOUBLE;
+    if (Ty->isPointerTy())
+      return RuntimeConstantType::PTR;
+    return std::nullopt;
+  }


Regression: long double is dropped

ZwFink · 2026-07-10T17:39:36Z

 // Legacy PM implementation.
 struct LegacyProteusPass : public ModulePass {
  static char ID;
  LegacyProteusPass() : ModulePass(ID) {}
  bool runOnModule(Module &M) override {
    ProteusPassImpl PPI{M};
    bool Changed = PPI.run(M, false);
    return Changed;


Should the lambda analysis be run here too?

ggeorgakoudis reviewed Apr 15, 2026

View reviewed changes

johnbowen42 force-pushed the feature/bowen/lambda-analysis-redesign branch from 2cf03a2 to d5b97c4 Compare April 17, 2026 22:16

johnbowen42 marked this pull request as ready for review April 17, 2026 22:16

johnbowen42 force-pushed the feature/bowen/lambda-analysis-redesign branch from d5b97c4 to a56679a Compare April 17, 2026 22:17

johnbowen42 force-pushed the feature/bowen/lambda-analysis-redesign branch from c3d3d2e to 5daef48 Compare May 11, 2026 21:30

johnbowen42 requested review from ZwFink and davidbeckingsale May 12, 2026 21:49

johnbowen42 force-pushed the feature/bowen/lambda-analysis-redesign branch from 4441506 to 3202250 Compare June 18, 2026 19:00

johnbowen42 mentioned this pull request Jul 2, 2026

draft: Proteus lambda redesign RAJA-side support llnl/RAJA#2052

Draft

johnbowen42 force-pushed the feature/bowen/lambda-analysis-redesign branch from 8e43319 to f0ed85a Compare July 2, 2026 20:39

johnbowen42 requested a review from ggeorgakoudis July 6, 2026 16:22

ZwFink mentioned this pull request Jul 6, 2026

Fixup: Ci/matrix integration testing #503

Merged

johnbowen42 force-pushed the feature/bowen/lambda-analysis-redesign branch 2 times, most recently from f3e5a08 to f59b51e Compare July 8, 2026 16:08

johnbowen42 and others added 13 commits July 9, 2026 09:51

add lambda analysis pass

2f649ad

Cleanup after rebase, rename __jit_ to __proteus_

d5a7877

Fix unused function warnings

2eb914d

fix LLVM 20 build

e9c5169

fix unused function

c34f071

Fix unused variable compiler warnings and includes

3c5ecae

Fix lambda analysis failure

f077bb6

add runtime device side dispatch to lambda IR transformation

ad064c3

Get lambda factory test passing

9192287

add visitor for kernel arg analysis

12a2245

tmp

91a00c6

Add kernel args visitor. Get Lambda factory test working

2b1d610

johnbowen42 and others added 28 commits July 9, 2026 09:51

tmp: fix regressions discovered in CEED benchmark. Not clean commit: …

cc09857

…need to rework host lambda specialization

Rework host side jit specialization to reduce runtime overhead

efdf276

Fix all host-side unit tests after host lambda specialization redesign

e81d93e

Fix device specialization, move all kernel reading to C instrumentation

732ebd2

Fix cuda build, remove unused variables and methods

b9bf402

Fix CUDA tests

933db71

Fix test

c3d7680

fix deprecation warnings

1638086

remove storage cache checks from ptrtoint

e5a3eef

Fix test

460f9c9

Fix bug where GPU lambdas without jit_variables causes errors through…

00e6dd2

… expecting jit_variable schema for every lambda

Trim lambda llvm.global.annotations to avoid warnings

9bd5a91

init commit for handling RAJA callbases in memory use analysis

4bab577

tmp

3d23a62

Init indirect launch with lambdas support

ddc2935

Add new indirect kernel launch support for kernels with lambdas

1712340

initial SSA test

05be7f6

Initial support for RAJA IR shapes

5eb76b3

Add support for thread privatizer methods in RAJA by adding extract_v…

09fe556

…alue and insert_value analysis

add thread safety to lambda registry

c0ccc0b

Remove unused fields from pointer use analysis

da7138d

fix use of hiplaunchkernel instead of generic

07f3c08

fix matrix build

c49d196

Remove separate lambda pass. Register lambda analysis at pipeline sta…

35fd58e

…rt through ProteusPass

clang tidy

a72fc1e

clang format

a752e5e

Clang format

69af86e

Clang format

c462791

johnbowen42 force-pushed the feature/bowen/lambda-analysis-redesign branch from f59b51e to c462791 Compare July 9, 2026 16:53

ZwFink requested changes Jul 10, 2026

View reviewed changes

		DenseMap<Value , GlobalVariable > getKernelHostStubs(Module &M) {
		DenseMap<Value , GlobalVariable > StubToKernelMap;

		Value getStubGV([[maybe_unused]] Value Operand) {
		// NOTE: when called by isDeviceKernelHostStub, Operand may not be a global

		bool hasDeviceLaunchKernelCalls(Module &M) {
		Function *LaunchKernelFn = nullptr;

Uh oh!

Conversation

johnbowen42 commented Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ggeorgakoudis left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

johnbowen42 commented Apr 17, 2026

Uh oh!

johnbowen42 commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ZwFink left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

johnbowen42 commented Apr 14, 2026 •

edited

Loading

johnbowen42 commented Jul 2, 2026 •

edited

Loading

codecov Bot commented Jul 8, 2026 •

edited

Loading

ZwFink left a comment •

edited

Loading