AbstractFFTs.jl icon indicating copy to clipboard operation
AbstractFFTs.jl copied to clipboard

Add EnzymeRules

Open sethaxen opened this issue 2 years ago • 8 comments

Will fix #99

sethaxen avatar May 22 '23 14:05 sethaxen

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?

sethaxen avatar May 22 '23 14:05 sethaxen

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

sethaxen avatar May 22 '23 22:05 sethaxen

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.

codecov[bot] avatar May 22 '23 22:05 codecov[bot]

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).

sethaxen avatar Jun 30 '23 10:06 sethaxen

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.

sethaxen avatar Aug 26 '23 19:08 sethaxen

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 is Const (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.

sethaxen avatar Sep 13 '23 13:09 sethaxen

Any hope of having this PR revived? Enzyme has come a long way lately and FFT support would be another great step forward.

danielwe avatar Aug 15 '24 05:08 danielwe

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.

sethaxen avatar Aug 18 '24 19:08 sethaxen