Fix sub_materialize for GPU arrays

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "1.11.1"
+version = "1.12.0"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -17,6 +17,7 @@ ArrayLayoutsSparseArraysExt = "SparseArrays"
 Aqua = "0.8"
 FillArrays = "1.2.1"
 Infinities = "0.1"
+JLArrays = "0.2"
 LinearAlgebra = "1"
 Quaternions = "0.7"
 Random = "1"
@@ -29,6 +30,7 @@ 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"
@@ -37,4 +39,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Infinities", "Quaternions", "Random", "StableRNGs", "SparseArrays", "StaticArrays", "Test"]
+test = ["Aqua", "Infinities", "JLArrays", "Quaternions", "Random", "StableRNGs", "SparseArrays", "StaticArrays", "Test"]

src/ArrayLayouts.jl

@@ -126,7 +126,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), 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)