Use Adapt.jl (#57)

Author: vchuravy

Project.toml

@@ -2,7 +2,11 @@ name = "OffsetArrays"
 uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 version = "1.3.1"
 
+[deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+
 [compat]
+Adapt = "2"
 julia = "0.7, 1"
 
 [extras]
@@ -12,7 +16,8 @@ DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "CatIndices", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra", "EllipsisNotation"]
+test = ["Aqua", "CatIndices", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra", "EllipsisNotation", "StaticArrays"]

src/OffsetArrays.jl

@@ -198,6 +198,12 @@ parenttype(A::OffsetArray) = parenttype(typeof(A))
 
 Base.parent(A::OffsetArray) = A.parent
 
+# TODO: Ideally we would delegate to the parent's broadcasting implementation, but that
+#       is currently broken in sufficiently many implementation, namely RecursiveArrayTools, DistributedArrays
+#       and StaticArrays, that it will take concentrated effort to get this working across the ecosystem.
+#       The goal would be to have `OffsetArray(CuArray) .+ 1 == OffsetArray{CuArray}`.
+# Base.Broadcast.BroadcastStyle(::Type{<:OffsetArray{<:Any, <:Any, AA}}) where AA = Base.Broadcast.BroadcastStyle(AA)
+
 Base.eachindex(::IndexCartesian, A::OffsetArray) = CartesianIndices(axes(A))
 Base.eachindex(::IndexLinear, A::OffsetVector)   = axes(A, 1)
 
@@ -473,4 +479,10 @@ if VERSION < v"1.1.0-DEV.783"
     Base.copyfirst!(dest::OffsetArray, src::OffsetArray) = (maximum!(parent(dest), parent(src)); return dest)
 end
 
+##
+# Adapt allows for automatic conversion of CPU OffsetArrays to GPU OffsetArrays
+##
+import Adapt
+Adapt.adapt_structure(to, x::OffsetArray) = OffsetArray(Adapt.adapt(to, parent(x)), x.offsets)
+
 end # module

test/runtests.jl

@@ -6,6 +6,8 @@ using LinearAlgebra
 using DelimitedFiles
 using CatIndices: BidirectionalVector
 using EllipsisNotation
+using Adapt
+using StaticArrays
 
 # https://github.com/JuliaLang/julia/pull/29440
 if VERSION < v"1.1.0-DEV.389"
@@ -1378,4 +1380,14 @@ end
     @test searchsorted(o,  6) ==  3:2
 end
 
+@testset "Adapt" begin
+    # We need another storage type, CUDA.jl defines one but we can't use that for CI
+    # let's define an appropriate method for SArrays
+    Adapt.adapt_storage(::Type{SA}, xs::AbstractArray) where SA<:SArray = convert(SA, xs)
+    arr = OffsetArray(rand(3, 3), -1:1, -1:1)
+    s_arr = adapt(SMatrix{3,3}, arr)
+    @test parent(s_arr) isa SArray
+    @test arr == adapt(Array, s_arr)
+end
+
 include("origin.jl")