File indexing completed on 2026-07-26 08:22:17
0001
0002
0003
0004
0005 import argparse
0006 import sys
0007 import csv
0008 import git
0009 import pathlib
0010 import shutil
0011 import logging
0012 import tempfile
0013 import subprocess
0014 import os
0015 import time
0016
0017 from traccc_bench_tools import parse_profile, types, build, profile, benchmark
0018
0019 ALL_COMPUTE_CAPABILITIES = [
0020 "50",
0021 "52",
0022 "60",
0023 "61",
0024 "70",
0025 "75",
0026 "80",
0027 "86",
0028 "89",
0029 "90",
0030 "100",
0031 "120",
0032 ]
0033
0034
0035 log = logging.getLogger("traccc_benchmark")
0036
0037
0038 def main():
0039 parser = argparse.ArgumentParser()
0040
0041 parser.add_argument(
0042 "repo",
0043 type=pathlib.Path,
0044 help="the traccc build repository",
0045 )
0046
0047 parser.add_argument(
0048 "db",
0049 type=pathlib.Path,
0050 help="the CSV database",
0051 )
0052
0053 parser.add_argument(
0054 "data",
0055 type=pathlib.Path,
0056 help="the traccc dataset to use",
0057 )
0058
0059 parser.add_argument(
0060 "-v",
0061 "--verbose",
0062 help="enable verbose output",
0063 action="store_true",
0064 )
0065
0066 parser.add_argument(
0067 "-j",
0068 "--parallel",
0069 help="number of threads to use for compilation",
0070 default=1,
0071 type=int,
0072 )
0073
0074 parser.add_argument(
0075 "-e",
0076 "--events",
0077 help="number of events to process per commit",
0078 type=int,
0079 default=10,
0080 )
0081
0082 parser.add_argument(
0083 "--cc",
0084 help="Compute Capability of the modelled GPU",
0085 type=str,
0086 required=True,
0087 dest="cc",
0088 )
0089
0090 parser.add_argument(
0091 "--num-sm",
0092 help="number of SMs in the modelled GPU",
0093 type=int,
0094 required=True,
0095 dest="num_sm",
0096 )
0097
0098 parser.add_argument(
0099 "--num-threads-per-sm",
0100 help="number of thread slots per SM in the modelled GPU",
0101 type=int,
0102 required=True,
0103 dest="num_threads_per_sm",
0104 )
0105
0106 parser.add_argument(
0107 "--ncu-wrapper",
0108 help="wrapper to use around the ncu command",
0109 type=str,
0110 dest="ncu_wrapper",
0111 )
0112
0113 args = parser.parse_args()
0114
0115 logging.basicConfig(
0116 level=logging.DEBUG if (args.verbose or False) else logging.INFO,
0117 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
0118 )
0119
0120 log.info(
0121 "Using GPU with %d SMs and %d thread slots per SM (%d thread slots total)",
0122 getattr(args, "num_sm"),
0123 getattr(args, "num_threads_per_sm"),
0124 getattr(args, "num_sm") * getattr(args, "num_threads_per_sm"),
0125 )
0126
0127 repo = git.Repo(args.repo)
0128
0129 log.info("Using git repository at %s", repo.git_dir)
0130
0131 if repo.is_dirty():
0132 e = "Repository is dirty; please clean it before use!"
0133 log.fatal(e)
0134 raise RuntimeError(e)
0135
0136 if "TRACCC_TEST_DATA_DIR" not in os.environ:
0137 e = 'Environment variable "TRACCC_TEST_DATA_DIR" is not set; aborting!'
0138 log.fatal(e)
0139 raise RuntimeError(e)
0140
0141 for exec in ["ncu", "g++", "nvcc", "cmake"]:
0142 if shutil.which(exec) is None:
0143 e = 'Executable "%s" is not available; aborting' % exec
0144 log.fatal(e)
0145 raise RuntimeError(e)
0146
0147 old_commit_hash = repo.head.object.hexsha
0148 log.info("Current commit hash is %s", old_commit_hash)
0149
0150 if args.db.is_file():
0151 log.info('Database file "%s" already exists; creating a backup', args.db)
0152 shutil.copy(str(args.db), str(args.db) + ".bak")
0153 results = []
0154 with open(args.db, "r") as f:
0155 reader = csv.DictReader(f)
0156 for i in reader:
0157 results.append(i)
0158 log.info("Database contained %d pre-existing results", len(results))
0159 else:
0160 log.info('Database file "%s" does not exist; starting from scratch', args.db)
0161 results = []
0162
0163 known_ccs = set(x["cc"] for x in results)
0164
0165 log.info(
0166 "Currently have pre-existing results for %d Compute Capabilities",
0167 len(known_ccs),
0168 )
0169
0170 for progress_id, x in enumerate(ALL_COMPUTE_CAPABILITIES):
0171 log.info(
0172 "Considering Compute Capability %s.%s (%u out of %u)",
0173 x[:-1],
0174 x[-1],
0175 progress_id,
0176 len(ALL_COMPUTE_CAPABILITIES),
0177 )
0178
0179
0180 if x in known_ccs:
0181 log.info(
0182 "Compute capability %s.%s is already know; skipping", x[:-1], x[-1]
0183 )
0184 continue
0185
0186 if int(x) > int(args.cc):
0187 log.info(
0188 "Compute Capability %s.%s is newer than target %s.%s; skipping",
0189 x[:-1],
0190 x[-1],
0191 args.cc[:-1],
0192 args.cc[-1],
0193 )
0194 continue
0195
0196 try:
0197 log.info("Running benchmark for Compute Capability %s.%s", x[:-1], x[-1])
0198
0199 result_df = benchmark.run_benchmark(
0200 source_dir=args.repo,
0201 data_dir=args.data,
0202 commit=repo.head.object,
0203 gpu_spec=types.GpuSpec(
0204 getattr(args, "num_sm"), getattr(args, "num_threads_per_sm")
0205 ),
0206 parallel=args.parallel,
0207 events=args.events,
0208 ncu_wrapper=getattr(args, "ncu_wrapper", None),
0209 cc=x,
0210 )
0211
0212 for y in result_df.iloc:
0213 results.append(
0214 {
0215 "cc": str(x),
0216 "kernel": y["Kernel Name"],
0217 "throughput": y["ThroughputMean"],
0218 "throughput_dev": y["ThroughputStd"],
0219 "rec_throughput": y["RecThroughputMean"],
0220 "rec_throughput_dev": y["RecThroughputStd"],
0221 }
0222 )
0223
0224 except Exception as e:
0225 log.exception(e)
0226 except KeyboardInterrupt as e:
0227 log.info("Received keyboard interrupt; skipping to post-processing")
0228 break
0229
0230 log.info("Gathered a total of %d results (incl. pre-existing)", len(results))
0231
0232 log.info("Writing data to %s", args.db)
0233 with open(args.db, "w") as f:
0234 writer = csv.DictWriter(
0235 f,
0236 fieldnames=[
0237 "cc",
0238 "kernel",
0239 "throughput",
0240 "throughput_dev",
0241 "rec_throughput",
0242 "rec_throughput_dev",
0243 ],
0244 )
0245 writer.writeheader()
0246
0247 for i in results:
0248 writer.writerow(i)
0249
0250 log.info("Processing complete; goodbye!")
0251
0252
0253 if __name__ == "__main__":
0254 main()
