Add on_architecture method for XESMF.Regridder (#4871)

* Add on_architecture method for XESMF.Regridder to support architecture-specific configurations

* Add architecture-specific regridder tests for CPU and GPU

* Update test_xesmf.jl

---------

Co-authored-by: Navid C. Constantinou <[email protected]>

Author: xkykai

ext/OceananigansXESMFExt.jl

@@ -7,6 +7,7 @@ using Oceananigans.Fields: AbstractField, topology, location
 using Oceananigans.Grids: AbstractGrid, λnodes, φnodes, Center, Face, total_length
 
 import Oceananigans.Fields: regrid!
+import Oceananigans.Architectures: on_architecture
 import XESMF: Regridder
 
 node_array(ξ::AbstractMatrix, Nx, Ny) = view(ξ, 1:Nx, 1:Ny)
@@ -145,6 +146,11 @@ function Regridder(dst_field::AbstractField, src_field::AbstractField; method="c
     return XESMF.Regridder(method, weights, temp_src, temp_dst)
 end
 
+on_architecture(on, r::XESMF.Regridder) = XESMF.Regridder(on_architecture(on, r.method),
+                                                          on_architecture(on, r.weights),
+                                                          on_architecture(on, r.src_temp),
+                                                          on_architecture(on, r.dst_temp))
+
 """
     regrid!(dst_field, regrider::XESMF.Regridder, src_field)
 

test/test_xesmf.jl

@@ -48,6 +48,13 @@ for arch in archs
                 @test regridder.weights isa CUDA.CUSPARSE.CuSparseMatrixCSC
             end
 
+            if arch isa GPU
+                cpu_regridder = on_architecture(CPU(), regridder)
+                @test cpu_regridder.weights isa SparseMatrixCSC
+                gpu_regridder = on_architecture(GPU(), cpu_regridder)
+                @test gpu_regridder.weights isa CUDA.CUSPARSE.CuSparseMatrixCSC
+            end
+
             regrid!(dst_field, regridder, src_field)
 
             # ∫ dst_field dA ≈ ∫ src_field dA