Improve GPU runtime trace output

Author: sbryngelson

src/common/m_trace_runtime.c

@@ -71,6 +71,7 @@ static int trace_file_fd = -1;
 static void mfc_trace_initialize_once(void) __attribute__((no_instrument_function));
 static void mfc_trace_initialize(void) __attribute__((no_instrument_function));
 static int mfc_trace_process_rank(void) __attribute__((no_instrument_function));
+static void mfc_trace_executable_name(char *buffer, size_t buffer_size) __attribute__((no_instrument_function));
 static void mfc_trace_format_function_name(const char *name, char *buffer, size_t buffer_size) __attribute__((no_instrument_function));
 static void mfc_trace_pretty_file(const char *file, char *buffer, size_t buffer_size) __attribute__((no_instrument_function));
 static void mfc_trace_shell_quote(const char *text, char *buffer, size_t buffer_size) __attribute__((no_instrument_function));
@@ -150,6 +151,12 @@ static void mfc_trace_initialize_once(void) {
             (void)write(STDERR_FILENO, message, sizeof(message) - 1);
         }
     }
+
+    if (trace_enabled) {
+        char executable[MFC_TRACE_NAME_MAX];
+        mfc_trace_executable_name(executable, sizeof(executable));
+        mfc_trace_write_line("TRACE_RUN", executable, NULL, 0);
+    }
 }
 
 static void mfc_trace_initialize(void) {
@@ -181,6 +188,24 @@ static int mfc_trace_process_rank(void) {
     return 0;
 }
 
+static void mfc_trace_executable_name(char *buffer, size_t buffer_size) {
+    char path[1024];
+    const char *name;
+    ssize_t len;
+
+    if (buffer_size == 0) return;
+
+    len = readlink("/proc/self/exe", path, sizeof(path) - 1);
+    if (len > 0) {
+        path[len] = '\0';
+        name = strrchr(path, '/');
+        snprintf(buffer, buffer_size, "%s", name != NULL ? name + 1 : path);
+        return;
+    }
+
+    snprintf(buffer, buffer_size, "program");
+}
+
 static void mfc_trace_format_function_name(const char *name, char *buffer, size_t buffer_size) {
     const char *mod_marker;
     const char *mp_marker;
@@ -407,6 +432,10 @@ static int mfc_trace_skip_symbol(const char *name) {
            strstr(name, "mfc_trace_") != NULL ||
            strstr(name, "m_trace:") != NULL ||
            strstr(name, "m_trace_s_") != NULL ||
+           strstr(name, "m_nvtx_") != NULL ||
+           strncmp(name, "..acc_", 6) == 0 ||
+           strstr(name, "acc_cuda_funcreg_constructor") != NULL ||
+           strstr(name, "acc_data_constructor") != NULL ||
            strcmp(name, "s_trace_point_begin") == 0 ||
            strcmp(name, "s_trace_point_end") == 0;
 }

toolchain/mfc/run/run.py

@@ -61,6 +61,15 @@ def __profiler_prepend() -> typing.List[str]:
     return []
 
 
+def __runner_prepend() -> typing.List[str]:
+    profiler = __profiler_prepend()
+
+    if __runtime_trace_requested() and ARG("gpu") == gpuConfigOptions.ACC.value:
+        return ["python3", os.path.join(MFC_ROOT_DIR, "toolchain", "mfc", "run", "trace.py")] + profiler
+
+    return profiler
+
+
 def get_baked_templates() -> dict:
     return {os.path.splitext(os.path.basename(f))[0]: file_read(f) for f in glob(os.path.join(MFC_TEMPLATE_DIR, "*.mako"))}
 
@@ -90,6 +99,18 @@ def __get_template() -> Template:
 
 
 def __generate_job_script(targets, case: input.MFCInputFile):
+    # Compute GPU mode booleans for templates
+    gpu_mode = ARG("gpu")
+
+    # Validate gpu_mode is one of the expected values
+    valid_gpu_modes = {e.value for e in gpuConfigOptions}
+    if gpu_mode not in valid_gpu_modes:
+        raise MFCException(f"Invalid GPU mode '{gpu_mode}'. Must be one of: {', '.join(sorted(valid_gpu_modes))}")
+
+    gpu_enabled = gpu_mode != gpuConfigOptions.NONE.value
+    gpu_acc = gpu_mode == gpuConfigOptions.ACC.value
+    gpu_mp = gpu_mode == gpuConfigOptions.MP.value
+
     env = {}
     if ARG("gpus") is not None:
         gpu_ids = ",".join([str(_) for _ in ARG("gpus")])
@@ -109,21 +130,11 @@ def __generate_job_script(targets, case: input.MFCInputFile):
                 "MFC_TRACE_STDOUT": "1" if sys.stdout.isatty() else "0",
             }
         )
