KernelInterface

Project.toml

@@ -10,6 +10,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OpenCL_jll = "6cb37087-e8b6-5417-8430-1f242f1e46e4"
@@ -33,6 +34,7 @@ Adapt = "4"
 GPUArrays = "11.2.1"
 GPUCompiler = "2.7"
 KernelAbstractions = "0.9.38"
+KernelInterface = "0.2"
 LLVM = "9.6"
 LinearAlgebra = "1"
 OpenCL_jll = "=2024.10.24"

src/OpenCL.jl

@@ -11,6 +11,8 @@ using Preferences
 
 import KernelAbstractions: KernelAbstractions
 
+import KernelInterface
+
 using Core: LLVMPtr
 
 # library wrappers
@@ -48,7 +50,12 @@ include("mapreduce.jl")
 include("gpuarrays.jl")
 include("random.jl")
 
-include("OpenCLKernels.jl")
+include("OpenCLKernelsOld.jl")
 import .OpenCLKernels: OpenCLBackend
 export OpenCLBackend
+
+# KernelInterface - NOT PUBLIC. Use KernelInterface.get_backend on an CLArray to get the backend
+include("OpenCLKernels.jl")
+import .OpenCLInterface
+
 end

src/OpenCLKernels.jl

@@ -1,9 +1,11 @@
-module OpenCLKernels
+module OpenCLInterface
 
 using ..OpenCL
-using ..OpenCL: @device_override, method_table
+using ..OpenCL: @device_override, method_table, kernel_convert, clfunction
 
-import KernelAbstractions as KA
+import KernelInterface as KI
+
+import SPIRVIntrinsics
 
 import StaticArrays
 
@@ -12,12 +14,14 @@ import Adapt
 
 ## Back-end Definition
 
-export OpenCLBackend
+# export OpenCLBackend
 
-struct OpenCLBackend <: KA.GPU
+struct OpenCLBackend <: KI.GPU
 end
 
