Merge pull request #3792 from abergeron/gpuarray_prealloc

theano/configdefaults.py

@@ -4,7 +4,8 @@ import logging
 
 import theano
 from theano.configparser import (AddConfigVar, BoolParam, ConfigParam, EnumStr,
-                                 IntParam, StrParam, TheanoConfigParser)
+                                 FloatParam, IntParam, StrParam,
+                                 TheanoConfigParser)
 from theano.misc.cpucount import cpuCount
 from theano.misc.windows import call_subprocess_Popen
 
@@ -219,6 +220,14 @@ AddConfigVar('gpuarray.sync',
              BoolParam(False),
              in_c_key=True)
 
+AddConfigVar('gpuarray.preallocate',
+             """If 0 it doesn't do anything.  If between 0 and 1 it
+             will preallocate that fraction of the total GPU memory.
+             If 1 or greater it will preallocate that amount of memory
+             (in megabytes).""",
+             FloatParam(0, lambda i: i >= 0),
+             in_c_key=False)
+
 
 def safe_no_dnn_workmem(workmem):
     """

theano/sandbox/gpuarray/__init__.py

@@ -39,14 +39,26 @@ register_transfer(transfer)
 
 
 def init_dev(dev, name=None):
-    if pygpu.gpuarray.api_version() != (-10000, 0):
-        raise RuntimeError("Wrong API version for gpuarray:",
-                           pygpu.gpuarray.api_version(),
+    v = pygpu.gpuarray.api_version()
+    if v[0] != -10000:
+        raise RuntimeError("Wrong major API version for gpuarray:", v[0],
                            "Make sure Theano and libgpuarray/pygpu "
                            "are in sync.")
+    if v[1] < 0:
+        raise RuntimeError("Wrong minor API version for gpuarray:", v[1],
+                           "Please update libgpuarray/pygpu.")
     global pygpu_activated
     if dev not in init_dev.devmap:
-        init_dev.devmap[dev] = pygpu.init(dev)
+        ctx = pygpu.init(dev)
+        init_dev.devmap[dev] = ctx
+        if config.gpuarray.preallocate != 0:
+            if config.gpuarray.preallocate < 1:
+                gmem = config.gpuarray.preallocate * ctx.total_gmem
+            else:
+                gmem = config.gpuarray.preallocate * (1024*1024)
+            # This will allocate and immediatly free an object of size gmem
+            # which will reserve that amount of memory on the GPU.
+            pygpu.empty((gmem,), dtype='int8', context=ctx)
     context = init_dev.devmap[dev]
     # This will map the context name to the real context object.
     reg_context(name, context)

theano/sandbox/gpuarray/dnn_base.c

@@ -125,3 +125,7 @@ cudnnHandle_t APPLY_SPECIFIC(_handle);
   }
   cuda_exit(PARAMS->ctx);
 }
+
+#section cleanup_code_struct
+
+cudnnDestroy(APPLY_SPECIFIC(_handle));