add ka extension

Author: omlins

src/ParallelKernel/KernelAbstractionsExt/allocators.jl

@@ -0,0 +1,76 @@
+import ParallelStencil
+using ParallelStencil.ParallelKernel.Exceptions: @KeywordArgumentError
+
+const PK = ParallelStencil.ParallelKernel
+
+# Helper to map KernelAbstractions allocations to the runtime backend.
+function runtime_allocator_symbol(kind::Symbol, hardware::Union{Symbol,Nothing})
+    target_package, _, _ = PK.resolve_runtime_backend(PK.PKG_KERNELABSTRACTIONS, hardware)
+    suffix = PK.allocator_suffix_for(target_package)
+    if suffix === nothing
+        @KeywordArgumentError("$(PK.ERRMSG_UNSUPPORTED_PACKAGE) (obtained: $target_package).")
+    end
+    return PK.allocator_function_symbol(kind, suffix)
+end
+
+invoke_runtime_allocator(kind::Symbol, hardware::Union{Symbol,Nothing}, args...) =
+    getfield(PK, runtime_allocator_symbol(kind, hardware))(args...)
+
+
+## RUNTIME ALLOCATOR FUNCTIONS
+
+function PK.zeros_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Number}
+    return invoke_runtime_allocator(:zeros, hardware, T, blocklength, args...)
+end
+
+function PK.ones_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Number}
+    return invoke_runtime_allocator(:ones, hardware, T, blocklength, args...)
+end
+
+function PK.rand_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{Number,Enum}}
+    return invoke_runtime_allocator(:rand, hardware, T, blocklength, args...)
+end
+
+function PK.falses_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Bool}
+    return invoke_runtime_allocator(:falses, hardware, T, blocklength, args...)
+end
+
+function PK.trues_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Bool}
+    return invoke_runtime_allocator(:trues, hardware, T, blocklength, args...)
+end
+
+function PK.fill_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{Number,Enum}}
+    return invoke_runtime_allocator(:fill, hardware, T, blocklength, args...)
+end
+
+function PK.zeros_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{SArray,FieldArray}}
+    return invoke_runtime_allocator(:zeros, hardware, T, blocklength, args...)
+end
+
+function PK.ones_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{SArray,FieldArray}}
+    return invoke_runtime_allocator(:ones, hardware, T, blocklength, args...)
+end
+
+function PK.rand_kernelabstractions(::Type{T}, blocklength::Val{B}, dims; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{SArray,FieldArray},B}
+    return invoke_runtime_allocator(:rand, hardware, T, blocklength, dims)
+end
+
+function PK.rand_kernelabstractions(::Type{T}, blocklength, dims...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{SArray,FieldArray}}
+    return PK.rand_kernelabstractions(T, blocklength, dims; hardware=hardware)
+end
+
+function PK.falses_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{SArray,FieldArray}}
+    return invoke_runtime_allocator(:falses, hardware, T, blocklength, args...)
+end
+
+function PK.trues_kernelabstractions(::Type{T}, blocklength, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{SArray,FieldArray}}
+    return invoke_runtime_allocator(:trues, hardware, T, blocklength, args...)
+end
+
+function PK.fill_kernelabstractions(::Type{T}, blocklength::Val{B}, x, args...; hardware::Union{Symbol,Nothing}=nothing) where {T<:Union{SArray,FieldArray},B}
+    return invoke_runtime_allocator(:fill, hardware, T, blocklength, x, args...)
+end
+
+function PK.fill_kernelabstractions!(A, x; hardware::Union{Symbol,Nothing}=nothing)
+    return invoke_runtime_allocator(:fill!, hardware, A, x)
+end

src/ParallelKernel/KernelAbstractionsExt/defaults.jl

