Merge pull request #192 from niboshi/dump-on-compile-error

Dump CUDA source on compilation error

cupy/cuda/compiler.py

@@ -1,7 +1,9 @@
 import hashlib
+import math
 import os
 import re
 import shutil
+import sys
 import tempfile
 
 import six
@@ -40,6 +42,16 @@ def __exit__(self, exc_type, exc_value, traceback):
         os.rmdir(self.path)
 
 
+def _get_bool_env_variable(name, default):
+    val = os.environ.get(name)
+    if val is None or len(val) == 0:
+        return default
+    try:
+        return int(val) == 1
+    except ValueError:
+        return False
+
+
 def compile_using_nvrtc(source, options=(), arch=None):
     if not arch:
         arch = _get_arch()
@@ -54,13 +66,29 @@ def compile_using_nvrtc(source, options=(), arch=None):
             cu_file.write(source)
 
         prog = _NVRTCProgram(source, cu_path)
-        ptx = prog.compile(options)
+        try:
+            ptx = prog.compile(options)
+        except CompileException as e:
+            dump = _get_bool_env_variable(
+                'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
+            if dump:
+                e.dump(sys.stderr)
+            raise
 
         return ptx
 
 
 def preprocess(source, options=()):
-    pp_src = _NVRTCProgram(source, '').compile(options)
+    prog = _NVRTCProgram(source, '')
+    try:
+        pp_src = prog.compile(options)
+    except CompileException as e:
+        dump = _get_bool_env_variable(
+            'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
+        if dump:
+            e.dump(sys.stderr)
+        raise
+
     if isinstance(pp_src, six.binary_type):
         pp_src = pp_src.decode('utf-8')
     return re.sub('(?m)^#.*$', '', pp_src)
@@ -141,8 +169,11 @@ def compile_with_cache(source, options=(), arch=None, cache_dir=None):
 
 class CompileException(Exception):
 
-    def __init__(self, msg):
+    def __init__(self, msg, source, name, options):
         self._msg = msg
+        self.source = source
+        self.name = name
+        self.options = options
 
     def __repr__(self):
         return str(self)
@@ -153,6 +184,20 @@ def __str__(self):
     def get_message(self):
         return self._msg
 
+    def dump(self, f):
+        lines = self.source.split('\n')
+        digits = int(math.floor(math.log10(len(lines)))) + 1
+        linum_fmt = '{{:0{}d}} '.format(digits)
+        f.write('NVRTC compilation error: {}\n'.format(self))
+        f.write('-----\n')
+        f.write('Name: {}\n'.format(' '.join(self.name)))
+        f.write('Options: {}\n'.format(' '.join(self.options)))
+        f.write('CUDA source:\n')
+        for i, line in enumerate(lines):
+            f.write(linum_fmt.format(i + 1) + line.rstrip() + '\n')
+        f.write('-----\n')
+        f.flush()
+
 
 class _NVRTCProgram(object):
 
@@ -164,6 +209,9 @@ def __init__(self, src, name="default_program", headers=(),
             src = src.decode('UTF-8')
         if isinstance(name, six.binary_type):
             name = name.decode('UTF-8')
+
+        self.src = src
+        self.name = name
         self.ptr = nvrtc.createProgram(src, name, headers, include_names)
 
     def __del__(self):
@@ -176,4 +224,4 @@ def compile(self, options=()):
             return nvrtc.getPTX(self.ptr)
         except nvrtc.NVRTCError:
             log = nvrtc.getProgramLog(self.ptr)
-            raise CompileException(log)
+            raise CompileException(log, self.src, self.name, options)

docs/source/reference/environment.rst

@@ -4,11 +4,15 @@ Environment variables
 Here are the environment variables CuPy uses.
 
 
-+--------------------+----------------------------------------------------+
-| ``CUPY_CACHE_DIR`` | Path to the directory to store kernel cache.       |
-|                    | ``$(HOME)/.cupy.kernel_cache`` is used by default. |
-|                    | See :ref:`overview` for details.                   |
-+--------------------+----------------------------------------------------+
++------------------------------------+----------------------------------------------------+
+| ``CUPY_CACHE_DIR``                 | Path to the directory to store kernel cache.       |
+|                                    | ``$(HOME)/.cupy.kernel_cache`` is used by default. |
+|                                    | See :ref:`overview` for details.                   |
++------------------------------------+----------------------------------------------------+
+| ``CUPY_DUMP_CUDA_SOURCE_ON_ERROR`` | If set to 1, when CUDA kernel compilation fails,   |
+|                                    | CuPy dumps CUDA kernel code to standard error.     |
+|                                    | It is disabled by default.                         |
++------------------------------------+----------------------------------------------------+
 
 
 For install