Merge pull request #73 from rutsky/compile_fatbin

add option to compile nvcc's fatbin
inducer · Jun 19, 2015 · a7a242e · a7a242e
2 parents 52fe395 + e282410
commit a7a242e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/pycuda/compiler.py b/pycuda/compiler.py
@@ -69,7 +69,7 @@ def preprocess_source(source, options, nvcc):
 def compile_plain(source, options, keep, nvcc, cache_dir, target="cubin"):
     from os.path import join
 
-    assert target in ["cubin", "ptx"]
+    assert target in ["cubin", "ptx", "fatbin"]
 
     if cache_dir:
         checksum = _new_md5()
@@ -191,7 +191,7 @@ def compile(source, nvcc="nvcc", options=None, keep=False,
         no_extern_c=False, arch=None, code=None, cache_dir=None,
         include_dirs=[], target="cubin"):
 
-    assert target in ["cubin", "ptx"]
+    assert target in ["cubin", "ptx", "fatbin"]
 
     if not no_extern_c:
         source = 'extern "C" {\n%s\n}\n' % source
@@ -201,10 +201,11 @@ def compile(source, nvcc="nvcc", options=None, keep=False,
 
     options = options[:]
     if arch is None:
+        from pycuda.driver import Error
         try:
             from pycuda.driver import Context
             arch = "sm_%d%d" % Context.get_device().compute_capability()
-        except RuntimeError:
+        except Error:
             pass
 
     from pycuda.driver import CUDA_DEBUGGING