Benchmark code refactoring (#5632)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2583

Pull Request resolved: #5632

- Benchmark code refactoring

Reviewed By: henrylhtsang

Differential Revision: D100881375

fbshipit-source-id: 2463d5e280ab49738912e32b369e8457ae479212

Author: q10

fbgemm_gpu/fbgemm_gpu/bench/analysis/__init__.py

@@ -23,5 +23,6 @@
     print_table,
     shorten_kernel_name,
 )
+from fbgemm_gpu.bench.analysis.parsers import detect_phase  # noqa: F401
 from fbgemm_gpu.bench.analysis.trace import KinetoTrace  # noqa: F401
 from fbgemm_gpu.bench.analysis.types import KernelStats, StatsMap  # noqa: F401

fbgemm_gpu/fbgemm_gpu/bench/analysis/parsers.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-ignore-all-errors
+"""Generic trace parsing utilities."""
+
+from __future__ import annotations
+
+
+def detect_phase(filename: str) -> str:
+    """Determine benchmark phase from a trace filename.
+
+    Returns ``"fwd"``, ``"fwd_bwd"``, ``"bwd"``, or ``"unknown"``.
+
+    The check order matters: ``_benchphase_fwd_bwd_`` and ``_fwd_bwd_`` must be
+    tested before ``_fwd_`` to avoid misclassifying combined traces as
+    forward-only.
+    """
+    if "_benchphase_fwd_bwd_" in filename or "_fwd_bwd_" in filename:
+        return "fwd_bwd"
+    if "_benchphase_fwd_" in filename or "_fwd_" in filename:
+        return "fwd"
+    if "_bwd_" in filename or "bwd" in filename.lower():
+        return "bwd"
+    return "unknown"

fbgemm_gpu/fbgemm_gpu/bench/analysis/types.py

@@ -62,6 +62,23 @@ def stdev_us(self) -> float:
     # Instance helpers
     # ------------------------------------------------------------------
 
+    def summary_tuple(self) -> tuple[int, float, float, float, float, float]:
+        """Return ``(count, mean_us, median_us, stdev_us, min_us, max_us)``.
+
+        Returns all-NaN tuple when there are no durations.
+        """
+        if not self.durations_us:
+            _nan = float("nan")
+            return 0, _nan, _nan, _nan, _nan, _nan
+        return (
+            self.count,
+            self.mean_us,
+            self.median_us,
+            self.stdev_us,
+            self.min_us,
+            self.max_us,
+        )
+
     def get(self, attr: str, default: float = nan) -> float:
         """Safe attribute access returning *default* when count == 0."""
         if self.count > 0:
@@ -83,6 +100,48 @@ def print_detail(self) -> None:
         print(f"  Stdev:  {self.stdev_us:,.2f} us ({self.stdev_us / 1000:,.4f} ms)")
         print()
 
+    # ------------------------------------------------------------------
+    # Factory classmethods
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_trace_file(
+        cls,
+        trace_file: str,
+        kernel_pattern: str,
+        *,
+        match_mode: str = "contains",
+        device_type: str | None = None,
+        name: str | None = None,
+    ) -> KernelStats:
+        """Load a trace file, extract matching durations, and return aggregated stats.
+
+        This replaces the common 5-line pattern::
+
+            trace = KinetoTrace.from_file(path)
+            bucketed = trace.extract_durations(pattern)
+            durations = [d for durs in bucketed.values() for d in durs]
+            stats = KernelStats(name=pattern, durations_us=durations)
+
+        Args:
+            trace_file: Path to a ``.json`` or ``.json.gz`` Chrome trace.
+            kernel_pattern: Pattern passed to
+                :meth:`KinetoTrace.extract_durations`.
+            match_mode: ``"contains"`` | ``"exact"`` | ``"startswith"`` |
+                ``"regex"``.
+            device_type: ``"cuda"``, ``"cpu"``, or ``None``.
+            name: Name for the returned :class:`KernelStats`.  Defaults to
+                *kernel_pattern*.
+        """
+        from fbgemm_gpu.bench.analysis.trace import KinetoTrace
+
+        trace = KinetoTrace.from_file(trace_file)
+        bucketed = trace.extract_durations(
+            kernel_pattern, match_mode=match_mode, device_type=device_type
+        )
+        durations = [d for durs in bucketed.values() for d in durs]
+        return cls(name=name or kernel_pattern, durations_us=durations)
+
     # ------------------------------------------------------------------
     # Collection-level classmethods
     # ------------------------------------------------------------------