From ae802fdd5078c76456a79cb5f0739d85467f1b5c Mon Sep 17 00:00:00 2001
From: David Berard <dberard@fb.com>
Date: Fri, 20 May 2022 14:51:23 -0700
Subject: [PATCH] NVFuser userbenchmarks - use JSON output format (#920)

Summary:
Pull Request resolved: https://github.com/pytorch/benchmark/pull/920

Ran `python run_benchmark.py nvfuser --filter autogen-4`.

There's some stdout, plus the following output that was dumped into `.userbenchmark/nvfuser/metrics-20220519221324.json`:

```json
{
    "name": "nvfuser",
    "environ": {
        "pytorch_git_version": "fdd5462cd93af8a835e5e5616d6c85c7553b9be1"
    },
    "metrics": {
        "no_fuser:autogen-4": 0.360220073023811,
        "nnc-static:autogen-4": 0.34597328188829124,
        "nnc-dynamic:autogen-4": 0.3488327481318265,
        "nvfuser:autogen-4": 0.4082540248055011
    }
}
```

Test Plan: Imported from OSS

Reviewed By: eellison, xuzhao9

Differential Revision: D36553164

Pulled By: davidberard98

fbshipit-source-id: 5487be82e1466b50a3bd102c5fe54bb000334900
---
 userbenchmark/nvfuser/__init__.py | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/userbenchmark/nvfuser/__init__.py b/userbenchmark/nvfuser/__init__.py
index 9861da783a..aa9077be61 100644
--- a/userbenchmark/nvfuser/__init__.py
+++ b/userbenchmark/nvfuser/__init__.py
@@ -1,5 +1,11 @@
+import torch
+
 import argparse
+import json
+import os
+import time
 import torch.utils.jit.log_extract as log_extract
+from datetime import datetime
 from typing import Any, List
 
 def parse_fusers(extra_args: List[str]):
@@ -38,6 +44,21 @@ def get_inputs(self) -> List[Any]:
         return inputs
 
 
+def dump_metrics(metrics):
+    output = {
+        "name": "nvfuser",
+        "environ": {"pytorch_git_version": torch.version.git_version},
+        "metrics": metrics,
+    }
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    target_dir = os.path.normpath(os.path.join(current_dir, "../../.userbenchmark/nvfuser/"))
+    os.makedirs(target_dir, exist_ok=True)
+    fname = "metrics-{}.json".format(datetime.fromtimestamp(time.time()).strftime("%Y%m%d%H%M%S"))
+    full_fname = os.path.join(target_dir, fname)
+    with open(full_fname, 'w') as f:
+        json.dump(output, f, indent=4)
+
+
 def run_nvfuser_microbenchmarks(extra_args: List[str]):
     from userbenchmark.nvfuser.ir import ir_list
     benchmarks = [NVFuserBenchmark(name, ir) for name, ir in ir_list]
@@ -48,12 +69,17 @@ def run_nvfuser_microbenchmarks(extra_args: List[str]):
     if len(fusers) == 0:
         fusers = ["no_fuser", "nnc-static", "nnc-dynamic", "nvfuser"]
 
+    metrics = {}
     for b in benchmarks:
         outputs = []
         for fuser in fusers:
             inputs = b.get_inputs()
-            outputs.append((fuser, b.run_test(inputs, fuser)))
+            runtime = b.run_test(inputs, fuser)
+            outputs.append((fuser, runtime))
+            metrics[f"{fuser}:{b.name}"] = runtime
         print(f"{b.name}:", "; ".join(f"{name} = {time:.3f} ms" for name, time in outputs))
+    dump_metrics(metrics)
+
 
 def run(args: List[str]):
     run_nvfuser_microbenchmarks(extra_args=args)