Add bench_ctx_device and fix JSON output

Author: danielfrg

cuda_bindings/benchmarks/.gitignore

@@ -11,3 +11,7 @@ __pycache__/
 
 # Override root .gitignore *.cpp rule (which targets Cython-generated files)
 !benchmarks/cpp/*.cpp
+
+
+results-python.json
+results-cpp.json

cuda_bindings/benchmarks/benchmarks/bench_ctx_device.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import time
+
+from runner.runtime import ensure_context
+
+from cuda.bindings import driver as cuda
+
+ensure_context()
+
+
+def bench_ctx_get_current(loops: int) -> float:
+    _cuCtxGetCurrent = cuda.cuCtxGetCurrent
+
+    t0 = time.perf_counter()
+    for _ in range(loops):
+        _cuCtxGetCurrent()
+    return time.perf_counter() - t0

cuda_bindings/benchmarks/benchmarks/cpp/CMakeLists.txt

@@ -46,3 +46,7 @@ endif()
 add_executable(bench_pointer_attributes_cpp bench_pointer_attributes.cpp)
 target_include_directories(bench_pointer_attributes_cpp PRIVATE "${CUDA_DRIVER_INCLUDE_DIR}")
 target_link_libraries(bench_pointer_attributes_cpp PRIVATE "${CUDA_DRIVER_LIBRARY}")
+
+add_executable(bench_ctx_device_cpp bench_ctx_device.cpp)
+target_include_directories(bench_ctx_device_cpp PRIVATE "${CUDA_DRIVER_INCLUDE_DIR}")
+target_link_libraries(bench_ctx_device_cpp PRIVATE "${CUDA_DRIVER_LIBRARY}")

cuda_bindings/benchmarks/benchmarks/cpp/bench_ctx_device.cpp

@@ -0,0 +1,65 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <cuda.h>
+
+#include "bench_support.hpp"
+
+#include <cstdlib>
+#include <iostream>
+
+
+static void check_cu(CUresult status, const char* message) {
+    if (status != CUDA_SUCCESS) {
+        const char* error_name = nullptr;
+        cuGetErrorName(status, &error_name);
+        std::cerr << message << ": " << (error_name ? error_name : "unknown") << '\n';
+        std::exit(1);
+    }
+}
+
+
+int main(int argc, char** argv) {
+    bench::Options options = bench::parse_args(argc, argv);
+    if (options.benchmark_name.empty()) {
+        options.benchmark_name = "cpp.ctx_device.ctx_get_current";
+    }
+
+    // Setup: init CUDA and create a context
+    check_cu(cuInit(0), "cuInit failed");
+
+    CUdevice device;
+    check_cu(cuDeviceGet(&device, 0), "cuDeviceGet failed");
+
+    CUcontext ctx;
+    CUctxCreateParams ctxParams = {};
+    check_cu(cuCtxCreate(&ctx, &ctxParams, 0, device), "cuCtxCreate failed");
+
+    CUcontext current_ctx = nullptr;
+
+    // Run benchmark
+    auto results = bench::run_benchmark(options, [&]() {
+        check_cu(
+            cuCtxGetCurrent(&current_ctx),
+            "cuCtxGetCurrent failed"
+        );
+    });
+
+    // Sanity check: the call actually returned our context
+    if (current_ctx != ctx) {
+        std::cerr << "unexpected: cuCtxGetCurrent returned a different context\n";
+    }
+
+    // Cleanup
+    check_cu(cuCtxDestroy(ctx), "cuCtxDestroy failed");
+
+    // Output
+    bench::print_summary(options.benchmark_name, results);
+
+    if (!options.output_path.empty()) {
+        bench::write_pyperf_json(options.output_path, options.benchmark_name, options.loops, results);
+    }
+
+    return 0;
+}

cuda_bindings/benchmarks/pixi.toml

@@ -76,7 +76,7 @@ cmd = ["cmake", "--build", "$PIXI_PROJECT_ROOT/.build/cpp"]
 depends-on = [{ task = "bench-cpp-configure" }]
 
 [target.linux.tasks.bench-cpp]
-cmd = ["$PIXI_PROJECT_ROOT/.build/cpp/bench_pointer_attributes_cpp"]
+cmd = ["python", "$PIXI_PROJECT_ROOT/run_cpp.py"]
 depends-on = [{ task = "bench-cpp-build" }]
 
 [target.linux.tasks.lint]

cuda_bindings/benchmarks/run_cpp.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from runner.cpp import main
+
+if __name__ == "__main__":
+    main()

cuda_bindings/benchmarks/runner/cpp.py

@@ -0,0 +1,181 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import json
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+BUILD_DIR = PROJECT_ROOT / ".build" / "cpp"
+DEFAULT_OUTPUT = PROJECT_ROOT / "results-cpp.json"
+
+BINARY_PREFIX = "bench_"
+BINARY_SUFFIX = "_cpp"
+
+
+def discover_binaries() -> dict[str, Path]:
+    """Discover C++ benchmark binaries in the build directory """
+    if not BUILD_DIR.is_dir():
+        return {}
+
+    registry: dict[str, Path] = {}
+    for path in sorted(BUILD_DIR.iterdir()):
+        if not path.is_file() or not path.name.startswith(BINARY_PREFIX):
+            continue
+        if not path.name.endswith(BINARY_SUFFIX):
+            continue
+        name = path.name.removeprefix(BINARY_PREFIX).removesuffix(BINARY_SUFFIX)
+        registry[name] = path
+    return registry
+
+
+def strip_output_args(argv: list[str]) -> list[str]:
+    cleaned: list[str] = []
+    skip_next = False
+    for arg in argv:
+        if skip_next:
+            skip_next = False
+            continue
+        if arg in ("-o", "--output"):
+            skip_next = True
+            continue
+        if arg.startswith("-o=") or arg.startswith("--output="):
+            continue
+        cleaned.append(arg)
+    return cleaned
+
+
+def merge_pyperf_json(individual_files: list[Path], output_path: Path) -> int:
+    """Merge individual pyperf JSON files into a single BenchmarkSuite file.
+
+    Each C++ binary produces a file with structure:
+        {"version": "1.0", "metadata": {...}, "benchmarks": [{...}]}
+
+    We merge them by collecting all benchmark entries into one file.
+    """
+    all_benchmarks = []
+
+    for path in individual_files:
+        with open(path) as f:
+            data = json.load(f)
+
+        file_metadata = data.get("metadata", {})
+        bench_name = file_metadata.get("name", "")
+        loops = file_metadata.get("loops")
+        unit = file_metadata.get("unit", "second")
+
+        for bench in data.get("benchmarks", []):
+            for run in bench.get("runs", []):
+                run_meta = run.setdefault("metadata", {})
+                if bench_name:
+                    run_meta.setdefault("name", bench_name)
+                if loops is not None:
+                    run_meta.setdefault("loops", loops)
+                run_meta.setdefault("unit", unit)
+
+            all_benchmarks.append(bench)
+
+    merged = {
+        "version": "1.0",
+        "benchmarks": all_benchmarks,
+    }
+
+    with open(output_path, "w") as f:
+        json.dump(merged, f)
+
+    return len(all_benchmarks)
+
+
+def parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
+    parser = argparse.ArgumentParser(
+        description="Run C++ CUDA benchmarks",
+        add_help=False,
+    )
+    parser.add_argument(
+        "--benchmark",
+        action="append",
+        default=[],
+        help="Benchmark name to run (e.g. 'ctx_device'). Repeat for multiple. Defaults to all.",
+    )
+    parser.add_argument(
+        "--list",
+        action="store_true",
+        help="Print discovered benchmark names and exit.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=Path,
+        default=DEFAULT_OUTPUT,
+        help=f"JSON output file path (default: {DEFAULT_OUTPUT.name})",
+    )
+    parsed, remaining = parser.parse_known_args(argv)
+    return parsed, remaining
+
+
+def main() -> None:
+    parsed, remaining_argv = parse_args(sys.argv[1:])
+
+    registry = discover_binaries()
+    if not registry:
+        print(
+            f"No C++ benchmark binaries found in {BUILD_DIR}.\n"
+            "Run 'pixi run bench-cpp-build' first.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    if parsed.list:
+        for name in sorted(registry):
+            print(name)
+        return
+
+    if parsed.benchmark:
+        missing = sorted(set(parsed.benchmark) - set(registry))
+        if missing:
+            known = ", ".join(sorted(registry))
+            unknown = ", ".join(missing)
+            print(
+                f"Unknown benchmark(s): {unknown}. Known benchmarks: {known}",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        names = parsed.benchmark
+    else:
+        names = sorted(registry)
+
+    # Strip any --output args to avoid conflicts with our output handling
+    passthrough_argv = strip_output_args(remaining_argv)
+
+    output_path = parsed.output.resolve()
+    failed = False
+    individual_files: list[Path] = []
+
+    with tempfile.TemporaryDirectory(prefix="cuda_bench_cpp_") as tmpdir:
+        tmpdir_path = Path(tmpdir)
+
+        for name in names:
+            binary = registry[name]
+            tmp_json = tmpdir_path / f"{name}.json"
+            cmd = [str(binary), "-o", str(tmp_json), *passthrough_argv]
+            result = subprocess.run(cmd)
+            if result.returncode != 0:
+                print(f"FAILED: {name} (exit code {result.returncode})", file=sys.stderr)
+                failed = True
+            elif tmp_json.exists():
+                individual_files.append(tmp_json)
+
+        if individual_files:
+            count = merge_pyperf_json(individual_files, output_path)
+            print(f"\nResults saved to {output_path} ({count} benchmark(s))")
+
+    if failed:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

cuda_bindings/benchmarks/runner/main.py

@@ -12,7 +12,9 @@
 
 import pyperf
 
-BENCH_DIR = Path(__file__).resolve().parent.parent / "benchmarks"
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+BENCH_DIR = PROJECT_ROOT / "benchmarks"
+DEFAULT_OUTPUT = PROJECT_ROOT / "results-python.json"
 
 
 def load_module(module_path: Path) -> ModuleType:
@@ -54,6 +56,22 @@ def discover_benchmarks() -> dict[str, Callable[[int], float]]:
     return registry
 
 
+def strip_pyperf_output_args(argv: list[str]) -> list[str]:
+    cleaned: list[str] = []
+    skip_next = False
+    for i, arg in enumerate(argv):
+        if skip_next:
+            skip_next = False
+            continue
+        if arg in ("-o", "--output", "--append"):
+            skip_next = True
+            continue
+        if arg.startswith("-o=") or arg.startswith("--output=") or arg.startswith("--append="):
+            continue
+        cleaned.append(arg)
+    return cleaned
+
+
 def parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
     parser = argparse.ArgumentParser(add_help=False)
     parser.add_argument(
@@ -67,13 +85,19 @@ def parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
         action="store_true",
         help="Print discovered benchmark IDs and exit.",
     )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=Path,
+        default=DEFAULT_OUTPUT,
+        help=f"JSON output file path (default: {DEFAULT_OUTPUT.name})",
+    )
     parsed, remaining = parser.parse_known_args(argv)
     return parsed, remaining
 
 
 def main() -> None:
     parsed, remaining_argv = parse_args(sys.argv[1:])
-    sys.argv = [sys.argv[0], *remaining_argv]
 
     registry = discover_benchmarks()
     if not registry:
@@ -94,10 +118,24 @@ def main() -> None:
     else:
         benchmark_ids = sorted(registry)
 
+    # Strip any --output args to avoid conflicts with our output handling
+    output_path = parsed.output.resolve()
+    remaining_argv = strip_pyperf_output_args(remaining_argv)
+    is_worker = "--worker" in remaining_argv
+
+    # Delete the file so this run starts fresh
+    if not is_worker:
+        output_path.unlink(missing_ok=True)
+
+    sys.argv = [sys.argv[0], "--append", str(output_path), *remaining_argv]
+
     runner = pyperf.Runner()
     for bench_id in benchmark_ids:
         runner.bench_time_func(bench_id, registry[bench_id])
 
+    if not is_worker:
+        print(f"\nResults saved to {output_path}")
+
 
 if __name__ == "__main__":
     main()