From 2084b3d1d32ee2d3e173eac69698d127a48ea5ff Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 11:12:37 +0200 Subject: [PATCH 01/10] fix: make (multi)basis work on CuArray --- ...ifferentiationInterfaceGPUArraysCoreExt.jl | 45 +++++-------------- DifferentiationInterface/src/utils/basis.jl | 3 -- 2 files changed, 12 insertions(+), 36 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl index d9d9749ed..d1a2e3629 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl @@ -1,42 +1,21 @@ module DifferentiationInterfaceGPUArraysCoreExt import DifferentiationInterface as DI -using GPUArraysCore: AbstractGPUArray +using GPUArraysCore: @allowscalar, AbstractGPUArray -""" - OneElement - -Efficient storage for a one-hot array, aka an array in the standard Euclidean basis. -""" -struct OneElement{I,N,T,A<:AbstractArray{T,N}} <: AbstractArray{T,N} - ind::I - val::T - a::A - - function OneElement(ind::Integer, val::T, a::A) where {N,T,A<:AbstractArray{T,N}} - right_ind = eachindex(a)[ind] - return new{typeof(right_ind),N,T,A}(right_ind, val, a) - end - - function OneElement( - ind::CartesianIndex{N}, val::T, a::A - ) where {N,T,A<:AbstractArray{T,N}} - linear_ind = LinearIndices(a)[ind] - right_ind = eachindex(a)[linear_ind] - return new{typeof(right_ind),N,T,A}(right_ind, val, a) - end -end - -Base.size(oe::OneElement) = size(oe.a) -Base.IndexStyle(oe::OneElement) = Base.IndexStyle(oe.a) - -function Base.getindex(oe::OneElement{<:Integer}, ind::Integer) - return ifelse(ind == oe.ind, oe.val, zero(eltype(oe.a))) +function DI.basis(a::AbstractGPUArray{T}, i) where {T} + b = similar(a) + fill!(b, zero(T)) + @allowscalar b[i] = one(T) + return b end -function DI.basis(a::AbstractGPUArray{T}, i) where {T} - b = zero(a) - b .+= OneElement(i, one(T), a) +function DI.multibasis(a::AbstractGPUArray{T}, inds) where {T} + b = similar(a) + fill!(b, zero(T)) + for i in inds + @allowscalar b[i] = one(T) + end return b end diff --git a/DifferentiationInterface/src/utils/basis.jl b/DifferentiationInterface/src/utils/basis.jl index 46c7162fc..6d1b5a55e 100644 --- a/DifferentiationInterface/src/utils/basis.jl +++ b/DifferentiationInterface/src/utils/basis.jl @@ -18,9 +18,6 @@ end multibasis(a::AbstractArray, inds) Construct the sum of the `i`-th standard basis arrays in the vector space of `a` for all `i ∈ inds`. - -!!! warning - Does not work on GPU, since this is intended for sparse autodiff and SparseMatrixColorings.jl doesn't work on GPUs either. """ function multibasis(a::AbstractArray{T}, inds) where {T} b = similar(a) From 39a68e49121ba5d775ef4c0f431164e9fe09c13b Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 12:40:18 +0200 Subject: [PATCH 02/10] Add CUDA tests --- DifferentiationInterface/Project.toml | 25 ++++++++++++++++--- .../test/GPU/CUDA/simple.jl | 22 ++++++++++++++++ 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index 486b91a5b..5896be5b2 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -1,7 +1,7 @@ name = "DifferentiationInterface" uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" authors = ["Guillaume Dalle", "Adrian Hill"] -version = "0.7.0" +version = "0.7.1" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -41,7 +41,9 @@ DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"] DifferentiationInterfaceGPUArraysCoreExt = "GPUArraysCore" DifferentiationInterfaceGTPSAExt = "GTPSA" DifferentiationInterfaceMooncakeExt = "Mooncake" -DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"] +DifferentiationInterfacePolyesterForwardDiffExt = [ + "PolyesterForwardDiff", "ForwardDiff", "DiffResults" +] DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"] DifferentiationInterfaceSparseArraysExt = "SparseArrays" DifferentiationInterfaceSparseConnectivityTracerExt = "SparseConnectivityTracer" @@ -121,4 +123,21 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["ADTypes", "Aqua", "ComponentArrays", "DataFrames", "ExplicitImports", "JET", "JLArrays", "JuliaFormatter", "Pkg", "Random", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test"] +test = [ + "ADTypes", + "Aqua", + "ComponentArrays", + "DataFrames", + "ExplicitImports", + "JET", + "JLArrays", + "JuliaFormatter", + "Pkg", + "Random", + "SparseArrays", + "SparseConnectivityTracer", + "SparseMatrixColorings", + "StableRNGs", + "StaticArrays", + "Test", +] diff --git a/DifferentiationInterface/test/GPU/CUDA/simple.jl b/DifferentiationInterface/test/GPU/CUDA/simple.jl index efc38541f..65a347079 100644 --- a/DifferentiationInterface/test/GPU/CUDA/simple.jl +++ b/DifferentiationInterface/test/GPU/CUDA/simple.jl @@ -2,4 +2,26 @@ using Pkg Pkg.add("CUDA") using CUDA +using DifferentiationInterface +import DifferentiationInterface as DI +using LinearAlgebra +using Test + CUDA.versioninfo() + +@testset "Basis" begin + x = CuVector(rand(Float32, 3)) + b = basis(x, 2) + @test Array(b) == [0, 1, 0] + + X = CuMatrix(rand(Float32, 2, 2)) + B = DI.multibasis(X, [2, 3]) + @test Array(B) == [0 1; 1 0] +end + +@testset "Jacobian" begin + x = CuVector(rand(Float32, 3)) + backend = DI.AutoSimpleFiniteDiff() + J = jacobian(identity, backend, x) + @test J == I +end From 29882c50ef2d482b7b09cfe7d34d4c0a03b718e7 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 12:42:19 +0200 Subject: [PATCH 03/10] Changelog --- DifferentiationInterface/CHANGELOG.md | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/DifferentiationInterface/CHANGELOG.md b/DifferentiationInterface/CHANGELOG.md index a3da3a10c..7dee7f9ba 100644 --- a/DifferentiationInterface/CHANGELOG.md +++ b/DifferentiationInterface/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.7.1] + +### Fixed + +- Make basis work for `CuArray` ([#810]) + ## [0.7.0] ### Changed @@ -27,11 +33,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Allocate Enzyme shadow memory during preparation ([#782]) -[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.0...main +[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.1...main +[0.7.1]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.0...DifferentiationInterface-v0.7.1 [0.7.0]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.54...DifferentiationInterface-v0.7.0 [0.6.54]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.53...DifferentiationInterface-v0.6.54 [0.6.53]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.52...DifferentiationInterface-v0.6.53 +[#810]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/810 [#799]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/799 [#795]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/795 [#790]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/790 From d037b77fd7011f77d890f8c8cb72819ba0186797 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 12:45:00 +0200 Subject: [PATCH 04/10] Test drafts on GPU --- .buildkite/pipeline.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 68aa085ee..d5fc78136 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1,8 +1,6 @@ steps: - label: "DI GPU tests" - if: | - !build.pull_request.draft && - build.pull_request.labels includes "gpu" + if: build.pull_request.labels includes "gpu" plugins: - JuliaCI/julia#v1: version: "1" From c1ac7ce9c3249a4fb3a02d34f701b32c3a3a7880 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 12:52:36 +0200 Subject: [PATCH 05/10] Install DifferentiationInterface --- DifferentiationInterface/test/GPU/CUDA/main.jl | 8 ++++++++ DifferentiationInterface/test/GPU/CUDA/simple.jl | 3 --- 2 files changed, 8 insertions(+), 3 deletions(-) create mode 100644 DifferentiationInterface/test/GPU/CUDA/main.jl diff --git a/DifferentiationInterface/test/GPU/CUDA/main.jl b/DifferentiationInterface/test/GPU/CUDA/main.jl new file mode 100644 index 000000000..094f3c841 --- /dev/null +++ b/DifferentiationInterface/test/GPU/CUDA/main.jl @@ -0,0 +1,8 @@ +@info "Testing on CUDA" +using Pkg +Pkg.add(["CUDA", "DifferentiationInterface"]) +using Test + +@testset verbose = true "Simple" begin + include("simple.jl") +end diff --git a/DifferentiationInterface/test/GPU/CUDA/simple.jl b/DifferentiationInterface/test/GPU/CUDA/simple.jl index 65a347079..9dc4f6edc 100644 --- a/DifferentiationInterface/test/GPU/CUDA/simple.jl +++ b/DifferentiationInterface/test/GPU/CUDA/simple.jl @@ -1,6 +1,3 @@ -@info "Testing on CUDA" -using Pkg -Pkg.add("CUDA") using CUDA using DifferentiationInterface import DifferentiationInterface as DI From 363cc18a7950b9c7c8934377e1693e62adfd2361 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 14:31:37 +0200 Subject: [PATCH 06/10] Run main --- .buildkite/pipeline.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index d5fc78136..70bec5f65 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -5,7 +5,7 @@ steps: - JuliaCI/julia#v1: version: "1" command: | - julia ./DifferentiationInterface/test/GPU/CUDA/simple.jl + julia ./DifferentiationInterface/test/GPU/CUDA/main.jl agents: queue: "juliagpu" cuda: "*" From 184e6577c9959ceb9a7dceccaa11bc4a1215f885 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 14:49:03 +0200 Subject: [PATCH 07/10] Fix --- DifferentiationInterface/test/GPU/CUDA/simple.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DifferentiationInterface/test/GPU/CUDA/simple.jl b/DifferentiationInterface/test/GPU/CUDA/simple.jl index 9dc4f6edc..81ff68eb7 100644 --- a/DifferentiationInterface/test/GPU/CUDA/simple.jl +++ b/DifferentiationInterface/test/GPU/CUDA/simple.jl @@ -8,7 +8,7 @@ CUDA.versioninfo() @testset "Basis" begin x = CuVector(rand(Float32, 3)) - b = basis(x, 2) + b = DI.basis(x, 2) @test Array(b) == [0, 1, 0] X = CuMatrix(rand(Float32, 2, 2)) From 72e715468335320cda18f45bb412a31bdc4511b2 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 14:54:52 +0200 Subject: [PATCH 08/10] Dev right version of DI --- DifferentiationInterface/test/GPU/CUDA/main.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/DifferentiationInterface/test/GPU/CUDA/main.jl b/DifferentiationInterface/test/GPU/CUDA/main.jl index 094f3c841..f7d1fabfb 100644 --- a/DifferentiationInterface/test/GPU/CUDA/main.jl +++ b/DifferentiationInterface/test/GPU/CUDA/main.jl @@ -1,6 +1,7 @@ @info "Testing on CUDA" using Pkg -Pkg.add(["CUDA", "DifferentiationInterface"]) +Pkg.add("CUDA") +Pkg.develop(PackageSpec(; path="./DifferentiationInterface")) using Test @testset verbose = true "Simple" begin From cdc4995d4df64bd297ad3733c5b9f717fd3ff9dc Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 17:10:40 +0200 Subject: [PATCH 09/10] Fix --- DifferentiationInterface/test/GPU/CUDA/simple.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DifferentiationInterface/test/GPU/CUDA/simple.jl b/DifferentiationInterface/test/GPU/CUDA/simple.jl index 81ff68eb7..6fac85c9d 100644 --- a/DifferentiationInterface/test/GPU/CUDA/simple.jl +++ b/DifferentiationInterface/test/GPU/CUDA/simple.jl @@ -20,5 +20,5 @@ end x = CuVector(rand(Float32, 3)) backend = DI.AutoSimpleFiniteDiff() J = jacobian(identity, backend, x) - @test J == I + @test (J .!= 0) == I end From 34019b50aa74d12457961f6598b24e33733e9946 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 10 Jun 2025 17:58:24 +0200 Subject: [PATCH 10/10] Cov --- DifferentiationInterface/test/Core/Internals/basis.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/DifferentiationInterface/test/Core/Internals/basis.jl b/DifferentiationInterface/test/Core/Internals/basis.jl index b2e1245e5..681d71bab 100644 --- a/DifferentiationInterface/test/Core/Internals/basis.jl +++ b/DifferentiationInterface/test/Core/Internals/basis.jl @@ -1,4 +1,4 @@ -using DifferentiationInterface: basis +using DifferentiationInterface: basis, multibasis using LinearAlgebra using StaticArrays, JLArrays using Test @@ -8,6 +8,9 @@ using Test @test basis(rand(3), 2) isa Vector @test basis(rand(3), 2) == b_ref @test basis(jl(rand(3)), 2) isa JLArray + @test Array(basis(jl(rand(3)), 2)) == [0, 1, 0] + @test multibasis(jl(rand(3)), [1, 2]) isa JLArray + @test Array(multibasis(jl(rand(3)), [1, 2])) == [1, 1, 0] @test all(basis(jl(rand(3)), 2) .== b_ref) @test basis(@SVector(rand(3)), 2) isa SVector @test basis(@SVector(rand(3)), 2) == b_ref