Bringing GPU programming to DFTK

[GVigne]

trafficstars

This is a work in progress and probably shouldn't be merged immediately into DFTK. As part of my GSoC (you can see some additional information here ), I am working on a GPU version of DFTK. One of the goal is to keep overall code changes as low as possible (I do not aim to build an other DFTK package using GPUs). This is the first step in my project: so far I managed to implement the Kinetic, Local and NonLocal terms and am running computations through the self_consistent_field function. I have not yet managed to make the SCF solvers work, so I disabled them for now.

Here is a MWE:

using DFTK
using CUDA
a = 10.263141334305942  # 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];
terms_LDA = [Kinetic(), AtomicLocal(), AtomicNonlocal()]

# Setup an LDA model and discretize using
# a single k-point and a small `Ecut` of 5 Hartree.
mod = Model(lattice, atoms, positions; terms=terms_LDA,symmetries=false)
basis = PlaneWaveBasis(mod; Ecut=30, kgrid=(1, 1, 1))
basis_gpu = PlaneWaveBasis(mod; Ecut=30, kgrid=(1, 1, 1), array_type = CuArray)

scfres = self_consistent_field(basis; tol=1e-3, solver=scf_damping_solver(1.0))
scfres_gpu = self_consistent_field(basis_gpu; tol=1e-3, solver=scf_damping_solver(1.0))

Any feedback is appreciated, be it on the code in itself or the implementation of new features!

[mfherbst]

Super awesome! I have not had a look at the code, but the interface is pretty much exactly how I'd like to have it.

[antoine-levitt]

It'd be good to merge some bits before we get the full GPU story, so that it's easier for you to stay and sync and for us to review. Eg you could have in separate PRs the uncontroversial fixes (eg zeros->similar, etc) that we can just merge easily now, and the LOBPCG fixes to not use blockarrays