← all records

NVIDIA + AMD kernel optimization

pitch-370008 · in Runtime improvements

Problem

In the last cycle we tried to profile and apply different optimizations to the NVIDIA and AMD kernels for dycore however due to the indeterministic code generation we couldn't come up to definitive answers. Since the code generation with specific branches in GT4Py and DaCe works much better we can focus on these optimizations this cycle. The goal is to improve performane for global and local grids on NVIDIA GH200 and AMD MI300A GPU.

Appetite

Solution

The current ideas we have are:

  • Do a sweep for thread block sizes in a bunch of GT4Py programs and try to figure out which of their kernels benefit from which thread block size. Based on their characteristics we could find ways to set certain thread block sizes to certain kernels depending on these characteristis.
  • Try vertical and horizontal loop blocking especially on AMD but NVIDIA as well by applying the above logic as well
  • See if limiting kernel fusion is beneficial since smaller kernels can better utilize the cache on AMD MI300A GPU

Rabbit holes

No-gos

Progress

  • Check whether the latest changes in https://github.com/C2SM/icon4py/pull/1367, https://github.com/GridTools/gt4py/tree/dace_toolchain_deterministic and https://github.com/GridTools/dace/tree/dace_toolchain_deterministic produce deterministic code every time
    • Currently there is an issue in only one program as described here
  • Add heuristics in the GPUSetBlockSize pass based on the range of the Y dimension (vertical) and the number of independent loads from memory that are indepent of the vertical dimension
    • See changes to src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py in https://github.com/GridTools/gt4py/compare/dace_toolchain_deterministic...threadblock_heuristics?expand=1
      • The above changes are based on a sweep in vertically_implicit_solver_at_predictor_step and compute_rho_theta_pgrad_and_update_vn for the following thread block sizes: [(32,8), (64,4), (128,2), (256,1)] on GH200 and MI300
      • The heuristic is tuned in such way that these two stencils give the best results on both GPUs
  • Run experiments
    • On GH200
      • Run the JW experiment on 1 node with R02B07 but didn't notice an improvement (difference in the best vs worse versions in GH200 of the above stencils is ~5%)
    • On MI300
      • Run the JW experiment on 1 node with R02B07 but noticed worse results (~3x) than the ones measured a month ago for PASC
      • Figure out if there's some issue with JW and communication or the kernels are that much slower
      • Need to check also with the optimal grid/thread block configuration for GCDs on AMD GPUs in DaCe: https://github.com/GridTools/dace/tree/chiplet_old_codegen (PASC measurements were done with this change)
  • Try to extend the heuristic across the whole dycore while dumping the independent and dependent bytes
  • Try to extend the heuristic to the icon_benchmark_global grid while dumping the independent and dependent bytes
  • Check horizontal/vertical unrolling based on https://github.com/GridTools/gt4py/pull/2578

As summarised at the review of 11.08.2026: benchmarked vertically_implicit_solver_at_predictor_step and compute_rho_theta_pgrad_and_update_vn for different thread block sizes on GH200 and collected runtimes for each kernel; tried to come up with heuristics to set the thread block sizes based on the data read and written to each stencil and whether they depend on the vertical dimension or not; didn't manage to see any big difference on GH200 (~5% in total with R02B07 in JW experiment on 1 node) — maybe the heuristic isn't general enough for all the stencils, need to continue looking into this, and didn't have time to do the same for AMD.

Since the review

The NVIDIA side of the bet closed with #1446, the thread block configuration for GH200, merged on 26.08.2026 and reviewed by edopao rather than by the Philip the table put on this row. That is the day this bet landed, two weeks past the cycle-37 review, and it is the end date above.

The AMD side did not fit the one week this was bet at. It carries on as pitch-3b71d9 in cycle 38, around #1444.