From c3f7647ade0de379925f6baaaaa854feccbbfe15 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 24 Jun 2026 22:35:51 +0900 Subject: [PATCH 01/32] [Docs] C++ trace pipeline design (runtime-tag pairing, ABI) --- PyTorchSimFrontend/extension_codecache.py | 9 +- docs/design/togsim_cpp_trace.md | 558 ++++++++++++++++++++++ 2 files changed, 561 insertions(+), 6 deletions(-) create mode 100644 docs/design/togsim_cpp_trace.md diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 492133a3..3fbb6c49 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -207,12 +207,9 @@ def load(cls, source_code, lock = FileLock(get_lock_path(write_path), timeout=LOCK_TIMEOUT) with lock: try: - # mlir-opt now runs only loop-padding/dma-fine-grained/pytorchsim-to-vcix - # and writes the post-vcix IR. The tile-operation-graph pass is ported - # to Python: run_tog reads that IR, writes the TOG (_tog.py) and the - # mutated IR (_custom.mlir: sample-mode step rewrite + compute markers), - # replacing the C++ -test-tile-operation-graph pass. - # loop-padding(timing, mlir-opt) -> Python fine-grained + vcix (one parse/print) + # mlir-opt now runs only loop-padding and writes the post-vcix IR; the + # tile-operation-graph pass is ported to Python. run_tog reads that IR and + # writes the TOG plus the mutated IR (step rewrite + compute markers). subprocess.check_call(gem5_pad_cmd) run_module_passes(sample_mlir_path + "_padded.mlir", sample_mlir_path + "_postvcix.mlir", diff --git a/docs/design/togsim_cpp_trace.md b/docs/design/togsim_cpp_trace.md new file mode 100644 index 00000000..179163a8 --- /dev/null +++ b/docs/design/togsim_cpp_trace.md @@ -0,0 +1,558 @@ +# TOGSim C++ Trace Generation + +TOGSim's timing path consumes a **Tile-Operation Graph (TOG)**: the stream of +modeled instructions a kernel executes (tile loads, tile computes, tile stores) +together with the dependency edges between them. The timing core turns that +graph into cycles. + +This document describes how the TOG is produced: the kernel's post-vcix MLIR is +compiled into a small C++ program, and *running* that program emits the graph. +It replaces an earlier producer that materialized a flattened graph at compile +time and serialized it as ONNX. + +Vocabulary used throughout: + +- **Producer** — the generated `.so`. Shape-parametric C++ compiled from the + kernel's MLIR. It contains no timing model and no functional compute; running + it emits a trace and nothing else. +- **Trace record** — one `togsim_*` callback invocation, i.e. one modeled + instruction. +- **Bridge** — the TOGSim-side code that turns the recorded trace into a + `TileGraph` of `Instruction`s with dependency edges. +- **Core** — TOGSim's existing timing model (systolic arrays, VPU, DMA engine, + SRAM, DRAM/NoC). It is unchanged by this pipeline. + +## 1. Motivation + +The legacy TOG producer (`MLIR -> Python dict -> ONNX -> C++ TileGraphParser`) +had four structural problems: + +1. **"ONNX in name only."** The graph was serialized as ONNX, but every op was a + custom `torchsim_*` attribute. That paid ONNX's costs (rigid schema, + protobuf, stringly-typed attributes) for none of its interop value. The + schema lived in three places — Python (`extension_op.py`), ONNX + (`AsmParser/onnx_utility.py`), C++ (`TileGraphParser`) — and drifted. + +2. **Synchronization was ad-hoc and DMA-specific.** One concept ("an async op + completed; a consumer may proceed") was expressed two different ways: a + content-addressed `tag_table` with overloaded magic values (`0` pending, `1` + signaled, `>1` consumed-count, `-1` sparse) plus a separate + `Instruction::ready_counter` / `child_inst` edge mechanism. Only the former + worked for DMA. + +3. **Static shape was baked in.** Loop bounds were resolved to constants and the + graph fully materialized per shape, so dynamic shape forced a recompile per + shape — pathological for LLM decode (a new `seq_len` every step) and MoE + (variable expert load). + +4. **Loop-flattening hacks.** `loop_end` tricks, the `calc_tag` content hash, + dedup-by-skip and magic offsets existed only to flatten loop nests into a + static graph. + +See [Appendix A](#appendix-a-legacy-path-references) for the file references. + +## 2. Model: trace-driven, not graph-materialized + +Rather than materializing a flattened graph, **the TOG becomes a stream emitted +by running a shape-parametric producer.** The producer keeps loops as loops +(symbolic bounds become C++ function parameters) and calls a small set of +callbacks. Each call emits one trace record. TOGSim `dlopen`s the producer, +injects a callback context, and records the stream. + +| Problem | Resolution | +|---|---| +| ONNX-in-name-only, 3-place schema | The C ABI is the single contract. No ONNX. | +| DMA-only, ad-hoc sync | Dependencies are explicit dataflow edges (§10). Async-DMA data arrival is one explicit barrier op keyed by the runtime tag slot. | +| Static shape | Loop bounds pass through from MLIR verbatim; symbolic bounds become native C++ loop bounds, so trip count is dynamic. | +| Loop-flatten hacks | Loops stay loops. `calc_tag` and dedup disappear. | + +This is **not** a dynamic hardware scheduler. Control flow is still statically +emitted by the compiler; the `.so` is a deterministic *trace generator*, not a +timing model. The trace-as-data boundary is preserved, so the timing core is +untouched. + +## 3. Trace op vocabulary + +Four primitives. Everything else is composition. + +- `dma(dir, arg_id, offset, shape, is_async, tag_id, tag_slot, read_bufs, + write_bufs)` — `dir ∈ {LOAD, STORE}`. A **synchronous** DMA is blocking: it + finishes when its data arrives, and consumers depend on it directly. An + **async** DMA returns control immediately and signals its tag at data arrival + (DMA response-complete); a later `memory_barrier` is the explicit point that + waits on it. +- `compute(tile_id, compute_type, dims, read_bufs, write_bufs)` — one fixed-size + tile kernel. Its cost is looked up (§6), never computed here. +- `memory_barrier(tag_id, tag_slot, write_bufs)` — the explicit async-DMA sync. + It waits until the async DMA carrying the same `(tag_id, tag_slot)` has + delivered its data, then becomes the last writer of the loaded buffer so + consumers gate on arrival. It is the source IR's `togsim.wait` mapped through, + not a synthesized barrier. +- `dispatch(tile_fn, iv, n_iv)` — run one parallel work-item (§9.3). + +Control flow lives in the producer: ordinary `for`/`if` with runtime bounds. + +Two things share the word "tag", and the pairing key is **both together**: + +- **`tag_id`** — compile-time identity of a DMA's tag memref (which logical + channel: the A-load vs the B-load). +- **`tag_slot`** — the runtime SRAM tile slot the loaded tile occupies (the + double-buffer / SRAM-capacity index). + +The key must include a runtime component: one static `togsim.dma` op executes +once per loop iteration, each iteration writing a different slot, so a +compile-time id alone cannot express the per-iteration pairing. + +`lower_to_vcix` writes the wait's tag index with a `-acc_iv` term per +accumulation (reduction) loop var — a sentinel marking the reduction axis, not +an arithmetic offset — and `build_skeleton` strips those terms so a +`memory_barrier` waits on the same slot its async load wrote. (The legacy +`TileGraphParser` mirrors this by skipping stride `-1`.) Without the strip, the +producer evaluates `-acc_iv` to a negative slot at reduction iteration > 0 and +the pairing fails on subtile + multi-tile-K. + +## 4. Decisions + +| Axis | Decision | +|---|---| +| Input MLIR | Use the **given MLIR as-is**. Do not touch inductor, the MLIR templates, or shape plumbing. Whatever bounds the MLIR carries (const or symbolic) pass through verbatim. | +| MLIR -> C++ | **EmitC dialect + `mlir-translate --mlir-to-cpp`** (upstream). | +| `.so` <-> TOGSim | **`dlopen` + an opaque `EmitCtx` callback context.** The ABI boundary is the main design surface. | +| `.so` role | **Timing trace only.** Functional correctness stays on the Spike/LLVM path. Every op without a timing dependency is stripped; the loop skeleton, the API ops, and the ops feeding bounds/addresses remain. | +| Compute cycle | A separate annotation pass reuses **gem5 sample-mode** to build a precomputed `tile_id -> cycle` table, looked up at runtime. | +| Dynamic shape | Falls out of symbolic loop bounds. Per-tile cost is static (tiles are fixed-size); only trip count is dynamic. | + +## 5. Architecture + +### 5.1 Artifacts (per kernel) + +- **Trace `.so`** — compiled from the skeleton MLIR. Shape-parametric: symbolic + bounds become C++ function parameters. +- **Cycle table** — `tile_id -> (cycle, overlapping_cycle)`, a TSV sidecar. + +Both are written next to the kernel's `tile_graph.onnx`. TOGSim picks the `.so` +up automatically; `TORCHSIM_LEGACY_TOG=1` forces the deprecated ONNX path. + +### 5.2 Pipeline + +``` +post-vcix MLIR (affine/scf.for + togsim.transfer/wait + vcix/vector compute) +| ++-- Branch A (trace): +| build_skeleton.py loops kept, bounds as-is (symbolic preserved) +| togsim.transfer -> togsim.dma(..., tag_id, %tag[%idx], is_async) +| togsim.wait -> togsim.memory_barrier(tag_id, %tag[%idx], write_bufs) +| compute body -> togsim.compute(tile_id, dims) +| DCE everything with no path to a loop bound, an +| address, or an API operand +| dep_analysis.py per-op read/write SRAM buffer sets (SSA) + the vcix +| preload/matmul pairing (§10.2) +| lower_to_emitc.py togsim.* -> emitc.call_opaque; drive the upstream +| lower-affine / convert-{scf,arith,func}-to-emitc +| mlir-translate --mlir-to-cpp -> C++ +| g++ -shared -> trace.so +| ++-- Branch B (cost): + cycle_table.py reuse the gem5 sample-mode cycle_list already computed + in extension_codecache -> tile_id -> (cycle, overlapping) + +TOGSim: + run_producer() dlopen(trace.so), resolve togsim_kernel, inject EmitCtx + { core pool; record sink; cycle table }, run it + -> a TraceRec stream + trace_to_tilegraph() TraceRec -> TileGraph (Instruction DAG, §10) + Simulator / Core cycles, DRAM traffic (unchanged) +``` + +### 5.3 Components + +- **`togsim_ops.py`** — the op vocabulary. The ops are kept *unregistered* (like + the existing `togsim.transfer`), so no C++ dialect registration is needed and + the togsim->emitc step is a Python rewrite rather than a registered + ConversionPass. +- **`build_skeleton.py`** — reduces the kernel to loops + API ops. Preserves + `is_async`; maps `togsim.wait` through to an explicit + `togsim.memory_barrier`. The IR verifies across sibling prefetch/compute loop + nests because the DMA/barrier pairing is by runtime tag slot, not a + cross-region SSA edge. +- **`dep_analysis.py`** — the read/write SRAM buffer sets per op (§10.2). +- **`lower_to_emitc.py`** — rewrites `togsim.*` to `emitc.call_opaque`, outlines + the work-item body (§9.4), then drives the upstream conversion passes. +- **`cycle_table.py`** — the `tile_id -> (cycle, overlapping_cycle)` sidecar. +- **`togsim_runtime.{h,cc}`** — the C ABI and the `EmitCtx` callback + implementations. +- **`togsim_loader.h`** — `run_producer`: `dlopen`, ABI-version check, run, + record. +- **`togsim_trace_bridge.{h,cc}`** — `TraceRec` stream -> `TileGraph`. + +### 5.4 ABI (v12) + +`mlir-translate --mlir-to-cpp` lowers `emitc.call_opaque` to *free function* +calls, so the contract is a set of `extern "C"` functions taking an opaque +`EmitCtx*` first argument. The loaded `.so` links back into the Simulator binary +(built with `ENABLE_EXPORTS`), so the symbols resolve without an explicit +function table. `togsim_abi_version()` guards against a producer built against a +stale header. + +```c +typedef struct EmitCtx EmitCtx; +typedef void (*togsim_tile_fn)(EmitCtx*, int64_t* iv, int32_t n_iv); + +int32_t togsim_abi_version(void); + +void togsim_dma(EmitCtx*, int32_t dir, int32_t arg_id, uint64_t offset, + int32_t ndim, const int64_t* dims, const int64_t* strides, + int32_t elem_bits, int32_t is_async, + int32_t tag_id, uint64_t tag_slot, + const int64_t* read_bufs, int32_t n_read, + const int64_t* write_bufs, int32_t n_write); + +void togsim_compute(EmitCtx*, uint64_t tile_id, int32_t compute_type, + int32_t ndim, const int64_t* dims, + const int64_t* read_bufs, int32_t n_read, + const int64_t* write_bufs, int32_t n_write); + +void togsim_memory_barrier(EmitCtx*, int32_t tag_id, uint64_t tag_slot, + const int64_t* write_bufs, int32_t n_write); + +void togsim_dispatch(EmitCtx*, togsim_tile_fn fn, int64_t* iv, int32_t n_iv); + +// entry point the loader resolves: +void togsim_kernel(EmitCtx*, int64_t* shape_args, int32_t n_shape_args); +``` + +`offset` is an **element** offset within tensor `arg_id`, computed by the +producer from the loop indices; only the runtime knows the tensors' allocation +bases, so it forms `base[arg_id] + offset * elem_bytes`. `compute_type` is +`0` vector / `1` matmul / `2` preload, and routes the op to the VPU or the +systolic array. + +## 6. Compute cost model + +A `togsim.compute(tile_id=...)` says *which* tile to compute, not how long it +takes. Because tiles are fixed-size (`TILE_M/N/K`), each tile's cost is +invariant — only the trip count varies with shape — so it is sampled once and +stored, keyed by `tile_id`. Two numbers per tile, mirroring the legacy TOG: + +- `cycle` — full compute latency, sampled by gem5 sample-mode (the existing + `cycle_list` measurement, reused so both paths stay cycle-consistent). +- `overlapping_cycle` — the portion that overlaps the previous instruction in + the systolic pipeline. The timing core uses it as + `finish = prev.finish + cycle - overlapping`. Derived exactly as the legacy + path does: vector -> `0`, matmul -> `max(cycle - x_offset, 0)`, preload -> + `max(cycle - w_offset, 0)`. + +`togsim_compute` looks both up and sets them on the `Instruction`. + +**Remainder tiles.** When a dimension is not divisible by the tile size, the +edge tile is partial and its true cost differs from the table entry. Today it is +charged the full-tile cost. Sampling a separate `tile_id` for the remainder is +the alternative; see §11. + +## 7. EmitC lowering notes + +`lower_to_emitc` drives the upstream `lower-affine`, `convert-scf-to-emitc`, +`convert-arith-to-emitc` and `convert-func-to-emitc` passes. One gap in this +LLVM 20 build: `convert-scf-to-emitc` emits `emitc.for` with `index` bounds, so +`convert-arith-to-emitc` leaves `builtin.unrealized_conversion_cast` on the +bounds (`emitc.size_t` <-> `index`) that `--reconcile-unrealized-casts` cannot +fold and `mlir-to-cpp` cannot print. + +`_retype_for_to_size_t` therefore retypes each `emitc.for` to `!emitc.size_t` +bounds and induction variable, and folds the residual casts. A `size_t` IV also +makes the lowered *address* arithmetic cast-free, which is what lets each +`togsim_dma` pass a real `(arg_id, element offset)` computed from the loop IVs. + +Unregistered ops (`togsim.*`) have no registered conversion patterns, which is +why the rewrite to `emitc.call_opaque` must be a custom pass and must run before +the upstream conversions. + +## 8. Validation + +The reproduction path for a single kernel: + +```sh +python -m PyTorchSimFrontend.mlir.passes.lower_to_emitc \ + --so trace.so [--emit-cpp trace.cpp] +bin/Simulator --config --trace_so trace.so \ + [--cycle_table trace_cycles.tsv] [--log_level trace] +``` + +`--log_level trace` prints the per-instruction issue/finish timeline, which is +how the dependency model is checked: on a 256^3 GEMM the trace shows the +preloads and matmuls pipelining across the systolic arrays, the store waiting +for the last matmul to drain, and each compute waiting for its async weight +load's *data arrival* rather than its issue. Compute work and DRAM traffic match +the legacy path on the same gem5 cycle table. + +## 9. Parallelism, reduction, and core dispatch + +### 9.1 Where the semantics come from + +Nothing has to be inferred. The post-vcix `affine.for` already carries the +mapping decision the frontend made, and `build_skeleton` preserves it: + +| attribute | meaning | role | +|---|---|---| +| `outer_loop` | PARALLEL axis (e.g. GEMM m, n) | independent output tiles -> distributable across cores | +| `accumulation_loop` | REDUCTION axis (e.g. GEMM k) | partial sums into one output tile -> ordered dependency | +| `inner_loop` | tile micro-loop | within one tile | + +### 9.2 Principle: bake intrinsic, parameterize extrinsic + +Two kinds of hardware dependence must be treated differently: + +- **Intrinsic** (vector lanes, `TILE_M/N/K`, systolic size) — defines the + *content and cost of each instruction*. Already baked into the IR. +- **Extrinsic** (`num_cores`) — defines only the *distribution* of an otherwise + fixed set of work-items. The tile set, the cost table, and the DMA tile shapes + are all `num_cores`-invariant. + +So `num_cores` is **not** baked into the producer. The producer is +**core-count transparent**: it never names a core or a core count. + +### 9.3 Core-transparent work function + dispatch hook + +The producer is two functions, split at the PARALLEL/ACCUMULATION boundary: + +```c +// WORK: the trace for ONE independent output tile. Takes the PARALLEL indices; +// names no core. Reduction (k) is program order -- the accumulator is +// core-local, so the ordering is implicit. +static void togsim_kernel_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + int64_t mi = iv[0], ni = iv[1]; + togsim_compute(ctx, /*tile_id=*/0, ...); // acc init + for (size_t ki = 0; ki < KT; ++ki) { // REDUCTION + togsim_dma(ctx, LOAD, A, offA(mi,ki), ..., /*is_async=*/1, /*tag_id=*/0, ki%D, ...); + togsim_dma(ctx, LOAD, B, offB(ki,ni), ..., /*is_async=*/1, /*tag_id=*/1, ki%D, ...); + togsim_memory_barrier(ctx, 1, ki%D, ...); togsim_compute(ctx, 1, ...); // preload + togsim_memory_barrier(ctx, 0, ki%D, ...); togsim_compute(ctx, 2, ...); // matmul + } + togsim_dma(ctx, STORE, C, offC(mi,ni), ...); +} + +// DISPATCH: enumerate the PARALLEL domain, one work-item per call. +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape, int32_t n) { + for (size_t mi = 0; mi < MT; ++mi) + for (size_t ni = 0; ni < NT; ++ni) { + int64_t iv[2] = {(int64_t)mi, (int64_t)ni}; + togsim_dispatch(ctx, togsim_kernel_tile, iv, 2); + } +} +``` + +Three orthogonal concepts: + +- **Parallel** = each `togsim_dispatch` call is an independent work-item. TOGSim + may place it on any core. +- **Reduction** = ordering *inside* one work-item: program order on its core. +- **Core assignment** = owned by `togsim_dispatch`, whose body lives in TOGSim. + It round-robins a core from the partition's pool and brackets the call with + `TILE_BEGIN`/`TILE_END` records, so the work-item's scope is exactly the + function call. A core is an assignment, not a held resource; there is no free. + +Because `togsim_dispatch` takes the work function as a pointer and forwards an +opaque `iv` array, one general dispatcher serves every kernel. The boundary +cannot be optimized away: TOGSim can only observe `togsim_*` callbacks across +the `dlopen` boundary, never a producer-internal call. + +### 9.4 Codegen and ABI + +`lower_to_emitc` outlines the innermost PARALLEL-loop body into an +`emitc.func togsim_kernel_tile(ctx, iv, n_iv)` and rewrites the dispatcher loop +to call `togsim_dispatch`. The `tile_id -> cycle` table is untouched by all of +this (it is `num_cores`-invariant). + +### 9.5 Stance and the split-K exception + +Refining "not a dynamic scheduler": **the per-work-item trace is static and +deterministic; only the work-item -> core binding is dynamic.** That is +independent-task distribution, not data-dependent control flow. + +The transparent model holds while work-items are independent (data-parallel over +output tiles). **Split-K** — a reduction split *across* cores — breaks +independence: the producer would have to emit `c` partials plus a combine, so +the instruction stream would depend on `num_cores`, and the cross-core +dependency would have to be a real dataflow edge rather than program order. +Split-K is a deliberate, scoped exception, not supported today. + +### 9.6 Work-items form a DAG + +Work-items are not always a flat independent set. A computation *between* +parallel loops can only run once the inner parallel region completes: + +``` +parallel for m: + parallel for n: A(m,n) # leaf work-items, each writes a tile of m's buffer + B(m) # join: reads that buffer -> depends on all n of this m +``` + +This needs **no new primitive**. It is the same dataflow-edge mechanism the +trace already uses (§10), at work-item granularity: the join op declares the +leaves' output buffer as an input, so the bridge makes it depend on every leaf +through the last-writer analysis. + +The general picture: **work-items form a DAG whose edges are buffer +producer -> consumer dependencies.** Independent data-parallel work is the +degenerate edge-less case; barriers, reduction across a parallel axis, and +split-K are the same DAG with real edges. + +### 9.7 Execution model: trace generation, not co-execution + +The producer is a pure trace generator. It never computes cycles, models +hardware, or schedules. Two consequences pin the model: + +- **What is an edge vs. what blocks.** Data dependencies are recorded *edges* — + the producer does not block on them. The only thing that can ever block the + producer is resource backpressure (finite cores, double-buffer slots, DMA + queue depth), which is flow control, not timing semantics. +- **Cores, double-buffering, DRAM and NoC are the timing core's job** — reused, + not reimplemented. The producer stays oblivious; depths and counts are + consumer-side config. + +Consumption is staged behind a swappable **sink**, so the choice touches neither +the producer nor the ABI: + +| sink | threads | when | +|---|---|---| +| *materializing* (today) — callbacks append to a `TraceRec` vector, which the bridge turns into a `TileGraph` | none | static shape | +| *streaming* — callbacks push to a bounded queue; the producer runs as a fiber and blocks on backpressure while the DES loop advances time and resumes it | producer fiber | when dynamic-shape trace size makes full materialization impractical | + +Even the streaming sink only blocks the producer on resource flow control, never +on timing-resolved data events. The single forward-compat requirement is that +the callback sink stays an interface. + +## 10. Dependency model + +The trace is an explicit dataflow DAG: every op declares the buffers it reads +and writes, and the bridge derives edges from that. There is no in-order chain, +no runtime content hash, and no op-pattern heuristic. + +### 10.1 Representation + +Each op carries the **SRAM buffer ids** it reads and writes (`read_bufs` / +`write_bufs`). The bridge maintains `writers(b)`, the set of current producers of +buffer `b`, and links each op against it. Resource scheduling — systolic-array +round-robin, double-buffering, SRAM capacity — stays entirely in the Core; the +trace only states producer -> consumer order. + +### 10.2 Two dependency sources + +A single "SRAM access" analysis is necessary but not sufficient: + +| dependency | source | visible in SRAM? | +|---|---|---| +| load -> compute (a DMA writes X_spad/W_spad; preload/matmul read it) | SRAM last-writer per buffer | yes | +| the accumulator chain (init writes Y_spad; the epilogue read-modify-writes it; the store reads it) | SRAM last-writer on Y_spad | yes | +| **preload -> matmul** (a preload loads weights into the systolic array's registers; the matmul consumes them) | **the vcix opcode FSM** (op1 = preload pairs with the following op0 = matmul) | **no — SA-internal, not a memref access** | + +On the 256^3 GEMM post-vcix, the SRAM buffers are `%0 = X_spad(A)`, +`%1 = W_spad(B)`, `%2 = Y_spad(acc)`. The matmul reads `%0` only; the preload +reads `%1`; the matmul does *not* read `%1`, because the weights come from the +systolic array. That is exactly why a memref-only analysis would let the matmul +run before the weight load. `dep_analysis` closes the gap by folding the +preload -> matmul pairing into a virtual `SA_WEIGHTS` buffer, so the FSM edge +becomes an ordinary last-writer edge. + +Both sources are available *before* `build_skeleton` collapses the compute +bodies, which is why the analysis runs on the post-vcix IR. + +### 10.3 Edge rules + +An instruction has two completion points. A systolic-array op **occupies** its +unit for `cycle - overlapping_cycle` (the initiation interval) and its **result** +is ready at `cycle`. `DepEvent::ISSUE` releases a successor at the former, +`DepEvent::DONE` at the latter. The bridge applies one rule per buffer `b`: + +- **Read `b`** — depend on every instruction in `writers(b)`. The edge is + `ISSUE` when consumer and producer are both systolic-array ops (a matmul + reading a preload or a matmul: they overlap on the pipeline), else `DONE`. +- **Write `b`** — replace: `writers(b) = {inst}`. +- **Exception — the commutative accumulator.** A matmul that both reads and + writes `b` is accumulating (`Y += X @ W`). Skip its read edge, and on the + write *union* rather than replace: it waits only for the non-matmul seed (the + init or bias) and joins `writers(b)` without ordering against its + co-matmuls. So the K matmuls do not chain through the accumulator, and a later + reader joins all of them. + +The last rule is what makes the store wait for the whole reduction while the +matmuls still pipeline: the store reads `Y_spad`, `writers(Y_spad)` is the union +of the K matmuls, and the store is not a systolic-array op, so it takes a `DONE` +edge from each. No explicit compute fence is needed. + +### 10.4 Resource models, not edges + +Write-after-read ordering is *not* expressed as an edge. Buffer reuse is a +resource question, and the Core already models it: + +- **SRAM capacity.** A coarse tile is one *version* of its buffer; the fine DMAs + that fill it share one allocation, freed once all of that version's consumers + have issued. A buffer reused by the next reduction iteration or work-item is a + new version that must wait for the old one to be freed — that is the + double-buffer / WAR constraint, enforced by capacity rather than by ordering + edges, so two versions may physically coexist. +- **Weight slots.** A preload takes a weight slot on a systolic array and holds + it until its matmuls have consumed it, which caps how many preloads can be in + flight per array. + +A latency WAR edge would instead force the new load to wait for the old tile's +readers to *finish*, defeating double-buffering. + +### 10.5 Async DMA: the memory barrier + +An async DMA's own finish is *issue*-complete; its data arrives later, at DMA +response-complete. A raw last-writer edge on the DMA would therefore release a +consumer before the data exists — a real bug this model had to fix. + +So an async load's last-writer edge is routed through its `MEMORY_BAR`. The DMA +registers its tag at issue; the barrier parks on `(tag_id, tag_slot)` in the +Core's existing tag table and is woken by `set_tag_finish` at response-complete; +the barrier then replaces `writers(b)`, so every consumer of the loaded buffer +gates on arrival. A synchronous DMA blocks to arrival itself and needs no +barrier. + +The bridge mints a per-record unique key so that successive reduction iterations +of one static DMA op get distinct tag-table keys, and the matching barrier reuses +that key. This works because the recorded stream is already per-iteration — the +producer ran the loops. + +## 11. Limitations and open work + +- **Per-iteration tags are reconstructed in the bridge.** `dma_fine_grained` + emits a fresh tag `memref.alloc` before each coarse load, but `build_skeleton` + DCEs it and keys `togsim.dma` by the alloc's static identity, so the bridge has + to re-derive per-iteration keys from record order. Threading the alloc identity + through as an SSA tag handle would remove that. +- **Preload occupancy.** A preload's `overlapping_cycle` equals its `cycle`, so + its occupancy is zero and concurrent preloads are not capped by the + systolic-array count. Giving the preload a non-zero occupancy (the weight-load + time) is a cycle-model input, not an edge-model change. This predates the trace + pipeline and affects the legacy path equally. +- **Dispatch granularity.** Work-items are enumerated per innermost parallel + loop. Distributing independent output sub-tiles across cores needs the sub-tile + axis roles, which the inner loops do not currently carry. +- **Remainder tiles** are charged the full-tile cost (§6). +- **Op coverage.** The dependency model has been characterized in detail on + GEMM. Other families go through the same rules but have not been + cycle-characterized. +- **Dynamic shape.** Symbolic bounds survive the pipeline, but end-to-end + `shape_args` plumbing and the streaming sink (§9.7) are not done. +- **Legacy path.** The ONNX TOG producer is deprecated and opt-in + (`TORCHSIM_LEGACY_TOG=1`). It is retired once the trace path is stable across + all op families. The cycle measurement (`cycle_list`, `x_offset`/`w_offset`) is + shared, so both paths stay cycle-consistent meanwhile. + +## Appendix A: legacy-path references + +- `TOGSim/include/DMA.h` — `tag_table` (overloaded `0/1/-1/>1`) + `waiters`; + `register_tag` / `set_tag_finish` / `register_tag_waiter` / `mark_tag_used` + (= init / signal / wait / consume). +- `TOGSim/src/Core.cc` — the async-DMA signal path and the barrier wait/consume + path over the tag table. +- `TOGSim/include/Instruction.h` — `ready_counter` / dependency edges and the tag + fields. +- `PyTorchSimFrontend/mlir/passes/build_tog.py` — `TogBuilder.print_operation` + dispatch and `_affine_for_bounds` (constant-bound resolution -> static shape). +- `AsmParser/tog_generator.py` — the ONNX serialization. +- `PyTorchSimFrontend/mlir/mlir_gemm_template.py` — the kernel template emitting + the `affine.for` nest, `linalg.matmul`, and the `togsim.transfer` DMA ops. From 1fe2cb7c40ac6bf7bc3f87c012ec4414d5aa2d4e Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 24 Jun 2026 22:35:51 +0900 Subject: [PATCH 02/32] [TOGSim] Add the C++ trace producer, runtime, loader and bridge Compile a kernel's post-vcix MLIR into a shape-parametric C++ trace producer (build_skeleton -> dep_analysis -> lower_to_emitc -> EmitC -> .so). TOGSim dlopens it, records the emitted instruction stream, and feeds the existing timing Core through togsim_trace_bridge. An async DMA and the memory_barrier that waits on its data pair at runtime by (tag_id, tag_slot): one static dma op runs once per loop iteration, so a compile-time id cannot express the pairing. That also covers multi-tile-K and conv, whose reduction is the kh*kw*C nest. --- AsmParser/tog_generator.py | 3 + PyTorchSimFrontend/extension_codecache.py | 35 +- PyTorchSimFrontend/mlir/passes/_mlir_util.py | 84 +++ .../mlir/passes/build_skeleton.py | 512 ++++++++++++++++++ PyTorchSimFrontend/mlir/passes/cycle_table.py | 102 ++++ .../mlir/passes/decompose_transfer.py | 10 +- .../mlir/passes/dep_analysis.py | 194 +++++++ .../mlir/passes/dma_fine_grained.py | 51 +- .../mlir/passes/lower_dma_to_gemmini.py | 11 +- .../mlir/passes/lower_to_emitc.py | 448 +++++++++++++++ .../mlir/passes/lower_to_vcix.py | 56 +- .../mlir/passes/lower_vlane_idx.py | 10 +- PyTorchSimFrontend/mlir/passes/togsim_ops.py | 101 ++++ TOGSim/include/Instruction.h | 11 +- TOGSim/include/togsim_loader.h | 51 ++ TOGSim/include/togsim_runtime.h | 69 +++ TOGSim/include/togsim_trace_bridge.h | 12 + TOGSim/src/CMakeLists.txt | 5 + TOGSim/src/Core.cc | 26 +- TOGSim/src/CoreTraceLog.cc | 2 +- TOGSim/src/Instruction.cc | 13 +- TOGSim/src/TileGraphParser.cc | 2 +- TOGSim/src/main.cc | 47 ++ TOGSim/src/togsim_runtime.cc | 182 +++++++ TOGSim/src/togsim_trace_bridge.cc | 190 +++++++ 25 files changed, 2144 insertions(+), 83 deletions(-) create mode 100644 PyTorchSimFrontend/mlir/passes/_mlir_util.py create mode 100644 PyTorchSimFrontend/mlir/passes/build_skeleton.py create mode 100644 PyTorchSimFrontend/mlir/passes/cycle_table.py create mode 100644 PyTorchSimFrontend/mlir/passes/dep_analysis.py create mode 100644 PyTorchSimFrontend/mlir/passes/lower_to_emitc.py create mode 100644 PyTorchSimFrontend/mlir/passes/togsim_ops.py create mode 100644 TOGSim/include/togsim_loader.h create mode 100644 TOGSim/include/togsim_runtime.h create mode 100644 TOGSim/include/togsim_trace_bridge.h create mode 100644 TOGSim/src/togsim_runtime.cc create mode 100644 TOGSim/src/togsim_trace_bridge.cc diff --git a/AsmParser/tog_generator.py b/AsmParser/tog_generator.py index a12460e3..8caa9df9 100644 --- a/AsmParser/tog_generator.py +++ b/AsmParser/tog_generator.py @@ -1,3 +1,6 @@ +# DEPRECATED (timing path): the legacy ONNX Tile-Operation-Graph producer, superseded +# by the C++ trace pipeline (build_skeleton + lower_to_emitc + cycle_table -> trace.so). +# Retired once the trace path is stable. See docs/design/togsim_cpp_trace.md. import os import sys import importlib.util diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 3fbb6c49..15c5f7e3 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -238,8 +238,13 @@ def load(cls, source_code, # Run cyclesim cyclesim = CycleSimulator() cycle_list = cyclesim.compile_and_simulate(os.path.join(write_path, cycle_binary_name), vectorlane_size, silent_mode=silent_mode) + # Snapshot for the P3-trace hook below: generate_tile_graph consumes + # cycle_list in place (cycle_list.pop(0) per tile), leaving it empty. + cycle_list_for_trace = list(cycle_list) - # Create TOG + # Create TOG -- DEPRECATED (timing path): the ONNX-TOG producer, superseded + # by the C++ trace pipeline. The cycle_list / x_offset / w_offset computed + # here are reused by cycle_table, so both paths stay cycle-consistent. w_offset, x_offset = vectorlane_size, vectorlane_size if kwargs['loop_size'] is not None and kwargs['loop_size'][-3] < vectorlane_size: x_offset = kwargs['loop_size'][-3] @@ -255,6 +260,34 @@ def load(cls, source_code, w_offset=w_offset, # FIXME. vector_lane=vectorlane_size ) + + # Trace pipeline (opt-in, TORCHSIM_DUMP_TRACE_SO=1): also emit the trace + # producer .so + cycle-table TSV from the SAME post-vcix IR and gem5 cycles, + # so it can be compared cycle-consistently. Best-effort: never breaks the build. + if os.environ.get("TORCHSIM_DUMP_TRACE_SO") == "1": + try: + import mlir.ir as ir + from PyTorchSimFrontend.mlir.passes import ( + build_skeleton as _bs, cycle_table as _ct, lower_to_emitc as _l2e) + pv = sample_mlir_path + "_postvcix.mlir" + _ctx = ir.Context(); _ctx.allow_unregistered_dialects = True + with _ctx: + _mod = ir.Module.parse(open(pv).read(), _ctx) + _bs.build_skeleton(_mod) + _ntiles = len(_ct._compute_types(_mod)) + # align lengths: gem5 gives one numCycles per compute node; + # pad with the last value / truncate if it disagrees. + _cl = list(cycle_list_for_trace) + if _cl and len(_cl) != _ntiles: + _cl = (_cl + [_cl[-1]] * _ntiles)[:_ntiles] + logger.info(f"[P3-trace] cycle_list={cycle_list_for_trace} -> {_cl} " + f"(#tiles={_ntiles}, x_off={x_offset}, w_off={w_offset})") + _tbl = _ct.build_cycle_table(_mod, _cl, x_offset, w_offset) + _ct.dump_cycle_table_tsv(_tbl, os.path.join(write_path, "trace_cycles.tsv")) + _l2e.build_trace_so(pv, os.path.join(write_path, "trace.so")) + logger.info(f"[P3-trace] wrote trace.so + trace_cycles.tsv in {write_path}") + except Exception as e: + logger.warning(f"[P3-trace] trace .so/sidecar dump skipped: {e}") return key class CustomAsyncCompile(AsyncCompile): diff --git a/PyTorchSimFrontend/mlir/passes/_mlir_util.py b/PyTorchSimFrontend/mlir/passes/_mlir_util.py new file mode 100644 index 00000000..a2e95b53 --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/_mlir_util.py @@ -0,0 +1,84 @@ +"""Small, dependency-light helpers shared across the MLIR passes. + +Every pass had its own copy of the same op-walk generator (named variously +`_iter_ops` / `_walk` / `_walk_ops`) and the same one-line attribute builders +(`_i32` / `_i64` / ...). This module is the single source for both. + +Import-safety: `walk_ops` is pure block/op attribute access and needs no MLIR +bindings, so this module does NOT import `mlir.ir` at top level -- some passes +(e.g. lower_vlane_idx, decompose_transfer) are deliberately importable without +the bindings present and only touch `mlir.ir` inside their run functions. The +attribute builders therefore import `mlir.ir` lazily; they require an active +MLIR context (the caller's `with ctx:`), exactly as the per-pass copies did. +""" + + +def walk_ops(block): + """Yield every op under `block` in program order, recursing into regions. + + Snapshots each block's operation list, so a caller may erase ops while + iterating (the strictest of the former copies; a superset of the rest).""" + for op in list(block.operations): + yield op + for region in op.operation.regions: + for b in region.blocks: + yield from walk_ops(b) + + +def _ir(): + import mlir.ir as ir + return ir + + +def i32(v): + """`i32` IntegerAttr for `v` (uses the active MLIR context).""" + ir = _ir() + return ir.IntegerAttr.get(ir.IntegerType.get_signless(32), int(v)) + + +def i64(v): + """`i64` IntegerAttr for `v`.""" + ir = _ir() + return ir.IntegerAttr.get(ir.IntegerType.get_signless(64), int(v)) + + +def i64_array(vals): + """ArrayAttr of `i64` IntegerAttrs for `vals`.""" + ir = _ir() + i = ir.IntegerType.get_signless(64) + return ir.ArrayAttr.get([ir.IntegerAttr.get(i, int(v)) for v in vals]) + + +def str_attr(v): + """StringAttr of `str(v)`.""" + ir = _ir() + return ir.StringAttr.get(str(v)) + + +# attribute readers -- accept an OpView or an Operation; `default` is returned when +# `key` is absent. +def _attrs(op): + return getattr(op, "operation", op).attributes + + +def attr_int(op, key, default=None): + """Integer value of `op`'s `key` attribute, or `default` if absent.""" + ir = _ir() + a = _attrs(op) + return ir.IntegerAttr(a[key]).value if key in a else default + + +def attr_bool(op, key, default=False): + """Bool value of `op`'s `key` attribute, or `default` if absent.""" + ir = _ir() + a = _attrs(op) + return bool(ir.BoolAttr(a[key]).value) if key in a else default + + +def attr_i64_array(op, key, default=None): + """`op`'s `key` ArrayAttr of integers as a Python list, or `default` if + absent (pass `default=[]` for the "missing -> empty" convention).""" + ir = _ir() + a = _attrs(op) + return ([ir.IntegerAttr(x).value for x in ir.ArrayAttr(a[key])] + if key in a else default) diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py new file mode 100644 index 00000000..df4c6046 --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -0,0 +1,512 @@ +"""build_skeleton pass (C2): reduce a kernel's post-vcix MLIR to the +*skeleton + API* form, in place. + +The trace pipeline (docs/design/togsim_cpp_trace.md) compiles a kernel to a +shape-parametric C++ trace producer. The producer is just the kernel's loop +skeleton with the data computation replaced by calls to the event-based runtime +API. This pass performs that reduction at the MLIR level: + + * `memref.dma_start` -> `togsim.dma(...) {tag_id, is_async, ...}` carrying the + runtime tag index operand (`%tag[%idx]`). + * `memref.dma_wait` -> `togsim.memory_barrier(tag_idx) {tag_id, write_bufs}`, + the explicit async-DMA sync. It pairs with its dma by + the RUNTIME tag slot (tag_id + the tag index), not a + compile-time id: one static dma op runs once per loop + iteration with a different `%tag[%idx]`, so only the + runtime slot can pair iteration i's dma with its wait. + * each compute node -> a single `togsim.compute {tile_id, compute_type}` + * everything else -> removed by a use-based DCE, keeping the loops and the + index/address arithmetic the survivors depend on. + +It reuses build_tog's traversal (`TogBuilder` / `_build`): loops, DMAs and +compute blocks are already identified there, each with a back-pointer to its +MLIR op(s), so this pass only adds the *rewrite*. Keeping a single traversal +guarantees the skeleton and the legacy TOG see the same structure. + +Counterpart to `build_tog.build_tog_and_mutate`. + +The DCE is safe by construction: it never erases an op whose results still have +uses, so at worst it leaves extra ops in the dump (visible for diagnosis) rather +than producing invalid IR. + +Requires the MLIR Python bindings (importing `build_tog` pulls in `mlir.ir`). +""" + +from . import togsim_ops as ts +from ._mlir_util import walk_ops, i32, i64, i64_array, str_attr +from .build_tog import ( + ir, + TogBuilder, + _build, + _reset_ids, + _find_kernel, + _value_key, + TOGDMANode, + TOGDMAWaitNode, + _COMPUTE_TYPE_NAME, +) + +#: Marker op names for the passes/__init__ fast-path (skip parsing if absent). +MARKERS = ("memref.dma_start", "memref.dma_wait") + +#: Ops the DCE must never remove (loops, terminators, our API ops). +_KEEP = { + "affine.for", "scf.for", "scf.while", + "affine.yield", "scf.yield", "func.return", + ts.DMA, ts.COMPUTE, ts.COMPUTE_BAR, ts.MEMORY_BAR, +} + + +def _kernel_block(module): + func_op = _find_kernel(module) + if func_op is None: + return None + return func_op.regions[0].blocks[0] + + +# --------------------------------------------------------------------------- +# op construction +# --------------------------------------------------------------------------- +def _arg_id_of(base_addr): + """Tensor func-arg ordinal from a build_tog base name ("arg3" -> 3); -1 if + it is not a plain block-arg base.""" + s = str(base_addr) + return int(s[3:]) if s.startswith("arg") and s[3:].isdigit() else -1 + + +def _emit_dma(ctx, dma_node, tag_id, dram_index, tag_index, read_bufs, write_bufs): + """Insert a `togsim.dma` before the original `memref.dma_start`. + + `tag_id` is the identity of this DMA's tag memref. An async DMA pairs with + its `togsim.memory_barrier` (the original dma_wait) by the RUNTIME tag slot + -- (tag_id, tag_index) -- not a compile-time identifier: one static dma op runs + once per loop iteration, each with a different runtime `%tag[%idx]` slot, so + only a runtime key can pair iteration i's dma with iteration i's wait. + + `dram_index` is the original linear DRAM index Value (the `affine.apply` + result that indexed the tensor in the `memref.dma_start`) -- carried as an + operand so the DCE keeps the address arithmetic live and the C4 lowering can + compute the real `base_addr = base[arg_id] + index*elem` (P3, approach A). + + `tag_index` is the original SRAM tag index Value (`%tag[%idx]`), carried as a + second operand: the runtime tag slot, used both to pair with the barrier and + for the double-buffer / SRAM-capacity (WAR) model. + Operand order: [dram_index, tag_index] (each omitted if absent).""" + op = dma_node.op + attrs = { + ts.ATTR_DIR: i32(ts.DIR_STORE if dma_node.is_write else ts.DIR_LOAD), + ts.ATTR_DIMS: i64_array(dma_node.tile_size), + ts.ATTR_STRIDES: i64_array(dma_node.tile_stride), + ts.ATTR_ELEM_BITS: i32(dma_node.element_size), + ts.ATTR_IS_ASYNC: ir.BoolAttr.get(bool(dma_node.is_async)), + ts.ATTR_TAG_ID: i32(tag_id), + ts.ATTR_ARG_ID: i32(_arg_id_of(dma_node.base_addr)), + "base": str_attr(dma_node.base_addr), + # SRAM spad this DMA touches (load writes it, store reads it) -- sec 10. + ts.ATTR_READ_BUFS: i64_array(read_bufs), + ts.ATTR_WRITE_BUFS: i64_array(write_bufs), + } + operands = [v for v in (dram_index, tag_index) if v is not None] + ir.Operation.create( + ts.DMA, + results=[], + operands=operands, + attributes=attrs, + loc=ir.Location.unknown(ctx), + ip=ir.InsertionPoint(op), + ) + + +def _emit_compute_bar(ctx, anchor_op): + """Insert a `togsim.compute_barrier` before `anchor_op` -- the fence that + drains in-flight async compute (the systolic-array matmuls) before a store + consumes their result (sec 10.7). + + FIXME: this is the one barrier still synthesized here rather than read from + the IR. Like the async-load memory barrier (now mapped 1:1 from the explicit + dma_wait), the compute fence should eventually appear explicitly in the input + MLIR and be mapped through, not auto-inserted -- no surprising insertion.""" + ir.Operation.create( + ts.COMPUTE_BAR, results=[], operands=[], attributes={}, + loc=ir.Location.unknown(ctx), ip=ir.InsertionPoint(anchor_op)) + + +def _emit_memory_bar(ctx, anchor_op, tag_id, tag_index, write_bufs): + """Insert a `togsim.memory_barrier` before `anchor_op` -- the explicit + async-DMA sync that was the original `memref.dma_wait`. It pairs with its + async `togsim.dma` by the RUNTIME tag slot (tag_id + tag_index), and carries + the SRAM buffer that dma loaded so consumers gate on data-arrival, not on the + async dma's issue-complete.""" + attrs = { + ts.ATTR_TAG_ID: i32(tag_id), + ts.ATTR_WRITE_BUFS: i64_array(write_bufs), + } + operands = [tag_index] if tag_index is not None else [] + ir.Operation.create( + ts.MEMORY_BAR, results=[], operands=operands, attributes=attrs, + loc=ir.Location.unknown(ctx), ip=ir.InsertionPoint(anchor_op)) + + +def _flatten_add(expr): + """Top-level additive summands of an AffineExpr (`.lhs`/`.rhs` come back typed + as the base AffineExpr, so use the `isinstance`/cast pattern, not Python + isinstance).""" + if ir.AffineAddExpr.isinstance(expr): + a = ir.AffineAddExpr(expr) + return _flatten_add(a.lhs) + _flatten_add(a.rhs) + return [expr] + + +def _neg_coeff_dim(summand): + """If `summand` is `dim * c` with a negative constant `c`, return that dim's + position; else None. lower_to_vcix tags each accumulation (reduction) loop var + with coefficient -1 in the dma_wait tag index -- a SENTINEL marking the + reduction axis, not an arithmetic offset (legacy TileGraphParser skips stride + -1 for the same reason).""" + if not ir.AffineMulExpr.isinstance(summand): + return None + mul = ir.AffineMulExpr(summand) + l, r = mul.lhs, mul.rhs + dim = l if ir.AffineDimExpr.isinstance(l) else (r if ir.AffineDimExpr.isinstance(r) else None) + con = l if ir.AffineConstantExpr.isinstance(l) else (r if ir.AffineConstantExpr.isinstance(r) else None) + if dim is None or con is None or ir.AffineConstantExpr(con).value >= 0: + return None + return ir.AffineDimExpr(dim).position + + +def _strip_accum_terms(ctx, tag_index, anchor_op): + """Return a tag-index Value with the accumulation-marked (-1 coefficient) terms + dropped, so a memory_barrier waits on the SAME subtile slot its async load + wrote. + + The wait tag index built by lower_to_vcix carries `-acc_iv` for each reduction + loop var; the matching load index (dma_fine_grained) is subtile-only. Without + this, at reduction iteration > 0 the producer EVALUATES `-acc_iv` to a negative + slot, so the recorded barrier slot diverges from the load slot and the runtime + tag pairing fails (TOGSim aborts with "Key does not exist in ... tag table"). + Dropping the -1 terms mirrors legacy TileGraphParser.cc, which skips stride -1 + and routes the reduction axis to a separate accum tag component; here the + per-iteration tag alloc (dma_fine_grained) already separates the reductions, so + the barrier only needs the subtile slot. + + Falls through (returns `tag_index` unchanged) for anything that is not an + affine.apply whose single result carries such a term -- e.g. the single-tile + case, whose index has no reduction term.""" + if tag_index is None: + return None + try: + apply_op = tag_index.owner + if apply_op.name != "affine.apply": + return tag_index + amap = ir.AffineMapAttr(apply_op.attributes["map"]).value + except Exception: + return tag_index + if amap.n_dims == 0 or amap.n_symbols != 0 or len(amap.results) != 1: + return tag_index + expr = amap.results[0] + dropped = sorted({p for p in (_neg_coeff_dim(s) for s in _flatten_add(expr)) + if p is not None}) + if not dropped: + return tag_index + n = amap.n_dims + kept = [i for i in range(n) if i not in dropped] + new_pos = {old: i for i, old in enumerate(kept)} + # compose the original expr with a selector that sends each dropped dim to 0 + # and renumbers the kept dims 0..k-1. + sel = [ir.AffineConstantExpr.get(0) if i in dropped + else ir.AffineDimExpr.get(new_pos[i]) for i in range(n)] + new_expr = expr.compose(ir.AffineMap.get(len(kept), 0, sel)) + new_map = ir.AffineMap.get(len(kept), 0, [new_expr]) + operands = list(apply_op.operands) + new_operands = [operands[i] for i in kept] + new_apply = ir.Operation.create( + "affine.apply", + results=[ir.IndexType.get(ctx)], + operands=new_operands, + attributes={"map": ir.AffineMapAttr.get(new_map)}, + loc=ir.Location.unknown(ctx), + ip=ir.InsertionPoint(anchor_op), + ) + return new_apply.results[0] + + +def _emit_compute(ctx, compute_node, tile_id, read_bufs, write_bufs): + front = compute_node.operations[0] + attrs = { + ts.ATTR_TILE_ID: i64(tile_id), + # int code (0 vector / 1 matmul / 2 preload) consumed by the C4 lowering; + # maps directly to the Core compute-unit enum. Keep the readable name too. + ts.ATTR_COMPUTE_TYPE: i32(int(compute_node.compute_type)), + "compute_type_name": str_attr(_COMPUTE_TYPE_NAME[compute_node.compute_type]), + # SRAM buffer ids read/written (sec 10 dataflow); the bridge builds the + # dependency DAG by last-writer per buffer. + ts.ATTR_READ_BUFS: i64_array(read_bufs), + ts.ATTR_WRITE_BUFS: i64_array(write_bufs), + } + ir.Operation.create( + ts.COMPUTE, + results=[], + operands=[], + attributes=attrs, + loc=ir.Location.unknown(ctx), + ip=ir.InsertionPoint(front), + ) + + +# --------------------------------------------------------------------------- +# DCE +# --------------------------------------------------------------------------- +def _has_nonempty_region(op): + for region in op.operation.regions: + for b in region.blocks: + if len(list(b.operations)) > 0: + return True + return False + + +def _results_unused(op): + for r in op.operation.results: + if len(list(r.uses)) > 0: + return False + return True + + +def _dce(block): + """Erase non-kept ops with no used results, to a fixed point. Safe: an op + with live SSA uses is never touched.""" + changed = True + while changed: + changed = False + victims = [] + for op in walk_ops(block): + name = op.operation.name + if name in _KEEP: + continue + if _has_nonempty_region(op): + continue + if _results_unused(op): + victims.append(op) + for op in victims: + try: + op.operation.erase() + changed = True + except Exception: + # Still referenced via something we will erase next round; retry. + pass + + +# --------------------------------------------------------------------------- +# driver +# --------------------------------------------------------------------------- +def _collect_dma_nodes(builder): + """Map op-identity -> DMA/DMAWait node, by walking the built tree.""" + by_op = {} + seen = set() + + def visit(n): + if id(n) in seen: + return + seen.add(id(n)) + if isinstance(n, (TOGDMANode, TOGDMAWaitNode)) and n.op is not None: + by_op[id(n.op.operation)] = n + for c in n.children: + visit(c) + + for ln in builder.loop_nodes: + visit(ln) + return by_op + + +class _BufferIds: + """Assigns each SRAM buffer name a stable small int id, shared by DMA and + compute so the bridge can match a reader to its buffer's writer (sec 10). + The virtual SA_WEIGHTS buffer (preload -> matmul) is numbered here too, on + first sight. `None` (a non-buffer base) is -1.""" + + def __init__(self): + self._ids = {} + + def of(self, name): + if name is None: + return -1 + return self._ids.setdefault(name, len(self._ids)) + + +class _TagIds: + """Identity of a DMA's tag memref -> stable small int, plus the SRAM buffer + that tag's async DMA loads. An async dma and its memory_barrier (the original + dma_wait) share a tag memref; this assigns it a tag_id (so the runtime can + pair them by the runtime tag slot) and remembers the loaded buffer so the + barrier can release it to consumers. Pairing is by tag, never a static id.""" + + def __init__(self): + self._ids = {} # tag value-key -> tag_id + self._buf = {} # tag value-key -> SRAM buffer id the dma loads + + def bind(self, key, buf): + tag_id = self._ids.setdefault(key, len(self._ids)) + self._buf[key] = buf + return tag_id + + def lookup(self, key): + """(tag_id, buffer) for a tag memref, or None if no dma used it.""" + if key not in self._ids: + return None + return self._ids[key], self._buf[key] + + +def _emit_computes(ctx, builder, bufs): + """Step 1: each compute node -> one togsim.compute carrying its tile_id and + the ids of the SRAM buffers it reads/writes. Returns the count.""" + from . import dep_analysis as dep # lazy: dep_analysis imports build_skeleton + n = 0 + for tile_id, cn in enumerate(builder.compute_nodes): + if not cn.operations: + continue + reads, writes = dep.compute_buffers(cn) + _emit_compute(ctx, cn, tile_id, + sorted(bufs.of(b) for b in reads), + sorted(bufs.of(b) for b in writes)) + n += 1 + return n + + +def _emit_one_dma(ctx, op, node, builder, bufs, tags): + """Rewrite one memref.dma_start as togsim.dma. A load reads DRAM and writes + its SRAM spad; a store reads the spad and writes DRAM -- which sets the + read/write buffer that drives the dependency edge (sec 10). The tag memref is + bound to a tag_id (with its loaded buffer) so the paired memory_barrier finds + it by the runtime tag slot.""" + from . import dep_analysis as dep # lazy: dep_analysis imports build_skeleton + f = builder._dma_start_fields(op) + dram_indices = f["dst_indices"] if node.is_write else f["src_indices"] + dram_index = dram_indices[0] if dram_indices else None + tag_indices = f["tag_indices"] + tag_index = tag_indices[0] if tag_indices else None + # the spad is the SRAM side of the copy: dst for a load, src for a store. + spad_id = bufs.of(dep._global_of(f["src"] if node.is_write else f["dst"])) + read_bufs = [spad_id] if node.is_write else [] + write_bufs = [] if node.is_write else [spad_id] + tag_id = tags.bind(_value_key(f["tag"]), spad_id) + if node.is_write: + _emit_compute_bar(ctx, op) # FIXME(sec10.7): auto-inserted; should be explicit in the IR. + _emit_dma(ctx, node, tag_id, dram_index, tag_index, read_bufs, write_bufs) + + +def _emit_one_wait(ctx, op, tags): + """Rewrite one memref.dma_wait as togsim.memory_barrier -- the explicit + async-DMA sync already in the IR. Paired with its dma by the tag memref + (tag_id) and the runtime tag index; carries the buffer the dma loaded. + Returns True iff emitted (a wait whose tag no dma used is dropped).""" + operands = list(op.operation.operands) + tag = operands[0] + tag_index = operands[1] if len(operands) >= 2 else None + binding = tags.lookup(_value_key(tag)) + if binding is None: + return False + tag_id, buf = binding + # honor lower_to_vcix's -1 accumulation marker: strip the reduction terms so + # the barrier slot equals the subtile slot the paired async load wrote. + tag_index = _strip_accum_terms(ctx, tag_index, op) + _emit_memory_bar(ctx, op, tag_id, tag_index, [buf]) + return True + + +def _emit_dmas_and_waits(ctx, block, builder, dma_by_op, bufs): + """Step 2: rewrite memref.dma_start -> togsim.dma and memref.dma_wait -> + togsim.memory_barrier in program order. An async dma and its barrier are + paired by the RUNTIME tag slot (tag_id + tag index), not a compile-time id: + one static dma op runs per loop iteration with a different `%tag[%idx]`, so + only the runtime slot can pair iteration i's dma with iteration i's wait. + Returns the original ops to erase and the (dma, wait) counts.""" + tags = _TagIds() + originals = [] + n_dma = n_wait = 0 + for op in list(walk_ops(block)): + name = op.operation.name + if name == "memref.dma_start": + node = dma_by_op.get(id(op.operation)) + if node is None: + continue + _emit_one_dma(ctx, op, node, builder, bufs, tags) + originals.append(op) + n_dma += 1 + elif name == "memref.dma_wait": + if _emit_one_wait(ctx, op, tags): + n_wait += 1 + originals.append(op) + return originals, n_dma, n_wait + + +def build_skeleton(module): + """Reduce `func.func @kernel` in `module` to the skeleton+API form, in place. + + Four steps: analyze the kernel into loop/compute/DMA nodes, emit a + togsim.compute per compute node, rewrite the DMAs/waits to togsim.dma/wait, + then DCE the leftover data computation. Returns a short text report (counts). + """ + _reset_ids() + builder = TogBuilder() + _build(module, builder) # populates loop/compute nodes + op back-pointers + + block = _kernel_block(module) + if block is None: + return "no @kernel found" + ctx = module.context + dma_by_op = _collect_dma_nodes(builder) + bufs = _BufferIds() + + n_compute = _emit_computes(ctx, builder, bufs) + originals, n_dma, n_wait = _emit_dmas_and_waits(ctx, block, builder, dma_by_op, bufs) + + # erase the now-replaced originals (result-less -> safe), then strip the + # leftover data computation. + for op in originals: + try: + op.operation.erase() + except Exception: + pass + _dce(block) + + return ("skeleton: compute=%d dma=%d wait=%d (unpaired waits dropped)" + % (n_compute, n_dma, n_wait)) + + +def run(module, vectorlane=128): + """passes/__init__ pass protocol entry (vectorlane unused; kept for parity).""" + build_skeleton(module) + + +def run_skeleton(in_path, out_path=None): + """Read post-vcix MLIR at `in_path`, reduce to skeleton+API, write it out. + + Requires the MLIR bindings. + """ + if out_path is None: + out_path = in_path + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(open(in_path).read(), ctx) + report = build_skeleton(module) + with open(out_path, "w") as fh: + fh.write(str(module)) + return report + + +def main(argv): + import argparse + + parser = argparse.ArgumentParser(prog="build_skeleton.py") + parser.add_argument("input") + parser.add_argument("--out", default=None) + args = parser.parse_args(argv[1:]) + report = run_skeleton(args.input, args.out) + import sys + sys.stderr.write(report + "\n") + return 0 + + +if __name__ == "__main__": + import sys + sys.exit(main(sys.argv)) diff --git a/PyTorchSimFrontend/mlir/passes/cycle_table.py b/PyTorchSimFrontend/mlir/passes/cycle_table.py new file mode 100644 index 00000000..5cb35801 --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/cycle_table.py @@ -0,0 +1,102 @@ +"""cycle_table: the precomputed tile_id -> (cycle, overlapping_cycle) table the +C++ trace pipeline looks up at runtime (docs/design/togsim_cpp_trace.md sec 6). + +A `togsim.compute(tile_id=...)` in the trace says *which* tile to compute, not +how long it takes. Because tiles are fixed size, each tile's cost is invariant +(only the trip count varies with shape), so it is sampled once and stored here, +keyed by `tile_id`. Two numbers per tile, mirroring the legacy TOG: + + * `cycle` -- full compute latency, sampled by gem5 sample-mode + (the existing measurement: `_rewrite_loop_steps` + + `_insert_compute_markers` in build_tog, run through + CycleSimulator -> the per-tile `cycle_list`). + * `overlapping_cycle` -- the portion that overlaps the previous instruction in + the systolic pipeline; the timing core uses it as + `finish = prev.finish + cycle - overlapped` (Core.cc). + Derived exactly as the legacy path does + (tog_generator.generate_tile_graph): + type 0 (VectorCompute) -> 0 + type 1 (MatmulCompute) -> max(cycle - x_offset, 0) + type 2 (MatmulPreload) -> max(cycle - w_offset, 0) + +This module only *builds/serializes* the table from a cycle_list; obtaining the +cycle_list reuses the existing sample-mode + gem5 path. The +`tile_id` order matches build_skeleton's `compute_nodes` order, which matches the +legacy TOG, so the same sampling keys both paths. + +Requires the MLIR Python bindings (to read the skeleton's togsim.compute ops). +""" + +import json + +from . import togsim_ops as ts +from ._mlir_util import walk_ops +from .build_tog import ( + ir, + VECTOR_COMPUTE, + MATMUL_COMPUTE, # noqa: F401 (documents the type enum used by the formula) + MATMUL_PRELOAD, +) + + +def overlapping_cycle(cycle, compute_type, x_offset, w_offset): + """Hideable (pipeline-overlapped) portion of `cycle`. Mirrors + tog_generator.generate_tile_graph.""" + if compute_type <= VECTOR_COMPUTE: # VectorCompute: no systolic overlap + return 0 + offset = w_offset if compute_type == MATMUL_PRELOAD else x_offset + return max(int(cycle) - int(offset), 0) + + +def _compute_types(skeleton_module): + """tile_id-ordered list of compute_type ints, from the skeleton's + togsim.compute ops.""" + items = [] + for op in walk_ops(skeleton_module.body): + if op.operation.name != ts.COMPUTE: + continue + tid = ir.IntegerAttr(op.operation.attributes[ts.ATTR_TILE_ID]).value + ct = ir.IntegerAttr(op.operation.attributes[ts.ATTR_COMPUTE_TYPE]).value + items.append((tid, ct)) + items.sort() + return [t for _, t in items] + + +def build_cycle_table(skeleton_module, cycle_list, x_offset, w_offset): + """Return `[(cycle, overlapping_cycle), ...]` indexed by tile_id. + + `cycle_list` is the per-tile gem5 measurement (compute_nodes order == + tile_id order). `x_offset`/`w_offset` are the systolic-fill offsets the + legacy path computes from the vector-lane size / loop size.""" + types = _compute_types(skeleton_module) + if len(cycle_list) != len(types): + raise ValueError( + "cycle_list (%d) does not match #compute tiles (%d)" + % (len(cycle_list), len(types))) + return [(int(c), overlapping_cycle(c, t, x_offset, w_offset)) + for c, t in zip(cycle_list, types)] + + +def dump_cycle_table(table, path, x_offset=None, w_offset=None): + """Serialize the table as a sidecar JSON next to the trace `.so`. The P3 C6 + loader reads it and sets compute_cycle + overlapping_cycle on each emitted + Instruction.""" + with open(path, "w") as fh: + json.dump({"x_offset": x_offset, "w_offset": w_offset, + "table": [list(e) for e in table]}, fh) + return path + + +def load_cycle_table(path): + with open(path) as fh: + return json.load(fh) + + +def dump_cycle_table_tsv(table, path): + """Plain `cycleoverlapping` per line, in tile_id order -- the trivial + format the C++ `--cycle_table` loader (main.cc, P3 trace pipeline) reads with + ifstream (no JSON dependency in TOGSim).""" + with open(path, "w") as fh: + for cycle, overlapping in table: + fh.write("%d\t%d\n" % (int(cycle), int(overlapping))) + return path diff --git a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py index c0e82b66..10b2edfb 100644 --- a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py +++ b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py @@ -32,13 +32,7 @@ OP_NAME = "togsim.transfer" MARKERS = (OP_NAME,) - -def _iter_ops(block): - for op in list(block.operations): - yield op - for region in op.operation.regions: - for b in region.blocks: - yield from _iter_ops(b) +from ._mlir_util import walk_ops def _int_array(attr): @@ -92,7 +86,7 @@ def run(module, vectorlane=128, **_): targets = [] for region in module.operation.regions: for b in region.blocks: - for op in _iter_ops(b): + for op in walk_ops(b): if op.operation.name == OP_NAME: targets.append(op.operation) diff --git a/PyTorchSimFrontend/mlir/passes/dep_analysis.py b/PyTorchSimFrontend/mlir/passes/dep_analysis.py new file mode 100644 index 00000000..fa4efc8a --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/dep_analysis.py @@ -0,0 +1,194 @@ +"""dep_analysis.py -- dependency-edge analysis for the C++ trace pipeline (sec 10). + +The current TOG pass does NO dependency analysis (it emits a lexical loop tree + +runtime tags). This module derives the producer->consumer edges that the explicit +dataflow trace needs, from two sources available on the post-vcix IR (before +build_skeleton collapses the compute regions): + + 1. SRAM access: each DMA/compute's read/write SRAM buffer(s), recovered by + following SSA (a vcix.iv's input vector -> its vector.transfer_read -> the + memref -> @global), and the DMA's spad operand. Edge: a reader depends on + the last node that wrote the same buffer. + 2. vcix preload/matmul pairing: a matmul (vcix opcode 0) consumes the weights a + preceding preload (opcode 1) loaded into the systolic array -- an SA-internal + dependency NOT visible as a memref access, so it comes from the opcode order. + +This is a node-level analysis (one node per build_tog compute/DMA node); the loops +replay the nodes, so loop-carried edges (the Y_spad accumulator) are materialized +per iteration downstream. First cut: buffer granularity (slot-level value matching +is a later refinement). Output is an edge list for validation / to drive emit. +""" +import sys +import os + +from .build_tog import TogBuilder, ir, _reset_ids +from . import build_skeleton as _bs + + +def _global_of(memref_val): + """memref SSA value -> @global symbol name (e.g. 'X_spad'), or None.""" + owner = memref_val.owner + op = owner if isinstance(owner, ir.Operation) else getattr(owner, "operation", None) + if op is None: + return None + if op.name == "memref.get_global": + return str(op.attributes["name"]).strip('@" ') + # walk through view-like ops (subview/cast) to their source + if op.operands: + try: + return _global_of(op.operands[0]) + except Exception: + return None + return None + + +def _read_buffers_of_compute(cn): + """SRAM buffers a compute node reads: (a) each vcix.iv input traced to its + vector.transfer_read source (activations/weights streamed into the SA), and + (b) any direct vector.transfer_read in the node (the epilogue's accumulator + read-modify-write of Y_spad).""" + bufs = set() + for op in cn.operations: + if op.name == "vector.transfer_read" and list(op.operands): + b = _global_of(op.operands[0]) + if b: + bufs.add(b) + elif op.name == "vcix.iv" and list(op.operands): + v = op.operands[0] + defop = v.owner if isinstance(v.owner, ir.Operation) else getattr(v.owner, "operation", None) + if defop is not None and defop.name == "vector.transfer_read" and list(defop.operands): + b = _global_of(defop.operands[0]) + if b: + bufs.add(b) + return bufs + + +def _write_buffers_of_compute(cn): + """SRAM buffers a compute node writes: vector.transfer_write / vector_store target.""" + bufs = set() + for op in cn.operations: + if op.name in ("vector.transfer_write", "affine.vector_store", "vector.store"): + # target memref is the last memref operand + for v in op.operands: + try: + if ir.MemRefType.isinstance(v.type): + b = _global_of(v) + if b: + bufs.add(b) + except Exception: + pass + return bufs + + +def _dma_buffer(builder, dma_node): + """The SRAM spad buffer a DMA touches (dst for load, src for store).""" + try: + f = builder._dma_start_fields(dma_node.op) + except Exception: + return None + val = f["dst"] if not dma_node.is_write else f["src"] + return _global_of(val) + + +# Virtual buffer for the systolic-array weight registers: a preload writes it, +# the following matmul reads it. This folds the SA-internal preload->matmul +# dependency (not a memref access) into the uniform "last-writer per buffer" rule. +SA_WEIGHTS = "__SA_WEIGHTS__" + + +def compute_buffers(cn): + """(read_buffers, write_buffers) for one compute node, including the virtual + SA_WEIGHTS edge (preload writes it, matmul reads it).""" + reads = set(_read_buffers_of_compute(cn)) + writes = set(_write_buffers_of_compute(cn)) + if cn.compute_type == 1: # MATMUL consumes the preloaded weights + reads.add(SA_WEIGHTS) + elif cn.compute_type == 2: # PRELOAD loads them + writes.add(SA_WEIGHTS) + return reads, writes + + +def analyze(module): + """Return (nodes, edges). nodes: list of dicts; edges: list of (consumer_idx, + producer_idx, reason).""" + _reset_ids() + builder = TogBuilder() + _bs._build(module, builder) + + nodes = [] + # DMA nodes only (the map also contains TOGDMAWaitNode; keep real DMAs). + dma_nodes = [dn for dn in dict.fromkeys(_bs._collect_dma_nodes(builder).values()) + if hasattr(dn, "is_write")] + for dn in dma_nodes: + buf = _dma_buffer(builder, dn) + nodes.append({ + "kind": "STORE" if dn.is_write else "LOAD", + "buf": buf, "arg": str(dn.base_addr), + "reads": {buf} if dn.is_write else set(), + "writes": {buf} if not dn.is_write else set(), + "node": dn, + }) + for cn in builder.compute_nodes: + if not cn.operations: + continue + ct = {0: "VECTOR", 1: "MATMUL", 2: "PRELOAD"}.get(cn.compute_type, f"c{cn.compute_type}") + nodes.append({ + "kind": ct, + "reads": _read_buffers_of_compute(cn), + "writes": _write_buffers_of_compute(cn), + "node": cn, + "compute_type": cn.compute_type, + }) + + # Order nodes by program position (last-writer needs program order: e.g. the + # store reads Y_spad written by the matmul, which lexically precedes it). + pos = {} + idx = [0] + def _index(op): + pos[op] = idx[0]; idx[0] += 1 + for r in op.regions: + for b in r.blocks: + for o in b.operations: + _index(o) + _index(module.operation) + def _key(n): + node = n["node"] + op = getattr(node, "op", None) or (node.operations[0] if getattr(node, "operations", None) else None) + return pos.get(op, 1 << 30) + nodes.sort(key=_key) + + # Edges: (1) buffer last-writer, (2) preload->matmul. + edges = [] + last_writer = {} # buffer -> node idx + prev_preload = None + for i, n in enumerate(nodes): + for b in sorted(n["reads"]): + if b in last_writer: + edges.append((i, last_writer[b], f"reads {b}")) + if n["kind"] == "MATMUL" and prev_preload is not None: + edges.append((i, prev_preload, "uses preloaded weights (vcix op1->op0)")) + for b in n["writes"]: + last_writer[b] = i + if n["kind"] == "PRELOAD": + prev_preload = i + return nodes, edges + + +def _main(): + path = sys.argv[1] + ctx = ir.Context(); ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(open(path).read(), ctx) + nodes, edges = analyze(module) + print("=== nodes ===") + for i, n in enumerate(nodes): + r = ",".join(sorted(n["reads"])) or "-" + w = ",".join(sorted(n["writes"])) or "-" + print(f" #{i:<2} {n['kind']:<8} reads[{r}] writes[{w}]") + print("=== edges (consumer -> producer) ===") + for c, p, why in edges: + print(f" #{c} ({nodes[c]['kind']}) -> #{p} ({nodes[p]['kind']}) [{why}]") + + +if __name__ == "__main__": + _main() diff --git a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py index 3f583ef2..f1872dca 100644 --- a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py +++ b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py @@ -30,6 +30,8 @@ import mlir.ir as ir # noqa: E402 +from ._mlir_util import walk_ops, attr_i64_array + MARKERS = ("subtile_size",) # only subtile DMAs are split MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 @@ -54,12 +56,6 @@ def _const_int(value, default=-1): return default -def _int_array_attr(op, key): - if key not in op.attributes: - return [] - return [ir.IntegerAttr(a).value for a in ir.ArrayAttr(op.attributes[key])] - - def _is_block_arg(v): return isinstance(v, ir.BlockArgument) @@ -106,13 +102,13 @@ def tile_shape(self): return list(mt.shape) def subtile_size(self): - return _int_array_attr(self.op, "subtile_size") + return attr_i64_array(self.op, "subtile_size", default=[]) def sram_stride(self): - return _int_array_attr(self.op, "sram_stride") + return attr_i64_array(self.op, "sram_stride", default=[]) def dram_stride(self): - return _int_array_attr(self.op, "dram_stride") + return attr_i64_array(self.op, "dram_stride", default=[]) def is_async(self): a = self.op.attributes @@ -244,6 +240,27 @@ def _const_index(v, ip): ir.IntegerAttr.get(ir.IndexType.get(), v), ip=ip).result +def _fresh_tag(dma): + """Give this DMA a fresh tag memref.alloc right BEFORE the (pre-split) coarse + dma_start, and rewire every use of the old tag -- the dma_start re-emitted + below AND its dma_wait -- to it. The coarse dma sits at the reduction-loop body + level (it has not been wrapped in a subtile load nest yet), so the alloc there + dominates both the load nest fine-grained is about to build and the sibling + wait nest. Each reduction iteration thus allocates its own tag -> successive + iterations are distinct (multi-tile-K / conv) and the per-iteration tag + semantics is in the IR, not reconstructed downstream. Old alloc becomes dead.""" + old = dma.tag + new_tag = ir.Operation.create("memref.alloc", results=[old.type], + operands=[], ip=ir.InsertionPoint(dma.op)).results[0] + old.replace_all_uses_with(new_tag) + dma.tag = new_tag + # the old (func-entry, per-tensor unique) alloc is now dead -- erase it. + try: + old.owner.erase() + except Exception: + pass + + # --------------------------------------------------------------------------- # Loop-nest construction # --------------------------------------------------------------------------- @@ -293,20 +310,12 @@ def _reaches(value, target): # --------------------------------------------------------------------------- # Pass driver # --------------------------------------------------------------------------- -def _iter_ops(block): - for op in list(block.operations): - yield op - for region in op.operation.regions: - for b in region.blocks: - yield from _iter_ops(b) - - def _run_func(func, vectorlane): from mlir.dialects import linalg # First matmul only. matmul = None dmas = [] - for op in _iter_ops(func.regions[0].blocks[0]): + for op in walk_ops(func.regions[0].blocks[0]): name = op.operation.name if name == "linalg.matmul" and matmul is None: matmul = op @@ -363,6 +372,12 @@ def _run_func(func, vectorlane): for d, f in enumerate(fuse["w_to_fused"]): bounds[f] = w_counts[d] + # Give each load a fresh per-iteration tag alloc just before its coarse dma + # (rewiring its dma_wait via the old tag's uses), so the tag is distinct per + # reduction iteration -- positioned to match the per-iteration tag semantics. + _fresh_tag(mvin_input) + _fresh_tag(mvin_weight) + # Insert the fused nest at the weight DMA (the later of the two): both DMAs' # original DRAM base indices (src_idx[0], computed in the enclosing loops) must # dominate the nest. Codegen emits input before weight, matching the C++ pass diff --git a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py index f5b841bb..998a6db5 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py @@ -22,6 +22,8 @@ WAIT_NAME = "memref.dma_wait" MARKERS = (OP_NAME, WAIT_NAME) +from ._mlir_util import attr_i64_array + # func7 instruction codes (CustomDMAAttribute.h) CONFIG, CONFIG2, CONFIG3, CONFIG4 = 0, 4, 5, 6 MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 @@ -124,8 +126,8 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): tile_shape = _subtile(op) if tile_shape is None: tile_shape = list(dst_ty.shape) if is_mvin else list(src_ty.shape) - dram_strides = _int_array(op, "dram_stride") - spad_strides = _int_array(op, "sram_stride") + dram_strides = attr_i64_array(op, "dram_stride") + spad_strides = attr_i64_array(op, "sram_stride") assert len(tile_shape) == len(dram_strides) == len(spad_strides), \ f"shape/stride rank mismatch: {tile_shape} {dram_strides} {spad_strides}" @@ -180,11 +182,6 @@ def _subtile(op): return [IntegerAttr(a).value for a in ArrayAttr(op.attributes["subtile_size"])] -def _int_array(op, name): - from mlir.ir import ArrayAttr, IntegerAttr - return [IntegerAttr(a).value for a in ArrayAttr(op.attributes[name])] - - def _elem_bytes(elem_type): from mlir.ir import IntegerType, FloatType bits = (IntegerType(elem_type).width if IntegerType.isinstance(elem_type) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py new file mode 100644 index 00000000..caca822e --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -0,0 +1,448 @@ +"""lower_to_emitc pass (C4): skeleton+API MLIR -> EmitC -> C++ -> trace `.so`. + +Second stage of the C++ trace pipeline (docs/design/togsim_cpp_trace.md, sec +5-7). Takes the skeleton+API module from `build_skeleton` (loop nest + +`togsim.*` ops) and produces an EmitC module whose single entry function + + extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) + +mirrors the loop skeleton, with every `togsim.*` op as an `emitc.call_opaque` +to the matching `togsim_runtime.h` free function (`togsim_ops.EMITC_CALLEE`). +`mlir-translate --mlir-to-cpp` renders it to C++, compiled to a `.so` that +exports `togsim_kernel` and leaves `togsim_dma/wait/compute/signal` undefined for +the TOGSim loader to resolve at `dlopen`. + +How the lowering is done -- it drives the *upstream* EmitC conversion passes and +adds only the glue they cannot do: + + 1. (python) Rewrite the unregistered `togsim.*` ops to `emitc.call_opaque`. + Unregistered ops have no registered conversion patterns, so this must be a + custom rewrite (design sec 8). Also rewrite the kernel's signature to the + ABI form (drop the memref tensor args -- the trace producer never touches + tensor data; base addresses are deferred to P3) and drop the aux + globals / wrapper func. + 2. (upstream passes, in-process PassManager) + func.func(lower-affine) -> convert-scf-to-emitc + -> convert-arith-to-emitc -> convert-func-to-emitc + This is the EmitC infrastructure: it lowers the affine/scf loop nest to + `emitc.for`, the index/arith (loop bounds, and in P3 the address + arithmetic) to EmitC, and the func to `emitc.func`. + 3. (python) Two small fixups the passes leave behind in this LLVM 20 build: + * `convert-scf-to-emitc` emits `emitc.for` with `index`-typed bounds, so + `convert-arith-to-emitc` (which makes constants `!emitc.size_t`) leaves + `builtin.unrealized_conversion_cast` on the bounds that nothing folds + and `mlir-to-cpp` cannot print (design sec 8 "EmitC coverage" risk). + `_fold_for_bound_casts` rewrites those casts away. + * add the `extern "C"` specifier so `dlsym` finds the entry unmangled. + +Requires the MLIR Python bindings (incl. `mlir.passmanager`); the .cpp/.so +steps additionally require `mlir-translate` (TORCHSIM_LLVM_PATH) and a host C++ +compiler. +""" + +import os +import subprocess + +from mlir.passmanager import PassManager + +from . import togsim_ops as ts +from ._mlir_util import walk_ops, i32, i64, attr_int, attr_i64_array +from .build_tog import ir, _find_kernel + +#: emitted entry symbol (== ts.ENTRY_SYMBOL == "togsim_kernel"). +ENTRY = ts.ENTRY_SYMBOL + +#: EmitC type of the opaque EmitCtx* threaded through every call. +CTX_TYPE = '!emitc.ptr>' + +#: upstream EmitC conversion pipeline (the infrastructure this pass drives). +_PIPELINE = ("builtin.module(" + "func.func(lower-affine)," + "convert-scf-to-emitc," + "convert-arith-to-emitc," + "convert-func-to-emitc)") + +#: prepended to the mlir-to-cpp output; pulls in size_t/intN_t and the ABI. +_PRELUDE = ( + "#include \n" + "#include \n" + "using std::size_t;\n" + '#include "togsim_runtime.h"\n' +) + + +# --------------------------------------------------------------------------- +# attribute builders / readers +# --------------------------------------------------------------------------- +def _idx(v): + return ir.IntegerAttr.get(ir.IndexType.get(), int(v)) + + +def _opaque(ctx, text): + return ir.Attribute.parse('#emitc.opaque<"%s">' % text, ctx) + + +def _arr(ctx, vals): + """A C compound-literal `(const int64_t[]){...}` arg, or `nullptr` if empty + (the call site decays it to a `const int64_t*`).""" + vals = list(vals) + if not vals: + return _opaque(ctx, "nullptr") + return _opaque(ctx, "(const int64_t[]){%s}" % ", ".join(str(int(v)) for v in vals)) + + +def _attr_bool(op, key): + return 1 if ir.BoolAttr(op.operation.attributes[key]).value else 0 + + +# --------------------------------------------------------------------------- +# step 1: rewrite signature + togsim.* ops (the unregistered-op glue) +# --------------------------------------------------------------------------- +def _strip_aux(module): + """Erase memref.global decls and every func except @kernel (the wrapper).""" + victims = [] + for op in module.body.operations: + name = op.operation.name + if name == "memref.global": + victims.append(op) + elif name == "func.func": + if ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": + victims.append(op) + for op in victims: + op.operation.erase() + + +def _rewrite_signature(kernel, ctx): + """Replace @kernel's memref tensor args with the ABI args + (EmitCtx*, int64_t* shape_args, int32_t n) and rename it to togsim_kernel. + Returns the ctx Value.""" + block = kernel.regions[0].blocks[0] + for arg in block.arguments: + if len(list(arg.uses)) > 0: + raise ValueError( + "kernel arg still used after build_skeleton; cannot drop it " + "(expected the DCE to have removed all tensor-data ops)") + # erase existing (memref) args high-to-low, then append the ABI args. + for i in reversed(range(len(block.arguments))): + block.erase_argument(i) + ptr = ir.Type.parse(CTX_TYPE, ctx) + i64ptr = ir.Type.parse("!emitc.ptr", ctx) + i32 = ir.IntegerType.get_signless(32) + loc = ir.Location.unknown(ctx) + block.add_argument(ptr, loc) + block.add_argument(i64ptr, loc) + block.add_argument(i32, loc) + kernel.operation.attributes["function_type"] = ir.TypeAttr.get( + ir.FunctionType.get([ptr, i64ptr, i32], [])) + kernel.operation.attributes["sym_name"] = ir.StringAttr.get(ENTRY) + return block.arguments[0] + + +def _call(ctx, ctx_val, op, callee, arg_attrs): + """Insert emitc.call_opaque (ctx) {args=[0:index, ...]} before `op`. + The leading `0 : index` references operand 0 (ctx); other entries are + literal C args (integer attr -> literal, #emitc.opaque -> verbatim).""" + ir.Operation.create( + "emitc.call_opaque", results=[], operands=[ctx_val], + attributes={"callee": ir.StringAttr.get(callee), + "args": ir.ArrayAttr.get([_idx(0)] + arg_attrs)}, + loc=ir.Location.unknown(ctx), ip=ir.InsertionPoint(op)) + + +def _innermost_outer_loop(block): + """Deepest `affine.for {outer_loop=true}` (the PARALLEL/ACCUMULATION + boundary). Returns the op or None if the kernel has no parallel loop.""" + found = [None] + + def is_outer(op): + a = op.operation.attributes + return "outer_loop" in a and ir.BoolAttr(a["outer_loop"]).value + + def walk(b): + for op in b.operations: + if op.operation.name == "affine.for" and is_outer(op): + found[0] = op # nested outer loops overwrite -> deepest wins + for r in op.operation.regions: + for bb in r.blocks: + walk(bb) + + walk(block) + return found[0] + + +def _insert_core_alloc(ctx, kernel, ctx_val): + """Insert `togsim_core_alloc(ctx)` at the start of each parallel work-item: + the first op of the innermost PARALLEL loop body (or the kernel entry if the + kernel has no parallel loop -> a single work-item). The runtime binds the + following ops to the returned core (sec 9.3); the producer never names + num_cores. The return value is discarded (no free -- a core is an assignment, + not a held resource).""" + block = kernel.regions[0].blocks[0] + target = _innermost_outer_loop(block) + body = target.operation.regions[0].blocks[0] if target is not None else block + ir.Operation.create( + "emitc.call_opaque", results=[], operands=[ctx_val], + attributes={"callee": ir.StringAttr.get(ts.CORE_ALLOC_CALLEE), + "args": ir.ArrayAttr.get([_idx(0)])}, + loc=ir.Location.unknown(ctx), + ip=ir.InsertionPoint.at_block_begin(body)) + + +def _rewrite_togsim_ops(ctx, kernel, ctx_val): + block = kernel.regions[0].blocks[0] + victims = [] + for op in walk_ops(block): + name = op.operation.name + ipo = ir.InsertionPoint(op) + if name == ts.DMA: + dims = attr_i64_array(op, ts.ATTR_DIMS) + # build_skeleton carries [dram_index, tag_index] (each optional) as + # operands. Pass each as a call operand so convert-arith-to-emitc lowers + # the address arithmetic into the producer; the runtime adds the base. + ins = list(op.operation.operands) + dram_operand = ins[0] if len(ins) >= 1 else None + tag_operand = ins[1] if len(ins) >= 2 else None + operands = [ctx_val] + offset_arg = i64(0) + tag_arg = i64(0) + if dram_operand is not None: + operands.append(dram_operand) + offset_arg = _idx(len(operands) - 1) + if tag_operand is not None: + operands.append(tag_operand) + tag_arg = _idx(len(operands) - 1) + args = [_idx(0), + i32(attr_int(op, ts.ATTR_DIR)), + i32(attr_int(op, ts.ATTR_ARG_ID)), + offset_arg, + i32(len(dims)), + _arr(ctx, dims), + _arr(ctx, attr_i64_array(op, ts.ATTR_STRIDES)), + i32(attr_int(op, ts.ATTR_ELEM_BITS)), + i32(_attr_bool(op, ts.ATTR_IS_ASYNC)), + i32(attr_int(op, ts.ATTR_TAG_ID)), + tag_arg] + _rb = attr_i64_array(op, ts.ATTR_READ_BUFS) + _wb = attr_i64_array(op, ts.ATTR_WRITE_BUFS) + args += [_arr(ctx, _rb), i32(len(_rb)), _arr(ctx, _wb), i32(len(_wb))] + # togsim_dma is void: the dma is paired with its barrier by the runtime + # (tag_id, tag_slot), not a returned handle. + ir.Operation.create( + "emitc.call_opaque", results=[], operands=operands, + attributes={"callee": ir.StringAttr.get(ts.EMITC_CALLEE[ts.DMA]), + "args": ir.ArrayAttr.get(args)}, + loc=ir.Location.unknown(ctx), ip=ipo) + victims.append(op) + elif name == ts.MEMORY_BAR: + # explicit async-DMA sync (the original dma_wait) -> + # togsim_memory_barrier(ctx, tag_id, tag_slot, write_bufs). The tag + # index operand (if any) is the runtime tag slot. + ins = list(op.operation.operands) + operands = [ctx_val] + tag_arg = i64(0) + if ins: + operands.append(ins[0]) + tag_arg = _idx(len(operands) - 1) + _wb = attr_i64_array(op, ts.ATTR_WRITE_BUFS) + ir.Operation.create( + "emitc.call_opaque", results=[], operands=operands, + attributes={"callee": ir.StringAttr.get(ts.EMITC_CALLEE[ts.MEMORY_BAR]), + "args": ir.ArrayAttr.get( + [_idx(0), i32(attr_int(op, ts.ATTR_TAG_ID)), tag_arg, + _arr(ctx, _wb), i32(len(_wb))])}, + loc=ir.Location.unknown(ctx), ip=ipo) + victims.append(op) + elif name == ts.COMPUTE: + # skeleton compute carries no dims (cost is keyed by tile_id) -> 0/null. + _rb = attr_i64_array(op, ts.ATTR_READ_BUFS) + _wb = attr_i64_array(op, ts.ATTR_WRITE_BUFS) + _call(ctx, ctx_val, op, ts.EMITC_CALLEE[ts.COMPUTE], + [i64(attr_int(op, ts.ATTR_TILE_ID)), + i32(attr_int(op, ts.ATTR_COMPUTE_TYPE)), + i32(0), _opaque(ctx, "nullptr"), + _arr(ctx, _rb), i32(len(_rb)), _arr(ctx, _wb), i32(len(_wb))]) + victims.append(op) + elif name == ts.COMPUTE_BAR: + # explicit compute fence -> togsim_compute_barrier(ctx) (sec 10.7). + ir.Operation.create( + "emitc.call_opaque", results=[], operands=[ctx_val], + attributes={"callee": ir.StringAttr.get(ts.EMITC_CALLEE[ts.COMPUTE_BAR]), + "args": ir.ArrayAttr.get([_idx(0)])}, + loc=ir.Location.unknown(ctx), ip=ipo) + victims.append(op) + for op in victims: + op.operation.erase() + + +# --------------------------------------------------------------------------- +# step 3: post-conversion fixups +# --------------------------------------------------------------------------- +def _retype_for_to_size_t(module): + """Make every `emitc.for` use `!emitc.size_t` bounds + induction variable, + then drop the `index`<->`!emitc.size_t` `unrealized_conversion_cast` ops that + `convert-scf-to-emitc` / `convert-arith-to-emitc` leave behind (mlir-to-cpp + cannot print them; --reconcile cannot fold them). + + `emitc.for` accepts `size_t` bounds with the explicit type, and a `size_t` IV + makes the lowered address arithmetic (`convert-arith-to-emitc`, which works + in `size_t`) cast-free. So: set each IV to size_t, then for every + index<->size_t cast replace its result with its source (every consumer here + -- `emitc.for` bounds, `emitc.call_opaque` operands, `emitc` arith -- accepts + either, and after the IV retype each such cast bridges equal types).""" + idx = ir.IndexType.get() + st = ir.Type.parse("!emitc.size_t", module.context) + + for op in list(walk_ops(module.body)): + if op.operation.name == "emitc.for": + op.operation.regions[0].blocks[0].arguments[0].set_type(st) + + dead = [] + for op in list(walk_ops(module.body)): + if op.operation.name != "builtin.unrealized_conversion_cast": + continue + res = op.results[0] + src = list(op.operation.operands)[0] + # idx<->size_t bridges (incl. the size_t->size_t identities left after + # the IV retype): every consumer here accepts either, so fold to source. + if src.type in (idx, st) and res.type in (idx, st): + res.replace_all_uses_with(src) + dead.append(op) + for d in dead: + try: + d.operation.erase() + except Exception: + pass + + +def _add_extern_c(module, ctx): + for op in module.body.operations: + if (op.operation.name == "emitc.func" + and ir.StringAttr(op.operation.attributes["sym_name"]).value == ENTRY): + op.operation.attributes["specifiers"] = ir.ArrayAttr.get( + [ir.StringAttr.get('extern "C"')]) + return + raise ValueError("emitc.func @%s not found after conversion" % ENTRY) + + +# --------------------------------------------------------------------------- +# driver +# --------------------------------------------------------------------------- +def lower_to_emitc(skeleton_module): + """Lower a skeleton+API module (in place) to an EmitC module with the + `togsim_kernel` entry function. Returns the same module.""" + ctx = skeleton_module.context + kernel = _find_kernel(skeleton_module) + if kernel is None: + raise ValueError("no @kernel found in skeleton module") + + _strip_aux(skeleton_module) + ctx_val = _rewrite_signature(kernel, ctx) + _insert_core_alloc(ctx, kernel, ctx_val) # core_alloc per work-item + _rewrite_togsim_ops(ctx, kernel, ctx_val) + + PassManager.parse(_PIPELINE, ctx).run(skeleton_module.operation) + + _retype_for_to_size_t(skeleton_module) + _add_extern_c(skeleton_module, ctx) + return skeleton_module + + +# --------------------------------------------------------------------------- +# C++ / .so backend +# --------------------------------------------------------------------------- +def _mlir_translate_bin(): + return os.path.join(os.environ.get("TORCHSIM_LLVM_PATH", "/usr/bin"), + "mlir-translate") + + +def emitc_to_cpp(emitc_module, mlir_translate=None): + """Render `emitc_module` to C++ source (prelude + mlir-to-cpp body).""" + mlir_translate = mlir_translate or _mlir_translate_bin() + proc = subprocess.run( + [mlir_translate, "--mlir-to-cpp"], + input=str(emitc_module), capture_output=True, text=True) + if proc.returncode != 0: + raise RuntimeError("mlir-translate --mlir-to-cpp failed:\n" + proc.stderr) + return _PRELUDE + proc.stdout + + +def compile_so(cpp_text, so_path, include_dir, cxx=None): + """Compile producer C++ to `so_path`. `include_dir` must hold + togsim_runtime.h. togsim_* symbols are left undefined (resolved at dlopen).""" + cxx = cxx or os.environ.get("CXX", "g++") + cpp_path = os.path.splitext(so_path)[0] + ".cpp" + with open(cpp_path, "w") as fh: + fh.write(cpp_text) + proc = subprocess.run( + [cxx, "-shared", "-fPIC", "-std=gnu++17", "-O2", + "-I", include_dir, cpp_path, "-o", so_path], + capture_output=True, text=True) + if proc.returncode != 0: + raise RuntimeError("%s failed:\n%s" % (cxx, proc.stderr)) + return so_path + + +def _default_include_dir(): + root = os.environ.get("TORCHSIM_DIR") + if not root: + root = os.path.dirname(os.path.dirname(os.path.dirname( + os.path.dirname(os.path.abspath(__file__))))) + return os.path.join(root, "TOGSim", "include") + + +def skeleton_to_so(skeleton_module, so_path, include_dir=None): + """skeleton module -> EmitC -> C++ -> compiled trace `.so`. Returns the + EmitC module text (for inspection / caching).""" + emitc = lower_to_emitc(skeleton_module) + cpp = emitc_to_cpp(emitc) + compile_so(cpp, so_path, include_dir or _default_include_dir()) + return str(emitc) + + +def build_trace_so(postvcix_path, so_path, include_dir=None): + """Full P2 path from a post-vcix kernel .mlir to a trace `.so`.""" + from . import build_skeleton as bs + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(open(postvcix_path).read(), ctx) + bs.build_skeleton(module) + return skeleton_to_so(module, so_path, include_dir) + + +def main(argv): + import argparse + + parser = argparse.ArgumentParser(prog="lower_to_emitc.py") + parser.add_argument("input", help="post-vcix kernel .mlir") + parser.add_argument("--so", required=True, help="output .so path") + parser.add_argument("--include-dir", default=None, + help="dir holding togsim_runtime.h (default: TOGSim/include)") + parser.add_argument("--emit-cpp", default=None, + help="also write the generated C++ here") + parser.add_argument("--emit-mlir", default=None, + help="also write the EmitC MLIR here") + args = parser.parse_args(argv[1:]) + + from . import build_skeleton as bs + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(open(args.input).read(), ctx) + bs.build_skeleton(module) + emitc = lower_to_emitc(module) + if args.emit_mlir: + open(args.emit_mlir, "w").write(str(emitc)) + cpp = emitc_to_cpp(emitc) + if args.emit_cpp: + open(args.emit_cpp, "w").write(cpp) + compile_so(cpp, args.so, args.include_dir or _default_include_dir()) + import sys + sys.stderr.write("wrote %s\n" % args.so) + return 0 + + +if __name__ == "__main__": + import sys + sys.exit(main(sys.argv)) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py index ac93ebc8..55bbae5a 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py @@ -29,6 +29,8 @@ import mlir.ir as ir # noqa: E402 +from ._mlir_util import walk_ops, i32, i64, attr_bool + MARKERS = ("linalg.matmul", "math.exp", "math.erf", "math.tanh", "math.sin", "math.cos") # math op name -> (opcode, imm) for the vcix.v.iv lowering (mirror Math*ToVCIX). @@ -80,20 +82,12 @@ def _legalize_vector_type(vt, vlen): return n, ir.VectorType.get([elt_count >> (n - 1)], elt_ty, scalable=[True]) -def _i64(v): - return ir.IntegerAttr.get(ir.IntegerType.get_signless(64), v) - - -def _i32(v): - return ir.IntegerAttr.get(ir.IntegerType.get_signless(32), v) - - def _viv(operand, result_ty, opcode, imm, rvl=None): """Create an unregistered vcix.v.iv (vcix::BinaryImmOp) op at the current IP.""" operands = [operand] if rvl is None else [operand, rvl] return ir.Operation.create( "vcix.v.iv", results=[result_ty], operands=operands, - attributes={"opcode": _i64(opcode), "imm": _i32(imm)}).results[0] + attributes={"opcode": i64(opcode), "imm": i32(imm)}).results[0] def _make_sf_vc_v_iv(vec, op_vt, n, legal_ty, opcode, imm): @@ -104,7 +98,7 @@ def _make_sf_vc_v_iv(vec, op_vt, n, legal_ty, opcode, imm): scalable = legal_ty.scalable rvl = None if scalable: - rvl = arith.ConstantOp(ir.IntegerType.get_signless(64), _i64(9)).result + rvl = arith.ConstantOp(ir.IntegerType.get_signless(64), i64(9)).result if n == 1: return _viv(vec, legal_ty, opcode, imm, rvl) elt_ty = legal_ty.element_type @@ -119,24 +113,16 @@ def _make_sf_vc_v_iv(vec, op_vt, n, legal_ty, opcode, imm): for i in range(total // elt_count): ext = vector.ExtractStridedSliceOp( legal_ty, vec, - ir.ArrayAttr.get([_i64(i * elt_count)]), - ir.ArrayAttr.get([_i64(elt_count)]), - ir.ArrayAttr.get([_i64(1)])).result + ir.ArrayAttr.get([i64(i * elt_count)]), + ir.ArrayAttr.get([i64(elt_count)]), + ir.ArrayAttr.get([i64(1)])).result v = _viv(ext, legal_ty, opcode, imm, rvl) res = vector.InsertStridedSliceOp( - v, res, ir.ArrayAttr.get([_i64(i * elt_count)]), - ir.ArrayAttr.get([_i64(1)])).result + v, res, ir.ArrayAttr.get([i64(i * elt_count)]), + ir.ArrayAttr.get([i64(1)])).result return res -def _iter_ops(block): - for op in list(block.operations): - yield op - for region in op.operation.regions: - for b in region.blocks: - yield from _iter_ops(b) - - # --------------------------------------------------------------------------- # matmul lowering helpers (mirror MatmulOpLowering) # --------------------------------------------------------------------------- @@ -146,11 +132,6 @@ def _elt_bits(elt_ty): return ir.FloatType(elt_ty).width -def _bool_attr_true(op, key): - a = op.attributes - return key in a and ir.BoolAttr(a[key]).value - - def _enclosing_loops(op): """Walk ancestor ops; return (accumulation, outer, inner) affine.for lists, outermost-first (mirror the C++ insert-at-begin).""" @@ -158,11 +139,11 @@ def _enclosing_loops(op): parent = op.operation.parent while parent is not None: if parent.name == "affine.for": - if _bool_attr_true(parent, "accumulation_loop"): + if attr_bool(parent, "accumulation_loop"): acc.insert(0, parent) - if _bool_attr_true(parent, "outer_loop"): + if attr_bool(parent, "outer_loop"): outer.insert(0, parent) - if _bool_attr_true(parent, "inner_loop"): + if attr_bool(parent, "inner_loop"): inner.insert(0, parent) parent = parent.parent return acc, outer, inner @@ -200,7 +181,7 @@ def _scan_conv_offsets(ow_loop, o_h, k_h, o_w, k_w): """Mirror the heuristic offset scan: find affine.apply(o_h,k_h)/(o_w,k_w) in the o_w loop and read the constant in its map (default 1).""" offset_h = offset_w = 1 - for o in _iter_ops(ow_loop.regions[0].blocks[0]): + for o in walk_ops(ow_loop.regions[0].blocks[0]): if o.operation.name != "affine.apply": continue ops = list(o.operation.operands) @@ -391,7 +372,7 @@ def _root(v): return owner.operands[0] return v rootA, rootB = _root(A), _root(B) - for o in _iter_ops(outer[-1].regions[0].blocks[0]): + for o in walk_ops(outer[-1].regions[0].blocks[0]): if o.operation.name == "affine.vector_store": dest = _root(o.operation.operands[1]) if dest == rootA: @@ -488,6 +469,9 @@ def _root(v): # --- B dma_wait --- nacc = len(acc) acc_ivs = [_loop_iv(l) for l in acc] + # LEGACY: coefficient -1 on an accumulation loop var is a SENTINEL for "this + # tag dim is the reduction axis", not an offset. The trace path strips it + # (build_skeleton._strip_accum_terms); remove here once legacy retires. bexpr = ir.AffineDimExpr.get(0) * -1 for i in range(1, nacc): bexpr = bexpr + ir.AffineDimExpr.get(i) * -1 @@ -544,6 +528,8 @@ def _root(v): with body_ip: # --- A dma_wait --- + # LEGACY: as for the B dma_wait above, the -1 coefficients mark the reduction + # axis; the trace path strips them. Remove once legacy retires. aexpr = ir.AffineDimExpr.get(0) * -1 for i in range(1, nacc): aexpr = aexpr + ir.AffineDimExpr.get(i) * -1 @@ -617,7 +603,7 @@ def run(module, vectorlane=128, vlen=128, **_): mms = [] for region in module.operation.regions: for b in region.blocks: - for o in _iter_ops(b): + for o in walk_ops(b): if o.operation.name == "linalg.matmul": mms.append(o.operation) for o in mms: @@ -625,7 +611,7 @@ def run(module, vectorlane=128, vlen=128, **_): targets = [] for region in module.operation.regions: for b in region.blocks: - for op in _iter_ops(b): + for op in walk_ops(b): if op.operation.name in _MATH_VIV: targets.append(op.operation) for op in targets: diff --git a/PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py b/PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py index 76e30cb3..3ed0a394 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py +++ b/PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py @@ -24,13 +24,7 @@ OP_NAME = "torchsim.vlane_idx" MARKERS = (OP_NAME,) - -def _iter_ops(block): - for op in list(block.operations): - yield op - for region in op.operation.regions: - for b in region.blocks: - yield from _iter_ops(b) +from ._mlir_util import walk_ops def run(module, **_): @@ -46,7 +40,7 @@ def run(module, **_): targets = [] for region in module.operation.regions: for b in region.blocks: - for op in _iter_ops(b): + for op in walk_ops(b): if op.operation.name == OP_NAME: targets.append(op.operation) diff --git a/PyTorchSimFrontend/mlir/passes/togsim_ops.py b/PyTorchSimFrontend/mlir/passes/togsim_ops.py new file mode 100644 index 00000000..85b89757 --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/togsim_ops.py @@ -0,0 +1,101 @@ +"""Shared vocabulary for the skeleton+API MLIR form (C1). + +The trace pipeline (docs/design/togsim_cpp_trace.md) reduces a kernel's MLIR to +a *loop skeleton + API calls*: native `affine.for`/`scf.for` loops (bounds kept +as-is, symbolic preserved) plus a handful of `togsim.*` ops that stand for the +runtime API. This module is the single source of truth for those op names and +attribute keys, shared by: + + * build_skeleton (C2) -- produces the skeleton+API MLIR, and + * togsim->emitc lowering (C4) -- rewrites each op to an `emitc.call_opaque`. + +The ops are kept *unregistered* (like the existing `togsim.transfer`), so there +is no C++ dialect to register; C4 is a custom rewrite, not a registered +ConversionPass. + +Grammar (each op lowers 1:1 to a `togsim_runtime.h` free function): + + "togsim.dma"(%dram_idx, %tag_idx) { -> togsim_dma(ctx, dir, arg_id, + dir = 0 | 1, # LOAD|STORE offset, ndim, dims, strides, + dims = [..], strides = [..], elem_bits, is_async, + elem_bits = i32, is_async = bool, tag_id, tag_slot, + tag_id = i32, arg_id = i32, read_bufs, write_bufs) + read_bufs = [..], write_bufs = [..] + } : (index, index) -> () + + "togsim.compute"() { -> togsim_compute(ctx, tile_id, + tile_id = i64, compute_type = i32, compute_type, ndim, dims, + read_bufs = [..], write_bufs = [..] read_bufs, write_bufs) + } : () -> () + + "togsim.memory_barrier"(%tag_idx) { -> togsim_memory_barrier(ctx, + tag_id = i32, write_bufs = [..] tag_id, tag_slot, write_bufs) + } : (index) -> () + + "togsim.compute_barrier"() : () -> () -> togsim_compute_barrier(ctx) + +How an async dma pairs with its sync point: NOT by a compile-time id. One static +`togsim.dma` op runs once per loop iteration, each with a different RUNTIME tag +slot `%tag[%idx]`, so the pairing must be a runtime key. `togsim.dma` carries a +`tag_id` (its tag memref identity) and the runtime `%tag[%idx]` operand; the +original `memref.dma_wait` becomes an explicit `togsim.memory_barrier` carrying +the same `tag_id` + tag index. They pair at runtime by `(tag_id, tag_slot)` via +the Core's tag table (the dma signals the tag at data-arrival; the barrier waits +it). `tag_id` (which tag memref) is distinct from `tag_slot` (the SRAM tile slot, +used for the double-buffer / capacity model). A sync (non-async) dma is blocking, +so it needs no barrier. The key must carry a runtime component: one static dma op +runs once per loop iteration, so a compile-time id cannot express the pairing. + +Keep this in lockstep with TOGSim/include/togsim_runtime.h (TOGSIM_ABI_VERSION). +""" + +# ---- op names ------------------------------------------------------------- +DMA = "togsim.dma" +COMPUTE = "togsim.compute" +COMPUTE_BAR = "togsim.compute_barrier" # fence: drain async compute before a consumer (sec 10.7) +MEMORY_BAR = "togsim.memory_barrier" # explicit async-DMA sync (the original dma_wait); tag-keyed + +#: every op this module owns (for matchers / DCE roots in C2). +OP_NAMES = (DMA, COMPUTE, COMPUTE_BAR, MEMORY_BAR) + +#: op name -> the togsim_runtime.h symbol C4 lowers it to. +EMITC_CALLEE = { + DMA: "togsim_dma", + COMPUTE: "togsim_compute", + COMPUTE_BAR: "togsim_compute_barrier", + MEMORY_BAR: "togsim_memory_barrier", +} + +#: producer entry-point symbol the TOGSim loader resolves (see togsim_runtime.h). +ENTRY_SYMBOL = "togsim_kernel" + +#: runtime callee emitted directly by lower_to_emitc (not a skeleton op): the +#: per-work-item core allocation. See togsim_cpp_trace.md sec 9.3. Kept in +#: lockstep with togsim_runtime.h. +CORE_ALLOC_CALLEE = "togsim_core_alloc" + +# ---- attribute keys ------------------------------------------------------- +ATTR_DIR = "dir" # i32: DIR_LOAD | DIR_STORE +ATTR_DIMS = "dims" # i64 array: tile extents +ATTR_STRIDES = "strides" # i64 array: tile strides +ATTR_ELEM_BITS = "elem_bits" # i32 +ATTR_IS_ASYNC = "is_async" # bool +ATTR_TILE_ID = "tile_id" # i64: key into the precomputed tile_id->cycle table +ATTR_COMPUTE_TYPE = "compute_type" # i32: 0 vector / 1 matmul / 2 preload (Core enum) +ATTR_READ_BUFS = "read_bufs" # i64 array: SRAM buffer ids this op reads (sec 10 dataflow) +ATTR_WRITE_BUFS = "write_bufs" # i64 array: SRAM buffer ids this op writes (sec 10 dataflow) +ATTR_TAG_ID = "tag_id" # i32: identity of the DMA's tag memref; pairs an async dma with + # its memory_barrier by the RUNTIME tag slot (tag_id + tag index) +ATTR_ARG_ID = "arg_id" # i32: which tensor (func arg) this DMA's base is + +# Must match togsim_dma_dir in togsim_runtime.h. +DIR_LOAD = 0 +DIR_STORE = 1 + + +def is_togsim_op(op): + """True if `op` (an Operation or a wrapping view) is one of ours.""" + name = getattr(op, "name", None) + if name is None: + name = getattr(getattr(op, "operation", None), "name", None) + return name in OP_NAMES diff --git a/TOGSim/include/Instruction.h b/TOGSim/include/Instruction.h index bb62a440..3740b53a 100644 --- a/TOGSim/include/Instruction.h +++ b/TOGSim/include/Instruction.h @@ -12,7 +12,10 @@ #include #include -enum class Opcode { MOVIN, MOVOUT, COMP, BAR, COUNT}; +// MEMORY_BAR waits a DMA tag in the tag table. COMPUTE_BAR waits the systolic-array +// compute pipelines to drain -- the explicit fence before a store consumes the +// results of async matmuls. +enum class Opcode { MOVIN, MOVOUT, COMP, MEMORY_BAR, COMPUTE_BAR, COUNT}; typedef uint64_t addr_type; typedef uint64_t cycle_type; @@ -29,6 +32,11 @@ class Instruction : public std::enable_shared_from_this { Instruction(Opcode opcode); void finish_instruction(); void add_child(std::shared_ptr child); + // Occupancy (SA-pipeline) dependency: the child is released when THIS op is + // ISSUED (enters the pipeline), not when it finishes -- so a preload/matmul + // successor overlaps it instead of waiting its full latency (sec 10.7). + void add_pipeline_child(std::shared_ptr child); + void release_pipeline_children(); bool check_ready() { return ready_counter == 0; } const Opcode get_opcode() { return opcode; } bool is_dma_read() { return opcode == Opcode::MOVIN; } @@ -103,6 +111,7 @@ class Instruction : public std::enable_shared_from_this { cycle_type overlapping_cycle; size_t ready_counter; std::set> child_inst; + std::set> _pipeline_children; // released at issue (sec 10.7) std::vector tile_size; std::vector tile_stride; size_t _tile_numel; diff --git a/TOGSim/include/togsim_loader.h b/TOGSim/include/togsim_loader.h new file mode 100644 index 00000000..52c1ac1e --- /dev/null +++ b/TOGSim/include/togsim_loader.h @@ -0,0 +1,51 @@ +#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. + +#include +#include + +#include "togsim_runtime.h" + +namespace togsim { + +// One modeled instruction recorded by the runtime callbacks. +struct TraceRec { + enum Kind { DISPATCH, DMA, COMPUTE, MEMORY_BAR, COMPUTE_BAR } kind; + int32_t core; // work-item -> core binding (set by togsim_core_alloc) + // DMA / MEMORY_BAR + int32_t dir; // togsim_dma_dir + int32_t arg_id; // tensor + int32_t elem_bits; + int32_t is_async; + uint64_t addr; // resolved DRAM byte address = base[arg_id] + off*bytes + int32_t tag_id; // DMA/MEMORY_BAR: tag memref identity; with tag_slot the + // runtime pairing key (an async dma <-> its memory_barrier) + uint64_t tag_slot; // SRAM tile slot (double-buffer / capacity model) + std::vector dims; // tile extents (DMA) + std::vector strides; // tile strides (DMA) + std::vector read_bufs; // SRAM buffer ids read (sec 10 dependency model) + std::vector write_bufs; // SRAM buffer ids written (MEMORY_BAR: released bufs) + // COMPUTE + uint64_t tile_id; + int32_t compute_type; // 0 vector / 1 matmul / 2 preload (Core unit enum) + int64_t cycle; // looked up from the cycle table + int64_t overlapping; // looked up from the cycle table +}; + +struct RunResult { + bool ok = false; + std::vector trace; +}; + +// 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 across `num_cores`. +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, + int32_t num_cores); + +} // namespace togsim diff --git a/TOGSim/include/togsim_runtime.h b/TOGSim/include/togsim_runtime.h new file mode 100644 index 00000000..f9c53c1c --- /dev/null +++ b/TOGSim/include/togsim_runtime.h @@ -0,0 +1,69 @@ +#pragma once +// togsim_runtime.h -- C ABI between a compiled, shape-parametric trace producer +// (`.so`, MLIR -> EmitC -> C++) and TOGSim: each call emits one modeled +// instruction, and the producer carries no timing model. See sec 5.4 of the doc. + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Producer/runtime ABI version. TOGSim refuses to load a producer whose +// embedded togsim_abi_version() does not match TOGSIM_ABI_VERSION. +#define TOGSIM_ABI_VERSION 11 +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. +typedef struct EmitCtx EmitCtx; + +// Direction for togsim_dma. +typedef enum { + TOGSIM_DMA_LOAD = 0, // DRAM -> SRAM (MOVIN) + TOGSIM_DMA_STORE = 1, // SRAM -> DRAM (MOVOUT) +} togsim_dma_dir; + +// Emit a DMA (sec 5.4). `offset` is an ELEMENT offset into tensor `arg_id`; null +// `strides` => contiguous. `is_async` => it finishes at ISSUE, so the barrier keyed +// by `(tag_id, tag_slot)` gates consumers. `read_bufs`/`write_bufs` -> sec 10. + +void togsim_dma(EmitCtx* ctx, int32_t dir, int32_t arg_id, + uint64_t offset, int32_t ndim, const int64_t* dims, + const int64_t* strides, int32_t elem_bits, + int32_t is_async, int32_t tag_id, uint64_t tag_slot, + const int64_t* read_bufs, int32_t n_read, + const int64_t* write_bufs, int32_t n_write); + +// Emit a fixed-size tile compute. Cost comes from the tile_id->cycle table (sec 6), +// not from `dims`. `compute_type` (0 vector / 1 matmul / 2 preload) routes the op to +// the VPU or the systolic array. +void togsim_compute(EmitCtx* ctx, uint64_t tile_id, int32_t compute_type, + int32_t ndim, const int64_t* dims, + const int64_t* read_bufs, int32_t n_read, + const int64_t* write_bufs, int32_t n_write); + +// The explicit async-DMA sync (sec 10.5). Pairs with its async togsim_dma by the +// runtime `(tag_id, tag_slot)` and becomes the last writer of `write_bufs`, so +// consumers gate on data arrival. A sync dma blocks to arrival and needs no barrier. +void togsim_memory_barrier(EmitCtx* ctx, int32_t tag_id, uint64_t tag_slot, + const int64_t* write_bufs, int32_t n_write); + +// Core allocation (sec 9.3): the producer calls this at each parallel work-item's +// start, and the ops that follow bind to the returned core. No free -- a core is an +// assignment. The producer never names num_cores; the runtime owns the pool. +int32_t togsim_core_alloc(EmitCtx* ctx); + +// Compute fence: drain in-flight async compute (the systolic-array matmuls) +// before the following op (a store) consumes their result. Explicit barrier in +// the trace; the loader turns it into a COMPUTE_BAR instruction (sec 10.7). +void togsim_compute_barrier(EmitCtx* ctx); + +// Entry point the loader resolves in the producer `.so`. `shape_args` carries +// the runtime values for the kernel's symbolic dimensions (in a kernel-specific +// order recorded alongside the cached `.so`); `n_shape_args` is their count. +void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n_shape_args); + +#ifdef __cplusplus +} // extern "C" +#endif diff --git a/TOGSim/include/togsim_trace_bridge.h b/TOGSim/include/togsim_trace_bridge.h new file mode 100644 index 00000000..f4b53c80 --- /dev/null +++ b/TOGSim/include/togsim_trace_bridge.h @@ -0,0 +1,12 @@ +#pragma once +// togsim_trace_bridge.h -- turn a recorded trace into a TileGraph the existing +// Simulator/Core can run: one Tile per work-item (a TILE_BEGIN/TILE_END span), +// dependency edges by last-writer per SRAM buffer (sec 10). +#include + +#include "TileGraph.h" +#include "togsim_loader.h" + +// Build a TileGraph from a recorded trace. `path`/`name` label the graph. +std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, + const std::string& name); diff --git a/TOGSim/src/CMakeLists.txt b/TOGSim/src/CMakeLists.txt index 65cd4dd4..d782d4d1 100644 --- a/TOGSim/src/CMakeLists.txt +++ b/TOGSim/src/CMakeLists.txt @@ -12,3 +12,8 @@ file(GLOB_RECURSE SRC_FILES # build add_executable(${LIB_NAME} ${SRC_FILES}) + +# Export the executable's dynamic symbols (-rdynamic) so a dlopen'd trace +# producer .so resolves the togsim_* runtime callbacks back into this binary +# (P3 trace pipeline). +set_target_properties(${LIB_NAME} PROPERTIES ENABLE_EXPORTS ON) diff --git a/TOGSim/src/Core.cc b/TOGSim/src/Core.cc index 9dad8597..8a728318 100644 --- a/TOGSim/src/Core.cc +++ b/TOGSim/src/Core.cc @@ -154,7 +154,7 @@ void Core::dma_cycle() { } else if(!finished_inst->is_dma_read()) { core_trace_log::log_error_dma_instruction_invalid(_core_cycle, _id); exit(EXIT_FAILURE); - } else if (finished_inst->get_opcode() == Opcode::BAR) { + } else if (finished_inst->get_opcode() == Opcode::MEMORY_BAR) { core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15(TraceLogTag::kInstructionFinished), @@ -265,6 +265,10 @@ void Core::cycle() { break; case Opcode::COMP: { + // sec 10.7: this op is now entering the pipeline -> release its + // occupancy (pipeline) dependents so a preload/matmul successor + // overlaps it instead of waiting its full latency. + inst->release_pipeline_children(); auto& target_pipeline = get_compute_pipeline(inst->get_compute_type()); if (target_pipeline.empty()) { inst->finish_cycle = _core_cycle + inst->get_compute_cycle(); @@ -297,7 +301,7 @@ void Core::cycle() { } } break; - case Opcode::BAR: + case Opcode::MEMORY_BAR: { auto& key = inst->get_tag_id(); uint32_t finished = _dma.get_tag_finish(inst->subgraph_id, key); @@ -324,6 +328,24 @@ void Core::cycle() { issued = true; } break; + case Opcode::COMPUTE_BAR: + { + // Compute fence: finish only once ALL compute pipelines have drained + // (every systolic array + the VPU empty). Until then it does not issue -- + // it stays in the ready queue and is re-checked next cycle. + bool drained = _vu_compute_pipeline.empty(); + for (int s = 0; s < _num_systolic_array_per_core; s++) + drained = drained && _sa_compute_pipeline.at(s).empty(); + if (drained) { + core_trace_log::trace_instruction_line(_core_cycle, _id, + TraceLogTag::pad15(TraceLogTag::kInstructionFinished), + inst->get_global_inst_id(), + core_trace_log::format_instruction_detail_line(*inst)); + finish_instruction(inst); + issued = true; + } + } + break; default: core_trace_log::log_error_undefined_opcode(); exit(EXIT_FAILURE); diff --git a/TOGSim/src/CoreTraceLog.cc b/TOGSim/src/CoreTraceLog.cc index ebc31de0..9761f9ec 100644 --- a/TOGSim/src/CoreTraceLog.cc +++ b/TOGSim/src/CoreTraceLog.cc @@ -70,7 +70,7 @@ std::string format_instruction_detail_line(Instruction& inst) { if (op == Opcode::MOVIN || op == Opcode::MOVOUT) { return fmt::format("{} (addr_name={})", opname, inst.get_addr_name()); } - if (op == Opcode::BAR) { + if (op == Opcode::MEMORY_BAR) { return fmt::format("{} (addr_name={} tag_id=[{}] tag_idx=[{}] tag_stride=[{}])", opname, inst.get_addr_name(), diff --git a/TOGSim/src/Instruction.cc b/TOGSim/src/Instruction.cc index f236d160..54e50511 100644 --- a/TOGSim/src/Instruction.cc +++ b/TOGSim/src/Instruction.cc @@ -23,7 +23,8 @@ std::string opcode_to_string(Opcode opcode) { case Opcode::MOVIN: return "MOVIN"; case Opcode::MOVOUT: return "MOVOUT"; case Opcode::COMP: return "COMP"; - case Opcode::BAR: return "BAR"; + case Opcode::MEMORY_BAR: return "MEMORY_BAR"; + case Opcode::COMPUTE_BAR: return "COMPUTE_BAR"; default: return "Unknown"; } } @@ -60,6 +61,16 @@ void Instruction::add_child(std::shared_ptr child) { child_inst.insert(child); } +void Instruction::add_pipeline_child(std::shared_ptr child) { + child->inc_ready_counter(); + _pipeline_children.insert(child); +} + +void Instruction::release_pipeline_children() { + for (auto& c : _pipeline_children) c->dec_ready_counter(); + _pipeline_children.clear(); +} + void Instruction::inc_waiting_request() { _nr_waiting_request++; } diff --git a/TOGSim/src/TileGraphParser.cc b/TOGSim/src/TileGraphParser.cc index 5060d336..572062e0 100644 --- a/TOGSim/src/TileGraphParser.cc +++ b/TOGSim/src/TileGraphParser.cc @@ -543,7 +543,7 @@ std::vector> TileLoopNode::get_tiles_from_iter(TileGraphPa fmt::join(new_tag_stride_list, ", ")); std::shared_ptr inst = std::make_shared( - Opcode::BAR, 0, + Opcode::MEMORY_BAR, 0, 0, base_addr, std::vector(), std::vector(), 0, tag_list, new_tag_stride_list, accum_tag_list diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 010826ef..8726cfdf 100644 --- a/TOGSim/src/main.cc +++ b/TOGSim/src/main.cc @@ -8,6 +8,8 @@ #include "Simulator.h" #include "TileGraphParser.h" #include "helper/CommandLineParser.h" +#include "togsim_loader.h" // P3 trace pipeline: run a compiled producer .so +#include "togsim_trace_bridge.h" // ... and bridge its trace to a TileGraph namespace fs = std::filesystem; namespace po = boost::program_options; @@ -104,6 +106,11 @@ int main(int argc, char** argv) { "models_list", "Path for the trace file (.trace)"); cmd_parser.add_command_line_option( "log_level", "Set for log level [trace, debug, info], default = info"); + cmd_parser.add_command_line_option( + "trace_so", "Path to a compiled trace producer .so (P3 trace pipeline)"); + cmd_parser.add_command_line_option( + "cycle_table", "Path to a 'cycleoverlapping' per-tile_id sidecar (TSV) " + "for --trace_so; falls back to a flat stub if omitted"); try { cmd_parser.parse(argc, argv); } catch (const CommandLineParser::ParsingError& e) { @@ -147,6 +154,46 @@ int main(int argc, char** argv) { exit(1); } + // P3 trace pipeline: if a compiled producer .so is given, run it, bridge the + // recorded trace to a TileGraph, and run the existing Simulator on it. + std::string trace_so_path; + cmd_parser.set_if_defined("trace_so", &trace_so_path); + if (!trace_so_path.empty()) { + const auto& cfg = simulator->get_hardware_config_yaml(); + int num_cores = cfg["num_cores"] ? cfg["num_cores"].as() : 1; + // First cut: stub tensor bases (real per-tensor addresses come later). + std::vector bases(16); + for (size_t i = 0; i < bases.size(); ++i) bases[i] = 0x100000ull * (i + 1); + // Cycle table: load the per-tile_id TSV sidecar if given, else a flat stub. + std::vector cyc, ovl; + std::string cycle_table_path; + cmd_parser.set_if_defined("cycle_table", &cycle_table_path); + if (!cycle_table_path.empty()) { + std::ifstream ct(cycle_table_path); + if (!ct.is_open()) { spdlog::error("[TOGSim] cannot open cycle_table {}", cycle_table_path); exit(1); } + int64_t c, o; + while (ct >> c >> o) { cyc.push_back(c); ovl.push_back(o); } + spdlog::info("[TOGSim-trace] loaded cycle table: {} tiles from {}", cyc.size(), cycle_table_path); + } else { + 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(), + num_cores); + if (!run.ok) { spdlog::error("[TOGSim] trace producer run failed"); exit(1); } + spdlog::info("[TOGSim-trace] recorded {} instructions", run.trace.size()); + auto tg = trace_to_tilegraph(run, "trace_kernel"); + tg->set_arrival_time(simulator->get_core_cycle()); + tg->set_kernel_id(0); + simulator->enqueue_graph(0, std::move(tg)); + simulator->run_simulator(); + spdlog::info("[TOGSim-trace] Total cycles: {}", simulator->get_core_cycle()); + simulator->print_core_stat(); + return 0; + } + // Get trace file path cmd_parser.set_if_defined("models_list", &trace_file_path); diff --git a/TOGSim/src/togsim_runtime.cc b/TOGSim/src/togsim_runtime.cc new file mode 100644 index 00000000..84cb2cd2 --- /dev/null +++ b/TOGSim/src/togsim_runtime.cc @@ -0,0 +1,182 @@ +// togsim_runtime.cc -- the producer ABI (togsim_runtime.h) and the loader +// (togsim_loader.h). The producer's calls each record a TraceRec on the opaque +// EmitCtx, resolving DRAM addresses and per-tile cycles as they go. + +#include "togsim_loader.h" + +#include +#include +#include +#include +#include + +// Full definition of the opaque handle from togsim_runtime.h. The producer holds +// only EmitCtx* and never dereferences it. +struct EmitCtx { + // inputs supplied by the loader + const uint64_t* tensor_base = nullptr; + int32_t n_tensors = 0; + const int64_t* cyc = nullptr; // tile_id -> cycle + const int64_t* ovl = nullptr; // tile_id -> overlapping_cycle + int32_t n_tiles = 0; + int32_t num_cores = 1; + // mutable run state + int32_t rr = 0; // round-robin core cursor + int32_t cur_core = -1; // current work-item's core + std::vector trace; +}; + +namespace { +inline togsim::TraceRec blank(togsim::TraceRec::Kind k, int32_t core) { + togsim::TraceRec r{}; + r.kind = k; + r.core = core; + return r; +} +} // namespace + +extern "C" { + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +int32_t togsim_core_alloc(EmitCtx* ctx) { + // Round-robin a core from the pool; the producer never sees num_cores. Binds + // it as the current core for the ops that follow (the work-item's reduction). + ctx->cur_core = ctx->num_cores > 0 ? (ctx->rr++ % ctx->num_cores) : 0; + ctx->trace.push_back(blank(togsim::TraceRec::DISPATCH, ctx->cur_core)); + return ctx->cur_core; +} + +void togsim_dma(EmitCtx* ctx, int32_t dir, int32_t arg_id, + uint64_t offset, int32_t ndim, const int64_t* dims, + const int64_t* strides, int32_t elem_bits, + int32_t is_async, int32_t tag_id, uint64_t tag_slot, + const int64_t* read_bufs, int32_t n_read, + const int64_t* write_bufs, int32_t n_write) { + 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); + 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; + 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); +} + +void togsim_compute(EmitCtx* ctx, uint64_t tile_id, int32_t compute_type, + int32_t ndim, const int64_t* dims, + 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); + r.tile_id = tile_id; + r.compute_type = compute_type; + 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); +} + +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); + r.tag_id = tag_id; r.tag_slot = tag_slot; + for (int32_t i = 0; i < n_write; ++i) r.write_bufs.push_back(write_bufs[i]); + ctx->trace.push_back(r); +} + +void togsim_compute_barrier(EmitCtx* ctx) { + ctx->trace.push_back(blank(togsim::TraceRec::COMPUTE_BAR, ctx->cur_core)); +} + +} // extern "C" + +namespace togsim { + +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, + int32_t num_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.num_cores = num_cores > 0 ? num_cores : 1; + emit(&ctx, (int64_t*)shape_args, n_shape); + + res.ok = true; + res.trace = std::move(ctx.trace); + return res; +} + +SimResult simulate(const RunResult& run, const TimingParams& params) { + SimResult out; + std::unordered_map dma_free; // DMA-engine free time, per core + std::unordered_map comp_free; // compute free time, per core + std::unordered_map prev_comp; // prev compute finish (overlap), per core + std::map, uint64_t> tag_finish; // (tag_id,tag_slot) -> finish + std::vector pending; // barrier-resolved deps since last compute + + for (const auto& t : run.trace) { + const int c = t.core; + switch (t.kind) { + case TraceRec::DMA: { + // DMAs serialize on the core's DMA engine (overlap compute -> separate + // timeline). finish = issue + latency, recorded under the runtime tag. + uint64_t start = dma_free[c]; + uint64_t fin = start + params.dma_latency; + dma_free[c] = fin; + tag_finish[{t.tag_id, t.tag_slot}] = fin; + out.n_dma++; + break; + } + case TraceRec::MEMORY_BAR: { + // the explicit async-DMA sync: gate the next compute on the paired dma's + // data-arrival, found by the runtime tag (tag_id, tag_slot). + auto it = tag_finish.find({t.tag_id, t.tag_slot}); + if (it != tag_finish.end()) pending.push_back(it->second); + break; + } + case TraceRec::COMPUTE: { + uint64_t deps = 0; + for (uint64_t f : pending) deps = std::max(deps, f); + pending.clear(); + uint64_t start = std::max(comp_free[c], deps); + uint64_t fin; + auto pit = prev_comp.find(c); + if (pit != prev_comp.end()) { + uint64_t prev = pit->second; + uint64_t tail = prev > start ? prev - start : 0; // prev still running + uint64_t overlapped = std::min(tail, (uint64_t)t.overlapping); + fin = std::max(start, prev) + (uint64_t)t.cycle - overlapped; + } else { + fin = start + (uint64_t)t.cycle; + } + comp_free[c] = fin; + prev_comp[c] = fin; + out.n_compute++; + break; + } + case TraceRec::DISPATCH: + case TraceRec::COMPUTE_BAR: + break; // work-item boundary / compute fence: no cost in this reference timer + } + } + for (auto& kv : dma_free) out.total_cycle = std::max(out.total_cycle, kv.second); + for (auto& kv : comp_free) out.total_cycle = std::max(out.total_cycle, kv.second); + return out; +} + +} // namespace togsim diff --git a/TOGSim/src/togsim_trace_bridge.cc b/TOGSim/src/togsim_trace_bridge.cc new file mode 100644 index 00000000..56e85a68 --- /dev/null +++ b/TOGSim/src/togsim_trace_bridge.cc @@ -0,0 +1,190 @@ +// togsim_trace_bridge.cc -- see togsim_trace_bridge.h +#include "togsim_trace_bridge.h" + +#include +#include +#include + +#include "Tile.h" +#include "Instruction.h" + +namespace { + +// `uniq` is a per-DMA-record Core tag key, so every reduction iteration of one +// static dma gets a distinct key (multi-tile-K, conv); its memory_barrier reuses +// it. `tag_idx` (the subtile slot) still drives the SRAM double-buffer model. + +// FIXME: `uniq` is reconstructed here from record order. build_skeleton should +// instead thread dma_fine_grained's per-iteration tag alloc through as an SSA +// handle on togsim.dma / togsim.memory_barrier (sec 11). +std::shared_ptr make_dma(const togsim::TraceRec& t, int64_t uniq) { + Opcode op = (t.dir == 1) ? Opcode::MOVOUT : Opcode::MOVIN; + std::vector tile_size(t.dims.begin(), t.dims.end()); + std::vector tile_stride(t.strides.begin(), t.strides.end()); + std::vector tag_idx{(int64_t)t.tag_slot}; + std::vector tag_stride{1}; + auto inst = std::make_shared( + op, /*compute_cycle=*/0, /*num_parents=*/0, /*dram_addr=*/t.addr, + tile_size, tile_stride, (size_t)t.elem_bits, tag_idx, tag_stride, + /*accum_tag_idx_list=*/std::vector{}); + inst->set_is_async(t.is_async != 0); + inst->set_addr_name("tag" + std::to_string(uniq), uniq); + inst->prepare_tag_key(); + return inst; +} + +// A MEMORY_BAR carrying the SAME `uniq` tag key as the async dma it gates -- the +// Core's tag table signals it at the dma's DATA-ready (resp-complete), unlike a +// raw add_child which the async dma releases at issue-complete. +std::shared_ptr make_mem_bar(const togsim::TraceRec& t, int64_t uniq) { + auto bar = std::make_shared( + Opcode::MEMORY_BAR, 0, 0, 0, + std::vector{}, std::vector{}, 0, + std::vector{(int64_t)t.tag_slot}, std::vector{1}, + std::vector{}); + bar->set_addr_name("tag" + std::to_string(uniq), uniq); + bar->prepare_tag_key(); + return bar; +} + +std::shared_ptr make_compute(const togsim::TraceRec& t) { + auto inst = std::make_shared( + Opcode::COMP, /*compute_cycle=*/(cycle_type)t.cycle, /*num_parents=*/0, + /*dram_addr=*/0, std::vector{}, std::vector{}, /*elem_bits=*/0, + std::vector{}, std::vector{}, std::vector{}); + inst->set_overlapping_cycle((cycle_type)t.overlapping); + inst->set_compute_type(t.compute_type); // route to VPU vs systolic array + 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({}); + + std::shared_ptr sg; + std::shared_ptr tile; + // Explicit dependency DAG (sec 10): a reader depends on the last writer of each + // SRAM buffer it reads. Scoped per work-item (reset at each dispatch) -- buffers + // are work-item-local, so distinct work-items are independent (-> parallel). + std::map> last_writer; // buffer id -> producer + // 1 load : N barriers, so track the CURRENT load per (tag_id, tag_slot), not a + // FIFO. Each load takes a fresh `uniq` Core key and its iteration's barriers reuse + // it. Correct only because a load nest and its consumers run in order. Per work-item. + std::map, + std::pair>> current_dma; + int64_t next_tag = 0; // mints a unique Core tag key per dma record + // Async compute (matmul/preload) pipelines on the systolic array. A store needs the + // drained result, so it FLUSHes -- one barrier before the store waits all outstanding + // async compute, with no per-op completion events. + std::vector> outstanding_async; + std::shared_ptr pending_bar; // last COMPUTE_BAR fence, awaited by the next store + auto is_async_compute = [](int ct) { return ct == 1 || ct == 2; }; // matmul / preload + + auto flush = [&]() { + if (sg && tile) { + sg->add_tile(tile); + tile->set_owner(sg); + tg->append_subgraph(sg); + } + sg.reset(); + tile.reset(); + last_writer.clear(); + current_dma.clear(); + next_tag = 0; + outstanding_async.clear(); + pending_bar.reset(); + }; + + // Edges from the recorded read/write buffer sets: a reader depends on the last writer + // of each buffer it reads. An SA-producer -> matmul edge is an OCCUPANCY dependency + // (released at ISSUE); every other edge is a LATENCY dependency (released at finish). + const int MATMUL_CT = 1, PRELOAD_CT = 2; + auto link = [&](std::shared_ptr inst, + const std::vector& reads, + const std::vector& writes) { + for (int64_t b : reads) { + auto it = last_writer.find(b); + if (it == last_writer.end()) continue; + int pct = it->second->get_compute_type(); + if (inst->get_compute_type() == MATMUL_CT && (pct == MATMUL_CT || pct == PRELOAD_CT)) + it->second->add_pipeline_child(inst); // SA pipeline -> occupancy (overlap) + else + it->second->add_child(inst); // data/result -> latency (full wait) + } + for (int64_t b : writes) last_writer[b] = inst; + tile->append_instuction(inst); + }; + + for (const auto& t : run.trace) { + if (t.kind == TraceRec::DISPATCH) { + // new work-item -> new subgraph (bound to its core) + tile. + flush(); + sg = std::make_shared(); + sg->set_core_id(t.core); + tile = std::make_shared(Tile::Status::INITIALIZED); + continue; + } + if (!tile) continue; // defensive: ops before the first core_alloc + + if (t.kind == TraceRec::DMA) { + int64_t uniq = next_tag++; // fresh Core tag key per dma record + auto inst = make_dma(t, uniq); + size_t numel = 1; // SRAM footprint (ready-tile ordering) + for (auto d : t.dims) numel *= (size_t)d; + tile->inc_required_sram_size(numel * (t.elem_bits / 8)); + if (t.dir == 1) { // STORE + if (pending_bar) { + // after a compute fence: wait it (drains the async matmuls) -- covers + // the accumulator read, so no per-buffer read edge. + pending_bar->add_child(inst); + pending_bar.reset(); + for (int64_t b : t.write_bufs) last_writer[b] = inst; + tile->append_instuction(inst); + } else { + link(inst, t.read_bufs, t.write_bufs); + } + } else { // LOAD + tile->append_instuction(inst); + // async load: the CURRENT load for this (tag_id, tag_slot), with a fresh uniq + // its barriers reuse. last_writer = the dma until its barrier overwrites it, 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) last_writer[b] = inst; + } + } else if (t.kind == TraceRec::MEMORY_BAR) { + // The explicit async-DMA sync. Pair with the CURRENT load for this (tag_id, + // tag_slot), reusing its uniq: the dma releases the bar at issue, the bar parks on + // the tag until resp-complete, and consumers of the buffer gate on the bar. + auto it = current_dma.find({t.tag_id, t.tag_slot}); + int64_t uniq = next_tag++; // fallback if unpaired + std::shared_ptr dma_inst; + if (it != current_dma.end()) { uniq = it->second.first; dma_inst = it->second.second; } + auto bar = make_mem_bar(t, uniq); + if (dma_inst) dma_inst->add_child(bar); + tile->append_instuction(bar); + for (int64_t b : t.write_bufs) last_writer[b] = bar; + } else if (t.kind == TraceRec::COMPUTE) { + auto inst = make_compute(t); + link(inst, t.read_bufs, t.write_bufs); + if (is_async_compute(t.compute_type)) outstanding_async.push_back(inst); + } else if (t.kind == TraceRec::COMPUTE_BAR) { + // explicit compute fence: ready once all outstanding async compute have + // ISSUED (pipeline-child release); the Core then waits the SA pipelines to + // drain before it finishes (-> the store it gates). + auto bar = std::make_shared(Opcode::COMPUTE_BAR); + for (auto& a : outstanding_async) a->add_pipeline_child(bar); + outstanding_async.clear(); + tile->append_instuction(bar); + pending_bar = bar; + } + } + flush(); + return tg; +} From 6bf8f614a6778e692eab1e5e8bd92618c87c27a2 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 24 Jun 2026 22:35:51 +0900 Subject: [PATCH 03/32] [TOGSim] Work-item outlining and ABI v12 dispatch togsim_dispatch with TILE_BEGIN/TILE_END; outline each work-item into togsim_kernel_tile. --- .../mlir/passes/lower_to_emitc.py | 195 ++++++++++++++++-- PyTorchSimFrontend/mlir/passes/togsim_ops.py | 12 +- TOGSim/include/togsim_loader.h | 4 +- TOGSim/include/togsim_runtime.h | 16 +- TOGSim/src/togsim_runtime.cc | 15 +- TOGSim/src/togsim_trace_bridge.cc | 12 +- 6 files changed, 216 insertions(+), 38 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index caca822e..30d3c796 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -170,22 +170,181 @@ def walk(b): return found[0] -def _insert_core_alloc(ctx, kernel, ctx_val): - """Insert `togsim_core_alloc(ctx)` at the start of each parallel work-item: - the first op of the innermost PARALLEL loop body (or the kernel entry if the - kernel has no parallel loop -> a single work-item). The runtime binds the - following ops to the returned core (sec 9.3); the producer never names - num_cores. The return value is discarded (no free -- a core is an assignment, - not a held resource).""" - block = kernel.regions[0].blocks[0] - target = _innermost_outer_loop(block) - body = target.operation.regions[0].blocks[0] if target is not None else block - ir.Operation.create( - "emitc.call_opaque", results=[], operands=[ctx_val], - attributes={"callee": ir.StringAttr.get(ts.CORE_ALLOC_CALLEE), - "args": ir.ArrayAttr.get([_idx(0)])}, - loc=ir.Location.unknown(ctx), - ip=ir.InsertionPoint.at_block_begin(body)) +def _is_outer(forop): + a = forop.operation.attributes + return "outer_loop" in a and ir.BoolAttr(a["outer_loop"]).value + + +def _parallel_loop_chain(block): + """The nested chain of `affine.for {outer_loop}` from `block` inward (one + work-item's parallel indices). Empty if the kernel has no parallel loop.""" + chain = [] + cur = block + while True: + nxt = None + for op in cur.operations: + if op.operation.name == "affine.for" and _is_outer(op): + nxt = op + break + if nxt is None: + break + chain.append(nxt) + cur = nxt.operation.regions[0].blocks[0] + return chain + + +def _const_op(value): + """The defining arith/emitc constant Operation if `value` is a constant + result, else None (block args / other ops).""" + owner = value.owner + if isinstance(owner, ir.Block): + return None + return owner if owner.name in ("arith.constant", "emitc.constant") else None + + +def _outline_work_item(ctx, kernel, ctx_val): + """Outline the innermost parallel work-item body into a uniform + `togsim_kernel_tile(ctx, iv, n)` func, replacing it with a + `togsim_dispatch(ctx, togsim_kernel_tile, iv, n)` call (sec 9.3). The + work-item SCOPE becomes the function body; the runtime wrapper owns the + core-alloc + the TILE_BEGIN/TILE_END boundary (a decorator). One uniform tile + signature -> a single general dispatcher serves every kernel. + + Runs after `_rewrite_togsim_ops`, so the moved body holds emitc.call_opaque + (not togsim.* ops). The only values captured from outside the body are ctx, + the enclosing parallel induction vars, and constants -- threaded via the iv + array (parallel IVs) / cloned (constants); anything else is unsupported + (dynamic shape -> P4).""" + kblk = kernel.regions[0].blocks[0] + chain = _parallel_loop_chain(kblk) + if chain: + L = chain[-1] + Lbody = L.operation.regions[0].blocks[0] + ivs = [c.operation.regions[0].blocks[0].arguments[0] for c in chain] + else: # no parallel loop -> the whole kernel body is one work-item + L = None + Lbody = kblk + ivs = [] + + i64 = ir.IntegerType.get_signless(64) + i32 = ir.IntegerType.get_signless(32) + idxty = ir.IndexType.get() + ctxty = ir.Type.parse(CTX_TYPE, ctx) + i64ptr = ir.Type.parse("!emitc.ptr", ctx) + loc = ir.Location.unknown(ctx) + + # --- the outlined tile function (before the kernel so C defines it first) --- + tile = ir.Operation.create( + "func.func", results=[], regions=1, + attributes={ + "function_type": ir.TypeAttr.get(ir.FunctionType.get([ctxty, i64ptr, i32], [])), + "sym_name": ir.StringAttr.get(ts.TILE_SYMBOL), + "sym_visibility": ir.StringAttr.get("private")}, + loc=loc, ip=ir.InsertionPoint(kernel)) + with loc: + tblk = tile.regions[0].blocks.append(ctxty, i64ptr, i32) + ctx2, iv2, _n2 = tblk.arguments + with ir.InsertionPoint(tblk): + tret = ir.Operation.create("func.return", results=[], operands=[], loc=loc) + + # in the tile fn: recover each parallel index = index_cast(iv[k]). + idx_vals = [] + with ir.InsertionPoint(tret): + for k in range(len(ivs)): + kc = ir.Operation.create("emitc.constant", results=[i64], + attributes={"value": ir.IntegerAttr.get(i64, k)}, loc=loc).results[0] + elem = ir.Operation.create("emitc.subscript", results=[i64], + operands=[iv2, kc], loc=loc).results[0] + idx_vals.append(ir.Operation.create("arith.index_cast", results=[idxty], + operands=[elem], loc=loc).results[0]) + + # move the work-item body into the tile fn (terminators stay behind). + for op in [o for o in Lbody.operations + if o.operation.name not in ("affine.yield", "func.return")]: + op.operation.move_before(tret) + + # remap captures (Value `==` is identity): ctx -> ctx2, each parallel IV -> + # its index_cast, each external constant -> a clone inside the tile fn. A + # constant defined inside the tile fn (moved/read) is internal -> left alone. + caps = [(ctx_val, ctx2)] + list(zip(ivs, idx_vals)) + internal_consts = [] + def _collect_internal(block): + for op in block.operations: + c = _const_op(op.operation.results[0]) if len(op.operation.results) == 1 else None + if c is not None: + internal_consts.append(op.operation.results[0]) + for rg in op.operation.regions: + for b in rg.blocks: + _collect_internal(b) + _collect_internal(tblk) + const_clones = [] + ext_consts = [] + def _find_ext_consts(block): + for op in block.operations: + for opnd in op.operation.operands: + if _const_op(opnd) is None: + continue + if any(opnd == ic for ic in internal_consts): + continue + if any(opnd == e for e in ext_consts): + continue + ext_consts.append(opnd) + for rg in op.operation.regions: + for b in rg.blocks: + _find_ext_consts(b) + _find_ext_consts(tblk) + top = ir.InsertionPoint(tblk.operations[0]) + for e in ext_consts: + c = _const_op(e) + clone = ir.Operation.create(c.name, results=[e.type], + attributes={"value": c.attributes["value"]}, loc=loc, ip=top).results[0] + const_clones.append((e, clone)) + + allcaps = caps + const_clones + def _remap(block): + for op in block.operations: + for i in range(len(op.operation.operands)): + cur = op.operation.operands[i] + for orig, new in allcaps: + if cur == orig: + op.operation.operands[i] = new + break + for rg in op.operation.regions: + for b in rg.blocks: + _remap(b) + _remap(tblk) + + # --- the dispatcher: marshal the IVs and hand the tile fn to togsim_dispatch --- + term = [o for o in Lbody.operations + if o.operation.name in ("affine.yield", "func.return")][0] + fn_ref = _opaque(ctx, ts.TILE_SYMBOL) # function name -> verbatim pointer in C + with ir.InsertionPoint(term): + if ivs: + arrty = ir.Type.parse("!emitc.array<%dxi64>" % len(ivs), ctx) + arr = ir.Operation.create("emitc.variable", results=[arrty], + attributes={"value": _opaque(ctx, "")}, loc=loc).results[0] + for k, iv in enumerate(ivs): + kc = ir.Operation.create("emitc.constant", results=[i64], + attributes={"value": ir.IntegerAttr.get(i64, k)}, loc=loc).results[0] + v64 = ir.Operation.create("arith.index_cast", results=[i64], + operands=[iv], loc=loc).results[0] + sub = ir.Operation.create("emitc.subscript", results=[i64], + operands=[arr, kc], loc=loc).results[0] + # emitc.assign operands are (lvalue dest, value). + ir.Operation.create("emitc.assign", results=[], operands=[sub, v64], loc=loc) + ir.Operation.create( + "emitc.call_opaque", results=[], operands=[ctx_val, arr], + attributes={"callee": ir.StringAttr.get(ts.DISPATCH_CALLEE), + "args": ir.ArrayAttr.get( + [_idx(0), fn_ref, _idx(1), ir.IntegerAttr.get(i32, len(ivs))])}, + loc=loc) + else: + ir.Operation.create( + "emitc.call_opaque", results=[], operands=[ctx_val], + attributes={"callee": ir.StringAttr.get(ts.DISPATCH_CALLEE), + "args": ir.ArrayAttr.get( + [_idx(0), fn_ref, _opaque(ctx, "nullptr"), ir.IntegerAttr.get(i32, 0)])}, + loc=loc) def _rewrite_togsim_ops(ctx, kernel, ctx_val): @@ -337,8 +496,8 @@ def lower_to_emitc(skeleton_module): _strip_aux(skeleton_module) ctx_val = _rewrite_signature(kernel, ctx) - _insert_core_alloc(ctx, kernel, ctx_val) # core_alloc per work-item - _rewrite_togsim_ops(ctx, kernel, ctx_val) + _rewrite_togsim_ops(ctx, kernel, ctx_val) # togsim.* -> emitc.call_opaque + _outline_work_item(ctx, kernel, ctx_val) # work-item body -> togsim_kernel_tile + dispatch PassManager.parse(_PIPELINE, ctx).run(skeleton_module.operation) diff --git a/PyTorchSimFrontend/mlir/passes/togsim_ops.py b/PyTorchSimFrontend/mlir/passes/togsim_ops.py index 85b89757..38104d15 100644 --- a/PyTorchSimFrontend/mlir/passes/togsim_ops.py +++ b/PyTorchSimFrontend/mlir/passes/togsim_ops.py @@ -69,10 +69,14 @@ #: producer entry-point symbol the TOGSim loader resolves (see togsim_runtime.h). ENTRY_SYMBOL = "togsim_kernel" -#: runtime callee emitted directly by lower_to_emitc (not a skeleton op): the -#: per-work-item core allocation. See togsim_cpp_trace.md sec 9.3. Kept in -#: lockstep with togsim_runtime.h. -CORE_ALLOC_CALLEE = "togsim_core_alloc" +#: outlined per-work-item function the dispatcher hands to togsim_dispatch +#: (uniform signature (ctx, int64* iv, i32 n); see togsim_cpp_trace.md sec 9.3). +TILE_SYMBOL = "togsim_kernel_tile" + +#: runtime callees emitted directly by lower_to_emitc (not skeleton ops), kept in +#: lockstep with togsim_runtime.h. DISPATCH_CALLEE is the per-work-item wrapper the +#: dispatcher loop calls, with TILE_SYMBOL as its function pointer. +DISPATCH_CALLEE = "togsim_dispatch" # ---- attribute keys ------------------------------------------------------- ATTR_DIR = "dir" # i32: DIR_LOAD | DIR_STORE diff --git a/TOGSim/include/togsim_loader.h b/TOGSim/include/togsim_loader.h index 52c1ac1e..df8063a5 100644 --- a/TOGSim/include/togsim_loader.h +++ b/TOGSim/include/togsim_loader.h @@ -12,8 +12,8 @@ namespace togsim { // One modeled instruction recorded by the runtime callbacks. struct TraceRec { - enum Kind { DISPATCH, DMA, COMPUTE, MEMORY_BAR, COMPUTE_BAR } kind; - int32_t core; // work-item -> core binding (set by togsim_core_alloc) + enum Kind { TILE_BEGIN, TILE_END, DMA, COMPUTE, MEMORY_BAR, COMPUTE_BAR } kind; + int32_t core; // work-item -> core binding (set by togsim_dispatch) // DMA / MEMORY_BAR int32_t dir; // togsim_dma_dir int32_t arg_id; // tensor diff --git a/TOGSim/include/togsim_runtime.h b/TOGSim/include/togsim_runtime.h index f9c53c1c..fe069e64 100644 --- a/TOGSim/include/togsim_runtime.h +++ b/TOGSim/include/togsim_runtime.h @@ -11,7 +11,7 @@ extern "C" { // Producer/runtime ABI version. TOGSim refuses to load a producer whose // embedded togsim_abi_version() does not match TOGSIM_ABI_VERSION. -#define TOGSIM_ABI_VERSION 11 +#define TOGSIM_ABI_VERSION 12 int32_t togsim_abi_version(void); // Opaque per-invocation context owned by TOGSim. Holds the record sink and the @@ -49,10 +49,16 @@ void togsim_compute(EmitCtx* ctx, uint64_t tile_id, int32_t compute_type, void togsim_memory_barrier(EmitCtx* ctx, int32_t tag_id, uint64_t tag_slot, const int64_t* write_bufs, int32_t n_write); -// Core allocation (sec 9.3): the producer calls this at each parallel work-item's -// start, and the ops that follow bind to the returned core. No free -- a core is an -// assignment. The producer never names num_cores; the runtime owns the pool. -int32_t togsim_core_alloc(EmitCtx* ctx); +// A parallel work-item body, outlined by the producer (sec 9.3): `iv` holds the +// packed parallel loop indices (e.g. the (m,n) output-tile indices). One uniform +// signature => one general dispatcher serves every kernel. The runtime only reads iv. +typedef void (*togsim_tile_fn)(EmitCtx* ctx, int64_t* iv, int32_t n_iv); + +// Dispatch one work-item (sec 9.3): round-robin a core, bracket `fn` with +// TILE_BEGIN/TILE_END, and invoke it -- so the work-item scope IS the call. Core +// choice is runtime-owned; the producer never names num_cores or a core. +void togsim_dispatch(EmitCtx* ctx, togsim_tile_fn fn, + int64_t* iv, int32_t n_iv); // Compute fence: drain in-flight async compute (the systolic-array matmuls) // before the following op (a store) consumes their result. Explicit barrier in diff --git a/TOGSim/src/togsim_runtime.cc b/TOGSim/src/togsim_runtime.cc index 84cb2cd2..df952a15 100644 --- a/TOGSim/src/togsim_runtime.cc +++ b/TOGSim/src/togsim_runtime.cc @@ -39,12 +39,14 @@ extern "C" { int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } -int32_t togsim_core_alloc(EmitCtx* ctx) { - // Round-robin a core from the pool; the producer never sees num_cores. Binds - // it as the current core for the ops that follow (the work-item's reduction). +void togsim_dispatch(EmitCtx* ctx, togsim_tile_fn fn, int64_t* iv, int32_t n_iv) { + // Work-item wrapper (sec 9.3): round-robin a core (the producer never sees + // num_cores), bracket the work-item with TILE_BEGIN/TILE_END, and run its body. The + // work-item scope is exactly this call; ops emit records under ctx->cur_core. ctx->cur_core = ctx->num_cores > 0 ? (ctx->rr++ % ctx->num_cores) : 0; - ctx->trace.push_back(blank(togsim::TraceRec::DISPATCH, ctx->cur_core)); - return ctx->cur_core; + ctx->trace.push_back(blank(togsim::TraceRec::TILE_BEGIN, ctx->cur_core)); + fn(ctx, iv, n_iv); + ctx->trace.push_back(blank(togsim::TraceRec::TILE_END, ctx->cur_core)); } void togsim_dma(EmitCtx* ctx, int32_t dir, int32_t arg_id, @@ -169,7 +171,8 @@ SimResult simulate(const RunResult& run, const TimingParams& params) { out.n_compute++; break; } - case TraceRec::DISPATCH: + case TraceRec::TILE_BEGIN: + case TraceRec::TILE_END: case TraceRec::COMPUTE_BAR: break; // work-item boundary / compute fence: no cost in this reference timer } diff --git a/TOGSim/src/togsim_trace_bridge.cc b/TOGSim/src/togsim_trace_bridge.cc index 56e85a68..2417fd4a 100644 --- a/TOGSim/src/togsim_trace_bridge.cc +++ b/TOGSim/src/togsim_trace_bridge.cc @@ -123,15 +123,21 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, }; for (const auto& t : run.trace) { - if (t.kind == TraceRec::DISPATCH) { - // new work-item -> new subgraph (bound to its core) + tile. + 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 + // brackets the tile fn call), not until the next begin. flush(); sg = std::make_shared(); sg->set_core_id(t.core); tile = std::make_shared(Tile::Status::INITIALIZED); continue; } - if (!tile) continue; // defensive: ops before the first core_alloc + if (t.kind == TraceRec::TILE_END) { + flush(); // close the work-item explicitly (scope = the tile fn call) + continue; + } + if (!tile) continue; // defensive: ops before the first TILE_BEGIN if (t.kind == TraceRec::DMA) { int64_t uniq = next_tag++; // fresh Core tag key per dma record From 8f9e389e31e7200cc190379667e47a37f9fe7cc6 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 24 Jun 2026 22:35:51 +0900 Subject: [PATCH 04/32] [TOGSim] SRAM-capacity and SA weight-buffer throttle for the trace path DMA-capacity throttle and frozen-state guard, per-core VMEM in the configs, and the SA weight-buffer throttle. --- TOGSim/include/Core.h | 31 +++ TOGSim/include/Instruction.h | 28 +++ TOGSim/include/SimulationConfig.h | 8 + TOGSim/include/Simulator.h | 3 + TOGSim/src/Common.cc | 4 + TOGSim/src/Core.cc | 110 +++++++++- TOGSim/src/Simulator.cc | 31 +++ TOGSim/src/togsim_trace_bridge.cc | 38 ++++ .../systolic_ws_128x128_c1_booksim_tpuv2.yml | 3 + .../systolic_ws_128x128_c1_booksim_tpuv3.yml | 3 + ...ystolic_ws_128x128_c1_simple_noc_tpuv2.yml | 3 + ...ystolic_ws_128x128_c1_simple_noc_tpuv3.yml | 3 + ...ic_ws_128x128_c1_simple_noc_tpuv3_half.yml | 3 + ...28x128_c1_simple_noc_tpuv3_timing_only.yml | 3 + ...ystolic_ws_128x128_c1_simple_noc_tpuv4.yml | 3 + .../systolic_ws_128x128_c2_booksim_tpuv3.yml | 3 + ...ws_128x128_c2_booksim_tpuv3_bw_quarter.yml | 3 + .../systolic_ws_128x128_c2_chiplet_tpuv3.yml | 3 + ...olic_ws_128x128_c2_chiplet_tpuv3_xnuma.yml | 3 + ...ystolic_ws_128x128_c2_simple_noc_tpuv2.yml | 3 + ...ystolic_ws_128x128_c2_simple_noc_tpuv3.yml | 3 + ...lic_ws_128x128_c2_simple_noc_tpuv3_ils.yml | 3 + ..._128x128_c2_simple_noc_tpuv3_partition.yml | 3 + ...ystolic_ws_128x128_c2_simple_noc_tpuv4.yml | 3 + configs/systolic_ws_8x8_c1_booksim.yml | 3 + configs/systolic_ws_8x8_c1_simple_noc.yml | 3 + scripts/trace_timeline.py | 199 ++++++++++++++++++ 27 files changed, 505 insertions(+), 1 deletion(-) create mode 100644 scripts/trace_timeline.py diff --git a/TOGSim/include/Core.h b/TOGSim/include/Core.h index 286feb5f..b611eff1 100644 --- a/TOGSim/include/Core.h +++ b/TOGSim/include/Core.h @@ -1,6 +1,7 @@ #pragma once #include #include +#include #include #include #include @@ -24,6 +25,10 @@ class Core { Core(uint32_t id, SimulationConfig config); ~Core()=default; virtual bool running(); + // True if this core has work actively in flight (DMA / compute pipeline / queues) + // that will produce a future finish event -- i.e. running() minus "tiles waiting". + // Used by the frozen-state (spad-too-small) guard. + bool has_inflight(); virtual bool can_issue(const std::shared_ptr& op); virtual void issue(std::shared_ptr tile); virtual std::shared_ptr pop_finished_tile(); @@ -55,6 +60,16 @@ class Core { void sa_cycle(); bool can_issue_compute(std::shared_ptr& inst); void update_stats(); + // SRAM-capacity throttle (sec 10.4): a consumer frees the buffer-versions it + // read (refcount -> 0 releases the spad bytes). Called when COMP/MOVOUT issue. + void release_sram(const std::shared_ptr& inst); + // SA weight-buffer throttle (sec 10.4): pick a systolic array that has a free + // weight slot (round-robin among free); -1 if all full -> the preload stalls. + int pick_free_weight_sa(); + // Free weight slots due this cycle: a matmul releases its slot at its + // streaming-end (finish - overlapping, when it stops reading the weight), + // scheduled at issue in _weight_release_q. Last consumer frees it. + void process_weight_releases(); /* Core id & config file */ const uint32_t _id; @@ -103,4 +118,20 @@ class Core { std::queue _request_queue; std::queue _response_queue; uint32_t _waiting_write_reqs; + + // SRAM-capacity throttle (sec 10.4). _sram_used = current per-core spad bytes; + // _sram_capacity = limit (0 = disabled); _sram_allocs maps a buffer-version id + // to its accumulated footprint bytes (freed when its last reader issues). + size_t _sram_used = 0; + size_t _sram_capacity = 0; + std::unordered_map _sram_allocs; + + // SA weight-buffer throttle (sec 10.4). _weight_slots_used[s] = weights resident + // on SA s (loaded by a preload, not yet freed by their last matmul); + // _weight_slot_depth = per-SA capacity (0 = disabled -> plain round-robin). + std::vector _weight_slots_used; + uint32_t _weight_slot_depth = 0; + // Pending weight-slot releases keyed by cycle (each matmul's streaming-end); + // process_weight_releases() drains those due and decrements the token. + std::multimap> _weight_release_q; }; \ No newline at end of file diff --git a/TOGSim/include/Instruction.h b/TOGSim/include/Instruction.h index 3740b53a..dc237580 100644 --- a/TOGSim/include/Instruction.h +++ b/TOGSim/include/Instruction.h @@ -17,6 +17,11 @@ // results of async matmuls. enum class Opcode { MOVIN, MOVOUT, COMP, MEMORY_BAR, COMPUTE_BAR, COUNT}; +// One weight slot on systolic array `sa` (sec 10.4). A preload sets refcount = +// the matmuls reusing the weight; each frees it at its streaming-end, the last +// one releases the slot. Shared (shared_ptr) by the preload's matmul consumers. +struct WeightToken { int sa; int refcount; }; + typedef uint64_t addr_type; typedef uint64_t cycle_type; @@ -37,6 +42,13 @@ class Instruction : public std::enable_shared_from_this { // successor overlaps it instead of waiting its full latency (sec 10.7). void add_pipeline_child(std::shared_ptr child); void release_pipeline_children(); + // SA weight-buffer model: the SA this op is pinned to (a preload picks it, its + // matmul consumers inherit it) and the shared weight slot the matmuls release. + const std::set>& get_pipeline_children() { return _pipeline_children; } + void set_assigned_sa(int s) { _assigned_sa = s; } + int get_assigned_sa() const { return _assigned_sa; } + void set_weight_token(const std::shared_ptr& t) { _weight_token = t; } + const std::shared_ptr& get_weight_token() const { return _weight_token; } bool check_ready() { return ready_counter == 0; } const Opcode get_opcode() { return opcode; } bool is_dma_read() { return opcode == Opcode::MOVIN; } @@ -94,6 +106,16 @@ class Instruction : public std::enable_shared_from_this { std::set>& get_child_inst() { return child_inst; } uint64_t get_global_inst_id() const { return _global_inst_id; } + // SRAM-capacity model (sec 10.4), filled by the bridge and enforced by Core: a + // load fills buffer-version `_sram_alloc_id` (-1 = untracked), and a version's + // 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); } + const std::vector& get_sram_release() const { return _sram_release_allocs; } + // bytes this load occupies in the spad (from the tile it moves in). + size_t sram_footprint() const { return _tile_numel * (_elem_bits / 8); } + cycle_type start_cycle; cycle_type finish_cycle; cycle_type bubble_cycle=0; @@ -132,4 +154,10 @@ class Instruction : public std::enable_shared_from_this { bool _is_indirect_mode=false; bool _is_sparse_inst=false; std::string _indirect_index_path=""; + // SRAM-capacity model (see the setters above). + int64_t _sram_alloc_id = -1; + std::vector _sram_release_allocs; + // SA weight-buffer model (see the setters above). + int _assigned_sa = -1; + std::shared_ptr _weight_token; }; \ No newline at end of file diff --git a/TOGSim/include/SimulationConfig.h b/TOGSim/include/SimulationConfig.h index 2ef08618..c099d057 100644 --- a/TOGSim/include/SimulationConfig.h +++ b/TOGSim/include/SimulationConfig.h @@ -27,6 +27,14 @@ struct SimulationConfig { uint32_t num_systolic_array_per_core = 1; uint32_t num_stonne_per_core = 1; uint32_t num_stonne_port = 1; + // Per-core VMEM/spad capacity (KB) for the trace-path DMA throttle (sec 10.4): a + // load that would overflow the spad waits for a consumer to free a tile. 0 = unset + // -> disabled. Legacy TileGraphParser insts have alloc id -1 and are never gated. + uint32_t core_spad_size_kb = 0; + // SA weight-buffer depth (sec 10.4): weight tiles a systolic array holds; a + // preload stalls until a slot frees (its matmuls finished). 2 = weight + // double-buffer (convention default, tunable). 0 = disabled. + uint32_t sa_weight_buffer_depth = 2; /* DRAM config */ DramType dram_type; diff --git a/TOGSim/include/Simulator.h b/TOGSim/include/Simulator.h index e3542d51..91baf5b5 100644 --- a/TOGSim/include/Simulator.h +++ b/TOGSim/include/Simulator.h @@ -48,6 +48,9 @@ class Simulator { void dram_cycle(); void icnt_cycle(); bool running(); + // Spad-too-small guard: if the sim stays frozen (running() but nothing in + // flight) past kWedgeThreshold cycles, error out and exit. Called each cycle. + void check_frozen(); void set_cycle_mask(); uint32_t get_dest_node(mem_fetch *access); SimulationConfig _config; diff --git a/TOGSim/src/Common.cc b/TOGSim/src/Common.cc index 3f84d885..6f9a74d7 100644 --- a/TOGSim/src/Common.cc +++ b/TOGSim/src/Common.cc @@ -64,6 +64,10 @@ SimulationConfig initialize_config(const YAML::Node& config, parsed_config.core_freq_mhz = get_config_value(config, "core_freq_mhz"); if (config["num_systolic_array_per_core"]) parsed_config.num_systolic_array_per_core = config["num_systolic_array_per_core"].as(); + if (config["core_spad_size_kb"]) + parsed_config.core_spad_size_kb = config["core_spad_size_kb"].as(); + if (config["sa_weight_buffer_depth"]) + parsed_config.sa_weight_buffer_depth = config["sa_weight_buffer_depth"].as(); if (config["num_stonne_per_core"]) parsed_config.num_stonne_per_core = config["num_stonne_per_core"].as(); if (config["num_stonne_port"]) diff --git a/TOGSim/src/Core.cc b/TOGSim/src/Core.cc index 8a728318..dfe40745 100644 --- a/TOGSim/src/Core.cc +++ b/TOGSim/src/Core.cc @@ -17,6 +17,42 @@ Core::Core(uint32_t id, SimulationConfig config) _stat_sa_compute_idle_cycle.resize(_num_systolic_array_per_core); _stat_inst_count.resize(static_cast(Opcode::COUNT), 0); _stat_tot_skipped_inst.resize(static_cast(Opcode::COUNT), 0); + _sram_capacity = (size_t)config.core_spad_size_kb * 1024; // 0 = throttle disabled + _weight_slot_depth = config.sa_weight_buffer_depth; // 0 = disabled (plain rr) + _weight_slots_used.resize(_num_systolic_array_per_core, 0); +} + +// Round-robin a systolic array that still has a free weight slot; -1 if all full +// (the preload must stall). Advances _systolic_array_rr past the chosen SA. +int Core::pick_free_weight_sa() { + for (uint32_t i = 0; i < _num_systolic_array_per_core; i++) { + uint32_t s = (_systolic_array_rr + i) % _num_systolic_array_per_core; + if (_weight_slots_used[s] < (int)_weight_slot_depth) { + _systolic_array_rr = (s + 1) % _num_systolic_array_per_core; + return (int)s; + } + } + return -1; +} + +void Core::process_weight_releases() { + while (!_weight_release_q.empty() && _weight_release_q.begin()->first <= _core_cycle) { + auto tok = _weight_release_q.begin()->second; + _weight_release_q.erase(_weight_release_q.begin()); + if (--tok->refcount <= 0) _weight_slots_used[tok->sa]--; // last reader frees the slot + } +} + +// The LAST reader of a buffer-version issued (bridge tags only that consumer): +// free the version's bytes back to the per-core spad. +void Core::release_sram(const std::shared_ptr& inst) { + if (!_sram_capacity) return; + for (int64_t id : inst->get_sram_release()) { + auto it = _sram_allocs.find(id); + if (it == _sram_allocs.end()) continue; + _sram_used -= it->second; + _sram_allocs.erase(it); + } } bool Core::can_issue(const std::shared_ptr& op) { @@ -200,6 +236,8 @@ void Core::cycle() { /* Increase core cycle counter */ _core_cycle++; + process_weight_releases(); // free weight slots due this cycle before dispatch + /* Iterate tile while an instruction is issued */ bool issued = false; @@ -240,6 +278,19 @@ void Core::cycle() { _stat_tot_skipped_inst.at(static_cast(inst->get_opcode()))++; break; } else { + // SRAM-capacity gate: a load that would overflow the per-core spad does + // not issue this cycle -- it stays in the ready queue until a consumer + // frees a tile. On issue it occupies its buffer-version allocation. + if (_sram_capacity && inst->get_sram_alloc() >= 0) { + size_t F = inst->sram_footprint(); + // Stall if the tile does not fit the free spad now. If it can never + // fit, Simulator::cycle() detects the frozen state and exits with a + // "spad too small" error rather than looping forever. + if (_sram_used + F > _sram_capacity) + break; // not issued -> retry next cycle + _sram_used += F; + _sram_allocs[inst->get_sram_alloc()] += F; // accumulate version footprint + } core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15( @@ -254,6 +305,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, _id, TraceLogTag::pad15(TraceLogTag::kInstructionIssued), @@ -265,11 +317,46 @@ void Core::cycle() { break; case Opcode::COMP: { + const int ct = inst->get_compute_type(); + // SA selection + weight-buffer gate: a preload picks a systolic array with + // a free weight slot and pins its matmul consumers to it (they free the slot + // on finish). This bounds preload run-ahead and keeps matmuls on their SA. + 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_pipeline_children()) + if (c->get_compute_type() == MATMUL) n_consumers++; + if (_weight_slot_depth > 0 && n_consumers > 0) { + 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_pipeline_children()) + if (c->get_compute_type() == MATMUL) { + c->set_assigned_sa(sa_idx); + c->set_weight_token(tok); + } + } else { // disabled / no consumers -> plain rr + sa_idx = _systolic_array_rr; + _systolic_array_rr = (_systolic_array_rr + 1) % _num_systolic_array_per_core; + } + } else { // MATMUL + sa_idx = inst->get_assigned_sa(); + if (sa_idx < 0) { // no preload pinned it -> rr fallback + sa_idx = _systolic_array_rr; + _systolic_array_rr = (_systolic_array_rr + 1) % _num_systolic_array_per_core; + } + } + inst->set_assigned_sa(sa_idx); // record the SA actually used (for the trace) + } + release_sram(inst); // consumer issued -> free the tiles it read // sec 10.7: this op is now entering the pipeline -> release its // occupancy (pipeline) dependents so a preload/matmul successor // overlaps it instead of waiting its full latency. inst->release_pipeline_children(); - auto& target_pipeline = get_compute_pipeline(inst->get_compute_type()); + auto& target_pipeline = (ct == VECTOR_UNIT) ? _vu_compute_pipeline + : _sa_compute_pipeline.at(sa_idx); if (target_pipeline.empty()) { inst->finish_cycle = _core_cycle + inst->get_compute_cycle(); inst->bubble_cycle = inst->get_overlapping_cycle(); @@ -280,6 +367,14 @@ void Core::cycle() { inst->bubble_cycle = bubble_cycle; } + // Release this matmul's weight slot at its streaming-end (finish - + // overlapping), not at full finish (the drain tail does not read it). + if (ct == MATMUL && inst->get_weight_token()) { + cycle_type rel = inst->finish_cycle > inst->get_overlapping_cycle() + ? inst->finish_cycle - inst->get_overlapping_cycle() : _core_cycle; + _weight_release_q.emplace(rel, inst->get_weight_token()); + } + if (inst->get_compute_cycle() == 0) { inst->finish_instruction(); static_cast(inst->get_owner())->inc_finished_inst(); @@ -409,6 +504,19 @@ void Core::finish_instruction(std::shared_ptr& inst, InstFinishTrac core_trace_log::format_instruction_detail_line(*inst)); } +bool Core::has_inflight() { + // running() without the "_tiles.size() > 0" term: work that will produce a + // finish event on its own (so the sim is NOT frozen). If this is false but + // tiles remain, only stalled ready instructions are left. + if (!_vu_compute_pipeline.empty()) return true; + for (int i = 0; i < _num_systolic_array_per_core; i++) + if (!_sa_compute_pipeline.at(i).empty()) return true; + if (!_dma_waiting_queue.empty() || !_dma_finished_queue.empty()) return true; + if (!_dma.empty()) return true; + if (!_ld_inst_queue.empty() || !_st_inst_queue.empty()) return true; + return false; +} + bool Core::running() { bool running = false; running = running || _tiles.size() > 0; diff --git a/TOGSim/src/Simulator.cc b/TOGSim/src/Simulator.cc index d987d787..17320fa0 100644 --- a/TOGSim/src/Simulator.cc +++ b/TOGSim/src/Simulator.cc @@ -184,6 +184,35 @@ void Simulator::icnt_cycle() { _icnt->cycle(); } +// Consecutive frozen cycles tolerated before declaring the sim wedged (spad too +// small). Generous so transient idle never false-fires; a true freeze is constant. +static constexpr uint64_t kWedgeThreshold = 5000; + +// Frozen-state guard: work remains but nothing is in flight to advance it, because +// the kernel's working set exceeds the whole per-core spad (core_spad_size_kb too +// small). The state repeats every cycle, so error out instead of looping forever. +void Simulator::check_frozen() { + static uint64_t stuck = 0; + // In flight = anything that will produce a future state change: icnt/dram busy, + // a core with DMA/compute pending, or a tile still schedulable. + bool inflight = _icnt->running() || _dram->running(); + for (int id = 0; id < _n_cores && !inflight; id++) { + if (_cores[id]->has_inflight()) inflight = true; + else if (!get_partition_scheduler(id)->empty(id)) inflight = true; + } + if (running() && !inflight) { + if (++stuck > kWedgeThreshold) { + spdlog::error("[Simulator] simulation wedged at cycle {}: work remains but " + "nothing is in flight -- the per-core spad (core_spad_size_kb) " + "is too small to hold a kernel's working set. Increase it.", + _core_cycles); + exit(EXIT_FAILURE); + } + } else { + stuck = 0; + } +} + void Simulator::cycle() { while (running() || _core_cycles < 1) { set_cycle_mask(); @@ -198,6 +227,8 @@ void Simulator::cycle() { // Interconnect cycle if (IS_ICNT_CYCLE(_cycle_mask)) icnt_cycle(); + + check_frozen(); // spad-too-small guard (errors out if wedged) } for (auto &core: _cores) { core->check_tag(); diff --git a/TOGSim/src/togsim_trace_bridge.cc b/TOGSim/src/togsim_trace_bridge.cc index 2417fd4a..48f1faf5 100644 --- a/TOGSim/src/togsim_trace_bridge.cc +++ b/TOGSim/src/togsim_trace_bridge.cc @@ -122,6 +122,40 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, tile->append_instuction(inst); }; + // SRAM-capacity tracking: 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. Only DMA-loaded buffers are tracked. v1 is single-core. + int64_t next_alloc = 0; + std::map cur_alloc; // buf -> current version id + std::map open_ver; // buf -> version still accepting loads + struct Ver { std::vector> loads, readers; }; + std::map vers; + 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); + }; + 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); + 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) { if (t.kind == TraceRec::TILE_BEGIN) { // togsim_dispatch opened a work-item -> new subgraph (bound to its core) + @@ -156,6 +190,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, } else { link(inst, t.read_bufs, t.write_bufs); } + for (int64_t b : t.read_bufs) sram_on_read(b, inst); // store frees what it drains } else { // LOAD tile->append_instuction(inst); // async load: the CURRENT load for this (tag_id, tag_slot), with a fresh uniq @@ -163,6 +198,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, // 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) last_writer[b] = inst; + for (int64_t b : t.write_bufs) sram_on_load(b, inst); // occupy spad } } else if (t.kind == TraceRec::MEMORY_BAR) { // The explicit async-DMA sync. Pair with the CURRENT load for this (tag_id, @@ -179,6 +215,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, } else if (t.kind == TraceRec::COMPUTE) { auto inst = make_compute(t); link(inst, t.read_bufs, t.write_bufs); + for (int64_t b : t.read_bufs) sram_on_read(b, inst); // frees the tiles it consumes if (is_async_compute(t.compute_type)) outstanding_async.push_back(inst); } else if (t.kind == TraceRec::COMPUTE_BAR) { // explicit compute fence: ready once all outstanding async compute have @@ -192,5 +229,6 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, } } flush(); + sram_finalize(); // readers per version are now final -> set each version's refcount return tg; } diff --git a/configs/systolic_ws_128x128_c1_booksim_tpuv2.yml b/configs/systolic_ws_128x128_c1_booksim_tpuv2.yml index 6d2537d9..7fea374b 100644 --- a/configs/systolic_ws_128x128_c1_booksim_tpuv2.yml +++ b/configs/systolic_ws_128x128_c1_booksim_tpuv2.yml @@ -22,3 +22,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c1_booksim_tpuv3.yml b/configs/systolic_ws_128x128_c1_booksim_tpuv3.yml index f830419b..3a96b588 100644 --- a/configs/systolic_ws_128x128_c1_booksim_tpuv3.yml +++ b/configs/systolic_ws_128x128_c1_booksim_tpuv3.yml @@ -26,3 +26,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c1_simple_noc_tpuv2.yml b/configs/systolic_ws_128x128_c1_simple_noc_tpuv2.yml index 1a8c60f6..41e267b6 100644 --- a/configs/systolic_ws_128x128_c1_simple_noc_tpuv2.yml +++ b/configs/systolic_ws_128x128_c1_simple_noc_tpuv2.yml @@ -25,3 +25,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.yml b/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.yml index ff976784..397f0fb7 100644 --- a/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.yml +++ b/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.yml @@ -26,3 +26,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c1_simple_noc_tpuv3_half.yml b/configs/systolic_ws_128x128_c1_simple_noc_tpuv3_half.yml index 2ed1bb12..f080fc69 100644 --- a/configs/systolic_ws_128x128_c1_simple_noc_tpuv3_half.yml +++ b/configs/systolic_ws_128x128_c1_simple_noc_tpuv3_half.yml @@ -26,3 +26,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c1_simple_noc_tpuv3_timing_only.yml b/configs/systolic_ws_128x128_c1_simple_noc_tpuv3_timing_only.yml index 1bcc9bb3..f89661b8 100644 --- a/configs/systolic_ws_128x128_c1_simple_noc_tpuv3_timing_only.yml +++ b/configs/systolic_ws_128x128_c1_simple_noc_tpuv3_timing_only.yml @@ -26,3 +26,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 8 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c1_simple_noc_tpuv4.yml b/configs/systolic_ws_128x128_c1_simple_noc_tpuv4.yml index 39d195b0..ca69d930 100644 --- a/configs/systolic_ws_128x128_c1_simple_noc_tpuv4.yml +++ b/configs/systolic_ws_128x128_c1_simple_noc_tpuv4.yml @@ -28,3 +28,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_booksim_tpuv3.yml b/configs/systolic_ws_128x128_c2_booksim_tpuv3.yml index bf01913b..b7b03e7a 100644 --- a/configs/systolic_ws_128x128_c2_booksim_tpuv3.yml +++ b/configs/systolic_ws_128x128_c2_booksim_tpuv3.yml @@ -26,3 +26,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_booksim_tpuv3_bw_quarter.yml b/configs/systolic_ws_128x128_c2_booksim_tpuv3_bw_quarter.yml index 8c71c528..903ffcbc 100644 --- a/configs/systolic_ws_128x128_c2_booksim_tpuv3_bw_quarter.yml +++ b/configs/systolic_ws_128x128_c2_booksim_tpuv3_bw_quarter.yml @@ -34,3 +34,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_chiplet_tpuv3.yml b/configs/systolic_ws_128x128_c2_chiplet_tpuv3.yml index d058f188..6a234017 100644 --- a/configs/systolic_ws_128x128_c2_chiplet_tpuv3.yml +++ b/configs/systolic_ws_128x128_c2_chiplet_tpuv3.yml @@ -28,3 +28,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_chiplet_tpuv3_xnuma.yml b/configs/systolic_ws_128x128_c2_chiplet_tpuv3_xnuma.yml index 019a0f0f..f0546e56 100644 --- a/configs/systolic_ws_128x128_c2_chiplet_tpuv3_xnuma.yml +++ b/configs/systolic_ws_128x128_c2_chiplet_tpuv3_xnuma.yml @@ -27,3 +27,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_simple_noc_tpuv2.yml b/configs/systolic_ws_128x128_c2_simple_noc_tpuv2.yml index 348babae..08ec26ac 100644 --- a/configs/systolic_ws_128x128_c2_simple_noc_tpuv2.yml +++ b/configs/systolic_ws_128x128_c2_simple_noc_tpuv2.yml @@ -25,3 +25,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_simple_noc_tpuv3.yml b/configs/systolic_ws_128x128_c2_simple_noc_tpuv3.yml index a0985aec..a6e073e9 100644 --- a/configs/systolic_ws_128x128_c2_simple_noc_tpuv3.yml +++ b/configs/systolic_ws_128x128_c2_simple_noc_tpuv3.yml @@ -26,3 +26,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_ils.yml b/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_ils.yml index 166e2e25..5436b3e8 100644 --- a/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_ils.yml +++ b/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_ils.yml @@ -29,3 +29,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_partition.yml b/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_partition.yml index 6119e83d..d928f9d3 100644 --- a/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_partition.yml +++ b/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_partition.yml @@ -30,3 +30,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_128x128_c2_simple_noc_tpuv4.yml b/configs/systolic_ws_128x128_c2_simple_noc_tpuv4.yml index 9100c22a..dd9dfac7 100644 --- a/configs/systolic_ws_128x128_c2_simple_noc_tpuv4.yml +++ b/configs/systolic_ws_128x128_c2_simple_noc_tpuv4.yml @@ -28,3 +28,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core VMEM (vector/scratchpad) size: TPUv2/v3/v4 = 16 MB. +core_spad_size_kb: 16384 diff --git a/configs/systolic_ws_8x8_c1_booksim.yml b/configs/systolic_ws_8x8_c1_booksim.yml index f46d380e..1593e148 100644 --- a/configs/systolic_ws_8x8_c1_booksim.yml +++ b/configs/systolic_ws_8x8_c1_booksim.yml @@ -23,3 +23,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core spad: 8x8 array, 128 KB x 8 = 1 MB. +core_spad_size_kb: 1024 diff --git a/configs/systolic_ws_8x8_c1_simple_noc.yml b/configs/systolic_ws_8x8_c1_simple_noc.yml index 1be24b85..b2d16c6a 100644 --- a/configs/systolic_ws_8x8_c1_simple_noc.yml +++ b/configs/systolic_ws_8x8_c1_simple_noc.yml @@ -24,3 +24,6 @@ codegen_external_mapping_file: '' codegen_autotune_max_retry: 10 codegen_autotune_template_topk: 4 codegen_compiler_optimization: all + +# Per-core spad: 8x8 array, 128 KB x 8 = 1 MB. +core_spad_size_kb: 1024 diff --git a/scripts/trace_timeline.py b/scripts/trace_timeline.py new file mode 100644 index 00000000..f7d2672e --- /dev/null +++ b/scripts/trace_timeline.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""Convert a TOGSim `--log_level trace` log into a Chrome Trace Event JSON that +opens in Perfetto (https://ui.perfetto.dev) or chrome://tracing as an interactive +timeline (Gantt). + +Each instruction becomes one duration slice on one of 3 per-core lanes: + dma -- MOVIN / MOVOUT + sa -- COMP compute_type 1 (matmul) / 2 (preload) + vector -- COMP compute_type 0 (vector) +grouped per core (pid). Time unit = core cycles. Barriers (MEMORY_BAR/COMPUTE_BAR) +are not drawn. A compute slice's width is its compute_cycle (the op's own latency), +not issue->finish (which balloons under pipeline backlog); a DMA slice is the +actual transfer ASYNC_DMA_ISSUE -> data-ready. + +Usage: + bin/Simulator --config --trace_so --cycle_table --log_level trace \ + 2>&1 | python scripts/trace_timeline.py -o timeline.json + # or + python scripts/trace_timeline.py trace.log -o timeline.json +Then drag timeline.json into https://ui.perfetto.dev . +""" +import argparse +import json +import re +import sys + +# [cycle][Core C][TAG ][INST_ID=N] OPCODE (detail...) +_LINE = re.compile( + r"\[(\d+)\]\[Core (\d+)\]\[([A-Z_]+)\s*\](?:\[INST_ID=(-?\d+)\])?\s*(\w+)?(.*)") + +# Only 3 lanes per core. Barriers are dropped (see _HIDE). +_LANE = {"MOVIN": "dma", "MOVOUT": "dma"} +_HIDE = {"MEMORY_BAR", "COMPUTE_BAR", "TILE_BEGIN", "TILE_END"} +_CT_NAME = {0: "vector", 1: "matmul", 2: "preload"} + + +def _label(opcode, detail): + if opcode == "COMP": + m = re.search(r"compute_type=(\d+)", detail) + ct = int(m.group(1)) if m else -1 + return _CT_NAME.get(ct, "comp") + m = re.search(r"addr_name=(\w+)", detail) + return f"{opcode} {m.group(1)}" if m else opcode + + +def _lane(opcode, detail): + if opcode == "COMP": + m = re.search(r"compute_type=(\d+)", detail) + ct = int(m.group(1)) if m else -1 + return "vector" if ct == 0 else "sa" + return _LANE.get(opcode, "dma") + + +def parse(lines): + # key = (core, inst_id) -> record + insts = {} + for ln in lines: + m = _LINE.search(ln) + if not m: + continue + cyc, core, tag, iid, opcode, detail = m.groups() + if iid is None or opcode is None: + continue + cyc, core, iid = int(cyc), int(core), int(iid) + key = (core, iid) + r = insts.setdefault(key, { + "core": core, "iid": iid, "opcode": opcode, "detail": detail, + "issued": None, "finished": None, "resp": None, "dma_issue": None}) + if not r["opcode"] or r["opcode"] == opcode: + r["opcode"] = opcode + if detail.strip(): + r["detail"] = detail + if tag == "INST_ISSUED" and r["issued"] is None: + r["issued"] = cyc + elif tag == "INST_FINISHED": + r["finished"] = cyc + elif tag == "DRAM_RESP_DONE": + r["resp"] = cyc + elif tag == "ASYNC_DMA_ISSUE": # actual transfer start (DMA engine busy) + r["dma_issue"] = cyc + return insts + + +def _occ(detail): + """(compute_cycle, overlapping_cycle) from a COMP detail string.""" + cc = re.search(r"compute_cycle=(\d+)", detail) + ov = re.search(r"overlapping_cycle=(\d+)", detail) + return (int(cc.group(1)) if cc else 0, int(ov.group(1)) if ov else 0) + + +def to_chrome(insts, num_sa=1): + """Model each hardware unit as a server and replay its ops in issue order, so + real idle gaps (bubbles) show and slices don't nest: + dma : MOVIN/MOVOUT -- 1 DMA engine; slice = actual transfer + (ASYNC_DMA_ISSUE -> data-ready). + vector : COMP type 0 -- 1 VPU. + sa : COMP type 1/2 -- num_sa systolic arrays, round-robin by issue order. + A compute slice's width is compute_cycle - overlapping_cycle (its occupancy = + latency minus the tail that overlaps the next op), starting when the unit + actually picks it up: start = max(issue, unit_free). num_sa>1 -> lanes sa0.. .""" + by_core = {} + for r in insts.values(): + op, detail, core = r["opcode"], r["detail"], r["core"] + if op in _HIDE: + continue + u = by_core.setdefault(core, {"dma": [], "vector": [], "sa": []}) + if op == "COMP": + m = re.search(r"compute_type=(\d+)", detail) + ct = int(m.group(1)) if m else -1 + u["vector" if ct == 0 else "sa"].append(r) + else: + u["dma"].append(r) + + events, lanes, cores = [], set(), set() + + def add(core, lane, ts, dur, name, r): + lanes.add((core, lane)) + cores.add(core) + events.append({"name": name, "cat": lane, "ph": "X", "ts": ts, + "dur": max(dur, 1), "pid": core, "tid": lane, + "args": {"inst_id": r["iid"], "issued": r["issued"], + "finished": r["finished"], "data_ready": r["resp"]}}) + + def issue_key(r): + return r["issued"] if r["issued"] is not None else 0 + + nsa = max(num_sa, 1) + for core, u in sorted(by_core.items()): + # DMA engine: one server, serialized. A load occupies it only while INJECTING + # requests -- [INST_ISSUED, ASYNC_DMA_ISSUE] -- not the response tail. So a load + # blocked on a full spad leaves a real idle gap = the SRAM throttle stalling it. + free = 0 + for r in sorted(u["dma"], key=issue_key): + start = r["issued"] if r["issued"] is not None else r["dma_issue"] + end = r["dma_issue"] + if end is None: # sync dma / store: no async-issue marker + end = r["resp"] if r["resp"] is not None else r["finished"] + if start is None: + continue + if end is None or end < start: + end = start + 1 + start = max(start, free) + free = max(end, start + 1) + add(core, "dma", start, free - start, _label(r["opcode"], r["detail"]), r) + # VPU: one server; slice = occupancy (compute_cycle - overlapping_cycle). + free = 0 + for r in sorted(u["vector"], key=issue_key): + if r["issued"] is None: + continue + cc, ov = _occ(r["detail"]) + dur = max(cc - ov, 1) + start = max(r["issued"], free) + free = start + dur + add(core, "vector", start, dur, "vector", r) + # SA: num_sa servers, round-robin in issue order (mirrors the Core's rr). + sa_free = [0] * nsa + for i, r in enumerate(sorted(u["sa"], key=issue_key)): + if r["issued"] is None: + continue + s = i % nsa + cc, ov = _occ(r["detail"]) + dur = max(cc - ov, 1) + start = max(r["issued"], sa_free[s]) + sa_free[s] = start + dur + lane = "sa" if nsa == 1 else f"sa{s}" + add(core, lane, start, dur, _label(r["opcode"], r["detail"]), r) + + for c in sorted(cores): + events.append({"name": "process_name", "ph": "M", "pid": c, "tid": 0, + "args": {"name": f"Core {c}"}}) + order = {"dma": 0, "sa": 1, "sa0": 1, "sa1": 2, "sa2": 3, "sa3": 4, "vector": 8} + for c, lane in sorted(lanes, key=lambda x: (x[0], order.get(x[1], 5))): + events.append({"name": "thread_name", "ph": "M", "pid": c, "tid": lane, + "args": {"name": lane}}) + events.append({"name": "thread_sort_index", "ph": "M", "pid": c, "tid": lane, + "args": {"sort_index": order.get(lane, 5)}}) + return {"traceEvents": events, "displayTimeUnit": "ns"} + + +def main(argv): + ap = argparse.ArgumentParser() + ap.add_argument("input", nargs="?", help="trace log file (default: stdin)") + ap.add_argument("-o", "--out", default="timeline.json") + ap.add_argument("-s", "--num-sa", type=int, default=1, + help="systolic arrays per core (num_systolic_array_per_core); " + ">1 splits into sa0..saN-1 lanes") + a = ap.parse_args(argv[1:]) + src = open(a.input) if a.input else sys.stdin + insts = parse(src) + trace = to_chrome(insts, a.num_sa) + with open(a.out, "w") as fh: + json.dump(trace, fh) + n = sum(1 for e in trace["traceEvents"] if e["ph"] == "X") + sys.stderr.write(f"wrote {a.out}: {n} slices -> open in https://ui.perfetto.dev\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) From 9a39ad92daedcb708a601d9ee4dfacccf783362c Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 24 Jun 2026 22:35:51 +0900 Subject: [PATCH 05/32] [Tooling] TOGSim trace timeline (Perfetto) and the trace emits it needs trace_timeline.py with per-work-item grouping and resource-centric DMA lanes; the trace logs the first DRAM response and the assigned systolic array, and scopes the compute barrier to its dispatch. --- TOGSim/include/Instruction.h | 14 ++++ TOGSim/include/TraceLogTags.h | 1 + TOGSim/src/Core.cc | 26 ++++--- TOGSim/src/CoreTraceLog.cc | 23 +++--- TOGSim/src/Instruction.cc | 7 +- TOGSim/src/togsim_trace_bridge.cc | 6 ++ scripts/trace_timeline.py | 123 ++++++++++++++++++++---------- 7 files changed, 140 insertions(+), 60 deletions(-) diff --git a/TOGSim/include/Instruction.h b/TOGSim/include/Instruction.h index dc237580..3dfdb796 100644 --- a/TOGSim/include/Instruction.h +++ b/TOGSim/include/Instruction.h @@ -49,6 +49,14 @@ class Instruction : public std::enable_shared_from_this { int get_assigned_sa() const { return _assigned_sa; } void set_weight_token(const std::shared_ptr& t) { _weight_token = t; } const std::shared_ptr& get_weight_token() const { return _weight_token; } + // Trace-only: which work-item (togsim_dispatch tile) this op belongs to, for + // grouping/coloring in the timeline. Set by the bridge per TILE_BEGIN. + void set_tile_group(int g) { _tile_group = g; } + int get_tile_group() const { return _tile_group; } + // COMPUTE_BAR fence: the max finish_cycle of the async computes it gates (its + // own dispatch only), so it drains those instead of every SA pipeline. + void update_fence_finish(cycle_type c) { if (c > _fence_finish) _fence_finish = c; } + cycle_type get_fence_finish() const { return _fence_finish; } bool check_ready() { return ready_counter == 0; } const Opcode get_opcode() { return opcode; } bool is_dma_read() { return opcode == Opcode::MOVIN; } @@ -71,6 +79,9 @@ class Instruction : public std::enable_shared_from_this { void inc_waiting_request(); void dec_waiting_request(); size_t get_waiting_request() { return _nr_waiting_request; } + // trace: log only the FIRST DRAM response of a load (when data starts arriving). + bool got_first_response() const { return _got_first_response; } + void mark_first_response() { _got_first_response = true; } std::vector& get_tile_size() { return tile_size; } std::vector& get_tile_stride() { return tile_stride; } void set_overlapping_cycle(cycle_type cycle) { overlapping_cycle = cycle; } @@ -138,6 +149,7 @@ class Instruction : public std::enable_shared_from_this { std::vector tile_stride; size_t _tile_numel; size_t _nr_waiting_request=0; + bool _got_first_response=false; size_t _elem_bits = 0; addr_type dram_addr; uint32_t _numa_id = 0; // For DMA instruction @@ -160,4 +172,6 @@ class Instruction : public std::enable_shared_from_this { // SA weight-buffer model (see the setters above). int _assigned_sa = -1; std::shared_ptr _weight_token; + int _tile_group = -1; // trace-only work-item id (see set_tile_group) + cycle_type _fence_finish = 0; // COMPUTE_BAR: drain target (see update_fence_finish) }; \ No newline at end of file diff --git a/TOGSim/include/TraceLogTags.h b/TOGSim/include/TraceLogTags.h index 6c158099..759a4fdb 100644 --- a/TOGSim/include/TraceLogTags.h +++ b/TOGSim/include/TraceLogTags.h @@ -24,6 +24,7 @@ inline constexpr const char* kInstructionFinished = "INST_FINISHED"; inline constexpr const char* kInstructionSkipped = "INST_SKIP"; inline constexpr const char* kAsyncDmaAllRequestsIssued = "ASYNC_DMA_ISSUE"; +inline constexpr const char* kFirstDramResponse = "DRAM_RESP_FIRST"; inline constexpr const char* kAllDramResponsesReceived = "DRAM_RESP_DONE"; inline constexpr const char* kL2CacheableStatusForAddress = "L2CACHE_STAT"; diff --git a/TOGSim/src/Core.cc b/TOGSim/src/Core.cc index dfe40745..c78afe5c 100644 --- a/TOGSim/src/Core.cc +++ b/TOGSim/src/Core.cc @@ -351,10 +351,6 @@ void Core::cycle() { inst->set_assigned_sa(sa_idx); // record the SA actually used (for the trace) } release_sram(inst); // consumer issued -> free the tiles it read - // sec 10.7: this op is now entering the pipeline -> release its - // occupancy (pipeline) dependents so a preload/matmul successor - // overlaps it instead of waiting its full latency. - inst->release_pipeline_children(); auto& target_pipeline = (ct == VECTOR_UNIT) ? _vu_compute_pipeline : _sa_compute_pipeline.at(sa_idx); if (target_pipeline.empty()) { @@ -366,6 +362,10 @@ void Core::cycle() { inst->finish_cycle = target_pipeline.back()->finish_cycle + inst->get_compute_cycle() - overlapped_cycle; inst->bubble_cycle = bubble_cycle; } + // sec 10.7: release the occupancy (pipeline) dependents so a successor + // overlaps this op. finish_cycle is set first so release can feed it to + // a COMPUTE_BAR child's per-dispatch fence (see release_pipeline_children). + inst->release_pipeline_children(); // Release this matmul's weight slot at its streaming-end (finish - // overlapping), not at full finish (the drain tail does not read it). @@ -425,13 +425,10 @@ void Core::cycle() { break; case Opcode::COMPUTE_BAR: { - // Compute fence: finish only once ALL compute pipelines have drained - // (every systolic array + the VPU empty). Until then it does not issue -- - // it stays in the ready queue and is re-checked next cycle. - bool drained = _vu_compute_pipeline.empty(); - for (int s = 0; s < _num_systolic_array_per_core; s++) - drained = drained && _sa_compute_pipeline.at(s).empty(); - if (drained) { + // Compute fence: finish once THIS dispatch's async computes have drained + // (their max finish is fed in by update_fence_finish at issue). Scoped to its + // own dispatch, so an unrelated tile's matmuls do not delay it. + if (_core_cycle >= inst->get_fence_finish()) { core_trace_log::trace_instruction_line(_core_cycle, _id, TraceLogTag::pad15(TraceLogTag::kInstructionFinished), inst->get_global_inst_id(), @@ -542,6 +539,13 @@ void Core::push_memory_response(mem_fetch* response) { Instruction* owner_inst = static_cast(response->get_custom_data()); assert(owner_inst->get_waiting_request()); + 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, + TraceLogTag::pad15(TraceLogTag::kFirstDramResponse), + owner_inst->get_global_inst_id(), + core_trace_log::format_instruction_detail_line(*owner_inst)); + } owner_inst->dec_waiting_request(); if (!owner_inst->get_waiting_request()) { auto it = _dma_waiting_queue.find(owner_inst); diff --git a/TOGSim/src/CoreTraceLog.cc b/TOGSim/src/CoreTraceLog.cc index 9761f9ec..7086893e 100644 --- a/TOGSim/src/CoreTraceLog.cc +++ b/TOGSim/src/CoreTraceLog.cc @@ -31,7 +31,7 @@ std::string format_dma_inst_issued_detail(Instruction& inst) { } return fmt::format( "addr_name={} dram=0x{:016x} rank={} elem_bits={} async={} indirect={} tag=0x{:016x} stride=[{}] size=[{}] " - "tag_idx=[{}]", + "tag_idx=[{}] tile={}", inst.get_addr_name(), static_cast(inst.get_base_dram_address()), rank, @@ -41,7 +41,8 @@ std::string format_dma_inst_issued_detail(Instruction& inst) { tag_hex, fmt::join(inst.get_tile_stride(), ","), fmt::join(ts, ","), - fmt::join(tidx, ",")); + fmt::join(tidx, ","), + inst.get_tile_group()); } std::string format_dma_inst_issued_trace_line(Instruction& inst) { @@ -52,31 +53,35 @@ std::string format_instruction_detail_line(Instruction& inst) { const Opcode op = inst.get_opcode(); const std::string opname = opcode_to_string(op); if (op == Opcode::COMP) { - return fmt::format("{} (compute_type={} compute_cycle={} overlapping_cycle={})", + return fmt::format("{} (compute_type={} compute_cycle={} overlapping_cycle={} sa={} tile={})", opname, inst.get_compute_type(), inst.get_compute_cycle(), - inst.get_overlapping_cycle()); + inst.get_overlapping_cycle(), + inst.get_assigned_sa(), + inst.get_tile_group()); } if ((op == Opcode::MOVIN || op == Opcode::MOVOUT) && inst.is_async_dma()) { - return fmt::format("{} (ASYNC subgraph_id={} addr_name={} tag_id=[{}] tag_idx=[{}] tag_stride=[{}])", + return fmt::format("{} (ASYNC subgraph_id={} addr_name={} tag_id=[{}] tag_idx=[{}] tag_stride=[{}] tile={})", opname, inst.subgraph_id, inst.get_addr_name(), format_tag_key_list_hex(inst.get_tag_id()), fmt::join(inst.get_tag_idx_list(), ","), - fmt::join(inst.get_tag_stride_list(), ",")); + fmt::join(inst.get_tag_stride_list(), ","), + inst.get_tile_group()); } if (op == Opcode::MOVIN || op == Opcode::MOVOUT) { - return fmt::format("{} (addr_name={})", opname, inst.get_addr_name()); + return fmt::format("{} (addr_name={} tile={})", opname, inst.get_addr_name(), inst.get_tile_group()); } if (op == Opcode::MEMORY_BAR) { - return fmt::format("{} (addr_name={} tag_id=[{}] tag_idx=[{}] tag_stride=[{}])", + return fmt::format("{} (addr_name={} tag_id=[{}] tag_idx=[{}] tag_stride=[{}] tile={})", opname, inst.get_addr_name(), format_tag_key_list_hex(inst.get_tag_id()), fmt::join(inst.get_tag_idx_list(), ","), - fmt::join(inst.get_tag_stride_list(), ",")); + fmt::join(inst.get_tag_stride_list(), ","), + inst.get_tile_group()); } return opname; } diff --git a/TOGSim/src/Instruction.cc b/TOGSim/src/Instruction.cc index 54e50511..d0471226 100644 --- a/TOGSim/src/Instruction.cc +++ b/TOGSim/src/Instruction.cc @@ -67,7 +67,12 @@ void Instruction::add_pipeline_child(std::shared_ptr child) { } void Instruction::release_pipeline_children() { - for (auto& c : _pipeline_children) c->dec_ready_counter(); + for (auto& c : _pipeline_children) { + // a COMPUTE_BAR child fences only its own dispatch -> it drains the max + // finish of the computes it gates, fed here as each one issues. + if (c->get_opcode() == Opcode::COMPUTE_BAR) c->update_fence_finish(finish_cycle); + c->dec_ready_counter(); + } _pipeline_children.clear(); } diff --git a/TOGSim/src/togsim_trace_bridge.cc b/TOGSim/src/togsim_trace_bridge.cc index 48f1faf5..1d67d1e1 100644 --- a/TOGSim/src/togsim_trace_bridge.cc +++ b/TOGSim/src/togsim_trace_bridge.cc @@ -80,6 +80,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, std::map, std::pair>> current_dma; 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) // Async compute (matmul/preload) pipelines on the systolic array. A store needs the // drained result, so it FLUSHes -- one barrier before the store waits all outstanding // async compute, with no per-op completion events. @@ -165,6 +166,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, sg = std::make_shared(); sg->set_core_id(t.core); tile = std::make_shared(Tile::Status::INITIALIZED); + cur_tile_group++; continue; } if (t.kind == TraceRec::TILE_END) { @@ -176,6 +178,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, if (t.kind == TraceRec::DMA) { int64_t uniq = next_tag++; // fresh Core tag key per dma record auto inst = make_dma(t, uniq); + inst->set_tile_group(cur_tile_group); size_t numel = 1; // SRAM footprint (ready-tile ordering) for (auto d : t.dims) numel *= (size_t)d; tile->inc_required_sram_size(numel * (t.elem_bits / 8)); @@ -209,11 +212,13 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, std::shared_ptr dma_inst; if (it != current_dma.end()) { uniq = it->second.first; dma_inst = it->second.second; } auto bar = make_mem_bar(t, uniq); + bar->set_tile_group(cur_tile_group); if (dma_inst) dma_inst->add_child(bar); tile->append_instuction(bar); for (int64_t b : t.write_bufs) last_writer[b] = bar; } else if (t.kind == TraceRec::COMPUTE) { auto inst = make_compute(t); + inst->set_tile_group(cur_tile_group); link(inst, t.read_bufs, t.write_bufs); for (int64_t b : t.read_bufs) sram_on_read(b, inst); // frees the tiles it consumes if (is_async_compute(t.compute_type)) outstanding_async.push_back(inst); @@ -222,6 +227,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, // ISSUED (pipeline-child release); the Core then waits the SA pipelines to // drain before it finishes (-> the store it gates). auto bar = std::make_shared(Opcode::COMPUTE_BAR); + bar->set_tile_group(cur_tile_group); for (auto& a : outstanding_async) a->add_pipeline_child(bar); outstanding_async.clear(); tile->append_instuction(bar); diff --git a/scripts/trace_timeline.py b/scripts/trace_timeline.py index f7d2672e..0dce5243 100644 --- a/scripts/trace_timeline.py +++ b/scripts/trace_timeline.py @@ -3,14 +3,16 @@ opens in Perfetto (https://ui.perfetto.dev) or chrome://tracing as an interactive timeline (Gantt). -Each instruction becomes one duration slice on one of 3 per-core lanes: - dma -- MOVIN / MOVOUT - sa -- COMP compute_type 1 (matmul) / 2 (preload) +Each instruction becomes one duration slice, grouped per core (pid). Lanes: + dram-rd -- loads crossing the DRAM bus (read bandwidth) + dram-wr -- stores crossing the DRAM bus (write bandwidth) + sa / sa0.. -- COMP compute_type 1 (matmul) / 2 (preload) vector -- COMP compute_type 0 (vector) -grouped per core (pid). Time unit = core cycles. Barriers (MEMORY_BAR/COMPUTE_BAR) -are not drawn. A compute slice's width is its compute_cycle (the op's own latency), -not issue->finish (which balloons under pipeline backlog); a DMA slice is the -actual transfer ASYNC_DMA_ISSUE -> data-ready. +Time unit = core cycles. Barriers (MEMORY_BAR/COMPUTE_BAR) are not drawn. A DMA bar +runs from the op's first DRAM response (DRAM_RESP_FIRST, logged by the Core -- so it +captures data moving even while still injecting) to its completion (load: data-ready; +store: finished), serialized per direction so each is one visible bar (packed row = +saturated bus). A compute slice's width is its occupancy (compute_cycle - overlapping). Usage: bin/Simulator --config --trace_so --cycle_table --log_level trace \ @@ -33,14 +35,39 @@ _HIDE = {"MEMORY_BAR", "COMPUTE_BAR", "TILE_BEGIN", "TILE_END"} _CT_NAME = {0: "vector", 1: "matmul", 2: "preload"} +# Perfetto/catapult reserved color names. Slices are tinted per work-item tile, so one +# tile's ops share a color across lanes. 16 names, because a core's tiles stride by +# num_cores and an 8-name palette collapses to 4 colors per core on 2 cores. +_TILE_PALETTE = ["good", "bad", "terrible", "yellow", "olive", "rail_response", + "rail_load", "rail_animation", "rail_idle", "thread_state_running", + "thread_state_runnable", "thread_state_iowait", + "thread_state_uninterruptible", "generic_work", "startup", + "vsync_highlight_color"] + + +def _tile_color(detail): + m = re.search(r"\btile=(\d+)", detail or "") + return _TILE_PALETTE[int(m.group(1)) % len(_TILE_PALETTE)] if m else None + + +_DMA_SHORT = {"MOVIN": "MVIN", "MOVOUT": "MVOUT"} + + +def _tile_of(detail): + m = re.search(r"\btile=(-?\d+)", detail or "") + return m.group(1) if m else "?" + def _label(opcode, detail): if opcode == "COMP": m = re.search(r"compute_type=(\d+)", detail) ct = int(m.group(1)) if m else -1 - return _CT_NAME.get(ct, "comp") - m = re.search(r"addr_name=(\w+)", detail) - return f"{opcode} {m.group(1)}" if m else opcode + return f"T{_tile_of(detail)} {_CT_NAME.get(ct, 'comp')}" + # DMA: keep each load's OWN identity (addr_name) so the input/weight/K-panel + # loads stay distinct; tile is conveyed by color (and args), not the name. + m = re.search(r"addr_name=(\w+)", detail or "") + who = m.group(1) if m else "?" + return f"{who} (T{_tile_of(detail)} {_DMA_SHORT.get(opcode, opcode)})" def _lane(opcode, detail): @@ -65,7 +92,8 @@ def parse(lines): key = (core, iid) r = insts.setdefault(key, { "core": core, "iid": iid, "opcode": opcode, "detail": detail, - "issued": None, "finished": None, "resp": None, "dma_issue": None}) + "issued": None, "finished": None, "resp": None, "dma_issue": None, + "first_resp": None}) if not r["opcode"] or r["opcode"] == opcode: r["opcode"] = opcode if detail.strip(): @@ -76,7 +104,9 @@ def parse(lines): r["finished"] = cyc elif tag == "DRAM_RESP_DONE": r["resp"] = cyc - elif tag == "ASYNC_DMA_ISSUE": # actual transfer start (DMA engine busy) + elif tag == "DRAM_RESP_FIRST" and r["first_resp"] is None: # first data arrived + r["first_resp"] = cyc + elif tag == "ASYNC_DMA_ISSUE": # all requests injected (engine done) r["dma_issue"] = cyc return insts @@ -94,7 +124,8 @@ def to_chrome(insts, num_sa=1): dma : MOVIN/MOVOUT -- 1 DMA engine; slice = actual transfer (ASYNC_DMA_ISSUE -> data-ready). vector : COMP type 0 -- 1 VPU. - sa : COMP type 1/2 -- num_sa systolic arrays, round-robin by issue order. + sa : COMP type 1/2 -- each op on the SA the Core reports (`sa=` field; + weight-pinned), so lanes auto-split sa0..; rr fallback if absent. A compute slice's width is compute_cycle - overlapping_cycle (its occupancy = latency minus the tail that overlaps the next op), starting when the unit actually picks it up: start = max(issue, unit_free). num_sa>1 -> lanes sa0.. .""" @@ -116,32 +147,36 @@ def to_chrome(insts, num_sa=1): def add(core, lane, ts, dur, name, r): lanes.add((core, lane)) cores.add(core) - events.append({"name": name, "cat": lane, "ph": "X", "ts": ts, - "dur": max(dur, 1), "pid": core, "tid": lane, - "args": {"inst_id": r["iid"], "issued": r["issued"], - "finished": r["finished"], "data_ready": r["resp"]}}) + args = {"inst_id": r["iid"], "tile": _tile_of(r["detail"]), + "issued": r["issued"], "first_data": r["first_resp"], + "finished": r["finished"], "data_ready": r["resp"]} + am = re.search(r"addr_name=(\w+)", r["detail"] or "") + if am: + args["addr"] = am.group(1) + ev = {"name": name, "cat": lane, "ph": "X", "ts": ts, + "dur": max(dur, 1), "pid": core, "tid": lane, "args": args} + cname = _tile_color(r["detail"]) + if cname: + ev["cname"] = cname + events.append(ev) def issue_key(r): return r["issued"] if r["issued"] is not None else 0 nsa = max(num_sa, 1) for core, u in sorted(by_core.items()): - # DMA engine: one server, serialized. A load occupies it only while INJECTING - # requests -- [INST_ISSUED, ASYNC_DMA_ISSUE] -- not the response tail. So a load - # blocked on a full spad leaves a real idle gap = the SRAM throttle stalling it. - free = 0 - for r in sorted(u["dma"], key=issue_key): - start = r["issued"] if r["issued"] is not None else r["dma_issue"] - end = r["dma_issue"] - if end is None: # sync dma / store: no async-issue marker - end = r["resp"] if r["resp"] is not None else r["finished"] - if start is None: - continue - if end is None or end < start: - end = start + 1 - start = max(start, free) - free = max(end, start + 1) - add(core, "dma", start, free - start, _label(r["opcode"], r["detail"]), r) + # DMA data on the DRAM bus, split by direction: a LOAD's data returns on the + # response, so its bar is [first response, data-ready]; a STORE's goes out with + # the request, so its bar is [issued, finished]. Serialized per direction. + for lane, op, sk, ek in (("dram-rd", "MOVIN", "first_resp", "resp"), + ("dram-wr", "MOVOUT", "issued", "finished")): + free = 0 + rows = [r for r in u["dma"] if r["opcode"] == op + and r[sk] is not None and r[ek] is not None and r[ek] > r[sk]] + for r in sorted(rows, key=lambda r: r[ek]): + start = max(r[sk], free) + free = max(r[ek], start + 1) + add(core, lane, start, free - start, _label(r["opcode"], r["detail"]), r) # VPU: one server; slice = occupancy (compute_cycle - overlapping_cycle). free = 0 for r in sorted(u["vector"], key=issue_key): @@ -152,23 +187,33 @@ def issue_key(r): start = max(r["issued"], free) free = start + dur add(core, "vector", start, dur, "vector", r) - # SA: num_sa servers, round-robin in issue order (mirrors the Core's rr). - sa_free = [0] * nsa - for i, r in enumerate(sorted(u["sa"], key=issue_key)): + # SA: each op runs on the systolic array the Core reports (the `sa=` field + # = its weight-pinned / round-robin assignment); fall back to round-robin + # by issue order for older logs without the field. Each SA is one server. + rows = sorted(u["sa"], key=issue_key) + + def _sa_of(r, i): + m = re.search(r"\bsa=(-?\d+)", r["detail"]) + return int(m.group(1)) if (m and int(m.group(1)) >= 0) else (i % nsa) + + max_sa = max([nsa] + [_sa_of(r, i) + 1 for i, r in enumerate(rows)]) + sa_free = [0] * max_sa + for i, r in enumerate(rows): if r["issued"] is None: continue - s = i % nsa + s = _sa_of(r, i) cc, ov = _occ(r["detail"]) dur = max(cc - ov, 1) start = max(r["issued"], sa_free[s]) sa_free[s] = start + dur - lane = "sa" if nsa == 1 else f"sa{s}" + lane = "sa" if max_sa == 1 else f"sa{s}" add(core, lane, start, dur, _label(r["opcode"], r["detail"]), r) for c in sorted(cores): events.append({"name": "process_name", "ph": "M", "pid": c, "tid": 0, "args": {"name": f"Core {c}"}}) - order = {"dma": 0, "sa": 1, "sa0": 1, "sa1": 2, "sa2": 3, "sa3": 4, "vector": 8} + order = {"dram-rd": 0, "dram-wr": 1, + "sa": 2, "sa0": 2, "sa1": 3, "sa2": 4, "sa3": 5, "vector": 7} for c, lane in sorted(lanes, key=lambda x: (x[0], order.get(x[1], 5))): events.append({"name": "thread_name", "ph": "M", "pid": c, "tid": lane, "args": {"name": lane}}) From 36b0b35b81dce79318c6142815c386dd7c00a9cb Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 24 Jun 2026 22:35:51 +0900 Subject: [PATCH 06/32] [TOGSim] Make the C++ trace path the default and stabilize it The trace .so is now picked up automatically next to the kernel's tile_graph.onnx; TORCHSIM_LEGACY_TOG=1 falls back to the deprecated ONNX producer. Also make the codegen cache replay safe: FxGraphCache can skip PyTorchSim codegen entirely, so the wrapper must carry what the .so needs, and concurrent compiles of one source must not observe a half-written .mlir. --- PyTorchSimFrontend/extension_codecache.py | 100 ++++++++------- PyTorchSimFrontend/mlir/mlir_autotune.py | 7 +- .../mlir/mlir_codegen_backend.py | 23 ++-- PyTorchSimFrontend/mlir/mlir_scheduling.py | 5 + PyTorchSimFrontend/mlir/mlir_template.py | 18 +-- PyTorchSimFrontend/mlir/passes/__init__.py | 7 +- .../mlir/passes/build_skeleton.py | 70 ++++++++++- .../mlir/passes/dep_analysis.py | 118 ++++++++++++------ .../mlir/passes/dma_fine_grained.py | 19 ++- .../mlir/passes/lower_to_emitc.py | 2 + Simulator/simulator.py | 13 +- TOGSim/include/Instruction.h | 17 +-- TOGSim/include/togsim_loader.h | 4 +- TOGSim/src/Core.cc | 7 +- TOGSim/src/Instruction.cc | 13 +- TOGSim/src/Simulator.cc | 2 +- TOGSim/src/main.cc | 84 +++++++++---- TOGSim/src/togsim_runtime.cc | 18 +-- TOGSim/src/togsim_trace_bridge.cc | 20 +++ 19 files changed, 368 insertions(+), 179 deletions(-) diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 15c5f7e3..964af004 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -5,7 +5,7 @@ import torch from PyTorchSimFrontend import extension_config -from torch._inductor.codecache import get_hash, write +from torch._inductor.codecache import get_hash, write, write_atomic from torch._inductor.async_compile import AsyncCompile from AsmParser.tog_generator import tog_generator from PyTorchSimFrontend.mlir.mlir_caller_codegen import MLIRKernelCallerCodeGen @@ -23,6 +23,13 @@ def get_write_path(src_code): return os.path.join(extension_config.get_dump_path(), hash_prefix(get_hash(src_code.strip()))) +_HEADER_BY_HASH = {} +def store_header(src_code, spike_header, gem5_header): + _HEADER_BY_HASH[get_hash(src_code.strip())] = (spike_header, gem5_header) +def get_header(src_code): + return _HEADER_BY_HASH.get(get_hash(src_code.strip())) + + def get_lock_path(write_path): """Return lock file path for the given write_path (per-source_code lock).""" return os.path.join(write_path, ".compile.lock") @@ -128,40 +135,50 @@ def load(cls, source_code, vlen = kwargs['vlen'] vlenb = vlen // 8 write_path = get_write_path(source_code) - key, input_path = write(source_code, "mlir", specified_dir=write_path) - # Run the Python out-of-line MLIR passes (MLIR bindings) on the kernel - # .mlir in place, before mlir-opt. Currently lowers torchsim.vlane_idx - # (replaces the old C++ -global-idx pass); add more in passes/__init__.py. + os.makedirs(write_path, exist_ok=True) + global_var_header = kwargs.get("global_var_header") + if global_var_header is not None: + write_atomic(os.path.join(write_path, "global_var.h"), global_var_header) + gem5_global_var_header = kwargs.get("gem5_global_var_header") + if gem5_global_var_header is not None: + write_atomic(os.path.join(write_path, "gem5_global_var.h"), gem5_global_var_header) + # The compile rewrites the kernel .mlir in place and reads it back, and two + # compiles of the same source share a write_path. Hold the per-path lock across + # the build, and skip it when a prior build finished (its tile_graph.onnx exists). + from filelock import FileLock from PyTorchSimFrontend.mlir.passes import ( run_python_passes, run_module_passes, POST_OPT_PASSES, run_standard_lowering, run_tog, ) - run_python_passes(input_path, vectorlane=vectorlane_size) - new_input_path = os.path.splitext(input_path)[0] - raw_tog_path = new_input_path + "_tog.py" tog_path = os.path.join(write_path, "tile_graph.onnx") - sample_mlir_path = new_input_path + "_sample" - validation_binary_path = os.path.join(write_path, validation_binary_name) - gem5_cmds = mlir_gem5_compile_command(new_input_path, sample_mlir_path, raw_tog_path, vectorlane_size) - - from filelock import FileLock - os.makedirs(write_path, exist_ok=True) lock = FileLock(get_lock_path(write_path), timeout=LOCK_TIMEOUT) - - if spad_info is not None: - link_option = f"-Wl,--section-start=.spad=0x{spad_info['spad_vaddr']:x}" - else: - link_option = "" - # Generate LLVM kernel calller and binary for validation - if extension_config.pytorchsim_functional_mode: - # Use custom malloc to avoid size error - new_link_option = link_option + " -Wl,--wrap=malloc -Wl,--wrap=free" - cmds = mlir_compile_command(new_input_path, vectorlane_size, vlen=vlen) - opt_pad_cmd = shlex.split(cmds[0]) - translate_cmd = shlex.split(cmds[1]) - llc_cmd = shlex.split(cmds[2]) - llc_asm_cmd = shlex.split(cmds[3]) - with lock: + with lock: + key, input_path = write(source_code, "mlir", specified_dir=write_path) + if os.path.isfile(tog_path): + return key + # Run the Python out-of-line MLIR passes (MLIR bindings) on the kernel + # .mlir in place, before mlir-opt. Currently lowers torchsim.vlane_idx + # (replaces the old C++ -global-idx pass); add more in passes/__init__.py. + run_python_passes(input_path, vectorlane=vectorlane_size) + new_input_path = os.path.splitext(input_path)[0] + raw_tog_path = new_input_path + "_tog.py" + sample_mlir_path = new_input_path + "_sample" + validation_binary_path = os.path.join(write_path, validation_binary_name) + gem5_cmds = mlir_gem5_compile_command(new_input_path, sample_mlir_path, raw_tog_path, vectorlane_size) + + if spad_info is not None: + link_option = f"-Wl,--section-start=.spad=0x{spad_info['spad_vaddr']:x}" + else: + link_option = "" + # Generate LLVM kernel calller and binary for validation + if extension_config.pytorchsim_functional_mode: + # Use custom malloc to avoid size error + new_link_option = link_option + " -Wl,--wrap=malloc -Wl,--wrap=free" + cmds = mlir_compile_command(new_input_path, vectorlane_size, vlen=vlen) + opt_pad_cmd = shlex.split(cmds[0]) + translate_cmd = shlex.split(cmds[1]) + llc_cmd = shlex.split(cmds[2]) + llc_asm_cmd = shlex.split(cmds[3]) try: # loop-padding (mlir-opt) -> Python fine-grained + vcix (one parse/print) subprocess.check_call(opt_pad_cmd) @@ -195,17 +212,11 @@ def load(cls, source_code, ) raise SpadOverflowError() - # Skip if TOG file already exists - if os.path.isfile(tog_path): - return key + # Launch tile graph generator + gem5_pad_cmd = shlex.split(gem5_cmds[0]) + gem5_translate_cmd = shlex.split(gem5_cmds[1]) + gem5_llc_cmd = shlex.split(gem5_cmds[2]) - # Launch tile graph generator - gem5_pad_cmd = shlex.split(gem5_cmds[0]) - gem5_translate_cmd = shlex.split(gem5_cmds[1]) - gem5_llc_cmd = shlex.split(gem5_cmds[2]) - - lock = FileLock(get_lock_path(write_path), timeout=LOCK_TIMEOUT) - with lock: try: # mlir-opt now runs only loop-padding and writes the post-vcix IR; the # tile-operation-graph pass is ported to Python. run_tog reads that IR and @@ -261,10 +272,10 @@ def load(cls, source_code, vector_lane=vectorlane_size ) - # Trace pipeline (opt-in, TORCHSIM_DUMP_TRACE_SO=1): also emit the trace - # producer .so + cycle-table TSV from the SAME post-vcix IR and gem5 cycles, - # so it can be compared cycle-consistently. Best-effort: never breaks the build. - if os.environ.get("TORCHSIM_DUMP_TRACE_SO") == "1": + # Trace pipeline (DEFAULT): emit the trace producer .so + cycle-table TSV + # from the post-vcix IR and gem5 cycles. TORCHSIM_LEGACY_TOG=1 opts back into + # the ONNX TOG, in which case the .so is unused. Never breaks the compile. + if os.environ.get("TORCHSIM_LEGACY_TOG") != "1": try: import mlir.ir as ir from PyTorchSimFrontend.mlir.passes import ( @@ -280,12 +291,9 @@ def load(cls, source_code, _cl = list(cycle_list_for_trace) if _cl and len(_cl) != _ntiles: _cl = (_cl + [_cl[-1]] * _ntiles)[:_ntiles] - logger.info(f"[P3-trace] cycle_list={cycle_list_for_trace} -> {_cl} " - f"(#tiles={_ntiles}, x_off={x_offset}, w_off={w_offset})") _tbl = _ct.build_cycle_table(_mod, _cl, x_offset, w_offset) _ct.dump_cycle_table_tsv(_tbl, os.path.join(write_path, "trace_cycles.tsv")) _l2e.build_trace_so(pv, os.path.join(write_path, "trace.so")) - logger.info(f"[P3-trace] wrote trace.so + trace_cycles.tsv in {write_path}") except Exception as e: logger.warning(f"[P3-trace] trace .so/sidecar dump skipped: {e}") return key diff --git a/PyTorchSimFrontend/mlir/mlir_autotune.py b/PyTorchSimFrontend/mlir/mlir_autotune.py index 396396f3..e4876b5b 100644 --- a/PyTorchSimFrontend/mlir/mlir_autotune.py +++ b/PyTorchSimFrontend/mlir/mlir_autotune.py @@ -54,7 +54,7 @@ def __str__(self) -> str: def make_run_fn( self, input_tensors: torch.Tensor, output_tensors: torch.Tensor ) -> Callable[[], None]: - from PyTorchSimFrontend.extension_codecache import CustomAsyncCompile + from PyTorchSimFrontend.extension_codecache import CustomAsyncCompile, get_header custom_async_compile = CustomAsyncCompile() # Check already cached result. @@ -80,12 +80,15 @@ def cached_run_fn(*args, autotune_subprocess_timeout_sec=None, **kwargs): return cached_run_fn # Run a candidate code + _headers = get_header(self.source_code) + _header_kwargs = {} if _headers is None else { + "global_var_header": _headers[0], "gem5_global_var_header": _headers[1]} run_method = custom_async_compile.mlir( self.source_code, vectorlane_size=self.extra_args["vector_lane"], loop_size=self.extra_args["loop_size"], spad_info=self.extra_args["spad_info"], vlen=self.extra_args["vlen"], arg_attributes=self.extra_args["arg_attributes"], origins=self.extra_args["origins"], silent_mode=True, - autotune=self.extra_args['autotune']) + autotune=self.extra_args['autotune'], **_header_kwargs) args = [ tensor diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 725e0dc6..2360542c 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -17,7 +17,6 @@ from torch._inductor.codegen import cpp, wrapper, common, memory_planning from torch._inductor.ir import GraphPartitionSignature from torch._inductor.virtualized import V, _ops as ops -from torch._inductor.codecache import write_atomic from torch._inductor.utils import ( IndentedBuffer, is_welford_reduction, @@ -1120,28 +1119,22 @@ def codegen_nodes(self, nodes, kernel_name): src_code, meta_code = super().codegen_nodes(nodes, kernel_name) self._prepare_simulator_headers(src_code) if "autotune" in extension_config.codegen_mapping_strategy and extension_config.pytorchsim_timing_mode: - optimal_src_code, meta_code = self.autotune(nodes, kernel_name)[:2] + # Use temporaries: autotune returns [None, None, None] when it cannot autotune + # (a size-1 pointwise kernel with ranges == [1]), and unpacking into meta_code + # would clobber the valid arg_attributes the fall-through below returns. + optimal_src_code, optimal_meta_code = self.autotune(nodes, kernel_name)[:2] if optimal_src_code is not None: - return optimal_src_code, meta_code + return optimal_src_code, optimal_meta_code return src_code, meta_code def _prepare_simulator_headers(self, src_code): - from filelock import FileLock - - write_path = extension_codecache.get_write_path(src_code) - os.makedirs(write_path, exist_ok=True) - - spike_write_path = os.path.join(write_path, "global_var.h") - gem5_write_path = os.path.join(write_path, "gem5_global_var.h") - spad_end_symbol = "int spad_end[0] __attribute__ ((section(\".spad\")));\n" spad_section_end_symbol = ( f"int spad_section_end[0] __attribute__ ((section(\".spad\"), aligned({self.spad_info['spad_size']*self.vector_lane})));" ) - lock = FileLock(extension_codecache.get_lock_path(write_path), timeout=extension_codecache.LOCK_TIMEOUT) - with lock: - write_atomic(spike_write_path, self.header.getvalue() + spad_end_symbol + spad_section_end_symbol) - write_atomic(gem5_write_path, self.gem5_header.getvalue()) + spike_content = self.header.getvalue() + spad_end_symbol + spad_section_end_symbol + gem5_content = self.gem5_header.getvalue() + extension_codecache.store_header(src_code, spike_content, gem5_content) def get_arg_info(self, name): arg_info = dict() diff --git a/PyTorchSimFrontend/mlir/mlir_scheduling.py b/PyTorchSimFrontend/mlir/mlir_scheduling.py index 41ec61af..8520596c 100644 --- a/PyTorchSimFrontend/mlir/mlir_scheduling.py +++ b/PyTorchSimFrontend/mlir/mlir_scheduling.py @@ -5,6 +5,7 @@ import operator from sympy import symbols, sympify from PyTorchSimFrontend import extension_config +from PyTorchSimFrontend import extension_codecache from PyTorchSimFrontend.mlir.mlir_codegen_backend import MLIRKernel from torch.utils._ordered_set import OrderedSet @@ -333,6 +334,10 @@ def define_kernel(self, src_code, meta_code, kernel_name, vector_lane, spad_info codecache_def.writeline(f"spad_info={spad_info},") codecache_def.writeline(f"origins={origins},") codecache_def.writeline(f"arg_attributes={meta_code},") + headers = extension_codecache.get_header(src_code) + if headers is not None: + codecache_def.writeline(f"global_var_header='''{headers[0]}''',") + codecache_def.writeline(f"gem5_global_var_header='''{headers[1]}''',") codecache_def.writeline(f"vlen={extension_config.vpu_vector_length_bits})") wrapper.define_kernel(kernel_name, codecache_def.getvalue(), gpu=False) return kernel_name diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index 529a49b5..2b8a0676 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -21,7 +21,6 @@ from torch._inductor.autotune_process import TensorMeta from torch._inductor.virtualized import V, NullHandler, _ops as ops from torch._inductor.utils import IndentedBuffer -from torch._inductor.codecache import write_atomic import PyTorchSimFrontend.extension_codecache as extension_codecache from PyTorchSimFrontend.mlir.mlir_autotune import MLIRBenchmarkRequest @@ -613,22 +612,11 @@ def codegen_nodes(self, tile_candidates, render, template_node, prologue_nodes, return src_code, meta_code def _prepare_simulator_headers(self, src_code): - from filelock import FileLock - spad_end_symbol = f"int spad_end[0] __attribute__ ((section(\".spad\")));\n" spad_section_end_symbol = f"int spad_section_end[0] __attribute__ ((section(\".spad\"), aligned({self.spad_info['spad_size']*self.vector_lane})));" - - write_path = extension_codecache.get_write_path(src_code) - os.makedirs(write_path, exist_ok=True) - spike_write_path = os.path.join(write_path, "global_var.h") - gem5_write_path = os.path.join(write_path, "gem5_global_var.h") - - lock = FileLock(extension_codecache.get_lock_path(write_path), timeout=extension_codecache.LOCK_TIMEOUT) - with lock: - if not os.path.exists(spike_write_path): - write_atomic(spike_write_path, self.header.getvalue()+spad_end_symbol+spad_section_end_symbol) - if not os.path.exists(gem5_write_path): - write_atomic(gem5_write_path, self.gem5_header.getvalue()) + spike_content = self.header.getvalue()+spad_end_symbol+spad_section_end_symbol + gem5_content = self.gem5_header.getvalue() + extension_codecache.store_header(src_code, spike_content, gem5_content) def codegen_prologue_body(self): body = IndentedBuffer() diff --git a/PyTorchSimFrontend/mlir/passes/__init__.py b/PyTorchSimFrontend/mlir/passes/__init__.py index 82cadc2f..d96035bb 100644 --- a/PyTorchSimFrontend/mlir/passes/__init__.py +++ b/PyTorchSimFrontend/mlir/passes/__init__.py @@ -76,8 +76,11 @@ def run_module_passes(in_path, out_path, passes, **opts): p.run(module, **opts) out = str(module) - with open(out_path, "w") as f: - f.write(out) + # Atomic write: this rewrites the kernel .mlir in place outside load()'s FileLock, + # and a concurrent compile must never see a truncated file -- mlir-opt would parse + # it to an empty module and silently drop the kernel. + from torch._inductor.codecache import write_atomic + write_atomic(out_path, out) return True diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index df4c6046..8124d757 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -271,6 +271,72 @@ def _results_unused(op): return True +def _strip_loop_iter_args(block): + """Drop loop-carried values (iter_args) from every affine.for/scf.for. + + The skeleton only needs the loop STRUCTURE (iteration counts) and the + togsim.* markers -- not the data flowing through the loop. Reduction kernels + carry a *vector* accumulator as an iter_arg; EmitC/C++ cannot represent a + loop carrying a vector, so the trace .so emission fails. Since the trace is + timing-only (values come from the recorded run), we rebuild each loop without + iter_args: body uses of an iter_arg become its init value, the loop result + becomes its init, and the now-orphaned accumulate ops are removed by _dce. + """ + # Only strip a loop whose RESULTS are unused: a loop whose result still feeds a + # kept op (an index accumulator behind a togsim.dma address) is left alone. Run + # after _dce, so nested reductions free up inner results round by round. + while True: + tgt = None + for op in walk_ops(block): + n = op.operation.name + if (n in ("affine.for", "scf.for") and len(op.operation.results) > 0 + and _results_unused(op)): + tgt = op + break + if tgt is None: + return + _rebuild_loop_no_iter(tgt) + + +def _rebuild_loop_no_iter(op): + o = op.operation + nres = len(o.results) + n_in = len(o.operands) + inits = [o.operands[n_in - nres + i] for i in range(nres)] + keep_operands = [o.operands[i] for i in range(n_in - nres)] # bound operands only + old_block = o.regions[0].blocks[0] + oargs = list(old_block.arguments) # [iv, *iter_args] + + attrs = {na.name: na.attr for na in o.attributes} + # affine.for tags its operand groups; zero the iter-arg group (last entry). + if "operandSegmentSizes" in attrs: + seg = [int(x) for x in str(attrs["operandSegmentSizes"]).split(":")[1].strip(" >").split(",")] + seg[-1] = 0 + attrs["operandSegmentSizes"] = ir.Attribute.parse( + "array") + + loc = ir.Location.unknown(o.context) + with loc: # default loc for new block args + new = ir.Operation.create(o.name, results=[], operands=keep_operands, + attributes=attrs, regions=1, loc=loc, + ip=ir.InsertionPoint(o)) + nb = new.regions[0].blocks.append(oargs[0].type) # block with the iv arg only + + oargs[0].replace_all_uses_with(nb.arguments[0]) # iv + for ba, ini in zip(oargs[1:], inits): # iter-arg uses -> init + ba.replace_all_uses_with(ini) + for res, ini in zip(o.results, inits): # loop result -> init + res.replace_all_uses_with(ini) + + term_name = "affine.yield" if o.name == "affine.for" else "scf.yield" + with ir.InsertionPoint(nb): + ir.Operation.create(term_name, results=[], operands=[], loc=loc) + new_term = list(nb.operations)[0] + for bop in list(old_block.operations)[:-1]: # move body (drop old yield) + bop.operation.move_before(new_term) + o.erase() + + def _dce(block): """Erase non-kept ops with no used results, to a fixed point. Safe: an op with live SSA uses is never touched.""" @@ -466,7 +532,9 @@ def build_skeleton(module): op.operation.erase() except Exception: pass - _dce(block) + _dce(block) # drop dead consumers (e.g. the result store) first, + _strip_loop_iter_args(block) # so a now-unused loop result lets us strip its iter_args + _dce(block) # then clean the orphaned accumulate ops return ("skeleton: compute=%d dma=%d wait=%d (unpaired waits dropped)" % (n_compute, n_dma, n_wait)) diff --git a/PyTorchSimFrontend/mlir/passes/dep_analysis.py b/PyTorchSimFrontend/mlir/passes/dep_analysis.py index fa4efc8a..36c1d724 100644 --- a/PyTorchSimFrontend/mlir/passes/dep_analysis.py +++ b/PyTorchSimFrontend/mlir/passes/dep_analysis.py @@ -42,42 +42,80 @@ def _global_of(memref_val): return None -def _read_buffers_of_compute(cn): - """SRAM buffers a compute node reads: (a) each vcix.iv input traced to its - vector.transfer_read source (activations/weights streamed into the SA), and - (b) any direct vector.transfer_read in the node (the epilogue's accumulator - read-modify-write of Y_spad).""" - bufs = set() - for op in cn.operations: - if op.name == "vector.transfer_read" and list(op.operands): - b = _global_of(op.operands[0]) - if b: - bufs.add(b) - elif op.name == "vcix.iv" and list(op.operands): - v = op.operands[0] - defop = v.owner if isinstance(v.owner, ir.Operation) else getattr(v.owner, "operation", None) - if defop is not None and defop.name == "vector.transfer_read" and list(defop.operands): - b = _global_of(defop.operands[0]) - if b: - bufs.add(b) - return bufs - - -def _write_buffers_of_compute(cn): - """SRAM buffers a compute node writes: vector.transfer_write / vector_store target.""" - bufs = set() - for op in cn.operations: - if op.name in ("vector.transfer_write", "affine.vector_store", "vector.store"): - # target memref is the last memref operand - for v in op.operands: - try: - if ir.MemRefType.isinstance(v.type): - b = _global_of(v) - if b: - bufs.add(b) - except Exception: - pass - return bufs +# Ops that touch SRAM-buffer DATA, by category. A view op only computes an address, +# so it is skipped; the real access is the load/store using it. Anything else with a +# memref operand raises, catching a new fusion pattern at compile time. +_LOAD_OPS = {"vector.transfer_read", "affine.vector_load", "vector.load", + "memref.load", "affine.load"} +_STORE_OPS = {"vector.transfer_write", "affine.vector_store", "vector.store", + "memref.store", "affine.store"} +_IGNORE_OPS = {"memref.dealloc"} # lifetime, not a data access + + +def _is_memref(v): + try: + return ir.MemRefType.isinstance(v.type) + except Exception: + return False + + +def _walk_compute_ops(cn): + """Every op in the compute node, recursing into nested regions (loop bodies). A + fused epilogue (BatchNorm/ReLU) keeps its ops inside an un-unrolled affine.for, so + a top-level-only scan (cn.operations) sees just the loop and misses every access.""" + for top in cn.operations: + stack = [top] + while stack: + op = stack.pop() + yield op + for region in op.operation.regions: + for block in region.blocks: + stack.extend(block.operations) + + +def _rw_buffers_of_compute(cn): + """(reads, writes): the @global SRAM buffers a compute node reads/writes, walking + nested regions and classifying each op that touches a memref.""" + reads, writes = set(), set() + def rd(v): + b = _global_of(v) + if b: + reads.add(b) + def wr(v): + b = _global_of(v) + if b: + writes.add(b) + for op in _walk_compute_ops(cn): + if any(_is_memref(r) for r in op.results): + continue # view/cast/alloc -- address only + mrefs = [v for v in op.operands if _is_memref(v)] + if not mrefs: + continue + name = op.name + if name in _LOAD_OPS: + for v in mrefs: + rd(v) + elif name in _STORE_OPS: + for v in mrefs: + wr(v) # the store target memref + elif name == "memref.copy": + rd(mrefs[0]) + wr(mrefs[-1]) + elif name.startswith("linalg."): # DPS: ins read, outs read+write + for v in op.inputs: + if _is_memref(v): + rd(v) + for v in op.outputs: + if _is_memref(v): + rd(v) + wr(v) + elif name in _IGNORE_OPS: + continue + else: + raise RuntimeError( + f"dep_analysis: unclassified memref op '{name}' in a compute node -- " + f"it touches an SRAM buffer; classify it in _LOAD_OPS/_STORE_OPS") + return reads, writes def _dma_buffer(builder, dma_node): @@ -99,8 +137,7 @@ def _dma_buffer(builder, dma_node): def compute_buffers(cn): """(read_buffers, write_buffers) for one compute node, including the virtual SA_WEIGHTS edge (preload writes it, matmul reads it).""" - reads = set(_read_buffers_of_compute(cn)) - writes = set(_write_buffers_of_compute(cn)) + reads, writes = _rw_buffers_of_compute(cn) if cn.compute_type == 1: # MATMUL consumes the preloaded weights reads.add(SA_WEIGHTS) elif cn.compute_type == 2: # PRELOAD loads them @@ -132,10 +169,11 @@ def analyze(module): if not cn.operations: continue ct = {0: "VECTOR", 1: "MATMUL", 2: "PRELOAD"}.get(cn.compute_type, f"c{cn.compute_type}") + creads, cwrites = _rw_buffers_of_compute(cn) nodes.append({ "kind": ct, - "reads": _read_buffers_of_compute(cn), - "writes": _write_buffers_of_compute(cn), + "reads": creads, + "writes": cwrites, "node": cn, "compute_type": cn.compute_type, }) diff --git a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py index f1872dca..d7571d2b 100644 --- a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py +++ b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py @@ -21,6 +21,7 @@ Pipeline entry point: run_fine_grained(in_path, out_path, vectorlane). """ +import itertools import os import sys @@ -383,11 +384,19 @@ def _run_func(func, vectorlane): # dominate the nest. Codegen emits input before weight, matching the C++ pass # which fuses after the weight subtile loop. ip = ir.InsertionPoint(mvin_weight.op) - fused_ivs, body_ip = _build_for_nest(bounds, ip) - in_ivs = [fused_ivs[fuse["in_to_fused"][d]] for d in range(rank)] - w_ivs = [fused_ivs[fuse["w_to_fused"][d]] for d in range(rank)] - _emit_dma(mvin_input, in_ivs, vectorlane, body_ip) - _emit_dma(mvin_weight, w_ivs, vectorlane, body_ip) + # Unroll the fused nest, emitting each distinct input/weight subtile ONCE (a load + # is invariant to the other operand's dims, so the cross-product re-emits it + # identically). Dedup by the operand's own coords; keep the fused issue order. + seen_in, seen_w = set(), set() + for it in itertools.product(*[range(b) for b in bounds]): + in_key = tuple(it[fuse["in_to_fused"][d]] for d in range(rank)) + if in_key not in seen_in: + seen_in.add(in_key) + _emit_dma(mvin_input, [_const_index(c, ip) for c in in_key], vectorlane, ip) + w_key = tuple(it[fuse["w_to_fused"][d]] for d in range(rank)) + if w_key not in seen_w: + seen_w.add(w_key) + _emit_dma(mvin_weight, [_const_index(c, ip) for c in w_key], vectorlane, ip) mvin_input.op.erase() mvin_weight.op.erase() diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index 30d3c796..4a54fbb9 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -57,7 +57,9 @@ #: upstream EmitC conversion pipeline (the infrastructure this pass drives). _PIPELINE = ("builtin.module(" + "convert-vector-to-scf{full-unroll=true}," "func.func(lower-affine)," + "func.func(lower-vector-multi-reduction)," "convert-scf-to-emitc," "convert-arith-to-emitc," "convert-func-to-emitc)") diff --git a/Simulator/simulator.py b/Simulator/simulator.py index 2b9f05be..3c191db8 100644 --- a/Simulator/simulator.py +++ b/Simulator/simulator.py @@ -560,7 +560,18 @@ def run_standalone( os.fsync(trace_file.fileno()) try: - cmd = f"{TOGSimulator.get_togsim_command(config_path, togsim_path)} --models_list {trace_file_path}" + # The C++ TOG (trace) path is the DEFAULT: drive the simulation from the + # emitted trace.so; TORCHSIM_LEGACY_TOG=1 opts into the legacy ONNX TOG. Each + # autotune candidate has its own trace.so. Fall back only if none was emitted. + trace_so = os.path.join(os.path.dirname(str(model_path)), "trace.so") + cycle_tsv = os.path.join(os.path.dirname(str(model_path)), "trace_cycles.tsv") + base_cmd = TOGSimulator.get_togsim_command(config_path, togsim_path) + use_trace = (os.environ.get("TORCHSIM_LEGACY_TOG") != "1" + and os.path.exists(trace_so)) + if use_trace: + cmd = f"{base_cmd} --trace_so {trace_so} --cycle_table {cycle_tsv}" + else: + cmd = f"{base_cmd} --models_list {trace_file_path}" if extension_config.CONFIG_TOGSIM_DEBUG_LEVEL: cmd += f" --log_level {extension_config.CONFIG_TOGSIM_DEBUG_LEVEL}" diff --git a/TOGSim/include/Instruction.h b/TOGSim/include/Instruction.h index 3dfdb796..2d17741e 100644 --- a/TOGSim/include/Instruction.h +++ b/TOGSim/include/Instruction.h @@ -127,12 +127,12 @@ class Instruction : public std::enable_shared_from_this { // bytes this load occupies in the spad (from the tile it moves in). size_t sram_footprint() const { return _tile_numel * (_elem_bits / 8); } - cycle_type start_cycle; - cycle_type finish_cycle; + cycle_type start_cycle = 0; + cycle_type finish_cycle = 0; cycle_type bubble_cycle=0; bool finished=false; - int subgraph_id; + int subgraph_id = 0; private: uint64_t _global_inst_id = 0; static uint64_t _next_global_inst_id; @@ -140,18 +140,19 @@ class Instruction : public std::enable_shared_from_this { void *_owner = nullptr; std::list>* _owner_ready_queue_ref = nullptr; Opcode opcode; - cycle_type compute_cycle; - cycle_type overlapping_cycle; - size_t ready_counter; + cycle_type compute_cycle = 0; + cycle_type overlapping_cycle = 0; + size_t ready_counter = 0; // parents not yet finished; the minimal Instruction(Opcode) + // ctor (barriers) relies on this default + inc_ready_counter std::set> child_inst; std::set> _pipeline_children; // released at issue (sec 10.7) std::vector tile_size; std::vector tile_stride; - size_t _tile_numel; + size_t _tile_numel = 0; size_t _nr_waiting_request=0; bool _got_first_response=false; size_t _elem_bits = 0; - addr_type dram_addr; + addr_type dram_addr = 0; uint32_t _numa_id = 0; // For DMA instruction int _compute_type = 0; std::vector _tag_idx_list; diff --git a/TOGSim/include/togsim_loader.h b/TOGSim/include/togsim_loader.h index df8063a5..2114055b 100644 --- a/TOGSim/include/togsim_loader.h +++ b/TOGSim/include/togsim_loader.h @@ -41,11 +41,11 @@ struct RunResult { // 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 across `num_cores`. +// 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, - int32_t num_cores); + const int32_t* partition_cores, int32_t n_partition_cores); } // namespace togsim diff --git a/TOGSim/src/Core.cc b/TOGSim/src/Core.cc index c78afe5c..c05857d8 100644 --- a/TOGSim/src/Core.cc +++ b/TOGSim/src/Core.cc @@ -56,8 +56,11 @@ void Core::release_sram(const std::shared_ptr& inst) { } bool Core::can_issue(const std::shared_ptr& op) { - /* Check SRAM is enough to run tile */ - return _tiles.size() < 4 && !op->is_stonne_tile(); + /* Bound concurrent dispatches so their combined spad working set fits: with the + * global @buffers each in-flight dispatch piles its own load versions, and too + * many at once overflow the spad (versions never free -> wedge). 2 keeps double- + * buffering overlap while leaving headroom. */ + return _tiles.size() < 2 && !op->is_stonne_tile(); } void Core::issue(std::shared_ptr op) { diff --git a/TOGSim/src/Instruction.cc b/TOGSim/src/Instruction.cc index d0471226..c5778a28 100644 --- a/TOGSim/src/Instruction.cc +++ b/TOGSim/src/Instruction.cc @@ -57,13 +57,18 @@ void Instruction::finish_instruction() { } void Instruction::add_child(std::shared_ptr child) { - child->inc_ready_counter(); - child_inst.insert(child); + // child_inst is a set (each child released exactly once at finish), so the + // ready_counter must be bumped only when the edge is NEW -- a producer that + // writes several buffers a single consumer reads (e.g. a sort tile reading the + // value+index buffers its predecessor wrote) links the same pair once per shared + // buffer; double-counting would leave ready_counter stuck above 0 -> deadlock. + if (child_inst.insert(child).second) + child->inc_ready_counter(); } void Instruction::add_pipeline_child(std::shared_ptr child) { - child->inc_ready_counter(); - _pipeline_children.insert(child); + if (_pipeline_children.insert(child).second) + child->inc_ready_counter(); } void Instruction::release_pipeline_children() { diff --git a/TOGSim/src/Simulator.cc b/TOGSim/src/Simulator.cc index 17320fa0..66992481 100644 --- a/TOGSim/src/Simulator.cc +++ b/TOGSim/src/Simulator.cc @@ -186,7 +186,7 @@ void Simulator::icnt_cycle() { // Consecutive frozen cycles tolerated before declaring the sim wedged (spad too // small). Generous so transient idle never false-fires; a true freeze is constant. -static constexpr uint64_t kWedgeThreshold = 5000; +static constexpr uint64_t kWedgeThreshold = 100000; // Frozen-state guard: work remains but nothing is in flight to advance it, because // the kernel's working set exceeds the whole per-core spad (core_spad_size_kb too diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 8726cfdf..4ce9aa18 100644 --- a/TOGSim/src/main.cc +++ b/TOGSim/src/main.cc @@ -1,3 +1,4 @@ +#include #include #include #include @@ -15,13 +16,61 @@ namespace fs = std::filesystem; namespace po = boost::program_options; +// Run a kernel's compiled trace producer (.so) and bridge it to a TileGraph for +// `partition_id`, whose cores its work-items round-robin over. The cycle-table TSV +// gives per-tile latency (a flat stub if absent). nullptr if the producer fails. +std::unique_ptr build_trace_tilegraph(Simulator* simulator, + const std::string& trace_so_path, + const std::string& cycle_table_path, + int partition_id) { + const auto& cfg = simulator->get_hardware_config_yaml(); + int num_cores = cfg["num_cores"] ? cfg["num_cores"].as() : 1; + std::vector partition_cores; + for (int c = 0; c < num_cores; c++) + if (simulator->get_partition_id(c) == partition_id) partition_cores.push_back(c); + if (partition_cores.empty()) partition_cores.push_back(0); + // First cut: stub tensor bases (real per-tensor addresses come later). + std::vector bases(16); + for (size_t i = 0; i < bases.size(); ++i) bases[i] = 0x100000ull * (i + 1); + // Cycle table: load the per-tile_id TSV sidecar if present, else a flat stub. + std::vector cyc, ovl; + std::ifstream ct(cycle_table_path); + if (ct.is_open()) { + int64_t c, o; + 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"); +} + 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) { - auto graph_praser = TileGraphParser(onnx_path, attribute_path, config_yaml); - std::unique_ptr& tile_graph = graph_praser.get_tile_graph(); + std::unique_ptr tile_graph; + std::string tog_path = onnx_path; // for the log line + // Prefer the C++ trace path: the kernel's trace.so / trace_cycles.tsv sit next to its + // tile_graph.onnx. Opt out with TORCHSIM_LEGACY_TOG=1, and fall back to the legacy ONNX + // parser when the .so is absent or fails to run. + const char* legacy = std::getenv("TORCHSIM_LEGACY_TOG"); + std::string dir = fs::path(onnx_path).parent_path().string(); + std::string trace_so = dir + "/trace.so"; + std::string cycle_tsv = dir + "/trace_cycles.tsv"; + if ((!legacy || std::string(legacy) != "1") && fs::exists(trace_so)) { + tile_graph = build_trace_tilegraph(simulator, trace_so, cycle_tsv, partition_id); + if (tile_graph) tog_path = trace_so; + else spdlog::warn("[TOGSim] trace.so run failed for {}; falling back to ONNX", trace_so); + } + if (!tile_graph) { + auto graph_praser = TileGraphParser(onnx_path, attribute_path, config_yaml); + tile_graph = std::move(graph_praser.get_tile_graph()); + } tile_graph->set_arrival_time(request_time ? request_time : simulator->get_core_cycle()); tile_graph->set_kernel_id(kernel_id); spdlog::info("[Scheduler {}] Enqueued kernel_id: {}, tog_path: {}, operation: {}, request_time_cycles: {}", - partition_id, kernel_id, onnx_path, tile_graph->get_name(), request_time); + partition_id, kernel_id, tog_path, tile_graph->get_name(), request_time); simulator->enqueue_graph(partition_id, std::move(tile_graph)); } @@ -159,37 +208,18 @@ int main(int argc, char** argv) { std::string trace_so_path; cmd_parser.set_if_defined("trace_so", &trace_so_path); if (!trace_so_path.empty()) { - const auto& cfg = simulator->get_hardware_config_yaml(); - int num_cores = cfg["num_cores"] ? cfg["num_cores"].as() : 1; - // First cut: stub tensor bases (real per-tensor addresses come later). - std::vector bases(16); - for (size_t i = 0; i < bases.size(); ++i) bases[i] = 0x100000ull * (i + 1); - // Cycle table: load the per-tile_id TSV sidecar if given, else a flat stub. - std::vector cyc, ovl; + // Standalone single-kernel trace run: enqueue to partition 0 (its work-items + // round-robin over partition 0's cores only; see build_trace_tilegraph). std::string cycle_table_path; cmd_parser.set_if_defined("cycle_table", &cycle_table_path); - if (!cycle_table_path.empty()) { - std::ifstream ct(cycle_table_path); - if (!ct.is_open()) { spdlog::error("[TOGSim] cannot open cycle_table {}", cycle_table_path); exit(1); } - int64_t c, o; - while (ct >> c >> o) { cyc.push_back(c); ovl.push_back(o); } - spdlog::info("[TOGSim-trace] loaded cycle table: {} tiles from {}", cyc.size(), cycle_table_path); - } else { - 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(), - num_cores); - if (!run.ok) { spdlog::error("[TOGSim] trace producer run failed"); exit(1); } - spdlog::info("[TOGSim-trace] recorded {} instructions", run.trace.size()); - auto tg = trace_to_tilegraph(run, "trace_kernel"); + auto tg = build_trace_tilegraph(simulator, trace_so_path, cycle_table_path, 0); + if (!tg) { spdlog::error("[TOGSim] trace producer run failed"); exit(1); } tg->set_arrival_time(simulator->get_core_cycle()); tg->set_kernel_id(0); simulator->enqueue_graph(0, std::move(tg)); simulator->run_simulator(); spdlog::info("[TOGSim-trace] Total cycles: {}", simulator->get_core_cycle()); + spdlog::info("Simulation finished"); simulator->print_core_stat(); return 0; } diff --git a/TOGSim/src/togsim_runtime.cc b/TOGSim/src/togsim_runtime.cc index df952a15..90ab9a9f 100644 --- a/TOGSim/src/togsim_runtime.cc +++ b/TOGSim/src/togsim_runtime.cc @@ -19,9 +19,9 @@ struct EmitCtx { const int64_t* cyc = nullptr; // tile_id -> cycle const int64_t* ovl = nullptr; // tile_id -> overlapping_cycle int32_t n_tiles = 0; - int32_t num_cores = 1; + std::vector cores{0}; // the partition's core ids; dispatch round-robins over these // mutable run state - int32_t rr = 0; // round-robin core cursor + int32_t rr = 0; // round-robin cursor into `cores` int32_t cur_core = -1; // current work-item's core std::vector trace; }; @@ -40,10 +40,11 @@ extern "C" { int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } void togsim_dispatch(EmitCtx* ctx, togsim_tile_fn fn, int64_t* iv, int32_t n_iv) { - // Work-item wrapper (sec 9.3): round-robin a core (the producer never sees - // num_cores), bracket the work-item with TILE_BEGIN/TILE_END, and run its body. The - // work-item scope is exactly this call; ops emit records under ctx->cur_core. - ctx->cur_core = ctx->num_cores > 0 ? (ctx->rr++ % ctx->num_cores) : 0; + // Work-item wrapper (sec 9.3): round-robin over THIS partition's cores only -- + // a work-item on another partition's core would sit in this partition's scheduler + // 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)); fn(ctx, iv, n_iv); ctx->trace.push_back(blank(togsim::TraceRec::TILE_END, ctx->cur_core)); @@ -105,7 +106,7 @@ 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, - int32_t num_cores) { + 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; } @@ -115,7 +116,8 @@ RunResult run_producer(const char* so_path, EmitCtx ctx; ctx.tensor_base = tensor_base; ctx.n_tensors = n_tensors; ctx.cyc = cyc; ctx.ovl = ovl; ctx.n_tiles = n_tiles; - ctx.num_cores = num_cores > 0 ? num_cores : 1; + 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); res.ok = true; diff --git a/TOGSim/src/togsim_trace_bridge.cc b/TOGSim/src/togsim_trace_bridge.cc index 1d67d1e1..5ecbe822 100644 --- a/TOGSim/src/togsim_trace_bridge.cc +++ b/TOGSim/src/togsim_trace_bridge.cc @@ -79,6 +79,11 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, // it. Correct only because a load nest and its consumers run in order. Per work-item. 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) // Async compute (matmul/preload) pipelines on the systolic array. A store needs the @@ -98,6 +103,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, tile.reset(); last_writer.clear(); current_dma.clear(); + bar_for_load.clear(); next_tag = 0; outstanding_async.clear(); pending_bar.reset(); @@ -111,6 +117,12 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, const std::vector& reads, const std::vector& writes) { for (int64_t b : reads) { + // A matmul reading its own accumulator imposes NO producer order: Y += X@W is + // commutative. Chaining them serializes the SA and DEADLOCKS the weight-slot + // pipeline. The store still waits every matmul via the COMPUTE_BAR fence. + bool is_accum = false; + for (int64_t w : writes) if (w == b) { is_accum = true; break; } + if (inst->get_compute_type() == MATMUL_CT && is_accum) continue; auto it = last_writer.find(b); if (it == last_writer.end()) continue; int pct = it->second->get_compute_type(); @@ -211,11 +223,19 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, int64_t uniq = next_tag++; // fallback if unpaired std::shared_ptr dma_inst; if (it != current_dma.end()) { uniq = it->second.first; dma_inst = it->second.second; } + // Identical wait (same slot, same load instance) already has a barrier -> reuse it + // 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) last_writer[b] = bf->second.second; + continue; + } auto bar = make_mem_bar(t, uniq); bar->set_tile_group(cur_tile_group); if (dma_inst) dma_inst->add_child(bar); tile->append_instuction(bar); for (int64_t b : t.write_bufs) last_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); inst->set_tile_group(cur_tile_group); From 5c4a7981ca82b75bd6550a8667e79b152af74b87 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 24 Jun 2026 20:36:13 +0900 Subject: [PATCH 07/32] [TOGSim] Redesign trace-bridge dependency, barrier, SRAM-version, and runtime model Replace the trace bridge's accumulated special cases with one dataflow rule and clean up the runtime that consumes it. Dependency rule: per SRAM buffer keep a writers SET; a reader depends on all current writers (occupancy=ISSUE when both are systolic-array ops, else latency=DONE); a writer REPLACEs the set. The only exception is is_mm_accum (a matmul that reads and writes the same buffer = a commutative accumulator): skip its read edge and UNION its write, waiting only the non-matmul init seed and not ordering co-matmuls. This drops the matmul-accumulator chain that deadlocked the SA weight-slot pipeline while keeping the init->matmul edge, and lets a vector epilogue or the store wait every K matmul (fixes the pure-vector store that an empty COMPUTE_BAR let slip). Remove COMPUTE_BAR entirely: a matmul is its own DONE-handle (finish == SA drain), so the store JOINs the matmul writers directly. The whole emit/loader chain is gone -- build_skeleton, lower_to_emitc, togsim.compute_barrier, the runtime symbol, the Opcode/case/_fence_finish, and TraceRec::COMPUTE_BAR -- so a stale producer fails loudly instead of emitting records the bridge would drop. Only MEMORY_BAR remains (an async load's DONE is its data arrival, not issue). Model compute-output spad footprint in the SRAM version/capacity machinery so buffer reuse (WAR) is capacity-modeled, not a hard edge. The output size comes from the DMA records that touch the same buffer (a buf_bytes pre-pass); an in-place buffer (accumulator, relu) is version-transparent so footprint is not double-counted. The occupy gate and version release sit in the MOVIN/MOVOUT/COMP issue points (release before the COMP skip path so a skipped matmul still frees). Runtime: collapse child_inst / _pipeline_children into one event-indexed _deps[ISSUE|DONE] with add_dep(c, on) and fire(e); collapse the weight-slot release queue and the async-load wakeup into one _due_events timed-effect table drained by process_due_events. Both are behavior-preserving (byte-identical). Require the weight-slot model: sa_weight_buffer_depth must be > 0 (errors at init), and the round-robin disable mode is removed. Degenerate traces (a consumer-less preload, an unpinned matmul) hit explicit error+exit guards rather than asserts that vanish under NDEBUG. Mark the legacy ONNX TOG path deprecated: it is superseded by the trace path, so TileGraphParser logs a deprecation warning and the TORCHSIM_LEGACY_TOG=1 opt-in warns at command build. --- CLAUDE.md | 1 + PyTorchSimFrontend/extension_codecache.py | 4 +- .../mlir/passes/build_skeleton.py | 18 +- .../mlir/passes/lower_to_emitc.py | 8 - PyTorchSimFrontend/mlir/passes/togsim_ops.py | 6 +- README.md | 3 +- Simulator/simulator.py | 4 +- TOGSim/include/Core.h | 22 ++- TOGSim/include/Instruction.h | 59 +++--- TOGSim/include/togsim_loader.h | 2 +- TOGSim/include/togsim_runtime.h | 5 - TOGSim/src/Core.cc | 129 +++++++------ TOGSim/src/Instruction.cc | 29 +-- TOGSim/src/TileGraphParser.cc | 6 +- TOGSim/src/main.cc | 7 +- TOGSim/src/togsim_runtime.cc | 63 ------- TOGSim/src/togsim_trace_bridge.cc | 173 ++++++++++-------- 17 files changed, 231 insertions(+), 308 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 12d48082..5a3a47cd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -85,6 +85,7 @@ Note: `TOGSIM_CONFIG` is **overwritten** while inside a `with TOGSimulator(confi Located under `configs/*.yml`: - `num_cores`, `core_freq_mhz`, `num_systolic_array_per_core` +- `sa_weight_buffer_depth` (per-SA resident weight slots; **must be > 0** — the simulator errors on 0. Raise it to effectively disable the preload run-ahead throttle. Defaults to 2 if the key is absent.) - `vpu_num_lanes`, `vpu_spad_size_kb_per_lane`, `vpu_vector_length_bits` - `dram_type` (`ramulator2` | `simple`), `dram_channels`, `dram_freq_mhz`, `ramulator_config_path` - `icnt_type` (`simple` | `booksim`), `icnt_latency_cycles`, `icnt_freq_mhz`, `icnt_config_path` diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 964af004..a520d9a2 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -273,8 +273,8 @@ def load(cls, source_code, ) # Trace pipeline (DEFAULT): emit the trace producer .so + cycle-table TSV - # from the post-vcix IR and gem5 cycles. TORCHSIM_LEGACY_TOG=1 opts back into - # the ONNX TOG, in which case the .so is unused. Never breaks the compile. + # from the post-vcix IR and gem5 cycles. The legacy ONNX TOG is DEPRECATED, + # an opt-in fallback via TORCHSIM_LEGACY_TOG=1. Never breaks the compile. if os.environ.get("TORCHSIM_LEGACY_TOG") != "1": try: import mlir.ir as ir diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index 8124d757..523b3015 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -53,7 +53,7 @@ _KEEP = { "affine.for", "scf.for", "scf.while", "affine.yield", "scf.yield", "func.return", - ts.DMA, ts.COMPUTE, ts.COMPUTE_BAR, ts.MEMORY_BAR, + ts.DMA, ts.COMPUTE, ts.MEMORY_BAR, } @@ -117,20 +117,6 @@ def _emit_dma(ctx, dma_node, tag_id, dram_index, tag_index, read_bufs, write_buf ) -def _emit_compute_bar(ctx, anchor_op): - """Insert a `togsim.compute_barrier` before `anchor_op` -- the fence that - drains in-flight async compute (the systolic-array matmuls) before a store - consumes their result (sec 10.7). - - FIXME: this is the one barrier still synthesized here rather than read from - the IR. Like the async-load memory barrier (now mapped 1:1 from the explicit - dma_wait), the compute fence should eventually appear explicitly in the input - MLIR and be mapped through, not auto-inserted -- no surprising insertion.""" - ir.Operation.create( - ts.COMPUTE_BAR, results=[], operands=[], attributes={}, - loc=ir.Location.unknown(ctx), ip=ir.InsertionPoint(anchor_op)) - - def _emit_memory_bar(ctx, anchor_op, tag_id, tag_index, write_bufs): """Insert a `togsim.memory_barrier` before `anchor_op` -- the explicit async-DMA sync that was the original `memref.dma_wait`. It pairs with its @@ -454,8 +440,6 @@ def _emit_one_dma(ctx, op, node, builder, bufs, tags): read_bufs = [spad_id] if node.is_write else [] write_bufs = [] if node.is_write else [spad_id] tag_id = tags.bind(_value_key(f["tag"]), spad_id) - if node.is_write: - _emit_compute_bar(ctx, op) # FIXME(sec10.7): auto-inserted; should be explicit in the IR. _emit_dma(ctx, node, tag_id, dram_index, tag_index, read_bufs, write_bufs) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index 4a54fbb9..d10651c8 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -423,14 +423,6 @@ def _rewrite_togsim_ops(ctx, kernel, ctx_val): i32(0), _opaque(ctx, "nullptr"), _arr(ctx, _rb), i32(len(_rb)), _arr(ctx, _wb), i32(len(_wb))]) victims.append(op) - elif name == ts.COMPUTE_BAR: - # explicit compute fence -> togsim_compute_barrier(ctx) (sec 10.7). - ir.Operation.create( - "emitc.call_opaque", results=[], operands=[ctx_val], - attributes={"callee": ir.StringAttr.get(ts.EMITC_CALLEE[ts.COMPUTE_BAR]), - "args": ir.ArrayAttr.get([_idx(0)])}, - loc=ir.Location.unknown(ctx), ip=ipo) - victims.append(op) for op in victims: op.operation.erase() diff --git a/PyTorchSimFrontend/mlir/passes/togsim_ops.py b/PyTorchSimFrontend/mlir/passes/togsim_ops.py index 38104d15..4130b576 100644 --- a/PyTorchSimFrontend/mlir/passes/togsim_ops.py +++ b/PyTorchSimFrontend/mlir/passes/togsim_ops.py @@ -32,8 +32,6 @@ tag_id = i32, write_bufs = [..] tag_id, tag_slot, write_bufs) } : (index) -> () - "togsim.compute_barrier"() : () -> () -> togsim_compute_barrier(ctx) - How an async dma pairs with its sync point: NOT by a compile-time id. One static `togsim.dma` op runs once per loop iteration, each with a different RUNTIME tag slot `%tag[%idx]`, so the pairing must be a runtime key. `togsim.dma` carries a @@ -52,17 +50,15 @@ # ---- op names ------------------------------------------------------------- DMA = "togsim.dma" COMPUTE = "togsim.compute" -COMPUTE_BAR = "togsim.compute_barrier" # fence: drain async compute before a consumer (sec 10.7) MEMORY_BAR = "togsim.memory_barrier" # explicit async-DMA sync (the original dma_wait); tag-keyed #: every op this module owns (for matchers / DCE roots in C2). -OP_NAMES = (DMA, COMPUTE, COMPUTE_BAR, MEMORY_BAR) +OP_NAMES = (DMA, COMPUTE, MEMORY_BAR) #: op name -> the togsim_runtime.h symbol C4 lowers it to. EMITC_CALLEE = { DMA: "togsim_dma", COMPUTE: "togsim_compute", - COMPUTE_BAR: "togsim_compute_barrier", MEMORY_BAR: "togsim_memory_barrier", } diff --git a/README.md b/README.md index f0bdc772..c2298376 100644 --- a/README.md +++ b/README.md @@ -385,6 +385,7 @@ num_cores: 1 core_freq_mhz: 940 core_stats_print_period_cycles: 10000 num_systolic_array_per_core: 2 +sa_weight_buffer_depth: 2 # per-SA resident weight slots; must be > 0 (default 2). Raise to loosen the preload throttle. # Optional: one entry per core, default ws_mesh # core_type: [ws_mesh, ws_mesh] # Optional STONNE cores: stonne_config_path, num_stonne_per_core, num_stonne_port @@ -453,7 +454,7 @@ codegen_compiler_optimization: all # all | none | list of option names One-line meaning for each group (details in the YAML block above). -- **Core (`num_cores`, `core_freq_mhz`, `core_stats_print_period_cycles`, `num_systolic_array_per_core`, optional `core_type`, STONNE keys)**: how many cores, their clock, stats cadence, systolic count per core, and optional non-default mesh vs STONNE mix. +- **Core (`num_cores`, `core_freq_mhz`, `core_stats_print_period_cycles`, `num_systolic_array_per_core`, `sa_weight_buffer_depth`, optional `core_type`, STONNE keys)**: how many cores, their clock, stats cadence, systolic count per core, the per-SA resident weight-slot count (must be > 0; bounds preload run-ahead—raise it to loosen the throttle), and optional non-default mesh vs STONNE mix. - **VPU (`vpu_*`)**: vector lane count, per-lane scratchpad (KB), and vector register width—**compiler** uses these for tiling/codegen. - **DRAM (`dram_type`, `dram_channels`, …)**: `ramulator2` uses `ramulator_config_path`; `simple` uses fixed latency and optional bandwidth caps (`dram_bandwidth_gbps_*`, `dram_freq_mhz` when capped). `dram_num_partitions` splits channels for NUMA-style addressing. - **Interconnect (`icnt_*`, `booksim_config_path`)**: `simple` adds fixed hop latency (`icnt_latency_cycles`); `booksim2` points at a BookSim2 topology file. diff --git a/Simulator/simulator.py b/Simulator/simulator.py index 3c191db8..d0a0df53 100644 --- a/Simulator/simulator.py +++ b/Simulator/simulator.py @@ -568,9 +568,11 @@ def run_standalone( base_cmd = TOGSimulator.get_togsim_command(config_path, togsim_path) use_trace = (os.environ.get("TORCHSIM_LEGACY_TOG") != "1" and os.path.exists(trace_so)) + if os.environ.get("TORCHSIM_LEGACY_TOG") == "1": + logger.warning("TORCHSIM_LEGACY_TOG=1 selects the DEPRECATED legacy ONNX TOG path") if use_trace: cmd = f"{base_cmd} --trace_so {trace_so} --cycle_table {cycle_tsv}" - else: + else: # DEPRECATED: legacy ONNX TOG path cmd = f"{base_cmd} --models_list {trace_file_path}" if extension_config.CONFIG_TOGSIM_DEBUG_LEVEL: cmd += f" --log_level {extension_config.CONFIG_TOGSIM_DEBUG_LEVEL}" diff --git a/TOGSim/include/Core.h b/TOGSim/include/Core.h index b611eff1..a8e7e54e 100644 --- a/TOGSim/include/Core.h +++ b/TOGSim/include/Core.h @@ -20,6 +20,13 @@ enum class InstFinishTraceTag { DmaRespComplete, }; +// A timed effect due at a cycle: free a weight slot, or wake a MEMORY_BAR. +struct DueAction { + enum Kind { FreeWeightSlot, WakeBar } kind; + std::shared_ptr token; + std::shared_ptr bar; +}; + class Core { public: Core(uint32_t id, SimulationConfig config); @@ -63,13 +70,14 @@ class Core { // SRAM-capacity throttle (sec 10.4): a consumer frees the buffer-versions it // read (refcount -> 0 releases the spad bytes). Called when COMP/MOVOUT issue. void release_sram(const std::shared_ptr& inst); + // Occupy inst's buffer-version footprint on issue; false if it would overflow + // the spad this cycle (the caller stalls it). True for untracked insts. + bool try_occupy_sram(const std::shared_ptr& inst); // SA weight-buffer throttle (sec 10.4): pick a systolic array that has a free // weight slot (round-robin among free); -1 if all full -> the preload stalls. int pick_free_weight_sa(); - // Free weight slots due this cycle: a matmul releases its slot at its - // streaming-end (finish - overlapping, when it stops reading the weight), - // scheduled at issue in _weight_release_q. Last consumer frees it. - void process_weight_releases(); + void process_due_events(); // drain _due_events due this cycle + void apply_due(const DueAction& a); /* Core id & config file */ const uint32_t _id; @@ -128,10 +136,8 @@ class Core { // SA weight-buffer throttle (sec 10.4). _weight_slots_used[s] = weights resident // on SA s (loaded by a preload, not yet freed by their last matmul); - // _weight_slot_depth = per-SA capacity (0 = disabled -> plain round-robin). + // _weight_slot_depth = per-SA weight-slot capacity (must be > 0). std::vector _weight_slots_used; uint32_t _weight_slot_depth = 0; - // Pending weight-slot releases keyed by cycle (each matmul's streaming-end); - // process_weight_releases() drains those due and decrements the token. - std::multimap> _weight_release_q; + std::multimap _due_events; }; \ No newline at end of file diff --git a/TOGSim/include/Instruction.h b/TOGSim/include/Instruction.h index 2d17741e..d4995eb9 100644 --- a/TOGSim/include/Instruction.h +++ b/TOGSim/include/Instruction.h @@ -6,16 +6,20 @@ #include #include +#include #include #include #include #include #include -// MEMORY_BAR waits a DMA tag in the tag table. COMPUTE_BAR waits the systolic-array -// compute pipelines to drain -- the explicit fence before a store consumes the -// results of async matmuls. -enum class Opcode { MOVIN, MOVOUT, COMP, MEMORY_BAR, COMPUTE_BAR, COUNT}; +// MEMORY_BAR: the DMA/memory barrier (waits a DMA tag in the tag table). +enum class Opcode { MOVIN, MOVOUT, COMP, MEMORY_BAR, COUNT}; + +// A dependency edge releases its consumer on one of the producer's lifecycle +// events: ISSUE (occupancy -- the consumer overlaps the producer on the SA +// pipeline) or DONE (latency -- the consumer needs the producer's result). +enum class DepEvent : uint8_t { ISSUE = 0, DONE = 1, COUNT = 2 }; // One weight slot on systolic array `sa` (sec 10.4). A preload sets refcount = // the matmuls reusing the weight; each frees it at its streaming-end, the last @@ -36,15 +40,20 @@ class Instruction : public std::enable_shared_from_this { std::vector accum_tag_idx_list); Instruction(Opcode opcode); void finish_instruction(); - void add_child(std::shared_ptr child); - // Occupancy (SA-pipeline) dependency: the child is released when THIS op is - // ISSUED (enters the pipeline), not when it finishes -- so a preload/matmul - // successor overlaps it instead of waiting its full latency (sec 10.7). - void add_pipeline_child(std::shared_ptr child); - void release_pipeline_children(); - // SA weight-buffer model: the SA this op is pinned to (a preload picks it, its - // matmul consumers inherit it) and the shared weight slot the matmuls release. - const std::set>& get_pipeline_children() { return _pipeline_children; } + // Subscribe `c` to this op's `on` event (ISSUE=occupancy, DONE=latency). The set + // dedups, so ready_counter is bumped only on a new edge (a producer writing + // several buffers one consumer reads links the pair once per buffer). + void add_dep(std::shared_ptr c, DepEvent on) { + if (_deps[static_cast(on)].insert(c).second) c->inc_ready_counter(); + } + // Release every subscriber of `e` (decrement its ready_counter) and clear. + void fire(DepEvent e) { + for (auto& c : _deps[static_cast(e)]) c->dec_ready_counter(); + _deps[static_cast(e)].clear(); + } + const std::set>& get_deps(DepEvent e) { + return _deps[static_cast(e)]; + } void set_assigned_sa(int s) { _assigned_sa = s; } int get_assigned_sa() const { return _assigned_sa; } void set_weight_token(const std::shared_ptr& t) { _weight_token = t; } @@ -53,10 +62,6 @@ class Instruction : public std::enable_shared_from_this { // grouping/coloring in the timeline. Set by the bridge per TILE_BEGIN. void set_tile_group(int g) { _tile_group = g; } int get_tile_group() const { return _tile_group; } - // COMPUTE_BAR fence: the max finish_cycle of the async computes it gates (its - // own dispatch only), so it drains those instead of every SA pipeline. - void update_fence_finish(cycle_type c) { if (c > _fence_finish) _fence_finish = c; } - cycle_type get_fence_finish() const { return _fence_finish; } bool check_ready() { return ready_counter == 0; } const Opcode get_opcode() { return opcode; } bool is_dma_read() { return opcode == Opcode::MOVIN; } @@ -114,7 +119,6 @@ class Instruction : public std::enable_shared_from_this { void prepare_tag_key(); bool is_sparse_inst() { return _is_sparse_inst; } void set_sparse_state(bool state) { _is_sparse_inst = state; } - std::set>& get_child_inst() { return child_inst; } uint64_t get_global_inst_id() const { return _global_inst_id; } // SRAM-capacity model (sec 10.4), filled by the bridge and enforced by Core: a @@ -124,10 +128,15 @@ class Instruction : public std::enable_shared_from_this { int64_t get_sram_alloc() const { return _sram_alloc_id; } void add_sram_release(int64_t id) { _sram_release_allocs.push_back(id); } const std::vector& get_sram_release() const { return _sram_release_allocs; } - // bytes this load occupies in the spad (from the tile it moves in). - size_t sram_footprint() const { return _tile_numel * (_elem_bits / 8); } + // 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 + // buffer's size is known from the DMA records that touch the same buffer). + void set_sram_footprint(size_t b) { _sram_footprint_override = b; } + size_t sram_footprint() const { + return _sram_footprint_override ? _sram_footprint_override + : _tile_numel * (_elem_bits / 8); + } - cycle_type start_cycle = 0; cycle_type finish_cycle = 0; cycle_type bubble_cycle=0; @@ -144,8 +153,10 @@ class Instruction : public std::enable_shared_from_this { cycle_type overlapping_cycle = 0; size_t ready_counter = 0; // parents not yet finished; the minimal Instruction(Opcode) // ctor (barriers) relies on this default + inc_ready_counter - std::set> child_inst; - std::set> _pipeline_children; // released at issue (sec 10.7) + // Per-event subscriber sets: _deps[ISSUE] released at issue (occupancy), + // _deps[DONE] at finish (latency). std::set dedups + fixes the release order. + std::array>, + static_cast(DepEvent::COUNT)> _deps; std::vector tile_size; std::vector tile_stride; size_t _tile_numel = 0; @@ -170,9 +181,9 @@ class Instruction : public std::enable_shared_from_this { // SRAM-capacity model (see the setters above). int64_t _sram_alloc_id = -1; std::vector _sram_release_allocs; + size_t _sram_footprint_override = 0; // SA weight-buffer model (see the setters above). int _assigned_sa = -1; std::shared_ptr _weight_token; int _tile_group = -1; // trace-only work-item id (see set_tile_group) - cycle_type _fence_finish = 0; // COMPUTE_BAR: drain target (see update_fence_finish) }; \ No newline at end of file diff --git a/TOGSim/include/togsim_loader.h b/TOGSim/include/togsim_loader.h index 2114055b..fef63e54 100644 --- a/TOGSim/include/togsim_loader.h +++ b/TOGSim/include/togsim_loader.h @@ -12,7 +12,7 @@ namespace togsim { // One modeled instruction recorded by the runtime callbacks. struct TraceRec { - enum Kind { TILE_BEGIN, TILE_END, DMA, COMPUTE, MEMORY_BAR, COMPUTE_BAR } kind; + enum Kind { TILE_BEGIN, TILE_END, DMA, COMPUTE, MEMORY_BAR } kind; int32_t core; // work-item -> core binding (set by togsim_dispatch) // DMA / MEMORY_BAR int32_t dir; // togsim_dma_dir diff --git a/TOGSim/include/togsim_runtime.h b/TOGSim/include/togsim_runtime.h index fe069e64..e00d51d9 100644 --- a/TOGSim/include/togsim_runtime.h +++ b/TOGSim/include/togsim_runtime.h @@ -60,11 +60,6 @@ typedef void (*togsim_tile_fn)(EmitCtx* ctx, int64_t* iv, int32_t n_iv); void togsim_dispatch(EmitCtx* ctx, togsim_tile_fn fn, int64_t* iv, int32_t n_iv); -// Compute fence: drain in-flight async compute (the systolic-array matmuls) -// before the following op (a store) consumes their result. Explicit barrier in -// the trace; the loader turns it into a COMPUTE_BAR instruction (sec 10.7). -void togsim_compute_barrier(EmitCtx* ctx); - // Entry point the loader resolves in the producer `.so`. `shape_args` carries // the runtime values for the kernel's symbolic dimensions (in a kernel-specific // order recorded alongside the cached `.so`); `n_shape_args` is their count. diff --git a/TOGSim/src/Core.cc b/TOGSim/src/Core.cc index c05857d8..23ee4d6d 100644 --- a/TOGSim/src/Core.cc +++ b/TOGSim/src/Core.cc @@ -18,7 +18,11 @@ Core::Core(uint32_t id, SimulationConfig config) _stat_inst_count.resize(static_cast(Opcode::COUNT), 0); _stat_tot_skipped_inst.resize(static_cast(Opcode::COUNT), 0); _sram_capacity = (size_t)config.core_spad_size_kb * 1024; // 0 = throttle disabled - _weight_slot_depth = config.sa_weight_buffer_depth; // 0 = disabled (plain rr) + _weight_slot_depth = config.sa_weight_buffer_depth; // per-SA weight slots (>0) + if (_weight_slot_depth == 0) { + spdlog::error("sa_weight_buffer_depth must be > 0 (raise it to loosen the preload throttle)"); + exit(EXIT_FAILURE); + } _weight_slots_used.resize(_num_systolic_array_per_core, 0); } @@ -35,11 +39,23 @@ int Core::pick_free_weight_sa() { return -1; } -void Core::process_weight_releases() { - while (!_weight_release_q.empty() && _weight_release_q.begin()->first <= _core_cycle) { - auto tok = _weight_release_q.begin()->second; - _weight_release_q.erase(_weight_release_q.begin()); - if (--tok->refcount <= 0) _weight_slots_used[tok->sa]--; // last reader frees the slot +void Core::apply_due(const DueAction& a) { + switch (a.kind) { + case DueAction::FreeWeightSlot: + if (--a.token->refcount <= 0) _weight_slots_used[a.token->sa]--; // last reader frees the slot + break; + case DueAction::WakeBar: { + auto bar = a.bar; // async load data arrived -> fire its MEMORY_BAR + finish_instruction(bar); + break; + } + } +} + +void Core::process_due_events() { + while (!_due_events.empty() && _due_events.begin()->first <= _core_cycle) { + apply_due(_due_events.begin()->second); + _due_events.erase(_due_events.begin()); } } @@ -55,6 +71,15 @@ void Core::release_sram(const std::shared_ptr& inst) { } } +bool Core::try_occupy_sram(const std::shared_ptr& inst) { + if (!_sram_capacity || inst->get_sram_alloc() < 0) return true; // untracked + size_t F = inst->sram_footprint(); + if (_sram_used + F > _sram_capacity) return false; // would overflow -> stall + _sram_used += F; + _sram_allocs[inst->get_sram_alloc()] += F; // accumulate version footprint + return true; +} + bool Core::can_issue(const std::shared_ptr& op) { /* Bound concurrent dispatches so their combined spad working set fits: with the * global @buffers each in-flight dispatch piles its own load versions, and too @@ -174,7 +199,7 @@ void Core::dma_cycle() { finish_instruction(instruction, InstFinishTraceTag::DmaRespComplete); for (auto & wait_inst : _dma.get_tag_waiter(instruction->subgraph_id, key)) { _dma.mark_tag_used(instruction->subgraph_id, key); - finish_instruction(wait_inst); + _due_events.emplace(_core_cycle, DueAction{DueAction::WakeBar, nullptr, wait_inst}); } } _dma_finished_queue.erase(_dma_finished_queue.begin()); @@ -239,7 +264,7 @@ void Core::cycle() { /* Increase core cycle counter */ _core_cycle++; - process_weight_releases(); // free weight slots due this cycle before dispatch + process_due_events(); // weight-slot frees + DMA-arrival wakeups due this cycle /* Iterate tile while an instruction is issued */ bool issued = false; @@ -248,9 +273,6 @@ void Core::cycle() { auto& instructions = _tiles[i]->get_ready_instructions(); for (auto it=instructions.begin(); it!=instructions.end();) { auto& inst = *it; - /* Skip instruction is not ready */ - //if (!inst->is_ready()) - // continue; switch (inst->get_opcode()) { case Opcode::MOVIN: @@ -281,19 +303,8 @@ void Core::cycle() { _stat_tot_skipped_inst.at(static_cast(inst->get_opcode()))++; break; } else { - // SRAM-capacity gate: a load that would overflow the per-core spad does - // not issue this cycle -- it stays in the ready queue until a consumer - // frees a tile. On issue it occupies its buffer-version allocation. - if (_sram_capacity && inst->get_sram_alloc() >= 0) { - size_t F = inst->sram_footprint(); - // Stall if the tile does not fit the free spad now. If it can never - // fit, Simulator::cycle() detects the frozen state and exits with a - // "spad too small" error rather than looping forever. - if (_sram_used + F > _sram_capacity) - break; // not issued -> retry next cycle - _sram_used += F; - _sram_allocs[inst->get_sram_alloc()] += F; // accumulate version footprint - } + // 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, _id, TraceLogTag::pad15( @@ -321,39 +332,38 @@ void Core::cycle() { case Opcode::COMP: { const int ct = inst->get_compute_type(); - // SA selection + weight-buffer gate: a preload picks a systolic array with - // a free weight slot and pins its matmul consumers to it (they free the slot - // on finish). This bounds preload run-ahead and keeps matmuls on their SA. + // a fresh-output compute occupies its spad bytes on issue; stall if full. + if (!try_occupy_sram(inst)) break; + // SA selection (sec 10.4): a preload picks an SA with a free weight slot + // and pins its matmul consumers there; a matmul runs on its pinned SA. 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_pipeline_children()) + for (auto& c : inst->get_deps(DepEvent::ISSUE)) if (c->get_compute_type() == MATMUL) n_consumers++; - if (_weight_slot_depth > 0 && n_consumers > 0) { - 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_pipeline_children()) - if (c->get_compute_type() == MATMUL) { - c->set_assigned_sa(sa_idx); - c->set_weight_token(tok); - } - } else { // disabled / no consumers -> plain rr - sa_idx = _systolic_array_rr; - _systolic_array_rr = (_systolic_array_rr + 1) % _num_systolic_array_per_core; + 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)) + if (c->get_compute_type() == MATMUL) { + c->set_assigned_sa(sa_idx); + c->set_weight_token(tok); + } } else { // MATMUL - sa_idx = inst->get_assigned_sa(); - if (sa_idx < 0) { // no preload pinned it -> rr fallback - sa_idx = _systolic_array_rr; - _systolic_array_rr = (_systolic_array_rr + 1) % _num_systolic_array_per_core; + sa_idx = inst->get_assigned_sa(); // pinned by its preload + if (sa_idx < 0) { // unpinned -> no preload set its SA + spdlog::error("matmul was not pinned to an SA by a preload (weight-slot model invariant)"); + exit(EXIT_FAILURE); } } inst->set_assigned_sa(sa_idx); // record the SA actually used (for the trace) } - release_sram(inst); // consumer issued -> free the tiles it read auto& target_pipeline = (ct == VECTOR_UNIT) ? _vu_compute_pipeline : _sa_compute_pipeline.at(sa_idx); if (target_pipeline.empty()) { @@ -365,19 +375,19 @@ void Core::cycle() { inst->finish_cycle = target_pipeline.back()->finish_cycle + inst->get_compute_cycle() - overlapped_cycle; inst->bubble_cycle = bubble_cycle; } - // sec 10.7: release the occupancy (pipeline) dependents so a successor - // overlaps this op. finish_cycle is set first so release can feed it to - // a COMPUTE_BAR child's per-dispatch fence (see release_pipeline_children). - inst->release_pipeline_children(); + // release the occupancy (ISSUE) dependents so a successor overlaps this op. + inst->fire(DepEvent::ISSUE); // Release this matmul's weight slot at its streaming-end (finish - // overlapping), not at full finish (the drain tail does not read it). if (ct == MATMUL && inst->get_weight_token()) { cycle_type rel = inst->finish_cycle > inst->get_overlapping_cycle() ? inst->finish_cycle - inst->get_overlapping_cycle() : _core_cycle; - _weight_release_q.emplace(rel, inst->get_weight_token()); + _due_events.emplace(rel, DueAction{DueAction::FreeWeightSlot, + inst->get_weight_token(), nullptr}); } + release_sram(inst); // free the tiles it read (before the skip path) if (inst->get_compute_cycle() == 0) { inst->finish_instruction(); static_cast(inst->get_owner())->inc_finished_inst(); @@ -404,7 +414,7 @@ void Core::cycle() { auto& key = inst->get_tag_id(); uint32_t finished = _dma.get_tag_finish(inst->subgraph_id, key); if (finished == -1) { - for (auto child_inst : inst->get_child_inst()) { + for (auto child_inst : inst->get_deps(DepEvent::DONE)) { if (child_inst->get_opcode() == Opcode::COMP && child_inst->get_compute_type() == MATMUL) { child_inst->set_compute_cycle(0); } @@ -426,21 +436,6 @@ void Core::cycle() { issued = true; } break; - case Opcode::COMPUTE_BAR: - { - // Compute fence: finish once THIS dispatch's async computes have drained - // (their max finish is fed in by update_fence_finish at issue). Scoped to its - // own dispatch, so an unrelated tile's matmuls do not delay it. - if (_core_cycle >= inst->get_fence_finish()) { - core_trace_log::trace_instruction_line(_core_cycle, _id, - TraceLogTag::pad15(TraceLogTag::kInstructionFinished), - inst->get_global_inst_id(), - core_trace_log::format_instruction_detail_line(*inst)); - finish_instruction(inst); - issued = true; - } - } - break; default: core_trace_log::log_error_undefined_opcode(); exit(EXIT_FAILURE); diff --git a/TOGSim/src/Instruction.cc b/TOGSim/src/Instruction.cc index c5778a28..ee184a1a 100644 --- a/TOGSim/src/Instruction.cc +++ b/TOGSim/src/Instruction.cc @@ -24,7 +24,6 @@ std::string opcode_to_string(Opcode opcode) { case Opcode::MOVOUT: return "MOVOUT"; case Opcode::COMP: return "COMP"; case Opcode::MEMORY_BAR: return "MEMORY_BAR"; - case Opcode::COMPUTE_BAR: return "COMPUTE_BAR"; default: return "Unknown"; } } @@ -51,36 +50,10 @@ Instruction::Instruction(Opcode opcode) } void Instruction::finish_instruction() { - for (auto& counter : child_inst) - counter->dec_ready_counter(); + fire(DepEvent::DONE); // latency consumers finished = true; } -void Instruction::add_child(std::shared_ptr child) { - // child_inst is a set (each child released exactly once at finish), so the - // ready_counter must be bumped only when the edge is NEW -- a producer that - // writes several buffers a single consumer reads (e.g. a sort tile reading the - // value+index buffers its predecessor wrote) links the same pair once per shared - // buffer; double-counting would leave ready_counter stuck above 0 -> deadlock. - if (child_inst.insert(child).second) - child->inc_ready_counter(); -} - -void Instruction::add_pipeline_child(std::shared_ptr child) { - if (_pipeline_children.insert(child).second) - child->inc_ready_counter(); -} - -void Instruction::release_pipeline_children() { - for (auto& c : _pipeline_children) { - // a COMPUTE_BAR child fences only its own dispatch -> it drains the max - // finish of the computes it gates, fed here as each one issues. - if (c->get_opcode() == Opcode::COMPUTE_BAR) c->update_fence_finish(finish_cycle); - c->dec_ready_counter(); - } - _pipeline_children.clear(); -} - void Instruction::inc_waiting_request() { _nr_waiting_request++; } diff --git a/TOGSim/src/TileGraphParser.cc b/TOGSim/src/TileGraphParser.cc index 572062e0..c252258e 100644 --- a/TOGSim/src/TileGraphParser.cc +++ b/TOGSim/src/TileGraphParser.cc @@ -584,7 +584,7 @@ std::vector> TileLoopNode::get_tiles_from_iter(TileGraphPa for (const auto& child_node: node->get_child()) { if (link_map.find(child_node) != link_map.end()) { std::shared_ptr child_inst = link_map[child_node]; - inst->add_child(child_inst); + inst->add_dep(child_inst, DepEvent::DONE); } } } @@ -606,7 +606,7 @@ std::vector> TileLoopNode::get_tiles_from_iter(TileGraphPa for (auto& inner_inst : inner_tile->get_instructions()) { tile_vec.back()->append_instuction(inner_inst); if (nr_inst) { - last_instruction->add_child(inner_inst); + last_instruction->add_dep(inner_inst, DepEvent::DONE); } } } @@ -662,7 +662,7 @@ std::vector> TileLoopNode::get_tiles_from_iter(TileGraphPa for (const auto& child_node: node->get_child()) { if (link_map.find(child_node) != link_map.end()) { std::shared_ptr child_inst = link_map[child_node]; - inst->add_child(child_inst); + inst->add_dep(child_inst, DepEvent::DONE); } } } diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 4ce9aa18..4ccaa4ae 100644 --- a/TOGSim/src/main.cc +++ b/TOGSim/src/main.cc @@ -51,9 +51,9 @@ std::unique_ptr build_trace_tilegraph(Simulator* simulator, 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) { std::unique_ptr tile_graph; std::string tog_path = onnx_path; // for the log line - // Prefer the C++ trace path: the kernel's trace.so / trace_cycles.tsv sit next to its - // tile_graph.onnx. Opt out with TORCHSIM_LEGACY_TOG=1, and fall back to the legacy ONNX - // parser when the .so is absent or fails to run. + // The C++ trace path is the supported one: the kernel's trace.so / trace_cycles.tsv + // sit next to its tile_graph.onnx (same write_path). The legacy ONNX parser below is + // DEPRECATED -- only used via TORCHSIM_LEGACY_TOG=1 or when the .so is absent / fails. const char* legacy = std::getenv("TORCHSIM_LEGACY_TOG"); std::string dir = fs::path(onnx_path).parent_path().string(); std::string trace_so = dir + "/trace.so"; @@ -64,6 +64,7 @@ void launchKernel(Simulator* simulator, unsigned int kernel_id, std::string onnx else spdlog::warn("[TOGSim] trace.so run failed for {}; falling back to ONNX", trace_so); } if (!tile_graph) { + spdlog::warn("[TOGSim] using the DEPRECATED legacy ONNX TOG path for {}", onnx_path); auto graph_praser = TileGraphParser(onnx_path, attribute_path, config_yaml); tile_graph = std::move(graph_praser.get_tile_graph()); } diff --git a/TOGSim/src/togsim_runtime.cc b/TOGSim/src/togsim_runtime.cc index 90ab9a9f..1a76231b 100644 --- a/TOGSim/src/togsim_runtime.cc +++ b/TOGSim/src/togsim_runtime.cc @@ -94,10 +94,6 @@ void togsim_memory_barrier(EmitCtx* ctx, int32_t tag_id, uint64_t tag_slot, ctx->trace.push_back(r); } -void togsim_compute_barrier(EmitCtx* ctx) { - ctx->trace.push_back(blank(togsim::TraceRec::COMPUTE_BAR, ctx->cur_core)); -} - } // extern "C" namespace togsim { @@ -125,63 +121,4 @@ RunResult run_producer(const char* so_path, return res; } -SimResult simulate(const RunResult& run, const TimingParams& params) { - SimResult out; - std::unordered_map dma_free; // DMA-engine free time, per core - std::unordered_map comp_free; // compute free time, per core - std::unordered_map prev_comp; // prev compute finish (overlap), per core - std::map, uint64_t> tag_finish; // (tag_id,tag_slot) -> finish - std::vector pending; // barrier-resolved deps since last compute - - for (const auto& t : run.trace) { - const int c = t.core; - switch (t.kind) { - case TraceRec::DMA: { - // DMAs serialize on the core's DMA engine (overlap compute -> separate - // timeline). finish = issue + latency, recorded under the runtime tag. - uint64_t start = dma_free[c]; - uint64_t fin = start + params.dma_latency; - dma_free[c] = fin; - tag_finish[{t.tag_id, t.tag_slot}] = fin; - out.n_dma++; - break; - } - case TraceRec::MEMORY_BAR: { - // the explicit async-DMA sync: gate the next compute on the paired dma's - // data-arrival, found by the runtime tag (tag_id, tag_slot). - auto it = tag_finish.find({t.tag_id, t.tag_slot}); - if (it != tag_finish.end()) pending.push_back(it->second); - break; - } - case TraceRec::COMPUTE: { - uint64_t deps = 0; - for (uint64_t f : pending) deps = std::max(deps, f); - pending.clear(); - uint64_t start = std::max(comp_free[c], deps); - uint64_t fin; - auto pit = prev_comp.find(c); - if (pit != prev_comp.end()) { - uint64_t prev = pit->second; - uint64_t tail = prev > start ? prev - start : 0; // prev still running - uint64_t overlapped = std::min(tail, (uint64_t)t.overlapping); - fin = std::max(start, prev) + (uint64_t)t.cycle - overlapped; - } else { - fin = start + (uint64_t)t.cycle; - } - comp_free[c] = fin; - prev_comp[c] = fin; - out.n_compute++; - break; - } - case TraceRec::TILE_BEGIN: - case TraceRec::TILE_END: - case TraceRec::COMPUTE_BAR: - break; // work-item boundary / compute fence: no cost in this reference timer - } - } - for (auto& kv : dma_free) out.total_cycle = std::max(out.total_cycle, kv.second); - for (auto& kv : comp_free) out.total_cycle = std::max(out.total_cycle, kv.second); - return out; -} - } // namespace togsim diff --git a/TOGSim/src/togsim_trace_bridge.cc b/TOGSim/src/togsim_trace_bridge.cc index 5ecbe822..08c40afb 100644 --- a/TOGSim/src/togsim_trace_bridge.cc +++ b/TOGSim/src/togsim_trace_bridge.cc @@ -1,6 +1,7 @@ // togsim_trace_bridge.cc -- see togsim_trace_bridge.h #include "togsim_trace_bridge.h" +#include #include #include #include @@ -35,7 +36,7 @@ std::shared_ptr make_dma(const togsim::TraceRec& t, int64_t uniq) { // A MEMORY_BAR carrying the SAME `uniq` tag key as the async dma it gates -- the // Core's tag table signals it at the dma's DATA-ready (resp-complete), unlike a -// raw add_child which the async dma releases at issue-complete. +// raw DONE edge that the async dma releases at issue-complete. std::shared_ptr make_mem_bar(const togsim::TraceRec& t, int64_t uniq) { auto bar = std::make_shared( Opcode::MEMORY_BAR, 0, 0, 0, @@ -70,13 +71,13 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, std::shared_ptr sg; std::shared_ptr tile; - // Explicit dependency DAG (sec 10): a reader depends on the last writer of each - // SRAM buffer it reads. Scoped per work-item (reset at each dispatch) -- buffers - // are work-item-local, so distinct work-items are independent (-> parallel). - std::map> last_writer; // buffer id -> producer + // 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` Core key and its iteration's barriers reuse - // it. Correct only because a load nest and its consumers run in order. Per work-item. + // 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, std::pair>> current_dma; // Dedup barriers on the CURRENT load of a (tag_id, tag_slot): a conv reads one @@ -86,12 +87,6 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, 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) - // Async compute (matmul/preload) pipelines on the systolic array. A store needs the - // drained result, so it FLUSHes -- one barrier before the store waits all outstanding - // async compute, with no per-op completion events. - std::vector> outstanding_async; - std::shared_ptr pending_bar; // last COMPUTE_BAR fence, awaited by the next store - auto is_async_compute = [](int ct) { return ct == 1 || ct == 2; }; // matmul / preload auto flush = [&]() { if (sg && tile) { @@ -101,48 +96,74 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, } sg.reset(); tile.reset(); - last_writer.clear(); + writers.clear(); current_dma.clear(); bar_for_load.clear(); next_tag = 0; - outstanding_async.clear(); - pending_bar.reset(); }; - // Edges from the recorded read/write buffer sets: a reader depends on the last writer - // of each buffer it reads. An SA-producer -> matmul edge is an OCCUPANCY dependency - // (released at ISSUE); every other edge is a LATENCY dependency (released at finish). + // The single dataflow rule (sec 10.3). READ b: depend on all writers(b), ISSUE + // when both are SA ops else DONE. WRITE b: replace writers(b), except a commutative + // matmul accumulator, which waits only the seed and unions. WAR: resource models. const int MATMUL_CT = 1, PRELOAD_CT = 2; + auto is_mm_accum = [&](const std::shared_ptr& inst, int64_t b, + const std::vector& writes) { + if (inst->get_compute_type() != MATMUL_CT) return false; + for (int64_t w : writes) if (w == b) return true; + return false; + }; auto link = [&](std::shared_ptr inst, const std::vector& reads, const std::vector& writes) { for (int64_t b : reads) { - // A matmul reading its own accumulator imposes NO producer order: Y += X@W is - // commutative. Chaining them serializes the SA and DEADLOCKS the weight-slot - // pipeline. The store still waits every matmul via the COMPUTE_BAR fence. - bool is_accum = false; - for (int64_t w : writes) if (w == b) { is_accum = true; break; } - if (inst->get_compute_type() == MATMUL_CT && is_accum) continue; - auto it = last_writer.find(b); - if (it == last_writer.end()) continue; - int pct = it->second->get_compute_type(); - if (inst->get_compute_type() == MATMUL_CT && (pct == MATMUL_CT || pct == PRELOAD_CT)) - it->second->add_pipeline_child(inst); // SA pipeline -> occupancy (overlap) - else - it->second->add_child(inst); // data/result -> latency (full wait) + if (is_mm_accum(inst, b, writes)) continue; // accumulator read -> handled in WRITE (UNION) + auto it = writers.find(b); + if (it != writers.end()) + for (auto& w : it->second) { + int pct = w->get_compute_type(); + // both SA ops -> occupancy (overlap on the SA pipeline); else latency. + DepEvent on = (inst->get_compute_type() == MATMUL_CT && + (pct == MATMUL_CT || pct == PRELOAD_CT)) + ? DepEvent::ISSUE : DepEvent::DONE; + w->add_dep(inst, on); + } + } + for (int64_t b : writes) { + if (is_mm_accum(inst, b, writes)) { // UNION (commutative accumulate) + auto it = writers.find(b); + if (it != writers.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 + writers[b].push_back(inst); // join; no reset, no co-matmul edge + } else { // REPLACE (normal output; resets the producer set) + writers[b] = { inst }; + } } - for (int64_t b : writes) last_writer[b] = inst; tile->append_instuction(inst); }; - // SRAM-capacity tracking: 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. Only DMA-loaded buffers are tracked. v1 is single-core. + // 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 loads + 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); + } 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++; @@ -152,6 +173,21 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, ld->set_sram_alloc(cur_alloc[b]); vers[cur_alloc[b]].loads.push_back(ld); }; + // 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. + auto sram_on_write = [&](int64_t b, const std::shared_ptr& w) { + auto bb = buf_bytes.find(b); + if (bb == buf_bytes.end()) return; // size unknown (never DMA'd) -> untracked + 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_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 @@ -191,34 +227,29 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, int64_t uniq = next_tag++; // fresh Core tag key per dma record auto inst = make_dma(t, uniq); inst->set_tile_group(cur_tile_group); - size_t numel = 1; // SRAM footprint (ready-tile ordering) - for (auto d : t.dims) numel *= (size_t)d; - tile->inc_required_sram_size(numel * (t.elem_bits / 8)); + tile->inc_required_sram_size(rec_bytes(t)); // SRAM footprint (ready-tile ordering) if (t.dir == 1) { // STORE - if (pending_bar) { - // after a compute fence: wait it (drains the async matmuls) -- covers - // the accumulator read, so no per-buffer read edge. - pending_bar->add_child(inst); - pending_bar.reset(); - for (int64_t b : t.write_bufs) last_writer[b] = inst; - tile->append_instuction(inst); - } else { - link(inst, t.read_bufs, t.write_bufs); - } + // store reads the result buffer(s) -> link() JOINs all their writers. + link(inst, t.read_bufs, t.write_bufs); for (int64_t b : t.read_bufs) sram_on_read(b, inst); // store frees what it drains } else { // LOAD tile->append_instuction(inst); - // async load: the CURRENT load for this (tag_id, tag_slot), with a fresh uniq - // its barriers reuse. last_writer = the dma until its barrier overwrites it, so - // consumers gate on arrival. A sync load blocks to arrival itself. + // async load: the CURRENT load for this (tag_id, tag_slot), with a fresh + // uniq its barriers reuse. writers = the dma until its barrier overwrites it, + // 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) last_writer[b] = inst; - for (int64_t b : t.write_bufs) sram_on_load(b, inst); // occupy spad + 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 }; + sram_on_load(b, inst); // occupy spad + } } } else if (t.kind == TraceRec::MEMORY_BAR) { // The explicit async-DMA sync. Pair with the CURRENT load for this (tag_id, - // tag_slot), reusing its uniq: the dma releases the bar at issue, the bar parks on - // the tag until resp-complete, and consumers of the buffer gate on the bar. + // tag_slot), reusing its uniq: the dma releases the bar at issue, the bar parks + // on the tag until resp-complete, and becomes the load's handle in writers(b). auto it = current_dma.find({t.tag_id, t.tag_slot}); int64_t uniq = next_tag++; // fallback if unpaired std::shared_ptr dma_inst; @@ -227,31 +258,29 @@ 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) last_writer[b] = bf->second.second; + for (int64_t b : t.write_bufs) writers[b] = { bf->second.second }; continue; } auto bar = make_mem_bar(t, uniq); bar->set_tile_group(cur_tile_group); - if (dma_inst) dma_inst->add_child(bar); + if (dma_inst) dma_inst->add_dep(bar, DepEvent::DONE); tile->append_instuction(bar); - for (int64_t b : t.write_bufs) last_writer[b] = 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 }; bar_for_load[{t.tag_id, t.tag_slot}] = {uniq, bar}; } else if (t.kind == TraceRec::COMPUTE) { auto inst = make_compute(t); inst->set_tile_group(cur_tile_group); link(inst, t.read_bufs, t.write_bufs); - for (int64_t b : t.read_bufs) sram_on_read(b, inst); // frees the tiles it consumes - if (is_async_compute(t.compute_type)) outstanding_async.push_back(inst); - } else if (t.kind == TraceRec::COMPUTE_BAR) { - // explicit compute fence: ready once all outstanding async compute have - // ISSUED (pipeline-child release); the Core then waits the SA pipelines to - // drain before it finishes (-> the store it gates). - auto bar = std::make_shared(Opcode::COMPUTE_BAR); - bar->set_tile_group(cur_tile_group); - for (auto& a : outstanding_async) a->add_pipeline_child(bar); - outstanding_async.clear(); - tile->append_instuction(bar); - pending_bar = bar; + // in-place buffers (read AND written) are version-transparent (accumulator, + // in-place vector): skip the self-read and the self-write so footprint is not + // double-counted. read_bufs/write_bufs are tiny, so a linear scan beats a set. + auto in = [](const std::vector& v, int64_t b) { + return std::find(v.begin(), v.end(), b) != v.end(); + }; + 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(); From 903dad210e9859693813392808a1cc54a1aa3f29 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 25 Jun 2026 16:25:35 +0900 Subject: [PATCH 08/32] [Frontend] Run the spad-overflow check in timing-only mode, budget at spad/2 The validation-binary spad-overflow check sat inside `if functional_mode:`, so in timing-only / autotune (non-functional) runs an over-spad tiling was never rejected and reached TOGSim, which wedged ("spad too small") and crashed the compile via assert 0. Move the compile + check out of the functional gate (the Spike execution itself stays gated, run_spike below) and budget per dispatch at spad/2 -- two work-items run concurrently (double buffer), so each must fit half the spad or they deadlock competing for it. This matches the GEMM tiling gate (max_spad_size = spad/2), which pointwise ops lacked. Fixes the resnet / test_scheduler wedge where a fused BatchNorm+ReLU tile exceeded spad/2. --- PyTorchSimFrontend/extension_codecache.py | 86 ++++++++++++----------- 1 file changed, 45 insertions(+), 41 deletions(-) diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index a520d9a2..6914bb16 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -170,47 +170,51 @@ def load(cls, source_code, link_option = f"-Wl,--section-start=.spad=0x{spad_info['spad_vaddr']:x}" else: link_option = "" - # Generate LLVM kernel calller and binary for validation - if extension_config.pytorchsim_functional_mode: - # Use custom malloc to avoid size error - new_link_option = link_option + " -Wl,--wrap=malloc -Wl,--wrap=free" - cmds = mlir_compile_command(new_input_path, vectorlane_size, vlen=vlen) - opt_pad_cmd = shlex.split(cmds[0]) - translate_cmd = shlex.split(cmds[1]) - llc_cmd = shlex.split(cmds[2]) - llc_asm_cmd = shlex.split(cmds[3]) - try: - # loop-padding (mlir-opt) -> Python fine-grained + vcix (one parse/print) - subprocess.check_call(opt_pad_cmd) - run_module_passes(new_input_path + "_padded.mlir", - new_input_path + "_custom.mlir", - POST_OPT_PASSES, vectorlane=vectorlane_size, vlen=vlen) - # Standard MLIR -> LLVM-dialect lowering (registered upstream - # passes) runs in-process via the bindings PassManager, picking - # up after the custom mlir-opt passes (memref-to-gemmini). - run_standard_lowering(new_input_path + "_custom.mlir", new_input_path + "_llvm.mlir") - subprocess.check_call(translate_cmd) - subprocess.check_call(llc_cmd) - subprocess.check_call(llc_asm_cmd) - except subprocess.CalledProcessError as e: - logger.error(f"Command failed with exit code {e.returncode}") - logger.error(f"Error output: {e.output.decode() if isinstance(e.output, bytes) else e.output}") - assert(0) - - val_llvm_caller = MLIRKernelCallerCodeGen(extension_config.pytorchsim_functional_mode, arg_attributes) - val_llvm_caller.generate_wrapper_file(write_path, validation_wrapper_name) - val_llvm_caller.compile_wih_kernel(write_path, key, validation_wrapper_name, - validation_binary_name, new_link_option) - - stack_size = val_llvm_caller.parse_stack_sizes(f"{write_path}/{key}.s", vlenb=vlenb) - spad_size = val_llvm_caller.get_spad_size(validation_binary_path) - spad_usage = stack_size + spad_size # Spad usage per lane - if extension_config.CONFIG_SPAD_INFO["spad_size"] < spad_usage: - logger.debug( - f"Scratchpad size exceeded: required {spad_usage} bytes, " - f"but only {extension_config.CONFIG_SPAD_INFO['spad_size']} bytes available." - ) - raise SpadOverflowError() + # Compile a validation binary and measure its .spad section to reject + # over-spad tilings, even in timing-only mode -- else the tiling wedges + # TOGSim. Spike *execution* stays gated on functional_mode (run_spike). + new_link_option = link_option + " -Wl,--wrap=malloc -Wl,--wrap=free" + cmds = mlir_compile_command(new_input_path, vectorlane_size, vlen=vlen) + opt_pad_cmd = shlex.split(cmds[0]) + translate_cmd = shlex.split(cmds[1]) + llc_cmd = shlex.split(cmds[2]) + llc_asm_cmd = shlex.split(cmds[3]) + try: + # loop-padding (mlir-opt) -> Python fine-grained + vcix (one parse/print) + subprocess.check_call(opt_pad_cmd) + run_module_passes(new_input_path + "_padded.mlir", + new_input_path + "_custom.mlir", + POST_OPT_PASSES, vectorlane=vectorlane_size, vlen=vlen) + # Standard MLIR -> LLVM-dialect lowering (registered upstream + # passes) runs in-process via the bindings PassManager, picking + # up after the custom mlir-opt passes (memref-to-gemmini). + run_standard_lowering(new_input_path + "_custom.mlir", new_input_path + "_llvm.mlir") + subprocess.check_call(translate_cmd) + subprocess.check_call(llc_cmd) + subprocess.check_call(llc_asm_cmd) + except subprocess.CalledProcessError as e: + logger.error(f"Command failed with exit code {e.returncode}") + logger.error(f"Error output: {e.output.decode() if isinstance(e.output, bytes) else e.output}") + assert(0) + + val_llvm_caller = MLIRKernelCallerCodeGen(extension_config.pytorchsim_functional_mode, arg_attributes) + val_llvm_caller.generate_wrapper_file(write_path, validation_wrapper_name) + val_llvm_caller.compile_wih_kernel(write_path, key, validation_wrapper_name, + validation_binary_name, new_link_option) + + stack_size = val_llvm_caller.parse_stack_sizes(f"{write_path}/{key}.s", vlenb=vlenb) + spad_size = val_llvm_caller.get_spad_size(validation_binary_path) + spad_usage = stack_size + spad_size # Spad usage per lane + # Budget per dispatch = half the spad: two work-items run concurrently + # (double-buffer), so each must fit in spad/2 or they deadlock competing for + # the shared spad. Matches the GEMM tiling gate (max_spad_size = spad/2). + spad_budget = extension_config.CONFIG_SPAD_INFO["spad_size"] // 2 + if spad_budget < spad_usage: + logger.debug( + f"Scratchpad size exceeded: required {spad_usage} bytes, but only " + f"{spad_budget} bytes (spad/2, double-buffer budget) available." + ) + raise SpadOverflowError() # Launch tile graph generator gem5_pad_cmd = shlex.split(gem5_cmds[0]) From 8164151ac39cb4b100770b816782737defa48d2f Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 25 Jun 2026 16:49:46 +0900 Subject: [PATCH 09/32] [Frontend] Generate the trace.cpp ABI/API banner from togsim_runtime.h Every generated trace.cpp now opens with a "GENERATED ... DO NOT EDIT" banner carrying the TOGSim ABI version and the togsim_* call formats, so a dumped trace is self-documenting. Both are read from togsim_runtime.h at codegen time -- the version from the TOGSIM_ABI_VERSION define, the call-format text from a marked block next to the declarations -- so they never drift when the ABI changes. --- .../mlir/passes/lower_to_emitc.py | 37 +++++++++++++++---- TOGSim/include/togsim_runtime.h | 11 ++++++ 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index d10651c8..5633769a 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -41,6 +41,7 @@ """ import os +import re import subprocess from mlir.passmanager import PassManager @@ -64,8 +65,8 @@ "convert-arith-to-emitc," "convert-func-to-emitc)") -#: prepended to the mlir-to-cpp output; pulls in size_t/intN_t and the ABI. -_PRELUDE = ( +#: includes every generated trace.cpp needs. +_INCLUDES = ( "#include \n" "#include \n" "using std::size_t;\n" @@ -73,6 +74,24 @@ ) +def _trace_banner(include_dir): + """Self-documenting header for a generated trace.cpp: the ABI version (the + TOGSIM_ABI_VERSION #define) and the call-format banner, both READ FROM + togsim_runtime.h so they never drift when the ABI changes. Falls back to a + bare banner if the header is unreadable.""" + head = "// GENERATED by PyTorchSim -- TOGSim trace producer. DO NOT EDIT.\n" + try: + text = open(os.path.join(include_dir, "togsim_runtime.h")).read() + except OSError: + return head + ver = re.search(r"#define\s+TOGSIM_ABI_VERSION\s+(\d+)", text) + blk = re.search(r"// --- BEGIN trace-producer call formats[^\n]*\n(.*?)" + r"// --- END trace-producer call formats[^\n]*\n", text, re.S) + head = ("// GENERATED by PyTorchSim -- TOGSim trace producer (ABI v%s). DO NOT EDIT.\n" + % (ver.group(1) if ver else "?")) + return head + (blk.group(1) if blk else "") + + # --------------------------------------------------------------------------- # attribute builders / readers # --------------------------------------------------------------------------- @@ -508,15 +527,16 @@ def _mlir_translate_bin(): "mlir-translate") -def emitc_to_cpp(emitc_module, mlir_translate=None): - """Render `emitc_module` to C++ source (prelude + mlir-to-cpp body).""" +def emitc_to_cpp(emitc_module, mlir_translate=None, include_dir=None): + """Render `emitc_module` to C++ source (banner + includes + mlir-to-cpp body). + `include_dir` locates togsim_runtime.h for the ABI banner (default resolves it).""" mlir_translate = mlir_translate or _mlir_translate_bin() proc = subprocess.run( [mlir_translate, "--mlir-to-cpp"], input=str(emitc_module), capture_output=True, text=True) if proc.returncode != 0: raise RuntimeError("mlir-translate --mlir-to-cpp failed:\n" + proc.stderr) - return _PRELUDE + proc.stdout + return _trace_banner(include_dir or _default_include_dir()) + _INCLUDES + proc.stdout def compile_so(cpp_text, so_path, include_dir, cxx=None): @@ -547,8 +567,9 @@ def skeleton_to_so(skeleton_module, so_path, include_dir=None): """skeleton module -> EmitC -> C++ -> compiled trace `.so`. Returns the EmitC module text (for inspection / caching).""" emitc = lower_to_emitc(skeleton_module) - cpp = emitc_to_cpp(emitc) - compile_so(cpp, so_path, include_dir or _default_include_dir()) + inc = include_dir or _default_include_dir() + cpp = emitc_to_cpp(emitc, include_dir=inc) + compile_so(cpp, so_path, inc) return str(emitc) @@ -587,7 +608,7 @@ def main(argv): emitc = lower_to_emitc(module) if args.emit_mlir: open(args.emit_mlir, "w").write(str(emitc)) - cpp = emitc_to_cpp(emitc) + cpp = emitc_to_cpp(emitc, include_dir=args.include_dir or _default_include_dir()) if args.emit_cpp: open(args.emit_cpp, "w").write(cpp) compile_so(cpp, args.so, args.include_dir or _default_include_dir()) diff --git a/TOGSim/include/togsim_runtime.h b/TOGSim/include/togsim_runtime.h index e00d51d9..c3183926 100644 --- a/TOGSim/include/togsim_runtime.h +++ b/TOGSim/include/togsim_runtime.h @@ -28,6 +28,17 @@ typedef enum { // `strides` => contiguous. `is_async` => it finishes at ISSUE, so the barrier keyed // by `(tag_id, tag_slot)` gates consumers. `read_bufs`/`write_bufs` -> sec 10. +// --- BEGIN trace-producer call formats (copied verbatim into generated trace.cpp) --- +// Each togsim_* call below lowers 1:1 to one of these free functions. Arg formats: +// togsim_dma(ctx, dir, arg_id, offset, ndim, dims[], strides[], elem_bits, +// is_async, tag_id, tag_slot, read_bufs[], n_read, write_bufs[], n_write) +// dir: 0=load (MOVIN), 1=store (MOVOUT) +// togsim_compute(ctx, tile_id, compute_type, ndim, dims[], read_bufs[], n_read, +// write_bufs[], n_write) compute_type: 0=vector, 1=matmul, 2=preload +// togsim_memory_barrier(ctx, tag_id, tag_slot, write_bufs[], n_write) +// togsim_dispatch(ctx, tile_fn, iv[], n_iv) // run one work-item +// togsim_kernel(ctx, shape_args[], n_shape_args) // producer entry point +// --- END trace-producer call formats --- void togsim_dma(EmitCtx* ctx, int32_t dir, int32_t arg_id, uint64_t offset, int32_t ndim, const int64_t* dims, const int64_t* strides, int32_t elem_bits, From 2c872700f9ca6803824efd07b2b7af877ccd0ec7 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 25 Jun 2026 17:44:39 +0900 Subject: [PATCH 10/32] [TOGSim] Pick 1- vs 2-dispatch concurrency by per-dispatch spad footprint The bridge sums each work-item's distinct-buffer footprint (each buffer once, so a reduction's reloads of the same section do not inflate it) onto its Tile. can_issue then admits two concurrent dispatches only when each fits half the spad; a dispatch larger than spad/2 runs alone with the whole spad, so two work-items never compete for the shared spad and deadlock -- a runtime safety net beneath the codegen spad/2 gate. The footprint and resulting max_dispatch are logged on the TILE_SCHEDULED trace line for debugging. --- TOGSim/include/CoreTraceLog.h | 3 ++- TOGSim/include/Tile.h | 4 ++++ TOGSim/include/togsim_trace_bridge.h | 2 +- TOGSim/src/Core.cc | 18 ++++++++++++------ TOGSim/src/CoreTraceLog.cc | 6 ++++-- TOGSim/src/togsim_trace_bridge.cc | 15 +++++++++++++++ 6 files changed, 38 insertions(+), 10 deletions(-) diff --git a/TOGSim/include/CoreTraceLog.h b/TOGSim/include/CoreTraceLog.h index e78c1ef2..2ce51012 100644 --- a/TOGSim/include/CoreTraceLog.h +++ b/TOGSim/include/CoreTraceLog.h @@ -18,7 +18,8 @@ std::string format_dma_inst_issued_trace_line(Instruction& inst); /** Opcode + (detail...) for COMP / BAR / MOVIN / MOVOUT finished or issued lines. */ std::string format_instruction_detail_line(Instruction& inst); -void trace_tile_scheduled(cycle_type core_cycle, uint32_t core_id, const std::string& tag15); +void trace_tile_scheduled(cycle_type core_cycle, uint32_t core_id, const std::string& tag15, + size_t spad_footprint, int max_dispatch); void trace_instruction_line(cycle_type core_cycle, uint32_t core_id, diff --git a/TOGSim/include/Tile.h b/TOGSim/include/Tile.h index d867a037..b11b9f8e 100644 --- a/TOGSim/include/Tile.h +++ b/TOGSim/include/Tile.h @@ -28,6 +28,9 @@ class Tile : public std::enable_shared_from_this { size_t get_required_sram_size() { return _required_sram_size; } void set_required_sram_size(size_t sram_size) { _required_sram_size=sram_size; } void inc_required_sram_size(size_t sram_size) { _required_sram_size+=sram_size; } + // Dispatch spad-buffer footprint (bytes, codegen .spad x lanes); Core picks 1- vs 2-dispatch by it. + size_t get_spad_footprint() { return _spad_footprint; } + void set_spad_footprint(size_t b) { _spad_footprint = b; } void append_instuction(std::shared_ptr& inst); void append_child(std::shared_ptr child); std::vector>& get_child_tile () { return _child_tiles; } @@ -52,6 +55,7 @@ class Tile : public std::enable_shared_from_this { std::shared_ptr _onwer_graph; Status _status = Status::EMPTY; size_t _required_sram_size=0; + size_t _spad_footprint=0; size_t _ready_counter=0; size_t _nr_insts = 0; size_t _nr_finished_insts = 0; diff --git a/TOGSim/include/togsim_trace_bridge.h b/TOGSim/include/togsim_trace_bridge.h index f4b53c80..212cd4d6 100644 --- a/TOGSim/include/togsim_trace_bridge.h +++ b/TOGSim/include/togsim_trace_bridge.h @@ -7,6 +7,6 @@ #include "TileGraph.h" #include "togsim_loader.h" -// Build a TileGraph from a recorded trace. `path`/`name` label the graph. +// 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); diff --git a/TOGSim/src/Core.cc b/TOGSim/src/Core.cc index 23ee4d6d..d3d174c8 100644 --- a/TOGSim/src/Core.cc +++ b/TOGSim/src/Core.cc @@ -81,17 +81,23 @@ bool Core::try_occupy_sram(const std::shared_ptr& inst) { } bool Core::can_issue(const std::shared_ptr& op) { - /* Bound concurrent dispatches so their combined spad working set fits: with the - * global @buffers each in-flight dispatch piles its own load versions, and too - * many at once overflow the spad (versions never free -> wedge). 2 keeps double- - * buffering overlap while leaving headroom. */ - return _tiles.size() < 2 && !op->is_stonne_tile(); + /* Bound concurrent dispatches so their combined spad working set fits. Two run + * concurrently (double-buffer) only if each fits half the spad; a dispatch whose + * footprint exceeds spad/2 runs alone with the whole spad (else two would compete + * for the shared spad and deadlock). spad_footprint = codegen .spad x lanes; 0 + * (unknown) falls back to 2. */ + size_t M = op->get_spad_footprint(); + int max_concurrent = (_sram_capacity && M > _sram_capacity / 2) ? 1 : 2; + return (int)_tiles.size() < max_concurrent && !op->is_stonne_tile(); } void Core::issue(std::shared_ptr op) { if (op->get_instructions().size()) { + size_t M = op->get_spad_footprint(); + int max_dispatch = (_sram_capacity && M > _sram_capacity / 2) ? 1 : 2; core_trace_log::trace_tile_scheduled(_core_cycle, _id, - TraceLogTag::pad15(TraceLogTag::kTileScheduled)); + TraceLogTag::pad15(TraceLogTag::kTileScheduled), + M, max_dispatch); } for (const auto& inst : op->get_instructions()) { if (inst->is_ready()) diff --git a/TOGSim/src/CoreTraceLog.cc b/TOGSim/src/CoreTraceLog.cc index 7086893e..1f6b720c 100644 --- a/TOGSim/src/CoreTraceLog.cc +++ b/TOGSim/src/CoreTraceLog.cc @@ -86,8 +86,10 @@ std::string format_instruction_detail_line(Instruction& inst) { return opname; } -void trace_tile_scheduled(cycle_type core_cycle, uint32_t core_id, const std::string& tag15) { - spdlog::trace("[{}][Core {}][{}]", core_cycle, core_id, tag15); +void trace_tile_scheduled(cycle_type core_cycle, uint32_t core_id, const std::string& tag15, + size_t spad_footprint, int max_dispatch) { + spdlog::trace("[{}][Core {}][{}] spad_footprint={} max_dispatch={}", + core_cycle, core_id, tag15, spad_footprint, max_dispatch); } void trace_instruction_line(cycle_type core_cycle, diff --git a/TOGSim/src/togsim_trace_bridge.cc b/TOGSim/src/togsim_trace_bridge.cc index 08c40afb..46e6f7b3 100644 --- a/TOGSim/src/togsim_trace_bridge.cc +++ b/TOGSim/src/togsim_trace_bridge.cc @@ -87,9 +87,12 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, 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 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); @@ -99,6 +102,8 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, writers.clear(); current_dma.clear(); bar_for_load.clear(); + cur_tile_bufs.clear(); + cur_tile_footprint = 0; next_tag = 0; }; @@ -164,6 +169,14 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, 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); } + // 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) + if (cur_tile_bufs.insert(b).second) { + auto it = buf_bytes.find(b); + if (it != buf_bytes.end()) cur_tile_footprint += it->second; + } + }; 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++; @@ -228,6 +241,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, auto inst = make_dma(t, uniq); inst->set_tile_group(cur_tile_group); tile->inc_required_sram_size(rec_bytes(t)); // SRAM footprint (ready-tile ordering) + note_bufs(t.read_bufs); note_bufs(t.write_bufs); // distinct-buffer footprint for 1- vs 2-dispatch if (t.dir == 1) { // STORE // store reads the result buffer(s) -> link() JOINs all their writers. link(inst, t.read_bufs, t.write_bufs); @@ -273,6 +287,7 @@ std::unique_ptr trace_to_tilegraph(const togsim::RunResult& run, auto inst = make_compute(t); inst->set_tile_group(cur_tile_group); link(inst, t.read_bufs, t.write_bufs); + note_bufs(t.read_bufs); note_bufs(t.write_bufs); // distinct-buffer footprint for 1- vs 2-dispatch // in-place buffers (read AND written) are version-transparent (accumulator, // in-place vector): skip the self-read and the self-write so footprint is not // double-counted. read_bufs/write_bufs are tiny, so a linear scan beats a set. From 1cdeb51d966a9683c76f4bea97c677d341667ae2 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 25 Jun 2026 18:07:58 +0900 Subject: [PATCH 11/32] [Frontend] Budget spad buffers honestly in GEMM tile selection Count the fused-epilogue buffers a tiling actually needs, and stop charging the kernel's stack frame against the spad budget -- it lives in DRAM, not the scratchpad, so including it rejected tilings that fit. --- PyTorchSimFrontend/extension_codecache.py | 5 ++--- PyTorchSimFrontend/mlir/mlir_gemm_template.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 6914bb16..18e0d3a9 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -202,9 +202,8 @@ def load(cls, source_code, val_llvm_caller.compile_wih_kernel(write_path, key, validation_wrapper_name, validation_binary_name, new_link_option) - stack_size = val_llvm_caller.parse_stack_sizes(f"{write_path}/{key}.s", vlenb=vlenb) - spad_size = val_llvm_caller.get_spad_size(validation_binary_path) - spad_usage = stack_size + spad_size # Spad usage per lane + # Only the .spad section consumes the scratchpad; the stack frame lives in main memory (sp in the -m region, not the scratchpad vaddr) so it is not charged against the per-lane spad budget. + spad_usage = val_llvm_caller.get_spad_size(validation_binary_path) # Budget per dispatch = half the spad: two work-items run concurrently # (double-buffer), so each must fit in spad/2 or they deadlock competing for # the shared spad. Matches the GEMM tiling gate (max_spad_size = spad/2). diff --git a/PyTorchSimFrontend/mlir/mlir_gemm_template.py b/PyTorchSimFrontend/mlir/mlir_gemm_template.py index 8a8cd585..871c244e 100644 --- a/PyTorchSimFrontend/mlir/mlir_gemm_template.py +++ b/PyTorchSimFrontend/mlir/mlir_gemm_template.py @@ -329,7 +329,7 @@ def select_tile(self, kernel, M, N, K, n_extra_node, n_extra_read, n_prologue_no else: # case 2: use heuristic mapping min_tile = (n_extra_node + n_prologue_node) == 0 - tile_candidates = kernel.gemm_combination_mapping(M, N, K, max(n_extra_read-2, 0), n_prologue_node, min_tile=True, precision_bytes=precision_bytes) + tile_candidates = kernel.gemm_combination_mapping(M, N, K, n_extra_node + n_extra_read, n_prologue_node, min_tile=True, precision_bytes=precision_bytes) # Edge case if (M == 0) or (N == 0) or (K == 0): From 6ae56f22e1db71bc5a905dc749ef46e6fd8f2ddb Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 25 Jun 2026 21:55:28 +0900 Subject: [PATCH 12/32] [CI] Bump spike pin to v1.0.3 Spike v1.0.3 zero-inits the MVIN/MVOUT DMA address buffer so ROUNDUP padding entries are skipped instead of dereferencing uninitialized garbage, fixing the host store segfault on 8-lane (wrapper3) configs where a tile's split axis is not a multiple of vlane_stride*n_vu. --- thirdparty/github-releases.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/github-releases.json b/thirdparty/github-releases.json index 4af6ccbd..390b68fb 100644 --- a/thirdparty/github-releases.json +++ b/thirdparty/github-releases.json @@ -13,7 +13,7 @@ }, "spike": { "repository": "PSAL-POSTECH/riscv-isa-sim", - "release_tag": "v1.0.2", + "release_tag": "v1.0.3", "asset_name": "spike-release.tar.gz" } } From 14302e18a340fdb5fa99853ef8652a0edaed971d Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 25 Jun 2026 21:11:52 +0900 Subject: [PATCH 13/32] [Frontend] Add a per-kernel CPU functional-verify sub-option pytorchsim_functional_verify_per_kernel compares every realized Spike buffer against a CPU golden and stops at the first divergent kernel, naming the op and the offending indices. --- CLAUDE.md | 7 + PyTorchSimFrontend/extension_config.py | 6 + .../extension_functional_verify.py | 165 ++++++++++++++++++ .../mlir/mlir_codegen_backend.py | 46 +++++ 4 files changed, 224 insertions(+) create mode 100644 PyTorchSimFrontend/extension_functional_verify.py diff --git a/CLAUDE.md b/CLAUDE.md index 5a3a47cd..46f07e1b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -58,6 +58,12 @@ export TORCHSIM_DUMP_MLIR_IR=1 export TORCHSIM_DUMP_LLVM_IR=1 ``` +**To find which op a wrong result first diverges at** (per-kernel CPU cross-check; +sub-option of functional mode). Set `pytorchsim_functional_verify_per_kernel: 1` +in the config YAML, clear the codegen cache, and re-run: each compiled kernel's +output is compared to a CPU golden and the run stops at the first divergent +kernel, naming the op and offending indices. + ## Key environment variables Read in `PyTorchSimFrontend/extension_config.py`: @@ -91,6 +97,7 @@ Located under `configs/*.yml`: - `icnt_type` (`simple` | `booksim`), `icnt_latency_cycles`, `icnt_freq_mhz`, `icnt_config_path` - `l2d_type` (e.g., `datacache`), `l2d_config` (AccelSim-format cache config string) - `pytorchsim_functional_mode` (Spike on/off), `pytorchsim_timing_mode` +- `pytorchsim_functional_verify_per_kernel` (debug: per-kernel CPU cross-check) - `codegen_mapping_strategy`: `heuristic` | `autotune` | `external-then-heuristic` | `external-then-autotune` - `codegen_external_mapping_file` (key `"M_N_K"` → `{TILE_M, TILE_K, TILE_N}` JSON) - `codegen_compiler_optimization`: `"all"` | `"none"` | a list from `{fusion, reduction_epilogue, reduction_reduction, prologue, single_batch_conv, multi_tile_conv, subtile}` diff --git a/PyTorchSimFrontend/extension_config.py b/PyTorchSimFrontend/extension_config.py index 0e1bc422..2d706bbc 100644 --- a/PyTorchSimFrontend/extension_config.py +++ b/PyTorchSimFrontend/extension_config.py @@ -57,6 +57,12 @@ def __getattr__(name): return config_yaml['pytorchsim_functional_mode'] if name == "pytorchsim_timing_mode": return config_yaml['pytorchsim_timing_mode'] + # Sub-option of functional mode: compare every realized Spike buffer against a CPU + # golden to localize the first kernel whose value diverges. Auto-disabled when + # functional mode is off (there are no Spike values to verify). + if name == "pytorchsim_functional_verify_per_kernel": + return bool(config_yaml.get('pytorchsim_functional_verify_per_kernel', False)) \ + and bool(config_yaml['pytorchsim_functional_mode']) # Mapping strategy if name == "codegen_mapping_strategy": diff --git a/PyTorchSimFrontend/extension_functional_verify.py b/PyTorchSimFrontend/extension_functional_verify.py new file mode 100644 index 00000000..de21be73 --- /dev/null +++ b/PyTorchSimFrontend/extension_functional_verify.py @@ -0,0 +1,165 @@ +"""Per-kernel CPU cross-check for the functional (Spike) path. + +This is a sub-option of ``pytorchsim_functional_mode``: enable it with the YAML +key ``pytorchsim_functional_verify_per_kernel: 1`` (auto-disabled when functional +mode is off -- there are no Spike values to verify otherwise). + +When enabled, the generated wrapper compares every realized buffer (the output of +each compiled, fused kernel that Spike just produced) against a CPU "golden" +reference. The golden is computed once per graph execution by running the +original aten graph (``V.graph.module``) on CPU with the same inputs. The first +buffer whose value diverges beyond tolerance pinpoints the fused op-cluster that +injected the error -- the finest granularity observable in a fused pipeline, +since intermediate aten ops inside a kernel are not separately materialized. + +Tolerances (read once at import): + TORCHSIM_FUNCTIONAL_VERIFY_RTOL relative tolerance (default 1e-4) + TORCHSIM_FUNCTIONAL_VERIFY_ATOL absolute tolerance (default 1e-4) + +The check raises FunctionalVerifyMismatch at the first divergent buffer +(stop-at-first), after logging the kernel, originating fx op, offending indices +and max abs diff. + +NOTE: codegen bakes the check calls into the wrapper only when enabled at +*compile* time. Toggle the option and clear the codegen cache together +(scripts/clear_codegen_cache.sh), or a cached wrapper without checks is replayed. +""" +import os + +import torch + +from PyTorchSimFrontend import extension_config +from PyTorchSimFrontend.extension_config import setup_logger + +logger = setup_logger(__name__) + +RTOL = float(os.environ.get("TORCHSIM_FUNCTIONAL_VERIFY_RTOL", "1e-4")) +ATOL = float(os.environ.get("TORCHSIM_FUNCTIONAL_VERIFY_ATOL", "1e-4")) + +# Codegen-time registry of runnable aten graphs, keyed by an int id baked into +# the generated wrapper. Persists in-process: the exec'd wrapper imports this +# very module, so the dict it sees is the one populated during codegen. +_GRAPHS = {} + +# Per-run state, reset by verify_init at the top of each call(args). +_STATE = {"golden": None, "n_checked": 0, "failed": False} + + +class FunctionalVerifyMismatch(RuntimeError): + """Raised at the first kernel buffer that diverges from the CPU golden.""" + + +def enabled(): + """True when functional mode AND the per-kernel verify sub-option are on. + + Read dynamically (the config can change inside a ``with TOGSimulator(...)`` + block); the config accessor already AND-gates this with functional mode. + """ + try: + return bool(extension_config.pytorchsim_functional_verify_per_kernel) + except Exception: + return False + + +def register_graph(gm): + """Register a runnable aten GraphModule, returning an id for the wrapper.""" + gid = len(_GRAPHS) + _GRAPHS[gid] = gm + return gid + + +class _GoldenInterpreter(torch.fx.Interpreter): + """Run the aten graph on CPU, recording each node's tensor output by name.""" + + def __init__(self, gm): + super().__init__(gm) + self.values = {} + + def run_node(self, n): + out = super().run_node(n) + # Move EVERY tensor output to CPU, preserving dtype, so a device-baked op + # (arange(device=npu)) or a consumer that needs co-located operands (aten.index) + # sees CPU tensors. Only the recorded golden is cast to float32 for allclose. + out = torch.utils._pytree.tree_map( + lambda x: x.detach().to("cpu") if isinstance(x, torch.Tensor) else x, out) + if isinstance(out, torch.Tensor): + self.values[n.name] = out.to(torch.float32) + return out + + +def _to_cpu(x): + return x.detach().cpu() if isinstance(x, torch.Tensor) else x + + +def verify_init(gid, inputs): + """Compute the CPU golden for this graph execution (top of call(args)).""" + _STATE["golden"] = None + _STATE["n_checked"] = 0 + _STATE["failed"] = False + gm = _GRAPHS.get(gid) + if gm is None: + logger.warning("[FuncVerify] no graph registered for id %s", gid) + return + try: + cpu_inputs = [_to_cpu(x) for x in inputs] + interp = _GoldenInterpreter(gm) + interp.run(*cpu_inputs) + _STATE["golden"] = interp.values + logger.info("[FuncVerify] golden ready: %d node tensors (rtol=%g atol=%g)", + len(interp.values), RTOL, ATOL) + except Exception as e: # never let the verify shadow break the real run + logger.warning("[FuncVerify] golden computation failed (%r); " + "per-kernel verify disabled for this graph", e) + _STATE["golden"] = None + + +def verify_check(value, buffer_name, node_name, op): + """Compare one realized buffer against its golden; raise on first divergence.""" + if _STATE["failed"]: + return + golden = _STATE["golden"] + if golden is None or not isinstance(value, torch.Tensor): + return + ref = golden.get(node_name) + if ref is None: + return # no reference captured for this fx node (non-tensor / folded) + val = value.detach().to("cpu", torch.float32) + if val.shape != ref.shape: + if val.numel() != ref.numel(): + return + val = val.reshape(ref.shape) + + _STATE["n_checked"] += 1 + if torch.allclose(val, ref, rtol=RTOL, atol=ATOL, equal_nan=True): + logger.debug("[FuncVerify] PASS %-14s %-28s %s", + buffer_name, op, tuple(ref.shape)) + return + + # ---- divergence ---- + _STATE["failed"] = True + diff = (val - ref).abs() + tol = ATOL + RTOL * ref.abs() + bad = diff > tol + n_bad = int(bad.sum()) + maxd = float(diff.max()) + idxs = bad.nonzero() + first = idxs[0].tolist() if idxs.numel() else None + sample = [f" {tuple(r.tolist())}: npu={val[tuple(r.tolist())].item():.6g} " + f"cpu={ref[tuple(r.tolist())].item():.6g}" + for r in idxs[:6]] + logger.error( + "\n================= PER-KERNEL FUNCTIONAL VERIFY: DIVERGENCE =================\n" + " first divergent buffer : %s\n" + " originating fx op : %s (node '%s')\n" + " shape : %s\n" + " elements over tol : %d / %d\n" + " max abs diff : %.6g (rtol=%g atol=%g)\n" + " first bad index : %s\n" + " buffers verified OK : %d\n" + " sample mismatches (npu vs cpu):\n%s\n" + "===========================================================================", + buffer_name, op, node_name, tuple(ref.shape), n_bad, ref.numel(), + maxd, RTOL, ATOL, first, _STATE["n_checked"] - 1, "\n".join(sample)) + raise FunctionalVerifyMismatch( + f"[FuncVerify] first divergence at buffer '{buffer_name}' (op {op}): " + f"max abs diff {maxd:.6g}, {n_bad}/{ref.numel()} elements over tol") diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 2360542c..40415f01 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -24,6 +24,7 @@ ) from torch.utils._sympy.functions import ModularIndexing, FloorDiv from PyTorchSimFrontend import extension_codecache +from PyTorchSimFrontend import extension_functional_verify as _func_verify from . import mlir_common from .mlir_common import LoopLevel, LoopNest from .mlir_ops import ExtensionOverrides @@ -103,6 +104,7 @@ def write_header(self): from PyTorchSimFrontend.extension_config import CONFIG_SRAM_BUFFER_PLAN, setup_logger from Simulator.simulator import TOGSimulator from PyTorchSimFrontend.extension_op import sparse_mm_dummy_stonne_outer + from PyTorchSimFrontend import extension_functional_verify as _fverify from torch._inductor.select_algorithm import extern_kernels # Configure logger for generated wrapper code @@ -162,6 +164,16 @@ def call(args): self.prefix.writeline(f"{lhs} = args") self.prefix.writeline("args.clear()") + # Per-kernel functional verify: register the runnable aten graph and + # emit a CPU golden build at the top of call(), passing graph inputs. + if _func_verify.enabled(): + gm = getattr(V.graph, "module", None) + if gm is not None: + gid = _func_verify.register_graph(gm) + in_names = list(V.graph.graph_inputs.keys()) + self.prefix.writeline( + f"_fverify.verify_init({gid}, [{', '.join(in_names)}])") + self.codegen_inputs() self.codegen_input_size_asserts() self.codegen_sram_plan_prefix() @@ -204,6 +216,7 @@ def generate(self, is_inference): result = IndentedBuffer() # result.splice(self.header) + self._fverify_seen = set() with contextlib.ExitStack() as stack: stack.enter_context(self.wrapper_call.indent()) self.memory_plan_reuse() @@ -220,6 +233,8 @@ def generate(self, is_inference): line.codegen(self.wrapper_call) elif isinstance(line, wrapper.KernelCallLine): self.wrapper_call.writeline(self.wrap_kernel_call(line.kernel_name, line.call_args)) + if _func_verify.enabled(): + self._fverify_emit_checks(line.call_args) else: if isinstance(line, wrapper.WrapperLine): line.codegen(self.wrapper_call) @@ -249,6 +264,37 @@ def generate(self, is_inference): self.kernel_declarations.getvaluewithlinemap(), ) + def _fverify_emit_checks(self, call_args): + """Emit per-kernel CPU verify calls for this kernel's output buffers. + + A buffer's value is produced by the first kernel that names it (producer + precedes consumers in topo order), so we check each bare-identifier buffer + arg the first time it is seen -- that occurrence is its output. The buffer + is mapped to its originating fx node (op) so the runtime check can compare + against the CPU golden keyed by that node. + """ + for a in call_args: + if not isinstance(a, str): + continue + name = a.strip() + if not name.isidentifier() or name in self._fverify_seen: + continue + self._fverify_seen.add(name) + if name in V.graph.graph_inputs: + continue # placeholders: golden == input, nothing to verify + try: + buf = V.graph.get_buffer(name) + except Exception: + buf = None + if buf is None: + continue + origin = getattr(buf, "origin_node", None) + if origin is None: + continue + op = str(getattr(origin, "target", "?")) + self.wrapper_call.writeline( + f'_fverify.verify_check({name}, "{name}", "{origin.name}", "{op}")') + def memory_plan(self): self.lines = memory_planning.MemoryPlanner(self).plan(self.lines) From a8b67c8f6ef21be2381bd64c2f95e0a113cc309f Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Fri, 26 Jun 2026 00:08:03 +0900 Subject: [PATCH 14/32] [Frontend] Budget fused-prologue spad buffers in BMM tile selection The GEMM path already budgeted its fused epilogue; BMM's fused prologue was unaccounted for, so a tiling that overflowed the scratchpad could be chosen. --- PyTorchSimFrontend/mlir/mlir_bmm_template.py | 40 +++++++++++++++++--- PyTorchSimFrontend/mlir/mlir_template.py | 10 ++--- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_bmm_template.py b/PyTorchSimFrontend/mlir/mlir_bmm_template.py index c5fd902f..5323fd7c 100644 --- a/PyTorchSimFrontend/mlir/mlir_bmm_template.py +++ b/PyTorchSimFrontend/mlir/mlir_bmm_template.py @@ -165,10 +165,10 @@ def render(self, prologue_nodes: Optional[List[IRNode]] = None, tile_info = None, **kwargs): - X, W, Y, Bias, W_tensor, X_tensor, B, M, N, K, n_extra_node, n_prologue_node = self.extract_info(template_buffer_node, epilogue_nodes, prologue_nodes) + X, W, Y, Bias, W_tensor, X_tensor, B, M, N, K, n_extra_node, n_prologue_node, n_extra_read = self.extract_info(template_buffer_node, epilogue_nodes, prologue_nodes) precision_bytes = mlir_common.get_dtype_nbytes(X.get_dtype()) if tile_info is None: - TILE_M, TILE_N, TILE_K, SUB_TILE_M, SUB_TILE_N, SUB_TILE_K = self.select_tile(kernel, M, N, K, n_extra_node, 0, n_prologue_node, precision_bytes)[0] + TILE_M, TILE_N, TILE_K, SUB_TILE_M, SUB_TILE_N, SUB_TILE_K = self.select_tile(kernel, M, N, K, n_extra_node, n_extra_read, n_prologue_node, precision_bytes)[0] else: TILE_M, TILE_N, TILE_K, SUB_TILE_M, SUB_TILE_N, SUB_TILE_K = tile_info @@ -342,7 +342,32 @@ def extract_info(self, template_buffer_node, epilogue_nodes, prologue_nodes): # Select tile size n_extra_node = len(epilogue_nodes) if epilogue_nodes is not None else 0 n_prologue_node = len(prologue_nodes) if prologue_nodes is not None else 0 - return X,W,Y,Bias,W_tensor,X_tensor,B,M,N,K,n_extra_node, n_prologue_node + n_extra_read = self.count_prologue_extra_buffers(prologue_nodes) + return X,W,Y,Bias,W_tensor,X_tensor,B,M,N,K,n_extra_node, n_prologue_node, n_extra_read + + def count_prologue_extra_buffers(self, prologue_nodes): + # Count prologue reads needing their own .spad global: the numel-matching main input reuses the matmul-operand buffer, every other read (e.g. softmax max/sum) gets a disjoint one (see codegen_template_code). + if not prologue_nodes: + return 0 + from functools import reduce + import operator + from torch._inductor.virtualized import V + buf_dict = {val.name: val for val in V.graph.buffers} + buf_dict.update(V.graph.graph_inputs) + prologue_outputs = {list(node.read_writes.writes)[0].name for node in prologue_nodes} + extra = set() + for node in prologue_nodes: + reads = sorted(i.name for i in node.read_writes.reads) + main_input = None + for candidate_read in reads: + if candidate_read in buf_dict and \ + reduce(operator.mul, buf_dict[candidate_read].get_size(), 1) == node.node.get_numel(): + main_input = candidate_read + break + for r in reads: + if r != main_input and r not in prologue_outputs: + extra.add(r) + return len(extra) def get_tile_candidates(self, kernel: MLIRTemplateKernel, @@ -350,12 +375,15 @@ def get_tile_candidates(self, epilogue_nodes: Optional[List[IRNode]] = None, prologue_nodes: Optional[List[IRNode]] = None, **kwargs): - X, W, Y, Bias, W_tensor, X_tensor, B, M, N, K, n_extra_node, n_prologue_node = self.extract_info(template_buffer_node, epilogue_nodes, prologue_nodes) + X, W, Y, Bias, W_tensor, X_tensor, B, M, N, K, n_extra_node, n_prologue_node, n_extra_read = self.extract_info(template_buffer_node, epilogue_nodes, prologue_nodes) precision_bytes = mlir_common.get_dtype_nbytes(X.get_dtype()) - return self.select_tile(kernel, M, N, K, n_extra_node, 0, n_prologue_node, precision_bytes) + return self.select_tile(kernel, M, N, K, n_extra_node, n_extra_read, n_prologue_node, precision_bytes) def select_tile(self, kernel, M, N, K, n_extra_node, n_extra_read, n_prologue_node, precision_bytes): - tile_candidates = kernel.gemm_combination_mapping(M, N, K, n_extra_node=n_extra_node, precision_bytes=precision_bytes) + # Budget the prologue extra-read globals as weight-tile buffers (their emitted size) so the chosen tile fits spad/2. + tile_candidates = kernel.gemm_combination_mapping( + M, N, K, n_extra_node=n_extra_node, n_prologue_extra_read=n_extra_read, + precision_bytes=precision_bytes) for idx, (TILE_M, TILE_N, TILE_K) in enumerate(tile_candidates): SUB_TILE_M = TILE_M if (TILE_M < kernel.vector_lane) or n_prologue_node else kernel.vector_lane SUB_TILE_N = TILE_N # if (TILE_N < kernel.vector_lane) or prologue_nodes else kernel.vector_lane diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index 2b8a0676..cd3ca018 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -204,7 +204,7 @@ def gemmini_gemm_mapping(self, M, N, K, precision_bytes=4): return inner_I, inner_J, inner_K - def gemm_combination_mapping(self, M, N, K, n_extra_node=0, n_prologue_node=0, pad_k=True, min_tile=False, is_conv=False, precision_bytes=4): + def gemm_combination_mapping(self, M, N, K, n_extra_node=0, n_prologue_node=0, n_prologue_extra_read=0, pad_k=True, min_tile=False, is_conv=False, precision_bytes=4): tile_candidates = [] spad_size_per_lane = self.spad_info["spad_size"] spad_size = spad_size_per_lane * self.vector_lane @@ -232,8 +232,8 @@ def gemm_combination_mapping(self, M, N, K, n_extra_node=0, n_prologue_node=0, p tile_M = i * self.vector_lane if M > self.vector_lane else M_padded for j in tile_N_range: tile_N = j * self.vector_lane if N > self.vector_lane else N_padded - used_spad_size = (tile_M * tile_K * (1 + n_prologue_node) + tile_K * tile_N + tile_M * tile_N * (1 + n_extra_node)) * precision_bytes - weight_size_per_lane = self.get_spad_size_per_lane(tile_K, tile_N) + used_spad_size = (tile_M * tile_K * (1 + n_prologue_node) + tile_K * tile_N * (1 + n_prologue_extra_read) + tile_M * tile_N * (1 + n_extra_node)) * precision_bytes + weight_size_per_lane = self.get_spad_size_per_lane(tile_K, tile_N) * (1 + n_prologue_extra_read) input_size_per_lane = self.get_spad_size_per_lane(tile_M * (1 + n_prologue_node), tile_K) output_size_per_lane = self.get_spad_size_per_lane(tile_M * (1 + n_extra_node), tile_N) used_spad_size_per_lane = (weight_size_per_lane + input_size_per_lane + output_size_per_lane) * precision_bytes @@ -258,8 +258,8 @@ def gemm_combination_mapping(self, M, N, K, n_extra_node=0, n_prologue_node=0, p tile_M = i * self.vector_lane if M > self.vector_lane else M_padded for j in tile_N_range: tile_N = j * self.vector_lane if N > self.vector_lane else N_padded - used_spad_size = (tile_M * tile_K * (1 + n_prologue_node) + tile_K * tile_N + tile_M * tile_N * (1 + n_extra_node)) * precision_bytes - weight_size_per_lane = self.get_spad_size_per_lane(tile_K, tile_N) + used_spad_size = (tile_M * tile_K * (1 + n_prologue_node) + tile_K * tile_N * (1 + n_prologue_extra_read) + tile_M * tile_N * (1 + n_extra_node)) * precision_bytes + weight_size_per_lane = self.get_spad_size_per_lane(tile_K, tile_N) * (1 + n_prologue_extra_read) input_size_per_lane = self.get_spad_size_per_lane(tile_M * (1 + n_prologue_node), tile_K) output_size_per_lane = self.get_spad_size_per_lane(tile_M * (1 + n_extra_node), tile_N) used_spad_size_per_lane = (weight_size_per_lane + input_size_per_lane + output_size_per_lane) * precision_bytes From 25491ad6dcab5e70ae12ec4aea1c647f54e727c3 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Fri, 26 Jun 2026 17:33:17 +0900 Subject: [PATCH 15/32] [TOGSim] Trim per-cycle overhead in the Core issue loop Skip the ready-queue scan when nothing can newly issue, skip building the per-instruction trace string when trace logging is off, and fix a use-after-erase in the zero-cycle COMP skip path. --- TOGSim/include/Core.h | 5 +++++ TOGSim/include/Instruction.h | 5 +++++ TOGSim/src/Core.cc | 21 +++++++++++++++++++-- TOGSim/src/CoreTraceLog.cc | 4 ++++ 4 files changed, 33 insertions(+), 2 deletions(-) diff --git a/TOGSim/include/Core.h b/TOGSim/include/Core.h index a8e7e54e..dbc91950 100644 --- a/TOGSim/include/Core.h +++ b/TOGSim/include/Core.h @@ -115,6 +115,11 @@ class Core { std::vector> _tiles; std::queue> _finished_tiles; + // Issue-scan re-arm (perf): cycle() skips the ready-queue scan unless this is set. + // EVERY event that can make a stalled instruction issuable must set it -- a new + // issue-gating throttle that forgets to will make cycle() skip the scan forever. + bool _issue_dirty = true; + std::queue> _vu_compute_pipeline; std::vector>> _sa_compute_pipeline; std::queue> _ld_inst_queue; diff --git a/TOGSim/include/Instruction.h b/TOGSim/include/Instruction.h index d4995eb9..74c5e339 100644 --- a/TOGSim/include/Instruction.h +++ b/TOGSim/include/Instruction.h @@ -77,6 +77,7 @@ class Instruction : public std::enable_shared_from_this { ready_counter--; if (!ready_counter && _owner_ready_queue_ref != nullptr) { _owner_ready_queue_ref->push_back(shared_from_this()); + if (_owner_dirty_ref) *_owner_dirty_ref = true; // ready set grew -> re-arm the owning Core's issue scan } } size_t get_tile_numel() { return _tile_numel; } @@ -103,6 +104,9 @@ class Instruction : public std::enable_shared_from_this { void* get_owner() { return _owner; } void set_owner(void *owner) { _owner = owner;} void set_owner_ready_queue(std::list>* q) { _owner_ready_queue_ref = q; } + // Points at the owning Core's _issue_dirty; set when the tile is issued to a core + // (Core::issue) so a dep-resolved enqueue re-arms only that core's issue scan. + void set_owner_dirty(bool* d) { _owner_dirty_ref = d; } void set_compute_type(int type) { _compute_type = type; } int get_compute_type() { return _compute_type; } void set_numa_id(int numa_id) { _numa_id = numa_id; } @@ -148,6 +152,7 @@ class Instruction : public std::enable_shared_from_this { void *_owner = nullptr; std::list>* _owner_ready_queue_ref = nullptr; + bool* _owner_dirty_ref = nullptr; // owning Core's _issue_dirty (re-arm gate) Opcode opcode; cycle_type compute_cycle = 0; cycle_type overlapping_cycle = 0; diff --git a/TOGSim/src/Core.cc b/TOGSim/src/Core.cc index d3d174c8..6097f287 100644 --- a/TOGSim/src/Core.cc +++ b/TOGSim/src/Core.cc @@ -42,7 +42,7 @@ int Core::pick_free_weight_sa() { void Core::apply_due(const DueAction& a) { switch (a.kind) { case DueAction::FreeWeightSlot: - if (--a.token->refcount <= 0) _weight_slots_used[a.token->sa]--; // last reader frees the slot + if (--a.token->refcount <= 0) { _weight_slots_used[a.token->sa]--; _issue_dirty = true; } // weight slot freed -> re-arm break; case DueAction::WakeBar: { auto bar = a.bar; // async load data arrived -> fire its MEMORY_BAR @@ -68,6 +68,7 @@ void Core::release_sram(const std::shared_ptr& inst) { if (it == _sram_allocs.end()) continue; _sram_used -= it->second; _sram_allocs.erase(it); + _issue_dirty = true; // freed spad bytes -> re-arm } } @@ -100,10 +101,12 @@ void Core::issue(std::shared_ptr op) { M, max_dispatch); } for (const auto& inst : op->get_instructions()) { + inst->set_owner_dirty(&_issue_dirty); // dep-resolved enqueues re-arm THIS core if (inst->is_ready()) op->enqueue_ready(inst); } _tiles.push_back(std::move(op)); + _issue_dirty = true; // new dispatch -> re-arm } std::shared_ptr Core::pop_finished_tile() { @@ -275,6 +278,12 @@ void Core::cycle() { /* Iterate tile while an instruction is issued */ bool issued = false; + // Re-arm gate: skip the scan unless _issue_dirty was set since the last scan (a + // ready-set grow or a resource free; else it re-walks the same blocked instructions + // and issues nothing). Issue-identical. + if (_issue_dirty) { + _issue_dirty = false; + for (int i=0; i<_tiles.size() && !issued; i++) { auto& instructions = _tiles[i]->get_ready_instructions(); for (auto it=instructions.begin(); it!=instructions.end();) { @@ -398,7 +407,9 @@ void Core::cycle() { inst->finish_instruction(); static_cast(inst->get_owner())->inc_finished_inst(); _stat_tot_skipped_inst.at(static_cast(inst->get_opcode()))++; - instructions.erase(it); + it = instructions.erase(it); // erase returns the next iterator; the + continue; // old code fell through to it++ on the + // erased (invalidated) iterator -> UB } else { core_trace_log::trace_instruction_line(_core_cycle, _id, @@ -456,6 +467,12 @@ void Core::cycle() { } } + // Keep dirty after an issue: the scan breaks at the first issue (!issued loop + // guard), so ready instructions in later tiles were not scanned this cycle. Needed + // for that early-break even when the issue woke no new dependent. + if (issued) _issue_dirty = true; + } // if (_issue_dirty) + /* Remove finshed tiles */ bool retry = true; while (retry) { diff --git a/TOGSim/src/CoreTraceLog.cc b/TOGSim/src/CoreTraceLog.cc index 1f6b720c..9a2a5120 100644 --- a/TOGSim/src/CoreTraceLog.cc +++ b/TOGSim/src/CoreTraceLog.cc @@ -46,10 +46,14 @@ std::string format_dma_inst_issued_detail(Instruction& inst) { } std::string format_dma_inst_issued_trace_line(Instruction& inst) { + // Built eagerly at the call site but only fed to spdlog::trace, so skip the format + // work when trace logging is off. + if (!spdlog::should_log(spdlog::level::trace)) return {}; return fmt::format("{} ({})", opcode_to_string(inst.get_opcode()), format_dma_inst_issued_detail(inst)); } std::string format_instruction_detail_line(Instruction& inst) { + if (!spdlog::should_log(spdlog::level::trace)) return {}; // see format_dma_inst_issued_trace_line const Opcode op = inst.get_opcode(); const std::string opname = opcode_to_string(op); if (op == Opcode::COMP) { From bf91b20ede4920c6a39d9a28951310fdee0e9714 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 25 Jun 2026 23:43:10 +0900 Subject: [PATCH 16/32] [Frontend] Make the C++ trace the sole main TOG path; drop legacy ONNX TOG The main compile/sim path no longer generates or selects the legacy ONNX Tile-Operation-Graph. extension_codecache emits only trace.so + trace_cycles.tsv (the build-skip now keys on trace.so), and TOGSimulator.run_standalone always drives TOGSim with --trace_so. The TORCHSIM_LEGACY_TOG opt-in is removed from the frontend. The ONNX --models_list branch is kept solely for the STONNE sparse path (extension_op.py); TOGSim's C++ ONNX parser is untouched (separate PR). origins (which FX nodes a kernel came from) is preserved: logged per kernel run and recorded as a trailing "# origins:" line in trace_cycles.tsv -- the legacy ONNX TOG carried this as node metadata, and the C++ cycle-table loader stops at the comment so the current parser is unaffected. Also drop the dead tog_file param from mlir_gem5_compile_command, migrate scripts/chiplet.sh to --trace_so/--cycle_table (the trace path stubs per-tensor addresses and --attributes_list is no longer a Simulator option), and refresh the CLAUDE.md TOG-generation notes. --- CLAUDE.md | 4 +- PyTorchSimFrontend/extension_codecache.py | 78 ++++++++----------- PyTorchSimFrontend/mlir/passes/cycle_table.py | 12 ++- Simulator/simulator.py | 24 +++--- scripts/chiplet.sh | 13 ++-- 5 files changed, 65 insertions(+), 66 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 46f07e1b..fb76c82d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -21,7 +21,7 @@ The pipeline runs in that order on every `torch.compile` invocation; you'll see | `Simulator/simulator.py` | Python drivers: `FunctionalSimulator` (Spike), `CycleSimulator` (Gem5), `TOGSimulator` (the cycle-accurate one + multi-tenant context manager) | | `Scheduler/scheduler.py` | Poisson arrival generator + scheduling utilities for multi-tenant runs | | `TOGSim/` | C++ TOGSim source. `src/Simulator.cc`, `Core.cc`, `Dram.cc`, `Interconnect.cc`, `L2Cache.cc`, `Tile.cc`, `TileGraph.cc` are the core models. Externals: ramulator2, booksim, stonneCore, onnx, protobuf, spdlog, yaml-cpp | -| `AsmParser/` | `tog_generator.py`, `onnx_utility.py` — TOG generation from ONNX/ASM | +| `AsmParser/` | `tog_generator.py`, `onnx_utility.py` — legacy ONNX TOG generation; now used only by the STONNE sparse path (the main path emits a C++ `trace.so` instead) | | `configs/` | TOGSim hardware configs (YAML). The default is `systolic_ws_128x128_c1_simple_noc_tpuv3.yml`. Naming pattern: `systolic_ws__c__.yml` | | `tests/` | Op- and model-level tests organized under `ops//` (elementwise, reduce, gemm, conv, attention, view, sort, sparsity, misc, fusion), `models//` (Llama, Mixtral8x7B, DeepSeek, Diffusion, MoE, MLP, MobileNet, Yolov5) plus single-file model tests (test_resnet, test_transformer, test_vit, test_mlp, test_single_perceptron), and `system/` (scheduler, eager, hetro, stonne, vectorops). Shared helper: `tests/_utils.py` | | `experiments/artifact/` | Paper reproduction scripts (`cycle_validation/run_cycle.sh`, `speedup/run_speedup.sh`) | @@ -130,7 +130,7 @@ Conan deps for TOGSim: `boost/1.79.0`, `robin-hood-hashing/3.11.5`, `spdlog/1.11 - **Adding a new op (Inductor lowering):** `PyTorchSimFrontend/mlir/mlir_ops.py`, `mlir_lowering.py`, plus a new `mlir__template.py` if it needs its own MLIR template. Decomposition rules: `mlir_decomposition.py`. Scheduling: `mlir_scheduling.py`. Autotune: `mlir_autotune.py`. - **Adding a PyTorch device op:** `PyTorchSimDevice/csrc/aten/native/*` (Minimal/Extra split mirrors `torch_openreg`). - **TOGSim hardware model changes:** `TOGSim/src/{Core,Dram,Interconnect,L2Cache,Tile,TileGraph}.cc` + matching `include/*.h`. -- **TOG generation:** `AsmParser/tog_generator.py` builds the raw graph and serializes it via `AsmParser/onnx_utility.py` to **ONNX, which is the on-disk TOG format** consumed by TOGSim. +- **TOG generation:** the main path compiles each kernel to a C++ **`trace.so`** (`mlir/passes/build_skeleton.py` + `lower_to_emitc.py`) plus a `trace_cycles.tsv` cycle table, which TOGSim turns into a TileGraph via `trace_to_tilegraph`. `AsmParser/tog_generator.py` + `onnx_utility.py` (the legacy ONNX TOG) remain only for the **STONNE sparse path** (`extension_op.py`). - **Eager fallback registration:** `torch.npu.register_eager_to_compile([...])` — see `tests/system/test_eager.py`. - **Per-run results:** `togsim_results/>.log` (stats) and `.trace` (instruction trace). The path is also printed at the end of every run. - **Wrapper codegen path:** printed as `Wrapper Codegen Path = /tmp/torchinductor_//...py` — useful for inspecting generated kernel code and tensor names for `SRAM_BUFFER_PLAN_PATH`. diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 18e0d3a9..308e0d6f 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -7,7 +7,6 @@ from PyTorchSimFrontend import extension_config from torch._inductor.codecache import get_hash, write, write_atomic from torch._inductor.async_compile import AsyncCompile -from AsmParser.tog_generator import tog_generator from PyTorchSimFrontend.mlir.mlir_caller_codegen import MLIRKernelCallerCodeGen from Simulator.simulator import FunctionalSimulator, CycleSimulator, TOGSimulator @@ -81,7 +80,7 @@ def mlir_compile_command(filename, vectorlane_size, vlen=256): """, ).strip()] -def mlir_gem5_compile_command(filename, sample_filename, tog_file, vectorlane_size, vlen=256): +def mlir_gem5_compile_command(filename, sample_filename, vectorlane_size, vlen=256): # See mlir_compile_command: -dma-fine-grained and -test-pytorchsim-to-vcix are # Python passes run in-process; mlir-opt runs only loop-padding here. return [re.sub(r"[ \n]+", " ", @@ -144,17 +143,17 @@ def load(cls, source_code, write_atomic(os.path.join(write_path, "gem5_global_var.h"), gem5_global_var_header) # The compile rewrites the kernel .mlir in place and reads it back, and two # compiles of the same source share a write_path. Hold the per-path lock across - # the build, and skip it when a prior build finished (its tile_graph.onnx exists). + # the build, and skip it when a prior build finished (its trace.so exists). from filelock import FileLock from PyTorchSimFrontend.mlir.passes import ( run_python_passes, run_module_passes, POST_OPT_PASSES, run_standard_lowering, run_tog, ) - tog_path = os.path.join(write_path, "tile_graph.onnx") + trace_so_path = os.path.join(write_path, "trace.so") lock = FileLock(get_lock_path(write_path), timeout=LOCK_TIMEOUT) with lock: key, input_path = write(source_code, "mlir", specified_dir=write_path) - if os.path.isfile(tog_path): + if os.path.isfile(trace_so_path): return key # Run the Python out-of-line MLIR passes (MLIR bindings) on the kernel # .mlir in place, before mlir-opt. Currently lowers torchsim.vlane_idx @@ -164,7 +163,7 @@ def load(cls, source_code, raw_tog_path = new_input_path + "_tog.py" sample_mlir_path = new_input_path + "_sample" validation_binary_path = os.path.join(write_path, validation_binary_name) - gem5_cmds = mlir_gem5_compile_command(new_input_path, sample_mlir_path, raw_tog_path, vectorlane_size) + gem5_cmds = mlir_gem5_compile_command(new_input_path, sample_mlir_path, vectorlane_size) if spad_info is not None: link_option = f"-Wl,--section-start=.spad=0x{spad_info['spad_vaddr']:x}" @@ -252,53 +251,39 @@ def load(cls, source_code, # Run cyclesim cyclesim = CycleSimulator() cycle_list = cyclesim.compile_and_simulate(os.path.join(write_path, cycle_binary_name), vectorlane_size, silent_mode=silent_mode) - # Snapshot for the P3-trace hook below: generate_tile_graph consumes - # cycle_list in place (cycle_list.pop(0) per tile), leaving it empty. cycle_list_for_trace = list(cycle_list) - # Create TOG -- DEPRECATED (timing path): the ONNX-TOG producer, superseded - # by the C++ trace pipeline. The cycle_list / x_offset / w_offset computed - # here are reused by cycle_table, so both paths stay cycle-consistent. + # Per-tile cycle offsets, shared with the trace cycle-table below. w_offset, x_offset = vectorlane_size, vectorlane_size if kwargs['loop_size'] is not None and kwargs['loop_size'][-3] < vectorlane_size: x_offset = kwargs['loop_size'][-3] if kwargs['loop_size'] is not None and kwargs['loop_size'][-1] < vectorlane_size: w_offset = kwargs['loop_size'][-1] w_offset = 0 # max(w_offset - x_offset, 0) - tile_graph_generator = tog_generator(origins) - tile_graph_generator.load_file(raw_tog_path) - tile_graph_generator.generate_tile_graph( - tog_path, - cycle_list=cycle_list, - x_offset=x_offset, # FIXME. - w_offset=w_offset, # FIXME. - vector_lane=vectorlane_size - ) - - # Trace pipeline (DEFAULT): emit the trace producer .so + cycle-table TSV - # from the post-vcix IR and gem5 cycles. The legacy ONNX TOG is DEPRECATED, - # an opt-in fallback via TORCHSIM_LEGACY_TOG=1. Never breaks the compile. - if os.environ.get("TORCHSIM_LEGACY_TOG") != "1": - try: - import mlir.ir as ir - from PyTorchSimFrontend.mlir.passes import ( - build_skeleton as _bs, cycle_table as _ct, lower_to_emitc as _l2e) - pv = sample_mlir_path + "_postvcix.mlir" - _ctx = ir.Context(); _ctx.allow_unregistered_dialects = True - with _ctx: - _mod = ir.Module.parse(open(pv).read(), _ctx) - _bs.build_skeleton(_mod) - _ntiles = len(_ct._compute_types(_mod)) - # align lengths: gem5 gives one numCycles per compute node; - # pad with the last value / truncate if it disagrees. - _cl = list(cycle_list_for_trace) - if _cl and len(_cl) != _ntiles: - _cl = (_cl + [_cl[-1]] * _ntiles)[:_ntiles] - _tbl = _ct.build_cycle_table(_mod, _cl, x_offset, w_offset) - _ct.dump_cycle_table_tsv(_tbl, os.path.join(write_path, "trace_cycles.tsv")) - _l2e.build_trace_so(pv, os.path.join(write_path, "trace.so")) - except Exception as e: - logger.warning(f"[P3-trace] trace .so/sidecar dump skipped: {e}") + + # Trace pipeline (sole sim path): emit the compiled trace producer .so + + # the cycle-table TSV from the post-vcix IR and gem5 cycle_list/offsets; + # TOGSim builds its C++ TOG from this via trace_to_tilegraph. + try: + import mlir.ir as ir + from PyTorchSimFrontend.mlir.passes import ( + build_skeleton as _bs, cycle_table as _ct, lower_to_emitc as _l2e) + pv = sample_mlir_path + "_postvcix.mlir" + _ctx = ir.Context(); _ctx.allow_unregistered_dialects = True + with _ctx: + _mod = ir.Module.parse(open(pv).read(), _ctx) + _bs.build_skeleton(_mod) + _ntiles = len(_ct._compute_types(_mod)) + # align lengths: gem5 gives one numCycles per compute node; + # pad with the last value / truncate if it disagrees. + _cl = list(cycle_list_for_trace) + if _cl and len(_cl) != _ntiles: + _cl = (_cl + [_cl[-1]] * _ntiles)[:_ntiles] + _tbl = _ct.build_cycle_table(_mod, _cl, x_offset, w_offset) + _ct.dump_cycle_table_tsv(_tbl, os.path.join(write_path, "trace_cycles.tsv"), origins=origins) + _l2e.build_trace_so(pv, os.path.join(write_path, "trace.so")) + except Exception as e: + logger.warning(f"[P3-trace] trace .so/sidecar dump skipped: {e}") return key class CustomAsyncCompile(AsyncCompile): @@ -323,6 +308,9 @@ def task(): def run_kernel_simulation(*args, autotune_subprocess_timeout_sec=None, **kwargs): # Wait for compilation key = future.result() + if not autotune and origins: + logger.info("[kernel %s] origins: %s", + hash_prefix(key), ", ".join(sorted(str(o) for o in origins))) from filelock import FileLock result_path = os.path.join(extension_config.get_dump_path(), hash_prefix(key)) lock = FileLock(get_lock_path(result_path), timeout=LOCK_TIMEOUT) diff --git a/PyTorchSimFrontend/mlir/passes/cycle_table.py b/PyTorchSimFrontend/mlir/passes/cycle_table.py index 5cb35801..73485d61 100644 --- a/PyTorchSimFrontend/mlir/passes/cycle_table.py +++ b/PyTorchSimFrontend/mlir/passes/cycle_table.py @@ -92,11 +92,19 @@ def load_cycle_table(path): return json.load(fh) -def dump_cycle_table_tsv(table, path): +def dump_cycle_table_tsv(table, path, origins=None): """Plain `cycleoverlapping` per line, in tile_id order -- the trivial format the C++ `--cycle_table` loader (main.cc, P3 trace pipeline) reads with - ifstream (no JSON dependency in TOGSim).""" + ifstream (no JSON dependency in TOGSim). + + `origins` (the FX nodes this kernel came from) is recorded as a trailing + `# origins: ...` comment after the data rows -- the legacy ONNX TOG carried + this as node metadata. The C++ loader's `while (ct >> c >> o)` stops at the + `#` once all (cycle, overlapping) rows are read, so the comment is safe with + the current parser; a future TOGSim change can promote it to a real field.""" with open(path, "w") as fh: for cycle, overlapping in table: fh.write("%d\t%d\n" % (int(cycle), int(overlapping))) + if origins: + fh.write("# origins: %s\n" % ", ".join(sorted(str(o) for o in origins))) return path diff --git a/Simulator/simulator.py b/Simulator/simulator.py index d0a0df53..9f1df4a2 100644 --- a/Simulator/simulator.py +++ b/Simulator/simulator.py @@ -319,7 +319,9 @@ def _send_command(self, command_type, device_index, stream_index, tog_path="", a command_type: Type of command ("LAUNCH_KERNEL" or "DEVICE_SYNC") device_index: Device index stream_index: Stream index - tog_path: Path to TOG file (ONNX model) - empty for DEVICE_SYNC + tog_path: kernel-dir handle; TOGSim derives trace.so/trace_cycles.tsv from + its directory (the ONNX file itself is only read on the STONNE sparse + path) - empty for DEVICE_SYNC attribute_path: Path to attribute file - empty for DEVICE_SYNC timestamp: Timestamp in nanoseconds (default: 0) @@ -410,7 +412,8 @@ def launch_kernel(self, device_index, stream_index, tog_path, attribute_path, ti Args: device_index: Device index stream_index: Stream index - tog_path: Path to TOG file (ONNX model) + tog_path: kernel-dir handle; TOGSim derives trace.so from its directory + (the ONNX file itself is only read on the STONNE sparse path) attribute_path: Path to attribute file timestamp: Timestamp in nanoseconds (default: 0) @@ -523,7 +526,8 @@ def run_standalone( For streaming multiple kernels, use launch_kernel() instead. Args: - model_path: Path to TOG file (ONNX model) + model_path: kernel-dir handle; trace.so/trace_cycles.tsv are derived from + its directory (the ONNX file itself is only read on the STONNE sparse path) attribute_path: Path to attribute file autotune_mode: If True, run in autotune mode (silent) config_path: Path to TOGSim config file (required) @@ -560,19 +564,15 @@ def run_standalone( os.fsync(trace_file.fileno()) try: - # The C++ TOG (trace) path is the DEFAULT: drive the simulation from the - # emitted trace.so; TORCHSIM_LEGACY_TOG=1 opts into the legacy ONNX TOG. Each - # autotune candidate has its own trace.so. Fall back only if none was emitted. + # Drive the simulation from the emitted trace.so (the C++ TOG path). The + # ONNX --models_list path remains only for callers without a trace.so (the + # STONNE sparse path); the normal compile always emits one. trace_so = os.path.join(os.path.dirname(str(model_path)), "trace.so") cycle_tsv = os.path.join(os.path.dirname(str(model_path)), "trace_cycles.tsv") base_cmd = TOGSimulator.get_togsim_command(config_path, togsim_path) - use_trace = (os.environ.get("TORCHSIM_LEGACY_TOG") != "1" - and os.path.exists(trace_so)) - if os.environ.get("TORCHSIM_LEGACY_TOG") == "1": - logger.warning("TORCHSIM_LEGACY_TOG=1 selects the DEPRECATED legacy ONNX TOG path") - if use_trace: + if os.path.exists(trace_so): cmd = f"{base_cmd} --trace_so {trace_so} --cycle_table {cycle_tsv}" - else: # DEPRECATED: legacy ONNX TOG path + else: # ONNX TOG path (STONNE sparse) cmd = f"{base_cmd} --models_list {trace_file_path}" if extension_config.CONFIG_TOGSIM_DEBUG_LEVEL: cmd += f" --log_level {extension_config.CONFIG_TOGSIM_DEBUG_LEVEL}" diff --git a/scripts/chiplet.sh b/scripts/chiplet.sh index e622874b..40aa77c4 100755 --- a/scripts/chiplet.sh +++ b/scripts/chiplet.sh @@ -35,7 +35,12 @@ for ATTRIBUTE in "$@"; do fi ATTRIBUTE_FILES+=("$ATTRIBUTE_FILE") done -MODELS_LIST="$GEMM_PATH/tile_graph.onnx" +# Trace (C++ TOG) path. NOTE: TOGSim currently stubs per-tensor addresses for the +# trace path (build_trace_tilegraph), so chiplet NoC/DRAM-partition accuracy is +# approximate until the trace path consumes real addresses; --attributes_list is +# no longer a Simulator option. +TRACE_SO="$GEMM_PATH/trace.so" +CYCLE_TABLE="$GEMM_PATH/trace_cycles.tsv" ATTRIBUTE_PATH="$GEMM_PATH/runtime_0000/attribute" for CONFIG in "${CONFIG_LIST[@]}"; do @@ -49,8 +54,7 @@ for CONFIG in "${CONFIG_LIST[@]}"; do OUTPUT_FILE="$RESULTS_DIR/${CONFIG_NAME}_result.txt" # Run Simulator - echo "$SIMULATOR_PATH" --config "$CONFIG" --models_list "$MODELS_LIST" --attributes_list "$ATTRIBUTE_PATH/$ATTRIBUTE_NAME" - "$SIMULATOR_PATH" --config "$CONFIG" --models_list "$MODELS_LIST" --log_level trace --attributes_list "$ATTRIBUTE_PATH/$ATTRIBUTE_NAME" > "$OUTPUT_FILE" & + echo "$SIMULATOR_PATH" --config "$CONFIG" --trace_so "$TRACE_SO" --cycle_table "$CYCLE_TABLE" "$SIMULATOR_PATH" --config "$CONFIG" --trace_so "$TRACE_SO" --cycle_table "$CYCLE_TABLE" --log_level trace --attributes_list "$ATTRIBUTE_PATH/$ATTRIBUTE_NAME" > "$OUTPUT_FILE" & echo "[TOGSim] for $CONFIG stored to \"$(pwd)/$OUTPUT_FILE\"" done done @@ -63,8 +67,7 @@ for CONFIG in "${CONFIG_LIST2[@]}"; do OUTPUT_FILE="$RESULTS_DIR/${CONFIG_NAME}_result.txt" # Run Simulator - # echo "$SIMULATOR_PATH" --config "$CONFIG" --models_list "$MODELS_LIST" --attributes_list "$ATTRIBUTE_PATH/$ATTRIBUTE_NAME" - "$SIMULATOR_PATH" --config "$CONFIG" --models_list "$MODELS_LIST" --log_level trace --attributes_list "$ATTRIBUTE_PATH/$ATTRIBUTE_NAME" > "$OUTPUT_FILE" & + # echo "$SIMULATOR_PATH" --config "$CONFIG" --trace_so "$TRACE_SO" --cycle_table "$CYCLE_TABLE" "$SIMULATOR_PATH" --config "$CONFIG" --trace_so "$TRACE_SO" --cycle_table "$CYCLE_TABLE" --log_level trace --attributes_list "$ATTRIBUTE_PATH/$ATTRIBUTE_NAME" > "$OUTPUT_FILE" & echo "[TOGSim] for $CONFIG stored to \"$(pwd)/$OUTPUT_FILE\"" done wait \ No newline at end of file From 88177e0e452141f2d1ad13a0fa50fd8f0b9108fb Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Fri, 26 Jun 2026 20:36:48 +0900 Subject: [PATCH 17/32] [Frontend] Carry gather/scatter offsets as an explicit transfer descriptor Represent a kernel body as ordered steps so a multi-step indirect access can be emitted, and hand the gather/scatter offset to the DMA as an explicit operand instead of inferring it from the operand count. --- .../mlir/mlir_codegen_backend.py | 244 +++++++++++------- .../mlir/passes/build_skeleton.py | 5 + .../mlir/passes/decompose_transfer.py | 8 +- .../mlir/passes/lower_dma_to_gemmini.py | 38 ++- tests/ops/misc/test_indirect_access.py | 27 ++ 5 files changed, 199 insertions(+), 123 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 40415f01..53235da1 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -310,6 +310,22 @@ def memory_plan(self): "MVOUT1": 3, } + +class Step: + """One load->compute->store unit of the kernel body (see codegen_loops). + + Bundles the DMA, mask, index and compute buffers so the body can be an + ordered list of steps; the formerly ad-hoc mask/index buffers are just + fields here. + """ + __slots__ = ("applys", "dma_loads", + "loads", "compute", "stores", "dma_stores") + + def __init__(self, **buffers): + for name, buf in buffers.items(): + setattr(self, name, buf) + + class MLIRKernel(mlir_common.BaseMLIRKernel): overrides = ExtensionOverrides newvar_prefix = "%" @@ -321,11 +337,15 @@ def __init__(self, kernel_group, reason=None): self.spad_buffer = IndentedBuffer() self.reduction_prefix = IndentedBuffer() self.reduction_suffix = IndentedBuffer() - self.applys = IndentedBuffer() - self.masks = IndentedBuffer() - self.dma_loads = IndentedBuffer() - self.dma_stores = IndentedBuffer() - self.indexed_buffer = IndentedBuffer() + # Kernel body = ordered load->compute->store steps; step 0 keeps the base + # loads/compute/stores (the CSE target default captured self.compute at init). + step0 = Step( + applys=IndentedBuffer(), + dma_loads=IndentedBuffer(), dma_stores=IndentedBuffer(), + loads=self.loads, compute=self.compute, stores=self.stores, + ) + self.steps = [step0] + self._bind_step(step0) self.global_vars = IndentedBuffer() self.header = IndentedBuffer() self.gem5_header = IndentedBuffer() @@ -364,6 +384,7 @@ def __init__(self, kernel_group, reason=None): self.welford_reduce_out = None self.reduce_iterator = {} self.spad_buffer_dict = dict() + self.indirect_symbols = set() # CSE-var names bound as indirect indices self.base_vector_initialized = False self.loop_size = None @@ -546,19 +567,9 @@ def parse_index_list(self, expr_list:list, offset=sympy.Number(0)) -> common.CSE return index def load(self, name: str, index: sympy.Expr): - index, comptute_depedency = self.convert_indirect_indexing(index) + index, offset_desc = self.convert_indirect_indexing(index) padding = self.get_padding_type() - # In case of special form of indirect access, we need to put load in dma_store buffer - if comptute_depedency: - apply_buffer = self.dma_stores - dma_buffer = self.dma_stores - load_buffer = self.dma_stores - else: - apply_buffer = None - dma_buffer = self.dma_loads - load_buffer = self.loads - # Extract dram info dram_var = self.kernel_group.args.input(name) dram_shape = mlir_common.MLIRKernelArgs.get_mlir_shape(self.buffer_types[name]) @@ -566,7 +577,7 @@ def load(self, name: str, index: sympy.Expr): mlir_dtype = mlir_common.DTYPE_TO_MLIR[dtype] # Extract sram info - local_tile_desc, index_var, dram_stride = self.get_dma_info(name, index, buffer=apply_buffer) + local_tile_desc, index_var, dram_stride = self.get_dma_info(name, index) vlane_split_axis = local_tile_desc.vmap.vlane_split_axis vlane_stride = local_tile_desc.vmap.vlane_stride tile_numel_per_lane = local_tile_desc.get_numel_per_lane() @@ -581,26 +592,21 @@ def load(self, name: str, index: sympy.Expr): compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) code = self.emit_transfer("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, int(padding)) - self.cse.generate(dma_buffer, code, assignment = False) # FIXME: assignment = False does not support caching + dram_shape, tile_shape, dram_stride, tile_stride, int(padding), offset=offset_desc) + self.cse.generate(self.dma_loads, code, assignment = False) # FIXME: assignment = False does not support caching - if not comptute_depedency: - # Generate vector load instruction - with self.override_buffer_cse(buffer=load_buffer): - out = ops._load(compute_vec_size, mlir_dtype, sram_var, compute_index_var, tile_shape) - else: - # FIXME. Any good idea? - out = sram_var - self.register_var_info(out, [compute_vec_size, mlir_dtype]) + with self.override_buffer_cse(buffer=self.loads): + out = ops._load(compute_vec_size, mlir_dtype, sram_var, compute_index_var, tile_shape) self.spad_buffer_dict[str(out)] = [sram_var, local_tile_desc.get_tile_size(), tile_numel_per_lane, sram_index_var, tile_shape, vshape] return out def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs): dtype = V.graph.get_dtype(name) mlir_dtype = mlir_common.DTYPE_TO_MLIR[dtype] + offset_desc = None # Handle scatter store - if "tmp" in str(index): + if self._has_indirect(index): # Convert the output buffer type to the inplace buffer arg_name = V.graph.scheduler.mutation_real_name.get(name, name) if arg_name not in self.kernel_group.args.inplace_buffers: @@ -609,7 +615,7 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) if mode == "atomic_add": loaded_value = ops.load(name, index) value = ops.add(loaded_value, value) - index, _ = self.convert_indirect_indexing(index) + index, offset_desc = self.convert_indirect_indexing(index) dram_var = self.kernel_group.args.output(name) # Prepare dma instruction @@ -651,7 +657,7 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) # Generate DMA instruction code = self.emit_transfer("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, 0) + dram_shape, tile_shape, dram_stride, tile_stride, 0, offset=offset_desc) self.dma_stores.writeline(common.DeferredLine(name, code)) def reduction(self, dtype, src_dtype, reduction_type, value): @@ -783,8 +789,12 @@ def store_reduction(self, name, index, value): self.reductions_suffix.writeline(common.DeferredLine(name, code)) def indirect_indexing(self, index_var, size, check=True, wrap_neg=True): + self.indirect_symbols.add(str(index_var)) # record the bound indirect symbol return str(index_var) + def _has_indirect(self, expr): + return any(s.name in self.indirect_symbols for s in expr.free_symbols) + def _index_expr(self, tile_desc, renamed_expression, index, base_vector_index): # In case of index expr, dimension size should be divisible by tile size if not self.kernel_group.tile_desc.is_dim_dividable(self.ranges): @@ -933,6 +943,30 @@ def index_expr(self, index, dtype): def codegen_global_init(self): return self.global_vars + def _bind_step(self, step): + # Make `step` the current emit sink: route the body buffers to its buffers + self.current_step = step + self.applys = step.applys + self.dma_loads = step.dma_loads + self.dma_stores = step.dma_stores + self.loads = step.loads + self.compute = step.compute + self.stores = step.stores + + def push_step(self): + # New load->compute->store step; later emits land here, steps bridge via spad + step = Step( + applys=IndentedBuffer(), + dma_loads=IndentedBuffer(), dma_stores=IndentedBuffer(), + loads=IndentedBuffer(), compute=IndentedBuffer(), stores=IndentedBuffer(), + ) + self.steps.append(step) + self._bind_step(step) + self.cse = self.cse.clone() # share name counter, fresh dedup cache (region-safe) + self.target_buffer_override.set(self.compute) + self.target_cse_override.set(self.cse) + return step + def codegen_loops(self): code = mlir_common.ParallelLoopBuffer() # Loop body part @@ -965,18 +999,18 @@ def codegen_loops(self): epilogue = reduction_loop.epilogue_line() code.writelines(reduction_lines) stack.enter_context(code.indent(attribute="{accumulation_loop=true}", suffix=epilogue)) - code.splice(self.applys) - code.splice(self.indexed_buffer) - code.splice(self.dma_loads) - # Compute body - code.writelines(self.compute_body_loop.lines()) - with contextlib.ExitStack() as stack: - stack.enter_context(code.indent(attribute="{inner_loop=false}",suffix=self.compute_body_loop.epilogue_line())) - code.splice(self.masks) - code.splice(self.loads) - code.splice(self.compute) - code.splice(self.stores) - code.splice(self.dma_stores) + for step in self.steps: + code.splice(step.applys) + code.splice(step.dma_loads) + # Compute body -- only steps that have one get the loop + epilogue + if any(b.getvalue() for b in (step.loads, step.compute, step.stores)): + code.writelines(self.compute_body_loop.lines()) + with contextlib.ExitStack() as stack: + stack.enter_context(code.indent(attribute="{inner_loop=false}",suffix=self.compute_body_loop.epilogue_line())) + code.splice(step.loads) + code.splice(step.compute) + code.splice(step.stores) + code.splice(step.dma_stores) code.splice(self.reductions_suffix) # Non-outerloop end code.writeline(f"return") @@ -1197,7 +1231,7 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe """ # Use loads as default if buffer is None: - buffer = self.applys if "tmp" not in str(index) else self.dma_loads + buffer = self.applys if not self._has_indirect(index) else self.dma_loads # TODO. kg_tile_desc = self.kernel_group.tile_desc @@ -1207,7 +1241,7 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe total_dims = [int(str(i)[5:]) for i in self.itervars] local_tile_desc = mlir_common.MLIRMultiDimTile([1], self.vector_lane) local_dims.sort() # Assume that smaller index is placed in the outer loop - indirect_syms = [s for s in index.free_symbols if "tmp" in s.name] + indirect_syms = [s for s in index.free_symbols if s.name in self.indirect_symbols] index = index.subs({s: 0 for s in indirect_syms}, simultaneous=True) indirect_dims = [f"{i}" for i in indirect_syms] @@ -1329,7 +1363,7 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtype, dram_var, dram_index_var, sram_var, sram_index_var, dram_shape, tile_shape, dram_stride, tile_stride, padding, - subtile_size=None, async_type=None): + subtile_size=None, async_type=None, offset=None): """Emit a generic togsim.transfer op for a DMA whose access exceeds the 4D Gemmini descriptor limit. Carries the full N-D access (dram/tile strides + shapes) plus the SSA operands a memref.dma_start needs @@ -1370,12 +1404,16 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp if subtile_size: av = int(async_type) if async_type is not None else 1 attrs += f', subtile_size = {list(subtile_size)}, async = {av} : i64' - # operands: dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_stride - return ( - f'"togsim.transfer"(%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' - f'%{tag}, %{dma_type}, %{vst}) {{{attrs}}} : ' - f'({dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index) -> ()' - ) + # operands: dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_stride [, offset spad] + operands = (f'%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' + f'%{tag}, %{dma_type}, %{vst}') + optypes = f'{dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index' + if offset is not None: # indirect: per-position offset spad (decompose lifts it to a symbol attr) + offset_buf, offset_type, offset_stride = offset + operands += f', %{offset_buf}' + optypes += f', {offset_type}' + attrs += f', indirect = true, offset_stride = {int(offset_stride)} : i64' + return f'"togsim.transfer"({operands}) {{{attrs}}} : ({optypes}) -> ()' def allocate_sram_buffer(self, dtype, dram_name, tile_desc, raw_index, buffer=None, forced_name=None): c_type = mlir_common.DTYPE_TO_C[dtype] @@ -1449,14 +1487,14 @@ def get_mask(self): upper_bound = ops.constant(self.compute_body_loop.size, "index") step_vec = ops.step(self.compute_body_loop.step, "index") - with self.override_buffer_cse(buffer=self.masks, cse=self.mask_cse): + with self.override_buffer_cse(buffer=self.compute, cse=self.mask_cse): gap = ops.sub(upper_bound, self.compute_idx) gap_vec = ops.broadcast(gap, self.compute_body_loop.step) mask_var = ops.lt(step_vec, gap_vec) return mask_shape, mask_var def convert_indirect_indexing(self, index :sympy.Expr): - if "tmp" not in str(index): + if not self._has_indirect(index): return index, None # Note: In case of indirect indexing, dimensions should be divisible by tile size @@ -1467,67 +1505,75 @@ def convert_indirect_indexing(self, index :sympy.Expr): raise mlir_common.RecompileSignal(f"Indirect access (tile size {self.kernel_group.tile_desc.get_tile_size()} is not divisible by {self.ranges})") # Process start - indirect_dims = [str(dim) for dim in index.free_symbols if "tmp" in str(dim)] + indirect_dims = [str(dim) for dim in index.free_symbols if str(dim) in self.indirect_symbols] indirect_dims.sort() first_dim = indirect_dims[0] spad_vars = dict() compute_dependecy = any([target_dim not in self.spad_buffer_dict for target_dim in indirect_dims]) - target_dma_buffers = self.dma_stores if compute_dependecy else self.dma_loads - - # Load indirect operands + # Store each newly-produced indirect index into spad, in its producing step for target_dim in indirect_dims: if target_dim in self.spad_buffer_dict: - sram_var, _, tile_numel_per_lane, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[target_dim] - else: - # FIXME. - var_info = [v for k, v in self.var_info.items() if str(k) == target_dim][0] - dtype = mlir_common.MLIR_TO_DTYPE[var_info[1]] - - local_tile_desc = self.kernel_group.tile_desc - tile_numel_per_lane = local_tile_desc.get_numel_per_lane() - tile_shape = local_tile_desc.get_mlir_shape(var_info[1]) - tile_vec = local_tile_desc.get_compute_vec_size() - vshape = f"vector<{var_info[0]}x{var_info[1]}>" - sram_var, sram_index_var = self.get_scratchpad_buffer(dtype, target_dim, local_tile_desc, target_dim) - self.spad_buffer_dict[target_dim] = [sram_var, local_tile_desc.get_tile_size(), tile_numel_per_lane, sram_index_var, tile_shape, vshape] - - # Store the indirect index variable - target_var = self.cse.varname_map[target_dim] - compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) - with self.override_buffer_cse(buffer=self.stores): - ops._store(target_var, sram_var, compute_index_var, tile_shape) - mlir_dtype = vshape.split("x")[1][:-1] - with self.override_buffer_cse(buffer=target_dma_buffers): - out = ops._load(tile_numel_per_lane, mlir_dtype, sram_var, sram_index_var, tile_shape) - spad_vars[target_dim] = out + continue + var_info = [v for k, v in self.var_info.items() if str(k) == target_dim][0] + dtype = mlir_common.MLIR_TO_DTYPE[var_info[1]] + local_tile_desc = self.kernel_group.tile_desc + tile_numel_per_lane = local_tile_desc.get_numel_per_lane() + tile_shape = local_tile_desc.get_mlir_shape(var_info[1]) + tile_vec = local_tile_desc.get_compute_vec_size() + vshape = f"vector<{var_info[0]}x{var_info[1]}>" + sram_var, sram_index_var = self.get_scratchpad_buffer(dtype, target_dim, local_tile_desc, target_dim) + self.spad_buffer_dict[target_dim] = [sram_var, local_tile_desc.get_tile_size(), tile_numel_per_lane, sram_index_var, tile_shape, vshape] + target_var = self.cse.varname_map[target_dim] + compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) + with self.override_buffer_cse(buffer=self.stores): + ops._store(target_var, sram_var, compute_index_var, tile_shape) + + # Offset build runs after the index is in spad -> own step when just produced + if compute_dependecy: + self.push_step() + + # Single indirect dim: the raw index IS the offset; the MVIN applies offset_stride (CONFIG4) + if len(indirect_dims) == 1: + offset_stride = 1 + for arg in list(index.args): + if not self._has_indirect(arg): + continue + if arg.is_Mul and arg.args[0].is_number: + offset_stride = int(arg.args[0]) + index = index.replace(arg, 0) + sram_var, _, _, _, tile_shape, _ = self.spad_buffer_dict[first_dim] + return index, (sram_var, tile_shape, offset_stride) - with self.override_buffer_cse(buffer=target_dma_buffers): - # Apply stride + # Multi indirect dim: sum the strided indices in the compute loop (chunked by compute_vec_size) + local_tile_desc = self.kernel_group.tile_desc + compute_vec_size = local_tile_desc.get_compute_vec_size() + for target_dim in indirect_dims: + sram_var, _, _, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[target_dim] + mlir_dtype = vshape.split("x")[1][:-1] + compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) + with self.override_buffer_cse(buffer=self.loads): + spad_vars[target_dim] = ops._load(compute_vec_size, mlir_dtype, sram_var, compute_index_var, tile_shape) + with self.override_buffer_cse(buffer=self.compute): for arg in index.args: - if "tmp" not in str(arg): + if not self._has_indirect(arg): continue if arg.is_Mul and arg.args[0].is_number: coeff_dtype = self.var_info[spad_vars[str(arg.args[1])]][1] coeff = self.get_const_cse(int(arg.args[0]), coeff_dtype) spad_vars[str(arg.args[1])] = ops.mul(spad_vars[str(arg.args[1])], coeff) index = index.replace(arg, 0) - - # Sum for dim, var in spad_vars.items(): if dim == first_dim: continue spad_vars[first_dim] = ops.add(spad_vars[first_dim], var) - - # Store index var - sram_var, _, tile_numel_per_lane, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[first_dim] - mlir_dtype = vshape.split("x")[1][:-1] - with self.override_buffer_cse(buffer=target_dma_buffers): - ops._store(spad_vars[first_dim], sram_var, sram_index_var, tile_shape) # FIXME. Maybe require fine grain compute... - - # Conversion - mlir_dtype = self.var_info[spad_vars[first_dim]][1] - with self.override_buffer_cse(buffer=target_dma_buffers): - out = ops._load(1, mlir_dtype, sram_var, sram_index_var, tile_shape) - if mlir_dtype != "index": - out = ops.index_cast(out, "index") - return index + sympy.Symbol(str(out)), compute_dependecy + # Summed offset goes to a DEDICATED spad (not an index buffer) to avoid clobbering a live index + var_info = [v for k, v in self.var_info.items() if str(k) == first_dim][0] + dtype = mlir_common.MLIR_TO_DTYPE[var_info[1]] + off_shape = local_tile_desc.get_mlir_shape(var_info[1]) + off_sram, off_index = self.get_scratchpad_buffer( + dtype, "indirect_offset_" + first_dim, local_tile_desc, "indirect_offset_" + first_dim) + off_compute_index = ",".join(off_index.split(",")[:-1] + [f"%{self.compute_idx}"]) + with self.override_buffer_cse(buffer=self.stores): + ops._store(spad_vars[first_dim], off_sram, off_compute_index, off_shape) + self.push_step() # offset-build compute loop must finish before the gather reads it + return index, (off_sram, off_shape, 1) diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index 523b3015..8d4cfd1b 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -439,6 +439,11 @@ def _emit_one_dma(ctx, op, node, builder, bufs, tags): spad_id = bufs.of(dep._global_of(f["src"] if node.is_write else f["dst"])) read_bufs = [spad_id] if node.is_write else [] write_bufs = [] if node.is_write else [spad_id] + if "indirect_offset" in op.attributes: # gather/scatter reads the offset spad -> dep on its build + from mlir.ir import FlatSymbolRefAttr + off_id = bufs.of(FlatSymbolRefAttr(op.attributes["indirect_offset"]).value) + if off_id not in read_bufs: + read_bufs = read_bufs + [off_id] tag_id = tags.bind(_value_key(f["tag"]), spad_id) _emit_dma(ctx, node, tag_id, dram_index, tag_index, read_bufs, write_bufs) diff --git a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py index 10b2edfb..0bf04e30 100644 --- a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py +++ b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py @@ -91,7 +91,10 @@ def run(module, vectorlane=128, **_): targets.append(op.operation) for op in targets: - dram, dram_idx, sram, sram_idx, tag, dma_type, vst = op.operands + op_operands = list(op.operands) + dram, dram_idx, sram, sram_idx, tag, dma_type, vst = op_operands[:7] + # indirect: offset spad operand -> lift to a symbol attr (memref.dma_start can't take the operand) + offset_sym = op_operands[7].owner.attributes["name"] if len(op_operands) > 7 else None kind = op.attributes["dma_kind"].value # StringAttr -> "MVIN"/"MVOUT" vlane_axis = IntegerAttr(op.attributes["vlane_split_axis"]).value dram_stride = _int_array(op.attributes["dram_stride"]) @@ -127,6 +130,9 @@ def _emit(sram_mem, sram_indices, dram_idx_val, vsa_val, dr_attr, tl_attr, st_at if st_attr is not None: attrs["subtile_size"] = st_attr attrs["async"] = async_attr + if offset_sym is not None: + attrs["indirect_offset"] = offset_sym + attrs["offset_stride"] = op.attributes["offset_stride"] Operation.create( "memref.dma_start", results=[], operands=operands, attributes=attrs) diff --git a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py index 998a6db5..5ca842c1 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py @@ -58,12 +58,18 @@ def run(module, timing=False): memref.dma_wait is erased in both modes (matches C++ DmaWaitOpLowering). """ from mlir.ir import (InsertionPoint, Operation, IntegerType, IndexType, - IntegerAttr, MemRefType) + IntegerAttr, MemRefType, FlatSymbolRefAttr, TypeAttr) from mlir.dialects import llvm, arith, memref i64 = IntegerType.get_signless(64) idx = IndexType.get() + # memref.global symbol -> type, to resolve the indirect_offset spad + sym2type = {} + for g in module.operation.regions[0].blocks[0].operations: + if g.operation.name == "memref.global": + sym2type[g.attributes["sym_name"].value] = MemRefType(TypeAttr(g.attributes["type"]).value) + def const_int(val): return IntegerAttr(val.owner.attributes["value"]).value @@ -119,9 +125,8 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): is_mvin = dma_type in (MVIN, MVIN2, MVIN3) elem_bytes = _elem_bytes(src_ty.element_type) - # Indirect (gather): the gather-side indices are src for mvin, dst for mvout. - gather_idx = src_idx if is_mvin else dst_idx - indirect, indirect_memref = _find_indirect(gather_idx) + # Indirect (gather): offset spad referenced by the indirect_offset symbol attr + indirect = "indirect_offset" in op.attributes tile_shape = _subtile(op) if tile_shape is None: @@ -155,10 +160,14 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): i64_const((spad4[2] << 32) | (spad4[3] & 0xFFFFFFFF))) if indirect: # CONFIG4: rs1 = indirect index-spad base address, rs2 = (elem_size<<16)|stride(1) + offset_sym = FlatSymbolRefAttr(op.attributes["indirect_offset"]).value + off_ty = sym2type[offset_sym] + indirect_memref = memref.GetGlobalOp(off_ty, offset_sym).result ind_base = memref.ExtractAlignedPointerAsIndexOp(indirect_memref).result ind_addr = arith.IndexCastOp(i64, ind_base).result - ind_esize = _elem_bytes(MemRefType(indirect_memref.type).element_type) - asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (1 & 0xFFFF))) + ind_esize = _elem_bytes(off_ty.element_type) + off_stride = IntegerAttr(op.attributes["offset_stride"]).value + asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (off_stride & 0xFFFF))) asm(dma_type, dram_addr, spad_addr) op.erase() @@ -189,23 +198,6 @@ def _elem_bytes(elem_type): return max(bits, 8) // 8 -def _find_indirect(indices): - """If a gather index is an affine.apply{indirect_access} whose operands include - index_cast(affine.load(%spad)), return (True, %spad memref); else (False, None).""" - for idx in indices: - ap = idx.owner - if getattr(ap, "name", None) != "affine.apply" or "indirect_access" not in ap.attributes: - continue - for operand in ap.operands: - ic = operand.owner - if getattr(ic, "name", None) != "arith.index_cast": - continue - ld = ic.operands[0].owner - if getattr(ld, "name", None) == "affine.load": - return True, ld.operands[0] # affine.load operand 0 == the index spad memref - return False, None - - def lower_text(text): if OP_NAME not in text: return text diff --git a/tests/ops/misc/test_indirect_access.py b/tests/ops/misc/test_indirect_access.py index f64fe50d..ae3a4ed4 100644 --- a/tests/ops/misc/test_indirect_access.py +++ b/tests/ops/misc/test_indirect_access.py @@ -52,6 +52,31 @@ def scatter_only(out, token_indices, weighted_output): res = opt_fn(out, token_indices, weighted_output) test_result("ScatterAdd(index_add_)", res, cpu_out) +def test_multidim_indirect(device, size=(64, 64), n=256): + torch.manual_seed(0) + def gather2d(x, ix, iy): + return x[ix, iy] + 1.0 + x = torch.randn(size, dtype=torch.float32).to(device=device) + ix = torch.randint(0, size[0], [n]).to(device=device) + iy = torch.randint(0, size[1], [n]).to(device=device) + opt_fn = torch.compile(dynamic=False)(gather2d) + res = opt_fn(x, ix, iy) + out = gather2d(x.cpu(), ix.cpu(), iy.cpu()) + test_result("Multi-dim Indirect (x[ix,iy])", res, out) + +def test_multidim_indirect_index_reuse(device, size=(64, 64), n=256): + torch.manual_seed(0) + def gather_reuse(x, ix, iy): + # ix is reused after the gather -> the offset spad must not clobber the index spad + return x[ix, iy] + ix.float() + x = torch.randn(size, dtype=torch.float32).to(device=device) + ix = torch.randint(0, size[0], [n]).to(device=device) + iy = torch.randint(0, size[1], [n]).to(device=device) + opt_fn = torch.compile(dynamic=False)(gather_reuse) + res = opt_fn(x, ix, iy) + out = gather_reuse(x.cpu(), ix.cpu(), iy.cpu()) + test_result("Multi-dim Indirect index reuse (x[ix,iy]+ix)", res, out) + def test_scatter_full(device, size=(128, 128)): def vectoradd(a, idx, b): a[idx, :] = b @@ -71,4 +96,6 @@ def vectoradd(a, idx, b): test_scatter_full(device, size=(2048, 2048)) test_scatter_add(device) test_indirect_vectoradd(device) + test_multidim_indirect(device) + test_multidim_indirect_index_reuse(device) #test_embedding(device, 1024, 2048) \ No newline at end of file From 8a2d2570a540d22dc6c1001e671c1cc5d322682d Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 15:35:41 +0900 Subject: [PATCH 18/32] [Frontend] Drop memref.dma_start: lower togsim.transfer directly to Gemmini togsim.transfer is unregistered, so it can carry every runtime descriptor as an operand -- including the masked-clamp low/high vectors -- which a registered memref.dma_start cannot. The trace path (build_skeleton) reads togsim.transfer/togsim.wait, and the indirect offset operand is detected by its `indirect` attribute rather than by counting operands. --- PyTorchSimFrontend/extension_codecache.py | 4 +- .../mlir/mlir_codegen_backend.py | 6 +- PyTorchSimFrontend/mlir/passes/__init__.py | 3 +- .../mlir/passes/build_skeleton.py | 74 +++-- PyTorchSimFrontend/mlir/passes/build_tog.py | 93 ++++--- .../mlir/passes/dma_fine_grained.py | 96 ++++--- .../mlir/passes/lower_to_llvm.py | 8 +- .../mlir/passes/lower_to_vcix.py | 45 +-- .../mlir/passes/lower_transfer_to_gemmini.py | 257 ++++++++++++++++++ 9 files changed, 456 insertions(+), 130 deletions(-) create mode 100644 PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 308e0d6f..9618b8e8 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -51,7 +51,7 @@ def mlir_compile_command(filename, vectorlane_size, vlen=256): return [re.sub(r"[ \n]+", " ", f""" {extension_config.CONFIG_TORCHSIM_LLVM_PATH}/mlir-opt \ - -test-loop-padding \ + -test-loop-padding --allow-unregistered-dialect \ {'--mlir-print-ir-after-all' if extension_config.CONFIG_TORCHSIM_DUMP_MLIR_IR else ''} \ {filename}.mlir -o {filename}_padded.mlir """, @@ -86,7 +86,7 @@ def mlir_gem5_compile_command(filename, sample_filename, vectorlane_size, vlen=2 return [re.sub(r"[ \n]+", " ", f""" {extension_config.CONFIG_TORCHSIM_LLVM_PATH}/mlir-opt \ - -test-loop-padding='timing_mode=1' \ + -test-loop-padding='timing_mode=1' --allow-unregistered-dialect \ {'--mlir-print-ir-after-all' if extension_config.CONFIG_TORCHSIM_DUMP_MLIR_IR else ''} \ {filename}.mlir -o {sample_filename}_padded.mlir """, diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 53235da1..ee7d9f7a 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -1404,10 +1404,10 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp if subtile_size: av = int(async_type) if async_type is not None else 1 attrs += f', subtile_size = {list(subtile_size)}, async = {av} : i64' - # operands: dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_stride [, offset spad] + # operands: dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vlane_stride [, offset spad] operands = (f'%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' - f'%{tag}, %{dma_type}, %{vst}') - optypes = f'{dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index' + f'%{tag}, %{zero_cse}, %{dma_type}, %{vst}') + optypes = f'{dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index, index' if offset is not None: # indirect: per-position offset spad (decompose lifts it to a symbol attr) offset_buf, offset_type, offset_stride = offset operands += f', %{offset_buf}' diff --git a/PyTorchSimFrontend/mlir/passes/__init__.py b/PyTorchSimFrontend/mlir/passes/__init__.py index d96035bb..98033fe8 100644 --- a/PyTorchSimFrontend/mlir/passes/__init__.py +++ b/PyTorchSimFrontend/mlir/passes/__init__.py @@ -41,9 +41,8 @@ def _ensure_mlir_bindings_on_path(): # Module rewrite passes around the one remaining mlir-opt pass (-test-loop-padding). # Each exposes MARKERS + run(module, **opts); run_module_passes parses once per phase. -# decompose_transfer first: togsim.transfer -> memref.dma_start (downstream expects it). +# togsim.transfer survives to lower_transfer_to_gemmini; loop-padding runs opaquely. PRE_OPT_PASSES = [ - decompose_transfer, lower_vlane_idx, ] # fine-grained first: splits the matmul DMAs that the vcix lowering then reads. diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index 8d4cfd1b..f4ed7d0d 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -6,9 +6,9 @@ skeleton with the data computation replaced by calls to the event-based runtime API. This pass performs that reduction at the MLIR level: - * `memref.dma_start` -> `togsim.dma(...) {tag_id, is_async, ...}` carrying the + * `togsim.transfer` -> `togsim.dma(...) {tag_id, is_async, ...}` carrying the runtime tag index operand (`%tag[%idx]`). - * `memref.dma_wait` -> `togsim.memory_barrier(tag_idx) {tag_id, write_bufs}`, + * `togsim.wait` -> `togsim.memory_barrier(tag_idx) {tag_id, write_bufs}`, the explicit async-DMA sync. It pairs with its dma by the RUNTIME tag slot (tag_id + the tag index), not a compile-time id: one static dma op runs once per loop @@ -47,7 +47,7 @@ ) #: Marker op names for the passes/__init__ fast-path (skip parsing if absent). -MARKERS = ("memref.dma_start", "memref.dma_wait") +MARKERS = ("togsim.transfer", "togsim.wait") #: Ops the DCE must never remove (loops, terminators, our API ops). _KEEP = { @@ -75,16 +75,16 @@ def _arg_id_of(base_addr): def _emit_dma(ctx, dma_node, tag_id, dram_index, tag_index, read_bufs, write_bufs): - """Insert a `togsim.dma` before the original `memref.dma_start`. + """Insert a `togsim.dma` before the original `togsim.transfer`. `tag_id` is the identity of this DMA's tag memref. An async DMA pairs with - its `togsim.memory_barrier` (the original dma_wait) by the RUNTIME tag slot + its `togsim.memory_barrier` (the original togsim.wait) by the RUNTIME tag slot -- (tag_id, tag_index) -- not a compile-time identifier: one static dma op runs once per loop iteration, each with a different runtime `%tag[%idx]` slot, so only a runtime key can pair iteration i's dma with iteration i's wait. `dram_index` is the original linear DRAM index Value (the `affine.apply` - result that indexed the tensor in the `memref.dma_start`) -- carried as an + result that indexed the tensor in the `togsim.transfer`) -- carried as an operand so the DCE keeps the address arithmetic live and the C4 lowering can compute the real `base_addr = base[arg_id] + index*elem` (P3, approach A). @@ -119,7 +119,7 @@ def _emit_dma(ctx, dma_node, tag_id, dram_index, tag_index, read_bufs, write_buf def _emit_memory_bar(ctx, anchor_op, tag_id, tag_index, write_bufs): """Insert a `togsim.memory_barrier` before `anchor_op` -- the explicit - async-DMA sync that was the original `memref.dma_wait`. It pairs with its + async-DMA sync that was the original `togsim.wait`. It pairs with its async `togsim.dma` by the RUNTIME tag slot (tag_id + tag_index), and carries the SRAM buffer that dma loaded so consumers gate on data-arrival, not on the async dma's issue-complete.""" @@ -146,7 +146,7 @@ def _flatten_add(expr): def _neg_coeff_dim(summand): """If `summand` is `dim * c` with a negative constant `c`, return that dim's position; else None. lower_to_vcix tags each accumulation (reduction) loop var - with coefficient -1 in the dma_wait tag index -- a SENTINEL marking the + with coefficient -1 in the togsim.wait tag index -- a SENTINEL marking the reduction axis, not an arithmetic offset (legacy TileGraphParser skips stride -1 for the same reason).""" if not ir.AffineMulExpr.isinstance(summand): @@ -387,7 +387,7 @@ def of(self, name): class _TagIds: """Identity of a DMA's tag memref -> stable small int, plus the SRAM buffer that tag's async DMA loads. An async dma and its memory_barrier (the original - dma_wait) share a tag memref; this assigns it a tag_id (so the runtime can + togsim.wait) share a tag memref; this assigns it a tag_id (so the runtime can pair them by the runtime tag slot) and remembers the loaded buffer so the barrier can release it to consumers. Pairing is by tag, never a static id.""" @@ -423,25 +423,51 @@ def _emit_computes(ctx, builder, bufs): return n +def _transfer_fields(op): + """Decode a `togsim.transfer`'s fixed operands by position. + + Layout (see mlir_codegen_backend.emit_transfer / lower_transfer_to_gemmini): + operands: dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst + [, offset_spad] # 8 or 9 operands + Unlike the old `memref.dma_start`, dram/sram are FIXED (not direction-swapped): + the DRAM side is always operand[0]/[1], the SRAM spad always operand[2], the + runtime tag slot always operand[4] (tag memref) + operand[5] (tag_idx). The + optional indirect-offset spad is operand[8]; its owning `memref.get_global` + carries the offset symbol name in its "name" attribute (matching + lower_transfer_to_gemmini's offset_sym derivation).""" + operands = list(op.operands) + offset = operands[8] if "indirect" in op.attributes else None + return { + "dram": operands[0], "dram_idx": operands[1], + "sram": operands[2], "sram_idx": operands[3], + "tag": operands[4], "tag_idx": operands[5], + "dma_type": operands[6], "vst": operands[7], + "offset": offset, + } + + def _emit_one_dma(ctx, op, node, builder, bufs, tags): - """Rewrite one memref.dma_start as togsim.dma. A load reads DRAM and writes + """Rewrite one togsim.transfer as togsim.dma. A load reads DRAM and writes its SRAM spad; a store reads the spad and writes DRAM -- which sets the read/write buffer that drives the dependency edge (sec 10). The tag memref is bound to a tag_id (with its loaded buffer) so the paired memory_barrier finds it by the runtime tag slot.""" from . import dep_analysis as dep # lazy: dep_analysis imports build_skeleton - f = builder._dma_start_fields(op) - dram_indices = f["dst_indices"] if node.is_write else f["src_indices"] - dram_index = dram_indices[0] if dram_indices else None - tag_indices = f["tag_indices"] - tag_index = tag_indices[0] if tag_indices else None - # the spad is the SRAM side of the copy: dst for a load, src for a store. - spad_id = bufs.of(dep._global_of(f["src"] if node.is_write else f["dst"])) + f = _transfer_fields(op) + # dram/sram are fixed operands now (not direction-swapped): the DRAM index is + # always operand[1], the SRAM spad always operand[2]. Direction (read/write) + # comes from the node (dma_kind attr / dma_type value). + dram_index = f["dram_idx"] + tag_index = f["tag_idx"] # single runtime tag slot operand (%tag[%idx]) + spad_id = bufs.of(dep._global_of(f["sram"])) read_bufs = [spad_id] if node.is_write else [] write_bufs = [] if node.is_write else [spad_id] - if "indirect_offset" in op.attributes: # gather/scatter reads the offset spad -> dep on its build - from mlir.ir import FlatSymbolRefAttr - off_id = bufs.of(FlatSymbolRefAttr(op.attributes["indirect_offset"]).value) + if f["offset"] is not None: # gather/scatter reads the offset spad -> dep on its build + # the offset symbol name lives on the offset operand's owning get_global + # ("name" attr), the same place lower_transfer_to_gemmini reads it. + off_owner = f["offset"].owner + off_sym = str(off_owner.attributes["name"]).strip('@" ') + off_id = bufs.of(off_sym) if off_id not in read_bufs: read_bufs = read_bufs + [off_id] tag_id = tags.bind(_value_key(f["tag"]), spad_id) @@ -449,7 +475,7 @@ def _emit_one_dma(ctx, op, node, builder, bufs, tags): def _emit_one_wait(ctx, op, tags): - """Rewrite one memref.dma_wait as togsim.memory_barrier -- the explicit + """Rewrite one togsim.wait as togsim.memory_barrier -- the explicit async-DMA sync already in the IR. Paired with its dma by the tag memref (tag_id) and the runtime tag index; carries the buffer the dma loaded. Returns True iff emitted (a wait whose tag no dma used is dropped).""" @@ -468,7 +494,7 @@ def _emit_one_wait(ctx, op, tags): def _emit_dmas_and_waits(ctx, block, builder, dma_by_op, bufs): - """Step 2: rewrite memref.dma_start -> togsim.dma and memref.dma_wait -> + """Step 2: rewrite togsim.transfer -> togsim.dma and togsim.wait -> togsim.memory_barrier in program order. An async dma and its barrier are paired by the RUNTIME tag slot (tag_id + tag index), not a compile-time id: one static dma op runs per loop iteration with a different `%tag[%idx]`, so @@ -479,14 +505,14 @@ def _emit_dmas_and_waits(ctx, block, builder, dma_by_op, bufs): n_dma = n_wait = 0 for op in list(walk_ops(block)): name = op.operation.name - if name == "memref.dma_start": + if name == "togsim.transfer": node = dma_by_op.get(id(op.operation)) if node is None: continue _emit_one_dma(ctx, op, node, builder, bufs, tags) originals.append(op) n_dma += 1 - elif name == "memref.dma_wait": + elif name == "togsim.wait": if _emit_one_wait(ctx, op, tags): n_wait += 1 originals.append(op) diff --git a/PyTorchSimFrontend/mlir/passes/build_tog.py b/PyTorchSimFrontend/mlir/passes/build_tog.py index ae515010..5a40feec 100644 --- a/PyTorchSimFrontend/mlir/passes/build_tog.py +++ b/PyTorchSimFrontend/mlir/passes/build_tog.py @@ -608,11 +608,11 @@ def bool_true(k): self.print_operation(inner, loop_node) return - if name == "memref.dma_start": + if name == "togsim.transfer": self._handle_dma_start(op, node) return - if name == "memref.dma_wait": + if name == "togsim.wait": self._handle_dma_wait(op, node) return @@ -716,36 +716,60 @@ def _handle_compute(self, op, node): return self._append_vector_compute(node, op) - # ---- dma_start ---- - def _dma_start_fields(self, op): - """Decode memref.dma_start operands by memref ranks. + # ---- transfer (formerly memref.dma_start) ---- + @staticmethod + def _transfer_is_load(op, dma_type_operand): + """True if a `togsim.transfer` is a load (DRAM -> SRAM), False for a store. - Layout: src[srcIdx], dst[dstIdx], numElements, tag[tagIdx], stride, - numElementsPerStride. - """ + Prefer the `dma_kind` attr ("MVIN"/"MVIN2"/"MVIN3" load, "MVOUT" store); + fall back to the `dma_type` operand constant (MVIN=2/MVIN2=1/MVIN3=14 are + loads, MVOUT=3 is a store).""" + oper = op.operation + if "dma_kind" in oper.attributes: + try: + kind = ir.StringAttr(oper.attributes["dma_kind"]).value + except Exception: + kind = str(oper.attributes["dma_kind"]).strip('"') + if kind: + return kind.upper().startswith("MVIN") + c = _const_index_value(dma_type_operand) + if c is not None: + return c in (1, 2, 14) + return True + + def _dma_start_fields(self, op): + """Decode a `togsim.transfer` into the legacy src/dst view. + + togsim.transfer operand layout (mirrors build_skeleton._transfer_fields / + lower_transfer_to_gemmini): + dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst[, offset] + The DRAM side is always operand[0]/[1], the SRAM spad operand[2]/[3], the + runtime tag slot operand[4] (tag memref) + operand[5] (tag_idx). The + optional indirect-offset spad is operand[8]. + + Direction (from dma_kind / dma_type) decides the src/dst mapping so the + rest of build_tog keeps the old memref.dma_start convention: for a load + the DRAM side is the SOURCE and the SRAM spad the DEST; a store reverses + it. src/dst therefore still carry the right memory spaces (DRAM=0, + SRAM=1) that `_handle_dma_start` recovers is_write from.""" operands = list(op.operands) - i = 0 - src = operands[i] - src_type = src.type - src_rank = len(ir.MemRefType(src_type).shape) - i += 1 - src_indices = operands[i:i + src_rank] - i += src_rank - dst = operands[i] - dst_type = dst.type - dst_rank = len(ir.MemRefType(dst_type).shape) - i += 1 - dst_indices = operands[i:i + dst_rank] - i += dst_rank - i += 1 # numElements - tag = operands[i] - tag_rank = len(ir.MemRefType(tag.type).shape) - i += 1 - tag_indices = operands[i:i + tag_rank] + dram, dram_idx = operands[0], operands[1] + sram, sram_idx = operands[2], operands[3] + tag, tag_idx = operands[4], operands[5] + dma_type = operands[6] + offset = operands[8] if len(operands) > 8 else None + + if self._transfer_is_load(op, dma_type): # DRAM -> SRAM + src, src_idx = dram, dram_idx + dst, dst_idx = sram, sram_idx + else: # SRAM -> DRAM + src, src_idx = sram, sram_idx + dst, dst_idx = dram, dram_idx return { - "src": src, "src_type": src_type, "src_indices": src_indices, - "dst": dst, "dst_type": dst_type, "dst_indices": dst_indices, - "tag": tag, "tag_indices": tag_indices, + "src": src, "src_type": src.type, "src_indices": [src_idx], + "dst": dst, "dst_type": dst.type, "dst_indices": [dst_idx], + "tag": tag, "tag_indices": [tag_idx], + "dma_type": dma_type, "offset": offset, } def _handle_dma_start(self, op, node): @@ -776,7 +800,7 @@ def _handle_dma_start(self, op, node): tile_size = [int(x) for x in tile_shape] tile_stride = [] - ds = _int_array_attr(oper, "dram_stride") + ds = _int_array_attr(oper, "dram_stride") # TOGDMANode.tile_stride keeps the DRAM stride (as before) if ds: tile_stride = list(ds) @@ -857,10 +881,11 @@ def _handle_dma_start(self, op, node): # ---- dma_wait ---- def _handle_dma_wait(self, op, node): oper = op.operation + # togsim.wait operands: (tag[0], tag_idx[1], num_elements[2]). tag_idx is + # now a single operand (memref.dma_wait had a variadic [tag_idx]). operands = list(oper.operands) tag = operands[0] - tag_rank = len(ir.MemRefType(tag.type).shape) - tag_indices = operands[1:1 + tag_rank] + tag_indices = operands[1:2] tag_index_list = [] tag_stride_list = [] @@ -880,11 +905,11 @@ def _handle_dma_wait(self, op, node): tag_stride_list = _collect_coefficients(amap.results[0]) tag_divider_list = _collect_dividers(amap.results[0]) - # base address: scan users of tag memref for a dma_start. + # base address: scan users of tag memref for a togsim.transfer. address = "arg" for use in tag.uses: user = use.owner - if user.name == "memref.dma_start": + if user.name == "togsim.transfer": f = self._dma_start_fields(user) dst_space = _memref_space(f["dst_type"]) src_space = _memref_space(f["src_type"]) diff --git a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py index d7571d2b..22eb999e 100644 --- a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py +++ b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py @@ -11,13 +11,16 @@ IRMapping; the MLIR Python bindings expose no IRMapping, so this port builds the fused nest directly and emits each DMA inside it using the fused induction vars (equivalence target: same loop structure / counts, same offset maps, same -dma_start operands+attrs -- validated against mlir-opt -dma-fine-grained and the -end-to-end gemm/conv/model tests, not byte-exact SSA text). +togsim.transfer operands+attrs -- validated against mlir-opt -dma-fine-grained and +the end-to-end gemm/conv/model tests, not byte-exact SSA text). -Operates on the customized memref.dma_start convention (see lower_dma_to_gemmini): -operands = src, *src_idx, dst, *dst_idx, num_elements(dma_type), tag, *tag_idx, -stride(=vlane_split_axis), num_elements_per_stride(=vlane_stride). MVIN dma_type in -{2,1,14}; tile shape = dst shape for MVIN. +Operates on the togsim.transfer convention (see mlir_codegen_backend.emit_transfer +and lower_transfer_to_gemmini): operands = dram, dram_idx, sram, sram_idx, tag, +tag_idx, dma_type, vst(=vlane_stride)[, offset]; attrs = dma_kind, vlane_split_axis +(i64), dram_stride[], tile_stride[], padding, [subtile_size, async]. Direction is +derived from dma_kind / dma_type: MVIN => src=dram, dst=sram; MVOUT => src=sram, +dst=dram. tile shape = the sram memref shape for BOTH directions. MVIN dma_type in +{2,1,14}. Pipeline entry point: run_fine_grained(in_path, out_path, vectorlane). """ @@ -62,25 +65,37 @@ def _is_block_arg(v): class _Dma: - """Positional view of a customized memref.dma_start op.""" + """Positional view of a togsim.transfer op. + + operands: dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst[, offset] + Direction from dma_kind / dma_type: MVIN => src=dram, dst=sram (MVOUT swaps). + self.src_idx is a single-element list holding the base DRAM idx for MVIN (the + SRAM idx for MVOUT); tile_shape is always the sram memref shape. + """ def __init__(self, op): self.op = op operands = list(op.operands) - src_rank = len(ir.MemRefType(operands[0].type).shape) - i = 0 - self.src = operands[i]; i += 1 - self.src_idx = operands[i:i + src_rank]; i += src_rank - self.dst = operands[i]; i += 1 - dst_rank = len(ir.MemRefType(self.dst.type).shape) - self.dst_idx = operands[i:i + dst_rank]; i += dst_rank - self.num_elements = operands[i]; i += 1 - self.tag = operands[i]; i += 1 - tag_rank = len(ir.MemRefType(self.tag.type).shape) - self.tag_idx = operands[i:i + tag_rank]; i += tag_rank - self.stride = operands[i]; i += 1 # = vlane_split_axis - self.num_elements_per_stride = operands[i] # = vlane_stride - self.src_rank, self.dst_rank, self.tag_rank = src_rank, dst_rank, tag_rank + # dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst[, offset] + self.dram = operands[0] + self.dram_idx = operands[1] + self.sram = operands[2] + self.sram_idx = operands[3] + self.tag = operands[4] + self.tag_idx = operands[5] + self.num_elements = operands[6] # = dma_type const operand + self.num_elements_per_stride = operands[7] # = vlane_stride (vst) + + self.sram_rank = len(ir.MemRefType(self.sram.type).shape) + # Direction: MVIN reads dram -> sram; MVOUT writes sram -> dram. + if self.is_mvin: + self.src, self.dst = self.dram, self.sram + self.src_idx = [self.dram_idx] + else: + self.src, self.dst = self.sram, self.dram + self.src_idx = [self.sram_idx] + self.src_rank = len(ir.MemRefType(self.src.type).shape) + self.dst_rank = len(ir.MemRefType(self.dst.type).shape) @property def dma_type(self): @@ -92,21 +107,21 @@ def is_mvin(self): @property def vlane_split_axis(self): - return _const_int(self.stride) + return ir.IntegerAttr(self.op.attributes["vlane_split_axis"]).value @property def vlane_stride(self): return _const_int(self.num_elements_per_stride) & 0x7FFF def tile_shape(self): - mt = ir.MemRefType((self.dst if self.is_mvin else self.src).type) - return list(mt.shape) + return list(ir.MemRefType(self.sram.type).shape) def subtile_size(self): return attr_i64_array(self.op, "subtile_size", default=[]) def sram_stride(self): - return attr_i64_array(self.op, "sram_stride", default=[]) + # togsim.transfer names the spad stride "tile_stride". + return attr_i64_array(self.op, "tile_stride", default=[]) def dram_stride(self): return attr_i64_array(self.op, "dram_stride", default=[]) @@ -200,10 +215,13 @@ def _apply(map_, operands, ip): def _dma_attrs(dma): - """Mirror getDmaAttrs: keep subtile/sram/dram strides, set async + fine_grained.""" + """Build the emitted togsim.transfer's attrs: copy dma_kind, vlane_split_axis, + dram_stride, tile_stride (the spad stride), subtile_size and padding straight + from the source op; set async (BoolAttr) + fine_grained (BoolAttr true).""" attrs = {} op = dma.op - for k in ("subtile_size", "sram_stride", "dram_stride"): + for k in ("dma_kind", "vlane_split_axis", "dram_stride", "tile_stride", + "subtile_size", "padding"): if k in op.attributes: attrs[k] = op.attributes[k] attrs["async"] = ir.BoolAttr.get(dma.is_async()) @@ -212,26 +230,20 @@ def _dma_attrs(dma): def _emit_dma(dma, ivs, vectorlane, ip): - """Emit one fine-grained memref.dma_start at `ip`, indexed by `ivs` (the fused + """Emit one fine-grained togsim.transfer at `ip`, indexed by `ivs` (the fused induction vars for this DMA's dims, in dim order).""" - idx_ty = ir.IndexType.get() - zero = _const_index(0, ip) - dram_off = _apply(_build_dram_map(dma), ivs, ip) - src_idx0 = dma.src_idx[0] - dram_idx = _apply(_sum_map(), [dram_off, src_idx0], ip) + # DRAM base index = the original transfer's dram_idx operand. + dram_idx = _apply(_sum_map(), [dram_off, dma.dram_idx], ip) + # SRAM offset is a SINGLE linear sram_idx operand (row-major stride 1). sram_off = _apply(_build_sram_map(dma, vectorlane), ivs, ip) + # Per-subtile tag index (required for async DMA<->barrier pairing downstream). tag_idx = _apply(_build_tag_map(dma, list(range(len(dma.tile_shape())))), ivs, ip) - # SRAM indices: zeros except the last = sram offset (mirror sramIndices.back()). - sram_indices = [zero] * dma.dst_rank - sram_indices[-1] = sram_off - - operands = [dma.src, dram_idx, dma.dst, *sram_indices, - dma.num_elements, dma.tag, tag_idx, - dma.stride, dma.num_elements_per_stride] - ir.Operation.create("memref.dma_start", results=[], operands=operands, + operands = [dma.dram, dram_idx, dma.sram, sram_off, + dma.tag, tag_idx, dma.num_elements, dma.num_elements_per_stride] + ir.Operation.create("togsim.transfer", results=[], operands=operands, attributes=_dma_attrs(dma), ip=ip) @@ -320,7 +332,7 @@ def _run_func(func, vectorlane): name = op.operation.name if name == "linalg.matmul" and matmul is None: matmul = op - elif name == "memref.dma_start": + elif name == "togsim.transfer": dmas.append(op) if matmul is None: return diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py b/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py index ad287499..4629df29 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py @@ -50,15 +50,15 @@ def run_standard_lowering(in_path, out_path=None, timing=False): out_path = in_path from mlir.ir import Context, Module, Location from mlir.passmanager import PassManager - from . import lower_dma_to_gemmini + from . import lower_transfer_to_gemmini ctx = Context() ctx.allow_unregistered_dialects = True with ctx, Location.unknown(): with open(in_path) as f: module = Module.parse(f.read()) - # Imperative Python pass: memref.dma_start/dma_wait -> Gemmini asm (replaces - # the C++ test-memref-to-gemmini), then the registered standard lowering. - lower_dma_to_gemmini.run(module, timing=timing) + # Imperative Python pass: togsim.transfer -> Gemmini asm directly (no + # memref.dma_start intermediate), then the registered standard lowering. + lower_transfer_to_gemmini.run(module, timing=timing) PassManager.parse(STANDARD_PIPELINE, ctx).run(module.operation) with open(out_path, "w") as f: f.write(str(module)) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py index 55bbae5a..d11c886a 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py @@ -255,8 +255,9 @@ def _transfer_write(value, dest, indices): def _dma_wait(tag, idx, num_elements): - from mlir.dialects import memref - memref.DmaWaitOp(tag, [idx], num_elements) + # togsim-level counterpart of the async-DMA barrier (paired with togsim.transfer), + # instead of memref.dma_wait -- no memref.* DMA ops remain in the pipeline. + ir.Operation.create("togsim.wait", results=[], operands=[tag, idx, num_elements]) def _vcix(name, operands, result_tys, attrs): @@ -277,22 +278,29 @@ def _reaches(value, target): class _DmaView: - """Positional view of a customized memref.dma_start (see lower_dma_to_gemmini).""" + """Positional view of a togsim.transfer op (see emit_transfer / + lower_transfer_to_gemmini). + + New operand layout: (dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, + vst [, offset]). Direction (which of dram/sram is source vs dest) comes from + the `dma_kind` attr ("MVIN"|"MVOUT"), not from operand position. To keep the + downstream _dram_is_write(src, dst) / `sram = d.dst` logic working (it + distinguishes MVIN loads from MVOUT stores by memory space), src/dst are + exposed direction-swapped: MVIN -> src=dram, dst=sram; MVOUT -> src=sram, + dst=dram.""" def __init__(self, op): self.op = op operands = list(op.operands) - src_rank = len(ir.MemRefType(operands[0].type).shape) - i = 0 - self.src = operands[i]; i += 1 - i += src_rank - self.dst = operands[i]; i += 1 - dst_rank = len(ir.MemRefType(self.dst.type).shape) - i += dst_rank - i += 1 # num_elements - self.tag = operands[i]; i += 1 - tag_rank = len(ir.MemRefType(self.tag.type).shape) - self.tag_idx = operands[i:i + tag_rank] + dram = operands[0] + sram = operands[2] + self.tag = operands[4] + self.tag_idx = [operands[5]] + kind = ir.StringAttr(op.attributes["dma_kind"]).value + if kind == "MVOUT": + self.src, self.dst = sram, dram + else: # MVIN (load) + self.src, self.dst = dram, sram def subtile_size(self): a = self.op.attributes @@ -358,10 +366,9 @@ def a64(v): return ir.IntegerAttr.get(i64, v) BiasIdx = None subtileM, subtileN, subtileK = M, N, K a_subk = b_subk = None - # Mirror the C++ isAInitialized / isBInitialized flags: an operand is - # "initialized" either by an MVIN dma_start (tag found below) or by a - # preceding affine.vector_store into its root memref (the fused case, e.g. - # SDPA scores.V where B is the softmax output produced in-place, not DMAed). + # Mirror the C++ isAInitialized / isBInitialized flags: an operand is initialized + # either by an MVIN togsim.transfer (tag found below) or by a preceding + # affine.vector_store into its root memref (the fused case, e.g. SDPA scores.V). isAInit = isBInit = False def _root(v): @@ -380,7 +387,7 @@ def _root(v): elif dest == rootB: isBInit = True continue - if o.operation.name != "memref.dma_start": + if o.operation.name != "togsim.transfer": continue d = _DmaView(o.operation) dram, is_write = _dram_is_write(d.src, d.dst) diff --git a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py new file mode 100644 index 00000000..fcb0aa2a --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py @@ -0,0 +1,257 @@ +"""Lower togsim.transfer DIRECTLY to Gemmini RISC-V inline asm (no memref.dma_start). + +Merges decompose_transfer (the <=4D Gemmini-limit handling: drop unit dims / +collapse / >4D affine.for peel with lane-banked SRAM offset) with +lower_dma_to_gemmini (the CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm emission). +togsim.transfer is unregistered so it carries every runtime descriptor as an +operand -- including the future masked-clamp low/high vectors -- which a registered +memref.dma_start cannot. + +timing=False: emit the gemmini asm. timing=True: erase the transfer (the TOG carries +DMA timing; the cycle binary needs no asm). +""" + +OP_NAME = "togsim.transfer" +WAIT_NAME = "togsim.wait" +MARKERS = (OP_NAME, WAIT_NAME) + +from ._mlir_util import walk_ops +from .lower_dma_to_gemmini import _i64_signed, _row_major_strides, _elem_bytes, _asm, CONSTRAINTS +from .decompose_transfer import _int_array, _const_int, _squeeze_reassociation + +CONFIG, CONFIG2, CONFIG3, CONFIG4 = 0, 4, 5, 6 +MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 +CONFIG_TYPE = {MVIN: 0, MVIN2: 1, MVIN3: 2, MVOUT: 3} +MAX_TENSOR_DIM = 4 + + +def run(module, timing=False, vectorlane=128, **_): + from mlir.ir import (InsertionPoint, Operation, MemRefType, ArrayAttr, + IntegerAttr, IntegerType, IndexType, DenseI64ArrayAttr, + DenseI32ArrayAttr, StridedLayoutAttr, AffineMap, AffineMapAttr, + AffineExpr, BoolAttr, FlatSymbolRefAttr, TypeAttr) + from mlir.dialects import affine, llvm, arith, memref + i64 = IntegerType.get_signless(64) + idx_ty = IndexType.get() + + sym2type = {} + for g in module.operation.regions[0].blocks[0].operations: + if g.operation.name == "memref.global": + sym2type[g.attributes["sym_name"].value] = MemRefType(TypeAttr(g.attributes["type"]).value) + + def i64_const(value): + return arith.ConstantOp(i64, IntegerAttr.get(i64, _i64_signed(value))).result + + def asm(func7, rs1, rs2): + llvm.InlineAsmOp(None, [rs1, rs2], _asm(func7), CONSTRAINTS, + has_side_effects=True, asm_dialect=0) + + def elem_addr_i64(memref_val, indices, mtype, elem_bytes): + base = memref.ExtractAlignedPointerAsIndexOp(memref_val).result + strides = _row_major_strides(list(mtype.shape)) + off = None + for k, ival in enumerate(indices): + if strides[k] == 0: + continue + term = ival + if strides[k] != 1: + term = arith.MulIOp(ival, arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, strides[k])).result).result + off = term if off is None else arith.AddIOp(off, term).result + if off is not None: + byte = arith.MulIOp(off, arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, elem_bytes)).result).result + base = arith.AddIOp(base, byte).result + return arith.IndexCastOp(i64, base).result + + targets, waits = [], [] + for region in module.operation.regions: + for b in region.blocks: + for op in walk_ops(b): + if op.operation.name == OP_NAME: + targets.append(op.operation) + elif op.operation.name == WAIT_NAME: + waits.append(op.operation) + + for op in waits: # togsim.wait: erase in both modes (the barrier is a sync marker) + op.erase() + + for op in targets: + op_operands = list(op.operands) + dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst = op_operands[:8] + offset_sym = (op_operands[8].owner.attributes["name"] if "indirect" in op.attributes else None) + kind = op.attributes["dma_kind"].value + dma_type_val = _const_int(dma_type) # MVIN(2)/MVIN2(1)/MVIN3(14)/MVOUT(3) + is_mvin = dma_type_val in (MVIN, MVIN2, MVIN3) + vlane_axis = IntegerAttr(op.attributes["vlane_split_axis"]).value + dram_stride = _int_array(op.attributes["dram_stride"]) + tile_stride = _int_array(op.attributes["tile_stride"]) + vlane_stride = _const_int(vst, 1) + try: + subtile = _int_array(op.attributes["subtile_size"]) + except KeyError: + subtile = None + + if timing: + op.erase() + continue + + sram_ty = MemRefType(sram.type) + elem, space = sram_ty.element_type, sram_ty.memory_space + elem_bytes = _elem_bytes(sram_ty.element_type) + dram_ty = MemRefType(dram.type) + tile_shape = list(sram_ty.shape) + eff = [i for i, e in enumerate(tile_shape) if e > 1] + indirect = offset_sym is not None + + def _const(v): + return arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, v)).result + + def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, + desc_dram_strides, desc_spad_strides, subtile_shape): + cfg_shape = subtile_shape if subtile_shape is not None else desc_shape + expand = MAX_TENSOR_DIM - len(cfg_shape) + shape4 = [1] * expand + list(cfg_shape) + dram4 = [0] * expand + list(desc_dram_strides) + spad4 = [0] * expand + list(desc_spad_strides) + vsa4 = vsa_val + expand + config_type = CONFIG_TYPE[dma_type_val] + sram_c = MemRefType(sram_mem.type) + dram_addr = elem_addr_i64(dram, [dram_idx_val], dram_ty, elem_bytes) + spad_addr = elem_addr_i64(sram_mem, sram_indices, sram_c, elem_bytes) + cfg_rs1 = i64_const(((shape4[0] & 0xFFFF) << 48) | ((shape4[1] & 0xFFFF) << 32) + | ((shape4[2] & 0xFFFF) << 16) | (shape4[3] & 0xFFFF)) + cfg_rs2 = i64_const((vlane_stride << 32) | ((config_type & 0x3) << 17) + | ((1 if indirect else 0) << 16) + | ((vsa4 & 0x3) << 14) | elem_bytes) + asm(CONFIG, cfg_rs1, cfg_rs2) + asm(CONFIG2, i64_const((dram4[0] << 32) | (dram4[1] & 0xFFFFFFFF)), + i64_const((dram4[2] << 32) | (dram4[3] & 0xFFFFFFFF))) + asm(CONFIG3, i64_const((spad4[0] << 32) | (spad4[1] & 0xFFFFFFFF)), + i64_const((spad4[2] << 32) | (spad4[3] & 0xFFFFFFFF))) + if indirect: + sym = FlatSymbolRefAttr(offset_sym).value if not isinstance(offset_sym, str) else offset_sym + off_ty = sym2type[sym] + ind_base = memref.ExtractAlignedPointerAsIndexOp(memref.GetGlobalOp(off_ty, sym).result).result + ind_addr = arith.IndexCastOp(i64, ind_base).result + ind_esize = _elem_bytes(off_ty.element_type) + off_stride = IntegerAttr(op.attributes["offset_stride"]).value + asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (off_stride & 0xFFFF))) + asm(dma_type_val, dram_addr, spad_addr) + + if offset_sym is not None: + offset_sym = FlatSymbolRefAttr(offset_sym).value if not isinstance(offset_sym, str) else offset_sym + + if len(tile_shape) <= 4: + with InsertionPoint(op): + # sram offset is linear (row-major stride 1) -> last index only; others 0. + sidx = [_const(0)] * (len(tile_shape) - 1) + [sram_idx] + _emit_asm(sram, sidx, dram_idx, vlane_axis, + tile_shape, dram_stride, tile_stride, subtile) + op.erase() + continue + + if len(eff) <= 4: + groups, target = _squeeze_reassociation(tile_shape) + reassoc = ArrayAttr.get( + [ArrayAttr.get([IntegerAttr.get(i64, d) for d in g]) for g in groups]) + collapsed_ty = MemRefType.get(target, elem, memory_space=space) + keep = [next((d for d in g if tile_shape[d] > 1), g[-1]) for g in groups] + dr = [dram_stride[i] for i in keep] + tl = [tile_stride[i] for i in keep] + st = [subtile[i] for i in keep] if subtile is not None else None + new_vlane = next(gi for gi, g in enumerate(groups) if vlane_axis in g) + with InsertionPoint(op): + sram_c = Operation.create( + "memref.collapse_shape", results=[collapsed_ty], operands=[sram], + attributes={"reassociation": reassoc}).results[0] + sidx = [_const(0)] * (len(target) - 1) + [sram_idx] + _emit_asm(sram_c, sidx, dram_idx, new_vlane, target, dr, tl, st) + op.erase() + continue + + # >4 effective dims: affine.for peel (mirrors decompose_transfer peel path) + peeled, inner = eff[:-4], eff[-4:] + ndim = len(tile_shape) + inner_shape = [tile_shape[d] for d in inner] + inner_strides = [tile_stride[d] for d in inner] + dr = [dram_stride[d] for d in inner] + tl = [tile_stride[d] for d in inner] + st = [subtile[d] for d in inner] if subtile is not None else None + if vlane_axis in inner: + new_vlane = inner.index(vlane_axis) + elif vlane_axis in peeled: + raise NotImplementedError( + f"vlane split axis {vlane_axis} peeled into the outer loop nest") + else: + new_vlane = 0 + split_extent = tile_shape[vlane_axis] + nr_outerloop = max( + (split_extent + vectorlane * vlane_stride - 1) // (vectorlane * vlane_stride), 1) + new_size = nr_outerloop * vlane_stride + target_stride = tile_stride[vlane_axis] + + def _phys(d): + s = tile_stride[d] + return s // split_extent * new_size if s > target_stride else s + + static_sizes = [1] * ndim + for d in inner: + static_sizes[d] = tile_shape[d] + res_ty = MemRefType.get( + inner_shape, elem, + layout=StridedLayoutAttr.get(0, inner_strides), memory_space=space) + + cur_ip = InsertionPoint(op) + ivs = [] + for d in peeled: + floop = affine.AffineForOp(0, tile_shape[d], 1, ip=cur_ip) + floop.operation.attributes["inner_loop"] = BoolAttr.get(True) + ivs.append(floop.induction_variable) + with InsertionPoint(floop.body): + affine.AffineYieldOp([]) + cur_ip = InsertionPoint.at_block_terminator(floop.body) + + npeel = len(peeled) + with cur_ip: + sub = Operation.create( + "memref.subview", results=[res_ty], operands=[sram], + attributes={"static_offsets": DenseI64ArrayAttr.get([0] * ndim), + "static_sizes": DenseI64ArrayAttr.get(static_sizes), + "static_strides": DenseI64ArrayAttr.get([1] * ndim), + "operandSegmentSizes": DenseI32ArrayAttr.get([1, 0, 0, 0])} + ).results[0] + sram_expr = AffineExpr.get_dim(0) * _phys(peeled[0]) + for k in range(1, npeel): + sram_expr = sram_expr + AffineExpr.get_dim(k) * _phys(peeled[k]) + sram_off_val = Operation.create( + "affine.apply", results=[idx_ty], operands=list(ivs), + attributes={"map": AffineMapAttr.get(AffineMap.get(npeel, 0, [sram_expr]))} + ).results[0] + dram_expr = AffineExpr.get_dim(0) + for k in range(npeel): + dram_expr = dram_expr + AffineExpr.get_dim(k + 1) * dram_stride[peeled[k]] + dram_idx_val = Operation.create( + "affine.apply", results=[idx_ty], operands=[dram_idx, *ivs], + attributes={"map": AffineMapAttr.get(AffineMap.get(npeel + 1, 0, [dram_expr]))} + ).results[0] + zero = _const(0) + _emit_asm(sub, [zero, zero, zero, sram_off_val], dram_idx_val, new_vlane, + inner_shape, dr, tl, st) + op.erase() + + +def lower_text(text, timing=False): + if OP_NAME not in text: + return text + from mlir.ir import Context, Module, Location + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + m = Module.parse(text) + run(m, timing=timing) + return str(m) + + +if __name__ == "__main__": + import sys + out = lower_text(open(sys.argv[1]).read()) + (open(sys.argv[2], "w").write(out) if len(sys.argv) > 2 else sys.stdout.write(out)) From bfef916f3f48f5d84a6a7d2dbb361de9399a9611 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 17:18:58 +0900 Subject: [PATCH 19/32] [Frontend] Lower togsim.transfer to a DMA descriptor + CONFIG_DESC (no CONFIG/2/3/4) Instead of packing the transfer params into four CONFIG asm instructions, build a 144-byte DMA descriptor as an 18xi64 memref.global (.data, non-constant so runtime fields can be written) and emit a single CONFIG_DESC that hands its address to the MVIN/MVOUT (which read the struct, matching the Spike-side torchsim_config_desc). This is the vehicle for the masked-DMA low/high clamp: dim_low/dim_high/fill are descriptor fields (currently defaulted, so behavior is unchanged), not new special registers. - _pack_desc encodes the fields into the struct's i64 slots (dim_size/low/high/mm_stride/ spad_stride/element_size/vlane/config_type/flags/indirect stride+esize). - descriptor globals are deduped by content; the indirect offset spad address is a runtime value, stored into slot 15 before the transfer (flags.bit0 = indirect). Requires the Spike torchsim_config_desc instruction (riscv-isa-sim). Validated functional (Spike) + trace (TOGSim): add, matmul (subtile+async), softmax, gather/scatter (indirect runtime store), constant_pad all pass. --- .../mlir/passes/lower_transfer_to_gemmini.py | 80 +++++++++++++++---- 1 file changed, 65 insertions(+), 15 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py index fcb0aa2a..46575ba7 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py @@ -19,8 +19,9 @@ from .lower_dma_to_gemmini import _i64_signed, _row_major_strides, _elem_bytes, _asm, CONSTRAINTS from .decompose_transfer import _int_array, _const_int, _squeeze_reassociation -CONFIG, CONFIG2, CONFIG3, CONFIG4 = 0, 4, 5, 6 +CONFIG, CONFIG2, CONFIG3, CONFIG4, CONFIG_DESC = 0, 4, 5, 6, 7 MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 +DESC_SLOTS = 18 # 144-byte DMA descriptor as 18 x i64 (see project-dma-descriptor) CONFIG_TYPE = {MVIN: 0, MVIN2: 1, MVIN3: 2, MVOUT: 3} MAX_TENSOR_DIM = 4 @@ -29,16 +30,70 @@ def run(module, timing=False, vectorlane=128, **_): from mlir.ir import (InsertionPoint, Operation, MemRefType, ArrayAttr, IntegerAttr, IntegerType, IndexType, DenseI64ArrayAttr, DenseI32ArrayAttr, StridedLayoutAttr, AffineMap, AffineMapAttr, - AffineExpr, BoolAttr, FlatSymbolRefAttr, TypeAttr) + AffineExpr, BoolAttr, FlatSymbolRefAttr, TypeAttr, + DenseElementsAttr, RankedTensorType, StringAttr, UnitAttr) from mlir.dialects import affine, llvm, arith, memref i64 = IntegerType.get_signless(64) idx_ty = IndexType.get() + module_top = module.operation.regions[0].blocks[0] sym2type = {} - for g in module.operation.regions[0].blocks[0].operations: + for g in module_top.operations: if g.operation.name == "memref.global": sym2type[g.attributes["sym_name"].value] = MemRefType(TypeAttr(g.attributes["type"]).value) + desc_ty = MemRefType.get([DESC_SLOTS], i64) + desc_globals = {} # slots tuple -> global sym name (dedup) + desc_counter = [0] + + def _pack_desc(shape4, dram4, spad4, elem_bytes, vlstride, vsa4, cfg_type, indirect, + ind_stride, ind_esize): + # 18 i64 slots matching the C dma_descriptor byte layout (little-endian). + lo = [0, 0, 0, 0] + hi = list(shape4) + def s2(a, b): return (a & 0xFFFFFFFF) | ((b & 0xFFFFFFFF) << 32) + m = 0xFFFFFFFFFFFFFFFF + slots = [ + s2(shape4[0], shape4[1]), s2(shape4[2], shape4[3]), + s2(lo[0], lo[1]), s2(lo[2], lo[3]), + s2(hi[0], hi[1]), s2(hi[2], hi[3]), + dram4[0] & m, dram4[1] & m, dram4[2] & m, dram4[3] & m, + spad4[0] & m, spad4[1] & m, spad4[2] & m, spad4[3] & m, + (elem_bytes & 0xFFFF) | ((vlstride & 0xFFFF) << 16) | ((vsa4 & 0xFF) << 32) + | ((cfg_type & 0xFF) << 40) | (((1 if indirect else 0) & 0xFFFF) << 48), + 0, # +120 indirect_addr (runtime) + (ind_stride & 0xFFFF) | ((ind_esize & 0xFFFF) << 16), # +128 + 0, # +136 fill (runtime/step3) + ] + return slots + + def _desc_global(slots): + key = tuple(slots) + if key not in desc_globals: + name = f"dma_desc_{desc_counter[0]}" + desc_counter[0] += 1 + import numpy as np + init = DenseElementsAttr.get( + np.array([_i64_signed(v) for v in slots], dtype=np.int64), + type=RankedTensorType.get([DESC_SLOTS], i64)) + with InsertionPoint.at_block_begin(module_top): + Operation.create("memref.global", results=[], operands=[], attributes={ + "sym_name": StringAttr.get(name), + "sym_visibility": StringAttr.get("private"), + "type": TypeAttr.get(desc_ty), + "initial_value": init}) + desc_globals[key] = name + return desc_globals[key] + + def desc_ptr_and_store_indirect(slots, ind_addr_i64): + # get_global -> byte pointer; for indirect, store the runtime addr into slot 15. + name = _desc_global(slots) + g = memref.GetGlobalOp(desc_ty, name).result + if ind_addr_i64 is not None: + memref.StoreOp(ind_addr_i64, g, [arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, 15)).result]) + base = memref.ExtractAlignedPointerAsIndexOp(g).result + return arith.IndexCastOp(i64, base).result + def i64_const(value): return arith.ConstantOp(i64, IntegerAttr.get(i64, _i64_signed(value))).result @@ -117,24 +172,19 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, sram_c = MemRefType(sram_mem.type) dram_addr = elem_addr_i64(dram, [dram_idx_val], dram_ty, elem_bytes) spad_addr = elem_addr_i64(sram_mem, sram_indices, sram_c, elem_bytes) - cfg_rs1 = i64_const(((shape4[0] & 0xFFFF) << 48) | ((shape4[1] & 0xFFFF) << 32) - | ((shape4[2] & 0xFFFF) << 16) | (shape4[3] & 0xFFFF)) - cfg_rs2 = i64_const((vlane_stride << 32) | ((config_type & 0x3) << 17) - | ((1 if indirect else 0) << 16) - | ((vsa4 & 0x3) << 14) | elem_bytes) - asm(CONFIG, cfg_rs1, cfg_rs2) - asm(CONFIG2, i64_const((dram4[0] << 32) | (dram4[1] & 0xFFFFFFFF)), - i64_const((dram4[2] << 32) | (dram4[3] & 0xFFFFFFFF))) - asm(CONFIG3, i64_const((spad4[0] << 32) | (spad4[1] & 0xFFFFFFFF)), - i64_const((spad4[2] << 32) | (spad4[3] & 0xFFFFFFFF))) + ind_addr = None + ind_stride = ind_esize = 0 if indirect: sym = FlatSymbolRefAttr(offset_sym).value if not isinstance(offset_sym, str) else offset_sym off_ty = sym2type[sym] ind_base = memref.ExtractAlignedPointerAsIndexOp(memref.GetGlobalOp(off_ty, sym).result).result ind_addr = arith.IndexCastOp(i64, ind_base).result ind_esize = _elem_bytes(off_ty.element_type) - off_stride = IntegerAttr(op.attributes["offset_stride"]).value - asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (off_stride & 0xFFFF))) + ind_stride = IntegerAttr(op.attributes["offset_stride"]).value + slots = _pack_desc(shape4, dram4, spad4, elem_bytes, vlane_stride, vsa4, + config_type, indirect, ind_stride, ind_esize) + desc_ptr = desc_ptr_and_store_indirect(slots, ind_addr) + asm(CONFIG_DESC, desc_ptr, i64_const(0)) asm(dma_type_val, dram_addr, spad_addr) if offset_sym is not None: From 8de755c6c15e74b65c97c52d8accfa8d2d6ff0c3 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 17:41:21 +0900 Subject: [PATCH 20/32] [Build] Bump spike to v1.0.5 (DMA descriptor, masked clamp, MVOUT accumulate) --- thirdparty/github-releases.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/github-releases.json b/thirdparty/github-releases.json index 390b68fb..2618c236 100644 --- a/thirdparty/github-releases.json +++ b/thirdparty/github-releases.json @@ -13,7 +13,7 @@ }, "spike": { "repository": "PSAL-POSTECH/riscv-isa-sim", - "release_tag": "v1.0.3", + "release_tag": "v1.0.5", "asset_name": "spike-release.tar.gz" } } From ddfc86ead4743fde6e4d443c36b6cd893d7cfb23 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 7 Jul 2026 11:34:03 +0900 Subject: [PATCH 21/32] [Frontend] Masked DMA: drop tile divisibility for non-dividing shapes Clamp each tile axis to its real DRAM extent in the DMA descriptor, so a tile that runs past the extent is filled with the consuming reduction's identity instead of MAC-ing garbage, and a store skips the ragged tail instead of writing the systolic pad region. Covers pointwise, pad, reduction, gather/scatter, matmul and conv, plus the fused-epilogue output store. Adds a golden unit test for the lower_transfer_to_gemmini lowering and a non-dividing regression suite to CI. --- .github/workflows/pytorchsim_test.yml | 18 +++ .../mlir/mlir_codegen_backend.py | 152 +++++++++++++++--- PyTorchSimFrontend/mlir/mlir_common.py | 53 +----- .../mlir/mlir_conv_mt_template.py | 4 + .../mlir/mlir_conv_sb_template.py | 3 + .../mlir/mlir_conv_sbs_template.py | 3 + PyTorchSimFrontend/mlir/mlir_conv_template.py | 3 + PyTorchSimFrontend/mlir/mlir_template.py | 46 +++++- .../mlir/passes/dma_fine_grained.py | 31 +++- .../mlir/passes/lower_transfer_to_gemmini.py | 104 +++++++++--- tests/lowering/__init__.py | 0 .../test_lower_transfer_to_gemmini.py | 119 ++++++++++++++ tests/lowering/test_masked_fill.py | 30 ++++ tests/ops/misc/test_indirect_access.py | 2 +- tests/ops/misc/test_masked_nondividing.py | 77 +++++++++ 15 files changed, 544 insertions(+), 101 deletions(-) create mode 100644 tests/lowering/__init__.py create mode 100644 tests/lowering/test_lower_transfer_to_gemmini.py create mode 100644 tests/lowering/test_masked_fill.py create mode 100644 tests/ops/misc/test_masked_nondividing.py diff --git a/.github/workflows/pytorchsim_test.yml b/.github/workflows/pytorchsim_test.yml index 345c716e..7cc76a5b 100644 --- a/.github/workflows/pytorchsim_test.yml +++ b/.github/workflows/pytorchsim_test.yml @@ -633,6 +633,24 @@ jobs: -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/misc/test_indirect_access.py + test_masked_nondividing: + name: Run test_masked_nondividing + runs-on: ubuntu-latest + steps: + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Run test_masked_nondividing.py + run: | + echo "Running test_masked_nondividing.py" + docker run --rm \ + -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ + ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/misc/test_masked_nondividing.py + test_scheduler: name: Run test_scheduler runs-on: ubuntu-latest diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index ee7d9f7a..5778213b 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -393,12 +393,21 @@ def reset(self, reason): self.__init__(self.kernel_group, reason=reason) self.exit_stack, self._nested_context_depth = save + @staticmethod + def _origin_is_exp(o): + """True iff the FX origin node is an aten.exp op -- matched by the op target, not a + name substring (so it does not fire on expand / expm1 / experimental).""" + t = getattr(o, "target", None) + pkt = getattr(t, "_overloadpacket", None) + name = getattr(pkt, "__name__", None) or getattr(t, "__name__", None) or "" + return name == "exp" + # padding type 0: zero-padding 1: negative-padding(-inf) ... def get_padding_type(self): ops = self.current_node.node.origins if self.current_node.is_reduction(): for op in ops: - if "exp" in op.name: # exponential reduciton case + if self._origin_is_exp(op): # exponential reduciton case return 1 # for op in ops: # TODO: padding has some problem in the case of max_pool # if "max_pool" in op.args[0].name: @@ -577,7 +586,7 @@ def load(self, name: str, index: sympy.Expr): mlir_dtype = mlir_common.DTYPE_TO_MLIR[dtype] # Extract sram info - local_tile_desc, index_var, dram_stride = self.get_dma_info(name, index) + local_tile_desc, index_var, dram_stride, local_dims = self.get_dma_info(name, index) vlane_split_axis = local_tile_desc.vmap.vlane_split_axis vlane_stride = local_tile_desc.vmap.vlane_stride tile_numel_per_lane = local_tile_desc.get_numel_per_lane() @@ -591,8 +600,10 @@ def load(self, name: str, index: sympy.Expr): sram_var, sram_index_var = self.get_scratchpad_buffer(dtype, name, local_tile_desc, index) compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) + masked_bounds = self._masked_bounds(name, index, dram_stride, local_tile_desc, is_load=True, buffer=self.dma_loads, local_dims=local_dims) code = self.emit_transfer("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, int(padding), offset=offset_desc) + dram_shape, tile_shape, dram_stride, tile_stride, int(padding), offset=offset_desc, + masked_bounds=masked_bounds, masked_fill=self._masked_fill_bits(dtype, index)) self.cse.generate(self.dma_loads, code, assignment = False) # FIXME: assignment = False does not support caching with self.override_buffer_cse(buffer=self.loads): @@ -606,20 +617,22 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) offset_desc = None # Handle scatter store + accumulate = False if self._has_indirect(index): # Convert the output buffer type to the inplace buffer arg_name = V.graph.scheduler.mutation_real_name.get(name, name) if arg_name not in self.kernel_group.args.inplace_buffers: self.kernel_group.args.make_inplace(arg_name, arg_name) - if mode == "atomic_add": - loaded_value = ops.load(name, index) - value = ops.add(loaded_value, value) + # index_add: let the MVOUT do out[idx] += val. The DMA processes positions + # sequentially, so duplicate indices accumulate correctly -- unlike the compute + # gather-add-overwrite, which loses duplicates landing in the same tile. + accumulate = (mode == "atomic_add") index, offset_desc = self.convert_indirect_indexing(index) dram_var = self.kernel_group.args.output(name) # Prepare dma instruction - local_tile_desc, index_var, dram_stride = self.get_dma_info(name, index) + local_tile_desc, index_var, dram_stride, local_dims = self.get_dma_info(name, index) vlane_split_axis = local_tile_desc.vmap.vlane_split_axis vlane_stride = local_tile_desc.vmap.vlane_stride @@ -656,8 +669,10 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) sram_index_var = self.spad_buffer_dict[str(value)][3] # Generate DMA instruction + masked_bounds = self._masked_bounds(name, index, dram_stride, local_tile_desc, is_load=False, buffer=self.dma_stores, local_dims=local_dims) code = self.emit_transfer("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, 0, offset=offset_desc) + dram_shape, tile_shape, dram_stride, tile_stride, 0, offset=offset_desc, masked_bounds=masked_bounds, + accumulate=accumulate, acc_float=dtype.is_floating_point) self.dma_stores.writeline(common.DeferredLine(name, code)) def reduction(self, dtype, src_dtype, reduction_type, value): @@ -756,7 +771,7 @@ def store_reduction(self, name, index, value): with self.override_buffer_cse(cse=self.reduction_cse): # Tile is always reuduced in inner loop - local_tile_desc, index_var, dram_stride = self.get_dma_info(name, index, broadcast=False, store_reduction=True, buffer=self.reductions_suffix) + local_tile_desc, index_var, dram_stride, _ = self.get_dma_info(name, index, broadcast=False, store_reduction=True, buffer=self.reductions_suffix) vlane_split_axis = local_tile_desc.vmap.vlane_split_axis vlane_stride = local_tile_desc.vmap.vlane_stride @@ -796,14 +811,6 @@ def _has_indirect(self, expr): return any(s.name in self.indirect_symbols for s in expr.free_symbols) def _index_expr(self, tile_desc, renamed_expression, index, base_vector_index): - # In case of index expr, dimension size should be divisible by tile size - if not self.kernel_group.tile_desc.is_dim_dividable(self.ranges): - new_tile_size = self.kernel_group.tile_desc.adjust_tile_to_divisible(self.ranges) - prior_tile_size, prior_ranges = self.kernel_group.tile_desc.get_tile_size(), self.ranges - self.kernel_group.tile_desc.set_tile_size(new_tile_size) - self.reset("recompile") - raise mlir_common.RecompileSignal(f"Index access (tile size {prior_tile_size} is not divisible by {prior_ranges})") - tile_size_per_lane = tile_desc.get_tile_size_per_lane() compute_vec_size = tile_desc.get_compute_vec_size() strides = tile_desc.get_tile_stride_per_lane() @@ -1358,12 +1365,96 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe if len(self.itervars) == 1 and self.reduction_depth == 0: # In case of reduction loop only case, we will add dummy loop so shift it once dram_stride = [0] + dram_stride[:-1] - return local_tile_desc, index_var, dram_stride + + # Return the tile-axis -> loop-dim map (local_dims) so load()/store() can pass it to + # _masked_bounds for the per-dim [low, high) clamp (it needs each tile axis' loop iv). + return local_tile_desc, index_var, dram_stride, local_dims + + _FILL_BITVIEW = {torch.float32: torch.int32, torch.float16: torch.int16, + torch.bfloat16: torch.int16, torch.float64: torch.int64} + + def _masked_fill_bits(self, dtype, index): + """Raw bits for the masked-DMA tail fill = the consuming reduction's identity + (reduction_init; sum->0, max->-inf, ...) in the LOAD dtype, 0 for a non-reduction + load. Log-sum-exp exception: a sum reducing exp(input) fills its PRIMARY input's + tail with -inf (exp(-inf)=0); broadcast operands keep their finite identity.""" + node = getattr(self, "current_node", None) + if node is None or getattr(node, "node", None) is None or not node.node.get_reduction_type(): + return 0 + rtype = node.node.get_reduction_type() + is_primary = bool(set(self.itervars[self.reduction_depth:]) & index.free_symbols) + if rtype == "sum" and is_primary and any(self._origin_is_exp(o) for o in node.node.origins): + init = "-inf" + else: + init = reduction_init(rtype, dtype) + val = {"-inf": float("-inf"), "inf": float("inf")}.get(init, init) + if isinstance(val, str): # e.g. welford_reduce -> "0.0" + val = float(val) + t = torch.tensor(val, dtype=dtype) + view = self._FILL_BITVIEW.get(dtype) + bits = int(t.view(view).item()) if view is not None else int(t.item()) + return bits & ((1 << (t.element_size() * 8)) - 1) + + def _masked_bounds(self, name, index, dram_stride, local_tile_desc, is_load, buffer, local_dims): + """Per tile-axis [low, high) clamp for a masked DMA. Per axis the valid GLOBAL range + is [glo, ghi) (store: [0, output_extent); padded load: [pad, pad + input_extent)), and + _emit_clamp turns it into the tile-local low/high SSA vars. Returns [(tile_axis, low_var, high_var), ...]. + """ + tile_size = local_tile_desc.get_tile_size() + const = int(index.as_coeff_Add()[0]) if not index.is_number else int(index) + # index const = -sum(pad_d * dram_stride[d]); recover per-dim pad greedily (desc stride). + pad = [0] * len(tile_size) + if is_load and const < 0: + rem = -const + for d in sorted(range(len(dram_stride)), key=lambda x: -dram_stride[x]): + s = dram_stride[d] + if s > 0 and d < len(pad): + pad[d] = rem // s + rem -= pad[d] * s + in_shape, in_stride = self.buffer_types[name][2], self.buffer_types[name][3] + axes = [] + for d, k in enumerate(local_dims): + if d >= len(tile_size) or k >= len(self.ranges): + continue + iv = str(self.itervars[k]) + # A padded load (const < 0) reads a smaller input: clamp per-dim to + # [pad_d, pad_d + input_extent_d). Every other DMA is valid over the whole + # loop extent, so only the trailing tail needs clamping. + if is_load and const < 0: + ext = next((int(in_shape[j]) for j, st in enumerate(in_stride) + if int(st) == int(dram_stride[d])), None) + glo, ghi = (pad[d], pad[d] + ext) if ext is not None else (0, int(self.ranges[k])) + else: + glo, ghi = 0, int(self.ranges[k]) + axes.append((d, iv, glo, ghi, int(tile_size[d]))) + return self._emit_clamp(axes, buffer) + + def _emit_clamp(self, axes, buffer): + """Emit the per-axis dynamic clamp. axes: [(tile_axis, base_iv, glo, ghi, tile), ...] + -- low = max(0, glo - base), high = min(tile, ghi - base) as affine.max/affine.min of + the loop iv so the last partial tile and the pad borders fall out per iteration. + Returns [(tile_axis, low_var, high_var), ...] for the non-trivial axes only.""" + result = [] + for d, iv, glo, ghi, tile in axes: + if glo == 0 and ghi % tile == 0: # every tile fully valid -> no clamp + continue + high_var = self.apply_cse.generate( + buffer, f"affine.min affine_map<(d0) -> ({tile}, {ghi} - d0)>(%{iv})") + self.register_var_info(high_var, [1, "index"]) + if glo > 0: + low_var = self.apply_cse.generate( + buffer, f"affine.max affine_map<(d0) -> (0, {glo} - d0)>(%{iv})") + self.register_var_info(low_var, [1, "index"]) + else: + low_var = self.get_const_cse(0) + result.append((d, low_var, high_var)) + return result def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtype, dram_var, dram_index_var, sram_var, sram_index_var, dram_shape, tile_shape, dram_stride, tile_stride, padding, - subtile_size=None, async_type=None, offset=None): + subtile_size=None, async_type=None, offset=None, masked_bounds=None, masked_fill=0, + accumulate=False, acc_float=False): """Emit a generic togsim.transfer op for a DMA whose access exceeds the 4D Gemmini descriptor limit. Carries the full N-D access (dram/tile strides + shapes) plus the SSA operands a memref.dma_start needs @@ -1404,6 +1495,10 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp if subtile_size: av = int(async_type) if async_type is not None else 1 attrs += f', subtile_size = {list(subtile_size)}, async = {av} : i64' + if accumulate: # index_add: MVOUT does out[idx] += val (float or integer add) + attrs += f', accumulate = true' + if acc_float: + attrs += f', acc_float = true' # operands: dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vlane_stride [, offset spad] operands = (f'%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' f'%{tag}, %{zero_cse}, %{dma_type}, %{vst}') @@ -1413,6 +1508,18 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp operands += f', %{offset_buf}' optypes += f', {offset_type}' attrs += f', indirect = true, offset_stride = {int(offset_stride)} : i64' + # masked-DMA dynamic clamp: append (low, high) index operands per clamped tile axis; + # masked_axes names the tile axis of each pair so the lowering writes the runtime + # values into the descriptor's dim_low/dim_high before the DMA. See _masked_bounds. + if masked_bounds: + axes = [d for d, _lo, _hi in masked_bounds] + for _d, lo, hi in masked_bounds: + operands += f', %{lo}, %{hi}' + optypes += ', index, index' + attrs += f', masked_axes = {axes}' + # box-excluded positions are filled with the consuming reduction's identity + # (0/1/-inf/+inf, per dtype); 0 for non-reduction loads. See _masked_fill_bits. + attrs += f', masked_fill = {int(masked_fill)} : i64' return f'"togsim.transfer"({operands}) {{{attrs}}} : ({optypes}) -> ()' def allocate_sram_buffer(self, dtype, dram_name, tile_desc, raw_index, buffer=None, forced_name=None): @@ -1497,13 +1604,6 @@ def convert_indirect_indexing(self, index :sympy.Expr): if not self._has_indirect(index): return index, None - # Note: In case of indirect indexing, dimensions should be divisible by tile size - if not self.kernel_group.tile_desc.is_dim_dividable(self.ranges): - new_tile_size = self.kernel_group.tile_desc.adjust_tile_to_divisible(self.ranges) - self.kernel_group.tile_desc.set_tile_size(new_tile_size) - self.reset("recompile") - raise mlir_common.RecompileSignal(f"Indirect access (tile size {self.kernel_group.tile_desc.get_tile_size()} is not divisible by {self.ranges})") - # Process start indirect_dims = [str(dim) for dim in index.free_symbols if str(dim) in self.indirect_symbols] indirect_dims.sort() diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index a70d1c7d..99f22762 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -50,7 +50,8 @@ torch.int8: "i8", torch.uint8: "i8", torch.bool: "i8", - torch.bfloat16: "bf16", + torch.bfloat16: "bf16", # present only to keep the dtype table total; bf16 is not + # actually supported (Spike has no bf16 and op-select rejects it) } MLIR_TO_DTYPE = { @@ -324,41 +325,6 @@ def apply_divisor(self, axis: int, divisor: int, mode: str = "split"): raise ValueError(f"Unknown mode: {mode}. Supported: 'pad', 'split'.") - def is_dim_dividable(self, dim_sizes: list[int]) -> bool: - if len(dim_sizes) != len(self._tile_size): - raise ValueError("dim_sizes must match the tile size dimensions") - - dim_sizes_cpy = list(dim_sizes) - axis, stride = self.vmap.vlane_split_axis, self.vmap.vlane_stride - remain = dim_sizes_cpy[axis] % stride - if remain: - dim_sizes_cpy[axis] += stride - remain - - return all(d % t == 0 for d, t in zip(dim_sizes_cpy, self._tile_size)) - - def adjust_tile_to_divisible(self, dim_sizes: list[int]) -> list[int]: - """Adjust current tile to be divisible by given dimensions.""" - if len(dim_sizes) != len(self._tile_size): - raise ValueError("dim_sizes must match the tile size dimensions") - - def _adjust_one(dim_size, tile_size): - for candidate in range(tile_size, 0, -1): - if dim_size % candidate == 0: - return candidate - return 1 - - candidate_tile_size = [_adjust_one(d, t) for d, t in zip(dim_sizes, self._tile_size)] - for i in range(len(candidate_tile_size)): - self.tile_constraint[i].must_divide_dim = True - - axis, stride = self.vmap.vlane_split_axis, self.vmap.vlane_stride - remain = candidate_tile_size[axis] % stride - - if remain: - # #201: relax vlane_stride constraints - self.vmap.vlane_stride = 1 - return candidate_tile_size - def scale_tile_dim(self, axis, dim_sz, scale_factor=2): axis_constrinat = self.tile_constraint[axis] current_sz = self._tile_size[axis] @@ -395,8 +361,6 @@ def trim_large_tail(self, ranges: list[int]): constraint = self.tile_constraint[i] if constraint.fixed: continue - elif constraint.must_divide_dim: - BETA = 0 padding_ratio = TileAdjustMixin.get_padding_ratio(tile_range, dim_range) if padding_ratio < self.tail_ratio_threshold: @@ -469,23 +433,14 @@ def get_padding_ratio(tile_range: int, dim_range: int) -> float: @dataclass class TileConstraint: multiple_of: int = 1 - must_divide_dim: bool = False fixed: bool = False def adjust(self, old: int, new: int, dim: int) -> int: if self.fixed: return old # Fixed tile size - tail = new % self.multiple_of - new -= tail - if not self.must_divide_dim: - return max(new, self.multiple_of) - - while new > 0: - if dim % new == 0: - return new - new -= self.multiple_of - raise extension_codecache.TileSizeError("Cannot find suitable tile size under the given constraints.") + new -= new % self.multiple_of + return max(new, self.multiple_of) class MLIRMultiDimTile(TileAdjustMixin): def __init__(self, tile_size, vector_lane, vlane_split_axis=None, vlane_stride=None, forced_vec_size=None): diff --git a/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py b/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py index 8b8288a8..7964da0f 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py @@ -143,6 +143,10 @@ def render(self, TOG_latency = O_W if TILE_M > O_W else TILE_M TOG_latency = 8 if TOG_latency < 8 else TOG_latency kernel.loop_size = [TOG_latency, TILE_N, TILE_K] + # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op). + # This template fuses channel and kernel-width into tile_k (loops 0..I_C*K_W). + kernel.loop_extents = {"tile_m": BATCH, "tile_n": O_C, "o_h": O_H, "o_w": O_W, + "k_h": K_H, "tile_k": I_C * K_W} # Prepare tile descriptors vlane_stride = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py b/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py index 92efff66..dfca23ec 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py @@ -144,6 +144,9 @@ def render(self, TOG_latency = O_W if TILE_M > O_W else TILE_M TOG_latency = 8 if TOG_latency < 8 else TOG_latency kernel.loop_size = [TOG_latency, TILE_N, TILE_K] + # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op). + kernel.loop_extents = {"tile_n": O_C, "o_h": O_H, "tile_m": O_W, + "k_h": K_H, "k_w": K_W, "tile_k": I_C} # Prepare tile descriptors vlane_stride = 1 vlane_split_axis = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py b/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py index dfd418d9..f1a42964 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py @@ -144,6 +144,9 @@ def render(self, TOG_latency = O_W if TILE_M > O_W else TILE_M TOG_latency = 8 if TOG_latency < 8 else TOG_latency kernel.loop_size = [TOG_latency, TILE_N, TILE_K] + # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op). + kernel.loop_extents = {"tile_n": O_C, "o_h": O_H, "tile_m": O_W, + "k_h": K_H, "k_w": K_W, "tile_k": I_C} # Prepare tile descriptors vlane_stride = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_conv_template.py b/PyTorchSimFrontend/mlir/mlir_conv_template.py index 178ba7c6..8bb64d48 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_template.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_template.py @@ -147,6 +147,9 @@ def render(self, TOG_latency = BATCH if TILE_M > BATCH else TILE_M TOG_latency = 8 if TOG_latency < 8 else TOG_latency kernel.loop_size = [TOG_latency, TILE_N, TILE_K] + # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op). + kernel.loop_extents = {"tile_m": BATCH, "tile_n": O_C, "o_h": O_H, "o_w": O_W, + "k_h": K_H, "k_w": K_W, "tile_k": I_C} # Prepare tile descriptors vlane_stride = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index cd3ca018..051e69b6 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -661,7 +661,6 @@ def template_store(): with contextlib.ExitStack() as stack: stack.enter_context(compute_body.indent(attribute="{inner_loop=false}",suffix=self.compute_body_loop.epilogue_line())) if self.reduction_fusion: - compute_body.splice(self.masks) compute_body.writelines(self.reduction_body_loop.lines()) stack.enter_context(compute_body.indent(attribute="{inner_loop=false}")) compute_body.splice(self.loads) @@ -940,9 +939,35 @@ def generate_dma_code(): zero_cse = self.get_const_cse(0, "index") sram_index_var = ", ".join([f"%{str(zero_cse)}"]*tile_desc.get_nr_dim()) + # masked-DMA clamp: per tile dim, clamp to the real DRAM extent so a tile + # padding past it is zero-filled, not MAC-ing garbage. Only a dim indexed by + # a SINGLE iv is clampable; a sum of ivs (im2col) is pre-padded instead. + + # FIXME: loop_extents is declared by hand in each conv template and drifts + # from the affine.for bounds. Record each loop's (iv -> bound) as the + # affine.for is emitted, as pointwise does with ranges/itervars. + loop_ext = getattr(self, "loop_extents", None) + layout_sizes, layout_strides = node_layout.size, node_layout.stride + tile_sizes = tile_desc.get_tile_size() + clamp_axes = [] + for d, idx_expr in enumerate(index_list): + syms = list(idx_expr.free_symbols) + if len(syms) != 1 or d >= len(tile_sizes): + continue + base_iv = str(syms[0]) + if loop_ext: + extent = loop_ext.get(base_iv) # explicit: clamp only known ivs + else: + extent = next((int(layout_sizes[j]) for j, st in enumerate(layout_strides) + if layout_sizes[j].is_number and int(st) == int(_dram_stride[d])), None) + if extent is not None: + clamp_axes.append((d, base_iv, 0, int(extent), int(tile_sizes[d]))) + masked_bounds = self._emit_clamp(clamp_axes, local_code) + code = self.emit_transfer(dma_type, vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, dram_shape, tile_shape, _dram_stride, sram_strides, int(padding), - subtile_size=subtile_size if subtile_size else None, async_type=async_type) + subtile_size=subtile_size if subtile_size else None, async_type=async_type, + masked_bounds=masked_bounds, masked_fill=0) local_code.writeline(code) return textwrap.indent(local_code.getvalue(), " "*indent_size).strip() @@ -1075,9 +1100,22 @@ def store_epilogue(self, name: str, index: sympy.Expr, value, *args, **kwargs): with self.override_buffer_cse(buffer=self.stores): ops._store(value, sram_var, compute_index_var, tile_shape, buffer_name=buffer_name) - # Generate DMA instruction + # Generate the DMA. Clamp the output tail, else a non-dividing epilogue store + # writes the systolic pad region past the real output extent. Each iv's extent is + # ranges[k]; _emit_clamp skips dividing dims, so a dividing output is a no-op. + iv_extent = {} + for itervar, rng in zip(self.itervars, self.ranges): + try: + iv_extent[str(itervar)] = int(rng) + except (TypeError, ValueError): + pass + tile_sizes = self.kernel_group.tile_desc.get_tile_size() + clamp_axes = [(d, iv, 0, iv_extent[iv], int(tile_sizes[d])) + for d, iv in enumerate(self.dim_aliasing.values()) + if d < len(tile_sizes) and iv in iv_extent] + masked_bounds = self._emit_clamp(clamp_axes, self.dma_stores) code = self.emit_transfer("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, 0) + dram_shape, tile_shape, dram_stride, tile_stride, 0, masked_bounds=masked_bounds) self.dma_stores.writeline(DeferredLine(name, code)) def reduction_epilogue(self, dtype, src_dtype, reduction_type, value): diff --git a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py index 22eb999e..b9b482d6 100644 --- a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py +++ b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py @@ -85,6 +85,14 @@ def __init__(self, op): self.tag_idx = operands[5] self.num_elements = operands[6] # = dma_type const operand self.num_elements_per_stride = operands[7] # = vlane_stride (vst) + # trailing operands: [offset (indirect)] then (low, high) per masked axis. + _extra = operands[8:] + self.offset = None + if "indirect" in op.attributes: + self.offset = _extra[0] + _extra = _extra[1:] + self.masked_axes = attr_i64_array(op, "masked_axes", default=[]) + self.masked_ops = _extra # low0, high0, low1, high1, ... self.sram_rank = len(ir.MemRefType(self.sram.type).shape) # Direction: MVIN reads dram -> sram; MVOUT writes sram -> dram. @@ -221,7 +229,8 @@ def _dma_attrs(dma): attrs = {} op = dma.op for k in ("dma_kind", "vlane_split_axis", "dram_stride", "tile_stride", - "subtile_size", "padding"): + "subtile_size", "padding", "masked_axes", "masked_fill", "indirect", + "offset_stride", "accumulate", "acc_float"): if k in op.attributes: attrs[k] = op.attributes[k] attrs["async"] = ir.BoolAttr.get(dma.is_async()) @@ -229,6 +238,18 @@ def _dma_attrs(dma): return attrs +def _remap_bound(bound, iv, sub, is_high, ip): + """The masked low/high are full-tile-local; shift to this subtile: subtile position + p maps to full-tile p + iv*sub, so subtile-local high = min(sub, high - iv*sub) and + low = max(0, low - iv*sub). Out-of-window (neg high / low>sub) -> Spike skips all.""" + from mlir.dialects import affine + d0, d1 = ir.AffineDimExpr.get(0), ir.AffineDimExpr.get(1) + edge = ir.AffineConstantExpr.get(sub if is_high else 0) + m = ir.AffineMap.get(2, 0, [edge, d0 - d1 * sub]) + op = affine.AffineMinOp if is_high else affine.AffineMaxOp + return op(m, [bound, iv], ip=ip).result + + def _emit_dma(dma, ivs, vectorlane, ip): """Emit one fine-grained togsim.transfer at `ip`, indexed by `ivs` (the fused induction vars for this DMA's dims, in dim order).""" @@ -243,6 +264,14 @@ def _emit_dma(dma, ivs, vectorlane, ip): operands = [dma.dram, dram_idx, dma.sram, sram_off, dma.tag, tag_idx, dma.num_elements, dma.num_elements_per_stride] + if dma.offset is not None: + operands.append(dma.offset) + # masked low/high are full-tile-local -> remap each per THIS subtile's offset. + sub = dma.subtile_size() + for i, axis in enumerate(dma.masked_axes): + low, high = dma.masked_ops[2 * i], dma.masked_ops[2 * i + 1] + operands.append(_remap_bound(low, ivs[axis], sub[axis], False, ip)) + operands.append(_remap_bound(high, ivs[axis], sub[axis], True, ip)) ir.Operation.create("togsim.transfer", results=[], operands=operands, attributes=_dma_attrs(dma), ip=ip) diff --git a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py index 46575ba7..f3599000 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py @@ -2,7 +2,7 @@ Merges decompose_transfer (the <=4D Gemmini-limit handling: drop unit dims / collapse / >4D affine.for peel with lane-banked SRAM offset) with -lower_dma_to_gemmini (the CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm emission). +lower_dma_to_gemmini (the CONFIG_DESC/MVIN|MVOUT asm emission). togsim.transfer is unregistered so it carries every runtime descriptor as an operand -- including the future masked-clamp low/high vectors -- which a registered memref.dma_start cannot. @@ -19,7 +19,7 @@ from .lower_dma_to_gemmini import _i64_signed, _row_major_strides, _elem_bytes, _asm, CONSTRAINTS from .decompose_transfer import _int_array, _const_int, _squeeze_reassociation -CONFIG, CONFIG2, CONFIG3, CONFIG4, CONFIG_DESC = 0, 4, 5, 6, 7 +CONFIG_DESC = 7 # func7 for the TMA-style descriptor pointer (replaces CONFIG/2/3/4) MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 DESC_SLOTS = 18 # 144-byte DMA descriptor as 18 x i64 (see project-dma-descriptor) CONFIG_TYPE = {MVIN: 0, MVIN2: 1, MVIN3: 2, MVOUT: 3} @@ -42,15 +42,23 @@ def run(module, timing=False, vectorlane=128, **_): if g.operation.name == "memref.global": sym2type[g.attributes["sym_name"].value] = MemRefType(TypeAttr(g.attributes["type"]).value) + i32 = IntegerType.get_signless(32) desc_ty = MemRefType.get([DESC_SLOTS], i64) + desc_i32_ty = MemRefType.get([DESC_SLOTS * 2], i32) # same bytes, i32-addressable (masked) desc_globals = {} # slots tuple -> global sym name (dedup) desc_counter = [0] + def _i32_signed(v): + v &= 0xFFFFFFFF + return v - (1 << 32) if v >= (1 << 31) else v + def _pack_desc(shape4, dram4, spad4, elem_bytes, vlstride, vsa4, cfg_type, indirect, - ind_stride, ind_esize): + ind_stride, ind_esize, hi=None, accumulate=False, acc_float=False): # 18 i64 slots matching the C dma_descriptor byte layout (little-endian). lo = [0, 0, 0, 0] - hi = list(shape4) + if hi is None: + hi = list(shape4) + masked = any(h < s for h, s in zip(hi, shape4)) # dim_high clamps below extent def s2(a, b): return (a & 0xFFFFFFFF) | ((b & 0xFFFFFFFF) << 32) m = 0xFFFFFFFFFFFFFFFF slots = [ @@ -60,7 +68,9 @@ def s2(a, b): return (a & 0xFFFFFFFF) | ((b & 0xFFFFFFFF) << 32) dram4[0] & m, dram4[1] & m, dram4[2] & m, dram4[3] & m, spad4[0] & m, spad4[1] & m, spad4[2] & m, spad4[3] & m, (elem_bytes & 0xFFFF) | ((vlstride & 0xFFFF) << 16) | ((vsa4 & 0xFF) << 32) - | ((cfg_type & 0xFF) << 40) | (((1 if indirect else 0) & 0xFFFF) << 48), + | ((cfg_type & 0xFF) << 40) + | (((1 if indirect else 0) | (2 if masked else 0) + | (4 if accumulate else 0) | (8 if acc_float else 0)) << 48), 0, # +120 indirect_addr (runtime) (ind_stride & 0xFFFF) | ((ind_esize & 0xFFFF) << 16), # +128 0, # +136 fill (runtime/step3) @@ -85,14 +95,46 @@ def _desc_global(slots): desc_globals[key] = name return desc_globals[key] - def desc_ptr_and_store_indirect(slots, ind_addr_i64): - # get_global -> byte pointer; for indirect, store the runtime addr into slot 15. - name = _desc_global(slots) - g = memref.GetGlobalOp(desc_ty, name).result - if ind_addr_i64 is not None: - memref.StoreOp(ind_addr_i64, g, [arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, 15)).result]) - base = memref.ExtractAlignedPointerAsIndexOp(g).result - return arith.IndexCastOp(i64, base).result + def desc_ptr(slots, masked_pairs=(), fill=0, ind_addr_i64=None): + """Byte pointer to the DMA descriptor global. A fully static descriptor (no + masked clamp, no indirect address) is a deduped i64 global. Anything with a + runtime-written field -- per-dim dim_low/dim_high (masked) and/or the indirect + address -- gets a UNIQUE i32-view global (same byte layout, i32-addressable) so + those int32/i64 fields can be stored individually.""" + import numpy as np + if not masked_pairs and ind_addr_i64 is None: + g = memref.GetGlobalOp(desc_ty, _desc_global(slots)).result + return arith.IndexCastOp(i64, memref.ExtractAlignedPointerAsIndexOp(g).result).result + slots = list(slots) + if masked_pairs: + slots[14] |= (2 << 48) # masked flag (+118 bit1) + slots[17] = fill & 0xFFFFFFFFFFFFFFFF # +136 fill: box-excluded positions + words = [] + for v in slots: + v &= 0xFFFFFFFFFFFFFFFF + words.append(_i32_signed(v & 0xFFFFFFFF)) + words.append(_i32_signed((v >> 32) & 0xFFFFFFFF)) + name = f"dma_desc_{desc_counter[0]}" + desc_counter[0] += 1 + init = DenseElementsAttr.get(np.array(words, dtype=np.int32), + type=RankedTensorType.get([DESC_SLOTS * 2], i32)) + with InsertionPoint.at_block_begin(module_top): + Operation.create("memref.global", results=[], operands=[], attributes={ + "sym_name": StringAttr.get(name), + "sym_visibility": StringAttr.get("private"), + "type": TypeAttr.get(desc_i32_ty), + "initial_value": init}) + g = memref.GetGlobalOp(desc_i32_ty, name).result + def store_word(v32, i32_idx): + memref.StoreOp(v32, g, [arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, i32_idx)).result]) + for axis4, low_val, high_val in masked_pairs: # dim_low i32 idx 4 (+16B), dim_high 8 (+32B) + store_word(arith.IndexCastOp(i32, low_val).result, 4 + axis4) + store_word(arith.IndexCastOp(i32, high_val).result, 8 + axis4) + if ind_addr_i64 is not None: # indirect_addr (i64) at +120 -> words 30/31 + store_word(arith.TruncIOp(i32, ind_addr_i64).result, 30) + shamt = arith.ConstantOp(i64, IntegerAttr.get(i64, 32)).result + store_word(arith.TruncIOp(i32, arith.ShRUIOp(ind_addr_i64, shamt).result).result, 31) + return arith.IndexCastOp(i64, memref.ExtractAlignedPointerAsIndexOp(g).result).result def i64_const(value): return arith.ConstantOp(i64, IntegerAttr.get(i64, _i64_signed(value))).result @@ -144,6 +186,20 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): subtile = _int_array(op.attributes["subtile_size"]) except KeyError: subtile = None + # masked-DMA dynamic clamp: masked_axes lists the clamped tile axes; the trailing + # (low, high) index operands (2 per axis) are runtime-stored into dim_low/dim_high. + if "masked_axes" in op.attributes: + masked_axes = _int_array(op.attributes["masked_axes"]) + n_base = 9 if "indirect" in op.attributes else 8 + mvals = op_operands[n_base:] + masked_pairs = [(masked_axes[i], mvals[2 * i], mvals[2 * i + 1]) + for i in range(len(masked_axes))] + masked_fill = IntegerAttr(op.attributes["masked_fill"]).value + else: + masked_pairs = [] + masked_fill = 0 + accumulate = "accumulate" in op.attributes # index_add: MVOUT out[idx] += val + acc_float = "acc_float" in op.attributes if timing: op.erase() @@ -161,7 +217,7 @@ def _const(v): return arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, v)).result def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, - desc_dram_strides, desc_spad_strides, subtile_shape): + desc_dram_strides, desc_spad_strides, subtile_shape, desc_masked_pairs=None): cfg_shape = subtile_shape if subtile_shape is not None else desc_shape expand = MAX_TENSOR_DIM - len(cfg_shape) shape4 = [1] * expand + list(cfg_shape) @@ -182,9 +238,11 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, ind_esize = _elem_bytes(off_ty.element_type) ind_stride = IntegerAttr(op.attributes["offset_stride"]).value slots = _pack_desc(shape4, dram4, spad4, elem_bytes, vlane_stride, vsa4, - config_type, indirect, ind_stride, ind_esize) - desc_ptr = desc_ptr_and_store_indirect(slots, ind_addr) - asm(CONFIG_DESC, desc_ptr, i64_const(0)) + config_type, indirect, ind_stride, ind_esize, + accumulate=accumulate, acc_float=acc_float) + pairs4 = [(expand + d, lo, hi) for d, lo, hi in (desc_masked_pairs or ())] # 4D-expand axis + desc_ptr_val = desc_ptr(slots, pairs4, masked_fill, ind_addr) + asm(CONFIG_DESC, desc_ptr_val, i64_const(0)) asm(dma_type_val, dram_addr, spad_addr) if offset_sym is not None: @@ -195,7 +253,7 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, # sram offset is linear (row-major stride 1) -> last index only; others 0. sidx = [_const(0)] * (len(tile_shape) - 1) + [sram_idx] _emit_asm(sram, sidx, dram_idx, vlane_axis, - tile_shape, dram_stride, tile_stride, subtile) + tile_shape, dram_stride, tile_stride, subtile, masked_pairs) op.erase() continue @@ -208,13 +266,16 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, dr = [dram_stride[i] for i in keep] tl = [tile_stride[i] for i in keep] st = [subtile[i] for i in keep] if subtile is not None else None + # remap each masked tile axis to its collapsed group index. + mp = [(next(gi for gi, g in enumerate(groups) if d in g), lo, hi) + for d, lo, hi in masked_pairs] new_vlane = next(gi for gi, g in enumerate(groups) if vlane_axis in g) with InsertionPoint(op): sram_c = Operation.create( "memref.collapse_shape", results=[collapsed_ty], operands=[sram], attributes={"reassociation": reassoc}).results[0] sidx = [_const(0)] * (len(target) - 1) + [sram_idx] - _emit_asm(sram_c, sidx, dram_idx, new_vlane, target, dr, tl, st) + _emit_asm(sram_c, sidx, dram_idx, new_vlane, target, dr, tl, st, mp) op.erase() continue @@ -226,6 +287,9 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, dr = [dram_stride[d] for d in inner] tl = [tile_stride[d] for d in inner] st = [subtile[d] for d in inner] if subtile is not None else None + if any(d in peeled for d, _lo, _hi in masked_pairs): + raise NotImplementedError("masked-DMA clamp on a peeled (outer-loop) axis") + mp = [(inner.index(d), lo, hi) for d, lo, hi in masked_pairs if d in inner] if vlane_axis in inner: new_vlane = inner.index(vlane_axis) elif vlane_axis in peeled: @@ -285,7 +349,7 @@ def _phys(d): ).results[0] zero = _const(0) _emit_asm(sub, [zero, zero, zero, sram_off_val], dram_idx_val, new_vlane, - inner_shape, dr, tl, st) + inner_shape, dr, tl, st, mp) op.erase() diff --git a/tests/lowering/__init__.py b/tests/lowering/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/lowering/test_lower_transfer_to_gemmini.py b/tests/lowering/test_lower_transfer_to_gemmini.py new file mode 100644 index 00000000..969bc02d --- /dev/null +++ b/tests/lowering/test_lower_transfer_to_gemmini.py @@ -0,0 +1,119 @@ +"""Unit tests for the lower_transfer_to_gemmini pass (no torch, no Spike). + +Hand-write a `togsim.transfer`, run the pass in-process, and assert the lowered IR +matches an embedded golden -- the exact expected output (MLIR SSA numbering is +deterministic for a fixed input). If a lowering change is intended, regenerate the +golden and review the diff. Skipped without the MLIR bindings. +""" +import importlib.util +import textwrap + +import pytest + +_MLIR = importlib.util.find_spec("mlir") is not None +pytestmark = pytest.mark.skipif(not _MLIR, reason="MLIR Python bindings not installed") + + +def _lower_transfer(ir_text, timing=False): + from mlir.ir import Context, Module, Location + from PyTorchSimFrontend.mlir.passes import lower_transfer_to_gemmini as L + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + m = Module.parse(ir_text) + L.run(m, timing=timing) + return str(m).strip() + + +# INPUT: one masked MVOUT of a [1,2,8,8] tile, with positional togsim.transfer operands +# (see emit_transfer). masked_axes=[2] is a SPARSE overlay clamping tile axis 2 to +# [0, 7), so the ragged tail is skipped on store; unlisted axes default to no clamp. +_MASKED_MVOUT = """ +module { + memref.global @buf1_spad : memref<1x2x8x8xf32, 1> + func.func @k(%arg1: memref<2048xf32>) { + %c0 = arith.constant 0 : index + %c7 = arith.constant 7 : index + %alloc = memref.alloc() : memref<1xi32> + %s = memref.get_global @buf1_spad : memref<1x2x8x8xf32, 1> + %c3 = arith.constant 3 : index + %c1 = arith.constant 1 : index + "togsim.transfer"(%arg1, %c0, %s, %c0, %alloc, %c0, %c3, %c1, %c0, %c7) {dma_kind="MVOUT", dram_stride=[128,64,8,1], tile_stride=[128,64,8,1], vlane_split_axis=3:i64, masked_axes=[2], masked_fill=0:i64} : (memref<2048xf32>, index, memref<1x2x8x8xf32,1>, index, memref<1xi32>, index, index, index, index, index) -> () + return + } +}""" + +# OUTPUT (golden): a 36-i32 @dma_desc_0 (dim_size, dim_low, dim_high, dram/tile strides, +# a packed flags slot), the masked clamp written into it at RUN TIME, then two +# `.insn r CUSTOM_1` ops: func7=7 CONFIG_DESC hands over the descriptor, func7=3 MVOUT. +_MASKED_MVOUT_LOWERED = textwrap.dedent("""\ + module { + memref.global "private" @dma_desc_0 : memref<36xi32> = dense<[1, 2, 8, 8, 0, 0, 0, 0, 1, 2, 8, 8, 128, 0, 64, 0, 8, 0, 1, 0, 128, 0, 64, 0, 8, 0, 1, 0, 65540, 131843, 0, 0, 0, 0, 0, 0]> + memref.global @buf1_spad : memref<1x2x8x8xf32, 1> + func.func @k(%arg0: memref<2048xf32>) { + %c0 = arith.constant 0 : index + %c7 = arith.constant 7 : index + %alloc = memref.alloc() : memref<1xi32> + %0 = memref.get_global @buf1_spad : memref<1x2x8x8xf32, 1> + %c3 = arith.constant 3 : index + %c1 = arith.constant 1 : index + %c0_0 = arith.constant 0 : index + %intptr = memref.extract_aligned_pointer_as_index %arg0 : memref<2048xf32> -> index + %c4 = arith.constant 4 : index + %1 = arith.muli %c0, %c4 : index + %2 = arith.addi %intptr, %1 : index + %3 = arith.index_cast %2 : index to i64 + %intptr_1 = memref.extract_aligned_pointer_as_index %0 : memref<1x2x8x8xf32, 1> -> index + %c128 = arith.constant 128 : index + %4 = arith.muli %c0_0, %c128 : index + %c64 = arith.constant 64 : index + %5 = arith.muli %c0_0, %c64 : index + %6 = arith.addi %4, %5 : index + %c8 = arith.constant 8 : index + %7 = arith.muli %c0_0, %c8 : index + %8 = arith.addi %6, %7 : index + %9 = arith.addi %8, %c0 : index + %c4_2 = arith.constant 4 : index + %10 = arith.muli %9, %c4_2 : index + %11 = arith.addi %intptr_1, %10 : index + %12 = arith.index_cast %11 : index to i64 + %13 = memref.get_global @dma_desc_0 : memref<36xi32> + %14 = arith.index_cast %c0 : index to i32 + %c6 = arith.constant 6 : index + memref.store %14, %13[%c6] : memref<36xi32> + %15 = arith.index_cast %c7 : index to i32 + %c10 = arith.constant 10 : index + memref.store %15, %13[%c10] : memref<36xi32> + %intptr_3 = memref.extract_aligned_pointer_as_index %13 : memref<36xi32> -> index + %16 = arith.index_cast %intptr_3 : index to i64 + %c0_i64 = arith.constant 0 : i64 + llvm.inline_asm has_side_effects asm_dialect = att ".insn r CUSTOM_1, 0x3, 7, x0, $0, $1", "r,r,~{dirflag},~{fpsr},~{flags}" %16, %c0_i64 : (i64, i64) -> () + llvm.inline_asm has_side_effects asm_dialect = att ".insn r CUSTOM_1, 0x3, 3, x0, $0, $1", "r,r,~{dirflag},~{fpsr},~{flags}" %3, %12 : (i64, i64) -> () + return + } + }""") + + +def test_masked_transfer_lowering_golden(): + assert _lower_transfer(_MASKED_MVOUT) == _MASKED_MVOUT_LOWERED + + +# Same transfer with subtile_size=[1,1,4,8] (the H axis is subtiled 8 -> 4). The +# descriptor's dim_size is the SUBTILE (config) shape, not the full tile; the masked +# axis + clamp stores are unchanged (a 4D subtile -> expand 0 -> same idx 6/10). +_MASKED_MVOUT_SUBTILE = _MASKED_MVOUT.replace( + 'vlane_split_axis=3:i64,', 'vlane_split_axis=3:i64, subtile_size=[1,1,4,8],') + +_MASKED_MVOUT_SUBTILE_LOWERED = _MASKED_MVOUT_LOWERED.replace( + "dense<[1, 2, 8, 8, 0, 0, 0, 0, 1, 2, 8, 8,", + "dense<[1, 1, 4, 8, 0, 0, 0, 0, 1, 1, 4, 8,") + + +def test_masked_transfer_subtile_lowering_golden(): + assert _lower_transfer(_MASKED_MVOUT_SUBTILE) == _MASKED_MVOUT_SUBTILE_LOWERED + + +def test_timing_mode_erases_transfer(): + # In timing mode the TOG carries DMA timing -> the transfer + descriptor are erased. + out = _lower_transfer(_MASKED_MVOUT, timing=True) + assert "togsim.transfer" not in out and "dma_desc" not in out diff --git a/tests/lowering/test_masked_fill.py b/tests/lowering/test_masked_fill.py new file mode 100644 index 00000000..f2abf4af --- /dev/null +++ b/tests/lowering/test_masked_fill.py @@ -0,0 +1,30 @@ +"""Unit test for the masked-DMA fill's exp-reduction detection. + +A log-sum-exp reduction (softmax / log_softmax) fills its masked tail with -inf so that +exp(-inf) = 0; a plain sum fills with the sum identity 0. `MLIRKernel._origin_is_exp` +decides which, and it must match the op TARGET, not a name substring -- otherwise `expand` +/ `expm1` (whose names contain "exp") would wrongly trigger the -inf fill on an ordinary +non-dividing sum, corrupting it to -inf. Skipped if torch / the frontend is unavailable. +""" +import importlib.util + +import pytest + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("torch") is None, reason="torch not available") + + +def test_origin_is_exp_matches_target_not_name_substring(): + import torch + from types import SimpleNamespace as NS + from PyTorchSimFrontend.mlir.mlir_codegen_backend import MLIRKernel + is_exp = MLIRKernel._origin_is_exp + + # the genuine exp op -> True + assert is_exp(NS(target=torch.ops.aten.exp.default, name="exp")) is True + # names contain "exp" but are NOT exp -> must be False (the substring bug this guards) + assert is_exp(NS(target=torch.ops.aten.expand.default, name="expand")) is False + assert is_exp(NS(target=torch.ops.aten.expm1.default, name="expm1")) is False + # unrelated op / a non-call origin (placeholder target is a str) -> False + assert is_exp(NS(target=torch.ops.aten.log.default, name="log")) is False + assert is_exp(NS(target="primals_1", name="primals_1")) is False diff --git a/tests/ops/misc/test_indirect_access.py b/tests/ops/misc/test_indirect_access.py index ae3a4ed4..cf57ef66 100644 --- a/tests/ops/misc/test_indirect_access.py +++ b/tests/ops/misc/test_indirect_access.py @@ -98,4 +98,4 @@ def vectoradd(a, idx, b): test_indirect_vectoradd(device) test_multidim_indirect(device) test_multidim_indirect_index_reuse(device) - #test_embedding(device, 1024, 2048) \ No newline at end of file + #test_embedding(device, 1024, 2048) diff --git a/tests/ops/misc/test_masked_nondividing.py b/tests/ops/misc/test_masked_nondividing.py new file mode 100644 index 00000000..20d12c8a --- /dev/null +++ b/tests/ops/misc/test_masked_nondividing.py @@ -0,0 +1,77 @@ +"""Masked-DMA non-dividing regression tests. + +Dropping the tile-divisibility constraint (is_dim_dividable / must_divide_dim and the +_index_expr / convert_indirect_indexing RecompileSignal) means a non-dividing shape no +longer falls back to a dividing tile -- correctness now rests entirely on the masked-DMA +[low, high) clamp + reduction-identity fill. These cases pin that path: what used to be a +(slow but safe) recompile is now a silent wrong answer if the clamp regresses. +""" +import os +import sys + +import torch +import torch.nn.functional as F + +sys.path.insert(0, os.path.join(os.environ.get("TORCHSIM_DIR", default="/workspace/PyTorchSim"), "tests")) +from _pytorchsim_utils import test_result + + +def test_pad_2d(device, size=(5, 10), pad=(2, 3, 1, 1)): + # 2-D constant pad on a non-dividing shape -> exercises the per-dim greedy pad recovery + # in _masked_bounds (both padded axes must be attributed to the right tile dim). + def fn(x): + return F.pad(x, pad) + x = torch.randn(size).to(device=device) + res = torch.compile(dynamic=False)(fn)(x) + test_result(f"Pad2D {size} {pad}", res, fn(x.cpu())) + + +def test_sum_over_broadcast(device, size=(5, 13)): + # sum over a non-dividing dim with a broadcast operand -- the tail fill must be the sum + # identity 0, NOT -inf (regression guard for the exp-origin substring bug: an origin + # name containing "exp" like expand/expm1 must not trigger the log-sum-exp -inf fill). + def fn(a, b): + return (a * b).sum(dim=1) + a = torch.randn(size).to(device=device) + b = torch.randn(size[0], 1).to(device=device) + res = torch.compile(dynamic=False)(fn)(a, b) + test_result(f"SumOverBroadcast {size}", res, fn(a.cpu(), b.cpu())) + + +def test_softmax_nondividing(device, size=(4, 10)): + # genuine log-sum-exp reduction: the exp path must still fill -inf and stay correct. + def fn(a): + return F.softmax(a, dim=1) + a = torch.randn(size).to(device=device) + res = torch.compile(dynamic=False)(fn)(a) + test_result(f"Softmax {size}", res, fn(a.cpu())) + + +def test_elementwise_nondividing(device, size=(7, 13)): + def fn(a, b): + return torch.relu(a) + b + a = torch.randn(size).to(device=device) + b = torch.randn(size).to(device=device) + res = torch.compile(dynamic=False)(fn)(a, b) + test_result(f"Elementwise {size}", res, fn(a.cpu(), b.cpu())) + + +def test_gather_nondividing(device, rows=13, cols=10, n=17): + # indirect (gather) on a non-dividing shape -- the divisibility RecompileSignal is gone. + def fn(x, idx): + return x[idx] + x = torch.randn(rows, cols).to(device=device) + idx = torch.randint(0, rows, [n]).to(device=device) + res = torch.compile(dynamic=False)(fn)(x, idx) + test_result(f"Gather [{rows},{cols}]->{n}", res, fn(x.cpu(), idx.cpu())) + + +if __name__ == "__main__": + device = torch.device("npu:0") + test_pad_2d(device, (5, 10), (2, 3, 1, 1)) + test_pad_2d(device, (7, 13), (1, 1, 3, 2)) + test_pad_2d(device, (1, 3, 10, 10), (2, 1, 1, 2)) + test_sum_over_broadcast(device, (5, 13)) + test_softmax_nondividing(device, (4, 10)) + test_elementwise_nondividing(device, (7, 13)) + test_gather_nondividing(device) From c1ae7b531ac2df64401c3ae049530e0cdb7d4629 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 7 Jul 2026 20:02:18 +0900 Subject: [PATCH 22/32] [Frontend] Reduce >4D togsim.transfer to <=4D before build_tog (fix pixel_shuffle trace) Both the Gemmini descriptor and the TOGSim DMA model cap at 4 tile dims, but a logical tile can exceed 4D (e.g. pixel_shuffle splits two spatial axes -> a 5D tile [1,1,2,4,2]). lower_transfer_to_gemmini reduces >4D as part of emitting Gemmini asm, but that runs only on the Spike/gem5 lowering path. The trace producer (build_tog) reads togsim.transfer BEFORE that, so it emitted a 5D DMA and TOGSim rejected it: "issued tile is not supported format.. tile.size: 5". Exposed by the pixel_shuffle case newly added to the wrapper1 CI suite. Add a peel_transfer POST_OPT pass that reduces every >4D togsim.transfer up front (before build_tog) while KEEPING it a togsim.transfer, so both consumers see <=4D. It mirrors lower_transfer_to_gemmini's two reductions: collapse unit dims when the effective rank is <=4 (memref.collapse_shape), else peel the outer effective dims into an affine.for nest with the lane-banked physical SRAM offset. The gemmini lowering keeps its own reduction as a now-defensive no-op. Validated on 128x128: test_floormod_axis_split (pixel_shuffle + 9 others), test_masked_nondividing, test_matmul, test_conv2d all pass. --- PyTorchSimFrontend/mlir/passes/__init__.py | 3 + .../mlir/passes/peel_transfer.py | 216 ++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 PyTorchSimFrontend/mlir/passes/peel_transfer.py diff --git a/PyTorchSimFrontend/mlir/passes/__init__.py b/PyTorchSimFrontend/mlir/passes/__init__.py index 98033fe8..55d33dd4 100644 --- a/PyTorchSimFrontend/mlir/passes/__init__.py +++ b/PyTorchSimFrontend/mlir/passes/__init__.py @@ -34,6 +34,7 @@ def _ensure_mlir_bindings_on_path(): from . import decompose_transfer from . import dma_fine_grained from . import lower_to_vcix +from . import peel_transfer from .lower_to_llvm import run_standard_lowering # noqa: F401 (re-exported) from .build_tog import run_tog # noqa: F401 (re-exported; replaces C++ test-tile-operation-graph) from .dma_fine_grained import run_fine_grained # noqa: F401 (re-exported; standalone/CLI) @@ -46,9 +47,11 @@ def _ensure_mlir_bindings_on_path(): lower_vlane_idx, ] # fine-grained first: splits the matmul DMAs that the vcix lowering then reads. +# peel_transfer last: split any >4D togsim.transfer so build_tog (trace) also sees <=4D. POST_OPT_PASSES = [ dma_fine_grained, lower_to_vcix, + peel_transfer, ] diff --git a/PyTorchSimFrontend/mlir/passes/peel_transfer.py b/PyTorchSimFrontend/mlir/passes/peel_transfer.py new file mode 100644 index 00000000..6af54378 --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/peel_transfer.py @@ -0,0 +1,216 @@ +"""Reduce a >4D togsim.transfer to <=4D (drop unit dims / peel outer dims). + +Both the Gemmini DMA descriptor and the TOGSim DMA model cap at 4 tile dims. A +logical tile can exceed 4D (e.g. pixel_shuffle splits two spatial axes -> a 5D +tile like [1,1,2,4,2]). lower_transfer_to_gemmini reduces >4D as part of emitting +Gemmini asm, but that runs only on the Spike/gem5 lowering path; the trace producer +(build_tog) reads togsim.transfer BEFORE that and would emit a >4D DMA that TOGSim +rejects ("issued tile is not supported format.. tile.size: 5"). + +This pass runs up front (POST_OPT, before build_tog) while KEEPING the op as +togsim.transfer, so both consumers see <=4D. It mirrors lower_transfer_to_gemmini's +two reductions: if the non-unit (effective) rank is <=4, collapse the unit dims +away (memref.collapse_shape); otherwise peel the outer effective dims into an +affine.for nest with the lane-banked physical SRAM offset. The innermost op is a +fresh <=4D togsim.transfer carrying the surviving dims' shape/strides/vlane axis +and the original dma_kind / masked / offset / subtile attributes. +""" + +OP_NAME = "togsim.transfer" +MARKERS = (OP_NAME,) + +from ._mlir_util import walk_ops +from .decompose_transfer import _int_array, _const_int, _squeeze_reassociation + + +def run(module, vectorlane=128, **_): + """Reduce every >4D togsim.transfer in `module` to <=4D, in place. Context active. + + vectorlane (= systolic-array size / number of vector lanes) feeds the lane-banked + physical SRAM offset in the peel, matching lower_transfer_to_gemmini. + """ + from mlir.ir import (InsertionPoint, Operation, MemRefType, ArrayAttr, + IntegerAttr, IntegerType, IndexType, DenseI64ArrayAttr, + DenseI32ArrayAttr, StridedLayoutAttr, AffineMap, AffineMapAttr, + AffineExpr, BoolAttr) + from mlir.dialects import affine + i64 = IntegerType.get_signless(64) + idx_ty = IndexType.get() + + def _arr(vals): + return ArrayAttr.get([IntegerAttr.get(i64, int(v)) for v in vals]) + + targets = [] + for region in module.operation.regions: + for b in region.blocks: + for op in walk_ops(b): + if op.operation.name == OP_NAME: + targets.append(op.operation) + + for op in targets: + op_operands = list(op.operands) + dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst = op_operands[:8] + sram_ty = MemRefType(sram.type) + elem, space = sram_ty.element_type, sram_ty.memory_space + tile_shape = list(sram_ty.shape) + if len(tile_shape) <= 4: + continue # TOGSim / Gemmini accept <=4 raw tile dims -- nothing to do + eff = [i for i, e in enumerate(tile_shape) if e > 1] + + has_indirect = "indirect" in op.attributes + offset_operand = op_operands[8] if has_indirect else None + vlane_axis = IntegerAttr(op.attributes["vlane_split_axis"]).value + dram_stride = _int_array(op.attributes["dram_stride"]) + tile_stride = _int_array(op.attributes["tile_stride"]) + vlane_stride = _const_int(vst, 1) + subtile = _int_array(op.attributes["subtile_size"]) if "subtile_size" in op.attributes else None + if "masked_axes" in op.attributes: + masked_axes = _int_array(op.attributes["masked_axes"]) + n_base = 9 if has_indirect else 8 + mvals = op_operands[n_base:] + masked_pairs = [(masked_axes[i], mvals[2 * i], mvals[2 * i + 1]) + for i in range(len(masked_axes))] + else: + masked_pairs = [] + + def _emit_inner(sram_mem, sram_idx_val, dram_idx_val, new_vlane, dr, tl, st, mp): + operands = [dram, dram_idx_val, sram_mem, sram_idx_val, tag, tag_idx, dma_type, vst] + if has_indirect: + operands.append(offset_operand) + for _axis, lo, hi in mp: + operands += [lo, hi] + attrs = { + "dma_kind": op.attributes["dma_kind"], + "vlane_split_axis": IntegerAttr.get(i64, new_vlane), + "dram_stride": _arr(dr), + "tile_stride": _arr(tl), + "padding": op.attributes["padding"], + } + if st is not None: + attrs["subtile_size"] = _arr(st) + if "async" in op.attributes: + attrs["async"] = op.attributes["async"] + if mp: + attrs["masked_axes"] = _arr([a for a, _lo, _hi in mp]) + attrs["masked_fill"] = op.attributes["masked_fill"] + if has_indirect: + attrs["indirect"] = op.attributes["indirect"] + attrs["offset_stride"] = op.attributes["offset_stride"] + if "accumulate" in op.attributes: + attrs["accumulate"] = op.attributes["accumulate"] + if "acc_float" in op.attributes: + attrs["acc_float"] = op.attributes["acc_float"] + Operation.create(OP_NAME, results=[], operands=operands, attributes=attrs) + + # <=4 effective dims: collapse the unit dims so the raw rank reaches <=4. + if len(eff) <= 4: + groups, target = _squeeze_reassociation(tile_shape) + reassoc = ArrayAttr.get( + [ArrayAttr.get([IntegerAttr.get(i64, d) for d in g]) for g in groups]) + collapsed_ty = MemRefType.get(target, elem, memory_space=space) + keep = [next((d for d in g if tile_shape[d] > 1), g[-1]) for g in groups] + dr = [dram_stride[i] for i in keep] + tl = [tile_stride[i] for i in keep] + st = [subtile[i] for i in keep] if subtile is not None else None + mp = [(next(gi for gi, g in enumerate(groups) if d in g), lo, hi) + for d, lo, hi in masked_pairs] + new_vlane = next(gi for gi, g in enumerate(groups) if vlane_axis in g) + with InsertionPoint(op): + sram_c = Operation.create( + "memref.collapse_shape", results=[collapsed_ty], operands=[sram], + attributes={"reassociation": reassoc}).results[0] + _emit_inner(sram_c, sram_idx, dram_idx, new_vlane, dr, tl, st, mp) + op.erase() + continue + + # >4 effective dims: peel the outer ones into an affine.for nest, keep the + # innermost 4. The SRAM slice offset is the lane-banked physical offset, + # delivered as sram_idx; the DRAM offset folds into dram_idx. + peeled, inner = eff[:-4], eff[-4:] + ndim = len(tile_shape) + inner_shape = [tile_shape[d] for d in inner] + inner_strides = [tile_stride[d] for d in inner] + dr = [dram_stride[d] for d in inner] + tl = [tile_stride[d] for d in inner] + st = [subtile[d] for d in inner] if subtile is not None else None + if any(d in peeled for d, _lo, _hi in masked_pairs): + raise NotImplementedError("masked-DMA clamp on a peeled (outer-loop) axis") + mp = [(inner.index(d), lo, hi) for d, lo, hi in masked_pairs if d in inner] + if vlane_axis in inner: + new_vlane = inner.index(vlane_axis) + elif vlane_axis in peeled: + raise NotImplementedError( + f"vlane split axis {vlane_axis} peeled into the outer loop nest") + else: + new_vlane = 0 + + split_extent = tile_shape[vlane_axis] + nr_outerloop = max( + (split_extent + vectorlane * vlane_stride - 1) // (vectorlane * vlane_stride), 1) + new_size = nr_outerloop * vlane_stride + target_stride = tile_stride[vlane_axis] + + def _phys(d): + s = tile_stride[d] + return s // split_extent * new_size if s > target_stride else s + + static_sizes = [1] * ndim + for d in inner: + static_sizes[d] = tile_shape[d] + res_ty = MemRefType.get( + inner_shape, elem, + layout=StridedLayoutAttr.get(0, inner_strides), memory_space=space) + + cur_ip = InsertionPoint(op) + ivs = [] + for d in peeled: + floop = affine.AffineForOp(0, tile_shape[d], 1, ip=cur_ip) + floop.operation.attributes["inner_loop"] = BoolAttr.get(True) + ivs.append(floop.induction_variable) + with InsertionPoint(floop.body): + affine.AffineYieldOp([]) + cur_ip = InsertionPoint.at_block_terminator(floop.body) + + npeel = len(peeled) + with cur_ip: + sub = Operation.create( + "memref.subview", results=[res_ty], operands=[sram], + attributes={"static_offsets": DenseI64ArrayAttr.get([0] * ndim), + "static_sizes": DenseI64ArrayAttr.get(static_sizes), + "static_strides": DenseI64ArrayAttr.get([1] * ndim), + "operandSegmentSizes": DenseI32ArrayAttr.get([1, 0, 0, 0])} + ).results[0] + sram_expr = AffineExpr.get_dim(0) * _phys(peeled[0]) + for k in range(1, npeel): + sram_expr = sram_expr + AffineExpr.get_dim(k) * _phys(peeled[k]) + sram_off_val = Operation.create( + "affine.apply", results=[idx_ty], operands=list(ivs), + attributes={"map": AffineMapAttr.get(AffineMap.get(npeel, 0, [sram_expr]))} + ).results[0] + dram_expr = AffineExpr.get_dim(0) + for k in range(npeel): + dram_expr = dram_expr + AffineExpr.get_dim(k + 1) * dram_stride[peeled[k]] + dram_idx_val = Operation.create( + "affine.apply", results=[idx_ty], operands=[dram_idx, *ivs], + attributes={"map": AffineMapAttr.get(AffineMap.get(npeel + 1, 0, [dram_expr]))} + ).results[0] + _emit_inner(sub, sram_off_val, dram_idx_val, new_vlane, dr, tl, st, mp) + op.erase() + + +def peel_text(text): + if OP_NAME not in text: + return text + from mlir.ir import Context, Module, Location + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + m = Module.parse(text) + run(m) + return str(m) + + +if __name__ == "__main__": + import sys + out = peel_text(open(sys.argv[1]).read()) + (open(sys.argv[2], "w").write(out) if len(sys.argv) > 2 else sys.stdout.write(out)) From bd521841052d35a7c509ccd6f628c6e08ee214f0 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 8 Jul 2026 16:06:51 +0900 Subject: [PATCH 23/32] [Frontend] Fix non-dividing tile over-reads found by the small-array CI Three MVIN over-reads on shapes whose tiles do not divide the extent: a cat concat-dim split, a vlane index_expr row stride on multi-row splits, and a 1-D reduction tile that exceeded its extent. --- PyTorchSimFrontend/mlir/mlir_cat_template.py | 20 ++++++++++++++++++- .../mlir/mlir_codegen_backend.py | 4 +++- PyTorchSimFrontend/mlir/mlir_common.py | 3 ++- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_cat_template.py b/PyTorchSimFrontend/mlir/mlir_cat_template.py index 7abdfee6..b922e51b 100644 --- a/PyTorchSimFrontend/mlir/mlir_cat_template.py +++ b/PyTorchSimFrontend/mlir/mlir_cat_template.py @@ -209,6 +209,22 @@ def _compute_excluded_dims(self, tile_sizes: list) -> list: tile_sizes[idx] = 1 return excluded + @staticmethod + def _largest_divisor_leq(extent, cap): + """Largest divisor of ``extent`` that does not exceed ``cap`` (>= 1). + + The concat-dim copy loop steps by the per-input tile over ``[0, extent)``. + If the tile does not divide ``extent`` the final iteration is ragged and, + because the DMA carries no remainder mask, over-reads the input and + over-writes the output. Snapping the tile down to a divisor keeps every + iteration full-width and in-bounds. + """ + cap = min(cap, extent) + for d in range(cap, 0, -1): + if extent % d == 0: + return d + return 1 + def _calculate_input_tile_sizes(self, kernel, input_sizes, tile_sizes, num_inputs, rank, precision_bytes): """Calculate tile sizes along the concat dimension for each input.""" non_dim_tile_elements = math.prod(tile_sizes) if tile_sizes else 1 @@ -218,7 +234,9 @@ def _calculate_input_tile_sizes(self, kernel, input_sizes, tile_sizes, num_input input_tile_sizes_dim = [] for i in range(num_inputs): if extra_concat > 0 and non_dim_tile_elements > 0: - tile_dim = min(input_sizes[i][self.dim], extra_concat) + # Snap to a divisor of the input's concat extent so the copy loop + # never emits a ragged (unmasked) tail tile. + tile_dim = self._largest_divisor_leq(input_sizes[i][self.dim], extra_concat) extra_concat -= tile_dim else: tile_dim = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 5778213b..56f123ea 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -847,7 +847,9 @@ def _index_expr(self, tile_desc, renamed_expression, index, base_vector_index): if idx == tile_desc.vmap.vlane_split_axis: # Need to add vector lane offset stride_dim = ops.remainder(dim, vlane_stride_vec) outer_dim = ops.remainder(ops.truncdiv(dim, vlane_stride_vec), vlane_outer_vec) - dim = ops.add(stride_dim, ops.mul(outer_dim, nr_vector_lane_vec)) + # Next sublane-row stride is vector_lane*vlane_stride, not vector_lane alone. + row_stride_vec = ops.mul(nr_vector_lane_vec, vlane_stride_vec) + dim = ops.add(stride_dim, ops.mul(outer_dim, row_stride_vec)) with self.override_buffer_cse(buffer=self.const_buffer, cse=self.const_cse): vlane_offset = ops.vlane_offset(vlane_vec, vlane_vec, attributes={"vlane_offset": offset}, comment="vlane offset") diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index 99f22762..9d610bdc 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -411,7 +411,8 @@ def init_tile_size(ranges, vlane_stride, vector_lane): return [1] tile_size = [1] * nr_dim if nr_dim == 1: - tile_size[0] = 1 if ranges[0] == 1 else 2 * vlane_stride * vector_lane + # Cap to extent so the tile never over-reads the DRAM buffer (extent 128 vs tile 512 -> 384 garbage folded into the reduction). No-op when extent >= tile. + tile_size[0] = 1 if ranges[0] == 1 else min(2 * vlane_stride * vector_lane, ranges[0]) elif nr_dim == 2: tile_size[-1] = vlane_stride * vector_lane tile_size[-2] = 2 * vector_lane From 47c64e1554663db9fddac3f1a0f04bfc827e42fd Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 8 Jul 2026 21:10:36 +0900 Subject: [PATCH 24/32] [Frontend] Keep argsort's sort kernel alive so indices are written torch.sort's frontend lowering exposes the sorted VALUES as the sole scheduler-visible template output; the INDICES buffer is written in place by the same kernel (make_inplace) but is not registered as an output. For argsort the values output is dead, so Inductor DCEs the whole sort kernel and the indices buffer is returned uninitialized (all-zeros), silently corrupting any argsort result (e.g. DeepSeek MoE expert routing). Register the in-place indices write as a MutationOutput owned by the sort operation so the scheduler keeps the kernel alive whenever indices is consumed. OperationBuffer.get_outputs() hardcodes [self], so add a small MLIRTemplateBuffer subclass that can own extra MutationOutputs and have MLIRTemplateCaller.output_node build it; the sort lowering attaches the indices mutation via add_mutation_output. For templates with no mutation this is behaviorally identical to the base (get_outputs == [self]). Then gate the values MVOUT on the values output being live: in argsort it is pruned from the kernel args, so emitting its MVOUT referenced an undeclared %YV. The values scratchpad is still produced internally for the bitonic compare; only the dead DRAM store is skipped. Verified: t.argsort() now returns correct indices (was all-zeros); torch.sort with both outputs live is unchanged; test_sort, test_matmul, test_cat, test_indirect_access unaffected by the output_node change. Note: x[argsort] end to end additionally needs the indirect-gather codegen fix (separate). --- PyTorchSimFrontend/mlir/mlir_lowering.py | 11 ++++++ PyTorchSimFrontend/mlir/mlir_sort_template.py | 6 +++ PyTorchSimFrontend/mlir/mlir_template.py | 38 ++++++++++++++++++- 3 files changed, 53 insertions(+), 2 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_lowering.py b/PyTorchSimFrontend/mlir/mlir_lowering.py index 7f33d956..90aa62fc 100644 --- a/PyTorchSimFrontend/mlir/mlir_lowering.py +++ b/PyTorchSimFrontend/mlir/mlir_lowering.py @@ -324,6 +324,17 @@ def _mlir_custom_sort_default( stable=stable_required, ) sorted_values = mlir_template.generate(template_buffer_node=value).output_node() + + def _unwrap(t): + t = t.data if isinstance(t, ir.TensorBox) else t + t = t.data if isinstance(t, ir.StorageBox) else t + return t + # The sort kernel writes `indices` in place, but only `sorted_values` is a + # scheduler-visible output. Advertise the indices write as a MutationOutput owned by + # the sort op, so argsort (values dead) does not DCE the whole kernel. + sort_op = _unwrap(sorted_values) + idx_buf = _unwrap(indices) + sort_op.add_mutation_output(ir.MutationOutput(idx_buf.get_layout(), idx_buf, sort_op)) return sorted_values, indices diff --git a/PyTorchSimFrontend/mlir/mlir_sort_template.py b/PyTorchSimFrontend/mlir/mlir_sort_template.py index 24b3a460..338f9636 100644 --- a/PyTorchSimFrontend/mlir/mlir_sort_template.py +++ b/PyTorchSimFrontend/mlir/mlir_sort_template.py @@ -3,6 +3,7 @@ from torch._inductor.ir import Buffer, IRNode from torch._inductor.virtualized import _ops as ops +from torch._inductor.virtualized import V from torch._inductor.codegen import common from PyTorchSimFrontend.mlir import mlir_common @@ -38,7 +39,9 @@ {{ BITONIC_BODY }} {{ kernel.def_dma_op("MVOUT", "XI", [], XI_TILE_DESC, indent_size=INDENT_SIZE, dram_stride=XI_DRAM_STRIDE, dram_offset="xi_dram_offset") }} + {%- if YV_LIVE %} {{ kernel.def_dma_op("MVOUT", "YV", [], YV_TILE_DESC, indent_size=INDENT_SIZE, dram_stride=YV_DRAM_STRIDE, dram_offset="yv_dram_offset") }} + {%- endif %} {%- for d in range(RANK-1) %} } { outer_loop=true } {%- endfor %} @@ -433,6 +436,9 @@ def render( X=x, XI=xi, YV=yv, + # In argsort the sorted-values output is dead and dropped from the kernel args; + # skip its MVOUT so the body does not reference the pruned %YV (undeclared SSA). + YV_LIVE=yv.get_name() not in V.graph.removed_buffers, X_TILE_DESC=x_tile_desc, XI_TILE_DESC=xi_tile_desc, YV_TILE_DESC=yv_tile_desc, diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index 051e69b6..71eec07c 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -15,9 +15,9 @@ from PyTorchSimFrontend import extension_config from torch._inductor.codegen.common import KernelTemplate, CSE, DeferredLine -from torch._inductor.ir import Buffer, IRNode, TemplateBuffer, ChoiceCaller, ir_node_to_tensor +from torch._inductor.ir import Buffer, IRNode, TemplateBuffer, ChoiceCaller, ir_node_to_tensor, TensorBox from torch._inductor.select_algorithm import PartialRender -from torch._inductor.codegen.cuda.cuda_kernel import CUDATemplateCaller +from torch._inductor.codegen.cuda.cuda_kernel import CUDATemplateCaller, CUDATemplateBuffer from torch._inductor.autotune_process import TensorMeta from torch._inductor.virtualized import V, NullHandler, _ops as ops from torch._inductor.utils import IndentedBuffer @@ -1299,7 +1299,41 @@ def set_tile_size(self, template_fusion_info, prologue=False): self.compute_body_loop.step = tile_desc.get_compute_vec_size() return tile_desc +class MLIRTemplateBuffer(CUDATemplateBuffer): + """Template output buffer that can own extra MutationOutputs beyond its primary output. + + The sort kernel's primary scheduler-visible output is the sorted values, but it also + writes the indices buffer in place (mlir_sort_template make_inplace). Registering that + write as a MutationOutput owned here -- so get_outputs() advertises it -- lets the + scheduler keep the kernel alive when only the indices are consumed (argsort), instead of + DCE-ing the whole sort. OperationBuffer.get_outputs() otherwise hardcodes [self]. + """ + + def add_mutation_output(self, mutation: IRNode) -> None: + if not hasattr(self, "_extra_mutation_outputs"): + self._extra_mutation_outputs = [] + self._extra_mutation_outputs.append(mutation) + + def get_outputs(self) -> List[Buffer]: + return [self, *getattr(self, "_extra_mutation_outputs", [])] + + class MLIRTemplateCaller(CUDATemplateCaller): + def output_node(self) -> TensorBox: + # Same as CUDATemplateCaller.output_node but builds a mutation-aware buffer so + # callers (e.g. the sort lowering) can attach in-place-write MutationOutputs. + self.bmreq.update_workspace_size() + return TensorBox.create( + MLIRTemplateBuffer( + layout=self.layout, + inputs=self.input_nodes, + make_kernel_render=self.make_kernel_render, + workspace_size=self.bmreq.workspace_size, + supports_epilogue_fusion=self.supports_epilogue_fusion, + template=self.template, + ) + ) + def __init__(self, name, category, input_nodes, layout, make_kernel_render, supports_epilogue_fusion, template, info_kwargs, description): bmreq = MLIRBenchmarkRequest( kernel_name=name, From 1d7da69e3ec660dcfd9e644485fd41f8add3a7c5 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 8 Jul 2026 21:10:36 +0900 Subject: [PATCH 25/32] [Frontend] Fix the gather base offset and drop the ill-posed padded-load clamp A bare indirect index (x[idx]) left the loop-local symbol in the DMA base offset. Separately, reverse-engineering per-dim padding from a single flat offset is ill-posed -- under channels_last the stride-1 axis absorbs the remainder and yields an empty [low, high) box that zeroed the whole load. --- .../mlir/mlir_codegen_backend.py | 39 ++++++++----------- 1 file changed, 16 insertions(+), 23 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 56f123ea..590271c6 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -1398,36 +1398,25 @@ def _masked_fill_bits(self, dtype, index): return bits & ((1 << (t.element_size() * 8)) - 1) def _masked_bounds(self, name, index, dram_stride, local_tile_desc, is_load, buffer, local_dims): - """Per tile-axis [low, high) clamp for a masked DMA. Per axis the valid GLOBAL range - is [glo, ghi) (store: [0, output_extent); padded load: [pad, pad + input_extent)), and - _emit_clamp turns it into the tile-local low/high SSA vars. Returns [(tile_axis, low_var, high_var), ...]. + """Per tile-axis [low, high) clamp for a masked DMA -- ONLY the trailing tail of a + non-dividing loop extent (valid GLOBAL range per axis is [0, loop_extent)). _emit_clamp + turns it into the tile-local low/high SSA vars. Returns [(tile_axis, low_var, high_var)]. + + A padded load reads out-of-bounds positions, but we do NOT clamp those here: the + consumer already yields the correct value at pad positions (a compute-side arith.select + for a padded gather, or the pad op's own fill). Reverse-engineering per-dim padding from + the single flat index offset is ill-posed -- the offset mixes the tap shift with the + padding, and under channels_last the stride-1 (channel) axis absorbs the offset + remainder, yielding an impossible clamp (e.g. [4, 8) on a size-2 channel tile) that + zeroed the whole load and made channels_last depthwise conv 99.9% wrong. """ tile_size = local_tile_desc.get_tile_size() - const = int(index.as_coeff_Add()[0]) if not index.is_number else int(index) - # index const = -sum(pad_d * dram_stride[d]); recover per-dim pad greedily (desc stride). - pad = [0] * len(tile_size) - if is_load and const < 0: - rem = -const - for d in sorted(range(len(dram_stride)), key=lambda x: -dram_stride[x]): - s = dram_stride[d] - if s > 0 and d < len(pad): - pad[d] = rem // s - rem -= pad[d] * s - in_shape, in_stride = self.buffer_types[name][2], self.buffer_types[name][3] axes = [] for d, k in enumerate(local_dims): if d >= len(tile_size) or k >= len(self.ranges): continue iv = str(self.itervars[k]) - # A padded load (const < 0) reads a smaller input: clamp per-dim to - # [pad_d, pad_d + input_extent_d). Every other DMA is valid over the whole - # loop extent, so only the trailing tail needs clamping. - if is_load and const < 0: - ext = next((int(in_shape[j]) for j, st in enumerate(in_stride) - if int(st) == int(dram_stride[d])), None) - glo, ghi = (pad[d], pad[d] + ext) if ext is not None else (0, int(self.ranges[k])) - else: - glo, ghi = 0, int(self.ranges[k]) + glo, ghi = 0, int(self.ranges[k]) axes.append((d, iv, glo, ghi, int(tile_size[d]))) return self._emit_clamp(axes, buffer) @@ -1643,6 +1632,10 @@ def convert_indirect_indexing(self, index :sympy.Expr): if arg.is_Mul and arg.args[0].is_number: offset_stride = int(arg.args[0]) index = index.replace(arg, 0) + # A bare indirect index (x[idx]: index IS the symbol, so index.args is empty) + # escapes the loop above. Zero any remaining indirect symbol -- the per-position + # gather rides the offset spad -- else affine.apply uses %sym before it exists. + index = index.subs({s: 0 for s in index.free_symbols if str(s) in self.indirect_symbols}) sram_var, _, _, _, tile_shape, _ = self.spad_buffer_dict[first_dim] return index, (sram_var, tile_shape, offset_stride) From 67c26369d1a036a3338c409210470be314b777be Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 9 Jul 2026 13:51:18 +0900 Subject: [PATCH 26/32] [Frontend] Store the template buffer when it outlives the fused epilogue codegen_epilogue_body() decided whether to emit the template buffer's MVOUT by asking "did a fused epilogue already store something?": if len(self.stores._lines) == 0: template_store() That is a proxy for the real question, "does anything outside this kernel still read the template buffer?", and it is wrong in two of the four cases: - epilogue + live template buffer -> store skipped, buffer never written - no epilogue + dead template buffer -> store emitted for a pruned DRAM arg The first case corrupts DeepSeek-V3's MoE gate. The gate is sigmoid(mm(hidden, gate_w) + bias), so add+sigmoid fuse onto the GEMM template, but the raw mm result is also read by a later gather kernel. Its buffer was declared as a kernel output and never stored, so the gather read an uninitialized buffer (all zeros), scrambling expert routing. Every config was affected; only 8x8 crossed the test's loose atol=2e-1 (final Max abs diff 0.2546) while 32x32 landed just under it and passed. Ask the real question instead. Inductor already answers it: Kernel.remove_kernel_local_buffers() drops a buffer only when Scheduler.can_buffer_be_removed_through_fusion() finds every user inside this fused kernel. A buffer that survived removal therefore has an outside user and must be stored. Record the template buffer's name in def_kernel() and def_conv_kernel() (conv has its own implementation) and gate template_store() on it, in both the main and the reduction_fusion branches. The proxy is gone. Use the kernel-local self.removed_buffers, not V.graph.removed_buffers: the former is what remove_buffer() populates, and the latter is still empty when the store hook runs. Note this makes a previously missing MVOUT issue, so DRAM traffic and cycle counts rise for any model with a live template buffer under a fused epilogue. That is the correct cost, not a regression, but reported cycles will shift. Verified on the 8x8 config: test_deepseek_v3_base goes from Max abs diff 0.2546 to passing; a dead template buffer (relu(mm)) still prunes its DRAM arg and emits a single MVOUT; the trace/TOGSim timing path runs with the extra store. Regressions pass: test_bmm_reduction, test_matmul, test_bmm, test_conv2d, test_cnn, test_group_conv, test_pool, test_reduce, test_softmax, test_layernorm, test_mlp. --- PyTorchSimFrontend/mlir/mlir_template.py | 27 ++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index 71eec07c..2566d2e0 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -645,11 +645,22 @@ def template_store(): code = self.def_dma_op("MVOUT", dram_var, index_list, tile_desc, lazy_mode=False) self.cse.generate(self.dma_stores, code, assignment = False) + def template_buffer_live(): + # Inductor's remove_kernel_local_buffers() drops the template buffer only if + # every user lives inside this fused kernel. If it survived, some user is + # outside it, so store it to DRAM even if a fused epilogue also stores. + name = getattr(self, "template_buffer_name", None) + assert name is not None, ( + "store_output() used by a template that never declared its output buffer; " + "call def_kernel()/def_conv_kernel() with outputs=[...] first" + ) + return name not in self.removed_buffers + body = IndentedBuffer() with self.epilogue_buffer_group.as_local(): # Do dma store first to overlap epilogue nodes if self.reduction_fusion: - if len(self.stores._lines) == 0: + if template_buffer_live(): template_store() body.splice(self.dma_stores) self.dma_stores.clear() @@ -668,7 +679,7 @@ def template_store(): else: compute_body.splice(self.loads) compute_body.splice(self.compute) - if len(self.stores._lines) == 0: + if template_buffer_live(): template_store() compute_body.splice(self.stores) if (compute_body.getvalue()): @@ -690,6 +701,12 @@ def def_kernel( f"{len(inputs) + len(outputs)=} != {len(names)=}, {inputs=}, {outputs=}, {names=}" ) + # The template's own output buffer. codegen_epilogue_body() stores it to DRAM iff + # it survived Inductor's kernel-local buffer removal -- i.e. it still has a user + # outside this fused kernel. + if outputs: + self.template_buffer_name = outputs[0].get_name() + if input_reorder is not None: assert len(inputs) == len(input_reorder) else: @@ -758,6 +775,12 @@ def def_conv_kernel( f"{len(inputs) + len(outputs)=} != {len(names)=}, {inputs=}, {outputs=}, {names=}" ) + # The template's own output buffer. codegen_epilogue_body() stores it to DRAM iff + # it survived Inductor's kernel-local buffer removal -- i.e. it still has a user + # outside this fused kernel. + if outputs: + self.template_buffer_name = outputs[0].get_name() + if input_reorder is not None: assert len(inputs) == len(input_reorder) else: From 22b32ce211a6f284620707da9b363ee4f386330b Mon Sep 17 00:00:00 2001 From: Jagggged Date: Thu, 9 Jul 2026 12:55:09 +0900 Subject: [PATCH 27/32] [Frontend] Add pointwise op lowerings and their CI Lower the remaining pointwise ops, route math.log/atan to VCIX in the Python pass, and add tests/ops/elementwise/test_pointwise.py. --- .github/workflows/pytorchsim_test.yml | 18 + PyTorchSimFrontend/mlir/mlir_ops.py | 341 ++++++++++++++++-- .../mlir/passes/lower_to_vcix.py | 13 +- gem5_script/vpu_config.py | 2 + tests/ops/elementwise/test_pointwise.py | 214 +++++++++++ 5 files changed, 555 insertions(+), 33 deletions(-) create mode 100644 tests/ops/elementwise/test_pointwise.py diff --git a/.github/workflows/pytorchsim_test.yml b/.github/workflows/pytorchsim_test.yml index 7cc76a5b..318c7945 100644 --- a/.github/workflows/pytorchsim_test.yml +++ b/.github/workflows/pytorchsim_test.yml @@ -57,6 +57,24 @@ jobs: -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/elementwise/test_transcendental.py + test_pointwise: + name: Run test_pointwise.py + runs-on: ubuntu-latest + steps: + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Run test_pointwise.py + run: | + echo "Running test_pointwise.py" + docker run --rm \ + -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ + ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/elementwise/test_pointwise.py + test_activation: name: Run test_activation.py runs-on: ubuntu-latest diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index f1fb4186..d21cc1d0 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -145,7 +145,7 @@ def where(condition, operand1, operand2, *args, **kwargs): operand2 = ops.broadcast(operand2, cond_type[0]) tile_size, ret_type = V.kernel.var_info[operand1] shape = f"vector<{tile_size}x{ret_type}>" if tile_size > 1 else ret_type - cond_shape = f"vector<{tile_size}xi1>" if tile_size > 1 else "" + cond_shape = f"vector<{tile_size}xi1>" if tile_size > 1 else "i1" op_str = f"arith.select %{condition}, %{operand1}, %{operand2}" shape = f"{cond_shape}, {shape}" @@ -309,7 +309,10 @@ def binary_elementwise_common(operand1, operand2): @staticmethod def abs(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + shape = f"vector<{tile_size}x{dtype}>" if tile_size > 1 else dtype + opcode = "math.absf" if dtype.startswith("f") else "math.absi" + return format_mlir_op(f'{opcode} %{operand}', shape, **kwargs), [tile_size, dtype] @staticmethod def exp(operand, *args, **kwargs): @@ -463,68 +466,263 @@ def erf(operand, *args, **kwargs): @staticmethod def cosh(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + operand = ops.broadcast(operand, 4) + val = ops.cosh(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + exp_pos = ops.exp(operand) + exp_neg = ops.exp(ops.neg(operand)) + + sum_exp = ops.add(exp_pos, exp_neg) + half_const = ops.constant(0.5, dtype) + + res = ops.mul(sum_exp, half_const) + return res, V.kernel.var_info[res] @staticmethod def sinh(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + operand = ops.broadcast(operand, 4) + val = ops.sinh(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + exp_pos = ops.exp(operand) + exp_neg = ops.exp(ops.neg(operand)) + + sub_exp = ops.sub(exp_pos, exp_neg) + half_const = ops.constant(0.5, dtype) + + res = ops.mul(sub_exp, half_const) + return res, V.kernel.var_info[res] @staticmethod def tanh(operand, *args, **kwargs): op_type = V.kernel.var_info[operand] + tile_size = op_type[0] + dtype = op_type[1] # Check scalar - op_type = V.kernel.var_info[operand] if op_type[0] == 1: operand = ops.broadcast(operand, 4) val = ops.tanh(operand) result = ops.extractelement(val, 0) return result, V.kernel.var_info[result] - op_type = V.kernel.var_info[operand] - tile_size = op_type[0] - dtype = op_type[1] - # Type check & auto cast - if dtype.startswith("f"): + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): operand = ops.to_dtype(operand, "f32") + dtype = "f32" shape = f"vector<{tile_size}x{dtype}>" if tile_size > 1 else dtype return format_mlir_op(f'math.tanh %{operand}', shape, **kwargs), [tile_size, dtype] @staticmethod def acos(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + operand = ops.broadcast(operand, 4) + val = ops.acos(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + asin_val = ops.asin(operand) + res = ops.sub(ops.constant(math.pi / 2, dtype), asin_val) + return res, V.kernel.var_info[res] @staticmethod def acosh(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + operand = ops.broadcast(operand, 4) + val = ops.acosh(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + x2 = ops.square(operand) + val = ops.sub(x2, ops.constant(1.0, dtype)) + sqrt_val = ops.sqrt(val) + sum_val = ops.add(operand, sqrt_val) + + res = ops.log(sum_val) + return res, V.kernel.var_info[res] @staticmethod def asin(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + operand = ops.broadcast(operand, 4) + val = ops.asin(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + x2 = ops.square(operand) + denom = ops.sqrt(ops.sub(ops.constant(1.0, dtype), x2)) + res = ops.atan(ops.truediv(operand, denom)) + return res, V.kernel.var_info[res] @staticmethod def asinh(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + operand = ops.broadcast(operand, 4) + val = ops.asinh(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + x2 = ops.square(operand) + val = ops.add(x2, ops.constant(1.0, dtype)) + sqrt_val = ops.sqrt(val) + sum_val = ops.add(operand, sqrt_val) + + res = ops.log(sum_val) + return res, V.kernel.var_info[res] @staticmethod def atan2(operand1, operand2, *args, **kwargs): - raise NotImplementedError + tile_size, ret_type, y, x = ExtensionOverrides.binary_elementwise_common(operand1, operand2) + if not ret_type.startswith("f"): + y = ops.to_dtype(y, "f32") + x = ops.to_dtype(x, "f32") + ret_type = "f32" + + pi = ops.constant(math.pi, ret_type) + zero = ops.constant(0.0, ret_type) + + base = ops.atan(ops.truediv(y, x)) + corrected = ops.add(base, ops.copysign(pi, y)) + res = ops.where(ops.lt(x, zero), corrected, base) + return res, [tile_size, ret_type] @staticmethod def atan(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + operand = ops.broadcast(operand, 4) + val = ops.atan(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + shape = f"vector<{tile_size}x{dtype}>" if tile_size > 1 else dtype + return format_mlir_op(f'math.atan %{operand}', shape, **kwargs), [tile_size, dtype] @staticmethod def atanh(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + # Check scalar + if tile_size == 1: + operand = ops.broadcast(operand, 4) + val = ops.atanh(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): + operand = ops.to_dtype(operand, "f32") + dtype = "f32" + + one_const = ops.constant(1.0, dtype) + val1 = ops.add(one_const, operand) + val2 = ops.sub(one_const, operand) + div_val = ops.truediv(val1, val2) + + half_const = ops.constant(0.5, dtype) + res = ops.mul(ops.log(div_val), half_const) + return res, V.kernel.var_info[res] @staticmethod def copysign(operand1, operand2, *args, **kwargs): - raise NotImplementedError + tile_size, ret_type, operand1, operand2 = ExtensionOverrides.binary_elementwise_common(operand1, operand2) + if not ret_type.startswith("f"): + raise ValueError("copysign is only supported for floats") + shape = f"vector<{tile_size}x{ret_type}>" if tile_size > 1 else ret_type + op_str = f'math.copysign %{operand1}, %{operand2}' + return format_mlir_op(op_str, shape, **kwargs), [tile_size, ret_type] @staticmethod def erfc(operand, *args, **kwargs): - raise NotImplementedError + """ + There is no direct MLIR operation for erfc, so we can implement it using the relationship: + erfc(x) = 1 - erf(x) + """ + op_type = V.kernel.var_info[operand] + + # Check scalar + if op_type[0] == 1: + operand = ops.broadcast(operand, 4) + val = ops.erfc(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] + + tile_size = op_type[0] + dtype = op_type[1] + erf_val = ops.erf(operand) + one_const = ops.constant(1.0, dtype) + res = ops.sub(one_const, erf_val) + + return res, V.kernel.var_info[res] @staticmethod def erfinv(operand, *args, **kwargs): @@ -536,7 +734,15 @@ def frexp(operand, *args, **kwargs): @staticmethod def hypot(operand1, operand2, *args, **kwargs): - raise NotImplementedError + tile_size, ret_type, operand1, operand2 = ExtensionOverrides.binary_elementwise_common(operand1, operand2) + if not ret_type.startswith("f"): + raise ValueError("hypot is only supported for floats") + + x_sq = ops.square(operand1) + y_sq = ops.square(operand2) + sum_sq = ops.add(x_sq, y_sq) + res = ops.sqrt(sum_sq) + return res, V.kernel.var_info[res] @staticmethod def log10(operand, *args, **kwargs): @@ -564,12 +770,20 @@ def log2(operand, *args, **kwargs): @staticmethod def log(operand, *args, **kwargs): op_type = V.kernel.var_info[operand] + if op_type[0] == 1: + operand = ops.broadcast(operand, 4) + val = ops.log(operand) + result = ops.extractelement(val, 0) + return result, V.kernel.var_info[result] tile_size = op_type[0] dtype = op_type[1] # Type check & auto cast - if dtype.startswith("f"): + # Float-only instruction: promote non-float inputs (e.g. integers) to f32 + # to run it. Native float widths (f16/f64) are left untouched. + if not dtype.startswith("f"): operand = ops.to_dtype(operand, "f32") + dtype = "f32" shape = f"vector<{tile_size}x{dtype}>" if tile_size > 1 else dtype return format_mlir_op(f'math.log %{operand}', shape, **kwargs), [tile_size, dtype] @@ -660,11 +874,21 @@ def bitwise_xor(operand1, operand2, *args, **kwargs): @staticmethod def bitwise_left_shift(operand1, operand2, *args, **kwargs): - raise NotImplementedError + tile_size, ret_type, operand1, operand2 = ExtensionOverrides.binary_elementwise_common(operand1, operand2) + if ret_type.startswith("f"): + raise ValueError("Bitwise left shift not supported for floats") + shape = f"vector<{tile_size}x{ret_type}>" if tile_size > 1 else ret_type + op_str = f'arith.shli %{operand1}, %{operand2}' + return format_mlir_op(op_str, shape, **kwargs), [tile_size, ret_type] @staticmethod def bitwise_right_shift(operand1, operand2, *args, **kwargs): - raise NotImplementedError + tile_size, ret_type, operand1, operand2 = ExtensionOverrides.binary_elementwise_common(operand1, operand2) + if ret_type.startswith("f"): + raise ValueError("Bitwise right shift not supported for floats") + shape = f"vector<{tile_size}x{ret_type}>" if tile_size > 1 else ret_type + op_str = f'arith.shrsi %{operand1}, %{operand2}' + return format_mlir_op(op_str, shape, **kwargs), [tile_size, ret_type] @staticmethod def rsqrt(operand, *args, **kwargs): @@ -689,15 +913,51 @@ def sigmoid(operand, *args, **kwargs): @staticmethod def fmod(operand1, operand2, *args, **kwargs): - raise NotImplementedError + tile_size, ret_type, operand1, operand2 = ExtensionOverrides.binary_elementwise_common(operand1, operand2) + + if ret_type.startswith("f"): + div_val = ops.truediv(operand1, operand2) + trunc_val = ops.trunc(div_val) + mul_val = ops.mul(trunc_val, operand2) + res = ops.sub(operand1, mul_val) + return res, V.kernel.var_info[res] + else: + shape = f"vector<{tile_size}x{ret_type}>" if tile_size > 1 else ret_type + op_str = f'arith.remsi %{operand1}, %{operand2}' + return format_mlir_op(op_str, shape, **kwargs), [tile_size, ret_type] @staticmethod def isinf(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + if dtype.startswith("f"): + abs_val = ops.abs(operand) + inf_val = ops.constant("inf", dtype) + res = ops.eq(abs_val, inf_val) + return res, V.kernel.var_info[res] + else: + const_false = ops.constant(False, "i1") + if tile_size > 1: + const_false = ops.broadcast(const_false, tile_size) + return const_false, V.kernel.var_info[const_false] @staticmethod def isnan(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + if dtype.startswith("f"): + # Unordered comparison (uno) to detect NaN (uno returns true if either operand is NaN) + operand_shape = f"vector<{tile_size}x{dtype}>" if tile_size > 1 else dtype + op_str = f"arith.cmpf uno, %{operand}, %{operand}" + res = format_mlir_op(op_str, operand_shape, **kwargs) + + V.kernel.var_info[res] = [tile_size, "i1"] + return res, [tile_size, "i1"] + else: + # Integers cannot be NaN + const_false = ops.constant(False, "i1") + if tile_size > 1: + const_false = ops.broadcast(const_false, tile_size) + return const_false, V.kernel.var_info[const_false] @staticmethod def round(operand, *args, **kwargs): @@ -723,7 +983,30 @@ def floor(operand, *args, **kwargs): @staticmethod def sign(operand, *args, **kwargs): - raise NotImplementedError + tile_size, dtype = V.kernel.var_info[operand] + + if dtype.startswith("f"): + v_zero, v_one, v_neg_one = 0.0, 1.0, -1.0 + else: + v_zero, v_one, v_neg_one = 0, 1, -1 + + const_zero = ops.constant(v_zero, dtype) + const_one = ops.constant(v_one, dtype) + const_neg_one = ops.constant(v_neg_one, dtype) + + if tile_size > 1: + const_zero = ops.broadcast(const_zero, tile_size) + const_one = ops.broadcast(const_one, tile_size) + const_neg_one = ops.broadcast(const_neg_one, tile_size) + + is_pos = ops.gt(operand, const_zero) + is_neg = ops.lt(operand, const_zero) + + V.kernel.var_info[is_pos] = [tile_size, "i1"] + V.kernel.var_info[is_neg] = [tile_size, "i1"] + + res = ops.where(is_pos, const_one, ops.where(is_neg, const_neg_one, const_zero)) + return res, V.kernel.var_info[res] @staticmethod def trunc(operand, *args, **kwargs): @@ -933,11 +1216,11 @@ def xor(operand1, operand2, *args, **kwargs): @staticmethod def lshift(operand1, operand2, *args, **kwargs): - raise NotImplementedError + return ops.bitwise_left_shift(operand1, operand2) @staticmethod def rshift(operand1, operand2, *args, **kwargs): - raise NotImplementedError + return ops.bitwise_right_shift(operand1, operand2) @staticmethod def truncdiv(operand1, operand2, *args, **kwargs): diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py index d11c886a..5b1367ea 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py @@ -1,7 +1,7 @@ """Python port of the C++ `-test-pytorchsim-to-vcix` conversion pass (TestPyTorchSimToVCIXConversion.cpp). -Lowers `linalg.matmul` and the transcendental math ops (exp/erf/tanh/sin/cos) to +Lowers `linalg.matmul` and the transcendental math ops (exp/erf/tanh/sin/cos/log/atan) to VCIX dialect ops (RISC-V vector custom instructions). The C++ pass is a dialect-conversion (`applyPartialConversion`); the MLIR Python bindings expose no conversion framework, so each matchAndRewrite is reimplemented as imperative IR @@ -14,7 +14,7 @@ `allow_unregistered_dialects` -- so emitting generic vcix ops here is consistent with the current pipeline. -Covers all 6 C++ patterns: linalg.matmul (gemm + conv2d) and exp/erf/tanh/sin/cos. +Covers all 8 C++ patterns: linalg.matmul (gemm + conv2d) and exp/erf/tanh/sin/cos/log/atan. Wired into extension_codecache (run_to_vcix) after fine-grained, before the standard lowering; mlir-opt then runs only -test-loop-padding. Validated structurally against `mlir-opt -test-pytorchsim-to-vcix` (non-constant ops byte-identical incl. dma_wait tag @@ -31,15 +31,20 @@ from ._mlir_util import walk_ops, i32, i64, attr_bool -MARKERS = ("linalg.matmul", "math.exp", "math.erf", "math.tanh", "math.sin", "math.cos") +MARKERS = ("linalg.matmul", "math.exp", "math.erf", "math.tanh", "math.sin", + "math.cos", "math.log", "math.atan") -# math op name -> (opcode, imm) for the vcix.v.iv lowering (mirror Math*ToVCIX). +# math op name -> (opcode, imm) for the vcix.v.iv lowering (mirror Math*ToVCIX). The +# sf.vc.v.iv opcode field is 2 bits: erf(0)/tanh(1)/sin,cos,log,atan(2)/exp(3), and +# the shared opcode 2 is disambiguated by imm (sin=0, cos=1, log=2, atan=3). _MATH_VIV = { "math.exp": (0b000011, 0), "math.erf": (0b000000, 0), "math.tanh": (0b000001, 0), "math.sin": (0b000010, 0), "math.cos": (0b000010, 1), + "math.log": (0b000010, 2), + "math.atan": (0b000010, 3), } diff --git a/gem5_script/vpu_config.py b/gem5_script/vpu_config.py index 33d26b5f..2bed69b1 100644 --- a/gem5_script/vpu_config.py +++ b/gem5_script/vpu_config.py @@ -20,6 +20,8 @@ class SpecialFunctionUnit(MinorFU): "CustomVtanh", "CustomVsin", "CustomVcos", + "CustomVlog", + "CustomVatan", ]) opLat = 10 diff --git a/tests/ops/elementwise/test_pointwise.py b/tests/ops/elementwise/test_pointwise.py new file mode 100644 index 00000000..74e6656d --- /dev/null +++ b/tests/ops/elementwise/test_pointwise.py @@ -0,0 +1,214 @@ +import torch +import os + +def clear_caches(): + from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache + from torch._inductor.codecache import FxGraphCache + AOTAutogradCache.clear() + torch._dynamo.reset() + os.environ["TORCHINDUCTOR_CACHE"] = "0" + FxGraphCache.clear() + +def test_result(name, out, cpu_out, rtol=1e-4, atol=1e-4): + if torch.allclose(out.cpu(), cpu_out, rtol=rtol, atol=atol): + message = f"|{name} Test Passed|" + print("-" * len(message)) + print(message) + print("-" * len(message)) + else: + message = f"|{name} Test Failed|" + print("-" * len(message)) + print(message) + print("-" * len(message)) + print("custom out: ", out.cpu()) + print("cpu out: ", cpu_out) + exit(1) + +def run_test(name, device, fn, inputs, size_desc, rtol=1e-4, atol=1e-4): + """ + Harness function to compile, execute on NPU, compare with CPU, and print details. + inputs: single tensor or tuple/list of tensors (on CPU) + """ + if not isinstance(inputs, (tuple, list)): + inputs = [inputs] + + npu_inputs = [x.to(device=device) for x in inputs] + cpu_inputs = [x.clone() for x in inputs] + + clear_caches() + + opt_fn = torch.compile(dynamic=False)(fn) + res = opt_fn(*npu_inputs) + out = fn(*cpu_inputs) + + # Print input / output slices (up to 10 elements) + for idx, x in enumerate(inputs): + label = f"X" if len(inputs) == 1 else f"X{idx+1}" + val = x.flatten()[:10].tolist() if x.numel() > 1 else x.item() + print(f"[{size_desc}] {name} Input {label}: {val}") + + out_val = res.flatten()[:10].tolist() if res.numel() > 1 else res.item() + print(f"[{size_desc}] {name} Output: {out_val}") + + test_result(f"{name}_{size_desc}", res, out, rtol=rtol, atol=atol) + + +# aligned, scalar, small tail, large tail +STD_SIZES = [(128, 128), (1, 1), (15, 15), (129, 129)] + + +def run_op(name, device, fn, make_input, cases=None, sizes=STD_SIZES, rtol=1e-4, atol=1e-4): + torch.manual_seed(0) + for rows, cols in sizes: + run_test(name, device, fn, make_input(rows, cols), f"{rows}x{cols}", rtol=rtol, atol=atol) + for desc, inputs in (cases or []): + run_test(name, device, fn, inputs, desc, rtol=rtol, atol=atol) + + +def test_abs(device): + run_op("Abs", device, torch.abs, lambda r, c: torch.randn(r, c) * 10, + cases=[ + ("int", torch.randint(-100, 100, (128, 128), dtype=torch.int32)), + ("strided_transposed", (torch.randn(128, 128) * 10).t()), + ("strided_sliced", (torch.randn(128, 256) * 10)[:, ::2]), + ]) + +def test_sign(device): + def make(r, c): + x = torch.randn(r, c) + x[x.abs() < 0.3] = 0.0 + return x + run_op("Sign", device, torch.sign, make, + cases=[ + ("zero", torch.tensor([[0.0]])), + ("nonzero", torch.tensor([[-4.5]])), + ("int", torch.randint(-5, 5, (128, 128), dtype=torch.int32)), + ]) + +def test_isnan(device): + def make(r, c): + x = torch.randn(r, c) + x[x.abs() < 0.3] = float('nan') + return x + run_op("IsNaN", device, torch.isnan, make, + cases=[("scalar_nan", torch.tensor([[float('nan')]]))]) + +def test_isinf(device): + def make(r, c): + x = torch.randn(r, c) + x[x > 1.0] = float('inf') + x[x < -1.0] = float('-inf') + return x + run_op("IsInf", device, torch.isinf, make, + cases=[("scalar_inf", torch.tensor([[float('-inf')]]))]) + +def test_fmod(device): + run_op("Fmod", device, torch.fmod, + lambda r, c: (torch.randn(r, c) * 10, torch.randn(r, c) * 3 + 1), + cases=[("broadcast", (torch.randn(128, 1) * 10, torch.randn(1, 128) * 3 + 1))]) + +def test_lshift(device): + run_op("LShift", device, torch.bitwise_left_shift, + lambda r, c: (torch.randint(1, 100, (r, c), dtype=torch.int32), + torch.randint(1, 5, (r, c), dtype=torch.int32)), + cases=[("broadcast", (torch.randint(1, 100, (128, 1), dtype=torch.int32), + torch.randint(1, 5, (1, 128), dtype=torch.int32)))]) + +def test_rshift(device): + run_op("RShift", device, torch.bitwise_right_shift, + lambda r, c: (torch.randint(10, 1000, (r, c), dtype=torch.int32), + torch.randint(1, 5, (r, c), dtype=torch.int32)), + cases=[("broadcast", (torch.randint(10, 1000, (128, 1), dtype=torch.int32), + torch.randint(1, 5, (1, 128), dtype=torch.int32)))]) + +def test_copysign(device): + run_op("Copysign", device, torch.copysign, + lambda r, c: (torch.randn(r, c) * 10, torch.randn(r, c)), + cases=[ + ("negzero", (torch.tensor([[3.0]]), torch.tensor([[-0.0]]))), + ("broadcast", (torch.randn(128, 1) * 10, torch.randn(1, 128))), + ]) + +def test_erfc(device): + run_op("Erfc", device, torch.erfc, lambda r, c: torch.randn(r, c), + cases=[("large", torch.tensor([[4.5]]))]) + +def test_hypot(device): + run_op("Hypot", device, torch.hypot, + lambda r, c: (torch.randn(r, c), torch.randn(r, c)), + cases=[ + ("3-4-5", (torch.tensor([[3.0]]), torch.tensor([[4.0]]))), + ("broadcast", (torch.randn(128, 1), torch.randn(1, 128))), + ]) + +def test_cosh(device): + run_op("Cosh", device, torch.cosh, lambda r, c: torch.randn(r, c), + cases=[("large", torch.tensor([[4.5]]))]) + +def test_sinh(device): + run_op("Sinh", device, torch.sinh, lambda r, c: torch.randn(r, c), + cases=[("large", torch.tensor([[4.5]]))]) + +def test_log(device): + # domain: x > 0 + run_op("Log", device, torch.log, lambda r, c: torch.rand(r, c) + 1e-5, + cases=[("e", torch.tensor([[2.71828]]))]) + +def test_acosh(device): + # domain: x >= 1 + run_op("Acosh", device, torch.acosh, lambda r, c: torch.rand(r, c) + 1, + cases=[("one", torch.tensor([[1.0]]))]) + +def test_asinh(device): + run_op("Asinh", device, torch.asinh, lambda r, c: torch.randn(r, c)) + +def test_atanh(device): + # domain: (-1, 1) + run_op("Atanh", device, torch.atanh, lambda r, c: torch.rand(r, c) * 2 - 1) + +def test_atan(device): + # domain: all reals; boundary: zero, unit, large |x| (range-reduction path) + run_op("Atan", device, torch.atan, lambda r, c: torch.randn(r, c) * 5, + cases=[("boundary", torch.tensor([[0.0, 1.0, -1.0, 1e3, -1e3, 1e-4]]))]) + +def test_asin(device): + # domain: [-1, 1]; boundary includes the +/-1 endpoints + run_op("Asin", device, torch.asin, lambda r, c: torch.rand(r, c) * 2 - 1, + cases=[("boundary", torch.tensor([[0.0, 1.0, -1.0, 0.5, -0.5]]))]) + +def test_acos(device): + # domain: [-1, 1]; boundary includes the +/-1 endpoints + run_op("Acos", device, torch.acos, lambda r, c: torch.rand(r, c) * 2 - 1, + cases=[("boundary", torch.tensor([[0.0, 1.0, -1.0, 0.5, -0.5]]))]) + +def test_atan2(device): + # atan2(y, x); boundary: axes and diagonals across all four quadrants + # (0, 0) excluded: composition yields atan(0/0)=nan while torch returns 0 + y = torch.tensor([[1.0, 0.0, -1.0, 0.0, 1.0, -1.0, 1.0, -1.0]]) + x = torch.tensor([[0.0, 1.0, 0.0, -1.0, 1.0, -1.0, -1.0, 1.0]]) + run_op("Atan2", device, torch.atan2, lambda r, c: (torch.randn(r, c), torch.randn(r, c)), + cases=[("boundary", (y, x))]) + +if __name__ == "__main__": + device = torch.device("npu:0") + + test_abs(device) + test_sign(device) + test_isnan(device) + test_isinf(device) + test_fmod(device) + test_lshift(device) + test_rshift(device) + test_copysign(device) + test_erfc(device) + test_hypot(device) + test_cosh(device) + test_sinh(device) + test_log(device) + test_acosh(device) + test_asinh(device) + test_atanh(device) + test_atan(device) + test_asin(device) + test_acos(device) + test_atan2(device) \ No newline at end of file From 0cc908bf3162cd0ac6726dd5b0d9bcfa5eb5db79 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 9 Jul 2026 12:55:09 +0900 Subject: [PATCH 28/32] [Build] Pin spike v1.0.6 and gem5 v1.0.2 Bump the third-party pins to the releases shipping the new pointwise instructions (spike torchsim_vlog/vatan, gem5 CustomVlog/Vatan + decode). --- thirdparty/github-releases.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/thirdparty/github-releases.json b/thirdparty/github-releases.json index 2618c236..1265d380 100644 --- a/thirdparty/github-releases.json +++ b/thirdparty/github-releases.json @@ -3,7 +3,7 @@ "pytorch_image": "ubuntu:22.04", "gem5": { "repository": "PSAL-POSTECH/gem5", - "release_tag": "v1.0.1", + "release_tag": "v1.0.2", "asset_name": "gem5-release.tar.gz" }, "llvm_project": { @@ -13,7 +13,7 @@ }, "spike": { "repository": "PSAL-POSTECH/riscv-isa-sim", - "release_tag": "v1.0.5", + "release_tag": "v1.0.6", "asset_name": "spike-release.tar.gz" } } From b7698459320ec0cd4cba62b931196512bfeee9cd Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 8 Jul 2026 21:50:05 +0900 Subject: [PATCH 29/32] [Frontend] Op fixes behind SwinV2's shifted-window attention Make int max/min reductions dtype-aware, lower torch.roll to narrow+cat and realize it so the consumer reads a plain affine index, flatten nested single-variable floor/mod in axis-split, define empty_strided_cpu in the generated wrapper header, and place the reduction axis outermost in the per-lane tile layout. --- PyTorchSimFrontend/mlir/axis_split.py | 73 +++++++++++++++++++ .../mlir/mlir_codegen_backend.py | 32 +++++--- PyTorchSimFrontend/mlir/mlir_decomposition.py | 50 ++++++++++++- PyTorchSimFrontend/mlir/mlir_ops.py | 9 ++- tests/ops/reduce/test_reduce.py | 21 +++++- 5 files changed, 170 insertions(+), 15 deletions(-) diff --git a/PyTorchSimFrontend/mlir/axis_split.py b/PyTorchSimFrontend/mlir/axis_split.py index 71ec4809..2cf58c8a 100644 --- a/PyTorchSimFrontend/mlir/axis_split.py +++ b/PyTorchSimFrontend/mlir/axis_split.py @@ -29,6 +29,77 @@ def _as_int(x): return None +def _as_digit(expr): + """If ``expr`` is a single-variable *digit extractor* -- any nesting of FloorDiv / + ModularIndexing whose innermost argument is one symbol ``v`` -- return ``(v, div, mod)`` + meaning ``(v // div) % mod`` (``mod is None`` -> a pure ``v // div``). Otherwise None. + + Every such nesting collapses to a single (div, mod) by composing divisors, from four + algebraic identities (v >= 0, all constants positive integers): + FloorDiv(v//a, b) = v // (a*b) + FloorDiv((v//a)%m, b) = (v // (a*b)) % (m//b) if b | m + ModularIndexing(v//a, b,m2)= (v // (a*b)) % m2 + ModularIndexing((v//a)%m, b,m2)= (v // (a*b)) % m2 if b*m2| m + The divisibility guards make every rewrite provably equality-preserving. A multi- + variable inner argument (e.g. torch.roll's v+shift) is not a digit extractor -> None. + """ + if isinstance(expr, sympy.Symbol): + return (expr, 1, None) + if isinstance(expr, FloorDiv): + inner, b = expr.args + b, e = _as_int(b), _as_digit(expr.args[0]) + if e is None or b is None: + return None + v, a, m = e + if m is None: + return (v, a * b, None) + if m % b == 0: + return (v, a * b, m // b) + return None + if isinstance(expr, ModularIndexing): + inner, b, m2 = expr.args + b, m2, e = _as_int(b), _as_int(m2), _as_digit(expr.args[0]) + if e is None or b is None or m2 is None: + return None + v, a, m = e + if m is None: + return (v, a * b, m2) + if m % (b * m2) == 0: + return (v, a * b, m2) + return None + return None + + +def _rebuild_digit(v, a, m): + """Canonical single-level form of ``(v // a) % m``.""" + x = v if a == 1 else FloorDiv(v, a) + return x if m is None else ModularIndexing(v, a, m) + + +def flatten_nested_floormod(expr): + """Collapse nested single-variable FloorDiv/ModularIndexing to one level. + + A composition of aligned reshapes on one iteration variable leaves a nested index like + ModularIndexing(ModularIndexing(p, 1, 64), 1, 8) that neither sympy nor + simplify_with_ranges reduces, so collect_boundaries skips its cut points (the inner base + is not a bare var) and the affine-only DMA check later rejects it. Rewriting each nested + digit extractor to its single-level (v // A) % M form (via _as_digit) exposes those cut + points to axis-split. General and pattern-free -- no per-shape special cases. + """ + try: + atoms = expr.atoms(FloorDiv, ModularIndexing) + except AttributeError: + return expr + replace = {} + for atom in atoms: + e = _as_digit(atom) + if e is not None: + canon = _rebuild_digit(*e) + if canon != atom: + replace[atom] = canon + return expr.xreplace(replace) if replace else expr + + def collect_boundaries(exprs, var_to_axis, var_ranges): """{axis_index: set(boundary cut points)} for the given index expressions. @@ -39,6 +110,7 @@ def collect_boundaries(exprs, var_to_axis, var_ranges): import collections bset = collections.defaultdict(set) for expr in exprs: + expr = flatten_nested_floormod(expr) # nested digit extractors -> single level for fd in expr.atoms(FloorDiv): base, div = fd.args k = _as_int(div) @@ -196,6 +268,7 @@ def _fold_with_ranges(expr, var_ranges): Iterated to a fixpoint (folding a mod can expose a foldable floor). """ from torch.utils._sympy.value_ranges import bound_sympy, ValueRanges + expr = flatten_nested_floormod(expr) # collapse any residual single-var nested digit ranges = {} for v, sz in var_ranges.items(): e = _as_int(sz) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 590271c6..7705be4a 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -44,10 +44,16 @@ def reduction_init(reduction_type, dtype): return float(0) if dtype.is_floating_point else int(0) if reduction_type == "prod": return float(1) if dtype.is_floating_point else int(1) + # Integer reductions cannot use a +/-inf identity (invalid as an int constant and + # overflows torch.tensor(inf, dtype=int)); use the dtype's representable extreme. if reduction_type in {"max", "argmax"}: - return "-inf" + if dtype.is_floating_point: + return "-inf" + return 0 if dtype is torch.bool else torch.iinfo(dtype).min if reduction_type in {"min", "argmin"}: - return "inf" + if dtype.is_floating_point: + return "inf" + return 1 if dtype is torch.bool else torch.iinfo(dtype).max if reduction_type in {"welford_reduce"}: return f"0.0" raise AssertionError(reduction_type) @@ -114,6 +120,7 @@ def write_header(self): inductor_ops = torch.ops.inductor assert_size_stride = torch._C._dynamo.guards.assert_size_stride assert_alignment = torch._C._dynamo.guards.assert_alignment + empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu alloc_from_pool = torch.ops.inductor._alloc_from_pool reinterpret_tensor = torch.ops.inductor._reinterpret_tensor custom_async_compile = CustomAsyncCompile() @@ -1321,16 +1328,19 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride - # Case 4. Tile is 4-D tile (e.g., Convolution epilogue) - elif len(local_dims) == 4: - is_reduction = self.reduction_depth < 3 and not store_reduction - if is_reduction: - raise NotImplementedError("Currently not implemented... ;)") - local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) - local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis - local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride + # Case 4+. Tile is 4-D or higher (Convolution epilogue, gathered attention bias, + # var_mean over an axis whose batch dims got split into many loop vars). else: - local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) + # A reduction tile must place the reduction axis-group OUTERMOST in the + # per-lane layout, so the 2-D [reduction | batch] multi_reduction reduces the + # reduction axis rather than a batch axis left inner by row-major order. + is_reduction = any(d >= self.reduction_depth for d in local_dims) and not store_reduction + if is_reduction: + r = self.get_nr_rdim() + axis_order = list(range(r, len(local_dims))) + list(range(r - 1, -1, -1)) + local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims], axis_order) + else: + local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride diff --git a/PyTorchSimFrontend/mlir/mlir_decomposition.py b/PyTorchSimFrontend/mlir/mlir_decomposition.py index a0030e3b..f9ddbc31 100644 --- a/PyTorchSimFrontend/mlir/mlir_decomposition.py +++ b/PyTorchSimFrontend/mlir/mlir_decomposition.py @@ -371,4 +371,52 @@ def decompose_native_multi_head_attention( attn_weights_mean = attn_weights.mean(dim=1) # Average over heads return output, attn_weights_mean else: - return (output, None) \ No newline at end of file + return (output, None) + + +# Lower roll as narrow + cat, then REALIZE: torch's decomposition is a modular gather the +# affine-only DMA cannot express, and even narrow+cat fuses into a modular reshape. +# copy_input forces the buffer boundary .contiguous() does not. +from torch._inductor.decomposition import decompositions as _inductor_decompositions +from torch._inductor import lowering as _ind_lowering +from torch._inductor import ir as _ind_ir + +_inductor_decompositions.pop(aten.roll.default, None) + + +def _roll_lowering(x, shifts, dims=()): + if not isinstance(shifts, (list, tuple)): + shifts = (shifts,) + if not isinstance(dims, (list, tuple)): + dims = (dims,) + slice_l = _ind_lowering.lowerings[aten.slice.Tensor] + cat_l = _ind_lowering.lowerings[aten.cat.default] + view_l = _ind_lowering.lowerings[aten.view.default] + + # Realize the INPUT: roll's producer is often a reshaped view (e.g. swinv2's window_reverse), + # and our slice's offset would otherwise fuse into that reshape's modular index. Reading a + # materialized buffer makes each slice a plain contiguous read. + x = _ind_ir.ExternKernel.copy_input(x) + + def roll_dim(t, shift, dim): + n = int(t.get_size()[dim]) + s = int(shift) % n + if s == 0: + return t + front = slice_l(t, dim, n - s, n) # narrow(t, dim, n-s, s) + back = slice_l(t, dim, 0, n - s) + return cat_l([front, back], dim) + + if len(dims) == 0: + numel = 1 + for d in x.get_size(): + numel *= int(d) + result = view_l(roll_dim(view_l(x, [numel]), shifts[0], 0), list(x.get_size())) + else: + result = x + for shift, dim in zip(shifts, dims): + result = roll_dim(result, shift, dim) + return _ind_ir.ExternKernel.copy_input(result) + + +_ind_lowering.lowerings[aten.roll.default] = _roll_lowering \ No newline at end of file diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index d21cc1d0..98e0272b 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -13,10 +13,15 @@ def reduction_combine_vec(reduction_type, vector_value, init_value, axis, shape, return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" if reduction_type == "prod": return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" + # element type of the reduced vector (e.g. "vector<8x2xi64>" -> "i64"); int max/min use + # the signed-integer reduction kinds, not the float ones (maximumf/minimumf reject ints). + _is_int = shape.rsplit("x", 1)[-1].rstrip(">").strip().startswith("i") if reduction_type == "max": - return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" + kind = "maxsi" if _is_int else "maximumf" + return f"vector.multi_reduction <{kind}>, %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" if reduction_type == "min": - return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" + kind = "minsi" if _is_int else "minimumf" + return f"vector.multi_reduction <{kind}>, %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" if reduction_type == "any": return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" raise AssertionError(reduction_type) diff --git a/tests/ops/reduce/test_reduce.py b/tests/ops/reduce/test_reduce.py index d80e17b9..c829261b 100644 --- a/tests/ops/reduce/test_reduce.py +++ b/tests/ops/reduce/test_reduce.py @@ -25,6 +25,24 @@ def reduce_sum(a, dim, keepdim): out = reduce_sum(x.cpu(), dim, keepdim) test_result("ReduceMax", res, out) +def test_reduce_gather_bias(device, NW=4, H=3, Q=32, K=32, T=64): + """A reduction fused with an INDIRECT gather bias, as in SwinV2 cosine-window + attention: score[w,h,q,k] + table[idx[q,k], h] -> amax over the key axis. The gather + blocks the head*query dim-merge, so the reduction tile stays 4-D. The reduction axis + must be hoisted to the outermost in-lane position; the 4-D reduction tile path used to + skip that reorder and reduce a batch axis (head) instead of the key axis, so head 0's + max picked up head 1's values (needs H>=2 and NW>=2 to expose the head bleed).""" + def fn(score, idx, table): + bias = table[idx.reshape(-1)].reshape(Q, K, H).permute(2, 0, 1).unsqueeze(0) + return (score + bias).amax(dim=-1) + torch.manual_seed(0) + score = torch.randn(NW, H, Q, K).to(device=device) + idx = torch.randint(0, T, (Q, K)).to(device=device) + table = torch.randn(T, H).to(device=device) + res = torch.compile(dynamic=False)(fn)(score, idx, table) + out = fn(score.cpu(), idx.cpu(), table.cpu()) + test_result("ReduceGatherBias", res, out) + if __name__ == "__main__": import argparse @@ -38,4 +56,5 @@ def reduce_sum(a, dim, keepdim): test_reduce_sum(device, (17, 68), 0, keepdim=True) test_reduce_sum(device, (327, 447), 1, keepdim=True) test_reduce_sum(device, (327, 447), 0, keepdim=True) - test_reduce_sum2(device, shape) \ No newline at end of file + test_reduce_sum2(device, shape) + test_reduce_gather_bias(device) \ No newline at end of file From 77c73382f973e5f020809d12b55bf3cfd877b1e2 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 9 Jul 2026 15:19:50 +0900 Subject: [PATCH 30/32] [Model] Enable SwinV2 shifted-window attention (issue #251) SwinV2 batch>1 used to fail codegen with "Unlinearized floor/mod in DMA index" because the shifted-window path (torch.roll composed with window partition/reverse) produced views axis-split could not linearize. With the roll->slice+cat lowering, the nested/shifted mod handling, and the reduction-axis layout fix, the whole backbone now compiles and matches CPU end to end (max diff ~5e-6 for image 64 / window 8 / batch 2). Add tests/models/test_swinv2.py, wire it as a self-hosted CI job next to the other model tests, and list SwinV2 in the README model coverage. --- .github/workflows/pytorchsim_test.yml | 25 +++++++++- README.md | 1 + tests/models/test_swinv2.py | 69 +++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_swinv2.py diff --git a/.github/workflows/pytorchsim_test.yml b/.github/workflows/pytorchsim_test.yml index 318c7945..749f79d6 100644 --- a/.github/workflows/pytorchsim_test.yml +++ b/.github/workflows/pytorchsim_test.yml @@ -18,8 +18,9 @@ on: # Runner policy: the CPU-only CI image is small enough to pull on GitHub-hosted # runners, so op and model tests run on ubuntu-latest. The memory/time-intensive -# jobs stay on self-hosted: test_deepseek (largest model), test_diffusion (UNet2D -# simulation OOMs the hosted runner), and test_accuracy (accuracy + speedup). +# jobs stay on self-hosted: test_deepseek (largest model), test_swinv2 (SwinV2 +# shifted-window backbone), test_diffusion (UNet2D simulation OOMs the hosted +# runner), and test_accuracy (accuracy + speedup). jobs: test_add: name: Run test_add.py @@ -743,6 +744,26 @@ jobs: -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ ${{ inputs.image_name }} python3 PyTorchSim/tests/models/DeepSeek/test_deepseek_v3_base.py + test_swinv2: + name: Run test_swinv2.py + # SwinV2 backbone (shifted-window attention) is heavy; keep it on a + # self-hosted runner like the other model tests. + runs-on: self-hosted + steps: + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Run test_swinv2.py + run: | + echo "Running test_swinv2.py" + docker run --rm \ + -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ + ${{ inputs.image_name }} python3 PyTorchSim/tests/models/test_swinv2.py + test_eager: name: Run test_eager.py runs-on: ubuntu-latest diff --git a/README.md b/README.md index c2298376..256196f2 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,7 @@ PyTorchSim **supports**: | Stable-diffusion v1 | 🤗 | ✅ | | | Llama 2/3 | 🤗 | ✅ | `tests/models/Llama/` (blocks & decode-style paths) | | DeepSeek-V3 (base) | 🤗 | ✅ | `tests/models/DeepSeek/` — several ops(e.g., gate ops) are not cycle-modeled | +| SwinV2 | 🤗 | ✅ | `tests/models/test_swinv2.py` (shifted-window attention) | | Llama-4 | 🤗 | ⏳ | In development | | Broader model support | — | ⏳ | In development |