Skip to content

Commit 2aecac6

Browse files
committed
Add HeterogenousContainer
1 parent bedeacf commit 2aecac6

File tree

1 file changed

+32
-0
lines changed

1 file changed

+32
-0
lines changed

src/auxiliary/containers.jl

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,4 +314,36 @@ end
314314
function raw_copy!(c::AbstractContainer, from::Int, destination::Int)
315315
raw_copy!(c, c, from, from, destination)
316316
end
317+
318+
# Containers that support multiple storage backends
319+
# should be subtypes of this type.
320+
#
321+
# The first type parameter must be `Array` by default and if
322+
# `Adapt.adapt_structure(to, container)` is called then this must be
323+
# `typeof(to)`. This is used in downstream code, e.g. for calling
324+
# `wrap_array` with an appropriate type after `resize!`ing.
325+
abstract type AbstractHeterogeneousContainer{T} <: AbstractContainer end
326+
array_type(::AbstractHeterogeneousContainer{T}) where {T} = T
327+
328+
# Subtypes must implement a method for these functions
329+
function Adapt.adapt_structure(to, ::AbstractHeterogeneousContainer)
330+
error("required method not implemented")
331+
end
332+
333+
# For some storage backends like CUDA.jl, empty arrays do seem to simply be
334+
# null pointers which can cause `unsafe_wrap` to fail when calling
335+
# Adapt.adapt (ArgumentError, see
336+
# https://github.com/JuliaGPU/CUDA.jl/blob/v5.4.2/src/array.jl#L212-L229).
337+
# To circumvent this, on length zero arrays this allocates
338+
# a separate empty array instead of wrapping.
339+
# However, since zero length arrays are not used in calculations,
340+
# it should be okay if the underlying storage vectors and wrapped arrays
341+
# are not the same as long as they are properly wrapped when `resize!`d etc.
342+
function unsafe_wrap_or_alloc(to, vec, size)
343+
if length(vec) == 0
344+
return similar(vec, size)
345+
else
346+
return unsafe_wrap(to, pointer(vec), size)
347+
end
348+
end
317349
end # @muladd

0 commit comments

Comments
 (0)