Thread loop optimizations RAJA launch #1949

artv3 · 2025-11-25T14:02:33Z

This PR is a collaboration space for exploring optimization within RAJA launch and the loop abstraction

include/RAJA/pattern/launch/launch_core.hpp

include/RAJA/policy/hip/launch.hpp

artv3 · 2025-12-02T17:13:47Z

Hi @LLNL/raja-core, in collaboration with AMD staff we found a key optimization for nested loops within RAJA::launch.
This would be nice to have for the upcoming release, the downside is that it introduces another set of policies similar to the existing RAJA::hip_thread_loop_{x,y,z} policies. One thought is to retire the old policies in favor of these new ones, but for performance tracking are the old ones worth keeping?

MrBurmark · 2025-12-02T17:32:31Z

I wonder if we can't use the same policies but with an extra template argument to choose between the global and context versions of the variables.

rhornung67 · 2025-12-02T17:34:41Z

I don't think it's a problem to keep the old policies and have a lot of alternatives for folks to try. We do need to work on documenting policies better and have a comprehensive cookbook of examples that clearly show the differences between policies choices, including usage, performance, and how to choose.

MrBurmark · 2025-12-02T17:39:26Z

Do we want to make populating the context variables optional in case we find any overhead there?

MrBurmark · 2025-12-02T17:40:06Z

Any thoughts on blockIdx and gridDim?

artv3 · 2025-12-02T19:12:53Z

Do we want to make populating the context variables optional in case we find any overhead there?

oh for register heavy kernels? that makes sense

artv3 · 2025-12-02T19:14:45Z

Any thoughts on blockIdx and gridDim?

I see pro and cons, pro - for completeness could be handy, con - takes up more registers. Maybe we can do partial specializations of the launch context or something like that. This may be less common use cases though.

MrBurmark · 2025-12-02T20:01:38Z

Ya, I'm imagining the context and some policies both having a switch. Then in the loop implementation it checks that if the policy uses the switch then the context must have the same switch.

template < bool switch >
struct Policy;

template < bool switch >
struct Context;

template < bool policy_switch, bool context_switch >
void loop(Policy<policy_switch>, Context<context_switch>)
{
  static_assert(!policy_switch || (policy_switch  && context_switch),
                "If policy has switch then context must have switch");
}

…into artv3/launch-loop-opt

include/RAJA/pattern/launch/launch_core.hpp

include/RAJA/policy/cuda/launch.hpp

MrBurmark · 2025-12-11T16:37:40Z

include/RAJA/policy/cuda/launch.hpp

+    for (int i = ::RAJA::internal::CudaDimHelper<DIM>::get(ctx.thread_id);
+         i < len; i += ::RAJA::internal::CudaDimHelper<DIM>::get(ctx.block_dim))


If constexpr to get the values based on StoreDim3?

Can you share an example?

See line 165 above if constexpr (LaunchContextT<LaunchContextPolicy>::hasDim3)

I wanted to echo that I think this is a good idea

Not sure if possible though because ctx is a function parameter

We were able to get the argument type if there is one operator() that is not templated and has at least one argument. If that isn't true then it will use the default.

Co-authored-by: Jason Burmark <[email protected]>

include/RAJA/policy/sequential/policy.hpp

artv3 · 2025-12-18T16:08:54Z

@MrBurmark , do you have time to take a look? I think I pushed up the ideas we had yesterday

include/RAJA/pattern/launch/launch_context_policy.hpp

include/RAJA/pattern/launch/launch_core.hpp

…v3/launch-loop-opt

include/RAJA/policy/cuda/launch.hpp

include/RAJA/policy/hip/launch.hpp

include/RAJA/pattern/launch/launch_context_policy.hpp

johnbowen42 · 2025-12-18T22:49:15Z

include/RAJA/pattern/launch/launch_context_policy.hpp

+//is pointer a type just pass it through otherwise do give me the operator
+//static error if not a function type
+
+//template <typename Lambda>