-function KA.allocate(::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T
+KI.versioninfo(io::IO, ::OpenCLBackend) = OpenCL.versioninfo(io)
+
+function KI.allocate(::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T
     if unified
         memory_backend = cl.unified_memory_backend()
         if memory_backend === cl.USMBackend()
@@ -32,164 +36,187 @@ function KA.allocate(::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = fa
     end
 end
 
-KA.supports_unified(::OpenCLBackend) = cl.default_memory_backend(cl.device(); unified=true) !== nothing
+KI.supports_unified(::OpenCLBackend) = cl.default_memory_backend(cl.device(); unified=true) !== nothing
 
-KA.get_backend(::CLArray) = OpenCLBackend()
+KI.get_backend(::CLArray) = OpenCLBackend()
 # TODO should be non-blocking
-KA.synchronize(::OpenCLBackend) = cl.finish(cl.queue())
-KA.supports_float64(::OpenCLBackend) = in("cl_khr_fp64", cl.device().extensions)
+KI.synchronize(::OpenCLBackend) = cl.finish(cl.queue())
+KI.supports_float64(::OpenCLBackend) = in("cl_khr_fp64", cl.device().extensions)
 
 Adapt.adapt_storage(::OpenCLBackend, a::Array) = Adapt.adapt(CLArray, a)
 Adapt.adapt_storage(::OpenCLBackend, a::CLArray) = a
-Adapt.adapt_storage(::KA.CPU, a::CLArray) = convert(Array, a)
+# Adapt.adapt_storage(::KI.CPU, a::CLArray) = convert(Array, a)
 
+## Device Selection
 
-## Memory Operations
+# devices are numbered consecutively across all platforms, in enumeration order
 
-function KA.copyto!(::OpenCLBackend, A, B)
-    copyto!(A, B)
-    # TODO: Address device to host copies in jl being synchronizing
+function KI.ndevices(::OpenCLBackend)
+    n = 0
+    for p in cl.platforms()
+        n += length(cl.devices(p))
+    end
+    return n
 end
 
+function KI.device(::OpenCLBackend)
+    current = cl.device()
+    i = 0
+    for p in cl.platforms()
+        for d in cl.devices(p)
+            i += 1
+            d == current && return i
+        end
+    end
+    error("Active OpenCL device $current not found among the available platforms.")
+end
+
+function KI.device!(::OpenCLBackend, id::Int)
+    i = 0
+    for p in cl.platforms()
+        for d in cl.devices(p)
+            i += 1
+            if i == id
+                cl.device!(d)
+                return nothing
+            end
+        end
+    end
+    throw(ArgumentError("Device id $id out of bounds."))
+end
 
-## Kernel Launch
+## Memory Operations
 
-function KA.mkcontext(kernel::KA.Kernel{OpenCLBackend}, _ndrange, iterspace)
-    KA.CompilerMetadata{KA.ndrange(kernel), KA.DynamicCheck}(_ndrange, iterspace)
-end
-function KA.mkcontext(kernel::KA.Kernel{OpenCLBackend}, I, _ndrange, iterspace,
-                      ::Dynamic) where Dynamic
-    KA.CompilerMetadata{KA.ndrange(kernel), Dynamic}(I, _ndrange, iterspace)
+function KI.copyto!(::OpenCLBackend, A, B)
+    copyto!(A, B)
+    # TODO: Address device to host copies in jl being synchronizing
 end
 
-function KA.launch_config(kernel::KA.Kernel{OpenCLBackend}, ndrange, workgroupsize)
-    if ndrange isa Integer
-        ndrange = (ndrange,)
-    end
-    if workgroupsize isa Integer
-        workgroupsize = (workgroupsize, )
-    end
 
-    # partition checked that the ndrange's agreed
-    if KA.ndrange(kernel) <: KA.StaticSize
-        ndrange = nothing
-    end
-
-    iterspace, dynamic = if KA.workgroupsize(kernel) <: KA.DynamicSize &&
-        workgroupsize === nothing
-        # use ndrange as preliminary workgroupsize for autotuning
-        KA.partition(kernel, ndrange, ndrange)
-    else
-        KA.partition(kernel, ndrange, workgroupsize)
-    end
+## Kernel Launch
 
-    return ndrange, workgroupsize, iterspace, dynamic
-end
 
 function threads_to_workgroupsize(threads, ndrange)
-    total = 1
+    total = Ref(1)
     return map(ndrange) do n
-        x = min(div(threads, total), n)
-        total *= x
+        x = min(div(threads, total[]), n)
+        total[] *= x
         return x
     end
 end
 
-function (obj::KA.Kernel{OpenCLBackend})(args...; ndrange=nothing, workgroupsize=nothing)
-    ndrange, workgroupsize, iterspace, dynamic =
-        KA.launch_config(obj, ndrange, workgroupsize)
+KI.argconvert(::OpenCLBackend, arg) = kernel_convert(arg)
 
-    # this might not be the final context, since we may tune the workgroupsize
-    ctx = KA.mkcontext(obj, ndrange, iterspace)
-    kernel = @opencl launch=false obj.f(ctx, args...)
+function KI.kernel_function(::OpenCLBackend, f::F, tt::TT=Tuple{}; name = nothing, kwargs...) where {F,TT}
+    kern = clfunction(f, tt; name, kwargs...)
+    KI.Kernel{OpenCLBackend, typeof(kern)}(OpenCLBackend(), kern)
+end
 
-    # figure out the optimal workgroupsize automatically
-    if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing
-        wg_info = cl.work_group_info(kernel.fun, cl.device())
-        wg_size_nd = threads_to_workgroupsize(wg_info.size, ndrange)
-        iterspace, dynamic = KA.partition(obj, ndrange, wg_size_nd)
-        ctx = KA.mkcontext(obj, ndrange, iterspace)
-    end
+function (obj::KI.Kernel{OpenCLBackend})(args...; numworkgroups=(), workgroupsize=(), ndrange=(), max_work_group_size=typemax(Int))
+    KI.check_launch_args(numworkgroups, workgroupsize, ndrange)
+    prod(ndrange) == 0 && return nothing
 
-    groups = length(KA.blocks(iterspace))
-    items = length(KA.workitems(iterspace))
+    numworkgroups, workgroupsize = KI.auto_launch_sizes(obj, numworkgroups, workgroupsize, ndrange, max_work_group_size)
+    local_size = (workgroupsize..., ntuple(_ -> 1, 3 - length(workgroupsize))...)
+    numworkgroups = (numworkgroups..., ntuple(_ -> 1, 3 - length(numworkgroups))...)
+    global_size = local_size .* numworkgroups
 
-    if groups == 0
-        return nothing
-    end
+    obj.kern(args...; local_size, global_size)
+    return nothing
+end
 
-    # Launch kernel
-    global_size = groups * items
-    local_size = items
-    kernel(ctx, args...; global_size, local_size)
 
-    return nothing
+function KI.kernel_max_work_group_size(kernel::KI.Kernel{<:OpenCLBackend}; max_work_items::Int=typemax(Int))::Int
+    wginfo = cl.work_group_info(kernel.kern.fun, cl.device())
+    Int(min(wginfo.size, max_work_items))
+end
+function KI.max_work_group_size(::OpenCLBackend)::Int
+    Int(cl.device().max_work_group_size)
+end
+function KI.sub_group_size(::OpenCLBackend)::Int
+    cl.sub_group_size(cl.device())
+end
+function KI.multiprocessor_count(::OpenCLBackend)::Int
+    Int(cl.device().max_compute_units)
 end
 
+function KI.shfl_down_types(::OpenCLBackend)
+    backend_extensions = cl.device().extensions
+    "cl_khr_subgroup_shuffle" in backend_extensions || return DataType[]
 
-## Indexing Functions
+    res = copy(SPIRVIntrinsics.gentypes)
 
-@device_override @inline function KA.__index_Local_Linear(ctx)
-    return get_local_id(1)
+    if "cl_khr_fp64" ∉ backend_extensions
+        res = setdiff(res, [Float64])
+    end
+    if "cl_khr_fp16" ∉ backend_extensions
+        res = setdiff(res, [Float16])
+    end
+
+    return res
 end
 
-@device_override @inline function KA.__index_Group_Linear(ctx)
-    return get_group_id(1)
+## Indexing Functions
+## COV_EXCL_START
+
+@device_override @inline function KI.get_local_id(::Type{T}) where {T}
+    return (; x = T(get_local_id(1)), y = T(get_local_id(2)), z = T(get_local_id(3)))
 end
 
-@device_override @inline function KA.__index_Global_Linear(ctx)
-    #return get_global_id(1)    # JuliaGPU/OpenCL.jl#346
-    I = KA.__index_Global_Cartesian(ctx)
-    @inbounds LinearIndices(KA.__ndrange(ctx))[I]
+@device_override @inline function KI.get_group_id(::Type{T}) where {T}
+    return (; x = T(get_group_id(1)), y = T(get_group_id(2)), z = T(get_group_id(3)))
 end
 
-@device_override @inline function KA.__index_Local_Cartesian(ctx)
-    @inbounds KA.workitems(KA.__iterspace(ctx))[get_local_id(1)]
+@device_override @inline function KI.get_global_id(::Type{T}) where {T}
+    return (; x = T(get_global_id(1)), y = T(get_global_id(2)), z = T(get_global_id(3)))
 end
 
-@device_override @inline function KA.__index_Group_Cartesian(ctx)
-    @inbounds KA.blocks(KA.__iterspace(ctx))[get_group_id(1)]
+@device_override @inline function KI.get_local_size(::Type{T}) where {T}
+    return (; x = T(get_local_size(1)), y = T(get_local_size(2)), z = T(get_local_size(3)))
 end
 
-@device_override @inline function KA.__index_Global_Cartesian(ctx)
-    return @inbounds KA.expand(KA.__iterspace(ctx), get_group_id(1), get_local_id(1))
+@device_override @inline function KI.get_num_groups(::Type{T}) where {T}
+    return (; x = T(get_num_groups(1)), y = T(get_num_groups(2)), z = T(get_num_groups(3)))
 end
 
-@device_override @inline function KA.__validindex(ctx)
-    if KA.__dynamic_checkbounds(ctx)
-        I = KA.__index_Global_Cartesian(ctx)
-        return I in KA.__ndrange(ctx)
-    else
-        return true
-    end
+@device_override @inline function KI.get_global_size(::Type{T}) where {T}
+    return (; x = T(get_global_size(1)), y = T(get_global_size(2)), z = T(get_global_size(3)))
 end
 
+@device_override KI.get_sub_group_size() = get_sub_group_size() % UInt32
+
+@device_override KI.get_max_sub_group_size() = get_max_sub_group_size() % UInt32
+
+@device_override KI.get_num_sub_groups() = get_num_sub_groups() % UInt32
+
+@device_override KI.get_sub_group_id() = get_sub_group_id() % UInt32
+
+@device_override KI.get_sub_group_local_id() = get_sub_group_local_id() % UInt32
 
 ## Shared and Scratch Memory
 
-@device_override @inline function KA.SharedMemory(::Type{T}, ::Val{Dims}, ::Val{Id}) where {T, Dims, Id}
+@device_override @inline function KI.localmemory(::Type{T}, ::Val{Dims}) where {T, Dims}
     ptr = OpenCL.emit_localmemory(T, Val(prod(Dims)))
     CLDeviceArray(Dims, ptr)
 end
 
-@device_override @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims}
-    StaticArrays.MArray{KA.__size(Dims), T}(undef)
-end
-
-
 ## Synchronization and Printing
 
-@device_override @inline function KA.__synchronize()
+@device_override @inline function KI.barrier()
     work_group_barrier(OpenCL.LOCAL_MEM_FENCE | OpenCL.GLOBAL_MEM_FENCE)
 end
 
-@device_override @inline function KA.__print(args...)
-    OpenCL._print(args...)
+@device_override @inline function KI.sub_group_barrier()
+    sub_group_barrier(OpenCL.LOCAL_MEM_FENCE | OpenCL.GLOBAL_MEM_FENCE)
 end
 
+@device_override function KI.shfl_down(val::T, offset::Integer) where T
+    sub_group_shuffle(val, get_sub_group_local_id() + offset)
+end
 
-## Other
-
-KA.argconvert(::KA.Kernel{OpenCLBackend}, arg) = OpenCL.kernel_convert(arg)
+@device_override @inline function KI._print(args...)
+    OpenCL._print(args...)
+end
+## COV_EXCL_STOP
 
 end