Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable usage of multiple CUDA GPUs in PyKokkos #85

Closed
wants to merge 27 commits into from

Conversation

NaderAlAwar
Copy link
Contributor

@NaderAlAwar NaderAlAwar commented Sep 19, 2022

This PR enables using multiple CUDA GPUs in PyKokkos in the same process. It depends on another pykokkos-base PR which makes multiple copies of libpykokkos, libkokkoscore, and libkokkoscontainers. PyKokkos will then link kernels compiled for CUDA to a different set of libraries depending on the GPU specified by the user.

The concept of Multi-GPU Kokkos is based on William Ruys's work here

NaderAlAwar and others added 27 commits August 29, 2022 14:20
README: update installation instructions
Fixes kokkos#86

* prohibit negative values in `pk.uint32`/`pk.uint64`-specified
views, and match the "wrapping" behavior used by NumPy in these
cases
* we can now pass an additional testing module
in the array API suite, `test_utils.py`, which
has also been added to the CI workflow now

* `DataTypeClass` subclasses have been given
a `np_equiv` attribute to facilitate interopartion
with NumPy types; admittedly, this is probably
just a temporary adjustment as the type system
matures

* `View` instances now have slightly more robust
`0-D` array handling because of two adjustments:
- lean a bit more heavily on NumPy in these cases
- start using a custom `__eq__` for scalar comparisons (which
apparently also requires a custom `__hash__` for other
parts of `pykokkos`)

* my primary metric that this is actually "progress" instead
of just leaning more heavily on NumPy is that we can pass
more array API tests now; ultimately, this may just
result in deferring proper implementations for some things,
though being able to satisfy the array API with some parts
falling back to NumPy probably isn't a bad thing anyway, at
least at first

* no changes to the `pykokkos` test suite were needed
* `bool` is now aliased to a more appropriate
`pykokkos` type -- `uint8`

* `int8` and `uint8` types have been updated
to include appropriate `np_equiv` attributes

* fix a broken conditional chain in `from_numpy()`
* this splits off a few cleanups from kokkosgh-73

* there is some pointless type remapping
code in the `View` class that has been removed

* removing the above code causes the test suite
to fail because of incorrect `uint32`/`uint64`
`value` mappings in `data_types` module, so fix
those to allow tests to pass
* implement `pk.ones()` such that
`array_api_tests/test_creation_functions.py::test_ones` passes

* the Python array API standard `test_ones()` makes use
of `pk.equal()`, which wasn't implemented, so an early-stage
implementation is added here (which should help later on as well,
since the need to occasionally do `view1 == view2` should be
obvious)

* this is no doubt a bit of a "hack," especially in terms
of full-featured broadcasting for `pk.equal`, which is currently only
supported for "scalar" broadcasting in my work here

* I was a bit surprised that I only needed to add limited 1D
support for `pk.equal()`, but I imagine the code there will grow
substantially to support more scenarios/remove hacks in the future

* note that this branch also borrows some of the typing
changes from kokkosgh-67 (and improves on them slightly I think)

* special shims are present for some `0-D` type scenarios,
and for now I've doubled down on the adoption of an unsigned
integer type as a proxy for a `bool` array--I suggest we at least
use a shorter width type for this in the future though, once
the pykokkos-base bindings are merged

* although I'm leaning on the Python array API standard tests here,
I did check that i.e., this works:

```python
import pykokkos as pk

def main():
    a = pk.ones((3, 3))
    print(a)

if __name__ == "__main__":
    main()
```

```
[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]]
```

* note, however, that the hook-in of `pk.equal()` to `__equal__`
wasn't the focus of this work (since the array API tests simply
call `pk.equal()` for the isolated test I was working on), so
for now:

```python
import pykokkos as pk

def main():
    a = pk.ones((3,))
    b = pk.ones((3,))
    result = a == b
    print(result)
    result2 = pk.equal(a, b)
    print(result2)

if __name__ == "__main__":
    main()
```

Produces:

```
False
[1 1 1]
```

* what we really want is `True` and `[True True True]`,
but the latter isn't too bad at least--the `1` values are
a proxy for truth in the `pk.uint16` type we're using
in place of `bool` for now; I suggest we delay the hook-in
to `__equal__` for now...
* updated the `equal()` ufunc code in this branch
to correctly use the new `uint8` type instead of
`uint16` as a temporary proxy for a `bool` type

* removed extraneous shape handling code
in `ones()` and `View`, as pointed out by
reviewer

* `View` `shape` argument type adjusted
to list OR tuple based on reviewer feedback

* simplified the `equal()` ufunc by removing
some of the `np.ndarray()` support shims--it seems
we don't have testable evidence that this is currently
a requirement
* add an option to report the top `N` slowest
tests when running `runtests.py`, which basically
just passes through to the built-in pytest option:
https://docs.pytest.org/en/7.1.x/how-to/usage.html#profiling-test-execution-duration

