Passing `::Type{InconcreteType}` to CUDA kernel


I have some kernels that share a lot of code but specialize in the end based on some type of AbstractAction in the MWE below Addition or Multiplication. Those concrete actions have some nonbitstype metadata, which is not important in the kernel, so I thought I’d just pass the Type{<:AbstractAction} to the GPU kernel rather than the actual object.

Even though isbitstype(Type{someting}) == false, there seems to be special handling in the GPU compiler which allows me to pass those objects to the kernel if something is a concrete type, however it does not work for non-concrete partial types. Hard to explain but, see the MWE below.

Is there a clever way to circumvent this issue? I can only think of

  • making up meta data type parameters to “complete” the type even though they don’t do anything,
  • creating special dispatch types like AdditionDispatch without any type parameters/metadata,

both of which I don’t really like…

using Pkg
pkg"activate --temp"
pkg"add KernelAbstractions, CUDA"
using KernelAbstractions
using CUDA

abstract type AbstractAction end
struct Addition{M} <: AbstractAction
struct Multiplication{M} <: AbstractAction

@kernel function kernel!(::Type{T}, z, x, y) where {T<:AbstractAction}
    i = @index(Global)
    z[i] = apply(T, x[i], y[i])

@inline apply(::Type{<:Addition}, x, y) = x + y
@inline apply(::Type{<:Multiplication}, x, y) = x * y

x = CuArray(1:10)
y = CuArray(1:10)
z = CuArray(zeros(10))
kernel = kernel!(get_backend(x))

# works on concrete type
kernel(Addition{:foo}, z, x, y; ndrange=length(z))
kernel(Multiplication{:foo}, z, x, y; ndrange=length(z))
isbitstype(Addition{:foo}) # no bitstype!

# does not work on abstract type, tells me its forbidden to pass non-isbits to kernel
kernel(Addition, z, x, y; ndrange=length(z))
kernel(Multiplication, z, x, y; ndrange=length(z))
isbitstype(Addition) # no bitstype!

Arguments are allowed if they’re ghost types or const types:

julia> Core.Compiler.isconstType(Type{Int})

julia> Core.Compiler.isconstType(Type{Number})

julia> Core.Compiler.isconstType(Type{Array})

So that doesn’t hold for your incomplete type.

That said, at some point in the past we used to support passing other arguments, as long as they were unused. Some code supporting that remains, GPUCompiler.jl/src/validation.jl at 85316c7f9f962667121ca8c2455b7a1e5993fd6b · JuliaGPU/GPUCompiler.jl · GitHub and CUDA.jl/src/compiler/execution.jl at 71311afa35463f3fce09ed73bae66f8810b5e821 · JuliaGPU/CUDA.jl · GitHub, so if this would work for you, open an issue or have a look at the GPUCompiler.jl git history to add this support back. Looks like it got lost in Clean-up version checks. · JuliaGPU/GPUCompiler.jl@eec85d5 · GitHub.

1 Like

Thanks for the explanation! I tried to reenable that feature but turns out to be more difficult now to check if the argument is used inside the kernel. I opened an issue instead.