Merge pull request #58 from JuliaGPU/vc/lets

Ensure that constify doesn't cause arguments to be captured

Author: vchuravy

src/backends/cpu.jl

@@ -59,12 +59,12 @@ function (obj::Kernel{CPU})(args...; ndrange=nothing, workgroupsize=nothing, dep
         ndrange = nothing
     end
 
-    t = Threads.@spawn __run(obj, ndrange, iterspace, args, dependencies)
+    t = Threads.@spawn __run(obj, ndrange, iterspace, args, dependencies, Val(dynamic))
     return CPUEvent(t)
 end
 
 # Inference barriers
-function __run(obj, ndrange, iterspace, args, dependencies)
+function __run(obj, ndrange, iterspace, args, dependencies, ::Val{dynamic}) where dynamic
     __waitall(CPU(), dependencies, yield)
     N = length(iterspace)
     Nthreads = Threads.nthreads()
@@ -79,16 +79,16 @@ function __run(obj, ndrange, iterspace, args, dependencies)
         len, rem = 1, 0
     end
     if Nthreads == 1
-        __thread_run(1, len, rem, obj, ndrange, iterspace, args)
+        __thread_run(1, len, rem, obj, ndrange, iterspace, args, Val(dynamic))
     else
         @sync for tid in 1:Nthreads
-            Threads.@spawn __thread_run(tid, len, rem, obj, ndrange, iterspace, args)
+            Threads.@spawn __thread_run(tid, len, rem, obj, ndrange, iterspace, args, Val(dynamic))
         end
     end
     return nothing
 end
 
-function __thread_run(tid, len, rem, obj, ndrange, iterspace, args)
+function __thread_run(tid, len, rem, obj, ndrange, iterspace, args, ::Val{dynamic}) where dynamic
     # compute this thread's iterations
     f = 1 + ((tid-1) * len)
     l = f + len - 1
@@ -105,27 +105,16 @@ function __thread_run(tid, len, rem, obj, ndrange, iterspace, args)
     # run this thread's iterations
     for i = f:l
         block = @inbounds blocks(iterspace)[i]
-        # TODO: how do we use the information that the iteration space maps perfectly to
-        #       the ndrange without incurring a 2x compilation overhead
-        # if dynamic
-        ctx = mkcontextdynamic(obj, block, ndrange, iterspace)
+        ctx = mkcontext(obj, block, ndrange, iterspace, Val(dynamic))
         Cassette.overdub(ctx, obj.f, args...)
-        # else
-        #     ctx = mkcontext(obj, blocks, ndrange, iterspace)
-        #     Threads.@spawn Cassette.overdub(ctx, obj.f, args...)
     end
     return nothing
 end
 
 Cassette.@context CPUCtx
 
-function mkcontext(kernel::Kernel{CPU}, I, _ndrange, iterspace)
-    metadata = CompilerMetadata{ndrange(kernel), false}(I, _ndrange, iterspace)
-    Cassette.disablehooks(CPUCtx(pass = CompilerPass, metadata=metadata))
-end
-
-function mkcontextdynamic(kernel::Kernel{CPU}, I, _ndrange, iterspace)
-    metadata = CompilerMetadata{ndrange(kernel), true}(I, _ndrange, iterspace)
+function mkcontext(kernel::Kernel{CPU}, I, _ndrange, iterspace, ::Val{dynamic}) where dynamic
+    metadata = CompilerMetadata{ndrange(kernel), dynamic}(I, _ndrange, iterspace)
     Cassette.disablehooks(CPUCtx(pass = CompilerPass, metadata=metadata))
 end
 

src/macros.jl

@@ -57,20 +57,23 @@ end
 # The easy case, transform the function for GPU execution
 # - mark constant arguments by applying `constify`.
 function transform_gpu!(def, constargs)
-    new_stmts = Expr[]
+    let_constargs = Expr[]
     for (i, arg) in enumerate(def[:args])
         if constargs[i]
-            push!(new_stmts, :($arg = $constify($arg)))
+            push!(let_constargs, :($arg = $constify($arg)))
         end
     end
 
-    def[:body] = quote
+    body = quote
         if $__validindex()
-            $(new_stmts...)
             $(def[:body])
         end
         return nothing
     end
+    def[:body] = Expr(:let,
+        Expr(:block, let_constargs...),
+        body,
+    )
 end
 
 # The hard case, transform the function for CPU execution
@@ -81,19 +84,22 @@ end
 #   - hoist workgroup definitions
 #   - hoist uniform variables
 function transform_cpu!(def, constargs)
-    new_stmts = Expr[]
+    let_constargs = Expr[]
     for (i, arg) in enumerate(def[:args])
         if constargs[i]
-            push!(new_stmts, :($arg = $constify($arg)))
+            push!(let_constargs, :($arg = $constify($arg)))
         end
     end
-
+    new_stmts = Expr[]
     body = MacroTools.flatten(def[:body])
     push!(new_stmts, Expr(:aliasscope))
     append!(new_stmts, split(body.args))
     push!(new_stmts, Expr(:popaliasscope))
     push!(new_stmts, :(return nothing))
-    def[:body] = Expr(:block, new_stmts...)
+    def[:body] = Expr(:let,
+        Expr(:block, let_constargs...),
+        Expr(:block, new_stmts...)
+    )
 end
 
 struct WorkgroupLoop

test/test.jl

@@ -121,7 +121,7 @@ end
     let kernel = constarg(CPU(), 8, (1024,))
         # this is poking at internals
         iterspace = NDRange{1, StaticSize{(128,)}, StaticSize{(8,)}}();
-        ctx = KernelAbstractions.mkcontext(kernel, 1, nothing, iterspace)
+        ctx = KernelAbstractions.mkcontext(kernel, 1, nothing, iterspace, Val(false))
         AT = Array{Float32, 2}
         IR = sprint() do io
             code_llvm(io, KernelAbstractions.Cassette.overdub,