ENH: array API start

.github/workflows/array_api.yml

@@ -0,0 +1,46 @@
+name: Array API Testing
+
+on:
+  push:
+    branches:
+      - develop
+  pull_request:
+    branches:
+      - develop
+
+jobs:
+  test_array_api:
+    strategy:
+      matrix:
+        platform: [ubuntu-latest]
+        python-version: ["3.10"]
+    runs-on: ${{ matrix.platform }}
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install --upgrade numpy mypy cmake pytest pybind11 scikit-build patchelf
+      - name: Install pykokkos-base
+        run: |
+          cd /tmp
+          git clone https://github.com/kokkos/pykokkos-base.git
+          cd pykokkos-base
+          python setup.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF
+      - name: Install pykokkos
+        run: |
+          python -m pip install .
+      - name: Check Array API conformance
+        run: |
+          cd /tmp
+          git clone https://github.com/data-apis/array-api-tests.git
+          cd array-api-tests
+          git submodule update --init
+          pip install -r requirements.txt
+          export ARRAY_API_TESTS_MODULE=pykokkos
+          # only run a subset of the conformance tests to get started
+          pytest array_api_tests/meta/test_broadcasting.py array_api_tests/meta/test_equality_mapping.py array_api_tests/meta/test_signatures.py array_api_tests/meta/test_special_cases.py

mypy.ini

@@ -95,3 +95,5 @@ ignore_errors = True
 [mypy-pykokkos.lib.ufuncs]
 ignore_errors = True
 
+[mypy-pykokkos.lib.info]
+ignore_missing_imports = True

pykokkos/__init__.py

@@ -19,6 +19,9 @@
                                  log1p,
                                  sqrt,
                                  sign)
+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()

pykokkos/interface/__init__.py

@@ -14,7 +14,8 @@
     DataType, DataTypeClass,
     int16, int32, int64,
     uint16, uint32, uint64,
-    float, double, real
+    float, double, real,
+    float32, float64, bool,
 )
 from .decorators import (
     callback, classtype, Decorator, function, functor, main,
@@ -46,7 +47,8 @@
     ScratchView, ScratchView1D, ScratchView2D,
     ScratchView3D, ScratchView4D, ScratchView5D,
     ScratchView6D, ScratchView7D, ScratchView8D,
-    from_cupy, from_numpy
+    from_cupy, from_numpy,
+    asarray,
 )
 
 

pykokkos/interface/data_types.py

@@ -1,6 +1,7 @@
 from enum import Enum
 
 from pykokkos.bindings import kokkos
+import numpy as np
 
 
 class DataType(Enum):
@@ -12,44 +13,59 @@ class DataType(Enum):
     uint64 = kokkos.uint64
     float = kokkos.float
     double = kokkos.double
+    # https://data-apis.org/array-api/2021.12/API_specification/data_types.html
+    # A conforming implementation of the array API standard
+    # must provide and support the dtypes listed above; for
+    # now, we will use aliases and possibly even incorrect/empty
+    # implementations so that we can start testing with the array
+    # API standard suite, otherwise we won't even be able to import
+    # the tests let alone run them
+    float32 = kokkos.float
+    float64 = kokkos.double
     real = None
+    bool = np.bool_
 
 
 class DataTypeClass:
     pass
 
 
 class int16(DataTypeClass):
-    pass
-
+    value = kokkos.int16
 
 class int32(DataTypeClass):
-    pass
-
+    value = kokkos.int32
 
 class int64(DataTypeClass):
-    pass
-
+    value = kokkos.int64
 
 class uint16(DataTypeClass):
-    pass
+    value = kokkos.uint16
 
 
 class uint32(DataTypeClass):
-    pass
+    value = kokkos.int32
 
 
 class uint64(DataTypeClass):
-    pass
+    value = kokkos.int64
 
 
 class float(DataTypeClass):
-    pass
-
+    value = kokkos.float
 
 class double(DataTypeClass):
-    pass
+    value = kokkos.double
 
 
 class real(DataTypeClass):
-    pass
+    value = None
+
+class float32(DataTypeClass):
+    value = kokkos.float
+
+class float64(DataTypeClass):
+    value = kokkos.double
+
+class bool(DataTypeClass):
+    value = kokkos.int16

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,7 +299,10 @@ def _get_type(self, dtype: Union[DataType, type]) -> Optional[DataType]:
             if dtype is real:
                 return DataType[km.get_default_precision().__name__]
 
-            return DataType[dtype.__name__]
+            if dtype == DataType.int64:
+                dtype = int64
+
+            return dtype
 
         if dtype is int:
             return DataType["int32"]
@@ -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}")
 
@@ -417,7 +449,16 @@ def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: O
     if layout is None:
         layout = Layout.LayoutDefault
 
-    return View(list(array.shape), dtype, space=space, trait=Trait.Unmanaged, array=array, layout=layout)
+    # TODO: pykokkos support for 0-D arrays?
+    # temporary/terrible hack here for array API testing..
+    if array.ndim == 0:
+        ret_list = ()
+        array = np.array(())
+    else:
+        ret_list = list((array.shape))
+
+
+    return View(ret_list, dtype, space=space, trait=Trait.Unmanaged, array=array, layout=layout)
 
 def from_cupy(array) -> ViewType:
     """
@@ -464,6 +505,23 @@ def from_cupy(array) -> ViewType:
 
     return from_numpy(np_array, MemorySpace.CudaSpace, layout)
 
+
+# asarray is required for comformance with the array API:
+# https://data-apis.org/array-api/2021.12/API_specification/creation_functions.html#objects-in-api
+
+def asarray(obj, /, *, dtype=None, device=None, copy=None):
+    # TODO: proper implementation/design
+    # for now, let's cheat and use NumPy asarray() followed
+    # by pykokkos from_numpy()
+    if "bool" in str(dtype):
+        dtype = np.bool_
+    arr = np.asarray(obj, dtype=dtype)
+    ret = from_numpy(arr)
+    return ret
+
+
+
+
 T = TypeVar("T")
 
 class View1D(Generic[T]):

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)

pykokkos/lib/info.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass
+
+from pykokkos.bindings import kokkos
+
+# the integer and float type information functions appear
+# to be required by the array API standard
+# i.e.,
+# https://data-apis.org/array-api/2021.12/API_specification/data_type_functions.html#objects-in-api
+
+
+@dataclass
+class info_type_attrs:
+    """
+    Store machine limits for numeric data types.
+    """
+    bits: int
+    max: int
+    min: int
+
+def iinfo(type_or_arr):
+    # TODO: more correct implementation
+    # this is really just an initial hack
+    # so we can run the array API tests,
+    # and effectively just copies return
+    # values from the NumPy equivalent
+    if "int32" in str(type_or_arr):
+        return info_type_attrs(bits=32,
+                               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,
+                               max=9223372036854775807)
+
+
+def finfo(type_or_arr):
+    # TODO: more correct implementation
+    # this is really just an initial hack
+    # so we can run the array API tests,
+    # and effectively just copies return
+    # values from the NumPy equivalent
+    if "float" in str(type_or_arr) and not "float64" in str(type_or_arr):
+        return info_type_attrs(bits=32,
+                               min=-3.4028235e+38,
+                               max=3.4028235e+38,)
+    elif "double" in str(type_or_arr) or "float64" in str(type_or_arr):
+        return info_type_attrs(bits=64,
+                               min=-1.7976931348623157e+308,
+                               max=1.7976931348623157e+308)
+

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)))