WIP: rebasing @localbtime and friends from LocalScopeBenchmarks.jl

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"
 

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 #

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

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

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)")