Skip to content

Commit 6d1a205

Browse files
ocotsocots
committed
feat: Complete GPU/CPU parameter system with comprehensive tests
- Add full GPU/CPU parameter support to methods() returning 4-tuples - Implement complete strategy builders with ResolvedMethod support - Enhance registry with parameter-aware strategy mapping - Add comprehensive test coverage (422 tests total): * Component Checks: 26 tests * Component Completion: 29 tests * Descriptive Routing: 75 tests * Kwarg Extraction: 59 tests * Methods: 39 tests * Print: 46 tests * Registry: 65 tests * Strategy Builders: 83 tests - Fix all test failures and ensure 100% pass rate - Add proper dependency handling for strategy building - Support both provided and build paths for strategy construction
1 parent a80e51e commit 6d1a205

16 files changed

Lines changed: 1292 additions & 128 deletions

‎Project.toml‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,24 +21,24 @@ SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843"
2121

2222
[compat]
2323
ADNLPModels = "0.8"
24+
BenchmarkTools = "1"
2425
CTBase = "0.18"
2526
CTDirect = "1"
2627
CTFlows = "0.8"
2728
CTModels = "0.9"
2829
CTParser = "0.8"
29-
CTSolvers = "0.3"
30+
CTSolvers = "0.4"
3031
CUDA = "5"
3132
CommonSolve = "0.2"
3233
DifferentiationInterface = "0.7"
3334
DocStringExtensions = "0.9"
3435
ExaModels = "0.9"
3536
ForwardDiff = "0.10, 1.0"
3637
LinearAlgebra = "1"
37-
MadNCL = "0.1"
38-
MadNLP = "0.8"
39-
MadNLPGPU = "0.7"
40-
MadNLPMumps = "0.5"
41-
NLPModels = "0.21.7"
38+
MadNCL = "0.2"
39+
MadNLP = "0.9"
40+
MadNLPGPU = "0.8"
41+
NLPModels = "0.21"
4242
NLPModelsIpopt = "0.11"
4343
NonlinearSolve = "4"
4444
OrdinaryDiffEq = "6"
@@ -51,14 +51,14 @@ Test = "1"
5151
julia = "1.10"
5252

5353
[extras]
54+
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
5455
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
5556
DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
5657
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
5758
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
5859
MadNCL = "434a0bcb-5a7c-42b2-a9d3-9e3f760e7af0"
5960
MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6"
6061
MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598"
61-
MadNLPMumps = "3b83494e-c0a4-4895-918b-9157a7a085a1"
6262
NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71"
6363
NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
6464
OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
@@ -67,4 +67,4 @@ SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
6767
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
6868

6969
[targets]
70-
test = ["CUDA", "DifferentiationInterface", "ForwardDiff", "LinearAlgebra", "MadNCL", "MadNLP", "MadNLPGPU", "MadNLPMumps", "NLPModelsIpopt", "NonlinearSolve", "OrdinaryDiffEq", "Printf", "SplitApplyCombine", "Test"]
70+
test = ["BenchmarkTools", "CUDA", "DifferentiationInterface", "ForwardDiff", "LinearAlgebra", "MadNCL", "MadNLP", "MadNLPGPU", "NLPModelsIpopt", "NonlinearSolve", "OrdinaryDiffEq", "Printf", "SplitApplyCombine", "Test"]

