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

feat: use DI for structured Jacobians #470

Merged
merged 3 commits into from
Oct 2, 2024
Merged

Conversation

avik-pal
Copy link
Member

@avik-pal avik-pal commented Oct 2, 2024

Seems to work beautifully. Haven't benchmarked yet, but the important part here is to ensure structure is preserved, since the main cost in these workloads is that Linear solving will be bottlnecked if a banded matrix is converted into a sparse matrix.

@gdalle
Copy link
Collaborator

gdalle commented Oct 2, 2024

Structure is indeed preserved for decompression (on the SparseMatrixColorings PR branch at least).
However, with the current implementation, a sparse matrix and its transpose will be allocated from the structured Jacobian nonetheless, as a preparation step for decompression.
This seems futile when the decompression target still has the same structure, but when it doesn't, this sparse encoding is actually useful for fast decompression. The scenario here would be someone passing sparsity_pattern::Tridiagonal but then calling jacobian! on a jac_prototype::SparseMatrixCSC.
Ultimately, coloring is a one-time cost, whereas decompression is paid at every Jacobian computation. So in SparseMatrixColorings we choose to allocate a little more during coloring, and that way we speed up more decompression scenarios down the road.

end

function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
autodiff = nothing, kwargs...) where {F}
fu = f(u, p)
if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing, nothing)
return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same question about cases where DI preparation is skipped

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only where the user provides an analytic jacobian. This one is a special case for scalars when DI is skipped if jacobian/jvp/vjp is provided.

src/internal/jacobian.jl Show resolved Hide resolved
docs/src/basics/sparsity_detection.md Outdated Show resolved Hide resolved
docs/src/basics/sparsity_detection.md Outdated Show resolved Hide resolved
docs/src/basics/sparsity_detection.md Outdated Show resolved Hide resolved
docs/src/basics/sparsity_detection.md Outdated Show resolved Hide resolved
docs/src/basics/sparsity_detection.md Outdated Show resolved Hide resolved
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@avik-pal avik-pal merged commit 89b679c into ap/sparse_di Oct 2, 2024
7 of 10 checks passed
@avik-pal avik-pal deleted the ap/structured_di branch October 2, 2024 16:55
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

Successfully merging this pull request may close these issues.

3 participants