Experimental AMDGPU implementation

JuliaMolSim · Jul 24, 2023 · 171b437 · 171b437
1 parent bfbe7af
commit 171b437
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 5 deletions.
diff --git a/examples/amdgpu.jl 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 → examples/cuda.jl b/examples/gpu.jl → examples/cuda.jl
diff --git 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
@@ -0,0 +1 @@
+synchronize_device(::GPU{<:AMDGPU.ROCArray}) = AMDGPU.Device.sync_workgroup()