[Runtime][Python] Add pyright checks and fix errors (#24561)
Signed-off-by: Florian Walbroel <walbroel@roofline.ai>
diff --git a/build_tools/github_actions/ci_requirements.txt b/build_tools/github_actions/ci_requirements.txt
index 5648db8..81570e4 100644
--- a/build_tools/github_actions/ci_requirements.txt
+++ b/build_tools/github_actions/ci_requirements.txt
@@ -3,3 +3,4 @@
cmake>=3.26
ninja
packaging
+pyright==1.1.410
diff --git a/runtime/bindings/python/CMakeLists.txt b/runtime/bindings/python/CMakeLists.txt
index 087a0c4..3f4e357 100644
--- a/runtime/bindings/python/CMakeLists.txt
+++ b/runtime/bindings/python/CMakeLists.txt
@@ -295,6 +295,21 @@
# Tests
################################################################################
+if(IREE_BUILD_PYTHON_BINDINGS AND IREE_BUILD_TESTS)
+ find_program(PYRIGHT_EXECUTABLE NAMES pyright REQUIRED)
+ add_test(
+ NAME iree/runtime/bindings/pyright
+ COMMAND
+ "${PYRIGHT_EXECUTABLE}"
+ "-p"
+ "${CMAKE_CURRENT_SOURCE_DIR}/pyrightconfig.json"
+ WORKING_DIRECTORY
+ "${CMAKE_CURRENT_SOURCE_DIR}"
+ )
+ set_property(TEST iree/runtime/bindings/pyright PROPERTY LABELS "runtime;lint")
+ set_property(TEST iree/runtime/bindings/pyright PROPERTY TIMEOUT 120)
+endif()
+
iree_py_test(
NAME
array_interop_test
diff --git a/runtime/bindings/python/iree/_runtime/libs.pyi b/runtime/bindings/python/iree/_runtime/libs.pyi
new file mode 100644
index 0000000..de17185
--- /dev/null
+++ b/runtime/bindings/python/iree/_runtime/libs.pyi
@@ -0,0 +1,5 @@
+from typing import Any
+
+_runtime: Any
+library_path: str
+name: str
diff --git a/runtime/bindings/python/iree/runtime/_binding.pyi b/runtime/bindings/python/iree/runtime/_binding.pyi
index f2013e0..05abddd 100644
--- a/runtime/bindings/python/iree/runtime/_binding.pyi
+++ b/runtime/bindings/python/iree/runtime/_binding.pyi
@@ -1,4 +1,5 @@
import os
+import numpy as np
from typing import (
Any,
Callable,
@@ -7,6 +8,7 @@
Optional,
Sequence,
Tuple,
+ TypeVar,
Union,
overload,
)
@@ -14,6 +16,8 @@
from .typing import HalModuleBufferViewTraceCallback
+_T = TypeVar("_T")
+
def create_hal_module(
instance: VmInstance,
device: Optional[HalDevice] = None,
@@ -67,31 +71,49 @@
def __or__(self, other: BufferUsage) -> int: ...
class ExternalTimepointType(int):
- NONE = ClassVar[ExternalTimepointType] = ...
- ASYNC_PRIMITIVE = ClassVar[ExternalTimepointType] = ...
- CUDA_EVENT = ClassVar[ExternalTimepointType] = ...
- HIP_EVENT = ClassVar[ExternalTimepointType] = ...
+ NONE: ClassVar[ExternalTimepointType] = ...
+ ASYNC_PRIMITIVE: ClassVar[ExternalTimepointType] = ...
+ CUDA_EVENT: ClassVar[ExternalTimepointType] = ...
+ HIP_EVENT: ClassVar[ExternalTimepointType] = ...
def __int__(self) -> int: ...
class SemaphoreCompatibility(int):
- NONE = ClassVar[SemaphoreCompatibility] = ...
- HOST_WAIT = ClassVar[SemaphoreCompatibility] = ...
- DEVICE_WAIT = ClassVar[SemaphoreCompatibility] = ...
- HOST_SIGNAL = ClassVar[SemaphoreCompatibility] = ...
- DEVICE_SIGNAL = ClassVar[SemaphoreCompatibility] = ...
- HOST_ONLY = ClassVar[SemaphoreCompatibility] = ...
- DEVICE_ONLY = ClassVar[SemaphoreCompatibility] = ...
- ALL = ClassVar[SemaphoreCompatibility] = ...
+ NONE: ClassVar[SemaphoreCompatibility] = ...
+ HOST_WAIT: ClassVar[SemaphoreCompatibility] = ...
+ DEVICE_WAIT: ClassVar[SemaphoreCompatibility] = ...
+ HOST_SIGNAL: ClassVar[SemaphoreCompatibility] = ...
+ DEVICE_SIGNAL: ClassVar[SemaphoreCompatibility] = ...
+ HOST_ONLY: ClassVar[SemaphoreCompatibility] = ...
+ DEVICE_ONLY: ClassVar[SemaphoreCompatibility] = ...
+ ALL: ClassVar[SemaphoreCompatibility] = ...
def __and__(self, other: SemaphoreCompatibility) -> int: ...
def __or__(self, other: SemaphoreCompatibility) -> int: ...
def __int__(self) -> int: ...
class ExternalTimepointFlags(int):
- NONE = ClassVar[ExternalTimepointFlags] = ...
+ NONE: ClassVar[ExternalTimepointFlags] = ...
def __and__(self, other: ExternalTimepointFlags) -> int: ...
def __or__(self, other: ExternalTimepointFlags) -> int: ...
def __int__(self) -> int: ...
+class _InvokeStatics: ...
+
+_invoke_statics: _InvokeStatics
+
+class InvokeContext:
+ def __init__(self, device: HalDevice) -> None: ...
+
+class ArgumentPacker:
+ def __init__(
+ self, statics: _InvokeStatics, arg_descs: Optional[list[Any]] = None
+ ) -> None: ...
+ def pack(
+ self,
+ invoke_context: InvokeContext,
+ pos_args: Sequence[Any],
+ kwargs: dict[str, Any],
+ ) -> VmVariantList: ...
+
class FileHandle:
@staticmethod
def wrap_memory(
@@ -129,6 +151,12 @@
def allocate_host_staging_buffer_copy(
self, device: HalDevice, initial_contents: object
) -> HalBuffer: ...
+ def import_host_allocation(
+ self,
+ buffer_obj: Any,
+ memory_type: Union[MemoryType, int],
+ allowed_usage: Union[BufferUsage, int],
+ ) -> HalBuffer: ...
def query_buffer_compatibility(
self,
memory_type: Union[MemoryType, int],
@@ -202,7 +230,7 @@
@property
def __iree_vm_ref__(self) -> VmRef: ...
-HalSemaphoreList = List[Tuple[HalSemaphore, int]]
+HalSemaphoreList = List[Tuple[HalSemaphore, int]] | HalFence
class HalDevice:
def begin_profiling(self, mode: Optional[str] = None) -> None: ...
@@ -345,7 +373,7 @@
class HalFence:
@staticmethod
- def create_at(value: int) -> HalFence: ...
+ def create_at(sem: HalSemaphore, value: int) -> HalFence: ...
@staticmethod
def join(fences: Sequence[HalFence]) -> HalFence: ...
def __init__(self, capacity: int) -> None: ...
@@ -422,7 +450,9 @@
class MappedMemory:
def __init__(self, *args, **kwargs) -> None: ...
- def asarray(self, shape: Sequence[int], numpy_dtype_descr: object) -> object: ...
+ def asarray(
+ self, shape: Sequence[int], numpy_dtype_descr: object
+ ) -> np.ndarray[Any, np.dtype[Any]]: ...
class MemoryAccess(int):
ALL: ClassVar[MemoryAccess] = ...
@@ -481,7 +511,7 @@
...
class ParameterIndex:
- def __init__() -> None: ...
+ def __init__(self) -> None: ...
def __len__(self) -> int: ...
def __getitem__(self, i) -> ParameterIndexEntry: ...
def reserve(self, new_capacity: int) -> None: ...
@@ -564,7 +594,7 @@
class VmContext:
def __init__(
- self, instance: VmInstance, modules: Optional[list[VmModule]] = None
+ self, instance: VmInstance, modules: Optional[Sequence[VmModule]] = None
) -> None: ...
def invoke(
self, function: VmFunction, inputs: VmVariantList, outputs: VmVariantList
@@ -624,12 +654,12 @@
@property
def name(self) -> str: ...
@property
- def stashed_flatbuffer_blob(self) -> object: ...
+ def stashed_flatbuffer_blob(self) -> bytes | None: ...
@property
def version(self) -> int: ...
class VmRef:
- def deref(self, value: Any, optional: bool = False) -> object: ...
+ def deref(self, value: type[_T], optional: bool = False) -> _T | None: ...
def isinstance(self, ref_type: type) -> bool: ...
def __eq__(self, other) -> bool: ...
@property
@@ -638,7 +668,10 @@
class VmVariantList:
def __init__(self, capacity: int) -> None: ...
def get_as_list(self, index: int) -> VmVariantList: ...
+ @overload
def get_as_object(self, index: int) -> object: ...
+ @overload
+ def get_as_object(self, index: int, clazz: type[_T]) -> _T: ...
def get_as_ref(self, index: int) -> VmRef: ...
def get_serialized_trace_value(self, index: int) -> dict: ...
def get_variant(self, index: int) -> Any: ...
diff --git a/runtime/bindings/python/iree/runtime/array_interop.py b/runtime/bindings/python/iree/runtime/array_interop.py
index 63deb86..da2b21c 100644
--- a/runtime/bindings/python/iree/runtime/array_interop.py
+++ b/runtime/bindings/python/iree/runtime/array_interop.py
@@ -5,7 +5,7 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""BufferView and Python Array Protocol interop."""
-from typing import Optional, Tuple
+from typing import Literal, Optional, Tuple, cast
import logging
import ml_dtypes
import numpy as np
@@ -99,6 +99,18 @@
func, types, args, kwargs
) # pytype: disable=attribute-error
+ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+ def transfer_input(inp):
+ if isinstance(inp, DeviceArray):
+ return inp._transfer_to_host(True)
+ return inp
+
+ host_inputs = [transfer_input(inp) for inp in inputs]
+ if kwargs.get("out") is not None:
+ kwargs = dict(kwargs)
+ kwargs["out"] = tuple(transfer_input(inp) for inp in kwargs["out"])
+ return getattr(ufunc, method)(*host_inputs, **kwargs)
+
def __repr__(self):
return f"<IREE DeviceArray: shape={np.shape(self)}, dtype={self.dtype}>"
@@ -142,7 +154,9 @@
# TODO: When synchronization is enabled, need to block here.
raw_dtype = self._get_raw_dtype()
mapped_memory = self._buffer_view.map()
- host_array = mapped_memory.asarray(self._buffer_view.shape, raw_dtype)
+ host_array = cast(
+ np.ndarray, mapped_memory.asarray(self._buffer_view.shape, raw_dtype)
+ )
# Detect if we need to force an explicit conversion. This happens when
# we were requested to pretend that the array is in a specific dtype,
# even if that is not representable on the device. You guessed it:
@@ -216,7 +230,12 @@
def shape(self):
return np.shape(self)
- def astype(self, dtype, casting="unsafe", copy=True):
+ def astype(
+ self,
+ dtype,
+ casting: Literal["no", "equiv", "safe", "same_kind", "unsafe"] = "unsafe",
+ copy=True,
+ ):
if self.dtype == dtype and not copy:
return self
host_ary = self.to_host()
@@ -305,6 +324,8 @@
buffer=a,
element_type=element_type,
)
+ if not isinstance(buffer_view, HalBufferView):
+ raise TypeError("Expected allocate_buffer_copy to return a HalBufferView")
return DeviceArray(
device,
buffer_view,
diff --git a/runtime/bindings/python/iree/runtime/benchmark.py b/runtime/bindings/python/iree/runtime/benchmark.py
index e1598dc..ecccde0 100644
--- a/runtime/bindings/python/iree/runtime/benchmark.py
+++ b/runtime/bindings/python/iree/runtime/benchmark.py
@@ -19,7 +19,6 @@
from typing import Union
from os import PathLike
-import iree.runtime
from . import VmModule
import numpy
import os
@@ -105,7 +104,7 @@
v = kwargs[k]
args.append(f"--{k}={v}")
- for inp in inputs:
+ for inp in inputs or []:
if isinstance(inp, str):
args.append(f"--input={inp}")
continue
diff --git a/runtime/bindings/python/iree/runtime/build_requirements.txt b/runtime/bindings/python/iree/runtime/build_requirements.txt
index 3700a42..81384ac 100644
--- a/runtime/bindings/python/iree/runtime/build_requirements.txt
+++ b/runtime/bindings/python/iree/runtime/build_requirements.txt
@@ -10,3 +10,4 @@
wheel>=0.36.2
sympy>=1.14.0
ml_dtypes>=0.5.4
+pyright==1.1.410
diff --git a/runtime/bindings/python/iree/runtime/function.py b/runtime/bindings/python/iree/runtime/function.py
index 0de8e94..c104d24 100644
--- a/runtime/bindings/python/iree/runtime/function.py
+++ b/runtime/bindings/python/iree/runtime/function.py
@@ -8,6 +8,7 @@
import json
import logging
+from typing import cast
import numpy as np
@@ -364,7 +365,9 @@
# since this is higher level and it preserves layering. Note that
# the reflection case also does this conversion.
if isinstance(converted, VmRef):
- converted_buffer_view = converted.deref(HalBufferView, True)
+ converted_buffer_view = cast(
+ HalBufferView | None, converted.deref(HalBufferView, True)
+ )
if converted_buffer_view:
converted = DeviceArray(
inv.device, converted_buffer_view, implicit_host_transfer=True
@@ -382,7 +385,7 @@
raise
except Exception as e:
_raise_return_error(
- inv, f"exception converting from VM type to Python", e
+ inv, "exception converting from VM type to Python", e
)
results.append(converted)
return results
diff --git a/runtime/bindings/python/iree/runtime/io.py b/runtime/bindings/python/iree/runtime/io.py
index e9d54da..dcd5ce8 100644
--- a/runtime/bindings/python/iree/runtime/io.py
+++ b/runtime/bindings/python/iree/runtime/io.py
@@ -26,20 +26,22 @@
class SplatValue:
def __init__(
self,
- pattern: Union[array.array, np.ndarray],
+ pattern: Union[array.array, np.ndarray, np.generic],
count: Union[Sequence[int], int],
):
- if hasattr(pattern, "shape"):
+ if isinstance(pattern, np.ndarray):
shape = pattern.shape
if not shape:
total_elements = 1
else:
- total_elements = reduce(lambda x, y: x * y, pattern.shape)
+ total_elements = reduce(lambda x, y: x * y, shape)
+ elif isinstance(pattern, np.generic):
+ total_elements = 1
else:
total_elements = len(pattern)
item_size = pattern.itemsize
if total_elements != 1:
- raise ValueError(f"SplatValue requires an array of a single element")
+ raise ValueError("SplatValue requires an array of a single element")
self.pattern = pattern
if isinstance(count, int):
logical_length = count
@@ -162,23 +164,23 @@
_DTYPE_TO_NAME = (
- (np.float16, "float16"),
- (np.float32, "float32"),
- (np.float64, "float64"),
- (np.int32, "int32"),
- (np.int64, "int64"),
- (np.int16, "int16"),
- (np.int8, "int8"),
- (np.uint32, "uint32"),
- (np.uint64, "uint64"),
- (np.uint16, "uint16"),
- (np.uint8, "uint8"),
- (np.bool_, "bool"),
- (np.complex64, "complex64"),
- (np.complex128, "complex128"),
+ (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] = {
+_NAME_TO_DTYPE: dict[str, np.dtype[Any]] = {
name: np_dtype for np_dtype, name in _DTYPE_TO_NAME
}
diff --git a/runtime/bindings/python/iree/runtime/system_api.py b/runtime/bindings/python/iree/runtime/system_api.py
index c41251c..6bd0bad 100644
--- a/runtime/bindings/python/iree/runtime/system_api.py
+++ b/runtime/bindings/python/iree/runtime/system_api.py
@@ -31,7 +31,7 @@
import os
import sys
-from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, List, MutableMapping, Optional, Sequence, Tuple, Union
from . import _binding
from .function import FunctionInvoker
@@ -77,7 +77,9 @@
self.default_vm_modules = (hal_module,)
-def _bool_to_int8(array: Any) -> Optional[Union[np.ndarray, List[Any], Tuple[Any]]]:
+def _bool_to_int8(
+ array: Any,
+) -> Optional[Union[np.ndarray, List[Any], Tuple[Any], dict[Any, Any]]]:
if not isinstance(array, np.ndarray):
return array
@@ -88,7 +90,9 @@
return array
-def normalize_value(value: Any) -> Optional[Union[np.ndarray, List[Any], Tuple[Any]]]:
+def normalize_value(
+ value: Any,
+) -> Optional[Union[np.ndarray, List[Any], Tuple[Any], dict[Any, Any]]]:
"""Normalizes the given value for input to (or comparison with) IREE."""
if value is None:
# Exclude None from falling through to blanket np.asarray conversion.
@@ -112,7 +116,7 @@
def _convert_lists_to_tuples(pytree):
if isinstance(pytree, Sequence):
return tuple(_convert_lists_to_tuples(leaf) for leaf in pytree)
- elif isinstance(pytree, Mapping):
+ elif isinstance(pytree, MutableMapping):
for key in pytree:
pytree[key] = _convert_lists_to_tuples(pytree[key])
return pytree
@@ -180,12 +184,17 @@
class SystemContext:
"""Global system."""
- def __init__(self, vm_modules=None, config: Optional[Config] = None):
+ def __init__(
+ self,
+ vm_modules: Optional[Sequence[_binding.VmModule]] = None,
+ config: Optional[Config] = None,
+ ):
self._config = config if config is not None else Config()
self._is_dynamic = vm_modules is None
if self._is_dynamic:
init_vm_modules = None
else:
+ assert vm_modules is not None
init_vm_modules = self._config.default_vm_modules + tuple(vm_modules)
self._vm_context = _binding.VmContext(
@@ -201,6 +210,7 @@
]
)
else:
+ assert init_vm_modules is not None
self._bound_modules = BoundModules(
[(m.name, BoundModule(self, m)) for m in init_vm_modules]
)
diff --git a/runtime/bindings/python/iree/runtime/system_setup.py b/runtime/bindings/python/iree/runtime/system_setup.py
index c592e46..029af5f 100644
--- a/runtime/bindings/python/iree/runtime/system_setup.py
+++ b/runtime/bindings/python/iree/runtime/system_setup.py
@@ -49,6 +49,8 @@
if cache:
existing = _GLOBAL_DEVICES_BY_URI.get(device_uri)
if existing is not None:
+ if isinstance(existing, Exception):
+ raise existing
return existing
driver = get_driver(device_uri)
diff --git a/runtime/bindings/python/iree/runtime/version.py b/runtime/bindings/python/iree/runtime/version.py
index 65f20b6..71ce247 100644
--- a/runtime/bindings/python/iree/runtime/version.py
+++ b/runtime/bindings/python/iree/runtime/version.py
@@ -9,4 +9,4 @@
package.
"""
-from .._runtime_libs.version import *
+from .._runtime_libs.version import * # pyright: ignore[reportMissingImports]
diff --git a/runtime/bindings/python/iree/runtime/version.pyi b/runtime/bindings/python/iree/runtime/version.pyi
new file mode 100644
index 0000000..dad2728
--- /dev/null
+++ b/runtime/bindings/python/iree/runtime/version.pyi
@@ -0,0 +1,3 @@
+PACKAGE_SUFFIX: str
+REVISIONS: dict[str, str]
+VERSION: str
diff --git a/runtime/bindings/python/pyrightconfig.json b/runtime/bindings/python/pyrightconfig.json
new file mode 100644
index 0000000..8382457
--- /dev/null
+++ b/runtime/bindings/python/pyrightconfig.json
@@ -0,0 +1,17 @@
+{
+ "typeCheckingMode": "basic",
+ "pythonVersion": "3.10",
+ "include": [
+ "iree"
+ ],
+ "stubPath": "typings",
+ "reportMissingModuleSource": "none",
+ "extraPaths": [
+ "."
+ ],
+ "exclude": [
+ "**/__pycache__",
+ "**/.venv",
+ "**/build"
+ ]
+}
diff --git a/runtime/bindings/python/tests/array_interop_test.py b/runtime/bindings/python/tests/array_interop_test.py
index 77bc132..0d12c1d 100644
--- a/runtime/bindings/python/tests/array_interop_test.py
+++ b/runtime/bindings/python/tests/array_interop_test.py
@@ -102,6 +102,14 @@
with self.assertRaises(ValueError):
_ = np.asarray(ary)
+ def testIllegalImplicitHostTransferForUfunc(self):
+ init_ary = np.zeros([3, 4], dtype=np.int32) + 2
+ ary = iree.runtime.asdevicearray(self.device, init_ary)
+ with self.assertRaises(ValueError):
+ _ = ary + 1
+ with self.assertRaises(ValueError):
+ _ = np.add(ary, 1)
+
def testImplicitHostArithmetic(self):
init_ary = np.zeros([3, 4], dtype=np.int32) + 2
ary = iree.runtime.asdevicearray(
diff --git a/runtime/bindings/python/typings/iree/_runtime/libs.pyi b/runtime/bindings/python/typings/iree/_runtime/libs.pyi
new file mode 100644
index 0000000..de17185
--- /dev/null
+++ b/runtime/bindings/python/typings/iree/_runtime/libs.pyi
@@ -0,0 +1,5 @@
+from typing import Any
+
+_runtime: Any
+library_path: str
+name: str
diff --git a/runtime/bindings/python/typings/iree/_runtime_libs/__init__.pyi b/runtime/bindings/python/typings/iree/_runtime_libs/__init__.pyi
new file mode 100644
index 0000000..dbf2dc4
--- /dev/null
+++ b/runtime/bindings/python/typings/iree/_runtime_libs/__init__.pyi
@@ -0,0 +1,6 @@
+from typing import Any
+
+__path__: list[str]
+_runtime: Any
+library_path: str
+name: str
diff --git a/runtime/bindings/python/typings/iree/_runtime_libs/version.pyi b/runtime/bindings/python/typings/iree/_runtime_libs/version.pyi
new file mode 100644
index 0000000..dad2728
--- /dev/null
+++ b/runtime/bindings/python/typings/iree/_runtime_libs/version.pyi
@@ -0,0 +1,3 @@
+PACKAGE_SUFFIX: str
+REVISIONS: dict[str, str]
+VERSION: str
diff --git a/runtime/bindings/python/typings/iree/_runtime_libs_tracy/__init__.pyi b/runtime/bindings/python/typings/iree/_runtime_libs_tracy/__init__.pyi
new file mode 100644
index 0000000..dbf2dc4
--- /dev/null
+++ b/runtime/bindings/python/typings/iree/_runtime_libs_tracy/__init__.pyi
@@ -0,0 +1,6 @@
+from typing import Any
+
+__path__: list[str]
+_runtime: Any
+library_path: str
+name: str
diff --git a/runtime/bindings/python/typings/iree/runtime/version.pyi b/runtime/bindings/python/typings/iree/runtime/version.pyi
new file mode 100644
index 0000000..dad2728
--- /dev/null
+++ b/runtime/bindings/python/typings/iree/runtime/version.pyi
@@ -0,0 +1,3 @@
+PACKAGE_SUFFIX: str
+REVISIONS: dict[str, str]
+VERSION: str
diff --git a/runtime/bindings/python/typings/render/cli.pyi b/runtime/bindings/python/typings/render/cli.pyi
new file mode 100644
index 0000000..df152a5
--- /dev/null
+++ b/runtime/bindings/python/typings/render/cli.pyi
@@ -0,0 +1 @@
+def main(argv: list[str] | None = None) -> int: ...
