Skip to content

Commit 7e4b13b

Browse files
ocotsocots
authored
Merge pull request #499 from control-toolbox/feat/gpu-di-parameter
feat(differentiation): parameterize DifferentiationInterface{P<:Union{CPU,GPU}}
2 parents 268ac4d + 454e6b2 commit 7e4b13b

3 files changed

Lines changed: 223 additions & 13 deletions

File tree

‎src/Differentiation/default.jl‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,3 +16,25 @@ for differentiation unless specified otherwise.
1616
See also: [`CTBase.Differentiation.DifferentiationInterface`](@ref).
1717
"""
1818
__ad_backend() = ADTypes.AutoForwardDiff()
19+
20+
"""
21+
$(TYPEDSIGNATURES)
22+
23+
Default AD backend for CPU execution: `AutoForwardDiff()`.
24+
25+
See also: [`CTBase.Differentiation.DifferentiationInterface`](@ref), [`CTBase.Strategies.CPU`](@ref).
26+
"""
27+
__ad_backend(::Type{Strategies.CPU}) = ADTypes.AutoForwardDiff()
28+
29+
"""
30+
$(TYPEDSIGNATURES)
31+
32+
Default AD backend for GPU execution: `AutoZygote()`.
33+
34+
`AutoZygote` is a GPU-capable AD backend (measured correct on `CuArray` inputs). The marker
35+
type comes from `ADTypes` (a hard dependency) and needs no Zygote loaded to construct; Zygote
36+
is required only when a gradient is actually evaluated.
37+
38+
See also: [`CTBase.Differentiation.DifferentiationInterface`](@ref), [`CTBase.Strategies.GPU`](@ref).
39+
"""
40+
__ad_backend(::Type{Strategies.GPU}) = ADTypes.AutoZygote()

‎src/Differentiation/differentiation_interface.jl‎

Lines changed: 97 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,41 +11,76 @@ Concrete AD backend strategy wrapping DifferentiationInterface.jl backends.
1111
`AutoZygote()`) and uses it to compute gradients via the DifferentiationInterface.jl
1212
ecosystem.
1313
14+
Parameterized on the execution device `P`:
15+
- `DifferentiationInterface{CPU}`: default backend `AutoForwardDiff()` (default device);
16+
- `DifferentiationInterface{GPU}`: default backend `AutoZygote()` (GPU-capable AD).
17+
18+
`DifferentiationInterface(...)` builds a `DifferentiationInterface{CPU}` — the device
19+
parameterization is fully backward compatible with existing call sites.
20+
1421
# Arguments
15-
- `backend=AutoForwardDiff()`: The DifferentiationInterface.jl backend to use. Defaults
16-
to `AutoForwardDiff()` from `ADTypes.jl` (a hard dependency).
22+
- `backend`: The DifferentiationInterface.jl backend to use. Defaults to the device default
23+
(`AutoForwardDiff()` on CPU, `AutoZygote()` on GPU) from `ADTypes.jl` (a hard dependency).
1724
- `kwargs...`: Additional options passed to `StrategyOptions`.
1825
1926
# Notes
20-
- `ADTypes.jl` is a hard dependency, so `AutoForwardDiff()` is always available in core.
27+
- `ADTypes.jl` is a hard dependency, so the default backend markers are always available in core;
28+
`AutoZygote()` is a marker type and needs no Zygote loaded to construct.
2129
- Gradient computation requires the `CTBaseDifferentiationInterface` extension.
2230
- Without the extension, the gradient methods throw `NotImplemented` with a helpful message.
2331
2432
See also: [`CTBase.Differentiation.AbstractADBackend`](@ref),
2533
[`CTBase.Differentiation.hamiltonian_gradient`](@ref),
2634
[`CTBase.Differentiation.variable_gradient`](@ref).
2735
"""
28-
struct DifferentiationInterface{O<:Strategies.StrategyOptions} <: AbstractADBackend
36+
struct DifferentiationInterface{
37+
P<:Union{Strategies.CPU,Strategies.GPU},O<:Strategies.StrategyOptions
38+
} <: AbstractADBackend
2939
options::O
3040
end
3141

