Strategies for Optimizing GPU-based EOS & Opacity Table Lookups in AMReX

[lwJi]

I'm working on an AMReX-based simulation where our physics kernels, running on the GPU, require frequent lookups from large, read-only Equation of State (EOS) and opacity tables (e.g., 2D or 3D tables, ~500MB size for each table).

Getting this right is critical for performance, and I'm facing the classic challenge of balancing memory management, data locality, and efficient computation. The two main hurdles are:

1. Efficient Loading: What is the best way to get these large tables into GPU memory so they are globally accessible to all GPU threads running AMReX kernels?

2. Fast Interpolation: Once the data is on the GPU, what is the most performant way to perform interpolation (e.g., bilinear, trilinear) inside a kernel launched with amrex::launch or amrex::MFIter without stalling the pipeline?

I'm looking for advice or best practices from the community on a few specific points:

• Data Management: What is the recommended AMReX approach for this? Should the tables be stored in Gpu::ManagedVector, Gpu::DeviceVector, or perhaps allocated via amrex::The_Arena()? How do you handle initialization and copy-to-device?

• Kernel Access: What's the cleanest way to pass this table data to a kernel? Should I be passing raw pointers, or is it better to wrap the table in a struct or Gpu::Device object that also contains its metadata (dims, min/max bounds, etc.)?

• Texture Memory: Has anyone had success using GPU texture memory for this? It seems ideally suited for this problem (hardware-accelerated interpolation, optimized caching for spatial locality), but I'm unsure how to integrate it cleanly with AMReX's data paradigm. Is this a recommended approach?

• Interpolation Code: Are there standard, GPU-optimized interpolation utilities within AMReX that I should be using, or is this typically a case of "roll your own" device-side functions?

[WeiqunZhang]

If you want texture memory, you will need to roll your own. Otherwise, you might want to have a look at amrex::TableData https://amrex-codes.github.io/amrex/doxygen/classamrex_1_1TableData.html#details.

If you run your jobs on a lot of nodes, having every process read 500 MB of data all at the same time might be bad for file system performance. Here is an example of limiting the number of processes doing I/O at the same time.

amrex/Tests/FillBoundaryComparison/main.cpp

Line 33 in 3c1f1c7

int nAtOnce = std::min(ParallelDescriptor::NProcs(), 32);

[lwJi]

@WeiqunZhang , could you comment on the strategy which read the table on one process and then broadcast the table to the rest of the processes?

[zingale]

we do this for our equation of state table:

https://github.com/AMReX-Astro/Microphysics/blob/main/EOS/helmholtz/actual_eos.H#L1418

[WeiqunZhang]

Yes, it's a good solution. You don't have to do the parallel I/O way I showed you, given that it's a one-time cost.

[lwJi]

we do this for our equation of state table:

https://github.com/AMReX-Astro/Microphysics/blob/main/EOS/helmholtz/actual_eos.H#L1418

Thanks @zingale

[lwJi]

Yes, it's a good solution. You don't have to do the parallel I/O way I showed you, given that it's a one-time cost.

Thanks @WeiqunZhang