docker-python/tests/test_numba.py at main · ZahidSQLDBA/docker-python · GitHub

35 lines (27 loc) · 1.08 KB

import unittest

import numpy as np

from numba import jit, cuda

from common import gpu_test

class TestNumba(unittest.TestCase):

def test_jit(self):

x = np.arange(100).reshape(10, 10)

@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit

def go_fast(a): # Function is compiled to machine code when called the first time

trace = 0.0

for i in range(a.shape[0]): # Numba likes loops

trace += np.tanh(a[i, i]) # Numba likes NumPy functions

return a + trace # Numba likes NumPy broadcasting

self.assertEqual(10, go_fast(x).shape[0])

@gpu_test

def test_cuda_jit(self):

x = np.arange(10)

@cuda.jit

def increment_by_one(an_array):

pos = cuda.grid(1)

if pos < an_array.size:

an_array[pos] += 1

threadsperblock = 32

blockspergrid = (x.size + (threadsperblock - 1))

self.assertEqual(0, x[0])

increment_by_one[blockspergrid, threadsperblock](x)

self.assertEqual(1, x[0])

Expand file tree