3242
"""
3343
$(TYPEDSIGNATURES)
3444
35-
Constructor for `DifferentiationInterface` with a specific backend.
45+
Construct a `DifferentiationInterface{CPU}` (the default device). Equivalent to
46+
`DifferentiationInterface{CPU}(...)`; delegates through
47+
[`CTBase.Strategies.default_parameter`](@ref).
3648
3749
# Arguments
3850
- `mode::Symbol=:strict`: Validation mode forwarded to `build_strategy_options`.
3951
- `kwargs...`: Options passed to `StrategyOptions`. The AD backend is set through the
4052
`:ad_backend` option (aliases `backend` and `ad`), e.g. `backend=AutoForwardDiff()`;
41-
it defaults to `AutoForwardDiff()`.
53+
it defaults to the device default (`AutoForwardDiff()` on CPU).
4254
4355
# Returns
44-
- `DifferentiationInterface`: A new backend strategy instance.
56+
- `DifferentiationInterface{CPU}`: A new backend strategy instance.
4557
"""
4658
function DifferentiationInterface(; mode::Symbol=:strict, kwargs...)
47-
opts = Strategies.build_strategy_options(DifferentiationInterface; mode=mode, kwargs...)
48-
return DifferentiationInterface{typeof(opts)}(opts)
59+
P = Strategies.default_parameter(DifferentiationInterface)
60+
return DifferentiationInterface{P}(; mode=mode, kwargs...)
61+
end
62+
63+
"""
64+
$(TYPEDSIGNATURES)
65+
66+
Construct a parameterized `DifferentiationInterface{P}` for the execution device `P`
67+
(`CPU` or `GPU`). The default `:ad_backend` is device-aware (`AutoForwardDiff()` on CPU,
68+
`AutoZygote()` on GPU) and can be overridden through the `:ad_backend` option.
69+
70+
# Arguments
71+
- `mode::Symbol=:strict`: Validation mode forwarded to `build_strategy_options`.
72+
- `kwargs...`: Options passed to `StrategyOptions` (`:ad_backend`, aliases `backend`/`ad`).
73+
74+
# Returns
75+
- `DifferentiationInterface{P}`: A new backend strategy instance.
76+
"""
77+
function DifferentiationInterface{P}(;
78+
mode::Symbol=:strict, kwargs...
79+
) where {P<:Strategies.AbstractStrategyParameter}
80+
opts = Strategies.build_strategy_options(
81+
DifferentiationInterface{P}; mode=mode, kwargs...
82+
)
83+
return DifferentiationInterface{P,typeof(opts)}(opts)
4984
end
5085

