65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
"""
|
|
Test basic language features
|
|
|
|
"""
|
|
|
|
import numpy as np
|
|
from numba import cuda, float64
|
|
from numba.cuda.testing import unittest, CUDATestCase
|
|
|
|
|
|
class TestLang(CUDATestCase):
|
|
def test_enumerate(self):
|
|
tup = (1., 2.5, 3.)
|
|
|
|
@cuda.jit("void(float64[:])")
|
|
def foo(a):
|
|
for i, v in enumerate(tup):
|
|
a[i] = v
|
|
|
|
a = np.zeros(len(tup))
|
|
foo[1, 1](a)
|
|
self.assertTrue(np.all(a == tup))
|
|
|
|
def test_zip(self):
|
|
t1 = (1, 2, 3)
|
|
t2 = (4.5, 5.6, 6.7)
|
|
|
|
@cuda.jit("void(float64[:])")
|
|
def foo(a):
|
|
c = 0
|
|
for i, j in zip(t1, t2):
|
|
c += i + j
|
|
a[0] = c
|
|
|
|
a = np.zeros(1)
|
|
foo[1, 1](a)
|
|
b = np.array(t1)
|
|
c = np.array(t2)
|
|
self.assertTrue(np.all(a == (b + c).sum()))
|
|
|
|
def test_issue_872(self):
|
|
'''
|
|
Ensure that typing and lowering of CUDA kernel API primitives works in
|
|
more than one block. Was originally to ensure that macro expansion works
|
|
for more than one block (issue #872), but macro expansion has been
|
|
replaced by a "proper" implementation of all kernel API functions.
|
|
'''
|
|
|
|
@cuda.jit("void(float64[:,:])")
|
|
def cuda_kernel_api_in_multiple_blocks(ary):
|
|
for i in range(2):
|
|
tx = cuda.threadIdx.x
|
|
for j in range(3):
|
|
ty = cuda.threadIdx.y
|
|
sm = cuda.shared.array((2, 3), float64)
|
|
sm[tx, ty] = 1.0
|
|
ary[tx, ty] = sm[tx, ty]
|
|
|
|
a = np.zeros((2, 3))
|
|
cuda_kernel_api_in_multiple_blocks[1, (2, 3)](a)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|