CUDA/GPU compatible discretization from MethodOfLines.jl

Specific Domains GPU
, , ,
1

MethodOfLines.jl (MOL) exports discretize which gives an ODEProblem that can then be directly passed to OrdinaryDiffEq.solve() to get a solution. This works great if you want a CPU evaluation, but it seems not so much if you want to evaluate on the GPU. So I’m curious, is there some implementation I’ve missed on using MOL to produce a GPU compatible discretization?

For the type of problem I’m working with, DiffEqGPU.jl recommends using “Within-Method GPU Parallelism”. In this implementation, you can just wrap the initial conditions of the ODEProblem in a CUDA array cu(..) and pass them to the ODEProblem constructor, but MOL gives an already-constructed ODEProblem object. Upon naively trying to update the initial condition u0 field with

prob = discretize(pdesys, discretization)
prob.u0 = cu(prob.u0)

I get met with an error as the ODEProblem is an immutable struct.

To try to get around this I found Accessors.jl. So I tried again with

prob = discretize(pdesys, discretization)
@reset prob.u0 = cu(prob.u0)

to effectively mutate the immutable. This is met with a MethodError occurring from ODEProblem on execution the @reset macro. I’m not really sure how to even begin parsing the error message as there’s so much going on with ODEProblem.

Minimal Example

Using the Heat Equation example from MOL.

using OrdinaryDiffEq, ModelingToolkit, MethodOfLines, DomainSets # MethodOfLines dependencies
using Accessors, CUDA # Need for GPU

# Parameters, variables, and derivatives
@parameters t x
@variables u(..)
Dt = Differential(t)
Dxx = Differential(x)^2

# 1D PDE and boundary conditions
eq  = Dt(u(t, x)) ~ Dxx(u(t, x))
bcs = [u(0, x) ~ cos(x),
        u(t, 0) ~ exp(-t),
        u(t, 1) ~ exp(-t) * cos(1)]

# Space and time domains
domains = [t ∈ Interval(0.0, 1.0),
           x ∈ Interval(0.0, 1.0)]

# PDE system
@named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])

# Method of lines discretization
dx = 0.1
order = 2
discretization = MOLFiniteDifference([x => dx], t)

# Convert the PDE problem into an ODE problem
prob = discretize(pdesys,discretization)
@set prob.u0 = cu(prob.u0)

The error still gets thrown even without actually changing anything: @set prob.u0 = prob.u0.

Error Message

ERROR: MethodError: no method matching (ODEProblem{true})(::ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#k#545"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x9e30353b, 0x76a0822f, 0xdba6ca3e, 0x77285d92, 0xd69d0c46), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xd1da9c88, 0x271bcf06, 0x449fb7ad, 0x202f1261, 0x20654849), Nothing}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Any}, ModelingToolkit.var"#630#generated_observed#555"{Bool, ODESystem, Dict{Any, Any}, Vector{Any}}, Nothing, ODESystem}, ::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ::Tuple{Float64, Float64}, ::Vector{Float64}, ::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::MethodOfLines.MOLMetadata{Val{true}(), MethodOfLines.DiscreteSpace{1, 1, MethodOfLines.CenterAlignedGrid}, MOLFiniteDifference{MethodOfLines.CenterAlignedGrid, MethodOfLines.ScalarizedDiscretization}, PDESystem, Base.RefValue{Any}, MethodOfLines.ScalarizedDiscretization})

Closest candidates are:
  (ODEProblem{iip})(::SciMLBase.AbstractODEFunction{iip}, ::Any, ::Any, ::Any, ::Any; kwargs...) where iip
   @ SciMLBase .julia\packages\SciMLBase\szsYq\src\problems\ode_problems.jl:111
  (ODEProblem{true})(::ModelingToolkit.AbstractODESystem, ::Any...; kwargs...)
   @ ModelingToolkit .julia\packages\ModelingToolkit\BsHty\src\systems\diffeqs\abstractodesystem.jl:911
  (ODEProblem{iip})(::SciMLBase.AbstractODEFunction{iip}, ::Any, ::Any, ::Any) where iip  
   @ SciMLBase .julia\packages\SciMLBase\szsYq\src\problems\ode_problems.jl:111
1 Like
2

You can open an issue to track this. I can at least tell you that the current discretization is not GPU compatible, but @xtalax is working on a change to the discretization that is in vector form that hopefully down the line will be able to be GPU-compatible, but the symbolic tooling just isn’t ready for this yet. It’s a target for the next year though.

If you need GPUs in the meantime, I would suggest doing a manual DifferentialEquations.jl problem construction.

1 Like
3

Okeydoke. Thanks!