Update

[ghstack-poisoned]

Author: GregoryComer

backends/test/suite/reporting.py

@@ -1,6 +1,7 @@
 import csv
 from collections import Counter
 from dataclasses import dataclass
+from datetime import timedelta
 from enum import IntEnum
 from functools import reduce
 from typing import TextIO
@@ -108,6 +109,12 @@ class TestCaseSummary:
     a single output tensor.
     """
 
+    quantize_time: timedelta | None = None
+    """ The total runtime of the quantization stage, or none, if the test did not run the quantize stage. """
+
+    lower_time: timedelta | None = None
+    """ The total runtime of the to_edge_transform_and_lower stage, or none, if the test did not run the quantize stage. """
+
 
 class TestSessionState:
     test_case_summaries: list[TestCaseSummary]
@@ -190,6 +197,8 @@ def generate_csv_report(summary: RunSummary, output: TextIO):
         "Backend",
         "Flow",
         "Result",
+        "Quantize Time (s)",
+        "Lowering Time (s)",
     ]
 
     # Tests can have custom parameters. We'll want to report them here, so we need
@@ -230,6 +239,12 @@ def generate_csv_report(summary: RunSummary, output: TextIO):
             "Backend": record.backend,
             "Flow": record.flow,
             "Result": record.result.display_name(),
+            "Quantize Time (s)": (
+                record.quantize_time.total_seconds() if record.quantize_time else None
+            ),
+            "Lowering Time (s)": (
+                record.lower_time.total_seconds() if record.lower_time else None
+            ),
         }
         if record.params is not None:
             row.update({k.capitalize(): v for k, v in record.params.items()})

backends/test/suite/runner.py

@@ -1,8 +1,10 @@
 import argparse
 import importlib
 import re
+import time
 import unittest
 
+from datetime import timedelta
 from typing import Any
 
 import torch
@@ -43,6 +45,7 @@ def run_test(  # noqa: C901
     """
 
     error_statistics: list[ErrorStatistics] = []
+    extra_stats = {}
 
     # Helper method to construct the summary.
     def build_result(
@@ -57,6 +60,7 @@ def build_result(
             result=result,
             error=error,
             tensor_error_statistics=error_statistics,
+            **extra_stats,
         )
 
     # Ensure the model can run in eager mode.
@@ -71,11 +75,16 @@ def build_result(
         return build_result(TestResult.UNKNOWN_FAIL, e)
 
     if flow.quantize:
+        start_time = time.perf_counter()
         try:
             tester.quantize(
                 flow.quantize_stage_factory() if flow.quantize_stage_factory else None
             )
+            elapsed = time.perf_counter() - start_time
+            extra_stats["quantize_time"] = timedelta(seconds=elapsed)
         except Exception as e:
+            elapsed = time.perf_counter() - start_time
+            extra_stats["quantize_time"] = timedelta(seconds=elapsed)
             return build_result(TestResult.QUANTIZE_FAIL, e)
 
     try:
@@ -86,9 +95,14 @@ def build_result(
     except Exception as e:
         return build_result(TestResult.EXPORT_FAIL, e)
 
+    lower_start_time = time.perf_counter()
     try:
         tester.to_edge_transform_and_lower()
+        elapsed = time.perf_counter() - lower_start_time
+        extra_stats["lower_time"] = timedelta(seconds=elapsed)
     except Exception as e:
+        elapsed = time.perf_counter() - lower_start_time
+        extra_stats["lower_time"] = timedelta(seconds=elapsed)
         return build_result(TestResult.LOWER_FAIL, e)
 
     is_delegated = any(
@@ -183,6 +197,7 @@ def parse_args():
         "--report",
         nargs="?",
         help="A file to write the test report to, in CSV format.",
+        default="backend_test_report.csv",
     )
     return parser.parse_args()