gasnetex_internal.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.
 */
 
// GASNet-EX network module internals
 
#ifndef GASNETEX_INTERNAL_H
#define GASNETEX_INTERNAL_H
 
#include "realm/gasnetex/gasnetex_module.h"
 
#ifdef REALM_USE_GASNETEX_WRAPPER
// Disable the definitions and importing of symbols we're not going to use
#define GEX_NO_PROTOTYPES 1
#endif
 
#include "realm/gasnetex/gasnetex_wrapper/gasnetex_wrapper.h"
 
#include "realm/bgwork.h"
#include "realm/atomics.h"
#include "realm/activemsg.h"
#include "realm/lists.h"
 
namespace Realm {
 
#define REALM_GEX_MODULE_VERSION 20240301
 
  extern gex_wrapper_handle_t gex_wrapper_handle;
 
  // rdma pointers need to identify which endpoint they belong to
  struct GASNetEXRDMAInfo {
    uintptr_t base;
    gex_ep_index_t ep_index;
  };
 
  class OutbufManager;
 
  class OutbufMetadata {
  public:
    OutbufMetadata();
 
    enum State
    {
      STATE_IDLE,
      STATE_DATABUF,
      STATE_PKTBUF,
    };
 
    void dec_usecount();
 
    void databuf_close();
 
    enum PktType
    {
      PKTTYPE_INVALID,
      PKTTYPE_INLINE,
      PKTTYPE_INLINE_SHORT,
      PKTTYPE_LONG,
      PKTTYPE_RGET,
      PKTTYPE_PUT,
      PKTTYPE_CANCELLED,
      PKTTYPE_COPY_IN_PROGRESS,
    };
 
    static const int MAX_PACKETS = 256;
 
    bool pktbuf_reserve(size_t bytes, int &pktidx, uintptr_t &offset);
    void pktbuf_close();
    uintptr_t pktbuf_get_offset(int pktidx);
    bool pktbuf_commit(int pktidx, PktType pkttype, bool update_realbuf);
 
  protected:
    friend class OutbufManager;
    friend class XmitSrcDestPair;
 
    void set_state(State new_state);
 
    State state;
    OutbufManager *manager;
    OutbufMetadata *nextbuf;
    uintptr_t baseptr;
    size_t size;
 
    bool is_overflow;
    OutbufMetadata *next_overflow;
    atomic<OutbufMetadata *> realbuf;
 
    atomic<int> remain_count;
 
    // dbuf reservations are NOT thread-safe - external serialization is used
    size_t databuf_rsrv_offset;
    int databuf_use_count;
 
    // pbuf reservatsions are NOT thread-safe - external serialization is used
    atomic<int> pktbuf_total_packets; // unsynchronized read ok
    size_t pktbuf_rsrv_offset;
    size_t pktbuf_pkt_ends[MAX_PACKETS]; // stores where a packet _ends_
 
    // pbuf commits are lock-free
    atomic<PktType> pktbuf_pkt_types[MAX_PACKETS];
    atomic<int> pktbuf_ready_packets;
 
    // pbuf consumption is NOT thread-safe - one pusher at a time
    int pktbuf_sent_packets;
    size_t pktbuf_sent_offset;
    int pktbuf_use_count;
  };
 
  class OutbufUsecountDec {
  public:
    OutbufUsecountDec(OutbufMetadata *_md)
      : md(_md)
    {}
 
    void operator()() const { md->dec_usecount(); }
 
  protected:
    OutbufMetadata *md;
  };
 
  class OutbufManager : public BackgroundWorkItem {
  public:
    OutbufManager();
    ~OutbufManager();
 
    void init(size_t _outbuf_count, size_t _outbuf_size, uintptr_t _baseptr);
 
    OutbufMetadata *alloc_outbuf(OutbufMetadata::State state, bool overflow_ok,
                                 bool new_endpoint);
    void free_outbuf(OutbufMetadata *md);
 
    virtual bool do_work(TimeLimit work_until);
 
