1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
|
cimport openmp
from cython.parallel import parallel, prange
import numpy as np
IF USE_BLIS:
cdef extern from 'cblas.h' nogil:
ctypedef enum CBLAS_ORDER:
CblasRowMajor=101
CblasColMajor=102
ctypedef enum CBLAS_TRANSPOSE:
CblasNoTrans=111
CblasTrans=112
CblasConjTrans=113
void dgemm 'cblas_dgemm' (
CBLAS_ORDER Order, CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB, int M, int N,
int K, double alpha, double *A, int lda,
double *B, int ldb, double beta, double *C, int ldc)
ELSE:
from scipy.linalg.cython_blas cimport dgemm
from threadpoolctl import ThreadpoolController
def check_nested_prange_blas(double[:, ::1] A, double[:, ::1] B, int nthreads):
"""Run multithreaded BLAS calls within OpenMP parallel loop"""
cdef:
int m = A.shape[0]
int n = B.shape[0]
int k = A.shape[1]
double[:, ::1] C = np.empty((m, n))
int n_chunks = 100
int chunk_size = A.shape[0] // n_chunks
char* trans = 't'
char* no_trans = 'n'
double alpha = 1.0
double beta = 0.0
int i
int prange_num_threads
int *prange_num_threads_ptr = &prange_num_threads
inner_info = [None]
with nogil, parallel(num_threads=nthreads):
if openmp.omp_get_thread_num() == 0:
with gil:
inner_info[0] = ThreadpoolController().info()
prange_num_threads_ptr[0] = openmp.omp_get_num_threads()
for i in prange(n_chunks):
IF USE_BLIS:
dgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
chunk_size, n, k, alpha, &A[i * chunk_size, 0], k,
&B[0, 0], k, beta, &C[i * chunk_size, 0], n)
ELSE:
dgemm(trans, no_trans, &n, &chunk_size, &k,
&alpha, &B[0, 0], &k, &A[i * chunk_size, 0], &k,
&beta, &C[i * chunk_size, 0], &n)
return np.asarray(C), prange_num_threads, inner_info[0]
|
