feat(rocthrust): use primbench

projects/rocthrust/.gitlab/run_benchmarks.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2024-2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -21,100 +21,325 @@
 # THE SOFTWARE.
 
 import argparse
-from collections import namedtuple
+import json
 import os
 import re
 import stat
 import subprocess
 import sys
-import time
 
-BenchmarkContext = namedtuple('BenchmarkContext', ['gpu_architecture', 'benchmark_output_dir', 'benchmark_dir', 'benchmark_filename_regex', 'benchmark_filter_regex', 'seed'])
 
-def run_benchmarks(benchmark_context):
+def run_benchmarks(
+    benchmark_executables_dir,
+    gpu_architecture,
+    json_out_dir,
+    csv_out_dir,
+    benchmark_executable_filter,
+    skip_gathered,
+    size,
+    hot,
+    seed,
+    filter,
+    dry,
+    min_gpu_ms_per_batch,
+    min_secs,
+    noise_timeout_secs,
+    batch_window_size,
+    noise_tolerance_percent,
+    min_gpu_temp,
+    max_gpu_temp,
+    max_warming_secs,
+    max_cooling_secs,
+    output_batches,
+    spaces_per_indent,
+    stream_blocking_timeout_secs,
+):
     def is_benchmark_executable(filename):
-        if not re.match(benchmark_context.benchmark_filename_regex, filename):
+        if not re.search(benchmark_executable_filter, filename):
             return False
-        path = os.path.join(benchmark_context.benchmark_dir, filename)
+
+        path = os.path.join(benchmark_executables_dir, filename)
         st_mode = os.stat(path).st_mode
 
-        # we are not interested in permissions, just whether there is any execution flag set
-        # and it is a regular file (S_IFREG)
-        return (st_mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)) and (st_mode & stat.S_IFREG)
+        # Return True when any execution flag is set, AND it is a regular file.
+        # Doesn't check permissions.
+        return (st_mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)) and (
+            st_mode & stat.S_IFREG
+        )
+
+    def should_skip(json_out_path):
+        if not skip_gathered:
+            return False
+
+        try:
+            with open(json_out_path) as f:
+                json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            return False
+
+        return True
+
+    # Create output directories if they don't exist
+    os.makedirs(json_out_dir, exist_ok=True)
+    if csv_out_dir:
+        os.makedirs(csv_out_dir, exist_ok=True)
+
+    benchmark_names = [
+        name
+        for name in os.listdir(benchmark_executables_dir)
+        if is_benchmark_executable(name)
+    ]
+    print(
+        "The following benchmarks will be run:\n{}".format("\n".join(benchmark_names)),
+        file=sys.stderr,
+        flush=True,
+    )
 
     success = True
-    benchmark_names = [name for name in os.listdir(benchmark_context.benchmark_dir) if is_benchmark_executable(name)]
-    benchmark_names.sort()
-    print('The following benchmarks will be ran:\n{}'.format('\n'.join(benchmark_names)), file=sys.stderr, flush=True)
+
     for benchmark_name in benchmark_names:
