Tpetra_Core.cpp

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#include "Tpetra_Core.hpp"
#include "Tpetra_Details_mpiIsInitialized.hpp"

#ifdef HAVE_TPETRACORE_MPI
#  include <Teuchos_DefaultMpiComm.hpp> // this includes mpi.h too
#endif // HAVE_TPETRACORE_MPI
#include <Teuchos_DefaultSerialComm.hpp>

#include <Kokkos_Core.hpp>
#include "Tpetra_Details_checkLaunchBlocking.hpp"
#include "Tpetra_Details_KokkosTeuchosTimerInjection.hpp"
#include "Tpetra_Details_Behavior.hpp"
#include "KokkosKernels_EagerInitialize.hpp"

namespace Tpetra {

  namespace { // (anonymous)

    class HideOutputExceptOnProcess0 {
    public:
      HideOutputExceptOnProcess0 (std::ostream& stream,
                                  const int myRank) :
        stream_ (stream),
        originalBuffer_ (stream.rdbuf ())
      {
        if (myRank != 0) {
          stream.rdbuf (blackHole_.rdbuf ());
        }
      }

      ~HideOutputExceptOnProcess0 () {
        stream_.rdbuf (originalBuffer_);
      }
    private:
      std::ostream& stream_;
      decltype (std::cout.rdbuf ()) originalBuffer_;
      Teuchos::oblackholestream blackHole_;
    };

#if defined(HAVE_TPETRACORE_MPI)
    bool mpiIsInitializedAndNotFinalized ()
    {
      int isInitialized = 0;
      int isFinalized = 0;
      // Not sure if MPI_Initialized or MPI_Finalized meet the strong
      // exception guarantee.
      try {
        (void) MPI_Initialized (&isInitialized);
      }
      catch (...) {
        isInitialized = 0;
      }
      try {
        (void) MPI_Finalized (&isFinalized);
      }
      catch (...) {
        isFinalized = 0;
      }
      return isInitialized != 0 && isFinalized == 0;
    }

    int getRankHarmlessly (MPI_Comm comm)
    {
      int myRank = 0;
      if (mpiIsInitializedAndNotFinalized ()) {
        try {
          (void) MPI_Comm_rank (comm, &myRank);
        }
        catch (...) {
          // Not sure if MPI_Comm_rank meets strong exception guarantee
          myRank = 0;
        }
      }
      return myRank;
    }
#endif // defined(HAVE_TPETRACORE_MPI)


    // Whether one of the Tpetra::initialize() functions has been called before.
    bool tpetraIsInitialized_ = false;

    // Whether Tpetra initialized Kokkos.  Tpetra::finalize only
    // finalizes Kokkos if it initialized Kokkos.  Otherwise,
    // something else initialized Kokkos and is responsible for
    // finalizing it.
    bool tpetraInitializedKokkos_ = false;

#ifdef HAVE_TPETRACORE_MPI
    // Whether Tpetra initialized MPI.  Tpetra::finalize only
    // finalizes MPI if it initialized MPI.  Otherwise, something else
    // initialized MPI and is responsible for finalizing it.
    bool tpetraInitializedMpi_ = false;
#endif // HAVE_TPETRACORE_MPI

    // Tpetra's default communicator, wrapped in a Teuchos wrapper.
    // After Tpetra::finalize() is called, this GOES AWAY (is set to null).
    Teuchos::RCP<const Teuchos::Comm<int> > wrappedDefaultComm_;

