Unsuported call to an unknown function when calling Distributions

[FredericWantiez]

Hello,

I've been trying to use metal with distributions but it seems code like this one would return an unsupported call error:

function kernel(state)
    idx = thread_position_in_grid_1d()
    state[idx] = rand(Normal{Float32}(0.f0, 1.f0))
    return
end

n = 10
state = Metal.ones(n)
@metal threads=n kernel(state)

ERROR: InvalidIRError: compiling MethodInstance for kernel(::MtlDeviceVector{Float32, 1}) 
resulted in invalid LLVM IRReason: unsupported call to an unknown function (call to julia.get_pgcstack)

Whereas the broadcast version works fine:

Metal.@sync rand!(Normal{Float32}(0.f0, 1.f0), output)

[christiangnrd]

Your kernel is constructing a Normal{Float32}(0.f0, 1.f0) object for every state index. On top of being very inefficient, this code will not run on GPU as it creates allocations.

What if you tried:

function kernel(state, dist)
    idx = thread_position_in_grid_1d()
    state[idx] = rand(dist)
    return
end

n = 10
state = Metal.ones(n)
@metal threads=n kernel(state, Normal{Float32}(0f0,1f0))

[FredericWantiez]

In the actual code, the distribution depends on the state, something like Normal(state[idx], 1.f0). Tried running your example but that still raises a pgstack error.

julia> @metal threads=n kernel(state, Normal{Float32}(0f0,1f0))
ERROR: InvalidIRError: compiling MethodInstance for kernel(::MtlDeviceVector{Float32, 1}, ::Normal{Float32}) resulted in invalid LLVM IR
Reason: unsupported call to an unknown function (call to julia.get_pgcstack)
Stacktrace:
 [1] randn_unlikely
   @ ~/.julia/juliaup/julia-1.9.4+0.aarch64.apple.darwin14/share/julia/stdlib/v1.9/Random/src/normal.jl:81
 [2] multiple call sites
   @ unknown:0

These are the system info:

macOS 14.5.0, Darwin 23.5.0

Toolchain:
- Julia: 1.10.4
- LLVM: 15.0.7

Julia packages:
- Metal.jl: 1.2.0
- LLVMDowngrader_jll: 0.3.0+1

1 device:
- Apple M1 Max (384.000 KiB allocated)

[maleadt]

External packages not supporting GPU execution is not a bug in the GPU support package, so it's probably better to file an issue on Distributions.jl for GPU compatibility. If that reveals specific features are needed for that support, you can file them here.

For example, your stack trace seems to show that rand is called by Distributions, which is currently not supported by Metal.jl kernels. If you can confirm that and create a MWE of what's needed, feel free to open a new issue with that information (but know that RNG support in kernels is fairly involved, see e.g. the CUDA.jl implementation, https://github.com/JuliaGPU/CUDA.jl/blob/master/src/device/random.jl).

(For other questions like this, something in GitHub's discussions, or on Discourse or Slack, is probably better suited than filing a bug report.)