diff --git a/nvcc4jupyter/__init__.py b/nvcc4jupyter/__init__.py index 356eb20..87f5779 100644 --- a/nvcc4jupyter/__init__.py +++ b/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" diff --git a/nvcc4jupyter/parsers.py b/nvcc4jupyter/parsers.py index a35e49f..fb626c2 100644 --- a/nvcc4jupyter/parsers.py +++ b/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", diff --git a/nvcc4jupyter/plugin.py b/nvcc4jupyter/plugin.py index 1da4f63..56f32b5 100644 --- a/nvcc4jupyter/plugin.py +++ b/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. - Defaults to False. + Compute or NVIDIA Nsight Systems and the profiling output is + 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: