Restore cuda_parallel_launch_constant_memory. by pkestene · Pull Request #9133 · kokkos/kokkos

pkestene · 2026-04-28T15:48:44Z

This is related to the discussion on slack about performance regression observed in kokkos >= 5 compared to 4.7.02 when using the Cuda backend.

We observed that for some cuda kernels the actual register count increased by more than 50, leading to both smaller device occupancy and reduced time to solution by about 30%. This performance regression is not due to the introduction of __grid_constant__ kernel argument annotation, but to the removal of the constant memory launch mechanism. In this PR we restore this lauch mechanism.

List of modifications :

restore the use of cuda_parallel_launch_constant_memory
using cuda_parallel_launch_constant_memory is a user choice through the use of ~~a new execution policy trait canlled ImplForceConstantLaunch (introduced has a member of WorkItemProperty)~~ HintHeavyWeight

Example of use:

#ifdef KOKKOS_ENABLE_CUDA
  using Property = Kokkos::Experimental::WorkItemProperty::HintHeavyWeight_t;
#else
  using Property = Kokkos::Experimental::WorkItemProperty::None_t;
#endif

  Kokkos::parallel_for("MyFunctor",
                       Kokkos::RangePolicy<exec_space, Property>(0, nbIterations),
                       functor);

Changelog Entry

crtrott

I don't think I want this this way. My current thinking is that the HeavyWeight hint is what we want to reuse. Essentially if its heavy weight go back to constant launch. We also need to discuss if we want to keep the "require" mechanism, or simply add these as template parameters directly to the policy - that is an orthogonal discussion though.

pkestene · 2026-04-29T19:06:56Z

You mean re-use the HeavyWeight hint instead of introducing a new WorkItemProperty ?
I thought the name ForceConstantLaunch was more meaningful, as a user might want to opt-in to constant memory launch as a fallback to avoid a potential bug from the compiler (when using the grid_constant annotations) or for any other reason (just testing or comparing performance).

I can change and reuse HeavyWeight if you prefer.

Signed-off-by: Pierre Kestener <pierre.kestener@cea.fr>

crtrott · 2026-05-04T15:22:58Z

The reason I want to reuse the heavy weight one are multi-fold:

more options adds more complexity in maintenance, documentation and teaching.
if at all possible we are generally trying to avoid backend specific options which masquerade as general options - users need to understand what "ConstantLaunch" means - which only expert CUDA programmers would have a chance for. It doesn't mean anything for other backend, and couldn't. How are other launches "non-constant" (after all, we always say the functor is const)" etc. The HeavyWeight indicates: you as the user are telling us that your functor will be a "heavy" one (meaning something like lots of code, lots of state, long runtime per iteration or a combination thereof).

Does that make sense?

Copilot

Pull request overview

This PR restores the CUDA “constant memory” kernel launch path (cuda_parallel_launch_constant_memory) to address a reported performance regression in Kokkos >= 5 (vs 4.7.02) for certain CUDA kernels, and wires the selection so users can opt into constant-memory launch behavior via WorkItemProperty::HintHeavyWeight (including when KOKKOS_IMPL_CUDA_USE_GRID_CONSTANT is enabled).

Changes:

Re-enables/ungates the constant-memory launch buffer and kernel entry points in Kokkos_Cuda_KernelLaunch.hpp.
Updates CUDA launch-mechanism deduction so HintHeavyWeight can select constant-memory launch even under grid-constant builds.
Makes the ConstantMemory invoker/types available in grid-constant configurations (previously compiled out).

 // If KOKKOS_IMPL_CUDA_USE_GRID_CONSTANT is used we leverage implicit constant
 // cache use via an argument attribute in the "local launch" mechanism. At that
 // point we only need local and global launch - the latter for functors that
 // exceed the kernel argument limit which is now 32kB. Local launch is always
 // strictly better than global launch - which means the light weight/heavy
 // weight property can be ignored - the only thing that matters is the size of
 // the functor.


pkestene · 2026-05-04T19:52:07Z

@@ -42,7 +41,6 @@
    [Kokkos::Impl::CudaTraits::ConstantMemoryUsage / sizeof(unsigned long)];

 #endif


Yes, that is a good point. It's probably a bad idea to permanently allocate buffer kokkos_impl_cuda_constant_memory_buffer only for constant memory launch mechanism especially if this mechanism is not used, or only for a few kernels.

Maybe that ConstantMemoryUsage could be made configurable (from cmake) ? But this would introduce a backend specific parameter

It seems to me a bit complex; in my particular case, the kernel that benefits from the constant memory launch weights 1088 Bytes, so a buffer much smaller than 32KiB would be enough. Not sure there is a good solution here.

