Support both KA 0.9 and KA 0.10

Author: christiangnrd

lib/JLArrays/Project.toml

@@ -1,7 +1,7 @@
 name = "JLArrays"
 uuid = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 authors = ["Tim Besard <[email protected]>"]
-version = "0.3.0"
+version = "0.2.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -12,6 +12,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 [compat]
 Adapt = "2.0, 3.0, 4.0"
 GPUArrays = "11.1"
-KernelAbstractions = "0.10"
+KernelAbstractions = "0.9, 0.10"
 Random = "1"
 julia = "1.8"

lib/JLArrays/src/JLArrays.jl

@@ -13,7 +13,11 @@ using GPUArrays
 using Adapt
 
 import KernelAbstractions
-import KernelAbstractions: Adapt, StaticArrays, Backend, Kernel, StaticSize, DynamicSize, partition, blocks, workitems, launch_config, POCL
+import KernelAbstractions: Adapt, StaticArrays, Backend, Kernel, StaticSize, DynamicSize, partition, blocks, workitems, launch_config
+
+@static if isdefined(JLArrays.KernelAbstractions, :POCL) # KA v0.10
+    import KernelAbstractions: POCL
+end
 
 
 #
@@ -40,8 +44,29 @@ 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*Base.elsize(x)
+
+    # 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
 
 #
 # Host abstractions
@@ -314,9 +339,14 @@ end
 
 ## GPUArrays interfaces
 
-function Adapt.adapt_storage(::Adaptor, x::JLArray{T,N}) where {T,N}
-    arr = typed_data(x)
-    Adapt.adapt_storage(POCL.KernelAdaptor([pointer(arr)]), arr)
+@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
 end
 
 function GPUArrays.mapreducedim!(f, op, R::AnyJLArray, A::Union{AbstractArray,Broadcast.Broadcasted};
@@ -358,8 +388,18 @@ KernelAbstractions.allocate(::JLBackend, ::Type{T}, dims::Tuple) where T = JLArr
     return ndrange, workgroupsize, iterspace, dynamic
 end
 
-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)
+@static if !isdefined(JLArrays.KernelAbstractions, :isgpu) # KA v0.9
+    KernelAbstractions.isgpu(b::JLBackend) = false
+end
+
+@static if !isdefined(JLArrays.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
 end
 
 function (obj::Kernel{JLBackend})(args...; ndrange=nothing, workgroupsize=nothing)
@@ -370,6 +410,11 @@ end
 
 Adapt.adapt_storage(::JLBackend, a::Array) = Adapt.adapt(JLArrays.JLArray, a)
 Adapt.adapt_storage(::JLBackend, a::JLArrays.JLArray) = a
-Adapt.adapt_storage(::KernelAbstractions.POCLBackend, a::JLArrays.JLArray) = convert(Array, a)
+
+@static if !isdefined(JLArrays.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)
+end
 
 end