diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 6d5d5c17..36c88c21 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -47,7 +47,17 @@ steps: # SPIRVIntrinsics first; otherwise the Pkg.add below resolves # against the registry, which has no SPIRVIntrinsics 1. Pkg.test # then carries the in-tree copy into the test sandbox. - Pkg.develop(path="lib/intrinsics") + if VERSION < v"1.12" + # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 are not registered yet, + # so develop them from their branch, in the same Pkg operation + run(`git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka`) + Pkg.develop([PackageSpec(; path) for path in ("lib/intrinsics", "ka", "ka/lib/KernelInterface")]) + # Pkg < 1.12 turns the developed weak dependency into a strong one; restore + # the project, keeping the developed package in the manifest + run(`git checkout Project.toml`) + else + Pkg.develop(path="lib/intrinsics") + end Pkg.add("{{matrix.pocl}}_jll") Pkg.add("InteractiveUtils") diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index ed5a15bc..1db493ff 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -142,7 +142,19 @@ jobs: using Pkg # Julia 1.10 does not support [sources], so dev the in-tree # SPIRVIntrinsics; Pkg.test then carries it into the test sandbox. - Pkg.develop(path="lib/intrinsics")' + if VERSION < v"1.12" + # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 are not registered yet, + # and Julia 1.11 does not pick them up from the [sources] of the test project + # without workspaces, so develop them from their branch, in the same + # Pkg operation + run(`git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka`) + Pkg.develop([PackageSpec(; path) for path in ("lib/intrinsics", "ka", "ka/lib/KernelInterface")]) + # Pkg < 1.12 turns the developed weak dependency into a strong one; restore + # the project, keeping the developed package in the manifest + run(`git checkout Project.toml`) + else + Pkg.develop(path="lib/intrinsics") + end' - name: Test OpenCL.jl uses: julia-actions/julia-runtest@v1 diff --git a/Project.toml b/Project.toml index df46ed4b..ff57b4d6 100644 --- a/Project.toml +++ b/Project.toml @@ -10,7 +10,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55" GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc" -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" @@ -23,10 +23,16 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SPIRVIntrinsics = "71d1d633-e7e8-4a92-83a1-de8814b09ba8" SPIRV_LLVM_Backend_jll = "4376b9bf-cff8-51b6-bb48-39421dff0d0c" SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4" -StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" spirv2clc_jll = "f0274c0c-8c8a-59f1-85b7-f7d60330c5fb" +[weakdeps] +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" + +[extensions] +KernelAbstractionsExt = "KernelAbstractions" + [sources] +KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"} SPIRVIntrinsics = {path = "lib/intrinsics"} [compat] @@ -34,7 +40,8 @@ Adapt = "4" GPUArrays = "11.2.1" GPUCompiler = "2.12" GPUToolbox = "3.3.2" -KernelAbstractions = "0.9.38" +KernelAbstractions = "0.10" +KernelInterface = "0.4" LLVM = "10" LinearAlgebra = "1" OpenCL_jll = "=2024.10.24" @@ -47,6 +54,9 @@ Reexport = "1" SPIRVIntrinsics = "1.2" SPIRV_LLVM_Backend_jll = "23" SPIRV_Tools_jll = "2025.1" -StaticArrays = "1" julia = "1.10" spirv2clc_jll = "0.2.0" + +# XXX: lets Pkg < 1.12 develop the weak dependency (see CI) +[extras] +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index e88779f0..f4da4eb1 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -13,5 +13,5 @@ Pages = ["95-reference.md"] ``` ```@autodocs -Modules = [OpenCL, OpenCL.cl] +Modules = [OpenCL, OpenCL.cl, OpenCL.OpenCLKernels] ``` diff --git a/ext/KernelAbstractionsExt.jl b/ext/KernelAbstractionsExt.jl new file mode 100644 index 00000000..14e97632 --- /dev/null +++ b/ext/KernelAbstractionsExt.jl @@ -0,0 +1,16 @@ +module KernelAbstractionsExt + +using OpenCL + +import KernelAbstractions as KA + +import Adapt + +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) + +end diff --git a/lib/cl/kernel.jl b/lib/cl/kernel.jl index c914f0e2..04b4e585 100644 --- a/lib/cl/kernel.jl +++ b/lib/cl/kernel.jl @@ -260,14 +260,15 @@ function enqueue_task(k::Kernel; wait_for=nothing) return ret_event[] end +# the arguments are passed as a tuple, see `set_args!` function call( - k::Kernel, args...; global_size = (1,), local_size = nothing, + k::Kernel, args::Tuple; global_size = (1,), local_size = nothing, global_work_offset = nothing, wait_on::Vector{Event} = Event[], indirect_memory::Vector{AbstractMemory} = AbstractMemory[], rng_state=false, ) return Base.@lock k begin - set_args!(k, args...) + _set_args!(k, args) if !isempty(indirect_memory) svm_pointers = CLPtr{Cvoid}[] usm_pointers = CLPtr{Cvoid}[] @@ -360,7 +361,7 @@ clcall(f::F, types::Tuple, args::Vararg{Any,N}; kwargs...) where {N,F} = function clcall(k::Kernel, types::Type{T}, args::Vararg{Any,N}; kwargs...) where {T,N} call_closure = function (converted_args::Vararg{Any,N}) - call(k, converted_args...; kwargs...) + call(k, converted_args; kwargs...) end convert_arguments(call_closure, types, args...) end @@ -415,10 +416,12 @@ unlock_arguments(locks) = foreach(unlock, Iterators.reverse(locks)) return ex end -function set_args!(k::Kernel, args...) - for (i, a) in enumerate(args) - set_arg!(k, i, a) - end +set_args!(k::Kernel, args::Vararg{Any,N}) where {N} = _set_args!(k, args) + +# one call per argument: a splat of more than 32 arguments isn't a direct call +@inline @generated function _set_args!(k::Kernel, args::Tuple) + calls = (:(set_arg!(k, $i, args[$i])) for i in 1:fieldcount(args)) + return :($(calls...); nothing) end function set_arg!(k::Kernel, idx::Integer, arg::T) where {T} diff --git a/src/OpenCL.jl b/src/OpenCL.jl index 3e3ef109..2dd665b8 100644 --- a/src/OpenCL.jl +++ b/src/OpenCL.jl @@ -10,7 +10,7 @@ using GPUArrays using Random using Preferences -import KernelAbstractions: KernelAbstractions +import KernelInterface using Core: LLVMPtr diff --git a/src/OpenCLKernels.jl b/src/OpenCLKernels.jl index 98220c9f..48e5e987 100644 --- a/src/OpenCLKernels.jl +++ b/src/OpenCLKernels.jl @@ -1,11 +1,9 @@ module OpenCLKernels using ..OpenCL -using ..OpenCL: @device_override, method_table +using ..OpenCL: @device_override, method_table, kernel_convert, clfunction -import KernelAbstractions as KA - -import StaticArrays +import KernelInterface as KI import Adapt @@ -14,17 +12,38 @@ import Adapt export OpenCLBackend -Base.@kwdef struct OpenCLBackend <: KA.GPU +""" + OpenCLBackend(; platform=cl.platform()) + +KernelInterface back end for the OpenCL devices of `platform`. + +A backend works with the task's active device if that is on its platform. Otherwise, work +for the backend (allocations, copies, compilation and launches) first activates the default +device of its platform, as `KernelInterface.device!` would, so that the arrays it creates +can be used afterwards. +""" +Base.@kwdef struct OpenCLBackend <: KI.Backend platform::cl.Platform = cl.platform() end -@noinline function platform_mismatch_warning(expected::cl.Platform, active::cl.Platform) - @warn "OpenCLBackend platform \"$(expected.name)\" is not the active platform \"$(active.name)\"" - return nothing +KI.versioninfo(io::IO, ::OpenCLBackend) = OpenCL.versioninfo(io) + +# the device that `b` works with +function backend_device(b::OpenCLBackend) + cl.platform() == b.platform && return cl.device() + dev = cl.default_device(b.platform) + dev === nothing && throw(ArgumentError("OpenCL platform \"$(b.platform.name)\" has no devices")) + return dev +end + +# make the backend's device the task's active device +@inline function activate(b::OpenCLBackend) + cl.platform() == b.platform || cl.platform!(b.platform) + return end -function KA.allocate(b::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T - b.platform === cl.platform() || platform_mismatch_warning(b.platform, cl.platform()) +function KI.allocate(b::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T + activate(b) if unified memory_backend = cl.unified_memory_backend() if memory_backend === cl.USMBackend() @@ -39,193 +58,274 @@ 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 +# OpenCL.jl creates a context per device +context_device(ctx::cl.Context) = ctx == cl.context() ? cl.device() : only(ctx.devices) -KA.get_backend(::CLArray) = OpenCLBackend() -KA.synchronize(::OpenCLBackend) = OpenCL.synchronize() -KA.supports_float64(::OpenCLBackend) = in("cl_khr_fp64", cl.device().extensions) +function KI.get_backend(A::CLArray) + ctx = OpenCL.context(A) + ctx == cl.context() && return OpenCLBackend(cl.platform()) + return OpenCLBackend(context_device(ctx).platform) +end -Adapt.adapt_storage(::OpenCLBackend, a::Array) = Adapt.adapt(CLArray, a) +function KI.synchronize(b::OpenCLBackend) + activate(b) + OpenCL.synchronize() + return +end + +# queues are task-local, so work is ordered across tasks with a marker event on the +# recording task's queue +function KI.record_event(b::OpenCLBackend) + activate(b) + event = cl.enqueue_marker_with_wait_list(cl.AbstractEvent[]) + # the waiting queue only makes progress if this one is submitted + cl.flush(cl.queue()) + return event +end + +function event_context(event::cl.Event) + ctx = Ref{cl.cl_context}() + cl.clGetEventInfo(event, cl.CL_EVENT_CONTEXT, sizeof(cl.cl_context), ctx, C_NULL) + return ctx[] +end + +function KI.wait_event(b::OpenCLBackend, event::cl.Event) + activate(b) + # the event has to stay alive until the driver has retained it + GC.@preserve event begin + if event_context(event) == cl.context().id + cl.enqueue_barrier_with_wait_list(cl.AbstractEvent[event]) + else + # XXX: queues can only wait for events of their own context, and OpenCL.jl + # creates a context per device, so wait for other devices on the host. + wait(event) + end + end + return +end + +function Adapt.adapt_storage(b::OpenCLBackend, a::Array) + activate(b) + return Adapt.adapt(CLArray, a) +end 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) ## 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) - 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)\".") +function device_index(b::OpenCLBackend, dev::cl.Device) + id = findfirst(==(dev), cl.devices(b.platform)) + id === nothing && + throw(ArgumentError("OpenCL device $(dev.name) is not on the backend's platform \"$(b.platform.name)\"")) + return id end -function KA.device!(b::OpenCLBackend, id::Int) - 0 < id <= KA.ndevices(b) || throw(ArgumentError("Device id $id out of bounds.")) +KI.device(b::OpenCLBackend) = device_index(b, backend_device(b)) + +KI.device(b::OpenCLBackend, A::CLArray) = device_index(b, context_device(OpenCL.context(A))) + +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]) return nothing end + ## Memory Operations -function KA.copyto!(::OpenCLBackend, A, B) +function KI.copyto!(b::OpenCLBackend, A, B) + length(A) == length(B) || + throw(ArgumentError("Arrays must have the same length, got $(length(A)) and $(length(B))")) + activate(b) copyto!(A, B) - # TODO: Address device to host copies in jl being synchronizing + return A end +KI.unsafe_free!(A::CLArray) = OpenCL.unsafe_free!(A) + ## Kernel Launch -function KA.mkcontext(kernel::KA.Kernel{OpenCLBackend}, _ndrange, iterspace) - KA.CompilerMetadata{KA.ndrange(kernel), KA.DynamicCheck}(_ndrange, iterspace) +KI.argconvert(::OpenCLBackend, arg) = kernel_convert(arg) + +# `f` is the host-side callable: the kernel keeps it as its `source`, which is converted +# again at every launch, because the converted callable only holds pointers to the arrays +# it captures +function KI.kernel_function(backend::OpenCLBackend, f::F, tt::TT=Tuple{}; name = nothing, kwargs...) where {F,TT} + activate(backend) + check_sub_group_size(backend, kwargs) + kern = GC.@preserve f clfunction(kernel_convert(f), tt; source=f, name, kwargs...) + KI.Kernel{OpenCLBackend, typeof(kern)}(backend, kern) end -function KA.mkcontext(kernel::KA.Kernel{OpenCLBackend}, I, _ndrange, iterspace, - ::Dynamic) where Dynamic - KA.CompilerMetadata{KA.ndrange(kernel), Dynamic}(I, _ndrange, iterspace) + +# kernels have to execute with the sub-group width that `KI.sub_group_size` reports +function check_sub_group_size(backend::OpenCLBackend, kwargs) + haskey(kwargs, :sub_group_size) && KI.supports_subgroups(backend) || return + width = KI.sub_group_size(backend) + kwargs[:sub_group_size] == width || + throw(ArgumentError("KernelInterface kernels execute with sub-group width $width, got `sub_group_size=$(kwargs[:sub_group_size])`")) + return 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 +# the context that a kernel was compiled for +function kernel_context(kernel::KI.Kernel{OpenCLBackend}) + ctx = Ref{cl.cl_context}() + cl.clGetKernelInfo(kernel.kern.fun, cl.CL_KERNEL_CONTEXT, sizeof(cl.cl_context), ctx, C_NULL) + return ctx[] +end - # partition checked that the ndrange's agreed - if KA.ndrange(kernel) <: KA.StaticSize - ndrange = nothing - end +function kernel_device(kernel::KI.Kernel{OpenCLBackend}) + ctx = kernel_context(kernel) + ctx == cl.context().id && return cl.device() + return context_device(cl.Context(ctx; retain=true)) +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 +@noinline function throw_device_mismatch(kernel) + throw(ArgumentError("Cannot launch a kernel compiled for $(kernel_device(kernel).name) on $(cl.device().name)")) +end - return ndrange, workgroupsize, iterspace, dynamic +@noinline function throw_geometry_keyword() + throw(ArgumentError("KernelInterface kernels take `numgroups`, `workgroupsize` or `ndrange`, not `global_size` or `local_size`")) end -function threads_to_workgroupsize(threads, ndrange) - total = 1 - return map(ndrange) do n - x = min(div(threads, total), n) - total *= x - return x +# passes the arguments on as a tuple, like calling the `HostKernel` does +function KI.launch(kernel::KI.Kernel{OpenCLBackend}, groups::Dims{3}, items::Dims{3}, + args::Tuple; kwargs...) + # KernelInterface has validated the launch geometry + if haskey(kwargs, :global_size) || haskey(kwargs, :local_size) + throw_geometry_keyword() end + activate(kernel.backend) + kernel_context(kernel) == cl.context().id || throw_device_mismatch(kernel) + OpenCL.launch_tuple(kernel.kern, args; local_size = items, global_size = items .* groups, + kwargs...) + return end -function (obj::KA.Kernel{OpenCLBackend})(args...; ndrange=nothing, workgroupsize=nothing) - obj.backend.platform === cl.platform() || platform_mismatch_warning(obj.backend.platform, cl.platform()) - - ndrange, workgroupsize, iterspace, dynamic = - KA.launch_config(obj, ndrange, workgroupsize) +function KI.max_work_group_size(kernel::KI.Kernel{OpenCLBackend})::Int + wginfo = cl.work_group_info(kernel.kern.fun, kernel_device(kernel)) + Int(wginfo.size) +end - # 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...) - # 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) +## Device Properties + +# querying the device allocates, so cache what launches and kernels need. the cache is +# keyed on the device, because the task-local device can be switched. +const DeviceProperties = @NamedTuple{ + max_work_group_size::Int, max_work_group_dims::NTuple{3, Int}, compute_units::Int, + float64::Bool, unified::Bool, + # 0 if the device doesn't support sub-groups of a fixed width + sub_group_size::Int, shuffle_types::Vector{DataType}, +} +function device_properties(dev::cl.Device) + cache = get!(task_local_storage(), :CLDeviceProperties) do + Dict{cl.Device, DeviceProperties}() + end::Dict{cl.Device, DeviceProperties} + return get!(cache, dev) do + sizes = dev.max_work_item_size + # the sub-group width is only fixed for kernels that request it, which `clfunction` + # does for devices with `cl_intel_required_subgroup_size` + fixed_sub_groups = cl.sub_groups_supported(dev) && + "cl_intel_required_subgroup_size" in dev.extensions + (; max_work_group_size = Int(dev.max_work_group_size), + max_work_group_dims = ntuple(d -> d <= length(sizes) ? Int(sizes[d]) : 1, 3), + compute_units = Int(dev.max_compute_units), + float64 = "cl_khr_fp64" in dev.extensions, + unified = cl.default_memory_backend(dev; unified=true) !== nothing, + sub_group_size = fixed_sub_groups ? cl.sub_group_size(dev) : 0, + shuffle_types = fixed_sub_groups ? cl.sub_group_shuffle_supported_types(dev) : DataType[]) end +end - groups = length(KA.blocks(iterspace)) - items = length(KA.workitems(iterspace)) +device_properties(b::OpenCLBackend) = device_properties(backend_device(b)) - if groups == 0 - return nothing - end +KI.max_work_group_size(b::OpenCLBackend)::Int = device_properties(b).max_work_group_size +KI.max_work_group_dims(b::OpenCLBackend)::NTuple{3, Int} = device_properties(b).max_work_group_dims +# OpenCL doesn't limit the number of work-groups, only the global size (to `size_t`) +function KI.max_num_groups(b::OpenCLBackend)::NTuple{3, Int} + return typemax(Int) .รท KI.max_work_group_dims(b) +end +KI.multiprocessor_count(b::OpenCLBackend)::Int = device_properties(b).compute_units - # Launch kernel - global_size = groups * items - local_size = items - kernel(ctx, args...; global_size, local_size) +KI.supports_float64(b::OpenCLBackend) = device_properties(b).float64 +KI.supports_unified(b::OpenCLBackend) = device_properties(b).unified +# 32-bit integer atomics are core OpenCL; float atomics fall back to compare-and-swap +KI.supports_atomics(::OpenCLBackend) = true - return nothing -end +KI.supports_subgroups(b::OpenCLBackend) = device_properties(b).sub_group_size > 0 +KI.sub_group_size(b::OpenCLBackend)::Int = device_properties(b).sub_group_size +KI.supports_shuffle(b::OpenCLBackend, ::Type{T}) where {T} = T in device_properties(b).shuffle_types ## Indexing Functions -@device_override @inline function KA.__index_Local_Linear(ctx) - return get_local_id(1) -end +# computed with `% T`, which unlike `T(x)` has no error path. KernelInterface derives the +# global queries from these. -@device_override @inline function KA.__index_Group_Linear(ctx) - return get_group_id(1) +@device_override @inline function KI.get_local_id(::Type{T}) where {T} + return (; x = get_local_id(1) % T, y = get_local_id(2) % T, z = get_local_id(3) % T) 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 = get_group_id(1) % T, y = get_group_id(2) % T, z = get_group_id(3) % T) 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 = get_local_size(1) % T, y = get_local_size(2) % T, z = get_local_size(3) % T) 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 = get_num_groups(1) % T, y = get_num_groups(2) % T, z = get_num_groups(3) % T) 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)) -end +# OpenCL's sub-group queries already have KernelInterface's semantics: the last sub-group +# of a work-group can be partial, and `get_sub_group_size` counts the work-items present -@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(::Type{T}) where {T} = get_sub_group_size() % T + +@device_override KI.get_max_sub_group_size(::Type{T}) where {T} = get_max_sub_group_size() % T + +@device_override KI.get_num_sub_groups(::Type{T}) where {T} = get_num_sub_groups() % T + +@device_override KI.get_sub_group_id(::Type{T}) where {T} = get_sub_group_id() % T + +@device_override KI.get_sub_group_local_id(::Type{T}) where {T} = get_sub_group_local_id() % T -## Shared and Scratch Memory +## Shared 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 +# out-of-range source lanes give an unspecified value, as KernelInterface allows +@device_override function KI.shfl_down(val::T, offset::Integer) where T + sub_group_shuffle(val, get_sub_group_local_id() % UInt32 + offset % UInt32) +end -## Other - -KA.argconvert(::KA.Kernel{OpenCLBackend}, arg) = OpenCL.kernel_convert(arg) +@device_override @inline function KI._print(args...) + OpenCL._print(args...) +end end diff --git a/src/compiler/exceptions.jl b/src/compiler/exceptions.jl index ba3f4344..8092f668 100644 --- a/src/compiler/exceptions.jl +++ b/src/compiler/exceptions.jl @@ -169,11 +169,14 @@ function with_exception_mailbox(f, mailbox::ExceptionMailbox) end end +# `(x, t...)`, without splatting +@inline @generated prepend(x, t::Tuple) = :((x, $((:(t[$i]) for i in 1:fieldcount(t))...))) + # Launch while holding the mailbox lock, so a concurrent synchronization cannot finish the # queue between enqueue and recording it as pending (or map the mailbox during submission). # Keep this in an ordinary function: the generated `AbstractKernel` call cannot contain a # closure or `do` block on Julia 1.13. -function launch_with_exception_mailbox(kernel::cl.Kernel, args...; +function launch_with_exception_mailbox(kernel::cl.Kernel, args::Tuple; indirect_memory::Vector{cl.AbstractMemory}, rng_state::Bool, kwargs...) ctx, dev, queue = cl.context(), cl.device(), cl.queue() @@ -187,7 +190,7 @@ function launch_with_exception_mailbox(kernel::cl.Kernel, args...; launch_id = (Base.@atomic mailbox.launch_id) + UInt64(1) state = KernelState(rng_state ? Base.rand(UInt32) : UInt32(0), mailbox.address, launch_id) - result = cl.call(kernel, state, args...; indirect_memory, rng_state, kwargs...) + result = cl.call(kernel, prepend(state, args); indirect_memory, rng_state, kwargs...) # only publish the launch once it has been submitted, so that `check_exceptions` # can tell whether kernels have been submitted while it waited without the lock Base.@atomic mailbox.launch_id = launch_id diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 6bacd8d6..9f9a6c79 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -176,11 +176,24 @@ abstract type AbstractKernel{F, TT} end pass_arg(@nospecialize dt) = !(isghosttype(dt) || Core.Compiler.isconstType(dt)) -@inline @generated function (kernel::AbstractKernel{F,TT})(args...; - call_kwargs...) where {F,TT} +# The arguments are passed on as a tuple: Julia doesn't turn a splat of more than 32 +# elements into a direct call, and a method with both varargs and keyword arguments splats +# them into its body. So the keyword method is defined explicitly. +(kernel::AbstractKernel)(args::Vararg{Any,N}) where {N} = launch_tuple(kernel, args) +Core.kwcall(kwargs::NamedTuple, kernel::AbstractKernel, args::Vararg{Any,N}) where {N} = + launch_tuple(kernel, args; kwargs...) + +""" + OpenCL.launch_tuple(kernel, args::Tuple; kwargs...) + +Launch `kernel`, as calling it with the arguments `args...` and the launch keywords +`kwargs` does, without splatting the arguments. +""" +@inline @generated function launch_tuple(kernel::AbstractKernel{F,TT}, args::Tuple; + call_kwargs...) where {F,TT} sig = Tuple{F, TT.parameters...} # Base.signature_type with a function type args = (:(kernel_convert(source, indirect_memory, managed)), - (:(kernel_convert(args[$i], indirect_memory, managed)) for i in 1:length(args))...) + (:(kernel_convert(args[$i], indirect_memory, managed)) for i in 1:fieldcount(args))...) # filter out ghost arguments that shouldn't be passed to_pass = map(pass_arg, sig.parameters) @@ -212,7 +225,7 @@ pass_arg(@nospecialize dt) = !(isghosttype(dt) || Core.Compiler.isconstType(dt)) locked = lock_managed(managed) try foreach(take_ownership!, locked) - launch_with_exception_mailbox(kernel.fun, $(converted...); + launch_with_exception_mailbox(kernel.fun, ($(converted...),); indirect_memory, rng_state=kernel.rng_state, call_kwargs...) finally diff --git a/src/util.jl b/src/util.jl index 65eff264..117a4d1d 100644 --- a/src/util.jl +++ b/src/util.jl @@ -72,7 +72,8 @@ function versioninfo(io::IO=stdout) (pkg[2], get(Base.loaded_modules, id, nothing)) end - for pkg in [:GPUArrays, :GPUCompiler, ("63c18a36-062a-441e-b654-da1e3ab1ce7c", "KernelAbstractions"), + for pkg in [:GPUArrays, :GPUCompiler, :KernelInterface, + ("63c18a36-062a-441e-b654-da1e3ab1ce7c", "KernelAbstractions"), :LLVM, :SPIRVIntrinsics, ("627d6b7a-bbe6-5189-83e7-98cc0a5aeadd", "pocl_jll"), ("59abdad9-3cfc-5436-8271-411e8cad6b82", "pocl_next_jll")] name, mod = get_module(pkg) diff --git a/test/Project.toml b/test/Project.toml index 14eef6fd..9b186565 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -9,6 +9,7 @@ IOCapture = "b5f81e59-6552-4d32-b1f0-c071b021bf89" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" @@ -29,6 +30,8 @@ pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" pocl_next_jll = "59abdad9-3cfc-5436-8271-411e8cad6b82" [sources] +KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"} +KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"} OpenCL = {path = ".."} SPIRVIntrinsics = {path = "../lib/intrinsics"} diff --git a/test/kernelabstractions.jl b/test/kernelabstractions.jl index 4f2d3de7..db2594ee 100644 --- a/test/kernelabstractions.jl +++ b/test/kernelabstractions.jl @@ -64,8 +64,52 @@ const KATestSuite = let mod.Testsuite end +# the others run kernels on KernelAbstractions' PoCL-based CPU back end skip_tests=Set([ "sparse", - "Convert", # Need to opt out of i128 + "CPU synchronization", + "fallback test: callable types", ]) KATestSuite.testsuite(OpenCLBackend, "OpenCL", OpenCL, CLArray, CLDeviceArray; skip_tests) + +@kernel function store_global_linear!(A) + I = @index(Global, Linear) + @inbounds A[I] = I +end + +@kernel function store_last_index!(A) + I = @index(Global, Linear) + if I == prod(@ndrange()) + @inbounds A[1] = I + @inbounds A[2] = @index(Global, Cartesian)[2] + end +end + +@testset "launch configuration" begin + backend = OpenCLBackend() + function select(kernel, ndrange, workgroupsize=nothing) + ndrange, workgroupsize, iterspace, _ = KernelAbstractions.launch_config(kernel, ndrange, workgroupsize) + KernelAbstractions.select_launch(kernel, workgroupsize, iterspace) + end + + # kernels are launched on an N-d grid, computing indices in 32 bits + kernel = store_global_linear!(backend) + @test select(kernel, (64, 32, 16)) === KernelAbstractions.NDLaunch{Int32}() + @test select(kernel, (4, 4, 4, 4)) === KernelAbstractions.LinearLaunch{Int32}() + + # which doesn't need divisions to compute the index of a dynamic N-d range + A = OpenCL.zeros(Int, 64, 32, 16) + ir = sprint(io -> @device_code_llvm io=io kernel(A; ndrange=size(A))) + @test !occursin(r"\b[su](div|rem) ", ir) + @test Array(A) == LinearIndices(A) + + # iteration spaces that don't fit 32 bits use 64-bit indices + kernel = store_last_index!(backend) + A = OpenCL.zeros(Int, 2) + for (dims, launch) in (((2^16 + 1, 2^15), KernelAbstractions.NDLaunch{Int}()), + ((2^11 + 1, 2^10, 2^10, 1), KernelAbstractions.LinearLaunch{Int}())) + @test select(kernel, dims) === launch + kernel(A; ndrange=dims) + @test Array(A) == [prod(dims), dims[2]] + end +end diff --git a/test/kernelinterface.jl b/test/kernelinterface.jl new file mode 100644 index 00000000..8242cff2 --- /dev/null +++ b/test/kernelinterface.jl @@ -0,0 +1,143 @@ +import KernelInterface +import KernelInterface as KI +import Adapt + +include(joinpath(dirname(pathof(KernelInterface)), "..", "test", "testsuite.jl")) + +Testsuite.testsuite(OpenCLBackend(), CLArray) + +function ki_fill_kernel(a, val) + i = KI.get_global_id().x + if i <= length(a) + @inbounds a[i] = val + end + return +end + +@testset "backend platform" begin + backend = OpenCLBackend() + a = KI.zeros(backend, Int32, 4) + @test KI.get_backend(a) == backend + + # a kernel belongs to the device it was compiled for + kernel = KI.@launch backend launch=false ki_fill_kernel(a, Int32(1)) + cl.context!(cl.Context(cl.device())) do + @test_throws ArgumentError kernel(a, Int32(1); ndrange = 4) + end + kernel(a, Int32(1); ndrange = 4) + @test Array(a) == ones(Int32, 4) + + # a backend for another platform activates it + others = filter(!=(cl.platform()), cl.platforms()) + if !isempty(others) + platform, device = cl.platform(), cl.device() + try + other = OpenCLBackend(; platform = first(others)) + @test KI.device(other) == 1 + b = KI.zeros(other, Int32, 4) + @test cl.platform() == first(others) + @test KI.get_backend(b) == other + @test KI.device(other, b) == KI.device(other) + + # so does moving arrays to it + cl.platform!(platform) + c = Adapt.adapt(other, Int32[1, 2]) + @test cl.platform() == first(others) + @test KI.get_backend(c) == other + KI.@launch other ndrange = 4 ki_fill_kernel(b, Int32(2)) + @test Array(b) == fill(Int32(2), 4) + + # arrays keep their platform + cl.device!(device) + @test KI.get_backend(b) == other + @test KI.get_backend(a) == backend + finally + cl.platform!(platform) + cl.device!(device) + end + end +end + +struct KICapturedArray{A} + array::A +end +Adapt.@adapt_structure KICapturedArray +function (f::KICapturedArray)() + @inbounds f.array[1] = 7 + return +end + +@testset "captured arrays" begin + function captured_kernel() + array = CLArray(Int32[0]) + kernel = KI.@launch OpenCLBackend() launch=false KICapturedArray(array)() + return kernel, WeakRef(array) + end + + # the kernel keeps the callable alive, and converts it again at launch, which makes the + # launch's queue the owner of the arrays it captures + kernel, owner = captured_kernel() + GC.gc(true) + @test owner.value !== nothing + queue = cl.CmdQueue() + cl.queue!(queue) do + kernel() + @test owner.value.data[].queue === queue + cl.finish(queue) + end + @test Array(owner.value) == Int32[7] +end + +@testset "launch keywords" begin + a = KI.zeros(OpenCLBackend(), Int32, 4) + kernel = KI.@launch OpenCLBackend() launch=false ki_fill_kernel(a, Int32(1)) + + # OpenCL's launch options are passed on + kernel(a, Int32(3); ndrange = 4, wait_on = cl.Event[]) + @test Array(a) == fill(Int32(3), 4) + + # but not ones that would override the launch geometry + @test_throws ArgumentError kernel(a, Int32(1); ndrange = 4, global_size = 8) + @test_throws ArgumentError kernel(a, Int32(1); ndrange = 4, local_size = 2) +end + +@testset "sub-groups" begin + backend = OpenCLBackend() + dev = cl.device() + # kernels only execute with a fixed sub-group width if they can request one + if cl.sub_groups_supported(dev) && "cl_intel_required_subgroup_size" in dev.extensions + @test KI.supports_subgroups(backend) + @test KI.sub_group_size(backend) == cl.sub_group_size(dev) + @test KI.supports_shuffle(backend, Int32) == + ("cl_khr_subgroup_shuffle" in dev.extensions) + + # which kernels can't opt out of + a = KI.zeros(backend, Int32, 4) + width = KI.sub_group_size(backend) + KI.@launch backend launch=false sub_group_size=width ki_fill_kernel(a, Int32(1)) + @test_throws ArgumentError KI.@launch backend launch=false sub_group_size=nothing ki_fill_kernel(a, Int32(1)) + @test_throws ArgumentError KI.@launch backend launch=false sub_group_size=2width ki_fill_kernel(a, Int32(1)) + else + @test !KI.supports_subgroups(backend) + end + @test !KI.supports_shuffle(backend, Complex{Float32}) +end + +@testset "events" begin + backend = OpenCLBackend() + a = KI.zeros(backend, Int32, 4) + KI.@launch backend ndrange = 4 ki_fill_kernel(a, Int32(5)) + event = KI.record_event(backend) + @test event isa cl.Event + + # queues of another context can't wait for the event, so the host does + cl.context!(cl.Context(cl.device())) do + @test KI.wait_event(backend, event) === nothing + @test event.status == :complete + end + @test Array(a) == fill(Int32(5), 4) +end + +@testset "versioninfo" begin + @test occursin("OpenCL.jl version", sprint(KI.versioninfo, OpenCLBackend())) +end