Skip to content

Commit 67f0683

Browse files
charleskawczynskicharleskawczynski
authored
Merge pull request #1844 from CliMA/ck/test_mapreduce_cuda
Add mapreduce unit tests
2 parents 5ea5fea + 619d92c commit 67f0683

File tree

3 files changed

+146
-0
lines changed

3 files changed

+146
-0
lines changed

‎.buildkite/pipeline.yml

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,10 @@ steps:
8585
key: unit_data_copyto
8686
command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_copyto.jl"
8787

88+
- label: "Unit: mapreduce"
89+
key: unit_data_mapreduce
90+
command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_mapreduce.jl"
91+
8892
- label: "Unit: data_opt_similar"
8993
key: data_opt_similar
9094
command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/opt_similar.jl"
@@ -129,6 +133,28 @@ steps:
129133
agents:
130134
slurm_gpus: 1
131135

136+
- label: "Unit: data mapreduce"
137+
key: gpu_unit_data_mapreduce
138+
command:
139+
- "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
140+
- "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_mapreduce.jl"
141+
env:
142+
CLIMACOMMS_DEVICE: "CUDA"
143+
agents:
144+
slurm_gpus: 1
145+
146+
- label: "Unit: data mapreduce (2 gpus)"
147+
key: gpu_2_unit_data_mapreduce
148+
command:
149+
- "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
150+
- "srun julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_mapreduce.jl"
151+
env:
152+
CLIMACOMMS_DEVICE: "CUDA"
153+
CLIMACOMMS_CONTEXT: "MPI"
154+
agents:
155+
slum_ntasks: 2
156+
slurm_gpus: 2
157+
132158
- label: "Unit: data copyto"
133159
key: gpu_unit_data_copyto
134160
command:

