src/realm: allow for XferDes objects to be reset

src/realm/cuda/cuda_internal.cc

@@ -2186,6 +2186,15 @@ namespace Realm {
       return 0;
     }
 
+    void GPUfillXferDes::reset(const std::vector<off_t> &ib_offsets)
+    {
+      XferDes::reset(ib_offsets);
+      assert(input_control.control_port_idx == -1);
+      input_control.current_io_port = -1;
+      input_control.remaining_count = fill_total;
+      input_control.eos_received = true;
+    }
+
     ////////////////////////////////////////////////////////////////////////
     //
     // class GPUreduceXferDes

src/realm/cuda/cuda_internal.h

@@ -918,12 +918,15 @@ namespace Realm {
                      const std::vector<XferDesPortInfo> &outputs_info, int _priority,
                      const void *_fill_data, size_t _fill_size, size_t _fill_total);
 
+      void reset(const std::vector<off_t> &ib_offsets);
+
       long get_requests(Request **requests, long nr);
 
       bool progress_xd(GPUfillChannel *channel, TimeLimit work_until);
 
     protected:
       size_t reduced_fill_size;
+      size_t fill_total;
     };
 
     class GPUfillChannel : public SingleXDQChannel<GPUfillChannel, GPUfillXferDes> {

src/realm/transfer/address_list.cc

@@ -125,6 +125,14 @@ namespace Realm {
     field_block = _field_block;
   }
 
+  void AddressList::reset()
+  {
+    total_bytes = 0;
+    write_pointer = 0;
+    read_pointer = 0;
+    memset(data.data(), 0, max_entries * sizeof(size_t));
+  }
+
   size_t AddressList::bytes_pending() const { return total_bytes; }
 
   size_t AddressList::full_field_bytes()
@@ -349,6 +357,15 @@ namespace Realm {
     }
   }
 
