2D Gridded Interpolation with 1xN coefficient matrix in GPU

[pvillacorta]

I have made the following function in order to create an interpolation from a matrix of coefficients d:

function interpolate(d::AbstractArray{T}) where {T<:Real}
    Ns, Nt = size(d)
    ITPType = Gridded(Linear())
    id = similar(d, Ns); copyto!(id, collect(range(oneunit(T), T(Ns), Ns)))
    t = similar(d, Nt);  copyto!(t, collect(range(zero(T), oneunit(T), Nt)))
    return Interpolations.GriddedInterpolation{eltype(d), 2, typeof(d), typeof(ITPType), typeof((id, t))}((id, t), d, ITPType)
end

This function returns a 2D GriddedInterpolation object and is valid for both CPU and GPU arrays:

julia> d = [0.0 1.0; 0.0 1.0]  # 2D CPU array
2×2 Matrix{Float64}:
 0.0  1.0
 0.0  1.0

julia> itp = KomaMRIBase.interpolate(d)
2×2 interpolate((::Vector{Float64},::Vector{Float64}), ::Matrix{Float64}, Gridded(Linear())) with element type Float64:
 0.0  1.0
 0.0  1.0

julia> d = [0.0 1.0]  # 1xNt CPU array
1×2 Matrix{Float64}:
 0.0  1.0

julia> itp = interpolate(d)
1×2 interpolate((::Vector{Float64},::Vector{Float64}), ::Matrix{Float64}, Gridded(Linear())) with element type Float64:
 0.0  1.0

julia> d = CuArray([0.0 1.0; 0.0 1.0])  # 2D GPU Array
2×2 CuArray{Float64, 2, CUDA.DeviceMemory}:
 0.0  1.0
 0.0  1.0

julia> itp = interpolate(d)
2×2 interpolate((::CuArray{Float64, 1, CUDA.DeviceMemory},::CuArray{Float64, 1, CUDA.DeviceMemory}), ::CuArray{Float64, 2, CUDA.DeviceMemory}, Gridded(Linear())) with element type Float64:
 0.0  1.0
 0.0  1.0

The problem comes only when these two things happen together:

• Ns = 1 (d is 1xNt)
• d is a GPU array

julia> d = CuArray([0.0 1.0])
1×2 CuArray{Float64, 2, CUDA.DeviceMemory}:
 0.0  1.0

