Skip to content

perf: HVP with in-place gradient + inner preparation#743

Merged
gdalle merged 15 commits intomainfrom
gd/hvp_inplacegrad
Mar 15, 2025
Merged

perf: HVP with in-place gradient + inner preparation#743
gdalle merged 15 commits intomainfrom
gd/hvp_inplacegrad

Conversation

@gdalle
Copy link
Copy Markdown
Member

@gdalle gdalle commented Mar 14, 2025

Fixes #86 by allowing preparation of the inner gradient in a forward-over-reverse HVP

  • Define new trait inner_preparation_behavior with three modalities:
    • PrepareInnerSimple (for finite difference backends)
    • PrepareInnerOverload (for ForwardDiff)
    • DontPrepareInner (for everyone else at the moment, but can evolve)
  • Restructure context type hierarchy and add a PrepContext for passing a preparation object created on the correct input type (it cannot be passed as a Cache because some of its values may matter):
    • ForwardDiff will keep it as-is
    • Enzyme will Duplicate it
  • Split hvp logic depending on
    • whether the outer backend supports in-place functions (determines whether to use shuffled_gradient or shuffled_gradient!)
    • what type of inner_preparation_behavior the outer backend has
  • Define overloaded_input for ForwardDiff and PolyesterForwardDiff to prepare the inner gradient on the right input type
  • Fix Cache handling for PolyesterForwardDiff (was incorrectly advertised as supported until now)
  • Add more tests, including on a new AutoForwardFromPrimitive

@codecov
Copy link
Copy Markdown

codecov Bot commented Mar 14, 2025

Codecov Report

Attention: Patch coverage is 99.02913% with 2 lines in your changes missing coverage. Please review.

Project coverage is 97.19%. Comparing base (6ae9532) to head (923eca4).

Files with missing lines Patch % Lines
DifferentiationInterface/src/second_order/hvp.jl 97.26% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #743      +/-   ##
==========================================
- Coverage   97.45%   97.19%   -0.26%     
==========================================
  Files         125      128       +3     
  Lines        6556     6678     +122     
==========================================
+ Hits         6389     6491     +102     
- Misses        167      187      +20     
Flag Coverage Δ
DI 97.84% <99.02%> (-0.39%) ⬇️
DIT 95.74% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@gdalle gdalle marked this pull request as ready for review March 14, 2025 20:20
@gdalle gdalle marked this pull request as draft March 14, 2025 21:22
@gdalle gdalle changed the title perf: compute in-place HVP from in-place gradient perf: HVP with in-place gradient + inner preparation Mar 14, 2025
@gdalle gdalle marked this pull request as ready for review March 15, 2025 09:18
@gdalle gdalle added the skipci Don't run CI tests and docs label Mar 15, 2025
@gdalle gdalle removed the skipci Don't run CI tests and docs label Mar 15, 2025
@gdalle gdalle merged commit 1a359ae into main Mar 15, 2025
4 checks passed
@gdalle gdalle deleted the gd/hvp_inplacegrad branch March 15, 2025 10:28
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.

Preparation for second order

1 participant