+  void AddressListCursor::reset()
+  {
+    // Not touching the addrlist.
+    partial = false;
+    partial_dim = 0;
+    for(size_t i = 0; i < pos.size(); i++)
+      pos[i] = 0;
+  }
+
   const FieldBlock *AddressListCursor::field_block() const
   {
     return addrlist->field_block;

src/realm/transfer/address_list.h

@@ -65,6 +65,7 @@ namespace Realm {
     [[nodiscard]] size_t *begin_entry(int max_dim, bool wrap_around = true);
     void commit_entry(int act_dim, size_t bytes);
     void attach_field_block(const FieldBlock *_field_block);
+    void reset();
 
     [[nodiscard]] size_t bytes_pending() const;
     [[nodiscard]] size_t full_field_bytes();
@@ -112,6 +113,7 @@ namespace Realm {
     // ─── progress───────────────────────────────────────────────────────────────
     void advance(int dim, size_t amount, int f = 1);
     void skip_bytes(size_t bytes);
+    void reset();
 
     // ─── field accessors ──────────────────────────────────────────────────────
     [[nodiscard]] const FieldBlock *field_block() const;

src/realm/transfer/channel.cc

@@ -276,6 +276,13 @@ namespace Realm {
     }
   }
 
+  void SequenceAssembler::reset()
+  {
+    contig_amount_x2.store(0);
+    first_noncontig.store((size_t)-1);
+    spans.clear();
+  }
+
   ////////////////////////////////////////////////////////////////////////
   //
   // class ControlPort::Encoder
@@ -462,8 +469,6 @@ namespace Realm {
     , nb_update_pre_bytes_total_calls_received(0)
   {
     input_ports.resize(inputs_info.size());
-    int gather_control_port = -1;
-    int scatter_control_port = -1;
     for(size_t i = 0; i < inputs_info.size(); i++) {
       XferPort &p = input_ports[i];
       const XferDesPortInfo &ii = inputs_info[i];
@@ -489,6 +494,7 @@ namespace Realm {
       p.remote_bytes_total.store(size_t(-1));
       p.ib_offset = ii.ib_offset;
       p.ib_size = ii.ib_size;
+      p.ib_index = ii.ib_index;
       p.addrcursor.set_addrlist(&p.addrlist);
       switch(ii.port_type) {
       case XferDesPortInfo::GATHER_CONTROL_PORT:
@@ -555,6 +561,7 @@ namespace Realm {
       p.remote_bytes_total.store(size_t(-1));
       p.ib_offset = oi.ib_offset;
       p.ib_size = oi.ib_size;
+      p.ib_index = oi.ib_index;
       p.addrcursor.set_addrlist(&p.addrlist);
 
       // if we're writing into an IB, the first 'ib_size' byte
@@ -655,6 +662,83 @@ namespace Realm {
     }
   }
 
+  void XferDes::reset(const std::vector<off_t> &ib_offsets)
+  {
+    iteration_completed.store_release(false);
+    bytes_write_pending.store_release(0);
+    transfer_completed.store_release(false);
+    progress_counter.store_release(0);
+    nb_update_pre_bytes_total_calls_received.store_release(0);
+    for(auto &info : input_ports) {
+      info.iter->reset();
+      info.local_bytes_total = 0;
+      info.local_bytes_cons.store_release(0);
+      info.remote_bytes_total.store_release(size_t(-1));
+      info.needs_pbt_update.store(false);
+      info.seq_local.reset();
+      info.seq_remote.reset();
+      info.addrlist.reset();
+      info.addrcursor.reset();
+      // If this XD is using an IB, update its offset value
+      // to the latest one for this copy instantiation.
+      if(info.ib_size > 0) {
+        info.ib_offset = ib_offsets[info.ib_index];
+        // TODO (rohany): This is pretty hacky, but I don't know
+        //  a better way to do it...
+        auto wrapit = dynamic_cast<WrappingFIFOIterator *>(info.iter);
+        assert(wrapit);
+        wrapit->set_base(info.ib_offset);
+      }
+    }
+
+    if(gather_control_port >= 0) {
+      input_control.control_port_idx = gather_control_port;
+      input_control.current_io_port = 0;
+      input_control.remaining_count = 0;
+      input_control.eos_received = false;
+    } else {
+      input_control.control_port_idx = -1;
+      input_control.current_io_port = 0;
+      input_control.remaining_count = size_t(-1);
+      input_control.eos_received = false;
+    }
+
+    for(auto &info : output_ports) {
+      info.iter->reset();
+      info.needs_pbt_update.store_release(info.peer_guid != XFERDES_NO_GUID);
+      info.local_bytes_total = 0;
+      info.local_bytes_cons.store_release(0);
+      info.remote_bytes_total.store_release(size_t(-1));
+      info.seq_local.reset();
+      info.seq_remote.reset();
+      info.addrlist.reset();
+      info.addrcursor.reset();
+      if(info.ib_size > 0) {
+        info.ib_offset = ib_offsets[info.ib_index];
+        // TODO (rohany): This is pretty hacky, but I don't know
+        //  a better way to do it...
+        auto wrapit = dynamic_cast<WrappingFIFOIterator *>(info.iter);
+        assert(wrapit);
+        wrapit->set_base(info.ib_offset);
+        // Also mark the remote sequence assembler as capable of
+        // writing into the ib memory.
+        info.seq_remote.add_span(0, info.ib_size);
+      }
+    }
+
+    if(scatter_control_port >= 0) {
+      output_control.control_port_idx = scatter_control_port;
+      output_control.current_io_port = 0;
+      output_control.remaining_count = 0;
+      output_control.eos_received = false;
+    } else {
+      output_control.control_port_idx = -1;
+      output_control.current_io_port = 0;
+      output_control.remaining_count = size_t(-1);
+      output_control.eos_received = false;
+    }
+  }
+
 #define MAX_GEN_REQS 3
 
   bool support_2d_xfers(XferDesKind kind)
@@ -1996,6 +2080,7 @@ namespace Realm {
                                  size_t _fill_total)
     : XferDes(_dma_op, _channel, _launch_node, _guid, inputs_info, outputs_info,
               _priority, _fill_data, _fill_size)
+    , fill_total(_fill_total)
   {
     kind = XFER_MEM_FILL;
 
@@ -2007,6 +2092,15 @@ namespace Realm {
     input_control.eos_received = true;
   }
 
+  void MemfillXferDes::reset(const std::vector<off_t> &ib_offsets)
+  {
+    XferDes::reset(ib_offsets);
+    assert(input_control.control_port_idx == -1);
+    input_control.current_io_port = -1;
+    input_control.remaining_count = fill_total;
+    input_control.eos_received = true;
+  }
+
   long MemfillXferDes::get_requests(Request **requests, long nr)
   {
     // unused

src/realm/transfer/channel.h

@@ -141,6 +141,9 @@ namespace Realm {
     //  modified during this call
     void swap(SequenceAssembler &other);
 
+    // NOT thread-safe.
+    void reset();
+
     // imports data from this assembler into another (this is thread-safe
     //  on the `other` but assumes no changes being made on `this`)
     void import(SequenceAssembler &other) const;
@@ -203,6 +206,9 @@ namespace Realm {
     Memory mem;
     RegionInstance inst;
     size_t ib_offset, ib_size;
+    // ib_index holds what index in the ib_offsets vector
+    // corresponds to the ib data for this port.
+    unsigned ib_index;
     TransferIterator *iter;
     CustomSerdezID serdez_id;
   };
@@ -297,6 +303,9 @@ namespace Realm {
     // current input and output port mask
     uint64_t current_in_port_mask, current_out_port_mask;
     uint64_t current_in_port_remain, current_out_port_remain;
+    // Stored to help reset XD state.
+    int gather_control_port = -1;
+    int scatter_control_port = -1;
     struct XferPort {
       MemoryImpl *mem;
       TransferIterator *iter;
@@ -314,6 +323,9 @@ namespace Realm {
       //  to complete)
       Memory ib_mem;
       size_t ib_offset, ib_size;
+      // Used when resetting an XD, records which ib_offset
+      // should be used.
+      unsigned ib_index;
       AddressList addrlist;
       AddressListCursor addrcursor;
     };
@@ -397,6 +409,14 @@ namespace Realm {
 
     void add_update_pre_bytes_total_received(void);
 
+    // Used to reset an XD so that it can be reused for a new copy.
+    // reset accepts a vector of ib_offsets that correspond to the
+    // ib's to be used after the XD is reset. Resetting an XD retains
+    // all pointers to the backing instance data, and it is up to the
+    // user to ensure that the data referenced by the XD is valid when
+    // the XD is used again after it is reset.
+    virtual void reset(const std::vector<off_t> &ib_offsets);
+
   protected:
     virtual ~XferDes();
 
@@ -550,7 +570,12 @@ namespace Realm {
     virtual Request *dequeue_request();
     virtual void enqueue_request(Request *req);
 
+    void reset(const std::vector<off_t> &ib_offsets);
+
     bool progress_xd(MemfillChannel *channel, TimeLimit work_until);
+
+  protected:
+    size_t fill_total;
   };
 
   class MemreduceChannel;

src/realm/transfer/lowlevel_dma.h

@@ -177,6 +177,7 @@ namespace Realm {
 
     virtual void reset(void);
     virtual bool done(void);
+    void set_base(size_t _base) { base = _base; }
 
     virtual size_t get_base_offset(void) const;
 

src/realm/transfer/transfer.cc

@@ -947,6 +947,10 @@ namespace Realm {
   {
     TransferIteratorBase<N, T>::reset();
     piece_idx = 0;
+    addrs_in->reset();
+    addrs_in_offset = 0;
+    point_pos = 0;
+    num_points = 0;
   }
 
   template <int N, typename T>
@@ -1192,6 +1196,8 @@ namespace Realm {
   {
     TransferIteratorBase<N, T>::reset();
     addrs_in->reset();
+    point_pos = 0;
+    num_points = 0;
   }
 
   template <int N, typename T>
@@ -4714,6 +4720,7 @@ namespace Realm {
           ii.inst = RegionInstance::NO_INST;
           ii.ib_offset = ib_offsets[xdn.inputs[j].edge];
           ii.ib_size = tg.ib_edges[xdn.inputs[j].edge].size;
+          ii.ib_index = xdn.inputs[j].edge;
           ii.iter = new WrappingFIFOIterator(ii.ib_offset, ii.ib_size);
           ii.serdez_id = 0;
           break;
@@ -4845,6 +4852,7 @@ namespace Realm {
           oi.inst = RegionInstance::NO_INST;
           oi.ib_offset = ib_offsets[xdn.outputs[j].edge];
           oi.ib_size = tg.ib_edges[xdn.outputs[j].edge].size;
+          oi.ib_index = xdn.outputs[j].edge;
           oi.iter = new WrappingFIFOIterator(oi.ib_offset, oi.ib_size);
           oi.serdez_id = 0;
           break;

src/realm/transfer/transfer.inl

@@ -76,7 +76,7 @@ namespace Realm {
   {
     return ((s << i.port_type) && (s << i.peer_guid) && (s << i.peer_port_idx) &&
             (s << i.indirect_port_idx) && (s << i.mem) && (s << i.inst) &&
-            (s << i.ib_offset) && (s << i.ib_size) && (s << *i.iter) &&
+            (s << i.ib_offset) && (s << i.ib_size) && (s << i.ib_index) && (s << *i.iter) &&
             (s << i.serdez_id));
   }
 
@@ -85,7 +85,7 @@ namespace Realm {
   {
     if(!((s >> i.port_type) && (s >> i.peer_guid) && (s >> i.peer_port_idx) &&
          (s >> i.indirect_port_idx) && (s >> i.mem) && (s >> i.inst) &&
-         (s >> i.ib_offset) && (s >> i.ib_size))) {
+         (s >> i.ib_offset) && (s >> i.ib_size) && (s >> i.ib_index))) {
       return false;
     }
     i.iter = TransferIterator::deserialize_new(s);