Repository: numba/numba

Issue:
Extend `__cuda_array_interface__` to support masked arrays
Currently the `__cuda_array_interface__` assumes the array is a single dense buffer. This unfortunately doesn't support nulls which is why there have been other array formats such as numpy masked arrays (https://docs.scipy.org/doc/numpy/reference/maskedarray.generic.html) and apache arrow arrays (https://arrow.apache.org/docs/memory_layout.html). The common theme is that there's a secondary buffer of booleans represented by either a single bit per value or a byte per value to indicate whether the value in the data buffer is valid or not.

Projects looking to integrate into the GPU PyData ecosystem have started using `__cuda_array_interface__` as the standard attribute to look for to determine whether the object is on the GPU or not, so being unable to support it for masked arrays has been problematic.

Continuing following the numpy `__array_interface__`, I'm proposing extending the `__cuda_array_interface__` to add the optional `mask` attribute.

We may want to consider not making the `mask` attribute a Python object and instead making it a nested dictionary which would they open the doors to considering making this cross language with something like JSON or a simple struct in the future.


Discussion:
Adding mask seems a bit special-case-y to me.  I wonder if there is some way to generalize this to anticipate future use cases.  Are there other things like mask that would motivate generalization here?  Categories, dataframes, collections of arrays generally (dataframe or xarray style).  I wonder if this motivates a `__cuda_arrays_interface__` protocol instead (note the `s`).
> Adding mask seems a bit special-case-y to me. I wonder if there is some way to generalize this to anticipate future use cases. Are there other things like mask that would motivate generalization here? Categories, dataframes, collections of arrays generally (dataframe or xarray style). I wonder if this motivates a `__cuda_arrays_interface__` protocol instead (note the `s`).

A categorical column would essentially be two arrays, one for the categories without a mask, and one for the codes into the categories potentially with a mask if there's nulls, each of which would expose a `__cuda_array_interface__` that could be effectively used. The mask isn't special-case-y in my mind because it's needed to prevent reading garbage data from the array which you otherwise can't prevent without some form of mask.
We talked through this a bit more in the Numba meeting today, and we think it is probably going out of scope for this interface to expand to include xarray/dataframe-like things that go beyond the ndarray-like containers CuPy/PyTorch/etc are providing.  A different protocol for dealing with those things is probably also useful, but would need much more thinking and/or leaning on Arrow's metadata itself.

The question we have is what implementations are there for a masked array on the GPU (neither PyTorch and CuPy offer masked arrays, AFAIK), and are these optional attributes (mask and mask_dtype = bit/byte) sufficient for them?  I also imagine the expected Python API for working with 
 a masked array on the device in Numba would be to have a .mask attribute that can also be indexed just like the array itself?

Also, as specified, the `__cuda_array_interface__` is a dictionary of simple JSON-encodable types (integers, tuple/list, string, bool) so there is nothing stopping a Python wrapper from getting it ultimately from a lower level library in JSON form and exporting it to Python as a dictionary.  Note that I'm assuming the `mask` attribute is in the same form as `data`: `(device_pointer_as_integer, read_only_flag_as_bool)`
Thinking about this more, the other question is "Could/should the mask be treated as a separate array?"  I think the obstacle there is that the current `__cuda_array_interface__` doesn't have a way to represent a packed array of 1-bit elements, so one of the mask formats is not representable.
FWIW one of the issues NumPy masked arrays ran into was premature optimization. This isn’t to say it wasn’t justified in some sense. Just that it’s caused a lot of pain since. In particular I’m thinking of how the trivial/non-existent mask is handled. Packed bits has a similar feeling to me (though I’m predisposed to worrying about things like this). Though it does make me wonder how far we can get without that feature. Thoughts?
@seibert for clarification I'm not proposing extending Numba functionality to support masked types, I'm only proposing extending the interface with the mask attribute.

From Numba's perspective it would be easy enough to check if the mask exists and isn't `None` to raise an exception that masked arrays aren't supported. Other libraries without masked array support could do the same. For libraries like cuDF and Arrow this would allow using the interface directly instead of having to roll something else combining multiple of these interfaces.

We actually don't need the `mask_dtype` attribute if we follow the `__array_interface__` spec and have the `mask` attribute be a Python object that exposes its own `__cuda_array_interface__`.
I'd rather not nest an arbitrary object in the interface dict, which should be only basic Python types.  If we want to embed the description of the mask into the array description, it should be another copy of the `__cuda_array_interface__` structure, which would also allow specification of the striding of the mask, which is likely necessary if the the array data itself is strided.

