Kokkos support in 4C

[georghammerl]

I would like to set up a meeting to get to know each other with Jeremy, our new colleague at Hereon who will be joining us in July. One of his key priorities is to add GPU support to 4C. Significant initial progress has already been made toward enabling Kokkos (#2012 (comment)).
The goal will be to get to know each other and exchange ideas on the future integration of Kokkos, GPU, and OpenMP support into 4C.

The vote does not work as expected. I apologize for the inconvenience. Only single vote is possible. Please comment your favorite dates. Thank you.

[mayrmt]

@georghammerl I can only cast a single vote in above's poll.

I'd be available on July 7th, 8th, and 9th.

[PhilipOesterlePekrun]

July 6th, 7th, or 8th would work for me.

[jeremylt]

At this time, all of these dates are open for me. HR hasn't put anything on my schedule yet

[georghammerl]

Let us have the meeting on July 7th, 10:00 - 11:30
Here is a Teams-Link:
https://teams.microsoft.com/meet/339083954025355?p=29JCC958S8fBo252mf

[jeremylt]

For me, some potential points of clarification/discussion would be

• Which pieces of the codebase are best suited to GPU computation? In my experience, the best computation to shift to the GPU is loops where the loop body has a lot of work and a high degree of independence between loop iterations, such as loops with lots of physics over quadrature points or particles. Identifying the most computationally heavy and easily parallelizeable loops would let us see the strongest initial results.

• Related, are there portions of computation than can/should be refactored to create more of this independence between loop body iterations? Sometimes rearranging the exact same math can make it easier for the compiler to generate faster kernels.

• Which data needs to be managed by Kokkos? The nice thing about using Kokkos is that we don't have to manage data synchronization between the host and device ourselves; we put the relevant data into containers managed by Kokkos. One point that can trip people up - any data synchronization between the host and device incurs a performance penalty, so we will want to minimize the total number of those transfers. Setting up as much computation as possible on device once we pay the cost of moving the data once helps us see those attractive performance benefits.

• Are there getters/setters in place that we can hide the Kokkos based data access behind? The less intrusive we can make these changes, the easier we can make it for other developers to maintain and extend these changes moving forward.

• Related, is there developer facing documentation (vs user facing) where we can put a guide to help developers who are working on new features or material models write them in a way that is amicable to good Kokkos performance? If there isn't much developer facing documentation yet, can we start doing that?

• Speaking of performance, what do we want to use as representative 'benchmark' problems? (I would be careful selling the '100x speedup' claims with such benchmarks because its an apples to oranges comparison.) Such benchmark problems can be used for profiling and performance tuning, so it would be nice if they are representative of common or important workflows - not too simple and not too niche.

• What is the testing plan? Once we have converted portions of the code to using Kokkos, we'll want to be running those kernels in CI to make sure that code stays correct. As I understand it, the plan right now is to test with OpenMP in CI.

• Related, do we want/need longitudinal benchmarking results? Perhaps from a weekly performance job? Then we could monitor performance changes as we move more computation to Kokkos-enabled, and we would be able to detect if a change in dependency versions hurts or helps performance.

• What am I forgetting to consider?

I am still learning 4C as a codebase, so some of these questions, but I feel like too much communication is better than too little in this sort of situation.

[PhilipOesterlePekrun]

Yes, we can address these points in tomorrow's meeting. One important thing to keep in mind is that the hybrid programming design should guard against oversubscription (of, in the case of kokkos CUDA backend, each compute node's GPU) when using more than one MPI rank per node. This works quite easily for CUDA, but unfortunately does not work in the same way for OpenMP, which is a problem I've been stuck on for a while now.

[jeremylt]

Interesting, do you have a toy problem or a writeup that you could point me at to better understand that issue? I have much more experience with CUDA than OpenMP, but I would have assumed Kokkos provided tools to prevent that sort of issue.

[jeremylt]

Also, is multiple MPI ranks per physical node a common configuration when using OpenMP? With CUDA/HIP I've usually done 1 rank per physical GPU.

[georghammerl]

Meeting notes 07.07.2026
Participants: @jeremylt @mayrmt @PhilipOesterlePekrun @georghammerl

• Kokkos integration and portability

  • Due to portability requirements and the desire to avoid additional dependencies, Kokkos appears to be the preferred solution for performance portability in 4C
  • Two integration paths available
    • Direct use of Kokkos / Kokkos kernels
    • Use of Kokkos through Tpetra (in future)
  • Build a sustainable Kokkos concept for 4C based on the current open PR (Enable using Kokkos with CUDA backend for MIRCO and in general #2012)
  • Teuchos::SerialDenseMatrix (see 4C_linalg_serialdensematrix.hpp) will likely be replaced by Kokkos data structures
  • Portability should not be limited to NVIDIA hardware; compilation on AMD platforms should also be considered (in the long run)

• Parallelization strategy

  • A promising starting point is one MPI rank for one CPU core per available GPU with parallelization delegated to shared memory on the node or accelerator device
  • Investigate opportunities to identify and refactor loops for device execution, similar to vectorization approaches used in deal.II
  • Review the centralized element evaluation workflow (Core::FE::Discretization::evaluate with its element loop down to Gauss-point loops) as a potential target for device portability

• Candidate code to bring to the GPU

  • The particle code is a good initial candidate because it is self-contained and independent of Trilinos/Epetra.
  • Potential representative test cases:
    • SPH (e.g., stomach simulations or dam-break) (for particles)
    • Elasticity/plasticity problems (for solid)

• Toolchain and infrastructure

  • The current requirement for an nvcc compiler wrapper is not ideal; For now, Clang appears to be a valid solution as parts of 4C uses modern C++ features that are problematic for NVIDIA's toolchain
  • Establish build test until Hereon Kubernetes cluster is available as a CI runner including GPUs
  • MPI/OpenMP testing can be developed independently

• Documentation and benchmarking

  • Need for developer documentation/tutorials describing best practices for Kokkos integration in one's code
  • Evaluate whether the existing benchmarking framework is sufficient to properly detect potential performance degradation

[jeremylt]

Using Kokkos via Trilinos 16.2.0 has us using Kokkos 4.7.1. We can improve that by swapping to Trilinos 16.2.1, which brings us to Kokkos 4.7.4 which includes some fixes for newer compilers. It's a bugfix release, so I don't think it should break anything for 4C.

As we add Kokkos to 4C, we should make sure that we only use features that are supported by Kokkos 5.0+, since the long term goal is to move 4C to Trilinos 17.0+, which uses Kokkos 5.0+. Don't need to add more tech debt.

There's something about how the implementation of Kokkos Views changed between 4.1 and 5.0, and I'm working on figuring out what that means for us in terms of any potential for tech debt.

[PhilipOesterlePekrun]

Ah, I did not see that minor release between 16.2.0 and 17.0. Thanks for the info, I'll switch to that.

[PhilipOesterlePekrun]

Interesting, 16.2.1 was actually released after 17.0.0. I guess we are not the only ones still using Epetra :)

[jeremylt]

Yeah, I think it was a patch release to catch some Kokkos patches

[jeremylt]

Ok, just a small update - I am working on some research and hope to have a proposal next week on how I'd recommend we update an entire module(?), like the particle code, to use Kokkos. Since Kokkos ships with Trilinos and we have it included in 4C automatically, I wouldn't worry to much about directly including Kokkos code into files, as Kokkos will compile them for the CPU if no GPU device is available.

I think the big concern I have is designing the interactions with common core modules like linalg and utils so that memory synchronization isn't isn't something most users or developers have to think about. It might make the most sense to attach Kokkos views to any large allocation of memory and only use the Kokkos view in modules that are adding Kokkos code. Those modules would have to do something similar when interfacing with other modules that are not Kokkos enabled.

Once we have the memory transfer between Kokkos enabled and not enabled modules defined, then we have natural break points and can transition the entire innards of a module to Kokkos execution gradually without disrupting other work. Once two modules that touch are both Kokkos enabled, then we can make a version of the interface that just hands off Kokkos views instead of CPU only objects.

Transitioning the innards to Kokkos isn't conceptually difficult - I imagine we'll largely start out with Kokkos parallel_for loops - but we'll want to profile as we go and also decide upon any algorithmic re-ordering to make loop bodies more parallelizable. This should also make those loop bodies emit better AVX instructions on CPU only, I'd guess, though I'm not sure exactly how that all works with Kokkos (yet!).

Anyways, its Friday, so this will rumble around in my head over the weekend and I'll hopefully have more concrete of ideas late next week.

I'm playing around with learning Kokkos and some of the options for not disrupting existing functions here: https://github.com/jeremylt/legacy-kokkos-integration-demo

I also hope to research what other projects have done when binging on Kokkos like this, as this one of the use cases it was designed for.

[jeremylt]

Ok, stage 1 here will be getting the memory management to work correctly. I have sketched out the user facing part of my plan here:

