Merge pull request #643 from nouiz/pycuda

doc/library/config.txt

@@ -145,6 +145,20 @@ import theano and print the config variable, as in:
 
     This flag's value cannot be modified during the program execution.
 
+.. attribute:: config.pycuda.init
+
+    Bool value: either ``True`` or ``False``
+
+    Default: ``False``
+
+    If True, always initialize PyCUDA when Theano want to initialize
+    the GPU.  With PyCUDA version 2011.2.2 or earlier, PyCUDA must
+    initialize the GPU before Theano does it.  Setting
+    this flag to True, ensure that, but always import PyCUDA.  It can
+    be done manually by importing theano.misc.pycuda_init before
+    Theano initialize the GPU device.  Newer version of PyCUDA
+    (currently only in the trunk) don't have this restriction.
+
 .. attribute:: floatX
 
     String value: either 'float64' or 'float32'

theano/misc/pycuda_init.py

@@ -2,29 +2,28 @@ import os
 import warnings
 
 import theano
-import theano.sandbox.cuda as cuda
-cuda_ndarray = cuda.cuda_ndarray.cuda_ndarray
+import theano.sandbox.cuda
+from theano import config
+
 
 def set_gpu_from_theano():
     """
     This set the GPU used by PyCUDA to the same as the one used by Theano.
     """
-    #import pdb;pdb.set_trace()
-    if cuda.use.device_number is None:
-        cuda.use("gpu",
-                 force=False,
-                 default_to_move_computation_to_gpu=False,
-                 move_shared_float32_to_gpu=False,
-                 enable_cuda=True,
-                 test_driver=True)
-
-    assert cuda.use.device_number == cuda_ndarray.active_device_number()
+    # Transfer the theano gpu binding to pycuda, for consistency
+    if config.device.startswith("gpu") and len(config.device) > 3:
+        os.environ["CUDA_DEVICE"] = theano.config.device[3:]
+    elif (config.init_gpu_device.startswith("gpu") and
+          len(config.init_gpu_device) > 3):
+        os.environ["CUDA_DEVICE"] = theano.config.init_gpu_device[3:]
 
-#    os.environ["CUDA_DEVICE"] = str(cuda.use.device_number)
 
 set_gpu_from_theano()
 pycuda_available = False
-if False:
+# If theano.sandbox.cuda don't exist, it is because we are importing
+# it and it try to import this file! This mean we must init the device.
+if (not hasattr(theano.sandbox, 'cuda') or
+    theano.sandbox.cuda.use.device_number is None):
     try:
         import pycuda
         import pycuda.autoinit
@@ -33,3 +32,26 @@ if False:
         # presumably, the user wanted to use pycuda, else they wouldn't have
         # imported this module, so issue a warning that the import failed.
         warnings.warn("PyCUDA import failed in theano.misc.pycuda_init")
+    except pycuda._driver.LogicError:
+        if theano.config.force_device:
+            raise
+        else:
+            if "CUDA_DEVICE" in os.environ:
+                del os.environ["CUDA_DEVICE"]
+            import pycuda.autoinit
+            pycuda_available = True
+else:
+    import pycuda.driver
+    if hasattr(pycuda.driver.Context, "attach"):
+        pycuda.driver.Context.attach()
+    else:
+        # Now we always import this file when we call theano.sandbox.cuda.use
+        # So this should not happen normally.
+        # TODO: make this an error.
+        warnings.warn("For some unknow reason, theano.misc.pycuda_init was not"
+                      " imported before Theano initialized the GPU and"
+                      " your PyCUDA version is 2011.2.2 or earlier."
+                      " To fix the problem, import theano.misc.pycuda_init"
+                      " manually before using/initializing the GPU, use the"
+                      " Theano flag pycuda.init=True or use a"
+                      " more recent version of PyCUDA.")

theano/sandbox/cuda/__init__.py

@@ -11,7 +11,7 @@ import theano
 from theano.compile import optdb
 from theano.gof.cmodule import get_lib_extension
 from theano.gof.compilelock import get_lock, release_lock
-from theano.configparser import config, AddConfigVar, StrParam
+from theano.configparser import config, AddConfigVar, StrParam, BoolParam
 import nvcc_compiler
 
 _logger_name = 'theano.sandbox.cuda'
@@ -29,6 +29,16 @@ AddConfigVar('cuda.root',
         """,
         StrParam(os.getenv('CUDA_ROOT', "AUTO")))
 
+AddConfigVar('pycuda.init',
+        """If True, always initialize PyCUDA when Theano want to
+           initilize the GPU.  Currently, we must always initialize
+           PyCUDA before Theano do it.  Setting this flag to True,
+           ensure that, but always import PyCUDA.  It can be done
+           manually by importing theano.misc.pycuda_init before theano
+           initialize the GPU device.
+             """,
+        BoolParam(False))
+
 if config.cuda.root == "AUTO":
     # set nvcc_path correctly and get the version
     nvcc_compiler.set_cuda_root()
@@ -328,10 +338,15 @@ def use(device,
         # No successful call to use() has been made yet
         if device != 'gpu' and device < 0:
             return
-        if device in [None, ""]:
-            device = 0
+
+        # Has PyCUDA already initialized the GPU context
+        pycuda_init_dev = False
+        if config.pycuda.init:
+            import theano.misc.pycuda_init
+            pycuda_init_dev = theano.misc.pycuda_init.pycuda_available
+
         try:
-            if device != 'gpu':
+            if (device != 'gpu') and not pycuda_init_dev:
                 assert isinstance(device, int)
                 gpu_init(device)
                 use.device_number = device