Add option to choose between NSYS and NCU profilers

nvcc4jupyter/__init__.py

@@ -2,7 +2,7 @@
 nvcc4jupyter: CUDA C++ plugin for Jupyter Notebook
 """
 
-from .parsers import set_defaults  # noqa: F401
+from .parsers import Profiler, set_defaults  # noqa: F401
 from .plugin import NVCCPlugin, load_ipython_extension  # noqa: F401
 
 __version__ = "1.1.0"

nvcc4jupyter/parsers.py

@@ -3,14 +3,28 @@ Parsers for the CUDA magic commands.
 """
 
 import argparse
-from typing import Callable, Optional
+from enum import Enum
+from typing import Callable, Optional, Type, TypeVar
 
+
+class Profiler(Enum):
+    """Choice between Nsight Compute and Nsight Systems profilers."""
+
+    NCU = "ncu"
+    NSYS = "nsys"
+
+
+_default_profiler: Profiler = Profiler.NCU
 _default_profiler_args: str = ""
 _default_compiler_args: str = ""
 
+T = TypeVar("T")
+
 
 def set_defaults(
-    compiler_args: Optional[str] = None, profiler_args: Optional[str] = None
+    profiler: Optional[Profiler] = None,
+    compiler_args: Optional[str] = None,
+    profiler_args: Optional[str] = None,
 ) -> None:
     """
     Set the default values for various arguments of the magic commands. These
@@ -18,17 +32,22 @@ def set_defaults(
     to override this behaviour on a cell by cell basis.
 
     Args:
+        profiler: If not None, this value becomes the new default profiler.
+            Defaults to None.
         compiler_args: If not None, this value becomes the new default compiler
-            config. Defaults to "".
+            config. Defaults to None.
         profiler_args: If not None, this value becomes the new default profiler
-            config. Defaults to "".
+            config. Defaults to None.
     """
 
     # pylint: disable=global-statement
+    global _default_profiler
+    if profiler is not None:
+        _default_profiler = profiler
     global _default_compiler_args
-    global _default_profiler_args
     if compiler_args is not None:
         _default_compiler_args = compiler_args
+    global _default_profiler_args
     if profiler_args is not None:
         _default_profiler_args = profiler_args
 
@@ -38,6 +57,11 @@ def str_to_lambda(arg: str) -> Callable[[], str]:
     return lambda: arg
 
 
+def class_to_lambda(arg: str, cls: Type[T]) -> Callable[[], T]:
+    """Convert string value to class and then to lambda"""
+    return lambda: cls(arg)
+
+
 def get_parser_cuda() -> argparse.ArgumentParser:
     """
     %%cuda magic command parser.
@@ -52,8 +76,14 @@ def get_parser_cuda() -> argparse.ArgumentParser:
     parser.add_argument("-t", "--timeit", action="store_true")
     parser.add_argument("-p", "--profile", action="store_true")
 
-    # --profiler-args and --compiler-args values are lambda functions to allow
+    # the type of the following arguments is a lambda lambda function to allow
     # changing the default value at runtime
+    parser.add_argument(
+        "-l",
+        "--profiler",
+        type=lambda arg: class_to_lambda(arg, cls=Profiler),
+        default=lambda: _default_profiler,
+    )
     parser.add_argument(
         "-a",
         "--profiler-args",

nvcc4jupyter/plugin.py

@@ -135,11 +135,12 @@ class NVCCPlugin(Magics):
 
         return executable_fpath
 
-    def _run(
+    def _run(  # pylint: disable=too-many-arguments
         self,
         exec_fpath: str,
         timeit: bool = False,
         profile: bool = False,
+        profiler: parsers.Profiler = parsers.Profiler.NCU,
         profiler_args: str = "",
     ) -> str:
         """
@@ -150,8 +151,9 @@ class NVCCPlugin(Magics):
             timeit: If True, returns the result of the "timeit" magic instead
                 of the standard output of the CUDA process. Defaults to False.
             profile: If True, the executable is profiled with NVIDIA Nsight
-                Compute profiling tool and its output is added to stdout.
+                Compute or NVIDIA Nsight Systems and the profiling output is
-                Defaults to False.
+                added to stdout. Defaults to False.
+            profiler: The profiling tool to use.
             profiler_args: The profiler arguments used to customize the
                 information gathered by it and its overall behaviour. Defaults
                 to an empty string.
@@ -173,7 +175,7 @@ class NVCCPlugin(Magics):
         else:
             run_args = []
             if profile:
-                run_args.extend(["ncu"] + profiler_args.split())
+                run_args.extend([profiler.value] + profiler_args.split())
             run_args.append(exec_fpath)
             output = subprocess.check_output(
                 run_args, stderr=subprocess.STDOUT
@@ -194,6 +196,7 @@ class NVCCPlugin(Magics):
                 exec_fpath=exec_fpath,
                 timeit=args.timeit,
                 profile=args.profile,
+                profiler=args.profiler(),
                 profiler_args=args.profiler_args(),
             )
         except subprocess.CalledProcessError as e: