Enable usage of multiple CUDA GPUs in PyKokkos #85
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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
BUG: uint32/64 wrapping
* 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()`
ENH, TST: more 0-D handling
* 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
MAINT, BUG: typing cleanups
* 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 =============================================================================================================================================== ```
ENH: pk.ones() to API standard
TST, MAINT: test durations
* 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`
… kokkos paths when multi gpu is not enabled
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR enables using multiple CUDA GPUs in PyKokkos in the same process. It depends on another
pykokkos-basePR which makes multiple copies oflibpykokkos,libkokkoscore, andlibkokkoscontainers. 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