ElasticArrays.jl
ElasticArrays.jl copied to clipboard
CUDA support
I got an error, when trying to use ElasticArray as input data, for Flux model.
ElasticArray([1,2]) |> gpu |>(Dense(2,1)|>gpu)
ERROR: ArgumentError: cannot take the CPU address of a CuArray{Float32,2}
Stacktrace:
[1] unsafe_convert(::Type{Ptr{Float32}}, ::CuArray{Float32,2}) at D:\x\.julia\packages\CUDA\dZvbp\src\array.jl:237
[2] gemv!(::Char, ::Float32, ::CuArray{Float32,2}, ::ElasticArray{Float32,1,0,CuArray{Float32,1}}, ::Float32, ::ElasticArray{Float32,1,0,CuArray{Float32,1}}) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\blas.jl:626
[3] gemv!(::ElasticArray{Float32,1,0,CuArray{Float32,1}}, ::Char, ::CuArray{Float32,2}, ::ElasticArray{Float32,1,0,CuArray{Float32,1}}, ::Bool, ::Bool) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:470
[4] mul! at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:66 [inlined]
[5] mul! at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:208 [inlined]
[6] * at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:47 [inlined]
[7] (::Dense{typeof(identity),CuArray{Float32,2},CuArray{Float32,1}})(::ElasticArray{Int64,1,0,CuArray{Int64,1}}) at D:\x\.julia\packages\Flux\05b38\src\layers\basic.jl:123
Also
julia> using ElasticArrays, CUDA
julia> CUDA.allowscalar(false)
julia> cu(rand(2,2)) * cu(ElasticArray(rand(2,2)))
ERROR: Scalar indexing is disallowed.
I think fixing this would require custom methods for (at least) LinearAlgebra.mul!, Broadcast.copy(bc::Broadcasted{<:AbstractGPUArrayStyle, ...}) and Base.copyto!(dest::BroadcastGPUArray, bc::Broadcasted{Nothing, ...}).
Question: Do we want elastic arrays on the GPU? Or should adapt return a non-elastic GPU array?