mirror of
https://github.com/andreinechaev/nvcc4jupyter.git
synced 2026-06-13 18:50:47 +05:30
Ștefan-Cosmin Ciocan
887c809d07
* Use NVIDIA Nsight Compute CLI profiler * Add profile and profiler-args options to argument parser. * Add missing comma to profiler-args option. * Use profile args in version 1 of the plugin * Change profiler-args option to take all remaining arguments * Change profiler_args type from string to list of strings * Add profile option to version 2 of the plugin * Add profiler usage instructions
63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
import os
|
|
import subprocess
|
|
import tempfile
|
|
import uuid
|
|
|
|
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'
|
|
|
|
|
|
@magics_class
|
|
class NVCCPlugin(Magics):
|
|
|
|
def __init__(self, shell):
|
|
super(NVCCPlugin, self).__init__(shell)
|
|
|
|
self.argparser = helper.get_argparser()
|
|
|
|
@staticmethod
|
|
def compile(file_path):
|
|
subprocess.check_output(
|
|
[compiler, file_path + ext, "-o", file_path + ".out", '-Wno-deprecated-gpu-targets'], stderr=subprocess.STDOUT)
|
|
|
|
def run(self, file_path, timeit=False, profile=False, profiler_args=[]):
|
|
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)
|
|
output = str(output) # convert TimeitResult object to human readable string
|
|
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
|
|
|
|
@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, profile=args.profile, profiler_args=args.profiler_args)
|
|
except subprocess.CalledProcessError as e:
|
|
helper.print_out(e.output.decode("utf8"))
|
|
output = None
|
|
return output
|