import numpy as np
from numba import cuda
import unittest
from numba.cuda.testing import CUDATestCase
@cuda.jit
def foo(x):
i = cuda.grid(1)
if i < x.size:
x[i] += 1
class TestForAll(CUDATestCase):
def test_forall_1(self):
arr = np.arange(11)
orig = arr.copy()
foo.forall(arr.size)(arr)
np.testing.assert_array_almost_equal(arr, orig + 1)
def test_forall_2(self):
@cuda.jit("void(float32, float32[:], float32[:])")
def bar(a, x, y):
i = cuda.grid(1)
if i < x.size:
y[i] = a * x[i] + y[i]
x = np.arange(13, dtype=np.float32)
y = np.arange(13, dtype=np.float32)
oldy = y.copy()
a = 1.234
bar.forall(y.size)(a, x, y)
np.testing.assert_array_almost_equal(y, a * x + oldy, decimal=3)
def test_forall_no_work(self):
# Ensure that forall doesn't launch a kernel with no blocks when called
# with 0 elements. See Issue #5017.
arr = np.arange(11)
foo.forall(0)(arr)
def test_forall_negative_work(self):
# Ensure that forall doesn't allow the creation of a forall with a
# negative element count.
with self.assertRaises(ValueError) as raises:
foo.forall(-1)
self.assertIn("Can't create ForAll with negative task count",
str(raises.exception))
if __name__ == '__main__':
unittest.main()