yeah lets leave it like this, that was the previous behavior too.

crtrott · 2026-05-05T14:03:51Z


 //------------------------------------------------------------------------------
 // <editor-fold desc="Constant Memory"> {{{2
-#ifndef KOKKOS_IMPL_CUDA_USE_GRID_CONSTANT
 template <class DriverType, unsigned int MaxThreadsPerBlock,
          unsigned int MinBlocksPerSM>
 struct CudaParallelLaunchKernelFunc<


I didn't check whether this observation is correct (that we don't pre-populate it), but if we didn't we should.

This observation is true. I updated the branch to address this.

crtrott · 2026-05-05T14:04:33Z

  static constexpr CudaLaunchMechanism launch_mechanism =
 #ifdef KOKKOS_IMPL_CUDA_USE_GRID_CONSTANT
-      default_launch_mechanism;
+      (((property & heavy_weight) == heavy_weight) and
+       (sizeof(DriverType) < CudaTraits::ConstantMemoryUsage))
+          ? CudaLaunchMechanism::ConstantMemory
+          : default_launch_mechanism;


we do have tests for all sizes, and did decide against testing explicitly for the path back in the day.

JBludau · 2026-05-06T16:01:40Z

If this becomes our official way of working around this Cuda issue, we should have documentation for Kokkos::Experimental::WorkItemProperty

…ernal constructor.

We observed situations where __grid_constant attribute for the functor does NOT result in the same optimization behavior as when the functor is stored explicitly in constant cache. We are considering that a compiler issue. This allows opting into classical constant cache launch mechanism via the HeavyWeight kernel trait. Signed-off-by: Pierre Kestener <pierre.kestener@cea.fr>

crtrott requested changes Apr 29, 2026

View reviewed changes

Restore cuda_parallel_launch_constant_memory.

910cec6

Signed-off-by: Pierre Kestener <pierre.kestener@cea.fr>

pkestene force-pushed the feature/restore_cuda_parallel_launch_constant_memory branch from 8980a57 to 910cec6 Compare May 2, 2026 07:55

crtrott approved these changes May 4, 2026

View reviewed changes

crtrott requested a review from Copilot May 4, 2026 15:24

Copilot started reviewing on behalf of crtrott May 4, 2026 15:25 View session

Copilot AI reviewed May 4, 2026

View reviewed changes

Adrien-Tab self-requested a review May 6, 2026 15:27

Make sure constantMemMutexPerDevice's mutex is initialized in CudaInt…

2d13383

…ernal constructor.

Adrien-Tab approved these changes May 19, 2026

View reviewed changes

crtrott added the SNL-CI-APPROVAL Required for non-SNL contributions to run on SNL CI label May 20, 2026

crtrott merged commit dda3bac into kokkos:develop May 20, 2026
50 of 66 checks passed

nmm0 mentioned this pull request Jul 8, 2026

Changelog for 5.2.0 #9334

Merged

NamjaeChoi mentioned this pull request Jul 14, 2026

Add an option to disable CUDA constant memory dispatch #9352

Closed

Navigation Menu

Restore cuda_parallel_launch_constant_memory. - #9133

Restore cuda_parallel_launch_constant_memory.#9133
crtrott merged 2 commits into
kokkos:developfrom
pkestene:feature/restore_cuda_parallel_launch_constant_memory

pkestene commented Apr 28, 2026 •

edited

Loading

Uh oh!

crtrott left a comment

Uh oh!

pkestene commented Apr 29, 2026

Uh oh!

crtrott commented May 4, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

pkestene May 4, 2026

Uh oh!

crtrott May 5, 2026

Uh oh!

crtrott May 5, 2026

Uh oh!

pkestene May 18, 2026

Uh oh!

crtrott May 5, 2026

Uh oh!

JBludau commented May 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

		@@ -42,7 +41,6 @@
		[Kokkos::Impl::CudaTraits::ConstantMemoryUsage / sizeof(unsigned long)];

		#endif

Conversation

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

Uh oh!

Changelog Entry

Uh oh!

crtrott left a comment

Choose a reason for hiding this comment

Uh oh!

pkestene commented Apr 29, 2026

Uh oh!

crtrott commented May 4, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

pkestene May 4, 2026

Choose a reason for hiding this comment

Uh oh!

crtrott May 5, 2026

Choose a reason for hiding this comment

Uh oh!

crtrott May 5, 2026

Choose a reason for hiding this comment

Uh oh!

pkestene May 18, 2026

Choose a reason for hiding this comment

Uh oh!

crtrott May 5, 2026

Choose a reason for hiding this comment

Uh oh!

JBludau commented May 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

pkestene commented Apr 28, 2026 •

edited

Loading