ai-content-maker/.venv/Lib/site-packages/blis/cy.pyx

642 lines
16 KiB
Cython
Raw Normal View History

2024-05-03 04:18:51 +03:00
# cython: infer_types=True
# cython: boundscheck=False
# Copyright ExplsionAI GmbH, released under BSD.
import atexit
cdef extern from "blis.h" nogil:
enum blis_err_t "err_t":
pass
cdef struct blis_cntx_t "cntx_t":
pass
cdef struct blis_rntm_t "rntm_s":
pass
ctypedef enum blis_trans_t "trans_t":
BLIS_NO_TRANSPOSE
BLIS_TRANSPOSE
BLIS_CONJ_NO_TRANSPOSE
BLIS_CONJ_TRANSPOSE
ctypedef enum blis_conj_t "conj_t":
BLIS_NO_CONJUGATE
BLIS_CONJUGATE
ctypedef enum blis_side_t "side_t":
BLIS_LEFT
BLIS_RIGHT
ctypedef enum blis_uplo_t "uplo_t":
BLIS_LOWER
BLIS_UPPER
BLIS_DENSE
ctypedef enum blis_diag_t "diag_t":
BLIS_NONUNIT_DIAG
BLIS_UNIT_DIAG
char* bli_info_get_int_type_size_str()
blis_err_t bli_init()
blis_err_t bli_finalize()
blis_err_t bli_rntm_init(blis_rntm_t* rntm);
# BLAS level 3 routines
void bli_dgemm_ex(
blis_trans_t transa,
blis_trans_t transb,
dim_t m,
dim_t n,
dim_t k,
double* alpha,
double* a, inc_t rsa, inc_t csa,
double* b, inc_t rsb, inc_t csb,
double* beta,
double* c, inc_t rsc, inc_t csc,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
# BLAS level 3 routines
void bli_sgemm_ex(
blis_trans_t transa,
blis_trans_t transb,
dim_t m,
dim_t n,
dim_t k,
float* alpha,
float* a, inc_t rsa, inc_t csa,
float* b, inc_t rsb, inc_t csb,
float* beta,
float* c, inc_t rsc, inc_t csc,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_dger_ex(
blis_conj_t conjx,
blis_conj_t conjy,
dim_t m,
dim_t n,
double* alpha,
double* x, inc_t incx,
double* y, inc_t incy,
double* a, inc_t rsa, inc_t csa,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_sger_ex(
blis_conj_t conjx,
blis_conj_t conjy,
dim_t m,
dim_t n,
float* alpha,
float* x, inc_t incx,
float* y, inc_t incy,
float* a, inc_t rsa, inc_t csa,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_dgemv_ex(
blis_trans_t transa,
blis_conj_t conjx,
dim_t m,
dim_t n,
double* alpha,
double* a, inc_t rsa, inc_t csa,
double* x, inc_t incx,
double* beta,
double* y, inc_t incy,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_sgemv_ex(
blis_trans_t transa,
blis_conj_t conjx,
dim_t m,
dim_t n,
float* alpha,
float* a, inc_t rsa, inc_t csa,
float* x, inc_t incx,
float* beta,
float* y, inc_t incy,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_daxpyv_ex(
blis_conj_t conjx,
dim_t m,
double* alpha,
double* x, inc_t incx,
double* y, inc_t incy,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_saxpyv_ex(
blis_conj_t conjx,
dim_t m,
float* alpha,
float* x, inc_t incx,
float* y, inc_t incy,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_dscalv_ex(
blis_conj_t conjalpha,
dim_t m,
double* alpha,
double* x, inc_t incx,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_sscalv_ex(
blis_conj_t conjalpha,
dim_t m,
float* alpha,
float* x, inc_t incx,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_ddotv_ex(
blis_conj_t conjx,
blis_conj_t conjy,
dim_t m,
double* x, inc_t incx,
double* y, inc_t incy,
double* rho,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_sdotv_ex(
blis_conj_t conjx,
blis_conj_t conjy,
dim_t m,
float* x, inc_t incx,
float* y, inc_t incy,
float* rho,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_snorm1v_ex(
dim_t n,
float* x, inc_t incx,
float* norm,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_dnorm1v_ex(
dim_t n,
double* x, inc_t incx,
double* norm,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_snormfv_ex(
dim_t n,
float* x, inc_t incx,
float* norm,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_dnormfv_ex(
dim_t n,
double* x, inc_t incx,
double* norm,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_snormiv_ex(
dim_t n,
float* x, inc_t incx,
float* norm,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_dnormiv_ex(
dim_t n,
double* x, inc_t incx,
double* norm,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_srandv_ex(
dim_t m,
float* x, inc_t incx,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_drandv_ex(
dim_t m,
double* x, inc_t incx,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
)
void bli_ssumsqv_ex(
dim_t m,
float* x, inc_t incx,
float* scale,
float* sumsq,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
) nogil
void bli_dsumsqv_ex(
dim_t m,
double* x, inc_t incx,
double* scale,
double* sumsq,
blis_cntx_t* cntx,
blis_rntm_t* rntm,
) nogil
bli_init()
cdef blis_rntm_t rntm;
def init():
bli_init()
bli_rntm_init(&rntm);
assert BLIS_NO_TRANSPOSE == <blis_trans_t>NO_TRANSPOSE
assert BLIS_TRANSPOSE == <blis_trans_t>TRANSPOSE
assert BLIS_CONJ_NO_TRANSPOSE == <blis_trans_t>CONJ_NO_TRANSPOSE
assert BLIS_CONJ_TRANSPOSE == <blis_trans_t>CONJ_TRANSPOSE
assert BLIS_NO_CONJUGATE == <blis_conj_t>NO_CONJUGATE
assert BLIS_CONJUGATE == <blis_conj_t>CONJUGATE
assert BLIS_LEFT == <blis_side_t>LEFT
assert BLIS_RIGHT == <blis_side_t>RIGHT
assert BLIS_LOWER == <blis_uplo_t>LOWER
assert BLIS_UPPER == <blis_uplo_t>UPPER
assert BLIS_DENSE == <blis_uplo_t>DENSE
assert BLIS_NONUNIT_DIAG == <blis_diag_t>NONUNIT_DIAG
assert BLIS_UNIT_DIAG == <blis_diag_t>UNIT_DIAG
def get_int_type_size():
cdef char* int_size = bli_info_get_int_type_size_str()
return '%d' % int_size[0]
# BLAS level 3 routines
cdef void gemm(
trans_t trans_a,
trans_t trans_b,
dim_t m,
dim_t n,
dim_t k,
double alpha,
reals_ft a, inc_t rsa, inc_t csa,
reals_ft b, inc_t rsb, inc_t csb,
double beta,
reals_ft c, inc_t rsc, inc_t csc
) nogil:
cdef float alpha_f = alpha
cdef float beta_f = beta
cdef double alpha_d = alpha
cdef double beta_d = beta
if reals_ft is floats_t:
bli_sgemm_ex(
<blis_trans_t>trans_a, <blis_trans_t>trans_b,
m, n, k,
&alpha_f, a, rsa, csa, b, rsb, csb, &beta_f, c, rsc, csc, NULL, &rntm)
elif reals_ft is doubles_t:
bli_dgemm_ex(
<blis_trans_t>trans_a, <blis_trans_t>trans_b,
m, n, k,
&alpha_d, a, rsa, csa, b, rsb, csb, &beta_d, c, rsc, csc, NULL, &rntm)
elif reals_ft is float1d_t:
bli_sgemm_ex(
<blis_trans_t>trans_a, <blis_trans_t>trans_b,
m, n, k,
&alpha_f, &a[0], rsa, csa, &b[0], rsb, csb, &beta_f, &c[0],
rsc, csc, NULL, &rntm)
elif reals_ft is double1d_t:
bli_dgemm_ex(
<blis_trans_t>trans_a, <blis_trans_t>trans_b,
m, n, k,
&alpha_d, &a[0], rsa, csa, &b[0], rsb, csb, &beta_d, &c[0],
rsc, csc, NULL, &rntm)
else:
# Impossible --- panic?
pass
cdef void ger(
conj_t conjx,
conj_t conjy,
dim_t m,
dim_t n,
double alpha,
reals_ft x, inc_t incx,
reals_ft y, inc_t incy,
reals_ft a, inc_t rsa, inc_t csa
) nogil:
cdef float alpha_f = alpha
cdef double alpha_d = alpha
if reals_ft is floats_t:
bli_sger_ex(
<blis_conj_t>conjx, <blis_conj_t>conjy,
m, n,
&alpha_f,
x, incx, y, incy, a, rsa, csa, NULL, &rntm)
elif reals_ft is doubles_t:
bli_dger_ex(
<blis_conj_t>conjx, <blis_conj_t>conjy,
m, n,
&alpha_d,
x, incx, y, incy, a, rsa, csa, NULL, &rntm)
elif reals_ft is float1d_t:
bli_sger_ex(
<blis_conj_t>conjx, <blis_conj_t>conjy,
m, n,
&alpha_f,
&x[0], incx, &y[0], incy, &a[0], rsa, csa, NULL, &rntm)
elif reals_ft is double1d_t:
bli_dger_ex(
<blis_conj_t>conjx, <blis_conj_t>conjy,
m, n,
&alpha_d,
&x[0], incx, &y[0], incy, &a[0], rsa, csa, NULL, &rntm)
else:
# Impossible --- panic?
pass
cdef void gemv(
trans_t transa,
conj_t conjx,
dim_t m,
dim_t n,
real_ft alpha,
reals_ft a, inc_t rsa, inc_t csa,
reals_ft x, inc_t incx,
real_ft beta,
reals_ft y, inc_t incy
) nogil:
cdef float alpha_f = alpha
cdef double alpha_d = alpha
cdef float beta_f = alpha
cdef double beta_d = alpha
if reals_ft is floats_t:
bli_sgemv_ex(
<blis_trans_t>transa, <blis_conj_t>conjx,
m, n,
&alpha_f, a, rsa, csa,
x, incx, &beta_f,
y, incy, NULL, &rntm)
elif reals_ft is doubles_t:
bli_dgemv_ex(
<blis_trans_t>transa, <blis_conj_t>conjx,
m, n,
&alpha_d, a, rsa, csa,
x, incx, &beta_d,
y, incy, NULL, &rntm)
elif reals_ft is float1d_t:
bli_sgemv_ex(
<blis_trans_t>transa, <blis_conj_t>conjx,
m, n,
&alpha_f, &a[0], rsa, csa,
&x[0], incx, &beta_f,
&y[0], incy, NULL, &rntm)
elif reals_ft is double1d_t:
bli_dgemv_ex(
<blis_trans_t>transa, <blis_conj_t>conjx,
m, n,
&alpha_d, &a[0], rsa, csa,
&x[0], incx, &beta_d,
&y[0], incy, NULL, &rntm)
else:
# Impossible --- panic?
pass
cdef void axpyv(
conj_t conjx,
dim_t m,
real_ft alpha,
reals_ft x, inc_t incx,
reals_ft y, inc_t incy
) nogil:
cdef float alpha_f = alpha
cdef double alpha_d = alpha
if reals_ft is floats_t:
bli_saxpyv_ex(<blis_conj_t>conjx, m, &alpha_f, x, incx, y, incy, NULL, &rntm)
elif reals_ft is doubles_t:
bli_daxpyv_ex(<blis_conj_t>conjx, m, &alpha_d, x, incx, y, incy, NULL, &rntm)
elif reals_ft is float1d_t:
bli_saxpyv_ex(<blis_conj_t>conjx, m, &alpha_f, &x[0], incx, &y[0], incy, NULL, &rntm)
elif reals_ft is double1d_t:
bli_daxpyv_ex(<blis_conj_t>conjx, m, &alpha_d, &x[0], incx, &y[0], incy, NULL, &rntm)
else:
# Impossible --- panic?
pass
cdef void scalv(
conj_t conjalpha,
dim_t m,
real_ft alpha,
reals_ft x, inc_t incx
) nogil:
cdef float alpha_f = alpha
cdef double alpha_d = alpha
if reals_ft is floats_t:
bli_sscalv_ex(<blis_conj_t>conjalpha, m, &alpha_f, x, incx, NULL, &rntm)
elif reals_ft is doubles_t:
bli_dscalv_ex(<blis_conj_t>conjalpha, m, &alpha_d, x, incx, NULL, &rntm)
elif reals_ft is float1d_t:
bli_sscalv_ex(<blis_conj_t>conjalpha, m, &alpha_f, &x[0], incx, NULL, &rntm)
elif reals_ft is double1d_t:
bli_dscalv_ex(<blis_conj_t>conjalpha, m, &alpha_d, &x[0], incx, NULL, &rntm)
else:
# Impossible --- panic?
pass
cdef double norm_L1(
dim_t n,
reals_ft x, inc_t incx
) nogil:
cdef double dnorm = 0
cdef float snorm = 0
if reals_ft is floats_t:
bli_snorm1v_ex(n, x, incx, &snorm, NULL, &rntm)
dnorm = snorm
elif reals_ft is doubles_t:
bli_dnorm1v_ex(n, x, incx, &dnorm, NULL, &rntm)
elif reals_ft is float1d_t:
bli_snorm1v_ex(n, &x[0], incx, &snorm, NULL, &rntm)
dnorm = snorm
elif reals_ft is double1d_t:
bli_dnorm1v_ex(n, &x[0], incx, &dnorm, NULL, &rntm)
else:
# Impossible --- panic?
pass
return dnorm
cdef double norm_L2(
dim_t n,
reals_ft x, inc_t incx
) nogil:
cdef double dnorm = 0
cdef float snorm = 0
if reals_ft is floats_t:
bli_snormfv_ex(n, x, incx, &snorm, NULL, &rntm)
dnorm = snorm
elif reals_ft is doubles_t:
bli_dnormfv_ex(n, x, incx, &dnorm, NULL, &rntm)
elif reals_ft is float1d_t:
bli_snormfv_ex(n, &x[0], incx, &snorm, NULL, &rntm)
dnorm = snorm
elif reals_ft is double1d_t:
bli_dnormfv_ex(n, &x[0], incx, &dnorm, NULL, &rntm)
else:
# Impossible --- panic?
pass
return dnorm
cdef double norm_inf(
dim_t n,
reals_ft x, inc_t incx
) nogil:
cdef double dnorm = 0
cdef float snorm = 0
if reals_ft is floats_t:
bli_snormiv_ex(n, x, incx, &snorm, NULL, &rntm)
dnorm = snorm
elif reals_ft is doubles_t:
bli_dnormiv_ex(n, x, incx, &dnorm, NULL, &rntm)
elif reals_ft is float1d_t:
bli_snormiv_ex(n, &x[0], incx, &snorm, NULL, &rntm)
dnorm = snorm
elif reals_ft is double1d_t:
bli_dnormiv_ex(n, &x[0], incx, &dnorm, NULL, &rntm)
else:
# Impossible --- panic?
pass
return dnorm
cdef double dotv(
conj_t conjx,
conj_t conjy,
dim_t m,
reals_ft x,
reals_ft y,
inc_t incx,
inc_t incy,
) nogil:
cdef double rho_d = 0.0
cdef float rho_f = 0.0
if reals_ft is floats_t:
bli_sdotv_ex(<blis_conj_t>conjx, <blis_conj_t>conjy, m, x, incx, y, incy, &rho_f, NULL, &rntm)
return rho_f
elif reals_ft is doubles_t:
bli_ddotv_ex(<blis_conj_t>conjx, <blis_conj_t>conjy, m, x, incx, y, incy, &rho_d, NULL, &rntm)
return rho_d
elif reals_ft is float1d_t:
bli_sdotv_ex(<blis_conj_t>conjx, <blis_conj_t>conjy, m, &x[0], incx, &y[0], incy,
&rho_f, NULL, &rntm)
return rho_f
elif reals_ft is double1d_t:
bli_ddotv_ex(<blis_conj_t>conjx, <blis_conj_t>conjy, m, &x[0], incx, &y[0], incy,
&rho_d, NULL, &rntm)
return rho_d
else:
raise ValueError("Unhandled fused type")
cdef void randv(dim_t m, reals_ft x, inc_t incx) nogil:
if reals_ft is floats_t:
bli_srandv_ex(m, x, incx, NULL, &rntm)
elif reals_ft is float1d_t:
bli_srandv_ex(m, &x[0], incx, NULL, &rntm)
if reals_ft is doubles_t:
bli_drandv_ex(m, x, incx, NULL, &rntm)
elif reals_ft is double1d_t:
bli_drandv_ex(m, &x[0], incx, NULL, &rntm)
else:
with gil:
raise ValueError("Unhandled fused type")
cdef void sumsqv(dim_t m, reals_ft x, inc_t incx,
reals_ft scale, reals_ft sumsq) nogil:
if reals_ft is floats_t:
bli_ssumsqv_ex(m, &x[0], incx, scale, sumsq, NULL, &rntm)
elif reals_ft is float1d_t:
bli_ssumsqv_ex(m, &x[0], incx, &scale[0], &sumsq[0], NULL, &rntm)
if reals_ft is doubles_t:
bli_dsumsqv_ex(m, x, incx, scale, sumsq, NULL, &rntm)
elif reals_ft is double1d_t:
bli_dsumsqv_ex(m, &x[0], incx, &scale[0], &sumsq[0], NULL, &rntm)
else:
with gil:
raise ValueError("Unhandled fused type")
cdef void dgemm(bint transA, bint transB, int M, int N, int K,
double alpha, const double* A, int lda, const double* B,
int ldb, double beta, double* C, int ldc) nogil:
gemm(
TRANSPOSE if transA else NO_TRANSPOSE,
TRANSPOSE if transB else NO_TRANSPOSE,
M, N, K,
alpha, A, lda, 1,
B, ldb, 1,
beta, C, ldc, 1
)
cdef void sgemm(bint transA, bint transB, int M, int N, int K,
float alpha, const float* A, int lda, const float* B,
int ldb, float beta, float* C, int ldc) nogil:
gemm(
TRANSPOSE if transA else NO_TRANSPOSE,
TRANSPOSE if transB else NO_TRANSPOSE,
M, N, K,
alpha, A, lda, 1,
B, ldb, 1,
beta, C, ldc, 1
)
cdef void saxpy(int N, float alpha, const float* X, int incX,
float* Y, int incY) nogil:
axpyv(NO_CONJUGATE, N, alpha, X, incX, Y, incY)
cdef void daxpy(int N, double alpha, const double* X, int incX,
double* Y, int incY) nogil:
axpyv(NO_CONJUGATE, N, alpha, X, incX, Y, incY)
@atexit.register
def finalize():
bli_finalize()