(0.87.1) Make set! a little more ambitious (#3225)

* Make set! a little more ambitious

* Add a test to ensure that halo data is copied when its easy

* Bump version

* Fix bug in test

* Another test bug

* Bump project and keep trying to fix the test

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.87.0"
+version = "0.87.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Fields/set!.jl

@@ -41,14 +41,29 @@ function set!(u::Field, f::Union{Array, CuArray, OffsetArray})
 end
 
 function set!(u::Field, v::Field)
-    # Note: we only copy interior points.
-    # To copy halos use `parent(u) .= parent(v)`.
+    # We implement some niceities in here that attempt to copy halo data,
+    # and revert to copying just interior points if that fails.
 
     if architecture(u) === architecture(v)
-        interior(u) .= interior(v)
+        # Note: we could try to copy first halo point even when halo
+        # regions are a different size. That's a bit more complicated than
+        # the below so we leave it for the future.
+
+        try # to copy halo regions along with interior data
+            parent(u) .= parent(v)
+        catch # this could fail if the halo regions are different sizes?
+            # copy just the interior data
+            interior(u) .= interior(v)
+        end
     else
         v_data = arch_array(architecture(u), v.data)
-        interior(u) .= interior(v_data, location(v), v.grid, v.indices)
+
+        # As above, we permit ourselves a little ambition and try to copy halo data:
+        try
+            parent(u) .= parent(v_data)
+        catch
+            interior(u) .= interior(v_data, location(v), v.grid, v.indices)
+        end
     end
 
     return u

test/test_field.jl

@@ -338,6 +338,40 @@ end
             @test CUDA.@allowscalar v[1, 2, 3] ≈ f(xv[1], yv[2], zv[3])
             @test CUDA.@allowscalar w[1, 2, 3] ≈ f(xw[1], yw[2], zw[3])
             @test CUDA.@allowscalar c[1, 2, 3] ≈ f(xc[1], yc[2], zc[3])
+
+            # Test for Field-to-Field setting on same architecture, and cross architecture.
+            # The behavior depends on halo size: if the halos of two fields are the same, we can
+            # (easily) copy halo data over.
+            # Otherwise, we take the easy way out (for now) and only copy interior data.
+            big_halo = (3, 3, 3)
+            small_halo = (1, 1, 1)
+            domain = (; x=(0, 1), y=(0, 1), z=(0, 1))
+            sz = (1, 1, 1)
+
+            grid = RectilinearGrid(arch, FT; halo=big_halo, size=sz, domain...)
+            a = CenterField(grid)
+            b = CenterField(grid)
+            parent(a) .= 1
+            set!(b, a)
+            @test parent(b) == parent(a)
+
+            grid_with_smaller_halo = RectilinearGrid(arch, FT; halo=small_halo, size=sz, domain...)
+            c = CenterField(grid_with_smaller_halo)
+            set!(c, a)
+            @test interior(c) == interior(a)
+
+            # Cross-architecture setting should have similar behavior
+            if arch isa GPU
+                cpu_grid = RectilinearGrid(CPU(), FT; halo=big_halo, size=sz, domain...)
+                d = CenterField(cpu_grid)
+                set!(d, a)
+                @test parent(d) == Array(parent(a))
+
+                cpu_grid_with_smaller_halo = RectilinearGrid(CPU(), FT; halo=small_halo, size=sz, domain...)
+                e = CenterField(cpu_grid_with_smaller_halo)
+                set!(e, a)
+                @test Array(interior(e)) == Array(interior((a)))
+            end
         end
     end