GatherMessageInfo.h

// -*- tab-width: 2; indent-tabs-mode: nil; coding: utf-8-with-signature -*-
//-----------------------------------------------------------------------------
// Copyright 2000-2025 CEA (www.cea.fr) IFPEN (www.ifpenergiesnouvelles.com)
// See the top-level COPYRIGHT file for details.
// SPDX-License-Identifier: Apache-2.0
//-----------------------------------------------------------------------------
/*---------------------------------------------------------------------------*/
/* GatherMessageInfo.h                                         (C) 2000-2025 */
/*                                                                           */
/* Informations pour les messages 'gather'.                                  */
/*---------------------------------------------------------------------------*/
#ifndef ARCCORE_MESSAGEPASSING_GATHERMESSAGEINFO_H
#define ARCCORE_MESSAGEPASSING_GATHERMESSAGEINFO_H
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
#include "arccore/collections/Array.h"
#include "arccore/message_passing/MessageRank.h"
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
namespace Arcane::MessagePassing
{
class ARCCORE_MESSAGEPASSING_EXPORT GatherMessageInfoBase
{
 public:

  enum class Mode
  {
    Gather,
    GatherVariable,
    GatherVariableNeedComputeInfo,
    Null
  };
 
 public:

  GatherMessageInfoBase() = default;

  explicit GatherMessageInfoBase(MessageRank dest_rank)
  : m_destination_rank(dest_rank)
  {}

  GatherMessageInfoBase(MessageRank dest_rank, eBlockingType blocking_type)
  : m_destination_rank(dest_rank)
  , m_is_blocking(blocking_type == Blocking)
  {}
 
 public:
 
  void setBlocking(bool is_blocking)
  {
    m_is_blocking = is_blocking;
  }
  bool isBlocking() const { return m_is_blocking; }

  MessageRank destinationRank() const { return m_destination_rank; }

  void setDestinationRank(MessageRank rank)
  {
    m_destination_rank = rank;
  }

  Mode mode() const { return m_mode; }

  void print(std::ostream& o) const;
 
  friend std::ostream& operator<<(std::ostream& o, const GatherMessageInfoBase& pmessage)
  {
    pmessage.print(o);
    return o;
  }
 
 public:
 
  // Indique si le message est valide (i.e: il a été initialisé avec un message valide)
  bool isValid() const
  {
    if (m_mode == Mode::Null)
      return false;
    return true;
  }
 
 protected:
 
  void _setType(Mode t)
  {
    m_mode = t;
  }
 
 private:
 
  MessageRank m_destination_rank;
  bool m_is_blocking = true;
  Mode m_mode = Mode::Null;
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
template <typename DataType>
class GatherMessageInfo
: public GatherMessageInfoBase
{
 public:
 
  using BaseClass = GatherMessageInfoBase;
 
 public:

  GatherMessageInfo() = default;

  explicit GatherMessageInfo(MessageRank dest_rank)
  : BaseClass(dest_rank)
  {}

  GatherMessageInfo(MessageRank dest_rank, eBlockingType blocking_type)
  : BaseClass(dest_rank, blocking_type)
  {}
 
 public:

  void setGather(Span<const DataType> send_buf, Span<DataType> receive_buf)
  {
    _setType(Mode::Gather);
    m_receive_buf = receive_buf;
    m_send_buffer = send_buf;
  }

  void setGatherVariable(Span<const DataType> send_buf, Array<DataType>* receive_array)
  {
    _setType(Mode::GatherVariableNeedComputeInfo);
    m_local_reception_buffer = receive_array;
    m_send_buffer = send_buf;
  }

  void setGatherVariable(Span<const DataType> send_buf, Span<DataType> receive_buf,
                         Span<const Int32> receive_counts, Span<const Int32> receive_displacements)
  {
    _setType(Mode::GatherVariable);
    m_receive_buf = receive_buf;
    m_send_buffer = send_buf;
    m_receive_displacements = receive_displacements;
    m_receive_counts = receive_counts;
  }

  Array<DataType>* localReceptionBuffer() const { return m_local_reception_buffer; }

  Span<const DataType> sendBuffer() const { return m_send_buffer; }

  Span<DataType> receiveBuffer() const { return m_receive_buf; }

  Span<const Int32> receiveDisplacement() { return m_receive_displacements; }

  Span<const Int32> receiveCounts() const { return m_receive_counts; }
 
 private:
 
  Array<DataType>* m_local_reception_buffer = nullptr;
  Span<const DataType> m_send_buffer;
  Span<DataType> m_receive_buf;
  Span<const Int32> m_receive_displacements;
  Span<const Int32> m_receive_counts;
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
} // End namespace Arccore::MessagePassing
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
#endif