[TOGSim] Build the TileGraph on demand, one dispatch tile at a time

TOGSim/include/CoreTraceLog.h

@@ -3,6 +3,8 @@
 #include <cstdint>
 #include <string>
 
+#include <spdlog/spdlog.h>
+
 #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);

TOGSim/include/Instruction.h

@@ -6,6 +6,7 @@
 #include <list>
 #include <numeric>
 
+#include <algorithm>
 #include <array>
 #include <set>
 #include <cassert>
@@ -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<int64_t>& tag_keys);
 
+class Instruction;
+// Order dependents by creation order (_global_inst_id), NOT by heap address.
+// std::set<shared_ptr<...>>'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<Instruction>& a,
+                  const std::shared_ptr<Instruction>& b) const;
+};
+
 class Instruction : public std::enable_shared_from_this<Instruction> {
  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<Instruction> {
     for (auto& c : _deps[static_cast<size_t>(e)]) c->dec_ready_counter();
     _deps[static_cast<size_t>(e)].clear();
   }
-  const std::set<std::shared_ptr<Instruction>>& get_deps(DepEvent e) {
+  const std::set<std::shared_ptr<Instruction>, DepLess>& get_deps(DepEvent e) {
     return _deps[static_cast<size_t>(e)];
   }
   void set_assigned_sa(int s) { _assigned_sa = s; }
@@ -63,6 +77,12 @@ class Instruction : public std::enable_shared_from_this<Instruction> {
   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<Instruction> {
   // 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<int64_t>& 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<Instruction> {
   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<std::set<std::shared_ptr<Instruction>>,
+  // _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<std::set<std::shared_ptr<Instruction>, DepLess>,
              static_cast<size_t>(DepEvent::COUNT)> _deps;
   std::vector<size_t> tile_size;
   std::vector<int> tile_stride;
@@ -191,4 +222,5 @@ class Instruction : public std::enable_shared_from_this<Instruction> {
   int _assigned_sa = -1;
   std::shared_ptr<WeightToken> _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()
 };

TOGSim/include/TileGraph.h

@@ -4,6 +4,7 @@
 #include <map>
 #include <queue>
 #include <set>
+#include <functional>
 #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<TileSubGraph> 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<std::shared_ptr<TileSubGraph>()> 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<std::shared_ptr<TileSubGraph>()> _pull;   // empty -> legacy
+  bool _exhausted = false;
+  bool refill();
   std::string _path;
   std::string _name = "?";
   unsigned int _kernel_id = 0;
   std::vector<std::string> _loop_index_list;
   std::vector<std::tuple<int, int, int>> _ranges;
   std::vector<std::shared_ptr<TileSubGraph>> _subgraph_vec;
-  std::vector<std::shared_ptr<TileSubGraph>> _finished_subgraph_vec;
   std::map<int, std::map<int, std::shared_ptr<TileSubGraph>>> _cpu_graph_map;
   std::shared_ptr<IntervalTree<unsigned long long, int>> _cache_plan;
   cycle_type _arrival_time;

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 <cstdint>
+#include <functional>
 #include <vector>
 
 #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<void(const TraceRec&)>;
+// ---------------------------------------------------------------------------
+// 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<int64_t> 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<WorkItem> _items;
+  uint64_t _stray = 0;
+  std::vector<uint64_t> _bases;
+  std::vector<int64_t> _cyc, _ovl;
+};
+
 }  // namespace togsim

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<TileGraph> 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<TileGraph> 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);

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>(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<Instruction>& 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<Instruction>& 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));