NOT TO BE MERGED: Python script to benchmark mxr files - consumes mxr files with json encoding

src/targets/gpu/compile_ops.cpp

@@ -37,6 +37,7 @@
 #include <migraphx/builtin.hpp>
 #include <migraphx/load_save.hpp>
 #include <migraphx/filesystem.hpp>
+#include <migraphx/json.hpp>
 #include <migraphx/gpu/compiler.hpp>
 #include <migraphx/gpu/compile_ops.hpp>
 #include <migraphx/gpu/context.hpp>
@@ -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<std::string>{}(to_string(config->problem));
             auto mxr_file     = mxr_dir / (preop.name() + "_" + std::to_string(i) + "_" +

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 <mxr_dir>
+"""
+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]} <mxr_dir>")
+    main(sys.argv[1])