diff --git a/Project.toml b/Project.toml index 11d2e3edd..11b4c64d3 100644 --- a/Project.toml +++ b/Project.toml @@ -45,7 +45,7 @@ Atomix = "1.2.1" EnzymeCore = "0.7, 0.8.1" GPUCompiler = "2.7" InteractiveUtils = "1.6" -KernelInterface = "0.1" +KernelInterface = "0.2" LLVM = "9.9" LinearAlgebra = "1.6" MacroTools = "0.5" diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 5d13f2b9a..b4b486cf7 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -63,7 +63,7 @@ get_backend These are called from inside a kernel. A backend provides each one with ```julia -@device_override KI.get_global_id() = ... +@device_override KI.get_global_id(::Type{T}) where {T} = ... ``` along with the corresponding on-device functionality. @@ -71,7 +71,9 @@ along with the corresponding on-device functionality. ### Indexing All index queries are **1-based** and return a named tuple of `x`, `y` and `z` -components. +components. They take an optional element type `T` for the components, defaulting +to `Int`, so a kernel can request e.g. `Int32` indices with +`KI.get_global_id(Int32)`. ```@docs get_global_size diff --git a/lib/KernelInterface/Project.toml b/lib/KernelInterface/Project.toml index 4b5c12506..c6892206f 100644 --- a/lib/KernelInterface/Project.toml +++ b/lib/KernelInterface/Project.toml @@ -1,7 +1,7 @@ name = "KernelInterface" uuid = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" authors = ["Valentin Churavy and contributors"] -version = "0.1.1" +version = "0.2.0" [compat] julia = "1.10" diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 328dc62c5..60a62fa09 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -1,20 +1,22 @@ """ - get_global_size()::@NamedTuple{x::Int, y::Int, z::Int} + get_global_size([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Return the number of global work-items specified. +Return the number of global work-items specified as a tuple of type `T`. +`T` defaults to `Int`. !!! note Backend implementations **must** implement: ``` - @device_override get_global_size()::@NamedTuple{x::Int, y::Int, z::Int} + @device_override get_global_size(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` + The zero-argument form forwards to `get_global_size(Int)`. """ -function get_global_size end +get_global_size() = get_global_size(Int) """ - get_global_id()::@NamedTuple{x::Int, y::Int, z::Int} + get_global_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Returns the unique global work-item ID. +Returns the unique global work-item ID as a tuple of type `T`. `T` defaults to `Int`. !!! note 1-based. @@ -22,28 +24,31 @@ Returns the unique global work-item ID. !!! note Backend implementations **must** implement: ``` - @device_override get_global_id()::@NamedTuple{x::Int, y::Int, z::Int} + @device_override get_global_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` + The zero-argument form forwards to `get_global_id(Int)`. """ -function get_global_id end +get_global_id() = get_global_id(Int) """ - get_local_size()::@NamedTuple{x::Int, y::Int, z::Int} + get_local_size([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Return the number of local work-items specified. +Return the number of local work-items specified as a tuple of type `T`. +`T` defaults to `Int`. !!! note Backend implementations **must** implement: ``` - @device_override get_local_size()::@NamedTuple{x::Int, y::Int, z::Int} + @device_override get_local_size(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` + The zero-argument form forwards to `get_local_size(Int)`. """ -function get_local_size end +get_local_size() = get_local_size(Int) """ - get_local_id()::@NamedTuple{x::Int, y::Int, z::Int} + get_local_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Returns the unique local work-item ID. +Returns the unique local work-item ID as a tuple of type `T`. `T` defaults to `Int`. !!! note 1-based. @@ -51,28 +56,30 @@ Returns the unique local work-item ID. !!! note Backend implementations **must** implement: ``` - @device_override get_local_id()::@NamedTuple{x::Int, y::Int, z::Int} + @device_override get_local_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` + The zero-argument form forwards to `get_local_id(Int)`. """ -function get_local_id end +get_local_id() = get_local_id(Int) """ - get_num_groups()::@NamedTuple{x::Int, y::Int, z::Int} + get_num_groups([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Returns the number of groups. +Returns the number of groups as a tuple of type `T`. `T` defaults to `Int`. !!! note Backend implementations **must** implement: ``` - @device_override get_num_groups()::@NamedTuple{x::Int, y::Int, z::Int} + @device_override get_num_groups(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` + The zero-argument form forwards to `get_num_groups(Int)`. """ -function get_num_groups end +get_num_groups() = get_num_groups(Int) """ - get_group_id()::@NamedTuple{x::Int, y::Int, z::Int} + get_group_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Returns the unique group ID. +Returns the unique group ID as a tuple of type `T`. `T` defaults to `Int`. !!! note 1-based. @@ -80,10 +87,11 @@ Returns the unique group ID. !!! note Backend implementations **must** implement: ``` - @device_override get_group_id()::@NamedTuple{x::Int, y::Int, z::Int} + @device_override get_group_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` + The zero-argument form forwards to `get_group_id(Int)`. """ -function get_group_id end +get_group_id() = get_group_id(Int) """ get_sub_group_size()::UInt32 diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 079e141db..ff838302c 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -48,16 +48,50 @@ end # The interface documents a concrete return type for each device-side function; # these kernels record whether the backend honors them. -const WorkItemNT = @NamedTuple{x::Int, y::Int, z::Int} +const WorkItemNT{T} = @NamedTuple{x::T, y::T, z::T} function typecheck_kernel(results) @inbounds begin - results[1] = KI.get_global_size() isa WorkItemNT - results[2] = KI.get_global_id() isa WorkItemNT - results[3] = KI.get_local_size() isa WorkItemNT - results[4] = KI.get_local_id() isa WorkItemNT - results[5] = KI.get_num_groups() isa WorkItemNT - results[6] = KI.get_group_id() isa WorkItemNT + results[1] = KI.get_global_size() isa WorkItemNT{Int} + results[2] = KI.get_global_id() isa WorkItemNT{Int} + results[3] = KI.get_local_size() isa WorkItemNT{Int} + results[4] = KI.get_local_id() isa WorkItemNT{Int} + results[5] = KI.get_num_groups() isa WorkItemNT{Int} + results[6] = KI.get_group_id() isa WorkItemNT{Int} + end + return +end + +# The indexing queries take an element type; the result must use it. +function typed_typecheck_kernel(results, ::Type{T}) where {T} + @inbounds begin + results[1] = KI.get_global_size(T) isa WorkItemNT{T} + results[2] = KI.get_global_id(T) isa WorkItemNT{T} + results[3] = KI.get_local_size(T) isa WorkItemNT{T} + results[4] = KI.get_local_id(T) isa WorkItemNT{T} + results[5] = KI.get_num_groups(T) isa WorkItemNT{T} + results[6] = KI.get_group_id(T) isa WorkItemNT{T} + end + return +end + +# Records the typed indexing queries for every work-item, so the host can check +# that they agree with the default `Int` form across all three dimensions. +# `results` is `(work-items, 18)`: one row per work-item, holding the `x`, `y` +# and `z` components of each of the six queries in turn. +function typed_index_kernel(results, ::Type{T}) where {T} + i, j, k = KI.get_global_id(T) + ni, nj, _ = KI.get_global_size(T) + lin = (k - one(T)) * ni * nj + (j - one(T)) * ni + i + + if lin <= size(results, 1) + vals = ( + KI.get_global_size(T)..., KI.get_global_id(T)..., KI.get_local_size(T)..., + KI.get_local_id(T)..., KI.get_num_groups(T)..., KI.get_group_id(T)..., + ) + for q in 1:18 + @inbounds results[lin, q] = vals[q] + end end return end @@ -180,6 +214,41 @@ function interface_testsuite(backend, AT) KI.@kernel backend() typecheck_kernel(results) KI.synchronize(backend()) @test all(Array(results)) + + @testset "$T" for T in (Int32, Int64, UInt32, UInt64) + typed_results = KI.zeros(backend(), Bool, 6) + KI.@kernel backend() typed_typecheck_kernel(typed_results, T) + KI.synchronize(backend()) + @test all(Array(typed_results)) + end + end + + @testset "Typed indexing" begin + workgroupsize = (2, 2, 2) + numworkgroups = (3, 2, 1) + N = prod(workgroupsize) * prod(numworkgroups) + + # `Int` is the reference: it is what the zero-argument form returns. + function run_typed(::Type{T}) where {T} + results = KI.zeros(backend(), T, N, 18) + KI.@kernel backend() workgroupsize = workgroupsize numworkgroups = numworkgroups typed_index_kernel(results, T) + KI.synchronize(backend()) + return Array(results) + end + reference = run_typed(Int) + + global_size = workgroupsize .* numworkgroups + @test all(eachrow(reference[:, 1:3]) .== Ref(collect(global_size))) + @test all(eachrow(reference[:, 7:9]) .== Ref(collect(workgroupsize))) + @test all(eachrow(reference[:, 13:15]) .== Ref(collect(numworkgroups))) + # every global id is seen exactly once + @test sort(Tuple.(eachrow(reference[:, 4:6]))) == sort(vec(Tuple.(CartesianIndices(global_size)))) + + @testset "$T" for T in (Int32, UInt32, UInt64) + typed = run_typed(T) + @test typed isa AbstractMatrix{T} + @test typed == reference + end end @testset "Basic interface functionality" begin diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index a9513d41b..ddc92e273 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -29,9 +29,6 @@ end # These have no fallback on purpose: a backend that forgets to `@device_override` # them should get a MethodError rather than silently wrong behaviour. stubs = [ - KI.get_global_size, KI.get_global_id, - KI.get_local_size, KI.get_local_id, - KI.get_num_groups, KI.get_group_id, KI.get_sub_group_size, KI.get_max_sub_group_size, KI.get_num_sub_groups, KI.get_sub_group_id, KI.get_sub_group_local_id, @@ -44,6 +41,21 @@ end for stub in stubs @test isempty(methods(stub)) end + + # The indexing queries take an element type; only the zero-argument form has a + # (forwarding) method, and it must reach the typed stub rather than recurse. + indexing = [ + KI.get_global_size, KI.get_global_id, + KI.get_local_size, KI.get_local_id, + KI.get_num_groups, KI.get_group_id, + ] + for f in indexing + @test length(methods(f)) == 1 + @test hasmethod(f, Tuple{}) + @test !hasmethod(f, Tuple{Type{Int}}) + @test_throws MethodError f() + @test_throws MethodError f(Int32) + end end struct StubBackend <: KI.Backend end diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 98a606afe..a8dfcd0cf 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -285,28 +285,28 @@ end ## Indexing Functions -@device_override @inline function KI.get_local_id() - return (; x = Int(get_local_id(1)), y = Int(get_local_id(2)), z = Int(get_local_id(3))) +@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 KI.get_group_id() - return (; x = Int(get_group_id(1)), y = Int(get_group_id(2)), z = Int(get_group_id(3))) +@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 KI.get_global_id() - return (; x = Int(get_global_id(1)), y = Int(get_global_id(2)), z = Int(get_global_id(3))) +@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 KI.get_local_size() - return (; x = Int(get_local_size(1)), y = Int(get_local_size(2)), z = Int(get_local_size(3))) +@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 KI.get_num_groups() - return (; x = Int(get_num_groups(1)), y = Int(get_num_groups(2)), z = Int(get_num_groups(3))) +@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 KI.get_global_size() - return (; x = Int(get_global_size(1)), y = Int(get_global_size(2)), z = Int(get_global_size(3))) +@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