Skip to content

add KA.get_backend(dev) #2779

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

Closed
wants to merge 1 commit into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/CUDAKernels.jl
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ KA.ones(::CUDABackend, ::Type{T}, dims::Tuple) where T = CUDA.ones(T, dims)

KA.get_backend(::CuArray) = CUDABackend()
KA.get_backend(::CUSPARSE.AbstractCuSparseArray) = CUDABackend()
KA.get_backend(::CuDevice) = CUDABackend()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Eh that is a surprising need. How does one obtain a CuDevice only using KA primitives?
That's why KA.device returns an ordinal and not a complex object.

Copy link
Contributor Author

@michel2323 michel2323 May 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes me think about the ordinality of KA.device. Oceananigans has functions that are only called with devices() as an argument or dispatch on the device type. No data, no backend. In addition, they implement this API:

const GPUVar = Union{CuArray, CuContext, CuPtr, Ptr}
@inline UT.sync_device!(::CuDevice)  = UT.sync_device!(CUDABackend())
@inline UT.getdevice(cu::GPUVar, i)     = CUDA.device(cu)
@inline UT.getdevice(cu::GPUVar)        = CUDA..device(cu)
@inline UT.switch_device!(dev::CuDevice)            = CUDA.device!(dev)

Now, the solution would be to add the backend as an argument. However, CUDA provides device(::CuPtr) etc, that they use and might useful. Add a get_backend(::CuPtr) for those? What do you think?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Curiously, CUDA implements CUDA.device(::Ptr).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am ok, with get_backend(ptr) that extends well to other back ends (which is always the Crux).

device(data) is an odd and increasingly ill formed query. (An array can be spread across multiple devices) So I would look very closely at those usages and how they generalize.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CUDA implements CUDA.device(::Ptr).

FWIW, that (and CUDA.device(::CuPtr)) are fairly specific, almost internal APIs. Not exported from CUDAdrv, but because of historical reasons it aliases with the exported device function that's much more commonly used. I may deprecate it at some point and replace it with a properties dictionary.

All not too relevant here, I just want to say that I wouldn't let it drive any decision.

KA.synchronize(::CUDABackend) = synchronize()

KA.functional(::CUDABackend) = CUDA.functional()
Expand Down