fixup N support

Author: vchuravy

src/KernelAbstractions.jl

@@ -52,7 +52,7 @@ synchronize(backend)
 ```
 """
 macro kernel(expr)
-    __kernel(expr, #=generate_cpu=# true, #=force_inbounds=# false)
+    __kernel(DynamicSize(), expr, #=generate_cpu=# true, #=force_inbounds=# false)
 end
 
 """
@@ -70,10 +70,11 @@ This allows for two different configurations:
 """
 macro kernel(ex...)
     if length(ex) == 1
-        __kernel(ex[1], true, false)
+        __kernel(DynamicSize(), ex[1], true, false)
     else
         generate_cpu = true
         force_inbounds = false
+        N = DynamicSize() # TODO parse N
         for i in 1:(length(ex) - 1)
             if ex[i] isa Expr && ex[i].head == :(=) &&
                     ex[i].args[1] == :cpu && ex[i].args[2] isa Bool
@@ -90,7 +91,7 @@ macro kernel(ex...)
                 )
             end
         end
-        __kernel(ex[end], generate_cpu, force_inbounds)
+        __kernel(N, ex[end], generate_cpu, force_inbounds)
     end
 end
 
@@ -586,7 +587,7 @@ in a workgroup.
     ```
     As well as the on-device functionality.
 """
-struct Kernel{Backend, N, WorkgroupSize <: _Size, NDRange <: _Size, Fun}
+struct Kernel{Backend, N <: _Size, WorkgroupSize <: _Size, NDRange <: _Size, Fun}
     backend::Backend
     f::Fun
 end
@@ -595,8 +596,9 @@ function Base.similar(kernel::Kernel{D, N, WS, ND}, f::F) where {D, N, WS, ND, F
     Kernel{D, N, WS, ND, F}(kernel.backend, f)
 end
 
-workgroupsize(::Kernel{D, N, WorkgroupSize}) where {D, WorkgroupSize} = WorkgroupSize
-ndrange(::Kernel{D, N, WorkgroupSize, NDRange}) where {D, WorkgroupSize, NDRange} = NDRange
+workgroupsize(::Kernel{D, N, WorkgroupSize}) where {D, N, WorkgroupSize} = WorkgroupSize
+ndrange(::Kernel{D, N, WorkgroupSize, NDRange}) where {D, N, WorkgroupSize, NDRange} = NDRange
+ndims(::Kernel{D, N}) where {D, N} = N
 backend(kernel::Kernel) = kernel.backend
 
 """
@@ -605,6 +607,7 @@ Partition a kernel for the given ndrange and workgroupsize.
 @inline function partition(kernel, ndrange, workgroupsize)
     static_ndrange = KernelAbstractions.ndrange(kernel)
     static_workgroupsize = KernelAbstractions.workgroupsize(kernel)
+    static_ndims = KernelAbstractions.ndims(kernel)
 
     if ndrange === nothing && static_ndrange <: DynamicSize ||
             workgroupsize === nothing && static_workgroupsize <: DynamicSize
@@ -655,11 +658,16 @@ Partition a kernel for the given ndrange and workgroupsize.
         workgroupsize = CartesianIndices(workgroupsize)
     end
 
+    if static_ndims <: StaticSize
+        @assert get(static_ndims) == length(ndrange)
+    end
+
+    # TODO: Add static_ndims
     iterspace = NDRange{length(ndrange), static_blocks, static_workgroupsize}(blocks, workgroupsize)
     return iterspace, dynamic
 end
 
-function construct(backend::Backend, ::Val{N}, ::S, ::NDRange, xpu_name::XPUName) where {Backend <: Union{CPU, GPU}, N, S <: _Size, NDRange <: _Size, XPUName}
+function construct(backend::Backend, ::N, ::S, ::NDRange, xpu_name::XPUName) where {Backend <: Union{CPU, GPU}, N <: _Size, S <: _Size, NDRange <: _Size, XPUName}
     return Kernel{Backend, N, S, NDRange, XPUName}(backend, xpu_name)
 end
 

src/macros.jl

@@ -10,7 +10,7 @@ function find_return(stmt)
 end
 
 # XXX: Proper errors
-function __kernel(expr, generate_cpu = true, force_inbounds = false)
+function __kernel(N, expr, generate_cpu = true, force_inbounds = false)
     def = splitdef(expr)
     name = def[:name]
     args = def[:args]
@@ -57,10 +57,10 @@ function __kernel(expr, generate_cpu = true, force_inbounds = false)
             $name(dev, size, range) = $name(dev, $StaticSize(size), $StaticSize(range))
             function $name(dev::Dev, sz::S, range::NDRange) where {Dev, S <: $_Size, NDRange <: $_Size}
                 if $isgpu(dev)
-                    return $construct(dev, sz, range, $gpu_name)
+                    return $construct(dev, $(N), sz, range, $gpu_name)
                 else
                     if $generate_cpu
-                        return $construct(dev, sz, range, $cpu_name)
+                        return $construct(dev, $(N), sz, range, $cpu_name)
                     else
                         error("This kernel is unavailable for backend CPU")
                     end

test/test.jl

@@ -10,7 +10,7 @@ identity(x) = x
 function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; skip_tests = Set{String}())
     @conditional_testset "partition" skip_tests begin
         backend = Backend()
-        let kernel = KernelAbstractions.Kernel{typeof(backend), StaticSize{(64,)}, DynamicSize, typeof(identity)}(backend, identity)
+        let kernel = KernelAbstractions.Kernel{typeof(backend), DynamicSize, StaticSize{(64,)}, DynamicSize, typeof(identity)}(backend, identity)
             iterspace, dynamic = KernelAbstractions.partition(kernel, (128,), nothing)
             @test length(blocks(iterspace)) == 2
             @test dynamic isa NoDynamicCheck
@@ -26,7 +26,7 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk
             @test_throws ErrorException KernelAbstractions.partition(kernel, (129,), (65,))
             @test KernelAbstractions.backend(kernel) == backend
         end
-        let kernel = KernelAbstractions.Kernel{typeof(backend), StaticSize{(64,)}, StaticSize{(128,)}, typeof(identity)}(backend, identity)
+        let kernel = KernelAbstractions.Kernel{typeof(backend), DynamicSize, StaticSize{(64,)}, StaticSize{(128,)}, typeof(identity)}(backend, identity)
             iterspace, dynamic = KernelAbstractions.partition(kernel, (128,), nothing)
             @test length(blocks(iterspace)) == 2
             @test dynamic isa NoDynamicCheck