diff --git a/TOGSim/include/CoreTraceLog.h b/TOGSim/include/CoreTraceLog.h index 2ce51012..47a3e6f8 100644 --- a/TOGSim/include/CoreTraceLog.h +++ b/TOGSim/include/CoreTraceLog.h @@ -3,6 +3,8 @@ #include #include +#include + #include "Instruction.h" #include "TraceLogTags.h" @@ -12,6 +14,14 @@ */ namespace core_trace_log { +/** Whether the instruction trace is actually being emitted. + * + * The trace_* helpers below take formatted strings, so their arguments are + * built whether or not spdlog will keep them: every issued instruction paid two + * std::string constructions (fmt::format over its tile dims/strides) that + * spdlog::trace then dropped. Guard the call sites with this. */ +inline bool trace_enabled() { return spdlog::should_log(spdlog::level::trace); } + std::string format_dma_inst_issued_detail(Instruction& inst); /** Opcode + (detail...) for DMA issue / skip traces. */ std::string format_dma_inst_issued_trace_line(Instruction& inst); diff --git a/TOGSim/include/Instruction.h b/TOGSim/include/Instruction.h index 74c5e339..cdce3161 100644 --- a/TOGSim/include/Instruction.h +++ b/TOGSim/include/Instruction.h @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -32,6 +33,19 @@ typedef uint64_t cycle_type; std::string opcode_to_string(Opcode opcode); std::string format_tag_key_list_hex(const std::vector& tag_keys); +class Instruction; +// Order dependents by creation order (_global_inst_id), NOT by heap address. +// std::set>'s default comparator orders by pointer value, which +// only looked stable because the whole TileGraph was allocated up front in one +// monotonically growing run: fire() then released dependents in creation order +// by accident. Under any other heap layout the release order -- and with it the +// issue order and the reported cycle count -- changes. The instruction id is +// intrinsic, so ordering by it makes the release order canonical. +struct DepLess { + bool operator()(const std::shared_ptr& a, + const std::shared_ptr& b) const; +}; + class Instruction : public std::enable_shared_from_this { public: Instruction(Opcode opcode, cycle_type compute_cycle, size_t num_parents, addr_type dram_addr, @@ -51,7 +65,7 @@ class Instruction : public std::enable_shared_from_this { for (auto& c : _deps[static_cast(e)]) c->dec_ready_counter(); _deps[static_cast(e)].clear(); } - const std::set>& get_deps(DepEvent e) { + const std::set, DepLess>& get_deps(DepEvent e) { return _deps[static_cast(e)]; } void set_assigned_sa(int s) { _assigned_sa = s; } @@ -63,6 +77,12 @@ class Instruction : public std::enable_shared_from_this { void set_tile_group(int g) { _tile_group = g; } int get_tile_group() const { return _tile_group; } bool check_ready() { return ready_counter == 0; } + // Number of MATMULs subscribed to this preload's ISSUE event -- the weight-slot + // refcount. It is a property of the graph, fixed once the tile is built, but + // Core's issue scan asks for it on every cycle a preload stalls waiting for a + // free weight slot. Counting it by walking the dep set each time dominated the + // simulation (83% of a conv2d's run: one preload can have 300k subscribers). + int matmul_consumers(); const Opcode get_opcode() { return opcode; } bool is_dma_read() { return opcode == Opcode::MOVIN; } bool is_dma_write() { return opcode == Opcode::MOVOUT; } @@ -130,7 +150,16 @@ class Instruction : public std::enable_shared_from_this { // last reader frees it on issue via `_sram_release_allocs`. void set_sram_alloc(int64_t id) { _sram_alloc_id = id; } int64_t get_sram_alloc() const { return _sram_alloc_id; } - void add_sram_release(int64_t id) { _sram_release_allocs.push_back(id); } + // Keep the list ascending by version id. sram_finalize() walks the version map + // in id order, so an instruction that frees several versions currently gets + // them ascending; making the invariant explicit lets a caller tag versions as + // they close (in buffer order) without changing the graph. The Core does not + // care about the order -- it frees each version -- but keeping it makes two + // builders byte-comparable. + void add_sram_release(int64_t id) { + auto& v = _sram_release_allocs; + v.insert(std::upper_bound(v.begin(), v.end(), id), id); + } const std::vector& get_sram_release() const { return _sram_release_allocs; } // bytes this instruction's buffer occupies in the spad. A DMA derives it from // the tile it moves; a compute output gets it set explicitly by the bridge (the @@ -159,8 +188,10 @@ class Instruction : public std::enable_shared_from_this { size_t ready_counter = 0; // parents not yet finished; the minimal Instruction(Opcode) // ctor (barriers) relies on this default + inc_ready_counter // Per-event subscriber sets: _deps[ISSUE] released at issue (occupancy), - // _deps[DONE] at finish (latency). std::set dedups + fixes the release order. - std::array>, + // _deps[DONE] released at finish (latency). std::set dedups + iterates in + // instruction-id order (DepLess), so the release order does not depend on the + // allocator. + std::array, DepLess>, static_cast(DepEvent::COUNT)> _deps; std::vector tile_size; std::vector tile_stride; @@ -191,4 +222,5 @@ class Instruction : public std::enable_shared_from_this { int _assigned_sa = -1; std::shared_ptr _weight_token; int _tile_group = -1; // trace-only work-item id (see set_tile_group) + int _n_matmul_consumers = -1; // lazily counted; see matmul_consumers() }; \ No newline at end of file diff --git a/TOGSim/include/TileGraph.h b/TOGSim/include/TileGraph.h index 4cad9355..14c08907 100644 --- a/TOGSim/include/TileGraph.h +++ b/TOGSim/include/TileGraph.h @@ -4,6 +4,7 @@ #include #include #include +#include #include "Tile.h" #include "IntervalTree.h" @@ -39,7 +40,22 @@ class TileGraph { public: TileGraph(std::string path, std::string name) : _path(path), _name(name), _subgraph_vec(), _cpu_graph_map() {} void append_subgraph(std::shared_ptr subgraph); + + // --- on-demand construction (single-threaded) -------------------------- + // `_pull` materializes the NEXT dispatch tile and returns its subgraph, or + // nullptr when the producer is exhausted. It is called from allocate_subgraph, + // i.e. only when a core actually needs new work, so at most a couple of tiles + // are alive at once: peak memory becomes O(tiles in flight), not O(dispatches). + // Safe because a dispatch tile is dependency-closed -- the bridge resets its + // writers/seeds/tag maps at every tile boundary, so no edge crosses tiles + // (measured: cross_tile_edges == 0). + void set_tile_source(std::function()> pull) { + _pull = std::move(pull); + } + bool on_demand() const { return (bool)_pull; } + bool empty(int core_id) { + if (_pull && !_exhausted) return false; // more tiles can still be made if (_vec_index != _subgraph_vec.size()) { return false; } @@ -132,13 +148,15 @@ class TileGraph { private: int _vec_index=0; + std::function()> _pull; // empty -> legacy + bool _exhausted = false; + bool refill(); std::string _path; std::string _name = "?"; unsigned int _kernel_id = 0; std::vector _loop_index_list; std::vector> _ranges; std::vector> _subgraph_vec; - std::vector> _finished_subgraph_vec; std::map>> _cpu_graph_map; std::shared_ptr> _cache_plan; cycle_type _arrival_time; diff --git a/TOGSim/include/togsim_loader.h b/TOGSim/include/togsim_loader.h index fef63e54..63851d58 100644 --- a/TOGSim/include/togsim_loader.h +++ b/TOGSim/include/togsim_loader.h @@ -4,6 +4,7 @@ // "materializing sink" of sec 5.3 / 9.7; the stream goes to togsim_trace_bridge.h. #include +#include #include #include "togsim_runtime.h" @@ -48,4 +49,64 @@ RunResult run_producer(const char* so_path, const int64_t* cyc, const int64_t* ovl, int32_t n_tiles, const int32_t* partition_cores, int32_t n_partition_cores); +// Streaming variant: feeds each emitted record to `sink` and retains NOTHING. +// The whole recorded stream is O(#tiles) and, for a small systolic array, that is +// millions of records -- materializing it in a RunResult AND then building the +// TileGraph from it means both live at peak (measured: ~equal halves of peak RSS, +// SIGKILL on large 8x8 convs). The TileGraph builder therefore consumes records +// through a sink, one work-item at a time (see LazyProducer below). +using TraceSink = std::function; +// --------------------------------------------------------------------------- +// On-demand (lazy) production of one work-item at a time. +// +// `togsim_dispatch(ctx, fn, iv, n)` hands us the work-item's function pointer +// and its parallel induction variables, and the tile body reads nothing but +// `ctx` and `iv`. So a single indexing pass -- which records (fn, iv, core) and +// SKIPS the call -- is enough to invoke any work-item later, on its own, with no +// replay of the others. That makes the TileGraph buildable one dispatch tile at +// a time, single-threaded: peak memory is O(tiles in flight), not O(dispatches). +// +// Legal because a dispatch tile is dependency-closed: the bridge resets its +// writers/seeds/tag maps at every tile boundary, so no dependency edge crosses +// tiles (measured: cross_tile_edges == 0). +struct WorkItem { + void* fn = nullptr; // togsim_tile_fn + std::vector iv; // the enclosing parallel loop indices + int32_t core = 0; // round-robin binding, fixed at index time +}; + +class LazyProducer { + public: + LazyProducer() = default; + ~LazyProducer(); + LazyProducer(const LazyProducer&) = delete; + LazyProducer& operator=(const LazyProducer&) = delete; + + // dlopen the .so and run togsim_kernel once, in INDEX mode: every + // togsim_dispatch is recorded (so it can be re-invoked alone later) AND run, + // with every record streamed to `sink`. That single run doubles as the + // footprint pre-pass, so records emitted outside a dispatch are not lost. + bool open(const char* so_path, const int64_t* shape_args, int32_t n_shape, + const uint64_t* tensor_base, int32_t n_tensors, + const int64_t* cyc, const int64_t* ovl, int32_t n_tiles, + const int32_t* partition_cores, int32_t n_partition_cores, + const TraceSink* sink = nullptr); + + size_t num_items() const { return _items.size(); } + // Emit work-item `i`'s record stream (TILE_BEGIN, body, TILE_END) into `sink`. + void run_item(size_t i, const TraceSink& sink); + // Records the producer emitted OUTSIDE any dispatch. The TileGraph builder + // drops these (they belong to no work-item) exactly as the eager builder did; + // they still reach the footprint pre-pass through open()'s sink. + uint64_t stray_records() const { return _stray; } + + private: + struct EmitCtx* _ctx = nullptr; // opaque; owned + void* _lib = nullptr; + std::vector _items; + uint64_t _stray = 0; + std::vector _bases; + std::vector _cyc, _ovl; +}; + } // namespace togsim diff --git a/TOGSim/include/togsim_trace_bridge.h b/TOGSim/include/togsim_trace_bridge.h index 212cd4d6..661bdecc 100644 --- a/TOGSim/include/togsim_trace_bridge.h +++ b/TOGSim/include/togsim_trace_bridge.h @@ -7,6 +7,24 @@ #include "TileGraph.h" #include "togsim_loader.h" -// Build a TileGraph from a recorded trace. `name` labels the graph. -std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, - const std::string& name); +// Build a TileGraph straight from the trace producer `so_path`. +// +// The graph is built ON DEMAND, one togsim_dispatch work-item at a time: an +// indexing pass records each dispatch's (tile fn, induction vars, core) without +// running it, and a tile's Instructions are materialized only when a core asks +// for that work-item (togsim_loader.h LazyProducer). Peak memory is therefore +// O(tiles in flight) rather than O(dispatches) -- a large 8x8 conv2d went from +// 12.4 GiB to ~0.5 GiB, with a bit-identical dependency DAG and cycle count. +// +// This is sound because a dispatch tile is dependency-closed: the bridge resets +// its writers/seeds/tag maps and finalizes its SRAM versions at every tile +// boundary, so no dependency edge and no buffer version crosses tiles. +// +// Args mirror run_producer. `name` labels the graph. Returns nullptr if the +// producer run fails, or if it emits any record outside a dispatch. +std::unique_ptr trace_to_tilegraph( + const char* so_path, const int64_t* shape_args, int32_t n_shape, + const uint64_t* tensor_base, int32_t n_tensors, + const int64_t* cyc, const int64_t* ovl, int32_t n_tiles, + const int32_t* partition_cores, int32_t n_partition_cores, + const std::string& name); diff --git a/TOGSim/src/Core.cc b/TOGSim/src/Core.cc index 6097f287..768d0b5c 100644 --- a/TOGSim/src/Core.cc +++ b/TOGSim/src/Core.cc @@ -228,7 +228,7 @@ void Core::dma_cycle() { core_trace_log::log_error_dma_instruction_invalid(_core_cycle, _id); exit(EXIT_FAILURE); } else if (finished_inst->get_opcode() == Opcode::MEMORY_BAR) { - core_trace_log::trace_instruction_line(_core_cycle, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15(TraceLogTag::kInstructionFinished), finished_inst->get_global_inst_id(), @@ -307,7 +307,7 @@ void Core::cycle() { finish_instruction(inst); else _dma.register_tag_waiter(inst->subgraph_id, key, inst); - core_trace_log::trace_instruction_line(_core_cycle, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15( TraceLogTag::kInstructionSkipped), @@ -320,7 +320,7 @@ void Core::cycle() { } else { // load occupies its spad bytes on issue; stall (retry next cycle) if full. if (!try_occupy_sram(inst)) break; - core_trace_log::trace_instruction_line(_core_cycle, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15( TraceLogTag::kInstructionIssued), @@ -335,7 +335,7 @@ void Core::cycle() { } case Opcode::MOVOUT: release_sram(inst); // store issued -> free the tiles it drained - core_trace_log::trace_instruction_line(_core_cycle, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15(TraceLogTag::kInstructionIssued), inst->get_global_inst_id(), @@ -354,15 +354,21 @@ void Core::cycle() { int sa_idx = -1; if (ct == MATMUL || ct == PRELOAD) { if (ct == PRELOAD) { - int n_consumers = 0; // matmuls reusing this weight - for (auto& c : inst->get_deps(DepEvent::ISSUE)) - if (c->get_compute_type() == MATMUL) n_consumers++; + // Ask for the slot FIRST. The issue scan revisits every stalled + // preload on every cycle, and with no free slot none of them can + // issue, so nothing else about them is worth computing. (Ordering + // is safe: try_occupy_sram above is a no-op for a preload -- the + // virtual SA-weights buffer is untracked -- so nothing was + // reserved that we would have to give back here.) + sa_idx = pick_free_weight_sa(); + if (sa_idx < 0) break; // all weight slots full -> stall (retry) + // counted once and cached: the ISSUE dep set of one preload can + // hold hundreds of thousands of matmuls. + const int n_consumers = inst->matmul_consumers(); if (n_consumers == 0) { // weight-slot model needs >=1 consumer spdlog::error("preload has no matmul consumer (weight-slot model invariant)"); exit(EXIT_FAILURE); } - sa_idx = pick_free_weight_sa(); - if (sa_idx < 0) break; // all weight slots full -> stall (retry) _weight_slots_used[sa_idx]++; auto tok = std::make_shared(WeightToken{sa_idx, n_consumers}); for (auto& c : inst->get_deps(DepEvent::ISSUE)) @@ -411,7 +417,7 @@ void Core::cycle() { continue; // old code fell through to it++ on the // erased (invalidated) iterator -> UB } else { - core_trace_log::trace_instruction_line(_core_cycle, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15( TraceLogTag::kInstructionIssued), @@ -443,7 +449,7 @@ void Core::cycle() { } else { _dma.register_tag_waiter(inst->subgraph_id, key, inst); } - core_trace_log::trace_instruction_line(_core_cycle, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15( TraceLogTag::kInstructionIssued), @@ -498,7 +504,7 @@ void Core::finish_instruction(std::shared_ptr& inst, InstFinishTrac core_trace_log::log_error_dram_responses_trace_not_finished(_core_cycle, _id); exit(EXIT_FAILURE); } - core_trace_log::trace_instruction_line(_core_cycle, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15(TraceLogTag::kAllDramResponsesReceived), inst->get_global_inst_id(), @@ -515,7 +521,7 @@ void Core::finish_instruction(std::shared_ptr& inst, InstFinishTrac const char* trace_tag = (tag == InstFinishTraceTag::DmaIssueComplete) ? TraceLogTag::kAsyncDmaAllRequestsIssued : TraceLogTag::kInstructionFinished; - core_trace_log::trace_instruction_line(_core_cycle, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15(trace_tag), inst->get_global_inst_id(), @@ -562,7 +568,7 @@ void Core::push_memory_response(mem_fetch* response) { if (!owner_inst->got_first_response()) { // first data of this load arrived owner_inst->mark_first_response(); - core_trace_log::trace_instruction_line(_core_cycle, _id, + if (core_trace_log::trace_enabled()) core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15(TraceLogTag::kFirstDramResponse), owner_inst->get_global_inst_id(), core_trace_log::format_instruction_detail_line(*owner_inst)); diff --git a/TOGSim/src/DMA.cc b/TOGSim/src/DMA.cc index 5d509953..ab8e05f1 100644 --- a/TOGSim/src/DMA.cc +++ b/TOGSim/src/DMA.cc @@ -33,7 +33,7 @@ std::shared_ptr> DMA::get_memory_access(cycle_type core_ bool is_cacheable = owner_subgraph->is_cacheable(base_daddr, base_daddr + _dram_req_size); - if (_l2_datacache_enabled) { + if (_l2_datacache_enabled && spdlog::should_log(spdlog::level::trace)) { spdlog::trace( "[{}][Core {}][{}][INST_ID={}] dram=0x{:016x} cacheable={}", core_cycle, @@ -43,6 +43,7 @@ std::shared_ptr> DMA::get_memory_access(cycle_type core_ base_daddr, is_cacheable); } + if (spdlog::should_log(spdlog::level::trace)) spdlog::trace( "[{}][Core {}][{}][INST_ID={}] core_id={} subgraph_id={} numa_id={} addr_name={} is_write={}", core_cycle, diff --git a/TOGSim/src/Instruction.cc b/TOGSim/src/Instruction.cc index ee184a1a..60375763 100644 --- a/TOGSim/src/Instruction.cc +++ b/TOGSim/src/Instruction.cc @@ -1,5 +1,7 @@ #include "Instruction.h" +#include + #include uint64_t Instruction::_next_global_inst_id = 0; @@ -32,14 +34,16 @@ Instruction::Instruction(Opcode opcode, cycle_type compute_cycle, size_t num_par addr_type dram_addr, std::vector tile_size, std::vector tile_stride, size_t elem_bits, std::vector tag_idx_list, std::vector tag_stride_list, std::vector accum_tag_idx_list) + // The vectors are taken by value, so move them into the members: copying them + // allocated a second buffer per vector for every DMA and barrier instruction. : opcode(opcode), compute_cycle(compute_cycle), ready_counter(num_parents), dram_addr(dram_addr), - tile_size(tile_size), tile_stride(tile_stride), _elem_bits(elem_bits), - _tag_idx_list(tag_idx_list), _tag_stride_list(tag_stride_list), - _accum_tag_idx_list(accum_tag_idx_list) { + tile_size(std::move(tile_size)), tile_stride(std::move(tile_stride)), _elem_bits(elem_bits), + _tag_idx_list(std::move(tag_idx_list)), _tag_stride_list(std::move(tag_stride_list)), + _accum_tag_idx_list(std::move(accum_tag_idx_list)) { _global_inst_id = _next_global_inst_id++; assert(_tag_idx_list.size()==_tag_stride_list.size()); _tile_numel = 1; - for (auto dim : tile_size) + for (auto dim : this->tile_size) // the parameter was moved from _tile_numel *= dim; } @@ -49,6 +53,22 @@ Instruction::Instruction(Opcode opcode) _tile_numel = 1; } +bool DepLess::operator()(const std::shared_ptr& a, + const std::shared_ptr& b) const { + return a->get_global_inst_id() < b->get_global_inst_id(); +} + +int Instruction::matmul_consumers() { + if (_n_matmul_consumers < 0) { + constexpr int kMatmul = 1; // Core's MATMUL compute-unit enum + int n = 0; + for (auto& c : _deps[static_cast(DepEvent::ISSUE)]) + if (c->get_compute_type() == kMatmul) n++; + _n_matmul_consumers = n; + } + return _n_matmul_consumers; +} + void Instruction::finish_instruction() { fire(DepEvent::DONE); // latency consumers finished = true; @@ -66,6 +86,9 @@ void Instruction::dec_waiting_request() { void Instruction::prepare_tag_key() { /* Calculate tag key */ int64_t key_offset = 0; + // exact size: the two unconditional pushes otherwise grow 0 -> 1 -> 2, i.e. an + // allocate + reallocate + free for every DMA and barrier. + _tag_key.reserve(2 + _accum_tag_idx_list.size()); _tag_key.push_back(_addr_id); for (size_t i = 0; i < _tag_idx_list.size(); i++) key_offset += _tag_idx_list.at(i) * _tag_stride_list.at(i); diff --git a/TOGSim/src/TileGraph.cc b/TOGSim/src/TileGraph.cc index 120d49e2..b1732d37 100644 --- a/TOGSim/src/TileGraph.cc +++ b/TOGSim/src/TileGraph.cc @@ -51,8 +51,20 @@ void TileGraph::append_subgraph(std::shared_ptr subgraph) { _subgraph_vec.push_back(std::move(subgraph)); } +// Materialize the next dispatch tile, on demand. Returns false once the producer +// is exhausted. Called only when a core has no work left, so the builder never +// runs ahead of the simulation. +bool TileGraph::refill() { + if (!_pull || _exhausted) return false; + auto sg = _pull(); + if (!sg) { _exhausted = true; return false; } + sg->init_cache_plan(_cache_plan); + _subgraph_vec.push_back(std::move(sg)); + return true; +} + bool TileGraph::is_finished() { - bool finished = _subgraph_vec.empty(); + bool finished = _subgraph_vec.empty() && (!_pull || _exhausted); /* Check all outer loop is allocated */ if (!finished) return finished; @@ -104,18 +116,24 @@ std::shared_ptr TileGraph::get_tile(int core_id, int slot_id) { } void TileGraph::allocate_subgraph(int core_id, int slot_id) { - if (_cpu_graph_map[core_id][slot_id] != nullptr) { - _finished_subgraph_vec.push_back(_cpu_graph_map[core_id][slot_id]); + // Drop the finished subgraph instead of parking it in a keep-everything vector: + // this is what actually releases a tile's Instructions once the Core is done. + if (_cpu_graph_map[core_id][slot_id] != nullptr) _cpu_graph_map[core_id][slot_id] = nullptr; - } - for (auto it = _subgraph_vec.begin(); it != _subgraph_vec.end(); ++it) { - if ((*it)->get_core_id() == -1 || (*it)->get_core_id() == core_id) { - std::shared_ptr subgraph = *it; - _cpu_graph_map[core_id][slot_id] = subgraph; - _subgraph_vec.erase(it); - return; + // Build tiles only now that one is needed. Work-items round-robin over the + // partition's cores, so the next tile for THIS core may be a few dispatches + // ahead; keep pulling until we find one (or the producer runs dry). With a + // single core this pulls exactly one. + for (;;) { + for (auto it = _subgraph_vec.begin(); it != _subgraph_vec.end(); ++it) { + if ((*it)->get_core_id() == -1 || (*it)->get_core_id() == core_id) { + std::shared_ptr subgraph = *it; + _cpu_graph_map[core_id][slot_id] = subgraph; + _subgraph_vec.erase(it); + return; + } } + if (!refill()) return; // exhausted: leave the slot empty } - return; } \ No newline at end of file diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 4ccaa4ae..0ef98eff 100644 --- a/TOGSim/src/main.cc +++ b/TOGSim/src/main.cc @@ -40,12 +40,11 @@ std::unique_ptr build_trace_tilegraph(Simulator* simulator, while (ct >> c >> o) { cyc.push_back(c); ovl.push_back(o); } } if (cyc.empty()) { cyc.assign(256, 128); ovl.assign(256, 0); } - auto run = togsim::run_producer(trace_so_path.c_str(), nullptr, 0, - bases.data(), (int)bases.size(), - cyc.data(), ovl.data(), (int)cyc.size(), - partition_cores.data(), (int32_t)partition_cores.size()); - if (!run.ok) return nullptr; - return trace_to_tilegraph(run, "trace_kernel"); + return trace_to_tilegraph(trace_so_path.c_str(), nullptr, 0, + bases.data(), (int)bases.size(), + cyc.data(), ovl.data(), (int)cyc.size(), + partition_cores.data(), (int32_t)partition_cores.size(), + "trace_kernel"); } void launchKernel(Simulator* simulator, unsigned int kernel_id, std::string onnx_path, std::string attribute_path, const YAML::Node& config_yaml, cycle_type request_time=0, int partition_id=0, int device_id=0) { diff --git a/TOGSim/src/togsim_runtime.cc b/TOGSim/src/togsim_runtime.cc index 1a76231b..57596880 100644 --- a/TOGSim/src/togsim_runtime.cc +++ b/TOGSim/src/togsim_runtime.cc @@ -23,16 +23,41 @@ struct EmitCtx { // mutable run state int32_t rr = 0; // round-robin cursor into `cores` int32_t cur_core = -1; // current work-item's core + // INDEX mode (LazyProducer::open): togsim_dispatch additionally records the + // work-item so it can be re-invoked on its own later. Records emitted outside + // any dispatch (depth == 0) are counted; the TileGraph builder drops them, so + // only the footprint pre-pass ever sees them. + bool capture = false; + std::vector* items = nullptr; + int depth = 0; + uint64_t stray = 0; + // Scratch record reused by every callback. The sink consumes a record before + // the next one is emitted (and run_producer copies it), so one buffer is + // enough -- and its vectors keep their capacity, which is the difference + // between four allocations per emitted record and none. + togsim::TraceRec scratch; + // Exactly one of these is active. With a sink the record is consumed and dropped + // (streaming, O(1) memory); without one it is appended (legacy run_producer). + const togsim::TraceSink* sink = nullptr; std::vector trace; }; namespace { -inline togsim::TraceRec blank(togsim::TraceRec::Kind k, int32_t core) { - togsim::TraceRec r{}; - r.kind = k; - r.core = core; +// Reset `ctx->scratch` for a new record. clear() keeps each vector's capacity. +inline togsim::TraceRec& blank(EmitCtx* ctx, togsim::TraceRec::Kind k, int32_t core) { + togsim::TraceRec& r = ctx->scratch; + r.dims.clear(); r.strides.clear(); r.read_bufs.clear(); r.write_bufs.clear(); + r.kind = k; r.core = core; + r.dir = 0; r.arg_id = 0; r.elem_bits = 0; r.is_async = 0; + r.addr = 0; r.tag_id = 0; r.tag_slot = 0; + r.tile_id = 0; r.compute_type = 0; r.cycle = 0; r.overlapping = 0; return r; } +inline void emit_rec(EmitCtx* ctx, const togsim::TraceRec& r) { + if (ctx->depth == 0) ++ctx->stray; // emitted outside any dispatch + if (ctx->sink) (*ctx->sink)(r); + else ctx->trace.push_back(r); +} } // namespace extern "C" { @@ -45,9 +70,18 @@ void togsim_dispatch(EmitCtx* ctx, togsim_tile_fn fn, int64_t* iv, int32_t n_iv) // forever. TILE_BEGIN/TILE_END bracket the ops `fn` emits under ctx->cur_core. ctx->cur_core = ctx->cores.empty() ? 0 : ctx->cores[ctx->rr++ % (int32_t)ctx->cores.size()]; - ctx->trace.push_back(blank(togsim::TraceRec::TILE_BEGIN, ctx->cur_core)); + if (ctx->capture) { // index pass: remember the work-item so it can be re-run alone + togsim::WorkItem w; + w.fn = (void*)fn; + w.core = ctx->cur_core; + if (iv && n_iv > 0) w.iv.assign(iv, iv + n_iv); + ctx->items->push_back(std::move(w)); + } + ++ctx->depth; + emit_rec(ctx, blank(ctx, togsim::TraceRec::TILE_BEGIN, ctx->cur_core)); fn(ctx, iv, n_iv); - ctx->trace.push_back(blank(togsim::TraceRec::TILE_END, ctx->cur_core)); + emit_rec(ctx, blank(ctx, togsim::TraceRec::TILE_END, ctx->cur_core)); + --ctx->depth; } void togsim_dma(EmitCtx* ctx, int32_t dir, int32_t arg_id, @@ -59,16 +93,20 @@ void togsim_dma(EmitCtx* ctx, int32_t dir, int32_t arg_id, uint64_t base = (arg_id >= 0 && arg_id < ctx->n_tensors) ? ctx->tensor_base[arg_id] : 0; uint64_t addr = base + offset * (uint64_t)(elem_bits / 8); - togsim::TraceRec r = blank(togsim::TraceRec::DMA, ctx->cur_core); + togsim::TraceRec& r = blank(ctx, togsim::TraceRec::DMA, ctx->cur_core); r.dir = dir; r.arg_id = arg_id; r.elem_bits = elem_bits; r.is_async = is_async; r.addr = addr; r.tag_id = tag_id; r.tag_slot = tag_slot; + if (dims) r.dims.reserve(ndim); + if (strides) r.strides.reserve(ndim); + r.read_bufs.reserve(n_read); + r.write_bufs.reserve(n_write); for (int32_t i = 0; i < ndim; ++i) { if (dims) r.dims.push_back(dims[i]); if (strides) r.strides.push_back(strides[i]); } for (int32_t i = 0; i < n_read; ++i) r.read_bufs.push_back(read_bufs[i]); for (int32_t i = 0; i < n_write; ++i) r.write_bufs.push_back(write_bufs[i]); - ctx->trace.push_back(r); + emit_rec(ctx, r); } void togsim_compute(EmitCtx* ctx, uint64_t tile_id, int32_t compute_type, @@ -76,49 +114,115 @@ void togsim_compute(EmitCtx* ctx, uint64_t tile_id, int32_t compute_type, const int64_t* read_bufs, int32_t n_read, const int64_t* write_bufs, int32_t n_write) { (void)ndim; (void)dims; - togsim::TraceRec r = blank(togsim::TraceRec::COMPUTE, ctx->cur_core); + togsim::TraceRec& r = blank(ctx, togsim::TraceRec::COMPUTE, ctx->cur_core); r.tile_id = tile_id; r.compute_type = compute_type; + r.read_bufs.reserve(n_read); + r.write_bufs.reserve(n_write); for (int32_t i = 0; i < n_read; ++i) r.read_bufs.push_back(read_bufs[i]); for (int32_t i = 0; i < n_write; ++i) r.write_bufs.push_back(write_bufs[i]); if (ctx->cyc && (int32_t)tile_id < ctx->n_tiles) r.cycle = ctx->cyc[tile_id]; if (ctx->ovl && (int32_t)tile_id < ctx->n_tiles) r.overlapping = ctx->ovl[tile_id]; - ctx->trace.push_back(r); + emit_rec(ctx, r); } void togsim_memory_barrier(EmitCtx* ctx, int32_t tag_id, uint64_t tag_slot, const int64_t* write_bufs, int32_t n_write) { - togsim::TraceRec r = blank(togsim::TraceRec::MEMORY_BAR, ctx->cur_core); + togsim::TraceRec& r = blank(ctx, togsim::TraceRec::MEMORY_BAR, ctx->cur_core); r.tag_id = tag_id; r.tag_slot = tag_slot; + r.write_bufs.reserve(n_write); for (int32_t i = 0; i < n_write; ++i) r.write_bufs.push_back(write_bufs[i]); - ctx->trace.push_back(r); + emit_rec(ctx, r); } } // extern "C" namespace togsim { +namespace { +void init_ctx(EmitCtx& ctx, const uint64_t* tensor_base, int32_t n_tensors, + const int64_t* cyc, const int64_t* ovl, int32_t n_tiles, + const int32_t* partition_cores, int32_t n_partition_cores) { + ctx.tensor_base = tensor_base; ctx.n_tensors = n_tensors; + ctx.cyc = cyc; ctx.ovl = ovl; ctx.n_tiles = n_tiles; + ctx.cores.assign(partition_cores, partition_cores + (n_partition_cores > 0 ? n_partition_cores : 0)); + if (ctx.cores.empty()) ctx.cores.push_back(0); +} +// dlopen the producer and run its togsim_kernel against `ctx`. +bool load_and_run(const char* so_path, const int64_t* shape_args, int32_t n_shape, + EmitCtx& ctx) { + void* lib = dlopen(so_path, RTLD_NOW | RTLD_GLOBAL); + if (!lib) { fprintf(stderr, "togsim: dlopen failed: %s\n", dlerror()); return false; } + auto emit = (void (*)(EmitCtx*, int64_t*, int32_t))dlsym(lib, "togsim_kernel"); + if (!emit) { fprintf(stderr, "togsim: dlsym togsim_kernel failed: %s\n", dlerror()); return false; } + emit(&ctx, (int64_t*)shape_args, n_shape); + return true; +} +} // namespace + RunResult run_producer(const char* so_path, const int64_t* shape_args, int32_t n_shape, const uint64_t* tensor_base, int32_t n_tensors, const int64_t* cyc, const int64_t* ovl, int32_t n_tiles, const int32_t* partition_cores, int32_t n_partition_cores) { RunResult res; - void* lib = dlopen(so_path, RTLD_NOW | RTLD_GLOBAL); - if (!lib) { fprintf(stderr, "togsim: dlopen failed: %s\n", dlerror()); return res; } - auto emit = (void (*)(EmitCtx*, int64_t*, int32_t))dlsym(lib, "togsim_kernel"); - if (!emit) { fprintf(stderr, "togsim: dlsym togsim_kernel failed: %s\n", dlerror()); return res; } - EmitCtx ctx; - ctx.tensor_base = tensor_base; ctx.n_tensors = n_tensors; - ctx.cyc = cyc; ctx.ovl = ovl; ctx.n_tiles = n_tiles; - ctx.cores.assign(partition_cores, partition_cores + (n_partition_cores > 0 ? n_partition_cores : 0)); - if (ctx.cores.empty()) ctx.cores.push_back(0); - emit(&ctx, (int64_t*)shape_args, n_shape); - + init_ctx(ctx, tensor_base, n_tensors, cyc, ovl, n_tiles, partition_cores, n_partition_cores); + if (!load_and_run(so_path, shape_args, n_shape, ctx)) return res; res.ok = true; res.trace = std::move(ctx.trace); return res; } +LazyProducer::~LazyProducer() { + delete _ctx; // _lib intentionally left open: the tile fns must stay callable +} + +bool LazyProducer::open(const char* so_path, const int64_t* shape_args, int32_t n_shape, + const uint64_t* tensor_base, int32_t n_tensors, + const int64_t* cyc, const int64_t* ovl, int32_t n_tiles, + const int32_t* partition_cores, int32_t n_partition_cores, + const TraceSink* sink) { + // own the tables: the caller's buffers are its locals, and the tile fns are + // invoked long after it returns. + if (tensor_base && n_tensors > 0) _bases.assign(tensor_base, tensor_base + n_tensors); + if (cyc && n_tiles > 0) _cyc.assign(cyc, cyc + n_tiles); + if (ovl && n_tiles > 0) _ovl.assign(ovl, ovl + n_tiles); + + _ctx = new EmitCtx(); + init_ctx(*_ctx, _bases.empty() ? nullptr : _bases.data(), (int32_t)_bases.size(), + _cyc.empty() ? nullptr : _cyc.data(), _ovl.empty() ? nullptr : _ovl.data(), + (int32_t)_cyc.size(), partition_cores, n_partition_cores); + + _lib = dlopen(so_path, RTLD_NOW | RTLD_GLOBAL); + if (!_lib) { fprintf(stderr, "togsim: dlopen failed: %s\n", dlerror()); return false; } + auto kernel = (void (*)(EmitCtx*, int64_t*, int32_t))dlsym(_lib, "togsim_kernel"); + if (!kernel) { fprintf(stderr, "togsim: dlsym togsim_kernel failed: %s\n", dlerror()); return false; } + + // One full producer run: it both indexes the dispatches and streams every + // record (including any emitted outside a dispatch) to `sink`. This is the + // footprint pre-pass the eager builder used to run, so nothing is missed. + _ctx->capture = true; + _ctx->items = &_items; + _ctx->sink = sink; + kernel(_ctx, (int64_t*)shape_args, n_shape); + _ctx->capture = false; + _ctx->items = nullptr; + _ctx->sink = nullptr; + _stray = _ctx->stray; + return true; +} + +void LazyProducer::run_item(size_t i, const TraceSink& sink) { + if (i >= _items.size()) return; + WorkItem& w = _items[i]; + _ctx->sink = &sink; + _ctx->cur_core = w.core; // the binding fixed at index time + _ctx->depth = 1; + emit_rec(_ctx, blank(_ctx, TraceRec::TILE_BEGIN, w.core)); + ((togsim_tile_fn)w.fn)(_ctx, w.iv.empty() ? nullptr : w.iv.data(), (int32_t)w.iv.size()); + emit_rec(_ctx, blank(_ctx, TraceRec::TILE_END, w.core)); + _ctx->depth = 0; + _ctx->sink = nullptr; +} } // namespace togsim diff --git a/TOGSim/src/togsim_trace_bridge.cc b/TOGSim/src/togsim_trace_bridge.cc index 46e6f7b3..ca6639be 100644 --- a/TOGSim/src/togsim_trace_bridge.cc +++ b/TOGSim/src/togsim_trace_bridge.cc @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -58,48 +59,113 @@ std::shared_ptr make_compute(const togsim::TraceRec& t) { return inst; } -} // namespace -std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, - const std::string& name) { - using togsim::TraceRec; - auto tg = std::make_unique(name, name); - // Empty cache plan (no L2/CMEM persistence) -- append_subgraph propagates it - // to each subgraph, and DMA::is_cacheable dereferences it, so it must be a - // valid (if empty) IntervalTree rather than null. - tg->init_cache_plan({}); +// All builder state for one kernel. Heap-allocated and owned by the TileGraph's +// tile source, so it survives between on-demand materializations of a tile. +struct BuildState { + togsim::LazyProducer prod; + size_t next = 0; // next work-item to materialize + std::map buf_bytes; + std::shared_ptr sg_out; // the tile just built std::shared_ptr sg; std::shared_ptr tile; - // The explicit dependency DAG (sec 10): per SRAM buffer, writers(b) is the SET of - // current producers' DONE-handles and readers(b) its consumers. Scoped per - // work-item, so distinct work-items are independent (-> parallel). - std::map>> writers; // buffer id -> current producers (DONE-handles) - // 1 load : N barriers, so track the CURRENT load per (tag_id, tag_slot), not a - // FIFO. Each load takes a fresh `uniq` and its iteration's barriers reuse it. - // Correct only because a load nest and its consumers run in order. Per work-item. + std::map>> writers; + std::map>> seeds; std::map, std::pair>> current_dma; - // Dedup barriers on the CURRENT load of a (tag_id, tag_slot): a conv reads one - // loaded subtile from many matmuls, so its per-consumer waits collapse to one bar. - // A new load bumps uniq, so a genuine new wait still gets its own bar. std::map, std::pair>> bar_for_load; - int64_t next_tag = 0; // mints a unique Core tag key per dma record - int cur_tile_group = -1; // work-item index, bumped per TILE_BEGIN (trace grouping) - std::set cur_tile_bufs; // distinct spad buffers this tile touches - size_t cur_tile_footprint = 0; // their footprint sum = the tile's resident set + int64_t next_tag = 0; + int cur_tile_group = -1; + std::set cur_tile_bufs; + size_t cur_tile_footprint = 0; + + // SRAM buffer versions. A version is NOT scoped to a work-item -- the spad is + // one physical resource, so a buffer reused by the next tile is a new version + // that must wait for the old one to free. The version *schedule* is therefore + // precomputed by sram_schedule(), which allocates nothing: the builder can tag + // a version's last reader as it goes instead of retaining every reader until + // end-of-stream (which is what forced the whole graph to stay in memory). + int64_t next_alloc = 0; + std::map cur_alloc; // buf -> current version id + std::map open_ver; // buf -> version still accepting writes + std::vector has_readers; // version -> ever read? + std::vector> last_reader; // version -> (work-item, record) + size_t item = 0, rec = 0; // position of the record being fed +}; + +// Walk the record stream once, allocating nothing, and work out for each SRAM +// buffer version (a) whether anything ever reads it and (b) where its LAST +// reader sits. Mirrors sram_on_load / sram_on_write / sram_on_read exactly, so +// the version ids it mints are the ones the builder will mint. +void sram_schedule(BuildState& S) { + using togsim::TraceRec; + int64_t next_alloc = 0; + std::map cur_alloc; + std::map open_ver; + size_t item = 0, rec = 0, pos = 0; + auto on_open = [&](int64_t b) { + if (!cur_alloc.count(b) || !open_ver[b]) { + cur_alloc[b] = next_alloc++; + open_ver[b] = true; + S.has_readers.push_back(0); + S.last_reader.emplace_back((size_t)-1, (size_t)-1); + } + }; + auto on_read = [&](int64_t b) { + auto it = cur_alloc.find(b); + if (it == cur_alloc.end()) return; + S.has_readers[it->second] = 1; + S.last_reader[it->second] = {item, pos}; + open_ver[b] = false; + }; + auto in = [](const std::vector& v, int64_t b) { + return std::find(v.begin(), v.end(), b) != v.end(); + }; + togsim::TraceSink sink = [&](const TraceRec& t) { + pos = rec++; + if (t.kind == TraceRec::DMA) { + if (t.dir == 1) { for (int64_t b : t.read_bufs) on_read(b); } // store drains the spad + else { for (int64_t b : t.write_bufs) on_open(b); } // load fills it + } else if (t.kind == TraceRec::COMPUTE) { + for (int64_t b : t.read_bufs) if (!in(t.write_bufs, b)) on_read(b); + for (int64_t b : t.write_bufs) if (!in(t.read_bufs, b) && S.buf_bytes.count(b)) on_open(b); + } + }; + for (item = 0; item < S.prod.num_items(); item++) { rec = 0; S.prod.run_item(item, sink); } +} + +// Materialize exactly one dispatch tile (work-item `S.next`) and return its +// subgraph; nullptr once the producer is exhausted. +std::shared_ptr build_one_tile(BuildState& S) { + using togsim::TraceRec; + if (S.next >= S.prod.num_items()) return nullptr; + auto& sg = S.sg; auto& tile = S.tile; + auto& writers = S.writers; auto& seeds = S.seeds; + auto& current_dma = S.current_dma; auto& bar_for_load = S.bar_for_load; + auto& next_tag = S.next_tag; auto& cur_tile_group = S.cur_tile_group; + auto& cur_tile_bufs = S.cur_tile_bufs; auto& cur_tile_footprint = S.cur_tile_footprint; + auto& next_alloc = S.next_alloc; auto& cur_alloc = S.cur_alloc; + auto& open_ver = S.open_ver; + auto& buf_bytes = S.buf_bytes; + auto rec_bytes = [](const TraceRec& t) { + size_t numel = 1; + for (auto d : t.dims) numel *= (size_t)d; + return numel * (t.elem_bits / 8); + }; auto flush = [&]() { if (sg && tile) { tile->set_spad_footprint(cur_tile_footprint); // distinct-buffer resident set (1- vs 2-dispatch) sg->add_tile(tile); tile->set_owner(sg); - tg->append_subgraph(sg); + S.sg_out = sg; // hand this tile to the consumer } sg.reset(); tile.reset(); writers.clear(); + seeds.clear(); current_dma.clear(); bar_for_load.clear(); cur_tile_bufs.clear(); @@ -117,6 +183,13 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, for (int64_t w : writes) if (w == b) return true; return false; }; + // REPLACE writers(b) and recompute its seed set in one place. + auto set_writer = [&](int64_t b, const std::shared_ptr& w) { + writers[b] = { w }; + auto& sd = seeds[b]; + sd.clear(); + if (w->get_compute_type() != MATMUL_CT) sd.push_back(w); + }; auto link = [&](std::shared_ptr inst, const std::vector& reads, const std::vector& writes) { @@ -135,40 +208,27 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, } for (int64_t b : writes) { if (is_mm_accum(inst, b, writes)) { // UNION (commutative accumulate) - auto it = writers.find(b); - if (it != writers.end()) + auto it = seeds.find(b); + if (it != seeds.end()) for (auto& s : it->second) - if (s->get_compute_type() != MATMUL_CT) - s->add_dep(inst, DepEvent::DONE); // wait the init/bias seed only + s->add_dep(inst, DepEvent::DONE); // wait the init/bias seed only writers[b].push_back(inst); // join; no reset, no co-matmul edge } else { // REPLACE (normal output; resets the producer set) - writers[b] = { inst }; + set_writer(b, inst); } } tile->append_instuction(inst); }; - // SRAM-capacity tracking (sec 10.4): a coarse tile is one version of its buffer, - // freed once all its consumers have issued. NOT reset in flush() -- the spad is a - // physical per-core resource. v1 is single-core; multi-core would key by (core, buf). - int64_t next_alloc = 0; - std::map cur_alloc; // buf -> current version id - std::map open_ver; // buf -> version still accepting writes - struct Ver { std::vector> loads, readers; }; - std::map vers; - // Spad bytes per buffer id, from the DMA records that touch it (a compute output - // takes its footprint from its store). Pre-pass, so it is known before the compute. - auto rec_bytes = [](const TraceRec& t) { // single source of the tile footprint - size_t numel = 1; - for (auto d : t.dims) numel *= (size_t)d; - return numel * (t.elem_bits / 8); - }; - std::map buf_bytes; - for (const auto& t : run.trace) { - if (t.kind != TraceRec::DMA) continue; - const auto& bs = (t.dir == 1) ? t.read_bufs : t.write_bufs; // store reads spad, load writes spad - for (int64_t b : bs) buf_bytes[b] = rec_bytes(t); - } + // --- SRAM-capacity tracking (buffer-version allocations, sec 10.4) --- + // A coarse tile = one version of its buffer; the fine DMAs that fill it share + // one allocation, freed once all the version's consumers have issued (refcount + // -> 0). NOT reset in flush(): the spad is one physical per-core resource, so a + // buffer reused by the next reduction iter / work-item is a NEW version that + // must wait for the old one to free (WAR / double-buffer). Both DMA-loaded + // buffers AND compute outputs (the accumulator, vector epilogue results) are + // tracked; the virtual SA-weights are not (weight slots model them). (v1: + // single-core; multi-core would key cur_alloc/vers by (core, buf).) // Add each buffer once to the current tile's footprint (reloads in a K-loop reuse the same id). auto note_bufs = [&](const std::vector& bufs) { for (int64_t b : bufs) @@ -177,14 +237,14 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, if (it != buf_bytes.end()) cur_tile_footprint += it->second; } }; + // a version nothing ever reads is never freed -> leave it untracked + auto tracked = [&](int64_t a) { return S.has_readers[a] ? a : (int64_t)-1; }; auto sram_on_load = [&](int64_t b, const std::shared_ptr& ld) { if (!cur_alloc.count(b) || !open_ver[b]) { // a read closed it -> new version cur_alloc[b] = next_alloc++; open_ver[b] = true; - vers[cur_alloc[b]] = {}; } - ld->set_sram_alloc(cur_alloc[b]); - vers[cur_alloc[b]].loads.push_back(ld); + ld->set_sram_alloc(tracked(cur_alloc[b])); }; // A compute that freshly produces b opens a version like a load, carrying b's // footprint; its last reader frees it -- identical lifecycle to a load. @@ -194,31 +254,23 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, if (!cur_alloc.count(b) || !open_ver[b]) { // a consuming read closed it -> new version cur_alloc[b] = next_alloc++; open_ver[b] = true; - vers[cur_alloc[b]] = {}; - w->set_sram_alloc(cur_alloc[b]); + w->set_sram_alloc(tracked(cur_alloc[b])); w->set_sram_footprint(bb->second); - vers[cur_alloc[b]].loads.push_back(w); } // already-open version (further producing writes): same physical bytes, no re-add. }; auto sram_on_read = [&](int64_t b, const std::shared_ptr& rd) { auto it = cur_alloc.find(b); if (it == cur_alloc.end()) return; // not a load buffer -> untracked - vers[it->second].readers.push_back(rd); + // The version's LAST reader frees it. sram_schedule() already found where + // that reader sits, so tag it now rather than retaining every reader. + if (S.last_reader[it->second] == std::make_pair(S.item, S.rec)) + rd->add_sram_release(it->second); open_ver[b] = false; // next write starts a new version }; - auto sram_finalize = [&]() { // tag only each version's LAST reader - for (auto& kv : vers) { - auto& v = kv.second; - if (v.readers.empty()) { // no consumer -> never freed: untrack - for (auto& ld : v.loads) ld->set_sram_alloc(-1); - continue; - } - v.readers.back()->add_sram_release(kv.first); // it frees the whole version on issue - } - }; - for (const auto& t : run.trace) { + auto feed = [&](const TraceRec& t) { // build, one record at a time + struct RecTick { size_t& r; ~RecTick() { ++r; } } tick{S.rec}; if (t.kind == TraceRec::TILE_BEGIN) { // togsim_dispatch opened a work-item -> new subgraph (bound to its core) + // tile. The scope runs until the matching TILE_END (the dispatch wrapper @@ -228,13 +280,13 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, sg->set_core_id(t.core); tile = std::make_shared(Tile::Status::INITIALIZED); cur_tile_group++; - continue; + return; } if (t.kind == TraceRec::TILE_END) { flush(); // close the work-item explicitly (scope = the tile fn call) - continue; + return; } - if (!tile) continue; // defensive: ops before the first TILE_BEGIN + if (!tile) return; // defensive: ops before the first TILE_BEGIN if (t.kind == TraceRec::DMA) { int64_t uniq = next_tag++; // fresh Core tag key per dma record @@ -253,10 +305,14 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, // so consumers gate on arrival. A sync load blocks to arrival itself. if (t.is_async) current_dma[{t.tag_id, t.tag_slot}] = {uniq, inst}; for (int64_t b : t.write_bufs) { - // No hard WAR edge: load-buffer reuse is modeled by the SRAM version/ - // capacity machinery (sec 10.4); an edge would force single-buffering. The - // accumulator is not a load buffer -- link()'s REPLACE branch handles it. - writers[b] = { inst }; + // No hard WAR edge here: load-buffer reuse (double-buffering, X_spad/ + // W_spad reloaded each reduction iter) is modeled by the SRAM + // version/capacity machinery (sram_on_load), which sizes how many + // versions physically coexist. A latency WAR edge would force + // single-buffering and kill the overlap the spad permits. (The + // accumulator Y is NOT a load buffer -> its cross-tile WAR is handled by + // the REPLACE branch of link() when the next tile's init overwrites it.) + set_writer(b, inst); sram_on_load(b, inst); // occupy spad } } @@ -272,8 +328,8 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, // so the buffer's consumers gate on it, instead of emitting a redundant barrier. auto bf = bar_for_load.find({t.tag_id, t.tag_slot}); if (bf != bar_for_load.end() && bf->second.first == uniq) { - for (int64_t b : t.write_bufs) writers[b] = { bf->second.second }; - continue; + for (int64_t b : t.write_bufs) set_writer(b, bf->second.second); + return; } auto bar = make_mem_bar(t, uniq); bar->set_tile_group(cur_tile_group); @@ -281,7 +337,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, tile->append_instuction(bar); // the bar is the load's DONE-handle: REPLACE writers(b) with it (no WAR -- the // load already WAR'd the prior readers when it wrote). - for (int64_t b : t.write_bufs) writers[b] = { bar }; + for (int64_t b : t.write_bufs) set_writer(b, bar); bar_for_load[{t.tag_id, t.tag_slot}] = {uniq, bar}; } else if (t.kind == TraceRec::COMPUTE) { auto inst = make_compute(t); @@ -297,8 +353,53 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, for (int64_t b : t.read_bufs) if (!in(t.write_bufs, b)) sram_on_read(b, inst); // consuming reads for (int64_t b : t.write_bufs) if (!in(t.read_bufs, b)) sram_on_write(b, inst); // fresh outputs } - } - flush(); - sram_finalize(); // readers per version are now final -> set each version's refcount + }; + + S.item = S.next; S.rec = 0; + S.prod.run_item(S.next++, feed); + auto out = std::move(S.sg_out); + S.sg_out.reset(); + return out; +} + +} // namespace + +std::unique_ptr trace_to_tilegraph( + const char* so_path, const int64_t* shape_args, int32_t n_shape, + const uint64_t* tensor_base, int32_t n_tensors, + const int64_t* cyc, const int64_t* ovl, int32_t n_tiles, + const int32_t* partition_cores, int32_t n_partition_cores, + const std::string& name) { + using togsim::TraceRec; + auto S = std::make_shared(); + + // Indexing pass, which doubles as the spad-footprint pre-pass: one full + // producer run that records each togsim_dispatch (tile fn, induction vars, + // core) so the work-item can be re-invoked on its own later, while streaming + // every record -- including any emitted outside a dispatch, which belong to no + // work-item and never enter the graph -- into the footprint sink. It builds no + // Instruction. + auto rec_bytes = [](const TraceRec& t) { // single source of the tile footprint + size_t numel = 1; + for (auto d : t.dims) numel *= (size_t)d; + return numel * (t.elem_bits / 8); + }; + togsim::TraceSink footprint = [&](const TraceRec& t) { + if (t.kind != TraceRec::DMA) return; + const auto& bs = (t.dir == 1) ? t.read_bufs : t.write_bufs; // store reads spad, load writes spad + for (int64_t b : bs) S->buf_bytes[b] = rec_bytes(t); + }; + if (!S->prod.open(so_path, shape_args, n_shape, tensor_base, n_tensors, + cyc, ovl, n_tiles, partition_cores, n_partition_cores, &footprint)) + return nullptr; + + sram_schedule(*S); // buffer-version lifetimes, materializing no Instruction + + auto tg = std::make_unique(name, name); + // Empty cache plan (no L2/CMEM persistence) -- the tile source propagates it + // to each subgraph, and DMA::is_cacheable dereferences it, so it must be a + // valid (if empty) IntervalTree rather than null. + tg->init_cache_plan({}); + tg->set_tile_source([S]() { return build_one_tile(*S); }); return tg; }