<filename>tests/test_gemmbatchedex.py
#!/usr/bin/env python

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

from ctypes import create_string_buffer
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
import time 
import sys 
sys.path.append('/home/hongy0a/seismicapp/scikit-cuda')

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)



def test_cublasCgemmBatched():
    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(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_cublasCgemmBatchedSwitch():
    l, m, k, n = 1, 4, 2, 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)
    A = (np.ones((l, m, k))+1j*np.ones((l, m, k))).astype(np.complex64)
    B = (np.ones((l, k, n))+1j*np.ones((l, k, n))).astype(np.complex64)
    print("A", A)
    print("B", B)
    C_res = np.einsum('nij,njk->nik', A, B)
    C_res = np.zeros((l,m,n), dtype=np.csingle)
    for i in range(C_res.shape[0]):
        C_res[i,:,:] = A[i,:,:] @ B[i,:,:]
    print("Real res",C_res)
    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(cublas_handle, 'n','n',
                                n, m, k, alpha,
                                b_arr.gpudata, n,
                                a_arr.gpudata, k,
                                beta, c_arr.gpudata, n, l)
    print("cgpu",c_gpu.get())
    assert np.allclose(C_res, c_gpu.get())

def test_cublasCgemmBatchedEx():
    l, m, k, n = 1, 3, 4, 2
    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)
    # C_res = np.zeros((l,m,n), dtype=np.np.csingle)
    # for i in range(C_res.shape[0]):
    #     C_res[i,:,:] = A[i,:,:] @ B[i,:,:]
    

    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.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    b_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_32F, n,
    a_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_32F, k,
    beta, c_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)
    print(C_res)
    print(c_gpu.get())
    assert np.allclose(C_res, c_gpu.get())

def test_cublasCgemmBatchedExFP32():
    l, m, k, n = 1, 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)
    rA = np.random.rand(l, m, k).astype(np.float32)
    iA = np.random.rand(l, m, k).astype(np.float32)
    rB = np.random.rand(l, k, n).astype(np.float32)
    iB = np.random.rand(l, k, n).astype(np.float32)
    A = (rA + 1j * iA).astype(np.complex64)
    B = (rB + 1j * iB).astype(np.complex64)
    C_res = np.einsum('nij,njk->nik', A, B)

    ra_gpu = gpuarray.to_gpu(rA)
    ia_gpu = gpuarray.to_gpu(iA)
    rb_gpu = gpuarray.to_gpu(rB)
    ib_gpu = gpuarray.to_gpu(iB)
    
    mid1_gpu = gpuarray.empty((l, m, n), np.float32)
    mid2_gpu = gpuarray.empty((l, m, n), np.float32)
    mid3_gpu = gpuarray.empty((l, m, n), np.float32)
    mid4_gpu = gpuarray.empty((l, m, n), np.float32)

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

    ra_arr = bptrs(ra_gpu)
    rb_arr = bptrs(rb_gpu)
    ia_arr = bptrs(ia_gpu)
    ib_arr = bptrs(ib_gpu)    

    mid1_arr = bptrs(mid1_gpu)
    mid2_arr = bptrs(mid2_gpu)
    mid3_arr = bptrs(mid3_gpu)
    mid4_arr = bptrs(mid4_gpu)

    cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    rb_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n,
    ra_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, k,
    beta, mid1_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)

    cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    ib_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n,
    ia_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, k,
    beta, mid2_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)

    cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    ib_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n,
    ra_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, k,
    beta, mid3_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)

    cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    rb_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n,
    ia_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, k,
    beta, mid4_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)
    mid1_gpu = mid1_gpu - mid2_gpu
    mid3_gpu = mid3_gpu + mid4_gpu
    c_gpu = mid1_gpu.get() + 1j * mid3_gpu.get()    

    assert np.allclose(C_res, c_gpu)

    
def test_cublasCgemmBatchedExFP16():
    l, m, k, n = 1, 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)
    A = (np.ones((l, m, k))+1j*np.ones((l, m, k))).astype(np.complex128)
    B = (np.ones((l, k, n))+1j*np.ones((l, k, n))).astype(np.complex128)
    print(A.real.dtype)
    rA = A.real.astype(np.float16)
    rB = B.real.astype(np.float16)
    print(rB.flags)
    # iB = B.imag.astype(np.float64)
    iA = A.imag.astype(np.float16)
    C_res = np.einsum('nij,njk->nik', A, B)

    ra_gpu = gpuarray.to_gpu(rA)
    rb_gpu = gpuarray.to_gpu(rB)
    ia_gpu = gpuarray.to_gpu(iA)
    # ib_gpu = gpuarray.to_gpu(iB)
    mid1_gpu = gpuarray.empty((l, m, n), np.float16)
    # mid2_gpu = gpuarray.empty((l, m, n), np.float16)
    # mid3_gpu = gpuarray.empty((l, m, n), np.float16)
    # mid4_gpu = gpuarray.empty((l, m, n), np.float16)


    print("before computation ", ra_gpu.get(), ra_gpu.dtype)
    print("before computation ", rb_gpu.get(), rb_gpu.dtype)
    alpha = np.float16(1.0)
    beta = np.float16(0.0)

    ra_arr = bptrs(ra_gpu)
    rb_arr = bptrs(rb_gpu)
    # ia_arr = bptrs(ia_gpu)
    # ib_arr = bptrs(ib_gpu)    

    mid1_arr = bptrs(mid1_gpu)
    # mid2_arr = bptrs(mid2_gpu)
    # mid3_arr = bptrs(mid3_gpu)
    # mid4_arr = bptrs(mid4_gpu)

    for i in range(10):
        s = time.time()
        cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
        rb_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_16F, n,
        ra_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_16F, k,
        beta, mid1_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_16F, n, l,
        cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,
        cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)
        e = time.time()
        print(e-s)
    print("Cres",C_res)
    
    print("mid1",mid1_gpu.get())

    # assert np.allclose(C_res, mid1_gpu.get())

    # 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(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())


    return 

    cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    ib_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n,
    ia_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, k,
    beta, mid2_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)

    cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    ib_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n,
    ra_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, k,
    beta, mid3_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)

    cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    rb_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n,
    ia_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, k,
    beta, mid4_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)
    mid1_gpu = mid1_gpu - mid2_gpu
    mid3_gpu = mid3_gpu + mid4_gpu
    c_gpu = mid1_gpu.get() + 1j * mid3_gpu.get()    

    assert np.allclose(C_res, c_gpu)

