From 47154cf2516e0a47c4290d3429fab27c2f5e2007 Mon Sep 17 00:00:00 2001 From: Emery Berger Date: Thu, 15 Aug 2024 16:22:14 -0400 Subject: [PATCH] Fixed GPU stats merging (#851) * Added handling for n_gpu_samples and gpu_mem_samples. * Include files with any GPU samples. --- scalene/scalene_json.py | 2 ++ scalene/scalene_statistics.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/scalene/scalene_json.py b/scalene/scalene_json.py index eb59c7517..862203c1f 100644 --- 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 index 002bb48c5..06742b7b3 100644 --- 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 )