diff --git a/pykokkos/__init__.py b/pykokkos/__init__.py index 55b743b2..ef4b20f8 100644 --- a/pykokkos/__init__.py +++ b/pykokkos/__init__.py @@ -19,6 +19,8 @@ log1p, sqrt) from pykokkos.lib.info import iinfo, finfo +from pykokkos.lib.create import zeros +from pykokkos.lib.util import all, any runtime_singleton.runtime = Runtime() defaults: Optional[CompilationDefaults] = runtime_singleton.runtime.compiler.read_defaults() diff --git a/pykokkos/interface/data_types.py b/pykokkos/interface/data_types.py index d6850554..0cc35108 100644 --- a/pykokkos/interface/data_types.py +++ b/pykokkos/interface/data_types.py @@ -30,41 +30,42 @@ class DataTypeClass: pass -int16 = DataType.int16 +class int16(DataTypeClass): + value = kokkos.int16 +class int32(DataTypeClass): + value = kokkos.int32 -int32 = DataType.int32 - - -int64 = DataType.int64 - +class int64(DataTypeClass): + value = kokkos.int64 class uint16(DataTypeClass): - pass + value = kokkos.uint16 class uint32(DataTypeClass): - pass + value = kokkos.int32 class uint64(DataTypeClass): - pass - + value = kokkos.int64 -float = DataType.float +class float(DataTypeClass): + value = kokkos.float -double = DataType.double -double.__name__ = "double" # type: ignore +class double(DataTypeClass): + value = kokkos.double class real(DataTypeClass): - pass + value = None class float32(DataTypeClass): - pass + value = kokkos.float class float64(DataTypeClass): - pass + value = kokkos.double -bool = DataType.bool +class bool(DataTypeClass): + value = kokkos.int16 diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index f5cdfab6..4b688246 100644 --- a/pykokkos/interface/views.py +++ b/pykokkos/interface/views.py @@ -1,5 +1,6 @@ from __future__ import annotations import ctypes +import os import math from enum import Enum import sys @@ -10,6 +11,7 @@ import numpy as np +import pykokkos as pk from pykokkos.bindings import kokkos import pykokkos.kokkos_manager as km @@ -131,6 +133,8 @@ def __len__(self) -> int: :returns: the length of the first dimension """ + if len(self.shape) == 0: + return 0 return self.shape[0] def __iter__(self) -> Iterator: @@ -240,6 +244,8 @@ def _init_view( """ self.shape: List[int] = shape + if self.shape == [0]: + self.shape = () self.size = math.prod(shape) self.dtype: Optional[DataType] = self._get_type(dtype) if self.dtype is None: @@ -259,9 +265,19 @@ def _init_view( self.layout: Layout = layout self.trait: Trait = trait + if self.dtype == pk.float: + self.dtype = DataType.float + elif self.dtype == pk.double: + self.dtype = DataType.double + elif self.dtype == pk.int32: + self.dtype = DataType.int32 + elif self.dtype == pk.int64: + pass if trait is trait.Unmanaged: self.array = kokkos.unmanaged_array(array, dtype=self.dtype.value, space=self.space.value, layout=self.layout.value) else: + if len(shape) == 0: + shape = [1] self.array = kokkos.array("", shape, None, None, self.dtype.value, space.value, layout.value, trait.value) self.data = np.array(self.array, copy=False) @@ -283,6 +299,9 @@ def _get_type(self, dtype: Union[DataType, type]) -> Optional[DataType]: if dtype is real: return DataType[km.get_default_precision().__name__] + if dtype == DataType.int64: + dtype = int64 + return dtype if dtype is int: @@ -326,10 +345,14 @@ def __init__(self, parent_view: Union[Subview, View], data_slice: Union[slice, T self.base_view: View = self._get_base_view(parent_view) self.data: np.ndarray = parent_view.data[data_slice] + self.dtype = parent_view.dtype self.array = kokkos.array( self.data, dtype=parent_view.dtype.value, space=parent_view.space.value, layout=parent_view.layout.value, trait=kokkos.Unmanaged) self.shape: List[int] = list(self.data.shape) + if self.data.shape == (0,): + self.data = np.array([], dtype=self.dtype) + self.shape = () self.parent_slice: List[Union[int, slice]] self.parent_slice = self._create_slice(data_slice) @@ -373,6 +396,13 @@ def _get_base_view(self, parent_view: Union[Subview, View]) -> View: return base_view + def __eq__(self, other): + if isinstance(other, View): + if len(self.data) == 0 and len(other.data) == 0: + return True + result_of_eq = self.data == other.data + return result_of_eq + def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: Optional[Layout] = None) -> ViewType: """ Create a PyKokkos View from a numpy array @@ -402,6 +432,8 @@ def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: O dtype = DataType.float # PyKokkos float elif np_dtype is np.float64: dtype = double + elif np_dtype is np.bool_: + dtype = int16 else: raise RuntimeError(f"ERROR: unsupported numpy datatype {np_dtype}") @@ -420,8 +452,8 @@ def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: O # TODO: pykokkos support for 0-D arrays? # temporary/terrible hack here for array API testing.. if array.ndim == 0: - ret_list = list((1,)) - array = [0] + ret_list = () + array = np.array(()) else: ret_list = list((array.shape)) diff --git a/pykokkos/lib/create.py b/pykokkos/lib/create.py new file mode 100644 index 00000000..6218dcd8 --- /dev/null +++ b/pykokkos/lib/create.py @@ -0,0 +1,4 @@ +import pykokkos as pk + +def zeros(shape, *, dtype=None, device=None): + return pk.View([*shape], dtype=dtype) diff --git a/pykokkos/lib/info.py b/pykokkos/lib/info.py index b90ca23c..bb8b088d 100644 --- a/pykokkos/lib/info.py +++ b/pykokkos/lib/info.py @@ -25,8 +25,13 @@ def iinfo(type_or_arr): # values from the NumPy equivalent if "int32" in str(type_or_arr): return info_type_attrs(bits=32, - min=2147483647, - max=-2147483648) + max=2147483647, + min=-2147483648) + elif "int16" in str(type_or_arr): + # iinfo(min=-32768, max=32767, dtype=int16) + return info_type_attrs(bits=16, + min=-32768, + max=32767) elif "int64" in str(type_or_arr): return info_type_attrs(bits=64, min=-9223372036854775808, diff --git a/pykokkos/lib/util.py b/pykokkos/lib/util.py new file mode 100644 index 00000000..7c7921a6 --- /dev/null +++ b/pykokkos/lib/util.py @@ -0,0 +1,21 @@ +import pykokkos as pk + +import numpy as np + +# TODO: add proper implementations rather +# than wrapping NumPy +# These are required for the array API: +# https://data-apis.org/array-api/2021.12/API_specification/utility_functions.html + +def all(x, /, *, axis=None, keepdims=False): + if x == True: + return True + elif x == False: + return False + np_result = np.all(x) + ret_val = pk.View(pk.from_numpy(np.all(x))) + return ret_val + + +def any(x, /, *, axis=None, keepdims=False): + return pk.View(pk.from_numpy(np.any(x)))