From 4fb77e26dde61793e3763c9f9c9590d4afa9022a Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 21 May 2025 13:39:20 +0000 Subject: [PATCH 1/3] Bump to use SciMLBase ImmutableNonlinearProblem --- docs/src/tutorials/nonlinear_solve_gpus.md | 4 +- lib/NonlinearSolveBase/Project.toml | 2 +- .../src/NonlinearSolveBase.jl | 4 +- .../src/immutable_problem.jl | 56 ------------------- 4 files changed, 4 insertions(+), 62 deletions(-) delete mode 100644 lib/NonlinearSolveBase/src/immutable_problem.jl diff --git a/docs/src/tutorials/nonlinear_solve_gpus.md b/docs/src/tutorials/nonlinear_solve_gpus.md index d594dae6e..a2f0871ae 100644 --- a/docs/src/tutorials/nonlinear_solve_gpus.md +++ b/docs/src/tutorials/nonlinear_solve_gpus.md @@ -142,7 +142,7 @@ end p = @SVector [@SVector(rand(Float32, 4)) for _ in 1:1024] u0 = SA[1f0, 2f0, 3f0, 4f0] -prob = NonlinearSolveBase.ImmutableNonlinearProblem{false}(p2_f, u0, p) +prob = SciMLBase.ImmutableNonlinearProblem{false}(p2_f, u0, p) ``` !!! note @@ -152,7 +152,7 @@ prob = NonlinearSolveBase.ImmutableNonlinearProblem{false}(p2_f, u0, p) In general, this means that you need to avoid any runtime operations in kernel code, such as allocating vectors, dynamic dispatch, type instabilities, etc. Thus to make this work, your `f` function should be non-allocating, your `u0` function should use - StaticArrays, and you must use `NonlinearSolveBase.ImmutableNonlinearProblem` + StaticArrays, and you must use `SciMLBase.ImmutableNonlinearProblem` (which is exactly the same as NonlinearProblem except it's immutable to satisfy the requirements of GPU kernels). Also, it's recommended that for most GPUs you use Float32 precision because many GPUs are much slower on 64-bit floating point operations. diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index ea872ba97..b762a6aa2 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -68,7 +68,7 @@ MaybeInplace = "0.1.4" Preferences = "1.4" Printf = "1.10" RecursiveArrayTools = "3" -SciMLBase = "2.89.1" +SciMLBase = "2.91" SciMLJacobianOperators = "0.1.1" SciMLOperators = "0.3.13, 0.4" SparseArrays = "1.10" diff --git a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl index d86718329..e6c61ef8d 100644 --- a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl +++ b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl @@ -20,7 +20,7 @@ using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinear AbstractNonlinearAlgorithm, AbstractNonlinearFunction, NonlinearProblem, NonlinearLeastSquaresProblem, StandardNonlinearProblem, NonlinearFunction, NullParameters, NLStats, LinearProblem, - LinearAliasSpecifier + LinearAliasSpecifier, ImmutableNonlinearProblem using SciMLJacobianOperators: JacobianOperator, StatefulJacobianOperator using SciMLOperators: AbstractSciMLOperator, IdentityOperator using SymbolicIndexingInterface: SymbolicIndexingInterface @@ -36,8 +36,6 @@ include("public.jl") include("utils.jl") include("abstract_types.jl") - -include("immutable_problem.jl") include("common_defaults.jl") include("termination_conditions.jl") diff --git a/lib/NonlinearSolveBase/src/immutable_problem.jl b/lib/NonlinearSolveBase/src/immutable_problem.jl deleted file mode 100644 index 2d429845f..000000000 --- a/lib/NonlinearSolveBase/src/immutable_problem.jl +++ /dev/null @@ -1,56 +0,0 @@ -struct ImmutableNonlinearProblem{uType, iip, P, F, K, PT} <: - AbstractNonlinearProblem{uType, iip} - f::F - u0::uType - p::P - problem_type::PT - kwargs::K - - SciMLBase.@add_kwonly function ImmutableNonlinearProblem{iip}( - f::AbstractNonlinearFunction{iip}, u0, p = NullParameters(), - problem_type = StandardNonlinearProblem(); kwargs...) where {iip} - if haskey(kwargs, :p) - error("`p` specified as a keyword argument `p = $(kwargs[:p])` to \ - `NonlinearProblem`. This is not supported.") - end - SciMLBase.warn_paramtype(p) - return new{ - typeof(u0), iip, typeof(p), typeof(f), typeof(kwargs), typeof(problem_type)}( - f, u0, p, problem_type, kwargs) - end - - """ - Define a steady state problem using the given function. - `isinplace` optionally sets whether the function is inplace or not. - This is determined automatically, but not inferred. - """ - function ImmutableNonlinearProblem{iip}( - f, u0, p = NullParameters(); kwargs...) where {iip} - return ImmutableNonlinearProblem{iip}(NonlinearFunction{iip}(f), u0, p; kwargs...) - end -end - -""" -Define a nonlinear problem using an instance of [`AbstractNonlinearFunction`](@ref). -""" -function ImmutableNonlinearProblem( - f::AbstractNonlinearFunction, u0, p = NullParameters(); kwargs...) - return ImmutableNonlinearProblem{SciMLBase.isinplace(f)}(f, u0, p; kwargs...) -end - -function ImmutableNonlinearProblem(f, u0, p = NullParameters(); kwargs...) - return ImmutableNonlinearProblem(NonlinearFunction(f), u0, p; kwargs...) -end - -""" -Define a ImmutableNonlinearProblem problem from SteadyStateProblem. -""" -function ImmutableNonlinearProblem(prob::AbstractNonlinearProblem) - return ImmutableNonlinearProblem{SciMLBase.isinplace(prob)}(prob.f, prob.u0, prob.p) -end - -function Base.convert( - ::Type{ImmutableNonlinearProblem}, prob::T) where {T <: NonlinearProblem} - return ImmutableNonlinearProblem{SciMLBase.isinplace(prob)}( - prob.f, prob.u0, prob.p, prob.problem_type; prob.kwargs...) -end From bf01fa793325282bc7410692d421cc783adc6d1e Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 21 May 2025 14:02:14 +0000 Subject: [PATCH 2/3] Update lib/NonlinearSolveBase/Project.toml --- lib/NonlinearSolveBase/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index b762a6aa2..7c50484f9 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -68,7 +68,7 @@ MaybeInplace = "0.1.4" Preferences = "1.4" Printf = "1.10" RecursiveArrayTools = "3" -SciMLBase = "2.91" +SciMLBase = "2.92" SciMLJacobianOperators = "0.1.1" SciMLOperators = "0.3.13, 0.4" SparseArrays = "1.10" From edc2ab29a291005b752c51036f7b4d24ca93074a Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 21 May 2025 15:06:58 +0000 Subject: [PATCH 3/3] Update NonlinearSolveBase.jl --- lib/NonlinearSolveBase/src/NonlinearSolveBase.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl index e6c61ef8d..7d1260260 100644 --- a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl +++ b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl @@ -17,9 +17,9 @@ using EnzymeCore: EnzymeCore using MaybeInplace: @bb using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinearProblem, - AbstractNonlinearAlgorithm, AbstractNonlinearFunction, - NonlinearProblem, NonlinearLeastSquaresProblem, StandardNonlinearProblem, - NonlinearFunction, NullParameters, NLStats, LinearProblem, + AbstractNonlinearAlgorithm, + NonlinearProblem, NonlinearLeastSquaresProblem, + NonlinearFunction, NLStats, LinearProblem, LinearAliasSpecifier, ImmutableNonlinearProblem using SciMLJacobianOperators: JacobianOperator, StatefulJacobianOperator using SciMLOperators: AbstractSciMLOperator, IdentityOperator