`to_device` desynchronizes spatial data between CPU and GPU versions of array

[ray-chew]

Describe the bug

When arrays are initialized on the CPU and passed to the GPU via ClimaCore.to_device, the spatial data (e.g., grid, grid.full_grid, etc.) appears to be copied independently, resulting in the loss of internal consistency. Consequently, data layouts that should be equivalent are evaluated as false, leading to bugs during subsequent operations or comparisons.

In particular, when layout-dependent operations like column_reduce! and column_accumulate! are used with the GPU version of the input arrays.

To Reproduce

The following code reproduces the issue and shows how the spatial data becomes inconsistent after calling to_device. In this minimal example, Fields.level(axes(...), 1) == axes(...) evaluates to true on the CPU but fails post-transfer, which should not occur.

I could also push this as a unit test on ClimaAtmos, if that is how I should do it.

using ClimaCore:
    Fields, Geometry, Meshes, Quadratures

using ClimaCore
using ClimaCore.CommonSpaces
import ClimaAtmos as CA
import ClimaComms

using CUDA

const FT = Float64

comms_ctx = ClimaComms.SingletonCommsContext()

(; config_file, job_id) = CA.commandline_kwargs()
config = CA.AtmosConfig(config_file; job_id, comms_ctx)

config.parsed_args["topography"] = "Earth";
config.parsed_args["topo_smoothing"] = false;
config.parsed_args["mesh_warp_type"] = "Linear";
(; parsed_args) = config

# Create meshes and spaces
h_elem = 16
nh_poly = 3
z_max = 42e3
z_elem = 33
dz_bottom = 300.0
radius = 6.371229e6

quad = Quadratures.GLL{nh_poly + 1}()
horizontal_mesh = CA.cubed_sphere_mesh(; radius, h_elem)
h_space = CA.make_horizontal_space(horizontal_mesh, quad, comms_ctx, false)
z_stretch = Meshes.HyperbolicTangentStretching(dz_bottom)
center_space, face_space =
    CA.make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args)

ᶜlocal_geometry = Fields.local_geometry_field(center_space)
ᶠlocal_geometry = Fields.local_geometry_field(face_space)

# create Y
Yc = map(ᶜlocal_geometry) do lg
    return (; ρ = FT(1.0), u_phy = FT(0), v_phy = FT(0), T = FT(0), qt = FT(0))
end
Yf = map(ᶠlocal_geometry) do lg
    return (; u₃ = Geometry.Covariant3Vector(FT(0), lg))
end
Y = Fields.FieldVector(c = Yc, f = Yf)

ᶜz = Fields.coordinate_field(Y.c).z

z_level = similar(Fields.level(ᶜz, 1), FT)

# Prints True
print("Axes check on the CPU: ")
println(Fields.level(axes(ᶜz),1) == axes(z_level))

A = ClimaCore.to_device(ClimaComms.CUDADevice(),ᶜz)
B = ClimaCore.to_device(ClimaComms.CUDADevice(),z_level)

# Should also print True, but prints False
# This has the implication that column_accumulate
# and column_reduce do not work on inputs with
# mixed VIJFH and IJFH fields.
print("Axes check on the GPU: ")
println(Fields.level(axes(A), 1) == axes(B))

# Let's look at what is wrong...
C = Fields.level(axes(A),1)
D = axes(B)

# Prints True as it should be...
print("C == D: ")
println(getfield(C.grid.full_grid.face_local_geometry, :array) == getfield(D.grid.full_grid.face_local_geometry, :array))

print("C === D: ")
# ... but physical equivalence evaluates to False
println(getfield(C.grid.full_grid.face_local_geometry, :array) === getfield(D.grid.full_grid.face_local_geometry, :array))

Project

I am using the .buildkite project:

# Project.toml
[deps]
ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
ArtifactWrappers = "a14bc488-3040-4b00-9dc1-f6467924858a"
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
ClimaAnalysis = "29b5916a-a76c-4e73-9657-3c8fd22e65e6"
ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
ClimaCoreSpectra = "c2caaa1d-32ae-4754-ba0d-80e7561362e9"
ClimaCoreTempestRemap = "d934ef94-cdd4-4710-83d6-720549644b70"
ClimaDiagnostics = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
ClimaReproducibilityTests = "e0c89595-00ba-42a9-9f9b-061ef3dc23a1"
ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
Debugger = "31a5f54b-26ea-5ae9-a837-f05ce5417438"
DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
NullBroadcasts = "0d71be07-595a-4f89-9529-4065a4ab43a6"
OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
Poppler_jll = "9c32591e-4766-534b-9725-b71a8799265b"
PrecompileCI = "76d61242-8ec2-4c91-8455-3234246697a2"
PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
ProfileCanvas = "efd6af41-a80b-495e-886c-e51b0c7d77a3"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
SnoopCompile = "aa65fe97-06da-5843-b5b1-d5d13cad87d2"
SnoopCompileCore = "e2b509da-e806-4183-be48-004708413034"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"

[compat]
CairoMakie = "0.11, 0.12"
ClimaCoreSpectra = "0.1"
ClimaCoreTempestRemap = "0.3"
JET = "0.9"
PrettyTables = "2"
ProfileCanvas = "0.1"
SnoopCompileCore = "3"
julia = "1.10"

