GPU Pipeline (#161)

* Initial buildkite script

* Dummy message

* Tweak pipeline

* Tweak

* Tweak

* Fix

* Tidy up labelling

* Add Adapt ext and CUDA + Adapt test-time deps

* CUDA ext

* Remove Adapt-specific stuff

* Fix issues

* More perf stuff

* Basic CUDA functionality

* Bump compat

.buildkite/pipeline.yml

@@ -0,0 +1,15 @@
+steps:
+  - label: "Julia v1"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    if: build.message !~ /\[skip tests\]/
+    timeout_in_minutes: 60
+    env:
+      TEST_GROUP: "gpu"

Project.toml

@@ -1,7 +1,7 @@
 name = "Tapir"
 uuid = "07d77754-e150-4737-8c94-cd238a1fb45b"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.2.15"
+version = "0.2.16"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -18,16 +18,19 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
+TapirCUDAExt = "CUDA"
 TapirLogDensityProblemsADExt = "LogDensityProblemsAD"
 TapirSpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
 ADTypes = "1.2"
 BenchmarkTools = "1"
+CUDA = "5"
 ChainRulesCore = "1"
 DiffRules = "1"
 DiffTests = "0.1"
@@ -49,6 +52,7 @@ julia = "1"
 [extras]
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
@@ -64,4 +68,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [targets]
-test = ["AbstractGPs", "BenchmarkTools", "DiffTests", "Distributions", "Documenter", "FillArrays", "KernelFunctions", "LogDensityProblemsAD", "PDMats", "ReverseDiff", "SpecialFunctions", "StableRNGs", "Test", "Turing", "TemporalGPs"]
+test = ["AbstractGPs", "BenchmarkTools", "CUDA", "DiffTests", "Distributions", "Documenter", "FillArrays", "KernelFunctions", "LogDensityProblemsAD", "PDMats", "ReverseDiff", "SpecialFunctions", "StableRNGs", "Test", "Turing", "TemporalGPs"]

ext/TapirCUDAExt.jl

@@ -0,0 +1,65 @@
+module TapirCUDAExt
+
+    using LinearAlgebra, Random, Tapir
+
+    using Base: IEEEFloat
+    using CUDA: CuArray, cu
+
+    import Tapir:
+        MinimalCtx,
+        rrule!!,
+        @is_primitive,
+        tangent_type,
+        zero_tangent,
+        randn_tangent,
+        increment!!,
+        set_to_zero!!,
+        _add_to_primal,
+        _diff,
+        _dot,
+        _scale,
+        TestUtils,
+        CoDual,
+        NoPullback
+
+    import Tapir.TestUtils: populate_address_map!, AddressMap, __increment_should_allocate
+
+    # Tell Tapir.jl how to handle CuArrays.
+
+    tangent_type(::Type{P}) where {P<:CuArray{<:IEEEFloat}} = P
+    zero_tangent(x::CuArray{<:IEEEFloat}) = zero(x)
+    function randn_tangent(rng::AbstractRNG, x::CuArray{Float32})
+        return cu(randn(rng, Float32, size(x)...))
+    end
+    TestUtils.has_equal_data(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x == y
+    increment!!(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x .+= y
+    __increment_should_allocate(::Type{<:CuArray{<:IEEEFloat}}) = true
+    set_to_zero!!(x::CuArray{<:IEEEFloat}) = x .= 0
+    _add_to_primal(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x + y
+    _diff(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x - y
+    _dot(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = Float64(dot(x, y))
+    _scale(x::Float64, y::P) where {T<:IEEEFloat, P<:CuArray{T}} = T(x) * y
+    function populate_address_map!(m::AddressMap, p::CuArray, t::CuArray)
+        k = pointer_from_objref(p)
+        v = pointer_from_objref(t)
+        haskey(m, k) && (@assert m[k] == v)
+        m[k] = v
+        return m
+    end
+
+    # Basic rules for operating on CuArrays.
+
+    @is_primitive(
+        MinimalCtx,
+        Tuple{Type{<:CuArray}, UndefInitializer, Vararg{Int, N}} where {N},
+    )
+    function rrule!!(
+        p::CoDual{Type{P}},
+        init::CoDual{UndefInitializer},
+        dims::CoDual{Int}...
+    ) where {P<:CuArray{<:Base.IEEEFloat}}
+        _dims = map(primal, dims)
+        y = CoDual(P(undef, _dims), P(undef, _dims))
+        return y, NoPullback(p, init, dims...)
+    end
+end

src/test_utils.jl

@@ -815,9 +815,9 @@ function test_tangent(
     perf && test_tangent_performance(rng, p)
 end
 
-function test_tangent(rng::AbstractRNG, p::P; interface_only=false) where {P}
+function test_tangent(rng::AbstractRNG, p::P; interface_only=false, perf=true) where {P}
     test_tangent_consistency(rng, p; interface_only)
-    test_tangent_performance(rng, p)
+    perf && test_tangent_performance(rng, p)
 end
 
 function test_equality_comparison(x)

test/integration_testing/cuda.jl

@@ -0,0 +1,18 @@
+using CUDA
+
+@testset "cuda" begin
+
+    # Check we can operate on CuArrays.
+    test_tangent(
+        Xoshiro(123456),
+        CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}(undef, 8, 8);
+        interface_only=false,
+    )
+
+    # Check we can instantiate a CuArray.
+    interp = Tapir.TapirInterpreter()
+    TestUtils.test_derived_rule(
+        sr(123456), CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, undef, 256;
+        interp, perf_flag=:none, interface_only=true, is_primitive=true,
+    )
+end

test/runtests.jl

@@ -63,6 +63,8 @@ include("front_matter.jl")
         include(joinpath("integration_testing", "temporalgps.jl"))
     elseif test_group == "interface"
         include("interface.jl")
+    elseif test_group == "gpu"
+        include(joinpath("integration_testing", "cuda.jl"))
     else
         throw(error("test_group=$(test_group) is not recognised"))
     end