Fix line number handling in @kernel

src/KernelAbstractions.jl

@@ -66,7 +66,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
 
 """
@@ -88,7 +88,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
@@ -112,7 +112,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
 
@@ -267,9 +267,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
 
 """
@@ -279,9 +277,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
 
 """
@@ -293,9 +289,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
 
 """
@@ -314,9 +308,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
 
 """
@@ -350,9 +342,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
 
 """
@@ -372,9 +362,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
 
 """
@@ -458,9 +446,7 @@ macro print(items...)
         end
     end
 
-    return quote
-        $__print($(map(esc, args)...))
-    end
+    return :($__print($(map(esc, args)...)))
 end
 
 """

src/macros.jl

@@ -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]
@@ -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
@@ -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
@@ -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_))
@@ -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 implemet 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_))
@@ -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

test/coverage.jl

@@ -0,0 +1,114 @@
+module Coverage
+
+using KernelAbstractions
+using Test
+
+# Script exercising the constructs whose line information used to be dropped or
+# misattributed by `@kernel`. Every line of both kernel bodies must show up as
+# tracked in the resulting coverage data.
+const SCRIPT = """
+using KernelAbstractions
+
+@kernel function mul2!(a)
+    i = @index(Global)
+    x = a[i]
+    a[i] = 2 * x
+end
+
+@kernel function sync2!(a)
+    i = @index(Local)
+    lm = @localmem Float64 (8,)
+    lm[i] = a[i]
+    @synchronize
+    a[i] = lm[i] + 1
+end
+
+function main()
+    a = ones(8)
+    mul2!(CPU(), 8)(a, ndrange = 8)
+    sync2!(CPU(), 8)(a, ndrange = 8)
+    synchronize(CPU())
+    a == fill(3.0, 8) || error("unexpected result: \$a")
+    return
+end
+
+main()
+"""
+
+const SCRIPT_NAME = "kernels.jl"
+
+"""
+Line numbers of `SCRIPT_NAME` that the LCOV tracefile reports as tracked.
+
+Records look like `SF:<path>`, followed by one `DA:<line>,<count>` per tracked
+line, terminated by `end_of_record`.
+"""
+function tracked_lines(tracefile)
+    lines = Set{Int}()
+    in_script = false
+    for line in eachline(tracefile)
+        if startswith(line, "SF:")
+            in_script = basename(line[4:end]) == SCRIPT_NAME
+        elseif line == "end_of_record"
+            in_script = false
+        elseif in_script && startswith(line, "DA:")
+            push!(lines, parse(Int, first(split(line[4:end], ','))))
+        end
+    end
+    return lines
+end
+
+function run_covered(dir)
+    script = joinpath(dir, SCRIPT_NAME)
+    write(script, SCRIPT)
+    log = joinpath(dir, "log.txt")
+    # Write to an LCOV tracefile rather than using `--code-coverage=user`: the
+    # latter drops a `.cov` file next to every user source file it tracks, which
+    # would litter both the checkout and the depot, and would perturb the outer
+    # coverage report when the suite itself runs under `Pkg.test(coverage=true)`.
+    tracefile = joinpath(dir, "lcov.info")
+
+    cmd = `$(Base.julia_cmd()) --startup-file=no --code-coverage=$tracefile
+        --project=$(Base.active_project()) $script`
+    proc = run(pipeline(ignorestatus(cmd); stdout = log, stderr = log))
+    if !success(proc)
+        @error "coverage subprocess failed" output = read(log, String)
+    end
+    @test success(proc)
+    @test isfile(tracefile)
+    isfile(tracefile) || return nothing
+
+    return tracked_lines(tracefile)
+end
+
+function coverage_testsuite()
+    # GPUCompiler records device coverage by visiting the source location of
+    # every `:code_coverage_effect` while compiling, so a kernel whose line
+    # information points into KernelAbstractions reads as untracked in the
+    # user's file. https://github.com/JuliaGPU/KernelAbstractions.jl/issues/732
+    mktempdir() do dir
+        tracked = run_covered(dir)
+        tracked === nothing && return
+
+        srclines = split(SCRIPT, '\n')
+        @testset "$needle" for needle in (
+                "@kernel function mul2!",
+                "i = @index(Global)",
+                "x = a[i]",
+                "a[i] = 2 * x",
+                "@kernel function sync2!",
+                "i = @index(Local)",
+                "lm = @localmem Float64 (8,)",
+                "lm[i] = a[i]",
+                "@synchronize",
+                "a[i] = lm[i] + 1",
+            )
+            line = findfirst(src -> occursin(needle, src), srclines)
+            @test line !== nothing
+            @test line in tracked
+        end
+    end
+    return
+end
+
+end # module