def cholesky(a): """Cholesky decomposition. Decompose a given two-dimensional square matrix into ``L * L.T``, where ``L`` is a lower-triangular matrix and ``.T`` is a conjugate transpose operator. Args: a (cupy.ndarray): The input matrix with dimension ``(N, N)`` Returns: cupy.ndarray: The lower-triangular matrix. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.cholesky` """ # TODO(Saito): Current implementation only accepts two-dimensional arrays util._assert_cupy_array(a) util._assert_rank2(a) util._assert_nd_squareness(a) if a.dtype.char == 'f' or a.dtype.char == 'd': dtype = a.dtype.char else: dtype = numpy.promote_types(a.dtype.char, 'f').char x = a.astype(dtype, order='C', copy=True) n = len(a) handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=numpy.int32) if dtype == 'f': potrf = cusolver.spotrf potrf_bufferSize = cusolver.spotrf_bufferSize elif dtype == 'd': potrf = cusolver.dpotrf potrf_bufferSize = cusolver.dpotrf_bufferSize elif dtype == 'F': potrf = cusolver.cpotrf potrf_bufferSize = cusolver.cpotrf_bufferSize else: # dtype == 'D': potrf = cusolver.zpotrf potrf_bufferSize = cusolver.zpotrf_bufferSize buffersize = potrf_bufferSize(handle, cublas.CUBLAS_FILL_MODE_UPPER, n, x.data.ptr, n) workspace = cupy.empty(buffersize, dtype=dtype) potrf(handle, cublas.CUBLAS_FILL_MODE_UPPER, n, x.data.ptr, n, workspace.data.ptr, buffersize, dev_info.data.ptr) cupy.linalg.util._check_cusolver_dev_info_if_synchronization_allowed( potrf, dev_info) util._tril(x, k=0) return x
def _batched_inv(a): assert (a.ndim >= 3) util._assert_cupy_array(a) util._assert_nd_squareness(a) if a.dtype == cupy.float32: getrf = cupy.cuda.cublas.sgetrfBatched getri = cupy.cuda.cublas.sgetriBatched elif a.dtype == cupy.float64: getrf = cupy.cuda.cublas.dgetrfBatched getri = cupy.cuda.cublas.dgetriBatched elif a.dtype == cupy.complex64: getrf = cupy.cuda.cublas.cgetrfBatched getri = cupy.cuda.cublas.cgetriBatched elif a.dtype == cupy.complex128: getrf = cupy.cuda.cublas.zgetrfBatched getri = cupy.cuda.cublas.zgetriBatched else: msg = ('dtype must be float32, float64, complex64 or complex128' ' (actual: {})'.format(a.dtype)) raise ValueError(msg) if 0 in a.shape: return cupy.empty_like(a) a_shape = a.shape # copy is necessary to present `a` to be overwritten. a = a.copy().reshape(-1, a_shape[-2], a_shape[-1]) handle = device.get_cublas_handle() batch_size = a.shape[0] n = a.shape[1] lda = n step = n * lda * a.itemsize start = a.data.ptr stop = start + step * batch_size a_array = cupy.arange(start, stop, step, dtype=cupy.uintp) pivot_array = cupy.empty((batch_size, n), dtype=cupy.int32) info_array = cupy.empty((batch_size, ), dtype=cupy.int32) getrf(handle, n, a_array.data.ptr, lda, pivot_array.data.ptr, info_array.data.ptr, batch_size) cupy.linalg.util._check_cublas_info_array_if_synchronization_allowed( getrf, info_array) c = cupy.empty_like(a) ldc = lda step = n * ldc * c.itemsize start = c.data.ptr stop = start + step * batch_size c_array = cupy.arange(start, stop, step, dtype=cupy.uintp) getri(handle, n, a_array.data.ptr, lda, pivot_array.data.ptr, c_array.data.ptr, ldc, info_array.data.ptr, batch_size) cupy.linalg.util._check_cublas_info_array_if_synchronization_allowed( getri, info_array) return c.reshape(a_shape)
def lsqr(A, b): """Solves linear system with QR decomposition. Find the solution to a large, sparse, linear system of equations. The function solves ``Ax = b``. Given two-dimensional matrix ``A`` is decomposed into ``Q * R``. Args: A (cupy.ndarray or cupy.sparse.csr_matrix): The input matrix with dimension ``(N, N)`` b (cupy.ndarray): Right-hand side vector. Returns: ret (tuple): Its length must be ten. It has same type elements as SciPy. Only the first element, the solution vector ``x``, is available and other elements are expressed as ``None`` because the implementation of cuSOLVER is different from the one of SciPy. You can easily calculate the fourth element by ``norm(b - Ax)`` and the ninth element by ``norm(x)``. .. seealso:: :func:`scipy.sparse.linalg.lsqr` """ if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') if not cupy.sparse.isspmatrix_csr(A): A = cupy.sparse.csr_matrix(A) util._assert_nd_squareness(A) util._assert_cupy_array(b) m = A.shape[0] if b.ndim != 1 or len(b) != m: raise ValueError('b must be 1-d array whose size is same as A') # Cast to float32 or float64 if A.dtype == 'f' or A.dtype == 'd': dtype = A.dtype else: dtype = numpy.find_common_type((A.dtype, 'f'), ()) handle = device.get_cusolver_sp_handle() nnz = A.nnz tol = 1.0 reorder = 1 x = cupy.empty(m, dtype=dtype) singularity = numpy.empty(1, numpy.int32) if dtype == 'f': csrlsvqr = cusolver.scsrlsvqr else: csrlsvqr = cusolver.dcsrlsvqr csrlsvqr(handle, m, nnz, A._descr.descriptor, A.data.data.ptr, A.indptr.data.ptr, A.indices.data.ptr, b.data.ptr, tol, reorder, x.data.ptr, singularity.ctypes.data) # The return type of SciPy is always float64. Therefore, x must be casted. x = x.astype(numpy.float64) ret = (x, None, None, None, None, None, None, None, None, None) return ret
def inv(a): """Computes the inverse of a matrix. This function computes matrix ``a_inv`` from n-dimensional regular matrix ``a`` such that ``dot(a, a_inv) == eye(n)``. Args: a (cupy.ndarray): The regular matrix Returns: cupy.ndarray: The inverse of a matrix. .. seealso:: :func:`numpy.linalg.inv` """ if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') # to prevent `a` to be overwritten a = a.copy() util._assert_cupy_array(a) util._assert_rank2(a) util._assert_nd_squareness(a) if a.dtype.char == 'f' or a.dtype.char == 'd': dtype = a.dtype.char else: dtype = numpy.find_common_type((a.dtype.char, 'f'), ()).char cusolver_handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=dtype) ipiv = cupy.empty((a.shape[0], 1), dtype=dtype) if dtype == 'f': getrf = cusolver.sgetrf getrf_bufferSize = cusolver.sgetrf_bufferSize getrs = cusolver.sgetrs else: # dtype == 'd' getrf = cusolver.dgetrf getrf_bufferSize = cusolver.dgetrf_bufferSize getrs = cusolver.dgetrs m = a.shape[0] buffersize = getrf_bufferSize(cusolver_handle, m, m, a.data.ptr, m) workspace = cupy.empty(buffersize, dtype=dtype) # LU factorization getrf(cusolver_handle, m, m, a.data.ptr, m, workspace.data.ptr, ipiv.data.ptr, dev_info.data.ptr) b = cupy.eye(m, dtype=dtype) # solve for the inverse getrs(cusolver_handle, 0, m, m, a.data.ptr, m, ipiv.data.ptr, b.data.ptr, m, dev_info.data.ptr) return b
def cholesky(a): """Cholesky decomposition. Decompose a given two-dimensional square matrix into ``L * L.T``, where ``L`` is a lower-triangular matrix and ``.T`` is a conjugate transpose operator. Note that in the current implementation ``a`` must be a real matrix, and only float32 and float64 are supported. Args: a (cupy.ndarray): The input matrix with dimension ``(N, N)`` Returns: cupy.ndarray: The lower-triangular matrix. .. seealso:: :func:`numpy.linalg.cholesky` """ if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') # TODO(Saito): Current implementation only accepts two-dimensional arrays util._assert_cupy_array(a) util._assert_rank2(a) util._assert_nd_squareness(a) # Cast to float32 or float64 if a.dtype.char == 'f' or a.dtype.char == 'd': dtype = a.dtype.char else: dtype = numpy.find_common_type((a.dtype.char, 'f'), ()).char x = a.astype(dtype, order='C', copy=True) n = len(a) handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=numpy.int32) if dtype == 'f': buffersize = cusolver.spotrf_bufferSize( handle, cublas.CUBLAS_FILL_MODE_UPPER, n, x.data.ptr, n) workspace = cupy.empty(buffersize, dtype=numpy.float32) cusolver.spotrf( handle, cublas.CUBLAS_FILL_MODE_UPPER, n, x.data.ptr, n, workspace.data.ptr, buffersize, dev_info.data.ptr) else: # dtype == 'd' buffersize = cusolver.dpotrf_bufferSize( handle, cublas.CUBLAS_FILL_MODE_UPPER, n, x.data.ptr, n) workspace = cupy.empty(buffersize, dtype=numpy.float64) cusolver.dpotrf( handle, cublas.CUBLAS_FILL_MODE_UPPER, n, x.data.ptr, n, workspace.data.ptr, buffersize, dev_info.data.ptr) status = int(dev_info[0]) if status > 0: raise linalg.LinAlgError( 'The leading minor of order {} ' 'is not positive definite'.format(status)) elif status < 0: raise linalg.LinAlgError( 'Parameter error (maybe caused by a bug in cupy.linalg?)') util._tril(x, k=0) return x
def solve(a, b): """Solves a linear matrix equation. It computes the exact solution of ``x`` in ``ax = b``, where ``a`` is a square and full rank matrix. Args: a (cupy.ndarray): The matrix with dimension ``(..., M, M)``. b (cupy.ndarray): The matrix with dimension ``(...,M)`` or ``(..., M, K)``. Returns: cupy.ndarray: The matrix with dimension ``(..., M)`` or ``(..., M, K)``. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.solve` """ # NOTE: Since cusolver in CUDA 8.0 does not support gesv, # we manually solve a linear system with QR decomposition. # For details, please see the following: # https://docs.nvidia.com/cuda/cusolver/index.html#qr_examples util._assert_cupy_array(a, b) util._assert_nd_squareness(a) if not ((a.ndim == b.ndim or a.ndim == b.ndim + 1) and a.shape[:-1] == b.shape[:a.ndim - 1]): raise ValueError( 'a must have (..., M, M) shape and b must have (..., M) ' 'or (..., M, K)') # Cast to float32 or float64 if a.dtype.char == 'f' or a.dtype.char == 'd': dtype = a.dtype else: dtype = numpy.find_common_type((a.dtype.char, 'f'), ()) cublas_handle = device.get_cublas_handle() cusolver_handle = device.get_cusolver_handle() a = a.astype(dtype) b = b.astype(dtype) if a.ndim == 2: return _solve(a, b, cublas_handle, cusolver_handle) x = cupy.empty_like(b) shape = a.shape[:-2] for i in six.moves.range(numpy.prod(shape)): index = numpy.unravel_index(i, shape) x[index] = _solve(a[index], b[index], cublas_handle, cusolver_handle) return x
def _slogdet_one(a): util._assert_rank2(a) util._assert_nd_squareness(a) dtype = a.dtype handle = device.get_cusolver_handle() m = len(a) ipiv = cupy.empty(m, 'i') info = cupy.empty((), 'i') # Need to make a copy because getrf works inplace a_copy = a.copy(order='F') if dtype == 'f': getrf_bufferSize = cusolver.sgetrf_bufferSize getrf = cusolver.sgetrf #<-- MODIFIED elif dtype == 'd': getrf_bufferSize = cusolver.dgetrf_bufferSize getrf = cusolver.dgetrf elif dtype == 'F': getrf_bufferSize = cusolver.cgetrf_bufferSize getrf = cusolver.cgetrf else: getrf_bufferSize = cusolver.zgetrf_bufferSize getrf = cusolver.zgetrf #<-- MODIFIED buffersize = getrf_bufferSize(handle, m, m, a_copy.data.ptr, m) workspace = cupy.empty(buffersize, dtype=dtype) getrf(handle, m, m, a_copy.data.ptr, m, workspace.data.ptr, ipiv.data.ptr, info.data.ptr) if info[()] == 0: diag = cupy.diag(a_copy) # ipiv is 1-origin non_zero = (cupy.count_nonzero(ipiv != cupy.arange(1, m + 1)) + cupy.count_nonzero(diag < 0)) # Note: sign == -1 ** (non_zero % 2) sign = (non_zero % 2) * -2 + 1 logdet = cupy.log(abs(diag)).sum() else: sign = cupy.array(0.0, dtype=dtype) #ORIGINAL # logdet = cupy.array(float('-inf'), dtype) #<-- MODIFIED if dtype in ['f', 'd']: logdet = cupy.array(float('-inf'), dtype) elif dtype == 'F': logdet = cupy.array(float('-inf'), cupy.float32) else: logdet = cupy.array(float('-inf'), cupy.float64) #<-- MODIFIED return sign, logdet
def solve(a, b): """Solves a linear matrix equation. It computes the exact solution of ``x`` in ``ax = b``, where ``a`` is a square and full rank matrix. Args: a (cupy.ndarray): The matrix with dimension ``(..., M, M)``. b (cupy.ndarray): The matrix with dimension ``(...,M)`` or ``(..., M, K)``. Returns: cupy.ndarray: The matrix with dimension ``(..., M)`` or ``(..., M, K)``. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.solve` """ if a.ndim > 2 and a.shape[-1] <= get_batched_gesv_limit(): # Note: There is a low performance issue in batched_gesv when matrix is # large, so it is not used in such cases. return batched_gesv(a, b) util._assert_cupy_array(a, b) util._assert_nd_squareness(a) if not ((a.ndim == b.ndim or a.ndim == b.ndim + 1) and a.shape[:-1] == b.shape[:a.ndim - 1]): raise ValueError( 'a must have (..., M, M) shape and b must have (..., M) ' 'or (..., M, K)') # Cast to float32 or float64 if a.dtype.char == 'f' or a.dtype.char == 'd': dtype = a.dtype else: dtype = numpy.promote_types(a.dtype.char, 'f') a = a.astype(dtype) b = b.astype(dtype) if a.ndim == 2: return cupy.cusolver.gesv(a, b) x = cupy.empty_like(b) shape = a.shape[:-2] for i in range(numpy.prod(shape)): index = numpy.unravel_index(i, shape) x[index] = cupy.cusolver.gesv(a[index], b[index]) return x
def lschol(A, b): """Solves linear system with cholesky decomposition. Find the solution to a large, sparse, linear system of equations. The function solves ``Ax = b``. Given two-dimensional matrix ``A`` is decomposed into ``L * L^*``. Args: A (cupy.ndarray or cupy.sparse.csr_matrix): The input matrix with dimension ``(N, N)``. Must be positive-definite input matrix. Only symmetric real matrix is supported currently. b (cupy.ndarray): Right-hand side vector. Returns: ret (cupy.ndarray): The solution vector ``x``. """ if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') if not cupy.sparse.isspmatrix_csr(A): A = cupy.sparse.csr_matrix(A) util._assert_nd_squareness(A) util._assert_cupy_array(b) m = A.shape[0] if b.ndim != 1 or len(b) != m: raise ValueError('b must be 1-d array whose size is same as A') # Cast to float32 or float64 if A.dtype == 'f' or A.dtype == 'd': dtype = A.dtype else: dtype = numpy.find_common_type((A.dtype, 'f'), ()) handle = device.get_cusolver_sp_handle() nnz = A.nnz tol = 1.0 reorder = 1 x = cupy.empty(m, dtype=dtype) singularity = numpy.empty(1, numpy.int32) if dtype == 'f': csrlsvchol = cusolver.scsrlsvchol else: csrlsvchol = cusolver.dcsrlsvchol csrlsvchol(handle, m, nnz, A._descr.descriptor, A.data.data.ptr, A.indptr.data.ptr, A.indices.data.ptr, b.data.ptr, tol, reorder, x.data.ptr, singularity.ctypes.data) # The return type of SciPy is always float64. x = x.astype(numpy.float64) return x
def solve(a, b): """Solves a linear matrix equation. It computes the exact solution of ``x`` in ``ax = b``, where ``a`` is a square and full rank matrix. Args: a (cupy.ndarray): The matrix with dimension ``(..., M, M)``. b (cupy.ndarray): The matrix with dimension ``(...,M)`` or ``(..., M, K)``. Returns: cupy.ndarray: The matrix with dimension ``(..., M)`` or ``(..., M, K)``. .. seealso:: :func:`numpy.linalg.solve` """ # NOTE: Since cusolver in CUDA 8.0 does not support gesv, # we manually solve a linear system with QR decomposition. # For details, please see the following: # https://docs.nvidia.com/cuda/cusolver/index.html#qr_examples if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') util._assert_cupy_array(a, b) util._assert_nd_squareness(a) if not ((a.ndim == b.ndim or a.ndim == b.ndim + 1) and a.shape[:-1] == b.shape[:a.ndim - 1]): raise ValueError( 'a must have (..., M, M) shape and b must have (..., M) ' 'or (..., M, K)') # Cast to float32 or float64 if a.dtype.char == 'f' or a.dtype.char == 'd': dtype = a.dtype else: dtype = numpy.find_common_type((a.dtype.char, 'f'), ()) cublas_handle = device.get_cublas_handle() cusolver_handle = device.get_cusolver_handle() a = a.astype(dtype) b = b.astype(dtype) if a.ndim == 2: return _solve(a, b, cublas_handle, cusolver_handle) x = cupy.empty_like(b) shape = a.shape[:-2] for i in six.moves.range(numpy.prod(shape)): index = numpy.unravel_index(i, shape) x[index] = _solve(a[index], b[index], cublas_handle, cusolver_handle) return x
def _slogdet_one(a): util._assert_rank2(a) util._assert_nd_squareness(a) dtype = a.dtype handle = device.get_cusolver_handle() m = len(a) ipiv = cupy.empty(m, dtype=numpy.int32) dev_info = cupy.empty((), dtype=numpy.int32) # Need to make a copy because getrf works inplace a_copy = a.copy(order='F') if dtype == 'f': getrf_bufferSize = cusolver.sgetrf_bufferSize getrf = cusolver.sgetrf else: getrf_bufferSize = cusolver.dgetrf_bufferSize getrf = cusolver.dgetrf buffersize = getrf_bufferSize(handle, m, m, a_copy.data.ptr, m) workspace = cupy.empty(buffersize, dtype=dtype) getrf(handle, m, m, a_copy.data.ptr, m, workspace.data.ptr, ipiv.data.ptr, dev_info.data.ptr) # dev_info < 0 means illegal value (in dimensions, strides, and etc.) that # should never happen even if the matrix contains nan or inf. # TODO(kataoka): assert dev_info >= 0 if synchronization is allowed for # debugging purposes. diag = cupy.diag(a_copy) # ipiv is 1-origin non_zero = (cupy.count_nonzero(ipiv != cupy.arange(1, m + 1)) + cupy.count_nonzero(diag < 0)) # Note: sign == -1 ** (non_zero % 2) sign = (non_zero % 2) * -2 + 1 logdet = cupy.log(abs(diag)).sum() singular = dev_info > 0 return ( cupy.where(singular, dtype.type(0), sign), cupy.where(singular, dtype.type('-inf'), logdet), )
def _slogdet_one(a): util._assert_rank2(a) util._assert_nd_squareness(a) dtype = a.dtype handle = device.get_cusolver_handle() m = len(a) ipiv = cupy.empty(m, dtype=numpy.int32) dev_info = cupy.empty(1, dtype=numpy.int32) # Need to make a copy because getrf works inplace a_copy = a.copy(order='F') if dtype == 'f': getrf_bufferSize = cusolver.sgetrf_bufferSize getrf = cusolver.sgetrf else: getrf_bufferSize = cusolver.dgetrf_bufferSize getrf = cusolver.dgetrf buffersize = getrf_bufferSize(handle, m, m, a_copy.data.ptr, m) workspace = cupy.empty(buffersize, dtype=dtype) getrf(handle, m, m, a_copy.data.ptr, m, workspace.data.ptr, ipiv.data.ptr, dev_info.data.ptr) try: cupy.linalg.util._check_cusolver_dev_info_if_synchronization_allowed( getrf, dev_info) diag = cupy.diag(a_copy) # ipiv is 1-origin non_zero = (cupy.count_nonzero(ipiv != cupy.arange(1, m + 1)) + cupy.count_nonzero(diag < 0)) # Note: sign == -1 ** (non_zero % 2) sign = (non_zero % 2) * -2 + 1 logdet = cupy.log(abs(diag)).sum() except linalg.LinAlgError: sign = cupy.array(0.0, dtype=dtype) logdet = cupy.array(float('-inf'), dtype) return sign, logdet
def inv(a): '''Computes the inverse of a matrix. This function computes matrix ``a_inv`` from n-dimensional regular matrix ``a`` such that ``dot(a, a_inv) == eye(n)``. Args: a (cupy.ndarray): The regular matrix Returns: cupy.ndarray: The inverse of a matrix. .. seealso:: :func:`numpy.linalg.inv` ''' if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') util._assert_cupy_array(a) util._assert_rank2(a) util._assert_nd_squareness(a) b = cupy.eye(len(a), dtype=a.dtype) return solve(a, b)
def inv(a): """Computes the inverse of a matrix. This function computes matrix ``a_inv`` from n-dimensional regular matrix ``a`` such that ``dot(a, a_inv) == eye(n)``. Args: a (cupy.ndarray): The regular matrix Returns: cupy.ndarray: The inverse of a matrix. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.inv` """ if a.ndim >= 3: return _batched_inv(a) # to prevent `a` to be overwritten a = a.copy() util._assert_cupy_array(a) util._assert_rank2(a) util._assert_nd_squareness(a) # support float32, float64, complex64, and complex128 if a.dtype.char in 'fdFD': dtype = a.dtype.char else: dtype = numpy.find_common_type((a.dtype.char, 'f'), ()).char cusolver_handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=numpy.int32) ipiv = cupy.empty((a.shape[0], 1), dtype=numpy.intc) if dtype == 'f': getrf = cusolver.sgetrf getrf_bufferSize = cusolver.sgetrf_bufferSize getrs = cusolver.sgetrs elif dtype == 'd': getrf = cusolver.dgetrf getrf_bufferSize = cusolver.dgetrf_bufferSize getrs = cusolver.dgetrs elif dtype == 'F': getrf = cusolver.cgetrf getrf_bufferSize = cusolver.cgetrf_bufferSize getrs = cusolver.cgetrs elif dtype == 'D': getrf = cusolver.zgetrf getrf_bufferSize = cusolver.zgetrf_bufferSize getrs = cusolver.zgetrs else: msg = ('dtype must be float32, float64, complex64 or complex128' ' (actual: {})'.format(a.dtype)) raise ValueError(msg) m = a.shape[0] buffersize = getrf_bufferSize(cusolver_handle, m, m, a.data.ptr, m) workspace = cupy.empty(buffersize, dtype=dtype) # LU factorization getrf(cusolver_handle, m, m, a.data.ptr, m, workspace.data.ptr, ipiv.data.ptr, dev_info.data.ptr) cupy.linalg.util._check_cusolver_dev_info_if_synchronization_allowed( getrf, dev_info) b = cupy.eye(m, dtype=dtype) # solve for the inverse getrs(cusolver_handle, 0, m, m, a.data.ptr, m, ipiv.data.ptr, b.data.ptr, m, dev_info.data.ptr) cupy.linalg.util._check_cusolver_dev_info_if_synchronization_allowed( getrs, dev_info) return b
def invh(a): """Compute the inverse of a Hermitian matrix. This function computes a inverse of a real symmetric or complex hermitian positive-definite matrix using Cholesky factorization. If matrix ``a`` is not positive definite, Cholesky factorization fails and it raises an error. Args: a (cupy.ndarray): Real symmetric or complex hermitian maxtix. Returns: cupy.ndarray: The inverse of matrix ``a``. """ # to prevent `a` from being overwritten a = a.copy() util._assert_cupy_array(a) util._assert_rank2(a) util._assert_nd_squareness(a) # support float32, float64, complex64, and complex128 if a.dtype.char in 'fdFD': dtype = a.dtype.char else: dtype = numpy.promote_types(a.dtype.char, 'f').char cusolver_handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=numpy.int32) if dtype == 'f': potrf = cusolver.spotrf potrf_bufferSize = cusolver.spotrf_bufferSize potrs = cusolver.spotrs elif dtype == 'd': potrf = cusolver.dpotrf potrf_bufferSize = cusolver.dpotrf_bufferSize potrs = cusolver.dpotrs elif dtype == 'F': potrf = cusolver.cpotrf potrf_bufferSize = cusolver.cpotrf_bufferSize potrs = cusolver.cpotrs elif dtype == 'D': potrf = cusolver.zpotrf potrf_bufferSize = cusolver.zpotrf_bufferSize potrs = cusolver.zpotrs else: msg = ('dtype must be float32, float64, complex64 or complex128' ' (actual: {})'.format(a.dtype)) raise ValueError(msg) m = a.shape[0] uplo = cublas.CUBLAS_FILL_MODE_LOWER worksize = potrf_bufferSize(cusolver_handle, uplo, m, a.data.ptr, m) workspace = cupy.empty(worksize, dtype=dtype) # Cholesky factorization potrf(cusolver_handle, uplo, m, a.data.ptr, m, workspace.data.ptr, worksize, dev_info.data.ptr) info = dev_info[0] if info != 0: if info < 0: msg = '\tThe {}-th parameter is wrong'.format(-info) else: msg = ('\tThe leading minor of order {} is not positive definite' .format(info)) raise RuntimeError('matrix inversion failed at potrf.\n' + msg) b = cupy.eye(m, dtype=dtype) # Solve: A * X = B potrs(cusolver_handle, uplo, m, m, a.data.ptr, m, b.data.ptr, m, dev_info.data.ptr) info = dev_info[0] if info > 0: assert False, ('Unexpected output returned by potrs (actual: {})' .format(info)) elif info < 0: raise RuntimeError('matrix inversion failed at potrs.\n' '\tThe {}-th parameter is wrong'.format(-info)) return b
def _batched_inv(a): assert (a.ndim >= 3) util._assert_cupy_array(a) util._assert_nd_squareness(a) if a.dtype == cupy.float32: getrf = cupy.cuda.cublas.sgetrfBatched getri = cupy.cuda.cublas.sgetriBatched elif a.dtype == cupy.float64: getrf = cupy.cuda.cublas.dgetrfBatched getri = cupy.cuda.cublas.dgetriBatched elif a.dtype == cupy.complex64: getrf = cupy.cuda.cublas.cgetrfBatched getri = cupy.cuda.cublas.cgetriBatched elif a.dtype == cupy.complex128: getrf = cupy.cuda.cublas.zgetrfBatched getri = cupy.cuda.cublas.zgetriBatched else: msg = ('dtype must be float32, float64, complex64 or complex128' ' (actual: {})'.format(a.dtype)) raise ValueError(msg) if 0 in a.shape: return cupy.empty_like(a) a_shape = a.shape # copy is necessary to present `a` to be overwritten. a = a.copy().reshape(-1, a_shape[-2], a_shape[-1]) handle = device.get_cublas_handle() batch_size = a.shape[0] n = a.shape[1] lda = n step = n * lda * a.itemsize start = a.data.ptr stop = start + step * batch_size a_array = cupy.arange(start, stop, step, dtype=cupy.uintp) pivot_array = cupy.empty((batch_size, n), dtype=cupy.int32) info_array = cupy.empty((batch_size, ), dtype=cupy.int32) getrf(handle, n, a_array.data.ptr, lda, pivot_array.data.ptr, info_array.data.ptr, batch_size) err = False err_detail = '' for i in range(batch_size): info = info_array[i] if info < 0: err = True err_detail += ('\tmatrix[{}]: illegal value at {}-the parameter.' '\n'.format(i, -info)) if info > 0: err = True err_detail += '\tmatrix[{}]: matrix is singular.\n'.format(i) if err: raise RuntimeError('matrix inversion failed at getrf.\n' + err_detail) c = cupy.empty_like(a) ldc = lda step = n * ldc * c.itemsize start = c.data.ptr stop = start + step * batch_size c_array = cupy.arange(start, stop, step, dtype=cupy.uintp) getri(handle, n, a_array.data.ptr, lda, pivot_array.data.ptr, c_array.data.ptr, ldc, info_array.data.ptr, batch_size) for i in range(batch_size): info = info_array[i] if info > 0: err = True err_detail += '\tmatrix[{}]: matrix is singular.\n'.format(i) if err: raise RuntimeError('matrix inversion failed at getri.\n' + err_detail) return c.reshape(a_shape)
def solve(a, b): '''Solves a linear matrix equation. It computes the exact solution of ``x`` in ``ax = b``, where ``a`` is a square and full rank matrix. Args: a (cupy.ndarray): The matrix with dimension ``(M, M)`` b (cupy.ndarray): The vector with ``M`` elements, or the matrix with dimension ``(M, K)`` Returns: cupy.ndarray: The vector with ``M`` elements, or the matrix with dimension ``(M, K)``. .. seealso:: :func:`numpy.linalg.solve` ''' # NOTE: Since cusolver in CUDA 8.0 does not support gesv, # we manually solve a linear system with QR decomposition. # For details, please see the following: # https://docs.nvidia.com/cuda/cusolver/index.html#qr_examples if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') # TODO(Saito): Current implementation only accepts two-dimensional arrays util._assert_cupy_array(a, b) util._assert_rank2(a) util._assert_nd_squareness(a) if 2 < b.ndim: raise linalg.LinAlgError('{}-dimensional array given. Array must be ' 'one or two-dimensional'.format(b.ndim)) if len(a) != len(b): raise linalg.LinAlgError('The number of rows of array a must be ' 'the same as that of array b') # Cast to float32 or float64 if a.dtype.char == 'f' or a.dtype.char == 'd': dtype = a.dtype.char else: dtype = numpy.find_common_type((a.dtype.char, 'f'), ()).char m, k = (b.size, 1) if b.ndim == 1 else b.shape a = a.transpose().astype(dtype, order='C', copy=True) b = b.transpose().astype(dtype, order='C', copy=True) cusolver_handle = device.get_cusolver_handle() cublas_handle = device.get_cublas_handle() dev_info = cupy.empty(1, dtype=numpy.int32) if dtype == 'f': geqrf = cusolver.sgeqrf geqrf_bufferSize = cusolver.sgeqrf_bufferSize ormqr = cusolver.sormqr trsm = cublas.strsm else: # dtype == 'd' geqrf = cusolver.dgeqrf geqrf_bufferSize = cusolver.dgeqrf_bufferSize ormqr = cusolver.dormqr trsm = cublas.dtrsm # 1. QR decomposition (A = Q * R) buffersize = geqrf_bufferSize(cusolver_handle, m, m, a.data.ptr, m) workspace = cupy.empty(buffersize, dtype=dtype) tau = cupy.empty(m, dtype=dtype) geqrf(cusolver_handle, m, m, a.data.ptr, m, tau.data.ptr, workspace.data.ptr, buffersize, dev_info.data.ptr) _check_status(dev_info) # 2. ormqr (Q^T * B) ormqr(cusolver_handle, cublas.CUBLAS_SIDE_LEFT, cublas.CUBLAS_OP_T, m, k, m, a.data.ptr, m, tau.data.ptr, b.data.ptr, m, workspace.data.ptr, buffersize, dev_info.data.ptr) _check_status(dev_info) # 3. trsm (X = R^{-1} * (Q^T * B)) trsm(cublas_handle, cublas.CUBLAS_SIDE_LEFT, cublas.CUBLAS_FILL_MODE_UPPER, cublas.CUBLAS_OP_N, cublas.CUBLAS_DIAG_NON_UNIT, m, k, 1, a.data.ptr, m, b.data.ptr, m) return b.transpose()
def batched_gesv(a, b): """Solves multiple linear matrix equations using cublas<t>getr[fs]Batched(). Computes the solution to system of linear equation ``ax = b``. Args: a (cupy.ndarray): The matrix with dimension ``(..., M, M)``. b (cupy.ndarray): The matrix with dimension ``(..., M)`` or ``(..., M, K)``. Returns: cupy.ndarray: The matrix with dimension ``(..., M)`` or ``(..., M, K)``. """ util._assert_cupy_array(a, b) util._assert_nd_squareness(a) if not ((a.ndim == b.ndim or a.ndim == b.ndim + 1) and a.shape[:-1] == b.shape[:a.ndim - 1]): raise ValueError( 'a must have (..., M, M) shape and b must have (..., M) ' 'or (..., M, K)') dtype = numpy.promote_types(a.dtype.char, 'f') if dtype == 'f': t = 's' elif dtype == 'd': t = 'd' elif dtype == 'F': t = 'c' elif dtype == 'D': t = 'z' else: raise TypeError('invalid dtype') getrf = getattr(cublas, t + 'getrfBatched') getrs = getattr(cublas, t + 'getrsBatched') bs = numpy.prod(a.shape[:-2]) if a.ndim > 2 else 1 n = a.shape[-1] nrhs = b.shape[-1] if a.ndim == b.ndim else 1 b_shape = b.shape a_data_ptr = a.data.ptr b_data_ptr = b.data.ptr a = cupy.ascontiguousarray(a.reshape(bs, n, n).transpose(0, 2, 1), dtype=dtype) b = cupy.ascontiguousarray(b.reshape(bs, n, nrhs).transpose(0, 2, 1), dtype=dtype) if a.data.ptr == a_data_ptr: a = a.copy() if b.data.ptr == b_data_ptr: b = b.copy() if n > get_batched_gesv_limit(): warnings.warn('The matrix size ({}) exceeds the set limit ({})'.format( n, get_batched_gesv_limit())) handle = device.get_cublas_handle() lda = n a_step = lda * n * a.itemsize a_array = cupy.arange(a.data.ptr, a.data.ptr + a_step * bs, a_step, dtype=cupy.uintp) ldb = n b_step = ldb * nrhs * b.itemsize b_array = cupy.arange(b.data.ptr, b.data.ptr + b_step * bs, b_step, dtype=cupy.uintp) pivot = cupy.empty((bs, n), dtype=numpy.int32) dinfo = cupy.empty((bs, ), dtype=numpy.int32) info = numpy.empty((1, ), dtype=numpy.int32) # LU factorization (A = L * U) getrf(handle, n, a_array.data.ptr, lda, pivot.data.ptr, dinfo.data.ptr, bs) util._check_cublas_info_array_if_synchronization_allowed(getrf, dinfo) # Solves Ax = b getrs(handle, cublas.CUBLAS_OP_N, n, nrhs, a_array.data.ptr, lda, pivot.data.ptr, b_array.data.ptr, ldb, info.ctypes.data, bs) if info[0] != 0: msg = 'Error reported by {} in cuBLAS. '.format(getrs.__name__) if info[0] < 0: msg += 'The {}-th parameter had an illegal value.'.format(-info[0]) raise linalg.LinAlgError(msg) return b.transpose(0, 2, 1).reshape(b_shape)
def lu_solve(lu_and_piv, b, trans=0, overwrite_b=False, check_finite=True): """Solve an equation system, ``a * x = b``, given the LU factorization of ``a`` Args: lu_and_piv (tuple): LU factorization of matrix ``a`` (``(M, M)``) together with pivot indices. b (cupy.ndarray): The matrix with dimension ``(M,)`` or ``(M, N)``. trans ({0, 1, 2}): Type of system to solve: ======== ========= trans system ======== ========= 0 a x = b 1 a^T x = b 2 a^H x = b ======== ========= overwrite_b (bool): Allow overwriting data in b (may enhance performance) check_finite (bool): Whether to check that the input matrices contain only finite numbers. Disabling may give a performance gain, but may result in problems (crashes, non-termination) if the inputs do contain infinities or NaNs. Returns: cupy.ndarray: The matrix with dimension ``(M,)`` or ``(M, N)``. .. seealso:: :func:`scipy.linalg.lu_solve` """ (lu, ipiv) = lu_and_piv util._assert_cupy_array(lu) util._assert_rank2(lu) util._assert_nd_squareness(lu) m = lu.shape[0] if m != b.shape[0]: raise ValueError('incompatible dimensions.') dtype = lu.dtype if dtype.char == 'f': getrs = cusolver.sgetrs elif dtype.char == 'd': getrs = cusolver.dgetrs else: raise NotImplementedError('Only float32 and float64 are supported.') if trans == 0: trans = cublas.CUBLAS_OP_N elif trans == 1: trans = cublas.CUBLAS_OP_T elif trans == 2: trans = cublas.CUBLAS_OP_C else: raise ValueError('unknown trans') lu = lu.astype(dtype, order='F', copy=False) ipiv = ipiv.astype(ipiv.dtype, order='F', copy=True) # cuSolver uses 1-origin while SciPy uses 0-origin ipiv += 1 b = b.astype(dtype, order='F', copy=(not overwrite_b)) if check_finite: if lu.dtype.kind == 'f' and not cupy.isfinite(lu).all(): raise ValueError( 'array must not contain infs or NaNs.\n' 'Note that when a singular matrix is given, unlike ' 'scipy.linalg.lu_factor, cupyx.scipy.linalg.lu_factor ' 'returns an array containing NaN.') if b.dtype.kind == 'f' and not cupy.isfinite(b).all(): raise ValueError('array must not contain infs or NaNs') n = 1 if b.ndim == 1 else b.shape[1] cusolver_handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=numpy.int32) # solve for the inverse getrs(cusolver_handle, trans, m, n, lu.data.ptr, m, ipiv.data.ptr, b.data.ptr, m, dev_info.data.ptr) if dev_info[0] < 0: raise ValueError('illegal value in %d-th argument of ' 'internal getrs (lu_solve)' % -dev_info[0]) return b
def inv(a): """Computes the inverse of a matrix. This function computes matrix ``a_inv`` from n-dimensional regular matrix ``a`` such that ``dot(a, a_inv) == eye(n)``. Args: a (cupy.ndarray): The regular matrix Returns: cupy.ndarray: The inverse of a matrix. .. seealso:: :func:`numpy.linalg.inv` """ if a.ndim >= 3: return _batched_inv(a) # to prevent `a` to be overwritten a = a.copy() util._assert_cupy_array(a) util._assert_rank2(a) util._assert_nd_squareness(a) # support float32, float64, complex64, and complex128 if a.dtype.char in 'fdFD': dtype = a.dtype.char else: dtype = numpy.find_common_type((a.dtype.char, 'f'), ()).char cusolver_handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=numpy.int32) ipiv = cupy.empty((a.shape[0], 1), dtype=numpy.intc) if dtype == 'f': getrf = cusolver.sgetrf getrf_bufferSize = cusolver.sgetrf_bufferSize getrs = cusolver.sgetrs elif dtype == 'd': getrf = cusolver.dgetrf getrf_bufferSize = cusolver.dgetrf_bufferSize getrs = cusolver.dgetrs elif dtype == 'F': getrf = cusolver.cgetrf getrf_bufferSize = cusolver.cgetrf_bufferSize getrs = cusolver.cgetrs elif dtype == 'D': getrf = cusolver.zgetrf getrf_bufferSize = cusolver.zgetrf_bufferSize getrs = cusolver.zgetrs else: msg = ('dtype must be float32, float64, complex64 or complex128' ' (actual: {})'.format(a.dtype)) raise ValueError(msg) m = a.shape[0] buffersize = getrf_bufferSize(cusolver_handle, m, m, a.data.ptr, m) workspace = cupy.empty(buffersize, dtype=dtype) # LU factorization getrf(cusolver_handle, m, m, a.data.ptr, m, workspace.data.ptr, ipiv.data.ptr, dev_info.data.ptr) b = cupy.eye(m, dtype=dtype) # solve for the inverse getrs(cusolver_handle, 0, m, m, a.data.ptr, m, ipiv.data.ptr, b.data.ptr, m, dev_info.data.ptr) return b
def lu_factor(a, overwrite_a=False, check_finite=True): """LU decomposition. Decompose a given two-dimensional square matrix into ``P * L * U``, where ``P`` is a permutation matrix, ``L`` lower-triangular with unit diagonal elements, and ``U`` upper-triangular matrix. Note that in the current implementation ``a`` must be a real matrix, and only :class:`numpy.float32` and :class:`numpy.float64` are supported. Args: a (cupy.ndarray): The input matrix with dimension ``(N, N)`` overwrite_a (bool): Allow overwriting data in ``a`` (may enhance performance) check_finite (bool): Whether to check that the input matrices contain only finite numbers. Disabling may give a performance gain, but may result in problems (crashes, non-termination) if the inputs do contain infinities or NaNs. Returns: tuple: ``(lu, piv)`` where ``lu`` is a :class:`cupy.ndarray` storing ``U`` in its upper triangle, and ``L`` without unit diagonal elements in its lower triangle, and `piv` is a :class:`cupy.ndarray` storing pivot indices representing permutation matrix ``P``. .. seealso:: :func:`scipy.linalg.lu_factor` .. note:: Current implementation returns result different from SciPy when the matrix singular. SciPy returns an array containing ``0.`` while the current implementation returns an array containing ``nan``. >>> import numpy as np >>> import scipy.linalg >>> scipy.linalg.lu_factor(np.array([[0, 1], [0, 0]], \ dtype=np.float32)) (array([[0., 1.], [0., 0.]], dtype=float32), array([0, 1], dtype=int32)) >>> import cupy as cp >>> import cupyx.scipy.linalg >>> cupyx.scipy.linalg.lu_factor(cp.array([[0, 1], [0, 0]], \ dtype=cp.float32)) (array([[ 0., 1.], [nan, nan]], dtype=float32), array([0, 1], dtype=int32)) """ if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') a = cupy.asarray(a) util._assert_rank2(a) util._assert_nd_squareness(a) dtype = a.dtype if dtype.char == 'f': getrf = cusolver.sgetrf getrf_bufferSize = cusolver.sgetrf_bufferSize elif dtype.char == 'd': getrf = cusolver.dgetrf getrf_bufferSize = cusolver.dgetrf_bufferSize else: raise NotImplementedError('Only float32 and float64 are supported.') a = a.astype(dtype, order='F', copy=(not overwrite_a)) if check_finite: if a.dtype.kind == 'f' and not cupy.isfinite(a).all(): raise ValueError('array must not contain infs or NaNs') cusolver_handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=numpy.intc) ipiv = cupy.empty((a.shape[0], ), dtype=numpy.intc) m = a.shape[0] buffersize = getrf_bufferSize(cusolver_handle, m, m, a.data.ptr, m) workspace = cupy.empty(buffersize, dtype=dtype) # LU factorization getrf(cusolver_handle, m, m, a.data.ptr, m, workspace.data.ptr, ipiv.data.ptr, dev_info.data.ptr) if dev_info[0] < 0: raise ValueError('illegal value in %d-th argument of ' 'internal getrf (lu_factor)' % -dev_info[0]) elif dev_info[0] > 0: warn('Diagonal number %d is exactly zero. Singular matrix.' % dev_info[0], RuntimeWarning, stacklevel=2) # cuSolver uses 1-origin while SciPy uses 0-origin ipiv -= 1 return (a, ipiv)
def slogdet(a): """Returns sign and logarithm of the determinant of an array. It calculates the natural logarithm of the determinant of a given value. Args: a (cupy.ndarray): The input matrix with dimension ``(..., N, N)``. Returns: tuple of :class:`~cupy.ndarray`: It returns a tuple ``(sign, logdet)``. ``sign`` represents each sign of the determinant as a real number ``0``, ``1`` or ``-1``. 'logdet' represents the natural logarithm of the absolute of the determinant. If the determinant is zero, ``sign`` will be ``0`` and ``logdet`` will be ``-inf``. The shapes of both ``sign`` and ``logdet`` are equal to ``a.shape[:-2]``. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. warning:: To produce the same results as :func:`numpy.linalg.slogdet` for singular inputs, set the `linalg` configuration to `raise`. .. seealso:: :func:`numpy.linalg.slogdet` """ if a.ndim < 2: msg = ('%d-dimensional array given. ' 'Array must be at least two-dimensional' % a.ndim) raise linalg.LinAlgError(msg) util._assert_nd_squareness(a) dtype = numpy.promote_types(a.dtype.char, 'f') real_dtype = numpy.dtype(dtype.char.lower()) if dtype not in (numpy.float32, numpy.float64, numpy.complex64, numpy.complex128): msg = ('dtype must be float32, float64, complex64, or complex128' ' (actual: {})'.format(a.dtype)) raise ValueError(msg) a_shape = a.shape shape = a_shape[:-2] n = a_shape[-2] if a.size == 0: # empty batch (result is empty, too) or empty matrices det([[]]) == 1 sign = cupy.ones(shape, dtype) logdet = cupy.zeros(shape, real_dtype) return sign, logdet lu, ipiv, dev_info = decomposition._lu_factor(a, dtype) # dev_info < 0 means illegal value (in dimensions, strides, and etc.) that # should never happen even if the matrix contains nan or inf. # TODO(kataoka): assert dev_info >= 0 if synchronization is allowed for # debugging purposes. diag = cupy.diagonal(lu, axis1=-2, axis2=-1) logdet = cupy.log(cupy.abs(diag)).sum(axis=-1) # ipiv is 1-origin non_zero = cupy.count_nonzero(ipiv != cupy.arange(1, n + 1), axis=-1) if dtype.kind == "f": non_zero += cupy.count_nonzero(diag < 0, axis=-1) # Note: sign == -1 ** (non_zero % 2) sign = (non_zero % 2) * -2 + 1 if dtype.kind == "c": sign = sign * cupy.prod(diag / cupy.abs(diag), axis=-1) singular = dev_info > 0 return ( cupy.where(singular, dtype.type(0), sign.astype(dtype)).reshape(shape), cupy.where(singular, real_dtype.type('-inf'), logdet).reshape(shape), )
def inv_core(a, cholesky=False): """Computes the inverse of a matrix. This function computes matrix ``a_inv`` from n-dimensional regular matrix ``a`` such that ``dot(a, a_inv) == eye(n)``. Args: a (cupy.ndarray): The regular matrix b (Boolean): Use cholesky decomposition Returns: cupy.ndarray: The inverse of a matrix. .. seealso:: :func:`numpy.linalg.inv` """ xp = cupy.get_array_module(a) if xp == numpy: if cholesky: warnings.warn( "Current fast-inv using cholesky doesn't support numpy.ndarray." ) return numpy.linalg.inv(a) if not cuda.cusolver_enabled: raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0') # to prevent `a` to be overwritten a = a.copy() util._assert_cupy_array(a) util._assert_rank2(a) util._assert_nd_squareness(a) if a.dtype.char == 'f' or a.dtype.char == 'd': dtype = a.dtype.char else: dtype = numpy.find_common_type((a.dtype.char, 'f'), ()).char cusolver_handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=cupy.int) m = a.shape[0] b = cupy.eye(m, dtype=dtype) if not cholesky: if dtype == 'f': getrf = cusolver.sgetrf getrf_bufferSize = cusolver.sgetrf_bufferSize getrs = cusolver.sgetrs else: # dtype == 'd' getrf = cusolver.dgetrf getrf_bufferSize = cusolver.dgetrf_bufferSize getrs = cusolver.dgetrs buffersize = getrf_bufferSize(cusolver_handle, m, m, a.data.ptr, m) # TODO(y1r): cache buffer to avoid malloc workspace = cupy.empty(buffersize, dtype=dtype) ipiv = cupy.empty((a.shape[0], 1), dtype=dtype) # LU Decomposition getrf(cusolver_handle, m, m, a.data.ptr, m, workspace.data.ptr, ipiv.data.ptr, dev_info.data.ptr) # TODO(y1r): check dev_info status # solve for the inverse getrs(cusolver_handle, 0, m, m, a.data.ptr, m, ipiv.data.ptr, b.data.ptr, m, dev_info.data.ptr) # TODO(y1r): check dev_info status else: if dtype == 'f': potrf = cusolver.spotrf potrf_bufferSize = cusolver.spotrf_bufferSize potrs = cusolver.spotrs else: # dtype == 'd' potrf = cusolver.dpotrf potrf_bufferSize = cusolver.dpotrf_bufferSize potrs = cusolver.dpotrs buffersize = potrf_bufferSize(cusolver_handle, cublas.CUBLAS_FILL_MODE_UPPER, m, a.data.ptr, m) # TODO(y1r): cache buffer to avoid malloc workspace = cupy.empty(buffersize, dtype=dtype) # Cholesky Decomposition potrf(cusolver_handle, cublas.CUBLAS_FILL_MODE_UPPER, m, a.data.ptr, m, workspace.data.ptr, buffersize, dev_info.data.ptr) # TODO(y1r): check dev_info status # solve for the inverse potrs(cusolver_handle, cublas.CUBLAS_FILL_MODE_UPPER, m, m, a.data.ptr, m, b.data.ptr, m, dev_info.data.ptr) # TODO(y1r): check dev_info status return b