  protected:
    OutbufMetadata *metadatas;
    size_t outbuf_size;
    Mutex mutex;
    OutbufMetadata *first_available;
    // we manage the copying over overflow bufs back to real outbufs - track
    //  both pending overflow bufs and reserved outbufs and request bgwork
    //  time when we've got at least one of each
    size_t num_overflow, num_reserved;
    // In order to detect the case of misconfigured objcounts given the
    // number of endpoints we count how many different endpoints we
    // need to allocate buffers for, if it is more than the total number
    // of buffers then we fail immediately since we'll never succeed
    // TODO: Remove this once we insert a slow insertion pathway
    // from dynamically allocated output buffers that are not
    // registered with GASNet
    // https://github.com/StanfordLegion/realm/issues/239
    size_t num_buffers, num_endpoints;
    OutbufMetadata *overflow_head;
    OutbufMetadata **overflow_tail;
    OutbufMetadata *reserved_head;
  };
 
  class PendingCompletionManager;
 
  class PendingCompletion {
  public:
    PendingCompletion();
 
    void *add_local_completion(size_t bytes, bool late_ok = false);
    void *add_remote_completion(size_t bytes);
 
    // marks ready and returns true if non-empty, else resets and returns false
    bool mark_ready(unsigned exp_local, unsigned exp_remote);
 
    bool has_local_completions();
    bool has_remote_completions();
 
    // these two calls can be concurrent, which complicates the determination of
    //  who can free the entry
    bool invoke_local_completions();
    bool invoke_remote_completions();
 
    int index;
    PendingCompletion *next_free;
    PendingCompletionManager *manager;
    // these three bits need to be in a single atomic location
    static const unsigned LOCAL_PENDING_BIT = 1;
    static const unsigned REMOTE_PENDING_BIT = 2;
    static const unsigned READY_BIT = 4;
    Realm::atomic<unsigned> state;
    Realm::atomic<unsigned> local_left, remote_left;
    size_t local_bytes, remote_bytes;
 
    static const size_t TOTAL_CAPACITY = 256;
    typedef char Storage_unaligned[TOTAL_CAPACITY];
    REALM_ALIGNED_TYPE_CONST(Storage_aligned, Storage_unaligned,
                             Realm::CompletionCallbackBase::ALIGNMENT);
    Storage_aligned storage;
  };
 
  struct PendingCompletionGroup {
    static const size_t LOG2_GROUPSIZE = 8; // 256 per group
 
    PendingCompletion entries[1 << LOG2_GROUPSIZE];
  };
 
  class PendingCompletionManager {
  public:
    PendingCompletionManager();
    ~PendingCompletionManager();
 
    PendingCompletion *get_available();
    void recycle_comp(PendingCompletion *comp);
 
    PendingCompletion *lookup_completion(int index);
    void invoke_completions(PendingCompletion *comp, bool do_local, bool do_remote);
 
    size_t num_completions_pending();
 
    bool over_pending_completion_soft_limit() const;
 
  protected:
    // NOTE: we stuff completion IDs, message IDs and 2 more bits into a
    //  32-bit word, so we're limited to 2^(30-12) = 256K completions
    static const size_t LOG2_MAXGROUPS = 10;
 
    // protects pops from the free list (to avoid A-B-A problem), but NOT pushes
    Realm::Mutex mutex;
    Realm::atomic<PendingCompletion *> first_free;
    Realm::atomic<int> num_groups; // number of groups currently allocated
    Realm::atomic<PendingCompletionGroup *> groups[1 << LOG2_MAXGROUPS];
    atomic<size_t> num_pending;
    size_t pending_soft_limit; // try to stall traffic above this threshold
  };
 
  template <typename T, unsigned CHUNK_SIZE>
  class ChunkedRecycler {
  public:
    ChunkedRecycler();
    ~ChunkedRecycler();
 
    T *alloc_obj();
    void free_obj(T *obj);
 
  protected:
    // we'll store a pointer back to the chunk header just past each object
    // each object needs enough space for a pointer back to the chunk header
    struct Chunk;
    struct WithPtr {
      typedef char Storage_unaligned[sizeof(T)];
      REALM_ALIGNED_TYPE_SAMEAS(Storage_aligned, Storage_unaligned, T);
      union {
        Storage_aligned raw_storage;
        uintptr_t nextptr;
      };
      Chunk *backptr;
    };
    struct Chunk {
      atomic<unsigned> remaining_count;
      Chunk *next_chunk;
      WithPtr elements[CHUNK_SIZE];
    };
 
