Add initial flamegraph support

Disabled for now because I'll need to implement Rust Analyzer-based path
resolution (converting function names to filenames and locations) in
order for the LLM to make use of the flamegraph.

Author: camelid

README.md

@@ -5,29 +5,26 @@
 1. Git clone
 2. Install [`uv`](https://github.com/astral-sh/uv) if not already installed
 3. To fix perf debuginfo issues: `cargo install addr2line --features="bin"`
+4. If you want support for sending flamegraphs to the LLM:
+    a. `cargo install flamegraph`
+    a. `cargo install resvg`
 
 ## Basic usage
 
-First, build your target program with optimizations on and debuginfo enabled, and then profile it with `perf` using something like the following:
-
-```console
-$ perf record -F99 --call-graph dwarf ./your-program
-```
-
-Then, in the `accelerant` repository, run:
+In the `accelerant` repository, run:
 
 ```console
 $ uv run accelerant_server.py
 ```
 
-Finally, in a separate terminal, run:
+In a separate terminal, run:
 
 ```console
-$ curl 'http://127.0.0.1:5000/optimize?project=PATH_TO_PROJECT_ROOT&perfDataPath=ABSOLUTE_PATH_TO_PERF_DATA'
+$ curl 'http://127.0.0.1:5000/optimize?project=PATH_TO_PROJECT_ROOT&targetBinary=target/release/REST_OF_PATH_TO_EXECUTABLE_TO_OPTIMIZE'
 ```
 
-Alternatively, you can ask to optimize a specific line without `perf` information using the following:
+Accelerant will automatically build, run, and profile your project using `cargo` and `perf`.
 
-```console
-$ curl 'http://127.0.0.1:5000/optimize?project=PATH_TO_PROJECT_ROOT&filename=RELATIVE_PATH_TO_FILE_IN_PROJECT&line=LINE_NUMBER_IN_FILE'
-```
+If you've already run the `perf` profiler and collected a `perf.data` file, you can give it to Accelerant by appending a `perfDataPath` query parameter with the path to the file.
+
+Also, if you know a particular line in your project is a hotspot, you can pass the (relative) path to its containing file in a `filename` paramater, with the line number in `line`.

accelerant/agent.py

@@ -35,6 +35,7 @@ def run_agent(
         tools.edit_code,
         tools.check_codebase_for_errors,
         tools.run_perf_profiler,
+        # tools.generate_flamegraph,
         tools.get_info,
         tools.get_references,
         tools.get_surrounding_code,

accelerant/flamegraph.py

@@ -0,0 +1,54 @@
+import base64
+from pathlib import Path
+import subprocess
+from tempfile import NamedTemporaryFile
+import re
+
+
+def make_flamegraph_png(perf_data_path: Path) -> bytes:
+    svg_str = make_flamegraph_svg(perf_data_path)
+    png_data = svg_to_png(svg_str)
+    return png_data
+
+
+def make_flamegraph_svg(perf_data_path: Path) -> str:
+    with NamedTemporaryFile(suffix=".svg") as output_svg_temp:
+        subprocess.run(
+            [
+                "flamegraph",
+                "--perfdata",
+                perf_data_path,
+                "--output",
+                output_svg_temp.name,
+            ],
+            check=True,
+        )
+        output_svg_temp.seek(0)
+        svg_data = output_svg_temp.read().decode()
+    return svg_data
+
+
+def svg_to_png(svg_str: str) -> bytes:
+    # HACK: resvg doesn't understand the monospace font-family, so replace it with concrete fonts
+    svg_str = re.sub(
+        "font-family: ?monospace",
+        "font-family: 'Fira Mono', 'DejaVu Sans Mono', 'Ubuntu Mono'",
+        svg_str,
+    )
+    with NamedTemporaryFile(suffix=".svg") as svg_temp:
+        svg_temp.write(svg_str.encode())
+        svg_temp.flush()
+        with NamedTemporaryFile(suffix=".png") as png_temp:
+            subprocess.run(
+                ["resvg", "--zoom=2", svg_temp.name, png_temp.name],
+                check=True,
+            )
+            png_temp.seek(0)
+            png_data = png_temp.read()
+    return png_data
+
+
+def png_to_data_url(png_data: bytes) -> str:
+    b64_encoded = base64.b64encode(png_data).decode()
+    data_url = f"data:image/png;base64,{b64_encoded}"
+    return data_url

accelerant/llm.py

@@ -16,11 +16,19 @@ def on_trace_end(self, trace):
         del self.active_traces[trace.trace_id]
 
     def on_span_start(self, span):
-        print(f"[blue]Starting span:[/blue] {span.span_data.export()}")
+        data = span.span_data.export()
+        if "data:image/png;base64" in str(data):
+            print("[blue]Starting span that includes image data:[/blue]")
+        else:
+            print(f"[blue]Starting span:[/blue] {data}")
         self.active_spans[span.span_id] = span
 
     def on_span_end(self, span):
-        print(f"[magenta]Ending span:[/magenta] {span.span_data.export()}")
+        data = span.span_data.export()
+        if "data:image/png;base64" in str(data):
+            print("[magenta]Ending span that includes image data:[/magenta]")
+        else:
+            print(f"[magenta]Ending span:[/magenta] {data}")
         del self.active_spans[span.span_id]
 
     def shutdown(self):

accelerant/perf.py

@@ -6,11 +6,16 @@
 
 
 class PerfData:
+    _path: Path
     _data: AttributedPerf
 
     def __init__(self, perf_data_path: Path, project_root: Path):
+        self._path = perf_data_path
         self._data = get_perf_data(str(perf_data_path), str(project_root))
 
+    def data_path(self) -> Path:
+        return self._path
+
     def lookup_pct_time(self, loc: LineLoc) -> Optional[float]:
         if loc not in self._data.hit_count:
             return None

accelerant/project.py

@@ -41,14 +41,21 @@ def lsp(self) -> LSP:
         return self._lsp
 
     def perf_data(self, version: Optional[FsVersion] = None) -> Optional[PerfData]:
+        perf_data_path = self.perf_data_path(version)
+        if perf_data_path is None:
+            return None
+
+        if perf_data_path not in self._perf_data_map:
+            self._perf_data_map[perf_data_path] = PerfData(perf_data_path, self._root)
+        return self._perf_data_map[perf_data_path]
+
+    def perf_data_path(self, version: Optional[FsVersion] = None) -> Optional[Path]:
         if version is None:
             version = self.fs_sandbox().version()
         if version not in self._perf_per_version:
             return None
         perf_data_path = self._perf_per_version[version]
-        if perf_data_path not in self._perf_data_map:
-            self._perf_data_map[perf_data_path] = PerfData(perf_data_path, self._root)
-        return self._perf_data_map[perf_data_path]
+        return perf_data_path
 
     def add_perf_data(self, version: FsVersion, perf_data_path: Path) -> None:
         self._perf_per_version[version] = perf_data_path

accelerant/tools.py

@@ -2,14 +2,16 @@
 from itertools import islice
 import shutil
 import subprocess
-from typing import Any, Optional
-from agents import RunContextWrapper, function_tool
+from typing import Optional
+from agents import RunContextWrapper, ToolOutputImage, function_tool
 from llm_utils import number_group_of_lines
 from perfparser import LineLoc
 
 from accelerant.chat_interface import CodeSuggestion
+from accelerant.flamegraph import make_flamegraph_png, png_to_data_url
 from accelerant.lsp import TOP_LEVEL_SYMBOL_KINDS, uri_to_relpath
 from accelerant.patch import apply_simultaneous_suggestions
+from accelerant.perf import PerfData
 from accelerant.util import find_symbol, truncate_for_llm
 from accelerant.project import Project
 
@@ -56,50 +58,65 @@ def check_codebase_for_errors(
     return "OK: Codebase has no errors!"
 
 
+def _shared_build_and_run_perf(project: Project) -> PerfData:
+    version = project.fs_sandbox().version()
+    perf_data = project.perf_data(version)
+    if perf_data is None:
+        project.build_for_profiling()
+        project.run_profiler()
+        perf_data = project.perf_data(version)
+    assert perf_data is not None, "perf data should be available after profiling"
+    return perf_data
+
+
 @function_tool
 def run_perf_profiler(
     ctx: RunContextWrapper[AgentContext],
-) -> list[dict[str, Any]]:
+) -> list[dict]:
     """Run a performance profiler on the target binary and return the top hotspots."""
-    try:
-        project = ctx.context.project
-        version = project.fs_sandbox().version()
-        perf_data = project.perf_data(version)
-        if perf_data is None:
-            project.build_for_profiling()
-            project.run_profiler()
-            perf_data = project.perf_data(version)
-        assert perf_data is not None, "perf data should be available after profiling"
-        perf_tabulated = perf_data.tabulate()
-        NUM_HOTSPOTS = 5
-
-        def get_parent_region(loc: LineLoc) -> Optional[str]:
-            parent_sym = project.lsp().syncexec(
-                project.lsp().request_nearest_parent_symbol(
-                    loc.path, loc.line - 1, TOP_LEVEL_SYMBOL_KINDS
-                ),
-            )
-            if parent_sym is None:
-                return None
-            return parent_sym["name"]
-
-        hotspots = list(
-            islice(
-                map(
-                    lambda x: {
-                        "parent_region": get_parent_region(x[0]) or "<unknown>",
-                        "loc": x[0],
-                        "pct_time": x[1] * 100,
-                    },
-                    filter(lambda x: x[0].line > 0, perf_tabulated),
-                ),
-                NUM_HOTSPOTS,
-            )
+    project = ctx.context.project
+    perf_data = _shared_build_and_run_perf(project)
+    perf_tabulated = perf_data.tabulate()
+    NUM_HOTSPOTS = 5
+
+    def get_parent_region(loc: LineLoc) -> Optional[str]:
+        parent_sym = project.lsp().syncexec(
+            project.lsp().request_nearest_parent_symbol(
+                loc.path, loc.line - 1, TOP_LEVEL_SYMBOL_KINDS
+            ),
         )
-        return hotspots
-    except Exception as e:
-        print("ERROR", e)
-        raise e
+        if parent_sym is None:
+            return None
+        return parent_sym["name"]
+
+    hotspots = list(
+        islice(
+            map(
+                lambda x: {
+                    "parent_region": get_parent_region(x[0]) or "<unknown>",
+                    "loc": x[0],
+                    "pct_time": x[1] * 100,
+                },
+                filter(lambda x: x[0].line > 0, perf_tabulated),
+            ),
+            NUM_HOTSPOTS,
+        )
+    )
+    return hotspots
+
+
+@function_tool
+def generate_flamegraph(
+    ctx: RunContextWrapper[AgentContext],
+) -> ToolOutputImage:
+    """Generate a flamegraph PNG image from the performance data, building the project and running the profiler if necessary."""
+    project = ctx.context.project
+    perf_data = _shared_build_and_run_perf(project)
+
+    flamegraph_data = make_flamegraph_png(perf_data.data_path())
+    flamegraph_data_url = png_to_data_url(flamegraph_data)
+    flamegraph_output = ToolOutputImage(image_url=flamegraph_data_url, detail="high")
+    return flamegraph_output
 
 
 @function_tool