benchmark: Add support for COMPARISON file

3 files changed, 91 insertions, 49 deletions

diff --git a/benchmark/COMPARISON b/benchmark/COMPARISON
index 55ab94b4..d8750e34 100644
--- a/benchmark/COMPARISON
+++ b/benchmark/COMPARISON
@@ -1,4 +1,4 @@
-CPU                                     | Mz   | exec/s | afl-*-config |
-========================================|======|========|==============|
-CPU 12th Gen Intel(R) Core(TM) i7-1270P | 4200 | 12750  | both         |
-AMD EPYC 7282 16-Core Processor         | 3190 | 10060  | both         |
+CPU                            | MHz   | singlecore | multicore | afl-*-config |
+===============================|=======|============|===========|==============|
+Apple Mac Studio M2 Ultra 2023 | 3500  | 174386     | 1112585   | both         |
+Intel(R) Core(TM) i9-9900K CPU | 4999  | 133706     | 1169989   | both         |
diff --git a/benchmark/benchmark-results.jsonl b/benchmark/benchmark-results.jsonl
index eef18384..1d52402d 100644
--- a/benchmark/benchmark-results.jsonl
+++ b/benchmark/benchmark-results.jsonl
@@ -546,3 +546,5 @@
 {"config": {"afl_persistent_config": true, "afl_system_config": true, "afl_version": "++4.09a", "comment": "AWS EC2 r6a.48xlarge spot instance", "compiler": "clang version 15.0.7 (Amazon Linux 15.0.7-3.amzn2023.0.1)", "target_arch": "x86_64-amazon-linux-gnu"}, "hardware": {"cpu_fastest_core_mhz": 3599.413, "cpu_model": "AMD EPYC 7R13 Processor", "cpu_threads": 192}, "targets": {"test-instr-persist-shmem": {"multicore": {"afl_execs_per_sec": 334191.98, "afl_execs_total": 102894843, "fuzzers_used": 198, "run_end": "2023-10-01 06:31:03.545925", "run_start": "2023-10-01 06:27:53.127882", "total_execs_per_sec": 270150.29, "total_run_time": 380.88}}}}
 {"config": {"afl_persistent_config": true, "afl_system_config": true, "afl_version": "++4.09a", "comment": "AWS EC2 r6a.48xlarge spot instance", "compiler": "clang version 15.0.7 (Amazon Linux 15.0.7-3.amzn2023.0.1)", "target_arch": "x86_64-amazon-linux-gnu"}, "hardware": {"cpu_fastest_core_mhz": 3599.521, "cpu_model": "AMD EPYC 7R13 Processor", "cpu_threads": 192}, "targets": {"test-instr-persist-shmem": {"multicore": {"afl_execs_per_sec": 332929.11, "afl_execs_total": 103414536, "fuzzers_used": 199, "run_end": "2023-10-01 06:37:27.645981", "run_start": "2023-10-01 06:34:15.843433", "total_execs_per_sec": 269442.01, "total_run_time": 383.81}}}}
 {"config": {"afl_persistent_config": true, "afl_system_config": true, "afl_version": "++4.09a", "comment": "AWS EC2 r6a.48xlarge spot instance", "compiler": "clang version 15.0.7 (Amazon Linux 15.0.7-3.amzn2023.0.1)", "target_arch": "x86_64-amazon-linux-gnu"}, "hardware": {"cpu_fastest_core_mhz": 3339.7, "cpu_model": "AMD EPYC 7R13 Processor", "cpu_threads": 192}, "targets": {"test-instr-persist-shmem": {"multicore": {"afl_execs_per_sec": 331957.22, "afl_execs_total": 103934201, "fuzzers_used": 200, "run_end": "2023-10-01 06:43:54.782368", "run_start": "2023-10-01 06:40:41.553700", "total_execs_per_sec": 268674.91, "total_run_time": 386.84}}}}
+{"config": {"afl_persistent_config": true, "afl_system_config": true, "afl_version": "++4.09a", "comment": "", "compiler": "Ubuntu clang version 16.0.6", "target_arch": "aarch64-unknown-linux-gnu"}, "hardware": {"cpu_fastest_core_mhz": 3500.0, "cpu_model": "Apple Mac Studio M2 Ultra 2023", "cpu_threads": 16}, "targets": {"test-instr-persist-shmem": {"multicore": {"afl_execs_per_sec": 1172198.55, "afl_execs_total": 12472081, "fuzzers_used": 16, "run_end": "2023-11-11 02:20:21.969027", "run_start": "2023-11-11 02:20:18.282561", "total_execs_per_sec": 1128695.11, "total_run_time": 11.05}, "singlecore": {"afl_execs_per_sec": 172105.52, "afl_execs_total": 779505, "fuzzers_used": 1, "run_end": "2023-11-11 02:20:10.920659", "run_start": "2023-11-11 02:20:09.459765", "total_execs_per_sec": 170570.02, "total_run_time": 4.57}}}}
+{"config": {"afl_persistent_config": true, "afl_system_config": true, "afl_version": "++4.09a", "comment": "", "compiler": "", "target_arch": "x86_64-pc-linux-gnu"}, "hardware": {"cpu_fastest_core_mhz": 4998.583, "cpu_model": "Intel(R) Core(TM) i9-9900K CPU @ 3.60GHz", "cpu_threads": 16}, "targets": {"test-instr-persist-shmem": {"multicore": {"afl_execs_per_sec": 1193793.35, "afl_execs_total": 12472080, "fuzzers_used": 16, "run_end": "2023-11-10 10:25:40.047364", "run_start": "2023-11-10 10:25:36.510399", "total_execs_per_sec": 1169988.74, "total_run_time": 10.66}, "singlecore": {"afl_execs_per_sec": 134426.26, "afl_execs_total": 779505, "fuzzers_used": 1, "run_end": "2023-11-10 10:25:29.381317", "run_start": "2023-11-10 10:25:27.420959", "total_execs_per_sec": 133705.83, "total_run_time": 5.83}}}}
diff --git a/benchmark/benchmark.py b/benchmark/benchmark.py
index e0dea299..6fde621b 100644
--- a/benchmark/benchmark.py
+++ b/benchmark/benchmark.py
@@ -42,8 +42,8 @@ class Config:
 
 @dataclass
 class Hardware:
-    cpu_fastest_core_mhz: Optional[float]
-    cpu_model: Optional[str]
+    cpu_fastest_core_mhz: float
+    cpu_model: str
     cpu_threads: int
 
 @dataclass
@@ -57,80 +57,85 @@ all_targets = [
     Target(source=Path("../utils/persistent_mode/test-instr.c").resolve(), binary=Path("test-instr-persist-shmem")),
     Target(source=Path("../test-instr.c").resolve(), binary=Path("test-instr"))
 ]
-mode_names = [mode.name for mode in all_modes]
-target_names = [str(target.binary) for target in all_targets]
+modes = [mode.name for mode in all_modes]
+targets = [str(target.binary) for target in all_targets]
 cpu_count = multiprocessing.cpu_count()
+env_vars = {
+    "AFL_BENCH_JUST_ONE": "1", "AFL_DISABLE_TRIM": "1", "AFL_I_DONT_CARE_ABOUT_MISSING_CRASHES": "1",
+    "AFL_NO_UI": "1", "AFL_TRY_AFFINITY": "1", "PATH": f'{str(Path("../").resolve())}:{os.environ["PATH"]}',
+}
 
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 parser.add_argument("-b", "--basedir", help="directory to use for temp files", type=str, default="/tmp/aflpp-benchmark")
 parser.add_argument("-d", "--debug", help="show verbose debugging output", action="store_true")
 parser.add_argument("-r", "--runs", help="how many runs to average results over", type=int, default=3)
 parser.add_argument("-f", "--fuzzers", help="how many afl-fuzz workers to use", type=int, default=cpu_count)
