This page was generated from unit-5.6-gpu/poisson_gpu.ipynb.

5.6.2 Solving the Poisson Equation on the common GPU layer

• After finite element discretization we obtain a (linear) system of equations.

• The common GPU layer (ngsolve.gpu) moves the linear operators to whichever accelerator this build has: an NVIDIA card through ngscuda, an Apple GPU through ngsmetal, or the host reference backend if neither is available.

• The host is steering, data stays on the device.

The kernels behind the device operators are written once in the common syntax of unit 5.6.1 and compiled at run-time for the backend in use. This notebook is the counterpart of 5.5.1, which uses the CUDA-only classes.

from ngsolve import *
from ngsolve.webgui import Draw
from time import time

mesh = Mesh(unit_square.GenerateMesh(maxh=0.1))
for l in range(5): mesh.Refine()
fes = H1(mesh, order=2, dirichlet=".*")
print ("ndof =", fes.ndof)

u, v = fes.TnT()
with TaskManager():
    a = BilinearForm(grad(u)*grad(v)*dx+u*v*dx).Assemble()
    f = LinearForm(x*v*dx).Assemble()

gfu = GridFunction(fes)

jac = a.mat.CreateSmoother(fes.FreeDofs())

with TaskManager():
    inv_host = CGSolver(a.mat, jac, maxiter=2000)
    ts = time()
    gfu.vec.data = inv_host * f.vec
    te = time()
    print ("steps =", inv_host.GetSteps(), ", time =", te-ts)

# Draw (gfu);

ndof = 460033
steps = 2000 , time = 4.217127561569214

Now we import the common GPU layer. It picks the backend, and tells us which one it took and whether the device computes in double precision. Apple GPUs do not, there the device operators are created in single precision.

• NGSolve matrices create their counterparts on the device with CreateDeviceMatrix. A sparse matrix becomes a DeviceSparseMatrix, a Jacobi smoother a DeviceDiagonalMatrix, a block smoother a DeviceBlockJacobi, a sparse Cholesky factorization a DeviceSparseCholesky. Sums, products and embeddings are translated recursively.

• Device operators create device vectors, which live in device memory and are transferred to the host only on demand.

• In the following, the conjugate gradients iteration runs on the host, but all operations involving big data are performed on the accelerator.

from ngsolve.gpu import *

dev = GetGPUDevice()
print ("backend =", backend, ", device =", dev.name, ", fp64 =", dev.has_float64)
ngsglobals.msg_level=1

backend = cuda , device = NVIDIA GeForce RTX 5090 , fp64 = True
CUDA Device Query...
There is 1 CUDA device.

[cusparse] handle created
[InitCuLinalg] cusparseSetStream bound to ngs_cuda_stream
[InitCuLinalg] callback wired, registering creators...

CUDA Device 0: NVIDIA GeForce RTX 5090, cap 12.0
Using device 0
Initializing cublas and cusparse.

adev = a.mat.CreateDeviceMatrix()
jacdev = jac.CreateDeviceMatrix()
print (adev.GetOperatorInfo())
print (jacdev.GetOperatorInfo())

fdev = adev.CreateColVector()
fdev.data = f.vec

inv = CGSolver(adev, jacdev, maxiter=2000, printrates=False)

ts = time()
res = (inv * fdev).Evaluate()
te = time()

print ("Time on device:", te-ts)
diff = Norm(gfu.vec - res)
print ("diff = ", diff)

DeviceSparseMatrix<double> (nze=5280769), h = 460033, w = 460033

DeviceDiagonalMatrix<double>, h = 460033, w = 460033

Time on device: 0.30512475967407227
diff =  2.1658092901716052e-14

The difference to the host solution reflects the precision of the device: about \(10^{-14}\) in double precision, about \(10^{-5}\) relative to the solution in single precision.

CGSolver computes its inner products on the host: every iteration waits for the device and reads two numbers back. DeviceCGSolver keeps the coefficients in device scalars and computes them with device kernels, so an iteration is a sequence of kernel launches without synchronisation; the residual is read back only every check iterations. On unified-memory machines, where the host is a strong competitor for sparse matrix-vector products, this is what makes the device iteration pay off.

With record=True the launches of one iteration are recorded once (ngsolve.gpu.Recording) and replayed afterwards: a CUDA graph on NVIDIA, a single command buffer on Metal. This removes the per-launch Python and driver cost, which dominates for small problems.

inv = DeviceCGSolver(adev, jacdev, maxiter=2000, check=20, record=True, printrates=False)

ts = time()
res = (inv * fdev).Evaluate()
te = time()

print ("Time on device:", te-ts, ", iterations =", inv.GetSteps())
print ("diff = ", Norm(gfu.vec - res))

Time on device: 0.29408884048461914 , iterations = 2000
diff =  2.281396443974453e-14

CG Solver with Block-Jacobi and exact low-order solver:

\[\begin{split}A = \left( \begin{array}{cc} A_{cc} & A_{cf} \\ A_{fc} & A_{ff} \end{array} \right)\end{split}\]

Additive Schwarz preconditioner:

\[C^{-1} = P A_{cc}^{-1} P^T + \sum_i E_i A_i^{-1} E_i^T\]

with

• \(P\) .. embedding of low-order space

• \(A_i\) .. blocks on edges/faces/cells

• \(E_i\) .. embedding matrices

fes = H1(mesh, order=5, dirichlet=".*")
print ("ndof =", fes.ndof)

