ai-content-maker/.venv/Lib/site-packages/numba/cuda/tests/cudadrv/test_pinned.py

38 lines
944 B
Python

import numpy as np
import platform
from numba import cuda
from numba.cuda.testing import unittest, ContextResettingTestCase
class TestPinned(ContextResettingTestCase):
def _run_copies(self, A):
A0 = np.copy(A)
stream = cuda.stream()
ptr = cuda.to_device(A, copy=False, stream=stream)
ptr.copy_to_device(A, stream=stream)
ptr.copy_to_host(A, stream=stream)
stream.synchronize()
self.assertTrue(np.allclose(A, A0))
def test_pinned(self):
machine = platform.machine()
if machine.startswith('arm') or machine.startswith('aarch64'):
count = 262144 # 2MB
else:
count = 2097152 # 16MB
A = np.arange(count)
with cuda.pinned(A):
self._run_copies(A)
def test_unpinned(self):
A = np.arange(2 * 1024 * 1024) # 16 MB
self._run_copies(A)
if __name__ == '__main__':
unittest.main()