Add tests for serialization and deserialization (#49)

Co-authored-by: Mark Saroufim <[email protected]>
Co-authored-by: Bert Maher <[email protected]>

Author: 3 people

BackendBench/torchbench_suite.py

@@ -2,68 +2,20 @@
 Load aten inputs from serialized txt files.
 """
 
-import math
 import re
 import tempfile
 from collections import defaultdict
 from pathlib import Path
 
 import requests
 import torch
-from torch.testing import make_tensor
+from BackendBench.utils import deserialize_args
 
 # the schema for this dataset is the one defined in tritonbench traces.
 # ie. https://github.com/pytorch-labs/tritonbench/blob/main/tritonbench/data/input_configs/hf_train/AlbertForMaskedLM_training.txt
 DEFAULT_HUGGINGFACE_URL = "https://huggingface.co/datasets/GPUMODE/huggingface_op_trace/resolve/main/tritonbench_op_trace.txt"
 
 
-dtype_abbrs = {
-    torch.bfloat16: "bf16",
-    torch.float64: "f64",
-    torch.float32: "f32",
-    torch.float16: "f16",
-    torch.complex32: "c32",
-    torch.complex64: "c64",
-    torch.complex128: "c128",
-    torch.int8: "i8",
-    torch.int16: "i16",
-    torch.int32: "i32",
-    torch.int64: "i64",
-    torch.bool: "b8",
-    torch.uint8: "u8",
-}
-
-dtype_abbrs_parsing = {value: key for key, value in dtype_abbrs.items()}
-
-_FLOATING_TYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
-
-
-def _deserialize_tensor(size, dtype, stride=None, device="cuda"):
-    kwargs = {}
-    if dtype in _FLOATING_TYPES:
-        kwargs.update({"low": 0, "high": 1})
-    if stride is not None:
-        extent = 1 + sum((size - 1) * stride for size, stride in zip(size, stride))
-        data = make_tensor(extent, dtype=dtype, device=device, **kwargs)
-        return data.as_strided(size, stride)
-    return make_tensor(size, dtype=dtype, device=device, **kwargs)
-
-
-def _deserialize_args(inps):
-    inps = inps.strip().strip("'")
-    global_vals = {
-        "T": _deserialize_tensor,
-        "th": torch,
-        "inf": math.inf,
-        "torch": torch,
-        **dtype_abbrs_parsing,
-    }
-    # f strings introduce quotations we dont want
-    for key in dtype_abbrs_parsing:
-        inps = inps.replace(f"'{key}'", key)
-    return eval(inps.strip().strip("'").strip('"'), global_vals)
-
-
 class TorchBenchTest:
     def __init__(self, *args, **kwargs):
         self.args = args
@@ -89,7 +41,7 @@ def __init__(self, op, inputs, topn):
     def tests(self):
         inputs_and_sizes = []
         for inp in self.inputs:
-            args, kwargs = _deserialize_args(inp)
+            args, kwargs = deserialize_args(inp)
             size = _args_size(args) + _args_size(list(kwargs.values()))
             inputs_and_sizes.append((size, inp))
         ret = [x[1] for x in sorted(inputs_and_sizes, reverse=True)]
@@ -98,13 +50,13 @@ def tests(self):
     @property
     def correctness_tests(self):
         for inp in self.tests():
-            args, kwargs = _deserialize_args(inp)
+            args, kwargs = deserialize_args(inp)
             yield TorchBenchTest(*args, **kwargs)
 
     @property
     def performance_tests(self):
         for inp in self.tests():
-            args, kwargs = _deserialize_args(inp)
+            args, kwargs = deserialize_args(inp)
             yield TorchBenchTest(*args, **kwargs)
 
 

BackendBench/utils.py

@@ -2,6 +2,29 @@
 import inspect
 import re
 import textwrap
+import math
+import torch
+from torch.testing import make_tensor
+
+dtype_abbrs = {
+    torch.bfloat16: "bf16",
+    torch.float64: "f64",
+    torch.float32: "f32",
+    torch.float16: "f16",
+    torch.complex32: "c32",
+    torch.complex64: "c64",
+    torch.complex128: "c128",
+    torch.int8: "i8",
+    torch.int16: "i16",
+    torch.int32: "i32",
+    torch.int64: "i64",
+    torch.bool: "b8",
+    torch.uint8: "u8",
+}
+
+dtype_abbrs_parsing = {value: key for key, value in dtype_abbrs.items()}
+
+_FLOATING_TYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
 
 
 def uses_cuda_stream(func) -> bool:
@@ -51,3 +74,82 @@ def visit_Call(self, node):
     finder = StreamCreationFinder()
     finder.visit(tree)
     return finder.found
+
+
+def _deserialize_tensor(size, dtype, stride=None, device="cuda"):
+    kwargs = {}
+    if dtype in _FLOATING_TYPES:
+        kwargs.update({"low": 0, "high": 1})
+
+    # Fall back to CPU if CUDA is not available
+    if device == "cuda" and not torch.cuda.is_available():
+        device = "cpu"
+
+    if stride is not None:
+        extent = 1 + sum((size - 1) * stride for size, stride in zip(size, stride))
+        data = make_tensor(extent, dtype=dtype, device=device, **kwargs)
+        return data.as_strided(size, stride)
+    return make_tensor(size, dtype=dtype, device=device, **kwargs)
+
+
+def _serialize_tensor(tensor):
+    """Helper function to serialize a tensor to string format"""
+    shape = list(tensor.shape)
+    dtype = dtype_abbrs[tensor.dtype]
+    stride = tensor.stride() if not tensor.is_contiguous() else None
+
+    if stride:
+        return f"T({shape}, {dtype}, {list(stride)})"
+    else:
+        return f"T({shape}, {dtype})"
+
+
+def _serialize_value(value):
+    """Helper function to serialize any value (tensor, list, primitive)"""
+    if isinstance(value, torch.Tensor):
+        return _serialize_tensor(value)
+    elif isinstance(value, list):
+        list_parts = [_serialize_value(item) for item in value]
+        return f"[{', '.join(list_parts)}]"
+    else:
+        return repr(value)
+
+
+def serialize_args(args, kwargs) -> str:
+    """Convert args and kwargs back to the BackendBench string format
+
+    Args:
+        args: List of arguments (can contain tensors, lists, primitives)
+        kwargs: Dict of keyword arguments
+
+    Returns:
+        Serialized string in format: (arg1, arg2, ..., key1=val1, key2=val2, ...)
+    """
+    if args is None or kwargs is None:
+        return "None"
+
+    # Process positional arguments
+    parts = [_serialize_value(arg) for arg in args]
+
+    # Process keyword arguments
+    kwargs_parts = [f"'{key}': {_serialize_value(val)}" for key, val in kwargs.items()]
+
+    # Handle empty args tuple properly
+    args_str = f"({', '.join(parts)},)" if parts else "()"
+
+    return f"({args_str}, {{{', '.join(kwargs_parts)}}})"
+
+
+def deserialize_args(inps):
+    inps = inps.strip().strip("'")
+    global_vals = {
+        "T": _deserialize_tensor,
+        "th": torch,
+        "inf": math.inf,
+        "torch": torch,
+        **dtype_abbrs_parsing,
+    }
+    # f strings introduce quotations we dont want
+    for key in dtype_abbrs_parsing:
+        inps = inps.replace(f"'{key}'", key)
+    return eval(inps.strip().strip("'").strip('"'), global_vals)