Skip to content

Handle FIRK expanded AD matvec fallbacks - #557

Draft
ChrisRackauckas-Claude wants to merge 4 commits into
SciML:masterfrom
ChrisRackauckas-Claude:fix-firk-pre-enzyme-fallback-clean
Draft

Handle FIRK expanded AD matvec fallbacks#557
ChrisRackauckas-Claude wants to merge 4 commits into
SciML:masterfrom
ChrisRackauckas-Claude:fix-firk-pre-enzyme-fallback-clean

Conversation

@ChrisRackauckas-Claude

@ChrisRackauckas-Claude ChrisRackauckas-Claude commented Jul 11, 2026

Copy link
Copy Markdown
Member

This PR should be ignored until reviewed by @ChrisRackauckas.

Stacked on #552. Until #552 is merged, GitHub's PR diff will include those prerequisite commits; the new change here is the top commit c72455db.

Summary

  • Use FIRK-local scalar interpolation matvecs for polynomial interpolation coefficient products used by interpolation boundary conditions.
  • On Julia 1.12+, route expanded FIRK Enzyme/Mooncake collocation and whole-residual AD modes through AutoForwardDiff, while leaving standard-problem BC AD modes intact.
  • Keep the fallback local to FIRK paths so ordinary Core/BLAS matvecs used by nested FIRK convergence tests are not redirected.
  • Derive FIRK interpolation matvec output types from multiplication instead of promote_type, avoiding BigFloat/SparseConnectivityTracer.Dual promotion ambiguity in Misc tests.

Fixes #554.

Local verification

Targeted NESTED LobattoIIIc4 convergence repro on Julia 1.12.6 after keeping Core matvecs on BLAS:

(i = 9, order = 5.851916060671115, errors = [5.10702591327572e-14, 2.8337332480532496e-12, 1.7036194677189087e-10])
(i = 10, order = 5.851916060671115, errors = [5.10702591327572e-14, 2.8319568912138493e-12, 1.7036194677189087e-10])
real 330.36

Full FIRK AD test group on Julia 1.12.6:

FIRK Expanded AD Tests | 9 pass, total 9, 9m42.8s
Testing BoundaryValueDiffEqFIRK tests passed
real 1135.61

Full root Misc test group on Julia 1.10.11 after the product-type fix:

BigFloat Tests | 7 pass, total 7, 3m59.1s
Testing BoundaryValueDiffEq tests passed
real 2904.12

Fresh-resolved targeted multipoint Enzyme repros on Julia 1.13.0-rc1 after the product-type fix:

array retcode=Success
interp retcode=Success
real 432.25

Formatting/checks:

Runic --check lib/BoundaryValueDiffEqFIRK/src/adaptivity.jl
Runic --check lib/BoundaryValueDiffEqCore/src/utils.jl
passed

git diff --check
passed

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the fix-firk-pre-enzyme-fallback-clean branch from a51bd47 to 344016e Compare July 11, 2026 21:02
@github-actions

github-actions Bot commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

Benchmark Results (Julia v1.11)

