# ---
# jupyter:
#   jupytext:
#     formats: py:percent
#   kernelspec:
#     display_name: Python 3
#     name: python3
#   language_info:
#     name: python
# ---

# %% [markdown]
# # Part G: pySDC on GPUs
#
# pySDC runs on GPUs through [CuPy](https://cupy.dev), and a problem class does not need a GPU twin to do so. It takes
# a `useGPU` flag, and a `setup_GPU` classmethod swaps what the class computes with: the array library, the sparse
# library and the data types. The body of the class then calls `self.xp.sin` where it would have called `numpy.sin`,
# and works either way. Everything above the problem class (sweepers, transfer operators, convergence controllers) is
# unchanged, so the same run that gives you SDC on a CPU gives you SDC on a GPU.
#
# This example solves one heat equation three ways, all of it on the device: with SDC on a single space level, with
# MLSDC on two, and with PFASST on two levels spread over several time ranks. The three differ only in what the
# controller is handed.
#
# ## The description
#
# The only difference to a CPU run is `useGPU`:

# %%
from pathlib import Path

from mpi4py import MPI

from pySDC.helpers.stats_helper import get_sorted
from pySDC.implementations.controller_classes.controller_MPI import controller_MPI
from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
from pySDC.implementations.problem_classes.HeatEquation_ND_FD import heatNd_unforced
from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
from pySDC.implementations.transfer_classes.TransferMesh import mesh_to_mesh


def get_description(useGPU, ml):
    """
    Set up the heat equation, on a GPU or not, on one space level or two.

    The only thing that differs between the CPU and the GPU version is `useGPU`. The problem class
    is the same one either way: `setup_GPU` swaps the array library, the sparse library and the
    datatypes, so the body of the class goes on calling `self.xp.sin` and does not care.

    Args:
        useGPU (bool): Run on a GPU
        ml (bool): Use two space levels rather than one

    Returns:
        dict: description of the problem to be solved
    """
    level_params = {'restol': 1e-10, 'dt': 1e-2}

    sweeper_params = {'quad_type': 'RADAU-RIGHT', 'num_nodes': 3, 'QI': 'LU'}

    problem_params = {
        'nu': 0.1,
        'freq': 2,
        'bc': 'periodic',
        'nvars': [128, 64] if ml else 128,
        'useGPU': useGPU,
    }

    step_params = {'maxiter': 20}

    description = {
        'problem_class': heatNd_unforced,
        'problem_params': problem_params,
        'sweeper_class': generic_implicit,
        'sweeper_params': sweeper_params,
        'level_params': level_params,
        'step_params': step_params,
    }

    if ml:
        # the space transfer works on GPU arrays as well: the interpolation and restriction
        # matrices are assembled with SciPy and then moved to the device once
        description['space_transfer_class'] = mesh_to_mesh
        description['space_transfer_params'] = {'rorder': 2, 'iorder': 4, 'periodic': True}

    return description


# %% [markdown]
# ## One run
#
# The serial controller for SDC and MLSDC, `controller_MPI` with a time communicator for PFASST:


# %%
def run(description, comm=None, num_procs=1, Tend=8e-2):
    """
    Run to `Tend` and report how it went.

    Args:
        description (dict): description of the problem to be solved
        comm (mpi4py.Intracomm): time communicator, for the parallel-in-time run
        num_procs (int): number of time steps to treat in parallel, for the serial controller
        Tend (float): time to run to

    Returns:
        float: error against the exact solution
        int: total number of iterations
    """
    controller_params = {'logger_level': 30}

    if comm is None:
        controller = controller_nonMPI(
            num_procs=num_procs, controller_params=controller_params, description=description
        )
        prob = controller.MS[0].levels[0].prob
    else:
        # the parallel-in-time controller sends the solution from one time rank to the next as a
        # GPU array, which needs MPI to have been told to expect device pointers -- see the note at the end
        controller = controller_MPI(controller_params=controller_params, description=description, comm=comm)
        prob = controller.S.levels[0].prob

    uinit = prob.u_exact(0.0)
    uend, stats = controller.run(u0=uinit, t0=0.0, Tend=Tend)

    error = abs(prob.u_exact(Tend) - uend)
    iterations = sum(count for _, count in get_sorted(stats, type='niter', comm=comm))

    return error, iterations


# %% [markdown]
# ## All three
#
# Run it with one rank per parallel time step, e.g. `mpirun -np 2 python G_pySDC_on_GPU.py`:


# %%
def main():
    """
    Solve the same heat equation with SDC, MLSDC and PFASST, all of it on GPUs.

    The three differ only in what they are given: one space level and one time step at a time is
    SDC, two space levels is MLSDC, and two space levels spread over several time ranks is PFASST.
    """
    comm = MPI.COMM_WORLD

    # every rank runs the two serial variants -- they are the reference the parallel one is judged
    # against, and running them everywhere keeps the ranks in step
    runs = [
        ('SDC   ', *run(get_description(useGPU=True, ml=False))),
        ('MLSDC ', *run(get_description(useGPU=True, ml=True))),
        ('PFASST', *run(get_description(useGPU=True, ml=True), comm=comm)),
    ]

    if comm.rank == 0:
        Path('data').mkdir(parents=True, exist_ok=True)
        with open('data/step_7_G_out.txt', 'a') as f:
            for name, error, iterations in runs:
                out = f'{name} on {comm.size} GPU(s): error {error:.4e}, {iterations} iterations in total'
                f.write(out + '\n')
                print(out)

    # all three solve the same problem, so they had better agree on the answer
    errors = [error for _, error, _ in runs]
    assert max(errors) < 1e-8, f'Some run was not accurate enough: {errors}'
    assert abs(errors[1] - errors[2]) < 1e-10, 'PFASST and MLSDC disagree, which they should not'


if __name__ == '__main__':
    main()

# %% [markdown]
# CuPy and GPUs are not available in the browser, nor in the environment this website is built in. Our CI runs this
# part on rented GPUs (NVIDIA T4, on [Modal](https://modal.com)), and checks that all three runs are accurate and
# that PFASST and MLSDC agree. These are the results of the run that built this page:
#
# :::{literalinclude} /../../data/step_7_G_out.txt
# :language: text
# :::
#
# :::{admonition} Important things to note
# - Space coarsening works on GPU arrays: `mesh_to_mesh` assembles its interpolation and restriction matrices with
#   SciPy and moves them to the device once, so the transfers themselves never leave it.
# - PFASST sends the solution from one time rank to the next as a GPU array. That needs an MPI built with CUDA
#   awareness, and told to use it. conda-forge's OpenMPI is built with it and ships it switched off, so export
#   `OMPI_MCA_opal_cuda_support=true` before launching.
# - `NCCLComm` in `pySDC/helpers/NCCL_communicator.py` routes the collectives through NCCL instead, which is worth
#   doing when a run is parallel in space as well.
# :::