From 926badc81760529132dd03c0fb6ca20caf21b8b1 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 2 Jan 2022 08:29:58 -0500 Subject: [PATCH 1/6] Fix all colon dispatch for GPUs --- .buildkite/pipeline.yml | 17 +++++++++++++++++ .github/workflows/CI.yml | 1 + src/RecursiveArrayTools.jl | 2 +- src/vector_of_array.jl | 7 +++++++ test/gpu/Project.toml | 2 ++ test/gpu/vectorofarray_gpu.jl | 9 +++++++++ test/runtests.jl | 22 ++++++++++++++++++---- 7 files changed, 55 insertions(+), 5 deletions(-) create mode 100644 .buildkite/pipeline.yml create mode 100644 test/gpu/Project.toml create mode 100644 test/gpu/vectorofarray_gpu.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml new file mode 100644 index 00000000..5fe3bef0 --- /dev/null +++ b/.buildkite/pipeline.yml @@ -0,0 +1,17 @@ +steps: + - label: "GPU" + plugins: + - JuliaCI/julia#v1: + version: "1" + - JuliaCI/julia-test#v1: + coverage: false # 1000x slowdown + agents: + queue: "juliagpu" + cuda: "*" + env: + GROUP: 'GPU' + JULIA_PKG_SERVER: "" # it often struggles with our large artifacts + # SECRET_CODECOV_TOKEN: "..." + timeout_in_minutes: 30 + # Don't run Buildkite if the commit message includes the text [skip tests] + if: build.message !~ /\[skip tests\]/ diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 85b20a26..dedf67cd 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -13,6 +13,7 @@ jobs: matrix: group: - Core + - Downstream version: - '1' - '1.6' diff --git a/src/RecursiveArrayTools.jl b/src/RecursiveArrayTools.jl index 88fb4240..05f53076 100644 --- a/src/RecursiveArrayTools.jl +++ b/src/RecursiveArrayTools.jl @@ -10,7 +10,7 @@ using Requires, RecipesBase, StaticArrays, Statistics, import ChainRulesCore import ChainRulesCore: NoTangent -import ZygoteRules +import ZygoteRules, Adapt using FillArrays diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 9e3129de..12a4b758 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -65,6 +65,13 @@ Base.@propagate_inbounds function Base.getindex(A::AbstractDiffEqArray{T, N}, I::Union{Int,AbstractArray{Int},CartesianIndex,Colon,BitArray,AbstractArray{Bool}}...) where {T, N} RecursiveArrayTools.VectorOfArray(A.u)[I...] end + +Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray{T, N}, + I::Colon...) where {T, N} + vecs = vec.(A.u) + return Adapt.adapt(__parameterless_type(T),reshape(reduce(hcat,vecs),size(A.u[1])...,length(A.u))) +end + Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, i::Int,::Colon) where {T, N} = [A.u[j][i] for j in 1:length(A)] Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, ::Colon,i::Int) where {T, N} = A.u[i] Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, i::Int,II::AbstractArray{Int}) where {T, N} = [A.u[j][i] for j in II] diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml new file mode 100644 index 00000000..c96c01b2 --- /dev/null +++ b/test/gpu/Project.toml @@ -0,0 +1,2 @@ +[deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" diff --git a/test/gpu/vectorofarray_gpu.jl b/test/gpu/vectorofarray_gpu.jl new file mode 100644 index 00000000..3c1a8e06 --- /dev/null +++ b/test/gpu/vectorofarray_gpu.jl @@ -0,0 +1,9 @@ +using RecursiveArrayTools, CUDA + +x = zeros(5) +y = VectorOfArray([x,x,x]) +y[:,:] + +x = CUDA.zeros(5) +y = VectorOfArray([x,x,x]) +y[:,:] diff --git a/test/runtests.jl b/test/runtests.jl index cb250429..6fe25d02 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,7 +11,15 @@ function activate_downstream_env() Pkg.instantiate() end +function activate_gpu_env() + Pkg.activate("gpu") + Pkg.develop(PackageSpec(path=dirname(@__DIR__))) + Pkg.instantiate() +end + @time begin + +if !is_APPVEYOR && GROUP == "Core" @time @testset "Utils Tests" begin include("utils_test.jl") end @time @testset "Partitions Tests" begin include("partitions_test.jl") end @time @testset "VecOfArr Indexing Tests" begin include("basic_indexing.jl") end @@ -20,9 +28,15 @@ end @time @testset "Linear Algebra Tests" begin include("linalg.jl") end @time @testset "Upstream Tests" begin include("upstream.jl") end @time @testset "Adjoint Tests" begin include("adjoints.jl") end +end + +if !is_APPVEYOR && GROUP == "Downstream" + activate_downstream_env() + @time @testset "DiffEqArray Indexing Tests" begin include("downstream/symbol_indexing.jl") end +end - if !is_APPVEYOR && GROUP == "Downstream" - activate_downstream_env() - @time @testset "DiffEqArray Indexing Tests" begin include("downstream/symbol_indexing.jl") end - end +if !is_APPVEYOR && GROUP == "GPU" + activate_gpu_env() + @time @testset "VectorOfArray GPU" begin include("gpu/vectorofarray_gpu.jl") end +end end From 9e7b24498e4db0d1a5e207c194323aa2331889d4 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 2 Jan 2022 08:47:43 -0500 Subject: [PATCH 2/6] fix dependencies --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 6652ff91..108bda4e 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ authors = ["Chris Rackauckas "] version = "2.21.1" [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" From eb04aee4544a94ef3c3447381d70b3ac2a0941c7 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Sun, 2 Jan 2022 21:05:26 -0500 Subject: [PATCH 3/6] Update vectorofarray_gpu.jl --- test/gpu/vectorofarray_gpu.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/gpu/vectorofarray_gpu.jl b/test/gpu/vectorofarray_gpu.jl index 3c1a8e06..6fd06c00 100644 --- a/test/gpu/vectorofarray_gpu.jl +++ b/test/gpu/vectorofarray_gpu.jl @@ -1,4 +1,5 @@ using RecursiveArrayTools, CUDA +CUDA.allowscalar(false) x = zeros(5) y = VectorOfArray([x,x,x]) From d7d0f415cd7f5cfb90d2c80fa3d40ca515ac3ce4 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Sun, 2 Jan 2022 21:06:13 -0500 Subject: [PATCH 4/6] Update vector_of_array.jl --- src/vector_of_array.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 12a4b758..a8376e0f 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -68,6 +68,7 @@ end Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray{T, N}, I::Colon...) where {T, N} + @assert length(I) == size(A.u[1])-1 vecs = vec.(A.u) return Adapt.adapt(__parameterless_type(T),reshape(reduce(hcat,vecs),size(A.u[1])...,length(A.u))) end From f6f09a6935365ad5bfc8416daf4e296a7a6e22c8 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Sun, 2 Jan 2022 21:40:16 -0500 Subject: [PATCH 5/6] Update vector_of_array.jl --- src/vector_of_array.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index a8376e0f..49827e9d 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -68,7 +68,7 @@ end Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray{T, N}, I::Colon...) where {T, N} - @assert length(I) == size(A.u[1])-1 + @assert length(I) == ndims(A.u[1])-1 vecs = vec.(A.u) return Adapt.adapt(__parameterless_type(T),reshape(reduce(hcat,vecs),size(A.u[1])...,length(A.u))) end From 8f5f185ef81d3f7189e09b804dc5195ba1de191a Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Sun, 2 Jan 2022 21:40:36 -0500 Subject: [PATCH 6/6] Update vector_of_array.jl --- src/vector_of_array.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 49827e9d..41c55ef4 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -68,7 +68,7 @@ end Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray{T, N}, I::Colon...) where {T, N} - @assert length(I) == ndims(A.u[1])-1 + @assert length(I) == ndims(A.u[1])+1 vecs = vec.(A.u) return Adapt.adapt(__parameterless_type(T),reshape(reduce(hcat,vecs),size(A.u[1])...,length(A.u))) end