Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RFC] Splitting up NonlinearSolve #456

Open
avik-pal opened this issue Aug 17, 2024 · 16 comments
Open

[RFC] Splitting up NonlinearSolve #456

avik-pal opened this issue Aug 17, 2024 · 16 comments

Comments

@avik-pal
Copy link
Member

avik-pal commented Aug 17, 2024

Similar to the treatment that was given to OrdinaryDiffEq, we want to split up this package into subpackages. This will be much easier to do here, since all the solvers here are extremely modular (that was one of the central points of our paper https://arxiv.org/abs/2403.16341). This is my first proposal:

  1. NonlinearSolveCore.jl: This houses most of what is present in timer_ouputs.jl, abstract_types.jl, and internal.jl internal/. The AD wrappers and such also go here via extensions
  2. LineSearch.jl (empty repo exists as of now): Let's move https://github.com/SciML/NonlinearSolve.jl/blob/master/src/globalization/line_search.jl to that package. We can iterate on a 2.0 design and move to native solvers later (is that okay @Vaibhavdixit02?)
  3. Next we do the main solver packages:
    1. NonlinearSolveApproximateJacobian.jl -- LBroyden, Broyden, Klement, and other solvers (and utilities) that reduce and are used in https://github.com/SciML/NonlinearSolve.jl/blob/master/src/core/approximate_jacobian.jl.
    2. NonlinearSolveFirstOrder.jl -- This is the biggest chunk housing TrustRegion, NewtonRaphson, etc.
    3. NonlinearSolveSpectralMethods.jl -- DFSane and contents of https://github.com/SciML/NonlinearSolve.jl/blob/master/src/core/spectral_methods.jl.
    4. Where do we put the extension solvers? NonlinearSolve.jl or another subpackage NonlinearSolveExtensionAlgs.jl?
  4. NonlinearSolve.jl -- We define the default algorithms here. We need to pull in all the packages (except Spectral methods) for this part.

To cut down more aggressively on compile times we need to finish NonlinearSolveBase.jl and use that in-place of DiffEqBase.jl.

Also this unblocks solvers like #404 that need to rely on TaylorDiff and heavier packages.

cc @oscardssmith @ChrisRackauckas

@avik-pal
Copy link
Member Author

Also as a part of this change should we move SimpleNonlinearSolve.jl into lib instead of a separate repo?

@ChrisRackauckas
Copy link
Member

Also as a part of this change should we move SimpleNonlinearSolve.jl into lib instead of a separate repo?

I'd entertain that.

@oscardssmith
Copy link
Contributor

IMO this split-up probably won't help much. The Nonlinear solvers are much simpler than the ODE ones and they use more of the same functionality. As such, I doubt this will lead to significant speed benefits.

@Vaibhavdixit02
Copy link
Member

LineSearch.jl (empty repo exists as of now): Let's move https://github.com/SciML/NonlinearSolve.jl/blob/master/src/globalization/line_search.jl to that package. We can iterate on a 2.0 design and move to native solvers later

Sounds perfect, probably the best way to get any activity started there 😅

@ChrisRackauckas
Copy link
Member

The 4 are rather different. LineSearch should be split anyways because of its use in optimization. We don't need to load spectral methods at all for any case users don't request it since they aren't in the defaults. And by default we don't need NonlinearSolveApproximateJacobian with the ODE solvers, we just want the first order methods. So we would probably cut out about half of what's loaded, and for everything else it would parallelize the precompilation.

Where do we put the extension solvers? NonlinearSolve.jl or another subpackage NonlinearSolveExtensionAlgs.jl?

I think NonlinearSolveExtensionAlgs.jl makes sense for that so it's a low dep.

IMO this split-up probably won't help much. The Nonlinear solvers are much simpler than the ODE ones and they use more of the same functionality. As such, I doubt this will lead to significant speed benefits.

I'm not sure that's really the case. What we're really get is parallelized precompilation. This package takes longer than OrdinaryDiffEq to precompile and is the biggest precompilation step right now. Splitting it to parallelize 4 way would already be a huge improvement if we just decrease it by 50%. I think we'd get even lower.

The other compile time decrease is to make some internal functions non-algorithm specific. That is another avenue to investigate, but I'd try to cut out the 3 portions and only depend on first order downstream first and see what we get. I wouldn't be surprised at something that's a good 3x in installation and precompilation time improvement, which would be a massive improvement.

@avik-pal
Copy link
Member Author

Also we should add a simple ASV benchmark script to all the repos to track load times. For example https://github.com/LuxDL/Lux.jl/blob/main/bench/asv.jl is very quick to run and adds comments like LuxDL/Lux.jl#848 (comment) with the load time

@ParamThakkar123
Copy link

ParamThakkar123 commented Aug 24, 2024

Should I take this one ? If nobody has taken it yet ?

@ParamThakkar123
Copy link

Should I take this one ? If nobody has taken it yet ?

@ChrisRackauckas @avik-pal ??

@avik-pal
Copy link
Member Author

We received an email from someone who wanted to take this on.

@ParamThakkar123
Copy link

I sent an email to sciml today for getting a claim on this project an hour ago

@ParamThakkar123
Copy link

Did the email come from [email protected]??

@ParamThakkar123
Copy link

@avik-pal ??

@ChrisRackauckas
Copy link
Member

ChrisRackauckas commented Aug 25, 2024

Let's work this out here. @Suavesito-Olimpiada put a claim in first so that has priority. However, it's unclear right now if that's for NonlinearSolve.jl or BoundaryValueDiffEq.jl. And it looks like @avik-pal already got started on this one himself #458 independently of that, so let's try to spread things out so we aren't overlapping 😅. I think we should hear from:

  1. @Suavesito-Olimpiada which one did you intend to claim?
  2. @avik-pal were you planning to do this one yourself? Or were you just starting the PR?

@ParamThakkar123
Copy link

Okay. No problem. I will put my step back from here 😅

@avik-pal
Copy link
Member Author

@avik-pal were you planning to do this one yourself? Or were you just starting the PR?

I was just setting up an initial sketch and PoC of the benefits of the split.

@Suavesito-Olimpiada
Copy link

Suavesito-Olimpiada commented Aug 26, 2024

I intended to claim both (as per here), but, started with this repo, as per the answer to my questions in email. 😅

I saw the PR in #458. From what I understand, that is independent of the plan exposed here by Avik.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

6 participants