Fixes and test for at_ka_code_typed

Author: mjulian31

src/backends/cuda.jl

@@ -134,28 +134,41 @@ end
 ###
 # Kernel launch
 ###
-function (obj::Kernel{CUDADevice})(args...; ndrange=nothing, dependencies=nothing, workgroupsize=nothing, progress=yield)
+function launch_config(kernel::Kernel{CUDADevice}, ndrange, workgroupsize)
     if ndrange isa Integer
         ndrange = (ndrange,)
     end
     if workgroupsize isa Integer
-        workgroupsize = (workgroupsize,)
+        workgroupsize = (workgroupsize, )
     end
 
-    if KernelAbstractions.workgroupsize(obj) <: DynamicSize && workgroupsize === nothing
+    if KernelAbstractions.workgroupsize(kernel) <: DynamicSize && workgroupsize === nothing
         # TODO: allow for NDRange{1, DynamicSize, DynamicSize}(nothing, nothing)
         #       and actually use CUDA autotuning
         workgroupsize = (256,)
     end
+
+    # partition checked that the ndrange's agreed
+    if KernelAbstractions.ndrange(kernel) <: StaticSize
+        ndrange = nothing
+    end
+
+    iterspace, dynamic = partition(kernel, ndrange, workgroupsize)
+
+    return ndrange, workgroupsize, iterspace, dynamic
+end
+
+function (obj::Kernel{CUDADevice})(args...; ndrange=nothing, dependencies=nothing, workgroupsize=nothing, progress=yield)
+
+    ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize)
+
     # If the kernel is statically sized we can tell the compiler about that
     if KernelAbstractions.workgroupsize(obj) <: StaticSize
         maxthreads = prod(get(KernelAbstractions.workgroupsize(obj)))
     else
         maxthreads = nothing
     end
 
-    iterspace, dynamic = partition(obj, ndrange, workgroupsize)
-
     nblocks = length(blocks(iterspace))
     threads = length(workitems(iterspace))
 

src/reflection.jl

@@ -1,13 +1,19 @@
 import InteractiveUtils
 export @ka_code_typed
+using CUDA
 
 function ka_code_typed(kernel, argtypes; ndrange=nothing, workgroupsize=nothing, dependencies=nothing, kwargs...)
     # get the iterspace and dynamic of a kernel
     ndrange, workgroupsize, iterspace, dynamic = KernelAbstractions.launch_config(kernel, ndrange, workgroupsize)
-    # get the first block
-    block = @inbounds KernelAbstractions.blocks(iterspace)[1]
-    # get a context of the kernel based on the first block
-    ctx = KernelAbstractions.mkcontext(kernel, block, ndrange, iterspace, dynamic)
+
+    if isa(kernel, Kernel{CPU})
+        # get the first block
+        block = @inbounds KernelAbstractions.blocks(iterspace)[1]
+        # get a context of the kernel based on the first block
+        ctx = KernelAbstractions.mkcontext(kernel, block, ndrange, iterspace, dynamic)
+    else
+        ctx = KernelAbstractions.mkcontext(kernel, ndrange, iterspace)
+    end
     # reformat
     if argtypes isa Type
         argtypes = argtypes.parameters
@@ -17,25 +23,63 @@ function ka_code_typed(kernel, argtypes; ndrange=nothing, workgroupsize=nothing,
 end
 
 
+"""
+Get the typed IR for a kernel
+
+# Examples
+```
+@ka_code_typed kernel(args. ndrange=...)
+@ka_code_typed kernel(args. ndrange=... workgroupsize=...)
+@ka_code_typed optimize=false kernel(args. ndrange=...)
+```
+If ndrange is statically defined, then you could call
+```
+@ka_code_typed kernel(args.)
+```
+Works for CPU or CUDA kernels, with static or dynamic declarations
+"""
 macro ka_code_typed(ex0...)
     ex = ()
+    args = gensym(:args)
+    old_args = nothing
+    kern = nothing
     for i = 1:length(ex0)
         if ex0[i].head == :call
+            # inside kernel() expr
             while length(ex0[i].args) > 2
-                kw = ex0[i].args[end]
-                @assert kw.head == :kw
-                kw.args[2] = esc(kw.args[2])
-                kw.head = Symbol("=")
-                resize!(ex0[i].args, length(ex0[i].args) - 1)
-                ex = (kw,)..., ex...
+                if isa(ex0[i].args[end], Expr)
+                    # at expr (like ndrange=10)
+                    kw = ex0[i].args[end]
+                    @assert kw.head == :kw
+                    kw.args[2] = esc(kw.args[2])
+                    kw.head = Symbol("=")
+                    resize!(ex0[i].args, length(ex0[i].args) - 1)
+                    ex = (kw,)..., ex...
+                else
+                    # only symbols left
+                    break
+                end
             end
+            # save kernel args
+            old_args = Expr(:tuple, map(esc, ex0[i].args[2:end])...)
+            resize!(ex0[i].args, 2)
+            ex0[i].args[2] = Expr(:..., args)
+            kern = esc(ex0[i].args[1])
         end
         ex = ex..., ex0[i]
     end
+    @assert(old_args != nothing)
+    @assert(kern != nothing)
 
     thecall = InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :ka_code_typed, ex)
 
     quote