‎test/DataLayouts/unit_mapreduce.jl

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
#=
2+
julia --project
3+
using Revise; include(joinpath("test", "DataLayouts", "unit_mapreduce.jl"))
4+
=#
5+
using Test
6+
using ClimaCore.DataLayouts
7+
import ClimaComms
8+
import Random
9+
ClimaComms.@import_required_backends
10+
11+
"""test mapreduce with data layouts with 1 component"""
12+
function test_mapreduce_1!(device, data)
13+
device isa ClimaComms.CUDADevice && return nothing
14+
Random.seed!(1234)
15+
ArrayType = ClimaComms.array_type(device)
16+
rand_data = ArrayType(rand(eltype(parent(data)), size(parent(data))))
17+
parent(data) .= rand_data
18+
@test minimum(data) == minimum(parent(data))
19+
@test maximum(data) == maximum(parent(data))
20+
end
21+
22+
"""test mapreduce with data layouts with 2 components"""
23+
function test_mapreduce_2!(device, data)
24+
device isa ClimaComms.CUDADevice && return nothing
25+
Random.seed!(1234)
26+
ArrayType = ClimaComms.array_type(device)
27+
rand_data = ArrayType(rand(eltype(parent(data)), size(parent(data))))
28+
parent(data) .= rand_data
29+
# mapreduce orders tuples lexicographically:
30+
# minimum(((2,3), (1,4))) # (1, 4)
31+
# minimum(((1,4), (2,3))) # (1, 4)
32+
# minimum(((4,1), (3,2))) # (3, 2)
33+
# so, for now, let's just assign the two components to match:
34+
parent(data.:2) .= parent(data.:1)
35+
@test minimum(data) == (minimum(parent(data.:1)), minimum(parent(data.:2)))
36+
@test maximum(data) == (maximum(parent(data.:1)), maximum(parent(data.:2)))
37+
end
38+
39+
@testset "mapreduce with Nf = 1" begin
40+
device = ClimaComms.device()
41+
device_zeros(args...) = ClimaComms.array_type(device)(zeros(args...))
42+
FT = Float64
43+
S = FT
44+
Nf = 1
45+
Nv = 4
46+
Nij = 3
47+
Nh = 5
48+
Nk = 6
49+
#! format: off
50+
data = DataF{S}(device_zeros(FT,Nf)); test_mapreduce_1!(device, data)
51+
data = IJFH{S, Nij}(device_zeros(FT,Nij,Nij,Nf,Nh)); test_mapreduce_1!(device, data)
52+
data = IFH{S, Nij}(device_zeros(FT,Nij,Nf,Nh)); test_mapreduce_1!(device, data)
53+
data = IJF{S, Nij}(device_zeros(FT,Nij,Nij,Nf)); test_mapreduce_1!(device, data)
54+
data = IF{S, Nij}(device_zeros(FT,Nij,Nf)); test_mapreduce_1!(device, data)
55+
data = VF{S, Nv}(device_zeros(FT,Nv,Nf)); test_mapreduce_1!(device, data)
56+
data = VIJFH{S, Nv, Nij}(device_zeros(FT,Nv,Nij,Nij,Nf,Nh)); test_mapreduce_1!(device, data)
57+
data = VIFH{S, Nv, Nij}(device_zeros(FT,Nv,Nij,Nf,Nh)); test_mapreduce_1!(device, data)
58+
#! format: on
59+
# data = DataLayouts.IJKFVH{S, Nij, Nk}(device_zeros(FT,Nij,Nij,Nk,Nf,Nv,Nh)); test_mapreduce_1!(device, data) # TODO: test
60+
# data = DataLayouts.IH1JH2{S, Nij}(device_zeros(FT,2*Nij,3*Nij)); test_mapreduce_1!(device, data) # TODO: test
61+
end
62+
63+
@testset "mapreduce with Nf > 1" begin
64+
device = ClimaComms.device()
65+
device_zeros(args...) = ClimaComms.array_type(device)(zeros(args...))
66+
FT = Float64
67+
S = Tuple{FT, FT}
68+
Nf = 2
69+
Nv = 4
70+
Nij = 3
71+
Nh = 5
72+
Nk = 6
73+
#! format: off
74+
data = DataF{S}(device_zeros(FT,Nf)); test_mapreduce_2!(device, data)
75+
data = IJFH{S, Nij}(device_zeros(FT,Nij,Nij,Nf,Nh)); test_mapreduce_2!(device, data)
76+
data = IFH{S, Nij}(device_zeros(FT,Nij,Nf,Nh)); test_mapreduce_2!(device, data)
77+
data = IJF{S, Nij}(device_zeros(FT,Nij,Nij,Nf)); test_mapreduce_2!(device, data)
78+
data = IF{S, Nij}(device_zeros(FT,Nij,Nf)); test_mapreduce_2!(device, data)
79+
data = VF{S, Nv}(device_zeros(FT,Nv,Nf)); test_mapreduce_2!(device, data)
80+
data = VIJFH{S, Nv, Nij}(device_zeros(FT,Nv,Nij,Nij,Nf,Nh)); test_mapreduce_2!(device, data)
81+
data = VIFH{S, Nv, Nij}(device_zeros(FT,Nv,Nij,Nf,Nh)); test_mapreduce_2!(device, data)
82+
#! format: on
83+
# TODO: test this
84+
# data = DataLayouts.IJKFVH{S, Nij, Nk}(device_zeros(FT,Nij,Nij,Nk,Nf,Nv,Nh)); test_mapreduce_2!(device, data) # TODO: test
85+
# data = DataLayouts.IH1JH2{S, Nij}(device_zeros(FT,2*Nij,3*Nij)); test_mapreduce_2!(device, data) # TODO: test
86+
end
87+
88+
@testset "mapreduce views with Nf > 1" begin
89+
device = ClimaComms.device()
90+
device_zeros(args...) = ClimaComms.array_type(device)(zeros(args...))
91+
data_view(data) = DataLayouts.rebuild(
92+
data,
93+
SubArray(
94+
parent(data),
95+
ntuple(i -> Base.OneTo(size(parent(data), i)), ndims(data)),
96+
),
97+
)
98+
FT = Float64
99+
S = Tuple{FT, FT}
100+
Nf = 2
101+
Nv = 4
102+
Nij = 3
103+
Nh = 5
104+
Nk = 6
105+
# Rather than using level/slab/column, let's just make views/SubArrays
106+
# directly so that we can easily test all cases:
107+
#! format: off
108+
data = IJFH{S, Nij}(device_zeros(FT,Nij,Nij,Nf,Nh)); test_mapreduce_2!(device, data_view(data))
109+
data = IFH{S, Nij}(device_zeros(FT,Nij,Nf,Nh)); test_mapreduce_2!(device, data_view(data))
110+
data = IJF{S, Nij}(device_zeros(FT,Nij,Nij,Nf)); test_mapreduce_2!(device, data_view(data))
111+
data = IF{S, Nij}(device_zeros(FT,Nij,Nf)); test_mapreduce_2!(device, data_view(data))
112+
data = VF{S, Nv}(device_zeros(FT,Nv,Nf)); test_mapreduce_2!(device, data_view(data))
113+
data = VIJFH{S, Nv, Nij}(device_zeros(FT,Nv,Nij,Nij,Nf,Nh)); test_mapreduce_2!(device, data_view(data))
114+
data = VIFH{S, Nv, Nij}(device_zeros(FT,Nv,Nij,Nf,Nh)); test_mapreduce_2!(device, data_view(data))
115+
#! format: on
116+
# TODO: test this
117+
# data = DataLayouts.IJKFVH{S, Nij, Nk}(device_zeros(FT,Nij,Nij,Nk,Nf,Nv,Nh)); test_mapreduce_2!(device, data_view(data)) # TODO: test
118+
# data = DataLayouts.IH1JH2{S, Nij}(device_zeros(FT,2*Nij,3*Nij)); test_mapreduce_2!(device, data_view(data)) # TODO: test
119+
end

‎test/runtests.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ UnitTest("DataLayouts 1D" ,"DataLayouts/data1d.jl"),
1616
UnitTest("DataLayouts 2D" ,"DataLayouts/data2d.jl"),
1717
UnitTest("DataLayouts 1dx" ,"DataLayouts/data1dx.jl"),
1818
UnitTest("DataLayouts 2dx" ,"DataLayouts/data2dx.jl"),
19+
UnitTest("DataLayouts mapreduce" ,"DataLayouts/unit_mapreduce.jl"),
1920
UnitTest("Geometry" ,"Geometry/geometry.jl"),
2021
UnitTest("rmul_with_projection" ,"Geometry/rmul_with_projection.jl"),
2122
UnitTest("AxisTensors" ,"Geometry/axistensors.jl"),

0 commit comments

Comments
 (0)