Merge pull request #25 from JuliaGPU/vc/val

Handle type parameters in kernel functions

Author: vchuravy

src/backends/cpu.jl

@@ -59,12 +59,12 @@ end
 Cassette.@context CPUCtx
 
 function mkcontext(kernel::Kernel{CPU}, I, _ndrange, _workgroupsize)
-    metadata = CompilerMetadata{workgroupsize(kernel), ndrange(kernel), false}(I, _ndrange, workgroupsize)
+    metadata = CompilerMetadata{workgroupsize(kernel), ndrange(kernel), false}(I, _ndrange, _workgroupsize)
     Cassette.disablehooks(CPUCtx(pass = CompilerPass, metadata=metadata))
 end
 
 function mkcontextdynamic(kernel::Kernel{CPU}, I, _ndrange, _workgroupsize)
-    metadata = CompilerMetadata{workgroupsize(kernel), ndrange(kernel), true}(I, _ndrange, workgroupsize)
+    metadata = CompilerMetadata{workgroupsize(kernel), ndrange(kernel), true}(I, _ndrange, _workgroupsize)
     Cassette.disablehooks(CPUCtx(pass = CompilerPass, metadata=metadata))
 end
 

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