mirror of
https://github.com/andreinechaev/nvcc4jupyter.git
synced 2026-06-13 18:50:47 +05:30
multi files working version + refactoring
This commit is contained in:
Andrei Nechaev
@@ -1 +1,25 @@
|
||||
## NVCC Plugin for Jupyter notebook
|
||||
|
||||
### V2 is available
|
||||
|
||||
V2 brings support of multiple source and header files.
|
||||
|
||||
##### Usage
|
||||
|
||||
- Load Extension
|
||||
> `%load_ext nvcc_plugin`
|
||||
|
||||
- Mark a cell to be treated as cuda cell
|
||||
> `%%cuda --name example.cu --compile false`
|
||||
>> NOTE: The cell must contain either code or comments to be run successfully.
|
||||
>> It accepts 2 arguments. `-n` | `--name` - which is the name of either CUDA source or Header
|
||||
>> The name parameter must have extension `.cu` or `.h`
|
||||
>> Second argument `-c` | `--compile`; default value is `false`. The argument is a flag to specify
|
||||
>> if the cell will be compiled and run right away or not. It might be usefull if you're playing in
|
||||
>> the `main` function
|
||||
|
||||
- To compile and run all CUDA files you need to run
|
||||
```
|
||||
%%cuda_run
|
||||
# This line just to bypass an exeption and can contain any text
|
||||
```
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
import argparse
|
||||
|
||||
|
||||
def get_argparser():
|
||||
parser = argparse.ArgumentParser(description='NVCCPlugin params')
|
||||
parser.add_argument("-t", "--timeit", action='store_true',
|
||||
help='flag to return timeit result instead of stdout')
|
||||
return parser
|
||||
+4
-129
@@ -1,135 +1,10 @@
|
||||
import os
|
||||
import uuid
|
||||
import timeit
|
||||
import argparse
|
||||
import tempfile
|
||||
import subprocess
|
||||
import IPython.core.magic as ipym
|
||||
from IPython.core.magic_arguments import (argument, magic_arguments,
|
||||
parse_argstring)
|
||||
|
||||
compiler = '/usr/local/cuda/bin/nvcc'
|
||||
ext = '.cu'
|
||||
|
||||
|
||||
def get_argparser():
|
||||
parser = argparse.ArgumentParser(description='NVCCPlugin params')
|
||||
parser.add_argument("-t", "--timeit", action='store_true',
|
||||
help='flag to return timeit result instead of stdout')
|
||||
return parser
|
||||
|
||||
|
||||
@ipym.magics_class
|
||||
class NVCCPlugin(ipym.Magics):
|
||||
|
||||
def __init__(self, shell):
|
||||
super(NVCCPlugin, self).__init__(shell)
|
||||
self.argparser = get_argparser()
|
||||
|
||||
@staticmethod
|
||||
def compile(file_path):
|
||||
subprocess.check_output([compiler, file_path + ext, "-o", file_path + ".out"], stderr=subprocess.STDOUT)
|
||||
|
||||
def run(self, file_path, timeit=False):
|
||||
if timeit:
|
||||
stmt = f"subprocess.check_output(['{file_path}.out'], stderr=subprocess.STDOUT)"
|
||||
output = self.shell.run_cell_magic(magic_name="timeit", line="-q -o import subprocess", cell=stmt)
|
||||
else:
|
||||
output = subprocess.check_output([file_path + ".out"], stderr=subprocess.STDOUT)
|
||||
output = output.decode('utf8')
|
||||
return output
|
||||
|
||||
@ipym.cell_magic
|
||||
def cu(self, line, cell):
|
||||
try:
|
||||
args = self.argparser.parse_args(line.split())
|
||||
except SystemExit as e:
|
||||
self.argparser.print_help()
|
||||
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)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(e.output.decode("utf8"))
|
||||
output = None
|
||||
return output
|
||||
|
||||
|
||||
out = "result.out"
|
||||
|
||||
|
||||
@ipym.magics_class
|
||||
class NVCCPluginV2(ipym.Magics):
|
||||
|
||||
def __init__(self, shell):
|
||||
super(NVCCPluginV2, self).__init__(shell)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self.output_dir = os.path.join(tmp, str(uuid.uuid4()))
|
||||
|
||||
@staticmethod
|
||||
def compile(file_path):
|
||||
subprocess.check_output([compiler, file_path, "-o", out], stderr=subprocess.STDOUT)
|
||||
|
||||
def run(self, file_path, timeit=False):
|
||||
if timeit:
|
||||
stmt = f"subprocess.check_output(['{out}'], stderr=subprocess.STDOUT)"
|
||||
output = self.shell.run_cell_magic(magic_name="timeit", line="-q -o import subprocess", cell=stmt)
|
||||
else:
|
||||
output = subprocess.check_output([file_path + ".out"], stderr=subprocess.STDOUT)
|
||||
output = output.decode('utf8')
|
||||
return output
|
||||
|
||||
@magic_arguments
|
||||
@argument('-n', '--name', type=str, help='file name that will be produced by the cell. must end with .cu extension')
|
||||
@argument('-c', '--compile', type=bool, help='Should be compiled?')
|
||||
@ipym.cell_magic
|
||||
def cuda(self, line='', cell=None):
|
||||
args = parse_argstring(self.cuda, line)
|
||||
if args.name[:-3] != '.cu':
|
||||
raise Exception('name must end with .cu')
|
||||
|
||||
file_path = os.path.join(self.tmp_dir, args.name)
|
||||
with open(file_path, "w") as f:
|
||||
f.write(cell)
|
||||
|
||||
if args.compile:
|
||||
try:
|
||||
self.compile(file_path)
|
||||
output = self.run(file_path, timeit=args.timeit)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(e.output.decode("utf8"))
|
||||
output = None
|
||||
else:
|
||||
output = f'File written in {file_path}'
|
||||
|
||||
return output
|
||||
|
||||
@ipym.cell_magic
|
||||
def cuda_run(self, line='', cell=None):
|
||||
try:
|
||||
args = self.argparser.parse_args(line.split())
|
||||
except SystemExit:
|
||||
self.argparser.print_help()
|
||||
return
|
||||
|
||||
try:
|
||||
self.compile('*.cu')
|
||||
output = self.run(out, timeit=args.timeit)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(e.output.decode("utf8"))
|
||||
output = None
|
||||
|
||||
return output
|
||||
from v1.v1 import NVCCPlugin as NVCC_V1
|
||||
from v2.v2 import NVCCPluginV2 as NVCC_V2
|
||||
|
||||
|
||||
def load_ipython_extension(ip):
|
||||
nvcc_plugin = NVCCPlugin(ip)
|
||||
nvcc_plugin = NVCC_V1(ip)
|
||||
ip.register_magics(nvcc_plugin)
|
||||
|
||||
nvcc_plugin_v2 = NVCCPluginV2(ip)
|
||||
nvcc_plugin_v2 = NVCC_V2(ip)
|
||||
ip.register_magics(nvcc_plugin_v2)
|
||||
|
||||
@@ -5,7 +5,7 @@ setup(
|
||||
version='0.0.2',
|
||||
author='Andrei Nechaev',
|
||||
author_email='lyfaradey@yahoo.com',
|
||||
py_modules=['nvcc_plugin'],
|
||||
py_modules=['nvcc_plugin', 'v2.v2', 'v1.v1', 'common.helper'],
|
||||
url='htpps://github.com/andreinechaev/nvcc4jupyter',
|
||||
license='LICENSE',
|
||||
description='Jupyter notebook plugin to run CUDA C/C++ code',
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import uuid
|
||||
|
||||
from IPython.core.magic import Magics, cell_magic, magics_class
|
||||
|
||||
compiler = '/usr/local/cuda/bin/nvcc'
|
||||
ext = '.cu'
|
||||
|
||||
|
||||
@magics_class
|
||||
class NVCCPlugin(Magics):
|
||||
|
||||
def __init__(self, shell):
|
||||
super(NVCCPlugin, self).__init__(shell)
|
||||
from common import helper
|
||||
self.argparser = helper.get_argparser()
|
||||
|
||||
@staticmethod
|
||||
def compile(file_path):
|
||||
subprocess.check_output([compiler, file_path + ext, "-o", file_path + ".out"], stderr=subprocess.STDOUT)
|
||||
|
||||
def run(self, file_path, timeit=False):
|
||||
if timeit:
|
||||
stmt = f"subprocess.check_output(['{file_path}.out'], stderr=subprocess.STDOUT)"
|
||||
output = self.shell.run_cell_magic(magic_name="timeit", line="-q -o import subprocess", cell=stmt)
|
||||
else:
|
||||
output = subprocess.check_output([file_path + ".out"], stderr=subprocess.STDOUT)
|
||||
output = output.decode('utf8')
|
||||
return output
|
||||
|
||||
@cell_magic
|
||||
def cu(self, line, cell):
|
||||
try:
|
||||
args = self.argparser.parse_args(line.split())
|
||||
except SystemExit as e:
|
||||
self.argparser.print_help()
|
||||
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)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(e.output.decode("utf8"))
|
||||
output = None
|
||||
return output
|
||||
@@ -0,0 +1,94 @@
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
from IPython.core.magic import Magics, cell_magic, magics_class
|
||||
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
|
||||
|
||||
from common import helper
|
||||
|
||||
compiler = '/usr/local/cuda/bin/nvcc'
|
||||
out = "result.out"
|
||||
|
||||
|
||||
@magics_class
|
||||
class NVCCPluginV2(Magics):
|
||||
|
||||
def __init__(self, shell):
|
||||
super(NVCCPluginV2, self).__init__(shell)
|
||||
self.argparser = helper.get_argparser()
|
||||
current_dir = os.getcwd()
|
||||
self.output_dir = os.path.join(current_dir, 'src')
|
||||
if not os.path.exists(self.output_dir):
|
||||
os.mkdir(self.output_dir)
|
||||
print(f'created output directory at {self.output_dir}')
|
||||
else:
|
||||
print(f'directory {self.output_dir} already exists')
|
||||
|
||||
@staticmethod
|
||||
def compile(output_dir, file_paths):
|
||||
res = subprocess.check_output([compiler, '-I' + output_dir, file_paths, "-o", out], stderr=subprocess.STDOUT)
|
||||
print(res)
|
||||
|
||||
def run(self, timeit=False):
|
||||
if timeit:
|
||||
stmt = f"subprocess.check_output(['{out}'], stderr=subprocess.STDOUT)"
|
||||
output = self.shell.run_cell_magic(magic_name="timeit", line="-q -o import subprocess", cell=stmt)
|
||||
else:
|
||||
output = subprocess.check_output([out], stderr=subprocess.STDOUT)
|
||||
output = output.decode('utf8')
|
||||
return output
|
||||
|
||||
@magic_arguments()
|
||||
@argument('-n', '--name', type=str, help='file name that will be produced by the cell. must end with .cu extension')
|
||||
@argument('-c', '--compile', type=bool, help='Should be compiled?')
|
||||
@cell_magic
|
||||
def cuda(self, line='', cell=None):
|
||||
args = parse_argstring(self.cuda, line)
|
||||
ex = args.name.split('.')[-1]
|
||||
if ex not in ['cu', 'h']:
|
||||
raise Exception('name must end with .cu or .h')
|
||||
|
||||
if not os.path.exists(self.output_dir):
|
||||
print(f'Output directory does not exist, creating')
|
||||
try:
|
||||
os.mkdir(self.output_dir)
|
||||
except OSError:
|
||||
print(f"Creation of the directory {self.output_dir} failed")
|
||||
else:
|
||||
print(f"Successfully created the directory {self.output_dir}")
|
||||
|
||||
file_path = os.path.join(self.output_dir, args.name)
|
||||
with open(file_path, "w") as f:
|
||||
f.write(cell)
|
||||
|
||||
if args.compile:
|
||||
try:
|
||||
self.compile(self.output_dir, file_path)
|
||||
output = self.run(timeit=args.timeit)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(e.output.decode("utf8"))
|
||||
output = None
|
||||
else:
|
||||
output = f'File written in {file_path}'
|
||||
|
||||
return output
|
||||
|
||||
@cell_magic
|
||||
def cuda_run(self, line='', cell=None):
|
||||
try:
|
||||
args = self.argparser.parse_args(line.split())
|
||||
except SystemExit:
|
||||
self.argparser.print_help()
|
||||
return
|
||||
|
||||
try:
|
||||
cuda_src = os.listdir(self.output_dir)
|
||||
cuda_src = [os.path.join(self.output_dir, x) for x in cuda_src if x[-3:] == '.cu']
|
||||
print(f'found sources: {cuda_src}')
|
||||
self.compile(self.output_dir, ' '.join(cuda_src))
|
||||
output = self.run(timeit=args.timeit)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(e.output.decode("utf8"))
|
||||
output = None
|
||||
|
||||
return output
|
||||
Reference in New Issue
Block a user