[JLArrays] Support KernelAbstractions 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:

docs/src/index.md

@@ -14,6 +14,6 @@ that builds on GPUArrays.jl such as [CUDA](https://github.com/JuliaGPU/CUDA.jl),
 This documentation is meant for users who might wish to implement a version of GPUArrays for another GPU backend and will cover the features you will need
 to implement, the functionality you gain by doing so, and the test suite that is available
 to verify your implementation. GPUArrays.jl also provides a reference implementation of
-these interfaces on the CPU: The `JLArray` array type uses Julia's parallel programming
-functionality to simulate GPU execution, and will be used throughout this documentation to
+these interfaces on the CPU: The `JLArray` array type runs its kernels on the CPU through
+KernelAbstractions' POCL back-end, and will be used throughout this documentation to
 illustrate functionality.

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.5"
+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"
@@ -42,10 +38,6 @@ include("device/array.jl")
 include("array.jl")
 include("sparse.jl")
 include("broadcast.jl")
-include("mapreduce.jl")
-include("sorting.jl")
-include("accumulate.jl")
-include("findall.jl")
 
 # KernelAbstractions
 include("JLKernels.jl")

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

@@ -57,11 +57,6 @@ 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
-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"
@@ -191,14 +186,16 @@ 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)
-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)
+# 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 find their selector bytes `maxsize` bytes past the data
+        POCL.CLDeviceArray{T,N,AS}(x.dims, ptr, x.tagoffset - x.offset)
+    else
+        POCL.CLDeviceArray{T,N,AS}(x.dims, ptr)
     end
 end
 

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[]))