GPU color autodiff

.gitlab-ci.yml

@@ -0,0 +1,61 @@
+image: "julia:1"
+
+variables:
+  JULIA_DEPOT_PATH: "$CI_PROJECT_DIR/.julia/"
+  JULIA_NUM_THREADS: '8'
+
+cache:
+  paths:
+    - .julia/
+
+build:
+  stage: build
+  tags:
+    - 'p6000'
+  script:
+    - curl https://julialang-s3.julialang.org/bin/linux/x64/1.1/julia-1.1.1-linux-x86_64.tar.gz -o julia.tar.gz
+    - unp julia.tar.gz
+    - export PATH="$(pwd)/julia-1.1.1/bin:$PATH"
+    - julia -e "using InteractiveUtils;
+                versioninfo()"
+    - julia --project -e "using Pkg;
+                          Pkg.update();
+                          Pkg.instantiate();
+                          Pkg.add(\"SparseDiffTools\");
+                          pkg\"precompile\";
+                          using SparseDiffTools;"
+  only:
+  - master
+  - tags
+  - external
+  - pushes
+  artifacts:
+    untracked: true
+    paths:
+      - .julia/**/*
+      - julia-1.1.1/**/*
+
+test-GPU:
+  stage: test
+  tags:
+    - 'p6000'
+  dependencies:
+    - build
+  variables:
+    GROUP: "GPU"
+  script:
+    - export PATH="$(pwd)/julia-1.1.1/bin:$PATH"
+    - julia -e "using InteractiveUtils;
+                versioninfo()"
+    - julia --project -e "using Pkg; Pkg.add(\"CuArrays\");
+                          Pkg.test(\"SparseDiffTools\"; coverage=true);"
+  only:
+  - master
+  - tags
+  - external
+  - pushes
+  artifacts:
+    untracked: true
+    paths:
+      - .julia/**/*
+      - julia-1.1.1/**/*

Project.toml

@@ -4,6 +4,8 @@ authors = ["Pankaj Mishra <[email protected]>", "Chris Rackauckas <cont
 version = "0.6.0"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 DiffEqDiffTools = "01453d9d-ee7c-5054-8395-0335cb756afa"
@@ -13,7 +15,6 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
-ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 
 [compat]
 ArrayInterface = "1.1"

src/SparseDiffTools.jl

@@ -7,6 +7,7 @@ using ForwardDiff
 using LightGraphs
 using Requires
 using VertexSafeGraphs
+using Adapt
 
 using LinearAlgebra
 using SparseArrays, ArrayInterface

src/differentiation/compute_jacobian_ad.jl

@@ -5,6 +5,7 @@ struct ForwardColorJacCache{T,T2,T3,T4,T5,T6}
     p::T4
     color::T5
     sparsity::T6
+    chunksize::Int
 end
 
 function default_chunk_size(maxcolor)
@@ -29,18 +30,19 @@ function ForwardColorJacCache(f,x,_chunksize = nothing;
         chunksize = _chunksize
     end
 
-    t = zeros(Dual{typeof(f), eltype(x), getsize(chunksize)},length(x))
+    p = adapt.(typeof(x),generate_chunked_partials(x,color,chunksize))
+    t = Dual{typeof(f)}.(x,first(p))
 
     if dx === nothing
         fx = similar(t)
         _dx = similar(x)
     else
-        fx = zeros(Dual{typeof(f), eltype(dx), getsize(chunksize)},length(dx))
+        fx = Dual{typeof(f)}.(dx,first(p))
         _dx = dx
     end
 
-    p = generate_chunked_partials(x,color,chunksize)
-    ForwardColorJacCache(t,fx,_dx,p,color,sparsity)
+
+    ForwardColorJacCache(t,fx,_dx,p,color,sparsity,getsize(chunksize))
 end
 
 generate_chunked_partials(x,color,N::Integer) = generate_chunked_partials(x,color,Val(N))
@@ -96,8 +98,8 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     p = jac_cache.p
     color = jac_cache.color
     sparsity = jac_cache.sparsity
+    chunksize = jac_cache.chunksize
     color_i = 1
-    chunksize = length(first(first(jac_cache.p)))
     fill!(J, zero(eltype(J)))
 
     for i in eachindex(p)

test/runtests.jl

@@ -1,10 +1,20 @@
 using SafeTestsets
 
-@time @safetestset "Exact coloring via contraction" begin include("test_contraction.jl") end
-@time @safetestset "Greedy distance-1 coloring" begin include("test_greedy_d1.jl") end
-@time @safetestset "Greedy star coloring" begin include("test_greedy_star.jl") end
-@time @safetestset "Matrix to graph conversion" begin include("test_matrix2graph.jl") end
-@time @safetestset "AD using color vector" begin include("test_ad.jl") end
-@time @safetestset "Integration test" begin include("test_integration.jl") end
-@time @safetestset "Special matrices" begin include("test_specialmatrices.jl") end
-@time @safetestset "Jac Vecs and Hes Vecs" begin include("test_jaches_products.jl") end
+const GROUP = get(ENV, "GROUP", "All")
+const is_APPVEYOR = ( Sys.iswindows() && haskey(ENV,"APPVEYOR") )
+const is_TRAVIS = haskey(ENV,"TRAVIS")
+
+if GROUP == "All"
+    @time @safetestset "Exact coloring via contraction" begin include("test_contraction.jl") end
+    @time @safetestset "Greedy distance-1 coloring" begin include("test_greedy_d1.jl") end
+    @time @safetestset "Greedy star coloring" begin include("test_greedy_star.jl") end
+    @time @safetestset "Matrix to graph conversion" begin include("test_matrix2graph.jl") end
+    @time @safetestset "AD using color vector" begin include("test_ad.jl") end
+    @time @safetestset "Integration test" begin include("test_integration.jl") end
+    @time @safetestset "Special matrices" begin include("test_specialmatrices.jl") end
+    @time @safetestset "Jac Vecs and Hes Vecs" begin include("test_jaches_products.jl") end
+end
+
+if GROUP == "GPU"
+    @time @safetestset "GPU AD" begin include("test_gpu_ad.jl") end
+end

test/test_gpu_ad.jl

@@ -0,0 +1,18 @@
+using SparseDiffTools, CuArrays, Test, LinearAlgebra
+using ArrayInterface: allowed_getindex, allowed_setindex!
+function f(dx,x)
+    dx[2:end-1] = x[1:end-2] - 2x[2:end-1] + x[3:end]
+    allowed_setindex!(dx,-2allowed_getindex(x,1) + allowed_getindex(x,2),1)
+    allowed_setindex!(dx,-2allowed_getindex(x,30) + allowed_getindex(x,29),30)
+    nothing
+end
+
+_J1 = similar(rand(30,30))
+_denseJ1 = cu(collect(_J1))
+x = cu(rand(30))
+CuArrays.allowscalar(false)
+forwarddiff_color_jacobian!(_denseJ1, f, x)
+forwarddiff_color_jacobian!(_denseJ1, f, x, sparsity = _J1)
+forwarddiff_color_jacobian!(_denseJ1, f, x, color = repeat(1:3,10), sparsity = _J1)
+_Jt = similar(Tridiagonal(_J1))
+@test_broken forwarddiff_color_jacobian!(_denseJ1, f, x, color = repeat(1:3,10), sparsity = _Jt)