Add option to snoop commands to control spawning separate process vs eval'ing in existing process

SnoopCompileCore/src/SnoopCompileCore.jl

@@ -21,6 +21,7 @@ end
 
 if VERSION >= v"1.6.0-DEV.1192"  # https://github.com/JuliaLang/julia/pull/37136
     include("snoopl.jl")
+    include("snoop_all.jl")
 end
 
 end

SnoopCompileCore/src/snoop_all.jl

@@ -0,0 +1,36 @@
+"""
+```
+SnoopCompileCore.@snoop_all "output_filename" begin
+    # Commands to execute
+end
+```
+
+Runs provided commands *in the current process*, and records snoop compilation results for
+all phases of the compiler supported by SnoopCompile: Inference, LLVM Optimization, Codegen.
+
+Returns a tuple of four values:
+- `tinf`: Results from `@snoopi_deep` - See https://timholy.github.io/SnoopCompile.jl/stable/snoopi_deep/#Viewing-the-results
+- `"snoopl.csv"`: File 1 of results for `@snoopl` - See https://timholy.github.io/SnoopCompile.jl/stable/reference/#SnoopCompile.read_snoopl
+- `"snoopl.yaml"`: File 2 of results for `@snoopl` - See https://timholy.github.io/SnoopCompile.jl/stable/reference/#SnoopCompile.read_snoopl
+- `"snoopc.csv"`: Results file for `@snoopc` - See https://timholy.github.io/SnoopCompile.jl/stable/snoopc/
+
+# Example
+```julia
+julia> tinf, snoopl_csv, snoopl_yaml, snoopc_csv = @snoop_all "foo" foo(1, 2, 3);
+```
+"""
+macro snoop_all(fname_prefix, commands)
+    snoopl_csv_f, snoopl_yaml_f, snoopc_csv_f =
+        "$fname_prefix.snoopl.csv", "$fname_prefix.snoopl.yaml", "$fname_prefix.snoopc.csv"
+    tinf = gensym(:tinf)
+    esc(quote
+        @snoopc pspawn=false $snoopc_csv_f begin
+            @snoopl pspawn=false $snoopl_csv_f $snoopl_yaml_f begin
+                global $tinf = @snoopi_deep begin
+                    @eval $commands
+                end
+            end
+        end;
+        $tinf, $snoopl_csv_f, $snoopl_yaml_f, $snoopc_csv_f
+    end)
+end

SnoopCompileCore/src/snoopc.jl

@@ -4,33 +4,50 @@ using Serialization
 
 """
 ```
-@snoopc "compiledata.csv" begin
-    # Commands to execute, in a new process
+@snoopc [pspawn=true] [String["<julia_flags...>"]] "compiledata.csv" begin
+    # Commands to execute, either in a new process if pspawn==true, or via eval if false.
 end
 ```
 causes the julia compiler to log all functions compiled in the course
 of executing the commands to the file "compiledata.csv". This file
 can be used for the input to `SnoopCompile.read`.
+
+Julia flags can optionally be passed as a string or an array of strings, which will be set
+on the newly spawned julia process if `pspawn=true`. If `pspawn=false`, setting
+`julia_flags` will be ignored.
+
+If `pspawn=false`, the commands will be run in the same julia process, via `eval()` in
+the current module. This will only report _new_ compilations, that haven't already been
+cached in the current julia process, which can be (carefully) used to prune down the
+results to only the code you are interested in.
 """
-macro snoopc(flags, filename, commands)
-    return :(snoopc($(esc(flags)), $(esc(filename)), $(QuoteNode(commands))))
+macro snoopc(args...)
+    @assert 4 >= length(args) >= 2 """Usage: @snoopl [args...] "snoopl.csv" "snoopl.yaml" commands"""
+    flags, (filename, commands) = args[1:end-2], args[end-1:end]
+    pspawn_expr, julia_flags = begin
+        if length(flags) == 2
+            flags[1], flags[2]
+        elseif flags[1].head == :(=) && flags[1].args[1] == :pspawn
+            flags[1], String[]
+        else
+            :(pspawn=false), flags[1]
+        end
+    end
+    return :(snoopc($(esc(julia_flags)), $(esc(filename)), $(QuoteNode(commands)), $__module__; $(esc(pspawn_expr))))
 end
 macro snoopc(filename, commands)
     return :(snoopc(String[], $(esc(filename)), $(QuoteNode(commands))))
 end
 
-function snoopc(flags, filename, commands)
-    println("Launching new julia process to run commands...")
-    # addprocs will run the unmodified version of julia, so we
-    # launch it as a command.
+function snoopc(flags, filename, commands, _module=nothing; pspawn=true)
+    if pspawn
+        println("Launching new julia process to run commands...")
+    end
     code_object = """
             using Serialization
-            while !eof(stdin)
-                Core.eval(Main, deserialize(stdin))
-            end
+            Core.eval(Main, deserialize(IOBuffer(read(stdin))))
             """
-    process = open(`$(Base.julia_cmd()) $flags --eval $code_object`, stdout, write=true)
-    serialize(process, quote
+    record_and_run_quote = quote
         let io = open($filename, "w")
             ccall(:jl_dump_compiles, Nothing, (Ptr{Nothing},), io.handle)
             try
@@ -40,9 +57,20 @@ function snoopc(flags, filename, commands)
                 close(io)
             end
         end
-        exit()
-    end)
-    wait(process)
-    println("done.")
+    end
+
+    if pspawn
+        # addprocs will run the unmodified version of julia, so we
+        # launch it as a command.
+        process = open(`$(Base.julia_cmd()) $flags --eval $code_object`, stdout, write=true)
+        serialize(process, quote
+            $record_and_run_quote
+        end)
+        close(process)
+        wait(process)
+        println("done.")
+    else
+        Core.eval(_module, record_and_run_quote)
+    end
     nothing
 end

