Remove KA 0.9 compat and clean up generic_matmatmul implementation

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

Project.toml

@@ -10,6 +10,7 @@ AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
@@ -31,7 +32,8 @@ AcceleratedKernels = "0.5"
 Adapt = "4.7.3"
 GPUArraysCore = "= 0.2.1"
 JLD2 = "0.4, 0.5, 0.6"
-KernelAbstractions = "0.9.43, 0.10"
+KernelAbstractions = "0.10"
+KernelInterface = "0.4.1"
 LLVM = "3.9, 4, 5, 6, 7, 8, 9, 10"
 LinearAlgebra = "1"
 Printf = "1"

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
 
@@ -44,4 +40,7 @@ include("findall.jl")
 include("JLKernels.jl")
 import .JLKernels: JLBackend
 
+# code reflection: JLArray kernels are compiled by POCL, for the device-side argument types
+import .POCL: code_typed, code_warntype, code_llvm, code_native
+
 end

lib/JLArrays/src/JLKernels.jl

@@ -3,8 +3,8 @@ 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 Adapt
 
@@ -13,73 +13,65 @@ 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
+## 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.
+
+KI.functional(::JLBackend) = true
 
-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 (: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)
+
+Adapt.adapt_storage(::JLBackend, x) = Adapt.adapt(JLArray, x)
 
 end

lib/JLArrays/src/array.jl

@@ -172,14 +172,9 @@ Adapt.adapt_storage(::Type{<:JLArray{T}}, xs::AbstractArray) where {T} =
 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
+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
 
 ## conversions

lib/JLArrays/src/device/array.jl

@@ -9,30 +9,3 @@ struct JlRefValue{T} <: Ref{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

lib/JLArrays/src/sparse.jl

@@ -153,6 +153,11 @@ 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))
+
+# dense to sparse, on the host
+SparseArrays.sparse(x::JLVector) = JLSparseVector(sparse(Array(x)))
+SparseArrays.sparse(x::JLMatrix) = JLSparseMatrixCSC(sparse(Array(x)))
+
 Base.length(x::JLSparseMatrixCSC) = prod(x.dims)
 Base.size(x::JLSparseMatrixCSC) = x.dims
 
@@ -209,9 +214,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