Code typed fixes (#144)

 fixes for code typed with complicated arguments to kernels

Author: mjulian31

src/reflection.jl

@@ -52,26 +52,7 @@ function ka_code_llvm(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=...)
-```
-To use interactive mode (with Cthulhu), call
-```
-@ka_code_typed interactive=true 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...)
+function format_ex(ex0)
     ex = ()
     args = gensym(:args)
     old_args = nothing
@@ -83,7 +64,10 @@ macro ka_code_typed(ex0...)
                 if isa(ex0[i].args[end], Expr)
                     # at expr (like ndrange=10)
                     kw = ex0[i].args[end]
-                    @assert kw.head == :kw
+                    if kw.head != :kw
+                        # if an expr in place of a variable, skip
+                        break
+                    end
                     kw.args[2] = esc(kw.args[2])
                     kw.head = Symbol("=")
                     resize!(ex0[i].args, length(ex0[i].args) - 1)
@@ -103,6 +87,31 @@ macro ka_code_typed(ex0...)
     end
     @assert(old_args != nothing)
     @assert(kern != nothing)
+    return ex, args, old_args, kern
+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=...)
+```
+To use interactive mode (with Cthulhu), call
+```
+@ka_code_typed interactive=true 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, old_args, kern = format_ex(ex0)
 
     thecall = InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :ka_code_typed, ex)
 
@@ -137,43 +146,18 @@ If ndrange is statically defined, then you could call
 Works for CPU kernels ONLY, with static or dynamic declarations
 """
 macro ka_code_llvm(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
-                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)
+    ex, args, old_args, kern = format_ex(ex0)
 
     thecall = InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :ka_code_llvm, ex)
 
     quote
         local $(esc(args)) = $(old_args)
 
+        if isa($kern, Kernel{CUDADevice})
+            # does not support CUDA kernels
+            error("@ka_code_llvm does not support CUDA kernels")
+        end
+
         local results = $thecall
     end
 end

test/reflection.jl

@@ -10,6 +10,11 @@ end
     A[I] = B[I] + C[I]
 end
 
+@kernel function addi(A, C, i)
+    I = @index(Global)
+    A[I] = i + C[I]
+end
+
 function test_typed_kernel_dynamic()
     A = ones(1024, 1024)
     kernel = mul2(CPU())
@@ -84,7 +89,26 @@ function test_typed_kernel_no_optimize()
     end
 end
 
+function test_expr_kernel()
+    A = ones(1024, 1024)
+    C = similar(A)
+    kernel = addi(CPU())
+    res = @ka_code_typed kernel(A, C, 1+2, 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)
+        C = similar(A)
+        kernel = addi(CUDADevice(), (32, 32))
+        res = @ka_code_typed kernel(A, C, 1+2, ndrange=size(A))
+        @test isa(res, Pair{Core.CodeInfo, DataType})
+        @test isa(res[1].code, Array{Any,1})
+    end
+end
+
 test_typed_kernel_dynamic()
 test_typed_kernel_dynamic_no_info()
 test_typed_kernel_static()
 test_typed_kernel_no_optimize()
+test_expr_kernel()