Cuda multi-GPU support: Allow execution space instance constructor to run #6706
+56
−6
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
masterleinad
force-pushed
the
cuda_multiple_devices_init_exec
branch
from
0e9fa7c
to
2a53772
Compare
ldh4
reviewed
Jan 9, 2024
dalg24
reviewed
Jan 9, 2024
Comment on lines
+27
to
+33
| KOKKOS_IMPL_CUDA_SAFE_CALL(cudaSetDevice(0)); | ||
| cudaStream_t stream0; | ||
| KOKKOS_IMPL_CUDA_SAFE_CALL(cudaStreamCreate(&stream0)); | ||
|
|
||
| KOKKOS_IMPL_CUDA_SAFE_CALL(cudaSetDevice(n_devices - 1)); | ||
| cudaStream_t stream; | ||
| KOKKOS_IMPL_CUDA_SAFE_CALL(cudaStreamCreate(&stream)); |
There was a problem hiding this comment.
Did you mean to mess with the CUDA runtime after Kokkos::initialize()?
There was a problem hiding this comment.
We are trying to make sure that we are always using the correct device even if the user changed it between Kokkos API calls.
dalg24
reviewed
Jan 10, 2024
dalg24
approved these changes
Jan 24, 2024
ldh4
approved these changes
Jan 24, 2024
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.
Part of #6091. This pull request allows running the Cuda execution space constructor with a device ID different from the default one. This basically means that we also handle the internal allocations correctly.
A caveat here (that wasn't exposed in the initial version of #6091) is that we need to keep track of the correct device ID for allocations since we might copy the allocation header back to the host. CallingcudaFreedoesn't seem to require using the correct device.This is handled in a hacky way for the internal allocations in this pull request by setting the singleton's device ID before deallocating the internal allocations and then resetting it to its previous value. This approach allows us to run the full unit test in #6091.
I decided to remove the guard we have for using a device ID different from the default one in this pull request already but I'm also open to introducing a macro that guards it and is only defined for the unit test.