-parser.add_argument("-m", "--mode", help="pick modes", action="append", default=["multicore"], choices=mode_names)
+parser.add_argument("-m", "--mode", help="pick modes", action="append", default=modes, choices=modes)
 parser.add_argument("-c", "--comment", help="add a comment about your setup", type=str, default="")
+parser.add_argument("--cpu", help="override the detected CPU model name", type=str, default="")
+parser.add_argument("--mhz", help="override the detected CPU MHz", type=str, default="")
 parser.add_argument(
-    "-t", "--target", help="pick targets", action="append", default=["test-instr-persist-shmem"], choices=target_names
+    "-t", "--target", help="pick targets", action="append", default=["test-instr-persist-shmem"], choices=targets
 )
 args = parser.parse_args()
 # Really unsatisfying argparse behavior: we want a default and to allow multiple choices, but if there's a manual choice
 # it should override the default.  Seems like we have to remove the default to get that and have correct help text?
-if len(args.target) > 1: args.target = args.target[1:]
-if len(args.mode) > 1: args.mode = args.mode[1:]
-
-targets = [target for target in all_targets if str(target.binary) in args.target]
-modes = [mode for mode in all_modes if mode.name in args.mode]
-results = Results(config=None, hardware=None, targets={str(t.binary): {m.name: None for m in modes} for t in targets})
+if len(args.target) > 1:
+    args.target = args.target[1:]
+if len(args.mode) > 2:
+    args.mode = args.mode[2:]
+
+chosen_modes = [mode for mode in all_modes if mode.name in args.mode]
+chosen_targets = [target for target in all_targets if str(target.binary) in args.target]
+results = Results(config=None, hardware=None, targets={
+    str(t.binary): {m.name: None for m in chosen_modes} for t in chosen_targets}
+)
 debug = lambda text: args.debug and print(blue(text))
-if Mode.multicore in modes:
+if Mode.multicore in chosen_modes:
     print(blue(f" [*] Using {args.fuzzers} fuzzers for multicore fuzzing "), end="")
     print(blue("(use --fuzzers to override)" if args.fuzzers == cpu_count else f"(the default is {cpu_count})"))
 
 async def clean_up_tempfiles() -> None:
     shutil.rmtree(f"{args.basedir}/in")
-    for target in targets:
+    for target in chosen_targets:
         target.binary.unlink()
-        for mode in modes:
+        for mode in chosen_modes:
             shutil.rmtree(f"{args.basedir}/out-{mode.name}-{str(target.binary)}")
 
 async def check_afl_persistent() -> bool:
-    with open("/proc/cmdline", "r") as cpuinfo:
-        return "mitigations=off" in cpuinfo.read().split(" ")
+    with open("/proc/cmdline", "r") as cmdline:
+        return "mitigations=off" in cmdline.read().strip().split(" ")
 
 async def check_afl_system() -> bool:
     sysctl = next((s for s in ["sysctl", "/sbin/sysctl"] if shutil.which(s)), None)
     if sysctl:
-        (returncode, stdout, _) = await run_command([sysctl, "kernel.randomize_va_space"], None)
+        (returncode, stdout, _) = await run_command([sysctl, "kernel.randomize_va_space"])
         return returncode == 0 and stdout.decode().rstrip().split(" = ")[1] == "0"
     return False
 
-async def prep_env() -> dict:
+async def prep_env() -> None:
     Path(f"{args.basedir}/in").mkdir(exist_ok=True, parents=True)
-    with open(f"{args.basedir}/in/in.txt", "wb") as seed: seed.write(b"\x00" * 10240)
-    return {
-        "AFL_BENCH_JUST_ONE": "1", "AFL_DISABLE_TRIM": "1", "AFL_I_DONT_CARE_ABOUT_MISSING_CRASHES": "1",
-        "AFL_NO_UI": "1", "AFL_TRY_AFFINITY": "1", "PATH": str(Path("../").resolve()),
-    }
+    with open(f"{args.basedir}/in/in.txt", "wb") as seed:
+        seed.write(b"\x00" * 10240)
 
 async def compile_target(source: Path, binary: Path) -> None:
     print(f" [*] Compiling the {binary} fuzzing harness for the benchmark to use.")
     (returncode, stdout, stderr) = await run_command(
-        [str(Path("../afl-clang-lto").resolve()), "-o", str(Path(binary.resolve())), str(Path(source).resolve())],
-        env={"AFL_LLVM_INSTRUMENT": "PCGUARD"},
+        [str(Path("../afl-clang-lto").resolve()), "-o", str(Path(binary.resolve())), str(Path(source).resolve())]
     )
