Simplify distributed_remapping.interpolate

I spent many hours tracking down
#2108 and could not find the
root issue.

I decided to take a different approach and simplify redefine
`interpolate` in terms of `interpolate!`.

Author: Sbozzolo

NEWS.md

@@ -2,6 +2,13 @@ ClimaCore.jl Release Notes
 ========================
 
 main
+-------
+
+- Fixed bug in distributed remapping with CUDA. Sometimes, `ClimaCore` would not
+  properly fill the output arrays with the correct values. This is now fixed. PR
+  [2169](https://github.com/CliMA/ClimaCore.jl/pull/2169)
+
+v0.14.24
 -------
 
  - A new `Adapt` wrapper was added, `to_device`, which allows users to adapt datalayouts, spaces, fields, and fieldvectors between the cpu and gpu. PR [2159](https://github.com/CliMA/ClimaCore.jl/pull/2159).

src/Remapping/distributed_remapping.jl

@@ -744,29 +744,6 @@ function _reset_interpolated_values!(remapper::Remapper)
     fill!(remapper._interpolated_values, 0)
 end
 
-"""
-    _collect_and_return_interpolated_values!(remapper::Remapper,
-                                             num_fields::Int)
-
-Perform an MPI call to aggregate the interpolated points from all the MPI processes and save
-the result in the local state of the `remapper`. Only the root process will return the
-interpolated data.
-
-`_collect_and_return_interpolated_values!` is type-unstable and allocates new return arrays.
-
-`num_fields` is the number of fields that have been interpolated in this batch.
-"""
-function _collect_and_return_interpolated_values!(
-    remapper::Remapper,
-    num_fields::Int,
-)
-    return ClimaComms.reduce(
-        remapper.comms_ctx,
-        remapper._interpolated_values[remapper.colons..., 1:num_fields],
-        +,
-    )
-end
-
 function _collect_interpolated_values!(
     dest,
     remapper::Remapper,
@@ -796,19 +773,6 @@ function _collect_interpolated_values!(
     return nothing
 end
 
-"""
-    batched_ranges(num_fields, buffer_length)
-
-Partition the indices from 1 to num_fields in such a way that no range is larger than
-buffer_length.
-"""
-function batched_ranges(num_fields, buffer_length)
-    return [
-        (i * buffer_length + 1):(min((i + 1) * buffer_length, num_fields)) for
-        i in 0:(div((num_fields - 1), buffer_length))
-    ]
-end
-
 """
    interpolate(remapper::Remapper, fields)
    interpolate!(dest, remapper::Remapper, fields)
@@ -860,58 +824,21 @@ int12 = interpolate(remapper, [field1, field2])
 ```
 """
 function interpolate(remapper::Remapper, fields)
-
+    ArrayType = ClimaComms.array_type(remapper.space)
+    FT = Spaces.undertype(remapper.space)
     only_one_field = fields isa Fields.Field
-    if only_one_field
-        fields = [fields]
-    end
-
-    for field in fields
-        axes(field) == remapper.space ||
-            error("Field is defined on a different space than remapper")
-    end
-
-    isa_vertical_space = remapper.space isa Spaces.FiniteDifferenceSpace
-
-    index_field_begin, index_field_end =
-        1, min(length(fields), remapper.buffer_length)
 
-    # Partition the indices in such a way that nothing is larger than
-    # buffer_length
-    index_ranges = batched_ranges(length(fields), remapper.buffer_length)
+    interpolated_values_dim..., _buffer_length =
+        size(remapper._interpolated_values)
 
-    cat_fn = (l...) -> cat(l..., dims = length(remapper.colons) + 1)
+    allocate_extra = only_one_field ? () : (length(fields),)
+    dest = ArrayType(zeros(FT, interpolated_values_dim..., allocate_extra...))
 
-    interpolated_values = mapreduce(cat_fn, index_ranges) do range
-        num_fields = length(range)
-
-        # Reset interpolated_values. This is needed because we collect distributed results
-        # with a + reduction.
-        _reset_interpolated_values!(remapper)
-        # Perform the interpolations (horizontal and vertical)
-        _set_interpolated_values!(
-            remapper,
-            view(fields, index_field_begin:index_field_end),
-        )
-
-        if !isa_vertical_space
-            # For spaces with an horizontal component, reshape the output so that it is a nice grid.
-            _apply_mpi_bitmask!(remapper, num_fields)
-        else
-            # For purely vertical spaces, just move to _interpolated_values
-            remapper._interpolated_values .= remapper._local_interpolated_values
-        end
-
-        # Finally, we have to send all the _interpolated_values to root and sum them up to
-        # obtain the final answer. Only the root will contain something useful.
-        return _collect_and_return_interpolated_values!(remapper, num_fields)
-    end
-
-    # Non-root processes
-    isnothing(interpolated_values) && return nothing
-
-    return only_one_field ? interpolated_values[remapper.colons..., begin] :
-           interpolated_values
+    # interpolate! has an MPI call, so it is important to return after it is
+    # called, not before!
+    interpolate!(dest, remapper, fields)
+    ClimaComms.iamroot(remapper.comms_ctx) || return nothing
+    return dest
 end
 
 # dest has to be allowed to be nothing because interpolation happens only on the root
@@ -927,13 +854,18 @@ function interpolate!(
     end
     isa_vertical_space = remapper.space isa Spaces.FiniteDifferenceSpace
 
+    for field in fields
+        axes(field) == remapper.space ||
+            error("Field is defined on a different space than remapper")
+    end
+
     if !isnothing(dest)
         # !isnothing(dest) means that this is the root process, in this case, the size have
         # to match (ignoring the buffer_length)
         dest_size = only_one_field ? size(dest) : size(dest)[1:(end - 1)]
 
         dest_size == size(remapper._interpolated_values)[1:(end - 1)] || error(
-            "Destination array is not compatible with remapper (size mismatch)",
+            "Destination array is not compatible with remapper (size mismatch), $dest_size",
         )
 
         expected_array_type =

test/Remapping/distributed_remapping.jl

@@ -31,14 +31,6 @@ atexit() do
     global_logger(prev_logger)
 end
 
-@testset "Utils" begin
-    # batched_ranges(num_fields, buffer_length)
-    @test Remapping.batched_ranges(1, 1) == [1:1]
-    @test Remapping.batched_ranges(1, 2) == [1:1]
-    @test Remapping.batched_ranges(2, 2) == [1:2]
-    @test Remapping.batched_ranges(3, 2) == [1:2, 3:3]
-end
-
 on_gpu = device isa ClimaComms.CUDADevice
 with_mpi = context isa ClimaComms.MPICommsContext