llvmpipe: FlameGraph profiling support

For use outside Linux.

Enable dumping JIT address mappings and assembly
using an environment variable. Add a script to map
JIT addresses in collapsed stacks, and annotate
assembly dumps with sample counts.

Reviewed-by: Konstantin Seurer <konstantin.seurer@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30626>

bin/flamegraph_map_lp_jit.py

@@ -0,0 +1,142 @@
+#
+# Copyright 2024 Autodesk, Inc.
+#
+# SPDX-License-Identifier: MIT
+#
+
+import argparse
+from bisect import bisect_left, bisect_right
+from dataclasses import dataclass
+from pathlib import Path
+import re
+
+
+@dataclass
+class Instruction:
+    address: int
+    assembly: str
+    samples: int = 0
+
+
+def mapping_address_key(mapping: tuple[int, int, str]):
+    return mapping[0]
+
+
+def instruction_address_key(instruction: Instruction):
+    return instruction.address
+
+
+def parse_mappings(map_file_path: Path):
+    mappings: list[tuple[int, int, str]] = []
+    with open(map_file_path) as map_file:
+        for mapping in map_file:
+            address_hex, size_hex, name = mapping.split(' ')
+            address = int(address_hex, base=16)
+            mappings.append((address, address + int(size_hex, base=16), name.strip()))
+
+    mappings.sort(key=mapping_address_key)
+    return mappings
+
+
+def parse_traces(trace_file_path: Path):
+    pattern = re.compile(r'((?:[^;]+;)*?[^;]+) (\d+)\n')
+
+    traces: list[tuple[list[str], int]] = []
+    with open(trace_file_path) as trace_file:
+        for trace in trace_file:
+            match = pattern.fullmatch(trace)
+            traces.append((match.group(1).split(';'), int(match.group(2))))
+
+    return traces
+
+
+def parse_asm(asm_file_path: Path):
+    symbol_pattern = re.compile(r'(\w+) ([0-9a-fA-F]+):\n')
+    instruction_pattern = re.compile(r' *([0-9a-fA-F]+):\t(.*?)\n')
+
+    asm: dict[tuple[int, str], list[Instruction]] = {}
+    with open(asm_file_path) as asm_file:
+        current_instructions = None
+        for line in asm_file:
+            if match := symbol_pattern.fullmatch(line):
+                symbol = (int(match.group(2), base=16), match.group(1))
+                current_instructions = asm[symbol] = []
+            elif match := instruction_pattern.fullmatch(line):
+                current_instructions.append(Instruction(int(match.group(1), base=16), match.group(2)))
+
+    return asm
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Map LLVMPipe JIT addresses in FlameGraph style '
+                                     'collapsed stack traces to their symbol name. Also optionally '
+                                     'annotate JIT assembly dumps with sample counts.')
+    parser.add_argument('jit_symbol_map', type=Path, help='JIT symbol map from LLVMPipe')
+    parser.add_argument('collapsed_traces', type=Path)
+    parser.add_argument('-a', '--asm', type=Path, nargs='?', const='', metavar='asm_path',
+                        help='JIT assembly dump from LLVMPipe. Defaults to "<jit_symbol_map>.asm"')
+    parser.add_argument('-o', '--out', type=Path, metavar='out_path')
+    arguments = parser.parse_args()
+
+    mappings = parse_mappings(arguments.jit_symbol_map)
+    traces = parse_traces(arguments.collapsed_traces)
+
+    asm = {}
+    asm_file_path: Path | None = arguments.asm
+    if asm_file_path:
+        if len(asm_file_path.parts) <= 0:
+            asm_file_path = Path(str(arguments.jit_symbol_map) + '.asm')
+            if asm_file_path.exists():
+                asm = parse_asm(asm_file_path)
+        else:
+            asm = parse_asm(asm_file_path)
+
+    merged_traces: dict[str, int] = {}
+    for stack, count in traces:
+        for i, function in enumerate(stack):
+            if not function.startswith('0x'):
+                continue
+
+            address = int(function, base=16)
+            mapping = mappings[bisect_right(mappings, address, key=mapping_address_key) - 1]
+            if address < mapping[0] or address >= mapping[1]:
+                continue
+
+            stack[i] = f'lp`{mapping[2]}@{mapping[0]:x}'
+
+            symbol = (mapping[0], mapping[2])
+            if symbol in asm:
+                instructions = asm[symbol]
+                instruction_address = address - symbol[0]
+                index = bisect_left(instructions, instruction_address, key=instruction_address_key)
+                if index < len(instructions) and instructions[index].address == instruction_address:
+                    instructions[index].samples += count
+
+        stack_key = ';'.join(stack)
+        if stack_key in merged_traces:
+            merged_traces[stack_key] += count
+        else:
+            merged_traces[stack_key] = count
+
+    out_file_path: Path | None = arguments.out
+    if not out_file_path:
+        out_file_path = arguments.collapsed_traces.with_stem(f'{arguments.collapsed_traces.stem}_mapped')
+    with open(out_file_path, 'w') as out:
+        for t, c in merged_traces.items():
+            print(f'{t} {c}', file=out)
+
+    if asm:
+        annotated_asm_file_path = asm_file_path.with_stem(f'{asm_file_path.stem}_annotated')
+        with open(annotated_asm_file_path, 'w') as out:
+            for symbol, instructions in asm.items():
+                print(f'{symbol[1]}: ;{symbol[0]:x}', file=out)
+                for instruction in instructions:
+                    print(f'\t{instruction.assembly}', end='', file=out)
+                    if instruction.samples:
+                        print(f' ;s {instruction.samples}', file=out)
+                    else:
+                        print(file=out)
+                print(file=out)
+
+if __name__ == '__main__':
+    main()