    Mutex mutex;
    uintptr_t free_head;
    uintptr_t *free_tail;
    Chunk *chunks_head;
    atomic<size_t> cur_alloc;
    size_t cur_capacity, max_alloc;
  };
 
  struct PendingPutHeader;
 
  struct PreparedMessage {
  protected:
    // should not be directly allocated
    template <typename T, unsigned CHUNK_SIZE>
    friend class ChunkedRecycler;
    PreparedMessage() {}
    ~PreparedMessage() {}
 
  public:
    enum Strategy
    {
      STRAT_UNKNOWN,
      STRAT_SHORT_IMMEDIATE,  // AM short, attempt immediate
      STRAT_SHORT_PBUF,       // AM short, deferred in pktbuf
      STRAT_MEDIUM_IMMEDIATE, // AM medium, attempt immediate
      STRAT_MEDIUM_PBUF,      // AM medium, header/data in pktbuf
      STRAT_MEDIUM_MALLOCSRC, // AM medium, malloc'd temp source (HACK)
      STRAT_MEDIUM_PREP,      // AM medium, using NPAM
      STRAT_LONG_IMMEDIATE,
      STRAT_LONG_PBUF,
      STRAT_RGET_IMMEDIATE,
      STRAT_RGET_PBUF,
      STRAT_PUT_IMMEDIATE,
      STRAT_PUT_PBUF,
    };
    Strategy strategy;
    gex_rank_t target;
    gex_ep_index_t source_ep_index, target_ep_index;
    unsigned short msgid;
    uintptr_t dest_payload_addr;
    void *temp_buffer;
    OutbufMetadata *databuf;
    OutbufMetadata *pktbuf;
    int pktidx;
    gex_am_src_desc_opaque_t srcdesc;
    PendingPutHeader *put;
  };
 
  // we'll keep separate transmit queues for each src/dst pair
  //  (i.e src_ep_index, dst_rank, dst_ep_index), hopefully avoiding any
  //  head of line blocking if networking resources are exhausted on a given
  //  path
  class XmitSrcDestPair {
  public:
    XmitSrcDestPair(GASNetEXInternal *_internal, gex_ep_index_t _src_ep_index,
                    gex_rank_t _tgt_rank, gex_ep_index_t _tgt_ep_index);
    ~XmitSrcDestPair();
 
    // indicates whether any packets are pending, including those that
    //  have not been committed yet - useful for quiescence detection but
    //  also for ensuring new packets don't jump the queue
    bool has_packets_queued() const;
 
    // a packet that is sent immediately is counted as if it was
    //  reserved and then sent in rapid succession
    void record_immediate_packet();
 
    // reserves space in a pbuf for an outbound packet (allocating and
    //  enqueuing a pbuf if needed)
    bool reserve_pbuf_inline(size_t hdr_bytes, size_t payload_bytes, bool overflow_ok,
                             OutbufMetadata *&pktbuf, int &pktidx, void *&hdr_base,
                             void *&payload_base);
    bool reserve_pbuf_long_rget(size_t hdr_bytes, bool overflow_ok,
                                OutbufMetadata *&pktbuf, int &pktidx, void *&hdr_base);
    bool reserve_pbuf_put(bool overflow_ok, OutbufMetadata *&pktbuf, int &pktidx);
    void commit_pbuf_inline(OutbufMetadata *pktbuf, int pktidx, const void *hdr_base,
                            gex_am_arg_t arg0, size_t act_payload_bytes);
    void commit_pbuf_long(OutbufMetadata *pktbuf, int pktidx, const void *hdr_base,
                          gex_am_arg_t arg0, const void *payload_base,
                          size_t payload_bytes, uintptr_t dest_addr,
                          OutbufMetadata *databuf);
    void commit_pbuf_rget(OutbufMetadata *pktbuf, int pktidx, const void *hdr_base,
                          gex_am_arg_t arg0, const void *payload_base,
                          size_t payload_bytes, uintptr_t dest_addr,
                          gex_ep_index_t src_ep_index, gex_ep_index_t tgt_ep_index);
    void commit_pbuf_put(OutbufMetadata *pktbuf, int pktidx, PendingPutHeader *put,
                         const void *payload_base, size_t payload_bytes,
                         uintptr_t dest_addr);
    void cancel_pbuf(OutbufMetadata *pktbuf, int pktidx);
 