    // This takes the same arguments as (the first two of) initialize().
    void initKokkosIfNeeded (int* argc, char*** argv, const int myRank)
    {
      if (! tpetraInitializedKokkos_) {
        // Kokkos doesn't have a global is_initialized().  However,
        // Kokkos::initialize() always initializes the default execution
        // space, so it suffices to check whether that was initialized.
        const bool kokkosIsInitialized =
          Kokkos::is_initialized ();
        if (! kokkosIsInitialized) {
          HideOutputExceptOnProcess0 hideCerr (std::cerr, myRank);
          HideOutputExceptOnProcess0 hideCout (std::cout, myRank);

          // Unlike MPI_Init, Kokkos promises not to modify argc and argv.
          Kokkos::initialize (*argc, *argv);
          tpetraInitializedKokkos_ = true;
        }
      }
      Details::checkOldCudaLaunchBlocking();
      const bool kokkosIsInitialized =
        Kokkos::is_initialized ();
      TEUCHOS_TEST_FOR_EXCEPTION
        (! kokkosIsInitialized, std::logic_error, "At the end of "
         "initKokkosIfNeeded, Kokkos is not initialized.  "
         "Please report this bug to the Tpetra developers.");
      // Now that the Kokkos backend(s) are initialized,
      // initialize all KokkosKernels TPLs.
      KokkosKernels::eager_initialize();
    }

#ifdef HAVE_TPETRACORE_MPI
    // This takes the same arguments as MPI_Init and the first two
    // arguments of initialize().
    void initMpiIfNeeded (int* argc, char*** argv)
    {
      // Both MPI_Initialized and MPI_Finalized report true after
      // MPI_Finalize has been called.  It's not legal to call
      // MPI_Init after MPI_Finalize has been called (see MPI 3.0
      // Standard, Section 8.7).  It would be unusual for users to
      // want to use Tpetra after MPI_Finalize has been called, but
      // there's no reason why we should forbid it.  It just means
      // that Tpetra will need to run without MPI.

      const bool mpiReady = mpiIsInitializedAndNotFinalized ();
      if (! mpiReady) {
        // Tpetra doesn't currently need to call MPI_Init_thread,
        // since with Tpetra, only one thread ever calls MPI
        // functions.  If we ever want to explore
        // MPI_THREAD_MULTIPLE, here would be the place to call
        // MPI_Init_thread.
        const int err = MPI_Init (argc, argv);
        TEUCHOS_TEST_FOR_EXCEPTION
          (err != MPI_SUCCESS, std::runtime_error, "MPI_Init failed with "
           "error code " << err << " != MPI_SUCCESS.  If MPI was set up "
           "correctly, then this should not happen, since we have already "
           "checked that MPI_Init (or MPI_Init_thread) has not yet been "
           "called.  This may indicate that your MPI library is corrupted "
           "or that it is incorrectly linked to your program.");
        tpetraInitializedMpi_ = true;
      }
    }
#endif // HAVE_TPETRACORE_MPI

  } // namespace (anonymous)

  bool isInitialized () {
    return tpetraIsInitialized_;
  }

  Teuchos::RCP<const Teuchos::Comm<int> > getDefaultComm ()
  {
    // It's technically not correct to call this function if Tpetra
    // has not yet been initialized, but requiring that may break some
    // downstream tests.
    //
    // This function initializes wrappedDefaultComm_ lazily.
    // Tpetra::initialize should not set it up.
    if (wrappedDefaultComm_.is_null ()) {
      Teuchos::RCP<const Teuchos::Comm<int> > comm;
#ifdef HAVE_TPETRACORE_MPI
      // Teuchos::MpiComm's constructor used to invoke MPI collectives.
      // It still reserves the right to do so.  This means MPI must be
      // initialized and not finalized.
      const bool mpiReady = mpiIsInitializedAndNotFinalized ();
      if (mpiReady) {
        comm = Teuchos::rcp (new Teuchos::MpiComm<int> (MPI_COMM_WORLD));
      }
      else {
        comm = Teuchos::rcp (new Teuchos::SerialComm<int> ());
      }
#else
      comm = Teuchos::rcp (new Teuchos::SerialComm<int> ());
#endif // HAVE_TPETRACORE_MPI
      wrappedDefaultComm_ = comm;
    }
    return wrappedDefaultComm_;
  }

