[TOGSim] Build the TileGraph on demand, one dispatch tile at a time#301
Open
YWHyuk wants to merge 9 commits into
Open
[TOGSim] Build the TileGraph on demand, one dispatch tile at a time#301YWHyuk wants to merge 9 commits into
YWHyuk wants to merge 9 commits into
Conversation
…rializing it The --trace_so path ran the producer into a std::vector<TraceRec> and then built the TileGraph from that vector, so the recorded stream and the graph both lived at peak. heaptrack on an 8x8 conv (28800 GEMM tiles) showed a 69.6M peak heap made of two nearly equal halves: 23.1M for the trace vector and 25.6M for the TileGraph. The stream is O(#tiles), and an 8x8 array has (128/8)^3 = 3136x the tiles of a 128x128 one, so a large conv (test_conv2d batch2 128->512 14x14 k7) outgrew the CI runner and was SIGKILLed. Feed the records to a sink instead. run_producer_stream runs the producer and hands each record to a callback, retaining nothing. trace_to_tilegraph now takes the .so directly and replays the producer twice: once to collect the per-buffer spad footprints (which must be known before the producing compute is processed), once to build the graph. togsim_kernel is a pure emitter, so the replay yields an identical stream. Measured on that same 8x8 conv: peak RSS 76.2M -> 54.5M (-28.5%), peak heap 69.6M -> 45.2M (-35%), and the trace vector no longer appears in the profile. Total execution cycles are unchanged (129461 before and after; 33000 on a smaller case), so the footprint prepass and the dependency DAG are built from exactly the same data as before. run_producer (materializing) is retained for simulate().
…ot O(K^2) link()'s is_mm_accum UNION branch rescanned all of writers(b) on every one of the K matmuls that accumulate into b, to add a dep edge to the non-MATMUL seed (the init/bias). But a UNION only ever appends MATMULs -- the seed set is fixed until the next REPLACE -- so the rescan found the same one seed every time and skipped everything else. Measured with a counter on an 8x8 conv whose K_tiles is 784 (50176 accumulating matmuls): the scan ran 1,258,840,576 times = N(N+1)/2, and only 50176 of those iterations produced an add_dep. 99.996 percent of the work was wasted. The real CI case (test_conv2d batch2 128->512 14x14 k7) has N=2458624, i.e. about 3e12 iterations, which is why it never finished. Track the non-MATMUL subset of writers(b) in `seeds` and scan only that; route every writers(b) REPLACE through set_writer() so the two stay consistent. The scan drops to 50176 on that case (25088x fewer iterations) and the set of add_dep edges is unchanged. Total execution cycles are identical on all three cases checked (33000, 129461, 202082), before and after.
Instruction::_deps is a std::set<std::shared_ptr<Instruction>>, whose default comparator orders by pointer value. That only looked deterministic because the whole TileGraph was allocated up front in one monotonically growing run, so addresses happened to track creation order and fire() released dependents in creation order. It is not a property of the code. Running the same binary on the same trace under a different allocator changes the reported cycle count: kernel glibc malloc tcmalloc 677alforuj6 7852 7846 6lxncgaxtnh 3490 3491 7ascx6t7uu7 2888 2889 Order the sets by _global_inst_id instead. The id is intrinsic to the instruction, so the release order -- and the issue order and the cycle count -- no longer depend on the heap layout. All three allocators above now agree. Five of twelve real kernels shift by 1..6 cycles; a two-tenant model shifts by 416 of 1.4M (0.03%). Those numbers were allocator artefacts, not the model. This is also a prerequisite for building the TileGraph on demand, which frees and reuses tiles and would otherwise perturb the release order.
sram_finalize() tags each buffer version's last reader while walking the version map, which is ordered by version id, so an instruction that frees several versions receives them ascending. Make that an invariant of add_sram_release instead of an accident of the caller's iteration order. No behaviour change today. It lets the builder tag a version the moment it closes -- in buffer order -- and still produce a byte-identical graph, which the on-demand construction that follows relies on.
togsim_dispatch(ctx, fn, iv, n) hands the runtime the work-item's function pointer and its parallel induction variables, and the tile body reads nothing but ctx and iv. So a single pass that records (fn, iv, core) for every dispatch is enough to invoke any one work-item later, on its own, with no replay of the others. LazyProducer::open() does that pass. It also runs the tiles and streams every record to an optional sink, so it doubles as the spad-footprint pre-pass the TileGraph builder already needed -- nothing in the stream is lost, including the records some producers emit outside a dispatch (those belong to no work-item and never entered the graph; stray_records() reports them). LazyProducer::run_item(i, sink) then replays exactly work-item i. Nothing uses it yet; the TileGraph builder switches over next.
The builder used to materialize every dispatch tile up front and hand the
finished graph to the Simulator, even though a Core only ever holds one or two
tiles (Core::max_concurrent). Peak memory therefore scaled with the number of
dispatches. An 8x8 conv2d(x[2,128,14,14], w[512,128,7,7], pad=3) -- 49 dispatch
tiles, 17.7M instructions, 37.0M dependency edges -- needed 12.7 GiB and was
SIGKILLed in CI, while the mathematically identical torch.mm needed 1.8 GiB.
Build a tile only when a Core asks for work:
* LazyProducer indexes the dispatches, so work-item k can be replayed alone.
* TileGraph::_pull materializes the next tile from allocate_subgraph(), i.e.
only when a core has run out; a finished subgraph is now dropped instead of
being parked in _finished_subgraph_vec, which is what actually releases a
tile's Instructions.
* sram_schedule() precomputes each buffer version's last reader in a pass that
allocates nothing, so the builder no longer has to retain every reader of
every version until end-of-stream (7.5M shared_ptrs on the conv above) just
to run sram_finalize().
This is sound because a dispatch tile is dependency-closed: flush() already
reset writers/seeds/tag maps at every tile boundary, so no dependency edge
crosses tiles (measured: 0 of 37.0M). Buffer versions do cross tiles, and still
do -- only their bookkeeping moved.
Peak resident memory, building the graph:
MM16 241,656 KB -> 44,032 KB (5.5x)
MM36 1,913,620 KB -> 488,484 KB (3.9x)
CI36 13,353,976 KB -> 549,948 KB (24.3x)
and it no longer grows with the dispatch count. Build time drops too (CI36:
23.8s -> 20.4s) -- fewer allocations and page faults.
Cycle counts are unchanged. Verified on 14 kernels (including a 2-core
2-partition config) and on the multi-tenant scheduler test: 1,401,595 cycles
before and after. run_producer_stream has no callers left; drop it.
Two redundant allocations per DMA and per barrier, on the graph-build path:
* Instruction's constructor takes its five vectors by value and then copied
them into the members, allocating a second buffer for each.
* prepare_tag_key() pushes two elements onto an empty vector, so it allocates
at capacity 1 and immediately reallocates at 2.
* the producer callbacks built a fresh TraceRec per record, growing its four
vectors from empty every time; the sink consumes a record before the next is
emitted, so one reused scratch record does.
Move instead of copy, reserve the exact size, and reuse the record. Building the
graph for an 8x8 conv2d with 17.7M instructions: 19.87s -> 19.60s of CPU. Small,
but these are 27M allocations that never had to happen.
No behaviour change: 297 kernels, identical cycle counts.
core_trace_log::trace_instruction_line() takes formatted strings, so its arguments were built at every call site whether or not spdlog would keep them: each issued instruction paid a TraceLogTag::pad15() and a fmt::format over its tile dims and strides, which spdlog::trace() then dropped. The same holds for two spdlog::trace() calls in DMA::cycle that format the instruction's addr_name. Guard them with core_trace_log::trace_enabled(). Output is unchanged when the trace is on; when it is off (the default) an 8x8 conv2d simulation drops from 77.59s to 76.84s of CPU. Modest -- the string work was never the bottleneck -- but it is pure waste.
…scan
Core's issue scan walks a tile's ready list every cycle and issues at most one
instruction. For a PRELOAD it counted the matmuls subscribed to its ISSUE event
by iterating the dependency set -- to size the weight-slot refcount -- before
asking whether a weight slot was even free. A preload that stalls waiting for a
slot is revisited on the next cycle, and the next, and recounted every time. One
preload's ISSUE set can hold hundreds of thousands of matmuls.
Measured on an 8x8 conv2d (CM16): that count was 83% of the whole simulation
(Core.cc:358-359 plus the std::set iteration inlined into it), and
matmul_consumers() was entered 762,064,838 times to perform 25,088 counts.
Two fixes:
* cache the count on the Instruction -- it is a property of the graph and
cannot change once the tile is built;
* ask pick_free_weight_sa() first, so a preload that cannot issue costs a
compare instead of a walk. Safe to reorder: try_occupy_sram() is a no-op for
a preload (the virtual SA-weights buffer is untracked), so nothing has been
reserved that would need giving back, and the call count -- hence the
round-robin cursor it advances -- is unchanged.
CM16 (conv, 803k cycles) 78.5s -> 29.3s (2.80x)
MM16 (gemm, 803k cycles) 55.6s -> 38.9s (1.43x)
MM36 (gemm, 9.8M cycles) ~1h56m -> 5m51s
Cycle counts unchanged: verified on 297 kernels.
The scan itself is now the top cost -- a stalled preload is still visited once
per cycle. Making the ready list resource-aware is the next step.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Build the TileGraph one dispatch tile at a time instead of all at once, and fix
the hot loops that a large conv then exposed. A big 8x8 conv2d that used to be
SIGKILLed in CI (peak 12.7 GiB, never finished) now completes at 1.66 GiB, and
the simulation is materially faster. Cycle counts are unchanged.
This targets develop and is nine commits, each buildable and
cycle-verified on its own.
The problem
A conv is a GEMM: M = batchHW, N = C_out, K = C_inkhkw. So the CI conv2d
(x[2,128,14,14], w[512,128,7,7], pad=3) is the identical 392x512x6272 multiply
as the matmul twin. Yet the conv had two problems:
17.7M instructions, 37.0M dependency edges), even though a core only ever
holds one or two tiles. The graph hit 12.7 GiB and the OOM killer took it.
What changed
Memory. The dispatch runtime already hands over each work-item's function
pointer and induction vars, and a tile body reads nothing else. So an indexing
pass records (fn, iv, core) per dispatch, and a tile's instructions are built
only when a core asks for that work-item -- single-threaded, no replay of the
others. Buffer-version lifetimes are precomputed in an allocation-free pass,
retiring the structure that held 7.5M instruction references alive just to run
one finalize at end-of-stream. Sound because a dispatch tile is dependency-
closed (measured: 0 of 37.0M edges cross a tile boundary).
MM16 full-sim peak RSS 232 MiB -> 71 MiB
MM36 graph build peak 1.83 GiB -> 477 MiB (3.9x)
CI36 graph build peak 12.74 GiB -> 537 MiB (24.3x)
Peak no longer scales with the dispatch count, so a bigger conv can't run out.
Speed. A CPU profile of a conv simulation pointed at one loop: the issue scan
counts the matmuls reusing a preload's weight by iterating its dependency set,
every cycle the preload stalls waiting for a weight slot. One preload's set can
hold hundreds of thousands of matmuls. Measured on CM16 it was 83% of the whole
simulation; matmul_consumers() was entered 762,064,838 times to perform 25,088
counts. Cache the count (it is a graph property) and check the weight slot first.
CM16 (conv, 803k cycles) 78.5s -> 29.3s CPU (2.8x)
MM16 (gemm, 803k cycles) 55.6s -> 38.9s CPU (1.4x)
MM36 (gemm, 9.8M cycles) ~1h56m -> 5m51s wall
Two smaller allocation cleanups (move-not-copy the instruction vectors, reserve
tag-key exactly, reuse one TraceRec) and gating the instruction trace strings
behind the log level are included.
A latent bug this surfaced
Instruction::_deps was a std::set<shared_ptr>, ordered by pointer value. That
only looked deterministic because the whole graph was allocated in one
monotonically growing run, so addresses tracked creation order by accident.
Running the same binary on the same trace under a different allocator moved the
cycle count (677alforuj6: glibc 7852 vs tcmalloc 7846). Ordering the set by the
intrinsic _global_inst_id makes the release order -- and the cycle count --
independent of heap layout. This is also a prerequisite for on-demand
construction, which frees and reuses tiles.
Verification
the multi-tenant scheduler run; 0 mismatches), including a 2-core 2-partition
config.
MACs / (882 arrays) = 39,335,938.
Not addressed here
The upstream cause of the conv's size is untouched: batch=2 pads to TILE_M=8, so
the conv still emits 4x the instructions of its matmul twin. That is a frontend
mapping fix, separate from this simulator-side work. The issue scan is now the
top remaining cost (a stalled preload is still visited once per cycle); making
the ready list resource-aware is a natural follow-up, but it touches issue order
directly and would need its own cycle-equivalence proof.
🤖 Generated with Claude Code
https://claude.ai/code/session_014PxddUzqSZzZBcKTC8uahV