docs/drivers/llvmpipe.rst

@@ -256,6 +256,25 @@ generated code annotated with the samples.
 You can obtain a call graph via
 `Gprof2Dot <https://github.com/jrfonseca/gprof2dot#linux-perf>`__.
 
+FlameGraph support
+~~~~~~~~~~~~~~~~~~~~~~
+
+Outside Linux, it is possible to generate a
+`FlameGraph https://github.com/brendangregg/FlameGraph`__:
+with resolved JIT symbols.
+
+Set the environment variable ``JIT_SYMBOL_MAP_DIR`` to a directory path,
+and run your LLVMpipe program. Follow the FlameGraph instructions:
+capture traces using a supported tool (for example DTrace),
+and fold the stacks using the associated script
+(``stackcollapse.pl`` for DTrace stacks).
+
+LLVMpipe will create a ``jit-symbols-XXXXX.map`` file containing the symbol
+address table inside the chosen directory. It will also dump the JIT
+disassemblies to ``jit-symbols-XXXXX.map.asm``. Run your folded traces and
+both output files through the ``bin/flamegraph_map_lp_jit.py`` script to map
+addresses to JIT symbols, and annotate the disassembly with the sample counts.
+
 Unit testing
 ------------
 

src/gallium/auxiliary/gallivm/lp_bld_debug.cpp

@@ -234,11 +234,14 @@ lp_disassemble(LLVMValueRef func, const void *code)
 extern "C" void
 lp_profile(LLVMValueRef func, const void *code)
 {
-#if defined(__linux__) && defined(PROFILE)
+#if defined(PROFILE)
    static std::ofstream perf_asm_file;
    static bool first_time = true;
    static FILE *perf_map_file = NULL;
    if (first_time) {
+      unsigned long long pid = (unsigned long long)getpid();
+      char filename[1024];
+#if defined(__linux__)
       /*
        * We rely on the disassembler for determining a function's size, but
        * the disassembly is a leaky and slow operation, so avoid running
@@ -246,13 +249,19 @@ lp_profile(LLVMValueRef func, const void *code)
        * by the PERF_BUILDID_DIR environment variable.
        */
       if (getenv("PERF_BUILDID_DIR")) {
-         pid_t pid = getpid();
-         char filename[256];
-         snprintf(filename, sizeof filename, "/tmp/perf-%llu.map", (unsigned long long)pid);
+         snprintf(filename, sizeof(filename), "/tmp/perf-%llu.map", pid);
          perf_map_file = fopen(filename, "wt");
-         snprintf(filename, sizeof filename, "/tmp/perf-%llu.map.asm", (unsigned long long)pid);
+         snprintf(filename, sizeof(filename), "/tmp/perf-%llu.map.asm", pid);
          perf_asm_file.open(filename);
       }
+#else
+      if (const char* output_dir = getenv("JIT_SYMBOL_MAP_DIR")) {
+         snprintf(filename, sizeof(filename), "%s/jit-symbols-%llu.map", output_dir, pid);
+         perf_map_file = fopen(filename, "wt");
+         snprintf(filename, sizeof(filename), "%s/jit-symbols-%llu.map.asm", output_dir, pid);
+         perf_asm_file.open(filename);
+      }
+#endif
       first_time = false;
    }
    if (perf_map_file) {