Support `ml_dtypes`

[leofang]

@benhg reported in the meeting that passing arrays with dtype from ml_dtypes would blow up the JIT compiler. Would be nice to know what happened.

[benhg]

It's not very exciting. I just get numba.core.errors.NumbaValueError: Unsupported array dtype: bfloat16

Full trace:

(venv) benjaming:~$ python3 ml_dtypes_test.py
<...path>/site-packages/numba_cuda/numba/cuda/dispatcher.py:697: NumbaPerformanceWarning: Grid size 2 will likely result in GPU under-utilization due to low occupancy.
  warn(NumbaPerformanceWarning(msg))
Traceback (most recent call last):
  File "<...path>/site-packages/numba/core/typing/typeof.py", line 264, in _typeof_ndarray
    dtype = numpy_support.from_dtype(val.dtype)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<...path>/site-packages/numba/np/numpy_support.py", line 137, in from_dtype
    raise errors.NumbaNotImplementedError(dtype)
numba.core.errors.NumbaNotImplementedError: bfloat16

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/benjaming/ml_dtypes_test.py", line 23, in <module>
    add_one_kernel[blocks_per_grid, threads_per_block](np_bf16)
  File "<...path>/site-packages/numba_cuda/numba/cuda/dispatcher.py", line 700, in __call__
    return self.dispatcher.call(
           ^^^^^^^^^^^^^^^^^^^^^
  File "<...path>/site-packages/numba_cuda/numba/cuda/dispatcher.py", line 1022, in call
    kernel = _dispatcher.Dispatcher._cuda_call(self, *args)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<...path>/site-packages/numba_cuda/numba/cuda/dispatcher.py", line 1029, in _compile_for_args
    argtypes = [self.typeof_pyval(a) for a in args]
                ^^^^^^^^^^^^^^^^^^^^
  File "<...path>/site-packages/numba_cuda/numba/cuda/dispatcher.py", line 1036, in typeof_pyval
    return typeof(val, Purpose.argument)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<...path>/site-packages/numba/core/typing/typeof.py", line 33, in typeof
    ty = typeof_impl(val, c)
         ^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/functools.py", line 909, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<...path>/site-packages/numba/core/typing/typeof.py", line 266, in _typeof_ndarray
    raise errors.NumbaValueError(f"Unsupported array dtype: {val.dtype}")
numba.core.errors.NumbaValueError: Unsupported array dtype: bfloat16

The code sample is:

import ml_dtypes
import numpy as np
from numba import cuda

# Allocate a CuPy array of BFloat16s using ml-dtypes
size = 16
# Create a NumPy array of bfloat16 using ml_dtypes
np_bf16 = np.arange(size, dtype=ml_dtypes.bfloat16)

# Define a simple Numba-CUDA kernel that adds 1 to each element
@cuda.jit
def add_one_kernel(arr):
    idx = cuda.grid(1)
    if idx < arr.size:
        arr[idx] = arr[idx] + 1

# Launch the kernel
threads_per_block = 8
blocks_per_grid = (size + threads_per_block - 1) // threads_per_block
add_one_kernel[blocks_per_grid, threads_per_block](np_bf16)

# Copy result back to host and print
result = np_bf16
print("Result (bfloat16):", result)
print("Result (float32 for display):", result.astype(np.float32))