[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
 
 module AS
 

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

@@ -171,15 +171,14 @@ 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.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)
-    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}
+    # POCL arrays of isbits unions read type tags after the data, which JLArrays don't store
+    Base.isbitsunion(T) &&
+        throw(ArgumentError("JLArrays with isbits-union elements cannot be used in kernels"))
+    AS = POCL.AS.CrossWorkgroup
+    POCL.CLDeviceArray{T,N,AS}(x.dims, reinterpret(Core.LLVMPtr{T,AS}, pointer(x)))
 end
 
 ## conversions

lib/JLArrays/src/device/array.jl

@@ -3,36 +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
-        dims::Dims{N}
-    end
-
-    Base.elsize(::Type{<:JLDeviceArray{T}}) where {T} = 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 Base.getindex(A::JLDeviceArray, index::Integer) = getindex(typed_data(A), index)
-    @inline Base.setindex!(A::JLDeviceArray, x, index::Integer) = setindex!(typed_data(A), x, index)
-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[]))

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

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"
 

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

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