Add some CUDA-specific tests for EnzymeTestUtils

.buildkite/pipeline.yml

@@ -24,6 +24,7 @@ steps:
 
       echo "+++ Run tests"
       julia --color=yes test/cuda.jl
+      julia --color=yes lib/EnzymeTestUtils/test/cuda_to_vec.jl
     env:
       JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
 

lib/EnzymeTestUtils/Project.toml

@@ -20,6 +20,7 @@ EnzymeTestUtilsGPUArraysCoreExt = ["Enzyme", "GPUArraysCore"]
 
 [compat]
 ConstructionBase = "1.4.1"
+CUDA = "5, 6"
 Enzyme = "0.13.78"
 EnzymeCore = "0.5, 0.6, 0.7, 0.8"
 FiniteDifferences = "0.12.33"
@@ -29,10 +30,11 @@ Quaternions = "0.7"
 julia = "1.10"
 
 [extras]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MetaTesting = "9e32d19f-1e4f-477a-8631-b16c78aa0f56"
 Quaternions = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"
 
 [targets]
-test = ["LinearAlgebra", "MetaTesting", "Quaternions"]
+test = ["LinearAlgebra", "MetaTesting", "Quaternions", "CUDA"]

lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl

@@ -15,7 +15,7 @@ function EnzymeTestUtils.to_vec(x::AbstractGPUArray{<:EnzymeTestUtils.ElementTyp
     has_seen = haskey(seen_vecs, x)
     is_const = Enzyme.Compiler.guaranteed_const(Core.Typeof(x))
     if has_seen || is_const
-        x_vec = Float32[]
+        x_vec = similar(x, Float32, 0)
     else
         x_vec = reshape(x, length(x))
         seen_vecs[x] = x_vec
@@ -38,11 +38,11 @@ function EnzymeTestUtils.to_vec(x::AbstractGPUArray{<:EnzymeTestUtils.ElementTyp
 end
 
 # basic containers: loop over defined elements, recursively converting them to vectors
-function to_vec(x::AbstractGPUArray{<:Complex{<:EnzymeTestUtils.ElementType}}, seen_vecs::EnzymeTestUtils.AliasDict)
+function EnzymeTestUtils.to_vec(x::AbstractGPUArray{<:Complex{<:EnzymeTestUtils.ElementType}}, seen_vecs::EnzymeTestUtils.AliasDict)
     has_seen = haskey(seen_vecs, x)
     is_const = Enzyme.Compiler.guaranteed_const(Core.Typeof(x))
     if has_seen || is_const
-        x_vec = Float32[]
+        x_vec = similar(x, Float32, 0)
     else
         y = reshape(x, length(x))
         x_vec = vcat(real.(y), imag.(y))
@@ -55,11 +55,11 @@ function to_vec(x::AbstractGPUArray{<:Complex{<:EnzymeTestUtils.ElementType}}, s
         end
         has_seen && return reshape(seen_xs[x], size(x))
         is_const && return x
-	x_new = Array{eltype(x)}(undef, sz)
-        @inbounds @simd for i in 1:length(x)
-            x_new[i] = eltype(x)(x_vec_new[i], x_vec_new[i + length(x)])
-        end
-	x_new = Core.Typeof(x)(x_new)
+        x_new = Core.Typeof(x)(undef, sz)
+        x_vec_new_real = view(x_vec_new, 1:length(x))
+        x_vec_new_imag = view(x_vec_new, (length(x) + 1):(2*length(x)))
+        x_vec_complex = reshape(complex.(x_vec_new_real, x_vec_new_imag), sz)
+        x_new .= x_vec_complex
         seen_xs[x] = x_new
         return x_new
     end
@@ -71,7 +71,7 @@ function to_vec(x::AbstractGPUArray, seen_vecs::EnzymeTestUtils.AliasDict)
     has_seen = haskey(seen_vecs, x)
     is_const = Enzyme.Compiler.guaranteed_const(Core.Typeof(x))
     if has_seen || is_const
-        x_vec = Float32[]
+        x_vec = similar(x, Float32, 0)
     else
         x_vecs = nothing
         from_vecs = []
@@ -87,7 +87,7 @@ function to_vec(x::AbstractGPUArray, seen_vecs::EnzymeTestUtils.AliasDict)
         end
 
         if x_vecs === nothing
-            x_vecs = (Float32[], true)
+            x_vecs = (similar(x, Float32, 0), true)
         end
         x_vec = x_vecs[1]
         seen_vecs[x] = x_vec
@@ -98,15 +98,15 @@ function to_vec(x::AbstractGPUArray, seen_vecs::EnzymeTestUtils.AliasDict)
         end
         has_seen && return reshape(seen_xs[x], size(x))
         is_const && return x
-	x_new = Array{eltype(x_vew_new)}(undef, size(x))
+        x_new = Array{eltype(x_vew_new)}(undef, size(x))
         k = 1
         for i in eachindex(x)
             isassigned(x, i) || continue
             xi = from_vecs[k](@view(x_vec_new[subvec_inds[k]]), seen_xs)
             x_new[i] = xi
             k += 1
         end
-	x_new = Core.Typeof(x)(x_new)
+        x_new = Core.Typeof(x)(x_new)
         seen_xs[x] = x_new
         return x_new
     end

lib/EnzymeTestUtils/test/cuda_to_vec.jl

@@ -0,0 +1,98 @@
+using EnzymeTestUtils
+using EnzymeTestUtils: to_vec
+using CUDA
+using Test
+
+include("helpers.jl")
+
+function test_to_vec(x)
+    x_vec, from_vec = to_vec(x)
+    @test x_vec isa CuVector{<:AbstractFloat}
+    x2 = from_vec(x_vec)
+    @test typeof(x2) === typeof(x)
+    return EnzymeTestUtils.test_approx(x2, x)
+end
+
+@testset "to_vec" begin
+    @testset "array of floats" begin
+        @testset for T in (Float32, Float64, ComplexF32, ComplexF64),
+                sz in (2, (2, 3), (2, 3, 4))
+
+            test_to_vec(CUDA.cuRAND.randn(T, sz))
+        end
+    end
+    #=
+    @testset "struct" begin
+        v = CUDA.cuRAND.randn(2, 3)
+        x = TestStruct(1, TestStruct("foo", v))
+        test_to_vec(x)
+        @test to_vec(x)[1] == vec(v)
+    end=# # doesn't work yet
+
+    @testset "incompletely initialized struct" begin
+        x = CUDA.cuRAND.randn(2, 3)
+        y = TestStruct2(x)
+        v, from_vec = to_vec(y)
+        @test v == vec(x)
+        v2 = CUDA.cuRAND.randn(size(v))
+        y2 = from_vec(v2)
+        @test y2.x == reshape(v2, size(x))
+        @test !isdefined(y2, :a)
+    end
+
+    @testset "mutable struct" begin
+        @testset for k in (:a, :x)
+            x = CUDA.cuRAND.randn(2, 3)
+            y = MutableTestStruct()
+            setfield!(y, k, x)
+            @test isdefined(y, k)
+            @test getfield(y, k) == x
+            v, from_vec = to_vec(y)
+            @test v == vec(x)
+            v2 = CUDA.cuRAND.randn(size(v)...)
+            y2 = from_vec(v2)
+            @test getfield(y2, k) == reshape(v2, size(x))
+            @test !isdefined(y2, k === :a ? :x : :a)
+        end
+    end
+
+    @testset "nested array" begin
+        @testset for T in (Float32, Float64, ComplexF32, ComplexF64),
+                sz in (2, (2, 3), (2, 3, 4))
+
+            test_to_vec([CUDA.cuRAND.randn(T, sz) for _ in 1:10])
+        end
+    end
+
+    @testset "dict" begin
+        x = Dict(:a => CUDA.cuRAND.randn(2), :b => CUDA.cuRAND.randn(3))
+        test_to_vec(x)
+    end
+
+    @testset "views of arrays" begin
+        x = CUDA.cuRAND.randn(2, 3)
+        test_to_vec(reshape(x, 3, 2))
+        test_to_vec(view(x, :, 1))
+    end
+
+    @testset "reshaped arrays share memory" begin
+        struct MyContainer1
+            a::Any
+            b::Any
+        end
+        mutable struct MyContainer2
+            a::Any
+            b::Any
+        end
+        @testset for T in (MyContainer1, MyContainer2)
+            x = CUDA.cuRAND.randn(2, 3)
+            x2 = vec(x)
+            y = T(x, x2)
+            test_to_vec(y)
+            v, from_vec = to_vec(y)
+            @test v == x2
+            y2 = from_vec(v)
+            @test Base.dataids(y2.a) == Base.dataids(y2.b)
+        end
+    end
+end

lib/EnzymeTestUtils/test/runtests.jl

@@ -1,6 +1,7 @@
 using EnzymeTestUtils
 using Random
 using Test
+using CUDA
 
 Random.seed!(0)
 
@@ -13,4 +14,5 @@ Random.seed!(0)
     include("test_forward.jl")
     include("test_reverse.jl")
     include("test_fd.jl")
+    CUDA.functional() && include("cuda_to_vec.jl")
 end