Skip to content

Fix line number handling in @kernel #734

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vchuravy merged 3 commits into from
Sep 13, 2026
+327 −42
Merged

Fix line number handling in @kernel #734

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
42c3339
Fix line number handling in `@kernel`
vchuravy Jul 27, 2026
e7e049f
Add end-to-end coverage test for kernel line attribution
vchuravy Jul 27, 2026
3084d86
Do not inherit the parent's coverage mode in the coverage test
vchuravy Sep 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 10 additions & 24 deletions src/KernelAbstractions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ synchronize(dev)
```
"""
macro kernel(expr)
return __kernel(expr, #=force_inbounds=# false, #=unsafe_indices=# false)
return __kernel(expr, __source__, #=force_inbounds=# false, #=unsafe_indices=# false)
end

"""
Expand All @@ -92,7 +92,7 @@ This allows for two different configurations:
"""
macro kernel(ex...)
if length(ex) == 1
return __kernel(ex[1], false, false)
return __kernel(ex[1], __source__, false, false)
else
unsafe_indices = false
force_inbounds = false
Expand All @@ -116,7 +116,7 @@ macro kernel(ex...)
)
end
end
return __kernel(ex[end], force_inbounds, unsafe_indices)
return __kernel(ex[end], __source__, force_inbounds, unsafe_indices)
end
end

Expand Down Expand Up @@ -168,9 +168,7 @@ a tuple corresponding to kernel configuration. In order to get
the total size you can use `prod(@groupsize())`.
"""
macro groupsize()
return quote
$groupsize($(esc(:__ctx__)))
end
return :($groupsize($(esc(:__ctx__))))
end

"""
Expand All @@ -180,9 +178,7 @@ Query the ndrange on the backend. This function returns
a tuple corresponding to kernel configuration.
"""
macro ndrange()
return quote
$size($ndrange($(esc(:__ctx__))))
end
return :($size($ndrange($(esc(:__ctx__)))))
end

"""
Expand All @@ -194,9 +190,7 @@ macro localmem(T, dims)
# Stay in sync with CUDAnative
id = gensym("static_shmem")

return quote
$SharedMemory($(esc(T)), Val($(esc(dims))), Val($(QuoteNode(id))))
end
return :($SharedMemory($(esc(T)), Val($(esc(dims))), Val($(QuoteNode(id)))))
end

"""
Expand All @@ -215,9 +209,7 @@ macro private(T, dims)
if dims isa Integer
dims = (dims,)
end
return quote
$Scratchpad($(esc(:__ctx__)), $(esc(T)), Val($(esc(dims))))
end
return :($Scratchpad($(esc(:__ctx__)), $(esc(T)), Val($(esc(dims)))))
end

"""
Expand Down Expand Up @@ -251,9 +243,7 @@ workgroup.
`@synchronize()` must be encountered by all workitems of a work-group executing the kernel or by none at all.
"""
macro synchronize()
return quote
$__synchronize()
end
return :($__synchronize())
end

"""
Expand All @@ -273,9 +263,7 @@ workgroup. `cond` is not allowed to have any visible sideffects.
Since v`0.9.34` this version of the macro is deprecated and lowers to `@synchronize()`
"""
macro synchronize(cond)
return quote
$__synchronize()
end
return :($__synchronize())
end

"""
Expand Down Expand Up @@ -340,9 +328,7 @@ macro print(items...)
end
end

return quote
$__print($(map(esc, args)...))
end
return :($__print($(map(esc, args)...)))
end

"""
Expand Down
78 changes: 60 additions & 18 deletions src/macros.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,23 @@ function find_return(stmt)
return result[]
end

# `quote` blocks insert `LineNumberNode`s pointing into this file. Rewriting them
# to the `@kernel` call site keeps coverage and profiling pointed at the user's
# code instead of at KernelAbstractions internals.
relocate_lines(expr, source::LineNumberNode) =
postwalk(x -> x isa LineNumberNode ? source : x, expr)

# `MacroTools.unblock` drops `LineNumberNode`s when it collapses a block down to
# its single remaining statement. Only unwrap blocks that carry no line
# information, so that we never discard it.
function unblock_lines(ex)
isexpr(ex, :block) || return ex
length(ex.args) == 1 || return ex
return unblock_lines(ex.args[1])
end

# XXX: Proper errors
function __kernel(expr, force_inbounds = false, unsafe_indices = false)
function __kernel(expr, __source__::LineNumberNode, force_inbounds = false, unsafe_indices = false)
def = splitdef(expr)
name = def[:name]
args = def[:args]
Expand Down Expand Up @@ -46,6 +61,7 @@ function __kernel(expr, force_inbounds = false, unsafe_indices = false)
$name(dev, size::$_Size, range::$_Size) = $_name(dev, size, range)
end
end
constructors = relocate_lines(constructors, __source__)

return Expr(:block, esc(gpu_function), esc(constructors))
end
Expand All @@ -60,7 +76,8 @@ function transform_gpu!(def, constargs, force_inbounds, unsafe_indices)
end
end
pushfirst!(def[:args], :__ctx__)
new_stmts = Expr[]
# `Any[]`, since `split` hands back `LineNumberNode`s alongside `Expr`s
new_stmts = Any[]
body = MacroTools.flatten(def[:body])
push!(new_stmts, Expr(:aliasscope))
if !unsafe_indices
Expand Down Expand Up @@ -91,6 +108,7 @@ struct WorkgroupLoop
stmts::Vector{Any}
allocations::Vector{Any}
terminated_in_sync::Bool
sync_line::Union{Nothing, LineNumberNode}
end

is_sync(expr) = @capture(expr, @synchronize() | @synchronize(a_))
Expand All @@ -109,48 +127,73 @@ function find_sync(stmt)
return result[]
end

# TODO proper handling of LineInfo
function split(stmts)
# 1. Split the code into blocks separated by `@synchronize`

current = Any[]
allocations = Any[]
new_stmts = Any[]
# `LineNumberNode` belonging to the statement currently being processed.
# Statements are moved between `current` and `allocations` and the two end
# up in different scopes of the emitted code, so instead of copying the line
# information over eagerly we attach it to whichever list the statement
# lands in. Otherwise hoisted allocations lose their source location.
line = nothing
# Flush the pending `LineNumberNode` into `stmts`.
function take_line!(stmts)
line === nothing && return
push!(stmts, line)
line = nothing
return
end

for stmt in stmts
if stmt isa LineNumberNode
line = stmt
continue
end

has_sync = find_sync(stmt)
if has_sync
loop = WorkgroupLoop(current, allocations, is_sync(stmt))
loop = WorkgroupLoop(current, allocations, is_sync(stmt), line)
push!(new_stmts, emit(loop))
allocations = Any[]
current = Any[]
is_sync(stmt) && continue
if is_sync(stmt)
# `emit` consumed `line` for the `@synchronize` itself
line = nothing
continue
end

# Recurse into scope constructs
# TODO: This currently implements hard scoping
# probably need to implement soft scoping
# by not deepcopying the environment.
recurse(x) = x
function recurse(expr::Expr)
expr = unblock(expr)
expr = unblock_lines(expr)
if is_scope_construct(expr) && any(find_sync, expr.args)
new_args = unblock(split(expr.args))
return Expr(expr.head, new_args...)
return Expr(expr.head, split(expr.args)...)
else
return Expr(expr.head, map(recurse, expr.args)...)
end
end
take_line!(new_stmts)
push!(new_stmts, recurse(stmt))
continue
end

if @capture(stmt, @uniform x_)
take_line!(allocations)
push!(allocations, stmt)
continue
elseif @capture(stmt, @private lhs_ = rhs_)
take_line!(allocations)
push!(allocations, :($lhs = $rhs))
continue
elseif @capture(stmt, lhs_ = rhs_ | (vs__, lhs_ = rhs_))
if @capture(rhs, @localmem(args__) | @uniform(args__))
take_line!(allocations)
push!(allocations, stmt)
continue
elseif @capture(rhs, @private(T_, dims_))
Expand All @@ -161,36 +204,35 @@ function split(stmts)
dims = (dims,)
end
alloc = :($Scratchpad(__ctx__, $T, Val($dims)))
take_line!(allocations)
push!(allocations, :($lhs = $alloc))
continue
end
end

take_line!(current)
push!(current, stmt)
end

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

function emit(loop)
# Note: built without `quote`, since that would splice `LineNumberNode`s
# pointing at this file into the middle of the user's kernel body.
stmts = Any[]

body = Expr(:block, loop.stmts...)
loopexpr = quote
$(loop.allocations...)
if __active_lane__
$(unblock(body))
end
end
push!(stmts, loopexpr)
append!(stmts, loop.allocations)
push!(stmts, Expr(:if, :__active_lane__, Expr(:block, loop.stmts...)))
if loop.terminated_in_sync
loop.sync_line === nothing || push!(stmts, loop.sync_line)
push!(stmts, :($__synchronize()))
end

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