提交 2d3e9ef7 authored 作者: notoraptor's avatar notoraptor

Wrap op params for theano.gpuarray.linalg.GpuMagmaMatrixInverse.

上级 1b46af51
......@@ -454,7 +454,8 @@ class GpuMagmaMatrixInverse(COp):
"""Computes the inverse of a matrix :math:`A` using magma library.
"""
__props__ = ('inplace', )
params_type = gpu_context_type
check_input = False
params_type = ParamsType(inplace=bool_t, context=gpu_context_type)
def __init__(self, inplace=False):
COp.__init__(self, ['magma_inv.c'], 'APPLY_SPECIFIC(magma_inv)')
......@@ -492,13 +493,7 @@ class GpuMagmaMatrixInverse(COp):
return theano.Apply(self, [x], [x.type()])
def get_params(self, node):
return node.inputs[0].type.context
def get_op_params(self):
if self.inplace:
return [('INPLACE', '1')]
else:
return []
return self.params_type.get_params(self, context=node.inputs[0].type.context)
def infer_shape(self, node, shapes):
return shapes
......
......@@ -5,7 +5,7 @@ setup_ext_cuda();
#section support_code_struct
int APPLY_SPECIFIC(magma_inv)(PyGpuArrayObject *A, PyGpuArrayObject **A_inv,
PyGpuContextObject *c) {
PARAMS_TYPE* params) {
const size_t *dims;
magma_int_t N, ldwork, info;
magma_int_t *piv = NULL;
......@@ -19,7 +19,7 @@ int APPLY_SPECIFIC(magma_inv)(PyGpuArrayObject *A, PyGpuArrayObject **A_inv,
}
// This is early to match the exit() in the fail label.
cuda_enter(c->ctx);
cuda_enter(params->context->ctx);
magma_init();
if (!GpuArray_IS_C_CONTIGUOUS(&A->ga)) {
......@@ -38,25 +38,25 @@ int APPLY_SPECIFIC(magma_inv)(PyGpuArrayObject *A, PyGpuArrayObject **A_inv,
"GpuMagmaMatrixInverse: matrix is not square");
goto fail;
}
#ifdef INPLACE
Py_XDECREF(*A_inv);
*A_inv = A;
Py_INCREF(*A_inv);
#else
*A_inv = theano_try_copy(*A_inv, A);
if (*A_inv == NULL) {
PyErr_SetString(
PyExc_RuntimeError,
"GpuMagmaMatrixInverse: failed to allocate memory for the output");
goto fail;
if (params->inplace) {
Py_XDECREF(*A_inv);
*A_inv = A;
Py_INCREF(*A_inv);
} else {
*A_inv = theano_try_copy(*A_inv, A);
if (*A_inv == NULL) {
PyErr_SetString(
PyExc_RuntimeError,
"GpuMagmaMatrixInverse: failed to allocate memory for the output");
goto fail;
}
}
#endif
// magma matrix inverse
N = dims[0];
ldwork = N * magma_get_sgetri_nb(N);
dwork = gpudata_alloc(c->ctx, ldwork * sizeof(float), NULL, 0, NULL);
dwork = gpudata_alloc(params->context->ctx, ldwork * sizeof(float), NULL, 0, NULL);
if (dwork == NULL) {
PyErr_SetString(PyExc_RuntimeError,
"GpuMagmaMatrixInverse: failed to allocate working memory");
......@@ -94,6 +94,6 @@ fail:
if (dwork != NULL)
gpudata_release(dwork);
magma_finalize();
cuda_exit(c->ctx);
cuda_exit(params->context->ctx);
return res;
}
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论