The Manifest.toml is too long. Including it as a `.txt` attachment.
[Manifest-v1.11.txt](https://github.com/user-attachments/files/20474817/Manifest-v1.11.txt)

System details

Any relevant system information:

• Julia version:

╰─$ julia --version
julia version 1.11.5

• operating system:

Description:	Manjaro Linux
Release:	25.0.3

• modules loaded on cluster (module list)

╰─$ module list
No Modulefiles Currently Loaded.

Related issues / PRs

Appears to be related to #2312 and #2260 .

I found the to_device function here: https://clima.github.io/ClimaCore.jl/dev/faq/#Moving-objects-between-the-CPU-and-GPU.

[charleskawczynski]

This is not a bug (this is what Adapt.jl does, too), we documented that to_device returns a copy (which doesn't satisfy ===):

"""
    out = to_device(device, x::Union{
        DataLayouts.AbstractData,
        Spaces.AbstractSpace,
        Fields.Field,
        Fields.FieldVector,
    })

Move `x` to the given `device`.

This is particularly useful to move different types of `Space.AbstractSpace`s,
`Fields.Field`s, and `Fields.FieldVector`s from CPUs to GPUs and viceversa.

If the input is already defined on the target device, returns a copy.

This means that `out === x` will not in general be satisfied.
"""

Also, this is not related to #2260, as that only pertains to our ability to move broadcasted objects between the host and device.

#2312 relates to our ability to perform operations on separate spaces, which is a design choice that we do have control over (as this check is performed inside our broadcast machinery).

For this issue, we really can't, as there is no way to preserve === between the host and device while using Adapt. So, I'd recommend closing this as not planned.

[ray-chew]

Thanks for your clarification, Charlie. I probably do not fully understand your response, as I am still new to ClimaCore and Julia.

However, this issue is not about preserving === between host and device, and regarding what you mentioned in your response, I fully understand why out === x will always evaluate as false. There is no way around that.

My concern is slightly different:
For example, if A and B are VIJFH and IJFH fields on the device, and they originated from the same mesh on the CPU, then I would expect Fields.level(axes(A),1).grid == axes(B).grid (or Fields.level(axes(A), 1) == axes(B)) to return true, as they do on the CPU. Currently, this is not the case after applying to_device, and it seems to stem from how layout metadata is deeply copied or reconstructed in a way that breaks logical equivalence, even when the structures should be identical.

I mentioned this briefly at the AMIP Council yesterday that I understand this might be a deliberate design decision. If so, is it reasonable to update the documentation to indicate that structural equality (==) between layouts or derived views may not be preserved across to_device?

[charleskawczynski]

However, this issue is not about preserving === between host and device, and regarding what you mentioned in your response, I fully understand why out === x will always evaluate as false. There is no way around that.
👍🏻

My concern is slightly different: For example, if A and B are VIJFH and IJFH fields on the device, and they originated from the same mesh on the CPU then I would expect Fields.level(axes(A),1).grid == axes(B).grid (or Fields.level(axes(A), 1) == axes(B)) to return true, as they do on the CPU. Currently, this is not the case after applying to_device, and it seems to stem from how layout metadata is deeply copied or reconstructed in a way that breaks logical equivalence, even when the structures should be identical.

They don't because they are copies and == falls back to ===, per the julia docs:

==(x, y)

  Generic equality operator. Falls back to ===. ...

I recommend that we instead write our code as:

A = ClimaCore.to_device(ClimaComms.CUDADevice(),ᶜz)
B = Fields.level(A, 1)

Which is also twice as memory efficient(!) as

A = ClimaCore.to_device(ClimaComms.CUDADevice(),ᶜz)
B = ClimaCore.to_device(ClimaComms.CUDADevice(),z_level)

If so, is it reasonable to update the documentation to indicate that structural equality (==) between layouts or derived views may not be preserved across to_device?

I feel that the docs are sufficient, but I don't suppose it would hurt to point out that == falls back to ===.

Also, here is a smaller reproducer:

using ClimaCore: Spaces, Fields, Geometry
using ClimaCore
using ClimaCore.CommonSpaces
import ClimaComms
ClimaComms.@import_required_backends

ᶜspace = ExtrudedCubedSphereSpace(;
    z_elem = 10,
    z_min = 0,
    z_max = 1,
    radius = 10,
    h_elem = 10,
    n_quad_points = 4,
    staggering = CellCenter()
);
ᶠspace = Spaces.face_space(ᶜspace);
ᶜlocal_geometry = Fields.local_geometry_field(ᶜspace);
ᶠlocal_geometry = Fields.local_geometry_field(ᶠspace);

# create Y
Yc = map(ᶜlocal_geometry) do lg
    FT = Geometry.undertype(typeof(lg))
    (; ρ = FT(1.0), u_phy = FT(0), v_phy = FT(0), T = FT(0), qt = FT(0))
end;
Yf = map(ᶠlocal_geometry) do lg
    FT = Geometry.undertype(typeof(lg))
    (; u₃ = Geometry.Covariant3Vector(FT(0), lg))
end;
Y = Fields.FieldVector(c = Yc, f = Yf);

ᶜz = Fields.coordinate_field(Y.c).z;

z_level = similar(Fields.level(ᶜz, 1));

using Test
@test Fields.level(axes(ᶜz),1) == axes(z_level) # passes
A = ClimaCore.to_device(ClimaComms.CPUSingleThreaded(),ᶜz)
B = ClimaCore.to_device(ClimaComms.CPUSingleThreaded(),z_level)
@test Fields.level(axes(A), 1) == axes(B) # fails

I hope that clarifies things!

By the way, the way I found this out was to use Julia's @which: @which Fields.level(axes(A), 1) == axes(B), which points to julia Base's == implementation

[charleskawczynski]

Shall we close this?

[ray-chew]

Thanks for your responses and tips, Charlie! Really helpful stuff.

Which is also twice as memory efficient(!) as

A = ClimaCore.to_device(ClimaComms.CUDADevice(),ᶜz)
B = ClimaCore.to_device(ClimaComms.CUDADevice(),z_level)

Thank you for pointing this out, I will try to avoid this in the implementation and just wanted to be clear about the behaviour of to_device in my MWE.

Shall we close this?

We could, but I will document the behaviour of to_device. Because this statement is not strictly true:

They don't because they are copies and == falls back to ===, per the julia docs: ...

In the above MWEs, we saw that to_device makes full/deep copies of our CPU arrays.

But if the arrays all belong to the same cache structure, to_device makes the expected and correct copying. I try to illustrate what I mean with a code below:

julia> AA = ClimaCore.to_device(ClimaComms.CUDADevice(),Y.c.u_phy)
julia> BB = ClimaCore.to_device(ClimaComms.CUDADevice(),Y.c,ρ)
julia> typeof(BB) == typeof(AA)
true
julia> CC = ClimaCore.to_device(ClimaComms.CUDADevice(),similar(Y.c.ρ,FT))
julia> typeof(BB) == typeof(CC)
false

I believe an extension of the existing documentation should inform the user that all arrays from derived views, e.g., obtained using Fields.level or similar, will not work correctly with to_device, but all "pure" ClimaCore arrays would work. Do you agree?

[charleskawczynski]

I'm not sure I understand-- if there is a difference in the types, then it's possible that there is an issue here. But you're example looks a bit broken (A is not defined, and Y.c,ρ is a typo), Can you provide a minimum working / breaking example? This should be a good starting point:

ENV["CLIMACOMMS_DEVICE"]="CUDA";
ClimaComms.@import_required_backends
using ClimaCore.CommonSpaces
using ClimaCore
using ClimaCore: Fields
using ClimaComms
space = ExtrudedCubedSphereSpace(;
    z_elem = 10,
    z_min = 0,
    z_max = 1,
    radius = 10,
    h_elem = 10,
    n_quad_points = 4,
    staggering = CellCenter()
)
y = Fields.ones(space);
x = Fields.ones(space);
A = ClimaCore.to_device(ClimaComms.CUDADevice(),x);
B = ClimaCore.to_device(ClimaComms.CUDADevice(),y);
C = ClimaCore.to_device(ClimaComms.CUDADevice(),similar(x));
typeof(A) == typeof(B) # true
typeof(B) == typeof(C) # true

[ray-chew]

Thanks for pointing out the errors in my example, and I apologise for shoddiness!

Edit: Give me awhile to think about my response below. I think I was conflating two issues and got confused.

Charlie's modified MWE

I tried modifying your MWE, and it behaves as expected:

Charlie's modified MWE

ENV["CLIMACOMMS_DEVICE"]="CPUSingleThreaded";
using ClimaCore.CommonSpaces
using ClimaCore
using ClimaCore: Fields
using ClimaComms
ClimaComms.@import_required_backends

using ClimaCore: Spaces, Utilities, Operators

space = ExtrudedCubedSphereSpace(;
    z_elem = 10,
    z_min = 0,
    z_max = 1,
    radius = 10,
    h_elem = 10,
    n_quad_points = 4,
    staggering = CellCenter()
)

# ᶜspace = Spaces.center_space(space)

x = Fields.ones(space);
y = Fields.ones(space);

Y = Fields.FieldVector(x = x, y = y)

A = ClimaCore.to_device(ClimaComms.CUDADevice(),Y.x);
B = ClimaCore.to_device(ClimaComms.CUDADevice(),Y.y);
C = ClimaCore.to_device(ClimaComms.CUDADevice(),similar(Y.x));
println(typeof(A) == typeof(B)) # true
println(typeof(B) == typeof(C)) # true

Ray's originals MWE

But below is the closer-to-real example reproduced with the MWE from the issue description above. Maybe I am initialising something wrongly? But in short, yes, the behaviour I'd expect is exactly like what your MWE showed. But the behaviour I see is the one in the MWE below.

using ClimaCore:
    Fields, Geometry, Meshes, Quadratures

using ClimaCore
using ClimaCore.CommonSpaces
import ClimaAtmos as CA
import ClimaComms

using CUDA

const FT = Float64

comms_ctx = ClimaComms.SingletonCommsContext()

(; config_file, job_id) = CA.commandline_kwargs()
config = CA.AtmosConfig(config_file; job_id, comms_ctx)

config.parsed_args["topography"] = "Earth";
config.parsed_args["topo_smoothing"] = false;
config.parsed_args["mesh_warp_type"] = "Linear";
(; parsed_args) = config

# Create meshes and spaces
h_elem = 16
nh_poly = 3
z_max = 42e3
z_elem = 33
dz_bottom = 300.0
radius = 6.371229e6

quad = Quadratures.GLL{nh_poly + 1}()
horizontal_mesh = CA.cubed_sphere_mesh(; radius, h_elem)
h_space = CA.make_horizontal_space(horizontal_mesh, quad, comms_ctx, false)
z_stretch = Meshes.HyperbolicTangentStretching(dz_bottom)
center_space, face_space =
    CA.make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args)

ᶜlocal_geometry = Fields.local_geometry_field(center_space)
ᶠlocal_geometry = Fields.local_geometry_field(face_space)

# create Y
Yc = map(ᶜlocal_geometry) do lg
    return (; ρ = FT(1.0), u_phy = FT(0), v_phy = FT(0), T = FT(0), qt = FT(0))
end
Yf = map(ᶠlocal_geometry) do lg
    return (; u₃ = Geometry.Covariant3Vector(FT(0), lg))
end
Y = Fields.FieldVector(c = Yc, f = Yf)

A = ClimaCore.to_device(ClimaComms.CUDADevice(),Y.c.ρ);
B = ClimaCore.to_device(ClimaComms.CUDADevice(),Y.c.T);
C = ClimaCore.to_device(ClimaComms.CUDADevice(),similar(Y.c.ρ));
println(typeof(A) == typeof(B)) # true
println(typeof(B) == typeof(C)) # false