looks like there might be some cruft here, we should delete before merging

johnbowen42 · 2025-12-18T22:50:01Z

include/RAJA/pattern/launch/launch_context_policy.hpp

+//  - If T is a function pointer, use function_traits<T> directly
+//  - Otherwise, assume it is a callable object and use &T::operator()
+template <typename T>
+using lambda_traits =


johnbowen42 · 2025-12-18T22:56:04Z

include/RAJA/policy/cuda/launch.hpp

+    for (int i = ::RAJA::internal::CudaDimHelper<DIM>::get(ctx.thread_id);
+         i < len; i += ::RAJA::internal::CudaDimHelper<DIM>::get(ctx.block_dim))


I wanted to echo that I think this is a good idea

johnbowen42 · 2025-12-18T22:56:40Z

include/RAJA/policy/cuda/launch.hpp

+    for (int i = ::RAJA::internal::CudaDimHelper<DIM>::get(ctx.thread_id);
+         i < len; i += ::RAJA::internal::CudaDimHelper<DIM>::get(ctx.block_dim))


Not sure if possible though because ctx is a function parameter

johnbowen42 · 2025-12-18T23:01:08Z

include/RAJA/policy/cuda/launch.hpp

-  ctx.shared_mem_ptr = raja_shmem_ptr;

-  RAJA::expt::invoke_body(reduce_params, body, ctx);
+  if constexpr (LaunchContextT<LaunchContextPolicy>::hasDim3)


I think we need to use your new type traits here as well? similar to the HIP backend

include/RAJA/policy/hip/launch.hpp

artv3 · 2025-12-19T10:56:38Z

include/RAJA/policy/hip/launch.hpp

