From 073cb86d1e7aab716ceb2afad3fa5be704d2e4dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hans=20W=C3=BCrfel?= Date: Thu, 20 Nov 2025 12:37:51 +0100 Subject: [PATCH 1/3] fix GPU compat of BrownFullBasicInit --- .../src/initialize_dae.jl | 6 +- test/gpu/Project.toml | 2 + test/gpu/simple_dae.jl | 65 +++++++++++++++++++ test/runtests.jl | 1 + 4 files changed, 71 insertions(+), 3 deletions(-) create mode 100644 test/gpu/simple_dae.jl diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl index e87963f362..1431893092 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl @@ -396,8 +396,8 @@ function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::OD M = integrator.f.mass_matrix M isa UniformScaling && return update_coefficients!(M, u, p, t) - algebraic_vars = [all(iszero, x) for x in eachcol(M)] - algebraic_eqs = [all(iszero, x) for x in eachrow(M)] + algebraic_vars = mapreduce(iszero, &, M, dims=1)[:] + algebraic_eqs = mapreduce(iszero, &, M, dims=2)[:] (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return tmp = get_tmp_cache(integrator)[1] @@ -456,7 +456,7 @@ function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::OD nlsolve = default_nlsolve(alg.nlsolve, isinplace, u, nlprob, isAD) nlsol = solve(nlprob, nlsolve; abstol = alg.abstol, reltol = integrator.opts.reltol) - alg_u .= nlsol + alg_u .= nlsol.u recursivecopy!(integrator.uprev, integrator.u) if alg_extrapolates(integrator.alg) diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml index e0e8db295f..fdf60e309f 100644 --- a/test/gpu/Project.toml +++ b/test/gpu/Project.toml @@ -6,6 +6,8 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8" +OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" [compat] diff --git a/test/gpu/simple_dae.jl b/test/gpu/simple_dae.jl new file mode 100644 index 0000000000..9c5a451d40 --- /dev/null +++ b/test/gpu/simple_dae.jl @@ -0,0 +1,65 @@ +using OrdinaryDiffEqRosenbrock +using OrdinaryDiffEqNonlinearSolve +using CUDA +using LinearAlgebra +using Adapt +using SparseArrays +using Test + +#= +du[1] = -u[1] +du[2] = -0.5*u[2] + 0 = u[1] + u[2] - u[3] + 0 = -u[1] + u[2] - u[4] +=# + +function dae!(du, u, p, t) + mul!(du, p, u) +end + +p = [-1 0 0 0 + 1 -0.5 0 0 + 1 1 -1 0 + -1 1 0 -1] + +# mass_matrix = [1 0 0 0 +# 0 1 0 0 +# 0 0 0 0 +# 0 0 0 0] +mass_matrix = Diagonal([1, 1, 0, 0]) +jac_prototype = sparse(map(x -> iszero(x) ? 0.0 : 1.0, p)) + +u0 = [1.0, 1.0, 0.5, 0.5] # force init +odef = ODEFunction(dae!, mass_matrix=mass_matrix, jac_prototype=jac_prototype) + +tspan = (0.0, 5.0) +prob = ODEProblem(odef, u0, tspan, p) +sol = solve(prob, Rodas5P()) + +# gpu version +mass_matrix_d = adapt(CuArray, mass_matrix) + +# TODO: jac_prototype fails +# jac_prototype_d = adapt(CuArray, jac_prototype) +# jac_prototype_d = CUDA.CUSPARSE.CuSparseMatrixCSR(jac_prototype) +jac_prototype_d = nothing + +u0_d = adapt(CuArray, u0) +p_d = adapt(CuArray, p) +odef_d = ODEFunction(dae!, mass_matrix=mass_matrix_d, jac_prototype=jac_prototype_d) +prob_d = ODEProblem(odef_d, u0_d, tspan, p_d) +sol_d = solve(prob_d, Rodas5P()) + +@testset "Test constraints in GPU sol" begin + for t in sol_d.t + u = Vector(sol_d(t)) + @test isapprox(u[1] + u[2], u[3]; atol=1e-6) + @test isapprox(-u[1] + u[2], u[4]; atol=1e-6) + end +end + +@testset "Compare GPU to CPU solution" begin + for t in tspan[begin]:0.1:tspan[end] + @test Vector(sol_d(t)) ≈ sol(t) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 09c0af2144..48f1ca0b17 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -205,6 +205,7 @@ end @time @safetestset "Linear Exponential GPU" include("gpu/linear_exp.jl") @time @safetestset "Reaction-Diffusion Stiff Solver GPU" include("gpu/reaction_diffusion_stiff.jl") @time @safetestset "Scalar indexing bug bypass" include("gpu/hermite_test.jl") + @time @safetestset "simple dae on GPU" include("gpu/simple_dae.jl") end if !is_APPVEYOR && GROUP == "QA" From 99ccab083108bb27e036de78650045c0aebf3bd3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hans=20W=C3=BCrfel?= Date: Thu, 20 Nov 2025 14:16:15 +0100 Subject: [PATCH 2/3] run formatter --- .../src/initialize_dae.jl | 4 ++-- test/gpu/simple_dae.jl | 16 ++++++++-------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl index 1431893092..048941c76b 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl @@ -396,8 +396,8 @@ function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::OD M = integrator.f.mass_matrix M isa UniformScaling && return update_coefficients!(M, u, p, t) - algebraic_vars = mapreduce(iszero, &, M, dims=1)[:] - algebraic_eqs = mapreduce(iszero, &, M, dims=2)[:] + algebraic_vars = mapreduce(iszero, &, M, dims = 1)[:] + algebraic_eqs = mapreduce(iszero, &, M, dims = 2)[:] (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return tmp = get_tmp_cache(integrator)[1] diff --git a/test/gpu/simple_dae.jl b/test/gpu/simple_dae.jl index 9c5a451d40..b8a73f0fb3 100644 --- a/test/gpu/simple_dae.jl +++ b/test/gpu/simple_dae.jl @@ -17,10 +17,10 @@ function dae!(du, u, p, t) mul!(du, p, u) end -p = [-1 0 0 0 - 1 -0.5 0 0 - 1 1 -1 0 - -1 1 0 -1] +p = [-1 0 0 0 + 1 -0.5 0 0 + 1 1 -1 0 + -1 1 0 -1] # mass_matrix = [1 0 0 0 # 0 1 0 0 @@ -30,7 +30,7 @@ mass_matrix = Diagonal([1, 1, 0, 0]) jac_prototype = sparse(map(x -> iszero(x) ? 0.0 : 1.0, p)) u0 = [1.0, 1.0, 0.5, 0.5] # force init -odef = ODEFunction(dae!, mass_matrix=mass_matrix, jac_prototype=jac_prototype) +odef = ODEFunction(dae!, mass_matrix = mass_matrix, jac_prototype = jac_prototype) tspan = (0.0, 5.0) prob = ODEProblem(odef, u0, tspan, p) @@ -46,15 +46,15 @@ jac_prototype_d = nothing u0_d = adapt(CuArray, u0) p_d = adapt(CuArray, p) -odef_d = ODEFunction(dae!, mass_matrix=mass_matrix_d, jac_prototype=jac_prototype_d) +odef_d = ODEFunction(dae!, mass_matrix = mass_matrix_d, jac_prototype = jac_prototype_d) prob_d = ODEProblem(odef_d, u0_d, tspan, p_d) sol_d = solve(prob_d, Rodas5P()) @testset "Test constraints in GPU sol" begin for t in sol_d.t u = Vector(sol_d(t)) - @test isapprox(u[1] + u[2], u[3]; atol=1e-6) - @test isapprox(-u[1] + u[2], u[4]; atol=1e-6) + @test isapprox(u[1] + u[2], u[3]; atol = 1e-6) + @test isapprox(-u[1] + u[2], u[4]; atol = 1e-6) end end From ed8a6f359a77d002b78c26c6cc0909626ed497d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hans=20W=C3=BCrfel?= Date: Thu, 20 Nov 2025 16:07:12 +0100 Subject: [PATCH 3/3] improve algebraic idx collection --- lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl index 048941c76b..cdd7014cdd 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl @@ -396,8 +396,9 @@ function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::OD M = integrator.f.mass_matrix M isa UniformScaling && return update_coefficients!(M, u, p, t) - algebraic_vars = mapreduce(iszero, &, M, dims = 1)[:] - algebraic_eqs = mapreduce(iszero, &, M, dims = 2)[:] + algebraic_vars = vec(all(iszero, M, dims = 1)) + algebraic_eqs = vec(all(iszero, M, dims = 2)) + (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return tmp = get_tmp_cache(integrator)[1]