Skip to content

Commit

Permalink
Forbid divergent execution of work-group barriers
Browse files Browse the repository at this point in the history
  • Loading branch information
vchuravy committed Feb 7, 2025
1 parent 90a10d7 commit 8072f4c
Show file tree
Hide file tree
Showing 2 changed files with 102 additions and 11 deletions.
10 changes: 9 additions & 1 deletion src/KernelAbstractions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down
103 changes: 93 additions & 10 deletions src/macros.jl
Original file line number Diff line number Diff line change
Expand Up @@ -58,22 +58,105 @@ 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__)))

Check warning on line 65 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L61-L65

Added lines #L61 - L65 were not covered by tests
if force_inbounds
body = quote
@inbounds $(body)
end
push!(new_stmts, Expr(:inbounds, true))

Check warning on line 67 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L67

Added line #L67 was not covered by tests
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))

Check warning on line 71 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L69-L71

Added lines #L69 - L71 were not covered by tests
end
push!(new_stmts, Expr(:popaliasscope))
push!(new_stmts, :(return nothing))

Check warning on line 74 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L73-L74

Added lines #L73 - L74 were not covered by tests
def[:body] = Expr(
:let,
Expr(:block, let_constargs...),
body,
Expr(:block, new_stmts...),
)
return
end

struct WorkgroupLoop
stmts::Vector{Any}
terminated_in_sync::Bool
end

is_sync(expr) = @capture(expr, @synchronize() | @synchronize(a_))

Check warning on line 88 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L88

Added line #L88 was not covered by tests

function is_scope_construct(expr::Expr)
return expr.head === :block # ||

Check warning on line 91 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L90-L91

Added lines #L90 - L91 were not covered by tests
# expr.head === :let
end

function find_sync(stmt)
result = false
postwalk(stmt) do expr
result |= is_sync(expr)
expr

Check warning on line 99 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L95-L99

Added lines #L95 - L99 were not covered by tests
end
return result

Check warning on line 101 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L101

Added line #L101 was not covered by tests
end

# TODO proper handling of LineInfo
function split(emit, stmts)

Check warning on line 105 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L105

Added line #L105 was not covered by tests
# 1. Split the code into blocks separated by `@synchronize`

current = Any[]
new_stmts = Any[]
for stmt in stmts
has_sync = find_sync(stmt)
if has_sync
loop = WorkgroupLoop(current, is_sync(stmt))
push!(new_stmts, emit(loop))
current = Any[]
is_sync(stmt) && continue

Check warning on line 116 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L108-L116

Added lines #L108 - L116 were not covered by tests

# 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))
return Expr(expr.head, new_args...)

Check warning on line 127 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L122-L127

Added lines #L122 - L127 were not covered by tests
else
return Expr(expr.head, map(recurse, expr.args)...)

Check warning on line 129 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L129

Added line #L129 was not covered by tests
end
end
push!(new_stmts, recurse(stmt))
continue

Check warning on line 133 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L132-L133

Added lines #L132 - L133 were not covered by tests
end

push!(current, stmt)
end

Check warning on line 137 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L136-L137

Added lines #L136 - L137 were not covered by tests

# everything since the last `@synchronize`
if !isempty(current)
loop = WorkgroupLoop(current, false)
push!(new_stmts, emit(loop))

Check warning on line 142 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L140-L142

Added lines #L140 - L142 were not covered by tests
end
return new_stmts

Check warning on line 144 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L144

Added line #L144 was not covered by tests
end

function emit_gpu(loop)
stmts = Any[]

Check warning on line 148 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L147-L148

Added lines #L147 - L148 were not covered by tests

body = Expr(:block, loop.stmts...)
loopexpr = quote
if __active_lane__
$(unblock(body))

Check warning on line 153 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L150-L153

Added lines #L150 - L153 were not covered by tests
end
end
push!(stmts, loopexpr)
if loop.terminated_in_sync
push!(stmts, :($__synchronize()))

Check warning on line 158 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L156-L158

Added lines #L156 - L158 were not covered by tests
end

return unblock(Expr(:block, stmts...))

Check warning on line 161 in src/macros.jl

View check run for this annotation

Codecov / codecov/patch

src/macros.jl#L161

Added line #L161 was not covered by tests
end

0 comments on commit 8072f4c

Please sign in to comment.