Coverage for pySDC/tutorial/step_7/G_pySDC_on_GPU.py: 100%
43 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1# ---
2# jupyter:
3# jupytext:
4# formats: py:percent
5# kernelspec:
6# display_name: Python 3
7# name: python3
8# language_info:
9# name: python
10# ---
12# %% [markdown]
13# # Part G: pySDC on GPUs
14#
15# pySDC runs on GPUs through [CuPy](https://cupy.dev), and a problem class does not need a GPU twin to do so. It takes
16# a `useGPU` flag, and a `setup_GPU` classmethod swaps what the class computes with: the array library, the sparse
17# library and the data types. The body of the class then calls `self.xp.sin` where it would have called `numpy.sin`,
18# and works either way. Everything above the problem class (sweepers, transfer operators, convergence controllers) is
19# unchanged, so the same run that gives you SDC on a CPU gives you SDC on a GPU.
20#
21# This example solves one heat equation three ways, all of it on the device: with SDC on a single space level, with
22# MLSDC on two, and with PFASST on two levels spread over several time ranks. The three differ only in what the
23# controller is handed.
24#
25# ## The description
26#
27# The only difference to a CPU run is `useGPU`:
29# %%
30from pathlib import Path
32from mpi4py import MPI
34from pySDC.helpers.stats_helper import get_sorted
35from pySDC.implementations.controller_classes.controller_MPI import controller_MPI
36from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
37from pySDC.implementations.problem_classes.HeatEquation_ND_FD import heatNd_unforced
38from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
39from pySDC.implementations.transfer_classes.TransferMesh import mesh_to_mesh
42def get_description(useGPU, ml):
43 """
44 Set up the heat equation, on a GPU or not, on one space level or two.
46 The only thing that differs between the CPU and the GPU version is `useGPU`. The problem class
47 is the same one either way: `setup_GPU` swaps the array library, the sparse library and the
48 datatypes, so the body of the class goes on calling `self.xp.sin` and does not care.
50 Args:
51 useGPU (bool): Run on a GPU
52 ml (bool): Use two space levels rather than one
54 Returns:
55 dict: description of the problem to be solved
56 """
57 level_params = {'restol': 1e-10, 'dt': 1e-2}
59 sweeper_params = {'quad_type': 'RADAU-RIGHT', 'num_nodes': 3, 'QI': 'LU'}
61 problem_params = {
62 'nu': 0.1,
63 'freq': 2,
64 'bc': 'periodic',
65 'nvars': [128, 64] if ml else 128,
66 'useGPU': useGPU,
67 }
69 step_params = {'maxiter': 20}
71 description = {
72 'problem_class': heatNd_unforced,
73 'problem_params': problem_params,
74 'sweeper_class': generic_implicit,
75 'sweeper_params': sweeper_params,
76 'level_params': level_params,
77 'step_params': step_params,
78 }
80 if ml:
81 # the space transfer works on GPU arrays as well: the interpolation and restriction
82 # matrices are assembled with SciPy and then moved to the device once
83 description['space_transfer_class'] = mesh_to_mesh
84 description['space_transfer_params'] = {'rorder': 2, 'iorder': 4, 'periodic': True}
86 return description
89# %% [markdown]
90# ## One run
91#
92# The serial controller for SDC and MLSDC, `controller_MPI` with a time communicator for PFASST:
95# %%
96def run(description, comm=None, num_procs=1, Tend=8e-2):
97 """
98 Run to `Tend` and report how it went.
100 Args:
101 description (dict): description of the problem to be solved
102 comm (mpi4py.Intracomm): time communicator, for the parallel-in-time run
103 num_procs (int): number of time steps to treat in parallel, for the serial controller
104 Tend (float): time to run to
106 Returns:
107 float: error against the exact solution
108 int: total number of iterations
109 """
110 controller_params = {'logger_level': 30}
112 if comm is None:
113 controller = controller_nonMPI(
114 num_procs=num_procs, controller_params=controller_params, description=description
115 )
116 prob = controller.MS[0].levels[0].prob
117 else:
118 # the parallel-in-time controller sends the solution from one time rank to the next as a
119 # GPU array, which needs MPI to have been told to expect device pointers -- see the note at the end
120 controller = controller_MPI(controller_params=controller_params, description=description, comm=comm)
121 prob = controller.S.levels[0].prob
123 uinit = prob.u_exact(0.0)
124 uend, stats = controller.run(u0=uinit, t0=0.0, Tend=Tend)
126 error = abs(prob.u_exact(Tend) - uend)
127 iterations = sum(count for _, count in get_sorted(stats, type='niter', comm=comm))
129 return error, iterations
132# %% [markdown]
133# ## All three
134#
135# Run it with one rank per parallel time step, e.g. `mpirun -np 2 python G_pySDC_on_GPU.py`:
138# %%
139def main():
140 """
141 Solve the same heat equation with SDC, MLSDC and PFASST, all of it on GPUs.
143 The three differ only in what they are given: one space level and one time step at a time is
144 SDC, two space levels is MLSDC, and two space levels spread over several time ranks is PFASST.
145 """
146 comm = MPI.COMM_WORLD
148 # every rank runs the two serial variants -- they are the reference the parallel one is judged
149 # against, and running them everywhere keeps the ranks in step
150 runs = [
151 ('SDC ', *run(get_description(useGPU=True, ml=False))),
152 ('MLSDC ', *run(get_description(useGPU=True, ml=True))),
153 ('PFASST', *run(get_description(useGPU=True, ml=True), comm=comm)),
154 ]
156 if comm.rank == 0:
157 Path('data').mkdir(parents=True, exist_ok=True)
158 with open('data/step_7_G_out.txt', 'a') as f:
159 for name, error, iterations in runs:
160 out = f'{name} on {comm.size} GPU(s): error {error:.4e}, {iterations} iterations in total'
161 f.write(out + '\n')
162 print(out)
164 # all three solve the same problem, so they had better agree on the answer
165 errors = [error for _, error, _ in runs]
166 assert max(errors) < 1e-8, f'Some run was not accurate enough: {errors}'
167 assert abs(errors[1] - errors[2]) < 1e-10, 'PFASST and MLSDC disagree, which they should not'
170if __name__ == '__main__':
171 main()
173# %% [markdown]
174# CuPy and GPUs are not available in the browser, nor in the environment this website is built in. Our CI runs this
175# part on rented GPUs (NVIDIA T4, on [Modal](https://modal.com)), and checks that all three runs are accurate and
176# that PFASST and MLSDC agree. These are the results of the run that built this page:
177#
178# :::{literalinclude} /../../data/step_7_G_out.txt
179# :language: text
180# :::
181#
182# :::{admonition} Important things to note
183# - Space coarsening works on GPU arrays: `mesh_to_mesh` assembles its interpolation and restriction matrices with
184# SciPy and moves them to the device once, so the transfers themselves never leave it.
185# - PFASST sends the solution from one time rank to the next as a GPU array. That needs an MPI built with CUDA
186# awareness, and told to use it. conda-forge's OpenMPI is built with it and ships it switched off, so export
187# `OMPI_MCA_opal_cuda_support=true` before launching.
188# - `NCCLComm` in `pySDC/helpers/NCCL_communicator.py` routes the collectives through NCCL instead, which is worth
189# doing when a run is parallel in space as well.
190# :::