handle type parameters in kernels

Author: vchuravy

src/macros.jl

@@ -7,7 +7,15 @@ function __kernel(expr)
     body = expr.args[2]
 
     # parse decl
-    @assert isexpr(decl, :call)
+    # `@kernel fname(::T) where {T}`
+    if isexpr(decl, :where) 
+        iswhere = true
+        whereargs = decl.args[2:end]
+        decl = decl.args[1]
+    else
+        iswhere = false
+    end
+    @assert isexpr(decl, :call) 
     name = decl.args[1]
 
     # List of tuple (Symbol, Bool) where the bool
@@ -37,6 +45,11 @@ function __kernel(expr)
     gpu_decl = Expr(:call, gpu_name, arglist...)
     cpu_decl = Expr(:call, cpu_name, arglist...)
 
+    if iswhere
+        gpu_decl = Expr(:where, gpu_decl, whereargs...)
+        cpu_decl = Expr(:where, cpu_decl, whereargs...)
+    end
+
     # Without the deepcopy we might accidentially modify expr shared between CPU and GPU
     gpu_body = transform_gpu(deepcopy(body), args)
     gpu_function = Expr(:function, gpu_decl, gpu_body)

test/test.jl

@@ -167,4 +167,18 @@ end
             @test occursin("@llvm.nvvm.ldg", IR)
         end
     end
-end
+end
+
+@kernel function kernel_val!(a, ::Val{m}) where {m}
+    I = @index(Global)
+    @inbounds a[I] = m
+end
+
+A = zeros(Int64, 1024)
+wait(kernel_val!(CPU())(A,Val(3), ndrange=size(A)))
+@test all((a)->a==3, A)
+if has_cuda_gpu()
+    A = CuArrays.zeros(Int64, 1024)
+    wait(kernel_val!(CUDA())(A,Val(3), ndrange=size(A)))
+    @test all((a)->a==3, A)
+end