enum MemorySpace {
  DefaultSpace, // Device data access
  HostSpace,    // Host data access
};

// ...

inline void init_dual_view(ParticleState state) {
  if (is_dual_valid_[state]) return;
  auto host_view = host_states_[state];
  Kokkos::View<int *> device_view(host_view.label(), host_view.size());
  // Note: need to copy to device, as DualView must start with both Views in sync
  Kokkos::deep_copy(device_view, host_view);
  dual_states_[state] = Kokkos::DualView<int *>(device_view, host_view);
  is_dual_valid_[state] = true;
}

// ...

inline const int* get_ptr_to_state(ParticleState state, MemorySpace space = DefaultSpace) {
  if (space == DefaultSpace) {
    if (!is_dual_valid_[state]) init_dual_view(state);
    dual_states_[state].sync_device();
    return dual_states_[state].view_device().data();
  } else {
    if (!is_dual_valid_[state]) return host_states_[state].data();
    dual_states_[state].sync_host();
    return dual_states_[state].view_host().data();
  }
}

inline int* get_ptr_to_state_writable(ParticleState state, MemorySpace space = DefaultSpace) {
  if (space == DefaultSpace) {
    if (!is_dual_valid_[state]) init_dual_view(state);
    dual_states_[state].sync_device();
    dual_states_[state].modify_device();
    return dual_states_[state].view_device().data();
  } else {
    if (!is_dual_valid_[state]) return host_states_[state].data();
    dual_states_[state].sync_host();
    dual_states_[state].modify_host();
    return dual_states_[state].view_host().data();
  }
}

Things to note

1. We need to distinguish between read only vs writable access. This is because Kokkos needs to know when to sync back and forth between the CPU memory and GPU memory. I'm open to name suggestions, but basically doing it this way I suggest above means the 4C developers are shielded from needing to manage flagging the modifications and calling the sync functions themselves, so I think this is the way to go.

2. I put in lazy initialization of the GPU side memory. I think its the right way to go as we are in the process of porting, as we will not need to have device memory for everything at first. We may want to revisit this in the future though. Also, I need to research the best way to handle this if we're running CPU only, as I don't want to effectively double our CPU allocations for no reason. More to come here.

3. Currently I set the default as the GPU side. This is somewhat aspirational and I was mirroring the Kokkos philosophy somewhat. We can flip this convention though easily if we want it to be an 'opt in' thing where we want the developer to affirm that they intend to use GPU functionality. I'm sorta leaning that way but want the opinions of others.

This basic set of changes will push changes through the internals of the particle container, but I think that it will largely be invisible to the user outside of needing to a) call specifically for writable access and 2) having to specific which memory space you want the pointer to be in.

[jeremylt]

Ok, updated to fix the risk of double allocation when Kokkos is configured CPU only and flip the convention to default to CPU for now. I'll start working on a PR for this idea, but feel free to throw thoughts as me, as this will take a while to get all the consequences straight in all of the methods and where they are called.

enum MemorySpace {
  DefaultSpace, // Device data access
  HostSpace,    // Host data access
};

// ...

inline void init_dual_view(ParticleState state) {
  if (is_dual_valid_[state]) return;
  auto host_view = host_states_[state];
  // Create mirror view and copy so both views are in sync before creating DualView
  Kokkos::View<double *> device_view =
      Kokkos::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace::memory_space(), host_view);
  dual_states_[state] = Kokkos::DualView<int *>(device_view, host_view);
  is_dual_valid_[state] = true;
}

// ...

inline const int* get_ptr_to_state(ParticleState state, MemorySpace space = HostSpace) {
  if (space == DefaultSpace) {
    if (!is_dual_valid_[state]) init_dual_view(state);
    dual_states_[state].sync_device();
    return dual_states_[state].view_device().data();
  } else {
    if (!is_dual_valid_[state]) return host_states_[state].data();
    dual_states_[state].sync_host();
    return dual_states_[state].view_host().data();
  }
}

inline int* get_ptr_to_state_writable(ParticleState state, MemorySpace space = HostSpace) {
  if (space == DefaultSpace) {
    if (!is_dual_valid_[state]) init_dual_view(state);
    dual_states_[state].sync_device();
    dual_states_[state].modify_device();
    return dual_states_[state].view_device().data();
  } else {
    if (!is_dual_valid_[state]) return host_states_[state].data();
    dual_states_[state].sync_host();
    dual_states_[state].modify_host();
    return dual_states_[state].view_host().data();
  }
}

[jeremylt]

See #2136