KernelInterface: per-dimension launch limits, inlined index queries

.buildkite/pipeline.yml

@@ -23,7 +23,7 @@ steps:
                 PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
       julia -e 'println("--- :julia: Developing CUDA")
                 using Pkg
-                url="https://github.com/christiangnrd/CUDA.jl"
+                url="https://github.com/JuliaGPU/CUDA.jl"
                 rev="intrinsics"
                 subdirs = ["CUDACore", "CUDATools", "lib/cusolver", "lib/curand", "lib/cublas", "lib/cufft", "lib/cusparse", "lib/cupti", "lib/nvml", "lib/custatevec", "lib/cudnn", "lib/cutensor", "lib/cutensornet"]
                 Pkg.add([(;url, rev); [(; url, rev, subdir) for subdir in subdirs]])'

.github/workflows/ci.yml

@@ -142,20 +142,22 @@ jobs:
       #         end
       #       end
       #     '
-      # Nightly tracks the next Julia release; don't fail CI on it
-      - uses: julia-actions/julia-buildpkg@v1
-        continue-on-error: ${{ matrix.version == 'nightly' }}
       # Known limitation: `[sources]` is only supported from Julia 1.11 on, so on 1.10
       # the `KernelInterface = {path = "lib/KernelInterface"}` entry in Project.toml is
       # silently ignored and Pkg resolves KernelInterface from the registry instead --
       # a stale release that need not have the features this repo's copy adds. Dev it
       # explicitly so 1.10 tests the copy in this repo, like 1.11+ already do.
       #
       # Required on every platform. Windows runs its tests de-escalated below, but that
-      # applies to the test run only, not to `Pkg.develop`.
+      # applies to the test run only, not to `Pkg.develop`. Do this before building, which
+      # already resolves the environment (and fails if this repo requires an unregistered
+      # version of KernelInterface).
       - name: Dev KernelInterface
         shell: bash
         run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface")'
+      # Nightly tracks the next Julia release; don't fail CI on it
+      - uses: julia-actions/julia-buildpkg@v1
+        continue-on-error: ${{ matrix.version == 'nightly' }}
       - uses: julia-actions/julia-runtest@v1
         if: runner.os != 'Windows'
         continue-on-error: ${{ matrix.version == 'nightly' }}
@@ -228,10 +230,13 @@ jobs:
         with:
           channel: ${{ matrix.version }}
       - uses: julia-actions/cache@v3
-      - uses: julia-actions/julia-buildpkg@v1
-        continue-on-error: ${{ matrix.version == 'nightly' }}
       # KernelInterface is dev'ed explicitly here for the same reason as in the CI job
       # above: Julia 1.10 does not support `[sources]`.
+      - name: Dev KernelInterface
+        shell: bash
+        run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface")'
+      - uses: julia-actions/julia-buildpkg@v1
+        continue-on-error: ${{ matrix.version == 'nightly' }}
       - name: "Instantiating project"
         run: |
           julia -e 'println("--- :julia: Instantiating project")

Project.toml

@@ -42,7 +42,7 @@ Adapt = "0.4, 1.0, 2.0, 3.0, 4"
 Atomix = "1.2.1"
 EnzymeCore = "0.7, 0.8.1"
 GPUCompiler = "2.7"
-KernelInterface = "0.2"
+KernelInterface = "0.2.3"
 LLVM = "9.9"
 LinearAlgebra = "1.6"
 MacroTools = "0.5"

docs/src/kernelinterface.md

@@ -171,6 +171,8 @@ supports_float64
 
 ```@docs
 max_work_group_size
+max_work_group_dims
+max_num_groups
 sub_group_size
 multiprocessor_count
 ```

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.2.2"
+version = "0.2.3"
 
 [compat]
 julia = "1.10"

lib/KernelInterface/src/device.jl

@@ -4,20 +4,24 @@
 Return the number of global work-items specified as a tuple of type `T`.
 `T` defaults to `Int`.
 
+The value is computed in `T`, and is undefined when it does not fit in `T`.
+
 !!! note
     Backend implementations **must** implement:
     ```
     @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)`.
 """
-get_global_size() = get_global_size(Int)
+@inline get_global_size() = get_global_size(Int)
 
 """
     get_global_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}
 
 Returns the unique global work-item ID as a tuple of type `T`. `T` defaults to `Int`.
 