SnoopCompileCore/src/snoopl.jl

@@ -4,8 +4,8 @@ using Serialization
 
 """
 ```
-@snoopl "func_names.csv" "llvm_timings.yaml" begin
-    # Commands to execute, in a new process
+@snoopl [jlflags="..."] [pspawn=true] "func_names.csv" "llvm_timings.yaml" begin
+    # Commands to execute, either in a new process if pspawn==true, or via eval if false.
 end
 ```
 causes the julia compiler to log timing information for LLVM optimization during the
@@ -14,41 +14,58 @@ be used for the input to `SnoopCompile.read_snoopl("func_names.csv", "llvm_timin
 
 The logs contain the amount of time spent optimizing each "llvm module", and information
 about each module, where a module is a collection of functions being optimized together.
+
+If `pspawn=false`, the commands will be run in the same julia process, via `eval()` in
+the current module. This will only report _new_ compilations, that haven't already been
+cached in the current julia process, which can be (carefully) used to prune down the
+results to only the code you are interested in.
 """
-macro snoopl(flags, func_file, llvm_file, commands)
-    return :(snoopl($(esc(flags)), $(esc(func_file)), $(esc(llvm_file)), $(QuoteNode(commands))))
+macro snoopl(args...)
+    @assert length(args) >= 3 """Usage: @snoopl [args...] "snoopl.csv" "snoopl.yaml" commands"""
+    flags, (func_file, llvm_file, commands) = args[1:end-3], args[end-2:end]
+    flags = [esc(e) for e in flags]
+    return :(snoopl($(esc(func_file)), $(esc(llvm_file)), $(QuoteNode(commands)), $__module__; $(flags...)))
 end
 macro snoopl(func_file, llvm_file, commands)
-    return :(snoopl(String[], $(esc(func_file)), $(esc(llvm_file)), $(QuoteNode(commands))))
+    return :(snoopl($(esc(func_file)), $(esc(llvm_file)), $(QuoteNode(commands)), $__module__))
 end
 
-function snoopl(flags, func_file, llvm_file, commands)
-    println("Launching new julia process to run commands...")
+function snoopl(func_file, llvm_file, commands, _module; pspawn=true, jlflags="")
+    if pspawn
+        println("Launching new julia process to run commands...")
+    end
     # addprocs will run the unmodified version of julia, so we
     # launch it as a command.
     code_object = """
             using Serialization
-            while !eof(stdin)
-                Core.eval(Main, deserialize(stdin))
-            end
+            Core.eval(Main, deserialize(IOBuffer(read(stdin))))
             """
-    process = open(`$(Base.julia_cmd()) $flags --eval $code_object`, stdout, write=true)
-    serialize(process, quote
+    record_and_run_quote = quote
         let func_io = open($func_file, "w"), llvm_io = open($llvm_file, "w")
             ccall(:jl_dump_emitted_mi_name, Nothing, (Ptr{Nothing},), func_io.handle)
             ccall(:jl_dump_llvm_opt, Nothing, (Ptr{Nothing},), llvm_io.handle)
             try
-                $commands
+                @eval $commands
             finally
                 ccall(:jl_dump_emitted_mi_name, Nothing, (Ptr{Nothing},), C_NULL)
                 ccall(:jl_dump_llvm_opt, Nothing, (Ptr{Nothing},), C_NULL)
                 close(func_io)
                 close(llvm_io)
             end
         end
-        exit()
-    end)
-    wait(process)
-    println("done.")
+    end
+
+    if pspawn
+        process = open(`$(Base.julia_cmd()) $jlflags --eval $code_object`, stdout, write=true)
+        @info process
+        serialize(process, quote
+            $record_and_run_quote
+        end)
+        close(process)
+        wait(process)
+        println("done.")
+    else
+        Core.eval(_module, record_and_run_quote)
+    end
     nothing
 end

test/runtests.jl

@@ -14,6 +14,9 @@ if VERSION >= v"1.6.0-DEV.1192"  # https://github.com/JuliaLang/julia/pull/37136
     @testset "snoopl" begin
         include("snoopl.jl")
     end
+    @testset "snoop_all" begin
+        include("snoop_all.jl")
+    end
 end
 
 using SnoopCompile

test/snoop_all.jl

@@ -0,0 +1,34 @@
+using Test
+
+using SnoopCompile
+
+f(x) = 2^x + 100
+
+@testset "basic snoop_all" begin
+    # First time not empty
+    tinf, snoopl_csv, snoopl_yaml, snoopc_csv =
+        SnoopCompileCore.@snoop_all "snoop_all-f" f(2)
+
+    @test length(collect(flatten(tinf))) > 1
+    @test filesize(snoopl_csv) != 0
+    @test filesize(snoopl_yaml) != 0
+    @test filesize(snoopc_csv) != 0
+
+    rm(snoopl_csv)
+    rm(snoopl_yaml)
+    rm(snoopc_csv)
+
+    # Second run is empty because f(x) is already compiled
+    tinf, snoopl_csv, snoopl_yaml, snoopc_csv =
+        SnoopCompileCore.@snoop_all "snoop_all-f" f(2)
+
+    @test length(collect(flatten(tinf))) == 1
+    @test filesize(snoopl_csv) == 0
+    @test filesize(snoopl_yaml) == 0
+    @test filesize(snoopc_csv) == 0
+
+    rm(snoopl_csv)
+    rm(snoopl_yaml)
+    rm(snoopc_csv)
+end
+