Add EnzymeRules
Will fix #99
For some reason, I can't seem to get the extension to work. Package precompilation fails with the error:
ERROR: The following 1 direct dependency failed to precompile:
AbstractFFTs [621f4979-c628-5d54-868e-fcf4e3e8185c]
Failed to precompile AbstractFFTs [621f4979-c628-5d54-868e-fcf4e3e8185c] to "/home/runner/.julia/compiled/v1.9/AbstractFFTs/jl_mYHZQL".
ERROR: LoadError: ArgumentError: Package AbstractFFTs does not have LinearAlgebra in its dependencies:
- You may have a partially installed environment. Try `Pkg.instantiate()`
to ensure all packages in the environment are installed.
- Or, if you have AbstractFFTs checked out for development and have
added LinearAlgebra as a dependency but haven't updated your primary
environment's manifest file, try `Pkg.resolve()`.
- Otherwise you may need to report an issue with AbstractFFTs
although LinearAlgebra is clearly listed as both a dep and a weak dep.
Weirder still, if I activate the project, it now says it's empty, whereas if I remove this extension, it shows the dependencies:
julia> using Pkg; Pkg.activate(".")
Activating project at `~/projects/AbstractFFTs.jl`
julia> Pkg.status()
Project AbstractFFTs v1.3.1
Status `~/projects/AbstractFFTs.jl/Project.toml` (empty project)
shell> head ./Project.toml
name = "AbstractFFTs"
uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
version = "1.3.1"
[deps]
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
[weakdeps]
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@KristofferC I've never had this problem with my extensions before. Do you know what could cause this?
Nevermind, it seems extensions cannot have weak deps that are also deps. In this case, the dep needs to be loaded within the extension from the main package, see e.g. https://github.com/JuliaStats/LogExpFunctions.jl/pull/63
Codecov Report
:x: Patch coverage is 0% with 20 lines in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 78.60%. Comparing base (a25656d) to head (859abf0).
:warning: Report is 34 commits behind head on master.
| Files with missing lines | Patch % | Lines |
|---|---|---|
| ext/AbstractFFTsEnzymeCoreExt.jl | 0.00% | 20 Missing :warning: |
:exclamation: There is a different number of reports uploaded between BASE (a25656d) and HEAD (859abf0). Click for more details.
HEAD has 2 uploads less than BASE
Flag BASE (a25656d) HEAD (859abf0) 8 6
Additional details and impacted files
@@ Coverage Diff @@
## master #103 +/- ##
==========================================
- Coverage 87.08% 78.60% -8.48%
==========================================
Files 3 4 +1
Lines 209 229 +20
==========================================
- Hits 182 180 -2
- Misses 27 49 +22
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
:rocket: New features to boost your workflow:
- :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
If #67 is merged, we could add rules for *(::Plan, ::StridedArray), so long as the plan is Const (if it's non-Const, then we would need the rule to support it being an in-place plan, which we can't do).
I've paused work on this until EnzymeTestUtils (https://github.com/EnzymeAD/Enzyme.jl/pull/782) is registered, which will make testing these rules reliably much more straightforward.
Coming back to this, I think Enzyme rules should only be defined here abstractly for cases where we know they will not be breaking downstream code that otherwise Enzyme would have handled fine. So I agree with the following restrictions:
- Restrict eltypes to BLAS types
- Restrict array types to
StridedArrays - only have rules for
fft,fft!, and other other variants. In general we cannot tell if a plan is in-place or not. If we can catch cases where it isConst(i.e. Enzyme has inferred it is not used to carry any derivative information) without breaking fall backs, then great, but otherwise we don't define the rule.
These rules are considerably stricter than the ChainRules and for good reason. ChainRules are by convention often defined to cover up indexing code and mutating code to help Zygote and Diffractor, but this comes at the cost of doing the wrong thing for lots of types, hence the ProjectTo mechanism. Enzyme, on the other hand, can in principle handle many more types well, so we want to avoid writing rules that do the wrong thing for any cases where with no rule Enzyme would have worked fine.
Rules for * with Plans can be define in packages like FFTW where the type informs the in-placeness of the plan.
Any hope of having this PR revived? Enzyme has come a long way lately and FFT support would be another great step forward.
Any hope of having this PR revived? Enzyme has come a long way lately and FFT support would be another great step forward.
I'm afraid I don't have the bandwidth now to revive this. I still think https://github.com/JuliaMath/AbstractFFTs.jl/pull/103#issuecomment-1717625945 is the right way forward, and this PR is a good starting point for someone who wants to take it on.