julia> itp = interpolate(d)
1×2 interpolate((::CuArray{Float64, 1, CUDA.DeviceMemory},::CuArray{Float64, 1, CUDA.DeviceMemory}), ::CuArray{Float64, 2, CUDA.DeviceMemory}, Gridded(Linear())) with element type Float64:
Error showing value of type Interpolations.GriddedInterpolation{Float64, 2, CuArray{Float64, 2, CUDA.DeviceMemory}, Interpolations.Gridded{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{CuArray{Float64, 1, CUDA.DeviceMemory}, CuArray{Float64, 1, CUDA.DeviceMemory}}}:
ERROR: CUDA error: invalid argument (code 1, ERROR_INVALID_VALUE)
Stacktrace:
  [1] throw_api_error(res::CUDA.cudaError_enum)
    @ CUDA ~/.julia/packages/CUDA/Tl08O/lib/cudadrv/libcuda.jl:30
  [2] check
    @ ~/.julia/packages/CUDA/Tl08O/lib/cudadrv/libcuda.jl:37 [inlined]
  [3] cuMemcpyDtoHAsync_v2
    @ ~/.julia/packages/CUDA/Tl08O/lib/utils/call.jl:34 [inlined]
  [4] unsafe_copyto!(dst::Ptr{Float64}, src::CuPtr{Float64}, N::Int64; stream::CuStream, async::Bool)
    @ CUDA ~/.julia/packages/CUDA/Tl08O/lib/cudadrv/memory.jl:414
  [5] unsafe_copyto!
    @ ~/.julia/packages/CUDA/Tl08O/lib/cudadrv/memory.jl:411 [inlined]
  [6] (::CUDA.var"#1127#1128"{Float64, Vector{Float64}, Int64, CuArray{Float64, 1, CUDA.DeviceMemory}, Int64, Int64})()
    @ CUDA ~/.julia/packages/CUDA/Tl08O/src/array.jl:558
  [7] #context!#990
    @ ~/.julia/packages/CUDA/Tl08O/lib/cudadrv/state.jl:168 [inlined]
  [8] context!
    @ ~/.julia/packages/CUDA/Tl08O/lib/cudadrv/state.jl:163 [inlined]
  [9] unsafe_copyto!(dest::Vector{Float64}, doffs::Int64, src::CuArray{Float64, 1, CUDA.DeviceMemory}, soffs::Int64, n::Int64)
    @ CUDA ~/.julia/packages/CUDA/Tl08O/src/array.jl:550
 [10] copyto!
    @ ~/.julia/packages/CUDA/Tl08O/src/array.jl:503 [inlined]
 [11] getindex
    @ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:52 [inlined]
 [12] weightedindex(fs::Tuple{typeof(Interpolations.value_weights)}, deg::Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}, knotvec::CuArray{Float64, 1, CUDA.DeviceMemory}, x::Float64, iclamp::Int64)
    @ Interpolations ~/.julia/packages/Interpolations/91PhN/src/gridded/indexing.jl:70
 [13] weightedindex_parts2
    @ ~/.julia/packages/Interpolations/91PhN/src/gridded/indexing.jl:62 [inlined]
 [14] weightedindex_parts
    @ ~/.julia/packages/Interpolations/91PhN/src/gridded/indexing.jl:56 [inlined]
 [15] map3argf
    @ ~/.julia/packages/Interpolations/91PhN/src/b-splines/indexing.jl:70 [inlined]
 [16] weightedindexes
    @ ~/.julia/packages/Interpolations/91PhN/src/b-splines/indexing.jl:66 [inlined]
 [17] GriddedInterpolation
    @ ~/.julia/packages/Interpolations/91PhN/src/gridded/indexing.jl:4 [inlined]
 [18] (::Interpolations.var"#42#43"{Interpolations.GriddedInterpolation{Float64, 2, CuArray{Float64, 2, CUDA.DeviceMemory}, Interpolations.Gridded{Interpolations.Linear{…}}, Tuple{CuArray{…}, CuArray{…}}}})(x::Tuple{Float64, Float64})
    @ Interpolations ./none:0
 [19] iterate
    @ ./generator.jl:47 [inlined]
 [20] collect(itr::Base.Generator{Base.Iterators.ProductIterator{Tuple{CuArray{…}, CuArray{…}}}, Interpolations.var"#42#43"{Interpolations.GriddedInterpolation{Float64, 2, CuArray{…}, Interpolations.Gridded{…}, Tuple{…}}}})
    @ Base ./array.jl:834
 [21] show_ranged(io::IOContext{Base.TTY}, X::Interpolations.GriddedInterpolation{Float64, 2, CuArray{…}, Interpolations.Gridded{…}, Tuple{…}}, knots::Tuple{CuArray{…}, CuArray{…}})
    @ Interpolations ~/.julia/packages/Interpolations/91PhN/src/io.jl:107
 [22] show(io::IOContext{Base.TTY}, ::MIME{Symbol("text/plain")}, X::Interpolations.GriddedInterpolation{Float64, 2, CuArray{…}, Interpolations.Gridded{…}, Tuple{…}})
    @ Interpolations ~/.julia/packages/Interpolations/91PhN/src/io.jl:117
 [23] (::REPL.var"#55#56"{REPL.REPLDisplay{REPL.LineEditREPL}, MIME{Symbol("text/plain")}, Base.RefValue{Any}})(io::Any)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:273
 [24] with_repl_linfo(f::Any, repl::REPL.LineEditREPL)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:569
 [25] display(d::REPL.REPLDisplay, mime::MIME{Symbol("text/plain")}, x::Any)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:259
 [26] display(d::REPL.REPLDisplay, x::Any)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:278
 [27] display(x::Any)
    @ Base.Multimedia ./multimedia.jl:340
 [28] #invokelatest#2
    @ ./essentials.jl:892 [inlined]
 [29] invokelatest
    @ ./essentials.jl:889 [inlined]
 [30] print_response(errio::IO, response::Any, show_value::Bool, have_color::Bool, specialdisplay::Union{Nothing, AbstractDisplay})
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:315
 [31] (::REPL.var"#57#58"{REPL.LineEditREPL, Pair{Any, Bool}, Bool, Bool})(io::Any)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:284
 [32] with_repl_linfo(f::Any, repl::REPL.LineEditREPL)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:569
 [33] print_response(repl::REPL.AbstractREPL, response::Any, show_value::Bool, have_color::Bool)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:282
 [34] (::REPL.var"#do_respond#80"{Bool, Bool, REPL.var"#93#103"{REPL.LineEditREPL, REPL.REPLHistoryProvider}, REPL.LineEditREPL, REPL.LineEdit.Prompt})(s::REPL.LineEdit.MIState, buf::Any, ok::Bool)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:911
 [35] (::VSCodeServer.var"#103#106"{REPL.var"#do_respond#80"{Bool, Bool, REPL.var"#93#103"{REPL.LineEditREPL, REPL.REPLHistoryProvider}, REPL.LineEditREPL, REPL.LineEdit.Prompt}})(mi::REPL.LineEdit.MIState, buf::IOBuffer, ok::Bool)
    @ VSCodeServer ~/.vscode/extensions/julialang.language-julia-1.79.2/scripts/packages/VSCodeServer/src/repl.jl:122
 [36] #invokelatest#2
    @ ./essentials.jl:892 [inlined]
 [37] invokelatest
    @ ./essentials.jl:889 [inlined]
 [38] run_interface(terminal::REPL.Terminals.TextTerminal, m::REPL.LineEdit.ModalInterface, s::REPL.LineEdit.MIState)
    @ REPL.LineEdit ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/LineEdit.jl:2656
 [39] run_frontend(repl::REPL.LineEditREPL, backend::REPL.REPLBackendRef)
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:1312
 [40] (::REPL.var"#62#68"{REPL.LineEditREPL, REPL.REPLBackendRef})()
    @ REPL ~/.julia/juliaup/julia-1.10.5+0.x64.linux.gnu/share/julia/stdlib/v1.10/REPL/src/REPL.jl:386
