Skip to content

Commit

Permalink
ENH: subtract() to API standard
Browse files Browse the repository at this point in the history
* bring the `subtract()` ufunc up to the API standard,
and turn the matching test on in the CI

* this is similar to kokkosgh-157, with similar broadcasting
memory drawback
  • Loading branch information
tylerjereddy committed Jan 16, 2023
1 parent ffd607c commit 5280ea5
Show file tree
Hide file tree
Showing 4 changed files with 426 additions and 33 deletions.
4 changes: 2 additions & 2 deletions .github/workflows/array_api.yml
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ jobs:
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
python setup.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF -DENABLE_VIEW_RANKS=5
- name: Install pykokkos
run: |
python -m pip install .
Expand All @@ -49,4 +49,4 @@ jobs:
# for hypothesis-driven test case generation
pytest $GITHUB_WORKSPACE/pre_compile_tools/pre_compile_ufuncs.py -s
# 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 array_api_tests/test_constants.py array_api_tests/meta/test_utils.py array_api_tests/test_creation_functions.py::test_ones array_api_tests/test_creation_functions.py::test_ones_like array_api_tests/test_data_type_functions.py::test_result_type array_api_tests/test_operators_and_elementwise_functions.py::test_log10 array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt array_api_tests/test_operators_and_elementwise_functions.py::test_isfinite array_api_tests/test_operators_and_elementwise_functions.py::test_log2 array_api_tests/test_operators_and_elementwise_functions.py::test_log1p array_api_tests/test_operators_and_elementwise_functions.py::test_isinf array_api_tests/test_operators_and_elementwise_functions.py::test_log array_api_tests/test_array_object.py::test_scalar_casting array_api_tests/test_operators_and_elementwise_functions.py::test_sign array_api_tests/test_operators_and_elementwise_functions.py::test_square array_api_tests/test_operators_and_elementwise_functions.py::test_cos array_api_tests/test_operators_and_elementwise_functions.py::test_round array_api_tests/test_operators_and_elementwise_functions.py::test_trunc array_api_tests/test_operators_and_elementwise_functions.py::test_ceil array_api_tests/test_operators_and_elementwise_functions.py::test_floor
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 array_api_tests/test_constants.py array_api_tests/meta/test_utils.py array_api_tests/test_creation_functions.py::test_ones array_api_tests/test_creation_functions.py::test_ones_like array_api_tests/test_data_type_functions.py::test_result_type array_api_tests/test_operators_and_elementwise_functions.py::test_log10 array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt array_api_tests/test_operators_and_elementwise_functions.py::test_isfinite array_api_tests/test_operators_and_elementwise_functions.py::test_log2 array_api_tests/test_operators_and_elementwise_functions.py::test_log1p array_api_tests/test_operators_and_elementwise_functions.py::test_isinf array_api_tests/test_operators_and_elementwise_functions.py::test_log array_api_tests/test_array_object.py::test_scalar_casting array_api_tests/test_operators_and_elementwise_functions.py::test_sign array_api_tests/test_operators_and_elementwise_functions.py::test_square array_api_tests/test_operators_and_elementwise_functions.py::test_cos array_api_tests/test_operators_and_elementwise_functions.py::test_round array_api_tests/test_operators_and_elementwise_functions.py::test_trunc array_api_tests/test_operators_and_elementwise_functions.py::test_ceil array_api_tests/test_operators_and_elementwise_functions.py::test_floor "array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[subtract(x1, x2)]"
2 changes: 1 addition & 1 deletion .github/workflows/main_ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ jobs:
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
python setup.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF -DENABLE_VIEW_RANKS=5
- name: Install pykokkos
run: |
python -m pip install .
Expand Down
350 changes: 350 additions & 0 deletions pykokkos/lib/ufunc_workunits.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,356 @@
import pykokkos as pk


