Skip to content

Commit 0df1d42

Browse files
amontoisonamontoison
authored
Use show_time for sparse AD backends (#326)
1 parent dec190a commit 0df1d42

File tree

4 files changed

+209
-149
lines changed

4 files changed

+209
-149
lines changed

‎src/ad.jl‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ function ADModelBackend(
8787
hessian_backend = if hessian_backend isa Union{AbstractNLPModel, ADBackend}
8888
hessian_backend
8989
else
90-
HB(nvar, f, ncon, c!; kwargs...)
90+
HB(nvar, f, ncon, c!; show_time, kwargs...)
9191
end
9292
end
9393
show_time && println("hessian backend $HB: $b seconds;")
@@ -170,7 +170,7 @@ function ADModelBackend(
170170
jacobian_backend = if jacobian_backend isa Union{AbstractNLPModel, ADBackend}
171171
jacobian_backend
172172
else
173-
JB(nvar, f, ncon, c!; kwargs...)
173+
JB(nvar, f, ncon, c!; show_time, kwargs...)
174174
end
175175
end
176176
show_time && println("jacobian backend $JB: $b seconds;")
@@ -180,7 +180,7 @@ function ADModelBackend(
180180
hessian_backend = if hessian_backend isa Union{AbstractNLPModel, ADBackend}
181181
hessian_backend
182182
else
183-
HB(nvar, f, ncon, c!; kwargs...)
183+
HB(nvar, f, ncon, c!; show_time, kwargs...)
184184
end
185185
end
186186
show_time && println("hessian backend $HB: $b seconds;")
@@ -263,7 +263,7 @@ function ADModelNLSBackend(
263263
hessian_backend = if hessian_backend isa Union{AbstractNLPModel, ADBackend}
264264
hessian_backend
265265
else
266-
HB(nvar, f, ncon, c!; kwargs...)
266+
HB(nvar, f, ncon, c!; show_time, kwargs...)
267267
end
268268
end
269269
show_time && println("hessian backend $HB: $b seconds;")
@@ -304,7 +304,7 @@ function ADModelNLSBackend(
304304
if jacobian_residual_backend isa Union{AbstractNLPModel, ADBackend}
305305
jacobian_residual_backend
306306
else
307-
JBLS(nvar, x -> zero(eltype(x)), nequ, F!; kwargs...)
307+
JBLS(nvar, x -> zero(eltype(x)), nequ, F!; show_time, kwargs...)
308308
end
309309
end
310310
show_time && println("jacobian_residual backend $JBLS: $b seconds;")
@@ -314,7 +314,7 @@ function ADModelNLSBackend(
314314
hessian_residual_backend = if hessian_residual_backend isa Union{AbstractNLPModel, ADBackend}
315315
hessian_residual_backend
316316
else
317-
HBLS(nvar, x -> zero(eltype(x)), nequ, F!; kwargs...)
317+
HBLS(nvar, x -> zero(eltype(x)), nequ, F!; show_time, kwargs...)
318318
end
319319
end
320320
show_time && println("hessian_residual backend $HBLS: $b seconds. \n")
@@ -410,7 +410,7 @@ function ADModelNLSBackend(
410410
jacobian_backend = if jacobian_backend isa Union{AbstractNLPModel, ADBackend}
411411
jacobian_backend
412412
else
413-
JB(nvar, f, ncon, c!; kwargs...)
413+
JB(nvar, f, ncon, c!; show_time, kwargs...)
414414
end
415415
end
416416
show_time && println("jacobian backend $JB: $b seconds;")
@@ -420,7 +420,7 @@ function ADModelNLSBackend(
420420
hessian_backend = if hessian_backend isa Union{AbstractNLPModel, ADBackend}
421421
hessian_backend
422422
else
423-
HB(nvar, f, ncon, c!; kwargs...)
423+
HB(nvar, f, ncon, c!; show_time, kwargs...)
424424
end
425425
end
426426
show_time && println("hessian backend $HB: $b seconds;")
@@ -471,7 +471,7 @@ function ADModelNLSBackend(
471471
if jacobian_residual_backend isa Union{AbstractNLPModel, ADBackend}
472472
jacobian_residual_backend
473473
else
474-
JBLS(nvar, x -> zero(eltype(x)), nequ, F!; kwargs...)
474+
JBLS(nvar, x -> zero(eltype(x)), nequ, F!; show_time, kwargs...)
475475
end
476476
end
477477
show_time && println("jacobian_residual backend $JBLS: $b seconds;")
@@ -481,7 +481,7 @@ function ADModelNLSBackend(
481481
hessian_residual_backend = if hessian_residual_backend isa Union{AbstractNLPModel, ADBackend}
482482
hessian_residual_backend
483483
else
484-
HBLS(nvar, x -> zero(eltype(x)), nequ, F!; kwargs...)
484+
HBLS(nvar, x -> zero(eltype(x)), nequ, F!; show_time, kwargs...)
485485
end
486486
end
487487
show_time && println("hessian_residual backend $HBLS: $b seconds. \n")

‎src/enzyme.jl‎

Lines changed: 67 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -112,11 +112,15 @@ function SparseEnzymeADJacobian(
112112
x0::AbstractVector = rand(nvar),
113113
coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
114114
detector::AbstractSparsityDetector = TracerSparsityDetector(),
115+
show_time::Bool = false,
115116
kwargs...,
116117
)
117-
output = similar(x0, ncon)
118-
J = compute_jacobian_sparsity(c!, output, x0, detector = detector)
119-
SparseEnzymeADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, kwargs...)
118+
timer = @elapsed begin
119+
output = similar(x0, ncon)
120+
J = compute_jacobian_sparsity(c!, output, x0, detector = detector)
121+
end
122+
show_time && println(" • Sparsity pattern detection of the Jacobian: $timer seconds.")
123+
SparseEnzymeADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, show_time, kwargs...)
120124
end
121125

122126
function SparseEnzymeADJacobian(
@@ -127,18 +131,26 @@ function SparseEnzymeADJacobian(
127131
J::SparseMatrixCSC{Bool, Int};
128132
x0::AbstractVector{T} = rand(nvar),
129133
coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
134+
show_time::Bool = false,
130135
kwargs...,
131136
) where {T}
132-
# We should support :row and :bidirectional in the future
133-
problem = ColoringProblem{:nonsymmetric, :column}()
134-
result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
137+
timer = @elapsed begin
138+
# We should support :row and :bidirectional in the future
139+
problem = ColoringProblem{:nonsymmetric, :column}()
140+
result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
141+
142+
rowval = J.rowval
143+
colptr = J.colptr
144+
nzval = T.(J.nzval)
145+
compressed_jacobian = similar(x0, ncon)
146+
end
147+
show_time && println(" • Coloring of the sparse Jacobian: $timer seconds.")
135148

136-
rowval = J.rowval
137-
colptr = J.colptr
138-
nzval = T.(J.nzval)
139-
compressed_jacobian = similar(x0, ncon)
140-
v = similar(x0)
141-
cx = zeros(T, ncon)
149+
timer = @elapsed begin
150+
v = similar(x0)
151+
cx = zeros(T, ncon)
152+
end
153+
show_time && println(" • Allocation of the AD buffers for the sparse Jacobian: $timer seconds.")
142154

143155
SparseEnzymeADJacobian(
144156
nvar,
@@ -177,10 +189,14 @@ function SparseEnzymeADHessian(
177189
x0::AbstractVector = rand(nvar),
178190
coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:substitution}(),
179191
detector::AbstractSparsityDetector = TracerSparsityDetector(),
192+
show_time::Bool = false,
180193
kwargs...,
181194
)
182-
H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector)
183-
SparseEnzymeADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, kwargs...)
195+
timer = @elapsed begin
196+
H = compute_hessian_sparsity(f, nvar, c!, ncon, detector = detector)
197+
end
198+
show_time && println(" • Sparsity pattern detection of the Hessian: $timer seconds.")
199+
SparseEnzymeADHessian(nvar, f, ncon, c!, H; x0, coloring_algorithm, show_time, kwargs...)
184200
end
185201

186202
function SparseEnzymeADHessian(
@@ -191,38 +207,48 @@ function SparseEnzymeADHessian(
191207
H::SparseMatrixCSC{Bool, Int};
192208
x0::AbstractVector{T} = rand(nvar),
193209
coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:substitution}(),
210+
show_time::Bool = false,
194211
kwargs...,
195212
) where {T}
196-
problem = ColoringProblem{:symmetric, :column}()
197-
result_coloring = coloring(H, problem, coloring_algorithm, decompression_eltype = T)
198-
199-
trilH = tril(H)
200-
rowval = trilH.rowval
201-
colptr = trilH.colptr
202-
nzval = T.(trilH.nzval)
203-
if coloring_algorithm isa GreedyColoringAlgorithm{:direct}
204-
coloring_mode = :direct
205-
compressed_hessian_icol = similar(x0)
206-
compressed_hessian = compressed_hessian_icol
207-
else
208-
coloring_mode = :substitution
209-
group = column_groups(result_coloring)
210-
ncolors = length(group)
211-
compressed_hessian_icol = similar(x0)
212-
compressed_hessian = similar(x0, (nvar, ncolors))
213+
timer = @elapsed begin
214+
problem = ColoringProblem{:symmetric, :column}()
215+
result_coloring = coloring(H, problem, coloring_algorithm, decompression_eltype = T)
216+
217+
trilH = tril(H)
218+
rowval = trilH.rowval
219+
colptr = trilH.colptr
220+
nzval = T.(trilH.nzval)
221+
if coloring_algorithm isa GreedyColoringAlgorithm{:direct}
222+
coloring_mode = :direct
223+
compressed_hessian_icol = similar(x0)
224+
compressed_hessian = compressed_hessian_icol
225+
else
226+
coloring_mode = :substitution
227+
group = column_groups(result_coloring)
228+
ncolors = length(group)
229+
compressed_hessian_icol = similar(x0)
230+
compressed_hessian = similar(x0, (nvar, ncolors))
231+
end
213232
end
214-
v = similar(x0)
215-
y = similar(x0, ncon)
216-
cx = similar(x0, ncon)
217-
grad = similar(x0)
218-
function ℓ(x, y, obj_weight, cx)
219-
res = obj_weight * f(x)
220-
if ncon != 0
221-
c!(cx, x)
222-
res += sum(cx[i] * y[i] for i = 1:ncon)
233+
show_time && println(" • Coloring of the sparse Hessian: $timer seconds.")
234+
235+
timer = @elapsed begin
236+
v = similar(x0)
237+
y = similar(x0, ncon)
238+
cx = similar(x0, ncon)
239+
grad = similar(x0)
240+
241+
function ℓ(x, y, obj_weight, cx)
242+
res = obj_weight * f(x)
243+
if ncon != 0
244+
c!(cx, x)
245+
res += sum(cx[i] * y[i] for i = 1:ncon)
246+
end
247+
return res
223248
end
224-
return res
225249
end
250+
show_time && println(" • Allocation of the AD buffers for the sparse Hessian: $timer seconds.")
251+
226252

227253
return SparseEnzymeADHessian(
228254
nvar,

0 commit comments

Comments
 (0)