[Runtime] Unify dtype handling in Python bindings (#24587)
Signed-off-by: Florian Walbroel <walbroel@roofline.ai>
diff --git a/runtime/bindings/python/CMakeLists.txt b/runtime/bindings/python/CMakeLists.txt
index 3f4e357..8668aa7 100644
--- a/runtime/bindings/python/CMakeLists.txt
+++ b/runtime/bindings/python/CMakeLists.txt
@@ -132,6 +132,7 @@
"iree/runtime/_binding.pyi"
"iree/runtime/array_interop.py"
"iree/runtime/benchmark.py"
+ "iree/runtime/dtypes.py"
"iree/runtime/flags.py"
"iree/runtime/function.py"
"iree/runtime/io.py"
@@ -319,6 +320,13 @@
iree_py_test(
NAME
+ dtypes_test
+ SRCS
+ "tests/dtypes_test.py"
+)
+
+iree_py_test(
+ NAME
flags_test
SRCS
"tests/flags_test.py"
diff --git a/runtime/bindings/python/iree/runtime/array_interop.py b/runtime/bindings/python/iree/runtime/array_interop.py
index da2b21c..a7256ba 100644
--- a/runtime/bindings/python/iree/runtime/array_interop.py
+++ b/runtime/bindings/python/iree/runtime/array_interop.py
@@ -7,7 +7,6 @@
from typing import Literal, Optional, Tuple, cast
import logging
-import ml_dtypes
import numpy as np
import numpy.lib.mixins
@@ -21,6 +20,7 @@
MemoryType,
HalFence,
)
+from .dtypes import map_dtype_to_hal_element_type
__all__ = [
"asdevicearray",
@@ -334,37 +334,5 @@
)
-# NOTE: Numpy dtypes are not hashable and exist in a hierarchy that should
-# be queried via isinstance checks. This should be done as a fallback but
-# this is a linear list for quick access to the most common. There may also
-# be a better way to do this.
-_DTYPE_TO_HAL_ELEMENT_TYPE = (
- (np.float16, HalElementType.FLOAT_16),
- (np.float32, HalElementType.FLOAT_32),
- (np.float64, HalElementType.FLOAT_64),
- (np.int32, HalElementType.SINT_32),
- (np.int64, HalElementType.SINT_64),
- (np.int16, HalElementType.SINT_16),
- (np.int8, HalElementType.SINT_8),
- (np.uint32, HalElementType.UINT_32),
- (np.uint64, HalElementType.UINT_64),
- (np.uint16, HalElementType.UINT_16),
- (np.uint8, HalElementType.UINT_8),
- (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),
-)
-
-
def map_dtype_to_element_type(dtype) -> Optional[HalElementType]:
- for match_dtype, element_type in _DTYPE_TO_HAL_ELEMENT_TYPE:
- if match_dtype == dtype:
- return element_type
- else:
- return None
+ return map_dtype_to_hal_element_type(dtype)
diff --git a/runtime/bindings/python/iree/runtime/benchmark.py b/runtime/bindings/python/iree/runtime/benchmark.py
index ecccde0..3a2d9da 100644
--- a/runtime/bindings/python/iree/runtime/benchmark.py
+++ b/runtime/bindings/python/iree/runtime/benchmark.py
@@ -23,6 +23,7 @@
import numpy
import os
import subprocess
+from .dtypes import DTYPE_TO_ABI_TYPE
__all__ = [
"benchmark_exe",
@@ -33,21 +34,6 @@
"BenchmarkResult", "benchmark_name time cpu_time iterations user_counters"
)
-DTYPE_TO_ABI_TYPE = {
- numpy.dtype(numpy.float16): "f16",
- numpy.dtype(numpy.float32): "f32",
- numpy.dtype(numpy.float64): "f64",
- numpy.dtype(numpy.int8): "i8",
- numpy.dtype(numpy.uint8): "i8",
- numpy.dtype(numpy.int16): "i16",
- numpy.dtype(numpy.uint16): "i16",
- numpy.dtype(numpy.int32): "i32",
- numpy.dtype(numpy.uint32): "i32",
- numpy.dtype(numpy.int64): "i64",
- numpy.dtype(numpy.uint64): "i64",
- numpy.dtype(numpy.bool_): "i1",
-}
-
class BenchmarkToolError(Exception):
"""Benchmark exception that preserves the command line and error output."""
@@ -95,7 +81,7 @@
args.append(f"--module={module}")
if entry_function is None:
raise ValueError(
- f"When specifying the module as a filepath the entry function must be specified."
+ "When specifying the module as a filepath the entry function must be specified."
)
args.append(f"--function={entry_function}")
diff --git a/runtime/bindings/python/iree/runtime/dtypes.py b/runtime/bindings/python/iree/runtime/dtypes.py
new file mode 100644
index 0000000..2d37b14
--- /dev/null
+++ b/runtime/bindings/python/iree/runtime/dtypes.py
@@ -0,0 +1,127 @@
+import dataclasses
+import numpy as np
+import ml_dtypes
+from typing import Any
+
+from ._binding import HalElementType
+
+
+@dataclasses.dataclass(frozen=True)
+class DTypeInfo:
+ dtype: np.dtype
+ name: str
+ abi_type: str
+ hal_type: HalElementType
+
+ @property
+ def is_complex(self) -> bool:
+ return np.issubdtype(self.dtype, np.complexfloating)
+
+ @property
+ def is_floating_point(self) -> bool:
+ return np.issubdtype(self.dtype, np.floating)
+
+ @property
+ def is_integer(self) -> bool:
+ return np.issubdtype(self.dtype, np.integer)
+
+ @property
+ def is_signed(self) -> bool:
+ return np.issubdtype(self.dtype, np.signedinteger)
+
+
+_ML_DTYPE_TO_HAL_ELEMENT_TYPE = {
+ np.dtype(ml_dtypes.bfloat16): HalElementType.BFLOAT_16,
+ np.dtype(ml_dtypes.float8_e4m3fn): HalElementType.FLOAT_8_E4M3_FN,
+ np.dtype(ml_dtypes.float8_e4m3fnuz): HalElementType.FLOAT_8_E4M3_FNUZ,
+ np.dtype(ml_dtypes.float8_e5m2): HalElementType.FLOAT_8_E5M2,
+ np.dtype(ml_dtypes.float8_e5m2fnuz): HalElementType.FLOAT_8_E5M2_FNUZ,
+ np.dtype(ml_dtypes.float8_e8m0fnu): HalElementType.FLOAT_8_E8M0_FNU,
+}
+
+
+_DTYPE_TO_INFO = {
+ info.dtype: info
+ for info in [
+ # Float
+ DTypeInfo(np.dtype(np.float16), "float16", "f16", HalElementType.FLOAT_16),
+ DTypeInfo(np.dtype(np.float32), "float32", "f32", HalElementType.FLOAT_32),
+ DTypeInfo(np.dtype(np.float64), "float64", "f64", HalElementType.FLOAT_64),
+ # Integer
+ DTypeInfo(np.dtype(np.int8), "int8", "i8", HalElementType.SINT_8),
+ DTypeInfo(np.dtype(np.uint8), "uint8", "i8", HalElementType.UINT_8),
+ DTypeInfo(np.dtype(np.int16), "int16", "i16", HalElementType.SINT_16),
+ DTypeInfo(np.dtype(np.uint16), "uint16", "i16", HalElementType.UINT_16),
+ DTypeInfo(np.dtype(np.int32), "int32", "i32", HalElementType.SINT_32),
+ DTypeInfo(np.dtype(np.uint32), "uint32", "i32", HalElementType.UINT_32),
+ DTypeInfo(np.dtype(np.int64), "int64", "i64", HalElementType.SINT_64),
+ DTypeInfo(np.dtype(np.uint64), "uint64", "i64", HalElementType.UINT_64),
+ # Boolean
+ DTypeInfo(np.dtype(np.bool_), "bool", "i1", HalElementType.BOOL_8),
+ # Complex
+ DTypeInfo(
+ np.dtype(np.complex64),
+ "complex64",
+ "complex<f32>",
+ HalElementType.COMPLEX_64,
+ ),
+ DTypeInfo(
+ np.dtype(np.complex128),
+ "complex128",
+ "complex<f64>",
+ HalElementType.COMPLEX_128,
+ ),
+ ]
+}
+
+_NAME_TO_INFO: dict[str, DTypeInfo] = {
+ info.name: info for info in _DTYPE_TO_INFO.values()
+}
+
+_ABI_TYPE_TO_INFO: dict[str, DTypeInfo] = {
+ info.abi_type: info
+ for info in _DTYPE_TO_INFO.values()
+ if not info.is_integer
+ or info.is_signed # Exclude unsigned integers since they share ABI types with signed integers.
+}
+
+
+def map_dtype_to_dtype_info(dtype: np.dtype | type) -> DTypeInfo:
+ dtype = np.dtype(dtype)
+ try:
+ return _DTYPE_TO_INFO[dtype]
+ except KeyError:
+ raise KeyError(f"Numpy dtype {dtype} not found.") from None
+
+
+def map_name_to_dtype_info(name: str) -> DTypeInfo:
+ try:
+ return _NAME_TO_INFO[name]
+ except KeyError:
+ raise KeyError(f"Numpy dtype name {name!r} not found.") from None
+
+
+def map_abi_type_to_dtype_info(abi_type: str) -> DTypeInfo:
+ try:
+ return _ABI_TYPE_TO_INFO[abi_type]
+ except KeyError:
+ raise KeyError(f"ABI type {abi_type!r} not found.") from None
+
+
+def map_dtype_to_hal_element_type(dtype: Any) -> HalElementType | None:
+ try:
+ normalized_dtype = np.dtype(dtype)
+ except TypeError:
+ return None
+ elem_ty = _ML_DTYPE_TO_HAL_ELEMENT_TYPE.get(normalized_dtype)
+ if elem_ty is not None:
+ return elem_ty
+ try:
+ return map_dtype_to_dtype_info(normalized_dtype).hal_type
+ except KeyError:
+ return None
+
+
+DTYPE_TO_ABI_TYPE: dict[np.dtype, str] = {
+ info.dtype: info.abi_type for info in _DTYPE_TO_INFO.values()
+}
diff --git a/runtime/bindings/python/iree/runtime/function.py b/runtime/bindings/python/iree/runtime/function.py
index c104d24..342dd94 100644
--- a/runtime/bindings/python/iree/runtime/function.py
+++ b/runtime/bindings/python/iree/runtime/function.py
@@ -4,7 +4,7 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-from typing import Dict, Optional
+from typing import Optional
import json
import logging
@@ -27,11 +27,11 @@
)
from .array_interop import (
- map_dtype_to_element_type,
DeviceArray,
)
-from .flags import (
- FUNCTION_INPUT_VALIDATION,
+
+from .dtypes import (
+ map_abi_type_to_dtype_info,
)
__all__ = [
@@ -204,7 +204,7 @@
buffer_view = vm_list.get_as_object(vm_index, HalBufferView)
dtype_str = desc[1]
try:
- dtype = ABI_TYPE_TO_DTYPE[dtype_str]
+ dtype = map_abi_type_to_dtype_info(dtype_str).dtype
except KeyError:
_raise_return_error(inv, f"unrecognized dtype '{dtype_str}'")
x = DeviceArray(
@@ -279,42 +279,12 @@
"bf16": _vm_to_scalar(float),
}
-ABI_TYPE_TO_DTYPE = {
- # TODO: Others.
- "f32": np.float32,
- "i32": np.int32,
- "i64": np.int64,
- "f64": np.float64,
- "i16": np.int16,
- "i8": np.int8,
- "i1": np.bool_,
-}
-
# When we get an ndarray as an argument and are implicitly converting it to a
# buffer view, flags for doing so.
IMPLICIT_BUFFER_ARG_MEMORY_TYPE = MemoryType.DEVICE_LOCAL
IMPLICIT_BUFFER_ARG_USAGE = BufferUsage.DEFAULT
-def _is_ndarray_descriptor(desc):
- return desc and desc[0] == "ndarray"
-
-
-def _is_0d_ndarray_descriptor(desc):
- # Example: ["ndarray", "f32", 0]
- return desc and desc[0] == "ndarray" and desc[2] == 0
-
-
-def _cast_scalar_to_ndarray(inv: Invocation, x, desc):
- # Example descriptor: ["ndarray", "f32", 0]
- dtype_str = desc[1]
- try:
- dtype = ABI_TYPE_TO_DTYPE[dtype_str]
- except KeyError:
- _raise_argument_error(inv, f"unrecognized dtype '{dtype_str}'")
- return dtype(x)
-
-
class ArgumentError(ValueError):
pass
diff --git a/runtime/bindings/python/iree/runtime/io.py b/runtime/bindings/python/iree/runtime/io.py
index dcd5ce8..bc430ea 100644
--- a/runtime/bindings/python/iree/runtime/io.py
+++ b/runtime/bindings/python/iree/runtime/io.py
@@ -13,6 +13,7 @@
from os import PathLike
from ._binding import ParameterIndex, ParameterIndexEntry
+from .dtypes import map_name_to_dtype_info, map_dtype_to_dtype_info
__all__ = [
"parameter_index_add_numpy_ndarray",
@@ -151,9 +152,11 @@
# Unpack/validate.
try:
- dtype = _NAME_TO_DTYPE[dtype_name]
- except KeyError:
- raise ValueError(f"Unknown dtype name '{dtype_name}'")
+ dtype = map_name_to_dtype_info(dtype_name).dtype
+ except KeyError as e:
+ raise ValueError(
+ f"Unsupported dtype for parameter entry {index_entry.key}"
+ ) from e
try:
shape = [int(d) for d in shape]
except ValueError as e:
@@ -163,36 +166,6 @@
return t.view(dtype=dtype).reshape(shape)
-_DTYPE_TO_NAME = (
- (np.dtype(np.float16), "float16"),
- (np.dtype(np.float32), "float32"),
- (np.dtype(np.float64), "float64"),
- (np.dtype(np.int32), "int32"),
- (np.dtype(np.int64), "int64"),
- (np.dtype(np.int16), "int16"),
- (np.dtype(np.int8), "int8"),
- (np.dtype(np.uint32), "uint32"),
- (np.dtype(np.uint64), "uint64"),
- (np.dtype(np.uint16), "uint16"),
- (np.dtype(np.uint8), "uint8"),
- (np.dtype(np.bool_), "bool"),
- (np.dtype(np.complex64), "complex64"),
- (np.dtype(np.complex128), "complex128"),
-)
-
-_NAME_TO_DTYPE: dict[str, np.dtype[Any]] = {
- name: np_dtype for np_dtype, name in _DTYPE_TO_NAME
-}
-
-
-def _map_dtype_to_name(dtype) -> str:
- for match_dtype, dtype_name in _DTYPE_TO_NAME:
- if match_dtype == dtype:
- return dtype_name
-
- raise KeyError(f"Numpy dtype {dtype} not found.")
-
-
_metadata_version = "TENSORv0"
"""Magic number to identify the format version.
The current version that will be used when adding tensors to a parameter index."""
@@ -209,20 +182,17 @@
def _make_tensor_metadata(t: np.ndarray) -> str:
"""Makes a tensor metadata blob that can be used to reconstruct the tensor."""
dtype = t.dtype
- dtype_name = _map_dtype_to_name(dtype)
- is_complex = np.issubdtype(dtype, np.complexfloating)
- is_floating_point = np.issubdtype(dtype, np.floating)
- is_signed = np.issubdtype(dtype, np.signedinteger)
+ dtype_info = map_dtype_to_dtype_info(dtype)
dtype_desc = {
"class_name": type(dtype).__name__,
- "is_complex": is_complex,
- "is_floating_point": is_floating_point,
- "is_signed": is_signed,
+ "is_complex": dtype_info.is_complex,
+ "is_floating_point": dtype_info.is_floating_point,
+ "is_signed": dtype_info.is_signed,
"itemsize": dtype.itemsize,
}
d = {
"type": "Tensor",
- "dtype": dtype_name,
+ "dtype": dtype_info.name,
"shape": list(t.shape),
"dtype_desc": dtype_desc,
}
diff --git a/runtime/bindings/python/tests/dtypes_test.py b/runtime/bindings/python/tests/dtypes_test.py
new file mode 100644
index 0000000..baa0400
--- /dev/null
+++ b/runtime/bindings/python/tests/dtypes_test.py
@@ -0,0 +1,108 @@
+# Copyright 2026 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import numpy as np
+import unittest
+
+from iree import runtime as rt
+import iree.runtime.dtypes as dtypes
+import ml_dtypes
+
+
+class DTypesTest(unittest.TestCase):
+ def testDTypeToInfo(self):
+ info = dtypes.map_dtype_to_dtype_info(np.dtype(np.bool_))
+ self.assertEqual(info.dtype, np.dtype(np.bool_))
+ self.assertEqual(info.name, "bool")
+ self.assertEqual(info.abi_type, "i1")
+ self.assertFalse(info.is_integer)
+ self.assertFalse(info.is_signed)
+ self.assertFalse(info.is_floating_point)
+ self.assertFalse(info.is_complex)
+
+ with self.assertRaises(KeyError):
+ dtypes.map_dtype_to_dtype_info(np.dtype("bfloat16"))
+
+ def testNameToInfo(self):
+ info = dtypes.map_name_to_dtype_info("float32")
+ self.assertEqual(info.dtype, np.dtype(np.float32))
+ self.assertEqual(info.name, "float32")
+ self.assertEqual(info.abi_type, "f32")
+ self.assertFalse(info.is_integer)
+ self.assertFalse(info.is_signed)
+ self.assertTrue(info.is_floating_point)
+ self.assertFalse(info.is_complex)
+
+ with self.assertRaises(KeyError):
+ dtypes.map_name_to_dtype_info("invalid")
+
+ def testAbiTypeToInfo(self):
+ info = dtypes.map_abi_type_to_dtype_info("i16")
+ self.assertEqual(info.dtype, np.dtype(np.int16))
+ self.assertEqual(info.name, "int16")
+ self.assertEqual(info.abi_type, "i16")
+ self.assertTrue(info.is_integer)
+ self.assertTrue(info.is_signed)
+ self.assertFalse(info.is_floating_point)
+ self.assertFalse(info.is_complex)
+
+ with self.assertRaises(KeyError):
+ dtypes.map_abi_type_to_dtype_info("invalid")
+
+ self.assertEqual(
+ dtypes.map_abi_type_to_dtype_info("i16").dtype, np.dtype(np.int16)
+ )
+ self.assertEqual(
+ dtypes.map_abi_type_to_dtype_info("i32").dtype, np.dtype(np.int32)
+ )
+ self.assertEqual(
+ dtypes.map_abi_type_to_dtype_info("i64").dtype, np.dtype(np.int64)
+ )
+
+ def testDTypeToAbiType(self) -> None:
+ self.assertEqual(dtypes.DTYPE_TO_ABI_TYPE[np.dtype(np.uint8)], "i8")
+ self.assertEqual(dtypes.DTYPE_TO_ABI_TYPE[np.dtype(np.uint16)], "i16")
+ self.assertEqual(dtypes.DTYPE_TO_ABI_TYPE[np.dtype(np.uint32)], "i32")
+
+ def testMapMlDtypeToHalElementTypeFromNumpyDType(self):
+ self.assertEqual(
+ dtypes.map_dtype_to_hal_element_type(np.dtype("bfloat16")),
+ rt.HalElementType.BFLOAT_16,
+ )
+ self.assertEqual(
+ dtypes.map_dtype_to_hal_element_type(np.dtype("float8_e4m3fn")),
+ rt.HalElementType.FLOAT_8_E4M3_FN,
+ )
+ self.assertEqual(
+ dtypes.map_dtype_to_hal_element_type(np.dtype("float8_e4m3fnuz")),
+ rt.HalElementType.FLOAT_8_E4M3_FNUZ,
+ )
+ self.assertEqual(
+ dtypes.map_dtype_to_hal_element_type(np.dtype("float8_e5m2")),
+ rt.HalElementType.FLOAT_8_E5M2,
+ )
+ self.assertEqual(
+ dtypes.map_dtype_to_hal_element_type(np.dtype("float8_e5m2fnuz")),
+ rt.HalElementType.FLOAT_8_E5M2_FNUZ,
+ )
+ self.assertEqual(
+ dtypes.map_dtype_to_hal_element_type(np.dtype("float8_e8m0fnu")),
+ rt.HalElementType.FLOAT_8_E8M0_FNU,
+ )
+
+ def testMapMlDtypeToHalElementTypeFromScalarType(self):
+ self.assertEqual(
+ dtypes.map_dtype_to_hal_element_type(ml_dtypes.bfloat16),
+ rt.HalElementType.BFLOAT_16,
+ )
+ self.assertEqual(
+ dtypes.map_dtype_to_hal_element_type(ml_dtypes.float8_e4m3fn),
+ rt.HalElementType.FLOAT_8_E4M3_FN,
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/runtime/bindings/python/tests/function_test.py b/runtime/bindings/python/tests/function_test.py
index 004c2dd..6010d51 100644
--- a/runtime/bindings/python/tests/function_test.py
+++ b/runtime/bindings/python/tests/function_test.py
@@ -56,6 +56,18 @@
self.assertEqual("[<VmVariantList(2): [1, 2]>]", vm_context.mock_arg_reprs)
self.assertEqual((3, 4), result)
+ def testNoReflectionBfloat16Ndarray(self):
+ def invoke(arg_list, ret_list):
+ ret_list.push_int(3)
+
+ vm_context = MockVmContext(invoke)
+ vm_function = MockVmFunction(reflection={})
+ invoker = FunctionInvoker(vm_context, self.device, vm_function)
+ arg = np.arange(4, dtype=np.dtype("bfloat16")).reshape(2, 2)
+ result = invoker(arg)
+ self.assertEqual(3, result)
+ self.assertIn("HalBufferView", vm_context.mock_arg_reprs)
+
def testKeywordArgs(self):
def invoke(arg_list, ret_list):
ret_list.push_int(3)
