Skip to content

Commit c798ff9

Browse files
efaulhaberefaulhaber
committed
Fix count_allocations 
1 parent 894867b commit c798ff9

File tree

5 files changed

+87
-85
lines changed

5 files changed

+87
-85
lines changed

‎test/count_allocations.jl‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ end
3737
end
3838

3939
# Count allocations of one call to the right-hand side (`kick!` + `drift!`)
40-
function count_rhs_allocations(sol, semi)
40+
function count_rhs_allocations(sol)
4141
t = sol.t[end]
4242
v_ode_, u_ode_ = sol.u[end].x
4343

@@ -48,7 +48,9 @@ function count_rhs_allocations(sol, semi)
4848
du_ode = similar(u_ode)
4949

5050
# Wrap neighborhood searches to avoid counting alloctations in the NHS update
51-
semi_no_nhs_update = copy_semi_with_no_update_nhs(semi)
51+
p = sol.prob.p
52+
semi_no_nhs_update = copy_semi_with_no_update_nhs(p.semi)
53+
p_no_update = TrixiParticles.@set p.semi = semi_no_nhs_update
5254

5355
try
5456
# Disable timers, which cause extra allocations
@@ -58,26 +60,26 @@ function count_rhs_allocations(sol, semi)
5860
# `TrixiParticles.timeit_debug_enabled()` is called, which is redefined in
5961
# `disable_debug_timings` above.
6062
return @invokelatest count_rhs_allocations_inner(dv_ode, du_ode, v_ode, u_ode,
61-
semi_no_nhs_update, t)
63+
p_no_update, t)
6264
finally
6365
# Enable timers again
6466
@invokelatest TrixiParticles.enable_debug_timings()
6567
end
6668
end
6769

68-
# Function barrier to avoid type instabilites with `semi_no_nhs_update`, which will
70+
# Function barrier to avoid type instabilites with `p_no_update`, which will
6971
# cause extra allocations.
7072
@inline function count_rhs_allocations_inner(dv_ode, du_ode, v_ode, u_ode,
71-
semi_no_nhs_update, t)
73+
p_no_update, t)
7274
# Run RHS once to avoid counting allocations from compilation
73-
TrixiParticles.kick!(dv_ode, v_ode, u_ode, semi_no_nhs_update, t)
74-
TrixiParticles.drift!(du_ode, v_ode, u_ode, semi_no_nhs_update, t)
75+
TrixiParticles.kick!(dv_ode, v_ode, u_ode, p_no_update, t)
76+
TrixiParticles.drift!(du_ode, v_ode, u_ode, p_no_update, t)
7577

7678
# Count allocations
7779
allocations_kick = @allocated TrixiParticles.kick!(dv_ode, v_ode, u_ode,
78-
semi_no_nhs_update, t)
80+
p_no_update, t)
7981
allocations_drift = @allocated TrixiParticles.drift!(du_ode, v_ode, u_ode,
80-
semi_no_nhs_update, t)
82+
p_no_update, t)
8183

8284
return allocations_kick + allocations_drift
8385
end

‎test/examples/dam_break_2d_corrections.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@
7474
density_diffusion=nothing)
7575

7676
@test sol.retcode == ReturnCode.Success
77-
@test count_rhs_allocations(sol, semi) == 0
77+
@test count_rhs_allocations(sol) == 0
7878
end
7979

8080
@testset verbose=true "$correction_name" for correction_name in keys(correction_dict)
@@ -107,6 +107,6 @@
107107
sol = solve(ode, RDPK3SpFSAL35(), save_everystep=false, callback=callbacks)
108108

109109
@test sol.retcode == ReturnCode.Success
110-
@test count_rhs_allocations(sol, semi) == 0
110+
@test count_rhs_allocations(sol) == 0
111111
end
112112
end

‎test/examples/examples.jl‎

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
if VERSION < v"1.12"
1616
# Older Julia versions produce allocations because `get_neighborhood_search`
1717
# is not type-stable with TLSPH.
18-
@test count_rhs_allocations(sol, semi) < 200
18+
@test count_rhs_allocations(sol) < 200
1919
else
20-
@test count_rhs_allocations(sol, semi) == 0
20+
@test count_rhs_allocations(sol) == 0
2121
end
2222
end
2323

@@ -34,9 +34,9 @@
3434
if VERSION < v"1.12"
3535
# Older Julia versions produce allocations because `get_neighborhood_search`
3636
# is not type-stable with TLSPH.
37-
@test count_rhs_allocations(sol, semi) < 200
37+
@test count_rhs_allocations(sol) < 200
3838
else
39-
@test count_rhs_allocations(sol, semi) == 0
39+
@test count_rhs_allocations(sol) == 0
4040
end
4141
end
4242

