SciML / SciML/RuntimeGeneratedFunctions.jl

GPU support

Open
#93 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
112
Forks
19
Avg merge
6h 11m
Merged PRs (30d)
9

Description

I've been using this project to circumvent the world age problems when generating functions and it has worked great so far.

Now I'm trying to also generate functions to run on the GPU, CUDA for now to be specific. There are two ways that I've tried to do that. The first is simply broadcasting the function on a CuVector, however, this does not work because RuntimeGeneratedFunctions are not isbits, which CUDA.jl does not like:

julia> func.(cu_inputs)
ERROR: GPU compilation of MethodInstance for (::GPUArrays.var"#34#36")(::CUDA.CuKernelContext, ::CuDeviceVector{ComplexF64, 1}, ::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{…}, Tuple{…}, RuntimeGeneratedFunction{…}, Tuple{…}}, ::Int64) failed
KernelError: passing and using non-bitstype argument

Argument 4 to your kernel function is of type Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1, CUDA.Mem.DeviceBuffer}, Tuple{Base.OneTo{Int64}}, RuntimeGeneratedFunction{(:data_input,), ComputableDAGs.var"#_RGF_ModTag", var"#_RGF_ModTag", (0x3d86ea6b, 0x46940344, 0xf0c50769, 0x501ce261, 0x5c3fa2fa), Expr}, Tuple{Base.Broadcast.Extruded{CuDeviceVector{<LongArgumentType...>, 1}, Tuple{Bool}, Tuple{Int64}}}}, which is not isbits:
  .f is of type RuntimeGeneratedFunction{(:data_input,), ComputableDAGs.var"#_RGF_ModTag", var"#_RGF_ModTag", (0x3d86ea6b, 0x46940344, 0xf0c50769, 0x501ce261, 0x5c3fa2fa), Expr} which is not isbits.
    .body is of type Expr which is not isbits.
      .head is of type Symbol which is not isbits.
      .args is of type Vector{Any} which is not isbits.

The other option is to use @cuda and call a (generated) kernel. However, this also does not work as RuntimeGeneratedFunctions seem to always return an Any, and GPU kernels are not allowed to return anything:

julia> @cuda kernel(cu_inputs, cu_outputs, 1000)
ERROR: GPU compilation of MethodInstance for (::RuntimeGeneratedFunction{(:input_vector, :output_vector, :n), CUDAExt.var"#_RGF_ModTag", var"#_RGF_ModTag", (0x26ea8e24, 0xf73e2d2a, 0xb1e0f07e, 0x53860e81, 0xc0771049), Expr})(::CuDeviceVector{PhaseSpacePoint{…}, 1}, ::CuDeviceVector{Float64, 1}, ::Int64) failed
KernelError: kernel returns a value of type `Any`

Make sure your kernel function ends in `return`, `return nothing` or `nothing`.

It would be great if compatibility for GPU broadcasting and kernels could be added. Looking at the code in this package, I don't understand how it is doing what it is doing so I don't think I can add support myself. It seems that fixing the return nothing problem with kernels might not be too hard though.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the two reported cases: broadcasting a RuntimeGeneratedFunction over a CuVector and calling a generated kernel with @cuda. Trace how RuntimeGeneratedFunction is lowered in each path; done means GPU broadcasting and kernels compile without non-isbits arguments or an Any return.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
hpc
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.