From 7ff7abb9e46ae76c392578eb7acea2d483faa752 Mon Sep 17 00:00:00 2001
From: Alexis Montoison <alexis.montoison@polymtl.ca>
Date: Thu, 7 Sep 2023 09:31:34 -0500
Subject: [PATCH] [documentation] AMDGPU.jl supports rocSPARSE

---
 .buildkite/pipeline.yml | 40 +++++++++++++++++++++-------------------
 docs/src/gpu.md         | 25 +++++++++++++++++++++++--
 test/gpu/amd.jl         | 21 ++++++++++++++++++++-
 3 files changed, 64 insertions(+), 22 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index f0b79523f..c22997428 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -19,25 +19,27 @@ steps:
         include("test/gpu/nvidia.jl")'
     timeout_in_minutes: 30
 
-  # - label: "AMD GPUs -- AMDGPU.jl"
-  #   plugins:
-  #     - JuliaCI/julia#v1:
-  #         version: 1.9
-  #   agents:
-  #     queue: "juliagpu"
-  #     rocm: "*"
-  #     rocmgpu: "gfx1031"
-  #   env:
-  #     JULIA_AMDGPU_CORE_MUST_LOAD: "1"
-  #     JULIA_AMDGPU_HIP_MUST_LOAD: "1"
-  #     JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
-  #   command: |
-  #     julia --color=yes --project -e '
-  #     using Pkg
-  #       Pkg.add("AMDGPU")
-  #       Pkg.instantiate()
-  #       include("test/gpu/amd.jl")'
-  #   timeout_in_minutes: 30
+  - label: "AMD GPUs -- AMDGPU.jl"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: 1.9
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "gfx1031"
+    env:
+      JULIA_NUM_THREADS: 4
+      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
+    command: |
+      julia --color=yes --project -e '
+        using Pkg
+        # Pkg.add("AMDGPU")
+        Pkg.add(url="https://github.com/JuliaGPU/AMDGPU.jl", rev="master")
+        Pkg.instantiate()
+        include("test/gpu/amd.jl")'
+    timeout_in_minutes: 30
 
   - label: "Intel GPUs -- oneAPI.jl"
     plugins:
diff --git a/docs/src/gpu.md b/docs/src/gpu.md
index cf5c80003..507d7add7 100644
--- a/docs/src/gpu.md
+++ b/docs/src/gpu.md
@@ -171,8 +171,29 @@ if AMDGPU.functional()
 end
 ```
 
-!!! info
-    The library `rocSPARSE` is not interfaced yet in AMDGPU.jl and only dense linear systems are supported.
+Sparse matrices have a specific storage on AMD GPUs (`ROCSparseMatrixCSC`, `ROCSparseMatrixCSR` or `ROCSparseMatrixCOO`):
+
+```julia
+using AMDGPU, Krylov
+using AMDGPU.rocSPARSE, SparseArrays
+
+if AMDGPU.functional()
+  # CPU Arrays
+  A_cpu = sprand(100, 200, 0.3)
+  b_cpu = rand(100)
+
+  # GPU Arrays
+  A_csc_gpu = ROCSparseMatrixCSC(A_cpu)
+  A_csr_gpu = ROCSparseMatrixCSR(A_cpu)
+  A_coo_gpu = ROCSparseMatrixCOO(A_cpu)
+  b_gpu = CuVector(b_cpu)
+
+  # Solve a rectangular and sparse system on an AMD GPU
+  x_csc, y_csc, stats_csc = lnlq(A_csc_gpu, b_gpu)
+  x_csr, y_csr, stats_csr = craig(A_csr_gpu, b_gpu)
+  x_coo, y_coo, stats_coo = craigmr(A_coo_gpu, b_gpu)
+end
+```
 
 ## Intel GPUs
 
diff --git a/test/gpu/amd.jl b/test/gpu/amd.jl
index 9fb6cdffd..6d9514ba4 100644
--- a/test/gpu/amd.jl
+++ b/test/gpu/amd.jl
@@ -1,4 +1,4 @@
-using AMDGPU
+using AMDGPU, AMDGPU.rocSPARSE
 
 include("gpu.jl")
 
@@ -14,6 +14,25 @@ include("gpu.jl")
     A_gpu = ROCMatrix(A_cpu)
     b_gpu = ROCVector(b_cpu)
     x, stats = minres(A_gpu, b_gpu)
+    r_gpu = b_gpu - A_gpu * x
+    @test norm(r_gpu) ≤ 1e-4
+
+    A_cpu = sprand(100, 200, 0.3)
+    b_cpu = rand(100)
+    A_csc_gpu = ROCSparseMatrixCSC(A_cpu)
+    A_csr_gpu = ROCSparseMatrixCSR(A_cpu)
+    A_coo_gpu = ROCSparseMatrixCOO(A_cpu)
+    b_gpu = ROXVector(b_cpu)
+    b_gpu = ROCVector(b_cpu)
+    x_csc, y_csc, stats_csc = lnlq(A_csc_gpu, b_gpu)
+    x_csr, y_csr, stats_csr = craig(A_csr_gpu, b_gpu)
+    x_coo, y_coo, stats_coo = craigmr(A_coo_gpu, b_gpu)
+    r_csc = b_gpu - A_csc_gpu * x_csc
+    r_csr = b_gpu - A_csr_gpu * x_csr
+    r_coo = b_gpu - A_coo_gpu * x_coo
+    @test norm(r_csc) ≤ 1e-4
+    @test norm(r_csr) ≤ 1e-4
+    @test norm(r_coo) ≤ 1e-4
   end
 
   for FC in (Float32, Float64, ComplexF32, ComplexF64)