@@ -60,9 +60,9 @@
6060
if VERSION < v"1.12"
6161
# Older Julia versions produce allocations because `get_neighborhood_search`
6262
# is not type-stable with TLSPH.
63-
@test count_rhs_allocations(sol, semi) < 200
63+
@test count_rhs_allocations(sol) < 200
6464
else
65-
@test count_rhs_allocations(sol, semi) == 0
65+
@test count_rhs_allocations(sol) == 0
6666
end
6767
end
6868
end
@@ -79,9 +79,9 @@
7979
if VERSION < v"1.12"
8080
# Older Julia versions produce allocations because `get_neighborhood_search`
8181
# is not type-stable with TLSPH.
82-
@test count_rhs_allocations(sol, semi) < 200
82+
@test count_rhs_allocations(sol) < 200
8383
else
84-
@test count_rhs_allocations(sol, semi) == 0
84+
@test count_rhs_allocations(sol) == 0
8585
end
8686
end
8787

@@ -98,9 +98,9 @@
9898
if VERSION < v"1.12"
9999
# Older Julia versions produce allocations because `get_neighborhood_search`
100100
# is not type-stable with TLSPH.
101-
@test count_rhs_allocations(sol, semi) < 200
101+
@test count_rhs_allocations(sol) < 200
102102
else
103-
@test count_rhs_allocations(sol, semi) == 0
103+
@test count_rhs_allocations(sol) == 0
104104
end
105105
end
106106

@@ -125,9 +125,9 @@
125125
if VERSION < v"1.12"
126126
# Older Julia versions produce allocations because `get_neighborhood_search`
127127
# is not type-stable with TLSPH.
128-
@test count_rhs_allocations(sol, semi) < 200
128+
@test count_rhs_allocations(sol) < 200
129129
else
130-
@test count_rhs_allocations(sol, semi) == 0
130+
@test count_rhs_allocations(sol) == 0
131131
end
132132

133133
# Use split integration and verify that we need fewer than 400 iterations
@@ -141,9 +141,9 @@
141141
if VERSION < v"1.12"
142142
# Older Julia versions produce allocations because `get_neighborhood_search`
143143
# is not type-stable with TLSPH.
144-
@test count_rhs_allocations(sol, semi) < 200
144+
@test count_rhs_allocations(sol) < 200
145145
else
146-
@test count_rhs_allocations(sol, semi) == 0
146+
@test count_rhs_allocations(sol) == 0
147147
end
148148

149149
# Use stage-level coupling and verify that it is not compatible with
@@ -191,9 +191,9 @@
191191
if VERSION < v"1.12"
192192
# Older Julia versions produce allocations because `get_neighborhood_search`
193193
# is not type-stable with TLSPH.
194-
@test count_rhs_allocations(sol, semi) < 200
194+
@test count_rhs_allocations(sol) < 200
195195
else
196-
@test count_rhs_allocations(sol, semi) == 0
196+
@test count_rhs_allocations(sol) == 0
197197
end
198198
end
199199

@@ -208,9 +208,9 @@
208208
if VERSION < v"1.12"
209209
# Older Julia versions produce allocations because `get_neighborhood_search`
210210
# is not type-stable with TLSPH.
211-
@test count_rhs_allocations(sol, semi) < 500
211+
@test count_rhs_allocations(sol) < 500
212212
else
213-
@test count_rhs_allocations(sol, semi) == 0
213+
@test count_rhs_allocations(sol) == 0
214214
end
215215
end
216216
end
@@ -221,7 +221,7 @@
221221
joinpath(examples_dir(), "n_body",
222222
"n_body_solar_system.jl"))
223223
@test sol.retcode == ReturnCode.Success
224-
@test count_rhs_allocations(sol, semi) == 0
224+
@test count_rhs_allocations(sol) == 0
225225
end
226226

227227
@trixi_testset "n_body/n_body_benchmark_trixi.jl" begin
@@ -306,7 +306,7 @@
306306
"rectangular_tank_2d.jl"),
307307
tspan=(0.0, 0.1))
308308
@test sol.retcode == ReturnCode.Success
309-
@test count_rhs_allocations(sol, semi) == 0
309+
@test count_rhs_allocations(sol) == 0
310310
end
311311
end
312312
@trixi_testset "dem/collapsing_sand_pile_3d.jl" begin
@@ -315,6 +315,6 @@
315315
"collapsing_sand_pile_3d.jl"),
316316
tspan=(0.0, 0.1))
317317
@test sol.retcode == ReturnCode.Success
318-
@test count_rhs_allocations(sol, semi) == 0
318+
@test count_rhs_allocations(sol) == 0
319319
end
320320
end

0 commit comments

Comments
 (0)