Time benchmarks
master c72455d... master / c72455d...
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK2() 0.66 ± 0.013 s 0.663 ± 0.007 s 0.996 ± 0.022
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK3() 13.4 ± 0.22 ms 13.4 ± 0.28 ms 0.998 ± 0.026
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK4() 2.84 ± 0.1 ms 2.89 ± 0.1 ms 0.982 ± 0.049
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK5() 3.46 ± 0.13 ms 3.48 ± 0.11 ms 0.994 ± 0.049
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK6() 1.65 ± 0.12 ms 1.67 ± 0.097 ms 0.984 ± 0.094
Simple Pendulum/IIP/MultipleShooting(10, Tsit5; grid_coarsening = false) 1.65 ± 0.16 ms 1.64 ± 0.11 ms 1.01 ± 0.12
Simple Pendulum/IIP/MultipleShooting(10, Tsit5; grid_coarsening = true) 3.32 ± 0.3 ms 3.4 ± 0.24 ms 0.979 ± 0.11
Simple Pendulum/IIP/MultipleShooting(100, Tsit5; grid_coarsening = false) 0.0437 ± 0.0048 s 0.0438 ± 0.005 s 0.998 ± 0.16
Simple Pendulum/IIP/MultipleShooting(100, Tsit5; grid_coarsening = true) 0.0694 ± 0.0022 s 0.0702 ± 0.002 s 0.988 ± 0.042
Simple Pendulum/IIP/Shooting(Tsit5()) 0.187 ± 0.085 ms 0.187 ± 0.088 ms 1 ± 0.65
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK2() 0.846 ± 0.012 s 0.874 ± 0.0076 s 0.968 ± 0.016
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK3() 18 ± 8.2 ms 18.1 ± 8.2 ms 0.997 ± 0.64
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK4() 3.56 ± 0.23 ms 3.62 ± 0.25 ms 0.984 ± 0.093
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK5() 4.37 ± 0.29 ms 4.42 ± 0.27 ms 0.988 ± 0.09
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK6() 2.03 ± 0.2 ms 2.08 ± 0.2 ms 0.974 ± 0.13
Simple Pendulum/OOP/MultipleShooting(10, Tsit5; grid_coarsening = false) 3.98 ± 1.1 ms 3.98 ± 3.2 ms 1 ± 0.85
Simple Pendulum/OOP/MultipleShooting(10, Tsit5; grid_coarsening = true) 7.71 ± 6.5 ms 7.57 ± 6.2 ms 1.02 ± 1.2
Simple Pendulum/OOP/MultipleShooting(100, Tsit5; grid_coarsening = false) 0.111 ± 0.0056 s 0.111 ± 0.0047 s 0.998 ± 0.066
Simple Pendulum/OOP/MultipleShooting(100, Tsit5; grid_coarsening = true) 0.172 ± 0.0087 s 0.175 ± 0.0086 s 0.984 ± 0.069
Simple Pendulum/OOP/Shooting(Tsit5()) 0.651 ± 0.12 ms 0.655 ± 0.14 ms 0.994 ± 0.28
time_to_load 8.06 ± 0.029 s 8.2 ± 0.017 s 0.983 ± 0.0041
Memory benchmarks
master c72455d... master / c72455d...
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK2() 0.389 M allocs: 0.0443 GB 0.389 M allocs: 0.0443 GB 1
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK3() 0.0436 M allocs: 4.88 MB 0.0436 M allocs: 4.88 MB 1
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK4() 15.7 k allocs: 1.65 MB 15.7 k allocs: 1.65 MB 1
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK5() 22.1 k allocs: 2.03 MB 22.1 k allocs: 2.03 MB 1
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK6() 12.8 k allocs: 1.05 MB 12.8 k allocs: 1.05 MB 1
Simple Pendulum/IIP/MultipleShooting(10, Tsit5; grid_coarsening = false) 25.5 k allocs: 1.82 MB 25.5 k allocs: 1.82 MB 1
Simple Pendulum/IIP/MultipleShooting(10, Tsit5; grid_coarsening = true) 0.049 M allocs: 3.39 MB 0.049 M allocs: 3.39 MB 1
Simple Pendulum/IIP/MultipleShooting(100, Tsit5; grid_coarsening = false) 0.553 M allocs: 0.0535 GB 0.553 M allocs: 0.0535 GB 1
Simple Pendulum/IIP/MultipleShooting(100, Tsit5; grid_coarsening = true) 0.833 M allocs: 0.0778 GB 0.833 M allocs: 0.0778 GB 1
Simple Pendulum/IIP/Shooting(Tsit5()) 4.63 k allocs: 0.224 MB 4.63 k allocs: 0.224 MB 1
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK2() 0.89 M allocs: 0.984 GB 0.89 M allocs: 0.984 GB 1
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK3() 0.0932 M allocs: 24.8 MB 0.0932 M allocs: 24.8 MB 1
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK4() 0.0324 M allocs: 3.96 MB 0.0324 M allocs: 3.96 MB 1
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK5() 0.045 M allocs: 4.99 MB 0.045 M allocs: 4.99 MB 1
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK6() 25.3 k allocs: 2.17 MB 25.3 k allocs: 2.17 MB 1
Simple Pendulum/OOP/MultipleShooting(10, Tsit5; grid_coarsening = false) 0.142 M allocs: 10.2 MB 0.142 M allocs: 10.2 MB 1
Simple Pendulum/OOP/MultipleShooting(10, Tsit5; grid_coarsening = true) 0.266 M allocs: 18.7 MB 0.266 M allocs: 18.7 MB 1
Simple Pendulum/OOP/MultipleShooting(100, Tsit5; grid_coarsening = false) 2.52 M allocs: 0.279 GB 2.52 M allocs: 0.279 GB 1
Simple Pendulum/OOP/MultipleShooting(100, Tsit5; grid_coarsening = true) 3.82 M allocs: 0.404 GB 3.82 M allocs: 0.404 GB 1
Simple Pendulum/OOP/Shooting(Tsit5()) 0.0373 M allocs: 1.69 MB 0.0373 M allocs: 1.69 MB 1
time_to_load 0.159 k allocs: 11.2 kB 0.159 k allocs: 11.2 kB 1

@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the fix-firk-pre-enzyme-fallback-clean branch from 344016e to c1c3376 Compare July 11, 2026 22:02
@ChrisRackauckas-Claude ChrisRackauckas-Claude changed the title Handle FIRK prerelease AD matvec fallbacks Handle FIRK expanded AD matvec fallbacks Jul 11, 2026
@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the fix-firk-pre-enzyme-fallback-clean branch from c1c3376 to 0f7b9f2 Compare July 12, 2026 00:26
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the fix-firk-pre-enzyme-fallback-clean branch from 0f7b9f2 to c72455d Compare July 12, 2026 01:42
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.

FIRK AD tests fail on Julia prerelease with EnzymeNoDerivativeError in BLAS gemv

2 participants