提交 0cb0d4ee authored 作者: Alexander Matyasko's avatar Alexander Matyasko

Add tests for float16 magma optimizations

上级 b1175d1e
...@@ -2183,7 +2183,8 @@ def local_gpu_magma_qr(op, context_name, inputs, outputs): ...@@ -2183,7 +2183,8 @@ def local_gpu_magma_qr(op, context_name, inputs, outputs):
return return
op = GpuMagmaQR(complete=True) op = GpuMagmaQR(complete=True)
if inputs[0].dtype == 'float16': if inputs[0].dtype == 'float16':
return op(inputs[0].astype('float32')).astype('float16') outputs = op(inputs[0].astype('float32'))
return [o.astype('float16') for o in outputs]
return op return op
......
...@@ -7,12 +7,14 @@ from numpy.linalg.linalg import LinAlgError ...@@ -7,12 +7,14 @@ from numpy.linalg.linalg import LinAlgError
import theano import theano
from theano import config from theano import config
from theano.gpuarray.linalg import (GpuCholesky, GpuMagmaMatrixInverse, from theano.gpuarray.linalg import (GpuCholesky, GpuMagmaCholesky,
GpuMagmaEigh, GpuMagmaMatrixInverse,
GpuMagmaQR, GpuMagmaSVD,
cusolver_available, gpu_matrix_inverse, cusolver_available, gpu_matrix_inverse,
gpu_solve, gpu_svd, GpuMagmaCholesky, gpu_solve, gpu_svd)
GpuMagmaQR, GpuMagmaEigh) from theano.tensor.nlinalg import (SVD, Eigh, MatrixInverse, QRFull,
from theano.tensor.nlinalg import matrix_inverse, qr, eigh QRIncomplete, eigh, matrix_inverse, qr)
from theano.tensor.slinalg import cholesky from theano.tensor.slinalg import Cholesky, cholesky
from theano.tests import unittest_tools as utt from theano.tests import unittest_tools as utt
from .. import gpuarray_shared_constructor from .. import gpuarray_shared_constructor
...@@ -216,6 +218,20 @@ class TestMagma(unittest.TestCase): ...@@ -216,6 +218,20 @@ class TestMagma(unittest.TestCase):
if not config.magma.enabled: if not config.magma.enabled:
self.skipTest('Magma is not enabled, skipping test') self.skipTest('Magma is not enabled, skipping test')
def test_magma_opt_float16(self):
ops_to_gpu = [(MatrixInverse(), GpuMagmaMatrixInverse),
(SVD(), GpuMagmaSVD),
(QRFull(mode='reduced'), GpuMagmaQR),
(QRIncomplete(mode='r'), GpuMagmaQR),
# TODO: add support for float16 to Eigh numpy
# (Eigh(), GpuMagmaEigh),
(Cholesky(), GpuMagmaCholesky)]
for op, gpu_op in ops_to_gpu:
A = theano.tensor.matrix("A", dtype="float16")
fn = theano.function([A], op(A), mode=mode_with_gpu.excluding('cusolver'))
assert any([isinstance(node.op, gpu_op)
for node in fn.maker.fgraph.toposort()])
def test_gpu_matrix_inverse(self): def test_gpu_matrix_inverse(self):
A = theano.tensor.fmatrix("A") A = theano.tensor.fmatrix("A")
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论