nvcc.py 1.74 KB
Newer Older
1
# pylint: disable=invalid-name
2
"""Utility to invoke nvcc compiler in the system"""
3
from __future__ import absolute_import as _abs
4 5
import sys
import subprocess
6
from . import util
7

8 9
def compile_cuda(code, target="ptx", arch=None,
                 options=None, path_target=None):
10 11 12 13 14 15 16
    """Compile cuda code with NVCC from env.

    Parameters
    ----------
    code : str
        The cuda code.

17
    target : str
18 19
        The target format

20 21 22
    arch : str
        The architecture

23 24 25
    options : str
        The additional options

26 27 28
    path_target : str, optional
        Output file.

29 30 31 32 33
    Return
    ------
    cubin : bytearray
        The bytearray of the cubin
    """
34
    temp = util.tempdir()
35 36
    if target not in ["cubin", "ptx", "fatbin"]:
        raise ValueError("target must be in cubin, ptx, fatbin")
37 38
    temp_code = temp.relpath("my_kernel.cu")
    temp_target = temp.relpath("my_kernel.%s" % target)
39

40
    with open(temp_code, "w") as out_file:
41
        out_file.write(code)
42 43
    if target == "cubin" and arch is None:
        raise ValueError("arch(sm_xy) must be passed for generating cubin")
44

45
    file_target = path_target if path_target else temp_target
46 47
    cmd = ["nvcc"]
    cmd += ["--%s" % target, "-O3"]
48 49
    if arch:
        cmd += ["-arch", arch]
50 51
    cmd += ["-o", file_target]

52 53
    if options:
        cmd += options
54
    cmd += [temp_code]
55 56 57 58 59 60 61 62 63 64
    args = ' '.join(cmd)

    proc = subprocess.Popen(
        args, shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT)
    (out, _) = proc.communicate()

    if proc.returncode != 0:
        sys.stderr.write("Compilation error:\n")
65
        sys.stderr.write(str(out))
66 67 68
        sys.stderr.flush()
        cubin = None
    else:
69
        cubin = bytearray(open(file_target, "rb").read())
70
    return cubin