Move CUDA stuff to an extension #4499
Conversation
|
Possibly, we should simply implement a CUDA extension in this PR with appropriate organization of the code and get on with the breaking change! tl;dr then after this is merged, anybody doing computations on nvidia GPU has to write using Oceananigans
using CUDA |
|
@simone-silvestri curious to hear your thoughts |
|
I think it's a good idea. It provides templates to add new architectures and makes the code completely architecture agnostic. the extra |
michel2323
force-pushed
the
ms/ka
branch
2 times, most recently
from
69eb545 to
59a441d
Compare
|
@michel2323 let us know when this is ready for prime time |
michel2323
force-pushed
the
ms/ka
branch
2 times, most recently
from
4bbb99e to
3e84145
Compare
| isnothing(devices) ? device!(node_rank % ndevices()) : device!(devices[node_rank+1]) | ||
| isnothing(devices) ? device!(child_architecture, node_rank % ndevices(child_architecture)) : device!(child_architecture, devices[node_rank+1]) |
|
@glwagner For the failing tests, we have 4 in total
|
|
I think the main problem is that I haven't figured out what |
To my knowledge that's an ill-formed query. |
@simone-silvestri @glwagner How do you want to proceed with these? Can this be rewritten to only use stuff from |
* versioninfo * dispatch fixes * Enable tests * Field broadcast fix
Finally ready. The documentation breaks due to something unrelated I think. buildkite/oceananigans-distributed isn't run because the code comes from a fork. |
| device(a::GPU) = a.device | ||
| device!(::CPU, i) = KA.device!(CPU(), i+1) | ||
| device!(::CPU) = nothing |
There was a problem hiding this comment.
There was a problem hiding this comment.
@vchuravy we may want to chat, but as far as I can tell there is basically an inconsistency in the concept of "device" in Oceananigans (we use the backend and "device id" interchangeably). It's a mess and we need to clean this up; I was thinking we would do this in a subsequent PR, but another option is to do it prior to this PR...
@michel2323 I don't see a call to device! in the code you linked. There is a call to switch_device! -- is that a synonym for device!?
| break | ||
| end | ||
| end | ||
| if backend isa Nothing |
There was a problem hiding this comment.
| if backend isa Nothing | |
| if backend === nothing |
| regional_return_values = Vector(undef, length(devs)) | ||
| for (r, dev) in enumerate(devs) | ||
| switch_device!(dev) | ||
| # switch_device!(dev) |
There was a problem hiding this comment.
You still need to switch devices don't you?
There was a problem hiding this comment.
we plan to discontinue support for multi-device MultiRegion (eg move towards the requirement that all multi-device code uses MPI). I don't think multi-device functionality is tested either...
| if isdefined(Main, :CUDA) | ||
| try | ||
| return versioninfo_with_gpu(GPU()) | ||
| catch | ||
| return "No GPU device found." | ||
| end | ||
| else | ||
| return "" |
There was a problem hiding this comment.
|
Seems like this is getting close which is exciting! |
|
@siddharthabishnu can you take a look at the cubed sphere / multi region stuff here, it will affect you |
This PR isolates CUDA into
src/arch_cuda.jl. This removes any direct CUDA calls in the remaining Oceananigans code base. That feel can either serve as a template for a new GPU architecture or for a future CUDA extension. @vchuravy