* I'm hoping not to build many more options in though,
I just wanted this one because we'll want a convenient way
to keep an eye on what tests are consuming time as our suite
grows

* some sample incantations/results:

- `python runtests.py -d 10`

```
===================================================================================================================================================
slowest 10 durations
===================================================================================================================================================
11.51s call     tests/test_ufuncs.py::test_caching
11.21s call
   tests/test_hierarchical.py::TestHierarchical::test_outer_for
9.23s call
  tests/test_AST_translator.py::TestASTTranslator::test_ann_assign
7.95s call
  tests/test_kokkosfunctions_translator.py::TestKokkosFunctionsTranslator::test_args_sum
7.83s call     tests/test_atomics.py::TestAtomic::test_atomic_add
6.27s call
  tests/test_ops_translator.py::TestOpsTranslator::test_add_op
5.36s call
  tests/test_hierarchical.py::TestHierarchical::test_yAx_vector
5.10s call
  tests/test_ufuncs.py::test_matmul_1d_exposed_ufuncs_vs_numpy[double-float64-matmul-matmul]
5.08s call
  tests/test_ufuncs.py::test_matmul_1d_exposed_ufuncs_vs_numpy[float-float32-matmul-matmul]
4.92s call
  tests/test_ufuncs.py::test_multi_array_1d_exposed_ufuncs_vs_numpy[float-float32-greater-greater]
============================================================================================================================
203 passed, 9 skipped, 9 xfailed, 16 warnings in 450.11s (0:07:30)
    ============================================================================================================================

```

- `python runtests.py -d 5 -t tests/test_views.py`

```
===================================================================================================================================================
slowest 5 durations
====================================================================================================================================================
0.02s call
  tests/test_views.py::test_asarray_consts_vs_numpy[None-None-nan]
0.02s call
  tests/test_views.py::test_sizes[input_arr2-view_dims2-View3D]
0.02s call
  tests/test_views.py::test_asarray_consts_vs_numpy[int32-int32-2.718281828459045]
0.02s call
  tests/test_views.py::test_sizes[input_arr1-view_dims1-View2D]
0.02s call
  tests/test_views.py::test_asarray_consts_vs_numpy[int64-int64-2.718281828459045]
==============================================================================================================================================
31 passed, 9 skipped in 0.78s
   ===============================================================================================================================================

```
* implement `pk.ones()` such that
`array_api_tests/test_creation_functions.py::test_ones` passes

* the Python array API standard `test_ones()` makes use
of `pk.equal()`, which wasn't implemented, so an early-stage
implementation is added here (which should help later on as well,
since the need to occasionally do `view1 == view2` should be
obvious)

* this is no doubt a bit of a "hack," especially in terms
of full-featured broadcasting for `pk.equal`, which is currently only
supported for "scalar" broadcasting in my work here

* I was a bit surprised that I only needed to add limited 1D
support for `pk.equal()`, but I imagine the code there will grow
substantially to support more scenarios/remove hacks in the future

* note that this branch also borrows some of the typing
changes from kokkosgh-67 (and improves on them slightly I think)

* special shims are present for some `0-D` type scenarios,
and for now I've doubled down on the adoption of an unsigned
integer type as a proxy for a `bool` array--I suggest we at least
use a shorter width type for this in the future though, once
the pykokkos-base bindings are merged

* although I'm leaning on the Python array API standard tests here,
I did check that i.e., this works:

```python
import pykokkos as pk

def main():
    a = pk.ones((3, 3))
    print(a)

if __name__ == "__main__":
    main()
```

```
[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]]
```

* note, however, that the hook-in of `pk.equal()` to `__equal__`
wasn't the focus of this work (since the array API tests simply
call `pk.equal()` for the isolated test I was working on), so
for now:

```python
import pykokkos as pk

def main():
    a = pk.ones((3,))
    b = pk.ones((3,))
    result = a == b
    print(result)
    result2 = pk.equal(a, b)
    print(result2)

if __name__ == "__main__":
    main()
```

Produces:

```
False
[1 1 1]
```

* what we really want is `True` and `[True True True]`,
but the latter isn't too bad at least--the `1` values are
a proxy for truth in the `pk.uint16` type we're using
in place of `bool` for now; I suggest we delay the hook-in
to `__equal__` for now...
* changes needed to pass
`array_api_tests/test_creation_functions.py::test_ones_like`

* this builds on kokkosgh-73, and is meant to be reviewed after
that one
* cleanups after rebase
ENH: changes needed to pass `array_api_tests/test_creation_functions.py::test_ones_like`
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants