reuse macro infrastructure

Author: vchuravy

src/macros.jl

@@ -94,29 +94,7 @@ function transform_gpu!(def, constargs, force_inbounds)
     if force_inbounds
         push!(new_stmts, Expr(:inbounds, true))
     end
-
-    # fix convergence
-    active_stmts = Any[]
-    for stmt in stmts
-        has_sync = find_sync(stmt)
-        if has_sync
-            push!(new_stmts, Expr(:if, :__active_lane__, Expr(:block, active_stmts...)))
-            empty!(active_stmts)
-            push!(new_stmts, stmt)
-            continue
-        end
-        if @capture(stmt, @uniform x_)
-            push!(new_stmts, stmt)
-            continue
-        elseif @capture(stmt, lhs_ = rhs_ | (vs__, lhs_ = rhs_))
-            if @capture(rhs, @localmem(args__) | @uniform(args__))
-                push!(new_stmts, stmt)
-                continue
-            end
-        end
-        push!(active_stmts, stmt)
-    end
-    push!(new_stmts, Expr(:if, :__active_lane__, Expr(:block, active_stmts...)))
+    append!(new_stmts, split(emit_gpu, body.args))
     if force_inbounds
         push!(new_stmts, Expr(:inbounds, :pop))
     end
@@ -151,7 +129,7 @@ function transform_cpu!(def, constargs, force_inbounds)
     if force_inbounds
         push!(new_stmts, Expr(:inbounds, true))
     end
-    append!(new_stmts, split(body.args))
+    append!(new_stmts, split(emit_cpu, body.args))
     if force_inbounds
         push!(new_stmts, Expr(:inbounds, :pop))
     end
@@ -191,6 +169,7 @@ end
 
 # TODO proper handling of LineInfo
 function split(
+        emit,
         stmts,
         indicies = Any[], private = Set{Symbol}(),
     )
@@ -221,7 +200,7 @@ function split(
             function recurse(expr::Expr)
                 expr = unblock(expr)
                 if is_scope_construct(expr) && any(find_sync, expr.args)
-                    new_args = unblock(split(expr.args, deepcopy(indicies), deepcopy(private)))
+                    new_args = unblock(split(emit, expr.args, deepcopy(indicies), deepcopy(private)))
                     return Expr(expr.head, new_args...)
                 else
                     return Expr(expr.head, map(recurse, expr.args)...)
@@ -270,7 +249,7 @@ function split(
     return new_stmts
 end
 
-function emit(loop)
+function emit_cpu(loop)
     idx = gensym(:I)
     for stmt in loop.indicies
         # splice index into the i = @index(Cartesian, $idx)
@@ -324,3 +303,37 @@ function emit(loop)
 
     return unblock(Expr(:block, stmts...))
 end
+
+function emit_gpu(loop)
+    stmts = Any[]
+    append!(stmts, loop.allocations)
+    for stmt in loop.private_allocations
+        if @capture(stmt, lhs_ = rhs_)
+            push!(stmts, :($lhs = $rhs))
+        else
+            error("@private $stmt not an assignment")
+        end
+    end
+
+    # don't emit empty loops
+    if !(isempty(loop.stmts) || all(s -> s isa LineNumberNode, loop.stmts))
+        body = Expr(:block, loop.stmts...)
+        body = postwalk(body) do expr
+            if @capture(expr, lhs_ = rhs_)
+                if lhs in loop.private
+                    error("Can't assign to variables marked private")
+                end
+            end
+            return expr
+        end
+        loopexpr = quote
+            $(loop.indicies...)
+            if __active_lane__
+                $(unblock(body))
+            end
+        end
+        push!(stmts, loopexpr)
+    end
+
+    return unblock(Expr(:block, stmts...))
+end