Fixed GPU stats merging (#851)

* Added handling for n_gpu_samples and gpu_mem_samples. * Include files with any GPU samples.
plasma-umass · Aug 15, 2024 · 47154cf · 47154cf
1 parent bdf9218
commit 47154cf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/scalene/scalene_json.py b/scalene/scalene_json.py
@@ -305,6 +305,7 @@ def output_profiles(
             not stats.total_cpu_samples
             and not stats.total_memory_malloc_samples
             and not stats.total_memory_free_samples
+            and not stats.total_gpu_samples
         ):
             # Nothing to output.
             return {}
@@ -315,6 +316,7 @@ def output_profiles(
                 + list(stats.cpu_samples_c.keys())
                 + list(stats.memory_free_samples.keys())
                 + list(stats.memory_malloc_samples.keys())
+                + list(stats.gpu_samples.keys())
             )
         )
         if not all_instrumented_files:

diff --git a/scalene/scalene_statistics.py b/scalene/scalene_statistics.py
@@ -362,6 +362,8 @@ def build_function_stats(self, filename: Filename) -> ScaleneStatistics:
         "function_map",
         "firstline_map",
         "gpu_samples",
+        "n_gpu_samples",
+        "gpu_mem_samples",
         "total_gpu_samples",
         "memory_malloc_count",
         "memory_free_count",
@@ -458,6 +460,9 @@ def merge_stats(self, the_dir_name: pathlib.Path) -> None:
                 self.increment_per_line_samples(
                     self.n_gpu_samples, x.n_gpu_samples
                 )
+                self.increment_per_line_samples(
+                    self.gpu_mem_samples, x.gpu_mem_samples
+                )
                 self.increment_per_line_samples(
                     self.memcpy_samples, x.memcpy_samples
                 )