@pk.workunit
def subtract_impl_1d_int64(tid: int, view1: pk.View1D[pk.int64], view2: pk.View1D[pk.int64], out: pk.View1D[pk.int64]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_int64(tid: int, view1: pk.View2D[pk.int64], view2: pk.View2D[pk.int64], out: pk.View2D[pk.int64]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_int64(tid: int, view1: pk.View3D[pk.int64], view2: pk.View3D[pk.int64], out: pk.View3D[pk.int64]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_int64(tid: int, view1: pk.View4D[pk.int64], view2: pk.View4D[pk.int64], out: pk.View4D[pk.int64]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_int64(tid: int, view1: pk.View5D[pk.int64], view2: pk.View5D[pk.int64], out: pk.View5D[pk.int64]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_int32(tid: int, view1: pk.View1D[pk.int32], view2: pk.View1D[pk.int32], out: pk.View1D[pk.int32]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_int32(tid: int, view1: pk.View2D[pk.int32], view2: pk.View2D[pk.int32], out: pk.View2D[pk.int32]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_int32(tid: int, view1: pk.View3D[pk.int32], view2: pk.View3D[pk.int32], out: pk.View3D[pk.int32]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_int32(tid: int, view1: pk.View4D[pk.int32], view2: pk.View4D[pk.int32], out: pk.View4D[pk.int32]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_int32(tid: int, view1: pk.View5D[pk.int32], view2: pk.View5D[pk.int32], out: pk.View5D[pk.int32]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_int16(tid: int, view1: pk.View1D[pk.int16], view2: pk.View1D[pk.int16], out: pk.View1D[pk.int16]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_int16(tid: int, view1: pk.View2D[pk.int16], view2: pk.View2D[pk.int16], out: pk.View2D[pk.int16]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_int16(tid: int, view1: pk.View3D[pk.int16], view2: pk.View3D[pk.int16], out: pk.View3D[pk.int16]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_int16(tid: int, view1: pk.View4D[pk.int16], view2: pk.View4D[pk.int16], out: pk.View4D[pk.int16]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_int16(tid: int, view1: pk.View5D[pk.int16], view2: pk.View5D[pk.int16], out: pk.View5D[pk.int16]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_int8(tid: int, view1: pk.View1D[pk.int8], view2: pk.View1D[pk.int8], out: pk.View1D[pk.int8]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_int8(tid: int, view1: pk.View2D[pk.int8], view2: pk.View2D[pk.int8], out: pk.View2D[pk.int8]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_int8(tid: int, view1: pk.View3D[pk.int8], view2: pk.View3D[pk.int8], out: pk.View3D[pk.int8]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_int8(tid: int, view1: pk.View4D[pk.int8], view2: pk.View4D[pk.int8], out: pk.View4D[pk.int8]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_int8(tid: int, view1: pk.View5D[pk.int8], view2: pk.View5D[pk.int8], out: pk.View5D[pk.int8]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_uint64(tid: int, view1: pk.View1D[pk.uint64], view2: pk.View1D[pk.uint64], out: pk.View1D[pk.uint64]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_uint64(tid: int, view1: pk.View2D[pk.uint64], view2: pk.View2D[pk.uint64], out: pk.View2D[pk.uint64]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_uint64(tid: int, view1: pk.View3D[pk.uint64], view2: pk.View3D[pk.uint64], out: pk.View3D[pk.uint64]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_uint64(tid: int, view1: pk.View4D[pk.uint64], view2: pk.View4D[pk.uint64], out: pk.View4D[pk.uint64]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_uint64(tid: int, view1: pk.View5D[pk.uint64], view2: pk.View5D[pk.uint64], out: pk.View5D[pk.uint64]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_uint32(tid: int, view1: pk.View1D[pk.uint32], view2: pk.View1D[pk.uint32], out: pk.View1D[pk.uint32]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_uint32(tid: int, view1: pk.View2D[pk.uint32], view2: pk.View2D[pk.uint32], out: pk.View2D[pk.uint32]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_uint32(tid: int, view1: pk.View3D[pk.uint32], view2: pk.View3D[pk.uint32], out: pk.View3D[pk.uint32]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_uint32(tid: int, view1: pk.View4D[pk.uint32], view2: pk.View4D[pk.uint32], out: pk.View4D[pk.uint32]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_uint32(tid: int, view1: pk.View5D[pk.uint32], view2: pk.View5D[pk.uint32], out: pk.View5D[pk.uint32]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_uint16(tid: int, view1: pk.View1D[pk.uint16], view2: pk.View1D[pk.uint16], out: pk.View1D[pk.uint16]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_uint16(tid: int, view1: pk.View2D[pk.uint16], view2: pk.View2D[pk.uint16], out: pk.View2D[pk.uint16]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_uint16(tid: int, view1: pk.View3D[pk.uint16], view2: pk.View3D[pk.uint16], out: pk.View3D[pk.uint16]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_uint16(tid: int, view1: pk.View4D[pk.uint16], view2: pk.View4D[pk.uint16], out: pk.View4D[pk.uint16]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_uint16(tid: int, view1: pk.View5D[pk.uint16], view2: pk.View5D[pk.uint16], out: pk.View5D[pk.uint16]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_uint8(tid: int, view1: pk.View1D[pk.uint8], view2: pk.View1D[pk.uint8], out: pk.View1D[pk.uint8]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_uint8(tid: int, view1: pk.View2D[pk.uint8], view2: pk.View2D[pk.uint8], out: pk.View2D[pk.uint8]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_uint8(tid: int, view1: pk.View3D[pk.uint8], view2: pk.View3D[pk.uint8], out: pk.View3D[pk.uint8]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_uint8(tid: int, view1: pk.View4D[pk.uint8], view2: pk.View4D[pk.uint8], out: pk.View4D[pk.uint8]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_uint8(tid: int, view1: pk.View5D[pk.uint8], view2: pk.View5D[pk.uint8], out: pk.View5D[pk.uint8]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_float(tid: int, view1: pk.View1D[pk.float], view2: pk.View1D[pk.float], out: pk.View1D[pk.float]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_float(tid: int, view1: pk.View2D[pk.float], view2: pk.View2D[pk.float], out: pk.View2D[pk.float]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_float(tid: int, view1: pk.View3D[pk.float], view2: pk.View3D[pk.float], out: pk.View3D[pk.float]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_float(tid: int, view1: pk.View4D[pk.float], view2: pk.View4D[pk.float], out: pk.View4D[pk.float]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_float(tid: int, view1: pk.View5D[pk.float], view2: pk.View5D[pk.float], out: pk.View5D[pk.float]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def subtract_impl_1d_double(tid: int, view1: pk.View1D[pk.double], view2: pk.View1D[pk.double], out: pk.View1D[pk.double]):
out[tid] = view1[tid] - view2[tid]


@pk.workunit
def subtract_impl_2d_double(tid: int, view1: pk.View2D[pk.double], view2: pk.View2D[pk.double], out: pk.View2D[pk.double]):
for i in range(view1.extent(1)):
out[tid][i] = view1[tid][i] - view2[tid][i]


@pk.workunit
def subtract_impl_3d_double(tid: int, view1: pk.View3D[pk.double], view2: pk.View3D[pk.double], out: pk.View3D[pk.double]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
out[tid][i][j] = view1[tid][i][j] - view2[tid][i][j]


@pk.workunit
def subtract_impl_4d_double(tid: int, view1: pk.View4D[pk.double], view2: pk.View4D[pk.double], out: pk.View4D[pk.double]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
out[tid][i][j][k] = view1[tid][i][j][k] - view2[tid][i][j][k]


@pk.workunit
def subtract_impl_5d_double(tid: int, view1: pk.View5D[pk.double], view2: pk.View5D[pk.double], out: pk.View5D[pk.double]):
for i in range(view1.extent(1)):
for j in range(view1.extent(2)):
for k in range(view1.extent(3)):
for l in range(view1.extent(4)):
out[tid][i][j][k][l] = view1[tid][i][j][k][l] - view2[tid][i][j][k][l]


@pk.workunit
def floor_impl_1d_double(tid: int, view: pk.View1D[pk.double], out: pk.View1D[pk.double]):
out[tid] = floor(view[tid])
Expand Down
Loading

0 comments on commit 5280ea5

Please sign in to comment.