u, v = fes.TnT()
with TaskManager():
    a = BilinearForm(grad(u)*grad(v)*dx+u*v*dx).Assemble()
    f = LinearForm(x*v*dx).Assemble()

gfu = GridFunction(fes)

jac = a.mat.CreateBlockSmoother(fes.CreateSmoothingBlocks())
lospace = fes.lospace
loinv = a.loform.mat.Inverse(inverse="sparsecholesky", freedofs=lospace.FreeDofs())
loemb = fes.loembedding

pre = jac + loemb@loinv@loemb.T
print ("mat", a.mat.GetOperatorInfo())
print ("preconditioner:")
print(pre.GetOperatorInfo())

with TaskManager():
    inv = CGSolver(a.mat, pre, maxiter=2000, printrates=False)
    ts = time()
    gfu.vec.data = inv * f.vec
    te = time()
    print ("iterations =", inv.GetSteps(), "time =", te-ts)

ndof = 2870401
mat SparseMatrixd (nze=88905601), h = 2870401, w = 2870401

preconditioner:
SumMatrix, h = 2870401, w = 2870401
  BlockJacobi-d, h = 2870401, w = 2870401
  EmbeddedTransposeMatrix, h = 2870401, w = 2870401
    EmbeddedMatrix, h = 2870401, w = 115329
      SparseCholesky-d, h = 115329, w = 115329

iterations = 70 time = 3.87028431892395

The whole preconditioner is translated by one call. The block inverses and the Cholesky factorization are computed on the host, the device applies them.

adev = a.mat.CreateDeviceMatrix()
predev = pre.CreateDeviceMatrix()
print (predev.GetOperatorInfo())

fdev = adev.CreateColVector()
fdev.data = f.vec

inv = CGSolver(adev, predev, maxiter=2000, printrates=False)
ts = time()
res = (inv * fdev).Evaluate()
te = time()
print ("iterations =", inv.GetSteps(), "time =", te-ts)
print ("diff =", Norm(gfu.vec-res) / Norm(gfu.vec))

SumMatrix, h = 2870401, w = 2870401
  DeviceBlockJacobi<double> (blocks=344705), h = 2870401, w = 2870401
  EmbeddedTransposeMatrix, h = 2870401, w = 2870401
    EmbeddedMatrix, h = 2870401, w = 115329
      DeviceSparseCholesky<double>, h = 115329, w = 115329

iterations = 70 time = 1.0113141536712646
diff = 2.963749294345198e-14

Using the BDDC preconditioner:

For the BDDC (balancing domain decomposition with constraints) preconditioning, we build a FEM system with relaxed connectivity:

This allows for static condensation of all local and interface degrees of freedom, only the wirebasket dofs enter the global solver. The resulting matrix \(\tilde A\) is much cheaper to invert.

The preconditioner is

\[P = R {\tilde A}^{-1} R^T\]

with an averaging operator \(R\).

mesh = Mesh(unit_square.GenerateMesh(maxh=0.1))
for l in range(3): mesh.Refine()
fes = H1(mesh, order=10, dirichlet=".*")
print ("ndof =", fes.ndof)

u, v = fes.TnT()
with TaskManager():
    a = BilinearForm(grad(u)*grad(v)*dx+u*v*dx)
    pre = Preconditioner(a, "bddc")
    a.Assemble()
    f = LinearForm(x*v*dx).Assemble()

gfu = GridFunction(fes)
with TaskManager():
    inv = CGSolver(a.mat, pre, maxiter=2000, printrates=False)
    ts = time()
    gfu.vec.data = (inv * f.vec).Evaluate()
    te = time()
    print ("iterations =", inv.GetSteps(), "time =", te-ts)

 Mesh bisection
 Bisection done
 Mesh bisection
 Bisection done
 Mesh bisection
 Bisection done
ndof = 718401
iterations = 79 time = 2.1538279056549072

predev = pre.mat.CreateDeviceMatrix()
print (predev.GetOperatorInfo())

ProductMatrix, h = 718401, w = 718401
  SumMatrix, h = 718401, w = 718401
    Identity, h = 718401, w = 718401
    DeviceSparseMatrix<double> (nze=2265228), h = 718401, w = 718401
  SumMatrix, h = 718401, w = 718401
    ProductMatrix, h = 718401, w = 718401
      DeviceSparseCholesky<double>, h = 718401, w = 718401
      SumMatrix, h = 718401, w = 718401
        Identity, h = 718401, w = 718401
        DeviceSparseMatrix<double> (nze=2265228), h = 718401, w = 718401
    DeviceSparseMatrix<double> (nze=54833760), h = 718401, w = 718401

adev = a.mat.CreateDeviceMatrix()
fdev = adev.CreateColVector()
fdev.data = f.vec

inv = CGSolver(adev, predev, maxiter=2000, printrates=False)
ts = time()
res = (inv * fdev).Evaluate()
te = time()
print ("iterations =", inv.GetSteps(), "time =", te-ts)
print ("diff =", Norm(gfu.vec-res) / Norm(gfu.vec))

iterations = 79 time = 0.5745899677276611
diff = 2.678107481725668e-14

The same notebook runs on the host reference backend, where the kernels are compiled with the system C++ compiler. This is slow, but a useful check that a device result is not an artifact of the backend:

# SetGPUDevice(GetCPUDevice())
# adev = a.mat.CreateDeviceMatrix(); predev = pre.mat.CreateDeviceMatrix()