profiler: Add option to generate perfetto url (#575)

* profiler: Add option to generate perfetto url

* remove sort

* Update api.md

Project.toml

@@ -20,6 +20,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
+Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 
 [weakdeps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -71,6 +72,7 @@ Random123 = "1.7"
 ReactantCore = "0.1.4"
 Reactant_jll = "0.0.42"
 Scratch = "1.2"
+Sockets = "1.10"
 SpecialFunctions = "2.4"
 Statistics = "1.10"
 YaoBlocks = "0.13"

docs/src/api/api.md

@@ -30,7 +30,7 @@ Reactant.@jit
 ## Profile XLA
 
 ```@docs
-Reactant.with_profiler
+Reactant.Profiler.with_profiler
 ```
 
 # Internal Functionality

src/Profiler.jl

@@ -1,10 +1,16 @@
+module Profiler
+
+import ..Reactant
+using Sockets: Sockets
+
 """
-    with_profiler(f, trace_output_dir::String; trace_device=true, trace_host=true)
+    with_profiler(f, trace_output_dir::String; trace_device=true, trace_host=true, create_perfetto_link=false)
 
 Runs the provided function under a profiler for XLA (similar to [JAX's profiler](https://jax.readthedocs.io/en/latest/profiling.html)).
 The traces will be exported in the provided folder and can be seen
 using tools like [perfetto.dev](https://ui.perfetto.dev). It will return the return values
-from the function.
+from the function. The `create_perfetto_link` parameter can be used
+to automatically generate a perfetto url to visualize the trace.
 
 ```julia
 with_profiler("./traces/") do
@@ -13,20 +19,113 @@ with_profiler("./traces/") do
 end
 ```
 """
-function with_profiler(f, trace_output_dir::String; trace_device=true, trace_host=true)
+function with_profiler(
+    f,
+    trace_output_dir::String;
+    trace_device=true,
+    trace_host=true,
+    create_perfetto_link=false,
+)
     # TODO: we should be able to inject traces from Julia to fill in the blank spots in the trace.
 
     device_tracer_level = UInt32(trace_device ? 1 : 0)
     host_tracer_level = UInt32(trace_host ? 2 : 0)
     profiler = @ccall Reactant.MLIR.API.mlir_c.CreateProfilerSession(
         device_tracer_level::UInt32, host_tracer_level::UInt32
     )::Ptr{Cvoid}
-    try
+
+    results = try
         f()
     finally
         @ccall Reactant.MLIR.API.mlir_c.ProfilerSessionCollectData(
             profiler::Ptr{Cvoid}, trace_output_dir::Cstring
         )::Cvoid
         @ccall Reactant.MLIR.API.mlir_c.ProfilerSessionDelete(profiler::Ptr{Cvoid})::Cvoid
     end
+
+    if create_perfetto_link
+        traces_path = joinpath(trace_output_dir, "plugins", "profile")
+        date = maximum(readdir(traces_path))
+        traces_path = joinpath(traces_path, date)
+
+        filename = first(f for f in readdir(traces_path) if endswith(f, ".trace.json.gz"))
+        serve_to_perfetto(joinpath(traces_path, filename))
+    end
+
+    return results
+end
+
+export with_profiler
+
+function serve_to_perfetto(path_to_trace_file)
+    port_hint = 9001
+    port, server = Sockets.listenany(port_hint)
+
+    try
+        url = "https://ui.perfetto.dev/#!/?url=http://127.0.0.1:$(port)/$(basename(path_to_trace_file))"
+        @info "Open $url"
+        # open_in_default_browser(url)
+
+        while true
+            isopen(server) || break
+
+            io = Sockets.accept(server)
+            @debug "Got connection"
+            msg = String(readuntil(io, UInt8['\r', '\n', '\r', '\n']))
+            @debug "Got request" msg
+            if startswith(msg, "OPTIONS")
+                isopen(io) || continue
+                write(
+                    io,
+                    """
+          HTTP/1.1 501
+          Server: Reactant.jl
+          Access-Control-Allow-Origin: *
+          Content-Length: 0
+
+          """,
+                )
+                close(io)
+                continue
+            end
+            if startswith(msg, "POST")
+                isopen(io) || continue
+                write(
+                    io,
+                    """
+          HTTP/1.1 404
+          Server: Reactant.jl
+          Access-Control-Allow-Origin: *
+          Content-Length: 0
+
+          """,
+                )
+                close(io)
+                continue
+            end
+
+            file = read(path_to_trace_file)
+            file_size = length(file)
+
+            isopen(io) || continue
+            write(
+                io,
+                """
+      HTTP/1.1 200
+      Server: Reactant.jl
+      Access-Control-Allow-Origin: *
+      Content-Length: $(file_size)
+      Content-Type: application/gzip
+
+      """,
+            )
+
+            write(io, file)
+            break
+        end
+    finally
+        isopen(server) && close(server)
+    end
 end
+
+end # module Profiler

src/Reactant.jl

@@ -249,6 +249,8 @@ include("Tracing.jl")
 include("Compiler.jl")
 include("Profiler.jl")
 
+const with_profiler = Profiler.with_profiler
+
 include("Overlay.jl")
 
 function Enzyme.make_zero(