‎src/helpers/component_completion.jl‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,15 +52,19 @@ function _complete_components(
5252
# Step 2: Complete the method description
5353
complete_description = _complete_description(partial_description)
5454

55+
# Step 2.5: Resolve method with parameter information
56+
families = _descriptive_families()
57+
resolved = CTSolvers.resolve_method(complete_description, families, registry)
58+
5559
# Step 3: Build or use strategies for each family
5660
final_discretizer = _build_or_use_strategy(
57-
complete_description, discretizer, CTDirect.AbstractDiscretizer, registry
61+
resolved, discretizer, :discretizer, families, registry
5862
)
5963
final_modeler = _build_or_use_strategy(
60-
complete_description, modeler, CTSolvers.AbstractNLPModeler, registry
64+
resolved, modeler, :modeler, families, registry
6165
)
6266
final_solver = _build_or_use_strategy(
63-
complete_description, solver, CTSolvers.AbstractNLPSolver, registry
67+
resolved, solver, :solver, families, registry
6468
)
6569

6670
return (discretizer=final_discretizer, modeler=final_modeler, solver=final_solver)

‎src/helpers/descriptive_routing.jl‎

Lines changed: 15 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ See also: [`_descriptive_families`](@ref), [`_descriptive_action_defs`](@ref),
166166
[`_build_components_from_routed`](@ref)
167167
"""
168168
function _route_descriptive_options(
169-
complete_description::Tuple{Symbol, Symbol, Symbol},
169+
complete_description::Tuple{Symbol, Symbol, Symbol, Symbol},
170170
registry::CTSolvers.StrategyRegistry,
171171
kwargs,
172172
)
@@ -192,7 +192,7 @@ $(TYPEDSIGNATURES)
192192
Build concrete strategy instances and extract action options from a routed options result.
193193
194194
Each strategy is constructed via
195-
[`CTSolvers.build_strategy_from_method`](@ref) using the options
195+
[`CTSolvers.build_strategy_from_resolved`](@ref) using the options
196196
that were routed to its family by [`_route_descriptive_options`](@ref).
197197
198198
Action options (`initial_guess`, `display`) are extracted from `routed.action`
@@ -220,31 +220,27 @@ true
220220
```
221221
222222
See also: [`_route_descriptive_options`](@ref),
223-
[`CTSolvers.build_strategy_from_method`](@ref)
223+
[`CTSolvers.build_strategy_from_resolved`](@ref)
224224
"""
225225
function _build_components_from_routed(
226226
ocp::CTModels.AbstractModel,
227-
complete_description::Tuple{Symbol, Symbol, Symbol},
227+
complete_description::Tuple{Symbol, Symbol, Symbol, Symbol},
228228
registry::CTSolvers.StrategyRegistry,
229229
routed::NamedTuple,
230230
)
231-
discretizer = CTSolvers.build_strategy_from_method(
232-
complete_description,
233-
CTDirect.AbstractDiscretizer,
234-
registry;
235-
routed.strategies.discretizer...,
231+
# Resolve method with parameter information as early as possible
232+
families = _descriptive_families()
233+
resolved = CTSolvers.resolve_method(complete_description, families, registry)
234+
235+
# Build strategies using resolved method
236+
discretizer = CTSolvers.build_strategy_from_resolved(
237+
resolved, :discretizer, families, registry; routed.strategies.discretizer...
236238
)
237-
modeler = CTSolvers.build_strategy_from_method(
238-
complete_description,
239-
CTSolvers.AbstractNLPModeler,
240-
registry;
241-
routed.strategies.modeler...,
239+
modeler = CTSolvers.build_strategy_from_resolved(
240+
resolved, :modeler, families, registry; routed.strategies.modeler...
242241
)
243-
solver = CTSolvers.build_strategy_from_method(
244-
complete_description,
245-
CTSolvers.AbstractNLPSolver,
246-
registry;
247-
routed.strategies.solver...,
242+
solver = CTSolvers.build_strategy_from_resolved(
243+
resolved, :solver, families, registry; routed.strategies.solver...
248244
)
249245

250246
# Extract and unwrap action options (OptionValue → raw value)

‎src/helpers/methods.jl‎

Lines changed: 40 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,63 @@
11
"""
22
$(TYPEDSIGNATURES)
33
4-
Return the tuple of available method triplets for solving optimal control problems.
4+
Return the tuple of available method quadruplets for solving optimal control problems.
55
6-
Each triplet consists of `(discretizer_id, modeler_id, solver_id)` where:
6+
Each quadruplet consists of `(discretizer_id, modeler_id, solver_id, parameter)` where:
77
- `discretizer_id`: Symbol identifying the discretization strategy
8-
- `modeler_id`: Symbol identifying the NLP modeling strategy
8+
- `modeler_id`: Symbol identifying the NLP modeling strategy
99
- `solver_id`: Symbol identifying the NLP solver
10+
- `parameter`: Symbol identifying the parameter (`:cpu` or `:gpu`)
11+
12+
GPU-capable methods use parameterized strategies that automatically get appropriate defaults:
13+
- `Exa{GPU}` gets `CUDA.CUDABackend()` by default
14+
- `MadNLP{GPU}` gets `MadNLPGPU.CUDSSSolver` by default
15+
- `MadNCL{GPU}` gets `MadNLPGPU.CUDSSSolver` by default
1016
1117
# Returns
12-
- `Tuple{Vararg{Tuple{Symbol, Symbol, Symbol}}}`: Available method combinations
18+
- `Tuple{Vararg{Tuple{Symbol, Symbol, Symbol, Symbol}}}`: Available method combinations
1319
1420
# Examples
1521
```julia
1622
julia> m = methods()
17-
((:collocation, :adnlp, :ipopt), (:collocation, :adnlp, :madnlp), ...)
23+
((:collocation, :adnlp, :ipopt, :cpu), (:collocation, :adnlp, :madnlp, :cpu), ...)
1824
1925
julia> length(m)
20-
8
26+
10 # CPU methods + GPU methods
27+
28+
julia> # CPU methods (existing behavior maintained)
29+
julia> methods()[1]
30+
(:collocation, :adnlp, :ipopt, :cpu)
31+
32+
julia> # GPU methods (new functionality)
33+
julia> methods()[9] # First GPU method
34+
(:collocation, :exa, :madnlp, :gpu)
2135
```
2236
37+
# Notes
38+
- All existing methods are now explicitly marked with `:cpu` parameter
39+
- GPU methods are available when CUDA.jl is loaded
40+
- Parameterized strategies provide smart defaults automatically
41+
2342
# See Also
2443
- [`solve`](@ref): Main solve function that uses these methods
2544
- [`CTBase.complete`](@ref): Completes partial method descriptions
45+
- [`get_strategy_registry`](@ref): Registry with parameterized strategies
2646
"""
27-
function Base.methods()::Tuple{Vararg{Tuple{Symbol, Symbol, Symbol}}}
47+
function Base.methods()::Tuple{Vararg{Tuple{Symbol, Symbol, Symbol, Symbol}}}
2848
return (
29-
(:collocation, :adnlp, :ipopt ),
30-
(:collocation, :adnlp, :madnlp),
31-
(:collocation, :exa, :ipopt ),
32-
(:collocation, :exa, :madnlp),
33-
(:collocation, :adnlp, :madncl),
34-
(:collocation, :exa, :madncl),
35-
(:collocation, :adnlp, :knitro),
36-
(:collocation, :exa, :knitro),
49+
# CPU methods (all existing methods now with :cpu parameter)
50+
(:collocation, :adnlp, :ipopt, :cpu),
51+
(:collocation, :adnlp, :madnlp, :cpu),
52+
(:collocation, :exa, :ipopt, :cpu),
53+
(:collocation, :exa, :madnlp, :cpu),
54+
(:collocation, :adnlp, :madncl, :cpu),
55+
(:collocation, :exa, :madncl, :cpu),
56+
(:collocation, :adnlp, :knitro, :cpu),
57+
(:collocation, :exa, :knitro, :cpu),
58+
59+
# GPU methods (only combinations that make sense)
60+
(:collocation, :exa, :madnlp, :gpu),
61+
(:collocation, :exa, :madncl, :gpu),
3762
)
3863
end

‎src/helpers/registry.jl‎

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,43 @@ $(TYPEDSIGNATURES)
33
44
Create and return the strategy registry for the solve system.
55
6-
The registry maps abstract strategy families to their concrete implementations:
6+
The registry maps abstract strategy families to their concrete implementations
7+
with their supported parameters:
78
- `CTDirect.AbstractDiscretizer` → Discretization strategies
8-
- `CTSolvers.AbstractNLPModeler` → NLP modeling strategies
9-
- `CTSolvers.AbstractNLPSolver` → NLP solver strategies
9+
- `CTSolvers.AbstractNLPModeler` → NLP modeling strategies (with CPU/GPU support)
10+
- `CTSolvers.AbstractNLPSolver` → NLP solver strategies (with CPU/GPU support)
11+
12+
Each strategy entry specifies which parameters it supports:
13+
- `CPU`: All strategies support CPU execution
14+
- `GPU`: Only GPU-capable strategies support GPU execution (Exa, MadNLP, MadNCL)
1015
1116
# Returns
12-
- `CTSolvers.StrategyRegistry`: Registry with all available strategies
17+
- `CTSolvers.StrategyRegistry`: Registry with all available strategies and their parameters
1318
1419
# Examples
1520
```julia
1621
julia> registry = OptimalControl.get_strategy_registry()
1722
StrategyRegistry with 3 families
23+
24+
julia> CTSolvers.strategy_ids(CTSolvers.AbstractNLPModeler, registry)
25+
(:adnlp, :exa)
26+
27+
julia> CTSolvers.strategy_ids(CTSolvers.AbstractNLPSolver, registry)
28+
(:ipopt, :madnlp, :madncl, :knitro)
29+
30+
julia> # Check which parameters a strategy supports
31+
julia> CTSolvers.available_parameters(:modeler, CTSolvers.Exa, registry)
32+
(CPU, GPU)
33+
34+
julia> CTSolvers.available_parameters(:solver, CTSolvers.Ipopt, registry)
35+
(CPU,)
1836
```
37+
38+
# Notes
39+
- GPU-capable strategies (Exa, MadNLP, MadNCL) support both CPU and GPU parameters
40+
- CPU-only strategies (ADNLP, Ipopt, Knitro) support only CPU parameter
41+
- Parameterization is handled at the method level in `methods()`
42+
- GPU strategies automatically get appropriate default configurations when parameterized
1943
"""
2044
function get_strategy_registry()::CTSolvers.StrategyRegistry
2145
return CTSolvers.create_registry(
@@ -24,14 +48,14 @@ function get_strategy_registry()::CTSolvers.StrategyRegistry
2448
# Add other discretizers as they become available
2549
),
2650
CTSolvers.AbstractNLPModeler => (
27-
CTSolvers.ADNLP,
28-
CTSolvers.Exa,
51+
(CTSolvers.ADNLP, [CTSolvers.CPU]),
52+
(CTSolvers.Exa, [CTSolvers.CPU, CTSolvers.GPU])
2953
),
3054
CTSolvers.AbstractNLPSolver => (
31-
CTSolvers.Ipopt,
32-
CTSolvers.MadNLP,
33-
CTSolvers.MadNCL,
34-
CTSolvers.Knitro,
35-
)
55+
(CTSolvers.Ipopt, [CTSolvers.CPU]),
56+
(CTSolvers.MadNLP, [CTSolvers.CPU, CTSolvers.GPU]),
57+
(CTSolvers.MadNCL, [CTSolvers.CPU, CTSolvers.GPU]),
58+
(CTSolvers.Knitro, [CTSolvers.CPU]),
59+
),
3660
)
3761
end

0 commit comments

Comments
 (0)