Launch_nd: Abstraction to switch between flat and nested loops by artv3 · Pull Request #2021 · llnl/RAJA

artv3 · 2026-04-30T18:20:24Z

I along with @tomstitt work an application where we have been exploring using a 1D GPU index for multi-level loops and have found that may be more performant than using hierarchical parallelism. This PR uses concepts from RAJA to create the RAJA::forall_nd convenience function.

Cases in which this approach is more performant:
it comes up when the threads per block are not high enough to saturate the GPU and using RAJA::forall + mods + divs allows us to increase the threads per block which ends up being more performant.

MrBurmark · 2026-04-30T19:45:05Z

I am concerned that people will use this without knowing that its doing expensive div/mod calculations and see that it performs more poorly than a native cuda/hip 2d or 3d kernel. Would it make sense to have policies that allow mapping to 2d or 3d kernels.

artv3 · 2026-04-30T20:39:49Z

I am concerned that people will use this without knowing that its doing expensive div/mod calculations and see that it performs more poorly than a native cuda/hip 2d or 3d kernel. Would it make sense to have policies that allow mapping to 2d or 3d kernels.

I think that can be addressed in the RAJA cookbook or examples explaining when this would be performant. In the case of @tomstitt and I, it comes up when the threads per block are not high enough to saturate the GPU and using RAJA::forall + mods + divs allows us to increase the threads per block which ends up being more performant. For 2D/3D gpu grids there are various ways to do that and perhaps we should direct developers to RAJA::launch or RAJA::kernel?

tomstitt · 2026-04-30T21:17:19Z

I am concerned that people will use this without knowing that its doing expensive div/mod calculations and see that it performs more poorly than a native cuda/hip 2d or 3d kernel. Would it make sense to have policies that allow mapping to 2d or 3d kernels.

I think that can be addressed in the RAJA cookbook or examples explaining when this would be performant. In the case of @tomstitt and I, it comes up when the threads per block are not high enough to saturate the GPU and using RAJA::forall + mods + divs allows us to increase the threads per block which ends up being more performant. For 2D/3D gpu grids there are various ways to do that and perhaps we should direct developers to RAJA::launch or RAJA::kernel?

I think my ideal is an interface where there is a choice of policy. We have 2d/3d kernels on 1d iteration spaces, using mod/div like Arturo said, to expose more parallelism. When we switch some of those to using our "true" 2d/3d grid launcher we lose performance because our block (16x16 , 8x8x8) doesn't map well onto the grid (because we just idle threads). It's of course not always true that our div/mod approach is going to be better, like Jason said, and if we had an easy way to pick we could put both behind our abstraction and correctly dispatch

artv3 · 2026-04-30T21:24:45Z

I am concerned that people will use this without knowing that its doing expensive div/mod calculations and see that it performs more poorly than a native cuda/hip 2d or 3d kernel. Would it make sense to have policies that allow mapping to 2d or 3d kernels.

I think that can be addressed in the RAJA cookbook or examples explaining when this would be performant. In the case of @tomstitt and I, it comes up when the threads per block are not high enough to saturate the GPU and using RAJA::forall + mods + divs allows us to increase the threads per block which ends up being more performant. For 2D/3D gpu grids there are various ways to do that and perhaps we should direct developers to RAJA::launch or RAJA::kernel?

I think my ideal is an interface where there is a choice of policy. We have 2d/3d kernels on 1d iteration spaces, using mod/div like Arturo said, to expose more parallelism. When we switch some of those to using our "true" 2d/3d grid launcher we lose performance because our block (16x16 , 8x8x8) doesn't map well onto the grid (because we just idle threads). It's of course not always true that our div/mod approach is going to be better, like Jason said, and if we had an easy way to pick we could put both behind our abstraction and correctly dispatch

@tomstitt , okay -- Let's get that imagination soaring, and cook up some ideas! #RAJA!

artv3 · 2026-04-30T23:04:34Z

+  RAJA::launch_nd(res, policy, RAJA::segments(cells, comps),
+                  [=] RAJA_HOST_DEVICE(int cell, int comp) {
+                    const int idx   = comp + num_comp * cell;
+                    values_ptr[idx] = 1000 * cell + comp;
+                  });


@tomstitt , @MrBurmark I invite you to take a look at this example, I think this may be what we are looking for.

MrBurmark · 2026-04-30T23:31:36Z

I am concerned that people will use this without knowing that its doing expensive div/mod calculations and see that it performs more poorly than a native cuda/hip 2d or 3d kernel. Would it make sense to have policies that allow mapping to 2d or 3d kernels.

I think that can be addressed in the RAJA cookbook or examples explaining when this would be performant. In the case of @tomstitt and I, it comes up when the threads per block are not high enough to saturate the GPU and using RAJA::forall + mods + divs allows us to increase the threads per block which ends up being more performant. For 2D/3D gpu grids there are various ways to do that and perhaps we should direct developers to RAJA::launch or RAJA::kernel?

I think my ideal is an interface where there is a choice of policy. We have 2d/3d kernels on 1d iteration spaces, using mod/div like Arturo said, to expose more parallelism. When we switch some of those to using our "true" 2d/3d grid launcher we lose performance because our block (16x16 , 8x8x8) doesn't map well onto the grid (because we just idle threads). It's of course not always true that our div/mod approach is going to be better, like Jason said, and if we had an easy way to pick we could put both behind our abstraction and correctly dispatch

@tomstitt , okay -- Let's get that imagination soaring, and cook up some ideas! #RAJA!

You certainly can use launch or teams to get some level of parallelism and then take those indices and do your own calculations with them. If we had the multiloop abstractions with a variety of policies that could take us a fair amount of the way.

artv3 · 2026-07-06T16:11:50Z

@llnl/raja-core , I received an approval from my project. Opening up the PR for review from the squad.

adrienbernede

@lc-hubcast approve

rhornung67 · 2026-07-08T15:36:19Z

@artv3 I merged in latest develop and ran clang-format so GHA CI check run. Is this ready to be merged when approved?

artv3 · 2026-07-08T15:57:25Z

@artv3 I merged in latest develop and ran clang-format so GHA CI check run. Is this ready to be merged when approved?

Not yet, I think we want @llnl/raja-core to take another look. We had a discussed a while potentially changing the name of the abstraction

MrBurmark · 2026-07-29T16:28:04Z

+stay written in logical multi-dimensional indices, but the best GPU mapping is
+not known from the source alone.


The GPU is part of this, but wouldn't this also be useful for doing things like collapsed openmp loops, or reordered sequential loops?

Underneath we just map to regular nested loops in launch, I think we would have to leave that for when we do the expanded version of this. One thing to explore is transitioning between flat loops to nested loops using threads in a GPU setting. We also would have to figure out the backend on how that would look like for OpenMP

I suppose it could, I haven't thought about that yet though actually, I think that could be a neat expansion. I think that goes beyond the targeted scope of this PR though, but something to definitely revisit as we expand capabilities. I would be interested in testing it some kernels in RAJAPerf

MrBurmark · 2026-07-29T16:41:23Z

+The example source defines backend-specific policy aliases. CUDA uses direct
+global loop policies for the true grid mapping:


It is confusing to have direct and loop in close proximity here.

good call, I updated the example

MrBurmark · 2026-07-29T16:42:08Z

+The flattened policy uses a forall-style execution policy such as
+``RAJA::cuda_exec<block_size_1d>`` or ``RAJA::hip_exec<block_size_1d>``. The
+grid policy uses ``RAJA::LaunchPolicy`` and direct global loop policies such as
+``RAJA::cuda_global_y_direct`` and ``RAJA::cuda_global_x_direct``.


Could you use RAJA::cuda_global_yx_direct?

If people want to do more complicated things like mixing direct and loop policies they could always use launch itself.

I agree actually, that would good but we would need to first update formalize support “multi-dim per-loop” policies. Something I've played around with but have not yet formalized