5186
# ==============================================================================
@@ -62,6 +97,35 @@ Strategies.id(::Type{<:DifferentiationInterface}) = :di
6297
"""
6398
$(TYPEDSIGNATURES)
6499
100+
Return the execution parameter type of a `DifferentiationInterface` strategy.
101+
102+
Extracts `P` from `DifferentiationInterface{P}` (`CPU` or `GPU`). Overrides the
103+
`AbstractADBackend` family default (`nothing`) since this strategy is device-parameterized.
104+
105+
# Returns
106+
- `Type{<:Union{CPU,GPU}}`: the execution parameter type.
107+
108+
See also: [`CTBase.Strategies.CPU`](@ref), [`CTBase.Strategies.GPU`](@ref)
109+
"""
110+
Strategies.parameter(
111+
::Type{<:DifferentiationInterface{P}}
112+
) where {P<:Union{Strategies.CPU,Strategies.GPU}} = P
113+
114+
"""
115+
$(TYPEDSIGNATURES)
116+
117+
Return the default execution parameter for `DifferentiationInterface` when none is specified.
118+
119+
Returns `CPU`, so `DifferentiationInterface(...)` builds a `DifferentiationInterface{CPU}` and
120+
every existing call site is unaffected by the device parameterization.
121+
122+
See also: [`CTBase.Strategies.CPU`](@ref)
123+
"""
124+
Strategies.default_parameter(::Type{<:DifferentiationInterface}) = Strategies.CPU
125+
126+
"""
127+
$(TYPEDSIGNATURES)
128+
65129
Return a human-readable description of the `DifferentiationInterface` strategy.
66130
"""
67131
function Strategies.description(::Type{<:DifferentiationInterface})
@@ -71,15 +135,22 @@ end
71135
"""
72136
$(TYPEDSIGNATURES)
73137
74-
Return metadata defining `DifferentiationInterface` options and their specifications.
138+
Return metadata defining `DifferentiationInterface{P}` options and their specifications.
139+
140+
The `:ad_backend` default is device-aware: `AutoForwardDiff()` on `CPU`, `AutoZygote()` on `GPU`
141+
(via [`CTBase.Differentiation.__ad_backend`](@ref)). The bare `metadata(DifferentiationInterface)`
142+
delegates here through `DifferentiationInterface{CPU}`.
75143
"""
76-
function Strategies.metadata(::Type{<:DifferentiationInterface})
144+
function Strategies.metadata(
145+
::Type{<:DifferentiationInterface{P}}
146+
) where {P<:Union{Strategies.CPU,Strategies.GPU}}
77147
return Strategies.StrategyMetadata(
78148
Strategies.OptionDefinition(;
79149
name=:ad_backend,
80150
type=ADTypes.AbstractADType,
81-
default=__ad_backend(),
82-
description="DifferentiationInterface.jl backend (e.g. AutoForwardDiff()).",
151+
default=__ad_backend(P),
152+
computed=true, # Default is computed from parameter P (CPU→ForwardDiff, GPU→Zygote)
153+
description="DifferentiationInterface.jl backend (e.g. AutoForwardDiff() on CPU, AutoZygote() on GPU).",
83154
aliases=(:backend, :ad),
84155
),
85156
)
@@ -88,6 +159,19 @@ end
88159
"""
89160
$(TYPEDSIGNATURES)
90161
162+
Fallback for the non-parameterized `DifferentiationInterface` type that delegates to
163+
`DifferentiationInterface{CPU}`. Preserves backward compatibility for
164+
`metadata(DifferentiationInterface)`.
165+
"""
166+
function Strategies.metadata(::Type{DifferentiationInterface})
167+
return Strategies.metadata(
168+
DifferentiationInterface{Strategies.default_parameter(DifferentiationInterface)}
169+
)
170+
end
171+
172+
"""
173+
$(TYPEDSIGNATURES)
174+
91175
Extract the AD backend from a `DifferentiationInterface` strategy.
92176
93177
# Arguments

‎test/suite/differentiation/test_di_parameter.jl‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
"""
2+
Contract tests for the device-parameterized `DifferentiationInterface{P}` strategy.
3+
"""
4+
5+
module TestDIParameter
6+
7+
using Test: Test
8+
import CTBase.Differentiation
9+
import CTBase.Strategies
10+
using ADTypes: ADTypes
11+
12+
const VERBOSE = isdefined(Main, :TestData) ? Main.TestData.VERBOSE : true
13+
const SHOWTIMING = isdefined(Main, :TestData) ? Main.TestData.SHOWTIMING : true
14+
15+
function test_di_parameter()
16+
DI = Differentiation.DifferentiationInterface
17+
Test.@testset "DifferentiationInterface{P} parameterization" verbose=VERBOSE showtiming=SHOWTIMING begin
18+
19+
# ====================================================================
20+
# CONTRACT - parameter / default_parameter / id
21+
# ====================================================================
22+
23+
Test.@testset "parameter contract" begin
24+
# Bare (device-unspecified) DI resolves to `nothing` via the AbstractADBackend
25+
# family default — required by consumers that register DI without a parameter
26+
# (e.g. CTFlows' flow registry) and query `parameter` on the bare type.
27+
Test.@test Strategies.parameter(DI) === nothing
28+
29+
Test.@test Strategies.parameter(DI{Strategies.CPU}) == Strategies.CPU
30+
Test.@test Strategies.parameter(DI{Strategies.GPU}) == Strategies.GPU
31+
Test.@test Strategies.default_parameter(DI) == Strategies.CPU
32+
33+
Test.@test Strategies.id(DI) === :di
34+
Test.@test Strategies.id(DI{Strategies.CPU}) === :di
35+
Test.@test Strategies.id(DI{Strategies.GPU}) === :di
36+
end
37+
38+
# ====================================================================
39+
# CONTRACT - per-parameter metadata + back-compat bare delegation
40+
# ====================================================================
41+
42+
Test.@testset "per-parameter metadata defaults" begin
43+
md_bare = Strategies.metadata(DI)
44+
md_cpu = Strategies.metadata(DI{Strategies.CPU})
45+
md_gpu = Strategies.metadata(DI{Strategies.GPU})
46+
47+
Test.@test md_cpu[:ad_backend].default === ADTypes.AutoForwardDiff()
48+
Test.@test md_gpu[:ad_backend].default === ADTypes.AutoZygote()
49+
50+
# The default is computed from P, so it is flagged `computed` (for `describe`/provenance).
51+
Test.@test md_cpu[:ad_backend].computed == true
52+
Test.@test md_gpu[:ad_backend].computed == true
53+
54+
# Bare `metadata(DI)` delegates to `metadata(DI{CPU})` (back-compat).
55+
Test.@test md_bare[:ad_backend].default === ADTypes.AutoForwardDiff()
56+
end
57+
58+
# ====================================================================
59+
# CONTRACT - construction per parameter (DI() ≡ DI{CPU})
60+
# ====================================================================
61+
62+
Test.@testset "construction per parameter" begin
63+
di_default = DI()
64+
Test.@test di_default isa DI{Strategies.CPU}
65+
Test.@test Strategies.parameter(typeof(di_default)) == Strategies.CPU
66+
Test.@test Differentiation.ad_backend(di_default) === ADTypes.AutoForwardDiff()
67+
68+
di_cpu = DI{Strategies.CPU}()
69+
Test.@test di_cpu isa DI{Strategies.CPU}
70+
71+
di_gpu = DI{Strategies.GPU}()
72+
Test.@test di_gpu isa DI{Strategies.GPU}
73+
Test.@test Strategies.parameter(typeof(di_gpu)) == Strategies.GPU
74+
# GPU default backend is AutoZygote (resolved from the per-device metadata).
75+
Test.@test Differentiation.ad_backend(di_gpu) === ADTypes.AutoZygote()
76+
77+
# Explicit override still wins on either device.
78+
di_gpu_fd = DI{Strategies.GPU}(; backend=ADTypes.AutoForwardDiff())
79+
Test.@test Differentiation.ad_backend(di_gpu_fd) === ADTypes.AutoForwardDiff()
80+
end
81+
82+
# ====================================================================
83+
# INTEGRATION - registry with [CPU, GPU] + global parameter extraction
84+
# ====================================================================
85+
86+
Test.@testset "registry with device parameters" begin
87+
r = Strategies.create_registry(
88+
Differentiation.AbstractADBackend => ((DI, [Strategies.CPU, Strategies.GPU]),),
89+
)
90+
91+
Test.@test :di in Strategies.strategy_ids(Differentiation.AbstractADBackend, r)
92+
93+
Test.@test Strategies.extract_global_parameter_from_method((:di, :cpu), r) ==
94+
Strategies.CPU
95+
Test.@test Strategies.extract_global_parameter_from_method((:di, :gpu), r) ==
96+
Strategies.GPU
97+
end
98+
end
99+
end
100+
101+
end # module
102+
103+
# CRITICAL: Redefine in outer scope for TestRunner
104+
test_di_parameter() = TestDIParameter.test_di_parameter()

0 commit comments

Comments
 (0)