activemsg.h

Go to the documentation of this file.

/*
 * Copyright 2025 Stanford University, NVIDIA Corporation
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
// hopefully a more user-friendly C++ template wrapper for GASNet active
//  messages...
 
#ifndef ACTIVEMSG_H
#define ACTIVEMSG_H
 
#include "realm/realm_config.h"
#include "realm/fragmented_message.h"
#include "realm/mutex.h"
#include "realm/serialize.h"
#include "realm/nodeset.h"
#include "realm/network.h"
#include "realm/atomics.h"
#include "realm/threads.h"
#include "realm/bgwork.h"
#include <type_traits>
#include <mutex>
 
#include <optional>
 
namespace Realm {
 
  namespace Config {
    // if true, the number and min/max/avg/stddev duration of handler per
    //  message type is recorded and printed
    extern bool profile_activemsg_handlers;
 
    // the maximum time we're willing to spend on inline message
    //  handlers
    extern long long max_inline_message_time;
  }; // namespace Config
 
  enum
  {
    PAYLOAD_NONE,    // no payload in packet
    PAYLOAD_KEEP,    // use payload pointer, guaranteed to be stable
    PAYLOAD_FREE,    // take ownership of payload, free when done
    PAYLOAD_COPY,    // make a copy of the payload
    PAYLOAD_SRCPTR,  // payload has been copied to the src data pool
    PAYLOAD_PENDING, // payload needs to be copied, but hasn't yet
    PAYLOAD_KEEPREG, // use payload pointer, AND it's registered!
    PAYLOAD_EMPTY,   // message can have payload, but this one is 0 bytes
  };
 
  class ActiveMessageImpl;
 
  template <typename T, size_t INLINE_STORAGE = 256>
  class ActiveMessage {
  public:
    // constructs an INACTIVE message object - call init(...) as needed
    ActiveMessage();
 
    // construct a new active message for either a single recipient or a mask
    //  of recipients
    // in addition to the header struct (T), a message can include a variable
    //  payload which can be delivered to a particular destination address
    ActiveMessage(NodeID _target, size_t _max_payload_size = 0);
    ActiveMessage(NodeID _target, size_t _max_payload_size,
                  const RemoteAddress &_dest_payload_addr);
    ActiveMessage(const Realm::NodeSet &_targets, size_t _max_payload_size = 0);
 
    // providing the payload (as a 1D or 2D reference, which must be PAYLOAD_KEEP)
    //  up front can avoid a copy if the source location is directly accessible
    //  by the networking hardware
    ActiveMessage(NodeID _target, const void *_data, size_t _datalen);
    ActiveMessage(NodeID _target, const LocalAddress &_src_payload_addr, size_t _datalen,
                  const RemoteAddress &_dest_payload_addr);
    ActiveMessage(const Realm::NodeSet &_targets, const void *_data, size_t _datalen);
    ActiveMessage(NodeID _target, const void *_data, size_t _bytes_per_line,
                  size_t _lines, size_t _line_stride);
    ActiveMessage(NodeID _target, const LocalAddress &_src_payload_addr,
                  size_t _bytes_per_line, size_t _lines, size_t _line_stride,
                  const RemoteAddress &_dest_payload_addr);
    ActiveMessage(const Realm::NodeSet &_targets, const void *_data,
                  size_t _bytes_per_line, size_t _lines, size_t _line_stride);
 
    ~ActiveMessage(void);
 
    // a version of `init` for each constructor above
    void init(NodeID _target, size_t _max_payload_size = 0);
    void init(NodeID _target, size_t _max_payload_size,
              const RemoteAddress &_dest_payload_addr);
    void init(const Realm::NodeSet &_targets, size_t _max_payload_size = 0);
    void init(NodeID _target, const void *_data, size_t _datalen);
    void init(NodeID _target, const LocalAddress &_src_payload_addr, size_t _datalen,
              const RemoteAddress &_dest_payload_addr);
    void init(const Realm::NodeSet &_targets, const void *_data, size_t _datalen);
    void init(NodeID _target, const void *_data, size_t _bytes_per_line, size_t _lines,
              size_t _line_stride);
    void init(NodeID _target, const LocalAddress &_src_payload_addr,
              size_t _bytes_per_line, size_t _lines, size_t _line_stride,
              const RemoteAddress &_dest_payload_addr);
    void init(const Realm::NodeSet &_targets, const void *_data, size_t _bytes_per_line,
              size_t _lines, size_t _line_stride);
 
    // large messages may need to be fragmented, so use cases that can
    //  handle the fragmentation at a higher level may want to know the
    //  largest size that is fragmentation-free - the answer can depend
    //  on whether the data is to be delivered to a known RemoteAddress
    //  and/or whether the source data location is known
    // a call that sets `with_congestion` may get a smaller value (maybe
    //  even 0) if the path to the named target(s) is getting full
    static size_t recommended_max_payload(NodeID target, bool with_congestion);
    static size_t recommended_max_payload(const NodeSet &targets, bool with_congestion);
    static size_t recommended_max_payload(NodeID target,
                                          const RemoteAddress &dest_payload_addr,
                                          bool with_congestion);
    static size_t recommended_max_payload(NodeID target, const void *data,
                                          size_t bytes_per_line, size_t lines,
                                          size_t line_stride, bool with_congestion);
    static size_t recommended_max_payload(const NodeSet &targets, const void *data,
                                          size_t bytes_per_line, size_t lines,
                                          size_t line_stride, bool with_congestion);
    static size_t
    recommended_max_payload(NodeID target, const LocalAddress &src_payload_addr,
                            size_t bytes_per_line, size_t lines, size_t line_stride,
                            const RemoteAddress &dest_payload_addr, bool with_congestion);
 
    // operator-> gives access to the header structure
    T *operator->(void);
    T &operator*(void);
 
    // variable payload can be written to in three ways:
    //  (a) Realm-style serialization (currently eager)
    template <typename T2>
    bool operator<<(const T2 &to_append);
    //  (b) old memcpy-like behavior (using the various payload modes)
    void add_payload(const void *data, size_t datalen, int payload_mode = PAYLOAD_COPY);
    void add_payload(const void *data, size_t bytes_per_line, size_t lines,
                     size_t line_stride, int payload_mode = PAYLOAD_COPY);
    //  (c) request for a pointer to write into (writes must be completed before
    //       call to commit or cancel)
    void *payload_ptr(size_t datalen);
 
    // register callbacks to be called upon:
    //  a) local completion - i.e. all source data has been read and can now
    //      be safely overwritten
    //  b) remote completion - message has been received AND HANDLED by target
    //
    // callbacks need to be "lightweight" - for heavier work, the message
    //  handler on the target can send an explicit response message
    template <typename CALLABLE>
    void add_local_completion(const CALLABLE &callable);
    template <typename CALLABLE>
    void add_remote_completion(const CALLABLE &callable);
 
    // every active message must eventually be commit()'ed or cancel()'ed
    void commit(void);
    void cancel(void);
 
  protected:
    ActiveMessageImpl *impl;
    T *header;
    Realm::Serialization::FixedBufferSerializer fbs;
    uint64_t inline_capacity[INLINE_STORAGE / sizeof(uint64_t)];
  };
 
  // type-erased wrappers for completion callbacks
  class REALM_INTERNAL_API_EXTERNAL_LINKAGE CompletionCallbackBase {
  public:
    virtual ~CompletionCallbackBase();
    virtual void invoke() = 0;
    virtual size_t size() const = 0;
    virtual CompletionCallbackBase *clone_at(void *p) const = 0;
 
    static const size_t ALIGNMENT = 8;
 
    // helper functions for invoking/cloning/destroying a collection of callbacks
    static void invoke_all(void *start, size_t bytes);
    static void clone_all(void *dst, const void *src, size_t bytes);
    static void destroy_all(void *start, size_t bytes);
  };
 
  template <typename CALLABLE>
  class CompletionCallback : public CompletionCallbackBase {
  public:
    CompletionCallback(const CALLABLE &_callable);
    virtual void invoke();
    virtual size_t size() const;
    virtual CompletionCallbackBase *clone_at(void *p) const;
 
  protected:
    CALLABLE callable;
  };
 
  // per-network active message implementations are mostly opaque, but a few
  //  fields are exposed to avoid virtual function calls
  class ActiveMessageImpl {
  public:
    virtual ~ActiveMessageImpl() {}
 
    // reserves space for a local/remote completion - caller will
    //  placement-new the completion at the provided address
    virtual void *add_local_completion(size_t size) = 0;
    virtual void *add_remote_completion(size_t size) = 0;
 
    virtual void commit(size_t act_payload_size) = 0;
    virtual void cancel() = 0;
 
    void *header_base;
    void *payload_base;
    size_t payload_size;
  };
 
  class ActiveMessageHandlerRegBase;
 
  struct ActiveMessageHandlerStats {
    atomic<size_t> count, sum, sum2, minval, maxval;
 
    ActiveMessageHandlerStats(void);
    void record(long long t_start, long long t_end);
  };
 
  struct FragmentInfo {
    uint32_t chunk_id{0};
    uint32_t total_chunks{0};
    uint64_t msg_id{0};
  };
 
  // singleton class that can convert message type->ID and ID->handler
  class ActiveMessageHandlerTable {
  public:
    ActiveMessageHandlerTable(void);
    ~ActiveMessageHandlerTable(void);
 
    typedef unsigned short MessageID;
    typedef void (*MessageHandler)(NodeID sender, const void *header, const void *payload,
                                   size_t payload_size, TimeLimit work_until);
    typedef void (*MessageHandlerNoTimeout)(NodeID sender, const void *header,
                                            const void *payload, size_t payload_size);
    typedef bool (*MessageHandlerInline)(NodeID sender, const void *header,
                                         const void *payload, size_t payload_size,
                                         TimeLimit work_until);
 
    template <typename T>
    MessageID lookup_message_id(void) const;
 
    const char *lookup_message_name(MessageID id);
    void record_message_handler_call(MessageID id, long long t_start, long long t_end);
    void report_message_handler_stats();
 
    static void append_handler_reg(ActiveMessageHandlerRegBase *new_reg);
 
    void construct_handler_table(void);
 
    typedef unsigned TypeHash;
 
    struct HandlerEntry {
      TypeHash hash;
      const char *name;
      MessageHandler handler;
      MessageHandlerNoTimeout handler_notimeout;
      MessageHandlerInline handler_inline;
      ActiveMessageHandlerStats stats;
 
      std::optional<const FragmentInfo &(*)(const void *)> extract_frag_info;
    };
 
    HandlerEntry *lookup_message_handler(MessageID id);
 
  protected:
    static ActiveMessageHandlerRegBase *pending_handlers;
 
    std::vector<HandlerEntry> handlers;
  };
 
  extern ActiveMessageHandlerTable activemsg_handler_table;
 
  class ActiveMessageHandlerRegBase {
  public:
    virtual ~ActiveMessageHandlerRegBase(void) {}
    virtual ActiveMessageHandlerTable::MessageHandler get_handler(void) const = 0;
    virtual ActiveMessageHandlerTable::MessageHandlerNoTimeout
    get_handler_notimeout(void) const = 0;
    virtual ActiveMessageHandlerTable::MessageHandlerInline
    get_handler_inline(void) const = 0;
 
    ActiveMessageHandlerTable::TypeHash hash;
    const char *name;
    bool must_free;
    ActiveMessageHandlerRegBase *next_handler;
    std::optional<const FragmentInfo &(*)(const void *)> extract_frag_info;
  };
 
  template <typename T, typename T2 = T>
  class ActiveMessageHandlerReg : public ActiveMessageHandlerRegBase {
  public:
    ActiveMessageHandlerReg(void);
    ~ActiveMessageHandlerReg(void);
 
    // when registering an active message handler, the following three methods
    //  are looked for in class T2
    // (a) void handle_message(NodeID, const T&, const void *, size_t, TimeLimit)
    // (b) void handle_message(NodeID, const T&, const void *, size_t)
    // (c) bool handle_inline(NodeID, const T&, const void *, size_t, TimeLimit)
    //
    // at least one of (a) or (b) must be present, with (a) being preferred
    //
    // if (c) is present, it will be used to attempt inline handling of
    //  active messages as they arrive, with the following constraints:
    //   (i) the handler must not block on any mutexes (trylocks are ok)
    //   (ii) the handler must not perform dynamic memory allocation/frees
    //   (iii) the handler must try very hard to stay within the specified
    //          time limit
    // if the inline handler is unable to satisfy these requirements, it should
    //  not attempt to handle the message, returning 'false' and letting it be
    //  queued as normal
 
    // returns either the requested kind of handler or a null pointer if
    //  it doesn't exist
    virtual ActiveMessageHandlerTable::MessageHandler get_handler(void) const;
    virtual ActiveMessageHandlerTable::MessageHandlerNoTimeout
    get_handler_notimeout(void) const;
    virtual ActiveMessageHandlerTable::MessageHandlerInline
    get_handler_inline(void) const;
 
    // this method does nothing, but can be called to force the instantiation
    //  of a handler registration object (needed when things are inside templates)
    void force_instantiation(void) {}
  };
 
  namespace ThreadLocal {
    // this flag will be true when we are running a message handler
    extern thread_local bool in_message_handler;
  }; // namespace ThreadLocal
 
  class REALM_INTERNAL_API_EXTERNAL_LINKAGE IncomingMessageManager
    : public BackgroundWorkItem {
  public:
    IncomingMessageManager(int _nodes, int _dedicated_threads,
                           Realm::CoreReservationSet &crs);
    ~IncomingMessageManager(void);
 
    typedef uintptr_t CallbackData;
    typedef void (*CallbackFnptr)(NodeID, CallbackData, CallbackData);
 
    // adds an incoming message to the queue
    // returns true if the call was handled immediately (in which case the
    //  callback, if present, will NOT be called), or false if the message
    //  will be processed later
    bool add_incoming_message(NodeID sender, ActiveMessageHandlerTable::MessageID msgid,
                              const void *hdr, size_t hdr_size, int hdr_mode,
                              const void *payload, size_t payload_size, int payload_mode,
                              CallbackFnptr callback_fnptr, CallbackData callback_data1,
                              CallbackData callback_data2, TimeLimit work_until);
 
    void start_handler_threads(size_t stack_size);
 
    // stalls caller until all incoming messages have been handled (and at
    //  least 'min_messages_handled' in total)
    void drain_incoming_messages(size_t min_messages_handled);
 
    void shutdown(void);
 
    virtual bool do_work(TimeLimit work_until);
 
    void handler_thread_loop(void);
 
  protected:
    struct MessageBlock;
 
    struct Message {
      MessageBlock *block;
      Message *next_msg;
      NodeID sender;
      ActiveMessageHandlerTable::HandlerEntry *handler;
      void *hdr;
      size_t hdr_size;
      bool hdr_needs_free;
      void *payload;
      size_t payload_size;
      bool payload_needs_free;
      CallbackFnptr callback_fnptr;
      CallbackData callback_data1, callback_data2;
    };
 
    struct MessageBlock {
      static MessageBlock *new_block(size_t _total_size);
      static void free_block(MessageBlock *block);
 
      void reset();
 
      // called with message manager lock held
      Message *append_message(size_t hdr_bytes_needed, size_t payload_bytes_needed);
 
      // called _without_ message manager lock held
      void recycle_message(Message *msg, IncomingMessageManager *manager);
 
      size_t total_size, size_used;
      atomic<unsigned> use_count;
      MessageBlock *next_free;
    };
 
    int get_messages(Message *&head, Message **&tail, bool wait);
    bool return_messages(int sender, size_t num_handled, Message *head, Message **tail);
 
    int nodes, dedicated_threads, sleeper_count;
    atomic<bool> bgwork_requested;
    int shutdown_flag;
    Message **heads;
    Message ***tails;
    bool *in_handler;
    int *todo_list; // list of nodes with non-empty message lists
    int todo_oldest, todo_newest;
    int handlers_active;
    bool drain_pending;
    size_t drain_min_count;
    size_t total_messages_handled;
    Mutex mutex;
    Mutex::CondVar condvar, drain_condvar;
    CoreReservation *core_rsrv;
    std::vector<Thread *> handler_threads;
    MessageBlock *current_block;
    MessageBlock *available_blocks;
    size_t num_available_blocks;
    size_t cfg_max_available_blocks, cfg_message_block_size;
 
    struct PairHash {
      std::size_t operator()(const std::pair<NodeID, uint64_t> &p) const
      {
        return std::hash<NodeID>()(p.first) ^ (std::hash<uint64_t>()(p.second) << 1);
      }
    };
 
    std::unordered_map<std::pair<NodeID, uint64_t>, std::unique_ptr<FragmentedMessage>,
                       PairHash>
        frag_message;
  };
 
  template <typename UserHdr>
  struct WrappedWithFragInfo {
    FragmentInfo frag_info;
    UserHdr user;
 
    UserHdr *operator->() { return &user; }
    const UserHdr *operator->() const { return &user; }
    UserHdr &operator*() { return user; }
    const UserHdr &operator*() const { return user; }
  };
 
  template <typename Hdr>
  using DefaultActiveMessageBuilder = ActiveMessage<WrappedWithFragInfo<Hdr>>;
 
  //
  // class ActiveMessageAuto
  //
  // A thin convenience wrapper that hides the decision of whether a payload
  // must be transmitted as a *single* ActiveMessage or split across multiple
  // FragmentedActiveMessages – hence the name "Auto".  The public interface
  // mirrors that of `ActiveMessage` so that existing call-sites can switch to
  // the automatic behaviour with only a type alias change.
  //
  // Behaviour
  // ---------
  // 1.  When `commit()` is called and the aggregated payload size is **less
  //     than or equal to** the `max_payload_size` provided to the
  //     constructor, the object internally constructs a single
  //     `ActiveMessage<WrappedWithFragInfo<UserHdr>>` (via the `Builder`
  //     template parameter) and sends it.
  // 2.  Otherwise the payload is automatically split into the minimum number
  //     of fragments.  Each fragment is sent with a compact `FragmentInfo`
  //     header and the remote side stitches them back together transparently
  //     via `FragmentedMessage`.
  //
  // Template Parameters
  // -------------------
  //  * `UserHdr` – the user-defined struct that forms the logical AM header.
  //  * `Builder` – meta-function that, given a *wire* header type, produces
  //               the concrete ActiveMessage builder to use.  The default
  //               simply wraps `UserHdr` with `FragmentInfo` so that the
  //               fragmentation metadata travels on the wire.
  //
  // Example
  // -------
  //   struct MyHdr { int value; };
  //   ActiveMessageAuto<MyHdr> msg(target_node,
  //                                 ActiveMessage<MyHdr>::recommended_max_payload(
  //                                     target_node, /*with_congestion=*/false));
  //   msg->value = 123;
  //   msg.add_payload(data_ptr, data_bytes);
  //   msg.commit();
  //
 
  template <typename UserHdr,
            template <typename> class Builder = DefaultActiveMessageBuilder>
  class ActiveMessageAuto {
  public:
    ActiveMessageAuto(NodeID target, size_t max_payload_size);
 
    UserHdr *operator->();
    UserHdr &operator*();
 
    void add_payload(const void *data, size_t size);
    void commit();
 
  private:
    using WireHdr = WrappedWithFragInfo<UserHdr>;
    uint64_t next_message_id(NodeID node_id);
 
    NodeID target_;
    size_t max_payload_size_;
    std::vector<char> payload_;
    UserHdr user_header_{};
  };
 
  template <typename UserHdr>
  struct AutoMessageRegistrar {
    ActiveMessageHandlerReg<WrappedWithFragInfo<UserHdr>, UserHdr> reg;
  };
 
} // namespace Realm
 
#include "realm/activemsg.inl"
 
#endif