add AD (Enzyme) support via MeshIntegralsEnzymeExt

[kylebeggs]

This is a fresh PR continuing the work in #129 because the rebase for that one was driving me insane:

Adds support to calculate the jacobian via Enzyme autodiff #35. This is done by adding a package extension for Enzyme, MeshIntegralsEnzymeExt.

[kylebeggs]

I simply created an error for a handful (4) of troublesome geometries with Enzyme so we can at least get this out for the rest of them

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (301bcff) to head (65dfe9c).
Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##              main      #152   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           18        18           
  Lines          165       184   +19     
=========================================
+ Hits           165       184   +19     

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

[github-actions]

Benchmark Results

	main	65dfe9c...	main/65dfe9c94c979d...
Differentials/Differential	0.24 ± 0.002 μs	0.203 ± 0.0019 μs	1.18
Differentials/Jacobian	0.205 ± 0.001 μs	0.167 ± 0.001 μs	1.23
Integrals/Segment/Scalar GaussKronrod	1.21 ± 0.007 μs	0.854 ± 0.012 μs	1.41
Integrals/Segment/Scalar GaussLegendre	4.11 ± 0.0086 μs	1.63 ± 0.0071 μs	2.52
Integrals/Segment/Scalar HAdaptiveCubature	1.44 ± 0.16 μs	1.16 ± 0.13 μs	1.24
Integrals/Segment/Vector GaussKronrod	3.43 ± 0.055 μs	3 ± 0.051 μs	1.14
Integrals/Segment/Vector GaussLegendre	19.8 ± 0.42 μs	17.2 ± 0.48 μs	1.15
Integrals/Segment/Vector HAdaptiveCubature	4.21 ± 0.082 μs	3.82 ± 0.078 μs	1.1
Integrals/Sphere/Scalar GaussKronrod	0.0823 ± 0.00078 ms	0.0726 ± 0.001 ms	1.13
Integrals/Sphere/Scalar GaussLegendre	2.25 ± 0.0086 ms	1.81 ± 0.0083 ms	1.24
Integrals/Sphere/Scalar HAdaptiveCubature	0.0578 ± 0.00012 ms	0.0472 ± 0.0001 ms	1.23
Integrals/Sphere/Vector GaussKronrod	0.116 ± 0.00096 ms	0.107 ± 0.00087 ms	1.09
Integrals/Sphere/Vector GaussLegendre	3.63 ± 0.071 ms	3.28 ± 0.067 ms	1.11
Integrals/Sphere/Vector HAdaptiveCubature	0.106 ± 0.00091 ms	0.0977 ± 0.00087 ms	1.09
Rules/GaussLegendre	21.8 ± 0.61 μs	21.8 ± 0.56 μs	1
Specializations/Scalar GaussLegendre/BezierCurve	0.247 ± 0.00067 ms	0.25 ± 0.0004 ms	0.985
Specializations/Scalar GaussLegendre/Line	8.96 ± 0.17 μs	7.43 ± 0.075 μs	1.21
Specializations/Scalar GaussLegendre/Plane	0.897 ± 0.0059 ms	0.738 ± 0.0023 ms	1.21
Specializations/Scalar GaussLegendre/Ray	7.79 ± 0.055 μs	5.97 ± 0.052 μs	1.3
Specializations/Scalar GaussLegendre/Rope	0.13 ± 0.0011 ms	0.0509 ± 0.00021 ms	2.55
Specializations/Scalar GaussLegendre/Tetrahedron	0.283 ± 0.004 s	0.281 ± 0.0028 s	1.01
Specializations/Scalar GaussLegendre/Triangle	0.729 ± 0.0054 ms	0.597 ± 0.0034 ms	1.22
time_to_load	1.51 ± 0.021 s	1.52 ± 0.02 s	0.992

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

[JoshuaLampert]

Thanks a lot! Please find below some suggestions and comments. The main suggestion is about refactoring the check_diff_method_support function.
I would be interested in some benchmarks. Could we maybe (at least temporarily) change the _default_diff_method to use AutoEnzyme to let the benchmarks script run for the AutoEnzyme backend, such that we get a comparison between this and the FiniteDifference backend?
This also leads to the question whether we want AutoEnzyme to be the _default_diff_backend permanently. We could define something like

function _default_diff_method(g::Type{G}) where {G <: Geometry}
  if supports_autoenzyme(g)
    return AutoEnzyme()
  else
    return FiniteDifference()
  end
end

