[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,15 @@ 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, NOT by heap address: std::set<shared_ptr>'s
+// default comparator sorts by pointer value, which makes fire()'s release order --
+// and with it the issue order and the reported cycle count -- allocator-dependent.
+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 +61,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 +73,10 @@ 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; }
+  // This preload's weight-slot refcount: the MATMULs subscribed to its ISSUE event.
+  // Fixed once the tile is built, but Core's issue scan asks on every stalled cycle
+  // and one preload can have 300k subscribers -- so count once and cache.
+  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 +144,13 @@ 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. The Core frees each version regardless
+  // of order, but the invariant lets a builder that tags versions as they close be
+  // byte-comparable with one that walks the version map at the end.
+  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 +179,9 @@ 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] at finish (latency). std::set dedups and, via DepLess, iterates in
+  // instruction-id order, 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 +212,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,16 @@ 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: `_pull` materializes the NEXT dispatch tile, or returns
+  // nullptr once the producer is exhausted. allocate_subgraph calls it only when a
+  // core needs work, so peak memory is O(tiles in flight), not O(dispatches).
+  void set_tile_source(std::function<std::shared_ptr<TileSubGraph>()> pull) {
+    _pull = std::move(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 +142,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

@@ -1,7 +1,7 @@
 #pragma once
 // togsim_loader.h -- the TOGSim half (not the producer ABI): `dlopen` a producer
-// `.so`, run its `togsim_kernel`, record the emitted instructions.  The
-// "materializing sink" of sec 5.3 / 9.7; the stream goes to togsim_trace_bridge.h.
+// `.so`, run its `togsim_kernel`, and record the emitted instructions as
+// TraceRecs (sec 5.3 / 9.7). togsim_trace_bridge.h turns them into a TileGraph.
 
 #include <cstdint>
 #include <vector>
@@ -34,18 +34,40 @@ struct TraceRec {
   int64_t  overlapping;   // looked up from the cycle table
 };
 
-struct RunResult {
-  bool ok = false;
-  std::vector<TraceRec> trace;
+// The producer is run ON DEMAND, one work-item at a time (materializing the whole
+// record stream cost 12.7 GiB on a large 8x8 conv). A tile body reads only ctx and
+// iv, so a work-item is (fn, iv, core) -- replayable on its own, whenever needed.
+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 when it is registered
 };
 
-// Load `so_path`, run its `togsim_kernel`, and return the recorded trace.
-// `tensor_base` gives each tensor argument's DRAM base, `cyc`/`ovl` the cycle table.
-// Work-items round-robin only over `partition_cores` (empty/null -> core 0).
-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);
+class LazyProducer {
+ public:
+  LazyProducer() = default;
+  ~LazyProducer();
+  LazyProducer(const LazyProducer&) = delete;
+  LazyProducer& operator=(const LazyProducer&) = delete;
+
+  // dlopen the .so and run togsim_kernel once: every togsim_dispatch registers
+  // its work-item and returns without running the tile body, so num_items() is
+  // known but no record is emitted yet.
+  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);
+
+  size_t num_items() const;
+  // Replay work-item `i` and return its record stream (TILE_BEGIN, body,
+  // TILE_END). The returned vector is a buffer reused by the next call.
+  const std::vector<TraceRec>& run_item(size_t i);
+
+ private:
+  struct EmitCtx* _ctx = nullptr;   // opaque; owns the work-item list
+  void* _lib = nullptr;
+  std::vector<uint64_t> _bases;
+  std::vector<int64_t> _cyc, _ovl;
+};
 
 }  // namespace togsim

TOGSim/include/togsim_runtime.h

@@ -14,8 +14,8 @@ extern "C" {
 #define TOGSIM_ABI_VERSION 12
 int32_t togsim_abi_version(void);
 
-// Opaque per-invocation context owned by TOGSim. Holds the record sink and the
-// tile_id->cycle lookup. Never dereferenced by the producer.
+// Opaque per-invocation context owned by TOGSim. Holds the recorded trace and
+// the tile_id->cycle lookup. Never dereferenced by the producer.
 typedef struct EmitCtx EmitCtx;
 
 // Direction for togsim_dma.

TOGSim/include/togsim_trace_bridge.h

@@ -7,6 +7,15 @@
 #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 (see togsim_loader.h): peak
+// memory is O(tiles in flight), not O(dispatches). Sound because a tile is
+// dependency-closed; SRAM buffer VERSIONS do cross tiles though.
+//
+// `name` labels the graph. Returns nullptr if the producer fails to load or run.
+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,16 @@ 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: with none free the preload cannot issue,
+                // so nothing else about it is worth computing. Safe to reorder --
+                // try_occupy_sram above is a no-op for a preload (weights untracked).
+                sa_idx = pick_free_weight_sa();
+                if (sa_idx < 0) break;              // all weight slots full -> stall (retry)
+                const int n_consumers = inst->matmul_consumers();   // cached
                 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 +412,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 +444,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 +499,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 +516,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 +563,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));

TOGSim/src/DMA.cc

@@ -33,7 +33,7 @@ std::shared_ptr<std::vector<mem_fetch*>> 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<std::vector<mem_fetch*>> 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,