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"

lib/cl/event.jl

@@ -238,11 +238,11 @@ function Base.getproperty(evt::AbstractEvent, s::Symbol)
     if s == :context
         ctx = Ref{cl_context}()
         clGetEventInfo(evt, CL_EVENT_CONTEXT, sizeof(cl_context), ctx, C_NULL)
-        return Context(ctx[])
+        return Context(ctx[], retain=true)
     elseif s == :command_queue
         cmd_q = Ref{cl_command_queue}()
         clGetEventInfo(evt, CL_EVENT_COMMAND_QUEUE, sizeof(cl_command_queue), cmd_q, C_NULL)
-        return CmdQueue(cmd_q[])
+        return CmdQueue(cmd_q[], retain=true)
     elseif s == :command_type
         cmd_t = Ref{Cint}()
         clGetEventInfo(evt, CL_EVENT_COMMAND_TYPE, sizeof(Cint), cmd_t, C_NULL)

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,9 +14,10 @@ import Adapt
 
 ## Back-end Definition
 
-export OpenCLBackend
+# export OpenCLBackend
+
 
-Base.@kwdef struct OpenCLBackend <: KA.GPU
+Base.@kwdef struct OpenCLBackend <: KI.GPU
     platform::cl.Platform = cl.platform()
 end
 
@@ -23,7 +26,9 @@ end
     return nothing
 end
 
-function KA.allocate(b::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T
+KI.versioninfo(io::IO, ::OpenCLBackend) = OpenCL.versioninfo(io)
+
+function KI.allocate(b::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T
     b.platform === cl.platform() || platform_mismatch_warning(b.platform, cl.platform())
     if unified
         memory_backend = cl.unified_memory_backend()
@@ -39,40 +44,31 @@ function KA.allocate(b::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = f
     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)
-
-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)
-
-# `@Const` applies `constify` inside the kernel, where arguments have already been
-# converted to device arrays, so the rule has to be registered for `CLDeviceArray`
-# rather than for `CLArray`.
-Adapt.adapt_storage(::KA.ConstAdaptor, a::CLDeviceArray) = Base.Experimental.Const(a)
+KI.synchronize(::OpenCLBackend) = cl.finish(cl.queue())
+KI.supports_float64(::OpenCLBackend) = in("cl_khr_fp64", cl.device().extensions)
 
 ## Device Selection
 
 # devices are numbered consecutively within the backend's platform, in enumeration order
 
-function KA.ndevices(b::OpenCLBackend)
+function KI.ndevices(b::OpenCLBackend)
     Int(cl.ndevices(b.platform))
 end
 
-function KA.device(b::OpenCLBackend)
+function KI.device(b::OpenCLBackend)
     current = cl.device()
     for (i, d) in enumerate(cl.devices(b.platform))
         d == current && return i
     end
     error("Active OpenCL device $current not found in the OpenCLBackend's platform \"$(b.platform.name)\".")
 end
 
-function KA.device!(b::OpenCLBackend, id::Int)
-    0 < id <= KA.ndevices(b) || throw(ArgumentError("Device id $id out of bounds."))
+function KI.device!(b::OpenCLBackend, id::Int)
+    0 < id <= KI.ndevices(b) || throw(ArgumentError("Device id $id out of bounds."))
     devs = cl.devices(b.platform)
 
     cl.device!(devs[id])
@@ -81,152 +77,160 @@ end
 
 ## Memory Operations
 
-function KA.copyto!(::OpenCLBackend, A, B)
+function KI.copyto!(::OpenCLBackend, A, B)
     copyto!(A, B)
     # TODO: Address device to host copies in jl being synchronizing
 end
 
 
 ## Kernel Launch
 
-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)
-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
-
-    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)
+KI.argconvert(::OpenCLBackend, arg) = kernel_convert(arg)
+
+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
+
+function (obj::KI.Kernel{OpenCLBackend})(args...; numworkgroups=(), workgroupsize=(), ndrange=(), max_work_group_size=typemax(Int))
     obj.backend.platform === cl.platform() || platform_mismatch_warning(obj.backend.platform, cl.platform())
+    KI.check_launch_args(numworkgroups, workgroupsize, ndrange)
+    prod(ndrange) == 0 && return nothing
 
-    ndrange, workgroupsize, iterspace, dynamic =
-        KA.launch_config(obj, ndrange, workgroupsize)
+    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
 
-    # 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...)
+    obj.kern(args...; local_size, global_size)
+    return nothing
+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
 
-    groups = length(KA.blocks(iterspace))
-    items = length(KA.workitems(iterspace))
+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
 
-    if groups == 0
-        return nothing
-    end
+function KI.shfl_down_types(::OpenCLBackend)
+    backend_extensions = cl.device().extensions
+    "cl_khr_subgroup_shuffle" in backend_extensions || return DataType[]
 
-    # Launch kernel
-    global_size = groups * items
-    local_size = items
-    kernel(ctx, args...; global_size, local_size)
+    res = copy(SPIRVIntrinsics.gentypes)
 
-    return nothing
-end
+    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
 
 ## Indexing Functions
+## COV_EXCL_START
 
-@device_override @inline function KA.__index_Local_Linear(ctx)
-    return get_local_id(1)
+@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_Group_Linear(ctx)
-    return get_group_id(1)
+@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_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_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_Local_Cartesian(ctx)
-    @inbounds KA.workitems(KA.__iterspace(ctx))[get_local_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_Group_Cartesian(ctx)
-    @inbounds KA.blocks(KA.__iterspace(ctx))[get_group_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.__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_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 @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
-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)
+## Events
+
+# The queue is task-local, so a marker with an empty wait list captures every command the
+# calling task has enqueued so far, on both in-order and out-of-order queues. The queue is
+# flushed so that other queues (or the host) waiting on the marker cannot deadlock on
+# commands that were never submitted to the device.
+function KI.record_event(::OpenCLBackend)
+    ev = cl.enqueue_marker_with_wait_list(cl.AbstractEvent[])
+    cl.flush(cl.queue())
+    return ev
 end
 
+function KI.wait_event(::OpenCLBackend, ev::cl.Event)
+    # if ev.context == cl.context()
+    #     # queues within a context can wait on each other's events without blocking the host
+    #     cl.enqueue_barrier_with_wait_list(cl.AbstractEvent[ev])
+    # else
+    #     # events cannot be shared across contexts (i.e. across devices), so wait on the host
+    #     wait(ev)
+    # end
+    return
+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