Enable generic kernel fusion
- Dominant language
- Julia
- Stars
- 117
- Forks
- 19
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 38
Description
## Purpose
As of #2466 and #2486, all stencil operations in ClimaCore are parallelized using similar CUDA kernels, with one column assigned to each block of threads, subdivided into one point per thread, and with shared memory used to propagate data between points. This now leaves us with five types of kernels in ClimaCore:
- Pointwise operations—one point per thread, each block handles points that are stored consecutively (block size chosen to maximize occupancy); i.e.,
- Applying a function to a `DataLayout` or `Field` (`f.(args...)` for a non-`AbstractOperator` `f`)
- Stencil operations—one point per thread, each block handles points in a column, with intermediate values cached in shared memory; i.e.,
- Applying a predefined stencil (`op.(args...)` for a `FiniteDifferenceOperator` `op`)
- Applying a generic stencil (`op .* arg` for a `ColumnwiseBandMatrixField` `op`)
- Inverting a generic stencil (`single_field_solve!`)
- Inverting multiple generic stencils (`multiple_field_solve!`)
- Spectral element operations—one point per thread, each block handles `Nvt` slabs that are stored consecutively (`Nvt` chosen so block size is roughly 256), with intermediate values cached in shared memory; i.e.,
- Applying a predefined spectral element operator (`op.(arg)` for a `SpectralElementOperator` `op`)
- Inter-element exchanges—one perimeter point per thread, each block handles perimeter points that are stored consecutively (block size is roughly 256), with intermediate values cached in global memory; i.e.,
- Applying DSS to a `DataLayout` (`dss!`) or `Field` (`weighted_dss!`), using a `DSSBuffer` as a cache
- Upward aggregations—one column per thread, each block handles columns that are stored consecutively (block size chosen to maximize occupancy); i.e.,
- Reducing each column to one value (`column_reduce!`)
- Accumulating values at each level along a column (`column_accumulate!`)
To minimize global memory reads, which currently use up the majority of our compute time, we can fuse sequential pointwise and stencil operations, and we can also fuse sequential pointwise and spectral element operations. If there are no unforeseen roadblocks, we could achieve the following without needing to modify the types of kernels we launch:
1. The implicit solve can be fused into a single kernel, since it only involves pointwise and stencil operations.
- This would singlehandedly eliminate most (~80%) of the reads that occur in each timestep.
2. The explicit tendency can be fused into a small collection of kernels that interleave stencil operations with spectral element operations.
- This would eliminate much of the remaining reads (~20%) that occur in each timestep, depending on how many interleaved kernels are required.
3. Other pointwise operations in the timestepper could also be fused with the kernels mentioned above.
- There is not much performance to be gained here following CliMA/ClimaTimeSteppers.jl#414, but this can be reassessed after more fusion has been achieved.
This SDI outlines a path to achieving a fully fused implicit solve and a mostly fused explicit tendency. Further elimination of global memory reads is still possible, but would require changing the types of kernels we launch. This work described here is a direct followup to #2522, which introduces the concept of a `DataScope`, implements a `foreach_point` communication primitive, and enables local caching of newly allocated `DataLayout`s with statically inferable sizes. The generic pointwise kernel fusion enabled by these changes must be extended to vertical and horizontal kernels before we can start to see significant performance benefits.
## Task Breakdown
1. Implement a `foreach_column!(f, args...; mask)` primitive to handle all stencil operations and upward aggregations.
- This will extend the `eachcol` function, like how `foreach_point` is used to extend `eachindex` in PR#2522.
- Can dispatch on the `DataScope` of its `args` to determine whether to launch a cuda kernel, thread pool, etc.
- Support any function `f(arg_columns...)` with arbitrary combinations of pointwise and stencil operations.
- Like `foreach_point`, this should terminate early and just call `f` when the `args` are already single columns.
- Will allow `foreach_column!` to get called from within itself, like from a stencil broadcast inside a fused kernel.
- Any `CuDeviceArray` passed to the kernel function may be replaced with a `CuStaticSharedArray` per column.
- The data for an `input_column` would get copied into its `CuStaticSharedArray` before evaluating `f`.
- The data from an `output_column` would get copied out of its `CuStaticSharedArray` after evaluating `f`.
2. Replace separate kernel launches with calls to `foreach_column!`.
- Verify that `copyto!` for a `StencilBroadcasted` can be replaced with `foreach_column!` without affecting performance.
- Similar to how `copyto!` for a `DataStyle` `Broadcasted` was implemented using `foreach_point` in PR#2522.
- Incrementally fuse operations in ClimaCore, ClimaAtmos, and ClimaTimeSteppers:
1) Fuse `ldiv!` for a `FieldMatrixWithSolver` in ClimaCore.
2) Fuse `implicit_tendency!` and `set_implicit_precomputed_quantities!` in ClimaAtmos.
3) Fuse `update_jacobian!` and `set_implicit_precomputed_quantities!` in ClimaAtmos.
4) Fuse `solve_implicit_equation!` in ClimaTimeSteppers.
- Eliminate cached values if hardware limit for shared memory is exceeded:
- On the L40S, occupancy will decrease after only 8 cached fields of `Float32` values (48 KB of shared memory per SM / 1536 max points per SM / 4 B per Float32 at each point).
- On the L40S, it will be impossible to launch a kernel with 188 or more cached fields of `Float32` values (48 KB of shared memory per SM / 64 min points per SM / 4 B per Float32 at each point).
- Fix GPU compatibility issues as they pop up (e.g., type-unstable variable manipulations in ClimaAtmos)
3. Implement a `foreach_slab!(f, args...; mask)` primitive to handle all spectral element operations.
- This will extend the `eachrow` function, like how `foreach_point`/`foreach_column` extend `eachindex`/`eachcol`.
- Support any function `f(arg_slabs...)` with arbitrary combinations of pointwise and spectral element operations.
- This should terminate early and just call `f` when the `args` are already single slabs.
4. Replace separate kernel launches with calls to `foreach_slab!`.
- Verify that `copyto!` for a `SpectralBroadcasted` can be replaced with `foreach_slab!` without affecting performance.
- Fuse sequential parts of `remaining_tendency!` and `set_precomputed_quantities!` in ClimaAtmos
- Rearrange internal operations as needed to improve fusion by reducing the number of sequential parts.
## Producers
This task will initially be assigned to ICCS (@Mikolaj-A-Kowalski et al). As the current refactoring work is completed, @imreddyTeja and @dennisYatunin will also be looped in.
## Reviewers
@dennisYatunin, @imreddyTeja, @tapios
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the DataScope and foreach_point work from #2522, then trace the existing stencil, upward aggregation, and spectral element kernel entry points named in the task. Implement and validate foreach_column! and foreach_slab!, replacing the listed separate launches and checking fusion in ClimaCore, ClimaAtmos, and ClimaTimeSteppers. Done means the planned implicit and explicit operations are fused without exceeding shared-memory limits or harming performance.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- hpc, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100