+        local $(esc(args)) = $(old_args)
+        if isa($kern, Kernel{CUDADevice})
+            # translate CuArray to CuDeviceArray
+            $(esc(args)) = map(CUDA.cudaconvert, $(esc(args)))
+        end
+
         local results = $thecall
         length(results) == 1 ? results[1] : results
     end

test/reflection.jl

@@ -1,24 +1,49 @@
-using KernelAbstractions, Test
+using KernelAbstractions, Test, CUDA
 
 @kernel function mul2(A)
     I = @index(Global)
     A[I] = 2 * A[I]
 end
 
+@kernel function add3(A, B, C)
+    I = @index(Global)
+    A[I] = B[I] + C[I]
+end
+
 function test_typed_kernel_dynamic()
     A = ones(1024, 1024)
     kernel = mul2(CPU())
     res = @ka_code_typed kernel(A, ndrange=size(A), workgroupsize=16)
     @test isa(res, Pair{Core.CodeInfo, DataType})
     @test isa(res[1].code, Array{Any,1})
+
+    if has_cuda_gpu()
+        A = CUDA.ones(1024, 1024)
+        kernel = mul2(CUDADevice())
+        res = @ka_code_typed kernel(A, ndrange=size(A), workgroupsize=(32, 32))
+        @test isa(res, Pair{Core.CodeInfo, DataType})
+        @test isa(res[1].code, Array{Any,1})
+    end
 end
 
 function test_typed_kernel_dynamic_no_info()
     A = ones(1024, 1024)
-    kernel = mul2(CPU())
-    res = @ka_code_typed kernel(A, ndrange=size(A))
+    B = similar(A)
+    C = similar(A)
+    kernel = add3(CPU())
+    res = @ka_code_typed kernel(A, B, C, ndrange=size(A))
     @test isa(res, Pair{Core.CodeInfo, DataType})
     @test isa(res[1].code, Array{Any,1})
+
+    if has_cuda_gpu()
+        A = CUDA.ones(1024, 1024)
+        B = similar(A)
+        C = similar(A)
+        kernel = add3(CUDADevice())
+        res = @ka_code_typed kernel(A, B, C, ndrange=size(A))
+        @test isa(res, Pair{Core.CodeInfo, DataType})
+        @test isa(res[1].code, Array{Any,1})
+    end
 end
 
 function test_typed_kernel_static()
@@ -27,15 +52,36 @@ function test_typed_kernel_static()
     res = @ka_code_typed kernel(A, ndrange=size(A))
     @test isa(res, Pair{Core.CodeInfo, DataType})
     @test isa(res[1].code, Array{Any,1})
+
+    if has_cuda_gpu()
+        A = CUDA.ones(1024, 1024)
+        kernel = mul2(CUDADevice(), (32, 32))
+        res = @ka_code_typed kernel(A, ndrange=size(A))
+        @test isa(res, Pair{Core.CodeInfo, DataType})
+        @test isa(res[1].code, Array{Any,1})
+    end
 end
 
 function test_typed_kernel_no_optimize()
     A = ones(1024, 1024)
-    kernel = mul2(CPU(), 16)
+    B = similar(A)
+    C = similar(A)
+    kernel = add3(CPU(), 16)
     res = @ka_code_typed optimize=false kernel(A, ndrange=size(A))
-    @test isa(res, Pair{Core.CodeInfo, DataType})
+    @test isa(res, Pair{Core.CodeInfo,Core.TypeofBottom})
     res_opt = @ka_code_typed kernel(A, ndrange=size(A))
     @test size(res[1].code) < size(res_opt[1].code)
+
+    if has_cuda_gpu()
+        A = CUDA.ones(1024, 1024)
+        B = similar(A)
+        C = similar(A)
+        kernel = add3(CUDADevice(), (32, 32))
+        res = @ka_code_typed optimize=false kernel(A, ndrange=size(A))
+        @test isa(res, Pair{Core.CodeInfo,Core.TypeofBottom})
+        res_opt = @ka_code_typed kernel(A, ndrange=size(A))
+        @test size(res[1].code) < size(res_opt[1].code)
+    end
 end
 
 test_typed_kernel_dynamic()