diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index ec098644bb..6f1e019541 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -3,21 +3,66 @@ steps: plugins: - JuliaCI/julia#v1: version: "1.6" - - JuliaCI/julia-test#v1: ~ agents: queue: "juliagpu" cuda: "*" if: build.message !~ /\[skip tests\]/ timeout_in_minutes: 60 + commands: | + echo "--- Setup Julia packages" + julia --color=yes --project=. -e ' + import Pkg + Pkg.develop(; path = joinpath(pwd(), "lib", "EnzymeCore")) + ' + + echo "+++ Run tests" + julia --color=yes --project=. -e ' + import Pkg + Pkg.test("Enzyme") + ' + +steps: + - label: "Julia v1.7" + plugins: + - JuliaCI/julia#v1: + version: "1.7" + agents: + queue: "juliagpu" + cuda: "*" + if: build.message !~ /\[skip tests\]/ + timeout_in_minutes: 60 + commands: | + echo "--- Setup Julia packages" + julia --color=yes --project=. -e ' + import Pkg + Pkg.develop(; path = joinpath(pwd(), "lib", "EnzymeCore")) + ' + + echo "+++ Run tests" + julia --color=yes --project=. -e ' + import Pkg + Pkg.test("Enzyme") + ' steps: - label: "Julia v1" plugins: - JuliaCI/julia#v1: version: "1" - - JuliaCI/julia-test#v1: ~ agents: queue: "juliagpu" cuda: "*" if: build.message !~ /\[skip tests\]/ - timeout_in_minutes: 60 \ No newline at end of file + timeout_in_minutes: 60 + commands: | + echo "--- Setup Julia packages" + julia --color=yes --project=. -e ' + import Pkg + Pkg.develop(; path = joinpath(pwd(), "lib", "EnzymeCore")) + ' + + echo "+++ Run tests" + julia --color=yes --project=. -e ' + import Pkg + Pkg.test("Enzyme") + ' \ No newline at end of file diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 80c889ca5c..72f9af6e59 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -73,7 +73,11 @@ jobs: ${{ runner.os }}-test-${{ env.cache-name }}- ${{ runner.os }}-test- ${{ runner.os }}- - - uses: julia-actions/julia-buildpkg@v1 + - name: add EnzymeCore + shell: julia --color=yes --project=. {0} + run: | + using Pkg + Pkg.develop(path="lib/EnzymeCore") - name: Build libEnzyme if: ${{ matrix.libEnzyme == 'local' && matrix.os != 'macOS-latest'}} run: | @@ -86,6 +90,7 @@ jobs: julia --project=deps -e 'using Pkg; Pkg.instantiate()' SDKROOT=`xcrun --show-sdk-path` julia --project=deps deps/build_local.jl cp LocalPreferences.toml test/ + - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 - uses: julia-actions/julia-processcoverage@v1 - uses: codecov/codecov-action@v1 @@ -102,6 +107,7 @@ jobs: - run: | julia --project=docs -e ' using Pkg + Pkg.develop(path="lib/EnzymeCore") Pkg.develop(PackageSpec(path=pwd())) Pkg.instantiate()' - run: | diff --git a/.github/workflows/scripts_deploy.yml b/.github/workflows/scripts_deploy.yml index b3bb1d7977..2ddeffada6 100644 --- a/.github/workflows/scripts_deploy.yml +++ b/.github/workflows/scripts_deploy.yml @@ -19,6 +19,7 @@ jobs: - run: | julia --project=docs -e ' using Pkg + Pkg.develop(path="lib/EnzymeCore") Pkg.develop(PackageSpec(path=pwd())) Pkg.instantiate()' - run: | diff --git a/Project.toml b/Project.toml index 6d557f3b73..6defb13344 100644 --- a/Project.toml +++ b/Project.toml @@ -4,8 +4,8 @@ authors = ["William Moses ", "Valentin Churavy ", "Valentin Churavy "] +version = "0.1.0" + +[deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" + +[compat] +Adapt = "3.3" +julia = "1.6" diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl new file mode 100644 index 0000000000..0b34b51743 --- /dev/null +++ b/lib/EnzymeCore/src/EnzymeCore.jl @@ -0,0 +1,129 @@ +module EnzymeCore + +using Adapt + +export Forward, Reverse +export Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplicatedNoNeed + +function batch_size end + +""" + abstract type Annotation{T} + +Abstract type for [`autodiff`](@ref) function argument wrappers like +[`Const`](@ref), [`Active`](@ref) and [`Duplicated`](@ref). +""" +abstract type Annotation{T} end +Base.eltype(::Type{<:Annotation{T}}) where T = T + +""" + Const(x) + +Mark a function argument `x` of [`autodiff`](@ref) as constant, +Enzyme will not auto-differentiate in respect `Const` arguments. +""" +struct Const{T} <: Annotation{T} + val::T +end +Adapt.adapt_structure(to, x::Const) = Const(adapt(to, x.val)) + +# To deal with Const(Int) and prevent it to go to `Const{DataType}(T)` +Const(::Type{T}) where T = Const{Type{T}}(T) + +""" + Active(x) + +Mark a function argument `x` of [`autodiff`](@ref) as active, +Enzyme will auto-differentiate in respect `Active` arguments. + +!!! note + + Enzyme gradients with respect to integer values are zero. + [`Active`](@ref) will automatically convert plain integers to floating + point values, but cannot do so for integer values in tuples and structs. +""" +struct Active{T} <: Annotation{T} + val::T +end +Adapt.adapt_structure(to, x::Active) = Active(adapt(to, x.val)) + +Active(i::Integer) = Active(float(i)) + +""" + Duplicated(x, ∂f_∂x) + +Mark a function argument `x` of [`autodiff`](@ref) as duplicated, Enzyme will +auto-differentiate in respect to such arguments, with `dx` acting as an +accumulator for gradients (so ``\\partial f / \\partial x`` will be *added to*) +`∂f_∂x`. +""" +struct Duplicated{T} <: Annotation{T} + val::T + dval::T +end +Adapt.adapt_structure(to, x::Duplicated) = Duplicated(adapt(to, x.val), adapt(to, x.dval)) + +""" + DuplicatedNoNeed(x, ∂f_∂x) + +Like [`Duplicated`](@ref), except also specifies that Enzyme may avoid computing +the original result and only compute the derivative values. +""" +struct DuplicatedNoNeed{T} <: Annotation{T} + val::T + dval::T +end +Adapt.adapt_structure(to, x::DuplicatedNoNeed) = DuplicatedNoNeed(adapt(to, x.val), adapt(to, x.dval)) + +""" + BatchDuplicated(x, ∂f_∂xs) + +Like [`Duplicated`](@ref), except contains several shadows to compute derivatives +for all at once. Argument `∂f_∂xs` should be a tuple of the several values of type `x`. +""" +struct BatchDuplicated{T,N} <: Annotation{T} + val::T + dval::NTuple{N,T} +end +Adapt.adapt_structure(to, x::BatchDuplicated) = BatchDuplicated(adapt(to, x.val), adapt(to, x.dval)) + +""" + BatchDuplicatedNoNeed(x, ∂f_∂xs) + +Like [`DuplicatedNoNeed`](@ref), except contains several shadows to compute derivatives +for all at once. Argument `∂f_∂xs` should be a tuple of the several values of type `x`. +""" +struct BatchDuplicatedNoNeed{T,N} <: Annotation{T} + val::T + dval::NTuple{N,T} +end +batch_size(::BatchDuplicated{T,N}) where {T,N} = N +batch_size(::BatchDuplicatedNoNeed{T,N}) where {T,N} = N +Adapt.adapt_structure(to, x::BatchDuplicatedNoNeed) = BatchDuplicatedNoNeed(adapt(to, x.val), adapt(to, x.dval)) + +""" + abstract type Mode + +Abstract type for what differentiation mode will be used. +""" +abstract type Mode end + +""" + struct Reverse <: Mode + +Reverse mode differentiation +""" +struct ReverseMode <: Mode +end +const Reverse = ReverseMode() + +""" + struct Forward <: Mode + +Forward mode differentiation +""" +struct ForwardMode <: Mode +end +const Forward = ForwardMode() + +end # module EnzymeCore diff --git a/src/Enzyme.jl b/src/Enzyme.jl index 753c87ffab..7a992d0568 100644 --- a/src/Enzyme.jl +++ b/src/Enzyme.jl @@ -1,136 +1,24 @@ module Enzyme +import EnzymeCore: Forward, Reverse export Forward, Reverse + +import EnzymeCore: Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplicatedNoNeed export Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplicatedNoNeed + +import EnzymeCore: batch_size +export batch_size + export autodiff, jacobian, gradient, gradient! export markType, batch_size, onehot, chunkedonehot using LinearAlgebra +import EnzymeCore: ReverseMode, ForwardMode # Independent code, must be loaded before "compiler.jl" include("pmap.jl") -using Adapt - -""" - abstract type Annotation{T} - -Abstract type for [`autodiff`](@ref) function argument wrappers like -[`Const`](@ref), [`Active`](@ref) and [`Duplicated`](@ref). -""" -abstract type Annotation{T} end -Base.eltype(::Type{<:Annotation{T}}) where T = T - -""" - Const(x) - -Mark a function argument `x` of [`autodiff`](@ref) as constant, -Enzyme will not auto-differentiate in respect `Const` arguments. -""" -struct Const{T} <: Annotation{T} - val::T -end -Adapt.adapt_structure(to, x::Const) = Const(adapt(to, x.val)) - -# To deal with Const(Int) and prevent it to go to `Const{DataType}(T)` -Const(::Type{T}) where T = Const{Type{T}}(T) - -""" - Active(x) - -Mark a function argument `x` of [`autodiff`](@ref) as active, -Enzyme will auto-differentiate in respect `Active` arguments. - -!!! note - - Enzyme gradients with respect to integer values are zero. - [`Active`](@ref) will automatically convert plain integers to floating - point values, but cannot do so for integer values in tuples and structs. -""" -struct Active{T} <: Annotation{T} - val::T -end -Adapt.adapt_structure(to, x::Active) = Active(adapt(to, x.val)) - -Active(i::Integer) = Active(float(i)) - -""" - Duplicated(x, ∂f_∂x) - -Mark a function argument `x` of [`autodiff`](@ref) as duplicated, Enzyme will -auto-differentiate in respect to such arguments, with `dx` acting as an -accumulator for gradients (so ``\\partial f / \\partial x`` will be *added to*) -`∂f_∂x`. -""" -struct Duplicated{T} <: Annotation{T} - val::T - dval::T -end -Adapt.adapt_structure(to, x::Duplicated) = Duplicated(adapt(to, x.val), adapt(to, x.dval)) - -""" - DuplicatedNoNeed(x, ∂f_∂x) - -Like [`Duplicated`](@ref), except also specifies that Enzyme may avoid computing -the original result and only compute the derivative values. -""" -struct DuplicatedNoNeed{T} <: Annotation{T} - val::T - dval::T -end -Adapt.adapt_structure(to, x::DuplicatedNoNeed) = DuplicatedNoNeed(adapt(to, x.val), adapt(to, x.dval)) - -""" - BatchDuplicated(x, ∂f_∂xs) - -Like [`Duplicated`](@ref), except contains several shadows to compute derivatives -for all at once. Argument `∂f_∂xs` should be a tuple of the several values of type `x`. -""" -struct BatchDuplicated{T,N} <: Annotation{T} - val::T - dval::NTuple{N,T} -end -Adapt.adapt_structure(to, x::BatchDuplicated) = BatchDuplicated(adapt(to, x.val), adapt(to, x.dval)) - -""" - BatchDuplicatedNoNeed(x, ∂f_∂xs) - -Like [`DuplicatedNoNeed`](@ref), except contains several shadows to compute derivatives -for all at once. Argument `∂f_∂xs` should be a tuple of the several values of type `x`. -""" -struct BatchDuplicatedNoNeed{T,N} <: Annotation{T} - val::T - dval::NTuple{N,T} -end -batch_size(::BatchDuplicated{T,N}) where {T,N} = N -batch_size(::BatchDuplicatedNoNeed{T,N}) where {T,N} = N -Adapt.adapt_structure(to, x::BatchDuplicatedNoNeed) = BatchDuplicatedNoNeed(adapt(to, x.val), adapt(to, x.dval)) - -""" - abstract type Mode - -Abstract type for what differentiation mode will be used. -""" -abstract type Mode end - -""" - struct Reverse <: Mode - -Reverse mode differentiation -""" -struct ReverseMode <: Mode -end -const Reverse = ReverseMode() guess_activity(::Type{T}, ::ReverseMode) where T = guess_activity(T) - -""" - struct Forward <: Mode - -Forward mode differentiation -""" -struct ForwardMode <: Mode -end -const Forward = ForwardMode() guess_activity(::Type{T}, ::ForwardMode) where T = guess_activity(T, API.DEM_ForwardMode) import LLVM