+        if gpu_acc:
+            env.update({"MFC_TRACE_ACC_NOTIFY": "1", "NV_ACC_NOTIFY": "1"})
         if ARG("mpi"):
             env.update({"MFC_TRACE_MPI": "1"})
 
-    # Compute GPU mode booleans for templates
-    gpu_mode = ARG("gpu")
-
-    # Validate gpu_mode is one of the expected values
-    valid_gpu_modes = {e.value for e in gpuConfigOptions}
-    if gpu_mode not in valid_gpu_modes:
-        raise MFCException(f"Invalid GPU mode '{gpu_mode}'. Must be one of: {', '.join(sorted(valid_gpu_modes))}")
-
-    gpu_enabled = gpu_mode != gpuConfigOptions.NONE.value
-    gpu_acc = gpu_mode == gpuConfigOptions.ACC.value
-    gpu_mp = gpu_mode == gpuConfigOptions.MP.value
-
     content = __get_template().render(
         **{**ARGS(), "targets": targets},
         ARG=ARG,
@@ -132,7 +143,7 @@ def __generate_job_script(targets, case: input.MFCInputFile):
         MFC_ROOT_DIR=MFC_ROOT_DIR,
         SIMULATION=SIMULATION,
         qsystem=queues.get_system(),
-        profiler=shlex.join(__profiler_prepend()),
+        profiler=shlex.join(__runner_prepend()),
         gpu_enabled=gpu_enabled,
         gpu_acc=gpu_acc,
         gpu_mp=gpu_mp,

toolchain/mfc/run/trace.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+import collections
+import os
+import re
+import subprocess
+import sys
+
+
+def _normalize_acc_line(line: str) -> str:
+    line = re.sub(r"file=[^ ]*/(src/[^ ]*)", r"file=\1", line)
+    line = re.sub(r"[^ ]*/(src/[^ ]*)", r"\1", line)
+    line = re.sub(r"devaddr=0x[0-9A-Fa-f]+", "devaddr=<addr>", line)
+    return line
+
+
+def _is_acc_noise(line: str) -> bool:
+    prefixes = (
+        "Enter ",
+        "Leave ",
+        "Wait ",
+        "Implicit wait",
+        "upload CUDA data",
+        "download CUDA data",
+        "create CUDA data",
+        "delete CUDA data",
+        "alloc",
+    )
+    return line.startswith(prefixes) or (".fpp function=" in line and " device=" in line)
+
+
+def main() -> int:
+    if len(sys.argv) < 2:
+        print("usage: trace.py <command> [args...]", file=sys.stderr)
+        return 2
+
+    command = sys.argv[1:]
+
+    if any(arg.endswith("/syscheck") for arg in command):
+        env = os.environ.copy()
+        env["MFC_TRACE"] = "0"
+        env.pop("MFC_TRACE_ACC_NOTIFY", None)
+        env["NV_ACC_NOTIFY"] = "0"
+        return subprocess.run(command, env=env, check=False).returncode
+
+    trace_file = os.environ.get("MFC_TRACE_FILE")
+    if not os.environ.get("MFC_TRACE_ACC_NOTIFY") or not trace_file:
+        return subprocess.run(command, check=False).returncode
+
+    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    assert process.stdout is not None
+
+    seen = set()
+    repeats = collections.Counter()
+
+    with open(trace_file, "a", encoding="utf-8") as trace:
+        for output_line in process.stdout:
+            line = output_line.rstrip("\n")
+            if line.startswith("launch CUDA kernel"):
+                short_line = _normalize_acc_line(line)
+                if short_line in seen:
+                    repeats[short_line] += 1
+                else:
+                    seen.add(short_line)
+                    trace.write(f"TRACE_ACC {short_line}\n")
+                    trace.flush()
+            elif not _is_acc_noise(line):
+                print(line, flush=True)
+
+        for line, count in sorted(repeats.items()):
+            trace.write(f"TRACE_ACC_REPEAT x{count} {line}\n")
+
+    return process.wait()
+
+
+if __name__ == "__main__":
+    sys.exit(main())