Skip to content

Commit 5204799

Browse files
committed
supports_unified
Add to docs
1 parent a73cf97 commit 5204799

File tree

3 files changed

+50
-13
lines changed

3 files changed

+50
-13
lines changed

docs/src/api.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ allocate
2121

2222
```@docs
2323
KernelAbstractions.zeros
24+
KernelAbstractions.supports_unified
2425
```
2526

2627
## Internal

src/KernelAbstractions.jl

Lines changed: 40 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -524,40 +524,67 @@ end
524524
# adapt_storage(::Backend, a::BackendArray) = a
525525

526526
"""
527-
allocate(::Backend, Type, dims...)::AbstractArray
527+
allocate(::Backend, Type, dims...; unified=false)::AbstractArray
528528
529-
Allocate a storage array appropriate for the computational backend.
529+
Allocate a storage array appropriate for the computational backend. `unified=true`
530+
allocates an array using unified memory if the backend supports it and throws otherwise.
531+
Use [`supports_unified`](@ref) to determine whether it is supported by a backend.
530532
531533
!!! note
532534
Backend implementations **must** implement `allocate(::NewBackend, T, dims::Tuple)`
533-
"""
534-
allocate(backend::Backend, T::Type, dims...) = allocate(backend, T, dims)
535-
allocate(backend::Backend, T::Type, dims::Tuple) = throw(MethodError(allocate, (backend, T, dims)))
535+
Backend implementations **should** implement `allocate(::NewBackend, T, dims::Tuple; unified::Bool=false)`
536+
"""
537+
allocate(backend::Backend, T::Type, dims...; kwargs...) = allocate(backend, T, dims; kwargs...)
538+
function allocate(backend::Backend, T::Type, dims::Tuple; unified::Union{Nothing, Bool} = nothing)
539+
if isnothing(unified)
540+
throw(MethodError(allocate, (backend, T, dims)))
541+
elseif unified
542+
throw(ArgumentError("`$(typeof(backend))` does not support unified memory. If you believe it does, please open a github issue."))
543+
else
544+
return allocate(backend, T, dims)
545+
end
546+
end
547+
536548

537549
"""
538-
zeros(::Backend, Type, dims...)::AbstractArray
550+
zeros(::Backend, Type, dims...; unified=false)::AbstractArray
539551
540552
Allocate a storage array appropriate for the computational backend filled with zeros.
553+
`unified=true` allocates an array using unified memory if the backend supports it and
554+
throws otherwise.
541555
"""
542-
zeros(backend::Backend, T::Type, dims...) = zeros(backend, T, dims)
543-
function zeros(backend::Backend, ::Type{T}, dims::Tuple) where {T}
544-
data = allocate(backend, T, dims...)
556+
zeros(backend::Backend, T::Type, dims...; kwargs...) = zeros(backend, T, dims; kwargs...)
557+
function zeros(backend::Backend, ::Type{T}, dims::Tuple; kwargs...) where {T}
558+
data = allocate(backend, T, dims...; kwargs...)
545559
fill!(data, zero(T))
546560
return data
547561
end
548562

549563
"""
550-
ones(::Backend, Type, dims...)::AbstractArray
564+
ones(::Backend, Type, dims...; unified=false)::AbstractArray
551565
552566
Allocate a storage array appropriate for the computational backend filled with ones.
567+
`unified=true` allocates an array using unified memory if the backend supports it and
568+
throws otherwise.
553569
"""
554-
ones(backend::Backend, T::Type, dims...) = ones(backend, T, dims)
555-
function ones(backend::Backend, ::Type{T}, dims::Tuple) where {T}
556-
data = allocate(backend, T, dims)
570+
ones(backend::Backend, T::Type, dims...; kwargs...) = ones(backend, T, dims; kwargs...)
571+
function ones(backend::Backend, ::Type{T}, dims::Tuple; kwargs...) where {T}
572+
data = allocate(backend, T, dims; kwargs...)
557573
fill!(data, one(T))
558574
return data
559575
end
560576

577+
"""
578+
supports_unified(::Backend)::Bool
579+
580+
Returns whether unified memory arrays are supported by the backend.
581+
582+
!!! note
583+
Backend implementations **must** implement this function
584+
only if they **do** support unified memory.
585+
"""
586+
supports_unified(::Backend) = false
587+
561588
"""
562589
supports_atomics(::Backend)::Bool
563590

test/test.jl

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,15 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk
7777
backendT = typeof(backend).name.wrapper # To look through CUDABackend{true, false}
7878
@test backend isa backendT
7979

80+
unified = KernelAbstractions.supports_unified(backend)
81+
@test unified isa Bool
82+
U = allocate(backend, Float32, 5; unified)
83+
if unified
84+
@test U[3] isa Float32
85+
else
86+
@test_throws ErrorException U[3]
87+
end
88+
8089
x = allocate(backend, Float32, 5)
8190
A = allocate(backend, Float32, 5, 5)
8291
@test @inferred(KernelAbstractions.get_backend(A)) isa backendT

0 commit comments

Comments
 (0)