Some type information was truncated. Use `show(err)` to see complete types.

This error happens when displaying. If I put a semicolon behind itp = interpolate(d), it does not appear, but when calling to the interpolation function:

julia> t = CuArray([0.0003 0.0003125 0.000325])
1×3 CuArray{Float64, 2, CUDA.DeviceMemory}:
 0.0003  0.0003125  0.000325

julia> itp.(CuArray([1.0]), t)
1×3 CuArray{Float64, 2, CUDA.DeviceMemory}:
 NaN  NaN  NaN

It returns an array of NaNs. When I tried to reproduce this error several times, sometimes it occurred and sometimes it did not.
I guess this is somehow related with #395 and with the fact that I am bypassing the GriddedInterpolation constructor, where check_gridded is called:

function GriddedInterpolation(::Type{TWeights}, knots::NTuple{N,GridIndex}, A::AbstractArray{TCoefs,N}, it::IT) where {N,TCoefs,TWeights<:Real,IT<:DimSpec{Gridded},pad}
    ...
    check_gridded(it, knots, axes(A))
    ...
    GriddedInterpolation{T,N,TCoefs,IT,typeof(knots)}(knots, A, it)
end

@inline function check_gridded(itpflag, knots, axs)
    flag, ax1, k1 = getfirst(itpflag), axs[1], knots[1]
   ...
    length(k1) == 1 && error("dimensions of length 1 not yet supported")  # FIXME <-- This is what I am bypassing
    ...
end

By the moment, I am solving it dispatching:

function GriddedInterpolation(nodes, A, ITP)
    return Interpolations.GriddedInterpolation{eltype(A), length(nodes), typeof(A), typeof(ITP), typeof(nodes)}(nodes, A, ITP)
end

function interpolate(d::AbstractArray{T}, Ns::Val{1}) where {T<:Real}
    _, Nt = size(d)
    ITPType = Gridded(Linear())
    t = similar(d, Nt); copyto!(t, collect(range(zero(T), oneunit(T), Nt)))
    return GriddedInterpolation((t, ), d[:], ITPType)
end

function interpolate(d::AbstractArray{T}, Ns::Val) where {T<:Real}
    Ns, Nt = size(d)
    ITPType = Gridded(Linear())
    id = similar(d, Ns); copyto!(id, collect(range(oneunit(T), T(Ns), Ns)))
    t = similar(d, Nt);  copyto!(t, collect(range(zero(T), oneunit(T), Nt)))
    return GriddedInterpolation((id, t), d, ITPType)
end

and, depending on Ns, calling the interpolator with:

• itp.(t) (Ns = 1)
• itp.(id, t) (Ns > 1)

But this solution is unnecessary complex.
Sorry for the long message and for not being able to reproduce the error in a more "general" way.
Thank you