Skip to content

Commit aa1b63e

Browse files
gdallegdalle
authored
Add allocation tests with AllocCheck (#577)
* Add allocation tests with AllocCheck * Reinsert no allocs bench * Bump * Fix
1 parent 4496997 commit aa1b63e

5 files changed

Lines changed: 288 additions & 19 deletions

File tree

‎DifferentiationInterfaceTest/Project.toml‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
name = "DifferentiationInterfaceTest"
22
uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
33
authors = ["Guillaume Dalle", "Adrian Hill"]
4-
version = "0.8.2"
4+
version = "0.8.3"
55

66
[deps]
77
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
8+
AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
89
Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
910
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
1011
DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
@@ -20,8 +21,8 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
2021
ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
2122
FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
2223
Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
23-
Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
2424
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
25+
Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
2526
JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
2627
Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
2728
LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"

‎DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ using ADTypes:
1919
ReverseMode,
2020
SymbolicMode,
2121
mode
22+
using AllocCheck: check_allocs
2223
using Chairmarks: @be, Benchmark, Sample
2324
using DataFrames: DataFrame
2425
using DifferentiationInterface
@@ -87,6 +88,8 @@ include("tests/correctness_eval.jl")
8788
include("tests/type_stability_eval.jl")
8889
include("tests/benchmark.jl")
8990
include("tests/benchmark_eval.jl")
91+
include("tests/allocs_eval.jl")
92+
9093
include("test_differentiation.jl")
9194

9295
export FIRST_ORDER, SECOND_ORDER

‎DifferentiationInterfaceTest/src/test_differentiation.jl‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,11 @@ This function always creates and runs a `@testset`, though its contents may vary
2121
2222
- `correctness=true`: whether to compare the differentiation results with the theoretical values specified in each scenario
2323
- `type_stability=:none`: whether (and how) to check type stability of operators with JET.jl.
24+
- `allocations=:none`: whether (and how) to check allocations inside operators with AllocCheck.jl
2425
- `benchmark=:none`: whether (and how) to benchmark operators with Chairmarks.jl
2526
26-
For `type_stability` and `benchmark`, the possible values are `:none`, `:prepared` or `:full`, each concerns a different subset of calls:
27+
For `type_stability`, `allocations` and `benchmark`, the possible values are `:none`, `:prepared` or `:full`.
28+
Each setting tests/benchmarks a different subset of calls:
2729
2830
| kwarg | prepared operator | unprepared operator | preparation |
2931
|---|---|---|---|
@@ -61,6 +63,7 @@ function test_differentiation(
6163
# test categories
6264
correctness::Bool=true,
6365
type_stability::Symbol=:none,
66+
allocations::Symbol=:none,
6467
benchmark::Symbol=:none,
6568
# misc options
6669
excluded::Vector{Symbol}=Symbol[],
@@ -79,11 +82,14 @@ function test_differentiation(
7982
else
8083
@nospecialize(f) -> f != Base.mapreduce_empty # fix for `mapreduce` in jacobian and hessian
8184
end,
85+
# allocs options
86+
skip_allocations::Bool=false, # private, only for code coverage
8287
# benchmark options
8388
count_calls::Bool=true,
8489
benchmark_test::Bool=true,
8590
)
8691
@assert type_stability in (:none, :prepared, :full)
92+
@assert allocations in (:none, :prepared, :full)
8793
@assert benchmark in (:none, :prepared, :full)
8894

8995
scenarios = filter(s -> !(operator(s) in excluded), scenarios)
@@ -149,6 +155,15 @@ function test_differentiation(
149155
)
150156
end
151157
yield()
158+
(allocations != :none) && @testset "Allocations" begin
159+
test_alloccheck(
160+
adapted_backend,
161+
scen;
162+
subset=allocations,
163+
skip=skip_allocations,
164+
)
165+
end
166+
yield()
152167
(benchmark != :none) && @testset "Benchmark" begin
153168
run_benchmark!(
154169
benchmark_data,
@@ -202,6 +217,7 @@ function benchmark_differentiation(
202217
scenarios;
203218
correctness=false,
204219
type_stability=:none,
220+
allocations=:none,
205221
benchmark,
206222
logging,
207223
excluded,

‎DifferentiationInterfaceTest/src/tests/allocs_eval.jl‎

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,222 @@
1+
function test_noallocs(skip::Bool, func, args...)
2+
possible_allocations = check_allocs(func, typeof.(args); ignore_throw=true)
3+
return @test isempty(possible_allocations) skip = skip
4+
end
5+
6+
for op in ALL_OPS
7+
op! = Symbol(op, "!")
8+
val_prefix = if op == :second_derivative
9+
"value_derivative_and_"
10+
elseif op in [:hessian, :hvp]
11+
"value_gradient_and_"
12+
else
13+
"value_and_"
14+
end
15+
val_and_op = Symbol(val_prefix, op)
16+
val_and_op! = Symbol(val_prefix, op!)
17+
prep_op = Symbol("prepare_", op)
18+
19+
S1out = Scenario{op,:out,:out}
20+
S1in = Scenario{op,:in,:out}
21+
S2out = Scenario{op,:out,:in}
22+
S2in = Scenario{op,:in,:in}
23+
24+
if op in [:derivative, :gradient, :jacobian]
25+
@eval function test_alloccheck(
26+
ba::AbstractADType, scen::$S1out; subset::Symbol, skip::Bool
27+
)
28+
(; f, x, contexts) = deepcopy(scen)
29+
prep = $prep_op(f, ba, x, contexts...)
30+
(subset == :full) && test_noallocs(skip, $prep_op, f, ba, x, contexts...)
31+
(subset == :full) && test_noallocs(skip, $op, f, ba, x, contexts...)
32+
(subset == :full) && test_noallocs(skip, $val_and_op, f, ba, x, contexts...)
33+
(subset != :none) && test_noallocs(skip, $op, f, prep, ba, x, contexts...)
34+
(subset != :none) &&
35+
test_noallocs(skip, $val_and_op, f, prep, ba, x, contexts...)
36+
return nothing
37+
end
38+
39+
@eval function test_alloccheck(
40+
ba::AbstractADType, scen::$S1in; subset::Symbol, skip::Bool
41+
)
42+
(; f, x, res1, contexts) = deepcopy(scen)
43+
res1_sim = mysimilar(res1)
44+
prep = $prep_op(f, ba, x, contexts...)
45+
(subset == :full) && test_noallocs(skip, $prep_op, f, ba, x, contexts...)
46+
(subset == :full) &&
47+
test_noallocs(skip, $op!, f, res1_sim, prep, ba, x, contexts...)
48+
(subset == :full) &&
49+
test_noallocs(skip, $val_and_op!, f, res1_sim, prep, ba, x, contexts...)
50+
(subset != :none) &&
51+
test_noallocs(skip, $op!, f, res1_sim, prep, ba, x, contexts...)
52+
(subset != :none) &&
53+
test_noallocs(skip, $val_and_op!, f, res1_sim, prep, ba, x, contexts...)
54+
return nothing
55+
end
56+
57+
op == :gradient && continue
58+
59+
@eval function test_alloccheck(
60+
ba::AbstractADType, scen::$S2out; subset::Symbol, skip::Bool
61+
)
62+
(; f, x, y, contexts) = deepcopy(scen)
63+
prep = $prep_op(f, y, ba, x, contexts...)
64+
(subset == :full) && test_noallocs(skip, $prep_op, f, y, ba, x, contexts...)
65+
(subset == :full) && test_noallocs(skip, $op, f, y, ba, x, contexts...)
66+
(subset == :full) && test_noallocs(skip, $val_and_op, f, y, ba, x, contexts...)
67+
(subset != :none) && test_noallocs(skip, $op, f, y, prep, ba, x, contexts...)
68+
(subset != :none) &&
69+
test_noallocs(skip, $val_and_op, f, y, prep, ba, x, contexts...)
70+
return nothing
71+
end
72+
73+
@eval function test_alloccheck(
74+
ba::AbstractADType, scen::$S2in; subset::Symbol, skip::Bool
75+
)
76+
(; f, x, y, res1, contexts) = deepcopy(scen)
77+
res1_sim = mysimilar(res1)
78+
prep = $prep_op(f, y, ba, x, contexts...)
79+
(subset == :full) && test_noallocs(skip, $prep_op, f, y, ba, x, contexts...)
80+
(subset == :full) &&
81+
test_noallocs(skip, $op!, f, y, res1_sim, ba, x, contexts...)
82+
(subset == :full) &&
83+
test_noallocs(skip, $val_and_op!, f, y, res1_sim, ba, x, contexts...)
84+
(subset != :none) &&
85+
test_noallocs(skip, $op!, f, y, res1_sim, prep, ba, x, contexts...)
86+
(subset != :none) &&
87+
test_noallocs(skip, $val_and_op!, f, y, res1_sim, prep, ba, x, contexts...)
88+
return nothing
89+
end
90+
91+
elseif op in [:second_derivative, :hessian]
92+
@eval function test_alloccheck(
93+
ba::AbstractADType, scen::$S1out; subset::Symbol, skip::Bool
94+
)
95+
(; f, x, contexts) = deepcopy(scen)
96+
prep = $prep_op(f, ba, x, contexts...)
97+
(subset == :full) && test_noallocs(skip, $prep_op, f, ba, x, contexts...)
98+
(subset == :full) && test_noallocs(skip, $op, f, ba, x, contexts...)
99+
(subset == :full) && test_noallocs(skip, $val_and_op, f, ba, x, contexts...)
100+
(subset != :none) && test_noallocs(skip, $op, f, prep, ba, x, contexts...)
101+
(subset != :none) &&
102+
test_noallocs(skip, $val_and_op, f, prep, ba, x, contexts...)
103+
return nothing
104+
end
105+
106+
@eval function test_alloccheck(
107+
ba::AbstractADType, scen::$S1in; subset::Symbol, skip::Bool
108+
)
109+
(; f, x, res1, res2, contexts) = deepcopy(scen)
110+
res1_sim, res2_sim = mysimilar(res1), mysimilar(res2)
111+
prep = $prep_op(f, ba, x, contexts...)
112+
(subset == :full) && test_noallocs(skip, $prep_op, f, ba, x, contexts...)
113+
(subset == :full) && test_noallocs(skip, $op!, f, res2_sim, ba, x, contexts...)
114+
(subset == :full) &&
115+
test_noallocs(skip, $val_and_op!, f, res1_sim, res2_sim, ba, x, contexts...)
116+
(subset != :none) &&
117+
test_noallocs(skip, $op!, f, res2_sim, prep, ba, x, contexts...)
118+
(subset != :none) && test_noallocs(
119+
skip, $val_and_op!, f, res1_sim, res2_sim, prep, ba, x, contexts...
120+
)
121+
return nothing
122+
end
123+
124+
elseif op in [:pushforward, :pullback]
125+
@eval function test_alloccheck(
126+
ba::AbstractADType, scen::$S1out; subset::Symbol, skip::Bool
127+
)
128+
(; f, x, tang, contexts) = deepcopy(scen)
129+
prep = $prep_op(f, ba, x, tang, contexts...)
130+
(subset == :full) && test_noallocs(skip, $prep_op, f, ba, x, tang, contexts...)
131+
(subset == :full) && test_noallocs(skip, $op, f, ba, x, tang, contexts...)
132+
(subset == :full) &&
133+
test_noallocs(skip, $val_and_op, f, ba, x, tang, contexts...)
134+
(subset != :none) && test_noallocs(skip, $op, f, prep, ba, x, tang, contexts...)
135+
(subset != :none) &&
136+
test_noallocs(skip, $val_and_op, f, prep, ba, x, tang, contexts...)
137+
return nothing
138+
end
139+
140+
@eval function test_alloccheck(
141+
ba::AbstractADType, scen::$S1in; subset::Symbol, skip::Bool
142+
)
143+
(; f, x, tang, res1, contexts) = deepcopy(scen)
144+
res1_sim = mysimilar(res1)
145+
prep = $prep_op(f, ba, x, tang, contexts...)
146+
(subset == :full) && test_noallocs(skip, $prep_op, f, ba, x, tang, contexts...)
147+
(subset == :full) &&
148+
test_noallocs(skip, $op!, f, res1_sim, ba, x, tang, contexts...)
149+
(subset == :full) &&
150+
test_noallocs(skip, $val_and_op!, f, res1_sim, ba, x, tang, contexts...)
151+
(subset != :none) &&
152+
test_noallocs(skip, $op!, f, res1_sim, prep, ba, x, tang, contexts...)
153+
(subset != :none) && test_noallocs(
154+
skip, $val_and_op!, f, res1_sim, prep, ba, x, tang, contexts...
155+
)
156+
return nothing
157+
end
158+
159+
@eval function test_alloccheck(
160+
ba::AbstractADType, scen::$S2out; subset::Symbol, skip::Bool
161+
)
162+
(; f, x, y, tang, contexts) = deepcopy(scen)
163+
prep = $prep_op(f, y, ba, x, tang, contexts...)
164+
(subset == :full) &&
165+
test_noallocs(skip, $prep_op, f, y, ba, x, tang, contexts...)
166+
(subset == :full) && test_noallocs(skip, $op, f, y, ba, x, tang, contexts...)
167+
(subset == :full) &&
168+
test_noallocs(skip, $val_and_op, f, y, ba, x, tang, contexts...)
169+
(subset != :none) &&
170+
test_noallocs(skip, $op, f, y, prep, ba, x, tang, contexts...)
171+
(subset != :none) &&
172+
test_noallocs(skip, $val_and_op, f, y, prep, ba, x, tang, contexts...)
173+
return nothing
174+
end
175+
176+
@eval function test_alloccheck(
177+
ba::AbstractADType, scen::$S2in; subset::Symbol, skip::Bool
178+
)
179+
(; f, x, y, tang, res1, contexts) = deepcopy(scen)
180+
res1_sim = mysimilar(res1)
181+
prep = $prep_op(f, y, ba, x, tang, contexts...)
182+
(subset == :full) &&
183+
test_noallocs(skip, $prep_op, f, y, ba, x, tang, contexts...)
184+
(subset == :full) &&
185+
test_noallocs(skip, $op!, f, y, res1_sim, ba, x, tang, contexts...)
186+
(subset == :full) &&
187+
test_noallocs(skip, $val_and_op!, f, y, res1_sim, ba, x, tang, contexts...)
188+
(subset != :none) &&
189+
test_noallocs(skip, $op!, f, y, res1_sim, prep, ba, x, tang, contexts...)
190+
(subset != :none) && test_noallocs(
191+
skip, $val_and_op!, f, y, res1_sim, prep, ba, x, tang, contexts...
192+
)
193+
return nothing
194+
end
195+
196+
elseif op in [:hvp]
197+
@eval function test_alloccheck(
198+
ba::AbstractADType, scen::$S1out; subset::Symbol, skip::Bool
199+
)
200+
(; f, x, tang, contexts) = deepcopy(scen)
201+
prep = $prep_op(f, ba, x, tang, contexts...)
202+
(subset == :full) && test_noallocs(skip, $prep_op, f, ba, x, tang, contexts...)
203+
(subset == :full) && test_noallocs(skip, $op, f, ba, x, tang, contexts...)
204+
(subset != :none) && test_noallocs(skip, $op, f, prep, ba, x, tang, contexts...)
205+
return nothing
206+
end
207+
208+
@eval function test_alloccheck(
209+
ba::AbstractADType, scen::$S1in; subset::Symbol, skip::Bool
210+
)
211+
(; f, x, tang, res1, res2, contexts) = deepcopy(scen)
212+
res1_sim, res2_sim = mysimilar(res1), mysimilar(res2)
213+
prep = $prep_op(f, ba, x, tang, contexts...)
214+
(subset == :full) && test_noallocs(skip, $prep_op, f, ba, x, tang, contexts...)
215+
(subset == :full) &&
216+
test_noallocs(skip, $op!, f, res2_sim, ba, x, tang, contexts...)
217+
(subset != :none) &&
218+
test_noallocs(skip, $op!, f, res2_sim, prep, ba, x, tang, contexts...)
219+
return nothing
220+
end
221+
end
222+
end

‎DifferentiationInterfaceTest/test/zero_backends.jl‎

Lines changed: 43 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -65,21 +65,48 @@ end
6565

6666
## Allocations
6767

68-
data_allocfree = vcat(
69-
benchmark_differentiation(
70-
AutoZeroForward(),
71-
allocfree_scenarios();
72-
excluded=[:pullback, :gradient],
73-
benchmark=:prepared,
74-
logging=LOGGING,
75-
),
76-
benchmark_differentiation(
77-
AutoZeroReverse(),
78-
allocfree_scenarios();
79-
excluded=[:pushforward, :derivative],
80-
benchmark=:prepared,
81-
logging=LOGGING,
82-
),
68+
@testset "Benchmark for zero allocations" begin
69+
data_allocfree = vcat(
70+
benchmark_differentiation(
71+
AutoZeroForward(),
72+
allocfree_scenarios();
73+
excluded=[:pullback, :gradient],
74+
benchmark=:prepared,
75+
logging=LOGGING,
76+
),
77+
benchmark_differentiation(
78+
AutoZeroReverse(),
79+
allocfree_scenarios();
80+
excluded=[:pushforward, :derivative],
81+
benchmark=:prepared,
82+
logging=LOGGING,
83+
),
84+
)
85+
@test all(iszero, data_allocfree[!, :allocs])
86+
end
87+
88+
test_differentiation(
89+
AutoZeroForward(),
90+
allocfree_scenarios();
91+
correctness=false,
92+
allocations=:prepared,
93+
excluded=[:pullback, :gradient, :jacobian],
94+
logging=LOGGING,
95+
)
96+
97+
test_differentiation(
98+
AutoZeroReverse(),
99+
allocfree_scenarios();
100+
correctness=false,
101+
allocations=:prepared,
102+
excluded=[:pushforward, :derivative, :jacobian],
103+
logging=LOGGING,
83104
)
84105

85-
@test all(iszero, data_allocfree[!, :allocs])
106+
test_differentiation(
107+
AutoZeroForward();
108+
correctness=false,
109+
allocations=:full,
110+
skip_allocations=true,
111+
logging=LOGGING,
112+
)

0 commit comments

Comments
 (0)