提交 6e9eb1f9 authored 作者: Alexander Matyasko's avatar Alexander Matyasko

Use cusolver first if matrix operation is available

上级 eaea1f87
...@@ -2124,12 +2124,12 @@ matrix_ops_db2 = LocalGroupDB(local_opt=theano.gof.opt.GraphToGPULocalOptGroup) ...@@ -2124,12 +2124,12 @@ matrix_ops_db2 = LocalGroupDB(local_opt=theano.gof.opt.GraphToGPULocalOptGroup)
matrix_ops_db2.__name__ = "matrix_ops_db2" matrix_ops_db2.__name__ = "matrix_ops_db2"
lifter = op_lifter([slinalg.Cholesky])(local_gpu_cholesky) lifter = op_lifter([slinalg.Cholesky])(local_gpu_cholesky)
matrix_ops_db.register("local_gpu_cholesky", lifter, matrix_ops_db.register("local_gpu_cholesky", lifter,
'gpuarray', 'fast_compile', 'fast_run', 'gpuarray', 'fast_compile', 'fast_run', 'cusolver',
position=1) position=0)
matrix_ops_db2.register("local_gpu_cholesky", matrix_ops_db2.register("local_gpu_cholesky",
local_optimizer([slinalg.Cholesky])(local_gpu_cholesky), local_optimizer([slinalg.Cholesky])(local_gpu_cholesky),
'gpuarray', 'fast_compile', 'fast_run', 'gpuarray', 'fast_compile', 'fast_run', 'cusolver',
position=1) position=0)
register_opt('fast_compile', name='matrix_ops_db')(matrix_ops_db) register_opt('fast_compile', name='matrix_ops_db')(matrix_ops_db)
register_opt2([slinalg.Solve], 'fast_compile', name='matrix_ops_db2')(matrix_ops_db2) register_opt2([slinalg.Solve], 'fast_compile', name='matrix_ops_db2')(matrix_ops_db2)
...@@ -2148,11 +2148,11 @@ def local_gpu_magma_cholesky(op, context_name, inputs, outputs): ...@@ -2148,11 +2148,11 @@ def local_gpu_magma_cholesky(op, context_name, inputs, outputs):
lifter = op_lifter([slinalg.Cholesky])(local_gpu_magma_cholesky) lifter = op_lifter([slinalg.Cholesky])(local_gpu_magma_cholesky)
matrix_ops_db.register("local_gpu_magma_cholesky", lifter, matrix_ops_db.register("local_gpu_magma_cholesky", lifter,
'gpuarray', 'fast_compile', 'fast_run', 'magma', 'gpuarray', 'fast_compile', 'fast_run', 'magma',
position=0) position=1)
matrix_ops_db2.register("local_gpu_magma_cholesky", matrix_ops_db2.register("local_gpu_magma_cholesky",
local_optimizer([slinalg.Cholesky])(local_gpu_magma_cholesky), local_optimizer([slinalg.Cholesky])(local_gpu_magma_cholesky),
'gpuarray', 'fast_compile', 'fast_run', 'magma', 'gpuarray', 'fast_compile', 'fast_run', 'magma',
position=0) position=1)
@register_inplace() @register_inplace()
......
...@@ -135,7 +135,7 @@ class TestGpuCholesky(unittest.TestCase): ...@@ -135,7 +135,7 @@ class TestGpuCholesky(unittest.TestCase):
cholesky_op = GpuCholesky(lower=lower, inplace=inplace) cholesky_op = GpuCholesky(lower=lower, inplace=inplace)
chol_A = cholesky_op(A) chol_A = cholesky_op(A)
return theano.function([A], chol_A, accept_inplace=inplace, return theano.function([A], chol_A, accept_inplace=inplace,
mode=mode_with_gpu.excluding('magma')) mode=mode_with_gpu)
def compare_gpu_cholesky_to_np(self, A_val, lower=True, inplace=False): def compare_gpu_cholesky_to_np(self, A_val, lower=True, inplace=False):
# Helper function to compare op output to np.cholesky output. # Helper function to compare op output to np.cholesky output.
...@@ -149,7 +149,7 @@ class TestGpuCholesky(unittest.TestCase): ...@@ -149,7 +149,7 @@ class TestGpuCholesky(unittest.TestCase):
def test_gpu_cholesky_opt(self): def test_gpu_cholesky_opt(self):
A = theano.tensor.matrix("A", dtype="float32") A = theano.tensor.matrix("A", dtype="float32")
fn = theano.function([A], cholesky(A), mode=mode_with_gpu.excluding('magma')) fn = theano.function([A], cholesky(A), mode=mode_with_gpu)
assert any([isinstance(node.op, GpuCholesky) assert any([isinstance(node.op, GpuCholesky)
for node in fn.maker.fgraph.toposort()]) for node in fn.maker.fgraph.toposort()])
...@@ -309,7 +309,7 @@ class TestMagma(unittest.TestCase): ...@@ -309,7 +309,7 @@ class TestMagma(unittest.TestCase):
def run_gpu_cholesky(self, A_val, lower=True): def run_gpu_cholesky(self, A_val, lower=True):
A = theano.tensor.fmatrix("A") A = theano.tensor.fmatrix("A")
f = theano.function([A], GpuMagmaCholesky(lower=lower)(A), f = theano.function([A], GpuMagmaCholesky(lower=lower)(A),
mode=mode_with_gpu) mode=mode_with_gpu.excluding('cusolver'))
return f(A_val) return f(A_val)
def check_cholesky(self, N, lower=True, rtol=None, atol=None): def check_cholesky(self, N, lower=True, rtol=None, atol=None):
...@@ -326,7 +326,7 @@ class TestMagma(unittest.TestCase): ...@@ -326,7 +326,7 @@ class TestMagma(unittest.TestCase):
def test_gpu_cholesky_opt(self): def test_gpu_cholesky_opt(self):
A = theano.tensor.matrix("A", dtype="float32") A = theano.tensor.matrix("A", dtype="float32")
fn = theano.function([A], cholesky(A), mode=mode_with_gpu) fn = theano.function([A], cholesky(A), mode=mode_with_gpu.excluding('cusolver'))
assert any([isinstance(node.op, GpuMagmaCholesky) assert any([isinstance(node.op, GpuMagmaCholesky)
for node in fn.maker.fgraph.toposort()]) for node in fn.maker.fgraph.toposort()])
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论