python-benchmarks/gemm/gemm_cython.pyx at 79d41430999764625e849f891976e55134e0c61f · numfocus/python-benchmarks · GitHub

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

# Authors: Jake Vanderplas, Olivier Grisel
# License: MIT

import numpy as np
cimport cython
from libc.math cimport sqrt

@cython.boundscheck(False)
@cython.wraparound(False)
def gemm_cython_for_loops(
        double alpha,
        double[:, ::1] A,
        double[:, ::1] B,
        double beta,
        double[:, ::1] C,
        ):
    cdef int M = C.shape[0]
    cdef int N = C.shape[1]
    cdef int K = A.shape[1]
    cdef double tmp, d
    for i in range(M):
        for j in range(N):
            d = 0.0
            for k in range(K):
                d += A[i, k] * B[k, j]
            C[i, j] = alpha * d + beta * C[i, j]


benchmarks = (
    gemm_cython_for_loops,
)