CliMA / CliMA/ClimaCore.jl

Unify operators into a single convolution kernel on the CUDA backend

Open
#2,497 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
117
Forks
19
Avg merge
3d 4h
Merged PRs (30d)
41

Description

## Purpose

Following #2522, the `DataLayout` interface will allow communication primitives to be implemented separately from specialized operators. That PR introduces the `foreach_point` primitive, and uses it to define the`fill!` and `copyto!` operators for `DataLayout`s and `LazyDataLayout`s (i.e., `DataStyle` broadcast expressions). #2514 lays out a plan for following this up with the `foreach_column` primitive, and later with the `foreach_slab` primitive.

This SDI outlines the other half of the operator simplification work: replacing all vertical stencil operations (specifically, functions called by `copyto!` for a `StencilBroadcasted`) with a single "convolution" operator. While this will not be as simple as the pointwise operators, it will still improve the maintainability of ClimaCore by centralizing all vertical operation logic in a single source code location, eliminating the need for code duplication in the CUDA extension. Since linear operator composition and inversion (also known as "deconvolution") are already operator-agnostic, as are `column_reduce!`/`column_accumulate!`, this is the final step needed to eliminiate redundant code for vertical operations.

## Task Breakdown

1. Allow `BandMatrixRow` to be applied as a standard function
- Define a method for `(::BandMatrixRow)(args...)`, where the number of arguments is the matrix bandwidth
- This function may either use `*`, or it may use `mul_with_projection`, whichever is easier
- Using `mul_with_projection` will require explicitly passing the `local_geometry` to every function
- Using `*` will require wrapping `Tensor` coefficients of `BandMatrixRow`s in some new `LinearForm` type that contains the `local_geometry`
- E.g., `*(x::LinearForm, y::Tensor)` can act equivalently to `mul_with_projection(x.tensor, y, x.local_geometry)`
2. Implement `column_convolve!(op, args...; bcs...)` to apply stencils, matching `column_reduce!`/`column_accumulate!`
- Support any generic convolution kernel `op(arg_points...)`, with linear kernels represented by `BandMatrixRow`s
- To represent the gradient as a `BandMatrixRow`, probably need a `LinearForm` type to encode "outer products"
- Nonlinear upwinding operators can use generic `Function`s or closures
- Support any generic boundary condition `bc_op(arg_points...)`, though we only need a few common cases
- Following #2521, all boundary conditions will be either extrapolating/padding arguments or setting the output value
- Constant output values at boundaries can be implemented as `bc_op = Returns(const)`
- Nonconstant output values can be implemented as `BandMatrixRow`s for linear operators
- Should work with any `DataScope` from #2522, not just the CUDA device scopes
- Same algorithm would be performant on CPUs if data were cached in `MArray`s instead of CUDA shmem arrays
3. Extend `column_convolve!` to compose stencils, instead of just applying them to input arguments
- Composing linear operators corresponds to matrix multiplication, which is also the case when they are applied
- If stencils and their boundary conditions are both `BandMatrixRow`s, composition is straightforward
- This is already implemented in ClimaCore, but with a lot of code duplication and unnecessary complexity
- Composing nonlinear operators might also be reasonable, if they are stored lazily as `op1 ∘ op2`
4. Further simplifications: delete old code once it is no longer necessary, use `foreach_column` primitive when it is available

## Producers

This task will initially be assigned to ICCS.

## Reviewers

@dennisYatunin, @imreddyTeja

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading the existing column_reduce!, column_accumulate!, BandMatrixRow, and StencilBroadcasted paths, then trace the duplicated vertical-operation logic in the CUDA extension. Define the scope of column_convolve! from the listed generic kernels, boundary conditions, DataScopes, and operator composition requirements. Done means vertical stencil operations use the unified convolution path and redundant old code is removed.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend, hpc, performance
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.