mirror of
https://github.com/andreinechaev/nvcc4jupyter.git
synced 2026-06-13 18:50:47 +05:30
Move source files to nvcc4jupyter directory
This commit is contained in:
Cosmin Ciocan
@@ -0,0 +1 @@
|
||||
__version__ = "1.0.0"
|
||||
@@ -0,0 +1,108 @@
|
||||
import argparse
|
||||
|
||||
|
||||
def get_parser_cuda() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"%%cuda magic that compiles and runs CUDA C++ code in this cell."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"-t",
|
||||
"--timeit",
|
||||
action="store_true",
|
||||
help=(
|
||||
'If set, returns the output of the "timeit" built-in ipython magic'
|
||||
" instead of stdout."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--profile",
|
||||
action="store_true",
|
||||
help=(
|
||||
"If set, runs the nvidia nsight compute profiler. Has no effect if"
|
||||
" used with --timeit."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"-a",
|
||||
"--profiler-args",
|
||||
type=str,
|
||||
default="",
|
||||
help=(
|
||||
"Extra options that can be passed to the nvidia nsight compute"
|
||||
" profiler. Must be the last option given to the argument parser"
|
||||
" so you can pass arguments with dashes."
|
||||
),
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def get_parser_cuda_group_run() -> argparse.ArgumentParser:
|
||||
parser = get_parser_cuda()
|
||||
parser.description = (
|
||||
"%%cuda_group_run magic that compiles and runs source files in a"
|
||||
" given group."
|
||||
)
|
||||
parser.add_argument(
|
||||
"-g",
|
||||
"--group",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The group whose files should be compiled and executed.",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def get_parser_cuda_group_save() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"%%cuda_group_save magic that saves CUDA C++ code in this cell for"
|
||||
" later compilation and execution with possibly more source files."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"-n",
|
||||
"--name",
|
||||
type=str,
|
||||
required=True,
|
||||
help=(
|
||||
'The name of the saved source file. Must have either the ".cu" or'
|
||||
' ".h" extension. In order to import a header file saved with this'
|
||||
" magic you can simply add '#include \"<name>\"'."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"-g",
|
||||
"--group",
|
||||
type=str,
|
||||
required=True,
|
||||
help=(
|
||||
"The group to which to add the saved source file. Groups are"
|
||||
" source files that get compiled together and do not interact with"
|
||||
" other groups. This allows you to have multiple unrelated CUDA"
|
||||
" programs within the same jupyter notebook. Adding files to a"
|
||||
' group named "shared" will make them available to all other'
|
||||
" source file groups. One use case for this is sharing error"
|
||||
" handling code."
|
||||
),
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def get_parser_cuda_group_delete() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"%%cuda_group_reset magic that deletes all files in a group."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"-g",
|
||||
"--group",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The group whose files should be deleted.",
|
||||
)
|
||||
return parser
|
||||
|
||||
@@ -0,0 +1,303 @@
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import uuid
|
||||
from typing import List, Optional
|
||||
|
||||
from IPython.core.interactiveshell import InteractiveShell
|
||||
from IPython.core.magic import Magics, cell_magic, line_magic, magics_class
|
||||
|
||||
from . import parsers
|
||||
|
||||
DEFAULT_EXEC_FNAME = "cuda_exec.out"
|
||||
SHARED_GROUP_NAME = "shared"
|
||||
|
||||
|
||||
def print_out(out: str):
|
||||
for l in out.split("\n"):
|
||||
print(l)
|
||||
|
||||
|
||||
@magics_class
|
||||
class NVCCPlugin(Magics):
|
||||
def __init__(self, shell: InteractiveShell):
|
||||
super(NVCCPlugin, self).__init__(shell)
|
||||
self.shell: InteractiveShell # type hint not provided by parent class
|
||||
|
||||
self.parser_cuda = parsers.get_parser_cuda()
|
||||
self.parser_cuda_group_save = parsers.get_parser_cuda_group_save()
|
||||
self.parser_cuda_group_delete = parsers.get_parser_cuda_group_delete()
|
||||
self.parser_cuda_group_run = parsers.get_parser_cuda_group_run()
|
||||
|
||||
self.workdir = tempfile.mkdtemp()
|
||||
print(f'Source files will be saved in "{self.workdir}".')
|
||||
|
||||
def _save_source(
|
||||
self, source_name: str, source_code: str, group_name: str
|
||||
) -> None:
|
||||
"""
|
||||
Save source code as a .cu or .h file in the group directory where
|
||||
files can be compiled together. Saving a source file to the group
|
||||
named "shared" will make those source files available when compiling
|
||||
any group.
|
||||
|
||||
Args:
|
||||
source_name: The name of the source file. Must end in ".cu" or
|
||||
".h".
|
||||
source_code: The source code to be written to the source file.
|
||||
group_name: The name of the group directory where the file will be
|
||||
saved.
|
||||
|
||||
Raises:
|
||||
ValueError: If the source name does not have a proper extension.
|
||||
"""
|
||||
_, ext = os.path.splitext(source_name)
|
||||
if ext != ".cu" and ext != ".h":
|
||||
raise ValueError(
|
||||
f'Given source name "{source_name}" must end in ".h" or ".cu".'
|
||||
)
|
||||
group_dirpath = os.path.join(self.workdir, group_name)
|
||||
os.makedirs(group_dirpath, exist_ok=True)
|
||||
source_fpath = os.path.join(group_dirpath, source_name)
|
||||
with open(source_fpath, "w", encoding="utf-8") as f:
|
||||
f.write(source_code)
|
||||
|
||||
def _delete_group(self, group_name: str) -> None:
|
||||
"""
|
||||
Removes all source files from the given group.
|
||||
|
||||
Args:
|
||||
group_name: The name of the source files group.
|
||||
"""
|
||||
group_dirpath = os.path.join(self.workdir, group_name)
|
||||
if os.path.exists(group_dirpath):
|
||||
shutil.rmtree(group_dirpath)
|
||||
|
||||
def _compile(
|
||||
self, group_name: str, executable_fname: str = DEFAULT_EXEC_FNAME
|
||||
) -> str:
|
||||
"""
|
||||
Compiles all source files in a given group together with all source
|
||||
files from the group named "shared".
|
||||
|
||||
Args:
|
||||
group_name: The name of the source file group to be compiled.
|
||||
executable_fname: The output executable file name. Defaults to
|
||||
"cuda_exec.out".
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the group does not exist or if does not have any
|
||||
source files associated with it.
|
||||
|
||||
Returns:
|
||||
The file path of the resulted executable file.
|
||||
"""
|
||||
shared_dirpath = os.path.join(self.workdir, SHARED_GROUP_NAME)
|
||||
group_dirpath = os.path.join(self.workdir, group_name)
|
||||
if not os.path.exists(group_dirpath):
|
||||
raise RuntimeError(f'Group "{group_name}" does not exist.')
|
||||
|
||||
source_files = list(glob.glob(os.path.join(group_dirpath, "*.cu")))
|
||||
if len(source_files) == 0:
|
||||
raise RuntimeError(
|
||||
f'Group "{group_name}" does not have any source files.'
|
||||
)
|
||||
source_files.extend(
|
||||
list(glob.glob(os.path.join(shared_dirpath, "*.cu")))
|
||||
)
|
||||
|
||||
executable_fpath = os.path.join(group_dirpath, executable_fname)
|
||||
|
||||
args = [
|
||||
"nvcc",
|
||||
"-I" + shared_dirpath + "," + group_dirpath,
|
||||
]
|
||||
args.extend(source_files)
|
||||
args.extend(
|
||||
[
|
||||
"-o",
|
||||
executable_fpath,
|
||||
"-Wno-deprecated-gpu-targets",
|
||||
]
|
||||
)
|
||||
subprocess.check_output(args, stderr=subprocess.STDOUT)
|
||||
|
||||
return executable_fpath
|
||||
|
||||
def _run(
|
||||
self,
|
||||
exec_fpath: str,
|
||||
timeit: bool = False,
|
||||
profile: bool = False,
|
||||
profiler_args: str = "",
|
||||
) -> str:
|
||||
"""
|
||||
Runs a CUDA executable.
|
||||
|
||||
Args:
|
||||
exec_fpath: The file path of the executable.
|
||||
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.
|
||||
profiler_args: The profiler arguments used to customize the
|
||||
information gathered by it and its overall behaviour. Defaults
|
||||
to an empty string.
|
||||
|
||||
Returns:
|
||||
The standard output of the CUDA process or the "timeit" magic
|
||||
output.
|
||||
"""
|
||||
if timeit:
|
||||
stmt = (
|
||||
f"subprocess.check_output(['{exec_fpath}'],"
|
||||
" stderr=subprocess.STDOUT)"
|
||||
)
|
||||
output = self.shell.run_cell_magic(
|
||||
magic_name="timeit", line="-q -o import subprocess", cell=stmt
|
||||
)
|
||||
# convert TimeitResult object to human readable string
|
||||
output = str(output)
|
||||
else:
|
||||
run_args = []
|
||||
if profile:
|
||||
run_args.extend(["ncu"] + profiler_args.split())
|
||||
run_args.append(exec_fpath)
|
||||
output = subprocess.check_output(
|
||||
run_args, stderr=subprocess.STDOUT
|
||||
)
|
||||
output = output.decode("utf8")
|
||||
|
||||
return output
|
||||
|
||||
def _compile_and_run(
|
||||
self, group_name: str, args: argparse.Namespace
|
||||
) -> str:
|
||||
try:
|
||||
exec_fpath = self._compile(group_name)
|
||||
output = self._run(
|
||||
exec_fpath=exec_fpath,
|
||||
timeit=args.timeit,
|
||||
profile=args.profile,
|
||||
profiler_args=args.profiler_args,
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
output = e.output.decode("utf8")
|
||||
return output
|
||||
|
||||
def _read_args(
|
||||
self, line: str, parser: argparse.ArgumentParser
|
||||
) -> Optional[argparse.Namespace]:
|
||||
"""
|
||||
Read arguments from the magic line. Makes sure to keep arguments
|
||||
between double quotes together for use with profiler arguments or
|
||||
compiler arguments.
|
||||
|
||||
Args:
|
||||
line: The arguments on the line of the magic call in the jupyter
|
||||
cell.
|
||||
parser: The parser which will process the arguments after they are
|
||||
correctly tokenized.
|
||||
|
||||
Returns:
|
||||
The parsed arguments.
|
||||
"""
|
||||
tokens = line.strip().split('"')
|
||||
args_tokenized: List[str] = []
|
||||
for index, tok in enumerate(tokens):
|
||||
if index % 2 == 0:
|
||||
# tokens found outside double quotes are split at whitespace
|
||||
args_tokenized.extend(tok.split(" "))
|
||||
else:
|
||||
# anything found between double quotes will not be split
|
||||
args_tokenized.append(tok)
|
||||
args_tokenized = [arg for arg in args_tokenized if len(arg) > 0]
|
||||
|
||||
try:
|
||||
return parser.parse_args(args_tokenized)
|
||||
except SystemExit:
|
||||
parser.print_help()
|
||||
return None
|
||||
|
||||
@cell_magic
|
||||
def cuda(self, line: str, cell: str) -> None:
|
||||
"""Compile and run the CUDA code in the cell.
|
||||
|
||||
Args:
|
||||
line: The arguments on the line of the magic call in the jupyter
|
||||
cell.
|
||||
cell: All of the lines in the jupyter cell besides the magic call
|
||||
itself. It should contain all of the source code to be
|
||||
compiled and run.
|
||||
"""
|
||||
args = self._read_args(line, self.parser_cuda)
|
||||
if args is None:
|
||||
return
|
||||
|
||||
group_name = str(uuid.uuid4())
|
||||
self._save_source(
|
||||
source_name="single_file.cu",
|
||||
source_code=cell,
|
||||
group_name=group_name,
|
||||
)
|
||||
|
||||
output = self._compile_and_run(group_name, args)
|
||||
print_out(output)
|
||||
|
||||
@cell_magic
|
||||
def cuda_group_save(self, line: str, cell: str) -> None:
|
||||
"""
|
||||
Save the CUDA code in the cell in a group of source files to be later
|
||||
compiled and executed by the "cuda_group_run" line magic.
|
||||
|
||||
Args:
|
||||
line: The arguments on the line of the magic call in the jupyter
|
||||
cell.
|
||||
cell: All of the lines in the jupyter cell besides the magic call
|
||||
itself. It should contain all of the source code to be
|
||||
saved.
|
||||
"""
|
||||
args = self._read_args(line, self.parser_cuda_group_save)
|
||||
if args is None:
|
||||
return
|
||||
|
||||
self._save_source(
|
||||
source_name=args.name,
|
||||
source_code=cell,
|
||||
group_name=args.group,
|
||||
)
|
||||
|
||||
@line_magic
|
||||
def cuda_group_run(self, line: str) -> None:
|
||||
"""
|
||||
Compile and run all source files inside a specific source file group.
|
||||
|
||||
Args:
|
||||
line: The arguments on the line of the magic call in the jupyter
|
||||
cell.
|
||||
"""
|
||||
args = self._read_args(line, self.parser_cuda_group_run)
|
||||
if args is None:
|
||||
return
|
||||
|
||||
output = self._compile_and_run(args.group, args)
|
||||
print_out(output)
|
||||
|
||||
@line_magic
|
||||
def cuda_group_delete(self, line: str) -> None:
|
||||
"""
|
||||
Remove all source files inside a specific source file group.
|
||||
|
||||
Args:
|
||||
line: The arguments on the line of the magic call in the jupyter
|
||||
cell.
|
||||
"""
|
||||
args = self._read_args(line, self.parser_cuda_group_delete)
|
||||
if args is None:
|
||||
return
|
||||
|
||||
self._delete_group(args.group)
|
||||
Reference in New Issue
Block a user