Merge pull request SciML#2941 from ChrisRackauckas-Claude/claude/rebase-hw-dae-gpu-pr2911

fix GPU compat of BrownFullBasicInit (rebased from SciML#2911)

Author: ChrisRackauckas

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 = [all(iszero, x) for x in eachcol(M)]
-    algebraic_eqs = [all(iszero, x) for x in eachrow(M)]
+    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]
 
@@ -456,7 +457,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)

test/gpu/Project.toml

@@ -7,13 +7,17 @@ FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 OrdinaryDiffEqRKIP = "a4daff8c-1d43-4ff3-8eff-f78720aeecdc"
+OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 
 [sources]
+OrdinaryDiffEqNonlinearSolve = {path = "../../lib/OrdinaryDiffEqNonlinearSolve"}
 OrdinaryDiffEqRKIP = {path = "../../lib/OrdinaryDiffEqRKIP"}
+OrdinaryDiffEqRosenbrock = {path = "../../lib/OrdinaryDiffEqRosenbrock"}
 
 [compat]
 Adapt = "4"
@@ -24,7 +28,9 @@ FastBroadcast = "0.3"
 FFTW = "1.8"
 FillArrays = "1"
 OrdinaryDiffEq = "6"
+OrdinaryDiffEqNonlinearSolve = "1"
 OrdinaryDiffEqRKIP = "1"
+OrdinaryDiffEqRosenbrock = "1"
 RecursiveArrayTools = "3"
 SciMLBase = "2.99"
 SciMLOperators = "1.3"

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

test/runtests.jl

@@ -206,6 +206,7 @@ end
         @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 "RKIP Semilinear PDE GPU" include("gpu/rkip_semilinear_pde.jl")
+        @time @safetestset "simple dae on GPU" include("gpu/simple_dae.jl")
     end
 
     if !is_APPVEYOR && GROUP == "QA"