diff --git a/lib/EnzymeTestUtils/Project.toml b/lib/EnzymeTestUtils/Project.toml index 39f84c88dd..ae4586010e 100644 --- a/lib/EnzymeTestUtils/Project.toml +++ b/lib/EnzymeTestUtils/Project.toml @@ -20,6 +20,7 @@ EnzymeTestUtilsGPUArraysCoreExt = ["Enzyme", "GPUArraysCore"] [compat] ConstructionBase = "1.4.1" +CUDA = "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"] diff --git a/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl b/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl index efd4baa80c..96f09c6dab 100644 --- a/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl +++ b/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 @@ -42,7 +42,7 @@ function to_vec(x::AbstractGPUArray{<:Complex{<:EnzymeTestUtils.ElementType}}, s 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)) diff --git a/lib/EnzymeTestUtils/test/cuda_to_vec.jl b/lib/EnzymeTestUtils/test/cuda_to_vec.jl new file mode 100644 index 0000000000..8beeded58d --- /dev/null +++ b/lib/EnzymeTestUtils/test/cuda_to_vec.jl @@ -0,0 +1,110 @@ +using EnzymeTestUtils +using EnzymeTestUtils: to_vec +using CUDA +using Test + +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 + + @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 "subarrays" begin + x = CUDA.cuRAND.randn(2, 3) + # note: bottom right 2x2 submatrix ommited from y but will be present in v + y = @views (x[:, 1], x[1, :]) + test_to_vec(y) + v, from_vec = to_vec(y) + @test v == vec(x) + v2 = CUDA.cuRAND.randn(size(v)) + y2 = from_vec(v2) + @test y2[1] == reshape(v2, size(x))[:, 1] + @test y2[2] == reshape(v2, size(x))[1, :] + @test Base.dataids(y2[1]) == Base.dataids(y2[2]) + 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 diff --git a/lib/EnzymeTestUtils/test/runtests.jl b/lib/EnzymeTestUtils/test/runtests.jl index 8883ee78ef..7468db5a14 100644 --- a/lib/EnzymeTestUtils/test/runtests.jl +++ b/lib/EnzymeTestUtils/test/runtests.jl @@ -1,6 +1,7 @@ using EnzymeTestUtils using Random using Test +using CUDA Random.seed!(0) @@ -12,4 +13,5 @@ Random.seed!(0) include("generate_tangent.jl") include("test_forward.jl") include("test_reverse.jl") + CUDA.functional() && include("cuda_to_vec.jl") end