Export input/output data to directory (#168)

Summary:
We need to pickle input/output data to directory for numeric checks across runs, such as CUDA version and Triton version.

Pull Request resolved: #168

Test Plan:
```
python run.py --op flash_attention --only triton_tutorial_flash_v2 --num-inputs 1 --export both --export-dir $PWD/flash_attention

python run.py --op flash_attention --only triton_tutorial_flash_v2 --num-inputs 1 --export both --export-dir .data/flash_attention/ --bwd --causal
```

```
ls .data/flash_attention

-rw-r--r-- 1 xzhao9 users 9438227 Feb 28 07:40 'x_(4, 48, 128, 64)-input.pkl'
-rw-r--r-- 1 xzhao9 users 9438257 Feb 28 07:39 'x_(4, 48, 128, 64)-triton_tutorial_flash_v2-bwd-grad.pkl'
-rw-r--r-- 1 xzhao9 users 3146151 Feb 28 07:40 'x_(4, 48, 128, 64)-triton_tutorial_flash_v2-fwd-output.pkl'
```

Fixes #84

Reviewed By: adamomainz

Differential Revision: D70394911

Pulled By: xuzhao9

fbshipit-source-id: 468980b0b50684a278f309afc8af0d0a55f8008f

Author: xuzhao9

tritonbench/components/export/__init__.py

@@ -0,0 +1 @@
+from .export import export_data

tritonbench/components/export/export.py

@@ -0,0 +1,50 @@
+"""
+Serialize pickled tensors to directory.
+"""
+
+import pickle
+from pathlib import Path
+
+from typing import Any, Callable
+
+from tritonbench.utils.input import input_cast
+
+
+def get_input_gradients(inputs):
+    all_input_grads = []
+    input_cast(lambda x: True, lambda y: all_input_grads.append(y.grad), inputs)
+    return all_input_grads
+
+
+def export_data(
+    x_val: str,
+    inputs: Any,
+    fn_mode: str,
+    fn: Callable,
+    export_type: str,
+    export_dir: str,
+):
+    # pickle naming convention
+    # x_<x_val>-input.pkl
+    # x_<x_val>-<fn_name>-fwd-output.pkl
+    # x_<x_val>-<fn_name>-bwd-grad.pkl
+    assert export_dir, "Export dir must be specified."
+    export_path = Path(export_dir)
+    assert export_path.exists(), f"Export path {export_dir} must exist."
+    if export_type == "input" or export_type == "both":
+        input_file_name = f"x_{x_val}-input.pkl"
+        input_file_path = export_path.joinpath(input_file_name)
+        with open(input_file_path, "wb") as ifp:
+            pickle.dump(inputs, ifp)
+    if export_type == "output" or export_type == "both":
+        if fn_mode == "fwd":
+            output_type = "output"
+            output = fn()
+        elif fn_mode == "bwd":
+            output_type = "grad"
+            # output of the backward pass are the input gradients
+            output = get_input_gradients(inputs)
+        output_file_name = f"x_{x_val}-{fn._name}-{fn_mode}-{output_type}.pkl"
+        output_file_path = export_path.joinpath(output_file_name)
+        with open(output_file_path, "wb") as ofp:
+            pickle.dump(output, ofp)

tritonbench/utils/parser.py

@@ -180,6 +180,18 @@ def get_parser(args=None):
         action="store_true",
         help="Include cold start time in compile_time and compile_trace.",
     )
+    parser.add_argument(
+        "--export",
+        default=None,
+        choices=["in", "out", "both"],
+        help="Export input or output. Must be used together with --export-dir.",
+    )
+    parser.add_argument(
+        "--export-dir",
+        default=None,
+        type=str,
+        help="The directory to store input or output.",
+    )
 
     if IS_FBCODE:
         parser.add_argument("--log-scuba", action="store_true", help="Log to scuba.")

tritonbench/utils/triton_op.py

@@ -24,6 +24,7 @@
 import triton
 
 from tritonbench.components.do_bench import do_bench_wrapper, Latency
+from tritonbench.components.export import export_data
 from tritonbench.components.ncu import ncu_analyzer, nsys_analyzer
 from tritonbench.utils.env_utils import apply_precision, set_env, set_random_seed
 from tritonbench.utils.input import input_cast
@@ -1371,6 +1372,15 @@ def _init_extra_metrics() -> Dict[str, Any]:
                     use_cuda_profiler_range=True,
                 )
                 metrics.extra_metrics["_nsys_rep_in_task"] = "success"
+            if self.tb_args.export:
+                export_data(
+                    x_val=self.get_x_val(self.example_inputs),
+                    inputs=self.example_inputs,
+                    fn_mode=self.mode.value,
+                    fn=fn,
+                    export_type=self.tb_args.export,
+                    export_dir=self.tb_args.export_dir,
+                )
             # generate customized metrics
             if self.name in REGISTERED_METRICS:
                 for metric_name in REGISTERED_METRICS[self.name]: