|
| 1 | +# Masks |
| 2 | + |
| 3 | +## Motivation |
| 4 | + |
| 5 | +ClimaCore spaces, `SpectralElementSpace2D`s in particular, support masks, where |
| 6 | +users can set horizontal nodal locations where operations are skipped. |
| 7 | + |
| 8 | +This is especially helpful for the land model, where they may have degrees of |
| 9 | +freedom over the ocean, but do not want to evaluate expressions in regions where |
| 10 | +data is missing. |
| 11 | + |
| 12 | +Masks in ClimaCore offer a solution to this by, ahead of time prescribing |
| 13 | +regions to skip. This helps both with the ergonomics, as well as performance. |
| 14 | + |
| 15 | +## User interface |
| 16 | + |
| 17 | +There are two user-facing parts for ClimaCore masks: |
| 18 | + |
| 19 | + - set the `enable_mask = true` keyword in the space constructor (when available), |
| 20 | + which is currently any constructor that returns/contains a `SpectralElementSpace2D`. |
| 21 | + - use `set_mask!` to set where the mask is `true` (where compute should occur) |
| 22 | + and `false` (where compute should be skipped) |
| 23 | + |
| 24 | +Here is an example |
| 25 | + |
| 26 | +```julia |
| 27 | +using ClimaComms |
| 28 | +ClimaComms.@import_required_backends |
| 29 | +import ClimaCore: Spaces, Fields |
| 30 | +using ClimaCore.CommonSpaces |
| 31 | +using Test |
| 32 | + |
| 33 | +FT = Float64 |
| 34 | +ᶜspace = ExtrudedCubedSphereSpace(FT; |
| 35 | + z_elem = 10, |
| 36 | + z_min = 0, |
| 37 | + z_max = 1, |
| 38 | + radius = 10, |
| 39 | + h_elem = 10, |
| 40 | + n_quad_points = 4, |
| 41 | + staggering = CellCenter(), |
| 42 | + enable_mask = true, |
| 43 | +) |
| 44 | + |
| 45 | +# How to set the mask |
| 46 | +Spaces.set_mask!(ᶜspace) do coords |
| 47 | + coords.lat > 0.5 |
| 48 | +end |
| 49 | +# Or |
| 50 | +mask = Fields.Field(FT, ᶜspace) |
| 51 | +mask .= map(cf -> cf.lat > 0.5 ? 0.0 : 1.0, Fields.coordinate_field(mask)) |
| 52 | +Spaces.set_mask!(ᶜspace, mask) |
| 53 | +``` |
| 54 | + |
| 55 | +Finally, operations over fields will be skipped where `mask == 0`, and applied |
| 56 | +where `mask == 1`: |
| 57 | + |
| 58 | +``` |
| 59 | +@. f = 1 # only applied where the mask is equal to 1 |
| 60 | +``` |
| 61 | + |
| 62 | +## Example script |
| 63 | + |
| 64 | +Here is a more complex script where the mask is used: |
| 65 | + |
| 66 | +```julia |
| 67 | +using ClimaComms |
| 68 | +ClimaComms.@import_required_backends |
| 69 | +import ClimaCore: Spaces, Fields, DataLayouts, Geometry, Operators |
| 70 | +using ClimaCore.CommonSpaces |
| 71 | +using Test |
| 72 | + |
| 73 | +FT = Float64 |
| 74 | +ᶜspace = ExtrudedCubedSphereSpace(FT; |
| 75 | + z_elem = 10, |
| 76 | + z_min = 0, |
| 77 | + z_max = 1, |
| 78 | + radius = 10, |
| 79 | + h_elem = 10, |
| 80 | + n_quad_points = 4, |
| 81 | + staggering = CellCenter(), |
| 82 | + enable_mask = true, |
| 83 | +) |
| 84 | +ᶠspace = Spaces.face_space(ᶜspace) |
| 85 | +ᶠcoords = Fields.coordinate_field(ᶠspace) |
| 86 | + |
| 87 | +# How to set the mask |
| 88 | +Spaces.set_mask!(ᶜspace) do coords |
| 89 | + coords.lat > 0.5 |
| 90 | +end |
| 91 | + |
| 92 | +# We also support the syntax `Spaces.set_mask!(::AbstractSpace, ::Field)` |
| 93 | + |
| 94 | +# We can check the mask directly: (internals, only for demonstrative purposes) |
| 95 | +mask = Spaces.get_mask(ᶜspace) |
| 96 | +@test count(parent(mask.is_active)) == 4640 |
| 97 | +@test length(parent(mask.is_active)) == 9600 |
| 98 | + |
| 99 | +# Let's skip operations that use fill! |
| 100 | +ᶜf = zeros(ᶜspace) # ignores mask |
| 101 | +@. ᶜf = 1 # tests fill! # abides by mask |
| 102 | + |
| 103 | +# Let's show that 4640 columns were impacted: |
| 104 | +@test count(x->x==1, parent(ᶜf)) == 4640 * Spaces.nlevels(axes(ᶜf)) |
| 105 | +@test length(parent(ᶜf)) == 9600 * Spaces.nlevels(axes(ᶜf)) |
| 106 | + |
| 107 | +# Let's skip operations that use copyto! |
| 108 | +ᶜz = Fields.coordinate_field(ᶜspace).z |
| 109 | +ᶜf = zeros(ᶜspace) |
| 110 | +@. ᶜf = 1 + 0 * ᶜz # tests copyto! |
| 111 | + |
| 112 | +# Let's again show that 4640 columns were impacted: |
| 113 | +@test count(x->x==1, parent(ᶜf)) == 4640 * Spaces.nlevels(axes(ᶜf)) |
| 114 | +@test length(parent(ᶜf)) == 9600 * Spaces.nlevels(axes(ᶜf)) |
| 115 | + |
| 116 | +# Let's skip operations in FiniteDifference operators |
| 117 | +ᶠf = zeros(ᶠspace) |
| 118 | +c = Fields.Field(FT, ᶜspace) |
| 119 | +div = Operators.DivergenceF2C() |
| 120 | +foo(f, cf) = cf.lat > 0.5 ? zero(f) : sqrt(-1) # results in NaN in masked out regions |
| 121 | +@. c = div(Geometry.WVector(foo(ᶠf, ᶠcoords))) |
| 122 | + |
| 123 | +# Check that this field should never yield NaNs |
| 124 | +@test count(isnan, parent(c)) == 0 |
| 125 | + |
| 126 | +# Doing the same thing with a space without a mask will yield NaNs: |
| 127 | +ᶜspace_no_mask = ExtrudedCubedSphereSpace(FT; |
| 128 | + z_elem = 10, |
| 129 | + z_min = 0, |
| 130 | + z_max = 1, |
| 131 | + radius = 10, |
| 132 | + h_elem = 10, |
| 133 | + n_quad_points = 4, |
| 134 | + staggering = CellCenter(), |
| 135 | +) |
| 136 | +ᶠspace_no_mask = Spaces.face_space(ᶜspace_no_mask) |
| 137 | +ᶠcoords_no_mask = Fields.coordinate_field(ᶠspace_no_mask) |
| 138 | +c_no_mask = Fields.Field(FT, ᶜspace_no_mask) |
| 139 | +ᶠf_no_mask = Fields.Field(FT, ᶠspace_no_mask) |
| 140 | +@. c_no_mask = div(Geometry.WVector(foo(ᶠf_no_mask, ᶠcoords_no_mask))) |
| 141 | +@test count(isnan, parent(c_no_mask)) == 49600 |
| 142 | +``` |
| 143 | + |
| 144 | +## Supported operations and caveats |
| 145 | + |
| 146 | +Currently, masked _operations_ are only supported for `Fields` (and not |
| 147 | +`DataLayouts`) with `SpectralElementSpace2D`s. We do not yet have support for |
| 148 | +masked `SpectralElement1DSpace`s, and we will likely never offer masked |
| 149 | +operation support for `DataLayouts`, as they do not have the space, and can |
| 150 | +therefore not use the mask. |
| 151 | + |
| 152 | +In addition, some operations with masked fields skip masked regions |
| 153 | +(i.e., mask-aware), and other operations execute everywhere |
| 154 | +(i.e., mask-unaware), effectively ignoring the mask. Here is a list of |
| 155 | +operations of mask-aware and mask-unaware: |
| 156 | + |
| 157 | + - `DataLayout` operations (`Fields.field_values(f) = 1`) mask-unaware (will likely never be mask-aware). |
| 158 | + - `fill!` (`@. f = 1`) mask-aware |
| 159 | + - point-wise `copyto!` (`@. f = 1 + z`) mask-aware |
| 160 | + - stencil `copyto!` (`@. ᶜf = 1 + DivergenceF2C()(Geometry.WVector(ᶠf))`) mask-aware (vertical derivatives and interpolations interpolations) |
| 161 | + - spectral element operations `copyto!` (`@. f = 1 + Operators.Divergence()(f)`), where `Operators.Divergence` carries out a divergence operation in horizontal directions. mask-unaware |
| 162 | + - fieldvector operations `copyto!` (`@. Y += 1`) mask-unaware |
| 163 | + - reductions: |
| 164 | + - `sum` (mask-unaware, warning is thrown) |
| 165 | + - `extrema` (mask-unaware, warning is thrown) |
| 166 | + - `min` (mask-unaware, warning is thrown) |
| 167 | + - `max` (mask-unaware, warning is thrown) |
| 168 | + - field constructors (`copy`, `Fields.Field`, `ones`, `zeros`) are mask-unaware. |
| 169 | + This was a design implementation detail, users should not generally depend on the results where `mask == 0`, in case this is changed in the future. |
| 170 | + - internal array operations (`fill!(parent(field), 0)`) mask-unaware. |
| 171 | + |
| 172 | +## Developer docs |
| 173 | + |
| 174 | +In order to support masks, we define their types in `DataLayouts`, since |
| 175 | +we need access to them from within kernels in `DataLayouts`. We could have made |
| 176 | +an API and kept them completely orthogonal, but that would have been a bit more |
| 177 | +complicated, also, it was convenient to make the masks themselves data layouts, |
| 178 | +so it seemed most natural for them to live there. |
| 179 | + |
| 180 | +We have a couple types: |
| 181 | + |
| 182 | + - abstract `AbstractMask` for subtyping masks and use for generic interface |
| 183 | + methods |
| 184 | + - `NoMask` (the default), which is a lazy object that should effectively result |
| 185 | + in a no-op, without any loss of runtime performance |
| 186 | + - `IJHMask` currently the only supported horizontal mask, which contains |
| 187 | + `is_active` (defined in `set_mask!`), `N` (the number of active columns), |
| 188 | + and maps containing indices to the `i, j, h` locations where `is_active` is |
| 189 | + true. The maps are defined in `set_mask_maps!`, allows us to launch cuda |
| 190 | + kernels to only target the active columns, and threads are not wasted on |
| 191 | + non-existent columns. The logic to handle this is relatively thin, and |
| 192 | + extends our current `ext/cuda/datalayouts_threadblock.jl` api |
| 193 | + (via `masked_partition` and `masked_universal_index`). |
| 194 | + |
| 195 | +An important note is that when we set the mask maps for active columns, the |
| 196 | +order that they are assigned can be permuted without impacting correctness, but |
| 197 | +this could have a big impact on performance on the gpu. We should investigate |
| 198 | +this. |
| 199 | + |
0 commit comments