Skip to content

icon_structured_slides.md

Latest commit

 

History

History
458 lines (337 loc) · 15.3 KB

icon_structured_slides.md

File metadata and controls

458 lines (337 loc) · 15.3 KB
Raw
marp false
theme cscs
paginate true
backgroundColor
backgroundImage url('slides-support/common/4k-slide-bg-white.png')
size 58140

ICON structured grid benchmark

bg cover

CSCS

Introduction

Motivation

  • ICON grid is partly structured
  • Investigate how expensive are indirect neighbor accesses compared to strided indexes of neighbors
  • Investigate optimizations based on strided indexes
  • Apply optimizations to current and future GT4Py development

ICON grid handling

ICON grid input

Manipulation of ICON grid

Torus grid characteristics

  • Torus grid has periodic boundaries
    • Adds complexity to compute strided indexes
    • Not applicable in production scenario of CH

Manipulation of ICON grid

Selection of computation domain

  • Filter torus grid to end up with a cartesian grid with 2 level halo
    • 1 would also be sufficient
    • Discard periodic edges
  • Calculate X and Y dimensions based on the distribution of vertices in space
    • x_dim equals the number of vertices with y == 0

Computation domain in red

Manipulation of ICON grid

Organization of edges in memory

Manipulation of ICON grid

Selection of computation domain

  • Fix e2v ordering of edges read by torus grid file to be by default per-vertex in GridManager
    • Some elements in e2v didn't follow this ordering by default
  • Select between per-vertex and per-orientation ordering in GridManager
  • Filter halo vertices and edges in benchmark python script

Computation domain in red

Selected kernels for exploration

  • Nabla4
    • Neighbor tables
      • e2c2v[4]
      • e2ecv[4]
    • Input fields
      • u_vert[VertexDim, KDim]
      • v_vert[VertexDim, KDim]
      • primal_normal_vert_v1[EdgeDim * ECVDim]
      • primal_normal_vert_v2[EdgeDim * ECVDim]
      • z_nabla2_e[OutputEdgeDim, KDim]
      • inv_vert_vert_length[OutputEdgeDim]
      • inv_primal_edge_length[OutputEdgeDim]
    • Output field
      • z_nabla4_e2[OutputEdgeDim, KDim]

Selected kernels for exploration

  • Nabla4 vertical (artificial)
    • Similar to nabla4 but all input fields are vertical
    • Neighbor tables
      • e2c2v[4]
      • e2ecv[4]
    • Input fields
      • u_vert[VertexDim, KDim]
      • v_vert[VertexDim, KDim]
      • primal_normal_vert_v1[EdgeDim * ECVDim, KDim]
      • primal_normal_vert_v2[EdgeDim * ECVDim, KDim]
      • z_nabla2_e[OutputEdgeDim, KDim]
      • inv_vert_vert_length[OutputEdgeDim, KDim]
      • inv_primal_edge_length[OutputEdgeDim, KDim]
    • Output field
      • z_nabla4_e2[OutputEdgeDim, KDim]

Selected kernels for exploration

  • Interpolate
      • Neighbor tables
      • v2e[6]
    • Input fields
      • p_e_in[EdgeDim, KDim]
      • ptr_coeff_1[OutputVertexDim, 6]
      • ptr_coeff_2[OutputVertexDim, 6]
    • Output fields
      • p_u_out[OutputVertexDim, KDim]
      • p_v_out[OutputVertexDim, KDim]

Nabla4 execution

Nabla4 execution

Interpolate execution

Nabla4 & Interpolate inlined v2e[e2c2v] execution

Nabla4 & Interpolate inlined v2e2c2v execution

Kernel implementations

  • Neighbor accesses
    • Indirect
      • Indirect accesses via neighbor tables
    • Strided
      • Neighbor accesses via strides
  • Iteration strategies
    • gpu_naive
      • Default iteration strategy in GridTools C++
      • One GPU thread calculates 1 element in horizontal and vertical axis
    • gpu_kloop
      • Optional iteration strategy in GridTools C++
      • One GPU thread calculates 1 element in horizontal axis but multiple vertical levels
      • Save neighbors and vertically independent fields in registers and iterate over multiple vertical fields

Kernel implementations

  • Separate
    • Execute nabla4 kernel and then interpolate
  • Inlined
    • Compute nabla4 output for every input of interpolate kernel
    • More computations
    • Less writes to device memory
  • Inlined v2v (indirect only)
    • Compress v2e[e2c2v] neighbor accesses to v2e2c2v
      • Read fields for 7 vertices instead of (6*4=) 24 vertices
    • Assumes certain order of vertices in e2c2v and edges in v2e

Kernel implementations

  • Inlined_cached
    • Save nabla4 output in shared memory and then use it in interpolate kernel
    • Reduces overcomputations compared to inlined implementation

