From 835aae2d22661c471d088ffbef3ac8b028ba79b3 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Wed, 25 Mar 2026 12:15:23 -0700 Subject: [PATCH] src/realm: allow for XferDes objects to be reset This commit ports the XferDes reset functionality from #391. --- src/realm/cuda/cuda_internal.cc | 9 +++ src/realm/cuda/cuda_internal.h | 3 + src/realm/transfer/address_list.cc | 17 ++++++ src/realm/transfer/address_list.h | 2 + src/realm/transfer/channel.cc | 98 +++++++++++++++++++++++++++++- src/realm/transfer/channel.h | 25 ++++++++ src/realm/transfer/lowlevel_dma.h | 1 + src/realm/transfer/transfer.cc | 8 +++ src/realm/transfer/transfer.inl | 4 +- 9 files changed, 163 insertions(+), 4 deletions(-) diff --git a/src/realm/cuda/cuda_internal.cc b/src/realm/cuda/cuda_internal.cc index 61f8e7ca5e3..03a3ff49ebb 100644 --- a/src/realm/cuda/cuda_internal.cc +++ b/src/realm/cuda/cuda_internal.cc @@ -2186,6 +2186,15 @@ namespace Realm { return 0; } + void GPUfillXferDes::reset(const std::vector &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 diff --git a/src/realm/cuda/cuda_internal.h b/src/realm/cuda/cuda_internal.h index dcc32a12fdb..5f9d3899881 100644 --- a/src/realm/cuda/cuda_internal.h +++ b/src/realm/cuda/cuda_internal.h @@ -918,12 +918,15 @@ namespace Realm { const std::vector &outputs_info, int _priority, const void *_fill_data, size_t _fill_size, size_t _fill_total); + void reset(const std::vector &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 { diff --git a/src/realm/transfer/address_list.cc b/src/realm/transfer/address_list.cc index e989f3469a6..23df9b33e4f 100644 --- a/src/realm/transfer/address_list.cc +++ b/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; diff --git a/src/realm/transfer/address_list.h b/src/realm/transfer/address_list.h index d702ccdf936..9c5f5e85ca2 100644 --- a/src/realm/transfer/address_list.h +++ b/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; diff --git a/src/realm/transfer/channel.cc b/src/realm/transfer/channel.cc index 9f2c609abf1..0968bb2ce09 100644 --- a/src/realm/transfer/channel.cc +++ b/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 &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(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(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 &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 diff --git a/src/realm/transfer/channel.h b/src/realm/transfer/channel.h index e693d0ff01a..3d7f5e51558 100644 --- a/src/realm/transfer/channel.h +++ b/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 &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 &ib_offsets); + bool progress_xd(MemfillChannel *channel, TimeLimit work_until); + + protected: + size_t fill_total; }; class MemreduceChannel; diff --git a/src/realm/transfer/lowlevel_dma.h b/src/realm/transfer/lowlevel_dma.h index 3209d130fce..2fa53e08e16 100644 --- a/src/realm/transfer/lowlevel_dma.h +++ b/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; diff --git a/src/realm/transfer/transfer.cc b/src/realm/transfer/transfer.cc index 54020a174fb..1276bf480a6 100644 --- a/src/realm/transfer/transfer.cc +++ b/src/realm/transfer/transfer.cc @@ -947,6 +947,10 @@ namespace Realm { { TransferIteratorBase::reset(); piece_idx = 0; + addrs_in->reset(); + addrs_in_offset = 0; + point_pos = 0; + num_points = 0; } template @@ -1192,6 +1196,8 @@ namespace Realm { { TransferIteratorBase::reset(); addrs_in->reset(); + point_pos = 0; + num_points = 0; } template @@ -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; diff --git a/src/realm/transfer/transfer.inl b/src/realm/transfer/transfer.inl index 7188130ebb6..62245492fe2 100644 --- a/src/realm/transfer/transfer.inl +++ b/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);