[KernelInterface 0.2] Take an element type in the indexing queries

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"

docs/src/kernelinterface.md

@@ -63,15 +63,17 @@ 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.
 
 ### 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

lib/KernelInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "KernelInterface"
 uuid = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 authors = ["Valentin Churavy <v.churavy@gmail.com> and contributors"]
-version = "0.1.1"
+version = "0.2.0"
 
 [compat]
 julia = "1.10"

lib/KernelInterface/src/device.jl

@@ -1,89 +1,97 @@
 """
-    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.
 
 !!! 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.
 
 !!! 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.
 
 !!! 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

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

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

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