Add API for device selection

[jkravs]

From a short conversation with @vchuravy I learned that there is currently no way of selecting between different GPUs by the same vendor on a single host using KernelAbstractions. It seems that KA.jl always defaults to the GPU with the lowest index.

CUDA.jl offers this feature already, see https://cuda.juliagpu.org/stable/usage/multigpu/, I'm not sure about AMDGPU.jl and oneAPI.jl. KA.jl should maybe offer some backend agnostic API for this.

[luraess]

I'm not sure about AMDGPU.jl

The closest to CUDA.device! you can get in AMDGPU is AMDGPU.device_id! (see https://github.com/JuliaGPU/AMDGPU.jl/commit/b80dc31ddb3097030aed253341a8ae4e4a943c36). Note that device indexing is 1-based in AMDGPU vs 0-based in CUDA.

I have no insights for oneAPI or Metal.

[stemann]

Something like the following pseudocode?

abstract type BackendDevice end

DeviceType(::Type{<:Backend})::Type{<:BackendDevice} # trait

devices(::TBackend)::AbstractVector{TBackendDevice} where {TBackend <: Backend, TBackendDevice = DeviceType(TBackend)}

device(::TBackend)::TBackendDevice where {TBackend <: Backend, TBackendDevice = DeviceType(TBackend)}

device!(dev::{<:BackendDevice})::Nothing # Selects dev device

function device(f::Function, dev::TBackendDevice)
    current_device = device()
    device!(dev)
    try
        f()
    finally
        device!(current_device)
    end
end

with e.g. CUDAKernels defining:

struct CUDAKernelDevice <: BackendDevice
    cu_device::CuDevice
end

KernelAbstractions.DeviceType(::Type{CUDABackend}) = CUDAKernelDevice

KernelAbstractions.devices(::CUDABackend) = CUDAKernelDevice.(devices())

KernelAbstractions.device(::CUDABackend) = CUDAKernelDevice(device())

KernelAbstractions.device!(dev::CUDAKernelDevice) = device!(dev.cu_device)

... and equivalent for ROCKernels, oneAPIKernels, and MetalKernels - seems like it can be implemented for all four backends.