    void enqueue_completion_reply(gex_am_arg_t comp_info);
 
    void enqueue_put_header(PendingPutHeader *put);
 
    // adds the xpair to the injector ready list or the poller critical
    //  pair list as appropriate, eventually resulting in a call to
    //  push_packets - MUST NOT be called until the push_packets
    //  that resulted from the previous enqueue is done-ish (i.e. not going
    //  to push anything else)
    void request_push(bool force_critical);
 
    void push_packets(bool immediate_mode, TimeLimit work_until);
 
    long long time_since_failure() const;
 
    // used when cfg_am_limit is nonzero
    bool try_consume_am_credit();
    void return_am_credits(int count);
 
    IntrusiveListLink<XmitSrcDestPair> xpair_list_link;
    REALM_PMTA_DEFN(XmitSrcDestPair, IntrusiveListLink<XmitSrcDestPair>, xpair_list_link);
    typedef IntrusiveList<XmitSrcDestPair,
                          REALM_PMTA_USE(XmitSrcDestPair, xpair_list_link), DummyLock>
        XmitPairList;
 
    struct LongRgetData {
      const void *payload_base;
      size_t payload_bytes;
      uintptr_t dest_addr;
      union {
        struct {
          OutbufMetadata *databuf;
        } l;
        struct {
          // rget needs to give both src and target ep index for data
          gex_ep_index_t src_ep_index, tgt_ep_index;
        } r;
      };
    };
 
    struct PutMetadata {
      const void *src_addr;
      uintptr_t dest_addr;
      size_t payload_bytes;
      PendingPutHeader *put;
    };
 
  protected:
    friend class GASNetEXInternal;
 
    bool reserve_pbuf_helper(size_t total_bytes, bool overflow_ok,
                             OutbufMetadata *&pktbuf, int &pktidx, uintptr_t &baseptr);
    bool commit_pbuf_helper(OutbufMetadata *pktbuf, int pktidx, const void *hdr_base,
                            uintptr_t &baseptr);
 
    GASNetEXInternal *internal;
    gex_ep_index_t src_ep_index;
    gex_rank_t tgt_rank;
    gex_ep_index_t tgt_ep_index;
    atomic<size_t> packets_reserved, packets_sent;
    Mutex mutex;
    // we don't hold the mutex while pushing packets, but we need definitely
    //  don't want multiple threads trying to push for the same src/dst pair
    MutexChecker push_mutex_check;
    atomic<OutbufMetadata *> first_pbuf; // read without mutex
    OutbufMetadata *cur_pbuf;
    atomic<unsigned> imm_fail_count;
    bool has_ready_packets;
    long long first_fail_time;
    atomic<PendingPutHeader *> put_head; // read without mutex
    atomic<PendingPutHeader *> *put_tailp;
    // circular queue of pending completion replys
    gex_am_arg_t *comp_reply_data;
    unsigned comp_reply_wrptr, comp_reply_rdptr;
    atomic<unsigned> comp_reply_count; // read without mutex
    unsigned comp_reply_capacity;
    // TODO: track packets in flight to avoid clogging?
    atomic<int> am_credits; // available end-to-end AM credits for this peer
  };
 
  class XmitSrc {
  public:
    XmitSrc(GASNetEXInternal *_internal, gex_ep_index_t _src_ep_index);
    ~XmitSrc();
 
    XmitSrcDestPair *lookup_pair(gex_rank_t tgt_rank, gex_ep_index_t tgt_ep_index);
 
  protected:
    friend class GASNetEXInternal;
 
    GASNetEXInternal *internal;
    gex_ep_index_t src_ep_index;
 
    // we'll allocate XmitSrcDestPair's on demand - atomics allow nonblocking
    //  lookup
    atomic<XmitSrcDestPair *> *pairs;
 
    // TODO: track in flight work at src level?
  };
 