  void initialize (int* argc, char*** argv)
  {
    if (! tpetraIsInitialized_) {
#if defined(HAVE_TPETRACORE_MPI)
      initMpiIfNeeded (argc, argv);
      // It's technically legal to initialize Tpetra after
      // MPI_Finalize has been called.  This means that we can't call
      // MPI_Comm_rank without first checking MPI_Finalized.
      const int myRank = getRankHarmlessly (MPI_COMM_WORLD);
#else
      const int myRank = 0;
#endif // defined(HAVE_TPETRACORE_MPI)
      initKokkosIfNeeded (argc, argv, myRank);

      // Add Kokkos calls to the TimeMonitor if the environment says so
      Tpetra::Details::AddKokkosDeepCopyToTimeMonitor();
      Tpetra::Details::AddKokkosFenceToTimeMonitor();
      Tpetra::Details::AddKokkosFunctionsToTimeMonitor();

      Tpetra::Details::Behavior::reject_unrecognized_env_vars();
    }
    tpetraIsInitialized_ = true;
  }

#ifdef HAVE_TPETRACORE_MPI
  void initialize (int* argc, char*** argv, MPI_Comm comm)
  {
    if (! tpetraIsInitialized_) {
#if defined(HAVE_TPETRACORE_MPI)
      initMpiIfNeeded (argc, argv);
      // It's technically legal to initialize Tpetra after
      // MPI_Finalize has been called.  This means that we can't call
      // MPI_Comm_rank without first checking MPI_Finalized.
      const int myRank = getRankHarmlessly (comm);
#else
      const int myRank = 0;
#endif // defined(HAVE_TPETRACORE_MPI)
      initKokkosIfNeeded (argc, argv, myRank);

      // Add Kokkos::deep calls to the TimeMonitor if the environment says so
      Tpetra::Details::AddKokkosDeepCopyToTimeMonitor();
      Tpetra::Details::AddKokkosFenceToTimeMonitor();
      Tpetra::Details::AddKokkosFunctionsToTimeMonitor();

      Tpetra::Details::Behavior::reject_unrecognized_env_vars();
    }
    tpetraIsInitialized_ = true;

    // Set the default communicator.  We set it here, after the above
    // initialize() call, just in case users have not yet initialized
    // MPI.  (This is legal if users pass in a predefined
    // communicator, like MPI_COMM_WORLD or MPI_COMM_SELF.)
    //
    // What if users have already called initialize() before, but with
    // a different default communicator?  There are two possible
    // things we could do here:
    //
    //   1. Test via MPI_Comm_compare whether comm differs from the
    //      raw MPI communicator in wrappedDefaultComm_ (if indeed it
    //      is an MpiComm).
    //   2. Accept that the user might want to change the default
    //      communicator, and let them do it.
    //
    // I prefer #2.  Perhaps it would be sensible to print a warning
    // here, but on which process?  Would we use the old or the new
    // communicator to find that process' rank?  We don't want to use
    // MPI_COMM_WORLD's Process 0, since neither communicator might
    // include that process.  Furthermore, in some environments, only
    // Process 0 in MPI_COMM_WORLD is allowed to do I/O.  Thus, we
    // just let the change go without a warning.
    wrappedDefaultComm_ = Teuchos::rcp (new Teuchos::MpiComm<int> (comm));
  }
#endif // HAVE_TPETRACORE_MPI

  void
  initialize (int* argc, char*** argv,
              const Teuchos::RCP<const Teuchos::Comm<int> >& comm)
  {
    if (! tpetraIsInitialized_) {
#if defined(HAVE_TPETRACORE_MPI)
      initMpiIfNeeded (argc, argv);
#endif // defined(HAVE_TPETRACORE_MPI)
      // It's technically legal to initialize Tpetra after
      // MPI_Finalize has been called.  This means that we can't call
      // MPI_Comm_rank without first checking MPI_Finalized.
      const int myRank = comm->getRank ();
      initKokkosIfNeeded (argc, argv, myRank);

      // Add Kokkos calls to the TimeMonitor if the environment says so
      Tpetra::Details::AddKokkosDeepCopyToTimeMonitor();
      Tpetra::Details::AddKokkosFenceToTimeMonitor();
      Tpetra::Details::AddKokkosFunctionsToTimeMonitor();

      Tpetra::Details::Behavior::reject_unrecognized_env_vars();
    }
    tpetraIsInitialized_ = true;
    wrappedDefaultComm_ = comm;
  }

  void finalize ()
  {
    if (! tpetraIsInitialized_) {
      return; // user didn't call initialize(), so do nothing at all
    }

    // Tpetra should only finalize Kokkos if it initialized Kokkos.
    // See Github Issue #434.
    if (tpetraInitializedKokkos_ && !Kokkos::is_finalized()) {
      Kokkos::finalize ();
    }

    // Make sure that no outstanding references to the communicator
    // remain.  If users gave initialize() an MPI_Comm, _they_ are
    // responsible for freeing it before calling finalize().
    wrappedDefaultComm_ = Teuchos::null;

#ifdef HAVE_TPETRACORE_MPI
    // Tpetra should only finalize MPI if it initialized MPI.
    // See Github Issue #434.
    if (tpetraInitializedMpi_) {
      // finalize() is a kind of destructor, so it's a bad idea to
      // throw an exception on error.  MPI implementations do have
      // the option to throw on error, so let's catch that here.
      try {
        if (Details::mpiIsInitialized ()) {
          // This must be called by the same thread that called
          // MPI_Init or MPI_Init_thread (possibly, but not
          // necessarily, in Tpetra::initialize()).
          (void) MPI_Finalize ();
        }
      }
      catch (...) {}
    }
#endif // HAVE_TPETRACORE_MPI

    tpetraIsInitialized_ = false; // it's not anymore.
  }

  ScopeGuard::ScopeGuard (int* argc, char*** argv)
  {
    ::Tpetra::initialize (argc, argv);
  }

#ifdef HAVE_TPETRA_MPI
  ScopeGuard::ScopeGuard (int* argc, char*** argv, MPI_Comm comm)
  {
    ::Tpetra::initialize (argc, argv, comm);
  }
#endif // HAVE_TPETRA_MPI

  ScopeGuard::~ScopeGuard ()
  {
    ::Tpetra::finalize ();
  }

} // namespace Tpetra