@@ -0,0 +1,19 @@
+const ERRMSG_KERNELABSTRACTIONSEXT_NOT_LOADED = "the KernelAbstractions extension was not loaded. Make sure to import KernelAbstractions before ParallelStencil."
+
+
+# shared.jl
+
+function get_priority_kastream end
+function get_kastream end
+function get_kernelabstractions_compute_capability end
+
+
+# allocators.jl
+
+zeros_kernelabstractions(arg...)  = @NotLoadedError(ERRMSG_KERNELABSTRACTIONSEXT_NOT_LOADED)
+ones_kernelabstractions(arg...)   = @NotLoadedError(ERRMSG_KERNELABSTRACTIONSEXT_NOT_LOADED)
+rand_kernelabstractions(arg...)   = @NotLoadedError(ERRMSG_KERNELABSTRACTIONSEXT_NOT_LOADED)
+falses_kernelabstractions(arg...) = @NotLoadedError(ERRMSG_KERNELABSTRACTIONSEXT_NOT_LOADED)
+trues_kernelabstractions(arg...)  = @NotLoadedError(ERRMSG_KERNELABSTRACTIONSEXT_NOT_LOADED)
+fill_kernelabstractions(arg...)   = @NotLoadedError(ERRMSG_KERNELABSTRACTIONSEXT_NOT_LOADED)
+fill_kernelabstractions!(arg...)  = @NotLoadedError(ERRMSG_KERNELABSTRACTIONSEXT_NOT_LOADED)

src/ParallelKernel/KernelAbstractionsExt/shared.jl

@@ -0,0 +1,58 @@
+import ParallelStencil
+import KernelAbstractions
+using ParallelStencil.ParallelKernel: PKG_KERNELABSTRACTIONS, PKG_CUDA, PKG_AMDGPU, PKG_METAL, PKG_THREADS, PKG_POLYESTER,
+    resolve_runtime_backend,
+    get_priority_custream, get_custream, get_priority_rocstream, get_rocstream,
+    get_priority_metalstream, get_metalstream,
+    get_cuda_compute_capability, get_amdgpu_compute_capability, get_metal_compute_capability, get_cpu_compute_capability,
+    ERRMSG_UNSUPPORTED_PACKAGE
+using ParallelStencil.ParallelKernel.Exceptions: @ArgumentError
+
+const PK = ParallelStencil.ParallelKernel
+
+
+## FUNCTIONS TO CHECK EXTENSION SUPPORT
+
+PK.is_loaded(::Val{:ParallelStencil_KernelAbstractionsExt}) = true
+
+
+## FUNCTIONS TO GET OR MANAGE KERNELABSTRACTIONS STREAMS
+
+function dispatch_kastream(kind::Symbol, id::Integer, hardware::Union{Symbol,Nothing})
+    target_package, symbol, _ = resolve_runtime_backend(PKG_KERNELABSTRACTIONS, hardware)
+    if target_package == PKG_CUDA
+        return (kind === :priority) ? get_priority_custream(id) : get_custream(id)
+    elseif target_package == PKG_AMDGPU
+        return (kind === :priority) ? get_priority_rocstream(id) : get_rocstream(id)
+    elseif target_package == PKG_METAL
+        return (kind === :priority) ? get_priority_metalstream(id) : get_metalstream(id)
+    elseif target_package == PKG_THREADS || target_package == PKG_POLYESTER
+        @ArgumentError("KernelAbstractions hardware $(symbol) does not expose GPU streams.")
+    else
+        @ArgumentError("$(ERRMSG_UNSUPPORTED_PACKAGE) (obtained: $target_package).")
+    end
+end
+
+function PK.get_priority_kastream(id::Integer; hardware::Union{Symbol,Nothing}=nothing)
+    return dispatch_kastream(:priority, id, hardware)
+end
+
+function PK.get_kastream(id::Integer; hardware::Union{Symbol,Nothing}=nothing)
+    return dispatch_kastream(:regular, id, hardware)
+end
+
+
+## FUNCTIONS TO QUERY DEVICE PROPERTIES
+
+function PK.get_kernelabstractions_compute_capability(default::VersionNumber; hardware::Union{Symbol,Nothing}=nothing)
+    target_package, _, _ = resolve_runtime_backend(PKG_KERNELABSTRACTIONS, hardware)
+    if target_package == PKG_CUDA
+        return get_cuda_compute_capability(default)
+    elseif target_package == PKG_AMDGPU
+        return get_amdgpu_compute_capability(default)
+    elseif target_package == PKG_METAL
+        return get_metal_compute_capability(default)
+    else
+        return get_cpu_compute_capability(default)
+    end
+end