#!/usr/bin/env python

"""
Unit tests for skcuda.cublas
"""

from unittest import main, makeSuite, TestCase, TestSuite

import pycuda.driver as drv
import pycuda.gpuarray as gpuarray
from pycuda.tools import clear_context_caches, make_default_context
import numpy as np

drv.init()

_SEPS = np.finfo(np.float32).eps
_DEPS = np.finfo(np.float64).eps

import skcuda.cublas as cublas
import skcuda.misc as misc

def bptrs(a):
    """
    Pointer array when input represents a batch of matrices.
    """

    return gpuarray.arange(a.ptr,a.ptr+a.shape[0]*a.strides[0],a.strides[0],
                dtype=cublas.ctypes.c_void_p)

class test_cublas(TestCase):
    @classmethod
    def setUpClass(cls):
        cls.ctx = make_default_context()
        cls.cublas_handle = cublas.cublasCreate()

    @classmethod
    def tearDownClass(cls):
        cublas.cublasDestroy(cls.cublas_handle)
        cls.ctx.pop()
        clear_context_caches()

    def setUp(self):
        np.random.seed(23)    # For reproducible tests.

    # ISAMAX, IDAMAX, ICAMAX, IZAMAX
    def test_cublasIsamax(self):
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasIsamax(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.argmax(x))

    def test_cublasIdamax(self):
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasIdamax(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.argmax(x))

    def test_cublasIcamax(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasIcamax(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.argmax(np.abs(x.real) + np.abs(x.imag)))

    def test_cublasIzamax(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasIzamax(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.argmax(np.abs(x.real) + np.abs(x.imag)))

    # ISAMIN, IDAMIN, ICAMIN, IZAMIN
    def test_cublasIsamin(self):
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasIsamin(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.argmin(x))

    def test_cublasIdamin(self):
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasIdamin(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.argmin(x))

    def test_cublasIcamin(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasIcamin(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.argmin(np.abs(x.real) + np.abs(x.imag)))

    def test_cublasIzamin(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasIzamin(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.argmin(np.abs(x.real) + np.abs(x.imag)))

    # SASUM, DASUM, SCASUM, DZASUM
    def test_cublasSasum(self):
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasSasum(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.sum(np.abs(x)))

    def test_cublasDasum(self):
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasDasum(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.sum(np.abs(x)))

    def test_cublasScasum(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasScasum(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.sum(np.abs(x.real)+np.abs(x.imag)))

    def test_cublasDzasum(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasDzasum(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.sum(np.abs(x.real)+np.abs(x.imag)))

    # SAXPY, DAXPY, CAXPY, ZAXPY
    def test_cublasSaxpy(self):
        alpha = np.float32(np.random.rand())
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        y = np.random.rand(5).astype(np.float32)
        y_gpu = gpuarray.to_gpu(y)
        cublas.cublasSaxpy(self.cublas_handle, x_gpu.size, alpha, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), alpha*x+y)

    def test_cublasDaxpy(self):
        alpha = np.float64(np.random.rand())
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        y = np.random.rand(5).astype(np.float64)
        y_gpu = gpuarray.to_gpu(y)
        cublas.cublasDaxpy(self.cublas_handle, x_gpu.size, alpha, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), alpha*x+y)

    def test_cublasCaxpy(self):
        alpha = np.complex64(np.random.rand()+1j*np.random.rand())
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        y = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        y_gpu = gpuarray.to_gpu(y)
        cublas.cublasCaxpy(self.cublas_handle, x_gpu.size, alpha, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), alpha*x+y)

    def test_cublasZaxpy(self):
        alpha = np.complex128(np.random.rand()+1j*np.random.rand())
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        y = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        y_gpu = gpuarray.to_gpu(y)
        cublas.cublasZaxpy(self.cublas_handle, x_gpu.size, alpha, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), alpha*x+y)

    # SCOPY, DCOPY, CCOPY, ZCOPY
    def test_cublasScopy(self):
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = gpuarray.zeros_like(x_gpu)
        cublas.cublasScopy(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), x_gpu.get())

    def test_cublasDcopy(self):
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = gpuarray.zeros_like(x_gpu)
        cublas.cublasDcopy(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), x_gpu.get())

    def test_cublasCcopy(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = misc.zeros_like(x_gpu)
        cublas.cublasCcopy(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), x_gpu.get())

    def test_cublasZcopy(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = misc.zeros_like(x_gpu)
        cublas.cublasZcopy(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), x_gpu.get())

    # SDOT, DDOT, CDOTU, CDOTC, ZDOTU, ZDOTC
    def test_cublasSdot(self):
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        y = np.random.rand(5).astype(np.float32)
        y_gpu = gpuarray.to_gpu(y)
        result = cublas.cublasSdot(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                                   y_gpu.gpudata, 1)
        assert np.allclose(result, np.dot(x, y))

    def test_cublasDdot(self):
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        y = np.random.rand(5).astype(np.float64)
        y_gpu = gpuarray.to_gpu(y)
        result = cublas.cublasDdot(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                                   y_gpu.gpudata, 1)
        assert np.allclose(result, np.dot(x, y))

    def test_cublasCdotu(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        y = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        y_gpu = gpuarray.to_gpu(y)
        result = cublas.cublasCdotu(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                                    y_gpu.gpudata, 1)
        assert np.allclose(result, np.dot(x, y))

    def test_cublasCdotc(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        y = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        y_gpu = gpuarray.to_gpu(y)
        result = cublas.cublasCdotc(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                                    y_gpu.gpudata, 1)
        assert np.allclose(result, np.dot(np.conj(x), y))

    def test_cublasZdotu(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        y = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        y_gpu = gpuarray.to_gpu(y)
        result = cublas.cublasZdotu(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                                    y_gpu.gpudata, 1)
        assert np.allclose(result, np.dot(x, y))

    def test_cublasZdotc(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        y = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        y_gpu = gpuarray.to_gpu(y)
        result = cublas.cublasZdotc(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                                    y_gpu.gpudata, 1)
        assert np.allclose(result, np.dot(np.conj(x), y))

    # SNRM2, DNRM2, SCNRM2, DZNRM2
    def test_cublasSrnm2(self):
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasSnrm2(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.linalg.norm(x))

    def test_cublasDrnm2(self):
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasDnrm2(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.linalg.norm(x))

    def test_cublasScrnm2(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasScnrm2(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.linalg.norm(x))

    def test_cublasDzrnm2(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        result = cublas.cublasDznrm2(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1)
        assert np.allclose(result, np.linalg.norm(x))

    # SSCAL, DSCAL, CSCAL, CSSCAL, ZSCAL, ZDSCAL
    def test_cublasSscal(self):
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        alpha = np.float32(np.random.rand())
        cublas.cublasSscal(self.cublas_handle, x_gpu.size, alpha,
                           x_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), alpha*x)

    def test_cublasCscal(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        alpha = np.complex64(np.random.rand()+1j*np.random.rand())
        cublas.cublasCscal(self.cublas_handle, x_gpu.size, alpha,
                           x_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), alpha*x)

    def test_cublasCsscal(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        alpha = np.float32(np.random.rand())
        cublas.cublasCscal(self.cublas_handle, x_gpu.size, alpha,
                           x_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), alpha*x)

    def test_cublasDscal(self):
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        alpha = np.float64(np.random.rand())
        cublas.cublasDscal(self.cublas_handle, x_gpu.size, alpha,
                           x_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), alpha*x)

    def test_cublasZscal(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        alpha = np.complex128(np.random.rand()+1j*np.random.rand())
        cublas.cublasZscal(self.cublas_handle, x_gpu.size, alpha,
                           x_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), alpha*x)

    def test_cublasZdscal(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        alpha = np.float64(np.random.rand())
        cublas.cublasZdscal(self.cublas_handle, x_gpu.size, alpha,
                           x_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), alpha*x)

    # SROT, DROT, CROT, CSROT, ZROT, ZDROT
    def test_cublasSrot(self):
        x = np.array([1, 2, 3]).astype(np.float32)
        y = np.array([4, 5, 6]).astype(np.float32)
        s = 2.0
        c = 3.0
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = gpuarray.to_gpu(x)
        cublas.cublasSrot(self.cublas_handle, x_gpu.size,
                          x_gpu.gpudata, 1,
                          y_gpu.gpudata, 1,
                          c, s)
        assert np.allclose(x_gpu.get(), [5, 10, 15])
        assert np.allclose(y_gpu.get(), [1, 2, 3])

    # SSWAP, DSWAP, CSWAP, ZSWAP
    def test_cublasSswap(self):
        x = np.random.rand(5).astype(np.float32)
        x_gpu = gpuarray.to_gpu(x)
        y = np.random.rand(5).astype(np.float32)
        y_gpu = gpuarray.to_gpu(y)
        cublas.cublasSswap(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), y)

    def test_cublasDswap(self):
        x = np.random.rand(5).astype(np.float64)
        x_gpu = gpuarray.to_gpu(x)
        y = np.random.rand(5).astype(np.float64)
        y_gpu = gpuarray.to_gpu(y)
        cublas.cublasDswap(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), y)

    def test_cublasCswap(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        x_gpu = gpuarray.to_gpu(x)
        y = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex64)
        y_gpu = gpuarray.to_gpu(y)
        cublas.cublasCswap(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), y)

    def test_cublasZswap(self):
        x = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        x_gpu = gpuarray.to_gpu(x)
        y = (np.random.rand(5)+1j*np.random.rand(5)).astype(np.complex128)
        y_gpu = gpuarray.to_gpu(y)
        cublas.cublasZswap(self.cublas_handle, x_gpu.size, x_gpu.gpudata, 1,
                           y_gpu.gpudata, 1)
        assert np.allclose(x_gpu.get(), y)

    # SGEMV, DGEMV, CGEMV, ZGEMV
    def test_cublasSgemv(self):
        a = np.random.rand(2, 3).astype(np.float32)
        x = np.random.rand(3, 1).astype(np.float32)
        a_gpu = gpuarray.to_gpu(a.T.copy())
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = gpuarray.empty((2, 1), np.float32)
        alpha = np.float32(1.0)
        beta = np.float32(0.0)
        cublas.cublasSgemv(self.cublas_handle, 'n', 2, 3, alpha,
                           a_gpu.gpudata, 2, x_gpu.gpudata,
                           1, beta, y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), np.dot(a, x))

    def test_cublasDgemv(self):
        a = np.random.rand(2, 3).astype(np.float64)
        x = np.random.rand(3, 1).astype(np.float64)
        a_gpu = gpuarray.to_gpu(a.T.copy())
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = gpuarray.empty((2, 1), np.float64)
        alpha = np.float64(1.0)
        beta = np.float64(0.0)
        cublas.cublasDgemv(self.cublas_handle, 'n', 2, 3, alpha,
                           a_gpu.gpudata, 2, x_gpu.gpudata,
                           1, beta, y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), np.dot(a, x))

    def test_cublasCgemv(self):
        a = (np.random.rand(2, 3)+1j*np.random.rand(2, 3)).astype(np.complex64)
        x = (np.random.rand(3, 1)+1j*np.random.rand(3, 1)).astype(np.complex64)
        a_gpu = gpuarray.to_gpu(a.T.copy())
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = gpuarray.empty((2, 1), np.complex64)
        alpha = np.complex64(1.0)
        beta = np.complex64(0.0)
        cublas.cublasCgemv(self.cublas_handle, 'n', 2, 3, alpha,
                           a_gpu.gpudata, 2, x_gpu.gpudata,
                           1, beta, y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), np.dot(a, x))

    def test_cublasZgemv(self):
        a = (np.random.rand(2, 3)+1j*np.random.rand(2, 3)).astype(np.complex128)
        x = (np.random.rand(3, 1)+1j*np.random.rand(3, 1)).astype(np.complex128)
        a_gpu = gpuarray.to_gpu(a.T.copy())
        x_gpu = gpuarray.to_gpu(x)
        y_gpu = gpuarray.empty((2, 1), np.complex128)
        alpha = np.complex128(1.0)
        beta = np.complex128(0.0)
        cublas.cublasZgemv(self.cublas_handle, 'n', 2, 3, alpha,
                           a_gpu.gpudata, 2, x_gpu.gpudata,
                           1, beta, y_gpu.gpudata, 1)
        assert np.allclose(y_gpu.get(), np.dot(a, x))

    # SGEAM, CGEAM, DGEAM, ZDGEAM
    def test_cublasSgeam(self):
        a = np.random.rand(2, 3).astype(np.float32)
        b = np.random.rand(2, 3).astype(np.float32)
        a_gpu = gpuarray.to_gpu(a.copy())
        b_gpu = gpuarray.to_gpu(b.copy())
        c_gpu = gpuarray.zeros_like(a_gpu)
        alpha = np.float32(np.random.rand())
        beta = np.float32(np.random.rand())
        cublas.cublasSgeam(self.cublas_handle, 'n', 'n', 2, 3,
                           alpha, a_gpu.gpudata, 2,
                           beta, b_gpu.gpudata, 2,
                           c_gpu.gpudata, 2)
        assert np.allclose(c_gpu.get(), alpha*a+beta*b)

    def test_cublasCgeam(self):
        a = (np.random.rand(2, 3)+1j*np.random.rand(2, 3)).astype(np.complex64)
        b = (np.random.rand(2, 3)+1j*np.random.rand(2, 3)).astype(np.complex64)
        a_gpu = gpuarray.to_gpu(a.copy())
        b_gpu = gpuarray.to_gpu(b.copy())
        c_gpu = gpuarray.zeros_like(a_gpu)
        alpha = np.complex64(np.random.rand()+1j*np.random.rand())
        beta = np.complex64(np.random.rand()+1j*np.random.rand())
        cublas.cublasCgeam(self.cublas_handle, 'n', 'n', 2, 3,
                           alpha, a_gpu.gpudata, 2,
                           beta, b_gpu.gpudata, 2,
                           c_gpu.gpudata, 2)
        assert np.allclose(c_gpu.get(), alpha*a+beta*b)

    def test_cublasDgeam(self):
        a = np.random.rand(2, 3).astype(np.float64)
        b = np.random.rand(2, 3).astype(np.float64)
        a_gpu = gpuarray.to_gpu(a.copy())
        b_gpu = gpuarray.to_gpu(b.copy())
        c_gpu = gpuarray.zeros_like(a_gpu)
        alpha = np.float64(np.random.rand())
        beta = np.float64(np.random.rand())
        cublas.cublasDgeam(self.cublas_handle, 'n', 'n', 2, 3,
                           alpha, a_gpu.gpudata, 2,
                           beta, b_gpu.gpudata, 2,
                           c_gpu.gpudata, 2)
        assert np.allclose(c_gpu.get(), alpha*a+beta*b)

    def test_cublasZgeam(self):
        a = (np.random.rand(2, 3)+1j*np.random.rand(2, 3)).astype(np.complex128)
        b = (np.random.rand(2, 3)+1j*np.random.rand(2, 3)).astype(np.complex128)
        a_gpu = gpuarray.to_gpu(a.copy())
        b_gpu = gpuarray.to_gpu(b.copy())
        c_gpu = gpuarray.zeros_like(a_gpu)
        alpha = np.complex128(np.random.rand()+1j*np.random.rand())
        beta = np.complex128(np.random.rand()+1j*np.random.rand())
        cublas.cublasZgeam(self.cublas_handle, 'n', 'n', 2, 3,
                           alpha, a_gpu.gpudata, 2,
                           beta, b_gpu.gpudata, 2,
                           c_gpu.gpudata, 2)
        assert np.allclose(c_gpu.get(), alpha*a+beta*b)

    # CgemmBatched, ZgemmBatched
    def test_cublasCgemmBatched(self):
        l, m, k, n = 11, 7, 5, 3
        A = (np.random.rand(l, m, k)+1j*np.random.rand(l, m, k)).astype(np.complex64)
        B = (np.random.rand(l, k, n)+1j*np.random.rand(l, k, n)).astype(np.complex64)

        C_res = np.einsum('nij,njk->nik', A, B)

        a_gpu = gpuarray.to_gpu(A)
        b_gpu = gpuarray.to_gpu(B)
        c_gpu = gpuarray.empty((l, m, n), np.complex64)

        alpha = np.complex64(1.0)
        beta = np.complex64(0.0)

        a_arr = bptrs(a_gpu)
        b_arr = bptrs(b_gpu)
        c_arr = bptrs(c_gpu)

        cublas.cublasCgemmBatched(self.cublas_handle, 'n','n',
                                  n, m, k, alpha,
                                  b_arr.gpudata, n,
                                  a_arr.gpudata, k,
                                  beta, c_arr.gpudata, n, l)

        assert np.allclose(C_res, c_gpu.get())

    def test_cublasZgemmBatched(self):
        l, m, k, n = 11, 7, 5, 3
        A = (np.random.rand(l, m, k)+1j*np.random.rand(l, m, k)).astype(np.complex128)
        B = (np.random.rand(l, k, n)+1j*np.random.rand(l, k, n)).astype(np.complex128)

        C_res = np.einsum('nij,njk->nik', A, B)

        a_gpu = gpuarray.to_gpu(A)
        b_gpu = gpuarray.to_gpu(B)
        c_gpu = gpuarray.empty((l, m, n), np.complex128)

        alpha = np.complex128(1.0)
        beta = np.complex128(0.0)

        a_arr = bptrs(a_gpu)
        b_arr = bptrs(b_gpu)
        c_arr = bptrs(c_gpu)

        cublas.cublasZgemmBatched(self.cublas_handle, 'n','n',
                                  n, m, k, alpha,
                                  b_arr.gpudata, n,
                                  a_arr.gpudata, k,
                                  beta, c_arr.gpudata, n, l)

        assert np.allclose(C_res, c_gpu.get())

    # SgemmBatched, DgemmBatched
    def test_cublasSgemmBatched(self):
        l, m, k, n = 11, 7, 5, 3
        A = np.random.rand(l, m, k).astype(np.float32)
        B = np.random.rand(l, k, n).astype(np.float32)

        C_res = np.einsum('nij,njk->nik', A, B)

        a_gpu = gpuarray.to_gpu(A)
        b_gpu = gpuarray.to_gpu(B)
        c_gpu = gpuarray.empty((l, m, n), np.float32)

        alpha = np.float32(1.0)
        beta = np.float32(0.0)

        a_arr = bptrs(a_gpu)
        b_arr = bptrs(b_gpu)
        c_arr = bptrs(c_gpu)

        cublas.cublasSgemmBatched(self.cublas_handle, 'n','n',
                                  n, m, k, alpha,
                                  b_arr.gpudata, n,
                                  a_arr.gpudata, k,
                                  beta, c_arr.gpudata, n, l)

        assert np.allclose(C_res, c_gpu.get())

    def test_cublasDgemmBatched(self):
        l, m, k, n = 11, 7, 5, 3
        A = np.random.rand(l, m, k).astype(np.float64)
        B = np.random.rand(l, k, n).astype(np.float64)

        C_res = np.einsum('nij,njk->nik',A,B)

        a_gpu = gpuarray.to_gpu(A)
        b_gpu = gpuarray.to_gpu(B)
        c_gpu = gpuarray.empty((l, m, n), np.float64)

        alpha = np.float64(1.0)
        beta = np.float64(0.0)

        a_arr = bptrs(a_gpu)
        b_arr = bptrs(b_gpu)
        c_arr = bptrs(c_gpu)

        cublas.cublasDgemmBatched(self.cublas_handle, 'n','n',
                                  n, m, k, alpha,
                                  b_arr.gpudata, n,
                                  a_arr.gpudata, k,
                                  beta, c_arr.gpudata, n, l)

        assert np.allclose(C_res, c_gpu.get())

    # StrsmBatched, DtrsmBatched
    def test_cublasStrsmBatched(self):
        l, m, n = 11, 7, 5
        A = np.random.rand(l, m, m).astype(np.float32)
        B = np.random.rand(l, m, n).astype(np.float32)

        A = np.array(list(map(np.triu, A)))
        X = np.array([np.linalg.solve(a, b) for a, b in zip(A, B)])

        alpha = np.float32(1.0)

        a_gpu = gpuarray.to_gpu(A)
        b_gpu = gpuarray.to_gpu(B)

        a_arr = bptrs(a_gpu)
        b_arr = bptrs(b_gpu)

        cublas.cublasStrsmBatched(self.cublas_handle, 'r', 'l', 'n', 'n',
                                  n, m, alpha,
                                  a_arr.gpudata, m,
                                  b_arr.gpudata, n, l)

        assert np.allclose(X, b_gpu.get(), 5)

    def test_cublasDtrsmBatched(self):
        l, m, n = 11, 7, 5
        A = np.random.rand(l, m, m).astype(np.float64)
        B = np.random.rand(l, m, n).astype(np.float64)

        A = np.array(list(map(np.triu, A)))
        X = np.array([np.linalg.solve(a, b) for a, b in zip(A, B)])

        alpha = np.float64(1.0)

        a_gpu = gpuarray.to_gpu(A)
        b_gpu = gpuarray.to_gpu(B)

        a_arr = bptrs(a_gpu)
        b_arr = bptrs(b_gpu)

        cublas.cublasDtrsmBatched(self.cublas_handle, 'r', 'l', 'n', 'n',
                                  n, m, alpha,
                                  a_arr.gpudata, m,
                                  b_arr.gpudata, n, l)

        assert np.allclose(X, b_gpu.get(), 5)

    # SgetrfBatched, DgetrfBatched
    def test_cublasSgetrfBatched(self):
        from scipy.linalg import lu_factor
        l, m = 11, 7
        A = np.random.rand(l, m, m).astype(np.float32)
        A = np.array([np.dot(a, a.T) for a in A])

        a_gpu = gpuarray.to_gpu(A)
        a_arr = bptrs(a_gpu)
        p_gpu = gpuarray.empty((l, m), np.int32)
        i_gpu = gpuarray.zeros(1, np.int32)
        X = np.array([ lu_factor(a)[0] for a in A])

        cublas.cublasSgetrfBatched(self.cublas_handle,
                                   m, a_arr.gpudata, m,
                                   p_gpu.gpudata, i_gpu.gpudata, l)

        X_ = np.array([a.T for a in a_gpu.get()])

        assert np.allclose(X, X_, atol=10*_SEPS)

    def test_cublasDgetrfBatched(self):
        from scipy.linalg import lu_factor
        l, m = 11, 7
        A = np.random.rand(l, m, m).astype(np.float64)
        A = np.array([np.dot(a, a.T) for a in A])

        a_gpu = gpuarray.to_gpu(A)
        a_arr = bptrs(a_gpu)
        p_gpu = gpuarray.empty((l, m), np.int32)
        i_gpu = gpuarray.zeros(1, np.int32)
        X = np.array([ lu_factor(a)[0] for a in A])

        cublas.cublasDgetrfBatched(self.cublas_handle,
                                   m, a_arr.gpudata, m,
                                   p_gpu.gpudata, i_gpu.gpudata, l)

        X_ = np.array([a.T for a in a_gpu.get()])

        assert np.allclose(X,X_)


