adding preliminary AMDGPU support #99
Draft
+328
−207
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.
I am still running the tests, so this PR is still a draft and might change in the next few days, but it adds preliminary support for AMDGPU devices.
There were a few small things to note:
CuArray
to use an Array TypeAT
value instead. There were a few places (such as inzygote.jl
with internal functions that were only used for the GPU. In these cases,AT
became a function argument. In most cases,AT
could be created without any change to the function calls, themselves.interactions/implicit_solver.jl
), I added a new conditional withROCArray
instead ofCuArray
System(...)
call was modified to try and figure out what device you wanted to use if you setgpu
to true, but will default toCUDA
if both GPUs are present and things are left unspecifiedrun_cuda_tests
andrun_rocm_tests
. It will then populate the arraygpu_array_types
withCuArray
,ROCArray
, or both. I tried to modify all the tests to use the array types from this array.Again, still messing around with things now, but thought I would put the draft PR up in the case that there are design decisions to talk about.