@@ -34,12 +34,12 @@ very large memory requirements at higher vertical resolutions.
3434
3535When the number of values in each column is very large, computing the entire
3636dense matrix in a single evaluation of `implicit_tendency!` can be too expensive
37- to compile and run. So, the dual number components are split into batches with a
38- maximum size of `max_simultaneous_derivatives`, and we call `implicit_tendency!`
39- once for each batch . That is, if the batch size is ``s``, then the first batch
40- evaluates the coefficients of ``ε₁`` through ``εₛ``, the second evaluates the
41- coefficients of ``εₛ₊₁`` through ``ε₂ₛ``, and so on until ``εₙ``. The default
42- batch size is 32.
37+ to compile and run. So, the dual number components are split into partitions
38+ with a maximum size of `max_simultaneous_derivatives`, and we call
39+ `implicit_tendency!` once for each partition . That is, if the partition size is
40+ ``s``, then the first partition evaluates the coefficients of ``ε₁`` through
41+ ``εₛ``, the second evaluates the coefficients of ``εₛ₊₁`` through ``ε₂ₛ``, and so
42+ on until ``εₙ``. The default partition size is 32.
4343"""
4444struct AutoDenseJacobian{S} <: JacobianAlgorithm end
4545AutoDenseJacobian (max_simultaneous_derivatives = 32 ) =
@@ -53,11 +53,11 @@ function jacobian_cache(alg::AutoDenseJacobian, Y, atmos)
5353 DA = ClimaComms. array_type (Y)
5454
5555 FT_dual = ForwardDiff. Dual{Jacobian, FT, max_simultaneous_derivatives (alg)}
56- Y_dual = replace_parent_type (Y, FT_dual)
57- Yₜ_dual = similar (Y_dual)
5856 precomputed_dual =
5957 replace_parent_type (implicit_precomputed_quantities (Y, atmos), FT_dual)
6058 scratch_dual = replace_parent_type (temporary_quantities (Y, atmos), FT_dual)
59+ Y_dual = replace_parent_type (Y, FT_dual)
60+ Yₜ_dual = similar (Y_dual)
6161
6262 N = length (Fields. column (Y, 1 , 1 , 1 ))
6363 n_columns = Fields. ncolumns (Y. c)
@@ -72,10 +72,10 @@ function jacobian_cache(alg::AutoDenseJacobian, Y, atmos)
7272 I_matrix = reshape (I_column_matrix, N, N, 1 )
7373
7474 return (;
75- Y_dual,
76- Yₜ_dual,
7775 precomputed_dual,
7876 scratch_dual,
77+ Y_dual,
78+ Yₜ_dual,
7979 column_matrices,
8080 column_lu_factors,
8181 column_lu_vectors,
@@ -85,7 +85,7 @@ function jacobian_cache(alg::AutoDenseJacobian, Y, atmos)
8585end
8686
8787function update_column_matrices! (alg:: AutoDenseJacobian , cache, Y, p, t)
88- (; Y_dual, Yₜ_dual, precomputed_dual, scratch_dual , column_matrices) = cache
88+ (; precomputed_dual, scratch_dual, Y_dual, Yₜ_dual , column_matrices) = cache
8989 device = ClimaComms. device (Y. c)
9090 column_indices = column_index_iterator (Y)
9191 scalar_names = scalar_field_names (Y)
@@ -99,48 +99,47 @@ function update_column_matrices!(alg::AutoDenseJacobian, cache, Y, p, t)
9999 for jacobian_index_to_Y_index_map_partition in
100100 ClimaComms. threadable (device, jacobian_index_to_Y_index_map_partitions)
101101
102- # Add a unique ε to Y for each Jacobian column index in this batch . With
102+ # Add a unique ε to each value in Y that is part of this partition . With
103103 # Y_col and Yᴰ_col denoting the columns of Y and Y_dual at column_index,
104- # set Yᴰ_col to Y_col + I[:, jacobian_indices_in_partition ] * εs, where
105- # I is the identity matrix for Y_col (i.e., the value of ∂Y_col/∂Y_col),
106- # εs is a vector of max_simultaneous_derivatives(alg) dual number
107- # components, and jacobian_indices_in_partition is equal to
104+ # set Yᴰ_col to Y_col + I[:, jacobian_column_indices ] * εs, where I is
105+ # the identity matrix for Y_col (i.e., the value of ∂Y_col/∂Y_col), εs
106+ # is a vector of max_simultaneous_derivatives(alg) dual number
107+ # components, and jacobian_column_indices is equal to
108108 # first.(jacobian_index_to_Y_index_map_partition).
109109 Y_dual .= Y
110110 ClimaComms. @threaded device begin
111111 # On multithreaded devices, assign one thread to each combination of
112- # spatial column index and Jacobian index in this batch .
112+ # spatial column index and Jacobian index in this partition .
113113 for column_index in column_indices,
114- (ε_index , (_, (scalar_index, level_index))) in
114+ (diagonal_ε_index , (_, (scalar_index, level_index))) in
115115 enumerate (jacobian_index_to_Y_index_map_partition)
116116
117- Y_partials =
118- ntuple (== (ε_index), Val (max_simultaneous_derivatives (alg)))
119- Y_dual_εs_value = ForwardDiff. Dual {Jacobian} (0 , Y_partials)
117+ n_εs_val = Val (max_simultaneous_derivatives (alg))
118+ Y_dual_ε_coefficients = ntuple (== (diagonal_ε_index), n_εs_val)
120119 unrolled_applyat (scalar_index, scalar_names) do name
121120 field = MatrixFields. get_field (Y_dual, name)
122121 @inbounds point (field, level_index, column_index... )[] +=
123- Y_dual_εs_value
122+ ForwardDiff . Dual {Jacobian} ( 0 , Y_dual_ε_coefficients)
124123 end
125124 end
126125 end
127126
128- # Compute this batch's portions of ∂p/∂Y and ∂Yₜ/∂Y.
127+ # Compute this partition of ∂p/∂Y and ∂Yₜ/∂Y.
129128 set_implicit_precomputed_quantities! (Y_dual, p_dual, t)
130129 implicit_tendency! (Yₜ_dual, Y_dual, p_dual, t)
131130
132- # Copy this batch's portion of ∂Yₜ/∂Y into column_matrices. With Yₜ_col
133- # and Yₜᴰ_col denoting the columns of Yₜ and Yₜ_dual at column_index, and
131+ # Copy this partition of ∂Yₜ/∂Y into column_matrices. With Yₜ_col and
132+ # Yₜᴰ_col denoting the columns of Yₜ and Yₜ_dual at column_index, and
134133 # with col_matrix denoting the matrix at the corresponding matrix_index
135134 # in column_matrices, copy the coefficients of the εs in Yₜᴰ_col into
136135 # col_matrix, where the previous steps have set Yₜᴰ_col to
137- # Yₜ_col + (∂Yₜ_col/∂Y_col)[:, jacobian_indices_in_batch ] * εs. In other
138- # words, set col_matrix[jacobian_row_index, jacobian_column_index] to
139- # ∂Yₜ_col[jacobian_row_index]/∂Y_col[jacobian_column_index], obtaining
140- # this derivative from the coefficient of εs[ε_index] in
141- # Yₜᴰ_col[jacobian_row_index], where ε_index is the index of
142- # jacobian_column_index in jacobian_indices_in_batch. After all batches
143- # are processed, col_matrix contains the full Jacobian ∂Yₜ_col/∂Y_col .
136+ # Yₜ_col + (∂Yₜ_col/∂Y_col)[:, jacobian_column_indices ] * εs. In
137+ # other words, set col_matrix[jacobian_row_index, jacobian_column_index]
138+ # to ∂Yₜ_col[jacobian_row_index]/∂Y_col[jacobian_column_index],
139+ # obtaining this derivative from the coefficient of
140+ # εs[jacobian_column_ε_index] in Yₜᴰ_col[jacobian_row_index], where
141+ # jacobian_column_ε_index is the index of jacobian_column_index in
142+ # jacobian_column_indices .
144143 ClimaComms. @threaded device begin
145144 # On multithreaded devices, assign one thread to each combination of
146145 # spatial column index and scalar level index.
@@ -153,16 +152,16 @@ function update_column_matrices!(alg::AutoDenseJacobian, cache, Y, p, t)
153152 field = MatrixFields. get_field (Yₜ_dual, name)
154153 @inbounds point (field, level_index, column_index... )[]
155154 end
156- Yₜ_partials = ForwardDiff. partials (Yₜ_dual_value)
157- for (ε_index , (jacobian_column_index, _)) in
155+ Yₜ_dual_ε_coefficients = ForwardDiff. partials (Yₜ_dual_value)
156+ for (jacobian_column_ε_index , (jacobian_column_index, _)) in
158157 enumerate (jacobian_index_to_Y_index_map_partition)
159158 cartesian_index = (
160159 jacobian_row_index,
161160 jacobian_column_index,
162161 matrix_index,
163162 )
164163 @inbounds column_matrices[cartesian_index... ] =
165- Yₜ_partials[ε_index ]
164+ Yₜ_dual_ε_coefficients[jacobian_column_ε_index ]
166165 end
167166 end
168167 end
0 commit comments