This page was generated from unit-5.5-cuda/poisson_cuda.ipynb.

5.5.1 Solving the Poisson Equation on devices

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

• The ngsolve.gpu module moves the linear operators to the device: CUDA cards via ngscuda, Apple GPUs via ngsmetal, and a host reference device otherwise.

• The host is steering, data stays on the device

• The modules are included in the NGSolve distributions and can be used whenever a supported accelerator and its runtime are available.

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.21914267539978

Now we import the NGSolve gpu module. It provides

• a DeviceVector, which keeps its data on the device and transfers to the host on demand (double precision on CUDA, single precision on Metal).

• NGSolve - matrices can create their counterparts on the device. 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 *
print ("gpu backend:", backend)
ngsglobals.msg_level=1
fdev = f.vec.CreateDeviceVector(copy=True)

CUDA Device Query...
gpu backend: cuda
There is 1 CUDA device.
CUDA Device 0: NVIDIA GeForce RTX 5090, cap 12.0
Using device 0

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

with TaskManager(pajetrace=10**8):
    inv = DeviceCGSolver(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 / Norm(gfu.vec))

Time on device: 0.8123941421508789
diff =  3.2528460741235983e-15

On an A-100 device I got (for 5 levels of refinement, ndof=476417):

Time on device: 0.4084775447845459 diff = 3.406979028373306e-12

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(pajetrace=10**8):
    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.844578504562378

adev = a.mat.CreateDeviceMatrix()
predev = pre.CreateDeviceMatrix()
fdev = f.vec.CreateDeviceVector()

with TaskManager(pajetrace=10**8):
    inv = CGSolver(adev, predev, maxiter=2000, printrates=False)
    ts = time()
    gfu.vec.data = (inv * fdev).Evaluate()
    te = time()
    print ("iterations =", inv.GetSteps(), "time =", te-ts)

iterations = 70 time = 0.2347126007080078

on the A-100:

SumMatrix, h = 2896001, w = 2896001 N4ngla20DevBlockJacobiMatrixE, h = 2896001, w = 2896001 EmbeddedTransposeMatrix, h = 2896001, w = 2896001 EmbeddedMatrix, h = 2896001, w = 116353 N4ngla17DevSparseCholeskyE, h = 116353, w = 116353 iterations = 37 time= 0.6766986846923828

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 averagingn 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.155575752258301

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()
predev = pre.mat.CreateDeviceMatrix()
fdev = f.vec.CreateDeviceVector()

with TaskManager(pajetrace=10**8):
    inv = CGSolver(adev, predev, maxiter=2000, printrates=False)
    ts = time()
    gfu.vec.data = (inv * fdev).Evaluate()
    te = time()

    print ("iterations =", inv.GetSteps() , "time =", te-ts)

iterations = 79 time = 0.12683534622192383

A100: iterations = 53 time= 0.16417622566223145

Vite - traces: