Skip to content

Commit

Permalink
Merge branch 'main' into tests
Browse files Browse the repository at this point in the history
  • Loading branch information
jgreener64 committed Oct 6, 2024
2 parents 06ef697 + 3c0871d commit cb387c7
Show file tree
Hide file tree
Showing 70 changed files with 15,295 additions and 6,752 deletions.
67 changes: 33 additions & 34 deletions .buildkite/pipeline.yml
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,10 @@ steps:
commands: |
echo "--- Setup Julia packages"
julia --color=yes -e '
import Pkg
Pkg.develop(; path = pwd())
Pkg.develop(; path = joinpath(pwd(), "lib", "EnzymeCore"))
Pkg.develop(; name = "CUDA")' || exit 3
using Pkg
pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")]
push!(pkgs, PackageSpec(; name="CUDA"))
Pkg.develop(pkgs)' || exit 3
echo "+++ Run tests"
julia --color=yes test/cuda.jl
Expand All @@ -41,40 +41,39 @@ steps:
commands: |
echo "--- Setup Julia packages"
julia --color=yes -e '
import Pkg
Pkg.develop(; path = pwd())
Pkg.develop(; path = joinpath(pwd(), "lib", "EnzymeCore"))
Pkg.develop(; name = "AMDGPU")' || exit 3
using Pkg
pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")]
push!(pkgs, PackageSpec(; name="AMDGPU"))
Pkg.develop(pkgs)' || exit 3
echo "+++ Run tests"
julia --color=yes test/amdgpu.jl
env:
JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager

# - label: "Metal Julia v{{matrix.version}}"
# matrix:
# setup:
# version:
# - "1.8"
# - "1.9"
# plugins:
# - JuliaCI/julia#v1:
# version: "{{matrix.version}}"
# agents:
# queue: "juliaecosystem"
# os: "macos"
# arch: "aarch64"
# if: build.message !~ /\[skip tests\]/
# timeout_in_minutes: 60
# commands: |
# echo "--- Setup Julia packages"
# julia --color=yes -e '
# import Pkg
# Pkg.develop(; path = pwd())
# Pkg.develop(; path = joinpath(pwd(), "lib", "EnzymeCore"))
# Pkg.develop(; name = "Metal")' || exit 3
- label: "Metal Julia v{{matrix.version}}"
matrix:
setup:
version:
- "1.10"
plugins:
- JuliaCI/julia#v1:
version: "{{matrix.version}}"
agents:
queue: "juliaecosystem"
os: "macos"
arch: "aarch64"
if: build.message !~ /\[skip tests\]/
timeout_in_minutes: 60
commands: |
echo "--- Setup Julia packages"
julia --color=yes -e '
using Pkg
pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")]
push!(pkgs, PackageSpec(; name="Metal"))
Pkg.develop(pkgs)' || exit 3
# echo "+++ Run tests"
# julia --color=yes test/metal.jl
# env:
# JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
echo "+++ Run tests"
julia --color=yes test/metal.jl
env:
JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
45 changes: 2 additions & 43 deletions .github/workflows/CI.yml
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,6 @@ jobs:
fail-fast: false
matrix:
version:
- '1.6'
- '1.7'
- '1.8'
- '1.9'
- '1.10'
- ~1.11.0-0
- 'nightly'
Expand All @@ -42,46 +38,11 @@ jobs:
arch: x64
libEnzyme: local
include:
- os: ubuntu-20.04
arch: x86
libEnzyme: packaged
version: '1.6'
assertions: false
- os: ubuntu-20.04
arch: x86
libEnzyme: packaged
version: '1.7'
assertions: false
- os: ubuntu-20.04
arch: x86
libEnzyme: packaged
version: '1.8'
assertions: false
- os: ubuntu-20.04
arch: x86
libEnzyme: packaged
version: '1.9'
assertions: false
- os: ubuntu-20.04
arch: x86
libEnzyme: packaged
version: '1.10'
assertions: false
- os: ubuntu-20.04
arch: x64
libEnzyme: packaged
version: '1.7'
assertions: true
- os: ubuntu-20.04
arch: x64
libEnzyme: packaged
version: '1.8'
assertions: true
- os: ubuntu-20.04
arch: x64
libEnzyme: packaged
version: '1.9'
assertions: true
- os: ubuntu-20.04
arch: x64
libEnzyme: packaged
Expand Down Expand Up @@ -125,7 +86,8 @@ jobs:
shell: julia --color=yes --project=. {0}
run: |
using Pkg
Pkg.develop(path="lib/EnzymeCore")
Pkg.develop([PackageSpec(; path) for path in ("lib/EnzymeCore", "lib/EnzymeTestUtils")])
Pkg.instantiate()
env:
JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
- name: Build libEnzyme
Expand Down Expand Up @@ -172,9 +134,6 @@ jobs:
fail-fast: false
matrix:
version:
- '1.7'
- '1.8'
- '1.9'
- '1.10'
- ~1.11.0-0
- 'nightly'
Expand Down
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
*.jl.cov
*.jl.mem
/Manifest.toml
/Manifest-v*.toml
/test/Manifest.toml
/docs/Manifest.toml
/docs/build/
Expand Down
10 changes: 6 additions & 4 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "Enzyme"
uuid = "7da242da-08ed-463a-9acd-ee780be4f1d9"
authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>"]
version = "0.12.35"
version = "0.13.7"

[deps]
CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
Expand All @@ -15,6 +15,7 @@ ObjectFile = "d8793406-e978-5875-9003-1fc021f44a92"
Preferences = "21216c6a-2e73-6563-6e65-726566657250"
Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[weakdeps]
BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
Expand All @@ -34,16 +35,17 @@ EnzymeStaticArraysExt = "StaticArrays"
BFloat16s = "0.2, 0.3, 0.4, 0.5"
CEnum = "0.4, 0.5"
ChainRulesCore = "1"
EnzymeCore = "0.7.8"
Enzyme_jll = "0.0.146"
EnzymeCore = "0.8.4"
Enzyme_jll = "0.0.153"
GPUCompiler = "0.21, 0.22, 0.23, 0.24, 0.25, 0.26, 0.27"
LLVM = "6.1, 7, 8, 9"
LogExpFunctions = "0.3"
ObjectFile = "0.4"
Preferences = "1.4"
SparseArrays = "1"
SpecialFunctions = "1, 2"
StaticArrays = "1"
julia = "1.6"
julia = "1.10"

[extras]
BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
Expand Down
3 changes: 2 additions & 1 deletion docs/Project.toml
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
[deps]
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[compat]
Literate = "2"
Documenter = "1"
Literate = "2"
22 changes: 14 additions & 8 deletions docs/src/faq.md
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,7 @@ That is why Enzyme provides a helper function `Enzyme.make_zero` that does this

```jldoctest sparse
Enzyme.make_zero(a)
Enzyme.gradient(Reverse, sum, a) # This calls make_zero(a)
Enzyme.gradient(Reverse, sum, a)[1] # This calls make_zero(a)
# output
Expand Down Expand Up @@ -268,7 +268,7 @@ Enzyme.autodiff(Reverse, f, Active(1.2), Const(Vector{Float64}(undef, 1)), Const
((0.0, nothing, nothing, nothing),)
```

Passing in a dupliacted (e.g. differentiable) variable for `tmp` now leads to the correct answer.
Passing in a duplicated (e.g. differentiable) variable for `tmp` now leads to the correct answer.

```jldoctest storage
Enzyme.autodiff(Reverse, f, Active(1.2), Duplicated(Vector{Float64}(undef, 1), zeros(1)), Const(1), Const(5)) # Correct (returns 10.367999999999999 == 1.2^4 * 5)
Expand All @@ -278,9 +278,11 @@ Enzyme.autodiff(Reverse, f, Active(1.2), Duplicated(Vector{Float64}(undef, 1), z
((10.367999999999999, nothing, nothing, nothing),)
```

However, even if we ignore the semantic guarantee provided by marking `tmp` as constant, another issue arises. When computing the original function, intermediate computations (like in `f` above) can use `tmp` for temporary storage. When computing the derivative, Enzyme also needs additional temporary storage space for the corresponding derivative variables as well. If `tmp` is marked as Const, Enzyme does not have any temporary storage space for the derivatives!
## Runtime Activity

Recent versions of Enzyme will attempt to error when they detect these latter types of situations, which we will refer to as `activity unstable`. This term is chosen to mirror the Julia notion of type-unstable code (e.g. where a type is not known at compile time). If an expression is activity unstable, it could either be constant, or active, depending on data not known at compile time. For example, consider the following:
When computing the derivative of mutable variables, Enzyme also needs additional temporary storage space for the corresponding derivative variables. If an argument `tmp` is marked as Const, Enzyme does not have any temporary storage space for the derivatives!

Enzyme will error when they detect these latter types of situations, which we will refer to as `activity unstable`. This term is chosen to mirror the Julia notion of type-unstable code (e.g. where a type is not known at compile time). If an expression is activity unstable, it could either be constant, or active, depending on data not known at compile time. For example, consider the following:

```julia
function g(cond, active_var, constant_var)
Expand All @@ -293,7 +295,7 @@ end
Enzyme.autodiff(Forward, g, Const(condition), Duplicated(x, dx), Const(y))
```

The returned value here could either by constant or duplicated, depending on the runtime-defined value of `cond`. If `cond` is true, Enzyme simply returns the shadow of `active_var` as the derivative. However, if `cond` is false, there is no derivative shadow for `constant_var` and Enzyme will throw a "Mismatched activity" error. For some simple types, e.g. a float Enzyme can circumvent this issue, for example by returning the float 0. Similarly, for some types like the Symbol type, which are never differentiable, such a shadow value will never be used, and Enzyme can return the original "primal" value as its derivative. However, for arbitrary data structures, Enzyme presently has no generic mechanism to resolve this.
The returned value here could either by constant or duplicated, depending on the runtime-defined value of `cond`. If `cond` is true, Enzyme simply returns the shadow of `active_var` as the derivative. However, if `cond` is false, there is no derivative shadow for `constant_var` and Enzyme will throw a `EnzymeRuntimeActivityError` error. For some simple types, e.g. a float Enzyme can circumvent this issue, for example by returning the float 0. Similarly, for some types like the Symbol type, which are never differentiable, such a shadow value will never be used, and Enzyme can return the original "primal" value as its derivative. However, for arbitrary data structures, Enzyme presently has no generic mechanism to resolve this.

For example consider a third function:
```julia
Expand All @@ -308,13 +310,17 @@ Enzyme provides a nice utility `Enzyme.make_zero` which takes a data structure a

If one created a new zero'd copy of each return from `g`, this would mean that the derivative `dresult` would have one copy made for the first element, and a second copy made for the second element. This could lead to incorrect results, and is unfortunately not a general resolution. However, for non-mutable variables (e.g. like floats) or non-differrentiable types (e.g. like Symbols) this problem can never arise.

Instead, Enzyme has a special mode known as "Runtime Activity" which can handle these types of situations. It can come with a minor performance reduction, and is therefore off by default. It can be enabled with `Enzyme.API.runtimeActivity!(true)` right after importing Enzyme for the first time.
Instead, Enzyme has a special mode known as "Runtime Activity" which can handle these types of situations. It can come with a minor performance reduction, and is therefore off by default. It can be enabled with by setting runtime activity to true in a desired differentiation mode.

The way Enzyme's runtime activity resolves this issue is to return the original primal variable as the derivative whenever it needs to denote the fact that a variable is a constant. As this issue can only arise with mutable variables, they must be represented in memory via a pointer. All addtional loads and stores will now be modified to first check if the primal pointer is the same as the shadow pointer, and if so, treat it as a constant. Note that this check is not saying that the same arrays contain the same values, but rather the same backing memory represents both the primal and the shadow (e.g. `a === b` or equivalently `pointer(a) == pointer(b)`).

Enabling runtime activity does therefore, come with a sharp edge, which is that if the computed derivative of a function is mutable, one must also check to see if the primal and shadow represent the same pointer, and if so the true derivative of the function is actually zero.

Generally, the preferred solution to these type of activity unstable codes should be to make your variables all activity-stable (e.g. always containing differentiable memory or always containing non-differentiable memory). However, with care, Enzyme does support "Runtime Activity" as a way to differentiate these programs without having to modify your code.
Generally, the preferred solution to these type of activity unstable codes should be to make your variables all activity-stable (e.g. always containing differentiable memory or always containing non-differentiable memory). However, with care, Enzyme does support "Runtime Activity" as a way to differentiate these programs without having to modify your code. One can enable runtime activity for your code by changing the mode, such as

```julia
Enzyme.autodiff(set_runtime_activity(Forward), h, Const(condition), Duplicated(x, dx), Const(y))
```

## Mixed activity

Expand Down Expand Up @@ -621,7 +627,7 @@ Presently Enzyme only considers floats as base types. As a result, Enzyme does n

```jldoctest types
f_int(x) = x * x
Enzyme.autodiff(Forward, f_int, DuplicatedNoNeed, Duplicated(3, 1))
Enzyme.autodiff(Forward, f_int, Duplicated, Duplicated(3, 1))
# output
Expand Down
Loading

0 comments on commit cb387c7

Please sign in to comment.