but that's something to discuss and maybe also depends on the benchmark results. (The code snippet above wouldn't run with my suggestions below because they define supports_autoenzyme(::Geometry) and not supports_autoenzyme(::Type{G}) where {G <: Geometry}, but these are details.)

[JoshuaLampert]

This looks like a pretty high compat (v0.13.22 is only a couple of days old). Does the compat bound need to be that high?

[kylebeggs]

I'm sure it doesn't, but I've always wondered the best approach to finding the lowest
compat other than sweeping the tests with lower versions until it breaks?

[JoshuaLampert]

I'm usually coming from the other side meaning I first try very low compat bounds. If it fails I see the reason why it fails and check what the smallest compat bound is such that this particular reason is not a problem anymore (e.g., because I use functionality from a specific version). Then I directly jump to this version. If it breaks again I repeat until it doesn't break. This avoids to try out many different successful compat bounds. And it's usually faster to wait for failing tests than to wait for successful tests 😉
But I would also be fine if we just pick a smaller compat bound without spending too much time in trying to find the minimal possible one.

[JoshuaLampert]

You would need to using Enzyme in the benchmark script to let the benchmarks load the package extension.

[kylebeggs]

Ahh thanks. Initial ideas about the Downgrade run?

[mikeingold]

If we're updating the lower bound on Meshes.jl to v0.51.20 then we can probably remove the version check in the `combinations.jl` ***@***.*** "ParametrizedCurve"` section.

…

On Fri, Dec 13, 2024 at 10:30 AM Joshua Lampert ***@***.***> wrote: @JoshuaLampert requested changes on this pull request. ________________________________ In benchmark/Project.toml: > LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Meshes = "eacbb407-ea5a-433e-ab97-5258b1ca43fa" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [compat] BenchmarkTools = "1.5" +Enzyme = "0.13" LinearAlgebra = "1" Meshes = "0.50, 0.51, 0.52" ⬇️ Suggested change -Meshes = "0.50, 0.51, 0.52" +Meshes = "0.51.20, 0.52 — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because your review was requested.Message ID: ***@***.***>

[JoshuaLampert]

The benchmark results look quite nice. Overall mostly a pretty solid speedup 👍

[mikeingold]

I just added some formatting tweaks and docstrings.

Is it possible to make AutoEnzyme default based on whether the module is loaded and the geometry/FP are supported?

What currently happens if a user tries to use AutoEnzyme without loading the package/extension first?

[JoshuaLampert]

I think we are almost done.

[kylebeggs]

Is it possible to make AutoEnzyme default based on whether the module is loaded and the geometry/FP are supported?

What currently happens if a user tries to use AutoEnzyme without loading the package/extension first?

I had this originally, but was fighting with it a bit so I dropped it. If the user tries to use it withoput loading Enzyme, you'll just see a method error, but it won't tell you why that method does not exist. One way around this is to write jacobian(..., ::AutoEnzyme) which returns a helpful error message in the main part of the code and then when you load Enzyme, it should overwrite this with the actual jacobian?

[kylebeggs]

@mikeingold thanks for the docstrings and sorry about removing the newer FiniteDifference constructors.
@JoshuaLampert did we want to add an option to use different modes for Enzyme? I think that could be pretty easy.

[kylebeggs]

How come the benchmarks don't print out as a comment here anymore? Does it just take some time?

[JoshuaLampert]

@JoshuaLampert did we want to add an option to use different modes for Enzyme? I think that could be pretty easy.

IMO, we can also add it later if needed or desired. But if you think it's easy and can be done also in this PR, I would also be fine.

How come the benchmarks don't print out as a comment here anymore? Does it just take some time?

The initial comment will always be updated for each (successful) run.

[kylebeggs]

Ok, agreed. I think this may be ready to merge then?

[JoshuaLampert]

Thanks a lot @kylebeggs! I would like to see a final approvement by @mikeingold. Then this is good to go!

[mikeingold]

@mikeingold thanks for the docstrings and sorry about removing the newer FiniteDifference constructors.

No worries. I was going to add a comment but realized I already had the file open locally, so it was a quick fix.

@JoshuaLampert did we want to add an option to use different modes for Enzyme? I think that could be pretty easy.

Concur with @JoshuaLampert here. I think this is definitely worth exploring. If it’s a quick change with obvious implementation I’d be up for seeing it here, but it might be worth completing this as-is and then exploring that in a new PR/branch with a tighter focus.

[mikeingold]

Approved. Just let us know what you want to do about the other modes.

Edit: just saw your last post.