Try using map instead of ntuple

Author: dennisYatunin

src/prognostic_equations/implicit/auto_dense_jacobian.jl

@@ -152,7 +152,7 @@ function update_column_matrices!(alg::AutoDenseJacobian, cache, Y, p, t)
                         field = MatrixFields.get_field(Yₜ_dual, name)
                         @inbounds point(field, level_index, column_index...)[]
                     end
-                ε_coefficients = ForwardDiff.partials(dual_number).values
+                ε_coefficients = ForwardDiff.partials(dual_number)
                 for (jacobian_column_ε_index, (jacobian_column_index, _)) in
                     enumerate(jacobian_index_to_Y_index_map_partition)
                     cartesian_index = (

src/prognostic_equations/implicit/auto_sparse_jacobian.jl

@@ -175,7 +175,7 @@ function jacobian_cache(alg::AutoSparseJacobian, Y, atmos; verbose = true)
         # making memory_for_I_matrix negligible compared to max_memory.
         n_εs_max = (max_memory - memory_for_I_matrix) ÷ memory_per_ε - 1
         n_εs_max >= 1 || error(
-            "Not enough free memory ($(cld(free_memory, 2^30)) GB) to allocate \
+            "Not enough free memory ($(cld(max_memory, 2^30)) GB) to allocate \
              Jacobian ($(cld(memory_for_I_matrix, 2^30)) GB for the identity \
              matrix and another $(cld(memory_per_ε, 2^30)) GB per ε component)",
         )
@@ -237,11 +237,11 @@ function jacobian_cache(alg::AutoSparseJacobian, Y, atmos; verbose = true)
             column_index in column_indices,
             index_pair in DA(collect(Y_index_to_diagonal_color_map))
 
-            ((scalar_index, level_index), diagonal_color) = index_pair
-            color_offset = (partition - 1) * n_εs
+            ((scalar_index, level_index), diagonal_entry_color) = index_pair
+            ε_offset = (partition - 1) * n_εs
             diagonal_ε_index =
-                color_offset < diagonal_color <= color_offset + n_εs ?
-                diagonal_color - color_offset : 0
+                ε_offset < diagonal_entry_color <= ε_offset + n_εs ?
+                diagonal_entry_color - ε_offset : 0
             n_εs_val = Val(ForwardDiff.npartials(eltype(I_matrix_partition_εs)))
             ε_coefficients = ntuple(==(diagonal_ε_index), n_εs_val)
             unrolled_applyat(scalar_index, scalar_names) do name
@@ -285,7 +285,7 @@ function jacobian_cache(alg::AutoSparseJacobian, Y, atmos; verbose = true)
         block_colors = last.(block_Y_index_to_color_map)
 
         map(block_Yₜ_indices) do (scalar_index, level_index)
-            band_colors = ntuple(colors_per_band_matrix_row) do band_index
+            entry_colors = ntuple(colors_per_band_matrix_row) do band_index
                 band = lower_band + band_index - 1
                 level_index_min = band < 0 ? 1 - band : 1
                 level_index_max =
@@ -296,7 +296,7 @@ function jacobian_cache(alg::AutoSparseJacobian, Y, atmos; verbose = true)
                     level_index_min <= level_index <= level_index_max
                 is_color_at_index ? block_colors[level_index + band] : 0
             end
-            ((block_index, level_index), (scalar_index, band_colors))
+            ((block_index, level_index), (scalar_index, entry_colors))
         end
     end
 
@@ -335,43 +335,41 @@ function update_jacobian!(::AutoSparseJacobian, cache, Y, p, dtγ, t)
         set_implicit_precomputed_quantities!(Y_dual, p_dual, t)
         implicit_tendency!(Yₜ_dual, Y_dual, p_dual, t)
 
-        # Move the entries of ∂Yₜ/∂Y * I_matrix_partition from Yₜ_dual into
-        # the blocks of autodiff_matrix. Drop the spatial data from every Field
+        # Move the entries of ∂Yₜ/∂Y * I_matrix_partition from Yₜ_dual into the
+        # blocks of autodiff_matrix. Drop spatial information from every Field
         # to ensure that this kernel stays below the GPU parameter memory limit.
-        Yₜ_dual_no_spaces =
+        Yₜ_dual_data =
             unrolled_map(Fields.field_values, Fields._values(Yₜ_dual))
-        matrix_fields_no_spaces =
+        matrix_fields_data =
             unrolled_map(Fields.field_values, values(autodiff_matrix))
         ClimaComms.@threaded device begin
             # On multithreaded devices, use one thread for each band matrix row.
             # TODO: Modify the map and use one thread for each dual number.
             for column_index in column_indices,
                 index_pair in band_matrix_row_index_to_colors_map
 
-                ((block_index, level_index), (scalar_index, band_colors)) =
+                ((block_index, level_index), (scalar_index, entry_colors)) =
                     index_pair
                 dual_number =
                     unrolled_applyat(scalar_index, scalar_names) do name
-                        data = MatrixFields.get_field(Yₜ_dual_no_spaces, name)
+                        data = MatrixFields.get_field(Yₜ_dual_data, name)
                         @inbounds point(data, level_index, column_index...)[]
                     end
-                ε_coefficients = ForwardDiff.partials(dual_number).values
+                ε_coefficients = ForwardDiff.partials(dual_number)
                 n_εs = length(ε_coefficients)
-                color_offset = (partition - 1) * n_εs
-                unrolled_applyat(block_index, matrix_fields_no_spaces) do data
-                    @inbounds band_entries =
-                        point(data.entries, level_index, column_index...)
-                    preexisting_band_entries = band_entries[]
-                    n_bands_val = Val(length(preexisting_band_entries))
-                    band_entries[] = ntuple(n_bands_val) do band_index
-                        # If this band's color is in the current partition, set
-                        # its entry to the ε coefficient for that color.
-                        # Otherwise, keep the entry already in the block's data.
-                        @inbounds band_color = band_colors[band_index]
-                        color_offset < band_color <= color_offset + n_εs ?
-                        (@inbounds ε_coefficients[band_color - color_offset]) :
-                        (@inbounds preexisting_band_entries[band_index])
-                    end
+                ε_offset = (partition - 1) * n_εs
+                unrolled_applyat(block_index, matrix_fields_data) do block_data
+                    @inbounds entries_data =
+                        point(block_data, level_index, column_index...).entries
+                    entries_data[] =
+                        map(entry_colors, entries_data[]) do entry_color, entry
+                            # If the entry has a color in the current partition,
+                            # set the entry to the ε coefficient for that color.
+                            # Otherwise, keep the value from the block's data.
+                            ε_offset < entry_color <= ε_offset + n_εs ?
+                            (@inbounds ε_coefficients[entry_color - ε_offset]) :
+                            entry
+                        end
                 end
             end
         end