# CUDA/GPU compatible discretization from MethodOfLines.jl

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

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

Okeydoke. Thanks!