  struct PendingReverseGet;
 
  class GASNetEXEvent {
  protected:
    // should not be directly allocated
    template <typename T, unsigned CHUNK_SIZE>
    friend class ChunkedRecycler;
    GASNetEXEvent();
    ~GASNetEXEvent() {}
 
  public:
    gex_event_opaque_t get_event() const;
 
    GASNetEXEvent &set_event(gex_event_opaque_t _event);
    GASNetEXEvent &set_local_comp(PendingCompletion *_local_comp);
    GASNetEXEvent &set_pktbuf(OutbufMetadata *_pktbuf);
    GASNetEXEvent &set_databuf(OutbufMetadata *_databuf);
    GASNetEXEvent &set_rget(PendingReverseGet *_rget);
    GASNetEXEvent &set_put(PendingPutHeader *_put);
    GASNetEXEvent &set_leaf(GASNetEXEvent *_leaf);
 
    void propagate_to_leaves();
 
    void trigger(GASNetEXInternal *internal);
 
    IntrusiveListLink<GASNetEXEvent> event_list_link;
    REALM_PMTA_DEFN(GASNetEXEvent, IntrusiveListLink<GASNetEXEvent>, event_list_link);
    typedef IntrusiveList<GASNetEXEvent, REALM_PMTA_USE(GASNetEXEvent, event_list_link),
                          DummyLock>
        EventList;
 
  protected:
    gex_ep_opaque_t event;
    PendingCompletion *local_comp;
    OutbufMetadata *pktbuf;
    OutbufMetadata *databuf;
    PendingReverseGet *rget;
    PendingPutHeader *put;
    GASNetEXEvent *leaf;
  };
 
  // an injector tries to send packets, but is not allowed to stall -
  //  all requests must use GEX_FLAG_IMMEDIATE and any failure results in
  //  handing the xpair off to the poller (which is allowed to block eventually)
  class GASNetEXInjector : public BackgroundWorkItem {
  public:
    GASNetEXInjector(GASNetEXInternal *_internal);
 
    void add_ready_xpair(XmitSrcDestPair *xpair);
 
    bool has_work_remaining();
 
    virtual bool do_work(TimeLimit work_until);
 
  protected:
    GASNetEXInternal *internal;
    Mutex mutex;
    XmitSrcDestPair::XmitPairList ready_xpairs;
  };
 
  class GASNetEXPoller : public BackgroundWorkItem {
  public:
    GASNetEXPoller(GASNetEXInternal *_internal);
 
    void begin_polling();
    void end_polling();
 
    void add_critical_xpair(XmitSrcDestPair *xpair);
 
    void add_pending_event(GASNetEXEvent *event);
 
    bool has_work_remaining();
 
    virtual bool do_work(TimeLimit work_until);
 
    // causes calling thread to wait for a full call to gasnet_AMPoll() to
    //  be performed by the poller
    void wait_for_full_poll_cycle();
 
  protected:
    bool started = false;
    GASNetEXInternal *internal;
    Mutex mutex;
    atomic<bool> shutdown_flag; // set/cleared inside mutex, but tested outside
    Mutex::CondVar shutdown_cond;
    atomic<bool> pollwait_flag; // set/cleared inside mutex, but tested outside
    Mutex::CondVar pollwait_cond;
    XmitSrcDestPair::XmitPairList critical_xpairs;
    GASNetEXEvent::EventList pending_events;
  };
 
  class GASNetEXCompleter : public BackgroundWorkItem {
  public:
    GASNetEXCompleter(GASNetEXInternal *_internal);
 
    void add_ready_events(GASNetEXEvent::EventList &newly_ready);
 
    bool has_work_remaining();
 
    virtual bool do_work(TimeLimit work_until);
 
  protected:
    GASNetEXInternal *internal;
    Mutex mutex;
    atomic<bool> has_work; // can be read without mutex
    GASNetEXEvent::EventList ready_events;
  };
 
  struct PendingPutHeader {
  protected:
    // should not be directly allocated
    template <typename T, unsigned CHUNK_SIZE>
    friend class ChunkedRecycler;
    PendingPutHeader() {}
    ~PendingPutHeader() {}
 
