WIP, ENH: array API start

* support the bare minimum required types per: https://data-apis.org/array-api/2021.12/API_specification/data_types.html (otherwise, we can't even import the array API tests, let alone run them--for now, some types are basically just "stubs" and/or aliases so we can get things rolling) * support `iinfo` and `finfo` functions required by the standard: https://data-apis.org/array-api/2021.12/API_specification/data_type_functions.html#objects-in-api * needed to add a (hacky) implemention of `asarray()` to comply with the array API standard per: https://data-apis.org/array-api/2021.12/API_specification/creation_functions.html#objects-in-api Note that the positional-only and keyword-only function signature is also mandatory. * the array API test suite detected multiple issue in our data type system; these mostly seemed to stem from having both `DataType(Enum)` and `DataTypeClass`, which is a model that is not consistent with the mappings expected by the standard; it is also mandatory to support more types than we currently do, so I've hacked around some of these issues for now, but with boolean type we'll currently fail the array API tests * we'll need to double check what to do for `0` dimensional arrays, but they are used in the array API test suite, so I've added a temporary hack around those until we support them "natively?" * add a CI job that starts to test for array API compliance--it will fail, but it should at least start running rather than erroring out at the import stage, which is a step forward from `develop` branch * the usual mypy ignores, at least for now..
kokkos · Aug 9, 2022 · a949831 · a949831
1 parent 6f6d199
commit a949831
Show file tree

Hide file tree

Showing 8 changed files with 162 additions and 16 deletions.
diff --git a/.github/workflows/array_api.yml b/.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/test_array_object.py::test_getitem
diff --git a/mypy.ini b/mypy.ini
@@ -92,3 +92,9 @@ ignore_errors = True
 [mypy-pykokkos.core.translators.bindings]
 ignore_errors = True
 
+[mypy-pykokkos.lib.ufuncs]
+ignore_errors = True
+
+[mypy-pykokkos.lib.info]
+ignore_missing_imports = True
+
diff --git a/pykokkos/__init__.py b/pykokkos/__init__.py
@@ -12,12 +12,13 @@
 )
 
 initialize()
-from pykokkos.lib.ufuncs import (reciprocal, # type: ignore
+from pykokkos.lib.ufuncs import (reciprocal,
                                  log,
                                  log2,
                                  log10,
                                  log1p,
                                  sqrt)
+from pykokkos.lib.info import iinfo, finfo
 
 runtime_singleton.runtime = Runtime()
 defaults: Optional[CompilationDefaults] = runtime_singleton.runtime.compiler.read_defaults()

diff --git a/pykokkos/interface/__init__.py b/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,
 )
 
 

diff --git a/pykokkos/interface/data_types.py b/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,23 +13,30 @@ 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
+int16 = DataType.int16
 
 
-class int32(DataTypeClass):
-    pass
+int32 = DataType.int32
 
 
-class int64(DataTypeClass):
-    pass
+int64 = DataType.int64
 
 
 class uint16(DataTypeClass):
@@ -43,13 +51,20 @@ class uint64(DataTypeClass):
     pass
 
 
-class float(DataTypeClass):
-    pass
+float = DataType.float
 
 
-class double(DataTypeClass):
-    pass
+double = DataType.double
+double.__name__ = "double" # type: ignore
 
 
 class real(DataTypeClass):
-    pass
+    pass
+
+class float32(DataTypeClass):
+    pass
+
+class float64(DataTypeClass):
+    pass
+
+bool = DataType.bool
diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py
@@ -283,7 +283,7 @@ 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__]
+            return dtype
 
         if dtype is int:
             return DataType["int32"]
@@ -417,7 +417,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 = list((1,))
+        array = [0]
+    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 +473,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]):

diff --git a/pykokkos/lib/__init__.py b/pykokkos/lib/__init__.py
diff --git a/pykokkos/lib/info.py b/pykokkos/lib/info.py
@@ -0,0 +1,50 @@
+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,
+                               min=2147483647,
+                               max=-2147483648)
+    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)
+