Add support for ml_dtypes to python runtime bindings (#21549)

Python runtime bindings should include support for atypical dtypes
provided by `ml_dtypes`. We also drop the requirement for `PYBUF_FORMAT`
as we never look at the format structure so this causes failures when
handling the atypical dtypes.

---------

Signed-off-by: Rob Suderman <[email protected]>

Author: rsuderman

runtime/bindings/python/hal.cc

@@ -180,7 +180,7 @@ py::object HalAllocator::AllocateBufferCopy(
   // only request that via PyBUF_ND. Long term, we should consult an
   // "oracle" in the runtime to determine the precise required format
   // and set flags accordingly (and fallback/copy on failure).
-  int flags = PyBUF_FORMAT | PyBUF_ND;
+  int flags = PyBUF_C_CONTIGUOUS | PyBUF_ND;
 
   // Acquire the backing buffer and setup RAII release.
   if (PyObject_GetBuffer(buffer.ptr(), &py_view, flags) != 0) {
@@ -243,7 +243,7 @@ HalBuffer HalAllocator::AllocateHostStagingBufferCopy(HalDevice& device,
   // only request that via PyBUF_ND. Long term, we should consult an
   // "oracle" in the runtime to determine the precise required format
   // and set flags accordingly (and fallback/copy on failure).
-  int flags = PyBUF_FORMAT | PyBUF_ND;
+  int flags = PyBUF_C_CONTIGUOUS | PyBUF_ND;
 
   // Acquire the backing buffer and setup RAII release.
   if (PyObject_GetBuffer(buffer.ptr(), &py_view, flags) != 0) {
@@ -2025,7 +2025,7 @@ void SetupHalBindings(nanobind::module_ m) {
              py::handle pattern, iree_device_size_t target_offset,
              std::optional<iree_device_size_t> length, bool end) {
             Py_buffer pattern_view;
-            int flags = PyBUF_FORMAT | PyBUF_ND;
+            int flags = PyBUF_C_CONTIGUOUS | PyBUF_ND;
             if (PyObject_GetBuffer(pattern.ptr(), &pattern_view, flags) != 0) {
               // The GetBuffer call is required to set an appropriate error.
               throw py::python_error();

runtime/bindings/python/iree/runtime/_binding.pyi

@@ -299,10 +299,11 @@ class HalElementType:
     BOOL_8: ClassVar[HalElementType] = ...
     COMPLEX_128: ClassVar[HalElementType] = ...
     COMPLEX_64: ClassVar[HalElementType] = ...
-    FLOAT_8_E4M3: ClassVar[HalElementType] = ...
+    FLOAT_8_E4M3_FN: ClassVar[HalElementType] = ...
     FLOAT_8_E4M3_FNUZ: ClassVar[HalElementType] = ...
     FLOAT_8_E5M2: ClassVar[HalElementType] = ...
     FLOAT_8_E5M2_FNUZ: ClassVar[HalElementType] = ...
+    FLOAT_8_E8M0_FNU: ClassVar[HalElementType] = ...
     FLOAT_16: ClassVar[HalElementType] = ...
     FLOAT_32: ClassVar[HalElementType] = ...
     FLOAT_64: ClassVar[HalElementType] = ...

runtime/bindings/python/iree/runtime/array_interop.py

@@ -7,6 +7,7 @@
 
 from typing import Optional, Tuple
 import logging
+import ml_dtypes
 import numpy as np
 import numpy.lib.mixins
 
@@ -305,6 +306,12 @@ def asdevicearray(
     (np.bool_, HalElementType.BOOL_8),
     (np.complex64, HalElementType.COMPLEX_64),
     (np.complex128, HalElementType.COMPLEX_128),
+    (ml_dtypes.bfloat16, HalElementType.BFLOAT_16),
+    (ml_dtypes.float8_e4m3fn, HalElementType.FLOAT_8_E4M3_FN),
+    (ml_dtypes.float8_e4m3fnuz, HalElementType.FLOAT_8_E4M3_FNUZ),
+    (ml_dtypes.float8_e5m2, HalElementType.FLOAT_8_E5M2),
+    (ml_dtypes.float8_e5m2fnuz, HalElementType.FLOAT_8_E5M2_FNUZ),
+    (ml_dtypes.float8_e8m0fnu, HalElementType.FLOAT_8_E8M0_FNU),
 )
 
 

runtime/bindings/python/iree/runtime/build_requirements.txt

@@ -9,3 +9,4 @@ numpy>=2.0.0b1
 requests>=2.28.0
 wheel>=0.36.2
 sympy==1.12.1
+ml_dtypes>=0.5.1

runtime/pyproject.toml

@@ -6,5 +6,6 @@ requires = [
     "ninja",
     "numpy>=2.0.0b1",
     "packaging",
+    "ml_dtypes>=0.5.1",
 ]
 build-backend = "setuptools.build_meta"

runtime/setup.py

@@ -672,5 +672,6 @@ def populate_built_package(abs_dir):
     },
     install_requires=[
         "numpy",
+        "ml_dtypes",
     ],
 )