[JLArrays] Support KA 0.10

.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:

lib/JLArrays/Project.toml

@@ -1,20 +1,22 @@
 name = "JLArrays"
 uuid = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
-version = "0.4.2"
+version = "0.5"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [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"
 
 [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"

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"

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

lib/JLArrays/src/array.jl

@@ -10,77 +10,79 @@ 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
 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