From c8e7c4e25ffa0c359aa419e33c9aacfa14feef79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mos=C3=A8=20Giordano?= <m.giordano@ucl.ac.uk>
Date: Mon, 28 Oct 2024 17:21:52 +0000
Subject: [PATCH] Automatically run a profiler with the executable

---
 benchmarks/modules/utils.py | 43 ++++++++++++++++++++++++++++++++++---
 1 file changed, 40 insertions(+), 3 deletions(-)

diff --git a/benchmarks/modules/utils.py b/benchmarks/modules/utils.py
index 89ef4473..d606cf11 100644
--- a/benchmarks/modules/utils.py
+++ b/benchmarks/modules/utils.py
@@ -7,7 +7,7 @@
 import pprint
 
 import reframe as rfm
-from reframe.core.exceptions import BuildSystemError
+from reframe.core.exceptions import BuildSystemError, CommandLineError
 from reframe.core.logging import getlogger
 from reframe.utility.osext import run_command
 import reframe.utility.osext as osext
@@ -244,6 +244,7 @@ class SpackTest(rfm.RegressionTest):
     build_system = 'Spack'
     spack_spec = variable(str, value='', loggable=True)
     spack_spec_dict = variable(str, value='', loggable=True)
+    profiler = variable(str, value='', loggable=True)
 
     @run_before('compile')
     def setup_spack_environment(self):
@@ -274,7 +275,7 @@ def setup_spack_environment(self):
                                 result_dict = {spec.name: {"compiler": {"name": spec.compiler.name, "version": str(spec.compiler.versions).lstrip("=")}, "variants": {key: str(spec.variants.dict[key].value) if isinstance(spec.variants.dict[key].value, bool) else "" if spec.variants.dict[key].value is None else list(spec.variants.dict[key].value) if isinstance(spec.variants.dict[key].value, tuple) else spec.variants.dict[key].value for key in key_list_for_each[i]},"mpi":str(spec["mpi"]) if "mpi" in spec else ""  } for i, spec in enumerate(spec_list)};\
                                 print(result_dict)'
         self.postrun_cmds.append(f'echo "spack_spec_dict: $(spack -e {self.build_system.environment} python -c \'{cmd_spack_spec_dict}\')"')
-        
+
         # Keep the `spack.lock` file in the output directory so that the Spack
         # environment can be faithfully reproduced later.
         self.keep_files.append(os.path.realpath(os.path.join(self.build_system.environment, 'spack.lock')))
@@ -287,13 +288,49 @@ def get_full_variants(self):
             # convert all single quotes to double quotes since JSON does not recognise it
             self.spack_spec_dict = self.spack_spec_dict.replace("'", "\"")
 
+
     @run_before('compile')
     def setup_build_system(self):
         # The `self.spack_spec` attribute is the user-facing and loggable
         # variable we use for setting the Spack spec, but then we need to
         # forward its value to `self.build_system.specs`, which is the way to
         # inform ReFrame which Spack specs to use.
-        self.build_system.specs = [self.spack_spec]
+        self.build_system.specs.append(self.spack_spec)
+
+
+    @run_before('compile')
+    def add_profiler(self):
+        if self.profiler:
+            if self.profiler == 'vtune':
+                # Spack package providing the profiler
+                self.build_system.specs.append('intel-oneapi-vtune')
+                # Name of output directory
+                vtune_out_dir = 'vtune-profiling'
+                # Prepend VTune call to the executable
+                self.executable = f'vtune -collect hotspots -r {vtune_out_dir} -- ' + self.executable
+                # Save the output directory
+                self.keep_files.append(vtune_out_dir)
+            elif self.profiler == 'nsight':
+                # Spack package providing the profiler
+                self.build_system.specs.append('nvidia-nsight-systems')
+                # Name of output file
+                nsys_out_file = 'nsys-trace'
+                # Prepend nsys call to the executable
+                self.executable = f'nsys profile --trace=cuda,mpi,nvtx,openmp,osrt,opengl,syscall --output {nsys_out_file} ' + self.executable
+                # Save the output file
+                self.keep_files.append(f'{nsys_out_file}.nsys-rep')
+            if self.profiler == 'advisor-roofline':
+                # Spack package providing the profiler
+                self.build_system.specs.append('intel-oneapi-advisor')
+                # Name of output directory
+                advisor_out_dir = 'advisor-roofline'
+                # Prepend advisor call to the executable
+                self.executable = f'advisor -collect roofline --project-dir={advisor_out_dir} -- ' + self.executable
+                # Save the output directory
+                self.keep_files.append(advisor_out_dir)
+            else:
+                raise CommandLineError(f'Unknown profiler {self.profiler}')
+
 
     @run_before('compile')
     def set_sge_num_slots(self):