From 200f31d5dfe37c7699dcbee539587b920562e12d Mon Sep 17 00:00:00 2001
From: Chris Rackauckas <accounts@chrisrackauckas.com>
Date: Wed, 31 Jul 2019 22:48:51 -0400
Subject: [PATCH 1/2] GPU color autodiff

---
 .gitlab-ci.yml                             | 61 ++++++++++++++++++++++
 Project.toml                               |  3 +-
 src/SparseDiffTools.jl                     |  1 +
 src/differentiation/compute_jacobian_ad.jl | 12 +++--
 test/test_gpu_ad.jl                        | 18 +++++++
 5 files changed, 89 insertions(+), 6 deletions(-)
 create mode 100644 .gitlab-ci.yml
 create mode 100644 test/test_gpu_ad.jl

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
new file mode 100644
index 00000000..677a4a2f
--- /dev/null
+++ b/.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/**/*
diff --git a/Project.toml b/Project.toml
index bebbbe47..6b0a833a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,6 +4,8 @@ authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>", "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"
diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl
index fc1e8a0e..c72b3053 100644
--- a/src/SparseDiffTools.jl
+++ b/src/SparseDiffTools.jl
@@ -7,6 +7,7 @@ using ForwardDiff
 using LightGraphs
 using Requires
 using VertexSafeGraphs
+using Adapt
 
 using LinearAlgebra
 using SparseArrays, ArrayInterface
diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl
index 468fe7d7..0c6e6729 100644
--- a/src/differentiation/compute_jacobian_ad.jl
+++ b/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)
diff --git a/test/test_gpu_ad.jl b/test/test_gpu_ad.jl
new file mode 100644
index 00000000..d2955e73
--- /dev/null
+++ b/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)

From b9bbc00ce00f5a7900f42e522014fadb5efdee11 Mon Sep 17 00:00:00 2001
From: Chris Rackauckas <accounts@chrisrackauckas.com>
Date: Wed, 31 Jul 2019 22:53:42 -0400
Subject: [PATCH 2/2] gpu CI

---
 test/runtests.jl | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index ee0a95c1..df43cf34 100644
--- a/test/runtests.jl
+++ b/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