Skip to content

draft: Proteus lambda redesign RAJA-side support#2052

Draft
johnbowen42 wants to merge 5 commits into
developfrom
feature/bowen/proteus-lambda-redesign-support
Draft

draft: Proteus lambda redesign RAJA-side support#2052
johnbowen42 wants to merge 5 commits into
developfrom
feature/bowen/proteus-lambda-redesign-support

Conversation

@johnbowen42

@johnbowen42 johnbowen42 commented Jul 2, 2026

Copy link
Copy Markdown
Contributor

We are redesigning lambda support to be more robust here Olympus-HPC/proteus#467, this PR changes how RAJA uses proteus to support the new, slightly changed interface. This PR also adds new examples of JIT compilation with launch and kernel, and also add tests to guarantee that runtime values are being turned into constants at JIT time

Comment thread examples/kernel-jit.cpp

auto Seg = RAJA::RangeSegment(0, N);

RAJA::kernel<kernel_policy>(RAJA::make_tuple(Seg), [=](int i) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

host device?

typename Resource,
typename... Bodies>
inline auto register_loop_bodies(
LoopData<SegmentTuple, ParamTuple, Resource, Bodies...> body)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

reference?

typename ParamTuple,
typename Resource,
typename... Bodies>
inline auto register_loop_bodies(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we want this to always return by value or return by reference without jit enabled?

Comment on lines +685 to +688
using registered_data_t = std::decay_t<decltype(registered_bodies)>;
using registered_launch_t =
CudaLaunchHelper<LaunchConfig, stmt_list_t, registered_data_t, Types>;
auto registered_function = registered_launch_t::get_func();

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a lot more code than it was.

@johnbowen42 johnbowen42 Jul 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah, I am debating simplifying this but shifitng the decision to JIT the lambda to the user. Basically, if a user wanted to JIT their lambda they could wrap it in RAJA::register_lambda before passing to kernel/launch/forall, and then we could get rid of all this boilerplate

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you think?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants