diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 3f1b0d7799..acc61d62ed 100644 --- a/.buildkite/pipeline.yml +++ b/.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 diff --git a/lib/EnzymeTestUtils/Project.toml b/lib/EnzymeTestUtils/Project.toml index b24a5e70f2..797b71bbad 100644 --- a/lib/EnzymeTestUtils/Project.toml +++ b/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"] diff --git a/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl b/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl index efd4baa80c..ab3231ef89 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 @@ -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,7 +98,7 @@ 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 @@ -106,7 +106,7 @@ function to_vec(x::AbstractGPUArray, seen_vecs::EnzymeTestUtils.AliasDict) 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 diff --git a/lib/EnzymeTestUtils/test/cuda_to_vec.jl b/lib/EnzymeTestUtils/test/cuda_to_vec.jl new file mode 100644 index 0000000000..a42f294fd3 --- /dev/null +++ b/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 diff --git a/lib/EnzymeTestUtils/test/runtests.jl b/lib/EnzymeTestUtils/test/runtests.jl index 2feed41cd4..25bbd75bc0 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) @@ -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