ai-content-maker/.venv/Lib/site-packages/numba/cuda/tests/cudapy/test_lang.py

65 lines
1.7 KiB
Python
Raw Permalink Normal View History

2024-05-03 04:18:51 +03:00
"""
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()