+The value is computed in `T`, and is undefined when it does not fit in `T`.
+
 !!! note
     1-based.
 
@@ -28,28 +32,32 @@ Returns the unique global work-item ID as a tuple of type `T`. `T` defaults to `
     ```
     The zero-argument form forwards to `get_global_id(Int)`.
 """
-get_global_id() = get_global_id(Int)
+@inline get_global_id() = get_global_id(Int)
 
 """
     get_local_size([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}
 
 Return the number of local work-items specified as a tuple of type `T`.
 `T` defaults to `Int`.
 
+The value is computed in `T`, and is undefined when it does not fit in `T`.
+
 !!! note
     Backend implementations **must** implement:
     ```
     @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)`.
 """
-get_local_size() = get_local_size(Int)
+@inline get_local_size() = get_local_size(Int)
 
 """
     get_local_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}
 
 Returns the unique local work-item ID as a tuple of type `T`. `T` defaults to `Int`.
 
+The value is computed in `T`, and is undefined when it does not fit in `T`.
+
 !!! note
     1-based.
 
@@ -60,27 +68,31 @@ Returns the unique local work-item ID as a tuple of type `T`. `T` defaults to `I
     ```
     The zero-argument form forwards to `get_local_id(Int)`.
 """
-get_local_id() = get_local_id(Int)
+@inline get_local_id() = get_local_id(Int)
 
 """
     get_num_groups([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}
 
 Returns the number of groups as a tuple of type `T`. `T` defaults to `Int`.
 
+The value is computed in `T`, and is undefined when it does not fit in `T`.
+
 !!! note
     Backend implementations **must** implement:
     ```
     @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)`.
 """
-get_num_groups() = get_num_groups(Int)
+@inline get_num_groups() = get_num_groups(Int)
 
 """
     get_group_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}
 
 Returns the unique group ID as a tuple of type `T`. `T` defaults to `Int`.
 
+The value is computed in `T`, and is undefined when it does not fit in `T`.
+
 !!! note
     1-based.
 
