Skip to content

Commit

Permalink
Add EnzymeRule for conv!/gather!/scatter!/dropout!/pool! (#536)
Browse files Browse the repository at this point in the history
* Add EnzymeRule for conv

* Fix

* Add missing file

* Change to enzymecore ext

* attempt fix

* Finish enzymecore rewrite

* Add missing file

* Also add gather

* Additional functions, tests, and fixes

* minor fixup

* Add pooling

* Add dropout

* Fix scatter bug

* fix pool

* More fixups

* fix up pool

* split conv/depth

* Cleanup rule

* Fix minor test bug

* Fix depthwise conv

* Fix typo

* Bound tests

* Failing to extension

* Add file

* Address review

* Remove inlining
  • Loading branch information
wsmoses authored Sep 28, 2023
1 parent d0a03f2 commit aea063c
Show file tree
Hide file tree
Showing 8 changed files with 542 additions and 8 deletions.
19 changes: 12 additions & 7 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,31 +16,37 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[weakdeps]
AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"

[extensions]
NNlibAMDGPUExt = "AMDGPU"
NNlibCUDAExt = "CUDA"
NNlibCUDACUDNNExt = ["CUDA", "cuDNN"]
NNlibCUDAExt = "CUDA"
NNlibEnzymeCoreExt = "EnzymeCore"

[compat]
AMDGPU = "0.5, 0.6"
Adapt = "3.2"
Atomix = "0.1"
ChainRulesCore = "1.13"
CUDA = "4, 5"
cuDNN = "1"
ChainRulesCore = "1.13"
EnzymeCore = "0.5, 0.6"
GPUArraysCore = "0.1"
KernelAbstractions = "0.9.2"
Requires = "1.0"
cuDNN = "1"
julia = "1.9"

[extras]
AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
Expand All @@ -52,6 +58,5 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"

[targets]
test = ["AMDGPU", "CUDA", "ChainRulesTestUtils", "Documenter",
"FiniteDifferences", "ForwardDiff", "Logging", "ReverseDiff",
"StableRNGs", "Test", "UnicodePlots", "Zygote", "cuDNN"]
test = ["AMDGPU", "CUDA", "ChainRulesTestUtils", "Documenter", "FiniteDifferences", "ForwardDiff", "Logging", "ReverseDiff", "StableRNGs", "Test", "UnicodePlots", "Zygote", "cuDNN",
"Enzyme", "EnzymeCore", "EnzymeTestUtils"]
Loading

0 comments on commit aea063c

Please sign in to comment.