Tpetra_Details_MpiTypeTraits.cpp

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

#include "Tpetra_Details_MpiTypeTraits.hpp"

#ifdef HAVE_TPETRACORE_MPI

namespace Teuchos {
namespace Details {
namespace Impl {

template<class T>
MPI_Datatype
computeKokkosComplexMpiDatatypeImpl (const ::Kokkos::complex<T>& z)
{
  static_assert (MpiTypeTraits<T>::isSpecialized, "This function only "
                 "works if MpiTypeTraits<T>::isSpecialized.");
  static_assert (! MpiTypeTraits<T>::needsFree, "This function requires "
                 "! MpiTypeTraits<T>::needsFree, since otherwise it would "
                 "leak memory.");
  // We assume here that every instance of T has the same
  // MPI_Datatype, i.e., has the same binary representation.
  MPI_Datatype innerDatatype = MpiTypeTraits<T>::getType (z.real ());
  MPI_Datatype outerDatatype; // return value

  // If Kokkos::complex<T> has the same layout as T[2], then we can
  // use a contiguous derived MPI_Datatype.  This is likely the only
  // code path that will execute.  Contiguous types are likely more
  // efficient for MPI to execute, and almost certainly more efficient
  // for MPI to set up.
  if (sizeof ( ::Kokkos::complex<T>) == 2 * sizeof (T)) {
    (void) MPI_Type_contiguous (2, innerDatatype, &outerDatatype);
  }
  else { // must use the general struct approach
    // I borrowed and adapted the code below from the MPICH
    // documentation:
    //
    // www.mpich.org/static/docs/v3.1/www3/MPI_Type_struct.html
    int blockLengths[3];
    MPI_Aint arrayOfDisplacements[3];
    MPI_Datatype arrayOfTypes[3];

    // See documentation of MyComplex (above) for explanation.
    static_assert (sizeof (MyComplex<T>) == sizeof ( ::Kokkos::complex<T>),
                   "Attempt to construct a struct of the same size and layout "
                   "as Kokkos::complex<T> failed.");
    ::Teuchos::Details::Impl::MyComplex<T> z2;

    // First entry in the struct.
    blockLengths[0] = 1;
    // Normally, &z2.re would equal &z2, but I'll be conservative and
    // actually compute the offset, even though it's probably just 0.
    //
    // Need the cast to prevent the compiler complaining about
    // subtracting addresses of different types.
    arrayOfDisplacements[0] = reinterpret_cast<uintptr_t> (&z2.re) - reinterpret_cast<uintptr_t> (&z2);
    arrayOfTypes[0] = innerDatatype;

    // Second entry in the struct.
    blockLengths[1] = 1;
    arrayOfDisplacements[1] = reinterpret_cast<uintptr_t> (&z2.im) - reinterpret_cast<uintptr_t> (&z2);
    arrayOfTypes[1] = innerDatatype;

#if MPI_VERSION < 2
    // Upper bound of the struct.
    blockLengths[2] = 1;
    arrayOfDisplacements[2] = sizeof (MyComplex<T>);
    arrayOfTypes[2] = MPI_UB; // "upper bound type"; signals end of struct
#endif // MPI_VERSION < 2

    // Define the MPI_Datatype.
#if MPI_VERSION < 2
    (void) MPI_Type_struct (3, blockLengths, arrayOfDisplacements,
                            arrayOfTypes, &outerDatatype);
#else
    // Don't include the upper bound with MPI_Type_create_struct.
    (void) MPI_Type_create_struct (2, blockLengths, arrayOfDisplacements,
                                   arrayOfTypes, &outerDatatype);
#endif // MPI_VERSION < 2
  }

  MPI_Type_commit (&outerDatatype);
  return outerDatatype;
}

MPI_Datatype
computeKokkosComplexMpiDatatype (const ::Kokkos::complex<double>& z)
{

  return computeKokkosComplexMpiDatatypeImpl<double> (z);
}

MPI_Datatype
computeKokkosComplexMpiDatatype (const ::Kokkos::complex<float>& z)
{
  return computeKokkosComplexMpiDatatypeImpl<float> (z);
}

} // namespace Impl

MPI_Datatype
MpiTypeTraits< ::Kokkos::complex<double> >::
getType (const ::Kokkos::complex<double>& z)
{
  if (hasMpi3) {
#if MPI_VERSION >= 3
    return MPI_C_DOUBLE_COMPLEX; // requires MPI 2.?
#else
    return MPI_DATATYPE_NULL; // FIXME (mfh 17 Nov 2016) Better to throw?
#endif // MPI_VERSION >= 3
  }
  else { // ! hasMpi3
    return Impl::computeKokkosComplexMpiDatatype (z);
  }
}

MPI_Datatype
MpiTypeTraits< ::Kokkos::complex<double> >::
getType ()
{
  if (hasMpi3) {
#if MPI_VERSION >= 3
    return MPI_C_DOUBLE_COMPLEX; // requires MPI 2.?
#else
    return MPI_DATATYPE_NULL; // FIXME (mfh 17 Nov 2016) Better to throw?
#endif // MPI_VERSION >= 3
  }
  else { // ! hasMpi3
    // Values are arbitrary.  The function just looks at the address
    // offsets of the class fields, not their contents.
    ::Kokkos::complex<double> z (3.0, 4.0);
    return Impl::computeKokkosComplexMpiDatatype (z);
  }
}

MPI_Datatype
MpiTypeTraits< ::Kokkos::complex<float> >::
getType (const ::Kokkos::complex<float>& z)
{
  if (hasMpi3) {
#if MPI_VERSION >= 3
    return MPI_C_FLOAT_COMPLEX; // requires MPI 2.?
#else
    return MPI_DATATYPE_NULL; // FIXME (mfh 17 Nov 2016) Better to throw?
#endif // MPI_VERSION >= 3
  }
  else { // ! hasMpi3
    return Impl::computeKokkosComplexMpiDatatype (z);
  }
}

MPI_Datatype
MpiTypeTraits< ::Kokkos::complex<float> >::
getType ()
{
  if (hasMpi3) {
#if MPI_VERSION >= 3
    return MPI_C_FLOAT_COMPLEX; // requires MPI 2.?
#else
    return MPI_DATATYPE_NULL; // FIXME (mfh 17 Nov 2016) Better to throw?
#endif // MPI_VERSION >= 3
  }
  else { // ! hasMpi3
    // Values are arbitrary.  The function just looks at the address
    // offsets of the class fields, not their contents.
    ::Kokkos::complex<float> z (3.0, 4.0);
    return Impl::computeKokkosComplexMpiDatatype (z);
  }
}

} // namespace Details
} // namespace Teuchos


#endif // HAVE_TPETRACORE_MPI