From 171b437fda886ab1342350f350fd310d892336d0 Mon Sep 17 00:00:00 2001
From: "Michael F. Herbst" <info@michael-herbst.com>
Date: Sat, 22 Apr 2023 22:34:22 +0200
Subject: [PATCH] Experimental AMDGPU implementation

---
 examples/amdgpu.jl            | 14 ++++++++++++++
 examples/{gpu.jl => cuda.jl}  |  0
 src/DFTK.jl                   |  9 ++++-----
 src/workarounds/roc_arrays.jl |  1 +
 4 files changed, 19 insertions(+), 5 deletions(-)
 create mode 100644 examples/amdgpu.jl
 rename examples/{gpu.jl => cuda.jl} (100%)
 create mode 100644 src/workarounds/roc_arrays.jl

diff --git a/examples/amdgpu.jl b/examples/amdgpu.jl
new file mode 100644
index 0000000000..20161682e9
--- /dev/null
+++ b/examples/amdgpu.jl
@@ -0,0 +1,14 @@
+using DFTK
+using AMDGPU
+
+a = 10.26  # Silicon lattice constant in Bohr
+lattice = a / 2 * [[0 1 1.];
+                   [1 0 1.];
+                   [1 1 0.]]
+Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
+atoms     = [Si, Si]
+positions = [ones(3)/8, -ones(3)/8]
+model = model_PBE(lattice, atoms, positions)
+
+basis  = PlaneWaveBasis(model; Ecut=30, kgrid=(5, 5, 5), architecture=DFTK.GPU(AMDGPU.ROCArray)
+scfres = self_consistent_field(basis; tol=1e-2, solver=scf_damping_solver())
diff --git a/examples/gpu.jl b/examples/cuda.jl
similarity index 100%
rename from examples/gpu.jl
rename to examples/cuda.jl
diff --git a/src/DFTK.jl b/src/DFTK.jl
index d1c4c048d4..211534c70b 100644
--- a/src/DFTK.jl
+++ b/src/DFTK.jl
@@ -230,15 +230,14 @@ function __init__()
     @require DoubleFloats="497a8b3b-efae-58df-a0af-a86822472b78" begin
         !isdefined(DFTK, :GENERIC_FFT_LOADED) && include("workarounds/fft_generic.jl")
     end
-    @require Plots="91a5bcdd-55d7-5caf-9e0b-520d859cae80" include("plotting.jl")
-    @require JLD2="033835bb-8acc-5ee8-8aae-3f567f8a3819"  include("external/jld2io.jl")
+    @require Plots="91a5bcdd-55d7-5caf-9e0b-520d859cae80"    include("plotting.jl")
+    @require JLD2="033835bb-8acc-5ee8-8aae-3f567f8a3819"     include("external/jld2io.jl")
     @require WriteVTK="64499a7a-5c06-52f2-abe2-ccb03c286192" include("external/vtkio.jl")
     @require wannier90_jll="c5400fa0-8d08-52c2-913f-1e3f656c1ce9" begin
         include("external/wannier90.jl")
     end
-    @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba"  begin
-        include("workarounds/cuda_arrays.jl")
-    end
+    @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba"   include("workarounds/cuda_arrays.jl")
+    @require AMDGPU="21141c5a-9bdb-4563-92ae-f87d6854732e" include("workarounds/roc_arrays.jl")
 end
 
 # Precompilation block with a basic workflow
diff --git a/src/workarounds/roc_arrays.jl b/src/workarounds/roc_arrays.jl
new file mode 100644
index 0000000000..4c34875c7d
--- /dev/null
+++ b/src/workarounds/roc_arrays.jl
@@ -0,0 +1 @@
+synchronize_device(::GPU{<:AMDGPU.ROCArray}) = AMDGPU.Device.sync_workgroup()