diff --git a/src/targets/gpu/compile_ops.cpp b/src/targets/gpu/compile_ops.cpp index 57f2e6e8211..520bd93a073 100644 --- a/src/targets/gpu/compile_ops.cpp +++ b/src/targets/gpu/compile_ops.cpp @@ -37,6 +37,7 @@ #include #include #include +#include #include #include #include @@ -489,8 +490,14 @@ struct compile_plan const auto& solution = config->solutions[i]; auto bench_prog = results[i]->make_program(); auto* mm = bench_prog.get_main_module(); - std::string comment_text = preop.name() + " problem=" + to_string(config->problem) + - " solution=" + to_string(solution); + + // Use json encoding for the comment used for benchmarking mxr files. + value comment_val = value::object{}; + comment_val["op"] = preop.name(); + comment_val["problem"] = config->problem; + comment_val["solution"] = solution; + std::string comment_text = to_json_string(comment_val); + mm->add_instruction(builtin::comment{comment_text}, {}); auto problem_hash = std::hash{}(to_string(config->problem)); auto mxr_file = mxr_dir / (preop.name() + "_" + std::to_string(i) + "_" + diff --git a/tools/benchmark_mxr.py b/tools/benchmark_mxr.py new file mode 100644 index 00000000000..92b82b21158 --- /dev/null +++ b/tools/benchmark_mxr.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +"""Benchmark dumped MIGraphX .mxr files and emit a problem cache. + +.mxr files dumped via MIGRAPHX_GPU_DUMP_BENCHMARK_MXR. + +After problem_cache.json file is generated by this script, +use MIGRAPHX_PROBLEM_CACHE=problem_cache.json to consume it. + +Usage: benchmark_mxr.py +""" +import sys +import json +import time +from pathlib import Path + +import migraphx + +WARMUP, ITERATIONS = 5, 50 + + +def benchmark(path): + p = migraphx.load(str(path)) + text = next(i.op().values()["text"] for i in p.get_main_module() + if i.name() == "@comment") + meta = dict(json.loads(text)) + p.finalize(migraphx.get_target("gpu")) + params = {k: migraphx.to_gpu(migraphx.generate_argument(s)) + for k, s in p.get_parameter_shapes().items()} + for _ in range(WARMUP): + p.run(params) + migraphx.gpu_sync() + start = time.perf_counter() + for _ in range(ITERATIONS): + p.run(params) + migraphx.gpu_sync() + avg_ms = (time.perf_counter() - start) / ITERATIONS * 1000 + return meta["op"], meta["problem"], meta["solution"], avg_ms + + +def main(mxr_dir): + best = {} + for path in sorted(Path(mxr_dir).glob("*.mxr")): + try: + op, problem, solution, avg_ms = benchmark(path) + except Exception as e: + print(f"SKIP {path.name}: {e}") + continue + print(f"{avg_ms:8.4f} ms {solution} ({path.name})") + key = (op, json.dumps(problem, sort_keys=True)) + if key not in best or avg_ms < best[key][0]: + best[key] = (avg_ms, op, problem, solution) + + cache = [[{"name": op, "problem": problem}, solution] + for _, op, problem, solution in best.values()] + out = Path(mxr_dir) / "problem_cache.json" + out.write_text(json.dumps(cache, indent=2)) + print(f"\nWrote {len(cache)} entries to {out}") + + +if __name__ == "__main__": + if len(sys.argv) != 2: + sys.exit(f"usage: {sys.argv[0]} ") + main(sys.argv[1])