Kernel implementations

  • gtfn
    • Only indirect
    • Based on GridTools C++
    • Improved GridTools C++
      • Memory loads via __ldg
      • gpu_kloop option
      • const neighbor tables and input fields
      • Kudos to Felix Thaler
  • CUDA
    • Plain cuda kernels
    • Launched by python script
    • Random input in benchmarks
    • Validated kernels with serialized data from GT4Py

General kernel optimizations

  • Occupancy
    • All kernels have been optimized for best occupancy
  • __launch_bounds__ and __maxnreg__
    • Launch bounds are applied to all kernels except some that were performing better with register limitation to a certain number of registers
  • Thread Block size
    • Specifically for gpu_naive implementations, increasing the vertical axis thread block size (ThreadBlockDim.y/z) was beneficial since there are more chances to find neighbor tables and vertically independent fields in cache
    • 4-8-9 most used number with 80 vertical levels in total. For some kernels 12 or 16
  • gpu_kloop
    • For this implementation vertical GridDim is set to 1. Iteration number is controlled by ThreadBlockDim.y/z and KDim. Exceptions are the inline_cached versions where the shared memory size is limiting the number of iterations possible
    • Optimized number of iterations per kernel. 5-40 iterations. 8-20 usually have the best performance

Notes for specific kernels

  • strided_gpu_{naive,kloop}_inlined_cached: Tried 2 different implementations, number of threads same as input but then deactivate for the output some of them and number of threads same as output where each thread caclulates multiple elements. Former is better
  • strided_*_inlined: calculate only necessary indexes, similar to v2e2c2v
  • indirect_*_inlined_v2v: pass v2e2c2v as input
  • indirect: nabla4 iterates on edges (per-orientation - 1 edge per thread)
  • strided: nabla4 iterates on vertices (per-vertex - 3 edges per vertex/thread)
  • Both strided and indirect versions operate on data with same ordering in memory
    • No SFC. Vertices are ordered per i and j coordinates and edges per orientation/vertices
  • strided: e2ecv is also computed
  • Smaller grids benefit by more threads and less k level iterations
  • Loop in k is done with stride blockDim.y/z * gridDim.y/z
    • It would be more beneficial for kernels that read data from adjacent k levels to do the looping with stride 1 in each thread
      • Current implementation in gtfn, not on the CUDA kernels since we don't evaluate such case

Results

  • GH200 GPU
  • Median runtime presented
    • 10 dry runs (not taken into account)
    • 101 runs to select median
  • 229758 edges
    • Close to the amount of edges that fit in a single GPU for ICON runs
  • 915948 edges
    • For exploration
  • Commit 979c2b9

gtfn improvements

  • strided separate kernels ~10% speed up
  • Inlining in gpu_naive doesn't help without extra optimizations due to overcomputations
  • cached approach in gpu_naive another ~20% speedup
  • gpu_kloop ~10-20% speedup compared to gpu_naive for separate kernels
  • indirect gpu_kloop inlined_v2v and strided gpu_kloop inlined ~2x faster
  • cached implementation not fastest in gpu_kloop
  • gpu_kloop not always beneficial since there are no vertical fields to save to registers
    • Only neighbors are saved in registers which are either loaded from memory or computed
  • strided gpu_kloop inlined_cached not better than more expensive overcomputations for other inlined implementations

Conclusions

  • For separate kernels strided is faster than indirect by 8-10%
  • For nabla4_interpolate where there are vertically independent fields gpu_kloop is beneficial
    • For nabla4_vertical_interpolate where there aren't many vertically independent fields gpu_kloop is not necessaarily beneficial due to more register usage and lower occupancy
  • strided is only slightly faster than indirect inlined_v2v in nabla4_interpolate but much faster than indirect
    • In nabla4_vertical_interpolate performance is different
      • Probably should understand it better but since strided_inlined_cached is not faster than strided_inlined maybe not worth it
  • inlined_cached using shared memory is not better than overcomputations/loading from memory for inlined version in strided implementations
    • Potentially room for improvement there

Next steps

  • Add interpolation stencil with c2v neighbor to see implact in inlined versions due to overcomputations
    • c2v neighbor will require computing 3 vertices per cell = 3 times the computations
    • Try c2v2e2c2v compressed neighbor in indirect version as well
  • Try cached approach for indirect implementation
    • Compute border coordinates for each Thread Block (should be the same as strided for our grid)
      • per-vertex ordering should be better due to smaller range of vertices/edges that need to be saved to shared memory
    • Load them to shared memory
  • Try TMA implementation
    • Tried cuda::pipeline to use LDGSTS/LDGSTS.BYPASS instructions for loading the input fields of nabla4 kernel
      • Big slowdown
      • Probably because computations are not enough to hide memory latencies
    • Try with TMA to see if it improves
    • More complex due to memory alignment requirements

Next steps

  • Compile time strides
    • x_dim is necessary to calculate the strides in strided version
      • Can be given in case of JIT compilation
  • Use cache hints for loads and non temporal stores

Questions?

bg cover

Resources