-        results_json_name = f'{benchmark_name}_{benchmark_context.gpu_architecture}.json'
-
-        benchmark_path = os.path.join(benchmark_context.benchmark_dir, benchmark_name)
-        results_json_path = os.path.join(benchmark_context.benchmark_output_dir, results_json_name)
-        args = [
-            benchmark_path,
-            '--name_format',
-            'json',
-            '--benchmark_out_format=json',
-            f'--benchmark_out={results_json_path}',
-            f'--benchmark_filter={benchmark_context.benchmark_filter_regex}'
-        ]
-        if benchmark_context.seed:
-            args += ['--seed', benchmark_context.seed]
+        results_json_name = f"{benchmark_name}_{gpu_architecture}.json"
+        json_out_path = os.path.join(json_out_dir, results_json_name)
+
+        if should_skip(json_out_path):
+            print(
+                f"Skipping {benchmark_name}, because its results have already been gathered at {json_out_path}",
+                file=sys.stderr,
+                flush=True,
+            )
+            continue
+
+        benchmark_path = os.path.join(benchmark_executables_dir, benchmark_name)
+
+        args = [benchmark_path]
+        args += ["--json-out", json_out_path]
+
+        if size:
+            args += ["--size", size]
+        if hot:
+            args += ["--hot"]
+        if seed:
+            args += ["--seed", seed]
+        if csv_out_dir:
+            results_csv_name = f"{benchmark_name}_{gpu_architecture}.csv"
+            csv_out_path = os.path.join(csv_out_dir, results_csv_name)
+            args += ["--csv-out", csv_out_path]
+        if filter:
+            args += ["--filter", filter]
+        if dry:
+            args += ["--dry"]
+        if min_gpu_ms_per_batch:
+            args += ["--min-gpu-ms-per-batch", min_gpu_ms_per_batch]
+        if min_secs:
+            args += ["--min-secs", min_secs]
+        if noise_timeout_secs:
+            args += ["--noise-timeout-secs", noise_timeout_secs]
+        if batch_window_size:
+            args += ["--batch-window-size", batch_window_size]
+        if noise_tolerance_percent:
+            args += ["--noise-tolerance-percent", noise_tolerance_percent]
+        if min_gpu_temp:
+            args += ["--min-gpu-temp", min_gpu_temp]
+        if max_gpu_temp:
+            args += ["--max-gpu-temp", max_gpu_temp]
+        if max_warming_secs:
+            args += ["--max-warming-secs", max_warming_secs]
+        if max_cooling_secs:
+            args += ["--max-cooling-secs", max_cooling_secs]
+        if output_batches:
+            args += ["--output-batches"]
+        if spaces_per_indent:
+            args += ["--spaces-per-indent", spaces_per_indent]
+        if stream_blocking_timeout_secs:
+            args += ["--stream-blocking-timeout-secs", stream_blocking_timeout_secs]
+
         try:
-            start_time = time.time()
             subprocess.check_call(args)
-            end_time = time.time()
-            duration = end_time - start_time
-
-            print(f'Benchmark {benchmark_name} took {duration:.3f} seconds to run', file=sys.stderr, flush=True)
         except subprocess.CalledProcessError as error:
-            print(f'Could not run benchmark at {benchmark_path}. Error: "{error}"', file=sys.stderr, flush=True)
+            print(
+                f'Could not run benchmark at {benchmark_path}. Error: "{error}"',
+                file=sys.stderr,
+                flush=True,
+            )
             success = False
-    return success
 
+    return success
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--benchmark_dir',
-        help='The local directory that contains the benchmark executables',
-        required=True)
-    parser.add_argument('--benchmark_gpu_architecture',
-        help='The architecture of the currently enabled GPU',
-        required=True)
-    parser.add_argument('--benchmark_output_dir',
-        help='The directory to write the benchmarks to',
-        required=True)
-    parser.add_argument('--benchmark_filename_regex',
-        help='Regular expression that controls the list of benchmark executables to run',
-        default=r'^benchmark',
-        required=False)
-    parser.add_argument('--benchmark_filter_regex',
-        help='Regular expression that controls the list of benchmarks to run in each benchmark executable',
-        default='',
-        required=False)
-    parser.add_argument('--seed',
-        help='Controls the seed for random number generation for each benchmark case',
-        default='',
-        required=False)
 
-    args = parser.parse_args()
+    # Python arguments
+    parser.add_argument(
+        "--benchmark-executables-dir",
+        help="Benchmark executables directory",
+        required=True,
+    )
+    parser.add_argument(
+        "--gpu-architecture",
+        help="Architecture of the currently enabled GPU",
+        required=True,
+    )
+    parser.add_argument(
+        "--json-out-dir",
+        help="Directory to write JSON result files to, creating parent directories if necessary",
+        required=True,
+    )
+    parser.add_argument(
+        "--csv-out-dir",
+        help="Directory to write CSV result files to, creating parent directories if necessary",
+        required=False,
+    )
+    parser.add_argument(
+        "--benchmark-executable-filter",
+        help="Regular expression that controls the list of benchmark executables to run",
+        default=r"^benchmark",
+        required=False,
+    )
+    parser.add_argument(
+        "--skip-gathered",
+        help="Skip running benchmarks whose JSON file has already been gathered",
+        default=False,
+        action="store_true",
+        required=False,
+    )
 
