From 61cc792e878464a2c05e10f098f38bfaa410f08d Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 31 Jan 2025 16:06:59 -0700 Subject: [PATCH] Clean up zeros --- Project.toml | 9 +++++---- .../distributed_architectures.jl | 2 +- src/Grids/latitude_longitude_grid.jl | 16 ++++++++-------- src/Grids/new_data.jl | 2 +- src/Grids/orthogonal_spherical_shell_grid.jl | 8 ++++---- src/Grids/zeros_and_ones.jl | 7 ++++--- src/OutputReaders/field_time_series.jl | 2 +- 7 files changed, 24 insertions(+), 22 deletions(-) diff --git a/Project.toml b/Project.toml index d2496d454b..9387c9da3f 100644 --- a/Project.toml +++ b/Project.toml @@ -33,17 +33,18 @@ SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" +TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a" [weakdeps] Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" MakieCore = "20f20a25-4f0e-4fdf-b5d1-57303727442b" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" [extensions] OceananigansEnzymeExt = "Enzyme" -OceananigansReactantExt = "Reactant" OceananigansMakieExt = ["MakieCore", "Makie"] +OceananigansReactantExt = "Reactant" [compat] Adapt = "4.1.1" @@ -55,8 +56,8 @@ Distances = "0.10" DocStringExtensions = "0.8, 0.9" Enzyme = "0.13.14" FFTW = "1" -Glob = "1.3" GPUArrays = "10.3" +Glob = "1.3" IncompleteLU = "0.2" InteractiveUtils = "1.9" IterativeSolvers = "0.9" @@ -83,8 +84,8 @@ julia = "1.9" CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a" diff --git a/src/DistributedComputations/distributed_architectures.jl b/src/DistributedComputations/distributed_architectures.jl index 4d9dea7916..8b75cd6563 100644 --- a/src/DistributedComputations/distributed_architectures.jl +++ b/src/DistributedComputations/distributed_architectures.jl @@ -306,7 +306,7 @@ ranks(arch::Distributed) = ranks(arch.partition) child_architecture(arch::Distributed) = arch.child_architecture device(arch::Distributed) = device(child_architecture(arch)) -zeros(FT, arch::Distributed, N...) = zeros(FT, child_architecture(arch), N...) +zeros(arch::Distributed, N...) = zeros(child_architecture(arch), FT, N...) array_type(arch::Distributed) = array_type(child_architecture(arch)) sync_device!(arch::Distributed) = sync_device!(arch.child_architecture) convert_to_device(arch::Distributed, arg) = convert_to_device(child_architecture(arch), arg) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 492e3186c9..8b3d5bb132 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -534,14 +534,14 @@ function allocate_metrics(grid::LatitudeLongitudeGrid) metric_size = (length(grid.Δλᶜᵃᵃ), length(grid.φᵃᶜᵃ)) end - Δxᶠᶜ = OffsetArray(zeros(FT, arch, metric_size...), offsets...) - Δxᶜᶠ = OffsetArray(zeros(FT, arch, metric_size...), offsets...) - Δxᶠᶠ = OffsetArray(zeros(FT, arch, metric_size...), offsets...) - Δxᶜᶜ = OffsetArray(zeros(FT, arch, metric_size...), offsets...) - Azᶠᶜ = OffsetArray(zeros(FT, arch, metric_size...), offsets...) - Azᶜᶠ = OffsetArray(zeros(FT, arch, metric_size...), offsets...) - Azᶠᶠ = OffsetArray(zeros(FT, arch, metric_size...), offsets...) - Azᶜᶜ = OffsetArray(zeros(FT, arch, metric_size...), offsets...) + Δxᶠᶜ = OffsetArray(zeros(arch, FT, metric_size...), offsets...) + Δxᶜᶠ = OffsetArray(zeros(arch, FT, metric_size...), offsets...) + Δxᶠᶠ = OffsetArray(zeros(arch, FT, metric_size...), offsets...) + Δxᶜᶜ = OffsetArray(zeros(arch, FT, metric_size...), offsets...) + Azᶠᶜ = OffsetArray(zeros(arch, FT, metric_size...), offsets...) + Azᶜᶠ = OffsetArray(zeros(arch, FT, metric_size...), offsets...) + Azᶠᶠ = OffsetArray(zeros(arch, FT, metric_size...), offsets...) + Azᶜᶜ = OffsetArray(zeros(arch, FT, metric_size...), offsets...) if grid isa YRegularLLG Δyᶠᶜ = Δyᶠᶜᵃ(1, 1, 1, grid) diff --git a/src/Grids/new_data.jl b/src/Grids/new_data.jl index 7b2bf1aac9..30072d73f0 100644 --- a/src/Grids/new_data.jl +++ b/src/Grids/new_data.jl @@ -64,7 +64,7 @@ with indices corresponding to a field on a `grid` of `size(grid)` and located at """ function new_data(FT::DataType, arch, loc, topo, sz, halo_sz, indices=default_indices(length(loc))) Tx, Ty, Tz = total_size(loc, topo, sz, halo_sz, indices) - underlying_data = zeros(FT, arch, Tx, Ty, Tz) + underlying_data = zeros(arch, FT, Tx, Ty, Tz) indices = validate_indices(indices, loc, topo, sz, halo_sz) return offset_data(underlying_data, loc, topo, sz, halo_sz, indices) end diff --git a/src/Grids/orthogonal_spherical_shell_grid.jl b/src/Grids/orthogonal_spherical_shell_grid.jl index 8814b0b66d..c1ae4b05af 100644 --- a/src/Grids/orthogonal_spherical_shell_grid.jl +++ b/src/Grids/orthogonal_spherical_shell_grid.jl @@ -860,10 +860,10 @@ function conformal_cubed_sphere_panel(filepath::AbstractString, architecture = C Tyᶠᶜ = total_length(loc_fc[2](), topology[2](), N[2], H[2]) Tyᶜᶠ = total_length(loc_cf[2](), topology[2](), N[2], H[2]) - λᶠᶜᵃ = offset_data(zeros(FT, architecture, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2]) - λᶜᶠᵃ = offset_data(zeros(FT, architecture, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2]) - φᶠᶜᵃ = offset_data(zeros(FT, architecture, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2]) - φᶜᶠᵃ = offset_data(zeros(FT, architecture, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2]) + λᶠᶜᵃ = offset_data(zeros(architecture, FT, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2]) + λᶜᶠᵃ = offset_data(zeros(architecture, FT, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2]) + φᶠᶜᵃ = offset_data(zeros(architecture, FT, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2]) + φᶜᶠᵃ = offset_data(zeros(architecture, FT, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2]) ## The vertical coordinates can come out of the regular rectilinear grid! Lz, z = generate_coordinate(FT, topology, (Nξ, Nη, Nz), halo, z, :z, 3, architecture) diff --git a/src/Grids/zeros_and_ones.jl b/src/Grids/zeros_and_ones.jl index d2352a5429..2bd8eb81f4 100644 --- a/src/Grids/zeros_and_ones.jl +++ b/src/Grids/zeros_and_ones.jl @@ -1,10 +1,11 @@ using Oceananigans.Architectures: device, AbstractArchitecture -import KernelAbstractions: zeros - -zeros(FT, arch::AbstractArchitecture, N...) = zeros(device(arch), FT, N...) +import KernelAbstractions +import Base: zeros +zeros(FT, arch::AbstractArchitecture, N...) = KernelAbstractions.zeros(device(arch), FT, N...) zeros(grid::AbstractGrid, N...) = zeros(eltype(grid), architecture(grid), N...) @inline Base.zero(grid::AbstractGrid) = zero(eltype(grid)) @inline Base.one(grid::AbstractGrid) = one(eltype(grid)) + diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 9d77c161c3..00b0dda9cb 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -356,7 +356,7 @@ new_data(FT, grid, loc, indices, ::Nothing) = nothing # https://github.com/CliMA/ClimaOcean.jl/actions/runs/8804916198/job/24166354095) function new_data(FT, grid, loc, indices, Nt::Union{Int, Int64}) space_size = total_size(grid, loc, indices) - underlying_data = zeros(FT, architecture(grid), space_size..., Nt) + underlying_data = zeros(architecture(grid), FT, space_size..., Nt) data = offset_data(underlying_data, grid, loc, indices) return data end