Add GPU pipeline

.buildkite/pipeline.yml

@@ -0,0 +1,62 @@
+steps:
+  - label: "CUDA with julia {{matrix.julia}}"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.julia}}"
+      - JuliaCI/julia-test#v1: 
+          test_args: "--quickfail"
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    timeout_in_minutes: 60
+    env:
+      GROUP: "GPU"
+      ADVANCEDVI_TEST_CUDA: "true"
+    matrix:
+      setup:
+        julia:
+           - "1.10"
+
+  - label: "Metal with julia {{matrix.julia}}"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.julia}}"
+      - JuliaCI/julia-test#v1:
+          test_args: "--quickfail"
+    agents:
+      queue: "juliaecosystem"
+      os: "macos"
+      arch: "aarch64"
+
+    if: build.message !~ /\[skip tests\]/
+    timeout_in_minutes: 60
+    env:
+      GROUP: "GPU"
+      ADVANCEDVI_TEST_METAL: "true"
+    matrix:
+      setup:
+        julia:
+           - "1.10"
+
+  - label: "AMD GPU with Julia {{matrix.julia}}"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.julia}}"
+      - JuliaCI/julia-test#v1:
+          test_args: "--quickfail"
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "*"
+    timeout_in_minutes: 60
+    env:
+      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
+      GROUP: "GPU"
+      ADVANCEDVI_TEST_AMDGPU: "true"
+      JULIA_NUM_THREADS: 4
+    matrix:
+      setup:
+        julia:
+           - "1.10"

src/objectives/elbo/repgradelbo.jl

@@ -93,9 +93,9 @@ function estimate_objective(obj::RepGradELBO, q, prob; n_samples::Int=obj.n_samp
 end
 
 function estimate_repgradelbo_ad_forward(params′, aux)
-    @unpack rng, obj, problem, adtype, restructure, q_stop = aux
+    @unpack device_rng, obj, problem, adtype, restructure, q_stop = aux
     q = restructure_ad_forward(adtype, restructure, params′)
-    samples, entropy = reparam_with_entropy(rng, q, q_stop, obj.n_samples, obj.entropy)
+    samples, entropy = reparam_with_entropy(device_rng, q, q_stop, obj.n_samples, obj.entropy)
     energy = estimate_energy_with_samples(problem, samples)
     elbo = energy + entropy
     return -elbo
@@ -109,11 +109,12 @@ function estimate_gradient!(
     prob,
     params,
     restructure,
-    state,
+    state;
+    device_rng::Random.AbstractRNG = rng,
 )
     q_stop = restructure(params)
     aux = (
-        rng=rng,
+        device_rng=device_rng,
         adtype=adtype,
         obj=obj,
         problem=prob,

src/optimize.jl

@@ -63,6 +63,7 @@ function optimize(
     prog=ProgressMeter.Progress(
         max_iter; desc="Optimizing", barlen=31, showspeed=true, enabled=show_progress
     ),
+    kwargs...
 )
     params, restructure = Optimisers.destructure(deepcopy(q_init))
     opt_st = maybe_init_optimizer(state_init, optimizer, params)
@@ -83,7 +84,8 @@ function optimize(
             params,
             restructure,
             obj_st,
-            objargs...,
+            objargs...;
+            kwargs...
         )
         stat = merge(stat, stat′)
 

test/gpu/amdgpu.jl

@@ -0,0 +1,5 @@
+
+@testset "amdgpu" begin
+    AMDGPU.ones(Float32, 16)
+    AMDGPU.randn(Float32, 16, 16)
+end

test/gpu/cuda.jl

@@ -0,0 +1,4 @@
+
+@testset "cuda" begin
+    CUDA.zeros(2, 3)
+end

test/gpu/metal.jl

@@ -0,0 +1,5 @@
+
+@testset "metal" begin
+    mtl(zeros(1,3); storage=Metal.SharedStorage)
+    Metal.randn(Float32, 16, 32)
+end

test/runtests.jl

@@ -61,3 +61,27 @@ if GROUP == "All" || GROUP == "Inference"
     include("inference/repgradelbo_locationscale.jl")
     include("inference/repgradelbo_locationscale_bijectors.jl")
 end
+
+if GROUP == "GPU"
+    if get(ENV, "ADVANCEDVI_TEST_CUDA", "false") == "true"
+        Pkg.add("CUDA")
+        using CUDA
+        include("gpu/cuda.jl")
+    end
+
+    if get(ENV, "ADVANCEDVI_TEST_METAL", "false") == "true"
+        Pkg.add("Metal")
+        using Metal
+        if Metal.functional()
+            include("gpu/metal.jl")
+        end
+    end
+
+    if get(ENV, "ADVANCEDVI_TEST_AMDGPU", "false") == "true"
+        Pkg.add("AMDGPU")
+        using AMDGPU
+        if AMDGPU.functional() && AMDGPU.functional(:MIOpen)
+            include("gpu/amdgpu.jl")
+        end
+    end
+end