Forbid divergent execution of work-group barriers

Author: vchuravy

src/KernelAbstractions.jl

@@ -297,6 +297,9 @@ end
 After a `@synchronize` statement all read and writes to global and local memory
 from each thread in the workgroup are visible in from all other threads in the
 workgroup.
+
+!!! note
+    `@synchronize()` must be encountered by all workitems of a work-group executing the kernel or by none at all.
 """
 macro synchronize()
     return quote
@@ -314,10 +317,15 @@ workgroup. `cond` is not allowed to have any visible sideffects.
 # Platform differences
   - `GPU`: This synchronization will only occur if the `cond` evaluates.
   - `CPU`: This synchronization will always occur.
+
+!!! warn
+    This variant of the `@synchronize` macro violates the requirement that `@synchronize` must be encountered
+    by all workitems of a work-group executing the kernel or by none at all.
+    Since v`0.9.34` this version of the macro is deprecated and lowers to `@synchronize()`
 """
 macro synchronize(cond)
     return quote
-        $(esc(cond)) && $__synchronize()
+        $__synchronize()
     end
 end
 

src/macros.jl

@@ -58,22 +58,165 @@ function transform_gpu!(def, constargs, force_inbounds)
         end
     end
     pushfirst!(def[:args], :__ctx__)
-    body = def[:body]
+    new_stmts = Expr[]
+    body = MacroTools.flatten(def[:body])
+    stmts = body.args
+    push!(new_stmts, Expr(:aliasscope))
+    push!(new_stmts, :(__active_lane__ = $__validindex(__ctx__)))
     if force_inbounds
-        body = quote
-            @inbounds $(body)
-        end
+        push!(new_stmts, Expr(:inbounds, true))
     end
-    body = quote
-        if $__validindex(__ctx__)
-            $(body)
-        end
-        return nothing
+    append!(new_stmts, split(emit_gpu, body.args))
+    if force_inbounds
+        push!(new_stmts, Expr(:inbounds, :pop))
     end
+    push!(new_stmts, Expr(:popaliasscope))
+    push!(new_stmts, :(return nothing))
     def[:body] = Expr(
         :let,
         Expr(:block, let_constargs...),
-        body,
+        Expr(:block, new_stmts...),
     )
     return
 end
+
+struct WorkgroupLoop
+    indicies::Vector{Any}
+    stmts::Vector{Any}
+    allocations::Vector{Any}
+    private_allocations::Vector{Any}
+    private::Set{Symbol}
+end
+
+is_sync(expr) = @capture(expr, @synchronize() | @synchronize(a_))
+
+function is_scope_construct(expr::Expr)
+    return expr.head === :block # ||
+    # expr.head === :let
+end
+
+function find_sync(stmt)
+    result = false
+    postwalk(stmt) do expr
+        result |= is_sync(expr)
+        expr
+    end
+    return result
+end
+
+# TODO proper handling of LineInfo
+function split(
+        emit,
+        stmts,
+        indicies = Any[], private = Set{Symbol}(),
+    )
+    # 1. Split the code into blocks separated by `@synchronize`
+    # 2. Aggregate `@index` expressions
+    # 3. Hoist allocations
+    # 4. Hoist uniforms
+
+    current = Any[]
+    allocations = Any[]
+    private_allocations = Any[]
+    new_stmts = Any[]
+    for stmt in stmts
+        has_sync = find_sync(stmt)
+        if has_sync
+            loop = WorkgroupLoop(deepcopy(indicies), current, allocations, private_allocations, deepcopy(private))
+            push!(new_stmts, emit(loop))
+            allocations = Any[]
+            private_allocations = Any[]
+            current = Any[]
+            is_sync(stmt) && continue
+
+            # Recurse into scope constructs
+            # TODO: This currently implements hard scoping
+            #       probably need to implemet soft scoping
+            #       by not deepcopying the environment.
+            recurse(x) = x
+            function recurse(expr::Expr)
+                expr = unblock(expr)
+                if is_scope_construct(expr) && any(find_sync, expr.args)
+                    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)...)
+                end
+            end
+            push!(new_stmts, recurse(stmt))
+            continue
+        end
+
+        if @capture(stmt, @uniform x_)
+            push!(allocations, stmt)
+            continue
+        elseif @capture(stmt, @private lhs_ = rhs_)
+            push!(private, lhs)
+            push!(private_allocations, :($lhs = $rhs))
+            continue
+        elseif @capture(stmt, lhs_ = rhs_ | (vs__, lhs_ = rhs_))
+            if @capture(rhs, @index(args__))
+                push!(indicies, stmt)
+                continue
+            elseif @capture(rhs, @localmem(args__) | @uniform(args__))
+                push!(allocations, stmt)
+                continue
+            elseif @capture(rhs, @private(T_, dims_))
+                # Implement the legacy `mem = @private T dims` as
+                # mem = Scratchpad(T, Val(dims))
+
+                if dims isa Integer
+                    dims = (dims,)
+                end
+                alloc = :($Scratchpad(__ctx__, $T, Val($dims)))
+                push!(allocations, :($lhs = $alloc))
+                push!(private, lhs)
+                continue
+            end
+        end
+
+        push!(current, stmt)
+    end
+
+    # everything since the last `@synchronize`
+    if !isempty(current)
+        loop = WorkgroupLoop(deepcopy(indicies), current, allocations, private_allocations, deepcopy(private))
+        push!(new_stmts, emit(loop))
+    end
+    return new_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
+            if __active_lane__
+                $(loop.indicies...)
+                $(unblock(body))
+            end
+            $__synchronize()
+        end
+        push!(stmts, loopexpr)
+    end
+
+    return unblock(Expr(:block, stmts...))
+end