Make profiler pre-extract CSV files needed by find_slowest (fairinternal/xformers#1413)

__original_commit__ = fairinternal/xformers@000d958

Author: lw

xformers/profiler/find_slowest.py

@@ -77,7 +77,10 @@ def sort_nccl_events(
     ]
 
 
-def parse_one_file(profile_trace_path):
+def read_one_file(profile_trace_path: str) -> pd.DataFrame:
+    if profile_trace_path.endswith(".csv"):
+        return pd.read_csv(profile_trace_path, names=["name", "dur"])
+
     jq_pipe = '.traceEvents[] | select(.cat == "kernel") | [.name, .dur] | @csv'
     if profile_trace_path.endswith(".gz"):
         cmd = (
@@ -97,6 +100,11 @@ def parse_one_file(profile_trace_path):
     finally:
         assert subp.wait() == 0
 
+    return kernel_events
+
+
+def parse_one_file(profile_trace_path: str) -> tuple[pd.DataFrame, pd.DataFrame]:
+    kernel_events = read_one_file(profile_trace_path)
     kernel_events["log_name"] = os.path.basename(profile_trace_path)
 
     communication_kernels = kernel_events[kernel_events.name.str.startswith("nccl")]
@@ -106,9 +114,13 @@ def parse_one_file(profile_trace_path):
 
 
 def print_profiling_info(cuda_profile_dir: str):
-    cuda_profile_path_name = f"{cuda_profile_dir}/*trace.json.gz"
+    cuda_profile_path_name = f"{cuda_profile_dir}/kernels_*.csv"
     profile_files = glob.glob(cuda_profile_path_name)
 
+    if len(profile_files) == 0:
+        cuda_profile_path_name = f"{cuda_profile_dir}/*.pt.trace.json.gz"
+        profile_files = glob.glob(cuda_profile_path_name)
+
     if len(profile_files) == 0:
         cuda_profile_path_name = f"{cuda_profile_dir}/*.json"
         profile_files = glob.glob(cuda_profile_path_name)

xformers/profiler/profiler.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import csv
 import logging
 import os
 import queue
@@ -86,12 +87,29 @@ def _on_trace(self, prof: torch.profiler.profiler.profile) -> None:
         os.makedirs(dir_name, exist_ok=True)
         file_name = f"{worker_name}.{time.time_ns()}.pt.trace.json.gz"
         prof.export_chrome_trace(os.path.join(dir_name, file_name))
+        csv_file_name = f"kernels_{worker_name}.{time.time_ns()}.csv"
+        self._preprocess_trace(prof, os.path.join(dir_name, csv_file_name))
         try:
             self._analyze_trace(prof)
         except Exception as exc:
             self.main_profiler.summary.append(("TraceAnalysis", "Error"))
             logger.warning("Exception analyzing kineto trace", exc_info=exc)
 
+    def _preprocess_trace(
+        self, prof: torch.profiler.profiler.profile, file_name: str
+    ) -> None:
+        if prof.profiler is None or prof.profiler.kineto_results is None:
+            return
+        with open(file_name, "w", newline="") as file:
+            writer = csv.writer(file)
+            for e in prof.profiler.kineto_results.events():
+                if (
+                    e.device_type().name == "CUDA"
+                    and not e.is_user_annotation()
+                    and e.duration_ns() > 0
+                ):
+                    writer.writerow([e.name(), f"{e.duration_ns() / 1_000}"])
+
     def _analyze_trace(self, prof: torch.profiler.profiler.profile) -> None:
         if prof.profiler is None or prof.profiler.kineto_results is None:
             return