From c8dd94e54f7816274385a390442f9a2d149a4973 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 14:06:41 +0800 Subject: [PATCH 1/2] fix(analysis): normalize memory metadata and traffic --- docs/spec/analysis.md | 263 +++++++----------- docs/spec/cli.md | 8 +- docs/spec/code-organization.md | 6 +- docs/spec/inspection.md | 5 +- docs/tutorial/optimize.ipynb | 8 +- docs/tutorial/optimize.md | 8 +- docs/tutorial/showcase.ipynb | 34 +-- docs/tutorial/showcase.md | 85 +++--- src/tilefoundry/analysis/__init__.py | 22 +- src/tilefoundry/analysis/access.py | 4 +- src/tilefoundry/analysis/check.py | 4 +- src/tilefoundry/analysis/compute_cost.py | 19 +- src/tilefoundry/analysis/iteration_scope.py | 127 +-------- src/tilefoundry/analysis/memory.py | 227 ++++++++------- src/tilefoundry/analysis/metadata.py | 125 +++------ src/tilefoundry/analysis/performance.py | 4 +- src/tilefoundry/analysis/registry.py | 5 +- src/tilefoundry/analysis/report.py | 41 +-- src/tilefoundry/analysis/roofline.py | 13 +- src/tilefoundry/cli/analyze.py | 12 +- src/tilefoundry/inspection/analysis_report.py | 14 +- src/tilefoundry/inspection/values.py | 41 +-- tests/analysis/test_analysis_families.py | 69 +++-- tests/analysis/test_analysis_invariants.py | 10 +- tests/analysis/test_analyze_at_a_size.py | 58 ++-- tests/analysis/test_analyze_by_hand.py | 201 +++++++++++++ tests/analysis/test_analyze_cross_module.py | 12 +- tests/cli/test_cli.py | 27 +- .../type_printer_sugar.analyzed.txt | 30 +- tests/fixtures/placed/hand_checked.py | 39 +++ tests/installed/models/contract.py | 14 +- tests/installed/smoke_analyze.py | 46 ++- tests/installed/smoke_target/smoke_v100.py | 21 +- tests/ops/ir/test_cache_update.py | 4 +- tests/ops/ir/test_local.py | 6 +- 35 files changed, 809 insertions(+), 803 deletions(-) create mode 100644 tests/analysis/test_analyze_by_hand.py create mode 100644 tests/fixtures/placed/hand_checked.py diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 98724821..87f987aa 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -41,7 +41,7 @@ Each owns its record types and declares its dependencies and output additions. | Selector | Requires | Owns | Attaches to | Rests on | Text summary adds | Annotates equations | |---|---|---|---|---|---|---| | `compute-cost` | - | `ComputeCostMetadata` | every measured Call and the Function | the authored program | `compute-cost` | every measured Call | -| `memory` | - | `MemoryMetadata`, `TrafficMetadata`, `LoopFootprintMetadata` | `MemoryMetadata` on the Function; `TrafficMetadata` on every measured Call and the Function; `LoopFootprintMetadata` on every `LoopRegion` | the authored program, `MemoryHierarchyFacts` | `peak-footprint`, `traffic`, `advisory` | none | +| `memory` | - | `MemoryMetadata`, `RegionMemoryMetadata` | `MemoryMetadata` on every measured Call; `RegionMemoryMetadata` on the Function | the authored program, `MemoryHierarchyFacts` | `memory`, `advisory` | every measured Call | | `roofline` | `compute-cost`, `memory` | `RooflineMetadata` | every measured Call and the Function | `ThroughputFacts` | `roofline` | every measured Call | | `performance` | `compute-cost`, `memory` | `PerformanceMetadata`, `PerformanceSummaryMetadata` | `PerformanceMetadata` on every Call with a modeled duration; `PerformanceSummaryMetadata` on the Function | `ThroughputFacts`, `ParallelCapacityFacts`, `MemoryHierarchyFacts` | `performance` | every Call with a modeled duration | @@ -63,11 +63,10 @@ The JSON report carries the same identity and selection in `target`, `module`, `function`, `topology`, `requested`, and `executed`. Whole-function projections are under `function_records`; `calls` is a value-ordered list whose entries have a `value` label and one key per selected family. `loops` is the -corresponding authored-loop list, labelled by induction variable. When memory is -selected, a loop whose backing storage has a same-scope implicit cache also has -`cache-pressure`: one target-aware row per cache, computed from the loop's -device-wide access footprint. `totals` appears when the selected view includes -compute cost or roofline's bounded work evidence. +corresponding authored-loop list, labelled by induction variable. Memory does not +attach an empty record to a loop for a conclusion it has not computed. `totals` +appears when the selected view includes compute cost or roofline's bounded work +evidence. One result is rendered once, and every surface reads that rendering: @@ -216,10 +215,11 @@ Each reported Call's JSON projection is under its `compute-cost` key: #### 1.2.2 `memory` -`memory` measures whole-Function value lifetimes and footprints, decides where -each value's bytes live, and states what every occurrence moves and at which -level. The movement is read off the Op's own registered evaluator and the -amounts its access relations reach. +`memory` measures whole-Function value lifetimes and placement peaks, and states +what every occurrence moves and at which level. The movement is read off the +Op's own registered evaluator and the amounts its access relations reach. The +records reserve an optional read-footprint conclusion, but this analysis does +not currently produce one: both attachment points report `footprint=None`. ```python class Spread[V]: @@ -251,23 +251,30 @@ class Breakdown[V]: kinds: tuple[tuple[str, Spread[V]], ...] = () -class TrafficMetadata(IRMetadata): - """What one Call moves, or what one Function moves over all its trips. +class Traffic: + """Movement grouped by storage and communication boundary.""" - Attributes: - topologies: attribute; The declared levels, in the order per_unit states them. - storage: attribute; TrafficBytes per storage level name. - communication: attribute; TrafficBytes per topology level name, for the boundary the bytes crossed. - operands: attribute; TrafficBytes per operand, positional against (*call.args, call); present only for a direct primitive call. - """ - - topologies: tuple[str, ...] = () storage: Breakdown[TrafficBytes] = Breakdown() communication: Breakdown[TrafficBytes] = Breakdown() + + +class Footprint: + """Unique read bytes, or the available lower bound when incomplete.""" + + buffers: tuple[tuple[str, Breakdown[int]], ...] = () + complete: bool = True + + +class MemoryMetadata(IRMetadata): + """One primitive Call's memory behavior for one occurrence.""" + + topologies: tuple[str, ...] = () + traffic: Traffic = Traffic() operands: tuple[TrafficBytes, ...] = () + footprint: Footprint | None = None -class MemoryLevelFootprint: +class MemoryLevelPeak: """How much of one memory level a function needs at its peak. Attributes: @@ -283,34 +290,6 @@ class MemoryLevelFootprint: persistent_bytes: int capacity_bytes: int | None = None -class BufferFootprint: - """Bytes one authored loop touches in one buffer at one storage level. - - Attributes: - buffer: attribute; The stable value name of the source buffer. - memory_level: attribute; The storage level containing that buffer. - bytes: attribute; Deduplicated bytes reached by one logical position. - device_bytes: attribute; Deduplicated bytes in the union across positions. - repeated_bytes: attribute; Per-position bytes without deduplicating repeated access. - """ - - buffer: str - memory_level: str - bytes: int - device_bytes: int - repeated_bytes: int - -class LoopFootprintMetadata(IRMetadata): - """Known buffer accesses or a lower bound within one authored LoopRegion. - - Attributes: - footprints: attribute; One row per source buffer and storage level. - known: attribute; Whether every access had a representable relation. - """ - - footprints: tuple[BufferFootprint, ...] - known: bool - class ValueLifetime: """One value's residency, as positions in the function's value order. @@ -330,31 +309,17 @@ class ValueLifetime: last_used_at: int persistent: bool = False -class AllocationMetadata: - """What showing this function's buffers fit came to. - - Attributes: - solver_status: attribute; `"feasible"` for the first validated placement. - """ +class RegionMemoryMetadata(IRMetadata): + """One Function's aggregate movement and placement conclusions.""" solver_status: str - -class MemoryMetadata(IRMetadata): - """One function's memory behaviour against one target's hierarchy. - - Attributes: - footprint: attribute; One row per level the function places values in. - lifetimes: attribute; One entry per value residency. - errors: attribute; Solved placement peaks that exceed stated capacity. - advisories: attribute; Capacity findings that do not invalidate the program. - allocation: attribute; What showing the addressable buffers fit came to. - """ - - footprint: tuple[MemoryLevelFootprint, ...] = () + topologies: tuple[str, ...] = () + traffic: Traffic = Traffic() + footprint: Footprint | None = None lifetimes: tuple[ValueLifetime, ...] = () + peaks: tuple[MemoryLevelPeak, ...] = () errors: tuple[str, ...] = () advisories: tuple[str, ...] = () - allocation: AllocationMetadata | None = None ``` Every traffic amount here is what a boundary's own relation reaches. The Op's @@ -398,10 +363,9 @@ the largest single projected logical value. MUST NOT make a program unplaceable, and a level owned per unit of a topology level other than the one being analysed MUST fail rather than be assumed. Domains holding the same buffers are one question, decided once. - - `allocation` MUST be absent only when no level could be projected against. A - function with nothing addressable MUST record a settled `allocation`: the - question was asked and there was nothing to decide. An attached - `solver_status` MUST be `"feasible"`; a domain that cannot be expressed or + - `RegionMemoryMetadata.solver_status` MUST be `"feasible"`, including when + the function has no addressable value: the question was asked and there was + nothing to place. A domain that cannot be expressed or does not settle in time MUST raise `AnalysisError` and leave no record. The solver MUST stop at its first feasible assignment rather than spend the remaining timeout proving a minimum. Its reported peak is that assignment's @@ -425,43 +389,36 @@ the largest single projected logical value. is longest, and MUST NOT sum them: one movement spends two resources over one span of time. A crossing at a level the target publishes no rate for MUST be stated and left untimed. - - A Call's `storage`, `communication` and `operands` MUST state one occurrence. Only - the Function record counts an occurrence as often as its authored loops - repeat it, and its `operands` MUST be empty: which operand moved what - belongs to the occurrence, not to the total. + - A Call's `traffic` and `operands` MUST state one occurrence. Only the Function + record counts an occurrence as often as its authored loops repeat it; which + operand moved what belongs to the occurrence, not to the region total. - A capacity conclusion MUST NOT correct or invent a movement number. What an - occurrence moves is counted once from its own boundaries, so a function with - no `allocation` still carries traffic -- a different question from whether a - time may be reported for it ([§1.2.4](#124-performance)) -- and a window + occurrence moves is counted once from its own boundaries, so placement and + traffic remain separate fields of one record. A window whose start arrives at run time reads that start rather than becoming a full read of its source and a write of its result. | Field | How it is computed | Reads the target | |---|---|---| -| `BufferFootprint.buffer` | The label value lifetimes use for that value, from the same derivation. Grouping stays by buffer identity, because two structurally equal buffers are distinct allocations; the label never reads an address. | No | -| `BufferFootprint.memory_level` | Read the source buffer's declared storage level. | No | -| `BufferFootprint.bytes` | Build relations from rank-preserving per-position Types, union the loop-prefixed access images, count the union's integer points, multiply by the dtype bit width, then round the whole buffer reading up to bytes. If the count is not an integer or exceeds `repeated_bytes`, that buffer reading is unavailable. | No | -| `BufferFootprint.device_bytes` | Repeat the same exact union measurement from authored Types without shard narrowing, giving the union across logical positions in bytes. | No | -| `BufferFootprint.repeated_bytes` | Multiply each operand's per-position element count by its enclosing trip counts, sum accesses to the same buffer, multiply by dtype bit width, then round the whole buffer reading up to bytes. | No | -| `LoopFootprintMetadata.footprints` | One `BufferFootprint` per known source buffer and storage level, sorted by buffer then level. When `known` is false these rows are the available lower bound rather than an empty replacement. | No | -| `LoopFootprintMetadata.known` | False when an access in the loop or a descendant loop lacks a representable forward relation, marking `footprints` as a lower bound; true otherwise. | No | | `ValueLifetime.binding` | Use the parameter or binding name, suffixed with `:` and the line of the value's source span when it has one. Repeated names already differ by the printer's numeric suffix in definition order; the line locates the row in authored source, which a suffix cannot. A value with neither name nor span is `` in definition order. | No | | `ValueLifetime.memory_level` | Emit one lifetime per storage level occupied by the value's Type. | No | | `ValueLifetime.bytes` | Project the Type through every authored split at or coarser than the explicit level's `owner`, then take its logical bytes; a target-owned or undeclared level remains global. | `MemoryHierarchyFacts.explicit_levels[].owner` | | `ValueLifetime.defined_at` | Definition event on the function-wide structured SSA timeline. | No | | `ValueLifetime.last_used_at` | Greatest ordinary-consumer, region-entry, loop-backedge, or region-exit use event; the final timeline event for a parameter. | No | | `ValueLifetime.persistent` | True for parameters and false for body allocations. | No | -| `MemoryLevelFootprint.memory_level` | Each storage level with at least one lifetime, sorted by name. | No | -| `MemoryLevelFootprint.peak_bytes` | For `gmem` and `smem`, the address high-water mark of the first feasible whole-Function placement. Exact pointwise relations and exact `insert_slice` partitions may permit overlap; widened or unknown relations do not. For `rmem`, the largest single projected logical value, without address placement or cross-value summation. | No | -| `MemoryLevelFootprint.persistent_bytes` | Sum of persistent lifetimes at that level. | No | -| `MemoryLevelFootprint.capacity_bytes` | Capacity of the matching explicit level, or `None` when it is unknown or undeclared. | `MemoryHierarchyFacts.explicit_levels[].capacity_bytes` | -| `MemoryMetadata.footprint` | One `MemoryLevelFootprint` per occupied storage level. | As above | -| `MemoryMetadata.lifetimes` | Every value residency except a `Reshape` or a `Transpose`, each of which describes bytes its operand already holds. | As above | -| `MemoryMetadata.errors` | One non-fatal error for each solved `gmem` or `smem` placement whose `peak_bytes` exceeds stated `capacity_bytes`; this includes a single value larger than capacity. | `MemoryHierarchyFacts.explicit_levels[].capacity_bytes` | -| `MemoryMetadata.advisories` | Cache/shared-capacity division and same-scope authored-loop access-footprint findings. | `MemoryHierarchyFacts` | -| `TrafficMetadata.storage` | One occurrence's per-boundary movement asked of the Op's access relations, charged to the storage levels its operand Types name. The total is asked in the whole program's window and each level's share in that level's, over Types projected through the authored `Split`s at or coarser than it. On a Function, summed over every reachable occurrence, each counted as often as its authored loops repeat it. A Type with leaves at several levels keeps those leaf bytes separate. A `UMAT` leaf has no residency of its own: when it appears in `Call.args`, charge its own bytes at the target's established `rmem` materialization level; when it appears only in an Op attribute, charge nothing. A Function Call takes the callee's grouped total. | No; projection reads resolved Mesh and effective Module topology extents. | -| `TrafficMetadata.communication` | What a Reshard sends off the unit it was on, when the shards on its two sides differ across a mesh axis that level owns. Zero where they agree. | No; the share each unit keeps follows from the mesh extents the shards name. | -| `TrafficMetadata.operands` | One occurrence's per-boundary movement in order `(*call.args, call)`, the same relation-derived amounts `storage` groups. Empty on a Function and on a Function Call, neither of which has a split. | No | +| `MemoryLevelPeak.memory_level` | Each storage level with at least one lifetime or traffic entry, sorted by name. | No | +| `MemoryLevelPeak.peak_bytes` | For `gmem` and `smem`, the address high-water mark of the first feasible whole-Function placement. Exact pointwise relations and exact `insert_slice` partitions may permit overlap; widened or unknown relations do not. For `rmem`, the largest single projected logical value. | No | +| `MemoryLevelPeak.persistent_bytes` | Sum of persistent lifetimes at that level. | No | +| `MemoryLevelPeak.capacity_bytes` | Capacity of the matching explicit level, or `None` when unknown. | `MemoryHierarchyFacts.explicit_levels[].capacity_bytes` | +| `MemoryMetadata.traffic` | One occurrence's per-boundary movement, grouped by storage and communication boundary. | No; projection reads resolved Mesh and topology extents. | +| `MemoryMetadata.operands` | One occurrence's movement in order `(*call.args, call)`. | No | +| `RegionMemoryMetadata.traffic` | Every reachable occurrence. `logical` multiplies only loops the value varies in; `total` and `per_unit` multiply every enclosing loop. | No | +| `RegionMemoryMetadata.footprint` | `None` until a read-footprint analysis has produced a conclusion; it MUST NOT use an empty `Footprint` to mean “not computed”. | No | +| `RegionMemoryMetadata.lifetimes` | Every value residency except a non-material view. | As above | +| `RegionMemoryMetadata.peaks` | One `MemoryLevelPeak` per occupied or moved storage level. | As above | +| `RegionMemoryMetadata.solver_status` | The validated whole-Function placement status. | No | +| `RegionMemoryMetadata.errors` | One non-fatal error per placement peak exceeding capacity. | `MemoryHierarchyFacts.explicit_levels[].capacity_bytes` | +| `RegionMemoryMetadata.advisories` | Lower-severity target-aware memory findings. | `MemoryHierarchyFacts` | One ordinary expression event uses its operands and defines its result. A region adds separate binding and exit events: a mesh argument is used before its @@ -472,25 +429,7 @@ exit, and a yielded value remains live through the backedge event. Event positions are monotonic across the whole Function, including nested and sibling regions. -The target-aware loop projection is report data rather than another metadata -record. `LoopFootprintMetadata` remains target-independent: - -```text -"cache-pressure": [{"cache_level": , "backing_level": , - "device_bytes": , "capacity_bytes": , - "status": "fits"|"exceeds"|"lower-bound"|"unknown"}, ...] -``` - -The projection MUST use `device_bytes`, sum rows at the cache's ultimate explicit -backing level, and compare only levels whose capacity scopes agree. A missing -backing level or a scope mismatch MUST emit no row and MUST NOT fail analysis. -`lower-bound` means an incomplete footprint has not yet exceeded capacity; -`exceeds` remains conclusive when the lower bound alone exceeds it. A cache with -no usable capacity emits `unknown`. A buffer with `device_bytes < bytes` MUST be -removed before projection and its `LoopFootprintMetadata` MUST be marked -incomplete. - -The family reads this target projection: +The family reads this target hierarchy: ```python class MemoryRelationKind(Enum): @@ -556,52 +495,42 @@ class MemoryHierarchyFacts: relations: tuple[MemoryLevelRelation, ...] ``` -Requesting memory adds the Function's own movement, one footprint line, and one -line per advisory: +Requesting memory adds one Function memory line and one line per finding: ```text -traffic traffic=:r/w@total,r/w@[,...][;:...] -peak-footprint=:[,:...] +memory traffic=:r/w@logical,r/w@total,r/w@[,...] peak=:[,...] error="" advisory="" ``` -An empty footprint states the family name alone; each error and advisory is its -own line and is quoted and escaped -([inspection §2.8](./inspection.md#28-record-comment-forms)). The record's own -comment form projects the footprint it holds, and `lifetimes` is read from JSON: +Each error and advisory is its own quoted and escaped line. Every measured Call +receives a `memory` annotation; its `operands` split is emitted only when asked +for ([cli Analyze](./cli.md#analyze)): ```text -memory peak=:[,...] persistent= errors= advisories= +memory traffic=:r/w@logical,r/w@total,r/w@[,...] [operands=:r/w[;:...]] ``` -Every measured Call also receives a `traffic` annotation, whose `operands` split -is emitted only when asked for ([cli Analyze](./cli.md#analyze)) and is absent -from a Function, which has no split: +Call and Function JSON projections are both under `memory`. The Function's full +projection is under `function_records.memory`: ```text -traffic traffic=:r/w@total,r/w@[,...][;:...] [operands=:r/w[;:...]] -``` - -Its JSON projection is under the reported value's `traffic` key, with `whole`, -`communication` and one `operands` entry per position carrying `read` and `write`. -The `analyze` equation printer emits no memory annotation because that record is -attached only to the Function. Its full JSON projection is under -`function_records.memory`: - -```text -{"footprint": [{"memory_level": , "peak_bytes": , - "persistent_bytes": , "capacity_bytes": }, ...], - "traffic": {: {"read": , "write": }}, +{"topologies": [, ...], + "traffic": {"storage": {: , ...}, + "communication": {: , ...}}, + "footprint": null, "lifetimes": [{"binding": , "memory_level": , "bytes": , "defined_at": , "last_used_at": , "persistent": }, ...], + "peaks": [{"memory_level": , "peak_bytes": , + "persistent_bytes": , "capacity_bytes": }, ...], + "solver_status": "feasible", "errors": [, ...], "advisories": [, ...]} ``` - constraints: - - `MemoryMetadata` MUST be attached per reachable `Function`; a peak spans its + - `RegionMemoryMetadata` MUST be attached per reachable `Function`; a peak spans its live ranges and belongs to no single expression. - `Reshape` and `Transpose` describe bytes their operand already holds and MUST NOT receive independent lifetimes. Every other result, a window and a @@ -636,8 +565,9 @@ attached only to the Function. Its full JSON projection is under capacity scope. - A solved explicit-level peak exceeding capacity, whether from one value or the aggregate placement, MUST produce a report `error` and MUST NOT fail the - call. An authored-loop access footprint exceeding an implicit cache capacity - MUST instead produce an advisory. + call. This contract produces no authored-loop cache conclusion; an + `advisory` MUST come from a recorded Function-level finding rather than a + report-layer reconstruction. #### 1.2.3 `roofline` @@ -744,11 +674,10 @@ Reported Call and Function records use the same projection under their ``` When roofline is requested without its dependencies being requested, `totals` -carries only exact `flops` and `traffic` sums, and `function_records.memory` -carries only `memory_level` and `peak_bytes` per footprint row. Persistent bytes, -capacities, advisories, lifetimes, operand splits, and dependency annotations do -not enter that view. Independently requesting a dependency selects its full form -as defined in that family's section. +carries only exact `flops`, `traffic`, and `communication` sums. Dependency +records remain on the semantic result but do not enter `function_records`, Call +annotations, or Call report rows. Independently requesting a dependency selects +its full form as defined in that family's section. - constraints: - `ThroughputFacts.peak_for` MUST return `None` for a dtype with no published @@ -761,17 +690,17 @@ as defined in that family's section. bandwidth rather than summing traffic across levels. - Performance local duration MUST divide one unit's share of `ComputeCostMetadata.flops` by `unit_flops`, of `service` by `unit_ops`, and - the `bandwidth_level` entry of `TrafficMetadata.storage` by + the `bandwidth_level` entry of `MemoryMetadata.traffic.storage` by `unit_bandwidth`, all at the level it was asked about. Compute and movement overlap within one occurrence, so its duration is the greater of the two sides rather than their sum. - Traffic at a level with no stated one-unit bandwidth MUST remain visible in - `TrafficMetadata` and MUST NOT enter a duration: an instruction throughput + `MemoryMetadata` and MUST NOT enter a duration: an instruction throughput standing in for a bandwidth prices a move as though it were arithmetic. - Having moved bytes and having work this can time are different questions. What decides the second is the quantities a rate exists for: a nonzero share of `flops` or of `service`, or nonzero - `TrafficMetadata.storage` at `bandwidth_level`. An occurrence with none of + `MemoryMetadata.traffic.storage` at `bandwidth_level`. An occurrence with none of them MUST take zero time, MUST NOT be required to carry an execution placement, and MUST still record its movement at any other level: it is untimed, not absent. Work of a dtype or kind the target states no one-unit @@ -933,12 +862,12 @@ model. An occurrence with no nonzero share of `flops`, none of `service` and no nonzero - `TrafficMetadata.storage` at `bandwidth_level` is structural to this + `MemoryMetadata.traffic.storage` at `bandwidth_level` is structural to this model: it needs no execution placement and MUST receive no record, because an empty interval reads as a measurement rather than as the absence of one. Movement at another level does not change that and MUST NOT be dropped from - `TrafficMetadata` because of it -- structural here means nothing to time, + `MemoryMetadata` because of it -- structural here means nothing to time, not nothing done. It still carries its producers' precedence to its consumers. Inputs MUST NOT supply placement for an unplaced occurrence. - The global total for an occurrence is its per-unit quantity multiplied by @@ -984,9 +913,9 @@ model. - `parallel_units` is compiler policy over hardware facts. It MUST NOT enter one-unit rates or the CTA-local layout, and is not a program rewrite. - The buffers a plan keeps live MUST have been placed by `memory` before a - time is reported for it, which a `MemoryMetadata` carrying an `allocation` - is the evidence of; one without it MUST fail with `AnalysisError`. A - placement that failed never reaches here, because `memory` refuses it. + time is reported for it. A successful dependency records + `RegionMemoryMetadata.solver_status="feasible"`; a placement that failed + never reaches performance because `memory` refuses it. Capacity therefore changes whether there is an answer, never which answer: two capacities that both admit a placement MUST produce the same intervals. - Performance is a modeled plan and MUST NOT be read as a guarantee about @@ -1209,14 +1138,16 @@ def analyze( ### 2.1 Shared IterationScope and Access The normalized HIR is visited once per `analyze()` call. That visit produces a -`IterationScope` tree parallel to Function/LoopRegion nesting and `Access` -relations for the narrow and device views. `IterationScope.domain` is the -accumulated authored loop domain; `IterationScope.accesses` and -`IterationScope.refused` are the only family inputs for -loop footprints, movement, and placement. An `Access` stores only its relation -and allocation expression; storage level and element width are read from the -allocation type. A refused descendant makes its owning scope unknown for that -view. Non-affine runtime indices retain the widest legal access approximation. +`IterationScope` tree parallel to Function/LoopRegion nesting and `Access` relations +for the narrow and device views. `IterationScope.domain` is the accumulated +authored loop domain; `IterationScope.accesses` and `IterationScope.refused` are +the shared family inputs for movement and any later read-footprint conclusion. +An input `Access` stores its original Call boundary index as well as its relation +and allocation expression; an output stores `input_index=None`. Failed boundaries +remain absent without shifting the index on later successful inputs. Storage level +and element width are read from the allocation type. A refused descendant makes +its owning scope unknown for that view. Non-affine runtime indices retain the +widest legal access approximation. Normalization clones each reached Function call site independently. Within one call site, source expressions shared by identity remain one shared expression in the clone; sharing never aliases the independently cloned body of another diff --git a/docs/spec/cli.md b/docs/spec/cli.md index 3cab0910..6ccc9991 100644 --- a/docs/spec/cli.md +++ b/docs/spec/cli.md @@ -388,10 +388,10 @@ explicit analysis; there is no ordinary `--target` option. - `--topology LEVEL` MUST be optional, passed through as the public analysis operation's `topology_level`, and name the unit for per-unit figures. Its help MUST state the default and, for every family, which figure changes with the level - and when to pass it, together with the global-traffic and observed-peak - assumptions. Compute cost MUST name `flops_per_unit` and `service_per_unit` - as its projected figures while keeping `flops` and `service` explicitly - global; movement MUST name `TrafficMetadata.per_unit` against `whole`. + and when to pass it, together with the logical/total traffic distinction and + the observed-peak assumption. Compute cost MUST distinguish `logical`, + `total`, and each topology level's per-unit share; movement MUST name the + `MemoryMetadata.traffic.storage` per-unit account against `total`. With no analysis flag it MUST be accepted and inert. - `--dim NAME=EXTENT` MUST bind one dimension the selection leaves open, and MUST be repeatable to bind several. One dimension MUST receive one extent; diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 5844ef96..e1444dca 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -37,16 +37,16 @@ truth for the directory's structure and invariants. | `analysis/errors.py` | [analysis](./analysis.md) | `AnalysisError`, the one diagnostic the whole analysis layer raises, so catching an analysis failure catches every analysis failure rather than the subset the caller happened to import. | | `analysis/visitor.py` | [analysis](./analysis.md) | The per-call `AnalyzeContext`, carrying the shared root/current `IterationScope` while a family traverses its work. | | `analysis/iteration_scope.py` | [analysis](./analysis.md) | The shared `IterationScope` tree built once from normalized HIR; families query it instead of constructing parallel structure. | -| `analysis/access.py` | [analysis](./analysis.md) | Access relations resolved against the authored iteration scopes. | +| `analysis/access.py` | [analysis](./analysis.md) | Access relations resolved against authored iteration scopes, retaining each input's original boundary index. | | `analysis/loop_domain.py` | [analysis](./analysis.md) | isl iteration domains built from authored loop bounds. | | `analysis/loop_terms.py` | [analysis](./analysis.md) | IR-only resolution of constants, authored loop axes, strides, and bounded invariant offsets into `LoopTerm`; it has no isl dependency. | | `analysis/footprint.py` | [analysis](./analysis.md) | Reserved for future target-independent authored-loop footprint policy; access-relation construction stays in `analysis/access.py`. | -| `analysis/report.py` | [analysis](./analysis.md) | Structured analysis report data, including record-family registration, field serialization, and target-aware report-only projections. It depends only on analysis/core modules; inspection consumes it to produce text and source annotations. | +| `analysis/report.py` | [analysis](./analysis.md) | Structured analysis report data, including record-family registration and field serialization. It depends only on analysis/core modules; inspection consumes it to produce text and source annotations. | | `analysis/check.py` | [analysis](./analysis.md) | The shared authored-program gate for analysis: authored-type re-derivation, authored validation, call-context validation, and checker-specific input checks. Established once per public call rather than per family. | | `analysis/facts.py` | [analysis](./analysis.md) | The narrow Facts aggregates the analysis families declare — the memory hierarchy graph, the throughput rates, and the parallel capacity. It is the record of how much hardware each measurement rests on, and names no backend; a Fact shared across consumer families belongs in `target/facts.py`. | | `analysis/metadata.py` | [analysis](./analysis.md) | The typed records the families leave on the IR, split by what each number depends on rather than by convenience. | | `analysis/compute_cost.py` | [analysis](./analysis.md) | The `compute-cost` family: logical flops per DType and bytes per storage level, from the authored program alone. | -| `analysis/memory.py` | [analysis](./analysis.md) | The `memory` family: value lifetimes, per-level peaks, and the capacity comparisons against a target's hierarchy — failing on an over-full addressable level and advising on an over-full cache. | +| `analysis/memory.py` | [analysis](./analysis.md) | The `memory` family: per-Call movement plus Function-wide movement, value lifetimes, per-level peaks, and non-fatal capacity errors for over-full addressable levels. | | `analysis/roofline.py` | [analysis](./analysis.md) | The `roofline` family: the recorded work divided by the target's published rates, per Call and aggregated per Function. Adds no count of its own. | | `analysis/performance.py` | [analysis](./analysis.md) | The `performance` family: occurrences projected from the shared `IterationScope` tree into flat timeline records and one function envelope, scaled by parallel capacity. It introduces no second scope tree. | | `visitor_registry/` | [visitor-registry](./visitor-registry.md) | Shared registry instances and derived visitors: access-relation construction, contexts, ISL helpers, relation building, shard propagation, type inference, verification, code generation, and cost evaluation. | diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index 3cf613fb..1b7b53e9 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -551,5 +551,6 @@ class DumpScope: A non-master worker appends `__` to the test-name leaf rather than adding a worker directory. `pytest.mark.no_dump` disables it. - A completed `analyze(...)` call MUST dump its complete report data as - `analysis.json` when `DumpFlags.ANALYSIS` is enabled. This includes - target-aware report-only projections such as authored-loop cache pressure. + `analysis.json` when `DumpFlags.ANALYSIS` is enabled. Every conclusion in + that dump MUST come from the selected typed records; the report layer MUST + NOT reconstruct target-aware findings that no record carries. diff --git a/docs/tutorial/optimize.ipynb b/docs/tutorial/optimize.ipynb index 67cb3a31..7bc9aa31 100644 --- a/docs/tutorial/optimize.ipynb +++ b/docs/tutorial/optimize.ipynb @@ -54,11 +54,11 @@ { "name": "stdout", "output_type": "stream", - "text": "# traffic traffic=gmem:r71680/w43456@total,r71680/w43456@cta,r336/w224@thread;rmem:r603496/w488344@total,r603496/w488344@cta,r2604/w2084@thread\n# roofline ideal-ns=24 bound-by=memory\n" + "text": "# memory traffic=gmem:r71680/w43456@logical,r71680/w43456@total,r71680/w43456@cta,r336/w224@thread;rmem:r603496/w488344@logical,r603496/w488344@total,r603496/w488344@cta,r2604/w2084@thread peak=gmem:71680;rmem:57344 persistent=43008\n# roofline ideal-ns=24 bound-by=memory\n" } ], "source": [ - "%%bash\nset -euo pipefail\ntilefoundry analyze rms_norm_quant.py:Naive Naive.txt --compute-cost --memory --roofline\ngrep -E '^# (traffic|roofline) ' Naive.txt\n" + "%%bash\nset -euo pipefail\ntilefoundry analyze rms_norm_quant.py:Naive Naive.txt --compute-cost --memory --roofline\ngrep -E '^# (memory|roofline) ' Naive.txt\n" ] }, { @@ -94,11 +94,11 @@ { "name": "stdout", "output_type": "stream", - "text": "# traffic traffic=gmem:r43008/w14784@total,r43008/w14784@cta,r224/w112@thread;rmem:r574824/w459672@total,r574824/w459672@cta,r2492/w1972@thread\n# roofline ideal-ns=13 bound-by=memory\n" + "text": "# memory traffic=gmem:r43008/w14784@logical,r43008/w14784@total,r43008/w14784@cta,r224/w112@thread;rmem:r574824/w459672@logical,r574824/w459672@total,r574824/w459672@cta,r2492/w1972@thread peak=gmem:57792;rmem:57344 persistent=43008\n# roofline ideal-ns=13 bound-by=memory\n" } ], "source": [ - "%%bash\nset -euo pipefail\ntilefoundry analyze rms_norm_quant.py:Fused Fused.txt --compute-cost --memory --roofline\ngrep -E '^# (traffic|roofline) ' Fused.txt\n" + "%%bash\nset -euo pipefail\ntilefoundry analyze rms_norm_quant.py:Fused Fused.txt --compute-cost --memory --roofline\ngrep -E '^# (memory|roofline) ' Fused.txt\n" ] }, { diff --git a/docs/tutorial/optimize.md b/docs/tutorial/optimize.md index a9aa7790..438dfc57 100644 --- a/docs/tutorial/optimize.md +++ b/docs/tutorial/optimize.md @@ -84,11 +84,11 @@ class Naive: ```bash set -euo pipefail tilefoundry analyze rms_norm_quant.py:Naive Naive.txt --compute-cost --memory --roofline -grep -E '^# (traffic|roofline) ' Naive.txt +grep -E '^# (memory|roofline) ' Naive.txt ``` ```text -# traffic traffic=gmem:r71680/w43456@total,r71680/w43456@cta,r336/w224@thread;rmem:r603496/w488344@total,r603496/w488344@cta,r2604/w2084@thread +# memory traffic=gmem:r71680/w43456@logical,r71680/w43456@total,r71680/w43456@cta,r336/w224@thread;rmem:r603496/w488344@logical,r603496/w488344@total,r603496/w488344@cta,r2604/w2084@thread peak=gmem:71680;rmem:57344 persistent=43008 # roofline ideal-ns=24 bound-by=memory ``` @@ -128,11 +128,11 @@ class Fused: ```bash set -euo pipefail tilefoundry analyze rms_norm_quant.py:Fused Fused.txt --compute-cost --memory --roofline -grep -E '^# (traffic|roofline) ' Fused.txt +grep -E '^# (memory|roofline) ' Fused.txt ``` ```text -# traffic traffic=gmem:r43008/w14784@total,r43008/w14784@cta,r224/w112@thread;rmem:r574824/w459672@total,r574824/w459672@cta,r2492/w1972@thread +# memory traffic=gmem:r43008/w14784@logical,r43008/w14784@total,r43008/w14784@cta,r224/w112@thread;rmem:r574824/w459672@logical,r574824/w459672@total,r574824/w459672@cta,r2492/w1972@thread peak=gmem:57792;rmem:57344 persistent=43008 # roofline ideal-ns=13 bound-by=memory ``` diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index 0647d479..a01411f3 100644 --- a/docs/tutorial/showcase.ipynb +++ b/docs/tutorial/showcase.ipynb @@ -90,7 +90,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage0_Naive function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta\n# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta\n# peak-footprint=gmem:1690380\n# roofline ideal-ns=632 bound-by=memory\n\n v0 = matmul(hidden, w_q, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; traffic traffic=gmem:r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory\n v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n v34 = matmul(v33, w_o, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; traffic traffic=gmem:r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage0_Naive function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta\n# memory traffic=gmem:r2225620/w806592@logical,r2225620/w806592@total,r2225620/w806592@cta peak=gmem:1690380 persistent=1409548\n# roofline ideal-ns=632 bound-by=memory\n\n v0 = matmul(hidden, w_q, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; memory traffic=gmem:r131584/w512@logical,r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory\n v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; memory traffic=gmem:r136/w128@logical,r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n v34 = matmul(v33, w_o, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; memory traffic=gmem:r131584/w512@logical,r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage0-128.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nprint()\nfor needle in (\"matmul(hidden, w_q\", \"cache_update(k_cache\", \"matmul(v33, w_o\"):\n line = next(line for line in annotated.splitlines() if needle in line)\n print(line.rstrip())\n" @@ -99,10 +99,10 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The number before `@` is global work or traffic. After it comes one share per\n", - "topology level the program names, each under that level's own name: what one CTA\n", - "of them does, what one thread does. With no authored split a level's share is the\n", - "whole. The annotated lines printed by the previous cell come from the same report\n", + "`@logical` is the authored request before loop replication; `@total` is the\n", + "whole execution. After them comes one share per topology level the program names,\n", + "each under that level's own name: what one CTA does, what one thread does. With no\n", + "authored split a level's share is the total. The annotated lines come from the same report\n", "file, so the call-level operands and roofline numbers stay tied to the command\n", "that produced them.\n" ] @@ -137,10 +137,10 @@ { "name": "stdout", "output_type": "stream", - "text": "| `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound |\n|---:|---:|---|---:|---:|---|\n| 128 | `200448@logical,200448@total,200448@cta` | `gmem:r2225620/w806592@total,r2225620/w806592@cta` | 1690380 | 632 | memory |\n| 512 | `799488@logical,799488@total,799488@cta` | `gmem:r4744660/w3202752@total,r4744660/w3202752@cta` | 2624268 | 1656 | memory |\n| 1024 | `1598208@logical,1598208@total,1598208@cta` | `gmem:r8103380/w6397632@total,r8103380/w6397632@cta` | 3869452 | 3022 | memory |\n| 2048 | `3195648@logical,3195648@total,3195648@cta` | `gmem:r14820820/w12787392@total,r14820820/w12787392@cta` | 6359820 | 5752 | memory |\n| 4096 | `6390528@logical,6390528@total,6390528@cta` | `gmem:r28255700/w25566912@total,r28255700/w25566912@cta` | 11340556 | 11214 | memory |\n| 8192 | `12780288@logical,12780288@total,12780288@cta` | `gmem:r55125460/w51125952@total,r55125460/w51125952@cta` | 21302028 | 22136 | memory |\n" + "text": "| `ctx_len` | f32 flops `logical@total@CTA` | traffic `logical@total@CTA` | peak gmem bytes | ideal ns | bound |\n|---:|---:|---|---:|---:|---|\n| 128 | `200448@logical,200448@total,200448@cta` | `gmem:r2225620/w806592@logical,r2225620/w806592@total,r2225620/w806592@cta` | 1690380 | 632 | memory |\n| 512 | `799488@logical,799488@total,799488@cta` | `gmem:r4744660/w3202752@logical,r4744660/w3202752@total,r4744660/w3202752@cta` | 2624268 | 1656 | memory |\n| 1024 | `1598208@logical,1598208@total,1598208@cta` | `gmem:r8103380/w6397632@logical,r8103380/w6397632@total,r8103380/w6397632@cta` | 3869452 | 3022 | memory |\n| 2048 | `3195648@logical,3195648@total,3195648@cta` | `gmem:r14820820/w12787392@logical,r14820820/w12787392@total,r14820820/w12787392@cta` | 6359820 | 5752 | memory |\n| 4096 | `6390528@logical,6390528@total,6390528@cta` | `gmem:r28255700/w25566912@logical,r28255700/w25566912@total,r28255700/w25566912@cta` | 11340556 | 11214 | memory |\n| 8192 | `12780288@logical,12780288@total,12780288@cta` | `gmem:r55125460/w51125952@logical,r55125460/w51125952@total,r55125460/w51125952@cta` | 21302028 | 22136 | memory |\n" } ], - "source": "import re\nfrom pathlib import Path\n\n\ndef metrics(ctx_len):\n report = Path(f\"tutorial-reports/stage0-{ctx_len}.txt\").read_text(encoding=\"utf-8\")\n lines = report.splitlines()\n compute = next(line for line in lines if line.startswith(\"# compute-cost \"))\n traffic = next(line for line in lines if line.startswith(\"# traffic \"))\n peak = next(line for line in lines if line.startswith(\"# peak-footprint=\"))\n roofline = next(line for line in lines if line.startswith(\"# roofline \"))\n f32 = re.search(r\"f32:([^ ]+)\", compute).group(1)\n traffic_value = traffic.removeprefix(\"# traffic traffic=\")\n gmem_peak = re.search(r\"gmem:([^,]+)\", peak).group(1)\n ideal, bound = re.search(r\"ideal-ns=([^ ]+) bound-by=([^ ]+)\", roofline).groups()\n return f32, traffic_value, gmem_peak, ideal, bound\n\n\nprint(\"| `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound |\")\nprint(\"|---:|---:|---|---:|---:|---|\")\nfor ctx_len in (128, 512, 1024, 2048, 4096, 8192):\n f32, traffic, peak, ideal, bound = metrics(ctx_len)\n print(f\"| {ctx_len} | `{f32}` | `{traffic}` | {peak} | {ideal} | {bound} |\")\n" + "source": "import re\nfrom pathlib import Path\n\n\ndef metrics(ctx_len):\n report = Path(f\"tutorial-reports/stage0-{ctx_len}.txt\").read_text(encoding=\"utf-8\")\n lines = report.splitlines()\n compute = next(line for line in lines if line.startswith(\"# compute-cost \"))\n memory = next(line for line in lines if line.startswith(\"# memory \"))\n roofline = next(line for line in lines if line.startswith(\"# roofline \"))\n f32 = re.search(r\"f32:([^ ]+)\", compute).group(1)\n traffic_value = re.search(r\"traffic=(.*?) peak=\", memory).group(1)\n gmem_peak = re.search(r\"peak=.*?gmem:([^;, ]+)\", memory).group(1)\n ideal, bound = re.search(r\"ideal-ns=([^ ]+) bound-by=([^ ]+)\", roofline).groups()\n return f32, traffic_value, gmem_peak, ideal, bound\n\n\nprint(\"| `ctx_len` | f32 flops `logical@total@CTA` | traffic `logical@total@CTA` | peak gmem bytes | ideal ns | bound |\")\nprint(\"|---:|---:|---|---:|---:|---|\")\nfor ctx_len in (128, 512, 1024, 2048, 4096, 8192):\n f32, traffic, peak, ideal, bound = metrics(ctx_len)\n print(f\"| {ctx_len} | `{f32}` | `{traffic}` | {peak} | {ideal} | {bound} |\")\n" }, { "cell_type": "markdown", @@ -188,7 +188,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:2833728@logical,2833728@total,354216@cta other-ops=special:14528@logical,14528@total,1816@cta\n# traffic traffic=gmem:r5563796/w3721856@total,r3936212/w3721408@cta;smem:r9597248/w9480000@total,r1199656/w1185000@cta\n# peak-footprint=gmem:3933836;smem:581312\n# error=\"smem placement peak 581312 B exceeds capacity 232448 B\"\n# roofline ideal-ns=1935 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:2833728@logical,2833728@total,354216@cta other-ops=special:14528@logical,14528@total,1816@cta\n# memory traffic=gmem:r5563796/w3721856@logical,r5563796/w3721856@total,r3936212/w3721408@cta;smem:r9597248/w9480000@logical,r9597248/w9480000@total,r1199656/w1185000@cta peak=gmem:3933836;smem:581312 persistent=1841676 errors=1\n# error=\"smem placement peak 581312 B exceeds capacity 232448 B\"\n# roofline ideal-ns=1935 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage2-1816.txt\").read_text(encoding=\"utf-8\")\nprint(report.partition(\"\\n\\n\")[0].rstrip())\n" @@ -223,7 +223,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:2839968@logical,2839968@total,354996@cta other-ops=special:14560@logical,14560@total,1820@cta\n# traffic traffic=gmem:r5573012/w3730048@total,r3941844/w3729600@cta;smem:r9618368/w9500864@total,r1202296/w1187608@cta\n# peak-footprint=gmem:3939468;smem:582592\n# error=\"smem placement peak 582592 B exceeds capacity 232448 B\"\n# roofline ideal-ns=1939 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:2839968@logical,2839968@total,354996@cta other-ops=special:14560@logical,14560@total,1820@cta\n# memory traffic=gmem:r5573012/w3730048@logical,r5573012/w3730048@total,r3941844/w3729600@cta;smem:r9618368/w9500864@logical,r9618368/w9500864@total,r1202296/w1187608@cta peak=gmem:3939468;smem:582592 persistent=1842700 errors=1\n# error=\"smem placement peak 582592 B exceeds capacity 232448 B\"\n# roofline ideal-ns=1939 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage2-1820.txt\").read_text(encoding=\"utf-8\")\nprint(report.partition(\"\\n\\n\")[0].rstrip())\n" @@ -274,10 +274,10 @@ { "name": "stdout", "output_type": "stream", - "text": "# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta\n# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta\n# peak-footprint=gmem:1690380\n# roofline ideal-ns=632 bound-by=memory\n" + "text": "# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta\n# memory traffic=gmem:r2225620/w806592@logical,r2225620/w806592@total,r2225620/w806592@cta peak=gmem:1690380 persistent=1409548\n# roofline ideal-ns=632 bound-by=memory\n" } ], - "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage0-128-summary.txt\").read_text(encoding=\"utf-8\")\nfor line in report.splitlines():\n if line.startswith((\"# compute-cost \", \"# traffic \", \"# peak-footprint=\", \"# roofline \")):\n print(line)\n" + "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage0-128-summary.txt\").read_text(encoding=\"utf-8\")\nfor line in report.splitlines():\n if line.startswith((\"# compute-cost \", \"# memory \", \"# roofline \")):\n print(line)\n" }, { "cell_type": "markdown", @@ -309,10 +309,10 @@ { "name": "stdout", "output_type": "stream", - "text": "# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:200448@logical,200448@total,25056@cta other-ops=special:1024@logical,1024@total,128@cta\n# traffic traffic=gmem:r1674644/w264832@total,r1559508/w264384@cta;smem:r684608/w675392@total,r85576/w84424@cta\n# peak-footprint=gmem:1557132;smem:41152\n# roofline ideal-ns=405 bound-by=memory\n" + "text": "# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:200448@logical,200448@total,25056@cta other-ops=special:1024@logical,1024@total,128@cta\n# memory traffic=gmem:r1674644/w264832@logical,r1674644/w264832@total,r1559508/w264384@cta;smem:r684608/w675392@logical,r684608/w675392@total,r85576/w84424@cta peak=gmem:1557132;smem:41152 persistent=1409548\n# roofline ideal-ns=405 bound-by=memory\n" } ], - "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage2-128-summary.txt\").read_text(encoding=\"utf-8\")\nfor line in report.splitlines():\n if line.startswith((\"# compute-cost \", \"# traffic \", \"# peak-footprint=\", \"# roofline \")):\n print(line)\n" + "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage2-128-summary.txt\").read_text(encoding=\"utf-8\")\nfor line in report.splitlines():\n if line.startswith((\"# compute-cost \", \"# memory \", \"# roofline \")):\n print(line)\n" }, { "cell_type": "markdown", @@ -360,7 +360,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage3_Fused function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,4392896@total,328672@cta;f32:3239808@logical,6418944@total,200592@cta other-ops=integer:9@logical,288@total,9@cta;special:33056@logical,33152@total,1036@cta\n# traffic traffic=gmem:r3476884/w4196992@total,r2558932/w4196544@cta;rmem:r656/w72@total,r656/w72@cta;smem:r5839296/w5662784@total,r682464/w672676@cta\n# peak-footprint=gmem:7145228;rmem:8;smem:42312\n# roofline ideal-ns=1599 bound-by=memory\n\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage3_Fused function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,4392896@total,328672@cta;f32:3239808@logical,6418944@total,200592@cta other-ops=integer:9@logical,288@total,9@cta;special:33056@logical,33152@total,1036@cta\n# memory traffic=gmem:r3476884/w4196992@logical,r3480468/w4196992@total,r2559380/w4196544@cta;rmem:r656/w72@logical,r768/w128@total,r768/w128@cta;smem:r5839296/w5662784@logical,r5871552/w5695040@total,r686496/w676708@cta peak=gmem:7145228;rmem:8;smem:42312 persistent=2425356\n# roofline ideal-ns=1600 bound-by=memory\n\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; memory traffic=gmem:r136/w128@logical,r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage3-4096.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nprint()\nprint(next(line.rstrip() for line in annotated.splitlines() if \"cache_update(k_cache\" in line))\n" @@ -411,7 +411,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,337408@total,42176@cta;f32:6390528@logical,51124224@total,6390528@cta other-ops=special:32768@logical,262144@total,32768@cta\n# traffic traffic=gmem:r28254676/w25566912@total,r27967956/w25565792@cta;smem:r331008/w329984@total,r43168/w42144@cta\n# peak-footprint=gmem:11340556;smem:16960\n# roofline ideal-ns=11213 bound-by=memory\n\n v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@total,r16384/w0@cta;smem:r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; traffic traffic=smem:r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute\n\n v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@total,r16384/w0@cta;smem:r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; traffic traffic=smem:r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,337408@total,42176@cta;f32:6390528@logical,51124224@total,6390528@cta other-ops=special:32768@logical,262144@total,32768@cta\n# memory traffic=gmem:r28254676/w25566912@logical,r28254676/w25566912@total,r27967956/w25565792@cta;smem:r331008/w329984@logical,r331008/w329984@total,r43168/w42144@cta peak=gmem:11340556;smem:16960 persistent=2425356\n# roofline ideal-ns=11213 bound-by=memory\n\n v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; memory traffic=gmem:r131072/w0@logical,r131072/w0@total,r16384/w0@cta;smem:r0/w131072@logical,r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; memory traffic=smem:r131584/w512@logical,r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute\n\n v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; memory traffic=gmem:r131072/w0@logical,r131072/w0@total,r16384/w0@cta;smem:r0/w131072@logical,r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; memory traffic=smem:r131584/w512@logical,r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage4-4096.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nlines = annotated.splitlines()\nfor needle in (\"reshard(w_q\", \"reshard(w_o\"):\n start = next(index for index, line in enumerate(lines) if needle in line)\n end = start\n while end + 1 < len(lines):\n end += 1\n if end > start and \" # \" in lines[end]:\n break\n print()\n print(\"\\n\".join(line.rstrip() for line in lines[start : end + 1]))\n" @@ -462,7 +462,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,1246400@total,328672@cta;f32:6410280@logical,6410280@total,801285@cta other-ops=integer:32@logical,256@total,32@cta;special:33040@logical,33040@total,4130@cta\n# traffic traffic=gmem:r7672476/w4198272@total,r4001116/w4197824@cta;rmem:r2560/w0@total,r2560/w0@cta;smem:r21788704/w21486432@total,r2723588/w2685804@cta\n# peak-footprint=gmem:3475212;rmem:0;smem:41920\n# roofline ideal-ns=2474 bound-by=memory\n\n v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; traffic traffic=rmem:r32/w0@total,r32/w0@cta operands=0:r0/w0;1:r32/w0;result:r0/w0; roofline\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,1246400@total,328672@cta;f32:6410280@logical,6410280@total,801285@cta other-ops=integer:32@logical,256@total,32@cta;special:33040@logical,33040@total,4130@cta\n# memory traffic=gmem:r7672476/w4198272@logical,r7672476/w4198272@total,r4001116/w4197824@cta;rmem:r2560/w0@logical,r2560/w0@total,r2560/w0@cta;smem:r21788704/w21486432@logical,r21883936/w21549920@total,r2735492/w2693740@cta peak=gmem:3475212;rmem:0;smem:41920 persistent=2425356\n# roofline ideal-ns=2474 bound-by=memory\n\n v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; memory traffic=rmem:r32/w0@logical,r32/w0@total,r32/w0@cta operands=0:r0/w0;1:r32/w0;result:r0/w0; roofline\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; memory traffic=gmem:r136/w128@logical,r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage5-4096.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nprint()\nfor needle in (\"slice(k_cache\", \"cache_update(k_cache\"):\n print(next(line.rstrip() for line in annotated.splitlines() if needle in line))\n" @@ -470,7 +470,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "```text\nweight staging: static tensor -> one output slice -> reusable for the step\ncache staging: growing context -> one block -> compute -> next block\n```\n\n## Feature ledger\n\nThe page uses this embedded ladder for features orthogonal to GQA:\n\n| feature | live program |\n|---|---|\n| `@module(entry/target/topologies)` | `Stage0_Naive` |\n| `@func`, `Tensor`, `ConstTensor`, `DimVar` | `Stage0_Naive` |\n| `pass` prototype and `@f.specialize(DimVarRangePat)` | `Stage1_Specialized` |\n| single `Mesh`, shard sugar `X @ m.axis`, `reshard` to smem/gmem | `Stage2_Sharded` |\n| split-K worker mesh and online softmax state | `Stage3_Fused` |\n| weight staging and output gather | `Stage4_WeightPrepared` |\n| cache update, block scan, `matmul`, `rope`, `reduce`, `cast` | `Stage5_CachePrepared` |\n| nested Mesh, `rmem`, rank-changing `reshard`, multi-level `Topology` | not in this embedded ladder |\n| runtime weight converter | [migrate](migrate.md) |\n\nThe command surface used by this page is:\n\n```text\n--compute-cost logical work and traffic\n--memory residency and peak footprint\n--roofline ideal bound and limiting resource\n--performance per-level execution projection\n--operands operand split in annotated call lines\n--dim bind ctx_len for one static analysis run\n--json write the same report data as JSON\n```\n\nFor example, the JSON form writes to a path just like the text form. The Bash cell runs the\nCLI with `--json`; the following Python cell loads the JSON report and prints a stable summary.\n" + "source": "```text\nweight staging: static tensor -> one output slice -> reusable for the step\ncache staging: growing context -> one block -> compute -> next block\n```\n\n## Feature ledger\n\nThe page uses this embedded ladder for features orthogonal to GQA:\n\n| feature | live program |\n|---|---|\n| `@module(entry/target/topologies)` | `Stage0_Naive` |\n| `@func`, `Tensor`, `ConstTensor`, `DimVar` | `Stage0_Naive` |\n| `pass` prototype and `@f.specialize(DimVarRangePat)` | `Stage1_Specialized` |\n| single `Mesh`, shard sugar `X @ m.axis`, `reshard` to smem/gmem | `Stage2_Sharded` |\n| split-K worker mesh and online softmax state | `Stage3_Fused` |\n| weight staging and output gather | `Stage4_WeightPrepared` |\n| cache update, block scan, `matmul`, `rope`, `reduce`, `cast` | `Stage5_CachePrepared` |\n| nested Mesh, `rmem`, rank-changing `reshard`, multi-level `Topology` | not in this embedded ladder |\n| runtime weight converter | [migrate](migrate.md) |\n\nThe command surface used by this page is:\n\n```text\n--compute-cost logical work and executed totals\n--memory movement, residency, and placement peaks\n--roofline ideal bound and limiting resource\n--performance per-level execution projection\n--operands operand split in annotated call lines\n--dim bind ctx_len for one static analysis run\n--json write the same report data as JSON\n```\n\nFor example, the JSON form writes to a path just like the text form. The Bash cell runs the\nCLI with `--json`; the following Python cell loads the JSON report and prints a stable summary.\n" }, { "cell_type": "code", diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index ab318d21..69f98b63 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -198,19 +198,18 @@ for needle in ("matmul(hidden, w_q", "cache_update(k_cache", "matmul(v33, w_o"): # analysis target=nvidia.h200_sxm module=Stage0_Naive function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline # compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta -# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta -# peak-footprint=gmem:1690380 +# memory traffic=gmem:r2225620/w806592@logical,r2225620/w806592@total,r2225620/w806592@cta peak=gmem:1690380 persistent=1409548 # roofline ideal-ns=632 bound-by=memory - v0 = matmul(hidden, w_q, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; traffic traffic=gmem:r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory - v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory - v34 = matmul(v33, w_o, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; traffic traffic=gmem:r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory + v0 = matmul(hidden, w_q, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; memory traffic=gmem:r131584/w512@logical,r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory + v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; memory traffic=gmem:r136/w128@logical,r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory + v34 = matmul(v33, w_o, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; memory traffic=gmem:r131584/w512@logical,r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory ``` -The number before `@` is global work or traffic. After it comes one share per -topology level the program names, each under that level's own name: what one CTA -of them does, what one thread does. With no authored split a level's share is the -whole. The annotated lines printed by the previous cell come from the same report +`@logical` is the authored request before loop replication; `@total` is the +whole execution. After them comes one share per topology level the program names, +each under that level's own name: what one CTA does, what one thread does. With no +authored split a level's share is the total. The annotated lines come from the same report file, so the call-level operands and roofline numbers stay tied to the command that produced them. @@ -240,31 +239,30 @@ def metrics(ctx_len): report = Path(f"tutorial-reports/stage0-{ctx_len}.txt").read_text(encoding="utf-8") lines = report.splitlines() compute = next(line for line in lines if line.startswith("# compute-cost ")) - traffic = next(line for line in lines if line.startswith("# traffic ")) - peak = next(line for line in lines if line.startswith("# peak-footprint=")) + memory = next(line for line in lines if line.startswith("# memory ")) roofline = next(line for line in lines if line.startswith("# roofline ")) f32 = re.search(r"f32:([^ ]+)", compute).group(1) - traffic_value = traffic.removeprefix("# traffic traffic=") - gmem_peak = re.search(r"gmem:([^,]+)", peak).group(1) + traffic_value = re.search(r"traffic=(.*?) peak=", memory).group(1) + gmem_peak = re.search(r"peak=.*?gmem:([^;, ]+)", memory).group(1) ideal, bound = re.search(r"ideal-ns=([^ ]+) bound-by=([^ ]+)", roofline).groups() return f32, traffic_value, gmem_peak, ideal, bound -print("| `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound |") +print("| `ctx_len` | f32 flops `logical@total@CTA` | traffic `logical@total@CTA` | peak gmem bytes | ideal ns | bound |") print("|---:|---:|---|---:|---:|---|") for ctx_len in (128, 512, 1024, 2048, 4096, 8192): f32, traffic, peak, ideal, bound = metrics(ctx_len) print(f"| {ctx_len} | `{f32}` | `{traffic}` | {peak} | {ideal} | {bound} |") ``` -| `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound | +| `ctx_len` | f32 flops `logical@total@CTA` | traffic `logical@total@CTA` | peak gmem bytes | ideal ns | bound | |---:|---:|---|---:|---:|---| -| 128 | `200448@logical,200448@total,200448@cta` | `gmem:r2225620/w806592@total,r2225620/w806592@cta` | 1690380 | 632 | memory | -| 512 | `799488@logical,799488@total,799488@cta` | `gmem:r4744660/w3202752@total,r4744660/w3202752@cta` | 2624268 | 1656 | memory | -| 1024 | `1598208@logical,1598208@total,1598208@cta` | `gmem:r8103380/w6397632@total,r8103380/w6397632@cta` | 3869452 | 3022 | memory | -| 2048 | `3195648@logical,3195648@total,3195648@cta` | `gmem:r14820820/w12787392@total,r14820820/w12787392@cta` | 6359820 | 5752 | memory | -| 4096 | `6390528@logical,6390528@total,6390528@cta` | `gmem:r28255700/w25566912@total,r28255700/w25566912@cta` | 11340556 | 11214 | memory | -| 8192 | `12780288@logical,12780288@total,12780288@cta` | `gmem:r55125460/w51125952@total,r55125460/w51125952@cta` | 21302028 | 22136 | memory | +| 128 | `200448@logical,200448@total,200448@cta` | `gmem:r2225620/w806592@logical,r2225620/w806592@total,r2225620/w806592@cta` | 1690380 | 632 | memory | +| 512 | `799488@logical,799488@total,799488@cta` | `gmem:r4744660/w3202752@logical,r4744660/w3202752@total,r4744660/w3202752@cta` | 2624268 | 1656 | memory | +| 1024 | `1598208@logical,1598208@total,1598208@cta` | `gmem:r8103380/w6397632@logical,r8103380/w6397632@total,r8103380/w6397632@cta` | 3869452 | 3022 | memory | +| 2048 | `3195648@logical,3195648@total,3195648@cta` | `gmem:r14820820/w12787392@logical,r14820820/w12787392@total,r14820820/w12787392@cta` | 6359820 | 5752 | memory | +| 4096 | `6390528@logical,6390528@total,6390528@cta` | `gmem:r28255700/w25566912@logical,r28255700/w25566912@total,r28255700/w25566912@cta` | 11340556 | 11214 | memory | +| 8192 | `12780288@logical,12780288@total,12780288@cta` | `gmem:r55125460/w51125952@logical,r55125460/w51125952@total,r55125460/w51125952@cta` | 21302028 | 22136 | memory | The table says: @@ -426,8 +424,7 @@ print(report.partition("\n\n")[0].rstrip()) # analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline # compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:2833728@logical,2833728@total,354216@cta other-ops=special:14528@logical,14528@total,1816@cta -# traffic traffic=gmem:r5563796/w3721856@total,r3936212/w3721408@cta;smem:r9597248/w9480000@total,r1199656/w1185000@cta -# peak-footprint=gmem:3933836;smem:581312 +# memory traffic=gmem:r5563796/w3721856@logical,r5563796/w3721856@total,r3936212/w3721408@cta;smem:r9597248/w9480000@logical,r9597248/w9480000@total,r1199656/w1185000@cta peak=gmem:3933836;smem:581312 persistent=1841676 errors=1 # error="smem placement peak 581312 B exceeds capacity 232448 B" # roofline ideal-ns=1935 bound-by=memory ``` @@ -453,8 +450,7 @@ print(report.partition("\n\n")[0].rstrip()) # analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline # compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:2839968@logical,2839968@total,354996@cta other-ops=special:14560@logical,14560@total,1820@cta -# traffic traffic=gmem:r5573012/w3730048@total,r3941844/w3729600@cta;smem:r9618368/w9500864@total,r1202296/w1187608@cta -# peak-footprint=gmem:3939468;smem:582592 +# memory traffic=gmem:r5573012/w3730048@logical,r5573012/w3730048@total,r3941844/w3729600@cta;smem:r9618368/w9500864@logical,r9618368/w9500864@total,r1202296/w1187608@cta peak=gmem:3939468;smem:582592 persistent=1842700 errors=1 # error="smem placement peak 582592 B exceeds capacity 232448 B" # roofline ideal-ns=1939 bound-by=memory ``` @@ -546,14 +542,13 @@ from pathlib import Path report = Path("tutorial-reports/stage0-128-summary.txt").read_text(encoding="utf-8") for line in report.splitlines(): - if line.startswith(("# compute-cost ", "# traffic ", "# peak-footprint=", "# roofline ")): + if line.startswith(("# compute-cost ", "# memory ", "# roofline ")): print(line) ``` ```text # compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta -# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta -# peak-footprint=gmem:1690380 +# memory traffic=gmem:r2225620/w806592@logical,r2225620/w806592@total,r2225620/w806592@cta peak=gmem:1690380 persistent=1409548 # roofline ideal-ns=632 bound-by=memory ``` @@ -573,14 +568,13 @@ from pathlib import Path report = Path("tutorial-reports/stage2-128-summary.txt").read_text(encoding="utf-8") for line in report.splitlines(): - if line.startswith(("# compute-cost ", "# traffic ", "# peak-footprint=", "# roofline ")): + if line.startswith(("# compute-cost ", "# memory ", "# roofline ")): print(line) ``` ```text # compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:200448@logical,200448@total,25056@cta other-ops=special:1024@logical,1024@total,128@cta -# traffic traffic=gmem:r1674644/w264832@total,r1559508/w264384@cta;smem:r684608/w675392@total,r85576/w84424@cta -# peak-footprint=gmem:1557132;smem:41152 +# memory traffic=gmem:r1674644/w264832@logical,r1674644/w264832@total,r1559508/w264384@cta;smem:r684608/w675392@logical,r684608/w675392@total,r85576/w84424@cta peak=gmem:1557132;smem:41152 persistent=1409548 # roofline ideal-ns=405 bound-by=memory ``` @@ -726,11 +720,10 @@ print(next(line.rstrip() for line in annotated.splitlines() if "cache_update(k_c # analysis target=nvidia.h200_sxm module=Stage3_Fused function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline # compute-cost flops=bf16:328896@logical,4392896@total,328672@cta;f32:3239808@logical,6418944@total,200592@cta other-ops=integer:9@logical,288@total,9@cta;special:33056@logical,33152@total,1036@cta -# traffic traffic=gmem:r3476884/w4196992@total,r2558932/w4196544@cta;rmem:r656/w72@total,r656/w72@cta;smem:r5839296/w5662784@total,r682464/w672676@cta -# peak-footprint=gmem:7145228;rmem:8;smem:42312 -# roofline ideal-ns=1599 bound-by=memory +# memory traffic=gmem:r3476884/w4196992@logical,r3480468/w4196992@total,r2559380/w4196544@cta;rmem:r656/w72@logical,r768/w128@total,r768/w128@cta;smem:r5839296/w5662784@logical,r5871552/w5695040@total,r686496/w676708@cta peak=gmem:7145228;rmem:8;smem:42312 persistent=2425356 +# roofline ideal-ns=1600 bound-by=memory - v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory + v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; memory traffic=gmem:r136/w128@logical,r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory ``` The embedded `Stage3_Fused` program is the split-K example for this page. @@ -863,15 +856,14 @@ for needle in ("reshard(w_q", "reshard(w_o"): # analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline # compute-cost flops=bf16:328896@logical,337408@total,42176@cta;f32:6390528@logical,51124224@total,6390528@cta other-ops=special:32768@logical,262144@total,32768@cta -# traffic traffic=gmem:r28254676/w25566912@total,r27967956/w25565792@cta;smem:r331008/w329984@total,r43168/w42144@cta -# peak-footprint=gmem:11340556;smem:16960 +# memory traffic=gmem:r28254676/w25566912@logical,r28254676/w25566912@total,r27967956/w25565792@cta;smem:r331008/w329984@logical,r331008/w329984@total,r43168/w42144@cta peak=gmem:11340556;smem:16960 persistent=2425356 # roofline ideal-ns=11213 bound-by=memory - v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@total,r16384/w0@cta;smem:r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory - v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; traffic traffic=smem:r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute + v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; memory traffic=gmem:r131072/w0@logical,r131072/w0@total,r16384/w0@cta;smem:r0/w131072@logical,r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory + v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; memory traffic=smem:r131584/w512@logical,r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute - v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@total,r16384/w0@cta;smem:r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory - v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; traffic traffic=smem:r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute + v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; memory traffic=gmem:r131072/w0@logical,r131072/w0@total,r16384/w0@cta;smem:r0/w131072@logical,r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory + v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; memory traffic=smem:r131584/w512@logical,r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute ``` ## 6. Stream the KV cache @@ -1015,12 +1007,11 @@ for needle in ("slice(k_cache", "cache_update(k_cache"): # analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline # compute-cost flops=bf16:328896@logical,1246400@total,328672@cta;f32:6410280@logical,6410280@total,801285@cta other-ops=integer:32@logical,256@total,32@cta;special:33040@logical,33040@total,4130@cta -# traffic traffic=gmem:r7672476/w4198272@total,r4001116/w4197824@cta;rmem:r2560/w0@total,r2560/w0@cta;smem:r21788704/w21486432@total,r2723588/w2685804@cta -# peak-footprint=gmem:3475212;rmem:0;smem:41920 +# memory traffic=gmem:r7672476/w4198272@logical,r7672476/w4198272@total,r4001116/w4197824@cta;rmem:r2560/w0@logical,r2560/w0@total,r2560/w0@cta;smem:r21788704/w21486432@logical,r21883936/w21549920@total,r2735492/w2693740@cta peak=gmem:3475212;rmem:0;smem:41920 persistent=2425356 # roofline ideal-ns=2474 bound-by=memory - v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=rmem:r32/w0@total,r32/w0@cta operands=0:r0/w0;1:r32/w0;result:r0/w0; roofline - v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory + v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; memory traffic=rmem:r32/w0@logical,r32/w0@total,r32/w0@cta operands=0:r0/w0;1:r32/w0;result:r0/w0; roofline + v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; memory traffic=gmem:r136/w128@logical,r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory ``` ```text @@ -1047,8 +1038,8 @@ The page uses this embedded ladder for features orthogonal to GQA: The command surface used by this page is: ```text ---compute-cost logical work and traffic ---memory residency and peak footprint +--compute-cost logical work and executed totals +--memory movement, residency, and placement peaks --roofline ideal bound and limiting resource --performance per-level execution projection --operands operand split in annotated call lines diff --git a/src/tilefoundry/analysis/__init__.py b/src/tilefoundry/analysis/__init__.py index 51165145..35d1ae66 100644 --- a/src/tilefoundry/analysis/__init__.py +++ b/src/tilefoundry/analysis/__init__.py @@ -16,16 +16,15 @@ ThroughputFacts, ) from .metadata import ( - BufferFootprint, - ComputeCostMetadata, - LoopFootprintMetadata, - AllocationMetadata, - MemoryLevelFootprint, MemoryMetadata, + ComputeCostMetadata, + Footprint, + MemoryLevelPeak, PerformanceMetadata, PerformanceSummaryMetadata, + RegionMemoryMetadata, RooflineMetadata, - TrafficMetadata, + Traffic, Breakdown, Spread, TimelineMetadata, @@ -41,23 +40,22 @@ "Analyzer", "AnalysisError", "AnalysisResult", - "BufferFootprint", + "MemoryMetadata", "ComputeCostMetadata", "ExplicitMemoryLevelFacts", "ImplicitMemoryLevelFacts", - "LoopFootprintMetadata", + "Footprint", "MemoryHierarchyFacts", - "MemoryLevelFootprint", + "MemoryLevelPeak", "MemoryLevelRelation", - "AllocationMetadata", - "MemoryMetadata", + "RegionMemoryMetadata", "MemoryRelationKind", "ParallelCapacityFacts", "PerformanceServiceFacts", "PerformanceMetadata", "PerformanceSummaryMetadata", "RooflineMetadata", - "TrafficMetadata", + "Traffic", "Breakdown", "Spread", "ThroughputFacts", diff --git a/src/tilefoundry/analysis/access.py b/src/tilefoundry/analysis/access.py index 1ad05fe1..e54232fc 100644 --- a/src/tilefoundry/analysis/access.py +++ b/src/tilefoundry/analysis/access.py @@ -40,6 +40,7 @@ class AccessPrecision(Enum): class Access: """One relation from an iteration scope to the allocation it reaches.""" + input_index: int | None relation: isl.map buffer: Expr precision: AccessPrecision = AccessPrecision.EXACT @@ -155,6 +156,7 @@ def resolve_access( scope: "IterationScope", ctx: TypeInferContext, *, + input_index: int | None, narrow: bool, ) -> Access | None: """Resolve one declared boundary into an access from its iteration scope.""" @@ -194,7 +196,7 @@ def resolve_access( precision = AccessPrecision.WIDENED if precision is AccessPrecision.EXACT and has_unbounded_param(relation): precision = AccessPrecision.UNKNOWN - return Access(relation, operand, precision) + return Access(input_index, relation, operand, precision) __all__ = [ diff --git a/src/tilefoundry/analysis/check.py b/src/tilefoundry/analysis/check.py index a3874ee2..9fb167de 100644 --- a/src/tilefoundry/analysis/check.py +++ b/src/tilefoundry/analysis/check.py @@ -58,8 +58,8 @@ MemoryMetadata, PerformanceMetadata, PerformanceSummaryMetadata, + RegionMemoryMetadata, RooflineMetadata, - TrafficMetadata, ) _INLINE_NODES = 10_000 @@ -69,8 +69,8 @@ PerformanceMetadata, PerformanceSummaryMetadata, RangeMetadata, + RegionMemoryMetadata, RooflineMetadata, - TrafficMetadata, } _ResourceKey = tuple[str, str] diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index eeb32427..46799ec8 100644 --- a/src/tilefoundry/analysis/compute_cost.py +++ b/src/tilefoundry/analysis/compute_cost.py @@ -18,7 +18,7 @@ from .errors import AnalysisError from .facts import PerformanceServiceFacts, ThroughputFacts -from .metadata import Breakdown, ComputeCostMetadata, breakdown, shares +from .metadata import Breakdown, ComputeCostMetadata, MemoryMetadata, breakdown, shares from .visitor import AnalyzeContext SELECTOR = "compute-cost" @@ -34,7 +34,7 @@ def _at(held: Breakdown[int], topologies: tuple[str, ...], level: str | None): def _is_structural_occurrence( cost: ComputeCostMetadata, - moved: "TrafficMetadata | None" = None, + moved: MemoryMetadata | None = None, *, unit: str, bandwidth_level: str | None = None, @@ -44,7 +44,7 @@ def _is_structural_occurrence( all(not value for _name, value in _at(cost.flops, cost.topologies, unit)) and all(not value for _kind, value in _at(cost.other_ops, cost.topologies, unit)) and not ( - _bytes(moved.storage, moved.topologies, bandwidth_level, unit) + _bytes(moved.traffic.storage, moved.topologies, bandwidth_level, unit) if moved is not None and bandwidth_level is not None else 0 ) @@ -56,7 +56,7 @@ def local_duration_ns( facts: ThroughputFacts, services: PerformanceServiceFacts, *, - moved: "TrafficMetadata | None" = None, + moved: MemoryMetadata | None = None, topology_level: str | None = None, level: str | None = None, scale: int = 1, @@ -102,7 +102,8 @@ def local_duration_ns( compute_ns += -(-(value * scale * 1_000_000_000) // throughput) crossed = ( - _bytes(moved.storage, moved.topologies, facts.bandwidth_level, topology_level) * scale + _bytes(moved.traffic.storage, moved.topologies, facts.bandwidth_level, topology_level) + * scale if moved is not None else 0 ) @@ -117,7 +118,13 @@ def local_duration_ns( memory_ns = -(-(crossed * 1_000_000_000) // throughput) sent = ( - _bytes(moved.communication, moved.topologies, topology_level, topology_level) * scale + _bytes( + moved.traffic.communication, + moved.topologies, + topology_level, + topology_level, + ) + * scale if moved is not None else 0 ) diff --git a/src/tilefoundry/analysis/iteration_scope.py b/src/tilefoundry/analysis/iteration_scope.py index 7a8d1f9e..0b8ee8c2 100644 --- a/src/tilefoundry/analysis/iteration_scope.py +++ b/src/tilefoundry/analysis/iteration_scope.py @@ -7,7 +7,7 @@ import isl -from tilefoundry.ir.core import Call, Expr, value_label, value_labels +from tilefoundry.ir.core import Call, Expr from tilefoundry.ir.core.module import Module from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion @@ -24,14 +24,12 @@ access_relation_registry, projected, relations_of, - static_bytes, ) from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext from .access import Access, resolve_access from .errors import AnalysisError from .loop_domain import induction_name, iteration_domain -from .metadata import BufferFootprint, LoopFootprintMetadata @dataclass(eq=False) @@ -119,113 +117,6 @@ def trips(self) -> int: self._trips_cache = result return result - def elements_per_trip(self, access: Access) -> int: - """Count elements reached while this scope's loop axes are held still.""" - cache = getattr(self, "_elements_per_trip_cache", {}) - cached = cache.get(id(access)) - if cached is not None: - return cached - standing = self.domain.insert_dims( - isl.dim_type.SET, - self.depth, - access.relation.dim(isl.dim_type.IN) - self.depth, - ) - relation = access.relation.intersect_domain(standing) - try: - points = param_points(relation.params()) - except UnboundedParameterBox as error: - label = value_label(access.buffer) or type(access.buffer).__name__ - raise AnalysisError( - f"scope access to {label!r} still has unbound parameter {error.parameter!r}" - ) from error - except ParameterBoxTooLarge as error: - raise AnalysisError( - f"scope access parameter box exceeds the {PARAM_POINT_LIMIT}-point analysis limit" - ) from error - amounts = [] - for point in points: - fixed = relation.intersect_params(point) - fixed_standing = standing.intersect_params(point) - for axis in range(self.depth): - low = fixed_standing.dim_min_val(axis) - if not low.is_int(): - raise AnalysisError("scope access has no finite one-pass extent") - fixed_standing = fixed_standing.fix_si( - isl.dim_type.SET, - axis, - low.get_num_si(), - ) - amount = cardinality(fixed.intersect_domain(fixed_standing).range()) - if amount is None: - raise AnalysisError("scope access has no finite one-pass extent") - amounts.append(amount) - result = max(amounts, default=0) - cache[id(access)] = result - self._elements_per_trip_cache = cache - return result - - def accesses_in(self, view: str) -> Iterator[Access]: - """Yield accesses owned by this scope and all descendant scopes.""" - for _call, values in self.accesses.get(view, {}).values(): - yield from values - for child in self.children: - yield from child.accesses_in(view) - - def is_complete(self, view: str) -> bool: - """Whether this scope and every descendant answered every access.""" - if self.refused.get(view): - return False - return all(child.is_complete(view) for child in self.children) - - def footprint(self) -> LoopFootprintMetadata: - """Summarize device and per-unit access bytes for this scope. - - Two structurally equal buffers are distinct allocations, so identity - groups the rows. It does not order or name them: an address is whatever - the allocator handed out this run, and a report exists to be compared - against another run. - """ - rows: dict[tuple[int, str], tuple[Expr, int, int, int]] = {} - for view, scale in (("narrow", "bytes"), ("device", "device_bytes")): - for access in self.accesses_in(view): - try: - amount = self.elements_per_trip(access) - except AnalysisError: - continue - size = static_bytes(access.buffer.type) - if size is None: - continue - device_amount = amount * max(1, self.trips()) - key = (id(access.buffer), str(getattr(access.buffer.type, "storage", "unknown"))) - current = rows.get(key, (access.buffer, len(rows), 0, 0)) - rows[key] = ( - current[0], - current[1], - current[2] + (amount * size if scale == "bytes" else 0), - current[3] + (device_amount * size if scale == "device_bytes" else 0), - ) - entries = list(rows.items()) - labels = value_labels(buffer for _, (buffer, _, _, _) in entries) - ordered = sorted( - (label, memory_level, local, device) - for label, ((_, memory_level), (_, _, local, device)) in zip(labels, entries) - ) - footprints = tuple( - BufferFootprint( - buffer=label, - memory_level=memory_level, - bytes=local, - device_bytes=device, - repeated_bytes=local * self.trips(), - ) - for label, memory_level, local, device in ordered - ) - return LoopFootprintMetadata( - footprints=footprints, - known=self.is_complete("narrow") and self.is_complete("device"), - ) - - class ScopeBuilder: """Build one IterationScope tree and its access views for a Function.""" @@ -267,14 +158,26 @@ def _record_accesses(self, expr: Call, scope: IterationScope) -> None: if index >= len(expr.args): continue access = resolve_access( - expr.args[index], boundary, scope, self.type_ctx, narrow=narrow + expr.args[index], + boundary, + scope, + self.type_ctx, + input_index=index, + narrow=narrow, ) if access is not None: built.append(access) scope.accesses.setdefault(view, {})[id(expr)] = (expr, tuple(built)) written: list[Access] = [] for boundary in local_relations.outputs: - access = resolve_access(expr, boundary, scope, self.type_ctx, narrow=narrow) + access = resolve_access( + expr, + boundary, + scope, + self.type_ctx, + input_index=None, + narrow=narrow, + ) if access is not None: written.append(access) scope.outputs.setdefault(view, {})[id(expr)] = (expr, tuple(written)) diff --git a/src/tilefoundry/analysis/memory.py b/src/tilefoundry/analysis/memory.py index 4c4e2ceb..778bf332 100644 --- a/src/tilefoundry/analysis/memory.py +++ b/src/tilefoundry/analysis/memory.py @@ -37,14 +37,13 @@ from .facts import MemoryHierarchyFacts from .liveness import Liveness, analyze_liveness from .metadata import ( - AllocationMetadata, Breakdown, - LoopFootprintMetadata, - MemoryLevelFootprint, + MemoryLevelPeak, MemoryMetadata, + RegionMemoryMetadata, Spread, + Traffic, TrafficBytes, - TrafficMetadata, ValueLifetime, ) from .visitor import AnalyzeContext @@ -210,7 +209,7 @@ def call_traffic( locals_by_unit: "dict[str, CostContext]", stated_relations: AccessRelations | None = None, asked: "str | None" = None, -) -> TrafficMetadata: +) -> MemoryMetadata: """What one Call moves, whole and for one participant. The same registered evaluator the work half reads, projected onto its @@ -220,10 +219,12 @@ def call_traffic( The Type of the leaf it reached names the level those bytes are charged at, and an allocation does not correct either answer. """ - storage: dict[str, dict[str, TrafficBytes]] = {} - crossing: dict[str, dict[str, TrafficBytes]] = {} + storage_whole: dict[str, TrafficBytes] = {} + storage_per_unit: dict[str, dict[str, TrafficBytes]] = {} + crossing_whole: dict[str, TrafficBytes] = {} + crossing_per_unit: dict[str, dict[str, TrafficBytes]] = {} operands: tuple[TrafficBytes, ...] = () - for key, ctx in (("", whole), *locals_by_unit.items()): + for unit, ctx in ((None, whole), *locals_by_unit.items()): try: cost = CostEvaluator().visit(expr, ctx) except (ValueError, VerifyError) as error: @@ -234,17 +235,27 @@ def call_traffic( ) levels, positional = _movement(expr, cost, ctx, types, stated_relations) for memory_level, moved in levels: - storage.setdefault(memory_level, {})[key] = moved + if unit is None: + storage_whole[memory_level] = moved + else: + storage_per_unit.setdefault(memory_level, {})[unit] = moved for boundary, moved in cost.sent: - if key and _finer_than(key, boundary, locals_by_unit): + if unit is not None and _finer_than(unit, boundary, locals_by_unit): continue - crossing.setdefault(boundary, {})[key] = moved - if not key: + if unit is None: + crossing_whole[boundary] = moved + else: + crossing_per_unit.setdefault(boundary, {})[unit] = moved + if unit is None: operands = positional - return TrafficMetadata( + return MemoryMetadata( topologies=tuple(locals_by_unit), - storage=_shares(storage, tuple(locals_by_unit)), - communication=_shares(crossing, tuple(locals_by_unit)), + traffic=Traffic( + storage=_occurrence_shares(storage_whole, storage_per_unit, tuple(locals_by_unit)), + communication=_occurrence_shares( + crossing_whole, crossing_per_unit, tuple(locals_by_unit) + ), + ), operands=operands, ) @@ -262,11 +273,10 @@ def _finer_than(unit: str, boundary: str, ordered: "dict[str, CostContext]") -> return names.index(unit) > names.index(boundary) -_WHOLE = "" - - -def _shares( - held: "dict[str, dict[str, TrafficBytes]]", topologies: tuple[str, ...] +def _occurrence_shares( + whole: dict[str, TrafficBytes], + per_unit: dict[str, dict[str, TrafficBytes]], + topologies: tuple[str, ...], ) -> "Breakdown[TrafficBytes]": """One entry per level, each carrying the whole and every level's share. @@ -279,34 +289,75 @@ def _shares( ( name, Spread( - logical=shares.get(_WHOLE, TrafficBytes()), - total=shares.get(_WHOLE, TrafficBytes()), - per_unit=tuple(shares.get(unit, TrafficBytes()) for unit in topologies), + logical=whole.get(name, TrafficBytes()), + total=whole.get(name, TrafficBytes()), + per_unit=tuple( + per_unit.get(name, {}).get(unit, TrafficBytes()) for unit in topologies + ), ), ) - for name, shares in sorted(held.items()) + for name in sorted({*whole, *per_unit}) ) ) +@dataclass +class _TrafficAccounts: + """Function traffic kept separate in its three counting domains.""" + + logical: dict[str, TrafficBytes] = field(default_factory=dict) + total: dict[str, TrafficBytes] = field(default_factory=dict) + per_unit: dict[str, dict[str, TrafficBytes]] = field(default_factory=dict) + + +def _account_shares( + accounts: _TrafficAccounts, topologies: tuple[str, ...] +) -> Breakdown[TrafficBytes]: + """Build a breakdown without encoding an account as a topology key.""" + names = sorted({*accounts.logical, *accounts.total, *accounts.per_unit}) + return Breakdown( + tuple( + ( + name, + Spread( + logical=accounts.logical.get(name, TrafficBytes()), + total=accounts.total.get(name, TrafficBytes()), + per_unit=tuple( + accounts.per_unit.get(name, {}).get(unit, TrafficBytes()) + for unit in topologies + ), + ), + ) + for name in names + ) + ) + + +def _accumulate(into: dict[str, TrafficBytes], name: str, moved: TrafficBytes, trips: int) -> None: + running = into.get(name, TrafficBytes()) + into[name] = TrafficBytes( + running.read + moved.read * trips, + running.write + moved.write * trips, + ) + + def add_traffic( - whole: "dict[str, dict[str, TrafficBytes]]", - per_unit: "dict[str, dict[str, TrafficBytes]]", - record: TrafficMetadata, - trips: int, + storage: _TrafficAccounts, + communication: _TrafficAccounts, + record: MemoryMetadata, + logical_trips: int, + total_trips: int, ) -> None: - """Add one occurrence's bytes to a function's, as often as it happens.""" - for into, stated in ((whole, record.storage), (per_unit, record.communication)): + """Add one occurrence to each independent Function counting domain.""" + for into, stated in ( + (storage, record.traffic.storage), + (communication, record.traffic.communication), + ): for name, spread in stated.kinds: - for key, moved in ( - (_WHOLE, spread.total), - *zip(record.topologies, spread.per_unit, strict=False), - ): - running = into.setdefault(name, {}).get(key, TrafficBytes()) - into[name][key] = TrafficBytes( - running.read + moved.read * trips, - running.write + moved.write * trips, - ) + _accumulate(into.logical, name, spread.logical, logical_trips) + _accumulate(into.total, name, spread.total, total_trips) + for unit, moved in zip(record.topologies, spread.per_unit, strict=False): + _accumulate(into.per_unit.setdefault(name, {}), unit, moved, total_trips) def _resident_value_ids(function: Function, liveness: Liveness) -> frozenset[int]: @@ -388,12 +439,12 @@ class MemoryContext(AnalyzeContext): whole: CostContext | None = None locals_by_unit: dict[str, CostContext] = field(default_factory=dict) - totals: dict[str, dict[str, TrafficBytes]] = field(default_factory=dict) - shares: dict[str, dict[str, TrafficBytes]] = field(default_factory=dict) + storage: _TrafficAccounts = field(default_factory=_TrafficAccounts) + communication: _TrafficAccounts = field(default_factory=_TrafficAccounts) class MemoryVisitor(ExprVisitor[None]): - """Attach per-Call traffic and loop footprints.""" + """Attach per-Call traffic and aggregate it over each enclosing loop.""" def visit_LoopRegion(self, expr: LoopRegion, ctx: MemoryContext) -> None: child = next(item for item in ctx.current.children if item.owner is expr) @@ -403,7 +454,6 @@ def visit_LoopRegion(self, expr: LoopRegion, ctx: MemoryContext) -> None: self.visit(expr.body, inner) for operand in expr.yield_values: self.visit(operand, inner) - attach(expr, child.footprint()) def default_visit_leaf( self, expr: Expr, _operands: tuple[None, ...], ctx: MemoryContext @@ -422,22 +472,31 @@ def default_visit_leaf( ctx.topology_level, ) if recorded - else TrafficMetadata() + else MemoryMetadata() ) attach(expr, moved) if not recorded: return - repeats = 1 + logical_repeats = 1 + total_repeats = 1 cursor = ctx.current while cursor.parent is not None: + trips = max(1, cursor.trips()) + total_repeats *= trips if cursor.is_variant(expr): - repeats *= max(1, cursor.trips()) + logical_repeats *= trips cursor = cursor.parent - add_traffic(ctx.totals, ctx.shares, moved, repeats) + add_traffic( + ctx.storage, + ctx.communication, + moved, + logical_repeats, + total_repeats, + ) def analyze_memory(function: Function, context: AnalyzeContext) -> None: - """Attach traffic and per-loop footprints from the shared IterationScope tree.""" + """Attach Call movement and Function-wide movement and placement.""" module = context.module topology_level = context.topology_level facts = context.target.get_facts(MemoryHierarchyFacts) @@ -465,14 +524,6 @@ def analyze_memory(function: Function, context: AnalyzeContext) -> None: locals_by_unit=locals_by_unit, ) MemoryVisitor().visit(function.body, memory_context) - attach( - function, - TrafficMetadata( - topologies=tuple(locals_by_unit), - storage=_shares(memory_context.totals, tuple(locals_by_unit)), - communication=_shares(memory_context.shares, tuple(locals_by_unit)), - ), - ) liveness = analyze_liveness(function) placement = CostContext( scope=FunctionScope(module, function), @@ -490,8 +541,10 @@ def analyze_memory(function: Function, context: AnalyzeContext) -> None: solver_options = ( context.options if isinstance(context.options, MemoryOptions) else MemoryOptions() ) - levels_list: list[MemoryLevelFootprint] = [] - for name in sorted({item.memory_level for item in lifetimes} | set(memory_context.totals)): + levels_list: list[MemoryLevelPeak] = [] + for name in sorted( + {item.memory_level for item in lifetimes} | set(memory_context.storage.total) + ): declared = facts.explicit(name) values = tuple(item for item in allocation_values if item.lifetime.memory_level == name) rows = [item.lifetime for item in values] @@ -518,7 +571,7 @@ def analyze_memory(function: Function, context: AnalyzeContext) -> None: ), ) levels_list.append( - MemoryLevelFootprint( + MemoryLevelPeak( memory_level=name, peak_bytes=peak, persistent_bytes=sum(item.bytes for item in rows if item.persistent), @@ -532,65 +585,25 @@ def analyze_memory(function: Function, context: AnalyzeContext) -> None: for item in levels if item.exceeds_capacity ) - allocation = AllocationMetadata(solver_status="feasible") attach( function, - MemoryMetadata( - footprint=levels, + RegionMemoryMetadata( + topologies=tuple(locals_by_unit), + traffic=Traffic( + storage=_account_shares(memory_context.storage, tuple(locals_by_unit)), + communication=_account_shares(memory_context.communication, tuple(locals_by_unit)), + ), lifetimes=lifetimes, + peaks=levels, + solver_status="feasible", errors=errors, - allocation=allocation, ), ) -def cache_pressure( - record: LoopFootprintMetadata, - facts: MemoryHierarchyFacts, - peaks: dict[str, int], -) -> tuple[dict[str, object], ...]: - """Compare one scope's device footprint with same-scope implicit caches.""" - rows: list[dict[str, object]] = [] - for cache in facts.implicit_levels: - backing_name = facts.backing_level(cache.name) - backing = facts.explicit(backing_name) - if backing is None or backing.scope != cache.scope: - continue - accesses = tuple(item for item in record.footprints if item.memory_level == backing_name) - if not accesses or any(item.device_bytes < item.bytes for item in accesses): - continue - working_set = sum(item.device_bytes for item in accesses) - capacity = cache.capacity_bytes - for peer, shared_bytes in facts.capacity_sharers(cache.name): - if shared_bytes is None: - continue - remaining = shared_bytes - peaks.get(peer, 0) - capacity = remaining if capacity is None else min(capacity, remaining) - status = ( - "unknown" - if capacity is None - else "exceeds" - if working_set > capacity - else "fits" - if record.known - else "lower-bound" - ) - rows.append( - { - "cache_level": cache.name, - "backing_level": backing_name, - "device_bytes": working_set, - "capacity_bytes": capacity, - "status": status, - } - ) - return tuple(rows) - - __all__ = [ "MemoryOptions", "SELECTOR", "analyze_memory", "analyze_value_lifetimes", - "cache_pressure", ] diff --git a/src/tilefoundry/analysis/metadata.py b/src/tilefoundry/analysis/metadata.py index f63f369b..3c61c968 100644 --- a/src/tilefoundry/analysis/metadata.py +++ b/src/tilefoundry/analysis/metadata.py @@ -92,37 +92,43 @@ class ComputeCostMetadata(IRMetadata): @dataclass(frozen=True) -class TrafficMetadata(IRMetadata): - """The bytes one occurrence moves, in the two coordinates a move has. +class Traffic: + """The bytes one occurrence or region moves in its two coordinates. ``storage`` says where the bytes are; ``communication`` says whose boundary they crossed, which a storage level cannot answer: data handed from one card to another is global memory at both ends and has still gone somewhere. One movement is counted in both, because it spends both. - ``operands`` is positional against ``(*call.args, call)`` on a Call and - empty on a Function, whose totals count each occurrence as often as its - loops repeat it. """ - topologies: tuple[str, ...] = () storage: Breakdown[TrafficBytes] = Breakdown() communication: Breakdown[TrafficBytes] = Breakdown() - operands: tuple[TrafficBytes, ...] = () - @property - def whole(self) -> tuple[tuple[str, TrafficBytes], ...]: - return tuple((name, spread.total) for name, spread in self.storage.kinds) - @property - def per_unit(self) -> tuple[tuple[str, TrafficBytes], ...]: - return tuple( - (name, spread.per_unit[-1] if spread.per_unit else spread.total) - for name, spread in self.storage.kinds - ) +@dataclass(frozen=True) +class Footprint: + """Unique read bytes by source buffer and memory-level counting domain.""" + + buffers: tuple[tuple[str, Breakdown[int]], ...] = () + complete: bool = True @dataclass(frozen=True) -class MemoryLevelFootprint: +class MemoryMetadata(IRMetadata): + """One Call's memory behavior for one occurrence. + + ``operands`` is positional against ``(*call.args, call)``. ``footprint`` is + absent until a read-footprint analysis has actually produced a conclusion. + """ + + topologies: tuple[str, ...] = () + traffic: Traffic = Traffic() + operands: tuple[TrafficBytes, ...] = () + footprint: Footprint | None = None + + +@dataclass(frozen=True) +class MemoryLevelPeak: """One level's solved high-water mark or largest logical value. ``persistent_bytes`` is the part that cannot be reclaimed within the @@ -134,44 +140,12 @@ class MemoryLevelFootprint: persistent_bytes: int capacity_bytes: int | None = None - @property - def level(self) -> str: - """Compatibility spelling for consumers rendering a memory level.""" - return self.memory_level - @property def exceeds_capacity(self) -> bool: """Whether the peak does not fit the stated capacity.""" return self.capacity_bytes is not None and self.peak_bytes > self.capacity_bytes -@dataclass(frozen=True) -class BufferFootprint: - """Per-position, device-wide, and repeated bytes touched in one buffer.""" - - buffer: str - memory_level: str - bytes: int - device_bytes: int - repeated_bytes: int - - @property - def level(self) -> str: - return self.memory_level - - -@dataclass(frozen=True) -class LoopFootprintMetadata(IRMetadata): - """Buffer bytes touched by one authored loop, grouped by storage level. - - ``known`` is false when some access has no representable relation; the - retained footprints are then a lower bound over the accesses that are known. - """ - - footprints: tuple[BufferFootprint, ...] - known: bool - - @dataclass(frozen=True) class ValueLifetime: """One value's residency on the function's structured SSA event timeline. @@ -191,47 +165,28 @@ class ValueLifetime: last_used_at: int persistent: bool = False - @property - def level(self) -> str: - """Compatibility spelling for consumers rendering a memory level.""" - return self.memory_level - @dataclass(frozen=True) -class AllocationMetadata: - """What showing this function's addressable buffers fit took. +class RegionMemoryMetadata(IRMetadata): + """Record one region's memory behavior against a target hierarchy. - Where any of them would sit is the solver's business and appears nowhere - here. ``feasible`` means the first validated placement was returned without - claiming that its high-water mark is minimal. + Today this is attached to a Function, because traffic totals and placement + span the whole function. A footprint is absent until one has been computed; + no empty record is attached to a LoopRegion merely to reserve the type. """ solver_status: str - - -@dataclass(frozen=True) -class MemoryMetadata(IRMetadata): - """Record one function's memory behavior against a target hierarchy. - - Function attachment reflects that peaks span all live ranges. ``errors`` - reports a solved placement whose high-water exceeds stated capacity without - suppressing the rest of the analysis result. Advisories carry lower-severity - capacity findings. - - ``allocation`` is absent when the function has no addressable buffer to - place at the level being analysed, which is a different answer from having - placed one: nothing was decided, so nothing is claimed. - """ - - footprint: tuple[MemoryLevelFootprint, ...] = () + topologies: tuple[str, ...] = () + traffic: Traffic = Traffic() + footprint: Footprint | None = None lifetimes: tuple[ValueLifetime, ...] = () + peaks: tuple[MemoryLevelPeak, ...] = () errors: tuple[str, ...] = () advisories: tuple[str, ...] = () - allocation: "AllocationMetadata | None" = None - def memory_level(self, name: str) -> MemoryLevelFootprint | None: - """The footprint recorded for *name*, if the function touches it.""" - return next((item for item in self.footprint if item.memory_level == name), None) + def peak_for(self, name: str) -> MemoryLevelPeak | None: + """The peak recorded for *name*, if this region touches it.""" + return next((item for item in self.peaks if item.memory_level == name), None) @dataclass(frozen=True) @@ -297,18 +252,18 @@ class PerformanceSummaryMetadata(IRMetadata): __all__ = [ - "AllocationMetadata", "Breakdown", - "BufferFootprint", - "ComputeCostMetadata", - "LoopFootprintMetadata", - "MemoryLevelFootprint", "MemoryMetadata", + "ComputeCostMetadata", + "Footprint", + "MemoryLevelPeak", "PerformanceMetadata", "PerformanceSummaryMetadata", + "RegionMemoryMetadata", "RooflineMetadata", "Spread", "TimelineMetadata", + "Traffic", "TrafficBytes", "ValueLifetime", "breakdown", diff --git a/src/tilefoundry/analysis/performance.py b/src/tilefoundry/analysis/performance.py index ad1f1cfb..f3963269 100644 --- a/src/tilefoundry/analysis/performance.py +++ b/src/tilefoundry/analysis/performance.py @@ -19,11 +19,11 @@ from .iteration_scope import IterationScope from .metadata import ( ComputeCostMetadata, + MemoryMetadata, PerformanceMetadata, PerformanceSummaryMetadata, RooflineMetadata, TimelineMetadata, - TrafficMetadata, ) from .visitor import AnalyzeContext @@ -58,7 +58,7 @@ def default_visit_leaf( return scope = ctx.current if id(expr) in ctx.current.accesses["narrow"] else ctx.root cost = get_metadata(expr, ComputeCostMetadata) - moved = get_metadata(expr, TrafficMetadata) + moved = get_metadata(expr, MemoryMetadata) if cost is None or moved is None: raise AnalysisError(f"performance: missing compute/memory record for {expr!r}") if ctx.facts is None or ctx.services is None: diff --git a/src/tilefoundry/analysis/registry.py b/src/tilefoundry/analysis/registry.py index 77436b3b..0a67e1c0 100644 --- a/src/tilefoundry/analysis/registry.py +++ b/src/tilefoundry/analysis/registry.py @@ -29,15 +29,14 @@ def builtin_analyzer(selector: str) -> Analyzer | None: if selector == "memory": from tilefoundry.analysis.memory import analyze_memory # noqa: PLC0415 from tilefoundry.analysis.metadata import ( # noqa: PLC0415 - LoopFootprintMetadata, MemoryMetadata, - TrafficMetadata, + RegionMemoryMetadata, ) return Analyzer( "memory", analyze_memory, - produces=(MemoryMetadata, LoopFootprintMetadata, TrafficMetadata), + produces=(MemoryMetadata, RegionMemoryMetadata), ) if selector == "roofline": from tilefoundry.analysis.metadata import RooflineMetadata # noqa: PLC0415 diff --git a/src/tilefoundry/analysis/report.py b/src/tilefoundry/analysis/report.py index d345da0a..e60204fd 100644 --- a/src/tilefoundry/analysis/report.py +++ b/src/tilefoundry/analysis/report.py @@ -9,18 +9,15 @@ from types import UnionType from typing import Union, get_args, get_origin, get_type_hints -from tilefoundry.analysis.facts import MemoryHierarchyFacts -from tilefoundry.analysis.memory import cache_pressure from tilefoundry.analysis.metadata import ( Breakdown, ComputeCostMetadata, - LoopFootprintMetadata, MemoryMetadata, PerformanceMetadata, PerformanceSummaryMetadata, + RegionMemoryMetadata, RooflineMetadata, Spread, - TrafficMetadata, ) from tilefoundry.ir.core import Call, IRMetadata, binding_name, get_metadata from tilefoundry.ir.core.module import Module @@ -136,13 +133,12 @@ def _pair_types(declared: object) -> tuple[object, object] | None: for _record_type in ( ComputeCostMetadata, - TrafficMetadata, - LoopFootprintMetadata, - MemoryMetadata, RooflineMetadata, ): declare_record(_record_type) +declare_record(MemoryMetadata, family="memory") +declare_record(RegionMemoryMetadata, family="memory") declare_record(PerformanceMetadata, family="performance") declare_record(PerformanceSummaryMetadata, family="performance") @@ -169,7 +165,7 @@ def _operand_name(operand: object) -> str: return type(operand).__name__.lower() -def _operands(record: TrafficMetadata, expr: object) -> list[dict[str, object]] | None: +def _operands(record: MemoryMetadata, expr: object) -> list[dict[str, object]] | None: """Each recorded amount, against the operand it was charged to.""" if not isinstance(expr, Call) or not record.operands: return None @@ -186,7 +182,7 @@ def _operands(record: TrafficMetadata, expr: object) -> list[dict[str, object]] ] -expr_field(TrafficMetadata, "operands", _operands) +expr_field(MemoryMetadata, "operands", _operands) def _records_of(expr: object, selected: frozenset[type[IRMetadata]]) -> dict[str, object]: @@ -242,10 +238,10 @@ def report_data( "executed": list(executed), "function_records": function_records, "calls": _call_records(function, selected, call_labels), - "loops": _loop_records(function, selected, target), + "loops": _loop_records(function, selected), } available = set(metadata_types) - asked = {ComputeCostMetadata, TrafficMetadata} + asked = {ComputeCostMetadata, MemoryMetadata, RegionMemoryMetadata} if asked & selected or ("roofline" in analyses and asked & available): data["totals"] = _work_totals(function) return data @@ -277,20 +273,8 @@ def _call_records( def _loop_records( function: Function, selected: frozenset[type[IRMetadata]], - target, ) -> list[dict[str, object]]: """Every selected record attached to an authored loop.""" - memory = get_metadata(function, MemoryMetadata) - facts = ( - target.get_facts(MemoryHierarchyFacts) - if memory is not None and MemoryMetadata in selected - else None - ) - peaks = ( - {item.memory_level: item.peak_bytes for item in memory.footprint} - if memory is not None - else {} - ) rows: list[dict[str, object]] = [] for expr in collect_exprs(function.body): if not isinstance(expr, LoopRegion): @@ -298,11 +282,6 @@ def _loop_records( records = _records_of(expr, selected) if not records: continue - record = get_metadata(expr, LoopFootprintMetadata) - if record is not None and facts is not None: - pressure = cache_pressure(record, facts, peaks) - if pressure: - records["cache-pressure"] = list(pressure) rows.append({"value": expr.induction_var.name, **records}) return rows @@ -328,11 +307,11 @@ def _work_totals(function: Function) -> dict[str, object]: a program asks of the machine beside what it moves through it. """ record = get_metadata(function, ComputeCostMetadata) - moved = get_metadata(function, TrafficMetadata) + moved = get_metadata(function, RegionMemoryMetadata) return { "flops": _totals_of(None if record is None else record.flops), - "traffic": _totals_of(None if moved is None else moved.storage), - "communication": _totals_of(None if moved is None else moved.communication), + "traffic": _totals_of(None if moved is None else moved.traffic.storage), + "communication": _totals_of(None if moved is None else moved.traffic.communication), } diff --git a/src/tilefoundry/analysis/roofline.py b/src/tilefoundry/analysis/roofline.py index 6f24a22f..dbf10373 100644 --- a/src/tilefoundry/analysis/roofline.py +++ b/src/tilefoundry/analysis/roofline.py @@ -18,9 +18,10 @@ from .facts import ThroughputFacts from .metadata import ( ComputeCostMetadata, + MemoryMetadata, + RegionMemoryMetadata, RooflineMetadata, TrafficBytes, - TrafficMetadata, ) from .visitor import AnalyzeContext @@ -94,7 +95,9 @@ def _bound(compute_ns: int, memory_ns: int, *, has_work: bool) -> RooflineMetada def _cost_bound( - cost: ComputeCostMetadata, moved: TrafficMetadata, facts: ThroughputFacts + cost: ComputeCostMetadata, + moved: MemoryMetadata | RegionMemoryMetadata, + facts: ThroughputFacts, ) -> RooflineMetadata: """Bound one occurrence from the work it does and the bytes it moves. @@ -105,7 +108,7 @@ def _cost_bound( nanosecond is what this could have priced, so a dtype whose rate is missing still owes one and a level nobody rated does not. """ - reached = moved.storage.of(facts.bandwidth_level) + reached = moved.traffic.storage.of(facts.bandwidth_level) crossed = reached.total if reached is not None else TrafficBytes() return _bound( _compute_ns(_totals(cost.flops), facts), @@ -132,7 +135,7 @@ def default_visit_leaf( f"{describe_expr(expr)}: roofline needs the compute-cost record " "this call was never given" ) - moved = get_metadata(expr, TrafficMetadata) + moved = get_metadata(expr, MemoryMetadata) if moved is None: raise AnalysisError( f"{describe_expr(expr)}: roofline needs the traffic record the " @@ -155,7 +158,7 @@ def analyze_roofline( f"function {function.name!r}: roofline needs the compute-cost root " "record this function was never given" ) - moved = get_metadata(function, TrafficMetadata) + moved = get_metadata(function, RegionMemoryMetadata) if moved is None: raise AnalysisError( f"function {function.name!r}: roofline needs the traffic root record " diff --git a/src/tilefoundry/cli/analyze.py b/src/tilefoundry/cli/analyze.py index d964e5c0..0eb4119f 100644 --- a/src/tilefoundry/cli/analyze.py +++ b/src/tilefoundry/cli/analyze.py @@ -27,7 +27,7 @@ EVIDENCE: dict[str, str] = { "compute-cost": "the logical work and traffic of every value: flops by dtype, bytes moved", - "memory": "where that traffic lands, and the footprint it holds live against the capacity", + "memory": "where traffic lands and what storage remains live against capacity", "roofline": "which of compute or memory limits each value, and the limit in time", "performance": "when each value runs, where its buffers fit, and the time that takes", } @@ -85,17 +85,17 @@ def guidance() -> str: compute-cost nothing. Every kind states its total never and every level's per-unit share memory nothing for traffic, which states every the program shards - level. Footprint follows its owner + level; placement remains Function-wide roofline nothing. The bound is the machine's never and is unchanged by program splits performance which level's parallel capacity the the program shards plan is issued against Two assumptions the reported numbers rest on: - global traffic is the device's and counted once, so units reading one operand - in common are assumed to read it from memory once. Whether they can is - a residency question, reported as an advisory. - a reported peak footprint holds under the order this walk took. Which + logical traffic omits loop replication that does not change an access; + total traffic counts every executed occurrence. Per-unit traffic is + the selected topology unit's share of that executed total. + a reported placement peak holds under the order this walk took. Which order the program really takes is settled by scheduling, so the peak is an observation, not a bound. diff --git a/src/tilefoundry/inspection/analysis_report.py b/src/tilefoundry/inspection/analysis_report.py index 8510972d..1b732a34 100644 --- a/src/tilefoundry/inspection/analysis_report.py +++ b/src/tilefoundry/inspection/analysis_report.py @@ -7,10 +7,9 @@ from tilefoundry.analysis.api import AnalysisResult from tilefoundry.analysis.metadata import ( ComputeCostMetadata, - MemoryMetadata, PerformanceSummaryMetadata, + RegionMemoryMetadata, RooflineMetadata, - TrafficMetadata, ) from tilefoundry.analysis.report import _type_text as _type_text from tilefoundry.analysis.report import ( @@ -24,12 +23,10 @@ from tilefoundry.inspection.values import ( AdvisorySummary, ErrorSummary, - MemorySummary, PerformanceSummaryView, Prose, ReportIdentity, ReportSelection, - peak_footprint, render_comment, ) from tilefoundry.ir.core import IRMetadata, get_metadata @@ -101,12 +98,11 @@ def _summary( ] if "totals" in data and "compute-cost" in data["executed"]: views.append(get_metadata(function, ComputeCostMetadata) or ComputeCostMetadata()) - if "traffic" in function_records: - views.append(get_metadata(function, TrafficMetadata) or TrafficMetadata()) if "memory" in function_records: - memory = get_metadata(function, MemoryMetadata) - views.append(MemorySummary(peak_footprint(memory))) - if MemoryMetadata in selected: + memory = get_metadata(function, RegionMemoryMetadata) + assert memory is not None + views.append(memory) + if RegionMemoryMetadata in selected: views.extend(ErrorSummary(Prose(note)) for note in memory.errors) views.extend(AdvisorySummary(Prose(note)) for note in memory.advisories) if "roofline" in function_records: diff --git a/src/tilefoundry/inspection/values.py b/src/tilefoundry/inspection/values.py index 09b72b44..5be5e907 100644 --- a/src/tilefoundry/inspection/values.py +++ b/src/tilefoundry/inspection/values.py @@ -4,7 +4,7 @@ import json from collections.abc import Mapping -from dataclasses import dataclass, field, fields, is_dataclass +from dataclasses import dataclass, fields, is_dataclass from tilefoundry.ir.core.metadata import IRMetadata from tilefoundry.ir.core.values import TripInterval @@ -41,11 +41,6 @@ class ReportSelection(IRMetadata): executed: tuple[str, ...] = () -@dataclass(frozen=True) -class MemorySummary(IRMetadata): - peak_bytes: dict[str, int] = field(default_factory=dict) - - @dataclass(frozen=True) class AdvisorySummary(IRMetadata): text: Prose @@ -144,8 +139,8 @@ def print_ComputeCostMetadata(self, record, **_): ), ) - def print_TrafficMetadata(self, record, *, opt_in=frozenset()): - traffic = self._breakdown(record.storage, record.topologies, logical=False) + def print_MemoryMetadata(self, record, *, opt_in=frozenset()): + traffic = self._breakdown(record.traffic.storage, record.topologies) values = [("traffic", traffic)] if "operands" in opt_in: last = len(record.operands) - 1 @@ -154,31 +149,20 @@ def print_TrafficMetadata(self, record, *, opt_in=frozenset()): for index, moved in enumerate(record.operands) } values.append(("operands", operands)) - return self._record("traffic", values) + return self._record("memory", values) - def print_MemoryMetadata(self, record, **_): + def print_RegionMemoryMetadata(self, record, **_): return self._record( "memory", ( - ("peak", {item.memory_level: item.peak_bytes for item in record.footprint}), - ("persistent", sum(item.persistent_bytes for item in record.footprint), 0), + ("traffic", self._breakdown(record.traffic.storage, record.topologies)), + ("peak", {item.memory_level: item.peak_bytes for item in record.peaks}), + ("persistent", sum(item.persistent_bytes for item in record.peaks), 0), ("errors", len(record.errors), 0), ("advisories", len(record.advisories), 0), ), ) - def print_LoopFootprintMetadata(self, record, **_): - footprints = { - f"{item.buffer}@{item.memory_level}": PAIR.join( - str(value) for value in (item.bytes, item.device_bytes, item.repeated_bytes) - ) - for item in record.footprints - } - return self._record( - "loop-footprint", - (("footprints", footprints), ("status", "complete" if record.known else "lower-bound")), - ) - def print_RooflineMetadata(self, record, **_): return self._record( "roofline", (("ideal_ns", record.ideal_ns, 0), ("bound_by", record.bound_by, "none")) @@ -216,9 +200,6 @@ def print_ReportSelection(self, record, **_): "selection", (("requested", record.requested, ()), ("executed", record.executed, ())) ) - def print_MemorySummary(self, record, **_): - return self._single("peak-footprint", record.peak_bytes, {}) - def print_AdvisorySummary(self, record, **_): return self._single("advisory", record.text) @@ -244,10 +225,6 @@ def render_comment(record, *, opt_in=frozenset()): return method(record, opt_in=opt_in) if method else None -def peak_footprint(record): - return {item.memory_level: item.peak_bytes for item in record.footprint} - - __all__ = [ "CommentPrinter", "Prose", @@ -262,9 +239,7 @@ def peak_footprint(record): "TRIPS", "ReportIdentity", "ReportSelection", - "MemorySummary", "AdvisorySummary", "ErrorSummary", "PerformanceSummaryView", - "peak_footprint", ] diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 9d023f9c..fb6f65b1 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -19,6 +19,7 @@ _LiteralStoreOffset, _SymbolicStoreOffset, ) +from tests.models.access_footprint.model import TiledQKVProjection from tilefoundry import func, module from tilefoundry.analysis import ( Breakdown, @@ -28,10 +29,11 @@ PerformanceMetadata, PerformanceServiceFacts, PerformanceSummaryMetadata, + RegionMemoryMetadata, RooflineMetadata, Spread, ThroughputFacts, - TrafficMetadata, + Traffic, ) from tilefoundry.analysis.api import analyze from tilefoundry.analysis.compute_cost import ( @@ -53,6 +55,23 @@ from tilefoundry.visitor_registry.contexts import TrafficBytes _ROUNDING_M = 14_593 + + +def test_invariant_gemm_operands_repeat_in_total_traffic() -> None: + """The QKV tiles load again even when one operand is invariant in one loop.""" + result = analyze( + TiledQKVProjection, + TiledQKVProjection.entry_function(), + analysis="memory", + ) + record = get_metadata(result.function, RegionMemoryMetadata) + gmem = record.traffic.storage.of("gmem") + + assert gmem is not None + assert gmem.logical == TrafficBytes(read=27_262_976, write=16_777_216) + assert gmem.total == TrafficBytes(read=142_606_336, write=16_777_216) + + _ROUNDING_N = 11_489 _ROUNDING_K = 298_224_413 _H200 = CudaTarget("nvidia.h200_sxm") @@ -246,7 +265,7 @@ def test_a_symbolic_store_stride_preserves_the_literal_control_result() -> None: literal_service, literal_local, literal_roofline, literal_performance = observed["literal"] symbolic_service, symbolic_local, symbolic_roofline, symbolic_performance = observed["symbolic"] - assert literal_roofline == symbolic_roofline == 139_407 + assert literal_roofline == symbolic_roofline == 398_459 assert symbolic_local - literal_local == 6 assert symbolic_service - literal_service == 6 * 128 assert symbolic_performance - literal_performance == 6 @@ -351,7 +370,7 @@ def test_a_program_whose_peak_exceeds_capacity_reports_an_error() -> None: roomy = replace(_SharedTile, target=_RoomyShared("nvidia.h200_sxm")) tight_memory = analyze(tight, split, analysis="memory") - tight_record = get_metadata(tight_memory.function, MemoryMetadata) + tight_record = get_metadata(tight_memory.function, RegionMemoryMetadata) assert tight_record.errors == ("smem placement peak 211200 B exceeds capacity 105600 B",) assert '# error="smem placement peak 211200 B exceeds capacity 105600 B"' in render_text( render_analysis(tight_memory) @@ -361,7 +380,9 @@ def test_a_program_whose_peak_exceeds_capacity_reports_an_error() -> None: ] tight_performance = analyze(tight, split, analysis="performance") - assert get_metadata(tight_performance.function, MemoryMetadata).errors == tight_record.errors + assert ( + get_metadata(tight_performance.function, RegionMemoryMetadata).errors == tight_record.errors + ) unrestated = next(item for item in _SharedTile.functions if item.name == "split") held = get_metadata( @@ -371,9 +392,9 @@ def test_a_program_whose_peak_exceeds_capacity_reports_an_error() -> None: analysis="memory", options=MemoryOptions(timeout_seconds=1.0), ).function, - MemoryMetadata, - ).footprint - assert next(item.peak_bytes for item in held if item.level == "smem") == 211_200 + RegionMemoryMetadata, + ).peaks + assert next(item.peak_bytes for item in held if item.memory_level == "smem") == 211_200 fits = analyze( roomy, @@ -382,8 +403,8 @@ def test_a_program_whose_peak_exceeds_capacity_reports_an_error() -> None: ) summary = get_metadata(fits.function, PerformanceSummaryMetadata) assert summary is not None - fits_memory = get_metadata(fits.function, MemoryMetadata) - assert fits_memory.allocation.solver_status == "feasible" + fits_memory = get_metadata(fits.function, RegionMemoryMetadata) + assert fits_memory.solver_status == "feasible" assert fits_memory.errors == () assert summary.timeline.end_ns > 0 @@ -394,11 +415,11 @@ def test_a_program_whose_peak_exceeds_capacity_reports_an_error() -> None: analysis="memory", ) assert [ - (item.binding, item.level, item.bytes, item.defined_at, item.last_used_at) - for item in get_metadata(fits.function, MemoryMetadata).lifetimes + (item.binding, item.memory_level, item.bytes, item.defined_at, item.last_used_at) + for item in get_metadata(fits.function, RegionMemoryMetadata).lifetimes ] == [ - (item.binding, item.level, item.bytes, item.defined_at, item.last_used_at) - for item in get_metadata(relieved.function, MemoryMetadata).lifetimes + (item.binding, item.memory_level, item.bytes, item.defined_at, item.last_used_at) + for item in get_metadata(relieved.function, RegionMemoryMetadata).lifetimes ] @@ -442,19 +463,21 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N level="cta", ) - crossed = TrafficMetadata( + crossed = MemoryMetadata( topologies=("cta",), - storage=Breakdown( - ( + traffic=Traffic( + storage=Breakdown( ( - throughput.bandwidth_level, - Spread( - TrafficBytes(read=4096), - TrafficBytes(read=4096), - (TrafficBytes(read=4096),), + ( + throughput.bandwidth_level, + Spread( + TrafficBytes(read=4096), + TrafficBytes(read=4096), + (TrafficBytes(read=4096),), + ), ), - ), - ) + ) + ), ), ) with pytest.raises( diff --git a/tests/analysis/test_analysis_invariants.py b/tests/analysis/test_analysis_invariants.py index 3c4e62fd..e13833a2 100644 --- a/tests/analysis/test_analysis_invariants.py +++ b/tests/analysis/test_analysis_invariants.py @@ -245,8 +245,8 @@ def measured(): both = measured() assert both.operands == (TrafficBytes(), TrafficBytes(read=12), TrafficBytes()) assert ( - both.storage.of("gmem").total, - both.storage.of("rmem").total, + both.traffic.storage.of("gmem").total, + both.traffic.storage.of("rmem").total, ) == ( TrafficBytes(read=4), TrafficBytes(read=8), @@ -262,11 +262,11 @@ def measured(): assert one.operands == (TrafficBytes(), TrafficBytes(read=8), TrafficBytes()), ( "the second number is eight bytes wide" ) - assert one.storage.of("rmem").total == TrafficBytes(read=8), ( + assert one.traffic.storage.of("rmem").total == TrafficBytes(read=8), ( "and it lives at rmem, so gmem was not touched at all" ) - assert one.storage.of("gmem") is None, "gmem was not touched at all" - assert one.storage.of("rmem").per_unit == (TrafficBytes(read=8),), ( + assert one.traffic.storage.of("gmem") is None, "gmem was not touched at all" + assert one.traffic.storage.of("rmem").per_unit == (TrafficBytes(read=8),), ( "one CTA is the only unit, so its share is the whole" ) diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index 1d19c498..06b42bd9 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -24,12 +24,11 @@ from tilefoundry.analysis import ( AnalysisResult, ComputeCostMetadata, - LoopFootprintMetadata, MemoryMetadata, PerformanceMetadata, PerformanceSummaryMetadata, + RegionMemoryMetadata, RooflineMetadata, - TrafficMetadata, analyze, ) from tilefoundry.analysis.access import Access, AccessPrecision @@ -89,6 +88,11 @@ class _PersistentScheduleExpectation: "gmem": 283_752, "rmem": 0, }, + "hand_checked.InvariantReuse.reuse[static]": { + "gmem": 80, + "rmem": 0, + "smem": 64, + }, "leaf_weights.Mod.entry[static]": { "gmem": 51_539_608_064, "rmem": 0, @@ -300,16 +304,16 @@ def assert_performance_contract(result: AnalysisResult) -> None: at the target's rates, and a solve that proved nothing says so. One a loop repeats is written once, so its interval is that many of its own durations and its last trip still lands inside the prediction that contains it. - A loop is not an occurrence and carries no timeline of its own, and still - states the buffers it touches. + A loop is not an occurrence and carries neither a timeline nor a placeholder + memory record of its own. """ fn = result.function summary = get_metadata(fn, PerformanceSummaryMetadata) assert summary is not None assert 0 <= summary.timeline.start_ns <= summary.timeline.end_ns - placement = get_metadata(fn, MemoryMetadata) - assert placement is not None and placement.allocation is not None - assert placement.allocation.solver_status in ("optimal", "feasible") + placement = get_metadata(fn, RegionMemoryMetadata) + assert placement is not None + assert placement.solver_status in ("optimal", "feasible") predicted_ns = summary.timeline.end_ns - summary.timeline.start_ns assert summary.waves > 0 and predicted_ns % summary.waves == 0 bound = get_metadata(fn, RooflineMetadata) @@ -329,7 +333,7 @@ def assert_performance_contract(result: AnalysisResult) -> None: cost, throughput, services, - moved=get_metadata(expr, TrafficMetadata), + moved=get_metadata(expr, MemoryMetadata), level=result.level, ) record = get_metadata(expr, PerformanceMetadata) @@ -369,7 +373,6 @@ def assert_performance_contract(result: AnalysisResult) -> None: continue assert get_metadata(expr, PerformanceMetadata) is None, describe_expr(expr) assert get_metadata(expr, PerformanceSummaryMetadata) is None, describe_expr(expr) - assert get_metadata(expr, LoopFootprintMetadata) is not None, describe_expr(expr) @pytest.mark.parametrize( @@ -401,7 +404,7 @@ def test_more_of_the_same_work_is_never_predicted_to_take_less_time(smaller, lar def _every_number_counts_something(result: AnalysisResult) -> None: """Every quantity these four families report is a count, so none is below zero. - Work, bytes, a footprint and a bound are all counts of something that + Work, moved bytes, placement peaks and a bound are all counts of something that happened or has to happen. A negative one is not a small answer but a derivation that ran backwards -- a projection dividing what it should have multiplied, or a difference taken the wrong way round -- and it would then be @@ -411,8 +414,8 @@ def _every_number_counts_something(result: AnalysisResult) -> None: for expr in (fn, *collect_exprs(fn.body)): for record, rows in ( (ComputeCostMetadata, ()), - (TrafficMetadata, ()), (MemoryMetadata, ()), + (RegionMemoryMetadata, ()), (RooflineMetadata, ()), (PerformanceMetadata, ()), ): @@ -432,18 +435,21 @@ def _every_number_counts_something(result: AnalysisResult) -> None: for name, spread in breakdown.kinds: for value in (spread.logical, spread.total, *spread.per_unit): assert value >= 0, f"{describe_expr(expr)}: {field}[{name}] = {value}" - if record is TrafficMetadata: - for field in ("whole", "per_unit"): - for level, moved in getattr(held, field): - assert moved.read >= 0 and moved.write >= 0, ( - f"{describe_expr(expr)}: {field}[{level}] = {moved}" - ) + if record in (MemoryMetadata, RegionMemoryMetadata): + for field in ("storage", "communication"): + breakdown = getattr(held.traffic, field) + for level, spread in breakdown.kinds: + for moved in (spread.logical, spread.total, *spread.per_unit): + assert moved.read >= 0 and moved.write >= 0, ( + f"{describe_expr(expr)}: {field}[{level}] = {moved}" + ) + if record is MemoryMetadata: for position, moved in enumerate(held.operands): assert moved.read >= 0 and moved.write >= 0, ( f"{describe_expr(expr)}: operand {position} = {moved}" ) - if record is MemoryMetadata: - for level in held.footprint: + if record is RegionMemoryMetadata: + for level in held.peaks: assert level.peak_bytes >= 0 and level.persistent_bytes >= 0 for item in held.lifetimes: assert item.bytes >= 0 and 0 <= item.defined_at <= item.last_used_at @@ -452,16 +458,6 @@ def _every_number_counts_something(result: AnalysisResult) -> None: assert held.ideal_ns >= 0 and held.compute_ns >= 0 and held.memory_ns >= 0 if record is PerformanceMetadata: assert 0 <= held.timeline.start_ns <= held.timeline.end_ns - for expr in collect_exprs(fn.body): - record = get_metadata(expr, LoopFootprintMetadata) - if record is None: - continue - rows = [(item.buffer, item.level) for item in record.footprints] - assert rows == sorted(rows), describe_expr(expr) - assert len(rows) == len(set(rows)), describe_expr(expr) - for item in record.footprints: - assert item.bytes >= 0 and item.device_bytes >= 0 and item.repeated_bytes >= 0 - assert " None: assert result.module is owner assert set(result.executed) == set(FAMILIES) assert_performance_contract(result) - placement = get_metadata(result.function, MemoryMetadata) + placement = get_metadata(result.function, RegionMemoryMetadata) assert placement is not None - observed = {item.level: item.peak_bytes for item in placement.footprint} + observed = {item.memory_level: item.peak_bytes for item in placement.peaks} assert observed == EXPECTED_MEMORY_PEAKS[case.id] expected_schedule = EXPECTED_PERSISTENT_SCHEDULES.get(case.id) if expected_schedule is not None: diff --git a/tests/analysis/test_analyze_by_hand.py b/tests/analysis/test_analyze_by_hand.py new file mode 100644 index 00000000..7d10926f --- /dev/null +++ b/tests/analysis/test_analyze_by_hand.py @@ -0,0 +1,201 @@ +"""Exact analysis values for programs small enough to compute on paper.""" + +from __future__ import annotations + +from tests.fixtures.placed.hand_checked import InvariantReuse +from tilefoundry.analysis import analyze +from tilefoundry.analysis.report import report_data + + +def _checked(data: dict) -> dict: + """Keep the movement and placement conclusions this fixture hand-checks.""" + function = data["function_records"]["memory"] + return { + "loops": data["loops"], + "calls": data["calls"], + "function_memory": { + "topologies": function["topologies"], + "traffic": function["traffic"], + "footprint": function["footprint"], + "peaks": function["peaks"], + "solver_status": function["solver_status"], + "errors": function["errors"], + "advisories": function["advisories"], + }, + } + + +INVARIANT_REUSE = { + "loops": [], + "calls": [ + { + "value": "v0", + "memory": { + "topologies": ["cta"], + "traffic": { + "storage": { + "smem": { + "logical": {"read": 0, "write": 16}, + "total": {"read": 0, "write": 16}, + "per_unit": [{"read": 0, "write": 16}], + } + }, + "communication": {}, + }, + "operands": [ + { + "arg": "result", + "name": "v0", + "type": "bf16[4,2] smem", + "read": 0, + "write": 16, + } + ], + "footprint": None, + }, + }, + { + "value": "v1", + "memory": { + "topologies": ["cta"], + "traffic": { + "storage": { + "rmem": { + "logical": {"read": 16, "write": 0}, + "total": {"read": 16, "write": 0}, + "per_unit": [{"read": 16, "write": 0}], + } + }, + "communication": {}, + }, + "operands": [ + { + "arg": 0, + "name": "x", + "type": "bf16[8,4] gmem", + "read": 0, + "write": 0, + }, + { + "arg": 1, + "name": "tuple", + "type": "i64[]+ umat", + "read": 16, + "write": 0, + }, + { + "arg": "result", + "name": "v1", + "type": "bf16[4,2] gmem", + "read": 0, + "write": 0, + }, + ], + "footprint": None, + }, + }, + { + "value": "v2", + "memory": { + "topologies": ["cta"], + "traffic": { + "storage": { + "gmem": { + "logical": {"read": 16, "write": 0}, + "total": {"read": 16, "write": 0}, + "per_unit": [{"read": 16, "write": 0}], + }, + "smem": { + "logical": {"read": 0, "write": 16}, + "total": {"read": 0, "write": 16}, + "per_unit": [{"read": 0, "write": 16}], + }, + }, + "communication": {}, + }, + "operands": [ + { + "arg": 0, + "name": "v1", + "type": "bf16[4,2] gmem", + "read": 16, + "write": 0, + }, + { + "arg": "result", + "name": "v2", + "type": "bf16[4,2] smem", + "read": 0, + "write": 16, + }, + ], + "footprint": None, + }, + }, + ], + "function_memory": { + "topologies": ["cta"], + "traffic": { + "storage": { + "gmem": { + "logical": {"read": 64, "write": 0}, + "total": {"read": 192, "write": 0}, + "per_unit": [{"read": 192, "write": 0}], + }, + "rmem": { + "logical": {"read": 64, "write": 0}, + "total": {"read": 192, "write": 0}, + "per_unit": [{"read": 192, "write": 0}], + }, + "smem": { + "logical": {"read": 0, "write": 80}, + "total": {"read": 0, "write": 208}, + "per_unit": [{"read": 0, "write": 208}], + }, + }, + "communication": {}, + }, + "footprint": None, + "peaks": [ + { + "memory_level": "gmem", + "peak_bytes": 80, + "persistent_bytes": 64, + "capacity_bytes": 141_000_000_000, + }, + { + "memory_level": "rmem", + "peak_bytes": 0, + "persistent_bytes": 0, + "capacity_bytes": 262_144, + }, + { + "memory_level": "smem", + "peak_bytes": 64, + "persistent_bytes": 0, + "capacity_bytes": 232_448, + }, + ], + "solver_status": "feasible", + "errors": [], + "advisories": [], + }, +} + + +def test_invariant_reuse_matches_the_written_arithmetic() -> None: + result = analyze( + InvariantReuse, + InvariantReuse.entry_function(), + analysis=("memory",), + ) + data = report_data( + module=result.module, + function=result.function, + analyses=result.analyses, + topology_level=result.topology_level, + executed=result.executed, + metadata_types=result.metadata_types, + ) + + assert _checked(data) == INVARIANT_REUSE diff --git a/tests/analysis/test_analyze_cross_module.py b/tests/analysis/test_analyze_cross_module.py index c9fd8c44..6eeef3bc 100644 --- a/tests/analysis/test_analyze_cross_module.py +++ b/tests/analysis/test_analyze_cross_module.py @@ -17,7 +17,7 @@ from tilefoundry import func, module from tilefoundry.analysis.api import analyze from tilefoundry.analysis.errors import AnalysisError -from tilefoundry.analysis.metadata import ComputeCostMetadata, TrafficMetadata +from tilefoundry.analysis.metadata import ComputeCostMetadata, MemoryMetadata from tilefoundry.dsl import ConstTensor, DimVar, Tensor, Topology, tf from tilefoundry.ir.core import Call, get_metadata from tilefoundry.ir.core.module import reachable_functions @@ -33,14 +33,14 @@ _CTA = (Topology("cta", 132),) -def _matmul_records(result) -> tuple[tuple[ComputeCostMetadata, TrafficMetadata], ...]: +def _matmul_records(result) -> tuple[tuple[ComputeCostMetadata, MemoryMetadata], ...]: """Both halves of the record on each inlined MatMul occurrence.""" records = [] for expr in collect_exprs(result.function.body): if not isinstance(expr, Call) or not isinstance(expr.target, MatMul): continue record = get_metadata(expr, ComputeCostMetadata) - moved = get_metadata(expr, TrafficMetadata) + moved = get_metadata(expr, MemoryMetadata) assert record is not None records.append((record, moved)) return tuple(records) @@ -60,7 +60,7 @@ def _traffic(records) -> dict[str, int]: total: dict[str, int] = {} for _record, moved in records: assert moved is not None, "traffic was asked of a run that did not measure it" - for name, spread in moved.storage.kinds: + for name, spread in moved.traffic.storage.kinds: total[name] = total.get(name, 0) + spread.total.total_bytes return total @@ -174,7 +174,7 @@ def fused(x: Tensor[(4, 8), "f32"]) -> Tensor[(4, 8), "f32"]: analysis=("compute-cost", "memory"), dims={"child_topology_extent": 132}, ) - assert set(result.metadata_types) >= {ComputeCostMetadata, TrafficMetadata} + assert set(result.metadata_types) >= {ComputeCostMetadata, MemoryMetadata} @pytest.mark.parametrize("name,root", REFERENCE_PROGRAMS, ids=[n for n, _ in REFERENCE_PROGRAMS]) @@ -182,7 +182,7 @@ def test_each_reference_program_is_one_analysable_kernel(name, root) -> None: """The shared programs measure as one root each; what a call means is not restated.""" result = analyze(root, root.entry_function(), analysis=("compute-cost", "memory")) - assert set(result.metadata_types) >= {ComputeCostMetadata, TrafficMetadata} + assert set(result.metadata_types) >= {ComputeCostMetadata, MemoryMetadata} def _placed_primitives(fn) -> list[tuple[str, object]]: diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 4dbf9d20..45623efd 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -168,7 +168,8 @@ def test_analyze_help_explains_topology_effects_and_assumptions(capsys) -> None: assert family in help_text assert "logical" in help_text and "per-unit share" in help_text assert "traffic" in help_text - assert "global traffic is the device's and counted once" in help_text + assert "logical traffic omits loop replication" in help_text + assert "total traffic counts every executed occurrence" in help_text assert "is an observation, not a bound" in help_text @@ -722,7 +723,6 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> ) assert set(payload["function_records"]) == { "compute-cost", - "traffic", "memory", "roofline", "performance", @@ -737,8 +737,8 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> summary = payload["function_records"]["performance"] cost = payload["function_records"]["compute-cost"] - moved = payload["function_records"]["traffic"] - peak = payload["function_records"]["memory"]["footprint"] + moved = payload["function_records"]["memory"] + peak = moved["peaks"] bound = payload["function_records"]["roofline"] assert header.splitlines() == [ "# analysis target=nvidia.h200_sxm module=MoEMegaKernel function=experts " @@ -749,12 +749,15 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> f"flops=f32:{cost['flops']['f32']['logical']}@logical," f"{cost['flops']['f32']['total']}@total," f"{cost['flops']['f32']['per_unit'][0]}@{payload['topology']}", - "# traffic " - f"traffic=gmem:r{moved['storage']['gmem']['total']['read']}" - f"/w{moved['storage']['gmem']['total']['write']}@total," - f"r{moved['storage']['gmem']['per_unit'][0]['read']}" - f"/w{moved['storage']['gmem']['per_unit'][0]['write']}@{payload['topology']}", - f"# peak-footprint=gmem:{peak[0]['peak_bytes']}", + "# memory " + f"traffic=gmem:r{moved['traffic']['storage']['gmem']['logical']['read']}" + f"/w{moved['traffic']['storage']['gmem']['logical']['write']}@logical," + f"r{moved['traffic']['storage']['gmem']['total']['read']}" + f"/w{moved['traffic']['storage']['gmem']['total']['write']}@total," + f"r{moved['traffic']['storage']['gmem']['per_unit'][0]['read']}" + f"/w{moved['traffic']['storage']['gmem']['per_unit'][0]['write']}" + f"@{payload['topology']} peak=gmem:{peak[0]['peak_bytes']} " + f"persistent={sum(item['persistent_bytes'] for item in peak)}", f"# roofline ideal-ns={bound['ideal_ns']} bound-by={bound['bound_by']}", "# performance root=MoEMegaKernel::experts " f"predicted-ns={summary['timeline']['end_ns']} " @@ -764,7 +767,7 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> name: spread["total"] for name, spread in cost["flops"].items() } assert payload["totals"]["traffic"] == { - name: value["total"] for name, value in moved["storage"].items() + name: value["total"] for name, value in moved["traffic"]["storage"].items() } hoisted = {line.split(" = ", 1)[0] for line in lines if " = Mesh((Topology(" in line} @@ -785,7 +788,7 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> rows = payload["calls"] assert len(rows) == 7 assert all( - set(row) - {"performance"} == {"value", "compute-cost", "traffic", "roofline"} + set(row) - {"performance"} == {"value", "compute-cost", "memory", "roofline"} for row in rows ) timed = [index for index, row in enumerate(rows) if "performance" in row] diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index b6ec3c9e..f5c373db 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -19,21 +19,21 @@ def composed_mesh_pipeline( ), "rmem"] ): with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: - v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r2048/w0@total,r512/w0@cta,r64/w0@thread;rmem:r0/w2048@total,r0/w512@cta,r0/w64@thread - v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@logical,512@total,128@cta,16@thread; traffic traffic=rmem:r2048/w2048@total,r512/w512@cta,r64/w64@thread - v2 = reshard(v1, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost; traffic traffic=rmem:r2048/w0@total,r512/w0@cta,r64/w0@thread;smem:r0/w2048@total,r0/w512@cta,r0/w64@thread - v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost flops=bf16:512@logical,512@total,128@cta,128@thread; traffic traffic=smem:r2048/w1024@total,r512/w256@cta,r512/w256@thread - v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta.tile, 2, 16, 4), (128, 64, 1, 16)), "smem"]; compute-cost; traffic traffic=smem:r1024/w1024@total,r256/w256@cta,r256/w256@thread - v5 = reshard(v4, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {cta.tile @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@total,r0/w256@cta,r0/w256@thread;smem:r1024/w0@total,r256/w0@cta,r256/w0@thread - v6 = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r64/w0@total,r64/w0@cta,r8/w0@thread;rmem:r0/w64@total,r0/w64@cta,r0/w8@thread - folded = reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic - summed = reshard(mixed, layout=((8, 16), {mesh.warp @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost; traffic - for _ in range(3): # loop-footprint footprints=folded@rmem:8192/196608/24576;summed@rmem:131072/393216/393216 status=complete - v10 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost flops=f32:128@logical,128@total,128@cta,128@thread; traffic traffic=rmem:r1024/w512@total,r1024/w512@cta,r1024/w512@thread - v11 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:128@logical,128@total,128@cta,16@thread; traffic traffic=rmem:r512/w512@total,r512/w512@cta,r64/w64@thread + v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost; memory traffic=gmem:r2048/w0@logical,r2048/w0@total,r512/w0@cta,r64/w0@thread;rmem:r0/w2048@logical,r0/w2048@total,r0/w512@cta,r0/w64@thread + v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@logical,512@total,128@cta,16@thread; memory traffic=rmem:r2048/w2048@logical,r2048/w2048@total,r512/w512@cta,r64/w64@thread + v2 = reshard(v1, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost; memory traffic=rmem:r2048/w0@logical,r2048/w0@total,r512/w0@cta,r64/w0@thread;smem:r0/w2048@logical,r0/w2048@total,r0/w512@cta,r0/w64@thread + v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost flops=bf16:512@logical,512@total,128@cta,128@thread; memory traffic=smem:r2048/w1024@logical,r2048/w1024@total,r512/w256@cta,r512/w256@thread + v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta.tile, 2, 16, 4), (128, 64, 1, 16)), "smem"]; compute-cost; memory traffic=smem:r1024/w1024@logical,r1024/w1024@total,r256/w256@cta,r256/w256@thread + v5 = reshard(v4, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {cta.tile @ B()})]; compute-cost; memory traffic=gmem:r0/w1024@logical,r0/w1024@total,r0/w256@cta,r0/w256@thread;smem:r1024/w0@logical,r1024/w0@total,r256/w0@cta,r256/w0@thread + v6 = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), "rmem"]; compute-cost; memory traffic=gmem:r64/w0@logical,r64/w0@total,r64/w0@cta,r8/w0@thread;rmem:r0/w64@logical,r0/w64@total,r0/w64@cta,r0/w8@thread + folded = reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; memory + summed = reshard(mixed, layout=((8, 16), {mesh.warp @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost; memory + for _ in range(3): + v10 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost flops=f32:128@logical,128@total,128@cta,128@thread; memory traffic=rmem:r1024/w512@logical,r1024/w512@total,r1024/w512@cta,r1024/w512@thread + v11 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:128@logical,128@total,128@cta,16@thread; memory traffic=rmem:r512/w512@logical,r512/w512@total,r512/w512@cta,r64/w64@thread folded = v11 summed = v10 - v13 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; traffic traffic=gmem:r0/w512@total,r0/w512@cta,r0/w64@thread;rmem:r512/w0@total,r512/w0@cta,r64/w0@thread - v15 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; traffic traffic=gmem:r0/w512@total,r0/w512@cta,r0/w512@thread;rmem:r512/w0@total,r512/w0@cta,r512/w0@thread - v17 = reshard(v6, layout=Layout((16,), (1,)), storage=gmem) # Tensor[(16,), "f32", Layout((16,), (1,))]; compute-cost; traffic traffic=gmem:r0/w64@total,r0/w64@cta,r0/w8@thread;rmem:r64/w0@total,r64/w0@cta,r8/w0@thread + v13 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; memory traffic=gmem:r0/w512@logical,r0/w512@total,r0/w512@cta,r0/w64@thread;rmem:r512/w0@logical,r512/w0@total,r512/w0@cta,r64/w0@thread + v15 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; memory traffic=gmem:r0/w512@logical,r0/w512@total,r0/w512@cta,r0/w512@thread;rmem:r512/w0@logical,r512/w0@total,r512/w0@cta,r512/w0@thread + v17 = reshard(v6, layout=Layout((16,), (1,)), storage=gmem) # Tensor[(16,), "f32", Layout((16,), (1,))]; compute-cost; memory traffic=gmem:r0/w64@logical,r0/w64@total,r0/w64@cta,r0/w8@thread;rmem:r64/w0@logical,r64/w0@total,r64/w0@cta,r8/w0@thread return (v5, v13, v15, v17) diff --git a/tests/fixtures/placed/hand_checked.py b/tests/fixtures/placed/hand_checked.py new file mode 100644 index 00000000..2bc5e2c3 --- /dev/null +++ b/tests/fixtures/placed/hand_checked.py @@ -0,0 +1,39 @@ +"""Small placed programs whose analysis results fit in handwritten arithmetic.""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +S, K, N = 8, 4, 6 +BM, BK, BN = 4, 2, 2 + +_H200 = CudaTarget("nvidia.h200_sxm") + + +@module(entry="reuse", target=_H200, topologies=(Topology("cta", 1),)) +class InvariantReuse: + """Expose an ``n``-invariant read without making it dead work. + + One ``x`` tile is ``BM * BK * sizeof(bf16) = 4 * 2 * 2 = 16 B``. + The authored nest executes it ``(S/BM) * (N/BN) * (K/BK) = 2 * 3 * 2`` + times, so its traffic is 192 B. Its logical account omits the invariant + ``n`` replication and is ``2 * 2 * 16 = 64 B``. + """ + + @func + def reuse( + x: Tensor[(S, K), "bf16"], + ): + with Mesh(("cta",), layout=(1,), names=("cta",)) as _cta: + result = tf.zeros(Tensor[(BM, BK), "bf16", (BM, BK), "smem"]) + for m in tile(S, BM): # noqa: F405 + for n in tile(N, BN): # noqa: F405 + for k in tile(K, BK): # noqa: F405 + result = tf.reshard(x[m, k], (BM, BK), "smem") + return result + + +__all__ = ["InvariantReuse"] diff --git a/tests/installed/models/contract.py b/tests/installed/models/contract.py index 9b74260f..72ab58b0 100644 --- a/tests/installed/models/contract.py +++ b/tests/installed/models/contract.py @@ -84,9 +84,7 @@ def analysed( ] if json_output: arguments.append("--json") - done, report_text = _run_with_report( - tf, arguments, suffix=".json" if json_output else ".py" - ) + done, report_text = _run_with_report(tf, arguments, suffix=".json" if json_output else ".py") assert done.returncode == 0, done.stderr return report_text @@ -131,9 +129,9 @@ def _memory_evidence(report: dict) -> str | None: as two. """ record = report["function_records"]["memory"] - placement = record.get("allocation") - if placement is None or placement["solver_status"] not in ("optimal", "feasible"): - return f"allocation is {placement!r}" + solver_status = record.get("solver_status") + if solver_status not in ("optimal", "feasible"): + return f"solver status is {solver_status!r}" gmem = report["totals"]["traffic"]["gmem"] if not gmem.get("read", 0) > 0: return f"reported no gmem read ({gmem!r})" @@ -312,7 +310,9 @@ def move(value): return type(value)(move(item) for item in value) return value - loaded = module.load(DictResource({name: value.to(device) for name, value in weights.items()})) + loaded = module.load( + DictResource({name: value.to(device) for name, value in weights.items()}) + ) got = loaded.forward(*(move(value) for value in activations)) want = move(expected[0] if len(expected) == 1 else tuple(expected)) expect = { diff --git a/tests/installed/smoke_analyze.py b/tests/installed/smoke_analyze.py index a4e924a8..1a71d5e3 100644 --- a/tests/installed/smoke_analyze.py +++ b/tests/installed/smoke_analyze.py @@ -17,7 +17,7 @@ def main(x: Tensor[(8,), "f32"]): return wrong """ -_OPEN_MODULE = ''' +_OPEN_MODULE = """ from tilefoundry import module from tilefoundry.dsl import DimVar, Tensor, Topology, func, tf from tilefoundry.target import CudaTarget @@ -29,7 +29,7 @@ class Open: @func def main(x: Tensor[(N,), "f32"]): return tf.add(x, x) -''' +""" def test_logical_analyses_run(tf, cmine, tmp_path) -> None: @@ -44,7 +44,7 @@ def test_logical_analyses_run(tf, cmine, tmp_path) -> None: assert done.returncode == 0, done.stderr assert done.stdout == "" report = (tmp_path / "logical.py").read_text(encoding="utf-8") - for conclusion in ("# compute-cost flops=", "# peak-footprint=", "# roofline ideal-ns="): + for conclusion in ("# compute-cost flops=", "# memory traffic=", "# roofline ideal-ns="): assert conclusion in report, conclusion @@ -81,7 +81,7 @@ def test_mega_kernel_reports_four_families_on_one_expanded_program(tf, tmp_path) families = ["compute-cost", "memory", "roofline", "performance"] assert payload["requested"] == payload["executed"] == families - assert set(payload["function_records"]) == {*families, "traffic"} + assert set(payload["function_records"]) == set(families) assert len(payload["calls"]) == 7 for row in payload["calls"]: name, line_text = row["value"].rsplit(":", 1) @@ -89,19 +89,21 @@ def test_mega_kernel_reports_four_families_on_one_expanded_program(tf, tmp_path) assert 1 <= line <= len(source_lines) assert f"{name} =" in source_lines[line - 1] assert all( - set(row) - {"performance"} == {"value", "compute-cost", "roofline", "traffic"} + set(row) - {"performance"} == {"value", "compute-cost", "memory", "roofline"} for row in payload["calls"] ) - assert [ - index for index, row in enumerate(payload["calls"]) if "performance" in row - ] == [1, 4, 6] + assert [index for index, row in enumerate(payload["calls"]) if "performance" in row] == [ + 1, + 4, + 6, + ] assert text.startswith( "# analysis target=nvidia.h200_sxm module=MoEMegaKernel function=experts" ) for conclusion in ( "# selection requested=compute-cost,memory,roofline,performance", "# compute-cost flops=f32:", - "# peak-footprint=", + "# memory traffic=", "# roofline ideal-ns=", "# performance root=MoEMegaKernel::experts predicted-ns=", ): @@ -114,7 +116,10 @@ def test_usage_errors_print_usage_before_error(tf) -> None: assert done.stdout == "" assert done.stderr.startswith("usage: tilefoundry analyze") - assert "tilefoundry analyze: error: the following arguments are required: SOURCE, PATH" in done.stderr + assert ( + "tilefoundry analyze: error: the following arguments are required: SOURCE, PATH" + in done.stderr + ) assert "SOURCE" in done.stderr assert "model.py[:Module[.child_module...][.function]]" not in done.stderr @@ -164,9 +169,7 @@ def test_a_bare_analyze_binds_every_open_dimension(tf, tmp_path) -> None: def test_performance_resolves_derived_execution_geometry(tf, derived_prefill, tmp_path) -> None: source = f"{derived_prefill}:DerivedPrefill.prefill" - unbound = tf( - "analyze", source, str(tmp_path / "unbound.json"), "--performance", "--json" - ) + unbound = tf("analyze", source, str(tmp_path / "unbound.json"), "--performance", "--json") assert unbound.returncode == 1 assert unbound.stdout == "" assert "prefill_n is declared as [1, 65)" in unbound.stderr @@ -204,23 +207,18 @@ def test_analyze_reports_only_the_analyses_that_were_requested(tf, cwide, tmp_pa assert done.stdout == "" report = (tmp_path / "report.py").read_text(encoding="utf-8") - assert ( - "# selection requested=roofline executed=compute-cost,memory,roofline" - in report - ) + assert "# selection requested=roofline executed=compute-cost,memory,roofline" in report assert "# compute-cost flops=" in report assert "# roofline ideal-ns=" in report - assert "# peak-footprint=" not in report + assert "# memory " not in report assert "# performance " not in report assert "; roofline ideal-ns=" in report - assert "; memory peak=" not in report + assert "; memory" not in report assert "; compute-cost" not in report assert "; performance=" not in report -def test_analyze_failure_reports_line_variable_and_reason( - tf, dynamic_trip_count, tmp_path -) -> None: +def test_analyze_failure_reports_line_variable_and_reason(tf, dynamic_trip_count, tmp_path) -> None: bad = tmp_path / "bad.py" bad.write_text(_BAD_MODULE, encoding="utf-8") @@ -233,9 +231,7 @@ def test_analyze_failure_reports_line_variable_and_reason( assert "dtype mismatch" in done.stderr report = tmp_path / "dynamic_report.py" - refused = tf( - "analyze", f"{dynamic_trip_count}:DynamicTripCount", str(report), "--memory" - ) + refused = tf("analyze", f"{dynamic_trip_count}:DynamicTripCount", str(report), "--memory") assert refused.returncode == 1 assert refused.stdout == "" assert "loop '_step'" in refused.stderr diff --git a/tests/installed/smoke_target/smoke_v100.py b/tests/installed/smoke_target/smoke_v100.py index 5b9f1c3c..cf382d59 100644 --- a/tests/installed/smoke_target/smoke_v100.py +++ b/tests/installed/smoke_target/smoke_v100.py @@ -25,10 +25,10 @@ def _v100_qwen(tf, tmp_path: Path) -> Path: ) .replace( 'CudaTarget("nvidia.h200_sxm")', - 'CudaTarget(\n' + "CudaTarget(\n" ' Path(__file__).parent / "hw" / "vendor_v100_sxm2_32gb.toml",\n' ' Path(__file__).parent / "hw" / "vendor_sm70.toml",\n' - ')', + ")", ), encoding="utf-8", ) @@ -39,9 +39,7 @@ def _v100_qwen(tf, tmp_path: Path) -> Path: return model -def test_external_v100_documents_analyse_a_copied_installed_model( - tf, tmp_path -) -> None: +def test_external_v100_documents_analyse_a_copied_installed_model(tf, tmp_path) -> None: model = _v100_qwen(tf, tmp_path) done = tf( "analyze", @@ -62,13 +60,14 @@ def test_external_v100_documents_analyse_a_copied_installed_model( assert report["totals"]["flops"]["f16"] > 0 assert report["function_records"]["roofline"]["ideal_ns"] > 0 gmem = next( - item for item in report["function_records"]["memory"]["footprint"] + item + for item in report["function_records"]["memory"]["peaks"] if item["memory_level"] == "gmem" ) assert gmem["peak_bytes"] < 32_000_000_000 record = report["function_records"]["performance"] assert record["waves"] == 2 - assert report["function_records"]["memory"]["allocation"]["solver_status"] in ( + assert report["function_records"]["memory"]["solver_status"] in ( "optimal", "feasible", ) @@ -76,14 +75,10 @@ def test_external_v100_documents_analyse_a_copied_installed_model( call["performance"]["timeline"] for call in report["calls"] if "performance" in call ] assert call_records - assert all( - set(call) == {"start_ns", "end_ns", "trips", "stride_ns"} - for call in call_records - ) + assert all(set(call) == {"start_ns", "end_ns", "trips", "stride_ns"} for call in call_records) assert all(call["end_ns"] > call["start_ns"] for call in call_records) per_wave_end = max( - call["end_ns"] + (call["trips"] - 1) * call["stride_ns"] - for call in call_records + call["end_ns"] + (call["trips"] - 1) * call["stride_ns"] for call in call_records ) per_wave_end = max(per_wave_end, report["function_records"]["roofline"]["ideal_ns"]) assert record["timeline"] == { diff --git a/tests/ops/ir/test_cache_update.py b/tests/ops/ir/test_cache_update.py index 6c36c768..4362b4a7 100644 --- a/tests/ops/ir/test_cache_update.py +++ b/tests/ops/ir/test_cache_update.py @@ -21,7 +21,7 @@ run_typeinfer_case, ) from tilefoundry import func, module -from tilefoundry.analysis import ComputeCostMetadata, TrafficMetadata +from tilefoundry.analysis import ComputeCostMetadata, MemoryMetadata from tilefoundry.analysis.api import analyze from tilefoundry.dsl import Mesh, Tensor, tf from tilefoundry.evaluator import evaluate @@ -228,7 +228,7 @@ def test_cache_update_function_analyzes_program_and_cta_cost() -> None: if isinstance(expr, Call) and isinstance(expr.target, CacheUpdate) ) record = get_metadata(analysed_update, ComputeCostMetadata) - moved = get_metadata(analysed_update, TrafficMetadata) + moved = get_metadata(analysed_update, MemoryMetadata) assert result.topology_level == "cta" assert record is not None assert record.flops.kinds == () diff --git a/tests/ops/ir/test_local.py b/tests/ops/ir/test_local.py index ce2347ea..847e045d 100644 --- a/tests/ops/ir/test_local.py +++ b/tests/ops/ir/test_local.py @@ -3,7 +3,7 @@ from __future__ import annotations from tilefoundry import func, module -from tilefoundry.analysis import ComputeCostMetadata, TrafficMetadata +from tilefoundry.analysis import ComputeCostMetadata, MemoryMetadata from tilefoundry.analysis.api import analyze from tilefoundry.dsl import Mesh, Tensor, Topology, tf from tilefoundry.ir.core import Call, get_metadata @@ -50,9 +50,9 @@ def test_local_analyzes_as_a_zero_traffic_topology_view() -> None: if isinstance(expr, Call) and isinstance(expr.target, Local) ) record = get_metadata(analysed_local, ComputeCostMetadata) - moved = get_metadata(analysed_local, TrafficMetadata) + moved = get_metadata(analysed_local, MemoryMetadata) assert result.topology_level == topology_level assert record is not None assert record.flops.kinds == () - assert moved.storage.kinds == () + assert moved.traffic.storage.kinds == () assert moved.operands == (TrafficBytes(), TrafficBytes()) From 79a1aa96f88b4c82b1dd33282dfdbda336863f76 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 15:03:26 +0800 Subject: [PATCH 2/2] test(analysis): focus invariant reuse regression --- tests/analysis/test_analyze_by_hand.py | 181 +------------------------ 1 file changed, 4 insertions(+), 177 deletions(-) diff --git a/tests/analysis/test_analyze_by_hand.py b/tests/analysis/test_analyze_by_hand.py index 7d10926f..4985a19c 100644 --- a/tests/analysis/test_analyze_by_hand.py +++ b/tests/analysis/test_analyze_by_hand.py @@ -7,182 +7,6 @@ from tilefoundry.analysis.report import report_data -def _checked(data: dict) -> dict: - """Keep the movement and placement conclusions this fixture hand-checks.""" - function = data["function_records"]["memory"] - return { - "loops": data["loops"], - "calls": data["calls"], - "function_memory": { - "topologies": function["topologies"], - "traffic": function["traffic"], - "footprint": function["footprint"], - "peaks": function["peaks"], - "solver_status": function["solver_status"], - "errors": function["errors"], - "advisories": function["advisories"], - }, - } - - -INVARIANT_REUSE = { - "loops": [], - "calls": [ - { - "value": "v0", - "memory": { - "topologies": ["cta"], - "traffic": { - "storage": { - "smem": { - "logical": {"read": 0, "write": 16}, - "total": {"read": 0, "write": 16}, - "per_unit": [{"read": 0, "write": 16}], - } - }, - "communication": {}, - }, - "operands": [ - { - "arg": "result", - "name": "v0", - "type": "bf16[4,2] smem", - "read": 0, - "write": 16, - } - ], - "footprint": None, - }, - }, - { - "value": "v1", - "memory": { - "topologies": ["cta"], - "traffic": { - "storage": { - "rmem": { - "logical": {"read": 16, "write": 0}, - "total": {"read": 16, "write": 0}, - "per_unit": [{"read": 16, "write": 0}], - } - }, - "communication": {}, - }, - "operands": [ - { - "arg": 0, - "name": "x", - "type": "bf16[8,4] gmem", - "read": 0, - "write": 0, - }, - { - "arg": 1, - "name": "tuple", - "type": "i64[]+ umat", - "read": 16, - "write": 0, - }, - { - "arg": "result", - "name": "v1", - "type": "bf16[4,2] gmem", - "read": 0, - "write": 0, - }, - ], - "footprint": None, - }, - }, - { - "value": "v2", - "memory": { - "topologies": ["cta"], - "traffic": { - "storage": { - "gmem": { - "logical": {"read": 16, "write": 0}, - "total": {"read": 16, "write": 0}, - "per_unit": [{"read": 16, "write": 0}], - }, - "smem": { - "logical": {"read": 0, "write": 16}, - "total": {"read": 0, "write": 16}, - "per_unit": [{"read": 0, "write": 16}], - }, - }, - "communication": {}, - }, - "operands": [ - { - "arg": 0, - "name": "v1", - "type": "bf16[4,2] gmem", - "read": 16, - "write": 0, - }, - { - "arg": "result", - "name": "v2", - "type": "bf16[4,2] smem", - "read": 0, - "write": 16, - }, - ], - "footprint": None, - }, - }, - ], - "function_memory": { - "topologies": ["cta"], - "traffic": { - "storage": { - "gmem": { - "logical": {"read": 64, "write": 0}, - "total": {"read": 192, "write": 0}, - "per_unit": [{"read": 192, "write": 0}], - }, - "rmem": { - "logical": {"read": 64, "write": 0}, - "total": {"read": 192, "write": 0}, - "per_unit": [{"read": 192, "write": 0}], - }, - "smem": { - "logical": {"read": 0, "write": 80}, - "total": {"read": 0, "write": 208}, - "per_unit": [{"read": 0, "write": 208}], - }, - }, - "communication": {}, - }, - "footprint": None, - "peaks": [ - { - "memory_level": "gmem", - "peak_bytes": 80, - "persistent_bytes": 64, - "capacity_bytes": 141_000_000_000, - }, - { - "memory_level": "rmem", - "peak_bytes": 0, - "persistent_bytes": 0, - "capacity_bytes": 262_144, - }, - { - "memory_level": "smem", - "peak_bytes": 64, - "persistent_bytes": 0, - "capacity_bytes": 232_448, - }, - ], - "solver_status": "feasible", - "errors": [], - "advisories": [], - }, -} - - def test_invariant_reuse_matches_the_written_arithmetic() -> None: result = analyze( InvariantReuse, @@ -198,4 +22,7 @@ def test_invariant_reuse_matches_the_written_arithmetic() -> None: metadata_types=result.metadata_types, ) - assert _checked(data) == INVARIANT_REUSE + traffic = data["function_records"]["memory"]["traffic"]["storage"]["gmem"] + assert traffic["logical"] == {"read": 64, "write": 0} + assert traffic["total"] == {"read": 192, "write": 0} + assert traffic["per_unit"] == [{"read": 192, "write": 0}]