JLD2 extension take 2 (#597)

Co-authored-by: lassepe <[email protected]>

Author: christiangnrd

Project.toml

@@ -15,9 +15,16 @@ ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
+[weakdeps]
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+
+[extensions]
+JLD2Ext = "JLD2"
+
 [compat]
 Adapt = "4.0"
 GPUArraysCore = "= 0.2.0"
+JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9.28"
 LLVM = "3.9, 4, 5, 6, 7, 8, 9"
 LinearAlgebra = "1"

ext/JLD2Ext.jl

@@ -0,0 +1,8 @@
+module JLD2Ext
+
+using GPUArrays: AbstractGPUArray
+using JLD2: JLD2
+
+JLD2.writeas(::Type{<:AbstractGPUArray{T, N}}) where {T, N} = Array{T, N}
+
+end

test/Project.toml

@@ -3,6 +3,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

test/testsuite.jl

@@ -10,6 +10,7 @@ using GPUArrays
 
 using KernelAbstractions
 using LinearAlgebra
+using JLD2
 using Random
 using Test
 
@@ -94,6 +95,7 @@ include("testsuite/random.jl")
 include("testsuite/uniformscaling.jl")
 include("testsuite/statistics.jl")
 include("testsuite/alloc_cache.jl")
+include("testsuite/jld2ext.jl")
 
 """
 Runs the entire GPUArrays test suite on array type `AT`

test/testsuite/jld2ext.jl

@@ -0,0 +1,25 @@
+@testsuite "ext/jld2" (AT, eltypes) -> begin
+    mktempdir() do dir
+        for ET in eltypes
+            @testset "$ET" begin
+                # Test with different array sizes and dimensions
+                for dims in ((2,), (2, 2), (2, 2, 2))
+                    # Create a random array
+                    x = AT(rand(ET, dims...))
+
+                    # Save to a temporary file
+                    file = joinpath(dir, "test.jld2")
+
+                    # Save and load
+                    JLD2.save_object(file, x)
+                    y = JLD2.load_object(file)
+
+                    # Verify the loaded array matches the original
+                    @test y isa AT{ET, length(dims)}
+                    @test size(y) == size(x)
+                    @test x == y
+                end
+            end
+        end
+    end
+end