Can ModelingToolkit be used with GPU arrays?

Specific Domains Modelling & Simulations
1

Hello, everyone. I’m trying to solve a complex system using ModelingToolkit, which involves an array variable. One of the equations is quite heavy computationally and I was trying to pass this array as an argument, then build a GPU array (CUDA in this case, I don’t have more GPUs now), do the calculation and return the value back in the original format. The problem is that symbolic variables are of type Num and they are not isbitstype.

I then decided to try a small example and build a simple model with a CuArray with the idea of knowing if I can build everything directly:

using CUDA, ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D

x0 = ones(2)
x0CU = CuArray(x0)
A = Float64[2 1; 1 2]
ACU = CuArray(A)

@mtkmodel Example begin
    @variables begin
        x(t)[1:2] = x0
    end
    @equations begin
        D(x) ~ A * x
    end
end

@mtkcompile example = Example()

@mtkmodel ExampleCU begin
    @variables begin
        x(t)[1:2] = x0CU
    end
    @equations begin
        D(x) ~ derivative(x, ACU)
    end
end

@mtkcompile exampleCU = ExampleCU()

The example model is compiled correctly on the CPU, but the exampleCU one for the GPU does not compile since it throws the scalar indexing error

ERROR: Scalar indexing is disallowed.

Note that something like

ACU * x0CU

outside of @mtkmodel works and returns a CuArray as a result.

Is there a way in ModelingToolkit for me to take a symbolic array, convert it to the desired numerical format for me to parallelize it and then return it back to MTK with the original type so the program can continue the evaluation? Or am I restricted to, at most, CPU parallelization?

2

I would try to not set x(t)[1:2] = x0CU in the @mtkmodel, but just declare @variables x(t)[1:2] without setting the default value. Instead set x when you create the ODEProblem.

I think the @mtkmodel step should be the same for the CPU and GPU. Just create one system and use that in both cases. The Num should not get transferred to the CPU; this is just a symbolic type that is used for the symbolic-to-numeric code generation. The CPU/GPU specialization happens first when you create and solve an ODEProblem.

I have not tried to run MTK code on the GPU myself, but you could have a look at this test and build from there. Maybe there also exists a fully documented and working example somewhere that I am not aware of?

3

I can start working from there and define the initial values like they do, thanks!

I guess I also have to rethink how I implement what I want. Because that expensive function makes use of StaticArrays (I’m using a vector of 3x3 SMatrix) and when I implement it into MTK it fails since it doesn’t understand how to multiply a Num with an SMatrix and then sum another SMatrix. I’ve tried using @register_array_symbolic without success.

4

Not yet, we’re working on a GPU codegen and DiffEqGPU support