@@ -91,7 +103,7 @@ Returns the unique group ID as a tuple of type `T`. `T` defaults to `Int`.
     ```
     The zero-argument form forwards to `get_group_id(Int)`.
 """
-get_group_id() = get_group_id(Int)
+@inline get_group_id() = get_group_id(Int)
 
 """
     get_sub_group_size()::UInt32

lib/KernelInterface/src/launch.jl

@@ -55,14 +55,26 @@ function check_launch_args(numworkgroups, workgroupsize, ndrange)
     return
 end
 
-function threads_to_workgroupsize(threads, ndrange)
-    total = Ref(1)
-    return map(ndrange) do n
-        # each dimension must be at least 1, even for a zero-sized ndrange
-        x = max(min(div(threads, total[]), n), 1)
-        total[] *= x
-        return x
-    end
+"""
+    threads_to_workgroupsize(threads, ndrange, [limits])
+
+Distribute `threads` work-items over the dimensions of `ndrange`, filling the first
+dimension first. Dimension `d` gets at most `limits[d]` work-items; dimensions past the end
+of `limits` are only bounded by `threads`.
+
+Every dimension gets at least one work-item, even for a zero-sized `ndrange`.
+"""
+threads_to_workgroupsize(threads, ndrange::Tuple, limits = ()) =
+    _threads_to_workgroupsize(threads, 1, ndrange, limits)
+threads_to_workgroupsize(threads, ndrange::Integer, limits = ()) =
+    only(threads_to_workgroupsize(threads, (ndrange,), limits))
+# written recursively, because a closure updating the running total would box it
+_threads_to_workgroupsize(threads, total, ::Tuple{}, limits) = ()
+function _threads_to_workgroupsize(threads, total, ndrange::Tuple, limits)
+    limit = isempty(limits) ? typemax(Int) : first(limits)
+    x = max(min(div(threads, total), first(ndrange), limit), 1)
+    rest = isempty(limits) ? () : Base.tail(limits)
+    return (x, _threads_to_workgroupsize(threads, total * x, Base.tail(ndrange), rest)...)
 end
 
 """
@@ -86,7 +98,7 @@ writing their own heuristic for calculating launch size.
     else
         workgroupsize = if workgroupsize == ()
             max_wgs = kernel_max_work_group_size(kernel; max_work_items = min(prod(ndrange), max_work_items))
-            threads_to_workgroupsize(max_wgs, ndrange)
+            threads_to_workgroupsize(max_wgs, ndrange, max_work_group_dims(kernel.backend))
         else
             workgroupsize
         end
@@ -130,6 +142,36 @@ kernel launch with too big a workgroup is attempted.
 """
 function max_work_group_size end
 
+"""
+    max_work_group_dims(backend)::NTuple{3, Int}
+
+The maximum number of work-items along each dimension of a workgroup, for the currently
+active device of `backend`. [`max_work_group_size`](@ref) bounds their product.
+
+!!! note
+    Backend implementations **should** implement:
+    ```
+    max_work_group_dims(backend::NewBackend)::NTuple{3, Int}
+    ```
+    The fallback does not limit individual dimensions.
+"""
+max_work_group_dims(::Backend) = (typemax(Int), typemax(Int), typemax(Int))
+
+"""
+    max_num_groups(backend)::NTuple{3, Int}
+
+The maximum number of workgroups along each dimension of a launch, for the currently
+active device of `backend`.
+
+!!! note
+    Backend implementations **should** implement:
+    ```
+    max_num_groups(backend::NewBackend)::NTuple{3, Int}
+    ```
+    The fallback does not limit the number of workgroups.
+"""
+max_num_groups(::Backend) = (typemax(Int), typemax(Int), typemax(Int))
+
 """
     sub_group_size(backend)::Int
 

lib/KernelInterface/test/interface.jl

@@ -186,6 +186,46 @@ function interface_testsuite(backend, AT)
         @test_throws ArgumentError (KI.@kernel backend() numworkgroups = (2, 2, 2) workgroupsize = (2, 2, 2, 2) launch_kernel3d(arr3d))
     end
 
+    @testset "Launch limits" begin
+        max_dims = KI.max_work_group_dims(backend())
+        max_groups = KI.max_num_groups(backend())
+        @test max_dims isa NTuple{3, Int} && all(>=(1), max_dims)
+        @test max_groups isa NTuple{3, Int} && all(>=(1), max_groups)
+
+        function fill_kernel(arr)
+            i, j, k = KI.get_global_id()
+            if i <= size(arr, 1) && j <= size(arr, 2) && k <= size(arr, 3)
+                @inbounds arr[i, j, k] = 1.0f0
+            end
+            return
+        end
+        kernel = KI.@kernel backend() launch = false fill_kernel(AT(zeros(Float32, 1, 1, 1)))
+        function fill_test(dims; kwargs...)
+            arr = AT(zeros(Float32, dims))
+            kernel(arr; kwargs...)
+            KI.synchronize(backend())
+            return all(Array(arr) .== 1)
+        end
+
+        # automatically chosen workgroup sizes respect the per-dimension limits
+        @testset "ndrange = $ndrange" for ndrange in ((1, 1, 5000), (1, 5000, 1), (1, 3, 2000))
+            @test fill_test(ndrange; ndrange)
+        end
+
+        # the reported limits can be launched
+        max_items = KI.kernel_max_work_group_size(kernel)
+        @testset "dimension $d" for d in 1:3
+            items = min(max_dims[d], max_items)
+            workgroupsize = ntuple(i -> i == d ? items : 1, 3)
+            @test fill_test(workgroupsize; workgroupsize, numworkgroups = (1, 1, 1))
+
+            # don't launch (practically) unlimited grids
+            groups = min(max_groups[d], 2^16)
+            numworkgroups = ntuple(i -> i == d ? groups : 1, 3)
+            @test fill_test(numworkgroups; workgroupsize = (1, 1, 1), numworkgroups)
+        end
+    end
+
     @testset "Host return types" begin
         b = backend()
 

lib/KernelInterface/test/runtests.jl

@@ -219,6 +219,18 @@ end
     @test KI.threads_to_workgroupsize(256, (0, 4)) == (1, 4)
     @test KI.threads_to_workgroupsize(256, (4, 0)) == (4, 1)
     @test KI.threads_to_workgroupsize(0, (5,)) == (1,)
+
+    # Per-dimension limits, as for CUDA's (1024, 1024, 64) blocks.
+    @test KI.threads_to_workgroupsize(1024, (1, 1, 5000), (1024, 1024, 64)) == (1, 1, 64)
+    @test KI.threads_to_workgroupsize(1024, (2000, 3), (512, 1024, 64)) == (512, 2)
+    # dimensions past the limits are only bounded by the thread budget
+    @test KI.threads_to_workgroupsize(64, (1, 1, 1, 100), (1024, 1024, 64)) == (1, 1, 1, 64)
+end
+
+@testset "per-dimension limits" begin
+    # Unlimited unless the backend says otherwise.
+    @test KI.max_work_group_dims(StubBackend()) == (typemax(Int), typemax(Int), typemax(Int))
+    @test KI.max_num_groups(StubBackend()) == (typemax(Int), typemax(Int), typemax(Int))
 end
 
 @testset "Kernel" begin
@@ -236,7 +248,17 @@ function KI.kernel_max_work_group_size(k::KI.Kernel{SizedBackend}; max_work_item
     return min(k.backend.maxThreads, max_work_items)
 end
 
+# ... and a fixed limit per workgroup dimension
+struct DimsBackend <: KI.Backend end
+KI.kernel_max_work_group_size(k::KI.Kernel{DimsBackend}; max_work_items::Int = typemax(Int)) =
+    min(1024, max_work_items)
+KI.max_work_group_dims(::DimsBackend) = (1024, 1024, 64)
+
 @testset "auto_launch_sizes" begin
+    # the per-dimension limit is respected
+    @test KI.auto_launch_sizes(KI.Kernel(DimsBackend(), nothing), (), (), (1, 1, 5000)) ===
+        ((1, 1, 79), (1, 1, 64))
+
     kernel = KI.Kernel(SizedBackend(256), nothing)
 
     # Without an ndrange the sizes pass through, defaulting to 1.

src/pocl/backend.jl

@@ -234,9 +234,19 @@ function KI.kernel_max_work_group_size(kernel::KI.Kernel{<:POCLBackend}; max_wor
     wginfo = cl.work_group_info(kernel.kern.fun, device())
     return Int(min(wginfo.size, max_work_items))
 end
-function KI.max_work_group_size(::POCLBackend)::Int
-    return Int(device().max_work_group_size)
+# querying the device allocates, so cache the limits that every launch needs
+function device_limits()
+    return get!(task_local_storage(), :POCLLimits) do
+        dev = device()
+        sizes = dev.max_work_item_size
+        (;
+            max_work_group_size = Int(dev.max_work_group_size),
+            max_work_group_dims = ntuple(d -> d <= length(sizes) ? sizes[d] : 1, 3),
+        )
+    end::@NamedTuple{max_work_group_size::Int, max_work_group_dims::NTuple{3, Int}}
 end
+KI.max_work_group_size(::POCLBackend)::Int = device_limits().max_work_group_size
+KI.max_work_group_dims(::POCLBackend)::NTuple{3, Int} = device_limits().max_work_group_dims
 function KI.sub_group_size(::POCLBackend)::Int
     # POCL can technically support any sub_group size.
     #  Check for common values used on GPUs then

test/runtests.jl

@@ -25,6 +25,18 @@ KernelAbstractions.versioninfo(POCLBackend())
 
 import KernelAbstractions.POCL: POCL, @opencl, @device_code_llvm
 
+module KernelInterfaceTests
+    import KernelInterface
+    using Test
+    include(joinpath(pkgdir(KernelInterface), "test", "testsuite.jl"))
+end
+@testset "POCL KernelInterface" begin
+    # POCL launches synchronously, so there's nothing for the events tests to order
+    KernelInterfaceTests.Testsuite.testsuite(
+        POCLBackend, "POCL", POCL, Array, POCL.CLDeviceArray; skip_tests = Set(["Events"])
+    )
+end
+
 @testset "POCL float atomics" begin
     # pocl's CPU device natively supports float add and min/max atomics in both global
     # and local memory, so the SPIR-V extensions guarding them must be permitted