Why can Flux not reduce this?

Ahmed_Salih · February 4, 2023, 11:11pm

Hello

I am trying:

using CUDA
S      = rand(SVector{3,Float32},5)
DST = zeros(SVector{3,Float32},5)
I       = [3,1,2,5,4]

# Works great on CPU

@CUDA.time NNlib.scatter!(+, DST,S,I)
0.000004 seconds
5-element Vector{SVector{3, Float32}}:
[1.3667885, 1.081403, 1.1134366]
[1.6543723, 0.50424564, 1.1448298]
[0.8695842, 1.8623418, 1.398939]
[1.0094697, 0.0052466393, 0.09184897]
[1.9389011, 0.3291297, 1.2308507]


# Bugs out on GPU

 @CUDA.time NNlib.scatter!(+, CuArray(DST),CuArray(S),CuArray(I))
ERROR: InvalidIRError: compiling kernel #scatter_kernel!(typeof(+), CuDeviceVector{SVector{3, Float32}, 1}, CuDeviceVector{SVector{3, Float32}, 1}, CuDeviceVector{Int64, 1}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to atomic_cas!)

Anyone knows why?

This is an example, I need it to work on GPU for a more complex case

Kind regards

maleadt · February 6, 2023, 9:21am

That seems like a bug in NNlibCUDA.jl. Please file an issue there.

Ahmed_Salih · February 6, 2023, 6:05pm

Thank you, I filed an issue 3 weeks ago:

github.com/FluxML/NNlib.jl

SVectors on GPU CuArray cannot use index from CuArray

opened 01:17AM - 14 Jan 23 UTC

AhmedSalih3d

NNlibCUDA

Hi! Using the package flux I want to scatter the following using NNlib: …```julia using Flux NNlib.scatter(+, [SVector(1,1,1),SVector(1,1,1),SVector(1,1,1)], [3,1,2]) 3-element Vector{SVector{3, Int64}}: [1, 1, 1] [1, 1, 1] [1, 1, 1] ```` Which works no problem. If I change the mid array to CuArray, then it works again, but tested that it is slow for large arrays (60k): ```julia NNlib.scatter(+, CuArray([SVector(1,1,1),SVector(1,1,1),SVector(1,1,1)]), [3,1,2]) 3-element CuArray{SVector{3, Int64}, 1, CUDA.Mem.DeviceBuffer}: [1, 1, 1] [1, 1, 1] [1, 1, 1] ``` If I try to do everything on GPU: ```julia NNlib.scatter(+, CuArray([SVector(1,1,1),SVector(1,1,1),SVector(1,1,1)]), CuArray([3,1,2])) ERROR: InvalidIRError: compiling kernel #scatter_kernel!(typeof(+), CuDeviceVector{SVector{3, Int64}, 1}, CuDeviceVector{SVector{3, Int64}, 1}, CuDeviceVector{Int64, 1}) resulted in invalid LLVM IR Reason: unsupported dynamic function invocation (call to atomic_cas!) ``` Which I think is an error? More info: https://discourse.julialang.org/t/how-to-reduce-an-array/92945/14 Kind regards

I got an answer on there stating why it doesn’t work. Personally, I just thought “it should work” since StaticArrays is such a core part of Julia imo. One can get it to work doing some code like:

  for i = 1:3 #Length of static array element i.e. SVector{**3**,Float32}
        o       = i - 1
        V_dst   = @view reinterpret(eltype(eltype(dst)), dst)[begin+o:3:end]

        V_src_  = @view reinterpret(eltype(eltype(src)), src)[begin+o:3:end]
        V_src   = @view V_src_[1:SYSTEM.MaxValidIndex[]]

        NNlib.scatter!(OP1,V_dst,V_src, @view(SYSTEM.I[1:SYSTEM.MaxValidIndex[]]))
        NNlib.scatter!(OP2,V_dst,V_src, @view(SYSTEM.J[1:SYSTEM.MaxValidIndex[]]))
    end

But I do not think this is the best fix, since I need to call NNlib.scatter!, 6 times to do this and possible lose out on some speed.

Kind regards

ToucheSir · February 7, 2023, 10:40pm

As I mentioned on that issue thread, you can turn this into 1 or 2 scatter! cols by using a 2D array instead.

Ahmed_Salih · February 8, 2023, 1:23pm

Yes, but I really wish to stay using static arrays since it allows me for a relatively fast CPU execution as well. Do you happen to know if there is a way to “view” a vector of static arrays as a 2d matrix?

I only know how to get a single column out as shown above

Kind regards

Ahmed_Salih · February 8, 2023, 7:57pm

Got it to work, thank you