Scalar indexing error while input view of cuarray to LSTMCell

[yuehhua]

MWE:

julia> using Flux, CUDA

julia> CUDA.allowscalar(false)

julia> l = Flux.LSTMCell(3, 5) |> gpu
LSTMCell(3 => 5)

julia> A = rand(3, 10) |> gpu;

julia> B = selectdim(A, 2, [1, 3, 5, 7, 9]);

julia> l(l.state0, B)
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore are only permitted from the REPL for prototyping purposes.
If you did intend to index this array, annotate the caller with @allowscalar.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/lojQM/src/GPUArraysCore.jl:87
  [3] getindex(::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, ::Int64, ::Int64)
    @ GPUArrays ~/.julia/packages/GPUArrays/fqD8z/src/host/indexing.jl:9
  [4] _generic_matmatmul!(C::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, tA::Char, tB::Char, A::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, B::SubArray{Float32, 2, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, Tuple{Base.Slice{Base.OneTo{Int64}}, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, false}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra ~/.julia/juliaup/julia-1.8.1+0.x64/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:876
  [5] generic_matmatmul!(C::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, tA::Char, tB::Char, A::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, B::SubArray{Float32, 2, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, Tuple{Base.Slice{Base.OneTo{Int64}}, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, false}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra ~/.julia/juliaup/julia-1.8.1+0.x64/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:844
  [6] mul!
    @ ~/.julia/juliaup/julia-1.8.1+0.x64/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:303 [inlined]
  [7] mul!
    @ ~/.julia/juliaup/julia-1.8.1+0.x64/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:276 [inlined]
  [8] *
    @ ~/.julia/juliaup/julia-1.8.1+0.x64/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:141 [inlined]
  [9] (::Flux.LSTMCell{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}})(::Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}, x::SubArray{Float32, 2, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, Tuple{Base.Slice{Base.OneTo{Int64}}, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, false})
    @ Flux ~/.julia/packages/Flux/EXOFx/src/layers/recurrent.jl:298
 [10] top-level scope
    @ REPL[7]:1
 [11] top-level scope
    @ ~/.julia/packages/CUDA/DfvRa/src/initialization.jl:52

[ToucheSir]

This is because no BLAS dispatches are hit on the CUDA.jl side. The reason is that selectdim generates a view in a way that CUDA.jl doesn't understand, see JuliaGPU/CUDA.jl#1542. Even then though, I'm not sure if cuBLAS is able to handle a discontinuous view, so this may still hit a fallback.

[CarloLucibello]

after the redesign #2500 the MWE is

using Flux, CUDA
CUDA.allowscalar(false)
l = Flux.LSTMCell(3 => 5) |> gpu
A = rand(3, 10) |> gpu;
B = selectdim(A, 2, [1, 3, 5, 7, 9])
h = zeros(Float32, 5) |> gpu
c = zeros(Float32, 5) |> gpu
l(B, (h, c))