From c661e71c9a42b0bafbace81ec52f14439a69ea2e Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 15 May 2026 11:44:19 +0200 Subject: [PATCH 1/7] Add some CUDA-specific tests for EnzymeTestUtils --- lib/EnzymeTestUtils/Project.toml | 4 +- .../ext/EnzymeTestUtilsGPUArraysCoreExt.jl | 4 +- lib/EnzymeTestUtils/test/cuda_to_vec.jl | 110 ++++++++++++++++++ lib/EnzymeTestUtils/test/runtests.jl | 2 + 4 files changed, 117 insertions(+), 3 deletions(-) create mode 100644 lib/EnzymeTestUtils/test/cuda_to_vec.jl diff --git a/lib/EnzymeTestUtils/Project.toml b/lib/EnzymeTestUtils/Project.toml index b24a5e70f2..c89d893dcc 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..4c56bf809f 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, 1) 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, 1) 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 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 From 178a6d3f0ec93e3ae967adaff9171612b16f5a2c Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 19 May 2026 01:07:41 -0400 Subject: [PATCH 2/7] Make the already seen x size zero --- lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl b/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl index 4c56bf809f..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 = similar(x, Float32, 1) + 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 = similar(x, Float32, 1) + x_vec = similar(x, Float32, 0) else y = reshape(x, length(x)) x_vec = vcat(real.(y), imag.(y)) From a4c6d898abed1dabf81a6af7d0c0efb4e4654553 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 27 May 2026 10:23:41 +0200 Subject: [PATCH 3/7] Respond to comments --- .buildkite/pipeline.yml | 1 + lib/EnzymeTestUtils/Project.toml | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 3f1b0d7799..6f64791807 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/cu2vec.jl env: JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager diff --git a/lib/EnzymeTestUtils/Project.toml b/lib/EnzymeTestUtils/Project.toml index c89d893dcc..797b71bbad 100644 --- a/lib/EnzymeTestUtils/Project.toml +++ b/lib/EnzymeTestUtils/Project.toml @@ -20,7 +20,7 @@ EnzymeTestUtilsGPUArraysCoreExt = ["Enzyme", "GPUArraysCore"] [compat] ConstructionBase = "1.4.1" -CUDA = "6" +CUDA = "5, 6" Enzyme = "0.13.78" EnzymeCore = "0.5, 0.6, 0.7, 0.8" FiniteDifferences = "0.12.33" From 7c71448a4241b866628750ccac948276177be162 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 28 May 2026 01:16:15 -0400 Subject: [PATCH 4/7] Fix bad filename --- .buildkite/pipeline.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 6f64791807..acc61d62ed 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -24,7 +24,7 @@ steps: echo "+++ Run tests" julia --color=yes test/cuda.jl - julia --color=yes lib/EnzymeTestUtils/test/cu2vec.jl + julia --color=yes lib/EnzymeTestUtils/test/cuda_to_vec.jl env: JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager From efa8fbff80f7ddade8ce91a2d2864813121f8bf5 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 28 May 2026 01:34:04 -0400 Subject: [PATCH 5/7] Fixes again --- lib/EnzymeTestUtils/test/cuda_to_vec.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lib/EnzymeTestUtils/test/cuda_to_vec.jl b/lib/EnzymeTestUtils/test/cuda_to_vec.jl index 8beeded58d..b83f9ac76e 100644 --- a/lib/EnzymeTestUtils/test/cuda_to_vec.jl +++ b/lib/EnzymeTestUtils/test/cuda_to_vec.jl @@ -3,9 +3,11 @@ 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} + @test x_vec isa Vector{<:AbstractFloat} x2 = from_vec(x_vec) @test typeof(x2) === typeof(x) return EnzymeTestUtils.test_approx(x2, x) From 7be171e92cfbc80af01ee510d3a0511cb5988174 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 28 May 2026 08:41:55 -0400 Subject: [PATCH 6/7] Incremental --- .../ext/EnzymeTestUtilsGPUArraysCoreExt.jl | 20 +++++++++---------- lib/EnzymeTestUtils/test/cuda_to_vec.jl | 18 ++--------------- 2 files changed, 12 insertions(+), 26 deletions(-) diff --git a/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl b/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl index 96f09c6dab..0bc9bf04bc 100644 --- a/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl +++ b/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl @@ -38,7 +38,7 @@ 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 @@ -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 = x_vec_new[1:length(x)] + x_vec_new_imag = 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 index b83f9ac76e..b97e9eac50 100644 --- a/lib/EnzymeTestUtils/test/cuda_to_vec.jl +++ b/lib/EnzymeTestUtils/test/cuda_to_vec.jl @@ -7,7 +7,7 @@ include("helpers.jl") function test_to_vec(x) x_vec, from_vec = to_vec(x) - @test x_vec isa Vector{<:AbstractFloat} + @test x_vec isa CuVector{<:AbstractFloat} x2 = from_vec(x_vec) @test typeof(x2) === typeof(x) return EnzymeTestUtils.test_approx(x2, x) @@ -49,7 +49,7 @@ end @test getfield(y, k) == x v, from_vec = to_vec(y) @test v == vec(x) - v2 = CUDA.cuRAND.randn(size(v)) + 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) @@ -75,20 +75,6 @@ end 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 From 894c622fb5d9704fa22379169e53e4a587cce70c Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 9 Jun 2026 04:53:41 -0400 Subject: [PATCH 7/7] Slight improvements --- lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl | 4 ++-- lib/EnzymeTestUtils/test/cuda_to_vec.jl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl b/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl index 0bc9bf04bc..ab3231ef89 100644 --- a/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl +++ b/lib/EnzymeTestUtils/ext/EnzymeTestUtilsGPUArraysCoreExt.jl @@ -56,8 +56,8 @@ function EnzymeTestUtils.to_vec(x::AbstractGPUArray{<:Complex{<:EnzymeTestUtils. has_seen && return reshape(seen_xs[x], size(x)) is_const && return x x_new = Core.Typeof(x)(undef, sz) - x_vec_new_real = x_vec_new[1:length(x)] - x_vec_new_imag = x_vec_new[(length(x) + 1):(2*length(x))] + 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 diff --git a/lib/EnzymeTestUtils/test/cuda_to_vec.jl b/lib/EnzymeTestUtils/test/cuda_to_vec.jl index b97e9eac50..a42f294fd3 100644 --- a/lib/EnzymeTestUtils/test/cuda_to_vec.jl +++ b/lib/EnzymeTestUtils/test/cuda_to_vec.jl @@ -21,13 +21,13 @@ end 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 + end=# # doesn't work yet @testset "incompletely initialized struct" begin x = CUDA.cuRAND.randn(2, 3)