+template<named_dim DIM, typename Dim3Like>
+RAJA_INLINE RAJA_DEVICE int get_dim(Dim3Like const& d)
+{
+  if constexpr (DIM == named_dim::x)


@MrBurmark @johnbowen42 , is this the usage of if constexpr you all had in mind?

No, I was thinking about doing if constexcpr (LaunchContextType::hasDim3) in the exec functions instead of adding new policies. Also there is already a function to get the named dims from a dim3.

ohh.. that makes sense now

artv3 added 2 commits November 25, 2025 05:59

initial commit for launch loop optimization

cd5065e

add structs to store gpu thread/block info in launch ctx

484ff1a

MrBurmark reviewed Nov 25, 2025

View reviewed changes

include/RAJA/pattern/launch/launch_core.hpp Outdated Show resolved Hide resolved

artv3 commented Nov 25, 2025

View reviewed changes

include/RAJA/policy/hip/launch.hpp Outdated Show resolved Hide resolved

artv3 added 2 commits December 2, 2025 09:05

add cuda variant and add build guards for cpu

18f332b

Merge branch 'develop' into artv3/launch-loop-opt

21f6184

artv3 marked this pull request as ready for review December 2, 2025 17:09

artv3 requested a review from a team December 2, 2025 17:09

artv3 changed the title ~~Exploring loop optimizations RAJA launch~~ Thread loop optimizations RAJA launch Dec 2, 2025

rhornung67 added this to the Dec 2025 Release milestone Dec 2, 2025

artv3 added 4 commits December 11, 2025 05:36

rework to support dim3 copy in ctx

73f224a

Merge branch 'artv3/launch-loop-opt' of https://github.com/LLNL/RAJA …

8a02fee

…into artv3/launch-loop-opt

minor clean up pass

1fbe50b

make format

672889e

artv3 commented Dec 11, 2025

View reviewed changes

include/RAJA/pattern/launch/launch_core.hpp Outdated Show resolved Hide resolved

MrBurmark reviewed Dec 11, 2025

View reviewed changes

include/RAJA/pattern/launch/launch_core.hpp Outdated Show resolved Hide resolved

MrBurmark reviewed Dec 11, 2025

View reviewed changes

include/RAJA/pattern/launch/launch_core.hpp Outdated Show resolved Hide resolved

MrBurmark reviewed Dec 11, 2025

View reviewed changes

include/RAJA/policy/cuda/launch.hpp Outdated Show resolved Hide resolved

MrBurmark reviewed Dec 11, 2025

View reviewed changes

artv3 and others added 3 commits December 11, 2025 09:05

Update include/RAJA/pattern/launch/launch_core.hpp

5908a20

Co-authored-by: Jason Burmark <[email protected]>

Merge branch 'develop' into artv3/launch-loop-opt

316e019

clean up pass

4d9f800

artv3 added 3 commits December 18, 2025 07:50

omp backend

d21c41f

seq backend + make style

40a5c1b

clean up pass

e0f4825

artv3 commented Dec 18, 2025

View reviewed changes

include/RAJA/policy/sequential/policy.hpp Outdated Show resolved Hide resolved

artv3 commented Dec 18, 2025

View reviewed changes

include/RAJA/pattern/launch/launch_context_policy.hpp Outdated Show resolved Hide resolved

Update include/RAJA/pattern/launch/launch_context_policy.hpp

96e99d5

MrBurmark reviewed Dec 18, 2025

View reviewed changes

include/RAJA/pattern/launch/launch_core.hpp Outdated Show resolved Hide resolved

MrBurmark reviewed Dec 18, 2025

View reviewed changes

include/RAJA/pattern/launch/launch_core.hpp Outdated Show resolved Hide resolved

artv3 and others added 4 commits December 18, 2025 08:53

minor clean up

a9f0cca

minor clean up

7d4595b

Merge branch 'artv3/launch-loop-opt' of github.com:LLNL/RAJA into art…

c23f76f

…v3/launch-loop-opt

revert changes to example

c990a4f

artv3 requested a review from a team December 18, 2025 18:40

MrBurmark reviewed Dec 18, 2025

View reviewed changes

include/RAJA/policy/cuda/launch.hpp Show resolved Hide resolved

remove specialization from launch policy

f7939fd

artv3 commented Dec 18, 2025

View reviewed changes

include/RAJA/policy/hip/launch.hpp Outdated Show resolved Hide resolved

MrBurmark reviewed Dec 18, 2025

View reviewed changes

include/RAJA/pattern/launch/launch_context_policy.hpp Outdated Show resolved Hide resolved

make work for function pointers

c24331c

johnbowen42 reviewed Dec 18, 2025

View reviewed changes

artv3 added 9 commits December 19, 2025 01:02

store dim3 based on launch context type - hip

0518138

rework omp backend

d5da29a

update sequential backend

af88dbb

get things building for cuda -- need a good clean up pass

21ad0a8

cuda clean up pass

646a95b

clean up ordering in hip launch

597641b

clean up ordering

5403737

make style

e41e970

use constexpt for getting dim values

7c95430

artv3 commented Dec 19, 2025

View reviewed changes

		for (int i = ::RAJA::internal::CudaDimHelper<DIM>::get(ctx.thread_id);
		i < len; i += ::RAJA::internal::CudaDimHelper<DIM>::get(ctx.block_dim))

Thread loop optimizations RAJA launch #1949

Are you sure you want to change the base?

Thread loop optimizations RAJA launch #1949

Conversation

artv3 commented Nov 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

artv3 commented Dec 2, 2025

Uh oh!

MrBurmark commented Dec 2, 2025

Uh oh!

rhornung67 commented Dec 2, 2025

Uh oh!

MrBurmark commented Dec 2, 2025

Uh oh!

MrBurmark commented Dec 2, 2025

Uh oh!

artv3 commented Dec 2, 2025

Uh oh!

artv3 commented Dec 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MrBurmark commented Dec 2, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

artv3 commented Dec 18, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MrBurmark Dec 19, 2025 • 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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

artv3 commented Nov 25, 2025 •

edited

Loading

artv3 commented Dec 2, 2025 •

edited

Loading

MrBurmark Dec 19, 2025 •

edited

Loading