Fix sub_materialize for GPU arrays (#261)

Currently, `sub_materialize` (through `sub_materialize_axes`) falls back
to materializing on CPU. This PR generalizes that logic by determining
the output destination with `similar`, which helps to support non-Array
types like GPU arrays. As a stand-in for other GPU arrays, I test this
using
[JLArrays.JLArray](https://github.com/JuliaGPU/GPUArrays.jl/tree/master/lib/JLArrays),
which is a reference implementation for the GPUArrays.jl interface that
runs on CPU.

An alternative design would be to define memory layouts for GPU arrays
(i.e. #9), which would allow more customizability for GPU array
backends, however I think it is helpful to have fallbacks that "just
work" if reasonable parts of the Base AbstractArray interface are
implemented.

I hit this issue because I was testing out
[`BlockArrays.BlockedArray`](https://juliaarrays.github.io/BlockArrays.jl/stable/man/blockedarrays)
wrapping a GPU array and noticed that calling `A[Block(1, 1)]` to access
a block instantiated the block on CPU, this PR fixes that issue.

---------

Co-authored-by: Sheehan Olver <[email protected]>

Author: mtfishman

Project.toml

@@ -1,12 +1,12 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "1.11.2"
+version = "1.12.0"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [weakdeps]
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -18,6 +18,7 @@ ArrayLayoutsSparseArraysExt = "SparseArrays"
 Aqua = "0.8"
 FillArrays = "1.2.1"
 Infinities = "0.1"
+JLArrays = "0.2"
 LinearAlgebra = "1"
 Quaternions = "0.7"
 Random = "1"
@@ -30,11 +31,12 @@ julia = "1.10"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Infinities = "e1ba4f0e-776d-440f-acd9-e1d2e9742647"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Quaternions = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Infinities", "Quaternions", "Random", "StableRNGs", "SparseArrays", "Test"]
+test = ["Aqua", "Infinities", "JLArrays", "Quaternions", "Random", "StableRNGs", "SparseArrays", "Test"]

src/ArrayLayouts.jl

@@ -127,7 +127,7 @@ include("triangular.jl")
 include("factorizations.jl")
 
 # Extend this function if you're only looking to dispatch on the axes
-@inline sub_materialize_axes(V, _) = Array(V)
+@inline sub_materialize_axes(V, _) = copyto!(similar(V, axes(V)), V)
 @inline sub_materialize(_, V, ax) = sub_materialize_axes(V, ax)
 @inline sub_materialize(L, V) = sub_materialize(L, V, axes(V))
 @inline sub_materialize(V::SubArray) = sub_materialize(MemoryLayout(V), V)

test/test_layouts.jl

@@ -1,6 +1,6 @@
 module TestLayouts
 
-using ArrayLayouts, LinearAlgebra, FillArrays, Test
+using ArrayLayouts, LinearAlgebra, FillArrays, JLArrays, Test
 import ArrayLayouts: MemoryLayout, DenseRowMajor, DenseColumnMajor, StridedLayout,
                         ConjLayout, RowMajor, ColumnMajor, UnitStride,
                         SymmetricLayout, HermitianLayout, UpperTriangularLayout,
@@ -404,6 +404,16 @@ struct FooNumber <: Number end
         @test ArrayLayouts.mul((1:11)', F) isa AbstractMatrix{Int}
     end
 
+    @testset "GPUArrays/JLArrays" begin
+        A = jl(randn(5,5))
+        @test MemoryLayout(A) == DenseColumnMajor()
+        @test ArrayLayouts.layout_getindex(A,1:3,1:3) == A[1:3,1:3]
+        @test ArrayLayouts.layout_getindex(A,1:3,1:3) isa JLArray{Float64}
+        V = view(A,1:3,1:3)
+        @test ArrayLayouts.sub_materialize(V) == A[1:3,1:3]
+        @test ArrayLayouts.sub_materialize(V) isa JLArray{Float64}
+    end
+
     @testset "Triangular col/rowsupport" begin
         A = randn(5,5)
         @test colsupport(UpperTriangular(A),3) ≡ Base.OneTo(3)