diff --git a/Project.toml b/Project.toml index b7464b31..d1f1aeb4 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,8 @@ version = "0.4.3" [deps] JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" diff --git a/src/BenchmarkTools.jl b/src/BenchmarkTools.jl index 619e493d..e3f964a1 100644 --- a/src/BenchmarkTools.jl +++ b/src/BenchmarkTools.jl @@ -5,9 +5,11 @@ using Base.Iterators using Statistics using Printf +using MacroTools: MacroTools, prewalk, postwalk, @capture +using OrderedCollections: OrderedDict -const BENCHMARKTOOLS_VERSION = v"0.4.3" +const BENCHMARKTOOLS_VERSION = v"0.4.4" ############## # Parameters # @@ -63,7 +65,10 @@ export tune!, @benchmark, @benchmarkable, @belapsed, - @btime + @btime, + @localbenchmark, + @localbelapsed, + @localbtime ################# # Serialization # diff --git a/src/execution.jl b/src/execution.jl index 9d4f7058..34fab648 100644 --- a/src/execution.jl +++ b/src/execution.jl @@ -400,3 +400,110 @@ macro btime(args...) $result end) end + +#################### +# Local Benchmarks # +#################### + +# Credit to Robin Deits @rdeits for introducing this code in LocalScopeBenchmarks.jl + +function collect_symbols(expr) + assignments = OrderedDict{Symbol, Expr}() + prewalk(expr) do x + if x isa Symbol + assignments[x] = Expr(:$, x) + return nothing + elseif x isa Expr && x.head == :$ + # Don't recurse inside $() interpolations, + # since those will already be interpolated + return nothing + else + return x + end + end + assignments +end + +function parse_setup(setup::Expr) + assignments = OrderedDict() + postwalk(setup) do x + if @capture(x, a_ = b_) + assignments[a] = b + end + x + end + assignments +end + +function lower_setup(assignments::AbstractDict) + Expr(:block, [Expr(:(=), k, v) for (k, v) in assignments]...) +end + +function parse_params(kwargs) + params_dict = OrderedDict((@assert x.head == :kw; x.args[1] => x.args[2]) for x in kwargs) +end + +function lower_params(params::AbstractDict) + [Expr(:kw, k, v) for (k, v) in params] +end + +function interpolate_locals_into_setup(args...) + core, kwargs = BenchmarkTools.prunekwargs(args...) + params = parse_params(kwargs) + setup_assignments = parse_setup(get(() -> Expr(:block), params, :setup)) + local_assignments = collect_symbols(core) + setup = merge(local_assignments, setup_assignments) + params[:setup] = lower_setup(setup) + core, lower_params(params) +end + + +macro localbenchmark(args...) + core, params = interpolate_locals_into_setup(args...) + quote + BenchmarkTools.@benchmark($(core), $(params...)) + end +end + +""" + @localbtime expression [other parameters...] + +Similar to the `@time` macro included with Julia, +this executes an expression, printing the time +it took to execute and the memory allocated before +returning the value of the expression. + +Unlike `@time`, it uses the `@benchmark` +macro, and accepts all of the same additional +parameters as `@benchmark`. The printed time +is the *minimum* elapsed time measured during the benchmark. + +This macro allows you to use local scoping with the expression called. +Please see the tests for further examples. +""" +macro localbtime(args...) + core, params = interpolate_locals_into_setup(args...) + quote + BenchmarkTools.@btime($(core), $(params...)) + end +end + +""" + @localbelapsed expression [other parameters...] + +Similar to the `@elapsed` macro included with Julia, +this returns the elapsed time (in seconds) to +execute a given expression. It uses the `@benchmark` +macro, however, and accepts all of the same additional +parameters as `@benchmark`. The returned time +is the *minimum* elapsed time measured during the benchmark. + +This macro allows you to use local scoping within the expression called. +Please see the tests for further examples. +""" +macro localbelapsed(args...) + core, params = interpolate_locals_into_setup(args...) + quote + BenchmarkTools.@belapsed($(core), $(params...)) + end +end diff --git a/test/LocalScopeBenchmarkTests.jl b/test/LocalScopeBenchmarkTests.jl new file mode 100644 index 00000000..5e3e60bf --- /dev/null +++ b/test/LocalScopeBenchmarkTests.jl @@ -0,0 +1,142 @@ +using Test +using BenchmarkTools +using Statistics + +function judge_loosely(t1, t2) + judge(ratio(mean(t1), mean(t2)), time_tolerance=0.2) +end + +global_x = 1.0 + +@testset "LocalScopeBenchmarks" begin + @testset "Basic benchmarks" begin + x = 1.0 + evals = 500 + t1 = @benchmark($sin($x), evals=500) + t2 = @localbenchmark(sin(x), evals=500) + j = judge_loosely(t1, t2) + @test isinvariant(j) + + t1 = @benchmark($sin($x), evals=500) + t2 = @localbenchmark(sin(x), evals=500) + j = judge_loosely(t1, t2) + @test isinvariant(j) + + f = sin + x = 1.0 + t1 = @benchmark($f($x), evals=500) + t2 = @localbenchmark(f(x), evals=500) + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + +# This test fails to run if copy/pasted into the REPL due to differing LineNumbers where vars get +# pulled from. + @testset "Generated code is identical" begin + x = 1.0 + ex1 = Meta.@lower(@benchmark($sin($x), evals=500)) + ex2 = Meta.@lower(@localbenchmark(sin(x), evals=500)) + end + + @testset "Benchmarks with setup" begin + @testset "Single setup" begin + x =1.0 + t1 = @benchmark sin(x) setup=(x = 2.0) + t2 = @localbenchmark sin(x) setup=(x = 2.0) + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + + @testset "Multiple setups" begin + t1 = @benchmark atan(x, y) setup=(x = 2.0; y = 1.5) + t2 = @localbenchmark atan(x, y) setup=(x = 2.0; y = 1.5) + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + + @testset "Setups override local vars" begin + x = 1.0 + t1 = @benchmark (@assert x == 2.0) setup=(x = 2.0) evals = 500 + t2 = @localbenchmark (@assert x == 2.0) setup=(x = 2.0) evals=500 + j = judge_loosely(t1,t2) + @test isinvariant(j) + end + + @testset "Mixed setup and local vars" begin + x = 1.0 + t1 = @benchmark atan($x, y) setup=(y = 2.0) + t2 = @localbenchmark atan(x, y) setup=(y = 2.0) + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + @testset "Simple generators and comprehensions" begin + x = [i for i in 1:1000] + t1 = @benchmark sum($x) + t2 = @localbenchmark sum(x) + j = judge_loosely(t1, t2) + @test isinvariant(j) + + x = (i for i in 1:1000) + t1 = @benchmark sum($x) + t2 = @localbenchmark sum(x) + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + @testset "Gens, comps, override local vars" begin + x = [1.0, 1.0, 1.0] + y = [2.0, 2.0, 2.0] + t1 = @benchmark atan.($x, y) setup=(y = [2.0 for i in 1:3]) + t2 = @localbenchmark atan.(x, y) setup=(y = [2.0 for i in 1:3]) + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + end + @testset "Additional kwargs" begin + @testset "evals kwarg" begin + x = 1.0 + t1 = @benchmark sin($x) evals=5 + t2 = @localbenchmark sin(x) evals=5 + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + + @testset "evals and setup kwargs" begin + x = 1.0 + t1 = @benchmark sin($x) setup=(x = 2.0) evals=500 + t2 = @localbenchmark sin(x) setup=(x = 2.0) evals=500 + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + @testset "kwargs, evals and gens and comprehension filters" begin + f(x) = x # define some generators based on local scope + i = π + N = 3 + x = [1, 3] + y = [f(i) for i = 1:N if f(i) % 2 != 0] + t1 = @benchmark atan.($x, y) setup=(y = [$f(i) for i in 1:$N if $f(i) % 2 != 0]) evals=100 + t2 = @localbenchmark atan.(x, y) setup=(x = [1, 3]) evals=100 + j = judge_loosely(t1, t2) + @test isinvariant(j) + end + end + + @testset "Test that local benchmarks are faster than globals" begin + t1 = @benchmark sin(global_x) evals=5 # note the lack of $ + t2 = @localbenchmark sin(global_x) evals=5 + j = judge_loosely(t1, t2) + @test isregression(j) + end + + @testset "Other macros" begin + x = 1.0 + t1 = @localbtime sin($x) + t2 = @localbelapsed sin(x) + end + + @testset "Interpolated values" begin + t1 = @benchmark sum($(rand(1000))) + t2 = @localbenchmark sum($(rand(1000))) + j = judge_loosely(t1, t2) + @test isinvariant(j) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index e7259482..9a1bd7ed 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -17,3 +17,7 @@ println("done (took ", took_seconds, " seconds)") print("Testing serialization...") took_seconds = @elapsed include("SerializationTests.jl") println("done (took ", took_seconds, " seconds)") + +print("Testing LocalScopeBenchmarks") +took_seconds = @elapsed include("LocalScopeBenchmarkTests.jl") +println("done (took ", took_seconds, " seconds)")