A more immediate issue is that there is no typestr for a bit array (there are only bit fields), so describing an Arrow style mask array with `__cuda_array_interface__` is not currently possible.  We would also need to establish a convention for that too.
@seibert I'm on board with nesting another dictionary for the `mask` attribute which follows the same interface. The numpy array interface allows specifying a bit field with the typestr `t` with "Bit field (following integer gives the number of bits in the bit field)." which should satisfy our use case with `t1` for a bitmask, whereas a bytemask can use `b1`.



OK, Siu pointed out that I had overlooked the definition of `mask` in the original `__array_interface__`, which does exactly what you proposed earlier.  In the interest of consistency with the NumPy interface, I'm reversing my suggestion and going with yours.  `mask` should be an object that exports the`__cuda_array_interface__`, rather than the dictionary itself.

Thanks for bearing with me on this discussion!  Feel free to open a PR to modify the documentation of `__cuda_array_interface__`.
> OK, Siu pointed out that I had overlooked the definition of `mask` in the original `__array_interface__`, which does exactly what you proposed earlier. In the interest of consistency with the NumPy interface, I'm reversing my suggestion and going with yours. `mask` should be an object that exports the`__cuda_array_interface__`, rather than the dictionary itself.
> 
> Thanks for bearing with me on this discussion! Feel free to open a PR to modify the documentation of `__cuda_array_interface__`.

Will do, as far as handling objects that specify the mask in Numba would you want to raise a ValueError type exception or return a tuple of two arrays? Either way will kick off the PR shortly and can continue discussion there!

Resolving patch:
diff --git a/numba/cuda/api.py b/numba/cuda/api.py
--- a/numba/cuda/api.py
+++ b/numba/cuda/api.py
@@ -30,6 +30,13 @@ def from_cuda_array_interface(desc, owner=None):
     The *owner* is the owner of the underlying memory.
     The resulting DeviceNDArray will acquire a reference from it.
     """
+    version = desc.get('version')
+    if version == 1:
+        mask = desc.get('mask')
+        # Would ideally be better to detect if the mask is all valid
+        if mask is not None:
+            raise NotImplementedError('Masked arrays are not supported')
+
     shape = desc['shape']
     strides = desc.get('strides')
     dtype = np.dtype(desc['typestr'])
diff --git a/numba/cuda/cudadrv/devicearray.py b/numba/cuda/cudadrv/devicearray.py
--- a/numba/cuda/cudadrv/devicearray.py
+++ b/numba/cuda/cudadrv/devicearray.py
@@ -120,7 +120,7 @@ def __cuda_array_interface__(self):
             'strides': tuple(self.strides),
             'data': (self.device_ctypes_pointer.value, False),
             'typestr': self.dtype.str,
-            'version': 0,
+            'version': 1,
         }
 
     def bind(self, stream=0):


Test patch:
diff --git a/numba/cuda/tests/cudapy/test_cuda_array_interface.py b/numba/cuda/tests/cudapy/test_cuda_array_interface.py
--- a/numba/cuda/tests/cudapy/test_cuda_array_interface.py
+++ b/numba/cuda/tests/cudapy/test_cuda_array_interface.py
@@ -205,6 +205,21 @@ def test_negative_strided_copy_to_host(self):
         expected_msg = 'D->H copy not implemented for negative strides'
         self.assertIn(expected_msg, str(raises.exception))
 
+    def test_masked_array(self):
+        h_arr = np.random.random(10)
+        h_mask = np.random.randint(2, size=10, dtype='bool')
+        c_arr = cuda.to_device(h_arr)
+        c_mask = cuda.to_device(h_mask)
+
+        # Manually create a masked CUDA Array Interface dictionary
+        masked_cuda_array_interface = c_arr.__cuda_array_interface__.copy()
+        masked_cuda_array_interface['mask'] = c_mask
+
+        with self.assertRaises(NotImplementedError) as raises:
+            cuda.from_cuda_array_interface(masked_cuda_array_interface)
+        expected_msg = 'Masked arrays are not supported'
+        self.assertIn(expected_msg, str(raises.exception))
+
 
 if __name__ == "__main__":
     unittest.main()