  public:
    gex_rank_t target;
    gex_ep_index_t src_ep_index, tgt_ep_index;
    gex_am_arg_t arg0;
    static const size_t MAX_HDR_SIZE = 128;
    size_t hdr_size;
    unsigned char hdr_data[MAX_HDR_SIZE];
    uintptr_t src_ptr, tgt_ptr;
    size_t payload_bytes;
    PendingCompletion *local_comp;
    XmitSrcDestPair *xpair;
    atomic<PendingPutHeader *> next_put;
  };
 
  // NOTE: ReverseGetter exists for now because we have to use RMAs instead
  //  of long AMs for transfers between segments bound to non-primordial
  //  endpoints (on the source and/or the target) - once GASNet supports
  //  AMs on all endpoints, this code should probably go away
  class ReverseGetter;
 
  struct PendingReverseGet {
  protected:
    // should not be directly allocated
    template <typename T, unsigned CHUNK_SIZE>
    friend class ChunkedRecycler;
    PendingReverseGet() {}
    ~PendingReverseGet() {}
 
  public:
    ReverseGetter *rgetter;
    PendingReverseGet *next_rget;
    gex_rank_t srcrank;
    gex_ep_index_t src_ep_index, tgt_ep_index;
    gex_am_arg_t arg0;
    static const size_t MAX_HDR_SIZE = 128;
    size_t hdr_size;
    unsigned char hdr_data[MAX_HDR_SIZE];
    uintptr_t src_ptr, tgt_ptr;
    size_t payload_bytes;
  };
 
  class ReverseGetter : public BackgroundWorkItem {
  public:
    ReverseGetter(GASNetEXInternal *_internal);
 
    void add_reverse_get(gex_rank_t srcrank, gex_ep_index_t src_ep_index,
                         gex_ep_index_t tgt_ep_index, gex_am_arg_t arg0, const void *hdr,
                         size_t hdr_bytes, uintptr_t src_ptr, uintptr_t tgt_ptr,
                         size_t payload_bytes);
 
    bool has_work_remaining();
 
    virtual bool do_work(TimeLimit work_until);
 
  protected:
    friend class GASNetEXEvent;
    void reverse_get_complete(PendingReverseGet *rget);
 
    GASNetEXInternal *internal;
    Mutex mutex;
    PendingReverseGet *head;
    PendingReverseGet **tailp;
    ChunkedRecycler<PendingReverseGet, 8> rget_alloc;
  };
 
  class GASNetEXInternal {
  public:
    GASNetEXInternal(GASNetEXModule *_module, RuntimeImpl *_runtime);
 
    ~GASNetEXInternal();
 
    bool init(int *argc, const char ***argv);
    uintptr_t attach(size_t size);
 
    bool attempt_binding(void *base, size_t size, NetworkSegmentInfo::MemoryType memtype,
                         NetworkSegmentInfo::MemoryTypeExtraData memextra,
                         gex_ep_index_t *ep_indexp);
    void publish_bindings();
 
    void detach();
 
    void get_shared_peers(Realm::NodeSet &shared_peers);
    void barrier();
    void broadcast(gex_rank_t root, const void *val_in, void *val_out, size_t bytes);
    void gather(gex_rank_t root, const void *val_in, void *vals_out, size_t bytes);
    void allgatherv(const char *val_in, size_t bytes, std::vector<char> &vals_out,
                    std::vector<size_t> &lengths);
 
    size_t sample_messages_received_count();
    bool check_for_quiescence(size_t sampled_receive_count);
 
    PendingCompletion *get_available_comp();
    PendingCompletion *early_local_completion(PendingCompletion *comp);
 
    size_t recommended_max_payload(gex_rank_t target, gex_ep_index_t target_ep_index,
                                   bool with_congestion, size_t header_size,
                                   uintptr_t dest_payload_addr);
    size_t recommended_max_payload(gex_rank_t target, gex_ep_index_t target_ep_index,
                                   const void *data, size_t bytes_per_line, size_t lines,
                                   size_t line_stride, bool with_congestion,
                                   size_t header_size, uintptr_t dest_payload_addr);
    size_t recommended_max_payload(bool with_congestion, size_t header_size);
 
