diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index fe2d8021..5bc92145 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -17,7 +17,7 @@ jobs: test: name: Julia ${{ matrix.version }} - ${{ matrix.os }} runs-on: ${{ matrix.os }} - timeout-minutes: 60 + timeout-minutes: 90 strategy: fail-fast: false matrix: diff --git a/lib/JLArrays/Project.toml b/lib/JLArrays/Project.toml index 698fb8eb..dff9d518 100644 --- a/lib/JLArrays/Project.toml +++ b/lib/JLArrays/Project.toml @@ -1,12 +1,13 @@ name = "JLArrays" uuid = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" -version = "0.4.2" +version = "0.5" authors = ["Tim Besard "] [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" @@ -14,7 +15,8 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" [compat] Adapt = "2.0, 3.0, 4.0" GPUArrays = "12.1" -KernelAbstractions = "0.9" +KernelAbstractions = "0.10" +KernelInterface = "0.4.2" LinearAlgebra = "1" Random = "1" SparseArrays = "1" diff --git a/lib/JLArrays/src/JLArrays.jl b/lib/JLArrays/src/JLArrays.jl index ff9d7c3e..997aaecd 100644 --- a/lib/JLArrays/src/JLArrays.jl +++ b/lib/JLArrays/src/JLArrays.jl @@ -16,11 +16,7 @@ using Random import GPUArrays: dense_array_type -import KernelAbstractions - -@static if isdefined(JLArrays.KernelAbstractions, :POCL) # KA v0.10 - import KernelAbstractions: POCL -end +import KernelAbstractions: KernelAbstractions, POCL # the size of an element in an array, including padding up to the element's alignment @static if VERSION < v"1.11" diff --git a/lib/JLArrays/src/JLKernels.jl b/lib/JLArrays/src/JLKernels.jl index ca4d84d0..44f68cbb 100644 --- a/lib/JLArrays/src/JLKernels.jl +++ b/lib/JLArrays/src/JLKernels.jl @@ -3,8 +3,10 @@ module JLKernels using ..JLArrays using ..JLArrays: jlconvert -import KernelAbstractions -import KernelAbstractions: Kernel, StaticSize, DynamicSize, partition, launch_config +import KernelAbstractions: KernelAbstractions, POCLBackend +import KernelInterface as KI + +import GPUArrays import Adapt @@ -13,73 +15,72 @@ import Adapt export JLBackend -const MAXTHREADS = 256 +struct JLBackend <: KI.Backend end -struct JLBackend <: KernelAbstractions.GPU - static::Bool - JLBackend(;static::Bool=false) = new(static) -end +KI.get_backend(a::JLA) where JLA <: JLArray = JLBackend() +KI.get_backend(a::JLA) where JLA <: Union{JLSparseMatrixCSC, JLSparseMatrixCSR, JLSparseVector} = JLBackend() -KernelAbstractions.get_backend(a::JLA) where JLA <: JLArray = JLBackend() -KernelAbstractions.get_backend(a::JLA) where JLA <: Union{JLSparseMatrixCSC, JLSparseMatrixCSR, JLSparseVector} = JLBackend() +KI.allocate(::JLBackend, ::Type{T}, dims::Tuple) where T = JLArray{T}(undef, dims) -function KernelAbstractions.mkcontext(kernel::Kernel{JLBackend}, I, _ndrange, iterspace, ::Dynamic) where Dynamic - return KernelAbstractions.CompilerMetadata{KernelAbstractions.ndrange(kernel), Dynamic}(I, _ndrange, iterspace) -end +KI.unsafe_free!(a::JLArray) = GPUArrays.unsafe_free!(a) + +## KernelInterface + +# JLArrays are executed by POCL: convert arguments to their POCL-compatible (CPU) +# counterparts, and forward compilation and launches to the POCL backend. KernelAbstractions +# kernels are launched through these methods too. -KernelAbstractions.allocate(::JLBackend, ::Type{T}, dims::Tuple) where T = JLArray{T}(undef, dims) - -@inline function launch_config(kernel::Kernel{JLBackend}, ndrange, workgroupsize) - if ndrange isa Integer - ndrange = (ndrange,) - end - if workgroupsize isa Integer - workgroupsize = (workgroupsize, ) - end - - if KernelAbstractions.workgroupsize(kernel) <: DynamicSize && workgroupsize === nothing - workgroupsize = (MAXTHREADS,) # Vectorization, 4x unrolling, minimal grain size - end - iterspace, dynamic = partition(kernel, ndrange, workgroupsize) - # partition checked that the ndrange's agreed - if KernelAbstractions.ndrange(kernel) <: StaticSize - ndrange = nothing - end - - return ndrange, workgroupsize, iterspace, dynamic +# POCL launches synchronously, and JLArrays live in host memory +KI.synchronize(::JLBackend) = nothing + +function KI.copyto!(::JLBackend, A, B) + length(A) == length(B) || + throw(ArgumentError("Arrays must match in length, got $(length(A)) and $(length(B))")) + Base.copyto!(A, B) + return A end -@static if isdefined(KernelAbstractions, :isgpu) # KA v0.9 - KernelAbstractions.isgpu(b::JLBackend) = false +KI.argconvert(::JLBackend, arg) = KI.argconvert(POCLBackend(), jlconvert(arg)) + +# the POCL kernel, and the callable it was compiled from, which has to be kept alive since +# its converted form only holds pointers to the arrays it captures +struct JLKernel{K, F} + kernel::K + f::F end -@static if !isdefined(KernelAbstractions, :POCL) # KA v0.9 - function convert_to_cpu(obj::Kernel{JLBackend, W, N, F}) where {W, N, F} - return Kernel{typeof(KernelAbstractions.CPU(; static = obj.backend.static)), W, N, F}(KernelAbstractions.CPU(; static = obj.backend.static), obj.f) - end -else - function convert_to_cpu(obj::Kernel{JLBackend, W, N, F}) where {W, N, F} - return Kernel{typeof(KernelAbstractions.POCLBackend()), W, N, F}(KernelAbstractions.POCLBackend(), obj.f) - end +function KI.kernel_function(backend::JLBackend, f::F, tt::TT = Tuple{}; kwargs...) where {F, TT} + kernel = KI.kernel_function(POCLBackend(), jlconvert(f), tt; kwargs...) + kern = JLKernel(kernel, f) + return KI.Kernel{JLBackend, typeof(kern)}(backend, kern) end -function (obj::Kernel{JLBackend})(args...; ndrange=nothing, workgroupsize=nothing) - ndrange, workgroupsize, _, _ = launch_config(obj, ndrange, workgroupsize) +function KI.launch(obj::KI.Kernel{JLBackend}, groups::Dims{3}, items::Dims{3}, args::Tuple; kwargs...) # Use `map` rather than `jlconvert.(args)` to skip the broadcast # machinery (broadcasted/materialize/ntuple) that would otherwise # be specialized per unique arg-tuple type. - device_args = map(jlconvert, args) - new_obj = convert_to_cpu(obj) - new_obj(device_args...; ndrange, workgroupsize) + f = obj.kern.f + GC.@preserve f args KI.launch(obj.kern.kernel, groups, items, map(jlconvert, args); kwargs...) end -Adapt.adapt_storage(::JLBackend, a::Array) = Adapt.adapt(JLArrays.JLArray, a) -Adapt.adapt_storage(::JLBackend, a::JLArrays.JLArray) = a - -@static if !isdefined(KernelAbstractions, :POCL) # KA v0.9 - Adapt.adapt_storage(::KernelAbstractions.CPU, a::JLArrays.JLArray) = convert(Array, a) -else - Adapt.adapt_storage(::KernelAbstractions.POCLBackend, a::JLArrays.JLArray) = convert(Array, a) +KI.max_work_group_size(kernel::KI.Kernel{JLBackend}) = KI.max_work_group_size(kernel.kern.kernel) +KI.launch_configuration(kernel::KI.Kernel{JLBackend}; kwargs...) = + KI.launch_configuration(kernel.kern.kernel; kwargs...) +for f in (:functional, + :max_work_group_size, :max_work_group_dims, :max_num_groups, :sub_group_size, + :multiprocessor_count, :supports_subgroups, :supports_linear_subgroups, + :supports_independent_subgroups, :supports_partitionable_subgroups, + :supports_float64, :supports_atomics) + @eval KI.$f(::JLBackend) = KI.$f(POCLBackend()) end +KI.supports_shuffle(::JLBackend, ::Type{T}) where {T} = KI.supports_shuffle(POCLBackend(), T) +KI.versioninfo(io::IO, ::JLBackend) = KI.versioninfo(io, POCLBackend()) + +# KernelAbstractions computes indices in `Int32` unless the back-end is POCL, where narrowing +# them keeps LLVM from vectorizing the loops over work-items. JLArray kernels run on POCL too. +KernelAbstractions.narrow_index_type(::JLBackend) = Int + +Adapt.adapt_storage(::JLBackend, a::Array) = Adapt.adapt(JLArray, a) +Adapt.adapt_storage(::JLBackend, a::JLArray) = a end diff --git a/lib/JLArrays/src/array.jl b/lib/JLArrays/src/array.jl index d3f708ae..be1bfc92 100644 --- a/lib/JLArrays/src/array.jl +++ b/lib/JLArrays/src/array.jl @@ -10,39 +10,40 @@ function check_eltype(T) end mutable struct JLArray{T, N} <: AbstractGPUArray{T, N} - data::DataRef{Vector{UInt8}} + # the buffer is a `Vector` of the element type the array was allocated with, which + # derived arrays (e.g. reinterpretations) share + data::DataRef{<:Vector} - offset::Int # offset of the data in the buffer, in bytes - - # arrays of isbits unions keep one selector byte per element past the data, like - # `Array` does: this is the offset of the array's first selector byte in the buffer - tagoffset::Int + # offset of the data in the buffer, in bytes. arrays of isbits unions are only derived + # as arrays of the same union, whose elements can take no bytes: for those, in elements + offset::Int dims::Dims{N} # allocating constructor function JLArray{T,N}(::UndefInitializer, dims::Dims{N}) where {T,N} check_eltype(T) - datasize = prod(dims) * aligned_sizeof(T) - bufsize = Base.isbitsunion(T) ? datasize + prod(dims) : datasize + len = prod(dims) - ref = GPUArrays.cached_alloc((JLArray, bufsize)) do - data = Vector{UInt8}(undef, bufsize) + # buffers of the same size are interchangeable, except for those of isbits unions, + # whose selector bytes Julia keeps + key = Base.isbitsunion(T) ? (JLArray, T, len) : (JLArray, len * aligned_sizeof(T)) + ref = GPUArrays.cached_alloc(key) do + data = Vector{T}(undef, len) DataRef(data) do data resize!(data, 0) end end - obj = new{T, N}(ref, 0, datasize, dims) + obj = new{T, N}(ref, 0, dims) finalizer(unsafe_free!, obj) return obj end # low-level constructor for wrapping existing data - function JLArray{T,N}(ref::DataRef{Vector{UInt8}}, dims::Dims{N}; offset::Int=0, - tagoffset::Int=offset + prod(dims) * aligned_sizeof(T)) where {T,N} + function JLArray{T,N}(ref::DataRef{<:Vector}, dims::Dims{N}; offset::Int=0) where {T,N} check_eltype(T) - obj = new{T,N}(ref, offset, tagoffset, dims) + obj = new{T,N}(ref, offset, dims) finalizer(unsafe_free!, obj) return obj end @@ -50,37 +51,38 @@ end GPUArrays.storage(a::JLArray) = a.data -GPUArrays.memory_location(a::JLArray) = (UInt(pointer(a.data[])), a.offset) +function GPUArrays.memory_location(a::JLArray) + base = UInt(buffer_pointer(a)) + return (base, Int(UInt(pointer(a)) - base)) +end GPUArrays.dense_array_type(a::JLArray{T, N}) where {T, N} = JLArray{T, N} GPUArrays.dense_array_type(::Type{JLArray{T, N}}) where {T, N} = JLArray{T, N} GPUArrays.dense_vector_type(a::JLArray{T, N}) where {T, N} = JLArray{T, 1} GPUArrays.dense_vector_type(::Type{JLArray{T, N}}) where {T, N} = JLArray{T, 1} -# conversion of untyped data to a typed Array +# the data as a host array: an `Array`, or for a derived array of isbits unions (which an +# `Array` can't wrap) a view of the buffer function typed_data(x::JLArray{T}) where {T} - unsafe_wrap(Array, pointer(x), x.dims) -end - -# the selector byte of element `i` of an array of isbits unions -typetagdata(a::JLArray, i=1) = pointer(a.data[]) + a.tagoffset + i - 1 -@static if VERSION >= v"1.11.0-DEV.753" - function typetagdata(a::Array, i=1) - # for memory of isbits unions, the reference holds the offset of the data in elements - ptr_or_offset = Int(a.ref.ptr_or_offset) - @ccall(jl_genericmemory_typetagdata(a.ref.mem::Any)::Ptr{UInt8}) + ptr_or_offset + i - 1 + if Base.isbitsunion(T) + buf = x.data[]::Vector{T} + len = prod(x.dims) + if x.offset == 0 && len == length(buf) + reshape(buf, x.dims) + else + reshape(view(buf, x.offset+1:x.offset+len), x.dims) + end + else + unsafe_wrap(Array, pointer(x), x.dims) end -else - typetagdata(a::Array, i=1) = ccall(:jl_array_typetagdata, Ptr{UInt8}, (Any,), a) + i - 1 end function GPUArrays.derive(::Type{T}, a::JLArray, dims::Dims{N}, offset::Int) where {T,N} ref = copy(a.data) # arrays of isbits unions are only derived as arrays of the same union - # (`reinterpret` rejects them), so their selector bytes advance by element - tagoffset = a.tagoffset + offset - offset = a.offset + offset * aligned_sizeof(T) - JLArray{T,N}(ref, dims; offset, tagoffset) + # (`reinterpret` rejects them), and are offset by element + offset = a.offset + (Base.isbitsunion(T) ? offset : offset * aligned_sizeof(T)) + JLArray{T,N}(ref, dims; offset) end @@ -136,9 +138,17 @@ StridedJLVector{T} = StridedJLArray{T,1} StridedJLMatrix{T} = StridedJLArray{T,2} StridedJLVecOrMat{T} = Union{StridedJLVector{T}, StridedJLMatrix{T}} +# (the buffer's element type isn't known) +buffer_pointer(x::JLArray) = convert(Ptr{UInt8}, pointer(x.data[]))::Ptr{UInt8} + # Pointer access is only available for callers that explicitly want a pointer. -Base.pointer(x::JLArray{T}) where {T} = - convert(Ptr{T}, pointer(x.data[])) + x.offset +function Base.pointer(x::JLArray{T}) where {T} + if Base.isbitsunion(T) + convert(Ptr{T}, pointer(x.data[]::Vector{T}, x.offset + 1)) + else + convert(Ptr{T}, buffer_pointer(x)) + x.offset + end +end @inline function Base.pointer(x::JLArray{T}, i::Integer) where T pointer(x) + (i - 1) * aligned_sizeof(T) end @@ -191,14 +201,28 @@ Adapt.adapt_storage(::Type{<:JLArray{T}}, xs::AbstractArray) where {T} = # adapt back to the CPU Adapt.adapt_storage(::Type{Array}, xs::JLArray) = convert(Array, xs) -# adapt for the device -@static if !isdefined(JLArrays.KernelAbstractions, :POCL) # KA v0.9 - Adapt.adapt_storage(::Adaptor, x::JLArray{T,N}) where {T,N} = - JLDeviceArray{T,N}(x.data[], x.offset, x.tagoffset, x.dims) +# the selector byte of element `i` of an `Array` of isbits unions +@static if VERSION >= v"1.11.0-DEV.753" + function typetagdata(a::Array, i=1) + # for memory of isbits unions, the reference holds the offset of the data in elements + ptr_or_offset = Int(a.ref.ptr_or_offset) + @ccall(jl_genericmemory_typetagdata(a.ref.mem::Any)::Ptr{UInt8}) + ptr_or_offset + i - 1 + end else - function Adapt.adapt_storage(::Adaptor, x::JLArray{T,N}) where {T,N} - arr = typed_data(x) - Adapt.adapt_storage(POCL.KernelAdaptor([pointer(arr)]), arr) + typetagdata(a::Array, i=1) = ccall(:jl_array_typetagdata, Ptr{UInt8}, (Any,), a) + i - 1 +end + +# adapt for the device, where JLArrays are POCL arrays (constructed directly, rather than +# by adapting a wrapping `Array`, to avoid allocating on every launch) +function Adapt.adapt_storage(::Adaptor, x::JLArray{T,N}) where {T,N} + AS = POCL.AS.CrossWorkgroup + ptr = reinterpret(Core.LLVMPtr{T,AS}, pointer(x)) + if Base.isbitsunion(T) + # POCL arrays of isbits unions read their selector bytes `maxsize` bytes past the data + tags = typetagdata(x.data[]::Vector{T}, x.offset + 1) + POCL.CLDeviceArray{T,N,AS}(x.dims, ptr, Int(UInt(tags) - UInt(pointer(x)))) + else + POCL.CLDeviceArray{T,N,AS}(x.dims, ptr) end end @@ -209,16 +233,20 @@ Base.convert(::Type{T}, x::T) where T <: JLArray = x ## memory operations -# copy the elements, and the selector bytes of isbits unions +host_data(x::Array) = x +host_data(x::JLArray) = typed_data(x) + function unsafe_copy_elements!(dest, d_offset, source, s_offset, amount) T = eltype(dest) GC.@preserve dest source begin - Base.unsafe_copyto!(convert(Ptr{UInt8}, pointer(dest, d_offset)), - convert(Ptr{UInt8}, pointer(source, s_offset)), - amount * aligned_sizeof(T)) if Base.isbitsunion(T) - Base.unsafe_copyto!(typetagdata(dest, d_offset), typetagdata(source, s_offset), - amount) + # (which also copies the selector bytes) + copyto!(host_data(dest), d_offset, host_data(source), s_offset, amount) + else + # (the data of reinterpreted arrays needn't be aligned for an `Array`) + Base.unsafe_copyto!(convert(Ptr{UInt8}, pointer(dest, d_offset)), + convert(Ptr{UInt8}, pointer(source, s_offset)), + amount * aligned_sizeof(T)) end end return dest @@ -254,7 +282,6 @@ function Base.resize!(a::DenseJLVector{T}, nl::Integer) where {T} a.data = copy(b.data) a.offset = b.offset - a.tagoffset = b.tagoffset a.dims = b.dims return a end diff --git a/lib/JLArrays/src/device/array.jl b/lib/JLArrays/src/device/array.jl index f6fa6597..c5bcd210 100644 --- a/lib/JLArrays/src/device/array.jl +++ b/lib/JLArrays/src/device/array.jl @@ -3,84 +3,6 @@ struct Adaptor end jlconvert(arg) = adapt(Adaptor(), arg) -# FIXME: add Ref to Adapt.jl (but make sure it doesn't cause ambiguities with CUDAnative's) -struct JlRefValue{T} <: Ref{T} - x::T -end -Base.getindex(r::JlRefValue) = r.x -Adapt.adapt_structure(to::Adaptor, r::Base.RefValue) = JlRefValue(adapt(to, r[])) - -## executed on-device - -# array type -@static if !isdefined(JLArrays.KernelAbstractions, :POCL) # KA v0.9 - struct JLDeviceArray{T, N} <: AbstractDeviceArray{T, N} - data::Vector{UInt8} - offset::Int - tagoffset::Int # see `JLArray` - dims::Dims{N} - end - - Base.elsize(::Type{<:JLDeviceArray{T}}) where {T} = aligned_sizeof(T) - - Base.size(x::JLDeviceArray) = x.dims - Base.sizeof(x::JLDeviceArray) = Base.elsize(x) * length(x) - - Base.unsafe_convert(::Type{Ptr{T}}, x::JLDeviceArray{T}) where {T} = - convert(Ptr{T}, pointer(x.data)) + x.offset - - # conversion of untyped data to a typed Array - function typed_data(x::JLDeviceArray{T}) where {T} - unsafe_wrap(Array, pointer(x), x.dims) - end - - @inline function Base.getindex(A::JLDeviceArray{T}, index::Integer) where {T} - if Base.isbitsunion(T) - @boundscheck checkbounds(A, index) - union_load(A, index) - else - getindex(typed_data(A), index) - end - end - @inline function Base.setindex!(A::JLDeviceArray{T}, x, index::Integer) where {T} - if Base.isbitsunion(T) - @boundscheck checkbounds(A, index) - union_store!(A, convert(T, x)::T, index) - return A - else - setindex!(typed_data(A), x, index) - end - end - - # elements of isbits unions are stored as the member type that their selector byte - # indexes in `Base.uniontypes(T)` - - @generated function union_load(A::JLDeviceArray{T}, index::Integer) where {T} - ex = :(error("invalid selector byte")) - for (sel, U) in Iterators.reverse(enumerate(Base.uniontypes(T))) - val = Base.issingletontype(U) ? U.instance : :(unsafe_load(convert(Ptr{$U}, ptr))) - ex = :(selector == $(UInt8(sel - 1)) ? $val : $ex) - end - quote - data = A.data - selector = @inbounds data[A.tagoffset + index] - GC.@preserve data begin - ptr = pointer(data) + A.offset + (index - 1) * $(Base.aligned_sizeof(T)) - $ex - end - end - end - - @generated function union_store!(A::JLDeviceArray{T}, x, index::Integer) where {T} - sel = findfirst(==(x), Base.uniontypes(T)) - quote - data = A.data - @inbounds data[A.tagoffset + index] = $(UInt8(sel - 1)) - GC.@preserve data begin - ptr = pointer(data) + A.offset + (index - 1) * $(Base.aligned_sizeof(T)) - $(Base.issingletontype(x) ? nothing : :(unsafe_store!(convert(Ptr{$x}, ptr), x))) - end - return - end - end -end +# convert what a `Ref` holds, but leave the `Ref` itself to POCL, which also handles `Ref`s +# of types (as passed by `convert.(T, A)`) +Adapt.adapt_structure(to::Adaptor, r::Base.RefValue) = Ref(adapt(to, r[])) diff --git a/lib/JLArrays/src/mapreduce.jl b/lib/JLArrays/src/mapreduce.jl index c1725e1b..5d5ccff5 100644 --- a/lib/JLArrays/src/mapreduce.jl +++ b/lib/JLArrays/src/mapreduce.jl @@ -1,13 +1,7 @@ # reductions struct ArrayNoCopy end -# (an `Array` can't wrap the storage of isbits unions, whose selector bytes are elsewhere) -Adapt.adapt_storage(::ArrayNoCopy, x::JLArray{T}) where {T} = - Base.isbitsunion(T) ? jlconvert(x) : typed_data(x) -@static if !isdefined(KernelAbstractions, :POCL) # KA v0.9 - # (AcceleratedKernels picks its host algorithms by the back-end) - KernelAbstractions.get_backend(::JLDeviceArray) = KernelAbstractions.CPU() -end +Adapt.adapt_storage(::ArrayNoCopy, x::JLArray) = typed_data(x) # GPUArrays reduces with AcceleratedKernels, whose kernels do not run on JLArrays' back-end. The # reference implementation runs AcceleratedKernels' host algorithms on the arrays' storage diff --git a/lib/JLArrays/src/sparse.jl b/lib/JLArrays/src/sparse.jl index 032bd422..2acca591 100644 --- a/lib/JLArrays/src/sparse.jl +++ b/lib/JLArrays/src/sparse.jl @@ -153,6 +153,7 @@ function JLSparseMatrixCSC(xs::SparseMatrixCSC{Tv, Ti}) where {Ti, Tv} return JLSparseMatrixCSC{Tv, Ti}(colPtr, rowVal, nzVal, (xs.m, xs.n)) end JLSparseMatrixCSC(xs::SparseVector) = JLSparseMatrixCSC(SparseMatrixCSC(xs)) + Base.length(x::JLSparseMatrixCSC) = prod(x.dims) Base.size(x::JLSparseMatrixCSC) = x.dims @@ -209,9 +210,15 @@ end ## adapt for the device -Adapt.adapt_structure(to::Adaptor, x::JLSparseMatrixCSC{Tv,Ti}) where {Tv,Ti} = -GPUSparseDeviceMatrixCSC{Tv,Ti,JLDeviceArray{Ti, 1}, JLDeviceArray{Tv, 1}, AS.Generic}(adapt(to, x.colPtr), adapt(to, x.rowVal), adapt(to, x.nzVal), x.dims, x.nnz) -Adapt.adapt_structure(to::Adaptor, x::JLSparseMatrixCSR{Tv,Ti}) where {Tv,Ti} = -GPUSparseDeviceMatrixCSR{Tv,Ti,JLDeviceArray{Ti, 1}, JLDeviceArray{Tv, 1}, AS.Generic}(adapt(to, x.rowPtr), adapt(to, x.colVal), adapt(to, x.nzVal), x.dims, x.nnz) -Adapt.adapt_structure(to::Adaptor, x::JLSparseVector{Tv,Ti}) where {Tv,Ti} = -GPUSparseDeviceVector{Tv,Ti,JLDeviceArray{Ti, 1}, JLDeviceArray{Tv, 1}, AS.Generic}(adapt(to, x.iPtr), adapt(to, x.nzVal), x.len, x.nnz) +function Adapt.adapt_structure(to::Adaptor, x::JLSparseMatrixCSC{Tv,Ti}) where {Tv,Ti} + colPtr, rowVal, nzVal = adapt(to, x.colPtr), adapt(to, x.rowVal), adapt(to, x.nzVal) + GPUSparseDeviceMatrixCSC{Tv,Ti,typeof(colPtr),typeof(nzVal),AS.Generic}(colPtr, rowVal, nzVal, x.dims, x.nnz) +end +function Adapt.adapt_structure(to::Adaptor, x::JLSparseMatrixCSR{Tv,Ti}) where {Tv,Ti} + rowPtr, colVal, nzVal = adapt(to, x.rowPtr), adapt(to, x.colVal), adapt(to, x.nzVal) + GPUSparseDeviceMatrixCSR{Tv,Ti,typeof(rowPtr),typeof(nzVal),AS.Generic}(rowPtr, colVal, nzVal, x.dims, x.nnz) +end +function Adapt.adapt_structure(to::Adaptor, x::JLSparseVector{Tv,Ti}) where {Tv,Ti} + iPtr, nzVal = adapt(to, x.iPtr), adapt(to, x.nzVal) + GPUSparseDeviceVector{Tv,Ti,typeof(iPtr),typeof(nzVal),AS.Generic}(iPtr, nzVal, x.len, x.nnz) +end diff --git a/test/Project.toml b/test/Project.toml index 7e816ed5..ab94d87b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,13 +1,17 @@ [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/runtests.jl b/test/runtests.jl index c2b9ea47..b43820e8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,6 +11,10 @@ const init_worker_code = quote TestSuite.sparse_types(::Type{<:JLArray}) = (JLSparseVector, JLSparseMatrixCSC, JLSparseMatrixCSR) TestSuite.sparse_types(::Type{<:Array}) = (SparseVector, SparseMatrixCSC) + + # JLArray kernels run on POCL, whose SPIR-V back-end lacks 128-bit integers + TestSuite.supported_eltypes(::Type{<:JLArray}) = + filter(!=(Int128), TestSuite.supported_eltypes()) end const init_code = quote @@ -26,6 +30,32 @@ for AT in (:JLArray, :Array), name in keys(TestSuite.tests) testsuite["$(AT)/$name"] = :(TestSuite.tests[$name]($AT)) end +testsuite["JLArray/kernelabstractions"] = quote + import KernelAbstractions + import KernelAbstractions as KA + using JLArrays.JLKernels + + include(joinpath(dirname(pathof(KernelAbstractions)), "..", "test", "testsuite.jl")) + + # JLArrays don't convert dense arrays to sparse ones, and the other two tests only run + # kernels on KernelAbstractions' own CPU back end + skip_tests = Set(["sparse", "CPU synchronization", "fallback test: callable types"]) + # JLArray kernels are compiled by POCL, like the CPU back end's, so check them the same way + Testsuite.testsuite(JLBackend, "CPU", KernelAbstractions.POCL, JLArray, + KernelAbstractions.POCL.CLDeviceArray; skip_tests) +end + +testsuite["JLArray/kernelinterface"] = quote + import KernelInterface + import KernelInterface as KI + using JLArrays.JLKernels + import Adapt + + include(joinpath(dirname(pathof(KernelInterface)), "..", "test", "testsuite.jl")) + + Testsuite.testsuite(JLBackend(), JLArray) +end + # captured variables that get boxed (`Core.Box`) lose their type if isdefined(Test, :detect_closure_boxes) testsuite["closure boxes"] = quote diff --git a/test/testsuite/broadcasting.jl b/test/testsuite/broadcasting.jl index 19e8b192..28d29412 100644 --- a/test/testsuite/broadcasting.jl +++ b/test/testsuite/broadcasting.jl @@ -44,6 +44,9 @@ function broadcasting(AT, eltypes) gres .= test_idx.(gidx, Base.RefValue(gy)) cres .= test_idx.(cidx, Base.RefValue(cy)) @test Array(gres) == cres + + # a `Ref` of a type, as passed by `convert.(T, A)` + @test compare(a -> convert.(ET, a), AT, rand(ET, N)) end @testset "Tuple" begin