-    benchmark_context = BenchmarkContext(
-        args.benchmark_gpu_architecture,
-        args.benchmark_output_dir,
-        args.benchmark_dir,
-        args.benchmark_filename_regex,
-        args.benchmark_filter_regex,
-        args.seed)
+    # primbench arguments
+    parser.add_argument(
+        "--size",
+        help="Input size. Benchmarks decide what this represents, but it is commonly the number of bytes or items",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--hot",
+        help="Skip clearing the GPU cache between batch iterations",
+        default=False,
+        action="store_true",
+        required=False,
+    )
+    parser.add_argument(
+        "--seed", help="Seed used for input generation", default="", required=False
+    )
+    parser.add_argument(
+        "--filter",
+        help="Regex filter of specialization names to benchmark",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--dry",
+        help="Perform a dry run. The benchmark setup is still run, and JSON and CSV files are still output, but state.run() immediately returns",
+        default=False,
+        action="store_true",
+        required=False,
+    )
+    parser.add_argument(
+        "--min-gpu-ms-per-batch",
+        help="Minimum duration of a batch in milliseconds (GPU time)",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--min-secs",
+        help="Minimum total benchmark duration in seconds (wall time)",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--noise-timeout-secs",
+        help="Maximum total benchmark duration in seconds before timing out a noisy run (wall time)",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--batch-window-size",
+        help="Number of batch times used in the noise (coefficient of variation) window to decide early stopping",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--noise-tolerance-percent",
+        help="Noise tolerance of batch times in percent, used to determine early benchmark stopping",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--min-gpu-temp",
+        help="Minimum GPU temperature in °C. Too low slows benchmarks; too high increases noise",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--max-gpu-temp",
+        help="Maximum GPU temperature in °C. Too low slows benchmarks; too high increases noise",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--max-warming-secs",
+        help="Maximum seconds allowed for GPU warming before an error is thrown",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--max-cooling-secs",
+        help="Maximum seconds allowed for GPU cooling before an error is thrown",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--output-batches",
+        help="Output a `batches` array for each specialization, containing per-batch details",
+        default=False,
+        action="store_true",
+        required=False,
+    )
+    parser.add_argument(
+        "--spaces-per-indent",
+        help="Number of spaces per indentation level in JSON output. Set to 0 for no indentation",
+        default="",
+        required=False,
+    )
+    parser.add_argument(
+        "--stream-blocking-timeout-secs",
+        help="Maximum stream blocking duration in seconds before timing out. Stream is blocked while queueing kernel calls. Use `primbench::flags::sync` if kernel is synchronous",
+        default="",
+        required=False,
+    )
+
+    args = parser.parse_args()
 
-    benchmark_run_successful = run_benchmarks(benchmark_context)
+    benchmark_run_successful = run_benchmarks(
+        args.benchmark_executables_dir,
+        args.gpu_architecture,
+        args.json_out_dir,
+        args.csv_out_dir,
+        args.benchmark_executable_filter,
+        args.skip_gathered,
+        args.size,
+        args.hot,
+        args.seed,
+        args.filter,
+        args.dry,
+        args.min_gpu_ms_per_batch,
+        args.min_secs,
+        args.noise_timeout_secs,
+        args.batch_window_size,
+        args.noise_tolerance_percent,
+        args.min_gpu_temp,
+        args.max_gpu_temp,
+        args.max_warming_secs,
+        args.max_cooling_secs,
+        args.output_batches,
+        args.spaces_per_indent,
+        args.stream_blocking_timeout_secs,
+    )
 
     return benchmark_run_successful
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     success = main()
     if success:
         exit(0)