RAJA::fornest : Abstraction to switch between collapse, nested loops, and tiling patterns 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-08-13T22:47:36Z

+
+The interface supports two mapping policy families:
+
+ * ``RAJA::fornest_flattened_policy<ExecPolicy, LayoutTag>`` maps the product


I thought we were going to call this collapse instead of a flatten as this is a multi-dimensional index space being mapped to a 1-D exec space?

Absolutely, we can call it a collapse. Thanks for catching that, I sometimes use the application terminology out of habit.

MrBurmark · 2026-08-14T15:12:25Z

+``RAJA::fornest`` runs a logical 2-D or 3-D loop body through selectable
+mappings. It is intended for kernels where the source code should stay written
+in logical multi-dimensional indices, but the best GPU mapping is not known from
+the source alone.


Should this dwell on the current support (2d and 3d) or be general and then note that currently only 2d and 3d are supported. The second sentence is a bit awkward and specific to a certain GPU use case when this API is intended to be fairly general for kernels with perfectly nested loops. Maybe something like "Use this for kernels with multi-dimensional indexing that can be expressed sequentially with perfectly nested loops". The "GPU mapping" statement seems like something too specific for the intro.

MrBurmark · 2026-08-14T15:13:06Z

+The interface supports two mapping policy families:
+
+ * ``RAJA::fornest_flattened_policy<ExecPolicy, LayoutTag>`` maps the product
+   of the logical dimensions to a 1-D iteration space. The ``ExecPolicy`` is a
+   regular forall-style policy such as ``RAJA::device_exec<256>`` (CUDA/HIP) or
+   ``RAJA::seq_exec`` (host). RAJA performs the linear-to-logical index
+   reconstruction internally.
+ * ``RAJA::fornest_mapping_policy<ExecPolicy, LoopPolicies...>`` maps the logical
+   dimensions directly using one loop policy per dimension. For CUDA/HIP this
+   allows explicit mapping to global, block, or thread spaces.


Regarding the naming and structure of these policies, are these aliases to a single underlying policy template? I am imagining something like RAJA::fornest_policy<OuterExecPolicy, InnerPolicyList=DefaultList> where in some cases an outer policy can have enough information for a basic case without specifying the inner policies.

artv3 · 2026-08-18T21:12:04Z

@MrBurmark @bechols97 , if you have some time can you take a look at the fornest-basic.cpp and dynamic-fornest.cpp ? I want to double check I captured the desired capabilities, then we can start wrapping this up.

MrBurmark · 2026-08-19T15:53:10Z

+
+The interface supports several policy families:
+
+ * ``RAJA::fornest_collapsed_policy<ExecPolicy, LayoutTag>`` maps the product


What does the LayoutTag do? Allow reordering essentially? Does it default to layout_right?

Also why collapsed with a d policy and not just collapse?

Do we still have permutation lists that we could imagine taking as the LayoutTag that could be used to reorder the loops?

MrBurmark · 2026-08-19T15:59:49Z

+   of the logical dimensions to a 1-D iteration space. The ``ExecPolicy`` is a
+   regular forall-style policy such as ``RAJA::seq_exec`` or
+   ``RAJA::device_exec<256>``. RAJA performs the linear-to-logical index
+   reconstruction internally. On host OpenMP builds, ``RAJA::fornest_omp_collapse_policy<...>``


The fornest_omp_collapse_policy policy is here to allow directly using omp collapse vs doing our own collapse with and using something like omp_parallel_for which can be done via fornest_collapsed_policy<omp_parallel_for, ...>.

MrBurmark · 2026-08-19T16:01:55Z

+   segment extents. Device builds can use ``RAJA::device_*`` mapping tags; on
+   CUDA/HIP these map to CUDA/HIP global, block, and thread indices, and on SYCL
+   they map to the corresponding work-group and work-item indices where
+   supported.


My wouldn't it be more consistent to use policies like RAJA::seq_exec, RAJA::device_exec<256>, etc as the ExecPolicy?

MrBurmark · 2026-08-19T16:04:00Z

