Add magma to theano config (by default not enabled)

theano/configdefaults.py

@@ -361,6 +361,23 @@ AddConfigVar('dnn.enabled',
              EnumStr("auto", "True", "False", "no_check"),
              in_c_key=False)
 
+AddConfigVar('magma.include_path',
+             "Location of the magma header",
+             StrParam(None),
+             in_c_key=False)
+
+AddConfigVar('magma.library_path',
+             "Location of the magma library",
+             StrParam(None),
+             in_c_key=False)
+
+AddConfigVar('magma.enabled',
+             " If True, use magma for matrix computation."
+             " If False, disable magma"
+             " If no_check, assume present and the version between header and library match (so less compilation at context init)",
+             BoolParam(False),
+             in_c_key=False)
+
 # This flag determines whether or not to raise error/warning message if
 # there is a CPU Op in the computational graph.
 AddConfigVar(

theano/gpuarray/linalg.py

@@ -8,7 +8,7 @@ import pkg_resources
 from numpy.linalg.linalg import LinAlgError
 
 import theano
-from theano import Op
+from theano import Op, config
 from theano.gof import COp
 from theano.gpuarray import GpuArrayType
 
@@ -365,11 +365,19 @@ class GpuMagmaSVD(COp):
                 'gpuarray_helper.h', 'magma.h']
 
     def c_header_dirs(self):
-        return [os.path.dirname(__file__), pygpu.get_include()]
+        dirs = [os.path.dirname(__file__), pygpu.get_include()]
+        if config.magma.include_path:
+            dirs.append(config.magma.include_path)
+        return dirs
 
     def c_libraries(self):
         return ['magma']
 
+    def c_lib_dirs(self):
+        if config.magma.library_path:
+            return [config.magma.library_path]
+        return []
+
     def make_node(self, A):
         if A.ndim != 2:
             raise LinAlgError("Matrix rank error")
@@ -430,11 +438,19 @@ class GpuMagmaMatrixInverse(COp):
                 'gpuarray_helper.h', 'magma.h']
 
     def c_header_dirs(self):
-        return [os.path.dirname(__file__), pygpu.get_include()]
+        dirs = [os.path.dirname(__file__), pygpu.get_include()]
+        if config.magma.include_path:
+            dirs.append(config.magma.include_path)
+        return dirs
 
     def c_libraries(self):
         return ['magma']
 
+    def c_lib_dirs(self):
+        if config.magma.library_path:
+            return [config.magma.library_path]
+        return []
+
     def make_node(self, x):
         if x.ndim != 2:
             raise LinAlgError("Matrix rank error")

theano/gpuarray/opt.py

@@ -2010,6 +2010,8 @@ def local_inplace_cholesky(node):
 @op_lifter([nlinalg.MatrixInverse])
 @register_opt2([theano.tensor.nlinalg.MatrixInverse], 'magma', 'fast_compile')
 def local_gpu_matrix_inverse(op, context_name, inputs, outputs):
+    if not config.magma.enabled:
+        return
     return GpuMagmaMatrixInverse()
 
 
@@ -2017,7 +2019,8 @@ def local_gpu_matrix_inverse(op, context_name, inputs, outputs):
 @op_lifter([nlinalg.SVD])
 @register_opt2([theano.tensor.nlinalg.SVD], 'magma', 'fast_compile')
 def local_gpu_svd(op, context_name, inputs, outputs):
-    import pudb; pudb.set_trace()
+    if not config.magma.enabled:
+        return
     return GpuMagmaSVD(full_matrices=op.full_matrices,
                        compute_uv=op.compute_uv)