From 234ae9df101011e40e2d0468af3431ff6df7e566 Mon Sep 17 00:00:00 2001 From: Olivier Cots Date: Sat, 18 Jul 2026 23:44:40 +0200 Subject: [PATCH 1/2] feat(integrators): parameterize SciML{P<:Union{CPU,GPU}} for GPU support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the SciML ODE integrator a device-parameterized strategy along the Exa/MadNLP pattern (GPU roadmap phase 1 / D1). SciML() ≡ SciML{CPU}, so every existing call site is unaffected. Adds per-parameter metadata seam, the __consistent_initial_condition device validator (CUDA extension), and Family-A parameter tests. Co-Authored-By: Claude Opus 4.8 --- ext/CTSolversCUDA.jl | 43 +++++- ext/CTSolversSciMLIntegrator.jl | 19 ++- src/Integrators/sciml.jl | 104 +++++++++++-- .../integrators/test_integrator_stubs.jl | 5 +- .../suite/integrators/test_sciml_parameter.jl | 140 ++++++++++++++++++ 5 files changed, 290 insertions(+), 21 deletions(-) create mode 100644 test/suite/integrators/test_sciml_parameter.jl diff --git a/ext/CTSolversCUDA.jl b/ext/CTSolversCUDA.jl index 9c0630d6..8190c9a3 100644 --- a/ext/CTSolversCUDA.jl +++ b/ext/CTSolversCUDA.jl @@ -1,13 +1,15 @@ """ CTSolversCUDA Extension -Extension providing CUDA backend functionality for Exa modeler. -Implements GPU-specific backend defaults and consistency validation. +Extension providing CUDA functionality for CTSolvers: +- `Modelers.Exa` GPU backend defaults and backend consistency validation; +- `Integrators.SciML{P}` initial-condition/device consistency validation. """ module CTSolversCUDA import CTSolvers.Modelers +import CTSolvers.Integrators import CTBase.Strategies using CUDA: CUDA using DocStringExtensions: DocStringExtensions @@ -91,4 +93,41 @@ function Modelers.__consistent_backend(::Type{Strategies.GPU}, backend::CUDA.CUD return true end +# ============================================================================ +# SciML integrator — initial-condition / device consistency +# ============================================================================ + +""" +$(DocStringExtensions.TYPEDSIGNATURES) + +A device `CuArray` initial condition is inconsistent with a `SciML{CPU}` integrator. + +Overrides the `true`-returning stub in `CTSolvers.Integrators`. Other host-array cases fall +through to the default (returns `true`). +""" +function Integrators.__consistent_initial_condition(::Type{Strategies.CPU}, u0::CUDA.AnyCuArray) + return false +end + +""" +$(DocStringExtensions.TYPEDSIGNATURES) + +A host `Array` initial condition is inconsistent with a `SciML{GPU}` integrator (the state +must live on the device). + +Overrides the `true`-returning stub in `CTSolvers.Integrators`. +""" +function Integrators.__consistent_initial_condition(::Type{Strategies.GPU}, u0::Array) + return false +end + +""" +$(DocStringExtensions.TYPEDSIGNATURES) + +A device `CuArray` initial condition is consistent with a `SciML{GPU}` integrator. +""" +function Integrators.__consistent_initial_condition(::Type{Strategies.GPU}, u0::CUDA.AnyCuArray) + return true +end + end # module CTSolversCUDA diff --git a/ext/CTSolversSciMLIntegrator.jl b/ext/CTSolversSciMLIntegrator.jl index 4d512277..afa06c97 100644 --- a/ext/CTSolversSciMLIntegrator.jl +++ b/ext/CTSolversSciMLIntegrator.jl @@ -57,12 +57,18 @@ end """ $(TYPEDSIGNATURES) -Return metadata defining `Integrators.SciML` options and their specifications. +Return metadata defining `Integrators.SciML{P}` options and their specifications. The `internalnorm` option defaults to `real_norm`, which extracts the primal (Float64) part of ForwardDiff dual numbers to ensure grid invariance (IND) when ForwardDiff is loaded. + +The metadata is specialized on the execution device `P` (`CPU`/`GPU`); the option set is +currently identical for both — `P` marks the seam where GPU-specific defaults/validators land +as they are discovered. The bare `metadata(SciML)` (core) delegates here through `SciML{CPU}`. """ -function Strategies.metadata(::Type{Integrators.SciML}) +function Strategies.metadata( + ::Type{Integrators.SciML{P}} +) where {P<:Union{Strategies.CPU,Strategies.GPU}} return Strategies.StrategyMetadata( Strategies.OptionDefinition(; name=:alg, @@ -297,6 +303,7 @@ cached dictionaries `options_point` (`:auto` → `false`) and `options_trajector # Arguments - `::Type{Integrators.SciMLTag}`: The SciML integrator tag type. +- `::Type{P}`: The execution device parameter (`CPU`/`GPU`); threaded into the built `SciML{P}`. - `mode::Symbol`: Validation mode for strategy options (`:strict` or `:permissive`). - `kwargs...`: User-provided option values. Explicit `true`/`false` override `:auto` resolution. @@ -310,9 +317,9 @@ cached dictionaries `options_point` (`:auto` → `false`) and `options_trajector See also: [`CTSolvers.Integrators.SciML`](@ref). """ function Integrators.build_sciml_integrator( - ::Type{Integrators.SciMLTag}; mode::Symbol=:strict, kwargs... -) - opts = Strategies.build_strategy_options(Integrators.SciML; mode=mode, kwargs...) + ::Type{Integrators.SciMLTag}, ::Type{P}; mode::Symbol=:strict, kwargs... +) where {P<:Strategies.AbstractStrategyParameter} + opts = Strategies.build_strategy_options(Integrators.SciML{P}; mode=mode, kwargs...) raw = Strategies.options_dict(opts) # Check if algorithm is missing and raise PreconditionError @@ -342,7 +349,7 @@ function Integrators.build_sciml_integrator( get(options_trajectory, key, :auto) === :auto && (options_trajectory[key] = true) end - return Integrators.SciML{typeof(opts),typeof(options_point),typeof(options_trajectory)}( + return Integrators.SciML{P,typeof(opts),typeof(options_point),typeof(options_trajectory)}( opts, options_point, options_trajectory ) end diff --git a/src/Integrators/sciml.jl b/src/Integrators/sciml.jl index 8c602514..4f3ee7e6 100644 --- a/src/Integrators/sciml.jl +++ b/src/Integrators/sciml.jl @@ -44,6 +44,13 @@ and solve) is provided by the `CTSolversSciMLIntegrator` package extension; this file declares the type and **stubs** that throw `ExtensionError` until the extension is loaded. +Parameterized on the execution device `P`: +- `SciML{CPU}`: CPU execution (default); +- `SciML{GPU}`: GPU execution (state on device arrays, e.g. `CuArray`). + +`SciML(...)` builds a `SciML{CPU}` — the device parameterization is fully backward +compatible with existing call sites. + To activate the extension, load any of: - `using OrdinaryDiffEqTsit5` (minimal) - `using OrdinaryDiffEq` @@ -53,7 +60,7 @@ To activate the extension, load any of: $(TYPEDFIELDS) """ -struct SciML{O<:Strategies.StrategyOptions,OP<:Dict{Symbol,Any},OT<:Dict{Symbol,Any}} <: +struct SciML{P<:Union{CPU,GPU},O<:Strategies.StrategyOptions,OP<:Dict{Symbol,Any},OT<:Dict{Symbol,Any}} <: AbstractSciMLIntegrator "Validated option bundle." options::O @@ -77,17 +84,29 @@ Strategies.id(::Type{<:SciML}) = :sciml """ $(TYPEDSIGNATURES) -Return the execution parameter type of the `SciML` integrator. +Return the execution parameter type of a `SciML` integrator. -Returns `nothing` because `SciML` is a non-parameterized strategy: it has no -execution parameter and does not require one. +Extracts the type parameter `P` from `SciML{P}`, which can be either `CPU` or `GPU` +since `SciML` supports both execution devices. # Returns -- `Nothing`: `SciML` is not parameterized. +- `Type{<:Union{CPU,GPU}}`: the execution parameter type. + +See also: [`CTSolvers.Integrators.SciML`](@ref), [`CTBase.Strategies.CPU`](@extref), [`CTBase.Strategies.GPU`](@extref) +""" +Strategies.parameter(::Type{<:SciML{P}}) where {P<:Union{CPU,GPU}} = P + +""" +$(TYPEDSIGNATURES) + +Return the default execution parameter for `SciML` when none is specified. + +Returns `CPU`, so `SciML(...)` builds a `SciML{CPU}` and every existing call site is +unaffected by the device parameterization. -See also: [`CTSolvers.Integrators.SciML`](@ref) +See also: [`CTSolvers.Integrators.SciML`](@ref), [`CTBase.Strategies.CPU`](@extref) """ -Strategies.parameter(::Type{<:SciML}) = nothing +Strategies.default_parameter(::Type{<:SciML}) = CPU """ $(TYPEDSIGNATURES) @@ -129,8 +148,8 @@ options_trajectory(integ::SciML) = integ.options_trajectory """ $(TYPEDSIGNATURES) -Construct a `SciML` integrator. Delegates to `build_sciml_integrator`, which -is overridden by the `CTSolversSciMLIntegrator` package extension. +Construct a `SciML{CPU}` integrator (the default device). Equivalent to +`SciML{CPU}(...)`; delegates through [`CTBase.Strategies.default_parameter`](@extref). # Arguments - `mode::Symbol=:strict`: Validation mode (`:strict` or `:permissive`). @@ -142,7 +161,28 @@ is overridden by the `CTSolversSciMLIntegrator` package extension. See also: [`CTSolvers.Integrators.SciML`](@ref), [`CTSolvers.Integrators.build_sciml_integrator`](@ref). """ function SciML(; mode::Symbol=:strict, kwargs...) - return build_sciml_integrator(SciMLTag; mode=mode, kwargs...) + P = Strategies.default_parameter(SciML) + return SciML{P}(; mode=mode, kwargs...) +end + +""" +$(TYPEDSIGNATURES) + +Construct a parameterized `SciML{P}` integrator for the execution device `P` +(`CPU` or `GPU`). Delegates to `build_sciml_integrator`, which is overridden by the +`CTSolversSciMLIntegrator` package extension. + +# Arguments +- `mode::Symbol=:strict`: Validation mode (`:strict` or `:permissive`). +- `kwargs...`: Options forwarded to the integrator builder (see extension documentation). + +# Throws +- `CTBase.Exceptions.ExtensionError`: If the `CTSolversSciMLIntegrator` extension is not loaded. + +See also: [`CTSolvers.Integrators.SciML`](@ref), [`CTSolvers.Integrators.build_sciml_integrator`](@ref). +""" +function SciML{P}(; mode::Symbol=:strict, kwargs...) where {P<:AbstractStrategyParameter} + return build_sciml_integrator(SciMLTag, P; mode=mode, kwargs...) end """ @@ -152,7 +192,9 @@ Stub builder for `SciML`. The real implementation is provided by `CTSolversSciMLIntegrator`; this stub throws `ExtensionError` until the extension is loaded. """ -function build_sciml_integrator(::Type{<:Core.AbstractTag}; kwargs...) +function build_sciml_integrator( + ::Type{<:Core.AbstractTag}, ::Type{<:AbstractStrategyParameter}; kwargs... +) return throw( Exceptions.ExtensionError( :OrdinaryDiffEqTsit5; @@ -188,6 +230,21 @@ end """ $(TYPEDSIGNATURES) +Fallback for the non-parameterized `SciML` type that delegates to `SciML{CPU}`. + +Preserves backward compatibility for `metadata(SciML)` once the extension defines only +the parameterized `metadata(SciML{P})`. Delegates through +[`CTBase.Strategies.default_parameter`](@extref). + +See also: [`CTSolvers.Integrators.SciML`](@ref), [`CTBase.Strategies.StrategyMetadata`](@extref). +""" +function Strategies.metadata(::Type{SciML}) + return Strategies.metadata(SciML{Strategies.default_parameter(SciML)}) +end + +""" +$(TYPEDSIGNATURES) + Return the default SciML ODE algorithm for the given tag type. This stub returns `missing` for the abstract tag type. The actual implementation @@ -201,3 +258,28 @@ See also: [`CTSolvers.Integrators.SciML`](@ref), [`CTSolvers.Integrators.Tsit5Ta function __default_sciml_algorithm(::Type{<:Core.AbstractTag}) return missing end + +""" +$(TYPEDSIGNATURES) + +Check whether an initial condition `u0` is consistent with the execution parameter `P` +of a `SciML{P}` integrator. + +Mirrors `Modelers.__consistent_backend`: the default returns `true` (all combinations +allowed); the `CTSolversCUDA` extension adds the device-aware methods that flag a `CuArray` +`u0` under `SciML{CPU}`, or a host `Array` `u0` under `SciML{GPU}`. The seam is defined here, +next to the integrator; the consuming package (e.g. CTFlows) calls it where a concrete `u0` +is available (problem construction / solve), since `SciML` does not receive `u0` at build time. + +# Arguments +- `parameter_type::Type{<:AbstractStrategyParameter}`: `CPU` or `GPU`. +- `u0`: The initial condition array to check. + +# Returns +- `Bool`: `true` if consistent, `false` otherwise. + +See also: [`CTSolvers.Integrators.SciML`](@ref), [`CTBase.Strategies.CPU`](@extref), [`CTBase.Strategies.GPU`](@extref). +""" +function __consistent_initial_condition(::Type{<:AbstractStrategyParameter}, u0) + return true +end diff --git a/test/suite/integrators/test_integrator_stubs.jl b/test/suite/integrators/test_integrator_stubs.jl index 32da0f53..efe5b659 100644 --- a/test/suite/integrators/test_integrator_stubs.jl +++ b/test/suite/integrators/test_integrator_stubs.jl @@ -3,6 +3,7 @@ module TestIntegratorExtensionStubs using Test: Test import CTBase.Core import CTBase.Exceptions +import CTBase.Strategies import CTSolvers.Integrators using CommonSolve: CommonSolve @@ -34,12 +35,12 @@ function test_integrator_stubs() Test.@testset "build_sciml_integrator stub" begin Test.@test_throws Exceptions.ExtensionError Integrators.build_sciml_integrator( - DummyTag + DummyTag, Strategies.CPU ) err = nothing try - Integrators.build_sciml_integrator(DummyTag) + Integrators.build_sciml_integrator(DummyTag, Strategies.CPU) catch e err = e end diff --git a/test/suite/integrators/test_sciml_parameter.jl b/test/suite/integrators/test_sciml_parameter.jl new file mode 100644 index 00000000..d3b046d4 --- /dev/null +++ b/test/suite/integrators/test_sciml_parameter.jl @@ -0,0 +1,140 @@ +module TestSciMLParameter + +using Test: Test +import CTBase.Core +import CTBase.Exceptions +import CTSolvers.Integrators +import CTBase.Strategies + +# Extensions: CTSolversSciMLIntegrator (DiffEqBase + SciMLBase), the Tsit5 default alg +# (OrdinaryDiffEqTsit5), and CTSolversCUDA (CUDA) for the device consistency validators. +using OrdinaryDiffEqTsit5: OrdinaryDiffEqTsit5, Tsit5 +using SciMLBase: SciMLBase +using DiffEqBase: DiffEqBase +using CUDA: CUDA + +const VERBOSE = isdefined(Main, :TestData) ? Main.TestData.VERBOSE : true +const SHOWTIMING = isdefined(Main, :TestData) ? Main.TestData.SHOWTIMING : true + +is_cuda_on() = CUDA.functional() + +""" + test_sciml_parameter() + +🧪 **Applying Testing Rule**: Contract tests for the device-parameterized `SciML{P}` integrator. + +Covers the `SciML{P<:Union{CPU,GPU}}` parameterization (GPU roadmap phase 1 / D1): the parameter +contract (`parameter`/`default_parameter`/`id`), per-parameter `metadata` and back-compat bare +delegation, per-parameter construction, registry `[CPU, GPU]` registration + global-parameter +extraction, and the `__consistent_initial_condition` device consistency validators. All assertions +run on CPU; only the `CuArray` cases are gated behind `is_cuda_on()`. +""" +function test_sciml_parameter() + Test.@testset "SciML{P} parameterization" verbose=VERBOSE showtiming=SHOWTIMING begin + + # ==================================================================== + # CONTRACT - parameter / default_parameter / id (no extension needed) + # ==================================================================== + + Test.@testset "parameter contract" begin + Test.@test Strategies.parameter(Integrators.SciML{Strategies.CPU}) == + Strategies.CPU + Test.@test Strategies.parameter(Integrators.SciML{Strategies.GPU}) == + Strategies.GPU + Test.@test Strategies.default_parameter(Integrators.SciML) == Strategies.CPU + + Test.@test Strategies.id(Integrators.SciML) === :sciml + Test.@test Strategies.id(Integrators.SciML{Strategies.CPU}) === :sciml + Test.@test Strategies.id(Integrators.SciML{Strategies.GPU}) === :sciml + end + + # ==================================================================== + # CONTRACT - per-parameter metadata + back-compat bare delegation + # ==================================================================== + + Test.@testset "per-parameter metadata" begin + md_bare = Strategies.metadata(Integrators.SciML) + md_cpu = Strategies.metadata(Integrators.SciML{Strategies.CPU}) + md_gpu = Strategies.metadata(Integrators.SciML{Strategies.GPU}) + + Test.@test md_bare isa Strategies.StrategyMetadata + Test.@test md_cpu isa Strategies.StrategyMetadata + Test.@test md_gpu isa Strategies.StrategyMetadata + + # Bare `metadata(SciML)` delegates to `metadata(SciML{CPU})` (back-compat). + Test.@test collect(keys(md_bare)) == collect(keys(md_cpu)) + + # Option set is currently identical for both devices (few defaults differ yet). + Test.@test collect(keys(md_cpu)) == collect(keys(md_gpu)) + for k in (:alg, :reltol, :abstol, :internalnorm) + Test.@test haskey(md_cpu, k) + Test.@test haskey(md_gpu, k) + end + end + + # ==================================================================== + # CONTRACT - construction per parameter (SciML() ≡ SciML{CPU}) + # ==================================================================== + + Test.@testset "construction per parameter" begin + integ_default = Integrators.SciML(; alg=Tsit5()) + Test.@test integ_default isa Integrators.SciML{Strategies.CPU} + Test.@test Strategies.parameter(typeof(integ_default)) == Strategies.CPU + + integ_cpu = Integrators.SciML{Strategies.CPU}(; alg=Tsit5()) + Test.@test integ_cpu isa Integrators.SciML{Strategies.CPU} + + # Building a SciML{GPU} needs no functional GPU — only the option bundle is built. + integ_gpu = Integrators.SciML{Strategies.GPU}(; alg=Tsit5()) + Test.@test integ_gpu isa Integrators.SciML{Strategies.GPU} + Test.@test Strategies.parameter(typeof(integ_gpu)) == Strategies.GPU + Test.@test Strategies.id(typeof(integ_gpu)) === :sciml + end + + # ==================================================================== + # INTEGRATION - registry with [CPU, GPU] + global parameter extraction + # ==================================================================== + + Test.@testset "registry with device parameters" begin + r = Strategies.create_registry( + Integrators.AbstractIntegrator => + ((Integrators.SciML, [Strategies.CPU, Strategies.GPU]),), + ) + + Test.@test :sciml in Strategies.strategy_ids(Integrators.AbstractIntegrator, r) + + Test.@test Strategies.extract_global_parameter_from_method((:sciml, :cpu), r) == + Strategies.CPU + Test.@test Strategies.extract_global_parameter_from_method((:sciml, :gpu), r) == + Strategies.GPU + end + + # ==================================================================== + # CONTRACT - __consistent_initial_condition device validator + # ==================================================================== + + Test.@testset "initial-condition consistency" begin + host = [1.0, 2.0] + + # Core default: a host array is consistent with CPU. + Test.@test Integrators.__consistent_initial_condition(Strategies.CPU, host) == true + + # CTSolversCUDA (loaded via `using CUDA`): a host array is inconsistent with GPU. + Test.@test Integrators.__consistent_initial_condition(Strategies.GPU, host) == false + + if is_cuda_on() + dev = CUDA.cu(host) + Test.@test Integrators.__consistent_initial_condition( + Strategies.GPU, dev + ) == true + Test.@test Integrators.__consistent_initial_condition( + Strategies.CPU, dev + ) == false + end + end + end +end + +end # module + +test_sciml_parameter() = TestSciMLParameter.test_sciml_parameter() From 9391006dde8ae9a6b81092fc76acec761209acbc Mon Sep 17 00:00:00 2001 From: Olivier Cots Date: Sun, 19 Jul 2026 00:48:32 +0200 Subject: [PATCH 2/2] fix(integrators): restore bare parameter(::Type{<:SciML})=nothing for registry back-compat A consumer that registers SciML without a device parameter (e.g. CTFlows' flow registry) queries parameter on the bare SciML type. The parameterization dropped the bare method, and AbstractIntegrator has no family-level fallback (unlike CTBase's AbstractADBackend), so bare parameter(SciML) threw NotImplemented, breaking every current CTFlows flow-construction path. Restore the bare method (nothing) alongside the parameterized one (P); add a regression guard test. Co-Authored-By: Claude Opus 4.8 --- src/Integrators/sciml.jl | 20 +++++++++++++++++-- .../suite/integrators/test_sciml_parameter.jl | 5 +++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/src/Integrators/sciml.jl b/src/Integrators/sciml.jl index 4f3ee7e6..eff0471e 100644 --- a/src/Integrators/sciml.jl +++ b/src/Integrators/sciml.jl @@ -84,10 +84,26 @@ Strategies.id(::Type{<:SciML}) = :sciml """ $(TYPEDSIGNATURES) -Return the execution parameter type of a `SciML` integrator. +Return the execution parameter of the non-parameterized `SciML` type: `nothing`. + +The bare `SciML` (device unspecified) resolves to `nothing`, matching the general +`AbstractStrategy` contract and preserving backward compatibility for consumers that +register `SciML` without a parameter (e.g. CTFlows' flow registry) and query +`parameter` on the bare type. A concrete `SciML{P}` resolves to `P` via the more +specific method below. + +See also: [`CTSolvers.Integrators.SciML`](@ref), [`CTBase.Strategies.default_parameter`](@extref) +""" +Strategies.parameter(::Type{<:SciML}) = nothing + +""" +$(TYPEDSIGNATURES) + +Return the execution parameter type of a parameterized `SciML{P}` integrator. Extracts the type parameter `P` from `SciML{P}`, which can be either `CPU` or `GPU` -since `SciML` supports both execution devices. +since `SciML` supports both execution devices. More specific than the bare +`parameter(::Type{<:SciML})` above, so it wins for a concrete `SciML{P}`. # Returns - `Type{<:Union{CPU,GPU}}`: the execution parameter type. diff --git a/test/suite/integrators/test_sciml_parameter.jl b/test/suite/integrators/test_sciml_parameter.jl index d3b046d4..0d518255 100644 --- a/test/suite/integrators/test_sciml_parameter.jl +++ b/test/suite/integrators/test_sciml_parameter.jl @@ -37,6 +37,11 @@ function test_sciml_parameter() # ==================================================================== Test.@testset "parameter contract" begin + # Bare (device-unspecified) SciML resolves to `nothing` — required by consumers + # that register SciML without a parameter (e.g. CTFlows' flow registry) and query + # `parameter` on the bare type. Regression guard for the removed-then-restored method. + Test.@test Strategies.parameter(Integrators.SciML) === nothing + Test.@test Strategies.parameter(Integrators.SciML{Strategies.CPU}) == Strategies.CPU Test.@test Strategies.parameter(Integrators.SciML{Strategies.GPU}) ==