-    if returncode != 0:
-        print(yellow(f" [*] afl-clang-lto was unable to compile; falling back to afl-cc."))
-
+    if returncode == 0:
+        return
+    print(yellow(f" [*] afl-clang-lto was unable to compile; falling back to afl-cc."))
     (returncode, stdout, stderr) = await run_command(
-        [str(Path("../afl-cc").resolve()), "-o", str(Path(binary.resolve())), str(Path(source).resolve())],
-        env={"AFL_LLVM_INSTRUMENT": "PCGUARD"},
+        [str(Path("../afl-cc").resolve()), "-o", str(Path(binary.resolve())), str(Path(source).resolve())]
     )
     if returncode != 0:
         sys.exit(red(f" [*] Error: afl-cc is unable to compile: {stderr.decode()} {stdout.decode()}"))
 
-async def run_command(cmd: list[str], env: Union[dict, None]) -> tuple[Union[int, None], bytes, bytes]:
-    debug(f"Launching command: {cmd} with env {env}")
+async def run_command(cmd: list[str]) -> tuple[Union[int, None], bytes, bytes]:
+    debug(f"Launching command: {cmd} with env {env_vars}")
     p = await asyncio.create_subprocess_exec(
-        *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=env
+        *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=env_vars
     )
     stdout, stderr = await p.communicate()
     debug(f"Output: {stdout.decode()} {stderr.decode()}")
@@ -142,13 +147,13 @@ async def check_deps() -> None:
         os.path.exists(Path("../SanitizerCoveragePCGUARD.so").resolve())):
         sys.exit(red(" [*] Compile AFL++: we need afl-fuzz, afl-clang-fast and SanitizerCoveragePCGUARD.so built."))
 
-    (returncode, stdout, stderr) = await run_command([str(Path("../afl-cc").resolve()), "-v"], env={})
+    (returncode, stdout, stderr) = await run_command([str(Path("../afl-cc").resolve()), "-v"])
     if returncode != 0:
         sys.exit(red(f" [*] Error: afl-cc -v returned: {stderr.decode()} {stdout.decode()}"))
     compiler = ""
     target_arch = ""
     for line in stderr.decode().split("\n"):
-        if m := re.match(r"^(clang version .*)", line):
+        if m := re.match(r"(.* clang version .*?)\s", line):
             compiler = m.group(1)
         elif m := re.match(r"^Target: (.*)", line):
             target_arch = m.group(1)
@@ -160,7 +165,6 @@ async def check_deps() -> None:
         print(yellow(f" [*] afl-persistent-config did not run; run it to improve performance (and decrease security)."))
     if not afl_sc:
         print(yellow(f" [*] afl-system-config did not run; run it to improve performance (and decrease security)."))
-
     results.config = Config(afl_persistent_config=afl_pc, afl_system_config=afl_sc, afl_version="",
                             comment=args.comment, compiler=compiler, target_arch=target_arch)
 
@@ -171,12 +175,41 @@ async def colon_values(filename: str, searchKey: str) -> list[str]:
         v_list = [v.rstrip() for k, v in kv_pairs if k.rstrip() == searchKey]
         return v_list
 
+async def describe_afl_config() -> str:
+   if results.config is None:
+       return "unknown"
+   elif results.config.afl_persistent_config and results.config.afl_system_config:
+       return "both"
+   elif results.config.afl_persistent_config:
+       return "persistent"
+   elif results.config.afl_system_config:
+       return "system"
+   else:
+       return "none"
+
 async def save_benchmark_results() -> None:
     """Append a single row to the benchmark results in JSON Lines format (which is simple to write and diff)."""
     with open("benchmark-results.jsonl", "a") as jsonfile:
         json.dump(asdict(results), jsonfile, sort_keys=True)
         jsonfile.write("\n")
         print(blue(f" [*] Results have been written to {jsonfile.name}"))
+    with open("COMPARISON", "a") as comparisonfile:
+        described_config = await describe_afl_config()
+        if results.hardware is None:
+            return
+        cpu_model = results.hardware.cpu_model.ljust(42)
+        cpu_mhz = str(round(results.hardware.cpu_fastest_core_mhz)).ljust(5)
+        if "test-instr-persist-shmem" in results.targets and "multicore" in results.targets["test-instr-persist-shmem"]:
+            if results.targets["test-instr-persist-shmem"]["singlecore"] is None or \
+               results.targets["test-instr-persist-shmem"]["multicore"] is None:
+                return
+            single = str(round(results.targets["test-instr-persist-shmem"]["singlecore"].total_execs_per_sec)).ljust(10)
+            multi = str(round(results.targets["test-instr-persist-shmem"]["multicore"].total_execs_per_sec)).ljust(9)
+            if len(cpu_model) > 30:
+                cpu_model = cpu_model[:30]
+            comparisonfile.write(f"{cpu_model} | {cpu_mhz} | {single} | {multi} | {described_config.ljust(12)} |\n")
+    with open("COMPARISON", "r") as comparisonfile:
+        print(comparisonfile.read())
 
 
 async def main() -> None:
@@ -185,17 +218,24 @@ async def main() -> None:
     except FileNotFoundError:
         pass
     await check_deps()
-    cpu_mhz_str = await colon_values("/proc/cpuinfo", "cpu MHz")
-    cpu_mhz = max([float(c) for c in cpu_mhz_str]) # use the fastest CPU MHz for now
-    cpu_model = await colon_values("/proc/cpuinfo", "model name")
-    # Only record the first core's speed for now, even though it can vary between cores.
+    if args.mhz:
+        cpu_mhz = float(args.mhz)
+    else:
+        cpu_mhz_str = await colon_values("/proc/cpuinfo", "cpu MHz")
+        if len(cpu_mhz_str) == 0:
+            cpu_mhz_str.append("0")
+        cpu_mhz = max([float(c) for c in cpu_mhz_str]) # use the fastest CPU MHz for now
+    if args.cpu:
+        cpu_model = [args.cpu]
+    else:
+        cpu_model = await colon_values("/proc/cpuinfo", "model name") or [""]
     results.hardware = Hardware(cpu_fastest_core_mhz=cpu_mhz, cpu_model=cpu_model[0], cpu_threads=cpu_count)
-    env_vars = await prep_env()
+    await prep_env()
     print(f" [*] Ready, starting benchmark...")
-    for target in targets:
+    for target in chosen_targets:
         await compile_target(target.source, target.binary)
         binary = str(target.binary)
-        for mode in modes:
+        for mode in chosen_modes:
             afl_execs_per_sec, execs_total, run_time_total = ([] for _ in range(3))
             for run_idx in range(0, args.runs):
                 print(gray(f" [*] {mode.name} {binary} run {run_idx+1} of {args.runs}, execs/s: "), end="", flush=True)
@@ -207,7 +247,7 @@ async def main() -> None:
                     cmds.append(["afl-fuzz", "-i", f"{args.basedir}/in"] + name + ["-s", "123", "-D", f"./{binary}"])
 
                 # Prepare the afl-fuzz tasks, and then block while waiting for them to finish.
-                fuzztasks = [run_command(cmds[cpu], env_vars) for cpu in fuzzers]
+                fuzztasks = [run_command(cmds[cpu]) for cpu in fuzzers]
                 start_time = datetime.datetime.now()
                 await asyncio.gather(*fuzztasks)
                 end_time = datetime.datetime.now()