MCPcopy Create free account
hub / github.com/apache/arrow / test_context_device_buffer

Function test_context_device_buffer

python/pyarrow/tests/test_cuda.py:123–247  ·  view source on GitHub ↗
(size)

Source from the content-addressed store, hash-verified

121
122@pytest.mark.parametrize("size", [0, 1, 1000])
123def test_context_device_buffer(size):
124 # Creating device buffer from host buffer;
125 arr, buf = make_random_buffer(size)
126 cudabuf = global_context.buffer_from_data(buf)
127 assert cudabuf.size == size
128 arr2 = np.frombuffer(cudabuf.copy_to_host(), dtype=np.uint8)
129 np.testing.assert_equal(arr, arr2)
130
131 # CudaBuffer does not support buffer protocol
132 with pytest.raises(BufferError):
133 memoryview(cudabuf)
134
135 # Creating device buffer from array:
136 cudabuf = global_context.buffer_from_data(arr)
137 assert cudabuf.size == size
138 arr2 = np.frombuffer(cudabuf.copy_to_host(), dtype=np.uint8)
139 np.testing.assert_equal(arr, arr2)
140
141 # Creating device buffer from bytes:
142 cudabuf = global_context.buffer_from_data(arr.tobytes())
143 assert cudabuf.size == size
144 arr2 = np.frombuffer(cudabuf.copy_to_host(), dtype=np.uint8)
145 np.testing.assert_equal(arr, arr2)
146
147 # Creating a device buffer from another device buffer, view:
148 cudabuf2 = cudabuf.slice(0, cudabuf.size)
149 assert cudabuf2.size == size
150 arr2 = np.frombuffer(cudabuf2.copy_to_host(), dtype=np.uint8)
151 np.testing.assert_equal(arr, arr2)
152
153 if size > 1:
154 cudabuf2.copy_from_host(arr[size//2:])
155 arr3 = np.frombuffer(cudabuf.copy_to_host(), dtype=np.uint8)
156 np.testing.assert_equal(np.concatenate((arr[size//2:], arr[size//2:])),
157 arr3)
158 cudabuf2.copy_from_host(arr[:size//2]) # restoring arr
159
160 # Creating a device buffer from another device buffer, copy:
161 cudabuf2 = global_context.buffer_from_data(cudabuf)
162 assert cudabuf2.size == size
163 arr2 = np.frombuffer(cudabuf2.copy_to_host(), dtype=np.uint8)
164 np.testing.assert_equal(arr, arr2)
165
166 cudabuf2.copy_from_host(arr[size//2:])
167 arr3 = np.frombuffer(cudabuf.copy_to_host(), dtype=np.uint8)
168 np.testing.assert_equal(arr, arr3)
169
170 # Slice of a device buffer
171 cudabuf2 = cudabuf.slice(0, cudabuf.size+10)
172 assert cudabuf2.size == size
173 arr2 = np.frombuffer(cudabuf2.copy_to_host(), dtype=np.uint8)
174 np.testing.assert_equal(arr, arr2)
175
176 cudabuf2 = cudabuf.slice(size//4, size+10)
177 assert cudabuf2.size == size - size//4
178 arr2 = np.frombuffer(cudabuf2.copy_to_host(), dtype=np.uint8)
179 np.testing.assert_equal(arr[size//4:], arr2)
180

Callers

nothing calls this directly

Calls 3

make_random_bufferFunction · 0.70
sliceMethod · 0.45
concatenateMethod · 0.45

Tested by

no test coverage detected