    PreparedMessage *prepare_message(gex_rank_t target, gex_ep_index_t target_ep_index,
                                     unsigned short msgid, void *&header_base,
                                     size_t header_size, void *&payload_base,
                                     size_t payload_size, uintptr_t dest_payload_addr);
    void commit_message(PreparedMessage *msg, PendingCompletion *comp, void *header_base,
                        size_t header_size, void *payload_base, size_t payload_size);
    void cancel_message(PreparedMessage *msg);
 
    gex_am_arg_t handle_short(gex_rank_t srcrank, gex_am_arg_t arg0, const void *hdr,
                              size_t hdr_bytes);
    gex_am_arg_t handle_medium(gex_rank_t srcrank, gex_am_arg_t arg0, const void *hdr,
                               size_t hdr_bytes, const void *data, size_t data_bytes);
    gex_am_arg_t handle_long(gex_rank_t srcrank, gex_am_arg_t arg0, const void *hdr,
                             size_t hdr_bytes, const void *data, size_t data_bytes);
    void handle_reverse_get(gex_rank_t srcrank, gex_ep_index_t src_ep_index,
                            gex_ep_index_t tgt_ep_index, gex_am_arg_t arg0,
                            const void *hdr, size_t hdr_bytes, uintptr_t src_ptr,
                            uintptr_t tgt_ptr, size_t payload_bytes);
    size_t handle_batch(gex_rank_t srcrank, gex_am_arg_t arg0, gex_am_arg_t cksum,
                        const void *data, size_t data_bytes, gex_am_arg_t *comps);
    void handle_completion_reply(gex_rank_t srcrank, const gex_am_arg_t *args,
                                 size_t nargs);
 
  protected:
    friend class ReverseGetter;
    friend class XmitSrc;
    friend class XmitSrcDestPair;
    friend class GASNetEXEvent;
    friend class GASNetEXPoller;
    friend class GASNetEXCompleter;
 
    // callbacks from IncomingMessageManager
    static void short_message_complete(NodeID sender, uintptr_t objptr,
                                       uintptr_t comp_info);
    static void medium_message_complete(NodeID sender, uintptr_t objptr,
                                        uintptr_t comp_info);
    static void long_message_complete(NodeID sender, uintptr_t objptr,
                                      uintptr_t comp_info);
 
    PendingCompletion *extract_arg0_local_comp(gex_am_arg_t &arg0);
 
    GASNetEXModule *module;
    RuntimeImpl *runtime;
    gex_client_opaque_t client;
    // order in 'eps' should match GASNet's indexing
    std::vector<gex_ep_opaque_t> eps;
    gex_tm_opaque_t prim_tm;
    gex_rank_t prim_rank, prim_size;
    gex_segment_opaque_t prim_segment;
    size_t prim_segsize;
 
    struct SegmentInfo {
      uintptr_t base, limit;
      gex_ep_index_t ep_index;
      gex_segment_opaque_t segment;
      NetworkSegmentInfo::MemoryType memtype;
      NetworkSegmentInfo::MemoryTypeExtraData memextra;
    };
    struct SegmentInfoSorter;
    // this list is sorted by address to enable quick address lookup
    std::vector<SegmentInfo> segments_by_addr;
 
    const SegmentInfo *find_segment(const void *srcptr) const;
 
    std::vector<XmitSrc *> xmitsrcs;
 
    GASNetEXPoller poller;
    GASNetEXInjector injector;
    GASNetEXCompleter completer;
    ReverseGetter rgetter;
    PendingCompletionManager compmgr;
    OutbufManager obmgr;
 
    // TODO: split counter into per-thread values to avoid contention?
    atomic<uint64_t> total_packets_received;
 
    // manage a single open databuf for all endpoints
    Mutex databuf_mutex;
    OutbufMetadata *databuf_md;
    gex_callback_handle_t gex_callback_handle;
 
    // allocator/managers for various objects we want to reuse
    ChunkedRecycler<GASNetEXEvent, 64> event_alloc;
    ChunkedRecycler<PreparedMessage, 32> prep_alloc;
    ChunkedRecycler<PendingPutHeader, 32> put_alloc;
  };
 
}; // namespace Realm
 
#endif