Skip to content

Commit 835845c

Browse files
authored
fix: empty array type mismatch between host and device (#612)
Fixes #483. As mentioned [here](numba/numba#9654 (comment)), the bug was due to: - empty arrays incorrectly identified as broadcast arrays - incorrect contiguity flags for empty arrays When running this: ``` from numba import cuda, typeof import numpy as np h_values = np.random.randint(low=0, high=2, size=(10, 0)) d_values = cuda.to_device(h_values) print(typeof(h_values)) print(typeof(d_values)) assert typeof(h_values) == typeof(d_values) ``` Below is the output: ``` array(int64, 2d, C) array(int64, 2d, C) ```
1 parent b06e183 commit 835845c

File tree

3 files changed

+67
-1
lines changed

3 files changed

+67
-1
lines changed

numba_cuda/numba/cuda/cudadrv/devicearray.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,8 @@ def _numba_type_(self):
178178
# of which will be 0, will not match those hardcoded in for 'C' or 'F'
179179
# layouts.
180180

181-
broadcast = 0 in self.strides
181+
broadcast = 0 in self.strides and (self.size != 0)
182+
182183
if self.flags["C_CONTIGUOUS"] and not broadcast:
183184
layout = "C"
184185
elif self.flags["F_CONTIGUOUS"] and not broadcast:

numba_cuda/numba/cuda/cudadrv/dummyarray.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,10 @@ def _compute_layout(self):
279279
if not self.dims:
280280
return {"C_CONTIGUOUS": True, "F_CONTIGUOUS": True}
281281

282+
# All 0-size arrays are considered contiguous, even if they are multidimensional
283+
if self.size == 0:
284+
return {"C_CONTIGUOUS": True, "F_CONTIGUOUS": True}
285+
282286
# If this is a broadcast array then it is not contiguous
283287
if any([dim.stride == 0 for dim in self.dims]):
284288
return {"C_CONTIGUOUS": False, "F_CONTIGUOUS": False}

numba_cuda/numba/cuda/tests/nocuda/test_dummyarray.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -387,5 +387,66 @@ def test_for_loop(self):
387387
x = val # noqa: F841
388388

389389

390+
@skip_on_cudasim("Tests internals of the CUDA driver device array")
391+
class TestEmptyArrays(unittest.TestCase):
392+
def test_empty_array_flags(self):
393+
test_shapes = [
394+
(0,),
395+
(10, 0),
396+
(0, 10),
397+
(0, 0),
398+
(5, 0, 3),
399+
(0, 5, 3),
400+
(5, 3, 0),
401+
(0, 0, 0),
402+
]
403+
for shape in test_shapes:
404+
with self.subTest(shape=shape):
405+
nparr = np.empty(shape)
406+
arr = Array.from_desc(
407+
0, nparr.shape, nparr.strides, nparr.dtype.itemsize
408+
)
409+
# Empty arrays should be both C and F contiguous
410+
self.assertEqual(
411+
arr.flags["C_CONTIGUOUS"],
412+
nparr.flags["C_CONTIGUOUS"],
413+
f"C_CONTIGUOUS mismatch for shape {shape}",
414+
)
415+
self.assertEqual(
416+
arr.flags["F_CONTIGUOUS"],
417+
nparr.flags["F_CONTIGUOUS"],
418+
f"F_CONTIGUOUS mismatch for shape {shape}",
419+
)
420+
self.assertTrue(arr.flags["C_CONTIGUOUS"])
421+
self.assertTrue(arr.flags["F_CONTIGUOUS"])
422+
423+
424+
@skip_on_cudasim("Tests CUDA device array type inference")
425+
class TestEmptyArrayTypeInference(unittest.TestCase):
426+
def test_empty_array_typeof(self):
427+
from numba import cuda, typeof
428+
429+
test_cases = [
430+
((0,), np.int64),
431+
((10, 0), np.int64),
432+
((0, 10), np.int64),
433+
((0, 0), np.float32),
434+
((5, 0, 3), np.float32),
435+
((0, 5, 3), np.int32),
436+
((5, 3, 0), np.float64),
437+
]
438+
439+
for shape, dtype in test_cases:
440+
with self.subTest(shape=shape, dtype=dtype):
441+
h_values = np.empty(shape, dtype=dtype)
442+
d_values = cuda.to_device(h_values)
443+
self.assertEqual(
444+
typeof(h_values),
445+
typeof(d_values),
446+
f"Type mismatch for shape {shape}, dtype {dtype}: "
447+
f"host={typeof(h_values)}, device={typeof(d_values)}",
448+
)
449+
450+
390451
if __name__ == "__main__":
391452
unittest.main()

0 commit comments

Comments
 (0)