Fix restarts with GPUs (#2309)

ClimaLand buildkite started failing with ClimaCore 0.14.31 with the
following error:
```julia
Got exception outside of a @test
  LoadError: ArgumentError: Illegal conversion of a CUDA.DeviceMemory to a Ptr{Bool}
  Stacktrace:
    [1] convert(T::Type{Ptr{Bool}}, mem::CUDA.DeviceMemory)
      @ CUDA /central/scratch/esm/slurm-buildkite/climaland-ci/depot/default/packages/CUDA/oymHm/lib/cudadrv/memory.jl:16
    [2] convert(::Type{Ptr{Bool}}, managed::CUDA.Managed{CUDA.DeviceMemory})
      @ CUDA /central/scratch/esm/slurm-buildkite/climaland-ci/depot/default/packages/CUDA/oymHm/src/memory.jl:583
    [3] unsafe_convert(typ::Type{Ptr{Bool}}, x::CUDA.CuArray{Bool, 4, CUDA.DeviceMemory})
      @ CUDA /central/scratch/esm/slurm-buildkite/climaland-ci/depot/default/packages/CUDA/oymHm/src/array.jl:448
    [4] unsafe_convert(::Type{Ptr{Nothing}}, a::CUDA.CuArray{Bool, 4, CUDA.DeviceMemory})
      @ Base ./pointer.jl:66
    [5] h5a_write(attr_hid::HDF5.Attribute, mem_type_id::HDF5.Datatype, buf::CUDA.CuArray{Bool, 4, CUDA.DeviceMemory})
      @ HDF5.API /central/scratch/esm/slurm-buildkite/climaland-ci/depot/default/packages/HDF5/Z859u/src/api/functions.jl:438
    [6] write_attribute(attr::HDF5.Attribute, memtype::HDF5.Datatype, x::CUDA.CuArray{Bool, 4, CUDA.DeviceMemory})
      @ HDF5 /central/scratch/esm/slurm-buildkite/climaland-ci/depot/default/packages/HDF5/Z859u/src/attributes.jl:154
    [7] write_attribute(parent::HDF5.Dataset, name::String, data::CUDA.CuArray{Bool, 4, CUDA.DeviceMemory}; pv::@kwargs{})
      @ HDF5 /central/scratch/esm/slurm-buildkite/climaland-ci/depot/default/packages/HDF5/Z859u/src/attributes.jl:205
    [8] write_attribute(parent::HDF5.Dataset, name::String, data::CUDA.CuArray{Bool, 4, CUDA.DeviceMemory})
      @ HDF5
      /central/scratch/esm/slurm-buildkite/climaland-ci/depot/default/packages/HDF5/Z859u/src/attributes.jl:202
```

I tracked this down to trying to write a CUArray as an attribute of an
HDF5 file. To fix this, I first move the array to the CPU.

As far as I understand, attributes are meant to be small (`h5py`
recommends keeping them under 64kb), so maybe the best solution would be
to write them as datasets instead.

This also exposed that InputOutput is not tested on GPU (or not tested
with a mask).

Add CUDA tests for InputOutput

Set mask to avoid integer division

Print debug

Add more debug printing

Fix mask initialization

Remove debug prints

Author: Sbozzolo

.buildkite/pipeline.yml

@@ -1179,6 +1179,14 @@ steps:
         key: unit_spectralelement2d
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/InputOutput/unit_spectralelement2d.jl"
 
+      - label: "Unit: spectralelement2d"
+        key: unit_spectralelement2d_gpu
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/InputOutput/unit_spectralelement2d.jl"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
       - label: "Unit: hybrid2dbox"
         key: unit_hybrid2dbox
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/InputOutput/unit_hybrid2dbox.jl"
@@ -1203,6 +1211,14 @@ steps:
         key: unit_hybrid3dcubedsphere_topography
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/InputOutput/unit_hybrid3dcubedsphere_topography.jl"
 
+      - label: "Unit: hybrid3dcubedsphere topography"
+        key: unit_hybrid3dcubedsphere_topography_gpu
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/InputOutput/unit_hybrid3dcubedsphere_topography.jl"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
       - label: "Unit: finitedifference"
         key: unit_finitedifference
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/InputOutput/unit_finitedifference.jl"

ext/cuda/data_layouts_threadblock.jl

@@ -59,6 +59,7 @@ end
 
 ##### linear partition
 @inline function linear_partition(nitems::Integer, n_max_threads::Integer)
+    @assert nitems > 0
     threads = min(nitems, n_max_threads)
     blocks = cld(nitems, threads)
     return (; threads, blocks)

src/DataLayouts/DataLayouts.jl

@@ -2340,16 +2340,13 @@ function IJHMask(is_active::Union{IJFH, IJHF})
     DA = unionall_type(typeof(parent(is_active)))
     (Ni, Nj, _, _, Nh) = size(is_active)
     Nijh = Ni * Nj * Nh
+    N = zeros(Int, 1)
     i_map = zeros(Int, Nijh)
     j_map = zeros(Int, Nijh)
     h_map = zeros(Int, Nijh)
-    return IJHMask(
-        rebuild(is_active, DA),
-        zeros(Int, 1), # N
-        DA(i_map),
-        DA(j_map),
-        DA(h_map),
-    )
+    mask = IJHMask(rebuild(is_active, DA), N, DA(i_map), DA(j_map), DA(h_map))
+    set_mask_maps!(mask)
+    return mask
 end
 
 """

src/InputOutput/writers.jl

@@ -604,7 +604,6 @@ function _write_mpi!(
         dxpl_mpio = :collective,
     )
     dataset[localidx...] = array
-    write_attribute(dataset, "array", array)
     write_attribute(dataset, "data_layout", string(nameof(typeof(values))))
     write_attribute(dataset, "data_eltype", string(eltype(values)))
     return name
@@ -626,7 +625,6 @@ function _write!(group, values::DataLayouts.AbstractData, name::AbstractString;)
     h_dim = DataLayouts.h_dim(DataLayouts.singleton(values))
     array = parent(values)
     dataset = write_plain_array!(group, array, name)
-    write_attribute(dataset, "array", array)
     write_attribute(dataset, "type", string(nameof(typeof(values))))
     write_attribute(dataset, "data_eltype", string(eltype(values)))
     return name