NVTX annotate some functions

Author: charleskawczynski

src/integrators.jl

@@ -159,7 +159,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, alg::Distribute
 end
 
 # either called directly (after init), or by DiffEqBase.solve (via __solve)
-function DiffEqBase.solve!(integrator::DistributedODEIntegrator)
+NVTX.@annotate function DiffEqBase.solve!(integrator::DistributedODEIntegrator)
     while !isempty(integrator.tstops) && integrator.step != integrator.stepstop
         __step!(integrator)
     end
@@ -230,7 +230,9 @@ function __step!(integrator)
 end
 
 # solvers need to define this interface
-step_u!(integrator) = step_u!(integrator, integrator.cache)
+NVTX.@annotate function step_u!(integrator)
+    step_u!(integrator, integrator.cache)
+end
 
 DiffEqBase.get_dt(integrator::DistributedODEIntegrator) = integrator._dt
 function set_dt!(integrator::DistributedODEIntegrator, dt)

src/nl_solvers/newtons_method.jl

@@ -428,7 +428,7 @@ function allocate_cache(alg::KrylovMethod, x_prototype)
     )
 end
 
-function solve_krylov!(alg::KrylovMethod, cache, Δx, x, f!, f, n, post_implicit!, j = nothing)
+NVTX.@annotate function solve_krylov!(alg::KrylovMethod, cache, Δx, x, f!, f, n, post_implicit!, j = nothing)
     (; jacobian_free_jvp, forcing_term, solve_kwargs) = alg
     (; disable_preconditioner, debugger) = alg
     type = solver_type(alg)
@@ -577,7 +577,7 @@ solve_newton!(
     post_implicit_last! = nothing,
 ) = nothing
 
-function solve_newton!(
+NVTX.@annotate function solve_newton!(
     alg::NewtonsMethod,
     cache,
     x,