MrBurmark · 2026-07-29T16:49:42Z

+``RAJA::launch_nd`` currently supports ``RAJA::TypedRangeSegment`` packs created
+with ``RAJA::segments``. The grid mapping supports 2-D and 3-D loops and
+requires one loop policy per segment. The flattened mapping uses
+``RAJA::layout_right`` by default, or ``RAJA::layout_left`` when that layout tag
+is supplied, to control which logical index is unit stride in the flattened
+space.


Why not just a tuple of segments?
The layouts in the policy are a way to reorder the loops?

In some sense it already is, struct TypedRangeSegmentPack { camp::tuple<RAJA::TypedRangeSegment<IdxTs>...> data; }; , the RAJA::Segments just enforces that we use pass in a RAJA::Segment type? Are you thinking of support for non-RAJA types?

Second question:
For the flatten policy it controls the linearization order (which logical index
varies fastest / is unit-stride)

For the grid policy:
it’s determined by the segment + loop-
policy order (one policy per segment)

rchen20 · 2026-07-30T23:03:43Z

+#if defined(RAJA_GPU_ACTIVE)
+template<typename PolicyList>
+RAJA_INLINE auto make_launch_nd_context(RAJA::resources::Resource resource,
+                                        std::string&& kernel_name)
+{
+  return resource.get_platform() == RAJA::Platform::host
+             ? util::make_context<typename PolicyList::host_policy_t>(
+                   std::move(kernel_name))


If we are checking for RAJA::Platform::host here, does this block need to be in RAJA_GPU_ACTIVE?

yes, because in GPU builds we can still dispatch to the CPU at run time.

rchen20 · 2026-08-05T23:50:02Z

+{
+
+template<typename... IdxTs>
+struct TypedRangeSegmentPack


Do we want TypedRangeSegmentPack to be in the outward facing RAJA API?

good, point that should be internal, since it is specific to this API

rchen20 · 2026-08-05T23:51:09Z

+};
+
+template<typename... IdxTs>
+RAJA_INLINE auto segments(RAJA::TypedRangeSegment<IdxTs> const&... segs)


Should we rename RAJA::segments to something else closer to its launch_nd-only usage? Or hide it better in a deeper namespace so that people don't accidentally use segments?

Agreed, renamed to RAJA::nd_segments(...) and updated docs/tests/examples accordingly

rchen20 · 2026-08-06T00:02:57Z

+    case ExecPlace::HOST:
+      return RAJA::launch_nd(
+          resource, launch_nd_flattened_policy<HostExecPolicy, LayoutTag> {},


Do we need to validate that the resource matches the ExecPlace?

yes, good call, added a helper to check

rchen20 · 2026-08-06T00:06:49Z

+template<typename HostExecPolicy,
+         typename DeviceExecPolicy,
+         typename LayoutTag = RAJA::layout_right>
+struct launch_nd_flattened_place_policy
+{
+  using host_exec_policy   = HostExecPolicy;
+  using device_exec_policy = DeviceExecPolicy;
+  using layout_tag         = LayoutTag;
+};
+
+/*!
+ * General runtime host/device launch_nd policy container.
+ *
+ * This allows callers to select different launch_nd policy *kinds* for host and
+ * device (e.g., a grid policy on host for nested loops and a flattened policy
+ * on device for a 1D mapping).
+ */
+template<typename HostPolicy, typename DevicePolicy>
+struct launch_nd_place_policy
+{
+  HostPolicy host;
+  DevicePolicy device;
+};


I know these serve two different overloads of launch_nd() below, but could these be combined for simplification? Maybe just have launch_nd_flattened_place_policy?

good idea, I reduced duplication by making launch_nd_flattened_place_policy just a convenience alias of the general launch_nd_place_policy, and removed the now-
redundant launch_nd(...) overloads.

artv3 · 2026-08-07T19:13:16Z

Thanks @rchen20 , I addressed your review comments

rchen20 · 2026-08-07T22:58:56Z

+{
+  switch (place)


We probably need validate_launch... here too. Would you mind checking where else we need to call it?

draft of forall_nd

8e84c53

artv3 requested a review from tomstitt April 30, 2026 18:20

artv3 added this to the June 2026 Release milestone Apr 30, 2026

draft of launch_nd

7e6c77a

artv3 changed the title ~~Draft: Forall_nd~~ Draft: launch_nd Apr 30, 2026

artv3 commented Apr 30, 2026

View reviewed changes

artv3 added 4 commits May 11, 2026 16:19

WIP benchmarking with RAJA launch_nd

65706f2

Merge branch 'develop' into artv3/feature/forall_nd

0a86357

minor changes for application usage

6349a4e

Merge branch 'develop' into artv3/feature/forall_nd

57870d4

artv3 requested a review from a team July 6, 2026 16:11

artv3 marked this pull request as ready for review July 6, 2026 16:11

Merge branch 'develop' into artv3/feature/forall_nd

fcdc683

artv3 changed the title ~~Draft: launch_nd~~ launch_nd Jul 6, 2026

artv3 changed the title ~~launch_nd~~ Launch_nd: Abstraction to switch between flat and nested loops Jul 6, 2026

adrienbernede reviewed Jul 7, 2026

View reviewed changes

rhornung67 added 2 commits July 8, 2026 08:29

Merge branch 'develop' into artv3/feature/forall_nd

0bbfaba

clang format

9329f4f

artv3 requested a review from a team July 14, 2026 16:22

artv3 added 2 commits July 14, 2026 09:23

Merge branch 'develop' into artv3/feature/forall_nd

ccd5925

Merge branch 'develop' into artv3/feature/forall_nd

47e3c8d

MrBurmark reviewed Jul 29, 2026

View reviewed changes

Merge branch 'develop' into artv3/feature/forall_nd

fc7f8de

rchen20 reviewed Jul 30, 2026

View reviewed changes

artv3 added 5 commits August 4, 2026 18:12

Merge branch 'develop' into artv3/feature/forall_nd

d719c1a

MR comments

b6a1ab0

clean up pass

6ddb464

fixes for cuda

b46b288

make style

f39d09e

rchen20 reviewed Aug 5, 2026

View reviewed changes

rchen20 reviewed Aug 6, 2026

View reviewed changes

rhornung67 modified the milestones: July 2026 Release, FY26 Development Aug 6, 2026

artv3 added 4 commits August 7, 2026 11:38

clean up pass

ffcbd82

clean up pass

463429d

clean up pass

e4a32bc

clean up pass

f644bb9

artv3 added 2 commits August 7, 2026 12:15

make style

17313a8

build fix for cuda

8c5b20d

rchen20 reviewed Aug 7, 2026

View reviewed changes

artv3 added 3 commits August 7, 2026 17:22

fix for nvcc

36ad2b7

simplify the internals

b9f872f

clean up pass

8829b2c

		stay written in logical multi-dimensional indices, but the best GPU mapping is
		not known from the source alone.

		The example source defines backend-specific policy aliases. CUDA uses direct
		global loop policies for the true grid mapping:

Uh oh!

Conversation

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

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

artv3 commented Apr 30, 2026

Uh oh!

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

Uh oh!

Uh oh!

artv3 commented Apr 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MrBurmark commented Apr 30, 2026

Uh oh!

artv3 commented Jul 6, 2026

Uh oh!

adrienbernede left a comment

Choose a reason for hiding this comment

Uh oh!

rhornung67 commented Jul 8, 2026

Uh oh!

artv3 commented Jul 8, 2026

Uh oh!

MrBurmark Jul 29, 2026 • 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!

MrBurmark Jul 29, 2026 • 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!

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!

artv3 commented Aug 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

artv3 commented Apr 30, 2026 •

edited

Loading

MrBurmark commented Apr 30, 2026 •

edited

Loading

tomstitt commented Apr 30, 2026 •

edited

Loading

MrBurmark Jul 29, 2026 •

edited

Loading

MrBurmark Jul 29, 2026 •

edited

Loading