Skip to content

Commit

Permalink
support multiple devices per backend (#554)
Browse files Browse the repository at this point in the history
  • Loading branch information
vchuravy authored Jan 21, 2025
1 parent dcb719e commit e5ef261
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 0 deletions.
28 changes: 28 additions & 0 deletions src/KernelAbstractions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -586,6 +586,34 @@ function priority!(::Backend, prio::Symbol)
return nothing
end

"""
device(::Backend)::Int
Returns the ordinal number of the currently active device starting at one.
"""
function device(::Backend)
return 1
end

"""
ndevices(::Backend)::Int
Returns the number of devices the backend supports.
"""
function ndevices(::Backend)
return 1
end

"""
device!(::Backend, id::Int)
"""
function device!(backend::Backend, id::Int)
if !(0 < id <= ndevices(backend))
throw(ArgumentError("Device id $id out of bounds."))
end
return nothing
end

"""
functional(::Backend)
Expand Down
14 changes: 14 additions & 0 deletions test/devices.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
function devices_testsuite(Backend)
backend = Backend()

current_device = KernelAbstractions.device(backend)
for i in KernelAbstractions.ndevices(backend)
KernelAbstractions.device!(backend, i)
@test KernelAbstractions.device(backend) == i
end

@test_throws ArgumentError KernelAbstractions.device!(backend, 0)
@test_throws ArgumentError KernelAbstractions.device!(backend, KernelAbstractions.ndevices(backend) + 1)
KernelAbstractions.device!(backend, current_device)
return nothing
end
5 changes: 5 additions & 0 deletions test/testsuite.jl
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ include("private.jl")
include("unroll.jl")
include("nditeration.jl")
include("copyto.jl")
include("devices.jl")
include("print_test.jl")
include("compiler.jl")
include("reflection.jl")
Expand Down Expand Up @@ -67,6 +68,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{
copyto_testsuite(backend, AT)
end

@conditional_testset "Devices" skip_tests begin
devices_testsuite(backend)
end

@conditional_testset "Printing" skip_tests begin
printing_testsuite(backend)
end
Expand Down

0 comments on commit e5ef261

Please sign in to comment.