[inductor] Improve operatorbench.py (#136808)

jansel · pytorchmergebot · commit 375921b755f0 · 2024-09-28T06:22:02.000Z
Pull Request resolved: #136808 Approved by: https://github.com/eellison
diff --git a/benchmarks/dynamo/microbenchmarks/operatorbench.py b/benchmarks/dynamo/microbenchmarks/operatorbench.py
@@ -1,10 +1,15 @@
 #!/usr/bin/env python3
-
+import csv
+import itertools
+import sys
+import time
+import warnings
 from contextlib import nullcontext
 
 import click
 import numpy as np
 from operator_inp_utils import OperatorInputsLoader
+from tqdm import tqdm
 
 import torch
 from torch._dynamo.backends.cudagraphs import cudagraphs_inner
@@ -21,10 +26,13 @@
 profile_enabled = False
 
 
+def maybe_record_function(name):
+    return torch.profiler.record_function(name) if profile_enabled else nullcontext()
+
+
 def compute_speedups(
     operator, models, example_inputs, repeats, accuracy_checking=False, device="cuda"
 ):
-    global profile_enabled
     expected = models[0](*example_inputs)
     if accuracy_checking:
         for model in models[1:]:
@@ -39,20 +47,10 @@ def compute_speedups(
 
     timings = np.zeros((repeats, len(models)), np.float64)
     for rep in range(repeats):
-        record_rep_context = (
-            torch.profiler.record_function(f"rep_{rep}")
-            if profile_enabled
-            else nullcontext()
-        )
-        with record_rep_context:
+        with maybe_record_function(f"rep_{rep}"):
             # interleave the runs to handle frequency scaling and load changes
             for m, model in enumerate(models):
-                record_model_context = (
-                    torch.profiler.record_function(f"model_{m}")
-                    if profile_enabled
-                    else nullcontext()
-                )
-                with record_model_context:
+                with maybe_record_function(f"model_{m}"):
                     if device == "cuda":
                         model(*example_inputs)
 
@@ -94,7 +92,7 @@ def to_channels_last(ten):
 
 
 def microbenchmark(
-    operator, args, kwargs, dtype, accuracy_checking, repeats, measure_nvfuser, device
+    operator, args, kwargs, accuracy_checking, repeats, measure_nvfuser, device
 ):
     gm, gm_args = gen_gm_and_inputs(operator, args, kwargs)
     torch.jit._builtins._register_builtin(
@@ -127,6 +125,13 @@ def microbenchmark(
     return medians
 
 
+quantiles_thresholds = (0.2, 0.5, 0.8)
+
+
+def quantiles(timings):
+    return np.quantile(timings, quantiles_thresholds).tolist()
+
+
 def skip_operator(operator):
     nyi_strings = (
         "aten.gather.default",
@@ -171,15 +176,17 @@ def skip_operator(operator):
     help="suite to load inps from: options: timm, huggingface, torchbench",
     default="torchbench",
 )
-@click.option("--op", help="operator overload to benchmark")
-@click.option("--dtype", help="dtype to benchmark")
+@click.option("--op", help="operator overload to benchmark", default="all")
+@click.option("--dtype", help="dtype to benchmark", default="float32")
 @click.option("--max-samples", help="max samples per op", default=15)
 @click.option("--accuracy-checking", help="check accuracy", default=False)
 @click.option(
     "--repeats", help="how many times to repeat for perf measurement", default=3
 )
 @click.option(
-    "--measure-nvfuser", help="default we only measure inductor", default=False
+    "--measure-nvfuser/--no-measure-nvfuser",
+    help="default we only measure inductor",
+    default=False,
 )
 @click.option("--device", help="cpu or cuda", default="cuda")
 @click.option("--inp-file", help="use custom input file instead of suite", default=None)
@@ -202,7 +209,10 @@ def benchmark(
     channels_last,
     profile,
 ):
+    warnings.filterwarnings("ignore", module="torch.jit._check")
+    torch.set_float32_matmul_precision("high")
     global profile_enabled
+
     if inp_file is not None:
         loader = OperatorInputsLoader(inp_file)
     else:
@@ -216,9 +226,35 @@ def benchmark(
 
     assert dtype in ("float16", "float32"), f"got {dtype}"
 
+    backend_names = ["inductor"]
+    if measure_nvfuser:
+        backend_names.append("nvfuser")
+
+    compare2 = len(backend_names) == 2
+    if compare2:
+        a, b = backend_names
+        backend_names.append(f"{a}/{b}")
+
+    output_fd = None
+    output_csv = None
     if op == "all":
-        filename = f"timings_{suite}_{op.replace('.', '_')}{dtype}.txt"
-        f = open(filename, "a")
+        filename = f"operatorbench_{suite}_{dtype}.csv"
+        output_fd = open(filename, "w")
+        output_csv = csv.writer(output_fd)
+        output_csv.writerow(
+            [
+                "operator",
+                *[
+                    f"{a} {b}"
+                    for a, b in itertools.product(
+                        backend_names,
+                        [f"{x * 100:.0f}th" for x in quantiles_thresholds],
+                    )
+                ],
+                "elapsed",
+                *map("{} abs".format, ["eager", *backend_names]),
+            ]
+        )
 
     dtype = torch.float16 if dtype == "float16" else torch.float32
 
@@ -233,8 +269,7 @@ def benchmark(
     for operator in ops:
         if skip_operator(operator):
             continue
-
-        print(f"Running {operator}")
+        start = time.perf_counter()
         inp_gen = loader.get_inputs_for_operator(operator, dtype=dtype, device=device)
         timings = []
         inputs_list = []
@@ -260,59 +295,62 @@ def benchmark(
             if profile_enabled
             else nullcontext()
         )
-        with profiler_context as prof:
-            for i, inps in enumerate(inputs_list):
+        with profiler_context:
+            for i, inps in enumerate(tqdm(inputs_list[start_idx:], desc=str(operator))):
                 if inps is None:
                     break
-                if i < start_idx:
-                    continue
-                print(f"Iter {i}")
                 args, kwargs = inps
                 if channels_last:
                     args, kwargs = tree_map_only(
                         torch.Tensor, to_channels_last, (args, kwargs)
                     )
                 try:
-                    iter_context = (
-                        torch.profiler.record_function(f"iter_{i}")
-                        if profile_enabled
-                        else nullcontext()
-                    )
-                    with iter_context:
+                    with maybe_record_function(f"iter_{i}"):
                         # aten, nvfuser, inductor
                         timings.append(
                             microbenchmark(
                                 operator,
                                 args,
                                 kwargs,
-                                dtype,
                                 accuracy_checking,
                                 repeats,
                                 measure_nvfuser,
                                 device,
                             )
                         )
-
                 except Exception as e:
-                    print(f"error {operator}")
-                    print(e)
+                    print(f"error {operator} input {i}: {type(e).__name__}: {e}")
                     # comment out this line to avoid blocking other tests
                     # raise e
 
         if not timings:
             continue
 
-        timings = torch.tensor(timings).T
-        q = torch.tensor([0.2, 0.5, 0.8], dtype=torch.float64)
-        output = f"{operator}:\nInductor Speedups : {(torch.quantile(timings[0] / timings[1], q)).tolist()}\n"
-        if measure_nvfuser:
-            output += f"NVFUSER Speedups :{(torch.quantile(timings[0] / timings[2], q)).tolist()}\n"
-        if op == "all":
-            f.write(output)
-        print(output)
-
-    if op == "all":
-        f.close()
+        timings = np.stack(timings)
+        speedups = [
+            quantiles(timings[:, 0] / timings[:, x]) for x in range(1, timings.shape[1])
+        ]
+        if compare2:
+            speedups.append(quantiles(timings[:, 1] / timings[:, 2]))
+        assert len(backend_names) == len(speedups)
+
+        row = [f"{operator}"]
+        sys.stdout.write(f"{operator}: ")
+        for backend, (low, mid, high) in zip(backend_names, speedups):
+            sys.stdout.write(f"{backend}={mid:.4f}x ({low:.4f}-{high:.4f}) ")
+            row.extend(map("{:.6f}".format, [low, mid, high]))
+        elapsed = time.perf_counter() - start
+        row.append(f"{elapsed:1f}")
+        row.extend(map("{:.8f}".format, np.mean(timings, axis=0).tolist()))
+        sys.stdout.write(f"took {elapsed:.0f}s\n")
+        sys.stdout.flush()
+        if output_csv:
+            output_csv.writerow(row)
+            output_fd.flush()
+
+    if output_fd:
+        print(f"Wrote {filename}")
+        output_fd.close()
 
 
 if __name__ == "__main__":