def suite():
    context = make_default_context()
    device = context.get_device()
    context.pop()

    s = TestSuite()
    s.addTest(test_cublas('test_cublasIsamax'))
    s.addTest(test_cublas('test_cublasIcamax'))
    s.addTest(test_cublas('test_cublasIsamin'))
    s.addTest(test_cublas('test_cublasIcamin'))
    s.addTest(test_cublas('test_cublasSasum'))
    s.addTest(test_cublas('test_cublasScasum'))
    s.addTest(test_cublas('test_cublasSaxpy'))
    s.addTest(test_cublas('test_cublasCaxpy'))
    s.addTest(test_cublas('test_cublasScopy'))
    s.addTest(test_cublas('test_cublasCcopy'))
    s.addTest(test_cublas('test_cublasSdot'))
    s.addTest(test_cublas('test_cublasCdotu'))
    s.addTest(test_cublas('test_cublasCdotc'))
    s.addTest(test_cublas('test_cublasSrnm2'))
    s.addTest(test_cublas('test_cublasScrnm2'))
    s.addTest(test_cublas('test_cublasSscal'))
    s.addTest(test_cublas('test_cublasCscal'))
    s.addTest(test_cublas('test_cublasSrot'))
    s.addTest(test_cublas('test_cublasSswap'))
    s.addTest(test_cublas('test_cublasCswap'))
    s.addTest(test_cublas('test_cublasSgemv'))
    s.addTest(test_cublas('test_cublasCgemv'))
    s.addTest(test_cublas('test_cublasSgeam'))
    s.addTest(test_cublas('test_cublasCgeam'))
    s.addTest(test_cublas('test_cublasSgemmBatched'))
    s.addTest(test_cublas('test_cublasCgemmBatched'))
    s.addTest(test_cublas('test_cublasStrsmBatched'))
    s.addTest(test_cublas('test_cublasSgetrfBatched'))
    if misc.get_compute_capability(device) >= 1.3:
        s.addTest(test_cublas('test_cublasIdamax'))
        s.addTest(test_cublas('test_cublasIzamax'))
        s.addTest(test_cublas('test_cublasIdamin'))
        s.addTest(test_cublas('test_cublasIzamin'))
        s.addTest(test_cublas('test_cublasDasum'))
        s.addTest(test_cublas('test_cublasDzasum'))
        s.addTest(test_cublas('test_cublasDaxpy'))
        s.addTest(test_cublas('test_cublasZaxpy'))
        s.addTest(test_cublas('test_cublasDcopy'))
        s.addTest(test_cublas('test_cublasZcopy'))
        s.addTest(test_cublas('test_cublasDdot'))
        s.addTest(test_cublas('test_cublasZdotu'))
        s.addTest(test_cublas('test_cublasZdotc'))
        s.addTest(test_cublas('test_cublasDrnm2'))
        s.addTest(test_cublas('test_cublasDzrnm2'))
        s.addTest(test_cublas('test_cublasDscal'))
        s.addTest(test_cublas('test_cublasZscal'))
        s.addTest(test_cublas('test_cublasZdscal'))
        s.addTest(test_cublas('test_cublasDswap'))
        s.addTest(test_cublas('test_cublasZswap'))
        s.addTest(test_cublas('test_cublasDgemv'))
        s.addTest(test_cublas('test_cublasZgemv'))
        s.addTest(test_cublas('test_cublasDgeam'))
        s.addTest(test_cublas('test_cublasZgeam'))
        s.addTest(test_cublas('test_cublasDgemmBatched'))
        s.addTest(test_cublas('test_cublasZgemmBatched'))
        s.addTest(test_cublas('test_cublasDtrsmBatched'))
        s.addTest(test_cublas('test_cublasDgetrfBatched'))
    return s

if __name__ == '__main__':
    main(defaultTest = 'suite')