def test_cublasCgemmBatchedExReal():
    l, m, k, n = 1, 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)
    rA = np.random.rand(l, m, k).astype(np.float32)
    iA = np.random.rand(l, m, k).astype(np.float32)
    rB = np.random.rand(l, k, n).astype(np.float32)
    iB = np.random.rand(l, k, n).astype(np.float32)
    A = (rA + 1j * iA).astype(np.complex64)
    B = (rB + 1j * iB).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.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
    b_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_32F, n,
    a_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_32F, k,
    beta, c_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_32F, n, l,
    cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)
    print(C_res)
    print(c_gpu.get())
    assert np.allclose(C_res, c_gpu.get())









    
def test_cublasCgemmBatchedExINT8():
    l, m, k, n = 1, 7, 5, 3

    A = (np.ones((l, m, k)) * 100 +1j*np.ones((l, m, k))).astype(np.complex128)
    B = (np.ones((l, k, n)) * 100 +1j*np.ones((l, k, n))).astype(np.complex128)
    
    rA = (np.random.rand(7,5)*127).astype(np.int8)
    rB = (np.random.rand(5,3)*127).astype(np.int8)
    C_res = np.einsum('nij,njk->nik', A, B)

    ra_gpu = gpuarray.to_gpu(rA)
    rb_gpu = gpuarray.to_gpu(rB)
    mid1_gpu = gpuarray.empty((l, m, n), np.float32)

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

    ra_arr = bptrs(ra_gpu)
    rb_arr = bptrs(rb_gpu)

    mid1_arr = bptrs(mid1_gpu)

    for i in range(10):
        s = time.time()
        cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha, 
        rb_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_8I, n,
        ra_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_8I, k,
        beta, mid1_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_R_32F, n, l,
        cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,
        cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)
        e = time.time()
        if i > 5:
            print("iter {} {:.3f} us.".format(i,(e-s)*1e6))
    
    print("rA", rA, "rB", rB)
    cpures = rA.astype(np.float32) @ rB.astype(np.float32)
    print("mid1",mid1_gpu.get())
    print("cpures ", cpures)


def test_cublasCgemmBatchedExC8():
    l, m, k, n = 1, 7, 5, 3

    A = (np.ones((l, m, k)) * 100 +1j*np.ones((l, m, k))).astype(np.complex64)
    B = (np.ones((l, k, n)) * 100 +1j*np.ones((l, k, n))).astype(np.complex64)
    
    rA = (np.random.rand(m,k*2)*127).astype(np.int8)
    rB = (np.random.rand(k,2*n)*127).astype(np.int8)
    C_res = np.einsum('nij,njk->nik', A, B)

    ra_gpu = gpuarray.to_gpu(rA)
    rb_gpu = gpuarray.to_gpu(rB)
    mid1_gpu = gpuarray.empty((l, m, n), np.float32)

    alpha = np.array([1.0,0.0]).astype(np.int8)
    beta = np.array([0.0,0.0]).astype(np.int8)
    alpha_arr = gpuarray.to_gpu(alpha)
    beta_arr = gpuarray.to_gpu(beta)
    ra_arr = bptrs(ra_gpu)
    rb_arr = bptrs(rb_gpu)

    mid1_arr = bptrs(mid1_gpu)

    for i in range(10):
        s = time.time()
        cublas.cublasGemmBatchedEx(cublas_handle, 'n', 'n', n, m, k, alpha_arr.gpudata, 
        rb_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_8I, n,
        ra_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_8I, k,
        beta_arr.gpudata, mid1_arr.gpudata, cublas.CUDA_DATATYPE.CUDA_C_32F, n, l,
        cublas.CUBLAS_COMPUTE_TYPE.CUBLAS_COMPUTE_32F,
        cublas.CUBLAS_GEMM_ALGO.CUBLAS_GEMM_DEFAULT)
        e = time.time()
        if i > 5:
            print("iter {} {:.3f} us.".format(i,(e-s)*1e6))
    
    # print("rA", rA, "rB", rB)
    # cpures = rA.astype(np.float32) @ rB.astype(np.float32)
    # print("mid1",mid1_gpu.get())
    # print("cpures ", cpures)


context = make_default_context()
device = context.get_device()
cublas_handle = cublas.cublasCreate()


# test_cublasCgemmBatchedExINT8()
test_cublasCgemmBatchedExC8()
context.pop()