mirror of
https://github.com/andreinechaev/nvcc4jupyter.git
synced 2026-06-13 18:50:47 +05:30
Add the functionality from v2 by reorganizing compilation into source file groups
This commit is contained in:
Cosmin Ciocan
@@ -1,62 +1,298 @@
|
||||
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 IPython.core.magic import Magics, cell_magic, magics_class
|
||||
from common import helper
|
||||
|
||||
compiler = '/usr/local/cuda/bin/nvcc'
|
||||
profiler = '/usr/local/cuda/bin/ncu'
|
||||
ext = '.cu'
|
||||
DEFAULT_EXEC_FNAME = "cuda_exec.out"
|
||||
SHARED_GROUP_NAME = "shared"
|
||||
|
||||
|
||||
@magics_class
|
||||
class NVCCPlugin(Magics):
|
||||
|
||||
def __init__(self, shell):
|
||||
def __init__(self, shell: InteractiveShell):
|
||||
super(NVCCPlugin, self).__init__(shell)
|
||||
self.shell: InteractiveShell # type hint not provided by parent class
|
||||
|
||||
self.argparser = helper.get_argparser()
|
||||
self.parser_cuda = helper.get_parser_cuda()
|
||||
self.parser_cuda_group_save = helper.get_parser_cuda_group_save()
|
||||
self.parser_cuda_group_delete = helper.get_parser_cuda_group_delete()
|
||||
self.parser_cuda_group_run = helper.get_parser_cuda_group_run()
|
||||
|
||||
@staticmethod
|
||||
def compile(file_path):
|
||||
subprocess.check_output(
|
||||
[compiler, file_path + ext, "-o", file_path + ".out", '-Wno-deprecated-gpu-targets'], stderr=subprocess.STDOUT)
|
||||
self.workdir = tempfile.mkdtemp()
|
||||
print(f'Source files will be saved in "{self.workdir}".')
|
||||
|
||||
def run(self, file_path, timeit=False, profile=False, profiler_args=[]):
|
||||
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(['{file_path}.out'], stderr=subprocess.STDOUT)"
|
||||
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)
|
||||
output = str(output) # convert TimeitResult object to human readable string
|
||||
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([profiler] + profiler_args)
|
||||
run_args.append(file_path + ".out")
|
||||
output = subprocess.check_output(run_args, stderr=subprocess.STDOUT)
|
||||
output = output.decode('utf8')
|
||||
|
||||
helper.print_out(output)
|
||||
return None
|
||||
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 cu(self, line, cell):
|
||||
try:
|
||||
args = self.argparser.parse_args(line.split())
|
||||
except SystemExit as e:
|
||||
self.argparser.print_help()
|
||||
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
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
file_path = os.path.join(tmp_dir, str(uuid.uuid4()))
|
||||
with open(file_path + ext, "w") as f:
|
||||
f.write(cell)
|
||||
try:
|
||||
self.compile(file_path)
|
||||
output = self.run(file_path, timeit=args.timeit, profile=args.profile, profiler_args=args.profiler_args)
|
||||
except subprocess.CalledProcessError as e:
|
||||
helper.print_out(e.output.decode("utf8"))
|
||||
output = None
|
||||
return output
|
||||
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)
|
||||
helper.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)
|
||||
helper.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