Skip to content

Commit

Permalink
ENH: support int8
Browse files Browse the repository at this point in the history
* add support for 8-bit signed integers with values that
should exist on the interval `[-128, 127]` for conformance
with the array API:
https://data-apis.org/array-api/latest/API_specification/data_types.html
kokkos/pykokkos#57

* add a regression test for kokkos<->NumPy type equivalencies that
includes the new `int8` type; I was hoping to do some simple
arithmetic testing as well, but not sure how practical that is
within `pykokkos-base` proper (maybe defer that kind of thing
to `pykokkos`?)
  • Loading branch information
tylerjereddy committed Aug 23, 2022
1 parent 87290b3 commit a38531a
Show file tree
Hide file tree
Showing 6 changed files with 41 additions and 3 deletions.
3 changes: 2 additions & 1 deletion include/fwd.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,8 @@ struct Device;
//----------------------------------------------------------------------------//

enum KokkosViewDataType {
Int16 = 0,
Int8 = 0,
Int16,
Int32,
Int64,
Uint16,
Expand Down
1 change: 1 addition & 0 deletions include/traits.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ VIEW_DATA_DIMS(8, T ********)
// the first string identifier is the "canonical name" (i.e. what gets encoded)
// and the remaining string entries are used to generate aliases
//
VIEW_DATA_TYPE(int8_t, Int8, "int8", "short_short")
VIEW_DATA_TYPE(int16_t, Int16, "int16", "short")
VIEW_DATA_TYPE(int32_t, Int32, "int32", "int")
VIEW_DATA_TYPE(int64_t, Int64, "int64", "long")
Expand Down
3 changes: 2 additions & 1 deletion kokkos/__init__.py.in
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,8 @@ try:
"read_dtype",
"initialize", # bindings
"finalize",
"int16", # data types
"int8", # data types
"int16",
"int32",
"int64",
"uint16",
Expand Down
33 changes: 33 additions & 0 deletions kokkos/test/test_types.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import kokkos

import numpy as np
import pytest


@pytest.mark.parametrize("type_val, expected_np_type", [
(kokkos.int8, np.int8),
(kokkos.int16, np.int16),
(kokkos.int32, np.int32),
(kokkos.int64, np.int64),
(kokkos.uint16, np.uint16),
(kokkos.uint32, np.uint32),
(kokkos.uint64, np.uint64),
(kokkos.float32, np.float32),
(kokkos.float64, np.float64),
(kokkos.float, np.float32),
(kokkos.double, np.float64),
(kokkos.short, np.int16),
(kokkos.int, np.int32),
(kokkos.long, np.int64),
])
def test_basic_type_equiv(type_val, expected_np_type):
# test some view to NumPy array type equivalencies
kokkos.initialize()
view = kokkos.array([2],
dtype=type_val,
space=kokkos.DefaultHostMemorySpace)

# NOTE: copy can still happen (attempt no copy,
# not guarantee)
arr = np.array(view, copy=False)
assert arr.dtype == expected_np_type
2 changes: 2 additions & 0 deletions kokkos/utility.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,8 @@ def read_dtype(_dtype):
try:
import numpy as np

if _dtype == np.int8:
return lib.int8
if _dtype == np.int16:
return lib.int16
elif _dtype == np.int32:
Expand Down
2 changes: 1 addition & 1 deletion src/variants/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ TARGET_LINK_LIBRARIES(libpykokkos-variants PUBLIC

SET(_types concrete dynamic)
SET(_variants layout memory_trait)
SET(_data_types Int16 Int32 Int64 Uint16 Uint32 Uint64 Float32 Float64)
SET(_data_types Int8 Int16 Int32 Int64 Uint16 Uint32 Uint64 Float32 Float64)

SET(layout_enums Right)
SET(memory_trait_enums Managed)
Expand Down

0 comments on commit a38531a

Please sign in to comment.