+ * ``RAJA::dynamic_fornest<camp::list<...>>(pol, ...)`` selects a policy from a
+   list at runtime.


Suggested change

* ``RAJA::dynamic_fornest<camp::list<...>>(pol, ...)`` selects a policy from a

list at runtime.

* ``RAJA::dynamic_fornest<camp::list<...>>(policy_index, ...)`` selects a policy from a

list at runtime.

Would this make it more clear that pol is mean to be an index into the list of policies?

MrBurmark · 2026-08-19T16:04:53Z

+ * ``RAJA::fornest_tiling_policy<...>`` adds per-dimension tiling on top of a
+   loop nest. Tile sizes may be fixed, runtime-provided, or chosen by RAJA.


Suggested change

* ``RAJA::fornest_tiling_policy<...>`` adds per-dimension tiling on top of a

loop nest. Tile sizes may be fixed, runtime-provided, or chosen by RAJA.

* ``RAJA::fornest_tiling_policy<ExecPolicy, ...>`` adds per-dimension tiling on top of a

loop nest. Tile sizes may be fixed, runtime-provided, or chosen by RAJA.

Don't we still want to have an ExecPolicy so we can do things like apply launch bounds?

MrBurmark · 2026-08-19T16:14:42Z

+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+
+#ifndef RAJA_pattern_fornest_HPP
+#define RAJA_pattern_fornest_HPP


I think we want to split this file up into multiple files.

MrBurmark · 2026-08-19T16:16:57Z

+{
+
+template<typename Policy>
+struct policy_dimensionality : std::integral_constant<camp::idx_t, 0>


If we write the policy we can have a static constexpr variable to make it easy to query the number of dimensions instead of writing a separate traits class.

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

artv3 added 4 commits August 12, 2026 08:44

clean up pass

d4e3e92

increase requested time on matrix

3c41d03

clean up pass

29a715a

make style

71c2bcf

MrBurmark reviewed Aug 13, 2026

View reviewed changes

flatten -> collapse

cfe619e

MrBurmark reviewed Aug 14, 2026

View reviewed changes

Comment thread include/RAJA/pattern/fornest.hpp Outdated

MrBurmark reviewed Aug 14, 2026

View reviewed changes

Comment thread include/RAJA/pattern/fornest.hpp Outdated

artv3 changed the title ~~RAJA::fornest : Abstraction to switch between flat and nested loops~~ RAJA::fornest : Abstraction to switch between collapse and nested loops Aug 14, 2026

artv3 added 3 commits August 17, 2026 21:20

update based on wishlist

8a10756

clean up pass

99f4ab9

build fix for sycl

1113732

artv3 changed the title ~~RAJA::fornest : Abstraction to switch between collapse and nested loops~~ RAJA::fornest : Abstraction to switch between collapse, nested loops, and tiling patterns Aug 18, 2026

clean up pass

5dfc51e

artv3 marked this pull request as draft August 18, 2026 18:23

artv3 added 4 commits August 18, 2026 12:35

clean up pass

ca6614f

clean up pass

36d3e07

clean up pass

e42d674

clean up pass

1904d1b

artv3 marked this pull request as ready for review August 18, 2026 21:10

MrBurmark reviewed Aug 19, 2026

View reviewed changes

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


		The interface supports two mapping policy families:

		* ``RAJA::fornest_flattened_policy<ExecPolicy, LayoutTag>`` maps the product


		The interface supports several policy families:

		* ``RAJA::fornest_collapsed_policy<ExecPolicy, LayoutTag>`` maps the product

		* ``RAJA::dynamic_fornest<camp::list<...>>(pol, ...)`` selects a policy from a
		list at runtime.

		* ``RAJA::fornest_tiling_policy<...>`` adds per-dimension tiling on top of a
		loop nest. Tile sizes may be fixed, runtime-provided, or chosen by RAJA.

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

artv3 commented Aug 18, 2026

Uh oh!

MrBurmark Aug 19, 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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

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 Aug 19, 2026 •

edited

Loading