gprMax/gprMax/contexts.py

# Copyright (C) 2015-2023: The University of Edinburgh, United Kingdom
#                 Authors: Craig Warren, Antonis Giannopoulos, and John Hartley
#
# This file is part of gprMax.
#
# gprMax is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# gprMax is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with gprMax.  If not, see <http://www.gnu.org/licenses/>.

import datetime
import logging
import sys

import humanize

import gprMax.config as config

from ._version import __version__, codename
from .model_build_run import ModelBuildRun
from .solvers import create_G, create_solver
from .utilities.host_info import (print_cuda_info, print_host_info,
                                  print_opencl_info)
from .utilities.utilities import get_terminal_width, logo, timer

logger = logging.getLogger(__name__)


class Context:
    """Standard context - models are run one after another and each model
        can exploit parallelisation using either OpenMP (CPU), CUDA (GPU), or
        OpenCL (CPU/GPU).
    """

    def __init__(self):
        self.model_range = range(config.sim_config.model_start, 
                                 config.sim_config.model_end)
        self.tsimend = None
        self.tsimstart = None        

    def run(self):
        """Run the simulation in the correct context.
        
        Returns:
            results: dict that can contain useful results/data from simulation.
        """

        results = {}
        self.tsimstart = timer()
        self.print_logo_copyright()
        print_host_info(config.sim_config.hostinfo)
        if config.sim_config.general['solver'] == 'cuda':
            print_cuda_info(config.sim_config.devices['devs'])
        elif config.sim_config.general['solver'] == 'opencl':
            print_opencl_info(config.sim_config.devices['devs'])

        # Clear list of model configs. It can be retained when gprMax is
        # called in a loop, and want to avoid this.
        config.model_configs = []

        for i in self.model_range:
            config.model_num = i
            model_config = config.ModelConfig()
            config.model_configs.append(model_config)

            # Always create a grid for the first model. The next model to run
            # only gets a new grid if the geometry is not re-used.
            if i != 0 and config.sim_config.args.geometry_fixed:
                config.get_model_config().reuse_geometry = True
            else:
                G = create_G()

            model = ModelBuildRun(G)
            model.build()
            
            if not config.sim_config.args.geometry_only:
                solver = create_solver(G)
                model.solve(solver)

        self.tsimend = timer()
        self.print_sim_time_taken()

        return results

    def print_logo_copyright(self):
        """Prints gprMax logo, version, and copyright/licencing information."""
        logo_copyright = logo(f'{__version__} ({codename})')
        logger.basic(logo_copyright)

    def print_sim_time_taken(self):
        """Prints the total simulation time based on context."""
        s = (f"\n=== Simulation completed in " +
             f"{humanize.precisedelta(datetime.timedelta(seconds=self.tsimend - self.tsimstart), format='%0.4f')}")
        logger.basic(f"{s} {'=' * (get_terminal_width() - 1 - len(s))}\n")


class MPIContext(Context):
    """Mixed mode MPI/OpenMP/CUDA context - MPI task farm is used to distribute
        models, and each model parallelised using either OpenMP (CPU), 
        CUDA (GPU), or OpenCL (CPU/GPU).
    """

    def __init__(self):
        super().__init__()
        from mpi4py import MPI

        from gprMax.mpi import MPIExecutor

        self.comm = MPI.COMM_WORLD
        self.rank = self.comm.rank
        self.MPIExecutor = MPIExecutor

    def _run_model(self, **work):
        """Process for running a single model.
        
        Args:
            work: dict of any additional information that is passed to MPI 
                    workers. By default only model number (i) is used.
        """

        # Create configuration for model
        config.model_num = work['i']
        model_config = config.ModelConfig()
        # Set GPU deviceID according to worker rank
        if config.sim_config.general['solver'] == 'cuda':
            model_config.device = {'dev': config.sim_config.devices['devs'][self.rank - 1],
                                   'snapsgpu2cpu': False}
        config.model_configs = model_config

        G = create_G()
        model = ModelBuildRun(G)
        model.build()
        
        if not config.sim_config.args.geometry_only:
            solver = create_solver(G)
            model.solve(solver)

    def run(self):
        """Specialise how the models are run.
        
        Returns:
            results: dict that can contain useful results/data from simulation.
        """

        if self.rank == 0:
            self.tsimstart = timer()
            self.print_logo_copyright()
            print_host_info(config.sim_config.hostinfo)
            if config.sim_config.general['solver'] == 'cuda':
                print_cuda_info(config.sim_config.devices['devs'])
            elif config.sim_config.general['solver'] == 'opencl':
                print_opencl_info(config.sim_config.devices['devs'])
            sys.stdout.flush()

        # Contruct MPIExecutor
        executor = self.MPIExecutor(self._run_model, comm=self.comm)

        # Check GPU resources versus number of MPI tasks
        if (executor.is_master() and
            config.sim_config.general['solver'] == 'cuda' and
            executor.size - 1 > len(config.sim_config.devices['devs'])):
            logger.exception('Not enough GPU resources for number of '
                             'MPI tasks requested. Number of MPI tasks '
                             'should be equal to number of GPUs + 1.')
            raise ValueError

        jobs = [{'i': i} for i in self.model_range]
        # Send the workers to their work loop
        executor.start()
        if executor.is_master():
            results = executor.submit(jobs)

        # Make the workers exit their work loop and join the main loop again
        executor.join()

        if executor.is_master():
            self.tsimend = timer()
            self.print_sim_time_taken()
            return results