ClimaCore.jl/ext/cuda/operators_integral.jl at 5b131e46542f6470a1ae0c3d63d9514d4d6a3c29 · CliMA/ClimaCore.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import ClimaCore: Spaces, Fields, level, column
import ClimaCore.Operators:
    left_idx,
    strip_space,
    column_reduce_device!,
    single_column_reduce!,
    column_accumulate_device!,
    single_column_accumulate!
import ClimaComms
using CUDA: @cuda

function column_reduce_device!(
    dev::ClimaComms.CUDADevice,
    f::F,
    transform::T,
    output,
    input,
    init,
    space,
) where {F, T}
    Ni, Nj, _, _, Nh = size(Fields.field_values(output))
    us = UniversalSize(Fields.field_values(output))
    mask = Spaces.get_mask(space)
    if !(mask isa DataLayouts.NoMask) && space isa Spaces.FiniteDifferenceSpace
        error("Masks not supported for FiniteDifferenceSpace")
    end
    cart_inds = cartesian_indices_columnwise(us)
    args = (
        single_column_reduce!,
        f,
        transform,
        strip_space(output, axes(output)), # The output space is irrelevant here
        strip_space(input, space),
        init,
        space,
        us,
        mask,
        cart_inds,
    )
    nitems = Ni * Nj * Nh
    threads = threads_via_occupancy(bycolumn_kernel!, args)
    n_max_threads = min(threads, nitems)
    p = linear_partition(nitems, n_max_threads)
    auto_launch!(
        bycolumn_kernel!,
        args;
        threads_s = p.threads,
        blocks_s = p.blocks,
    )
    call_post_op_callback() && post_op_callback(
        output,
        (dev, f, transform, output, input, init, space),
        (;),
    )
end

function column_accumulate_device!(
    ::ClimaComms.CUDADevice,
    f::F,
    transform::T,
    output,
    input,
    init,
    space,
    reverse,
) where {F, T}
    out_fv = Fields.field_values(output)
    mask = Spaces.get_mask(space)
    if !(mask isa DataLayouts.NoMask) && space isa Spaces.FiniteDifferenceSpace
        error("Masks not supported for FiniteDifferenceSpace")
    end
    us = UniversalSize(out_fv)
    cart_inds = cartesian_indices_columnwise(us)
    args = (
        single_column_accumulate!,
        f,
        transform,
        strip_space(output, space),
        strip_space(input, space),
        init,
        space,
        us,
        mask,
        cart_inds,
        reverse,
    )
    (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(us)
    nitems = Ni * Nj * Nh
    threads = threads_via_occupancy(bycolumn_kernel!, args)
    n_max_threads = min(threads, nitems)
    p = linear_partition(nitems, n_max_threads)
    auto_launch!(
        bycolumn_kernel!,
        args;
        threads_s = p.threads,
        blocks_s = p.blocks,
    )
end

function bycolumn_kernel!(
    single_column_function!::S,
    f::F,
    transform::T,
    output,
    input,
    init,
    space,
    us::DataLayouts.UniversalSize,
    mask,
    cart_inds,
    reverse,
) where {S, F, T}
    if space isa Spaces.FiniteDifferenceSpace
        single_column_function!(f, transform, output, input, init, space, reverse)
    else
        tidx = linear_thread_idx()
        if linear_is_valid_index(tidx, us) && tidx ≤ length(unval(cart_inds))
            I = unval(cart_inds)[tidx]
            (i, j, h) = I.I
            ui = CartesianIndex((i, j, 1, 1, h))
            DataLayouts.should_compute(mask, ui) || return nothing
            single_column_function!(
                f,
                transform,
                column(output, i, j, h),
                column(input, i, j, h),
                init,
                column(space, i, j, h),
                reverse,
            )
        end
    end
end