From d06e82b120e5f14336b999f6e637abb7b8d947c5 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Tue, 23 Jul 2024 09:59:44 -0400
Subject: [PATCH] Possible MPI driver bugfix

---
 examples/hybrid/driver.jl       | 16 ++++++++--------
 ext/cuda/data_layouts_copyto.jl |  2 +-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl
index 8c60b127bf..6b72273018 100644
--- a/examples/hybrid/driver.jl
+++ b/examples/hybrid/driver.jl
@@ -186,6 +186,10 @@ walltime = @elapsed sol = OrdinaryDiffEq.solve!(integrator)
 any(isnan, sol.u[end]) && error("NaNs found in result.")
 
 if is_distributed # replace sol.u on the root processor with the global sol.u
+    global_Y_c_1 =
+        DataLayouts.gather(comms_ctx, Fields.field_values(sol.u[1].c))
+    global_Y_f_1 =
+        DataLayouts.gather(comms_ctx, Fields.field_values(sol.u[1].f))
     if ClimaComms.iamroot(comms_ctx)
         global_h_space = make_horizontal_space(
             horizontal_mesh,
@@ -194,14 +198,10 @@ if is_distributed # replace sol.u on the root processor with the global sol.u
         )
         global_center_space, global_face_space =
             make_hybrid_spaces(global_h_space, z_max, z_elem; z_stretch)
-        global_Y_c_type = Fields.Field{
-            typeof(Fields.field_values(Y.c)),
-            typeof(global_center_space),
-        }
-        global_Y_f_type = Fields.Field{
-            typeof(Fields.field_values(Y.f)),
-            typeof(global_face_space),
-        }
+        global_Y_c_type =
+            Fields.Field{typeof(global_Y_c_1), typeof(global_center_space)}
+        global_Y_f_type =
+            Fields.Field{typeof(global_Y_f_1), typeof(global_face_space)}
         global_Y_type = Fields.FieldVector{
             FT,
             NamedTuple{(:c, :f), Tuple{global_Y_c_type, global_Y_f_type}},
diff --git a/ext/cuda/data_layouts_copyto.jl b/ext/cuda/data_layouts_copyto.jl
index 9a8308fed8..3e53d05c46 100644
--- a/ext/cuda/data_layouts_copyto.jl
+++ b/ext/cuda/data_layouts_copyto.jl
@@ -112,7 +112,7 @@ Base.copyto!(dest::IFH{S, Ni, Nh},        bc::DataLayouts.BroadcastedUnionIFH{S,
 Base.copyto!(dest::IJF{S, Nij},           bc::DataLayouts.BroadcastedUnionIJF{S, Nij}, ::ToCUDA) where {S, Nij} = cuda_copyto!(dest, bc)
 Base.copyto!(dest::IF{S, Ni},             bc::DataLayouts.BroadcastedUnionIF{S, Ni}, ::ToCUDA) where {S, Ni} = cuda_copyto!(dest, bc)
 Base.copyto!(dest::VIFH{S, Nv, Ni, Nh},   bc::DataLayouts.BroadcastedUnionVIFH{S, Nv, Ni, Nh}, ::ToCUDA) where {S, Nv, Ni, Nh} = cuda_copyto!(dest, bc)
-# Base.copyto!(dest::VIJFH{S, Nv, Nij},     bc::DataLayouts.BroadcastedUnionVIJFH{S, Nv, Nij, Nh}, ::ToCUDA) where {S, Nv, Nij, Nh} = cuda_copyto!(dest, bc)
+# Base.copyto!(dest::VIJFH{S, Nv, Nij, Nh}, bc::DataLayouts.BroadcastedUnionVIJFH{S, Nv, Nij, Nh}, ::ToCUDA) where {S, Nv, Nij, Nh} = cuda_copyto!(dest, bc)
 # Base.copyto!(dest::VF{S, Nv},             bc::DataLayouts.BroadcastedUnionVF{S, Nv}, ::ToCUDA) where {S, Nv} = cuda_copyto!(dest, bc)
 # Base.copyto!(dest::DataF{S},              bc::DataLayouts.BroadcastedUnionDataF{S}, ::ToCUDA) where {S} = cuda_copyto!(dest, bc)
 #! format: on