Skip to content

Commit 3c5f99e

Browse files
yebaiyebai
authored
Merge branch 'main' into develop
2 parents 0f0b9fc + eeb4c86 commit 3c5f99e

17 files changed

Lines changed: 168 additions & 75 deletions

File tree

‎.buildkite/pipeline.yml‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
11
steps:
22
- label: "DI GPU tests"
3-
if: |
4-
!build.pull_request.draft &&
5-
build.pull_request.labels includes "gpu"
3+
if: build.pull_request.labels includes "gpu"
64
plugins:
75
- JuliaCI/julia#v1:
86
version: "1"
97
command: |
10-
julia ./DifferentiationInterface/test/GPU/CUDA/simple.jl
8+
julia ./DifferentiationInterface/test/GPU/CUDA/main.jl
119
agents:
1210
queue: "juliagpu"
1311
cuda: "*"

‎DifferentiationInterface/CHANGELOG.md‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77

88
## [Unreleased]
99

10+
## [0.7.1]
11+
12+
### Feat
13+
14+
- Use Mooncake's internal copy utilities ([#809])
15+
16+
### Fixed
17+
18+
- Take `absstep` into account for FiniteDiff ([#812])
19+
- Make basis work for `CuArray` ([#810])
20+
1021
## [0.7.0]
1122

1223
### Changed
@@ -27,11 +38,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
2738

2839
- Allocate Enzyme shadow memory during preparation ([#782])
2940

30-
[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.0...main
41+
[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.1...main
42+
[0.7.1]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.0...DifferentiationInterface-v0.7.1
3143
[0.7.0]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.54...DifferentiationInterface-v0.7.0
3244
[0.6.54]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.53...DifferentiationInterface-v0.6.54
3345
[0.6.53]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.52...DifferentiationInterface-v0.6.53
3446

47+
[#812]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/812
48+
[#810]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/810
49+
[#809]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/809
3550
[#799]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/799
3651
[#795]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/795
3752
[#790]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/790

‎DifferentiationInterface/Project.toml‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "DifferentiationInterface"
22
uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
33
authors = ["Guillaume Dalle", "Adrian Hill"]
4-
version = "0.7.0"
4+
version = "0.7.1"
55

66
[deps]
77
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -41,7 +41,9 @@ DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
4141
DifferentiationInterfaceGPUArraysCoreExt = "GPUArraysCore"
4242
DifferentiationInterfaceGTPSAExt = "GTPSA"
4343
DifferentiationInterfaceMooncakeExt = "Mooncake"
44-
DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"]
44+
DifferentiationInterfacePolyesterForwardDiffExt = [
45+
"PolyesterForwardDiff", "ForwardDiff", "DiffResults"
46+
]
4547
DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
4648
DifferentiationInterfaceSparseArraysExt = "SparseArrays"
4749
DifferentiationInterfaceSparseConnectivityTracerExt = "SparseConnectivityTracer"
@@ -72,7 +74,7 @@ JET = "0.9"
7274
JLArrays = "0.2.0"
7375
JuliaFormatter = "1,2"
7476
LinearAlgebra = "1"
75-
Mooncake = "0.4.121"
77+
Mooncake = "0.4.122"
7678
Pkg = "1"
7779
PolyesterForwardDiff = "0.1.2"
7880
Random = "1"
@@ -121,4 +123,21 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
121123
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
122124

123125
[targets]
124-
test = ["ADTypes", "Aqua", "ComponentArrays", "DataFrames", "ExplicitImports", "JET", "JLArrays", "JuliaFormatter", "Pkg", "Random", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test"]
126+
test = [
127+
"ADTypes",
128+
"Aqua",
129+
"ComponentArrays",
130+
"DataFrames",
131+
"ExplicitImports",
132+
"JET",
133+
"JLArrays",
134+
"JuliaFormatter",
135+
"Pkg",
136+
"Random",
137+
"SparseArrays",
138+
"SparseConnectivityTracer",
139+
"SparseMatrixColorings",
140+
"StableRNGs",
141+
"StaticArrays",
142+
"Test",
143+
]

‎DifferentiationInterface/README.md‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ We support the following backends defined by [ADTypes.jl](https://github.com/Sci
3232

3333
- [ChainRulesCore.jl](https://github.com/JuliaDiff/ChainRulesCore.jl)
3434
- [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl) (currently broken)
35-
- [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl)
35+
- [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) (see below)
3636
- [FastDifferentiation.jl](https://github.com/brianguenter/FastDifferentiation.jl)
3737
- [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl)
3838
- [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl)
@@ -45,10 +45,8 @@ We support the following backends defined by [ADTypes.jl](https://github.com/Sci
4545
- [Tracker.jl](https://github.com/FluxML/Tracker.jl)
4646
- [Zygote.jl](https://github.com/FluxML/Zygote.jl)
4747

48-
Note that in some cases, going through DifferentiationInterface.jl might be slower than a direct call to the backend's API.
49-
This is mostly true for Enzyme.jl, whose handling of activities and multiple arguments unlocks additional performance.
50-
We are working on this challenge, and welcome any suggestions or contributions.
51-
Meanwhile, if differentiation fails or takes too long, consider using Enzyme.jl directly.
48+
> [!CAUTION]
49+
> Note that in some cases, going through DifferentiationInterface.jl might be slower or cause more errors than a direct call to the backend's API. This is especially true for Enzyme.jl, whose handling of activities and multiple arguments is not fully supported here. We are working on this challenge, and welcome any suggestions or contributions. Meanwhile, if differentiation fails or takes too long, consider using Enzyme.jl through its [native API](https://enzymead.github.io/Enzyme.jl/stable/) instead.
5250
5351
## Installation
5452

‎DifferentiationInterface/docs/make.jl‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@ links = InterLinks(
1616
"Symbolics" => "https://symbolics.juliasymbolics.org/stable/",
1717
)
1818

19-
cp(joinpath(@__DIR__, "..", "README.md"), joinpath(@__DIR__, "src", "index.md"); force=true)
19+
readme_str = read(joinpath(@__DIR__, "..", "README.md"), String)
20+
readme_str = replace(readme_str, "> [!CAUTION]\n> " => "!!! warning\n ")
21+
write(joinpath(@__DIR__, "src", "index.md"), readme_str)
2022

2123
makedocs(;
2224
modules=[DifferentiationInterface],

‎DifferentiationInterface/docs/src/explanation/backends.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,10 @@ We only implement `pushforward`.
120120
Depending on the `mode` attribute inside [`AutoEnzyme`](@extref ADTypes.AutoEnzyme), we implement either `pushforward` or `pullback` based on `Enzyme.autodiff`.
121121
When necessary, preparation chooses a number of chunks (for `gradient` and `jacobian` in forward mode, for `jacobian` only in reverse mode).
122122

123+
!!! warning
124+
Enzyme.jl's handling of activities and multiple arguments is not fully supported here, which can cause slowdowns or errors.
125+
If differentiation fails or takes too long, consider using Enzyme.jl through its [native API](https://enzymead.github.io/Enzyme.jl/stable/) instead.
126+
123127
### FastDifferentiation
124128

125129
For every operator, preparation generates an [executable function](https://brianguenter.github.io/FastDifferentiation.jl/stable/makefunction/) from the symbolic expression of the differentiated function.

‎DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ function DI.prepare_pushforward_nokwarg(
2727
absstep = if isnothing(backend.absstep)
2828
relstep
2929
else
30-
backend.relstep
30+
backend.absstep
3131
end
3232
dir = backend.dir
3333
return FiniteDiffOneArgPushforwardPrep(_sig, cache, relstep, absstep, dir)
@@ -144,7 +144,7 @@ function DI.prepare_derivative_nokwarg(
144144
absstep = if isnothing(backend.absstep)
145145
relstep
146146
else
147-
backend.relstep
147+
backend.absstep
148148
end
149149
dir = backend.dir
150150
return FiniteDiffOneArgDerivativePrep(_sig, cache, relstep, absstep, dir)
@@ -269,7 +269,7 @@ function DI.prepare_gradient_nokwarg(
269269
absstep = if isnothing(backend.absstep)
270270
relstep
271271
else
272-
backend.relstep
272+
backend.absstep
273273
end
274274
dir = backend.dir
275275
return FiniteDiffGradientPrep(_sig, cache, relstep, absstep, dir)
@@ -359,7 +359,7 @@ function DI.prepare_jacobian_nokwarg(
359359
absstep = if isnothing(backend.absstep)
360360
relstep
361361
else
362-
backend.relstep
362+
backend.absstep
363363
end
364364
dir = backend.dir
365365
return FiniteDiffOneArgJacobianPrep(_sig, cache, relstep, absstep, dir)
@@ -465,8 +465,16 @@ function DI.prepare_hessian_nokwarg(
465465
else
466466
backend.relstep
467467
end
468-
absstep_g = isnothing(backend.absstep) ? relstep_g : backend.absstep
469-
absstep_h = isnothing(backend.absstep) ? relstep_h : backend.absstep
468+
absstep_g = if isnothing(backend.absstep)
469+
relstep_g
470+
else
471+
backend.absstep
472+
end
473+
absstep_h = if isnothing(backend.absstep)
474+
relstep_h
475+
else
476+
backend.absstep
477+
end
470478
return FiniteDiffHessianPrep(
471479
_sig, gradient_cache, hessian_cache, relstep_g, absstep_g, relstep_h, absstep_h
472480
)

‎DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ function DI.prepare_pushforward_nokwarg(
3131
absstep = if isnothing(backend.absstep)
3232
relstep
3333
else
34-
backend.relstep
34+
backend.absstep
3535
end
3636
dir = backend.dir
3737
return FiniteDiffTwoArgPushforwardPrep(_sig, cache, relstep, absstep, dir)
@@ -175,7 +175,7 @@ function DI.prepare_derivative_nokwarg(
175175
absstep = if isnothing(backend.absstep)
176176
relstep
177177
else
178-
backend.relstep
178+
backend.absstep
179179
end
180180
dir = backend.dir
181181
return FiniteDiffTwoArgDerivativePrep(_sig, cache, relstep, absstep, dir)
@@ -295,7 +295,7 @@ function DI.prepare_jacobian_nokwarg(
295295
absstep = if isnothing(backend.absstep)
296296
relstep
297297
else
298-
backend.relstep
298+
backend.absstep
299299
end
300300
dir = backend.dir
301301
return FiniteDiffTwoArgJacobianPrep(_sig, cache, relstep, absstep, dir)

‎DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl‎

Lines changed: 12 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,42 +1,21 @@
11
module DifferentiationInterfaceGPUArraysCoreExt
22

33
import DifferentiationInterface as DI
4-
using GPUArraysCore: AbstractGPUArray
4+
using GPUArraysCore: @allowscalar, AbstractGPUArray
55

6-
"""
7-
OneElement
8-
9-
Efficient storage for a one-hot array, aka an array in the standard Euclidean basis.
10-
"""
11-
struct OneElement{I,N,T,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
12-
ind::I
13-
val::T
14-
a::A
15-
16-
function OneElement(ind::Integer, val::T, a::A) where {N,T,A<:AbstractArray{T,N}}
17-
right_ind = eachindex(a)[ind]
18-
return new{typeof(right_ind),N,T,A}(right_ind, val, a)
19-
end
20-
21-
function OneElement(
22-
ind::CartesianIndex{N}, val::T, a::A
23-
) where {N,T,A<:AbstractArray{T,N}}
24-
linear_ind = LinearIndices(a)[ind]
25-
right_ind = eachindex(a)[linear_ind]
26-
return new{typeof(right_ind),N,T,A}(right_ind, val, a)
27-
end
28-
end
29-
30-
Base.size(oe::OneElement) = size(oe.a)
31-
Base.IndexStyle(oe::OneElement) = Base.IndexStyle(oe.a)
32-
33-
function Base.getindex(oe::OneElement{<:Integer}, ind::Integer)
34-
return ifelse(ind == oe.ind, oe.val, zero(eltype(oe.a)))
6+
function DI.basis(a::AbstractGPUArray{T}, i) where {T}
7+
b = similar(a)
8+
fill!(b, zero(T))
9+
@allowscalar b[i] = one(T)
10+
return b
3511
end
3612

37-
function DI.basis(a::AbstractGPUArray{T}, i) where {T}
38-
b = zero(a)
39-
b .+= OneElement(i, one(T), a)
13+
function DI.multibasis(a::AbstractGPUArray{T}, inds) where {T}
14+
b = similar(a)
15+
fill!(b, zero(T))
16+
for i in inds
17+
@allowscalar b[i] = one(T)
18+
end
4019
return b
4120
end
4221

‎DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,12 @@ using Mooncake:
2121
zero_fcodual,
2222
MinimalCtx,
2323
NoRData,
24-
primal
24+
primal,
25+
_copy_output,
26+
_copy_to_output!!
2527

2628
DI.check_available(::AutoMooncake) = true
2729

28-
copyto!!(dst::Number, src::Number) = convert(typeof(dst), src)
29-
copyto!!(dst, src) = DI.ismutable_array(dst) ? copyto!(dst, src) : convert(typeof(dst), src)
30-
3130
get_config(::AutoMooncake{Nothing}) = Config()
3231
get_config(backend::AutoMooncake{<:Config}) = backend.config
3332

0 commit comments

Comments
 (0)