diff --git a/Project.toml b/Project.toml index b442c669..caa37aae 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorOperations" uuid = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" -version = "5.6.2" +version = "5.6.3" authors = ["Lukas Devos ", "Maarten Van Damme ", "Jutho Haegeman "] [deps] diff --git a/docs/src/man/implementation.md b/docs/src/man/implementation.md index d5bc5b33..50902012 100644 --- a/docs/src/man/implementation.md +++ b/docs/src/man/implementation.md @@ -77,13 +77,17 @@ objects. ```@docs TensorOperations._flatten -TensorOperations.removelinenumbernode TensorOperations.addtensoroperations TensorOperations.insertargument TensorOperations.insertbackend TensorOperations.insertallocator ``` +Finally, after all postprocessors have run, the parser strips the `LineNumberNode`s that it +synthesized itself, while preserving the ones that came from the user's code. As a result the +generated code stays attributable to the lines of the original `@tensor` expression, which is +what makes those lines show up in code coverage reports and in stacktraces. + ## Analysis of contraction graphs and optimizing contraction order The macro [`@tensoropt`](@ref) or the combination of [`@tensor`](@ref) with the keyword diff --git a/ext/TensorOperationsBumperExt.jl b/ext/TensorOperationsBumperExt.jl index 12791765..b6713666 100644 --- a/ext/TensorOperationsBumperExt.jl +++ b/ext/TensorOperationsBumperExt.jl @@ -27,16 +27,14 @@ function TensorOperations._butensor(src, ex...) buf_sym = gensym("buffer") # TODO: there is no check for doubled tensor kwargs - newex = quote - $buf_sym = $(Expr(:call, GlobalRef(Bumper, :default_buffer))) - $( - Expr( - :macrocall, GlobalRef(TensorOperations, Symbol("@tensor")), - src, :(allocator = $buf_sym), ex... - ) + return Expr( + :block, + Expr(:(=), buf_sym, Expr(:call, GlobalRef(Bumper, :default_buffer))), + Expr( + :macrocall, GlobalRef(TensorOperations, Symbol("@tensor")), + src, :(allocator = $buf_sym), ex... ) - end - return Base.remove_linenums!(newex) + ) end end diff --git a/src/indexnotation/contractiontrees.jl b/src/indexnotation/contractiontrees.jl index 7f442f4a..8d3bce05 100644 --- a/src/indexnotation/contractiontrees.jl +++ b/src/indexnotation/contractiontrees.jl @@ -161,7 +161,7 @@ function insertcontractiontrees!( end ) end - push!(postexprs, removelinenumbernode(costcompareex)) + push!(postexprs, costcompareex) return treeex end diff --git a/src/indexnotation/parser.jl b/src/indexnotation/parser.jl index 6afc0729..7a0b6c4a 100644 --- a/src/indexnotation/parser.jl +++ b/src/indexnotation/parser.jl @@ -13,7 +13,7 @@ mutable struct TensorParser contractiontreebuilder = defaulttreebuilder contractiontreesorter = defaulttreesorter contractioncostcheck = nothing - postprocessors = [_flatten, removelinenumbernode, addtensoroperations] + postprocessors = [_flatten, addtensoroperations] return new( preprocessors, contractiontreebuilder, contractiontreesorter, contractioncostcheck, @@ -24,6 +24,9 @@ end function (parser::TensorParser)(ex::Expr) verifytensorexpr(ex) + # any `LineNumberNode` present here belongs to the user's code: record its file so that the + # ones synthesized further down can be told apart and removed again at the very end + userfiles = linenumberfiles(ex) for p in parser.preprocessors ex = p(ex)::Expr end @@ -35,6 +38,8 @@ function (parser::TensorParser)(ex::Expr) for p in parser.postprocessors ex = p(ex)::Expr end + # this has to happen after all (possibly user-supplied) postprocessors have run + ex = removeinternallinenumbernodes(ex, userfiles)::Expr return ex end diff --git a/src/indexnotation/postprocessors.jl b/src/indexnotation/postprocessors.jl index 25b4297a..ee03b5e8 100644 --- a/src/indexnotation/postprocessors.jl +++ b/src/indexnotation/postprocessors.jl @@ -31,19 +31,60 @@ function _flatten(ex) end """ - removelinenumbernode(ex) + linenumberfiles(ex, files = Set{Symbol}()) -Remove all `LineNumberNode`s from an expression. +Collect the set of files referenced by the `LineNumberNode`s in `ex`. + +This is used on the expression that is handed to a [`TensorParser`](@ref), before any +processing takes place, to determine which `LineNumberNode`s belong to the user's code: +see [`removeinternallinenumbernodes`](@ref). +""" +function linenumberfiles(ex, files = Set{Symbol}()) + if ex isa LineNumberNode + push!(files, ex.file) + elseif ex isa Expr + foreach(e -> linenumberfiles(e, files), ex.args) + end + return files +end + +""" + removeinternallinenumbernodes(ex, userfiles) + +Remove the `LineNumberNode`s that were synthesized by the parser, i.e. the ones whose file is +not in `userfiles`, as obtained from [`linenumberfiles`](@ref) on the original expression. + +`LineNumberNode`s originating from user code are kept, so that the generated code remains +attributable to the user's source lines. This matters for code coverage: Julia only emits a +coverage counter for a line that a `LineNumberNode` points at, so stripping the user's +`LineNumberNode`s leaves every statement of a `@tensor begin ... end` block after the first one +without any coverage information at all. Conversely, a synthesized `LineNumberNode` would +re-attribute all statements that follow it to a line in the parser's own source, so both halves +are needed. """ -function removelinenumbernode(ex) +function removeinternallinenumbernodes(ex, userfiles) if isexpr(ex, :block) - args = [removelinenumbernode(e) for e in ex.args if !(e isa LineNumberNode)] + # within a block, `LineNumberNode`s are statement markers: drop the internal ones + args = Any[ + removeinternallinenumbernodes(e, userfiles) for e in ex.args + if !_isinternallinenumber(e, userfiles) + ] return Expr(:block, args...) + elseif isa(ex, Expr) + # elsewhere a `LineNumberNode` may be structurally required -- most notably as the + # mandatory 2nd argument of a `:macrocall` -- so keep all positions here and only + # recurse into nested blocks + return Expr( + ex.head, Any[removeinternallinenumbernodes(e, userfiles) for e in ex.args]... + ) else return ex end end +_isinternallinenumber(@nospecialize(x), userfiles) = + x isa LineNumberNode && x.file ∉ userfiles + # list of functions that are used in expressions produced by `@tensor` const tensoroperationsfunctions = ( :tensoralloc, :tensorfree!, diff --git a/test/butensor.jl b/test/butensor.jl index fefcb211..ae658d09 100644 --- a/test/butensor.jl +++ b/test/butensor.jl @@ -6,6 +6,20 @@ end using Bumper +@testset "@butensor preserves user line numbers (issue #280)" begin + # `@butensor` wraps the block in an inner `@tensor`, so the expansion has to be recursive + # here to reach the code that the parser generated. This also covers the extension itself: + # `_butensor` must not introduce `LineNumberNode`s pointing into `ext/`. + firstline = @__LINE__() + 2 + block = @macroexpand @butensor begin + T[a, b] := X[a, c] * Y[c, b] + Z[a, b] := T[a, c] * W[c, b] + end + lnns = statementlinenumbernodes(block) + @test all(l -> l.file === Symbol(@__FILE__), lnns) + @test sort!(unique(l.line for l in lnns)) == collect(firstline .+ (0:1)) +end + @testset "Bumper tests with eltype $T" for T in (Float32, ComplexF64) D1, D2, D3 = 30, 40, 20 d1, d2 = 2, 3 diff --git a/test/macro_kwargs.jl b/test/macro_kwargs.jl index db577e49..5d45f13f 100644 --- a/test/macro_kwargs.jl +++ b/test/macro_kwargs.jl @@ -92,6 +92,91 @@ end end end +# https://github.com/QuantumKitHub/TensorOperations.jl/issues/280: the generated code must keep +# the user's `LineNumberNode`s (so `@tensor` lines show up in code coverage) while dropping the +# parser's own, which would steal that attribution. See `statementlinenumbernodes` in +# `runtests.jl` for why only block-statement positions are inspected. `@macroexpand1` is used +# throughout so that we look at what `@tensor` itself produced, and not at the expansion of +# macros that it merely passes through (e.g. the `@warn` of the `costcheck` path). +@testset "line numbers (issue #280)" begin + thisfile = Symbol(@__FILE__) + + @testset "single-statement expressions carry no LineNumberNodes" begin + # nothing to preserve here: the input has no `LineNumberNode`s of its own, and the + # statement is attributed to the line of the `@tensor` call by the surrounding scope + exprs = [ + @macroexpand1(@tensor T[a, b] := A[a, c] * B[c, b]), + @macroexpand1(@tensor R[a, b] := A[a, c] * B[c, d] * C[d, e] * E[e, f] * F[f, b]), + @macroexpand1(@tensor s = X[a, b] * Y[a, b]), + @macroexpand1(@tensoropt R[a, b] := A[a, c] * B[c, d] * C[d, e] * E[e, b]), + @macroexpand1(@tensor allocator = alloc R[a, b] := A[a, c] * B[c, d] * C[d, b]), + @macroexpand1(@tensor costcheck = warn R[a, b] := A[a, c] * B[c, d] * C[d, b]), + @macroexpand1(@tensor contractcheck = true R[a, b] := A[a, c] * B[c, b]), + ] + for ex in exprs + @test isempty(statementlinenumbernodes(ex)) + end + end + + @testset "one user LineNumberNode per statement of a block" begin + # multi-statement block, including a scalar assignment; `s = ...` additionally + # exercises the `_flatten` path that hoists a block into the right hand side + firstline = @__LINE__() + 2 + block = @macroexpand1 @tensor begin + T[a, e] := A[a, c] * B[c, d] * C[d, e] + D[a, b] := T[a, e] * E[e, b] + s = D[a, b] * F[a, b] + end + lnns = statementlinenumbernodes(block) + @test all(l -> l.file === thisfile, lnns) + @test sort!(unique(l.line for l in lnns)) == collect(firstline .+ (0:2)) + + # dst-reuse: `tensorify` wraps this in a `quote` of its own + reuseline = @__LINE__() + 2 + reuseblock = @macroexpand1 @tensor begin + T[a, b] := A[a, c] * B[c, b] + T[a, b] := T[a, c] * C[c, b] + end + reuselnns = statementlinenumbernodes(reuseblock) + @test all(l -> l.file === thisfile, reuselnns) + @test sort!(unique(l.line for l in reuselnns)) == collect(reuseline .+ (0:1)) + + optline = @__LINE__() + 2 + optblock = @macroexpand1 @tensoropt begin + T[a, e] := A[a, c] * B[c, d] * C[d, e] + D[a, b] := T[a, e] * E[e, b] + end + optlnns = statementlinenumbernodes(optblock) + @test all(l -> l.file === thisfile, optlnns) + @test sort!(unique(l.line for l in optlnns)) == collect(optline .+ (0:1)) + end + + @testset "kwargs that generate extra code preserve user line numbers" begin + # `allocator` inserts a checkpoint `quote`, `costcheck` inserts a `@notensor` block + # containing a `@warn` whose own `LineNumberNode` is structurally required + allocline = @__LINE__() + 2 + allocblock = @macroexpand1 @tensor allocator = alloc begin + T[a, b] := A[a, c] * B[c, b] + Z[a, b] := T[a, c] * C[c, b] + end + alloclnns = statementlinenumbernodes(allocblock) + @test all(l -> l.file === thisfile, alloclnns) + @test sort!(unique(l.line for l in alloclnns)) == collect(allocline .+ (0:1)) + + costline = @__LINE__() + 2 + costblock = @macroexpand1 @tensor costcheck = warn begin + T[a, b] := A[a, c] * B[c, b] + Z[a, b] := T[a, c] * C[c, b] + end + costlnns = statementlinenumbernodes(costblock) + @test all(l -> l.file === thisfile, costlnns) + @test sort!(unique(l.line for l in costlnns)) == collect(costline .+ (0:1)) + # the `@warn` macrocall must have kept its (TensorOperations-internal) LineNumberNode, + # otherwise the expression cannot be expanded at all + @test macroexpand(@__MODULE__, costblock; recursive = true) isa Expr + end +end + @testset "opt" begin A = randn(5, 5, 5, 5) B = randn(5, 5, 5) diff --git a/test/runtests.jl b/test/runtests.jl index 79432edd..55edbced 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,6 +11,24 @@ using TensorOperations: DefaultAllocator, ManualAllocator, BufferAllocator precision(::Type{<:Union{Float32, Complex{Float32}}}) = 1.0e-2 precision(::Type{<:Union{Float64, Complex{Float64}}}) = 1.0e-8 +# https://github.com/QuantumKitHub/TensorOperations.jl/issues/280: the generated code has to +# keep the user's `LineNumberNode`s -- Julia only emits a coverage counter for lines that a +# `LineNumberNode` points at -- and drop the ones synthesized by the parser, which would +# otherwise re-attribute the statements that follow them to a line in TensorOperations itself. +# Only `LineNumberNode`s in block-statement position matter for that: elsewhere (most notably +# the mandatory 2nd argument of a `:macrocall`) they are structurally required and left alone. +function statementlinenumbernodes(ex, acc = LineNumberNode[]) + if ex isa Expr + if ex.head === :block + for e in ex.args + e isa LineNumberNode && push!(acc, e) + end + end + foreach(e -> statementlinenumbernodes(e, acc), ex.args) + end + return acc +end + # don't run all tests on GPU, only the GPU # specific ones is_buildkite = get(ENV, "BUILDKITE", "false") == "true"