38 lines
944 B
Python
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()
|