From 1ff2b6dca24f145e8e71e02ffa3393a0c881ae2b Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 23:37:02 +0800 Subject: [PATCH 1/7] feat(analysis): preserve logical work alongside projected cost --- docs/spec/analysis.md | 4 +++ src/tilefoundry/analysis/compute_cost.py | 36 ++++++++++++++++++++++-- src/tilefoundry/analysis/metadata.py | 2 ++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 60f48865f..0e1fd9a1e 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -171,12 +171,16 @@ class ComputeCostMetadata(IRMetadata): Attributes: topologies: attribute; The declared levels, in the order per_unit states them. flops: attribute; Flop counts per compute DType name. + flops_logical: attribute; Flop counts before execution-position replication. service: attribute; Result counts per service kind. + service_logical: attribute; Service counts before execution-position replication. """ topologies: tuple[str, ...] = () flops: Breakdown[int] = Breakdown() + flops_logical: tuple[tuple[str, int], ...] = () service: Breakdown[int] = Breakdown() + service_logical: tuple[tuple[str, int], ...] = () ``` | Field | How it is computed | Reads the target | diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index bd4f73609..530c4112e 100644 --- a/src/tilefoundry/analysis/compute_cost.py +++ b/src/tilefoundry/analysis/compute_cost.py @@ -156,6 +156,7 @@ def _call_cost_record( expr: Call, locals_by_unit: "dict[str, CostContext]", positions_by_unit: "dict[str, int]", + whole: CostContext, asked: "str | None" = None, ) -> ComputeCostMetadata: """Measure the work one Call asks for, without attaching the record. @@ -169,6 +170,10 @@ def _call_cost_record( service_by_unit: list[dict[str, int]] = [] whole_flops: dict[str, int] = {} whole_service: dict[str, int] = {} + try: + whole_cost = CostEvaluator().visit(expr, whole) + except (ValueError, VerifyError) as error: + raise AnalysisError(str(error)) from None for unit, local in locals_by_unit.items(): try: cost = CostEvaluator().visit(expr, local) @@ -183,7 +188,9 @@ def _call_cost_record( return ComputeCostMetadata( topologies=tuple(locals_by_unit), flops=breakdown(whole_flops, flops_by_unit, 0), + flops_logical=tuple(sorted(_named(whole_cost.flops).items())), service=breakdown(whole_service, service_by_unit, 0), + service_logical=tuple(sorted(whole_cost.service.items())), ) @@ -221,15 +228,21 @@ def _bytes( def _accumulate( flops: dict[str, int], + flops_logical: dict[str, int], service: dict[str, int], + service_logical: dict[str, int], by_unit: "dict[str, dict[str, dict[str, int]]]", record: ComputeCostMetadata, trips: int, ) -> None: for kind, spread in record.flops.kinds: flops[kind] = flops.get(kind, 0) + spread.total * trips + for kind, value in record.flops_logical: + flops_logical[kind] = flops_logical.get(kind, 0) + value * trips for kind, spread in record.service.kinds: service[kind] = service.get(kind, 0) + spread.total * trips + for kind, value in record.service_logical: + service_logical[kind] = service_logical.get(kind, 0) + value * trips for index, unit in enumerate(record.topologies): held = by_unit.setdefault(unit, {"flops": {}, "service": {}}) for kind, spread in record.flops.kinds: @@ -248,10 +261,13 @@ class ComputeCostContext(AnalyzeContext): """ locals_by_unit: dict[str, CostContext] = field(default_factory=dict) + whole: CostContext | None = None current_mesh: Mesh | None = None executing_positions: dict[str, int] = field(default_factory=dict) flops: dict[str, int] = field(default_factory=dict) + flops_logical: dict[str, int] = field(default_factory=dict) service: dict[str, int] = field(default_factory=dict) + service_logical: dict[str, int] = field(default_factory=dict) by_unit: dict[str, dict[str, dict[str, int]]] = field(default_factory=dict) call_count: list[int] = field(default_factory=lambda: [0]) @@ -291,7 +307,11 @@ def default_visit_leaf( if not ctx.locals_by_unit: raise AnalysisError("compute-cost: visitor context is missing its cost context") record = _call_cost_record( - expr, ctx.locals_by_unit, ctx.executing_positions, ctx.topology_level + expr, + ctx.locals_by_unit, + ctx.executing_positions, + ctx.whole, + ctx.topology_level, ) attach(expr, record) owner = ctx.current if id(expr) in ctx.current.accesses["narrow"] else ctx.root @@ -301,7 +321,15 @@ def default_visit_leaf( if cursor.is_variant(expr): repeats *= max(1, cursor.trips()) cursor = cursor.parent - _accumulate(ctx.flops, ctx.service, ctx.by_unit, record, repeats) + _accumulate( + ctx.flops, + ctx.flops_logical, + ctx.service, + ctx.service_logical, + ctx.by_unit, + record, + repeats, + ) def analyze_compute_cost( @@ -312,6 +340,7 @@ def analyze_compute_cost( module, topology_level = context.module, context.topology_level topologies = module.effective_topologies() scope = FunctionScope(module, function) + whole = CostContext(scope=scope) units = tuple(topology.name for topology in topologies) or (topology_level,) locals_by_unit = { unit: CostContext(scope=scope, topology_level=unit, topologies=topologies) @@ -326,6 +355,7 @@ def analyze_compute_cost( root=context.root, current=context.current, locals_by_unit=locals_by_unit, + whole=whole, executing_positions=dict.fromkeys(locals_by_unit, 1), ) ComputeCostVisitor().visit(function.body, cost_context) @@ -339,11 +369,13 @@ def analyze_compute_cost( [dict(held["flops"]) for held in cost_context.by_unit.values()], 0, ), + flops_logical=tuple(sorted(cost_context.flops_logical.items())), service=breakdown( dict(cost_context.service), [dict(held["service"]) for held in cost_context.by_unit.values()], 0, ), + service_logical=tuple(sorted(cost_context.service_logical.items())), ), ) diff --git a/src/tilefoundry/analysis/metadata.py b/src/tilefoundry/analysis/metadata.py index ed87f58f2..aefa1b4da 100644 --- a/src/tilefoundry/analysis/metadata.py +++ b/src/tilefoundry/analysis/metadata.py @@ -108,7 +108,9 @@ class ComputeCostMetadata(IRMetadata): topologies: tuple[str, ...] = () flops: Breakdown[int] = Breakdown() + flops_logical: tuple[tuple[str, int], ...] = () service: Breakdown[int] = Breakdown() + service_logical: tuple[tuple[str, int], ...] = () @dataclass(frozen=True) From c434c39ab0d529391e4f27c137241679ff93752c Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 15 Sep 2026 10:36:45 +0800 Subject: [PATCH 2/7] refactor: align follow-up with updated codegen baseline --- src/tilefoundry/analysis/compute_cost.py | 254 +++----- src/tilefoundry/analysis/facts.py | 2 +- src/tilefoundry/analysis/metadata.py | 172 ++---- src/tilefoundry/analysis/roofline.py | 2 +- src/tilefoundry/cli/analyze.py | 17 +- src/tilefoundry/codegen/cpu/module.py | 516 +++++++++++----- src/tilefoundry/codegen/cuda/module.py | 243 +++----- .../codegen/cuda/tir/prim_function.py | 294 ++++++--- .../codegen/cuda/tir/stmts/mesh_scope.py | 83 +-- src/tilefoundry/codegen/emitter.py | 158 +++++ src/tilefoundry/codegen/topology.py | 112 ---- src/tilefoundry/inspection/values.py | 580 +++++++----------- .../passes/transforms/host_entry.py | 216 ++----- src/tilefoundry/visitor_registry/contexts.py | 36 +- src/tilefoundry/visitor_registry/op_cost.py | 25 +- tests/analysis/test_analysis_families.py | 133 +--- tests/analysis/test_analyze_at_a_size.py | 63 +- tests/cli/test_cli.py | 39 +- tests/ops/ir/cost_utils.py | 4 +- tests/ops/ir/test_binary.py | 2 +- tests/ops/ir/test_unary.py | 2 +- tests/ops/ir/test_where.py | 2 +- 22 files changed, 1361 insertions(+), 1594 deletions(-) create mode 100644 src/tilefoundry/codegen/emitter.py delete mode 100644 src/tilefoundry/codegen/topology.py diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index 530c4112e..6a7787806 100644 --- a/src/tilefoundry/analysis/compute_cost.py +++ b/src/tilefoundry/analysis/compute_cost.py @@ -1,7 +1,7 @@ """How much work the authored program asks for. -This family reads the program and nothing else. Flops and typed service come -from each op's registered cost evaluator, so the record it leaves is the same on +This family reads the program and nothing else. Flops and typed machine ops +come from each op's registered cost evaluator, so the record it leaves is the same on every backend. What that work moves is the memory family's half of the same declaration, and what it costs in time is a separate question again, asked against a target's rates. @@ -17,19 +17,18 @@ from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.types import DType -from tilefoundry.ir.types.shard import Mesh, composed, topology_axes +from tilefoundry.ir.types.shard import Mesh, composed, level_axes from tilefoundry.ir.types.shard.mesh import _positions_layout from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.visitor_registry.contexts import ( CostContext, FunctionScope, - TrafficBytes, ) from tilefoundry.visitor_registry.visitors import CostEvaluator from .errors import AnalysisError from .facts import PerformanceServiceFacts, ThroughputFacts -from .metadata import Breakdown, ComputeCostMetadata, breakdown, shares +from .metadata import ComputeCostMetadata from .visitor import AnalyzeContext SELECTOR = "compute-cost" @@ -39,12 +38,11 @@ def _is_structural_occurrence( cost: ComputeCostMetadata, moved: "TrafficMetadata | None" = None, *, - unit: str, bandwidth_level: str | None = None, ) -> bool: """Whether an occurrence asks for nothing this model puts on a clock. - Only what could take time is counted: the flops, the typed service, and the + Only what could take time is counted: the flops, the typed machine ops, and the bytes at the one level a bandwidth is published for. Movement at any other level is still movement and still recorded -- what it is not is work this model can lay on a timeline, so it neither earns a duration nor asks for a @@ -52,10 +50,10 @@ def _is_structural_occurrence( different questions, and this is the second one. """ return ( - all(not value for _name, value in shares(cost.flops, cost.topologies, unit).items()) - and all(not value for _kind, value in shares(cost.service, cost.topologies, unit).items()) + all(not value for _name, value in cost.flops_per_unit) + and all(not value for _kind, value in cost.ops_per_unit) and not ( - _bytes(moved.storage, moved.topologies, bandwidth_level, unit) + moved.per_unit_at(bandwidth_level).total_bytes if moved is not None and bandwidth_level is not None else 0 ) @@ -68,32 +66,29 @@ def _local_duration_ns( services: PerformanceServiceFacts, *, moved: "TrafficMetadata | None" = None, - topology_level: str, + level: str, scale: int = 1, ) -> int: """Price one occurrence's projected work against one unit's throughputs. - Compute, movement and what leaves the unit overlap, so the duration is - whichever takes longest; the same bytes can spend two resources and still - take one span of time. Work with no stated throughput is refused rather - than priced at nothing. Movement at a level the target publishes no - bandwidth for is a different case -- stated and left untimed, because a - rate nobody published is not one this may invent. + Compute and movement overlap within one occurrence, so its duration is + whichever side takes longer. Work with no stated throughput is refused + rather than priced at nothing: a number with a hole in it reads as a program + that does less than it does. Movement at a level the target publishes no + bandwidth for is a different case -- it is stated and left untimed, because + a rate nobody published is not one this may invent. """ - if services.unit != topology_level: + if services.unit != level: raise AnalysisError( - f"performance: selected topology level {topology_level!r}, but the " - f"target's one-unit throughputs are stated for {services.unit!r}" + f"performance: selected topology level {level!r}, but the target's " + f"one-unit throughputs are stated for {services.unit!r}" ) - if _is_structural_occurrence( - cost, moved, unit=topology_level, bandwidth_level=facts.bandwidth_level - ): + if _is_structural_occurrence(cost, moved, bandwidth_level=facts.bandwidth_level): return 0 compute_ns = 0 - for name, share in shares(cost.flops, cost.topologies, topology_level).items(): - value = share + for name, value in cost.flops_per_unit: if not value: continue dtype = getattr(DType, name, None) @@ -102,27 +97,22 @@ def _local_duration_ns( throughput = services.flops(dtype) if throughput is None or throughput <= 0: raise AnalysisError( - f"performance: target states no one-unit throughput for dtype " - f"{name!r} at {topology_level!r}" + f"performance: target states no one-unit throughput for dtype {name!r} at {level!r}" ) compute_ns += -(-(value * scale * 1_000_000_000) // throughput) - for kind, share in shares(cost.service, cost.topologies, topology_level).items(): - value = share + for kind, value in cost.ops_per_unit: if not value: continue throughput = services.ops(kind) if throughput is None or throughput <= 0: raise AnalysisError( - f"performance: target states no one-unit throughput for {kind!r} " - f"work at {topology_level!r}" + f"performance: target states no one-unit throughput for {kind!r} work at {level!r}" ) compute_ns += -(-(value * scale * 1_000_000_000) // throughput) crossed = ( - _bytes(moved.storage, moved.topologies, facts.bandwidth_level, topology_level) * scale - if moved is not None - else 0 + moved.per_unit_at(facts.bandwidth_level).total_bytes * scale if moved is not None else 0 ) memory_ns = 0 if crossed: @@ -130,79 +120,56 @@ def _local_duration_ns( if throughput is None or throughput <= 0: raise AnalysisError( f"performance: target states no one-unit throughput for level " - f"{facts.bandwidth_level!r} at {topology_level!r}" + f"{facts.bandwidth_level!r} at {level!r}" ) memory_ns = -(-(crossed * 1_000_000_000) // throughput) - - sent = ( - _bytes(moved.communication, moved.topologies, topology_level, topology_level) * scale - if moved is not None - else 0 - ) - link_ns = 0 - if sent: - rate = services.bandwidth(topology_level) - if rate: - link_ns = -(-(sent * 1_000_000_000) // rate) - return max(compute_ns, memory_ns, link_ns) + return max(compute_ns, memory_ns) -def _named(flops: dict) -> dict[str, int]: - """The flop counts keyed by dtype name, as a record states them.""" - return {dtype.name: value for dtype, value in flops.items()} +def _flops(flops: dict) -> tuple[tuple[str, int], ...]: + return tuple(sorted((dtype.name, value) for dtype, value in flops.items())) def _call_cost_record( - expr: Call, - locals_by_unit: "dict[str, CostContext]", - positions_by_unit: "dict[str, int]", - whole: CostContext, - asked: "str | None" = None, + expr: Call, whole: CostContext, local: CostContext, executing_positions: int ) -> ComputeCostMetadata: """Measure the work one Call asks for, without attaching the record. - Work only: what an occurrence moves is the memory family's answer. One - unit's work is measured at every declared level rather than at one chosen - for the reader, and global work is any of them repeated over the positions - executing this scope -- the product is the same whichever level states it. + Work only: what an occurrence moves is the memory family's answer. Global + work is one selected unit's work repeated over the positions executing this + scope, so it is derived from the same registered evaluator rather than + independently recomputed from the authored type. """ - flops_by_unit: list[dict[str, int]] = [] - service_by_unit: list[dict[str, int]] = [] - whole_flops: dict[str, int] = {} - whole_service: dict[str, int] = {} try: whole_cost = CostEvaluator().visit(expr, whole) + local_cost = CostEvaluator().visit(expr, local) except (ValueError, VerifyError) as error: raise AnalysisError(str(error)) from None - for unit, local in locals_by_unit.items(): - try: - cost = CostEvaluator().visit(expr, local) - except (ValueError, VerifyError) as error: - raise AnalysisError(str(error)) from None - flops_by_unit.append(_named(cost.flops)) - service_by_unit.append(dict(cost.service)) - if unit == (asked or next(iter(locals_by_unit))): - repeats = positions_by_unit[unit] - whole_flops = {name: value * repeats for name, value in _named(cost.flops).items()} - whole_service = {kind: value * repeats for kind, value in cost.service.items()} return ComputeCostMetadata( - topologies=tuple(locals_by_unit), - flops=breakdown(whole_flops, flops_by_unit, 0), - flops_logical=tuple(sorted(_named(whole_cost.flops).items())), - service=breakdown(whole_service, service_by_unit, 0), - service_logical=tuple(sorted(whole_cost.service.items())), + flops=_flops( + {dtype: value * executing_positions for dtype, value in local_cost.flops.items()} + ), + flops_logical=sum(whole_cost.flops.values()), + flops_per_unit=_flops(local_cost.flops), + ops=tuple( + sorted( + (kind, value * executing_positions) for kind, value in local_cost.ops.items() + ) + ), + ops_logical=tuple(sorted(whole_cost.ops.items())), + ops_per_unit=tuple(sorted(local_cost.ops.items())), ) -def _scope_position_count(mesh: Mesh, topology_level: str | None, topologies: tuple) -> int: +def _scope_position_count(mesh: Mesh, level: str | None, topologies: tuple) -> int: """Count positions at or above the selected level within *mesh*.""" - if topology_level is None: + if level is None: return 1 declared = {topology.name: index for index, topology in enumerate(topologies)} - selected = declared[topology_level] + selected = declared[level] shape, _strides, _offset = _positions_layout(mesh) positions = 1 - for topology, axes in zip(mesh.topologies, topology_axes(mesh)): + for topology, axes in zip(mesh.topologies, level_axes(mesh)): if declared[topology.name] > selected: continue for axis in axes: @@ -215,60 +182,48 @@ def _scope_position_count(mesh: Mesh, topology_level: str | None, topologies: tu return positions -def _bytes( - held: "Breakdown[TrafficBytes]", - topologies: tuple[str, ...], - kind: str, - topology_level: "str | None", -) -> int: - """One kind's bytes for one unit of *topology_level*, read and written together.""" - moved = shares(held, topologies, topology_level).get(kind) - return moved.total_bytes if moved is not None else 0 - - def _accumulate( flops: dict[str, int], - flops_logical: dict[str, int], - service: dict[str, int], - service_logical: dict[str, int], - by_unit: "dict[str, dict[str, dict[str, int]]]", + flops_logical: list[int], + flops_per_unit: dict[str, int], + ops: dict[str, int], + ops_logical: dict[str, int], + ops_per_unit: dict[str, int], record: ComputeCostMetadata, trips: int, ) -> None: - for kind, spread in record.flops.kinds: - flops[kind] = flops.get(kind, 0) + spread.total * trips - for kind, value in record.flops_logical: - flops_logical[kind] = flops_logical.get(kind, 0) + value * trips - for kind, spread in record.service.kinds: - service[kind] = service.get(kind, 0) + spread.total * trips - for kind, value in record.service_logical: - service_logical[kind] = service_logical.get(kind, 0) + value * trips - for index, unit in enumerate(record.topologies): - held = by_unit.setdefault(unit, {"flops": {}, "service": {}}) - for kind, spread in record.flops.kinds: - held["flops"][kind] = held["flops"].get(kind, 0) + spread.at(index) * trips - for kind, spread in record.service.kinds: - held["service"][kind] = held["service"].get(kind, 0) + spread.at(index) * trips + for name, value in record.flops: + flops[name] = flops.get(name, 0) + value * trips + flops_logical[0] += record.flops_logical * trips + for name, value in record.flops_per_unit: + flops_per_unit[name] = flops_per_unit.get(name, 0) + value * trips + for name, value in record.ops: + ops[name] = ops.get(name, 0) + value * trips + for name, value in record.ops_logical: + ops_logical[name] = ops_logical.get(name, 0) + value * trips + for name, value in record.ops_per_unit: + ops_per_unit[name] = ops_per_unit.get(name, 0) + value * trips @dataclass class ComputeCostContext(AnalyzeContext): """State carried through the compute-cost expression walk. - ``executing_positions`` is how many positions of each declared level the - enclosing scope holds; it is the multiplier that turns one unit's work at - that level into total replicated work. + ``executing_positions`` is the number of positions in the enclosing scope + through the selected topology level; it is the D68 multiplier that turns + per-unit work into total replicated work. """ - locals_by_unit: dict[str, CostContext] = field(default_factory=dict) + local: CostContext | None = None whole: CostContext | None = None current_mesh: Mesh | None = None - executing_positions: dict[str, int] = field(default_factory=dict) + executing_positions: int = 1 flops: dict[str, int] = field(default_factory=dict) - flops_logical: dict[str, int] = field(default_factory=dict) - service: dict[str, int] = field(default_factory=dict) - service_logical: dict[str, int] = field(default_factory=dict) - by_unit: dict[str, dict[str, dict[str, int]]] = field(default_factory=dict) + flops_logical: list[int] = field(default_factory=lambda: [0]) + flops_per_unit: dict[str, int] = field(default_factory=dict) + ops: dict[str, int] = field(default_factory=dict) + ops_logical: dict[str, int] = field(default_factory=dict) + ops_per_unit: dict[str, int] = field(default_factory=dict) call_count: list[int] = field(default_factory=lambda: [0]) @@ -280,10 +235,7 @@ def visit_MeshRegion(self, expr: MeshRegion, ctx: ComputeCostContext) -> None: for arg in expr.args: self.visit(arg, ctx) mesh = composed((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh - topologies = ctx.module.effective_topologies() - positions = { - unit: _scope_position_count(mesh, unit, topologies) for unit in ctx.locals_by_unit - } + positions = _scope_position_count(mesh, ctx.level, ctx.module.effective_topologies()) self.visit( expr.body, replace(ctx, executing_positions=positions, current_mesh=mesh), @@ -304,15 +256,11 @@ def default_visit_leaf( if not isinstance(expr, Call): return ctx.call_count[0] += 1 - if not ctx.locals_by_unit: + if ctx.local is None: raise AnalysisError("compute-cost: visitor context is missing its cost context") - record = _call_cost_record( - expr, - ctx.locals_by_unit, - ctx.executing_positions, - ctx.whole, - ctx.topology_level, - ) + if ctx.whole is None: + raise AnalysisError("compute-cost: visitor context is missing its whole context") + record = _call_cost_record(expr, ctx.whole, ctx.local, ctx.executing_positions) attach(expr, record) owner = ctx.current if id(expr) in ctx.current.accesses["narrow"] else ctx.root repeats = 1 @@ -324,9 +272,10 @@ def default_visit_leaf( _accumulate( ctx.flops, ctx.flops_logical, - ctx.service, - ctx.service_logical, - ctx.by_unit, + ctx.flops_per_unit, + ctx.ops, + ctx.ops_logical, + ctx.ops_per_unit, record, repeats, ) @@ -337,45 +286,32 @@ def analyze_compute_cost( context: AnalyzeContext, ) -> None: """Attach one-trip work per Call and multiplicity-aware totals per Function.""" - module, topology_level = context.module, context.topology_level + module, level = context.module, context.level topologies = module.effective_topologies() scope = FunctionScope(module, function) whole = CostContext(scope=scope) - units = tuple(topology.name for topology in topologies) or (topology_level,) - locals_by_unit = { - unit: CostContext(scope=scope, topology_level=unit, topologies=topologies) - for unit in units - if unit is not None - } + local = CostContext(scope=scope, level=level, topologies=topologies) cost_context = ComputeCostContext( module=module, target=context.target, - topology_level=topology_level, + level=level, options=context.options, root=context.root, current=context.current, - locals_by_unit=locals_by_unit, + local=local, whole=whole, - executing_positions=dict.fromkeys(locals_by_unit, 1), ) ComputeCostVisitor().visit(function.body, cost_context) if cost_context.call_count[0] > 0: attach( function, ComputeCostMetadata( - topologies=tuple(cost_context.by_unit), - flops=breakdown( - dict(cost_context.flops), - [dict(held["flops"]) for held in cost_context.by_unit.values()], - 0, - ), - flops_logical=tuple(sorted(cost_context.flops_logical.items())), - service=breakdown( - dict(cost_context.service), - [dict(held["service"]) for held in cost_context.by_unit.values()], - 0, - ), - service_logical=tuple(sorted(cost_context.service_logical.items())), + flops=tuple(sorted(cost_context.flops.items())), + flops_logical=cost_context.flops_logical[0], + flops_per_unit=tuple(sorted(cost_context.flops_per_unit.items())), + ops=tuple(sorted(cost_context.ops.items())), + ops_logical=tuple(sorted(cost_context.ops_logical.items())), + ops_per_unit=tuple(sorted(cost_context.ops_per_unit.items())), ), ) diff --git a/src/tilefoundry/analysis/facts.py b/src/tilefoundry/analysis/facts.py index 6967cf6ec..1129fce17 100644 --- a/src/tilefoundry/analysis/facts.py +++ b/src/tilefoundry/analysis/facts.py @@ -147,7 +147,7 @@ class PerformanceServiceFacts: """Everything one unit gets through, by the kind of work it is asked for. Floating-point work by dtype, movement by the level it crosses, and the work - that is not floating point at all by the service it asks for. A comparison, + that is not floating point at all by the operation kind it asks for. A comparison, a select, an integer add and a local move all take a machine time, and none of them is a FLOP; a dtype is not a kind of work. """ diff --git a/src/tilefoundry/analysis/metadata.py b/src/tilefoundry/analysis/metadata.py index aefa1b4da..0254a63af 100644 --- a/src/tilefoundry/analysis/metadata.py +++ b/src/tilefoundry/analysis/metadata.py @@ -8,139 +8,69 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence from dataclasses import dataclass from tilefoundry.ir.core.metadata import IRMetadata from tilefoundry.visitor_registry.contexts import TrafficBytes -@dataclass(frozen=True) -class Spread[V]: - """One quantity, whole and as one unit of each topology level holds it. - - ``per_unit`` runs in the order the record's ``topologies`` states, one - entry per level, so a level's name is written once for the whole record - rather than once per quantity. A finer level's unit sits inside a coarser - one, so its share is the coarser share divided again by whatever the mesh - splits between them -- which is why every level is stated rather than only - the one an analysis was asked about. - """ - - total: V - per_unit: tuple[V, ...] = () - - def at(self, index: int) -> V: - """One level's share by position, or the total when there is none.""" - return self.per_unit[index] if index < len(self.per_unit) else self.total - - -@dataclass(frozen=True) -class Breakdown[V]: - """One category's quantities, split by the kind of thing each one is. - - The kind is what the rate pricing it is stated for: a dtype prices flops, - a service kind prices what is not floating point, a memory level prices - bytes. Two kinds are never summed, because two rates cannot be. - """ - - kinds: tuple[tuple[str, Spread[V]], ...] = () - - def of(self, kind: str) -> Spread[V] | None: - """This kind's quantity, or ``None`` when the record states none.""" - return next((value for name, value in self.kinds if name == kind), None) - - def names(self) -> tuple[str, ...]: - """Every kind this record states, in the order it states them.""" - return tuple(name for name, _ in self.kinds) - - -def breakdown[V]( - total: "Mapping[str, V]", per_unit: "Sequence[Mapping[str, V]]", zero: V -) -> Breakdown[V]: - """Gather one category's kinds, each with its total and every level's share. - - A kind any of them states appears in all of them, at *zero* where it was - not stated, so one kind's total and its shares stay one row. - """ - kinds = sorted({*total, *(kind for share in per_unit for kind in share)}) - return Breakdown( - tuple( - ( - kind, - Spread( - total.get(kind, zero), - tuple(share.get(kind, zero) for share in per_unit), - ), - ) - for kind in kinds - ) - ) - - -def shares[V]( - held: Breakdown[V], topologies: tuple[str, ...], topology_level: "str | None" = None -) -> dict[str, V]: - """Each kind's value for one unit of *topology_level*, or its total without one. - - The only place a level's name is turned back into a position, because - ``topologies`` is where the names are written and a ``Spread`` states its - shares in that order and carries none of its own. A level the record does - not state reads as the total, which is what a record over one unit says. - """ - index = topologies.index(topology_level) if topology_level in topologies else None - return { - kind: spread.total if index is None else spread.at(index) for kind, spread in held.kinds - } - - @dataclass(frozen=True) class ComputeCostMetadata(IRMetadata): """Record one occurrence's work, or one Function's total work. - ``service`` counts what is not floating point -- comparing, selecting, - whole-number arithmetic -- by the service it asks for, because a predicate - priced as a FLOP is a number about a pipe the work never went down. What an - occurrence moves is a separate record, kept by the family that knows where - values live. On a Call these state one occurrence; on a Function, loops - contribute their trip count. + ``flops`` and ``ops`` state global work; their ``*_per_unit`` partners + apply shard projection at the requested topology level. ``ops`` counts + non-floating-point machine operations by kind -- comparing, selecting, + whole-number arithmetic and special functions. What an occurrence moves is a separate + record, kept by the family that knows where values live. On a Call these + state one occurrence; on a Function, loops contribute their trip count. """ - topologies: tuple[str, ...] = () - flops: Breakdown[int] = Breakdown() - flops_logical: tuple[tuple[str, int], ...] = () - service: Breakdown[int] = Breakdown() - service_logical: tuple[tuple[str, int], ...] = () + flops: tuple[tuple[str, int], ...] = () + flops_logical: int = 0 + flops_per_unit: tuple[tuple[str, int], ...] = () + ops: tuple[tuple[str, int], ...] = () + ops_logical: tuple[tuple[str, int], ...] = () + ops_per_unit: tuple[tuple[str, int], ...] = () + + def ops_per_unit_of(self, kind: str) -> int: + """One unit's count of *kind*, zero when the call asks for none.""" + return next((value for name, value in self.ops_per_unit if name == kind), 0) @dataclass(frozen=True) class TrafficMetadata(IRMetadata): - """The bytes one occurrence moves, in the two coordinates a move has. - - ``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. + """The bytes one occurrence moves, whole and for one participant. + + Which way a boundary moves is its Op's evaluator's answer and how much is + its relation's; the family that decides where values live attaches the + record, and where they landed never corrects a crossing. ``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() + whole: tuple[tuple[str, TrafficBytes], ...] = () + per_unit: tuple[tuple[str, TrafficBytes], ...] = () operands: tuple[TrafficBytes, ...] = () + def at(self, level: str) -> TrafficBytes: + """Bytes moved at *level*, zero when the occurrence does not touch it.""" + return next((value for name, value in self.whole if name == level), TrafficBytes()) + + def per_unit_at(self, level: str) -> TrafficBytes: + """One unit's bytes at *level*, zero when it does not touch it.""" + return next((value for name, value in self.per_unit if name == level), TrafficBytes()) + @dataclass(frozen=True) -class MemoryLevelFootprint: +class LevelFootprint: """How much of one memory level a function needs at its peak. ``persistent_bytes`` is the part that cannot be reclaimed within the function, so it is the floor the peak can never fall below. """ - memory_level: str + level: str peak_bytes: int persistent_bytes: int capacity_bytes: int | None = None @@ -156,7 +86,7 @@ class BufferFootprint: """Per-position, device-wide, and repeated bytes touched in one buffer.""" buffer: str - memory_level: str + level: str bytes: int device_bytes: int repeated_bytes: int @@ -187,7 +117,7 @@ class ValueLifetime: """ binding: str - memory_level: str + level: str bytes: int defined_at: int last_used_at: int @@ -219,14 +149,14 @@ class MemoryMetadata(IRMetadata): placed one: nothing was decided, so nothing is claimed. """ - footprint: tuple[MemoryLevelFootprint, ...] = () + footprint: tuple[LevelFootprint, ...] = () lifetimes: tuple[ValueLifetime, ...] = () advisories: tuple[str, ...] = () allocation: "AllocationMetadata | None" = None - def memory_level(self, name: str) -> MemoryLevelFootprint | None: + def level(self, name: str) -> LevelFootprint | 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) + return next((item for item in self.footprint if item.level == name), None) @dataclass(frozen=True) @@ -265,6 +195,21 @@ class TimelineMetadata: stride_ns: int = 0 +@dataclass(frozen=True) +class TripInterval: + """A half-open interval, parameterized by the trip index of its loop. + + A single trip occurs once, so ``stride`` says nothing about it and the trip + index does not appear. Repeated trips state the first interval and the + stride every later one is offset by. + """ + + start: int + end: int + stride: int = 0 + trips: int = 1 + + @dataclass(frozen=True) class PerformanceMetadata(IRMetadata): """One occurrence's interval within one local wave of its Function. @@ -293,19 +238,16 @@ class PerformanceSummaryMetadata(IRMetadata): __all__ = [ "AllocationMetadata", - "Breakdown", "BufferFootprint", "ComputeCostMetadata", + "LevelFootprint", "LoopFootprintMetadata", - "MemoryLevelFootprint", "MemoryMetadata", "PerformanceMetadata", "PerformanceSummaryMetadata", "RooflineMetadata", - "Spread", "TimelineMetadata", + "TripInterval", "TrafficBytes", "ValueLifetime", - "breakdown", - "shares", ] diff --git a/src/tilefoundry/analysis/roofline.py b/src/tilefoundry/analysis/roofline.py index 695893c40..244335f8d 100644 --- a/src/tilefoundry/analysis/roofline.py +++ b/src/tilefoundry/analysis/roofline.py @@ -100,7 +100,7 @@ def _cost_bound( Whole-device work against whole-device rates: the flops the target publishes a peak for, and the bytes at the level it publishes a bandwidth for. Typed - service has no whole-device rate to divide by and so does not enter a bound, + typed ops have no whole-device rate to divide by and so do not enter a bound, and neither do bytes at a level with no published bandwidth: what asks for a nanosecond is what this could have priced, so a dtype whose rate is missing still owes one and a level nobody rated does not. diff --git a/src/tilefoundry/cli/analyze.py b/src/tilefoundry/cli/analyze.py index 144ac3b56..0f4cbd809 100644 --- a/src/tilefoundry/cli/analyze.py +++ b/src/tilefoundry/cli/analyze.py @@ -43,7 +43,8 @@ def _watch(limit: float) -> None: def stop() -> None: time.sleep(limit) sys.stderr.write( - f"tilefoundry: error: analysis too complex, timed out after {limit:.0f}s\n" + "tilefoundry: error: analysis too complex, " + f"timed out after {limit:.0f}s\n" ) sys.stderr.flush() os._exit(1) @@ -82,10 +83,10 @@ def guidance() -> str: family what --topology changes pass it when ------------ -------------------------------------- --------------------- - 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 + compute-cost flops_per_unit and ops_per_unit. the program shards + flops and ops stay global + memory per-unit traffic. Footprint follows the program shards + its owner for each storage level 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 @@ -155,7 +156,7 @@ def run_authored_analysis( Path(out_path).write_text(annotated, encoding="utf-8") return 0 - result = analyze(module, function, analysis=analyses, topology_level=topology, dims=dims) + result = analyze(module, function, analysis=analyses, level=topology, dims=dims) rendered = render_analysis(result, operands=operands and not as_json) if as_json: Path(out_path).write_text( @@ -164,7 +165,9 @@ def run_authored_analysis( ) return 0 - Path(out_path).write_text(f"{render_text(rendered)}\n\n{rendered.annotated}", encoding="utf-8") + Path(out_path).write_text( + f"{render_text(rendered)}\n\n{rendered.annotated}", encoding="utf-8" + ) return 0 diff --git a/src/tilefoundry/codegen/cpu/module.py b/src/tilefoundry/codegen/cpu/module.py index 54b36a38b..547ead958 100644 --- a/src/tilefoundry/codegen/cpu/module.py +++ b/src/tilefoundry/codegen/cpu/module.py @@ -1,31 +1,32 @@ -"""The host ``.cpp`` translation unit: one entry, and the launches it makes. +"""Emit the host module for the split code-generation pipeline. -The entry takes runtime tensors, checks each is on the device its parameter's -storage names, and calls a launch shim per ``Launch`` in its body. Which -arguments that call takes is the callee's answer, reached through the symbol -table, so nothing here spells a device type or a device ABI. Output uses only -TVM FFI, DLPack and standard C++. +Launch entries validate and bind runtime tensors before calling a device shim. +Dispatch entries use a first-match shape predicate and throw on fallback. Output +uses only TVM FFI, DLPack, and standard C++; CUDA syntax and types remain in the +device module and shims. """ from __future__ import annotations -from dataclasses import replace - -from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cpu.templates import render +from tilefoundry.codegen.cuda.module import shim_symbol +from tilefoundry.codegen.cuda.tir.prim_function import ( + _internal_wrapper_symbol, + _is_hidden_shape_scalar, + _parse_shape_param_name, +) from tilefoundry.codegen.linkable import LinkableFunction, LinkableModule from tilefoundry.codegen.registry import CodeGenerator -from tilefoundry.codegen.signature import ( - CallableSignature, - LaunchSignature, - tensor_signature_of, -) +from tilefoundry.codegen.signature import LaunchSignature from tilefoundry.ir.core import Call, Constant, Var -from tilefoundry.ir.core.module import Module +from tilefoundry.ir.core.module import Module, owning_module +from tilefoundry.ir.core.pattern import DimVarRangePat, locate_dim_var from tilefoundry.ir.tir.launch import Launch from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.tir.stmts import Evaluate, Sequential +from tilefoundry.ir.tir.symbol_ref import symbol_call +from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.dim import ( DimAdd, DimFloorDiv, @@ -45,6 +46,27 @@ StorageKind.HOST: "kDLCPU", } + +_LAUNCH_ABI_DECL = ["int", "int", "int", "int", "int", "int", "int", "void*"] + + +def _static_smem(value) -> int: + """Resolve ``dynamic_smem`` to a static int. + + Resolve ``dynamic_smem`` to a static int. Host index-expr codegen + covers the projected grid/block geometry only; a dynamic shared-memory size expression + is not supported. + """ + sv = static_dim_value(value) + if sv is not None: + return sv + raise ValueError( + "emit_host_module: dynamic_smem must be a static int/Constant; a " + f"dynamic shared-memory size expression is not supported, got " + f"{type(value).__name__}" + ) + + _DIM_BINOP_CXX = { DimAdd: "+", DimSub: "-", @@ -84,48 +106,49 @@ def default_visit(self, expr, ctx=None) -> str: raise ValueError(f"emit_host_module: unsupported launch-extent node {type(expr).__name__}") -def _extent(expr) -> str: - """One grid or block extent as a host C++ int expression. +def _emit_host_int_expr(expr) -> str: + """Lower a launch-extent Expr to a host C++ integer expression. - An ``int`` or integer ``Constant`` is written as a literal, a ``ShapeOf`` - reads the forwarded tensor's shape, and dim arithmetic over those composes. - Anything else raises -- there is no silent zero. + Accepts an ``int`` / integer ``Constant`` (emitted as a literal), a + ``ShapeOf`` (the forwarded tensor's ``shape()`` access), or a + dim-arithmetic ``Call`` over those. Any unsupported node raises — there + is no silent zero/default. """ - value = static_dim_value(expr) - if value is not None: - return str(value) - if not isinstance(expr, (ShapeOf, Call, Constant)): - raise ValueError(f"emit_host_module: unsupported launch-extent node {type(expr).__name__}") - return f"static_cast({_HostIntExprVisitor().visit(expr)})" + sv = static_dim_value(expr) + if sv is not None: + return str(sv) + if isinstance(expr, (ShapeOf, Call, Constant)): + return _HostIntExprVisitor().visit(expr) + raise ValueError(f"emit_host_module: unsupported launch-extent node {type(expr).__name__}") -def _static_smem(value) -> int: - """The dynamic shared-memory size, which the host states as a constant. +def _hidden_names(params) -> set: + return {p.name for p in params if _is_hidden_shape_scalar(p, params)} - Host index-expr codegen covers grid and block extents only; a shared-memory - size that is decided at run time has nowhere to be decided. - """ - static = static_dim_value(value) - if static is None: - raise ValueError( - "emit_host_module: dynamic_smem must be a static int/Constant; a " - f"dynamic shared-memory size expression is not supported, got " - f"{type(value).__name__}" - ) - return static +def _is_user_scalar(p, hidden: set) -> bool: + return p.name not in hidden and isinstance(p.type, TensorType) and not p.type.shape -def _reject_unsupported_config(cfg: Launch) -> None: - if cfg.cluster is not None: - raise NotImplementedError("emit_host_module: launch `cluster` is not supported yet") - if cfg.stream is not None: - raise NotImplementedError("emit_host_module: launch `stream` is not supported yet") - if cfg.attrs.entries: - raise NotImplementedError("emit_host_module: launch `attrs` are not supported yet") + +def _is_tensor(p, hidden: set) -> bool: + return p.name not in hidden and not _is_user_scalar(p, hidden) + + +def _extent(c): + value = static_dim_value(c) + return str(value) if value is not None else f"static_cast({_emit_host_int_expr(c)})" + + +def _call_arg(p, host_names, hidden): + name = host_names[p.name] + if _is_tensor(p, hidden): + return f"{name}.data_ptr()" + if p.name in hidden: + return name + return f"static_cast({name})" def _placement_line(name: str, storage) -> str: - """Refuse a tensor that is not where the parameter's storage says it is.""" device_type = _STORAGE_DEVICE_TYPE.get(storage) if device_type is None: raise ValueError( @@ -139,145 +162,314 @@ def _placement_line(name: str, storage) -> str: ) -def _launch_statements(entry: PrimFunction) -> tuple[Evaluate, ...]: - """The launches *entry* makes, which is the whole of what a host entry does.""" - body = entry.body - statements = body.body if isinstance(body, Sequential) else () - if not statements or not all( - isinstance(stmt, Evaluate) and isinstance(stmt.callable, Launch) for stmt in statements - ): - raise ValueError( - f"emit_host_module: entry {entry.name!r} body must be one or more Launch statements" - ) - return tuple(statements) +def _shim_decl(fn: PrimFunction) -> str: + """A types-only ``extern "C"`` forward declaration of *fn*'s launch shim.""" + hidden = _hidden_names(fn.params) + tokens = ["void*" if _is_tensor(p, hidden) else "long long" for p in fn.params] + tokens += _LAUNCH_ABI_DECL + return f'extern "C" void {shim_symbol(fn.name)}({", ".join(tokens)});' -def _as_this_scope_names_it( - shim: CallableSignature, forwarded: tuple[Var, ...], launch: Evaluate -) -> CallableSignature: - """*shim*'s parameters rewritten as the values this launch passes for them. +def emit_host_module( + module: Module, functions: tuple[PrimFunction, ...], target: Target +) -> LinkableModule: + """Emit the host ``.cpp`` linkable module for a CPU *entry*. - Order and kind are the callee's and do not change; every name does, because - a caller writes what its own scope calls the thing. The ids the target - states keep their names: the entry was told them under exactly those. + *module* is the enclosing ``Module``; the dispatch path resolves each + case's ``SymbolRef`` callee through ``module.lookup`` to read the variant's + parameters. """ - geometry = ( - *(_extent(arg) for arg in launch.args[1:7]), - str(_static_smem(launch.callable.dynamic_smem)), - "nullptr", + if len(functions) != 1: + raise ValueError("emit_host_module: expected exactly one CPU host entry") + entry = functions[0] + body = entry.body + if ( + isinstance(body, Sequential) + and body.body + and all( + isinstance(stmt, Evaluate) and isinstance(stmt.callable, Launch) for stmt in body.body + ) + ): + if len(body.body) > 1: + shim_decls, body_lines, sig = _lower_launches(entry, body.body, module) + else: + shim_decls, body_lines, sig = _lower_launch(entry, body.body[0], module) + elif entry.variants: + pat = entry.variants[0].specializations[0] + loc_idx = locate_dim_var(entry.params, pat.dim_var) + loc = (entry.params[loc_idx[0]], loc_idx[1]) if loc_idx is not None else None + if loc is None: + raise ValueError("emit_host_module: cannot derive specialization subject") + p, axis = loc + subject = ShapeOf(type=TensorType.scalar(DType.i32), param=p, axis=axis) + + def _variant_call(v): + args = [] + for vp in v.params: + parsed = _parse_shape_param_name(vp.name) + if parsed is None: + args.append(next(p for p in entry.params if p.name == vp.name)) + else: + base, ax = parsed + ep = next(p for p in entry.params if p.name == base) + args.append(ShapeOf(type=vp.type, param=ep, axis=ax)) + return symbol_call(v, tuple(args)) + + calls = tuple(_variant_call(v) for v in entry.variants) + shim_decls, body_lines, sig = _lower_dispatch( + entry, + module, + callee_name=entry.name, + subject=subject, + variants=entry.variants, + calls=calls, + ) + else: + raise ValueError(f"emit_host_module: entry {entry.name!r} body must be a single Launch") + source = render( + "cpu_module.cpp.j2", + shim_decls=shim_decls, + internal_host_symbol=_internal_wrapper_symbol(entry.name), + wrapper_params_sig=sig, + body_lines=body_lines, + entry_name=entry.name, ) - return replace( - shim, - params=tuple(tensor_signature_of(var) for var in forwarded), - trailing=tuple( - LaunchSignature(name=value, ctype=signature.ctype) - for signature, value in zip(shim.trailing, geometry) - ), + return LinkableModule( + target="cpu", + language="cpp", + source=source, + functions=(LinkableFunction(name=entry.name, source=source),), ) -def _forwarded(entry: PrimFunction, launch: Evaluate, callee: PrimFunction) -> tuple[Var, ...]: - """The entry parameters this launch hands the callee, in the callee's order.""" - args = launch.args[7:] - if len(args) != len(callee.params): - raise ValueError( - f"emit_host_module: launch passes {len(args)} args but device function " - f"{callee.name!r} has {len(callee.params)} parameters" +CPU_CODE_GENERATOR = CodeGenerator(emit_host_module) + + +def _lower_launch(entry: PrimFunction, evaluate, module): + device_fn = module.lookup(evaluate.args[0].name) + if device_fn.variants: + dim_name = device_fn.variants[0].specializations[0].dim_var + loc_idx = locate_dim_var(device_fn.params, dim_name) + loc = (device_fn.params[loc_idx[0]], loc_idx[1]) if loc_idx is not None else None + if loc is None: + raise ValueError(f"cannot derive specialization subject {dim_name!r}") + p, axis = loc + subject = ShapeOf(type=TensorType.scalar(DType.i32), param=p, axis=axis) + calls = tuple(symbol_call(v, tuple(evaluate.args[1:])) for v in device_fn.variants) + return _lower_dispatch( + entry, + module, + callee_name=device_fn.name, + subject=subject, + variants=device_fn.variants, + calls=calls, ) - held = {param.name: param for param in entry.params} - for arg in args: - if not isinstance(arg, Var): + return _lower_launches(entry, (evaluate,), module) + + +def _lower_launches(entry: PrimFunction, evaluates, module): + bindings = [] + kinds = {} + used = set() + for evaluate in evaluates: + launch_op = evaluate.callable + device_fn = module.lookup(evaluate.args[0].name) + if device_fn.variants: + raise NotImplementedError( + "emit_host_module: multi-launch entries do not support variants" + ) + dev_params = device_fn.params + _reject_unsupported_config(launch_op) + hidden = _hidden_names(dev_params) + visible_dev = [p for p in dev_params if p.name not in hidden] + args = evaluate.args[1:] + if len(args) != len(visible_dev): + raise ValueError( + f"emit_host_module: launch passes {len(args)} args but device " + f"function {device_fn.name!r} has {len(visible_dev)} host-visible " + "parameters (hidden shape scalars are derived from tensor shapes)" + ) + if not all(isinstance(a, Var) for a in args): raise ValueError( "emit_host_module: launch args must be host entry parameters (Var); " "expressions are not accepted" ) - if arg.name not in held: + bound = [] + for arg in args: + ep = next((p for p in entry.params if p.name == arg.name), None) + if ep is None: + raise ValueError( + f"emit_host_module: launch arg {arg.name!r} is not a parameter " + f"of entry {entry.name!r}" + ) + bound.append(ep) + local_seen = set() + host_name_of = {} + for vp, ep in zip(visible_dev, bound): + if id(ep) in local_seen: + raise ValueError( + f"emit_host_module: entry parameter {ep.name!r} is bound more " + "than once in one launch" + ) + local_seen.add(id(ep)) + used.add(id(ep)) + kind = _is_tensor(vp, hidden) + if id(ep) in kinds and kinds[id(ep)] != kind: + raise ValueError( + f"emit_host_module: entry parameter {ep.name!r} is bound to " + "incompatible device parameter types across launches" + ) + kinds[id(ep)] = kind + host_name_of[vp.name] = ep.name + bindings.append((evaluate, device_fn, hidden, host_name_of)) + for ep in entry.params: + if id(ep) not in used: raise ValueError( - f"emit_host_module: launch arg {arg.name!r} is not a parameter " - f"of entry {entry.name!r}" + f"emit_host_module: entry parameter {ep.name!r} is not used by any " + f"launch in {entry.name!r} -- an unused parameter has no device " + "type to give the wrapper's signature" ) - return tuple(held[arg.name] for arg in args) - - -def _one_launch( - entry: PrimFunction, launch: Evaluate, callee: PrimFunction, ctx: CpuCodegenContext -) -> list[str]: - """The lines one ``Launch`` becomes: the checks it owes, then the call.""" - _reject_unsupported_config(launch.callable) - forwarded = _forwarded(entry, launch, callee) - shim = ctx.signature_of(callee) - bound = _as_this_scope_names_it(shim, forwarded, launch) - lines = [ - _placement_line(var.name, param.type.storage) - for var, param in zip(forwarded, callee.params) - ] - lines.append(f"{shim.name}({ctx.arguments(bound, callee.target)});") - return lines - - -def _declare_callee(callee: PrimFunction, ctx: CpuCodegenContext) -> str: - """The forward declaration of a symbol another unit defines. - - Written by the callee's own target, so this unit and the one that defines - it cannot come to disagree about what the symbol takes. - """ - shim = ctx.signature_of(callee) - return ( - f'extern "C" void {shim.name}' - f"({ctx.parameters(shim, callee.target, exported=True)});" + body_lines = [] + shim_decls = [] + + for index, (evaluate, device_fn, hidden, dev_to_host) in enumerate(bindings): + dev_params = device_fn.params + launch_op = evaluate.callable + host_names = { + **dev_to_host, + **{p.name: f"l{index}__{p.name}" for p in dev_params if p.name in hidden}, + } + for p in dev_params: + if _is_tensor(p, hidden): + body_lines.append(_placement_line(host_names[p.name], p.type.storage)) + for p in dev_params: + if p.name in hidden: + base, axis = _parse_shape_param_name(p.name) + host_base = dev_to_host.get(base) + if host_base is None: + raise ValueError( + f"emit_host_module: hidden shape scalar {p.name!r} " + f"references unknown base parameter {base!r}" + ) + body_lines.append( + f"long long {host_names[p.name]} = static_cast(" + f"{host_base}.shape()[{axis}]);" + ) + + owner = owning_module(module, device_fn) or module + launch = LaunchSignature.from_topologies(owner.effective_topologies()) + call_args = [_call_arg(p, host_names, hidden) for p in dev_params] + [ + *(str(d) for d in (*launch.grid, *launch.block)), + str(_static_smem(launch_op.dynamic_smem)), + "nullptr", + ] + body_lines.append(f"{shim_symbol(device_fn.name)}({', '.join(call_args)});") + shim_decls.append(_shim_decl(device_fn)) + sig = ", ".join( + f"tvm::ffi::Tensor {ep.name}" if kinds[id(ep)] else f"int {ep.name}" for ep in entry.params ) + return shim_decls, body_lines, sig -def _host_entry( - entry: PrimFunction, - signature: CallableSignature, - body_lines: list[str], - ctx: CpuCodegenContext, -) -> LinkableFunction: - """*entry* in both of its positions, written from the one signature it is called by.""" - return LinkableFunction( - name=entry.name, - declaration=f"void {signature.name}({ctx.parameters(signature)});", - definition=render( - "cpu_entry.cpp.j2", - host_symbol=signature.name, - params=ctx.parameters(signature), - body_lines=body_lines, - entry_name=entry.name, - ), - ) +def _lower_dispatch(entry: PrimFunction, module, *, callee_name, subject, variants, calls): + if not isinstance(subject, ShapeOf): + raise NotImplementedError( + "emit_host_module: dispatch v1 expects exactly one ShapeOf subject" + ) + for variant in variants: + if len(variant.specializations) != 1 or not isinstance( + variant.specializations[0], DimVarRangePat + ): + raise NotImplementedError( + "emit_host_module: dispatch v1 expects exactly one DimVarRangePat per case" + ) + entry_params = entry.params + entry_names = {p.name for p in entry_params} + hidden = _hidden_names(entry_params) -def emit_host_module( - module: Module, - functions: tuple[PrimFunction, ...], - target: Target, - ctx: CpuCodegenContext, -) -> LinkableModule: - """Emit the host ``.cpp`` linkable module for a CPU *entry*. + def _host_name(ref) -> str: + nm = ref.name if isinstance(ref, Var) else ref.param.name + if nm not in entry_names: + raise ValueError( + f"emit_host_module: dispatch arg {nm!r} is not a parameter of entry {entry.name!r}" + ) + return nm + + wrapper_tokens, body_lines = [], [] + for p in entry_params: + if p.name in hidden: + continue + if _is_user_scalar(p, hidden): + wrapper_tokens.append(f"int {p.name}") + else: + wrapper_tokens.append(f"tvm::ffi::Tensor {p.name}") + body_lines.append(_placement_line(p.name, p.type.storage)) + + subj = subject + s = "__tf_dispatch_subject" + body_lines.append( + f"long long {s} = static_cast({_host_name(subj)}.shape()[{subj.axis}]);" + ) - *module* is the enclosing ``Module``, which resolves each launch's callee; - the FFI export names the symbol it exports, so it is written with the - definition and not in the preamble every declaration follows. - """ - if len(functions) != 1: - raise ValueError("emit_host_module: expected exactly one CPU host entry") - entry = functions[0] - declarations: dict[str, str] = {} - body_lines: list[str] = [] - for launch in _launch_statements(entry): - callee = module.lookup(launch.args[0].name) - declarations[ctx.signature_of(callee).name] = _declare_callee(callee, ctx) - body_lines += _one_launch(entry, launch, callee, ctx) - return LinkableModule( - target="cpu", - language="cpp", - preamble=render("cpu_preamble.cpp.j2", shim_decls=list(declarations.values())), - functions=(_host_entry(entry, ctx.signature_of(entry), body_lines, ctx),), + for variant, call in zip(variants, calls): + if len(call.args) != len(variant.params): + raise ValueError( + f"emit_host_module: dispatch call to {variant.name!r} passes " + f"{len(call.args)} args for {len(variant.params)} parameters" + ) + shim_decls: dict[str, str] = {} + for idx, (variant, call) in enumerate(zip(variants, calls)): + pat = variant.specializations[0] + variant_symbol = variant.name + shim_decls[shim_symbol(variant_symbol)] = _shim_decl(variant) + if len(call.args) != len(variant.params): + raise ValueError( + f"emit_host_module: dispatch call to {variant.name!r} passes " + f"{len(call.args)} args for {len(variant.params)} parameters" + ) + v_hidden = _hidden_names(variant.params) + shim_args = [] + for vp, arg in zip(variant.params, call.args): + if _is_tensor(vp, v_hidden): + shim_args.append(f"{_host_name(arg)}.data_ptr()") + elif vp.name in v_hidden: + if not isinstance(arg, ShapeOf): + raise NotImplementedError( + f"emit_host_module: hidden shape param {vp.name!r} expects a ShapeOf arg" + ) + shim_args.append(f"static_cast({_host_name(arg)}.shape()[{arg.axis}])") + else: + shim_args.append(f"static_cast({_host_name(arg)})") + owner = owning_module(module, variant) or module + launch = LaunchSignature.from_topologies(owner.effective_topologies()) + if launch.grid[0] is None: + raise ValueError( + f"emit_host_module: dispatch variant {variant.name!r} has a " + f"launch-provided (dynamic) CTA extent; the dispatch host path " + f"requires a static grid" + ) + shim_args += [str(d) for d in (*launch.grid, *launch.block, 0)] + shim_args.append("nullptr") + pred = f"(({pat.lo} <= {s}) && ({s} <= {pat.hi}))" + prefix = "if" if idx == 0 else "} else if" + body_lines.append(f"{prefix} ({pred}) {{") + body_lines.append(f" {shim_symbol(variant_symbol)}({', '.join(shim_args)});") + body_lines.append("} else {") + body_lines.append( + f' throw std::runtime_error("tilefoundry: no matching dispatch variant for {entry.name}");' ) + body_lines.append("}") + return list(shim_decls.values()), body_lines, ", ".join(wrapper_tokens) -CPU_CODE_GENERATOR = CodeGenerator(emit_host_module) +def _reject_unsupported_config(cfg) -> None: + if cfg.cluster is not None: + raise NotImplementedError("emit_host_module: launch `cluster` is not supported yet") + if cfg.stream is not None: + raise NotImplementedError("emit_host_module: launch `stream` is not supported yet") + if cfg.attrs.entries: + raise NotImplementedError("emit_host_module: launch `attrs` are not supported yet") __all__ = ["emit_host_module"] diff --git a/src/tilefoundry/codegen/cuda/module.py b/src/tilefoundry/codegen/cuda/module.py index 9be4c143b..0d982308f 100644 --- a/src/tilefoundry/codegen/cuda/module.py +++ b/src/tilefoundry/codegen/cuda/module.py @@ -1,176 +1,125 @@ -"""The device ``.cu`` translation unit: one kernel and one shim per function. - -Two passes over each function and nothing between them: one writes the -declaration a caller sees without entering the body, the other writes the -definition by walking it. What is not a function -- the includes, the program -dimensions this unit specializes, the device state something asked for while -it was being written -- settles only once every body has been walked, so it is -assembled last. -""" +"""CUDA device linkable-module emitter (split pipeline). +Emits the device ``.cu`` translation unit: the ``__global__`` kernels (identical +to the single-source path) plus, per kernel, an ``extern "C"`` launch shim that +performs the ``<<>>`` launch. Grid / block / smem / +stream arrive as plain C ABI arguments from the host module, so all CUDA types +stay inside this translation unit. +""" from __future__ import annotations -import math - -from tilefoundry.codegen.cuda.abi import ( - PROGRAM_META, - PROGRAM_META_CTYPE, - kernel_arguments, - kernel_signature, -) -from tilefoundry.codegen.cuda.context import CudaCodegenContext, topology_scope_str +from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.codegen.cuda.templates import render +from tilefoundry.codegen.cuda.tir.prim_function import ( + _compute_kernel_fields, +) from tilefoundry.codegen.linkable import LinkableFunction, LinkableModule from tilefoundry.codegen.registry import CodeGenerator -from tilefoundry.codegen.signature import CallableSignature -from tilefoundry.codegen.topology import Geometry, coarsest_topology -from tilefoundry.ir.core.module import Module +from tilefoundry.ir.core.module import Module, owning_module from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.target import Target -def _emit_function(fn: PrimFunction, ctx: CudaCodegenContext) -> LinkableFunction: - """*fn* in both of its positions, written from the one convention it is called by. +def shim_symbol(kernel_name: str) -> str: + """The ``extern "C"`` launch-shim symbol for *kernel_name*. - A caller can name the shim and nothing else, so that is what the - declaration states; the kernel is defined beside it and named only here. + The ``extern "C"`` launch-shim symbol for *kernel_name* (a plain C + identifier; mangled-variant ``$`` is replaced). """ - shim = ctx.signature_of(fn) - kernel = kernel_signature(shim, fn.name) - return LinkableFunction( - name=fn.name, - declaration=f'extern "C" void {shim.name}({ctx.parameters(shim, exported=True)});', - definition=ctx.capture(lambda scope: _define(fn, shim, kernel, scope)), - ) + return "tilefoundry_" + kernel_name.replace("$", "__") + "_launch" -def _define( - fn: PrimFunction, - shim: CallableSignature, - kernel: CallableSignature, - ctx: CudaCodegenContext, -) -> None: - """The kernel and the shim that launches it, in that order.""" - _define_kernel(fn, kernel, ctx) - ctx.blank() - _define_shim(shim, kernel, ctx) - - -def _define_kernel( - fn: PrimFunction, kernel: CallableSignature, ctx: CudaCodegenContext -) -> None: - """The ``__global__``: what it was told, then what it was written to do.""" - ctx.reset_barrier_ids() - ctx.emit(f"__global__ void {kernel.name}({ctx.parameters(kernel)}) {{") - ctx.indent() - if kernel.leading: - ctx.emit(f"tilefoundry::program_meta({PROGRAM_META.name});") - ctx.emit_node(fn) - ctx.dedent() - ctx.emit("}") - - -def _define_shim( - shim: CallableSignature, kernel: CallableSignature, ctx: CudaCodegenContext -) -> None: - """The ``extern "C"`` entry: the geometry it was handed, spent on one launch.""" - ctx.emit(f'extern "C" void {shim.name}({ctx.parameters(shim, exported=True)}) {{') - ctx.indent() - _tell_the_block(shim, ctx) - grid_x, grid_y, grid_z, block_x, block_y, block_z, smem, stream = ( - signature.name for signature in shim.trailing - ) - ctx.emit(f"dim3 grid({grid_x}, {grid_y}, {grid_z});") - ctx.emit(f"dim3 block({block_x}, {block_y}, {block_z});") - ctx.emit( - f"{kernel.name}<<>>" - f"({kernel_arguments(kernel, ctx)});" - ) - ctx.dedent() - ctx.emit("}") +_LAUNCH_ABI_PARAMS = ( + "int grid_x, int grid_y, int grid_z, " + "int block_x, int block_y, int block_z, " + "int dynamic_smem, void* stream" +) -def _tell_the_block(shim: CallableSignature, ctx: CudaCodegenContext) -> None: - """Gather the ids the shim was told into the block the kernel takes. - Which levels these are is the signature's answer, so the shim states no - level of its own: one line per id the convention put in front. - """ - if not shim.leading: - return - ctx.emit(f"{PROGRAM_META_CTYPE} {PROGRAM_META.name}{{}};") - for signature in shim.leading: - scope = topology_scope_str(signature.topology_level) - ctx.emit( - f"{PROGRAM_META.name}.program_id[int({scope})] = " - f"static_cast({signature.name});" - ) - - -def _one_geometry(kernels: tuple[PrimFunction, ...], ctx: CudaCodegenContext) -> Geometry: - """The geometry this unit runs at, which every kernel in it states the same. - - A unit specializes ``program_dim`` once, so two kernels launched at - different geometries cannot share one: they are two units. - """ - by_kernel = {fn.name: _launched_at(fn, ctx) for fn in kernels} - distinct = set(by_kernel.values()) - if len(distinct) > 1: - raise ValueError( - f"emit_cuda_module: kernels disagree on launch geometry ({by_kernel}); " - f"one translation unit runs every kernel in it at one geometry" - ) - return distinct.pop() - - -def _launched_at(fn: PrimFunction, ctx: CudaCodegenContext) -> Geometry: - """The geometry *fn*'s call states, which is the only thing that states one.""" - geometry = ctx.launches.get(id(fn)) - if geometry is None: - raise ValueError( - f"emit_cuda_module: nothing launches {fn.name!r}, so no call " - f"states the geometry its unit runs at" - ) - return geometry - - -def _program_dims(geometry: Geometry) -> list[dict[str, str]]: - """One instance count per level: what ``program_dim`` states for this unit. - - ``program_shape`` is derived from these, so a unit specializes the counts - and nothing else. A launch-provided grid has no count here; the runtime - states that one when it is asked. - """ - grid, block = geometry - dims = [] - if grid[0] is not None: - dims.append({"scope": topology_scope_str("cta"), "count": str(math.prod(grid))}) - dims.append({"scope": topology_scope_str("thread"), "count": str(math.prod(block))}) - return dims +def _emit_kernel_and_shim(fields) -> str: + shim_params = [] + shim_casts = [] + call_args = [] + for p in fields.params: + if fields.param_kinds[p.name] == "tensor": + cpp = fields.param_cpp_types[p.name] + shim_params.append(f"void* {p.name}") + shim_casts.append(f"{cpp}* {p.name}_p = static_cast<{cpp}*>({p.name});") + call_args.append(f"{p.name}_p") + else: + shim_params.append(f"long long {p.name}") + shim_casts.append(f"int {p.name}_i = static_cast({p.name});") + call_args.append(f"{p.name}_i") + shim_params_sig = ", ".join((*shim_params, _LAUNCH_ABI_PARAMS)) + return render( + "device_kernel.cu.j2", + kernel_name=fields.kernel_name, + kernel_params_sig=fields.kernel_params_sig, + param_wrappers=fields.param_wrappers, + kernel_body=fields.kernel_body, + shim_name=shim_symbol(fields.kernel_name), + shim_params_sig=shim_params_sig, + shim_casts=shim_casts, + kernel_call_args=", ".join(call_args), + ) def emit_cuda_module( - module: Module, cuda_fns: tuple[PrimFunction, ...], target: Target, ctx: CudaCodegenContext + module: Module, cuda_fns: tuple[PrimFunction, ...], target: Target ) -> LinkableModule: - """Emit the device ``.cu`` linkable module for one topology domain.""" - kernels = tuple(fn for fn in cuda_fns if fn.target == target) - if not kernels: - raise ValueError(f"emit_cuda_module: module {module.name!r} has no CUDA device kernels") - functions = tuple(_emit_function(fn, ctx) for fn in kernels) - geometry = _one_geometry(kernels, ctx) - preamble = render( - "cuda_preamble.cu.j2", - program_topology=topology_scope_str(coarsest_topology(module)), - program_dims=_program_dims(geometry), - dynamic_cta=geometry[0][0] is None, - grid_barrier_state=ctx.needs_grid_barrier_state, + """Emit the device ``.cu`` linkable module for *cuda_fns*. + + Emit the device ``.cu`` linkable module for *cuda_fns* (the CUDA-target + PrimFunctions of a module). + """ + from tilefoundry.codegen.cuda.emit import _topology_dim_specializations # noqa: PLC0415 + + kernel_texts = [] + all_fields = [] + expanded_fns = tuple( + {v.name: v for fn in cuda_fns for v in (fn.variants or (fn,))}.values() + ) + for fn in expanded_fns: + owner = owning_module(module, fn) or module + ctx = CudaCodegenContext(target, owner.effective_topologies()) + fields = _compute_kernel_fields(fn, ctx) + kernel_texts.append(_emit_kernel_and_shim(fields)) + all_fields.append(fields) + if not all_fields: + raise ValueError("emit_cuda_module: no CUDA device kernels") + + + + base = all_fields[0] + for f in all_fields[1:]: + if f.launch != base.launch: + raise ValueError( + "emit_cuda_module: kernels disagree on launch topology — " + f"{base.kernel_name!r} has {base.launch!r} " + f"but {f.kernel_name!r} has {f.launch!r}" + ) + specs = _topology_dim_specializations(base.launch.grid, base.launch.block) + + + uses_grid_barrier = any("tf_grid_bar_state" in text for text in kernel_texts) + source = render( + "cuda_module.cu.j2", + topology_dim_specializations=specs, + kernels="\n".join(kernel_texts), + dynamic_cta=base.launch.grid[0] is None, + uses_grid_barrier=uses_grid_barrier, + ) + functions = tuple( + LinkableFunction(name=f.kernel_name, source=text) + for f, text in zip(all_fields, kernel_texts) ) return LinkableModule( - target="cuda", language="cu", preamble=preamble, functions=functions + target="cuda", language="cu", source=source, functions=functions ) CUDA_CODE_GENERATOR = CodeGenerator(emit_cuda_module) -__all__ = ["CUDA_CODE_GENERATOR", "emit_cuda_module"] +__all__ = ["CUDA_CODE_GENERATOR", "emit_cuda_module", "shim_symbol"] diff --git a/src/tilefoundry/codegen/cuda/tir/prim_function.py b/src/tilefoundry/codegen/cuda/tir/prim_function.py index c3124c312..a623b2dec 100644 --- a/src/tilefoundry/codegen/cuda/tir/prim_function.py +++ b/src/tilefoundry/codegen/cuda/tir/prim_function.py @@ -1,107 +1,233 @@ -"""Emit the body of one CUDA kernel: its buffers, then its statements. +"""Emit a TIR primitive function's CUDA kernel and host wrapper. -A kernel is handed pointers and the extents its parameter types leave open, so -the body starts by wrapping each pointer in the cute tensor the ops index. A -specialization prototype has no body of its own: its variants are branches of -the one kernel, chosen on the extent the pattern ranges over, because every -kernel in a translation unit runs at the geometry that unit states once. +The wrapper accepts TVM tensors, extracts pointers and hidden integer shape +scalars, and derives launch dimensions from mesh scopes. A specialization entry +containing only dispatch emits no global kernel; its host wrapper selects and +calls variant wrappers because CUDA kernels cannot call host code. +See [codegen §1](docs/spec/codegen.md#1-pipeline). """ - from __future__ import annotations -from tilefoundry.codegen.cuda.context import CudaCodegenContext +from dataclasses import dataclass + +from tilefoundry.codegen.cuda.context import CodegenContext from tilefoundry.codegen.cuda.tir.memory.tensor_view import render_shard_layout_value -from tilefoundry.codegen.signature import TensorSignature, tensor_signature_of -from tilefoundry.ir.core.pattern import DimVarRangePat +from tilefoundry.codegen.emitter import CudaEmitter +from tilefoundry.codegen.signature import LaunchSignature from tilefoundry.ir.tir.prim_function import PrimFunction +from tilefoundry.ir.tir.shape import ( + is_hidden_shape_scalar as _is_hidden_shape_scalar, +) +from tilefoundry.ir.tir.shape import ( + parse_shape_var_name as _parse_shape_param_name, +) +from tilefoundry.ir.tir.shape import ( + shape_var_name, +) +from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.shape_helpers import shape_numel_upper_bound from tilefoundry.ir.types.shard.shard_layout import ShardLayout -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen -def _wrap_buffer(signature: TensorSignature, ctx: CudaCodegenContext) -> None: - """Wrap one parameter's pointer in the tensor the body indexes it through. +def _internal_wrapper_symbol(kernel_name: str) -> str: + """Map a user-facing kernel name to its internal C++ wrapper symbol. - The layout is the envelope the whole dispatch range fits in, so one kernel - covers every shape its type admits; a sharded parameter carries which part - of that envelope this instance holds. + The user-facing name may be ``main`` (collides with ``::main``) or + a mangled variant like ``main$S$1_4`` (``$`` is a GCC extension, + not portable). The internal symbol is always a plain C++ identifier + so the generated source compiles under strict toolchains. """ - name = signature.name - total = shape_numel_upper_bound(signature.type.shape) + return "__tilefoundry_" + kernel_name.replace("$", "__") + "_host" + + +def _param_wrapper(name: str, total: int, cpp_type: str) -> str: layout = f"cute::make_layout(cute::Shape>{{}})" - pointer = f"cute::make_tensor(cute::make_gmem_ptr({name}), {layout})" - shard = signature.type.layout - if not isinstance(shard, ShardLayout): - ctx.emit(f"auto {name}_tensor = {pointer};") - return - preamble, value = render_shard_layout_value( - f"{name}_tensor", shard, ctx.dynamic_extents, signature.type.storage, ctx + return ( + f"auto {name}_tensor = cute::make_tensor(" + f"cute::make_gmem_ptr({name}), {layout});" ) - for line in preamble: - ctx.emit(line) - ctx.emit( - f"auto {name}_tensor = tilefoundry::make_shard_tensor(" - f"{pointer}, {layout}, {value});" + + +def _param_cpp_types(params: tuple, ctx: CodegenContext) -> dict[str, str]: + """Map each param name → CUDA C++ type from its TensorType dtype.""" + result: dict[str, str] = {} + for p in params: + ty = p.type + if isinstance(ty, TensorType): + result[p.name] = ctx.dtype_to_cpp(ty.dtype.name) + else: + result[p.name] = "float" + return result + + +def _param_wrapper_shard( + name: str, total: int, shard_layout: ShardLayout, dim_var_runtime=None +) -> str: + """Emit make_shard_tensor wrapping a kernel param with ShardLayout value.""" + global_layout = f"cute::make_layout(cute::Shape>{{}})" + tensor_ref = f"{name}_tensor" + preamble, shard_value = render_shard_layout_value( + tensor_ref, shard_layout, dim_var_runtime, None, ctx + ) + wrapper = ( + f"auto {tensor_ref} = tilefoundry::make_shard_tensor(" + f"cute::make_tensor(cute::make_gmem_ptr({name}), {global_layout}), " + f"{global_layout}, {shard_value});" ) + return "\n".join([*preamble, wrapper]) -def _dispatch(fn: PrimFunction, ctx: CudaCodegenContext) -> None: - """Write the variants of *fn* as branches on the extent they range over. +@dataclass(frozen=True) +class _KernelFields: + """Represent KernelFields. - The extent is already a parameter of this kernel, so the choice costs one - comparison inside the launch the host would otherwise have had to make - before it -- and the host has nothing to gain by making it, since all the - variants run at one geometry anyway. + Per-PrimFunction codegen pieces shared by the single-source kernel + emitter and the split device / host fragment emitters. """ - ctx.bind_extents(tensor_signature_of(param) for param in fn.params) - for index, variant in enumerate(fn.variants): - pattern = _range_over_one_dimension(fn, variant) - subject = _subject(fn, pattern, ctx) - opening = "if" if index == 0 else "} else if" - ctx.emit(f"{opening} (({pattern.lo} <= {subject}) && ({subject} <= {pattern.hi})) {{") - ctx.indent() - ctx.emit_node(variant) - ctx.dedent() - ctx.emit("} else {") - ctx.indent() - ctx.emit("__trap();") - ctx.dedent() - ctx.emit("}") - - -def _range_over_one_dimension(fn: PrimFunction, variant: PrimFunction) -> DimVarRangePat: - """The pattern selecting *variant*, which a branch can only be a range.""" - pattern = variant.specializations[0] - if not isinstance(pattern, DimVarRangePat): - raise NotImplementedError( - f"CUDA dispatch: {fn.name!r} selects {variant.name!r} by " - f"{type(pattern).__name__}; only a range over one dimension is written" - ) - return pattern + kernel_name: str + internal_wrapper_name: str + params: tuple + param_cpp_types: dict + param_kinds: dict + kernel_params_sig: str + wrapper_params_sig: str + user_params: tuple + launch_args: str + param_wrappers: list + wrapper_locals: list + kernel_body: str + launch: LaunchSignature + entry_host_only: bool -def _subject(fn: PrimFunction, pattern: DimVarRangePat, ctx: CudaCodegenContext) -> str: - """Where this kernel reads the dimension *pattern* ranges over.""" - subject = ctx.dynamic_extents.get(pattern.dim_var) - if subject is None: + +def _compute_kernel_fields(node: PrimFunction, ctx: CodegenContext) -> _KernelFields: + if node.variants: raise ValueError( - f"CUDA dispatch: {fn.name!r} selects on {pattern.dim_var!r}, which " - f"no parameter's type leaves open, so the kernel is never told it" + f"PrimFunction {node.name!r} is a specialization prototype; " + "codegen must dispatch variants instead of emitting an empty kernel" ) - return subject - - -@register_codegen(CudaTarget, Role.EMIT, PrimFunction) -def _emit(fn: PrimFunction, ctx: CudaCodegenContext) -> None: - """Write what runs inside one ``__global__``: its buffers, then its statements.""" - if fn.variants: - _dispatch(fn, ctx) - return - params = tuple(tensor_signature_of(param) for param in fn.params) - for param in fn.params: - ctx.register_kernel_param(param) - ctx.bind_extents(params) - for signature in params: - _wrap_buffer(signature, ctx) - ctx.emit_node(fn.body) + for p in node.params: + ctx.register_kernel_param(p) + + + + + + + ctx._dim_var_runtime = {} + for p in node.params: + ty = p.type + if not isinstance(ty, TensorType): + continue + for axis, dim in enumerate(ty.shape): + if isinstance(dim, DimVar) and dim.name not in ctx._dim_var_runtime: + ctx._dim_var_runtime[dim.name] = shape_var_name(p.name, axis) + + ctx.reset_barrier_ids() + + entry_host_only = bool(node.variants) + + + + body = ctx.capture(lambda inner: CudaEmitter(context=inner).visit(node.body)) + + param_cpp_types = _param_cpp_types(node.params, ctx) + hidden_shape = tuple( + p for p in node.params if _is_hidden_shape_scalar(p, node.params) + ) + hidden_names = {p.name for p in hidden_shape} + user_params = tuple(p for p in node.params if p.name not in hidden_names) + + + + + def _is_user_scalar(p) -> bool: + return ( + p.name not in hidden_names + and isinstance(p.type, TensorType) + and not p.type.shape + ) + + def _kind(p) -> str: + if p.name in hidden_names: + return "hidden_scalar" + if _is_user_scalar(p): + return "user_scalar" + return "tensor" + + param_kinds = {p.name: _kind(p) for p in node.params} + + def _kernel_sig_token(p) -> str: + if param_kinds[p.name] != "tensor": + return f"int {p.name}" + return f"{param_cpp_types[p.name]}* {p.name}" + + kernel_params_sig = ", ".join(_kernel_sig_token(p) for p in node.params) + + + + + def _wrapper_param_token(p) -> str: + if param_kinds[p.name] == "user_scalar": + return f"int {p.name}" + return f"tvm::ffi::Tensor {p.name}" + + wrapper_params_sig = ", ".join(_wrapper_param_token(p) for p in user_params) + wrapper_locals = [] + for p in hidden_shape: + parsed = _parse_shape_param_name(p.name) + + assert parsed is not None + base, axis = parsed + wrapper_locals.append( + f"int {p.name} = static_cast({base}.shape()[{axis}]);" + ) + + + + def _launch_arg(p) -> str: + if param_kinds[p.name] != "tensor": + return p.name + return f"static_cast<{param_cpp_types[p.name]}*>({p.name}.data_ptr())" + + launch_args = ", ".join(_launch_arg(p) for p in node.params) + + + + buffer_params = tuple( + p for p in user_params if param_kinds[p.name] == "tensor" + ) + param_wrappers = [] + for p in buffer_params: + total = shape_numel_upper_bound(p.type.shape) + sl = getattr(p.type, "layout", None) + if isinstance(sl, ShardLayout): + param_wrappers.append( + _param_wrapper_shard(p.name, total, sl, ctx._dim_var_runtime) + ) + else: + param_wrappers.append( + _param_wrapper(p.name, total, param_cpp_types[p.name]) + ) + + launch = LaunchSignature.from_topologies(ctx.topologies) + + codegen_name = node.name + return _KernelFields( + kernel_name=codegen_name, + internal_wrapper_name=_internal_wrapper_symbol(codegen_name), + params=node.params, + param_cpp_types=param_cpp_types, + param_kinds=param_kinds, + kernel_params_sig=kernel_params_sig, + wrapper_params_sig=wrapper_params_sig, + user_params=user_params, + launch_args=launch_args, + param_wrappers=param_wrappers, + wrapper_locals=wrapper_locals, + kernel_body=body, + launch=launch, + entry_host_only=entry_host_only, + ) diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py index 0e6ef022c..d7f178633 100644 --- a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py +++ b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py @@ -1,21 +1,15 @@ -"""Emit CUDA mesh scopes. +"""Helpers for CUDA ``MeshScope`` lowering. -Emitter for `tir.MeshScope` — emits a C++ block + comment marker + -constexpr Mesh type alias ([runtime §2.3.2](docs/spec/runtime.md#232-layoutmeshcuh)). +The ``CudaEmitter`` owns Stmt traversal and invokes these stateless mesh +layout helpers while visiting a ``MeshScope``. """ - from __future__ import annotations -from tilefoundry.codegen.cuda.context import ( - CudaCodegenContext, - topology_scope_str, -) -from tilefoundry.ir.tir.stmts import MeshScope +from tilefoundry.codegen.cuda.context import topology_scope_str from tilefoundry.ir.tir.sync import participation from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout from tilefoundry.ir.types.shard.mesh import Mesh, Topology, positions_at -from tilefoundry.target import CudaTarget, validate_cuda_topology_levels -from tilefoundry.visitor_registry.registries import Role, register_codegen +from tilefoundry.target import validate_cuda_topology_levels def _resolved(topology: Topology | str) -> Topology: @@ -43,14 +37,14 @@ def _validate_topology(mesh: Mesh, target) -> None: validate_cuda_topology_levels(target, (_resolved(t).name for t in mesh.topologies)) -def _group_mesh_layout(mesh: Mesh, topos) -> str: +def _levelwise_layout(mesh: Mesh, topos) -> str: """One nest per topology level, each in that level's own numbering. A flat shape says nothing about which axes are whose, so the runtime's - ``get`` needs the boundary stated. ``topology_axes`` hands the - axes to the levels left to right and ``positions_at`` divides each level's - strides by what the levels under it contribute, so a nest reads as the - layout that level would have alone. The composite is a grouping, not a map: only the + ``get`` needs the boundary stated. ``level_axes`` hands the axes to + the levels left to right and ``positions_at`` divides each level's strides + by what the levels under it contribute, so a nest reads as the layout that + level would have alone. The composite is a grouping, not a map: only the nests are evaluated. A level every instance shares keeps a mode of one, so that ``get`` still has something to pick. """ @@ -59,9 +53,20 @@ def _group_mesh_layout(mesh: Mesh, topos) -> str: level_shape, level_strides = positions_at(mesh, topology.name) if not level_shape: level_shape, level_strides = (1,), (0,) - shapes.append("cute::Shape<" + ", ".join(f"cute::Int<{s}>" for s in level_shape) + ">") - strides.append("cute::Stride<" + ", ".join(f"cute::Int<{s}>" for s in level_strides) + ">") - return f"cute::Layout, cute::Stride<{', '.join(strides)}>>" + shapes.append( + "cute::Shape<" + + ", ".join(f"cute::Int<{s}>" for s in level_shape) + + ">" + ) + strides.append( + "cute::Stride<" + + ", ".join(f"cute::Int<{s}>" for s in level_strides) + + ">" + ) + return ( + f"cute::Layout, " + f"cute::Stride<{', '.join(strides)}>>" + ) def mesh_type(mesh: Mesh) -> str: @@ -83,7 +88,10 @@ def mesh_type(mesh: Mesh) -> str: if len(topos) == 1: shape_types = ", ".join(f"cute::Int<{s}>" for s in shape) stride_types = ", ".join(f"cute::Int<{s}>" for s in strides) - layout = f"cute::Layout, cute::Stride<{stride_types}>>" + layout = ( + f"cute::Layout, " + f"cute::Stride<{stride_types}>>" + ) else: if base: raise NotImplementedError( @@ -91,7 +99,7 @@ def mesh_type(mesh: Mesh) -> str: "be sliced; the slice and the level boundary would both be " "deciding which positions these are" ) - layout = _group_mesh_layout(mesh, topos) + layout = _levelwise_layout(mesh, topos) if base: layout = f"cute::ComposedLayout, {layout}>" return f"tilefoundry::Mesh<{layout}, {', '.join(topology_scope_str(t.name) for t in topos)}>" @@ -108,35 +116,4 @@ def _is_dynamic_mesh(mesh: Mesh) -> bool: return any(s is None for s in mesh.layout.shape) -@register_codegen(CudaTarget, Role.EMIT, MeshScope) -def _emit(node: MeshScope, ctx: CudaCodegenContext) -> None: - """Emit the block a mesh scope is, and the mesh object it states.""" - if ctx.target is None: - raise RuntimeError("CUDA MeshScope emission requires its Target") - _validate_topology(node.mesh, ctx.target) - name = ctx.name_for(node.binding) - ctx.emit(f"// mesh scope: {program_topologies(node.mesh)[0].name}") - - is_slice = isinstance(node.mesh.layout, ComposedLayout) - ctx.emit("{") - ctx.indent() - outer_aliases = ctx._mesh_aliases - ctx._mesh_aliases = dict(outer_aliases) - try: - if not _is_dynamic_mesh(node.mesh): - alias = f"{name}_mesh_t" - mesh_type_str = mesh_type(node.mesh) - ctx._mesh_aliases[id(node.mesh)] = (alias, mesh_type_str) - ctx.emit(f"using {alias} = {mesh_type_str};") - ctx.emit(f"constexpr {alias} {name}_mesh{{}};") - if is_slice: - ctx.emit(f"if (tilefoundry::contains({name}_mesh, tilefoundry::program_ids())) {{") - ctx.indent() - ctx.emit_node(node.body) - if is_slice: - ctx.dedent() - ctx.emit("}") - finally: - ctx._mesh_aliases = outer_aliases - ctx.dedent() - ctx.emit("}") +__all__ = ["mesh_type", "program_topologies"] diff --git a/src/tilefoundry/codegen/emitter.py b/src/tilefoundry/codegen/emitter.py new file mode 100644 index 000000000..ae40454bb --- /dev/null +++ b/src/tilefoundry/codegen/emitter.py @@ -0,0 +1,158 @@ +"""C++ source emitters shared by CUDA and CPU code generation.""" + +from __future__ import annotations + +from tilefoundry.codegen.cuda.tir.stmts.mesh_scope import ( + _is_dynamic_mesh, + _validate_topology, + mesh_type, + program_topologies, +) +from tilefoundry.codegen.cuda.tir.stmts.scalar_expr import render_scalar_expr +from tilefoundry.ir.core import Call +from tilefoundry.ir.tir.stmts import ( + Evaluate, + For, + If, + LetStmt, + MeshScope, + Return, + Sequential, + While, +) +from tilefoundry.ir.types import UnitType +from tilefoundry.ir.types.shard.layout import ComposedLayout +from tilefoundry.ir.visitor import StmtVisitor + + +class CppEmitter(StmtVisitor[None]): + """Write C++ from TIR; the context carries state, this class traverses.""" + + def __init__(self, *, context) -> None: + self.context = context + + def emit(self, line: str) -> None: + self.context.emit(line) + + def indent(self) -> None: + self.context.indent() + + def dedent(self) -> None: + self.context.dedent() + + def statement(self, text: str) -> None: + self.emit(f"{text};") + + def generic_visit(self, node) -> None: + raise RuntimeError(f"no emitter for {type(node).__name__}") + + def visit_Sequential(self, node: Sequential) -> None: + for stmt in node.body: + self.visit(stmt) + + def visit_If(self, node: If) -> None: + cond = self.render_scalar(node.cond) + self.emit(f"if ({cond}) {{") + self.indent() + self.visit(node.then_body) + self.dedent() + if getattr(node.else_body, "body", None): + self.emit("} else {") + self.indent() + self.visit(node.else_body) + self.dedent() + self.emit("}") + + def visit_For(self, node: For) -> None: + name = self.context.name_for(node.induction_var) + start = self.render_scalar(node.start) + stop = self.render_scalar(node.stop) + step = self.render_scalar(node.step) + self.emit(f"for (int {name} = {start}; {name} < {stop}; {name} += {step}) {{") + self.indent() + self.visit(node.body) + self.dedent() + self.emit("}") + + def visit_While(self, node: While) -> None: + self.emit(f"while ({self.render_scalar(node.cond)}) {{") + self.indent() + self.visit(node.body) + self.dedent() + self.emit("}") + + def visit_Return(self, node: Return) -> None: + self.statement("return") + + def visit_LetStmt(self, node: LetStmt) -> None: + if not isinstance(node.value, Call): + raise RuntimeError( + f"LetStmt.value must be a Call (TIR-owned Expr Op), " + f"got {type(node.value).__name__}" + ) + self.emit_let_op(node.value, node) + self.visit(node.body) + + def visit_Evaluate(self, node: Evaluate) -> None: + op = node.callable + self.emit_evaluate_op(Call(type=UnitType(), target=op, args=node.args)) + + def visit_MeshScope(self, node: MeshScope) -> None: + raise NotImplementedError + + def emit_evaluate_op(self, call: Call) -> None: + handler = self.context.lookup(type(call.target)) + if handler is None: + raise RuntimeError(f"no codegen handler for Op {type(call.target).__name__}") + handler(call, self.context) + + def emit_let_op(self, call: Call, stmt: LetStmt) -> None: + handler = self.context.lookup(type(call.target)) + if handler is None: + raise RuntimeError(f"no codegen handler for Op {type(call.target).__name__}") + handler(stmt, self.context) + + def render_scalar(self, expr) -> str: + return render_scalar_expr(expr, self.context) + + +class CudaEmitter(CppEmitter): + """CUDA C++ emitter, including mesh scope lowering.""" + + def visit_MeshScope(self, node: MeshScope) -> None: + ctx = self.context + if ctx.target is None: + raise RuntimeError("CUDA MeshScope emission requires its Target") + _validate_topology(node.mesh, ctx.target) + name = ctx.name_for(node.binding) + ctx.emit(f"// mesh scope: {program_topologies(node.mesh)[0].name}") + + is_slice = isinstance(node.mesh.layout, ComposedLayout) + ctx.emit("{") + ctx.indent() + outer_aliases = ctx._mesh_aliases + ctx._mesh_aliases = dict(outer_aliases) + try: + if not _is_dynamic_mesh(node.mesh): + alias = f"{name}_mesh_t" + mesh_type_str = mesh_type(node.mesh) + ctx._mesh_aliases[id(node.mesh)] = (alias, mesh_type_str) + ctx.emit(f"using {alias} = {mesh_type_str};") + ctx.emit(f"constexpr {alias} {name}_mesh{{}};") + if is_slice: + ctx.emit( + f"if (tilefoundry::contains({name}_mesh, " + "tilefoundry::program_ids())) {" + ) + ctx.indent() + self.visit(node.body) + if is_slice: + ctx.dedent() + ctx.emit("}") + finally: + ctx._mesh_aliases = outer_aliases + ctx.dedent() + ctx.emit("}") + + +__all__ = ["CppEmitter", "CudaEmitter"] diff --git a/src/tilefoundry/codegen/topology.py b/src/tilefoundry/codegen/topology.py deleted file mode 100644 index adc5df63d..000000000 --- a/src/tilefoundry/codegen/topology.py +++ /dev/null @@ -1,112 +0,0 @@ -"""The topology questions every emitter asks of a module and its target. - -Host module, device module and the entry signature describe three segments of -one call, so they read the answer here rather than each deciding for -themselves. -""" - -from __future__ import annotations - -from collections.abc import Iterator, Mapping - -from tilefoundry.ir.core.module import Module, module_functions, subtree -from tilefoundry.ir.tir.launch import Launch -from tilefoundry.ir.tir.stmts import Evaluate, Sequential -from tilefoundry.ir.types.shape_helpers import static_dim_value - -Geometry = tuple[tuple[int | None, int, int], tuple[int, int, int]] -"""The grid and block one launch runs at; a launch-provided CTA extent is ``None``.""" - - -def coarsest_topology(module) -> str: - """The coarsest topology level this module's program names. - - ``effective_topologies()`` is ordered coarsest first, so the first entry - says where this program's run of levels begins; a module that names none - is one CTA. - """ - declared = module.effective_topologies() - return declared[0].name if declared else "cta" - - -def topology_domains(root: Module) -> Iterator[tuple[Module, tuple]]: - """Each module that states topologies of its own, and the functions under it. - - A program states its instance counts once, so a module that states its own - is its own translation unit; one that states none runs in its owner's - program and its functions are emitted into the owner's unit. - """ - yield root, _domain_functions(root) - for node in subtree(root): - for child in node.modules: - if child.topologies is not None: - yield child, _domain_functions(child) - - -def _domain_functions(node: Module) -> tuple: - """*node*'s own functions and those of every child that states no topologies.""" - inherited = tuple( - function - for child in node.modules - if child.topologies is None - for function in _domain_functions(child) - ) - return (*node.functions, *inherited) - - -def launch_geometry(root: Module) -> Mapping[int, Geometry]: - """The geometry each launched function runs at, as its ``Launch`` states it. - - The host entry settled it when it was written, so the device side reads it - rather than walking a body to derive the same answer a second time. - """ - geometry: dict[int, Geometry] = {} - for function in module_functions(root): - body = function.body - if not isinstance(body, Sequential): - continue - for stmt in body.body: - if not (isinstance(stmt, Evaluate) and isinstance(stmt.callable, Launch)): - continue - callee = root.lookup(stmt.args[0].name) - geometry[id(callee)] = ( - _grid(stmt.args[1:4], callee.name), - _block(stmt.args[4:7], callee.name), - ) - return geometry - - -def _grid(args, name: str) -> tuple: - """The three grid extents, whose first may be the one a launch provides. - - A launch-provided CTA count is the one extent a translation unit cannot - state; the runtime states that one when it is asked. - """ - return ( - static_dim_value(args[0]), - *(_static(arg, name, "grid", axis + 1) for axis, arg in enumerate(args[1:])), - ) - - -def _block(args, name: str) -> tuple: - """The three block extents, which multiply into the thread count a unit states.""" - return tuple(_static(arg, name, "block", axis) for axis, arg in enumerate(args)) - - -def _static(arg, name: str, what: str, axis: int) -> int: - """One extent as the compile-time count the program dimensions need.""" - value = static_dim_value(arg) - if value is None: - raise ValueError( - f"launch of {name!r}: {what}[{axis}] is decided at run time, and " - f"only a launch-provided CTA count can be" - ) - return value - - -__all__ = [ - "Geometry", - "coarsest_topology", - "launch_geometry", - "topology_domains", -] diff --git a/src/tilefoundry/inspection/values.py b/src/tilefoundry/inspection/values.py index d78453163..0b8c390cf 100644 --- a/src/tilefoundry/inspection/values.py +++ b/src/tilefoundry/inspection/values.py @@ -1,304 +1,35 @@ -"""How a record renders as one line of comment. - -A record states typed fields ([core-ir §2](docs/spec/core-ir.md#2-expr)); what -they look like is decided here, so a family arrives by declaring what it has -rather than by writing a form string -- which is how five came to spell one value -five ways. A value earns a rendering of its own only where Python cannot already -read it: an ``int``, a token, and a mapping of them get none. The constants below -are the ladder of [inspection §2.8](docs/spec/inspection.md#28-record-comment-forms), -one each, and nothing else in the tree spells one. -""" +"""Canonical metadata comment printer.""" from __future__ import annotations import json -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field, fields -from typing import get_type_hints - -from tilefoundry.analysis.metadata import ( - Breakdown, - ComputeCostMetadata, - LoopFootprintMetadata, - MemoryMetadata, - PerformanceMetadata, - PerformanceSummaryMetadata, - RooflineMetadata, - TrafficMetadata, -) -from tilefoundry.analysis.report import ( - declare_record, - declared_records, - expr_field, - expr_fields, - family_of, - render_record, -) -from tilefoundry.ir.core.metadata import IRMetadata, SourceSpanMetadata -from tilefoundry.ir.core.values import TotalAndPerUnit, TripInterval -from tilefoundry.visitor_registry.contexts import TrafficBytes +from collections.abc import Mapping +from dataclasses import dataclass, field, fields, is_dataclass -PAIR = "/" -PER_UNIT = "@" -ENTRY = ":" -ENTRIES = "," -FIELD = "=" -FIELDS = " " -PARTS = "; " -TRIPS = "*" - - -class Prose(str): - """Text that is a sentence rather than a token. - - A token renders bare, which only works while it holds no separator. A - sentence holds several -- spaces, commas, a semicolon -- so it renders as a - quoted, escaped string literal, which brackets it: a reader splits a layer - outside the quotes ([inspection §2.8](docs/spec/inspection.md#28-record-comment-forms)). - """ - - -def _trip_interval(value: TripInterval, render: Callable[[object], str]) -> str: - """One interval, offset by the trip index when it repeats.""" - if value.trips <= 1: - return f"[{value.start}{ENTRIES}{value.end})" - offset = f"{value.stride}t+" - return f"[{offset}{value.start}{ENTRIES}{offset}{value.end}){TRIPS}{value.trips}" - - -RENDER: dict[type, Callable[..., str]] = { - Prose: lambda value, render: json.dumps(str(value)), - TrafficBytes: lambda value, render: f"r{value.read}{PAIR}w{value.write}", - TotalAndPerUnit: ( - lambda value, render: f"{render(value.total)}{PER_UNIT}{render(value.per_unit)}" - ), - TripInterval: _trip_interval, -} - - -def render_value(value: object) -> str: - """One value, rendered by its own type or as the entries it holds.""" - for value_type, render in RENDER.items(): - if isinstance(value, value_type): - return render(value, render_value) - entries = _entries_of(value) - if entries is not None: - return ENTRIES.join(f"{key}{ENTRY}{render_value(item)}" for key, item in entries) - if isinstance(value, tuple | list): - return ENTRIES.join(render_value(item) for item in value) - return str(value) - - -def _entries_of(value: object) -> tuple[tuple[object, object], ...] | None: - """*value* as key/value pairs, or ``None`` when it is not a mapping.""" - if isinstance(value, Mapping): - return tuple(value.items()) - if isinstance(value, tuple) and all( - isinstance(item, tuple) and len(item) == 2 for item in value - ): - return value - return None +from tilefoundry.analysis.metadata import TripInterval +from tilefoundry.ir.core.metadata import IRMetadata +from tilefoundry.visitor_registry.contexts import TrafficBytes +PAIR, PER_UNIT, ENTRY, ENTRIES, FIELD, FIELDS, PARTS, TRIPS = ( + "/", + "@", + ":", + ",", + "=", + " ", + "; ", + "*", +) _UNSET = object() -@dataclass(frozen=True) -class Projection: - """One key a comment emits: its name, its type, and where its value is read. - - A field is the identity projection of itself. Anything else -- a sum, a - count, several fields folded into one value -- is declared here, which is - what keeps the next one from appearing unannounced. - - *default* is the value this key says nothing by, so it is left out. Without - one, an empty mapping says nothing and every other value is emitted. - """ - - key: str - type: object - of: Callable[..., object] - default: object = _UNSET - opt_in: bool = False - - -@dataclass(frozen=True) -class RecordComment: - """What one record type emits, in order, under which family name.""" - - family: str - emissions: tuple[Projection, ...] - - -_COMMENTS: dict[type[IRMetadata], RecordComment] = {} - - -def _field_projections(record: type[IRMetadata], names: tuple[str, ...]) -> tuple[Projection, ...]: - hints = get_type_hints(record) - defaults = {item.name: item.default for item in fields(record)} - return tuple(Projection(name, hints[name], _read(name), defaults[name]) for name in names) - - -def _read(name: str) -> Callable[..., object]: - return lambda record: getattr(record, name) - - -def comment( - record: type[IRMetadata], - *emitted: str | Projection, - family: str | None = None, -) -> None: - """Declare that *record* renders as a comment, and what it emits. - - A field is named by its own name; anything else is a ``Projection``. Naming - nothing emits every field in declaration order. A record with no declaration - renders as ``None``, which is how metadata that is not a report stays out of - the comment. - - *family* is only for a record whose reported name is not the one its class - name states. - """ - declared = emitted or tuple(item.name for item in fields(record)) - emissions = tuple( - item if isinstance(item, Projection) else _field_projections(record, (item,))[0] - for item in declared - ) - declare_record(record, family=family) - _COMMENTS[record] = RecordComment(family=family_of(record), emissions=emissions) - - -def comment_of(record: type[IRMetadata]) -> RecordComment | None: - """What *record* declared it emits, if it declared anything.""" - return _COMMENTS.get(record) - - -def render_comment(record: IRMetadata, *, opt_in: frozenset[str] = frozenset()) -> str | None: - """One record as one comment, or ``None`` when it does not report. - - A key is its declared name with ``_`` written as ``-``; a unit belongs in - that name, which is where it is said once. A record declaring one key uses - the family name as that key, because for a record of one thing the family - and the key say the same thing twice. - """ - declared = _COMMENTS.get(type(record)) - if declared is None: - return None - alone = len(declared.emissions) == 1 - emitted: list[str] = [] - for emission in declared.emissions: - if emission.opt_in and emission.key not in opt_in: - continue - value = emission.of(record) - if _says_nothing(value, emission.default): - continue - key = declared.family if alone else emission.key.replace("_", "-") - emitted.append(f"{key}{FIELD}{render_value(value)}") - if alone and emitted: - return emitted[0] - return FIELDS.join([declared.family, *emitted]) - - -def _says_nothing(value: object, default: object) -> bool: - """Whether this value adds nothing to the comment it would appear in.""" - if default is not _UNSET: - return value == default - entries = _entries_of(value) - return entries is not None and not entries - - -def _paired[V]( - held: "Breakdown[V]", topologies: tuple[str, ...] -) -> dict[str, TotalAndPerUnit[dict[str, V]]]: - """Each kind's quantity, whole beside what one unit of each level holds. - - The level names come from the record's own ``topologies``, which is where - they are written once; a ``Spread`` states its shares in that order and - carries no names of its own. - """ - return { - kind: TotalAndPerUnit(spread.total, dict(zip(topologies, spread.per_unit, strict=False))) - for kind, spread in held.kinds - } - - -def _paired_flops(record: ComputeCostMetadata): - """Each dtype's work, whole and per unit at every declared level.""" - return _paired(record.flops, record.topologies) - - -def _paired_service(record: ComputeCostMetadata): - """Each service kind's work, whole and per unit at every declared level. - - What a machine is asked for that is not floating point: comparing, selecting, - integer arithmetic, a reciprocal, a local move. Reported beside the flops - rather than folded into them, because a predicate priced as a FLOP is a - number about a pipe the work never went down. - """ - return _paired(record.service, record.topologies) - - -def _paired_traffic(record: TrafficMetadata): - """Each storage level's traffic, whole and per unit at every declared level.""" - return _paired(record.storage, record.topologies) - - -def _by_operand(record: TrafficMetadata) -> dict[str, TrafficBytes]: - """What each operand moved, positional against ``(*call.args, call)``.""" - last = len(record.operands) - 1 - return { - ("result" if index == last else str(index)): moved - for index, moved in enumerate(record.operands) - } - - -def peak_footprint(record: MemoryMetadata) -> dict[str, int]: - """How much of each level the function holds at its peak.""" - return {item.memory_level: item.peak_bytes for item in record.footprint} - - -def _persistent_bytes(record: MemoryMetadata) -> int: - """The part of the peak the function cannot reclaim, across levels.""" - return sum(item.persistent_bytes for item in record.footprint) - - -def _advisory_count(record: MemoryMetadata) -> int: - """How many advisories there are; the comment is not where they are read.""" - return len(record.advisories) - - -def _loop_footprints(record: LoopFootprintMetadata) -> dict[str, str]: - return { - f"{item.buffer}@{item.memory_level}": ( - f"{item.bytes}/{item.device_bytes}/{item.repeated_bytes}" - ) - for item in record.footprints - } - - -def _loop_footprint_status(record: LoopFootprintMetadata) -> str: - return "complete" if record.known else "lower-bound" - - -def _interval(record: PerformanceMetadata) -> TripInterval: - """The occurrence's interval, with its repetition folded in.""" - timeline = record.timeline - return TripInterval(timeline.start_ns, timeline.end_ns, timeline.stride_ns, timeline.trips) - - -def _predicted_ns(record: PerformanceSummaryMetadata) -> int: - """How long the whole Function is predicted to take.""" - return record.timeline.end_ns - record.timeline.start_ns - - -def _source_span(record: SourceSpanMetadata) -> str: - """Where the expression was authored, as one location.""" - return f"{record.file}:{record.line}:{record.column}" +class Prose(str): + """Sentence text rendered as a quoted DSL string.""" @dataclass(frozen=True) class ReportIdentity(IRMetadata): - """Which program, on which machine, this report is about.""" - target: str = "" module: str = "" function: str = "" @@ -307,120 +38,227 @@ class ReportIdentity(IRMetadata): @dataclass(frozen=True) class ReportSelection(IRMetadata): - """What was asked for, and what running it took.""" - requested: tuple[str, ...] = () executed: tuple[str, ...] = () @dataclass(frozen=True) class MemorySummary(IRMetadata): - """What one function holds at its peak, per level.""" - peak_bytes: dict[str, int] = field(default_factory=dict) @dataclass(frozen=True) class AdvisorySummary(IRMetadata): - """One thing the memory walk observed that a reader should weigh.""" - text: Prose @dataclass(frozen=True) class PerformanceSummaryView(IRMetadata): - """One function's prediction, under the root the report is about. - - ``root`` is the report's own identity composed for a reader, not something - the family measured, which is why it lives here and not on - ``PerformanceSummaryMetadata``. ``waves`` is stated even when it is one: how - many passes over the machine a plan takes is a conclusion, and one wave is an - answer rather than nothing to say. It is declared with no value it says - nothing by, so the suppression rule itself stays one rule. - """ - root: str = "" predicted_ns: int = 0 waves: int = 1 -comment( - ComputeCostMetadata, - Projection("flops", dict[str, TotalAndPerUnit[dict[str, int]]], _paired_flops), - Projection("service", dict[str, TotalAndPerUnit[dict[str, int]]], _paired_service), -) -comment( - TrafficMetadata, - Projection( - "traffic", - dict[str, TotalAndPerUnit[dict[str, TrafficBytes]]], - _paired_traffic, - ), - Projection("operands", dict[str, TrafficBytes], _by_operand, opt_in=True), -) -comment( - MemoryMetadata, - Projection("peak", dict[str, int], peak_footprint), - Projection("persistent", int, _persistent_bytes, default=0), - Projection("advisories", int, _advisory_count, default=0), -) -comment( - LoopFootprintMetadata, - Projection("footprints", dict[str, str], _loop_footprints), - Projection("status", str, _loop_footprint_status), -) -comment(RooflineMetadata, "ideal_ns", "bound_by") -comment( - PerformanceMetadata, - Projection("interval", TripInterval, _interval), - family="performance", -) -comment( - PerformanceSummaryMetadata, - Projection("predicted_ns", int, _predicted_ns), - Projection("waves", int, _read("waves")), - family="performance", -) -comment(SourceSpanMetadata, Projection("span", str, _source_span), family="source") -comment(ReportIdentity, family="analysis") -comment(ReportSelection, family="selection") -comment(MemorySummary, family="peak-footprint") -comment(AdvisorySummary, family="advisory") -comment( - PerformanceSummaryView, - "root", - "predicted_ns", - Projection("waves", int, _read("waves")), - family="performance", -) +class CommentPrinter: + """Render metadata by explicit ``print_`` methods.""" + + def print(self, value): + method = getattr(self, f"print_{type(value).__name__}", None) + if method: + return method(value) + if is_dataclass(value): + return ENTRIES.join( + f"{item.name}{ENTRY}{self.print(getattr(value, item.name))}" + for item in fields(value) + ) + if isinstance(value, Mapping): + return ENTRIES.join(f"{k}{ENTRY}{self.print(v)}" for k, v in value.items()) + if isinstance(value, (tuple, list)): + return ENTRIES.join(self.print(v) for v in value) + return str(value) + + def print_Prose(self, value): + return json.dumps(str(value)) + + def print_TrafficBytes(self, value): + return f"r{value.read}{PAIR}w{value.write}" + + def print_TripInterval(self, value): + if value.trips <= 1: + return f"[{value.start},{value.end})" + offset = f"{value.stride}t+" + return f"[{offset}{value.start},{offset}{value.end}){TRIPS}{value.trips}" + + @staticmethod + def _empty_without_default(value): + if isinstance(value, Mapping | tuple | list): + return not value + return False + + def _record(self, family, values): + out = [] + for item in values: + key, value = item[:2] + default = item[2] if len(item) == 3 else _UNSET + if default is not _UNSET: + if value == default: + continue + elif self._empty_without_default(value): + continue + out.append(f"{key.replace('_', '-')}={self.print(value)}") + return FIELDS.join([family, *out]) if out else family + + def _single(self, family, value, default=_UNSET): + if default is not _UNSET and value == default: + return family + if default is _UNSET and self._empty_without_default(value): + return family + return f"{family}{FIELD}{self.print(value)}" + + def print_ComputeCostMetadata(self, record, **_): + flops_per_unit = dict(record.flops_per_unit) + ops_per_unit = dict(record.ops_per_unit) + flops = { + dtype: f"{total}{PER_UNIT}{flops_per_unit.get(dtype, 0)}" + for dtype, total in record.flops + } + ops = { + kind: f"{total}{PER_UNIT}{ops_per_unit.get(kind, 0)}" + for kind, total in record.ops + } + return self._record( + "compute-cost", + ( + ("flops", flops), + ("flops_logical", record.flops_logical), + ("ops", ops), + ("ops_logical", dict(record.ops_logical)), + ), + ) + + def print_TrafficMetadata(self, record, *, opt_in=frozenset()): + per_unit = dict(record.per_unit) + traffic = { + level: ( + f"{self.print(moved)}{PER_UNIT}" + f"{self.print(per_unit.get(level, TrafficBytes()))}" + ) + for level, moved in record.whole + } + values = [("traffic", traffic)] + if "operands" in opt_in: + last = len(record.operands) - 1 + operands = { + "result" if index == last else str(index): moved + for index, moved in enumerate(record.operands) + } + values.append(("operands", operands)) + return self._record("traffic", values) + + def print_MemoryMetadata(self, record, **_): + return self._record( + "memory", + ( + ("peak", {item.level: item.peak_bytes for item in record.footprint}), + ("persistent", sum(item.persistent_bytes for item in record.footprint), 0), + ("advisories", len(record.advisories), 0), + ), + ) + + def print_LoopFootprintMetadata(self, record, **_): + footprints = { + f"{item.buffer}@{item.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")) + ) + + def print_PerformanceMetadata(self, record, **_): + interval = TripInterval( + record.timeline.start_ns, + record.timeline.end_ns, + record.timeline.stride_ns, + record.timeline.trips, + ) + return self._single("performance", interval) + + def print_PerformanceSummaryMetadata(self, record, **_): + predicted = record.timeline.end_ns - record.timeline.start_ns + return self._record("performance", (("predicted_ns", predicted), ("waves", record.waves))) + + def print_SourceSpanMetadata(self, record, **_): + return self._single("source", f"{record.file}:{record.line}:{record.column}") + + def print_ReportIdentity(self, record, **_): + return self._record( + "analysis", + ( + ("target", record.target, ""), + ("module", record.module, ""), + ("function", record.function, ""), + ("topology", record.topology, "none"), + ), + ) + + def print_ReportSelection(self, record, **_): + return 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) + + def print_PerformanceSummaryView(self, record, **_): + return self._record( + "performance", + ( + ("root", record.root, ""), + ("predicted_ns", record.predicted_ns, 0), + ("waves", record.waves), + ), + ) + + +_PRINTER = CommentPrinter() + + +def render_comment(record, *, opt_in=frozenset()): + method = getattr(_PRINTER, f"print_{type(record).__name__}", None) + return method(record, opt_in=opt_in) if method else None + + +def peak_footprint(record): + return {item.level: item.peak_bytes for item in record.footprint} __all__ = [ - "AdvisorySummary", - "ENTRIES", + "CommentPrinter", + "Prose", + "render_comment", + "PAIR", + "PER_UNIT", "ENTRY", + "ENTRIES", "FIELD", "FIELDS", - "PAIR", "PARTS", - "PER_UNIT", - "Prose", - "RENDER", "TRIPS", - "MemorySummary", - "PerformanceSummaryView", - "Projection", - "RecordComment", "ReportIdentity", "ReportSelection", - "comment", - "comment_of", - "declared_records", - "expr_field", - "expr_fields", - "family_of", + "MemorySummary", + "AdvisorySummary", + "PerformanceSummaryView", "peak_footprint", - "render_comment", - "render_record", - "render_value", ] diff --git a/src/tilefoundry/passes/transforms/host_entry.py b/src/tilefoundry/passes/transforms/host_entry.py index b385a944b..378b8b7c6 100644 --- a/src/tilefoundry/passes/transforms/host_entry.py +++ b/src/tilefoundry/passes/transforms/host_entry.py @@ -1,27 +1,18 @@ -"""Give a device-only module the CPU entry that calls into it. +"""Synthesize a CPU entry that launches a lone CUDA ``PrimFunction``. -The entry is synthesized, never retargeted: which target a function records is -what its author wrote, and a pass that rewrites it decides in the wrong place. -Launch geometry is settled here and written into the ``Launch``, so nothing -downstream derives it a second time. A launch-provided (dynamic) grid is not -supported by the synthesized entry. +It copies device parameters and the existing static launch configuration; +dynamic grids remain unsupported. """ - from __future__ import annotations from dataclasses import replace -from tilefoundry.ir.core import Call, Constant, Tuple, Var +from tilefoundry.codegen.signature import LaunchSignature +from tilefoundry.ir.core import Var from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.tir.launch import launch_call from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.tir.shape import ShapeOf -from tilefoundry.ir.tir.stmts import LetStmt, MeshScope, Sequential -from tilefoundry.ir.tir.symbol_ref import SymbolRef -from tilefoundry.ir.types.shard.shard_layout import ShardLayout -from tilefoundry.ir.visitor import ExprWalker -from tilefoundry.passes.pass_base import ModulePass +from tilefoundry.ir.tir.stmts import Sequential from tilefoundry.target import CpuTarget, CudaTarget _DEFAULT_ENTRY_NAME = "main" @@ -38,18 +29,35 @@ def _fresh_entry_name(taken: set[str]) -> str: return candidate -def _device_callee(module: Module) -> PrimFunction: - """The one device function a synthesized entry can be written against. - - Refuses rather than guesses: a second device function would leave the - entry's parameters undecided, and a CPU function that is not the entry - means the author meant some other function to be called first. +def insert_default_host_entry(module: Module) -> Module: + """Return a module whose entry is host-callable (CPU-target). + + Return a module whose entry is host-callable (CPU-target): + - a CPU entry is already present → unchanged; + - the entry is a specialization prototype (host-only) → retarget it + to CPU in place of the original (``module.entry`` name unchanged); + - exactly one CUDA device kernel with no CPU entry → synthesize a CPU entry + whose body is a single ``Launch`` of that kernel. """ + entry_fn = module.entry_function() + if isinstance(entry_fn.target, CpuTarget): + return module + + + if any(isinstance(fn.target, CpuTarget) for fn in module.functions): raise ValueError( "insert_default_host_entry: module has a CPU function that is not " "the entry; refusing to guess the host entry" ) + + + if entry_fn.variants: + cpu_entry = replace(entry_fn, target=CpuTarget()) + new_functions = tuple( + cpu_entry if fn is entry_fn else fn for fn in module.functions + ) + return replace(module, functions=new_functions) device_fns = [ fn for fn in module.functions @@ -58,159 +66,21 @@ def _device_callee(module: Module) -> PrimFunction: if len(device_fns) != 1: raise ValueError( f"insert_default_host_entry: expected exactly one CUDA device " - f"prim_function and no CPU entry, found {len(device_fns)} device " - f"functions" + f"prim_function (or a dispatch entry) and no CPU entry, found " + f"{len(device_fns)} device functions" ) - return device_fns[0] - - -def _launch_config_from_body( - body: Sequential, -) -> tuple[tuple[int | None, int, int], tuple[int, int, int]]: - """Derive grid and block dimensions from body mesh topologies. - - CTA topology sizes multiply into ``grid.x`` and thread sizes into - ``block.x``; warp axes do not contribute. A launch-provided CTA extent - yields ``grid.x = None`` for static callers to reject. Other dimensions - remain one until a public convention exists. - """ - grid_x = 1 - block_x = 1 - cta_dynamic = False - - def _topo_dims(mesh) -> tuple[int, int]: - """Return grid and block contributions for *mesh*. - - Return ``(grid_size_contribution, block_size_contribution)`` - for *mesh*'s full topology list. - """ - nonlocal cta_dynamic - g = 1 - b = 1 - for t in mesh.topologies: - tname = t.name - size = t.size - if not isinstance(size, int): - if tname == "cta": - cta_dynamic = True - continue - raise ValueError( - f"_launch_config_from_body: topology {tname!r} has a " - f"dynamic/scalar extent ({size!r}) that cannot be converted " - f"to a static launch config; only a 'cta' level may be " - f"launch-provided" - ) - tsize = size - if tname == "cta": - g *= tsize - else: - b *= tsize - return g, b - - def _harvest_from_layout(layout) -> None: - nonlocal grid_x, block_x - if isinstance(layout, ShardLayout): - g, b = _topo_dims(layout.mesh) - grid_x = max(grid_x, g) - block_x = max(block_x, b) - - class _LayoutVisitor(ExprWalker[None]): - def _record(self, expr) -> None: - _harvest_from_layout(getattr(getattr(expr, "type", None), "layout", None)) - - def visit_Call(self, expr: Call, ctx=None) -> None: - self._record(expr) - self.visit_operands(expr, ctx) - - def visit_Tuple(self, expr: Tuple, ctx=None) -> None: - self._record(expr) - - def visit_LoopRegion(self, expr: LoopRegion, ctx=None) -> None: - self._record(expr) - - def visit_Var(self, expr: Var, ctx=None) -> None: - self._record(expr) - - def visit_Constant(self, expr: Constant, ctx=None) -> None: - self._record(expr) - - def visit_SymbolRef(self, expr: SymbolRef, ctx=None) -> None: - self._record(expr) - - def visit_ShapeOf(self, expr: ShapeOf, ctx=None) -> None: - self._record(expr) - - def default_visit(self, expr, ctx=None) -> None: - self._record(expr) - - def walk(stmt) -> None: - nonlocal grid_x, block_x - match stmt: - case MeshScope(): - g, b = _topo_dims(stmt.mesh) - grid_x = max(grid_x, g) - block_x = max(block_x, b) - walk(stmt.body) - case Sequential(): - for s in stmt.body: - walk(s) - case LetStmt(): - if hasattr(stmt, "value"): - _LayoutVisitor().visit(stmt.value) - if hasattr(stmt, "var") and getattr(stmt.var, "type", None) is not None: - _harvest_from_layout(getattr(stmt.var.type, "layout", None)) - walk(stmt.body) - - walk(body) - grid = (None, 1, 1) if cta_dynamic else (grid_x, 1, 1) - return grid, (block_x, 1, 1) - - -def _launch_config_of( - fn: PrimFunction, -) -> tuple[tuple[int | None, int, int], tuple[int, int, int]]: - """The geometry one launch of *fn* runs at. - - A specialization prototype has variants and no body of its own, so the - variants state it -- and they have to agree, because a translation unit - states its program dimensions once for every kernel it holds. - """ - if fn.variants and not fn.body.body: - by_variant = { - variant.name: _launch_config_from_body(variant.body) for variant in fn.variants - } - distinct = set(by_variant.values()) - if len(distinct) > 1: - raise ValueError( - f"insert_default_host_entry: prototype {fn.name!r} has variants " - f"launching at different geometries ({by_variant}); one " - f"translation unit runs every kernel in it at one geometry" - ) - return distinct.pop() - return _launch_config_from_body(fn.body) + device_fn = device_fns[0] -def insert_default_host_entry(module: Module) -> Module: - """Return a module whose entry is host-callable, synthesizing one if it is not. - - A CPU entry passes through. Otherwise the device entry gains a CPU entry - that launches it with its own parameters. Whether that call reads as one - launch or as a dispatch over variants is the callee's shape to state, so - a prototype and a lone kernel take the same route here. - """ - if isinstance(module.entry_function().target, CpuTarget): - return module - - device_fn = _device_callee(module) entry_params = tuple(Var(type=p.type, name=p.name) for p in device_fn.params) - grid, block = _launch_config_of(device_fn) - if grid[0] is None: + launch_signature = LaunchSignature.from_topologies(module.effective_topologies()) + if launch_signature.grid[0] is None: raise ValueError( f"insert_default_host_entry: device function {device_fn.name!r} has " f"a launch-provided (dynamic) CTA extent; the implicit host entry " f"cannot derive its grid — launch it from an explicit host entry" ) - launch = launch_call(device_fn, entry_params, grid, block) + launch = launch_call(device_fn, entry_params) name = _fresh_entry_name({fn.name for fn in module.functions}) entry = PrimFunction( name=name, @@ -222,18 +92,4 @@ def insert_default_host_entry(module: Module) -> Module: return replace(module, functions=(*module.functions, entry), entry=name) -class InsertHostEntryPass(ModulePass): - """Give a device-only module a host-callable entry. - - The only pass a TIR module needs: everything else it is made of was - settled before it became TIR. - """ - - name: str = "insert_host_entry" - requires: tuple[str, ...] = () - - def run(self, module: Module) -> Module: - return insert_default_host_entry(module) - - -__all__ = ["InsertHostEntryPass", "insert_default_host_entry"] +__all__ = ["insert_default_host_entry"] diff --git a/src/tilefoundry/visitor_registry/contexts.py b/src/tilefoundry/visitor_registry/contexts.py index 3705490bd..c641f259a 100644 --- a/src/tilefoundry/visitor_registry/contexts.py +++ b/src/tilefoundry/visitor_registry/contexts.py @@ -123,14 +123,14 @@ class VerifyContext(TypeInferContext): class CostContext(TypeInferContext): """Cost Evaluator context for one topology window. - ``topology_level=None`` exposes the types as written. A named level projects - them to what one unit of that level holds, letting the same registered - evaluator answer both global and per-unit questions. + ``level=None`` exposes the types as written. A named level projects them to + what one unit of that level holds, letting the same registered evaluator + answer both global and per-unit questions. """ selected_types: Mapping[int, Type] = field(default_factory=dict) selected_output_type: Type | None = None - topology_level: str | None = None + level: str | None = None topologies: tuple[Topology, ...] = () def type_of(self, expr: Expr) -> Type: @@ -141,18 +141,18 @@ def local_type_of(self, expr: Expr) -> Type: """Return ``expr``'s Type in this context's topology window.""" selected = self.selected_types.get(id(expr)) type_ = selected if selected is not None else expr.type - if self.topology_level is None: + if self.level is None: return type_ - return local_type_of(type_, topology_level=self.topology_level, topologies=self.topologies) + return local_type_of(type_, level=self.level, topologies=self.topologies) def local_output_type(self, call: Call) -> Type: """Return the selected candidate output in recursive-local form.""" output = self.selected_output_type if output is None: output = call.type - if self.topology_level is None: + if self.level is None: return output - return local_type_of(output, topology_level=self.topology_level, topologies=self.topologies) + return local_type_of(output, level=self.level, topologies=self.topologies) @dataclass(frozen=True) @@ -174,24 +174,16 @@ class Cost: ``flops`` groups leaf-local logical work by compute ``DType`` so one Op can report mixed work without selecting an ALU/TensorCore implementation. - ``service`` counts what is not floating point at all -- a comparison, a - select, an integer add -- by the service it asks for, because a dtype is not - a kind of work. ``traffic`` carries one entry per operand in call order with + ``ops`` counts what is not floating point at all -- a comparison, a + select, an integer add -- by the machine operation kind it asks for, because + a dtype is not a kind of work. ``traffic`` carries one entry per operand in call order with the result last, so an Op that reads part of an input says so where it knows it. No field names a hardware implementation or a memory level. """ flops: Mapping[DType, int] traffic: tuple[TrafficBytes, ...] - service: Mapping[str, int] = field(default_factory=dict) - sent: tuple[tuple[str, TrafficBytes], ...] = () - """Bytes that left the unit they were on, by the boundary they crossed. - - ``traffic`` says which storage level each operand touched. Data handed - between two units of one topology level touches the same storage at both - ends and has still gone somewhere, so it is stated against that level's - name instead, from the sending unit's own view. - """ + ops: Mapping[str, int] = field(default_factory=dict) @property def bytes(self) -> int: @@ -209,9 +201,9 @@ def __post_init__(self) -> None: raise ValueError("Cost traffic must be non-negative integers") if any( isinstance(value, bool) or not isinstance(value, int) or value < 0 - for value in self.service.values() + for value in self.ops.values() ): - raise ValueError("Cost service work must be non-negative integers") + raise ValueError("Cost ops work must be non-negative integers") __all__ = [ diff --git a/src/tilefoundry/visitor_registry/op_cost.py b/src/tilefoundry/visitor_registry/op_cost.py index 479774fc5..eb04ac85f 100644 --- a/src/tilefoundry/visitor_registry/op_cost.py +++ b/src/tilefoundry/visitor_registry/op_cost.py @@ -107,11 +107,11 @@ def _elementwise(call: Call, ctx: CostContext, *, dtype: DType | None = None) -> return Cost({result_dtype: numel(output)}, _traffic(inputs, output)) -def _serviced(call: Call, ctx: CostContext, kind: str) -> Cost: - """One result of *kind* per element, and no floating-point work at all.""" +def _ops(call: Call, ctx: CostContext, kind: str) -> Cost: + """One machine operation of *kind* per element, and no FLOP work.""" inputs = _input_types(call, ctx) output = _output_type(call, ctx) - return Cost({}, _traffic(inputs, output), {kind: numel(output)}) + return Cost({}, _traffic(inputs, output), ops={kind: numel(output)}) @register_cost_evaluator(MatMul) @@ -206,15 +206,15 @@ def _binary(call: Call, ctx: CostContext) -> Cost: A comparison over floats produces booleans, and calling that a bool FLOP asks the target for a rate no machine publishes. What it really asks for is - a predicate, which is a service a machine does state a throughput for. + a predicate, which is a machine operation with a separately stated throughput. ``_PREDICATES`` holds the kinds whose result is one, whatever the operands were. """ kind = call.target.kind if kind in _PREDICATES: - return _serviced(call, ctx, "predicate") + return _ops(call, ctx, "predicate") if _integral(call, ctx): - return _serviced(call, ctx, "integer") + return _ops(call, ctx, "integer") return _elementwise(call, ctx) @@ -236,17 +236,17 @@ def _unary(call: Call, ctx: CostContext) -> Cost: ``_SPECIAL`` holds the kinds the machine answers on its special-function unit, at a rate of its own -- a quarter of the scalar one here. Counting one of those as a single FLOP would put it on the float pipe - at four times the throughput the unit has, so it is a service rather than - arithmetic. What is left for ``flops`` is the arithmetic that really is a + at four times the throughput the unit has, so it is a separately priced + operation rather than arithmetic. What is left for ``flops`` is the arithmetic that really is a multiply or an add. """ kind = call.target.kind if kind in _SPECIAL: - return _serviced(call, ctx, "special") + return _ops(call, ctx, "special") if kind is UnaryKind.NOT: - return _serviced(call, ctx, "predicate") + return _ops(call, ctx, "predicate") if _integral(call, ctx): - return _serviced(call, ctx, "integer") + return _ops(call, ctx, "integer") return _elementwise(call, ctx) @@ -350,8 +350,7 @@ def _softmax(call: Call, ctx: CostContext) -> Cost: @register_cost_evaluator(Where) def _where(call: Call, ctx: CostContext) -> Cost: """Choosing between two values it already has is a select, not arithmetic.""" - return _serviced(call, ctx, "select") - + return _ops(call, ctx, "select") @register_cost_evaluator(LayerNorm) def _layer_norm(call: Call, ctx: CostContext) -> Cost: diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 0da3e23a7..8cd7805a9 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -14,19 +14,16 @@ import pytest -from tests.fixtures.placed.fused_boundary import FusedBoundary from tests.fixtures.placed.performance_findings import _CrossScopePerformance from tests.fixtures.placed.symbolic_offset import ( _LiteralStoreOffset, _SymbolicStoreOffset, ) -from tests.fixtures.placed.tp_all_to_all import GPUS, SENT_BYTES, TransposeShard from tilefoundry import func, module from tilefoundry.analysis import ( ComputeCostMetadata, MemoryHierarchyFacts, MemoryMetadata, - ParallelCapacityFacts, PerformanceMetadata, PerformanceServiceFacts, PerformanceSummaryMetadata, @@ -39,7 +36,6 @@ _local_duration_ns, ) from tilefoundry.analysis.errors import AnalysisError -from tilefoundry.analysis.metadata import Breakdown, Spread, shares from tilefoundry.dsl import ConstTensor, DimVar, Mesh, Tensor, Topology, tf from tilefoundry.ir.core import ( Call, @@ -178,7 +174,8 @@ def test_performance_orders_a_predecessor_materialized_in_a_child_scope() -> Non cross_scope = results["cross_scope"] timed = tuple( - (call, get_metadata(call, PerformanceMetadata)) for call in _calls(cross_scope.function) + (call, get_metadata(call, PerformanceMetadata)) + for call in _calls(cross_scope.function) ) scalar_indices = tuple( (call, record) @@ -236,14 +233,18 @@ def test_a_symbolic_store_stride_preserves_the_literal_control_result() -> None: summary = get_metadata(result.function, PerformanceSummaryMetadata) assert cost is not None and bound is not None and summary is not None observed[name] = ( - shares(cost.service, cost.topologies)["integer"], - shares(cost.service, cost.topologies, "cta")["integer"], + dict(cost.ops)["integer"], + dict(cost.ops_per_unit)["integer"], bound.ideal_ns, summary.timeline.end_ns - summary.timeline.start_ns, ) - literal_service, literal_local, literal_roofline, literal_performance = observed["literal"] - symbolic_service, symbolic_local, symbolic_roofline, symbolic_performance = observed["symbolic"] + 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 symbolic_local - literal_local == 6 assert symbolic_service - literal_service == 6 * 128 @@ -280,8 +281,12 @@ def last_axis( w: ConstTensor[(_SPLIT_HIDDEN, _SPLIT_OUT), "bf16"], ): with Mesh(("cta",), layout=(_SPLIT_GRID,), names=("unit",)) as mesh: - rows = tf.reshard(x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem") - strip = tf.reshard(w[0:_SPLIT_BLOCK, :], (_SPLIT_BLOCK, _SPLIT_OUT @ mesh.unit), "smem") + rows = tf.reshard( + x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem" + ) + strip = tf.reshard( + w[0:_SPLIT_BLOCK, :], (_SPLIT_BLOCK, _SPLIT_OUT @ mesh.unit), "smem" + ) return tf.matmul(rows, strip) @@ -295,7 +300,9 @@ def strip_major( w: ConstTensor[(_SPLIT_GRID, _SPLIT_HIDDEN, _SPLIT_PER), "bf16"], ): with Mesh(("cta",), layout=(_SPLIT_GRID,), names=("unit",)) as mesh: - rows = tf.reshard(x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem") + rows = tf.reshard( + x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem" + ) strip = tf.reshard( w[:, 0:_SPLIT_BLOCK, :], (_SPLIT_GRID @ mesh.unit, _SPLIT_BLOCK, _SPLIT_PER), @@ -324,8 +331,8 @@ def test_a_matmul_counts_its_rows_once_whichever_axis_the_mesh_split() -> None: cost = get_metadata(product, ComputeCostMetadata) summary = get_metadata(report.function, PerformanceSummaryMetadata) per_layout[name] = ( - shares(cost.flops, cost.topologies)["bf16"], - shares(cost.flops, cost.topologies, "cta")["bf16"], + dict(cost.flops)["bf16"], + dict(cost.flops_per_unit)["bf16"], summary.timeline.end_ns, ) @@ -359,7 +366,7 @@ def test_a_program_whose_buffers_have_nowhere_to_sit_is_refused() -> None: held = get_metadata( analyze(_SharedTile, unrestated, analysis="memory").function, MemoryMetadata ).footprint - assert next(item.peak_bytes for item in held if item.memory_level == "smem") == 422_400 + assert next(item.peak_bytes for item in held if item.level == "smem") == 422_400 fits = analyze( roomy, @@ -378,10 +385,10 @@ def test_a_program_whose_buffers_have_nowhere_to_sit_is_refused() -> None: analysis="memory", ) assert [ - (item.binding, item.memory_level, item.bytes, item.defined_at, item.last_used_at) + (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) + (item.binding, item.level, item.bytes, item.defined_at, item.last_used_at) for item in get_metadata(relieved.function, MemoryMetadata).lifetimes ] @@ -405,7 +412,7 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N work = next( record for record in records - if record is not None and any(shares(record.flops, record.topologies, "cta").values()) + if record is not None and any(v for _n, v in record.flops_per_unit) ) with pytest.raises( @@ -413,30 +420,18 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N match=r"^performance: selected topology level 'thread', but the target's " r"one-unit throughputs are stated for 'cta'$", ): - _local_duration_ns(work, throughput, services, topology_level="thread") + _local_duration_ns(work, throughput, services, level="thread") with pytest.raises(AnalysisError, match=r"unknown compute dtype 'f9e9m9'"): _local_duration_ns( - replace( - work, - topologies=("cta",), - flops=Breakdown((("f9e9m9", Spread(8, (8,))),)), - ), + replace(work, flops_per_unit=(("f9e9m9", 8),)), throughput, services, - topology_level="cta", + level="cta", ) crossed = TrafficMetadata( - topologies=("cta",), - storage=Breakdown( - ( - ( - throughput.bandwidth_level, - Spread(TrafficBytes(read=4096), (TrafficBytes(read=4096),)), - ), - ) - ), + per_unit=((throughput.bandwidth_level, TrafficBytes(read=4096)),) ) with pytest.raises( AnalysisError, @@ -447,73 +442,5 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N throughput, replace(services, unit_bandwidth=()), moved=crossed, - topology_level="cta", + level="cta", ) - - -def test_one_pass_states_the_work_at_every_level_and_prices_it_there() -> None: - """The numbers this program reports at each level it declares, not the relations. - - One analysis states one unit's work at every declared level, and what one - unit gets through is the device peak over however many of that unit the - device holds, so a thread is priced far slower than a CTA for the same - share. Nothing here is sharded across ``thread``, so a thread repeats its - CTA's share rather than dividing it: asking at ``thread`` unrolls that - repetition into the whole-program total and leaves the shares alone. - """ - function = FusedBoundary.entry_function() - reported = {} - for level in ("cta", "thread"): - result = analyze(FusedBoundary, function, analysis=("performance",), topology_level=level) - cost = get_metadata(result.function, ComputeCostMetadata) - summary = get_metadata(result.function, PerformanceSummaryMetadata) - reported[level] = { - "answered_at": result.topology_level, - "levels": cost.topologies, - "whole_program": shares(cost.flops, cost.topologies), - "per_cta": shares(cost.flops, cost.topologies, "cta"), - "per_thread": shares(cost.flops, cost.topologies, "thread"), - "end_ns": summary.timeline.end_ns, - } - - assert reported == { - "cta": { - "answered_at": "cta", - "levels": ("cta", "thread"), - "whole_program": {"bf16": 576, "f32": 128}, - "per_cta": {"bf16": 72, "f32": 16}, - "per_thread": {"bf16": 72, "f32": 16}, - "end_ns": 6, - }, - "thread": { - "answered_at": "thread", - "levels": ("cta", "thread"), - "whole_program": {"bf16": 576, "f32": 2112}, - "per_cta": {"bf16": 72, "f32": 16}, - "per_thread": {"bf16": 72, "f32": 16}, - "end_ns": 1849, - }, - } - capacity = FusedBoundary.resolve_target().get_facts(ParallelCapacityFacts, "gpu") - assert capacity.parallel_units == 1, "a target that never said how many cards is one card" - - -def test_a_reshard_that_changes_shards_says_what_left_the_card() -> None: - """Splitting one tensor two ways is an exchange, and it costs both resources. - - What crosses is counted against the boundary it crossed, which no storage - level names: the data is global memory at both ends and has still gone to - another card. A unit inside the boundary is not one of the parties, so it - states no share of the crossing. - """ - result = analyze( - TransposeShard, TransposeShard.entry_function(), analysis=("compute-cost", "memory") - ) - moved = get_metadata(result.function, TrafficMetadata) - - crossing = moved.communication.of("gpu") - assert crossing.total == TrafficBytes(read=SENT_BYTES * GPUS, write=SENT_BYTES * GPUS) - by_level = dict(zip(moved.topologies, crossing.per_unit, strict=True)) - assert by_level["gpu"] == TrafficBytes(read=SENT_BYTES, write=SENT_BYTES) - assert by_level["cta"] == TrafficBytes() - assert moved.communication.of("cta") is None diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index 56c13dcb4..8989cea65 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -92,7 +92,7 @@ def assert_performance_contract(result: AnalysisResult) -> None: module_target = result.module.resolve_target() throughput = module_target.get_facts(ThroughputFacts) - services = module_target.get_facts(PerformanceServiceFacts, result.topology_level) + services = module_target.get_facts(PerformanceServiceFacts) scopes = tuple(walk_scopes(build_scopes(result.module, fn))) timed = 0 for expr in collect_exprs(fn.body): @@ -105,7 +105,7 @@ def assert_performance_contract(result: AnalysisResult) -> None: throughput, services, moved=get_metadata(expr, TrafficMetadata), - topology_level=result.topology_level, + level=result.level, ) record = get_metadata(expr, PerformanceMetadata) if not duration: @@ -184,36 +184,39 @@ def _every_number_counts_something(result: AnalysisResult) -> None: """ fn = result.function for expr in (fn, *collect_exprs(fn.body)): - for record in ( - ComputeCostMetadata, - TrafficMetadata, - MemoryMetadata, - RooflineMetadata, - PerformanceMetadata, + for record, rows in ( + ( + ComputeCostMetadata, + ( + "flops", + "flops_logical", + "flops_per_unit", + "ops", + "ops_logical", + "ops_per_unit", + ), + ), + (TrafficMetadata, ()), + (MemoryMetadata, ()), + (RooflineMetadata, ()), + (PerformanceMetadata, ()), ): held = get_metadata(expr, record) if held is None: continue - if record is ComputeCostMetadata: - for field in ("flops", "service"): - for kind, spread in getattr(held, field).kinds: - for where, value in ( - ("total", spread.total), - *zip(held.topologies, spread.per_unit, strict=False), - ): - assert value >= 0, ( - f"{describe_expr(expr)}: {field}[{kind}].{where} = {value}" - ) + for field in rows: + value = getattr(held, field) + if field == "flops_logical": + assert value >= 0, f"{describe_expr(expr)}: {field} = {value}" + continue + for name, value in value: + assert value >= 0, f"{describe_expr(expr)}: {field}[{name}] = {value}" if record is TrafficMetadata: - for field in ("storage", "communication"): - for level, spread in getattr(held, field).kinds: - for where, moved in ( - ("total", spread.total), - *zip(held.topologies, spread.per_unit, strict=False), - ): - assert moved.read >= 0 and moved.write >= 0, ( - f"{describe_expr(expr)}: {field}[{level}].{where} = {moved}" - ) + 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}" + ) for position, moved in enumerate(held.operands): assert moved.read >= 0 and moved.write >= 0, ( f"{describe_expr(expr)}: operand {position} = {moved}" @@ -232,7 +235,7 @@ def _every_number_counts_something(result: AnalysisResult) -> None: record = get_metadata(expr, LoopFootprintMetadata) if record is None: continue - rows = [(item.buffer, item.memory_level) for item in record.footprints] + 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: @@ -270,9 +273,7 @@ def test_every_analysis_runs_at_a_stated_size(family: str) -> None: def _predicted_ns(module, dims=None) -> int: """What the four families together say one program takes.""" - result = analyze( - module, module.entry_function(), analysis=FAMILIES, topology_level="cta", dims=dims - ) + result = analyze(module, module.entry_function(), analysis=FAMILIES, level="cta", dims=dims) summary = get_metadata(result.function, PerformanceSummaryMetadata) assert summary is not None return summary.timeline.end_ns - summary.timeline.start_ns diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 7345d588a..147ea1e4a 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -166,7 +166,8 @@ def test_analyze_help_explains_topology_effects_and_assumptions(capsys) -> None: help_text = capsys.readouterr().out for family in ("compute-cost", "memory", "roofline", "performance"): assert family in help_text - assert "every level's per-unit share" in help_text + assert "flops_per_unit" in help_text + assert "per-unit traffic" in help_text assert "global traffic is the device's and counted once" in help_text assert "is an observation, not a bound" in help_text @@ -356,7 +357,7 @@ def test_named_provider_registers_only_decorated_targets_and_replays_them( "from tilefoundry.target import Target, register_target\n" "@dataclass(frozen=True)\n" "class _VendorBase(Target):\n" - " pass\n" + " topology_levels = ('core',)\n" "@register_target\n" "@dataclass(frozen=True)\n" "class VendorOne(_VendorBase):\n" @@ -679,18 +680,6 @@ def test_analyze_binds_an_extent_on_a_root_that_reaches_a_child(tmp_path, capsys assert "leaf(" not in expanded -def _shares_text(per_unit: list, topologies: list) -> str: - """One quantity's per-level shares, as the comment renders a dict.""" - return ",".join(f"{name}:{value}" for name, value in zip(topologies, per_unit)) - - -def _bytes_shares_text(per_unit: list, topologies: list) -> str: - """The same for read/write pairs.""" - return ",".join( - f"{name}:r{moved['read']}/w{moved['write']}" for name, moved in zip(topologies, per_unit) - ) - - def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> None: source = Path(__file__).parents[1] / "fixtures" / "placed" / "moe_mega_kernel.py" selector = f"{source}:MoEMegaKernel" @@ -757,27 +746,21 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> f"# selection requested={','.join(payload['requested'])} " f"executed={','.join(payload['executed'])}", "# compute-cost " - f"flops=f32:{cost['flops']['f32']['total']}" - f"@{_shares_text(cost['flops']['f32']['per_unit'], cost['topologies'])}", + f"flops=f32:{cost['flops']['f32']}@{cost['flops_per_unit']['f32']}" + f" flops-logical={cost['flops_logical']}", "# traffic " - f"traffic=gmem:r{moved['storage']['gmem']['total']['read']}" - f"/w{moved['storage']['gmem']['total']['write']}" - f"@{_bytes_shares_text(moved['storage']['gmem']['per_unit'], moved['topologies'])}", + f"traffic=gmem:r{moved['whole']['gmem']['read']}" + f"/w{moved['whole']['gmem']['write']}" + f"@r{moved['per_unit']['gmem']['read']}" + f"/w{moved['per_unit']['gmem']['write']}", f"# peak-footprint=gmem:{peak[0]['peak_bytes']}", f"# roofline ideal-ns={bound['ideal_ns']} bound-by={bound['bound_by']}", "# performance root=MoEMegaKernel::experts " f"predicted-ns={summary['timeline']['end_ns']} " f"waves={summary['waves']}", ] - assert payload["totals"]["flops"] == { - kind: value["total"] for kind, value in cost["flops"].items() - } - assert payload["totals"]["traffic"] == { - level: value["total"] for level, value in moved["storage"].items() - } - assert payload["totals"]["communication"] == { - level: value["total"] for level, value in moved["communication"].items() - } + assert payload["totals"]["flops"] == cost["flops"] + assert payload["totals"]["traffic"] == moved["whole"] hoisted = {line.split(" = ", 1)[0] for line in lines if " = Mesh((Topology(" in line} scoped = { diff --git a/tests/ops/ir/cost_utils.py b/tests/ops/ir/cost_utils.py index 73f469357..ca3f7a7cd 100644 --- a/tests/ops/ir/cost_utils.py +++ b/tests/ops/ir/cost_utils.py @@ -20,7 +20,7 @@ class CostCase: op: object inputs: tuple[TensorType, ...] flops: Mapping[DType, int] = field(default_factory=dict) - service: Mapping[str, int] = field(default_factory=dict) + ops: Mapping[str, int] = field(default_factory=dict) traffic: tuple[TrafficBytes, ...] = () topology_level: str | None = None topologies: tuple[Topology, ...] = () @@ -44,5 +44,5 @@ def run_cost_case(case: CostCase) -> None: cost = CostEvaluator().visit_Call(call, ctx) assert cost.flops == case.flops - assert cost.service == case.service + assert cost.ops == case.ops assert cost.traffic == case.traffic diff --git a/tests/ops/ir/test_binary.py b/tests/ops/ir/test_binary.py index 7ca0c1984..6c19623b7 100644 --- a/tests/ops/ir/test_binary.py +++ b/tests/ops/ir/test_binary.py @@ -173,7 +173,7 @@ def test_lower_rank_split_right_aligns(): "lt_over_reals_is_a_predicate", Binary(kind=BinaryKind.LT), (_REAL, _REAL), - service={"predicate": 8}, + ops={"predicate": 8}, traffic=(TrafficBytes(read=32), TrafficBytes(read=32), TrafficBytes(write=1)), ), CostCase( diff --git a/tests/ops/ir/test_unary.py b/tests/ops/ir/test_unary.py index 14756c8ec..8013e8032 100644 --- a/tests/ops/ir/test_unary.py +++ b/tests/ops/ir/test_unary.py @@ -101,7 +101,7 @@ def test_a_unary_evaluates_to_its_oracle(name, op, drawn, oracle): "exp_is_special_function_work", Unary(kind=UnaryKind.EXP), (_REAL,), - service={"special": 8}, + ops={"special": 8}, traffic=(TrafficBytes(read=32), TrafficBytes(write=32)), ), CostCase( diff --git a/tests/ops/ir/test_where.py b/tests/ops/ir/test_where.py index baf2045ac..facda266f 100644 --- a/tests/ops/ir/test_where.py +++ b/tests/ops/ir/test_where.py @@ -119,7 +119,7 @@ def test_where_cost_counts_selection_and_materialization(): make_tensor_type((2, 3, 4), DType.f32), make_tensor_type((2, 3, 4), DType.f32), ), - service={"select": 24}, + ops={"select": 24}, traffic=( TrafficBytes(read=3), TrafficBytes(read=96), From 9351eb091452e299f6658d8582e6e1098ab76bd8 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 15 Sep 2026 13:49:47 +0800 Subject: [PATCH 3/7] refactor(codegen): close review follow-ups --- docs/spec/analysis.md | 33 +- docs/tutorial/showcase.ipynb | 16 +- docs/tutorial/showcase.md | 30 +- src/tilefoundry/analysis/api.py | 9 + src/tilefoundry/analysis/compute_cost.py | 295 +++++----- src/tilefoundry/analysis/facts.py | 17 +- src/tilefoundry/analysis/memory.py | 5 +- src/tilefoundry/analysis/metadata.py | 176 ++++-- src/tilefoundry/analysis/report.py | 12 +- src/tilefoundry/analysis/roofline.py | 6 +- src/tilefoundry/cli/analyze.py | 17 +- src/tilefoundry/codegen/context.py | 45 +- src/tilefoundry/codegen/context_builder.py | 312 +++++++++++ src/tilefoundry/codegen/cpu/abi.py | 6 +- src/tilefoundry/codegen/cpu/context.py | 14 +- src/tilefoundry/codegen/cpu/module.py | 517 ++++++------------ src/tilefoundry/codegen/cuda/abi.py | 27 +- src/tilefoundry/codegen/cuda/context.py | 15 +- src/tilefoundry/codegen/cuda/module.py | 247 +++++---- .../codegen/cuda/tir/prim_function.py | 292 +++------- .../codegen/cuda/tir/stmts/abort.py | 11 - .../codegen/cuda/tir/stmts/for_.py | 24 - src/tilefoundry/codegen/cuda/tir/stmts/if_.py | 25 - .../codegen/cuda/tir/stmts/let_stmt.py | 30 - .../codegen/cuda/tir/stmts/return_.py | 13 - .../codegen/cuda/tir/stmts/sequential.py | 13 - .../codegen/cuda/tir/stmts/while_.py | 15 - src/tilefoundry/codegen/emitter.py | 16 +- src/tilefoundry/codegen/registry.py | 49 +- src/tilefoundry/codegen/signature.py | 27 +- src/tilefoundry/compile.py | 95 ++-- src/tilefoundry/inspection/values.py | 56 +- .../passes/transforms/host_entry.py | 107 ++-- src/tilefoundry/visitor_registry/contexts.py | 41 +- src/tilefoundry/visitor_registry/op_cost.py | 25 +- tests/analysis/test_analysis_families.py | 51 +- tests/analysis/test_analyze_at_a_size.py | 20 +- tests/analysis/test_analyze_cross_module.py | 4 +- tests/cli/test_cli.py | 23 +- tests/codegen/test_codegen_registry.py | 100 ---- tests/codegen/test_device_dispatch.py | 24 +- tests/codegen/test_host_multi_launch.py | 13 +- tests/codegen/test_linkable.py | 10 +- tests/codegen/test_symbol_table.py | 98 ---- tests/fixtures/tir/async_sync.py | 2 +- tests/fixtures/tir/square.py | 2 +- tests/fixtures/tir/sync.py | 2 +- tests/installed/models/smoke_qwen3_1_7b.py | 6 +- tests/ops/ir/test_cache_update.py | 2 +- tests/ops/ir/test_local.py | 2 +- tests/ops/tir/cuda/test_mma.py | 24 +- tests/ops/tir/cuda/test_tma.py | 37 +- tests/ops/tir/test_copy.py | 26 +- tests/ops/tir/test_sync.py | 8 +- tests/passes/test_host_entry.py | 28 - tests/target/test_target.py | 45 +- 56 files changed, 1459 insertions(+), 1706 deletions(-) create mode 100644 src/tilefoundry/codegen/context_builder.py delete mode 100644 src/tilefoundry/codegen/cuda/tir/stmts/abort.py delete mode 100644 src/tilefoundry/codegen/cuda/tir/stmts/for_.py delete mode 100644 src/tilefoundry/codegen/cuda/tir/stmts/if_.py delete mode 100644 src/tilefoundry/codegen/cuda/tir/stmts/let_stmt.py delete mode 100644 src/tilefoundry/codegen/cuda/tir/stmts/return_.py delete mode 100644 src/tilefoundry/codegen/cuda/tir/stmts/sequential.py delete mode 100644 src/tilefoundry/codegen/cuda/tir/stmts/while_.py delete mode 100644 tests/codegen/test_codegen_registry.py delete mode 100644 tests/codegen/test_symbol_table.py diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 0e1fd9a1e..9a4655001 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -166,35 +166,26 @@ target hardware facts. What an occurrence moves is the memory family's answer ```python class ComputeCostMetadata(IRMetadata): - """One Call's logical work, as the authored program states it. - - Attributes: - topologies: attribute; The declared levels, in the order per_unit states them. - flops: attribute; Flop counts per compute DType name. - flops_logical: attribute; Flop counts before execution-position replication. - service: attribute; Result counts per service kind. - service_logical: attribute; Service counts before execution-position replication. - """ + """Typed work in logical, expanded, and topology-unit domains.""" topologies: tuple[str, ...] = () - flops: Breakdown[int] = Breakdown() - flops_logical: tuple[tuple[str, int], ...] = () - service: Breakdown[int] = Breakdown() - service_logical: tuple[tuple[str, int], ...] = () + flops: Spread[tuple[tuple[str, int], ...]] = Spread((), (), ()) + other_ops: Spread[tuple[tuple[str, int], ...]] = Spread((), (), ()) ``` + | Field | How it is computed | Reads the target | |---|---|---| | `topologies` | The effective Module topology levels, coarsest first. | No | | `flops` | For a primitive Call, run its registered cost evaluator over operand and result Types as written; the total then multiplies by the enclosing recomputation factor and the number of positions in its execution scope, and each level's share is the same evaluator over Types projected through authored `Split`s at or coarser than that level. For a Function Call, take the callee's summed record and multiply by the call site's factor. | No; projection reads resolved Mesh and effective Module topology extents. | -| `service` | The same evaluator's service counts -- the results it asks a machine for that are not floating point -- totalled and shared the same way. | No; projection reads resolved Mesh and effective Module topology extents. | +| `other_ops` | The evaluator's non-floating-point operation counts (`integer`, `predicate`, `select`, `special`), totalled and shared the same way. These keys map directly to target one-unit operation-throughput keys; target-side service naming is unchanged. | No; projection reads resolved Mesh and effective Module topology extents. | Requesting this family adds one summary line, prefixed by `# `: the Function's own record, stated exactly as a Call's is. The whole program's work is not a second record. ```text -compute-cost flops=:@[,...] service=:@[,...] +compute-cost flops=:/@:[,...] other-ops=:/@:[,...] ``` Every measured Call receives this annotation. Each key pairs the whole quantity @@ -204,8 +195,10 @@ Each reported Call's JSON projection is under its `compute-cost` key: ```text {"topologies": [, ...], - "flops": {: {"total": , "per_unit": [, ...]}}, - "service": {: {"total": , "per_unit": [, ...]}}} + "flops": {"logical": {: }, "total": {: }, + "per_unit": [{: }, ...]}, + "other_ops": {"logical": {: }, "total": {: }, + "per_unit": [{: }, ...]}} ``` - constraints: @@ -238,10 +231,12 @@ class Spread[V]: and a reader who knew to ask. Attributes: - total: attribute; What the whole program asks for. - per_unit: attribute; What one unit of each level holds, in the record's topologies order. + logical: attribute; What the authored operation asks for before replication. + total: attribute; What the whole execution asks for. + per_unit: attribute; What one unit of each level holds, in topology order. """ + logical: V total: V per_unit: tuple[V, ...] = () diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index a7aa768fa..62e6201d5 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@cta:328896,f32:200448@cta:200448 service=special:1024@cta:1024\n# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592\n# peak-footprint=gmem:1675788\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@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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@cta:r136/w128 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@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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/328896@cta:328896,f32:200448/200448@cta:200448 other-ops=special:1024/1024@cta:1024\n# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592\n# peak-footprint=gmem:1675788\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/131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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@cta:r136/w128 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/131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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" @@ -137,7 +137,7 @@ { "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@cta:200448` | `gmem:r2225620/w806592@cta:r2225620/w806592` | 1675788 | 632 | memory |\n| 512 | `799488@cta:799488` | `gmem:r4744660/w3202752@cta:r4744660/w3202752` | 2572812 | 1656 | memory |\n| 1024 | `1598208@cta:1598208` | `gmem:r8103380/w6397632@cta:r8103380/w6397632` | 3768844 | 3022 | memory |\n| 2048 | `3195648@cta:3195648` | `gmem:r14820820/w12787392@cta:r14820820/w12787392` | 6160908 | 5752 | memory |\n| 4096 | `6390528@cta:6390528` | `gmem:r28255700/w25566912@cta:r28255700/w25566912` | 10945036 | 11214 | memory |\n| 8192 | `12780288@cta:12780288` | `gmem:r55125460/w51125952@cta:r55125460/w51125952` | 20513292 | 22136 | memory |\n" + "text": "| `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound |\n|---:|---:|---|---:|---:|---|\n| 128 | `200448/200448@cta:200448` | `gmem:r2225620/w806592@cta:r2225620/w806592` | 1675788 | 632 | memory |\n| 512 | `799488/799488@cta:799488` | `gmem:r4744660/w3202752@cta:r4744660/w3202752` | 2572812 | 1656 | memory |\n| 1024 | `1598208/1598208@cta:1598208` | `gmem:r8103380/w6397632@cta:r8103380/w6397632` | 3768844 | 3022 | memory |\n| 2048 | `3195648/3195648@cta:3195648` | `gmem:r14820820/w12787392@cta:r14820820/w12787392` | 6160908 | 5752 | memory |\n| 4096 | `6390528/6390528@cta:6390528` | `gmem:r28255700/w25566912@cta:r28255700/w25566912` | 10945036 | 11214 | memory |\n| 8192 | `12780288/12780288@cta:12780288` | `gmem:r55125460/w51125952@cta:r55125460/w51125952` | 20513292 | 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" @@ -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:2629376@cta:328672,f32:2833728@cta:354216 service=special:14528@cta:1816\n# traffic traffic=gmem:r5563796/w3721856@cta:r3936212/w3721408,smem:r9597248/w9480000@cta:r1199656/w1185000\n# peak-footprint=gmem:3701260,smem:472160\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/2629376@cta:328672,f32:2833728/2833728@cta:354216 other-ops=special:14528/14528@cta:1816\n# traffic traffic=gmem:r5563796/w3721856@cta:r3936212/w3721408,smem:r9597248/w9480000@cta:r1199656/w1185000\n# peak-footprint=gmem:3701260,smem:472160\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" @@ -274,7 +274,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# compute-cost flops=bf16:328896@cta:328896,f32:200448@cta:200448 service=special:1024@cta:1024\n# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592\n# peak-footprint=gmem:1675788\n# roofline ideal-ns=632 bound-by=memory\n" + "text": "# compute-cost flops=bf16:328896/328896@cta:328896,f32:200448/200448@cta:200448 other-ops=special:1024/1024@cta:1024\n# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592\n# peak-footprint=gmem:1675788\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" @@ -309,7 +309,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# compute-cost flops=bf16:2629376@cta:328672,f32:200448@cta:25056 service=special:1024@cta:128\n# traffic traffic=gmem:r1674644/w264832@cta:r1559508/w264384,smem:r684608/w675392@cta:r85576/w84424\n# peak-footprint=gmem:1540620,smem:33280\n# roofline ideal-ns=405 bound-by=memory\n" + "text": "# compute-cost flops=bf16:328896/2629376@cta:328672,f32:200448/200448@cta:25056 other-ops=special:1024/1024@cta:128\n# traffic traffic=gmem:r1674644/w264832@cta:r1559508/w264384,smem:r684608/w675392@cta:r85576/w84424\n# peak-footprint=gmem:1540620,smem:33280\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" @@ -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:4392896@cta:328672,f32:6418944@cta:200592 service=integer:288@cta:9,special:33152@cta:1036\n# traffic traffic=gmem:r3476884/w4196992@cta:r2558932/w4196544,rmem:r656/w72@cta:r656/w72,smem:r5839296/w5662784@cta:r682464/w672676\n# peak-footprint=gmem:5047436,rmem:16,smem:33408\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@cta:r136/w128 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/4392896@cta:328672,f32:3239808/6418944@cta:200592 other-ops=integer:9/288@cta:9,special:33056/33152@cta:1036\n# traffic traffic=gmem:r3476884/w4196992@cta:r2558932/w4196544,rmem:r656/w72@cta:r656/w72,smem:r5839296/w5662784@cta:r682464/w672676\n# peak-footprint=gmem:5047436,rmem:16,smem:33408\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@cta:r136/w128 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:337408@cta:42176,f32:51124224@cta:6390528 service=special:262144@cta:32768\n# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144\n# peak-footprint=gmem:10945036,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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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:337408@cta:42176,f32:51124224@cta:6390528 service=special:262144@cta:32768\n# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144\n# peak-footprint=gmem:10945036,smem:16960\n# roofline ideal-ns=11213 bound-by=memory\n\n v1 = reshard(w_q, layout=ShardLayout(\n layout=Layout((1, 256, 8, 32), None),\n attrs=(S(2),),\n mesh=cta,\n ), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n\n v42 = reshard(w_o, layout=ShardLayout(\n layout=Layout((1, 256, 8, 32), None),\n attrs=(S(2),),\n mesh=cta,\n ), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\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:1246400@cta:328672,f32:6410280@cta:801285 service=integer:256@cta:32,special:33040@cta:4130\n# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804\n# peak-footprint=gmem:2950412,rmem:0,smem:33536\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@cta:r32/w0 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@cta:r136/w128 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:1246400@cta:328672,f32:6410280@cta:801285 service=integer:256@cta:32,special:33040@cta:4130\n# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804\n# peak-footprint=gmem:2950412,rmem:0,smem:33536\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@cta:r32/w0 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@cta:r136/w128 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" diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index 45bb5856d..c9331af5b 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -197,14 +197,14 @@ for needle in ("matmul(hidden, w_q", "cache_update(k_cache", "matmul(v33, w_o"): ```text # 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@cta:328896,f32:200448@cta:200448 service=special:1024@cta:1024 +# compute-cost flops=bf16:328896/328896@cta:328896,f32:200448/200448@cta:200448 other-ops=special:1024/1024@cta:1024 # traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592 # peak-footprint=gmem:1675788 # 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@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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/131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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@cta:r136/w128 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@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=28 bound-by=memory + v34 = matmul(v33, w_o, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072/131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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 @@ -259,12 +259,12 @@ for ctx_len in (128, 512, 1024, 2048, 4096, 8192): | `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound | |---:|---:|---|---:|---:|---| -| 128 | `200448@cta:200448` | `gmem:r2225620/w806592@cta:r2225620/w806592` | 1675788 | 632 | memory | -| 512 | `799488@cta:799488` | `gmem:r4744660/w3202752@cta:r4744660/w3202752` | 2572812 | 1656 | memory | -| 1024 | `1598208@cta:1598208` | `gmem:r8103380/w6397632@cta:r8103380/w6397632` | 3768844 | 3022 | memory | -| 2048 | `3195648@cta:3195648` | `gmem:r14820820/w12787392@cta:r14820820/w12787392` | 6160908 | 5752 | memory | -| 4096 | `6390528@cta:6390528` | `gmem:r28255700/w25566912@cta:r28255700/w25566912` | 10945036 | 11214 | memory | -| 8192 | `12780288@cta:12780288` | `gmem:r55125460/w51125952@cta:r55125460/w51125952` | 20513292 | 22136 | memory | +| 128 | `200448/200448@cta:200448` | `gmem:r2225620/w806592@cta:r2225620/w806592` | 1675788 | 632 | memory | +| 512 | `799488/799488@cta:799488` | `gmem:r4744660/w3202752@cta:r4744660/w3202752` | 2572812 | 1656 | memory | +| 1024 | `1598208/1598208@cta:1598208` | `gmem:r8103380/w6397632@cta:r8103380/w6397632` | 3768844 | 3022 | memory | +| 2048 | `3195648/3195648@cta:3195648` | `gmem:r14820820/w12787392@cta:r14820820/w12787392` | 6160908 | 5752 | memory | +| 4096 | `6390528/6390528@cta:6390528` | `gmem:r28255700/w25566912@cta:r28255700/w25566912` | 10945036 | 11214 | memory | +| 8192 | `12780288/12780288@cta:12780288` | `gmem:r55125460/w51125952@cta:r55125460/w51125952` | 20513292 | 22136 | memory | The table says: @@ -426,7 +426,7 @@ print(report.partition("\n\n")[0].rstrip()) ```text # 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:2629376@cta:328672,f32:2833728@cta:354216 service=special:14528@cta:1816 +# compute-cost flops=bf16:328896/2629376@cta:328672,f32:2833728/2833728@cta:354216 other-ops=special:14528/14528@cta:1816 # traffic traffic=gmem:r5563796/w3721856@cta:r3936212/w3721408,smem:r9597248/w9480000@cta:r1199656/w1185000 # peak-footprint=gmem:3701260,smem:472160 # roofline ideal-ns=1935 bound-by=memory @@ -554,7 +554,7 @@ for line in report.splitlines(): ``` ```text -# compute-cost flops=bf16:328896@cta:328896,f32:200448@cta:200448 service=special:1024@cta:1024 +# compute-cost flops=bf16:328896/328896@cta:328896,f32:200448/200448@cta:200448 other-ops=special:1024/1024@cta:1024 # traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592 # peak-footprint=gmem:1675788 # roofline ideal-ns=632 bound-by=memory @@ -581,7 +581,7 @@ for line in report.splitlines(): ``` ```text -# compute-cost flops=bf16:2629376@cta:328672,f32:200448@cta:25056 service=special:1024@cta:128 +# compute-cost flops=bf16:328896/2629376@cta:328672,f32:200448/200448@cta:25056 other-ops=special:1024/1024@cta:128 # traffic traffic=gmem:r1674644/w264832@cta:r1559508/w264384,smem:r684608/w675392@cta:r85576/w84424 # peak-footprint=gmem:1540620,smem:33280 # roofline ideal-ns=405 bound-by=memory @@ -728,7 +728,7 @@ print(next(line.rstrip() for line in annotated.splitlines() if "cache_update(k_c ```text # 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:4392896@cta:328672,f32:6418944@cta:200592 service=integer:288@cta:9,special:33152@cta:1036 +# compute-cost flops=bf16:328896/4392896@cta:328672,f32:3239808/6418944@cta:200592 other-ops=integer:9/288@cta:9,special:33056/33152@cta:1036 # traffic traffic=gmem:r3476884/w4196992@cta:r2558932/w4196544,rmem:r656/w72@cta:r656/w72,smem:r5839296/w5662784@cta:r682464/w672676 # peak-footprint=gmem:5047436,rmem:16,smem:33408 # roofline ideal-ns=1599 bound-by=memory @@ -865,7 +865,7 @@ for needle in ("reshard(w_q", "reshard(w_o"): ```text # 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:337408@cta:42176,f32:51124224@cta:6390528 service=special:262144@cta:32768 +# compute-cost flops=bf16:328896/337408@cta:42176,f32:6390528/51124224@cta:6390528 other-ops=special:32768/262144@cta:32768 # traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144 # peak-footprint=gmem:10945036,smem:16960 # roofline ideal-ns=11213 bound-by=memory @@ -1017,7 +1017,7 @@ for needle in ("slice(k_cache", "cache_update(k_cache"): ```text # 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:1246400@cta:328672,f32:6410280@cta:801285 service=integer:256@cta:32,special:33040@cta:4130 +# compute-cost flops=bf16:328896/1246400@cta:328672,f32:6410280/6410280@cta:801285 other-ops=integer:32/256@cta:32,special:33040/33040@cta:4130 # traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804 # peak-footprint=gmem:2950412,rmem:0,smem:33536 # roofline ideal-ns=2474 bound-by=memory diff --git a/src/tilefoundry/analysis/api.py b/src/tilefoundry/analysis/api.py index f61e5df01..e1ef0ea25 100644 --- a/src/tilefoundry/analysis/api.py +++ b/src/tilefoundry/analysis/api.py @@ -46,6 +46,11 @@ class AnalysisResult: executed: tuple[str, ...] metadata_types: tuple[type[IRMetadata], ...] + @property + def level(self) -> str | None: + """The selected topology level.""" + return self.topology_level + def _algorithm(target: Target, selector: str, *, root: str) -> Analyzer: """The service selected by the resolved Target for *selector*.""" @@ -107,6 +112,7 @@ def analyze( *, analysis: str | Sequence[str], topology_level: str | None = None, + level: str | None = None, options: object | None = None, dims: "Mapping[str, int] | None" = None, ) -> AnalysisResult: @@ -118,6 +124,9 @@ def analyze( variant owned by the module, and the result identifies the concrete inlined view that received records. """ + if topology_level is not None and level is not None and topology_level != level: + raise ValueError("analyze: topology_level and level disagree") + topology_level = topology_level if topology_level is not None else level if not isinstance(module, Module): raise TypeError( f"analyze: expected a Module, got {type(module).__name__}. A Function " diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index 6a7787806..7c760d92f 100644 --- a/src/tilefoundry/analysis/compute_cost.py +++ b/src/tilefoundry/analysis/compute_cost.py @@ -1,11 +1,4 @@ -"""How much work the authored program asks for. - -This family reads the program and nothing else. Flops and typed machine ops -come from each op's registered cost evaluator, so the record it leaves is the same on -every backend. What that work moves is the memory family's half of the same -declaration, and what it costs in time is a separate question again, asked -against a target's rates. -""" +"""Count typed floating-point and other operations in authored HIR.""" from __future__ import annotations @@ -17,43 +10,42 @@ from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.types import DType -from tilefoundry.ir.types.shard import Mesh, composed, level_axes +from tilefoundry.ir.types.shard import Mesh, composed, topology_axes from tilefoundry.ir.types.shard.mesh import _positions_layout from tilefoundry.ir.visitor import ExprVisitor -from tilefoundry.visitor_registry.contexts import ( - CostContext, - FunctionScope, -) +from tilefoundry.visitor_registry.contexts import CostContext, FunctionScope, TrafficBytes from tilefoundry.visitor_registry.visitors import CostEvaluator from .errors import AnalysisError from .facts import PerformanceServiceFacts, ThroughputFacts -from .metadata import ComputeCostMetadata +from .metadata import Breakdown, ComputeCostMetadata, Spread, shares from .visitor import AnalyzeContext SELECTOR = "compute-cost" +def _counts(values: dict) -> tuple[tuple[str, int], ...]: + return tuple(sorted((getattr(kind, "name", kind), value) for kind, value in values.items())) + + +def _at(spread: Spread, topologies: tuple[str, ...], level: str | None): + index = topologies.index(level) if level in topologies else None + return spread.total if index is None else spread.at(index) + + def _is_structural_occurrence( cost: ComputeCostMetadata, moved: "TrafficMetadata | None" = None, *, + unit: str, bandwidth_level: str | None = None, ) -> bool: - """Whether an occurrence asks for nothing this model puts on a clock. - - Only what could take time is counted: the flops, the typed machine ops, and the - bytes at the one level a bandwidth is published for. Movement at any other - level is still movement and still recorded -- what it is not is work this - model can lay on a timeline, so it neither earns a duration nor asks for a - placement to be laid at. Having moved bytes and having timed work are - different questions, and this is the second one. - """ + """Whether an occurrence asks for nothing this model puts on a clock.""" return ( - all(not value for _name, value in cost.flops_per_unit) - and all(not value for _kind, value in cost.ops_per_unit) + 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 ( - moved.per_unit_at(bandwidth_level).total_bytes + _bytes(moved.storage, moved.topologies, bandwidth_level, unit) if moved is not None and bandwidth_level is not None else 0 ) @@ -66,29 +58,26 @@ def _local_duration_ns( services: PerformanceServiceFacts, *, moved: "TrafficMetadata | None" = None, - level: str, + topology_level: str | None = None, + level: str | None = None, scale: int = 1, ) -> int: - """Price one occurrence's projected work against one unit's throughputs. - - Compute and movement overlap within one occurrence, so its duration is - whichever side takes longer. Work with no stated throughput is refused - rather than priced at nothing: a number with a hole in it reads as a program - that does less than it does. Movement at a level the target publishes no - bandwidth for is a different case -- it is stated and left untimed, because - a rate nobody published is not one this may invent. - """ - if services.unit != level: + """Price one occurrence's projected work against one unit's throughputs.""" + topology_level = topology_level if topology_level is not None else level + if topology_level is None: + raise AnalysisError("performance: no topology level was selected") + if services.unit != topology_level: raise AnalysisError( - f"performance: selected topology level {level!r}, but the target's " - f"one-unit throughputs are stated for {services.unit!r}" + f"performance: selected topology level {topology_level!r}, but the " + f"target's one-unit throughputs are stated for {services.unit!r}" ) - - if _is_structural_occurrence(cost, moved, bandwidth_level=facts.bandwidth_level): + if _is_structural_occurrence( + cost, moved, unit=topology_level, bandwidth_level=facts.bandwidth_level + ): return 0 compute_ns = 0 - for name, value in cost.flops_per_unit: + for name, value in _at(cost.flops, cost.topologies, topology_level): if not value: continue dtype = getattr(DType, name, None) @@ -97,22 +86,26 @@ def _local_duration_ns( throughput = services.flops(dtype) if throughput is None or throughput <= 0: raise AnalysisError( - f"performance: target states no one-unit throughput for dtype {name!r} at {level!r}" + f"performance: target states no one-unit throughput for dtype " + f"{name!r} at {topology_level!r}" ) compute_ns += -(-(value * scale * 1_000_000_000) // throughput) - for kind, value in cost.ops_per_unit: + for kind, value in _at(cost.other_ops, cost.topologies, topology_level): if not value: continue throughput = services.ops(kind) if throughput is None or throughput <= 0: raise AnalysisError( - f"performance: target states no one-unit throughput for {kind!r} work at {level!r}" + f"performance: target states no one-unit throughput for {kind!r} " + f"work at {topology_level!r}" ) compute_ns += -(-(value * scale * 1_000_000_000) // throughput) crossed = ( - moved.per_unit_at(facts.bandwidth_level).total_bytes * scale if moved is not None else 0 + _bytes(moved.storage, moved.topologies, facts.bandwidth_level, topology_level) * scale + if moved is not None + else 0 ) memory_ns = 0 if crossed: @@ -120,56 +113,67 @@ def _local_duration_ns( if throughput is None or throughput <= 0: raise AnalysisError( f"performance: target states no one-unit throughput for level " - f"{facts.bandwidth_level!r} at {level!r}" + f"{facts.bandwidth_level!r} at {topology_level!r}" ) memory_ns = -(-(crossed * 1_000_000_000) // throughput) - return max(compute_ns, memory_ns) - -def _flops(flops: dict) -> tuple[tuple[str, int], ...]: - return tuple(sorted((dtype.name, value) for dtype, value in flops.items())) + sent = ( + _bytes(moved.communication, moved.topologies, topology_level, topology_level) * scale + if moved is not None + else 0 + ) + link_ns = 0 + if sent: + rate = services.bandwidth(topology_level) + if rate: + link_ns = -(-(sent * 1_000_000_000) // rate) + return max(compute_ns, memory_ns, link_ns) def _call_cost_record( - expr: Call, whole: CostContext, local: CostContext, executing_positions: int + expr: Call, + locals_by_unit: dict[str, CostContext], + positions_by_unit: dict[str, int], + whole: CostContext, + asked: str | None = None, ) -> ComputeCostMetadata: - """Measure the work one Call asks for, without attaching the record. - - Work only: what an occurrence moves is the memory family's answer. Global - work is one selected unit's work repeated over the positions executing this - scope, so it is derived from the same registered evaluator rather than - independently recomputed from the authored type. - """ + """Measure one Call in logical, expanded, and every topology-unit domain.""" + flops_by_unit: list[tuple[tuple[str, int], ...]] = [] + other_ops_by_unit: list[tuple[tuple[str, int], ...]] = [] + total_flops: tuple[tuple[str, int], ...] = () + total_other_ops: tuple[tuple[str, int], ...] = () try: - whole_cost = CostEvaluator().visit(expr, whole) - local_cost = CostEvaluator().visit(expr, local) + logical = CostEvaluator().visit(expr, whole) except (ValueError, VerifyError) as error: raise AnalysisError(str(error)) from None + for unit, local in locals_by_unit.items(): + try: + cost = CostEvaluator().visit(expr, local) + except (ValueError, VerifyError) as error: + raise AnalysisError(str(error)) from None + unit_flops = _counts(cost.flops) + unit_other_ops = _counts(cost.service) + flops_by_unit.append(unit_flops) + other_ops_by_unit.append(unit_other_ops) + if unit == (asked or next(iter(locals_by_unit))): + repeats = positions_by_unit[unit] + total_flops = tuple((name, value * repeats) for name, value in unit_flops) + total_other_ops = tuple((kind, value * repeats) for kind, value in unit_other_ops) return ComputeCostMetadata( - flops=_flops( - {dtype: value * executing_positions for dtype, value in local_cost.flops.items()} - ), - flops_logical=sum(whole_cost.flops.values()), - flops_per_unit=_flops(local_cost.flops), - ops=tuple( - sorted( - (kind, value * executing_positions) for kind, value in local_cost.ops.items() - ) - ), - ops_logical=tuple(sorted(whole_cost.ops.items())), - ops_per_unit=tuple(sorted(local_cost.ops.items())), + topologies=tuple(locals_by_unit), + flops=Spread(_counts(logical.flops), total_flops, tuple(flops_by_unit)), + other_ops=Spread(_counts(logical.service), total_other_ops, tuple(other_ops_by_unit)), ) -def _scope_position_count(mesh: Mesh, level: str | None, topologies: tuple) -> int: - """Count positions at or above the selected level within *mesh*.""" - if level is None: +def _scope_position_count(mesh: Mesh, topology_level: str | None, topologies: tuple) -> int: + if topology_level is None: return 1 declared = {topology.name: index for index, topology in enumerate(topologies)} - selected = declared[level] + selected = declared[topology_level] shape, _strides, _offset = _positions_layout(mesh) positions = 1 - for topology, axes in zip(mesh.topologies, level_axes(mesh)): + for topology, axes in zip(mesh.topologies, topology_axes(mesh)): if declared[topology.name] > selected: continue for axis in axes: @@ -182,48 +186,43 @@ def _scope_position_count(mesh: Mesh, level: str | None, topologies: tuple) -> i return positions -def _accumulate( - flops: dict[str, int], - flops_logical: list[int], - flops_per_unit: dict[str, int], - ops: dict[str, int], - ops_logical: dict[str, int], - ops_per_unit: dict[str, int], - record: ComputeCostMetadata, - trips: int, -) -> None: - for name, value in record.flops: - flops[name] = flops.get(name, 0) + value * trips - flops_logical[0] += record.flops_logical * trips - for name, value in record.flops_per_unit: - flops_per_unit[name] = flops_per_unit.get(name, 0) + value * trips - for name, value in record.ops: - ops[name] = ops.get(name, 0) + value * trips - for name, value in record.ops_logical: - ops_logical[name] = ops_logical.get(name, 0) + value * trips - for name, value in record.ops_per_unit: - ops_per_unit[name] = ops_per_unit.get(name, 0) + value * trips +def _bytes( + held: Breakdown[TrafficBytes], + topologies: tuple[str, ...], + kind: str, + topology_level: str | None, +) -> int: + moved = shares(held, topologies, topology_level).get(kind) + return moved.total_bytes if moved is not None else 0 -@dataclass -class ComputeCostContext(AnalyzeContext): - """State carried through the compute-cost expression walk. +def _add(target: dict[str, int], values, trips: int) -> None: + for name, value in values: + target[name] = target.get(name, 0) + value * trips + + +def _accumulate(ctx: "ComputeCostContext", record: ComputeCostMetadata, trips: int) -> None: + _add(ctx.flops_logical, record.flops.logical, trips) + _add(ctx.flops, record.flops.total, trips) + _add(ctx.other_ops_logical, record.other_ops.logical, trips) + _add(ctx.other_ops, record.other_ops.total, trips) + for index, unit in enumerate(record.topologies): + held = ctx.by_unit.setdefault(unit, {"flops": {}, "other_ops": {}}) + _add(held["flops"], record.flops.at(index), trips) + _add(held["other_ops"], record.other_ops.at(index), trips) - ``executing_positions`` is the number of positions in the enclosing scope - through the selected topology level; it is the D68 multiplier that turns - per-unit work into total replicated work. - """ - local: CostContext | None = None +@dataclass +class ComputeCostContext(AnalyzeContext): + locals_by_unit: dict[str, CostContext] = field(default_factory=dict) whole: CostContext | None = None current_mesh: Mesh | None = None - executing_positions: int = 1 + executing_positions: dict[str, int] = field(default_factory=dict) flops: dict[str, int] = field(default_factory=dict) - flops_logical: list[int] = field(default_factory=lambda: [0]) - flops_per_unit: dict[str, int] = field(default_factory=dict) - ops: dict[str, int] = field(default_factory=dict) - ops_logical: dict[str, int] = field(default_factory=dict) - ops_per_unit: dict[str, int] = field(default_factory=dict) + flops_logical: dict[str, int] = field(default_factory=dict) + other_ops: dict[str, int] = field(default_factory=dict) + other_ops_logical: dict[str, int] = field(default_factory=dict) + by_unit: dict[str, dict[str, dict[str, int]]] = field(default_factory=dict) call_count: list[int] = field(default_factory=lambda: [0]) @@ -231,15 +230,14 @@ class ComputeCostVisitor(ExprVisitor[None]): """Attach per-Call work and accumulate multiplicity-aware totals.""" def visit_MeshRegion(self, expr: MeshRegion, ctx: ComputeCostContext) -> None: - """Carry the region's execution multiplicity into each contained Call.""" for arg in expr.args: self.visit(arg, ctx) mesh = composed((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh - positions = _scope_position_count(mesh, ctx.level, ctx.module.effective_topologies()) - self.visit( - expr.body, - replace(ctx, executing_positions=positions, current_mesh=mesh), - ) + topologies = ctx.module.effective_topologies() + positions = { + unit: _scope_position_count(mesh, unit, topologies) for unit in ctx.locals_by_unit + } + self.visit(expr.body, replace(ctx, executing_positions=positions, current_mesh=mesh)) def visit_LoopRegion(self, expr: LoopRegion, ctx: ComputeCostContext) -> None: child = next(item for item in ctx.current.children if item.owner is expr) @@ -256,11 +254,11 @@ def default_visit_leaf( if not isinstance(expr, Call): return ctx.call_count[0] += 1 - if ctx.local is None: + if not ctx.locals_by_unit or ctx.whole is None: raise AnalysisError("compute-cost: visitor context is missing its cost context") - if ctx.whole is None: - raise AnalysisError("compute-cost: visitor context is missing its whole context") - record = _call_cost_record(expr, ctx.whole, ctx.local, ctx.executing_positions) + record = _call_cost_record( + expr, ctx.locals_by_unit, ctx.executing_positions, ctx.whole, ctx.topology_level + ) attach(expr, record) owner = ctx.current if id(expr) in ctx.current.accesses["narrow"] else ctx.root repeats = 1 @@ -269,49 +267,48 @@ def default_visit_leaf( if cursor.is_variant(expr): repeats *= max(1, cursor.trips()) cursor = cursor.parent - _accumulate( - ctx.flops, - ctx.flops_logical, - ctx.flops_per_unit, - ctx.ops, - ctx.ops_logical, - ctx.ops_per_unit, - record, - repeats, - ) + _accumulate(ctx, record, repeats) -def analyze_compute_cost( - function: Function, - context: AnalyzeContext, -) -> None: +def analyze_compute_cost(function: Function, context: AnalyzeContext) -> None: """Attach one-trip work per Call and multiplicity-aware totals per Function.""" - module, level = context.module, context.level + module, topology_level = context.module, context.topology_level topologies = module.effective_topologies() scope = FunctionScope(module, function) - whole = CostContext(scope=scope) - local = CostContext(scope=scope, level=level, topologies=topologies) + units = tuple(topology.name for topology in topologies) or (topology_level,) + locals_by_unit = { + unit: CostContext(scope=scope, topology_level=unit, topologies=topologies) + for unit in units + if unit is not None + } cost_context = ComputeCostContext( module=module, target=context.target, - level=level, + topology_level=topology_level, options=context.options, root=context.root, current=context.current, - local=local, - whole=whole, + locals_by_unit=locals_by_unit, + whole=CostContext(scope=scope), + executing_positions=dict.fromkeys(locals_by_unit, 1), ) ComputeCostVisitor().visit(function.body, cost_context) if cost_context.call_count[0] > 0: + units = tuple(cost_context.by_unit) attach( function, ComputeCostMetadata( - flops=tuple(sorted(cost_context.flops.items())), - flops_logical=cost_context.flops_logical[0], - flops_per_unit=tuple(sorted(cost_context.flops_per_unit.items())), - ops=tuple(sorted(cost_context.ops.items())), - ops_logical=tuple(sorted(cost_context.ops_logical.items())), - ops_per_unit=tuple(sorted(cost_context.ops_per_unit.items())), + topologies=units, + flops=Spread( + _counts(cost_context.flops_logical), + _counts(cost_context.flops), + tuple(_counts(cost_context.by_unit[u]["flops"]) for u in units), + ), + other_ops=Spread( + _counts(cost_context.other_ops_logical), + _counts(cost_context.other_ops), + tuple(_counts(cost_context.by_unit[u]["other_ops"]) for u in units), + ), ), ) diff --git a/src/tilefoundry/analysis/facts.py b/src/tilefoundry/analysis/facts.py index 1129fce17..b1756cd3f 100644 --- a/src/tilefoundry/analysis/facts.py +++ b/src/tilefoundry/analysis/facts.py @@ -89,15 +89,11 @@ class MemoryHierarchyFacts: def explicit(self, name: str) -> ExplicitMemoryLevelFacts | None: """The explicit level called *name*, if the target has one.""" - return next( - (level for level in self.explicit_levels if level.name == name), None - ) + return next((level for level in self.explicit_levels if level.name == name), None) def implicit(self, name: str) -> ImplicitMemoryLevelFacts | None: """The implicit level called *name*, if the target has one.""" - return next( - (level for level in self.implicit_levels if level.name == name), None - ) + return next((level for level in self.implicit_levels if level.name == name), None) def cached_level(self, name: str) -> str | None: """The level *name* caches, if it caches one.""" @@ -105,8 +101,7 @@ def cached_level(self, name: str) -> str | None: ( relation.far for relation in self.relations - if relation.kind is MemoryRelationKind.CACHES - and relation.near == name + if relation.kind is MemoryRelationKind.CACHES and relation.near == name ), None, ) @@ -122,9 +117,7 @@ def backing_level(self, name: str) -> str: current = name while (nearer := self.cached_level(current)) is not None: if nearer in seen: - raise ValueError( - f"memory hierarchy: caching cycle through {current!r}" - ) + raise ValueError(f"memory hierarchy: caching cycle through {current!r}") seen.add(nearer) current = nearer return current @@ -147,7 +140,7 @@ class PerformanceServiceFacts: """Everything one unit gets through, by the kind of work it is asked for. Floating-point work by dtype, movement by the level it crosses, and the work - that is not floating point at all by the operation kind it asks for. A comparison, + that is not floating point at all by the service it asks for. A comparison, a select, an integer add and a local move all take a machine time, and none of them is a FLOP; a dtype is not a kind of work. """ diff --git a/src/tilefoundry/analysis/memory.py b/src/tilefoundry/analysis/memory.py index 218f8ecd6..4660a9aff 100644 --- a/src/tilefoundry/analysis/memory.py +++ b/src/tilefoundry/analysis/memory.py @@ -277,8 +277,9 @@ def _shares( ( name, Spread( - shares.get(_WHOLE, TrafficBytes()), - tuple(shares.get(unit, TrafficBytes()) for unit in topologies), + logical=shares.get(_WHOLE, TrafficBytes()), + total=shares.get(_WHOLE, TrafficBytes()), + per_unit=tuple(shares.get(unit, TrafficBytes()) for unit in topologies), ), ) for name, shares in sorted(held.items()) diff --git a/src/tilefoundry/analysis/metadata.py b/src/tilefoundry/analysis/metadata.py index 0254a63af..1f85f9b15 100644 --- a/src/tilefoundry/analysis/metadata.py +++ b/src/tilefoundry/analysis/metadata.py @@ -8,6 +8,7 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from dataclasses import dataclass from tilefoundry.ir.core.metadata import IRMetadata @@ -15,66 +16,130 @@ @dataclass(frozen=True) -class ComputeCostMetadata(IRMetadata): - """Record one occurrence's work, or one Function's total work. - - ``flops`` and ``ops`` state global work; their ``*_per_unit`` partners - apply shard projection at the requested topology level. ``ops`` counts - non-floating-point machine operations by kind -- comparing, selecting, - whole-number arithmetic and special functions. What an occurrence moves is a separate - record, kept by the family that knows where values live. On a Call these - state one occurrence; on a Function, loops contribute their trip count. +class Spread[V]: + """One typed value in logical, expanded, and topology-unit domains. + + ``per_unit`` follows ``ComputeCostMetadata.topologies``. The value retains + its semantic structure while ``Spread`` states only its counting domain. """ - flops: tuple[tuple[str, int], ...] = () - flops_logical: int = 0 - flops_per_unit: tuple[tuple[str, int], ...] = () - ops: tuple[tuple[str, int], ...] = () - ops_logical: tuple[tuple[str, int], ...] = () - ops_per_unit: tuple[tuple[str, int], ...] = () + logical: V + total: V + per_unit: tuple[V, ...] = () - def ops_per_unit_of(self, kind: str) -> int: - """One unit's count of *kind*, zero when the call asks for none.""" - return next((value for name, value in self.ops_per_unit if name == kind), 0) + def at(self, index: int) -> V: + """One topology level's value by position, or ``total`` when absent.""" + return self.per_unit[index] if index < len(self.per_unit) else self.total @dataclass(frozen=True) -class TrafficMetadata(IRMetadata): - """The bytes one occurrence moves, whole and for one participant. +class Breakdown[V]: + """Typed quantities grouped by resource kind for traffic records.""" + + kinds: tuple[tuple[str, Spread[V]], ...] = () + + def of(self, kind: str) -> Spread[V] | None: + return next((value for name, value in self.kinds if name == kind), None) + + def names(self) -> tuple[str, ...]: + return tuple(name for name, _ in self.kinds) + + +def breakdown[V]( + total: "Mapping[str, V]", per_unit: "Sequence[Mapping[str, V]]", zero: V +) -> Breakdown[V]: + """Gather traffic kinds, filling missing values at each topology level.""" + kinds = sorted({*total, *(kind for share in per_unit for kind in share)}) + return Breakdown( + tuple( + ( + kind, + Spread( + logical=total.get(kind, zero), + total=total.get(kind, zero), + per_unit=tuple(share.get(kind, zero) for share in per_unit), + ), + ) + for kind in kinds + ) + ) + + +def shares[V]( + held: Breakdown[V] | Spread[tuple[tuple[str, V], ...]], + topologies: tuple[str, ...], + topology_level: "str | None" = None, +) -> dict[str, V]: + """Each traffic kind's value at one topology level, or its total.""" + index = topologies.index(topology_level) if topology_level in topologies else None + if isinstance(held, Spread): + value = held.total if index is None else held.at(index) + return dict(value) + return { + kind: spread.total if index is None else spread.at(index) for kind, spread in held.kinds + } + + +TypedCounts = tuple[tuple[str, int], ...] + + +@dataclass(frozen=True) +class ComputeCostMetadata(IRMetadata): + """Typed floating-point and other operation work in each counting domain.""" - Which way a boundary moves is its Op's evaluator's answer and how much is - its relation's; the family that decides where values live attaches the - record, and where they landed never corrects a crossing. ``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, ...] = () + flops: Spread[TypedCounts] = Spread((), (), ()) + other_ops: Spread[TypedCounts] = Spread((), (), ()) + + +@dataclass(frozen=True) +class TrafficMetadata(IRMetadata): + """The bytes one occurrence moves, in the two coordinates a move has. + + ``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. """ - whole: tuple[tuple[str, TrafficBytes], ...] = () - per_unit: tuple[tuple[str, TrafficBytes], ...] = () + topologies: tuple[str, ...] = () + storage: Breakdown[TrafficBytes] = Breakdown() + communication: Breakdown[TrafficBytes] = Breakdown() operands: tuple[TrafficBytes, ...] = () - def at(self, level: str) -> TrafficBytes: - """Bytes moved at *level*, zero when the occurrence does not touch it.""" - return next((value for name, value in self.whole if name == level), TrafficBytes()) + @property + def whole(self) -> tuple[tuple[str, TrafficBytes], ...]: + return tuple((name, spread.total) for name, spread in self.storage.kinds) - def per_unit_at(self, level: str) -> TrafficBytes: - """One unit's bytes at *level*, zero when it does not touch it.""" - return next((value for name, value in self.per_unit if name == level), TrafficBytes()) + @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 LevelFootprint: +class MemoryLevelFootprint: """How much of one memory level a function needs at its peak. ``persistent_bytes`` is the part that cannot be reclaimed within the function, so it is the floor the peak can never fall below. """ - level: str + memory_level: str peak_bytes: int 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.""" @@ -86,11 +151,15 @@ class BufferFootprint: """Per-position, device-wide, and repeated bytes touched in one buffer.""" buffer: str - level: 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): @@ -117,12 +186,17 @@ class ValueLifetime: """ binding: str - level: str + memory_level: str bytes: int defined_at: int 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: @@ -149,14 +223,14 @@ class MemoryMetadata(IRMetadata): placed one: nothing was decided, so nothing is claimed. """ - footprint: tuple[LevelFootprint, ...] = () + footprint: tuple[MemoryLevelFootprint, ...] = () lifetimes: tuple[ValueLifetime, ...] = () advisories: tuple[str, ...] = () allocation: "AllocationMetadata | None" = None - def level(self, name: str) -> LevelFootprint | 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.level == name), None) + return next((item for item in self.footprint if item.memory_level == name), None) @dataclass(frozen=True) @@ -195,21 +269,6 @@ class TimelineMetadata: stride_ns: int = 0 -@dataclass(frozen=True) -class TripInterval: - """A half-open interval, parameterized by the trip index of its loop. - - A single trip occurs once, so ``stride`` says nothing about it and the trip - index does not appear. Repeated trips state the first interval and the - stride every later one is offset by. - """ - - start: int - end: int - stride: int = 0 - trips: int = 1 - - @dataclass(frozen=True) class PerformanceMetadata(IRMetadata): """One occurrence's interval within one local wave of its Function. @@ -238,16 +297,19 @@ class PerformanceSummaryMetadata(IRMetadata): __all__ = [ "AllocationMetadata", + "Breakdown", "BufferFootprint", "ComputeCostMetadata", - "LevelFootprint", "LoopFootprintMetadata", + "MemoryLevelFootprint", "MemoryMetadata", "PerformanceMetadata", "PerformanceSummaryMetadata", "RooflineMetadata", + "Spread", "TimelineMetadata", - "TripInterval", "TrafficBytes", "ValueLifetime", + "breakdown", + "shares", ] diff --git a/src/tilefoundry/analysis/report.py b/src/tilefoundry/analysis/report.py index 2e78e83b1..d345da0a6 100644 --- a/src/tilefoundry/analysis/report.py +++ b/src/tilefoundry/analysis/report.py @@ -104,6 +104,7 @@ def _reported_value(value: object, declared: object) -> object: if origin is Spread: (item,) = get_args(declared) return { + "logical": _reported_value(value.logical, item), "total": _reported_value(value.total, item), "per_unit": [_reported_value(share, item) for share in value.per_unit], } @@ -135,6 +136,7 @@ def _pair_types(declared: object) -> tuple[object, object] | None: for _record_type in ( ComputeCostMetadata, + TrafficMetadata, LoopFootprintMetadata, MemoryMetadata, RooflineMetadata, @@ -305,14 +307,12 @@ def _loop_records( return rows -def _totals_of(held: "Breakdown | None") -> dict[str, object]: - """Each kind's whole-program amount, without any level's share. - - A total is what the whole program asks for, so this section states one - number per kind; the per-level shares are on the family's own record. - """ +def _totals_of(held: "Breakdown | Spread | None") -> dict[str, object]: + """Each kind's whole-program amount, without any level's share.""" if held is None: return {} + if isinstance(held, Spread): + return dict(held.total) return { kind: _reported_value(spread.total, type(spread.total)) if is_dataclass(type(spread.total)) diff --git a/src/tilefoundry/analysis/roofline.py b/src/tilefoundry/analysis/roofline.py index 244335f8d..00e88195a 100644 --- a/src/tilefoundry/analysis/roofline.py +++ b/src/tilefoundry/analysis/roofline.py @@ -31,8 +31,8 @@ def _totals(held) -> tuple[tuple[str, int], ...]: - """One category's kinds with their whole-program counts.""" - return tuple((kind, spread.total) for kind, spread in held.kinds) + """One typed work category's whole-program counts.""" + return held.total def _compute_ns(flops: tuple[tuple[str, int], ...], facts: ThroughputFacts) -> int: @@ -100,7 +100,7 @@ def _cost_bound( Whole-device work against whole-device rates: the flops the target publishes a peak for, and the bytes at the level it publishes a bandwidth for. Typed - typed ops have no whole-device rate to divide by and so do not enter a bound, + service has no whole-device rate to divide by and so does not enter a bound, and neither do bytes at a level with no published bandwidth: what asks for a nanosecond is what this could have priced, so a dtype whose rate is missing still owes one and a level nobody rated does not. diff --git a/src/tilefoundry/cli/analyze.py b/src/tilefoundry/cli/analyze.py index 0f4cbd809..144ac3b56 100644 --- a/src/tilefoundry/cli/analyze.py +++ b/src/tilefoundry/cli/analyze.py @@ -43,8 +43,7 @@ def _watch(limit: float) -> None: def stop() -> None: time.sleep(limit) sys.stderr.write( - "tilefoundry: error: analysis too complex, " - f"timed out after {limit:.0f}s\n" + f"tilefoundry: error: analysis too complex, timed out after {limit:.0f}s\n" ) sys.stderr.flush() os._exit(1) @@ -83,10 +82,10 @@ def guidance() -> str: family what --topology changes pass it when ------------ -------------------------------------- --------------------- - compute-cost flops_per_unit and ops_per_unit. the program shards - flops and ops stay global - memory per-unit traffic. Footprint follows the program shards - its owner for each storage level + 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 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 @@ -156,7 +155,7 @@ def run_authored_analysis( Path(out_path).write_text(annotated, encoding="utf-8") return 0 - result = analyze(module, function, analysis=analyses, level=topology, dims=dims) + result = analyze(module, function, analysis=analyses, topology_level=topology, dims=dims) rendered = render_analysis(result, operands=operands and not as_json) if as_json: Path(out_path).write_text( @@ -165,9 +164,7 @@ def run_authored_analysis( ) return 0 - Path(out_path).write_text( - f"{render_text(rendered)}\n\n{rendered.annotated}", encoding="utf-8" - ) + Path(out_path).write_text(f"{render_text(rendered)}\n\n{rendered.annotated}", encoding="utf-8") return 0 diff --git a/src/tilefoundry/codegen/context.py b/src/tilefoundry/codegen/context.py index ce521a77e..7f06cd047 100644 --- a/src/tilefoundry/codegen/context.py +++ b/src/tilefoundry/codegen/context.py @@ -12,14 +12,12 @@ from collections.abc import Callable, Mapping from tilefoundry.codegen.signature import CallableSignature, Signature, TensorSignature -from tilefoundry.ir.core.expr import Call from tilefoundry.ir.tir.stmts import Evaluate -from tilefoundry.ir.types import UnitType from tilefoundry.target.base import Target from tilefoundry.visitor_registry.registries import DispatchRegistry, Role, spelled -class CodegenContext: +class EmitContext: """One compile, from the first function read to the last line written.""" target_kind: type[Target] @@ -38,10 +36,14 @@ def __init__( *, symbols: Mapping[int, CallableSignature] | None = None, target: Target | None = None, + resolved_launches: Mapping[int, object] | None = None, + codegen_context: object | None = None, ) -> None: self.registry = registry + self.codegen_context = codegen_context self.symbols: Mapping[int, CallableSignature] = {} if symbols is None else symbols self.target = target + self.resolved_launches = {} if resolved_launches is None else resolved_launches self.exported = False self._lines: list[str] = [] self._indent = 0 @@ -49,6 +51,13 @@ def __init__( self._counter = 0 self._kernel_param_ids: set[int] = set() + def callee_of(self, launch): + """The callee root context resolved for one Launch statement.""" + callee = self.resolved_launches.get(id(launch)) + if callee is None: + raise KeyError("codegen: Launch was not resolved by the root context") + return callee + def signature_of(self, fn) -> CallableSignature: """How a caller writes a call to *fn*, as the symbol table settled it.""" signature = self.symbols.get(id(fn)) @@ -109,14 +118,24 @@ def capture(self, fn) -> str: self._lines = saved_lines def emit_node(self, node) -> None: - """Write *node*, dispatching an ``Evaluate`` by the Op it wraps. + """Write *node* through the target Visitor or function orchestrator.""" + from tilefoundry.ir.tir.abort import Abort # noqa: PLC0415 + from tilefoundry.ir.tir.stmts import ( # noqa: PLC0415 + For, + If, + LetStmt, + MeshScope, + Return, + Sequential, + While, + ) - An effect Op reaches codegen inside an ``Evaluate``, so the handler - that answers for it is registered against the Op's own class. - """ - if isinstance(node, Evaluate): - op = node.callable - self.handler_for(type(op))(Call(type=UnitType(), target=op, args=node.args), self) + if isinstance( + node, (Abort, Evaluate, For, If, LetStmt, MeshScope, Return, Sequential, While) + ): + from tilefoundry.codegen.emitter import CudaEmitter # noqa: PLC0415 + + CudaEmitter(context=self).visit(node) return self.handler_for(type(node))(node, self) @@ -174,9 +193,7 @@ def argument(self, signature: Signature, callee_target: Target) -> tuple[str, .. def arguments(self, signature: CallableSignature, callee_target: Target) -> str: """The whole argument list of one call, in the order the callee declared.""" return ", ".join( - token - for param in signature.all_params - for token in self.argument(param, callee_target) + token for param in signature.all_params for token in self.argument(param, callee_target) ) def local_value(self, signature: Signature) -> str: @@ -204,4 +221,4 @@ def _handler(self, key) -> Callable: return handler -__all__ = ["CodegenContext"] +__all__ = ["EmitContext"] diff --git a/src/tilefoundry/codegen/context_builder.py b/src/tilefoundry/codegen/context_builder.py new file mode 100644 index 000000000..00f827960 --- /dev/null +++ b/src/tilefoundry/codegen/context_builder.py @@ -0,0 +1,312 @@ +"""Build the compile-wide code-generation context before source emission.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, replace +from types import MappingProxyType + +from tilefoundry.codegen.signature import ( + LAUNCH_ABI, + CallableSignature, + called_as, + program_id_params, +) +from tilefoundry.ir.core.module import Module, module_functions, owning_module, subtree +from tilefoundry.ir.tir.launch import Launch +from tilefoundry.ir.tir.prim_function import PrimFunction +from tilefoundry.ir.tir.stmts import Evaluate, Sequential +from tilefoundry.target import CpuTarget, Target + + +@dataclass(frozen=True) +class FunctionSymbol: + """Resolved IR identity and generated ABI names for one callable function.""" + + function: PrimFunction + ir_name: str + kernel_name: str | None + shim_name: str | None + host_name: str | None + owner: Module + target: Target + callable: CallableSignature + + +@dataclass(frozen=True) +class FunctionSymbolTable: + """Separate IR and generated-name namespaces shared by every emitter.""" + + by_function: Mapping[int, FunctionSymbol] + by_ir_name: Mapping[str, FunctionSymbol] + by_kernel_name: Mapping[str, FunctionSymbol] + + @property + def callables(self) -> Mapping[int, CallableSignature]: + return MappingProxyType( + {identity: symbol.callable for identity, symbol in self.by_function.items()} + ) + + +@dataclass(frozen=True) +class LaunchEdge: + """A resolved host-to-device Launch and its domain-owned callable ABI.""" + + launch: Evaluate + caller: PrimFunction + callee: PrimFunction + symbol: FunctionSymbol + callable: CallableSignature + grid_x: object + block_x: object + + +@dataclass(frozen=True) +class EmissionGroup: + owner: Module + target: Target + functions: tuple[PrimFunction, ...] + + +@dataclass(frozen=True) +class EmitContext: + """The compile facts required by one target/topology emission group.""" + + codegen: CodegenContext + owner: Module + target: Target + functions: tuple[PrimFunction, ...] + symbols: Mapping[int, CallableSignature] + launches: Mapping[int, tuple[tuple[object, object, object], tuple[object, object, object]]] + resolved_launches: Mapping[int, PrimFunction] + + +@dataclass(frozen=True) +class CodegenContext: + """Complete compile-wide resolution performed once before emission.""" + + module: Module + functions: tuple[PrimFunction, ...] + ownership: Mapping[int, Module] + targets: Mapping[int, Target] + symbols: FunctionSymbolTable + launch_edges: tuple[LaunchEdge, ...] + topology_domains: Mapping[int, Module] + launch_callables: Mapping[int, CallableSignature] + groups: tuple[EmissionGroup, ...] + + def for_group(self, group: EmissionGroup) -> EmitContext: + identities = {id(fn) for fn in group.functions} + symbols = { + identity: symbol.callable + for identity, symbol in self.symbols.by_function.items() + if identity in identities or symbol.target != group.target + } + launches = { + id(edge.callee): ((edge.grid_x, 1, 1), (edge.block_x, 1, 1)) + for edge in self.launch_edges + if id(edge.callee) in identities + } + for edge in self.launch_edges: + if id(edge.callee) in identities or edge.caller in group.functions: + symbols[id(edge.callee)] = edge.callable + resolved_launches = { + id(edge.launch): edge.callee + for edge in self.launch_edges + if edge.caller in group.functions + } + return EmitContext( + codegen=self, + owner=group.owner, + target=group.target, + functions=group.functions, + symbols=MappingProxyType(symbols), + launches=MappingProxyType(launches), + resolved_launches=MappingProxyType(resolved_launches), + ) + + +def _domain_functions(node: Module) -> tuple[PrimFunction, ...]: + inherited = tuple( + function + for child in node.modules + if child.topologies is None + for function in _domain_functions(child) + ) + return (*node.functions, *inherited) + + +def _domains(root: Module) -> tuple[tuple[Module, tuple[PrimFunction, ...]], ...]: + domains = [(root, _domain_functions(root))] + for node in subtree(root): + domains.extend( + (child, _domain_functions(child)) + for child in node.modules + if child.topologies is not None + ) + return tuple(domains) + + +def _product(values: tuple[object, ...]) -> object: + result: object = 1 + for value in values: + if value is None: + raise ValueError( + "codegen: topology extent cannot be None; use a DimVar-backed " + "CPU-computable launch expression" + ) + result = result * value + return result + + +def _launch_extents(owner: Module) -> tuple[object, object]: + """Derive one-dimensional grid/block extents from declared topologies.""" + ctas = tuple( + topology.size for topology in owner.effective_topologies() if topology.name == "cta" + ) + threads = tuple( + topology.size for topology in owner.effective_topologies() if topology.name == "thread" + ) + return _product(ctas), _product(threads) + + +def _launch_statements(function: PrimFunction) -> tuple[Evaluate, ...]: + body = function.body + if not isinstance(body, Sequential): + return () + return tuple( + stmt + for stmt in body.body + if isinstance(stmt, Evaluate) and isinstance(stmt.callable, Launch) + ) + + +def build_codegen_context(module: Module) -> CodegenContext: + """Resolve functions, symbols, Launch edges, domains, and groups once.""" + functions = tuple(module_functions(module)) + if not all(isinstance(function, PrimFunction) for function in functions): + bad = next(function for function in functions if not isinstance(function, PrimFunction)) + raise TypeError( + f"codegen expects PrimFunction values, got {type(bad).__name__} {bad.name!r}" + ) + expanded_list: list[PrimFunction] = [] + seen_expanded: set[int] = set() + for function in functions: + for variant in function.variants or (function,): + if id(variant) not in seen_expanded: + seen_expanded.add(id(variant)) + expanded_list.append(variant) + expanded = tuple(expanded_list) + ownership: dict[int, Module] = {} + targets: dict[int, Target] = {} + for function in (*functions, *expanded): + owner = owning_module(module, function) + if owner is None: + raise ValueError(f"codegen: function {function.name!r} has no unique owning Module") + if function.target is None: + raise ValueError(f"codegen: function {function.name!r} has no resolved Target") + ownership[id(function)] = owner + targets[id(function)] = function.target + + device_targets = tuple( + dict.fromkeys( + function.target + for function in (*functions, *expanded) + if not isinstance(function.target, CpuTarget) + ) + ) + if len(device_targets) != 1: + raise ValueError(f"codegen: expected one device Target, found {len(device_targets)}") + program_ids = program_id_params(module, device_targets[0]) + by_function: dict[int, FunctionSymbol] = {} + by_ir_name: dict[str, FunctionSymbol] = {} + by_kernel_name: dict[str, FunctionSymbol] = {} + for function in functions: + signature = called_as(function, program_ids) + is_device = not isinstance(function.target, CpuTarget) + symbol = FunctionSymbol( + function=function, + ir_name=function.name, + kernel_name=function.name if is_device else None, + shim_name=signature.name if is_device else None, + host_name=signature.name if not is_device else None, + owner=ownership[id(function)], + target=function.target, + callable=signature, + ) + if function.name in by_ir_name: + raise ValueError(f"codegen: duplicate IR function name {function.name!r}") + by_function[id(function)] = symbol + by_ir_name[function.name] = symbol + if symbol.kernel_name is not None: + if symbol.kernel_name in by_kernel_name: + raise ValueError(f"codegen: duplicate kernel symbol {symbol.kernel_name!r}") + by_kernel_name[symbol.kernel_name] = symbol + + table = FunctionSymbolTable( + MappingProxyType(by_function), + MappingProxyType(by_ir_name), + MappingProxyType(by_kernel_name), + ) + edges: list[LaunchEdge] = [] + launch_callables: dict[int, CallableSignature] = {} + for caller in functions: + for stmt in _launch_statements(caller): + reference = stmt.args[0] + symbol = by_ir_name.get(reference.name) + if symbol is None: + raise ValueError( + f"codegen: launch from {caller.name!r} names unknown function {reference.name!r}" + ) + grid_x, block_x = _launch_extents(symbol.owner) + callable_signature = replace(symbol.callable, trailing=LAUNCH_ABI) + launch_callables[id(stmt)] = callable_signature + edges.append( + LaunchEdge( + stmt, + caller, + symbol.function, + symbol, + callable_signature, + grid_x, + block_x, + ) + ) + + domain_rows = _domains(module) + topology_domains = {id(function): owner for owner, owned in domain_rows for function in owned} + groups: list[EmissionGroup] = [] + for owner, owned in domain_rows: + by_target: dict[Target, list[PrimFunction]] = {} + for function in owned: + by_target.setdefault(function.target, []).append(function) + for target, grouped in by_target.items(): + if not isinstance(target, CpuTarget) and not any( + edge.callee in grouped for edge in edges + ): + missing = ", ".join(function.name for function in grouped) + raise ValueError(f"codegen: no Launch states geometry for {missing}") + groups.append(EmissionGroup(owner, target, tuple(grouped))) + + return CodegenContext( + module=module, + functions=functions, + ownership=MappingProxyType(ownership), + targets=MappingProxyType(targets), + symbols=table, + launch_edges=tuple(edges), + topology_domains=MappingProxyType(topology_domains), + launch_callables=MappingProxyType(launch_callables), + groups=tuple(groups), + ) + + +__all__ = [ + "CodegenContext", + "EmissionGroup", + "EmitContext", + "FunctionSymbol", + "FunctionSymbolTable", + "LaunchEdge", + "build_codegen_context", +] diff --git a/src/tilefoundry/codegen/cpu/abi.py b/src/tilefoundry/codegen/cpu/abi.py index 2b095a0e6..30414d083 100644 --- a/src/tilefoundry/codegen/cpu/abi.py +++ b/src/tilefoundry/codegen/cpu/abi.py @@ -7,7 +7,7 @@ from __future__ import annotations -from tilefoundry.codegen.context import CodegenContext +from tilefoundry.codegen.context import EmitContext from tilefoundry.codegen.signature import ( CallableSignature, ProgramIdSignature, @@ -49,13 +49,13 @@ def _called_as_entry( @register_codegen(CpuTarget, Role.CALLEE, TensorSignature) -def _declare_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ...]: +def _declare_tensor(sig: TensorSignature, ctx: EmitContext) -> tuple[str, ...]: """A runtime tensor, which carries its own extents, so nothing else declares them.""" return (f"{TENSOR_CTYPE} {sig.name}",) @register_codegen(CpuTarget, Role.CALLEE, ProgramIdSignature) -def _declare_program_id(sig: ProgramIdSignature, ctx: CodegenContext) -> tuple[str, ...]: +def _declare_program_id(sig: ProgramIdSignature, ctx: EmitContext) -> tuple[str, ...]: """The id of a level no device register answers for, passed in by value.""" return (f"{PROGRAM_ID_CTYPE} {sig.name}",) diff --git a/src/tilefoundry/codegen/cpu/context.py b/src/tilefoundry/codegen/cpu/context.py index 5abe92468..7e9ef7deb 100644 --- a/src/tilefoundry/codegen/cpu/context.py +++ b/src/tilefoundry/codegen/cpu/context.py @@ -10,14 +10,14 @@ from collections.abc import Mapping -from tilefoundry.codegen.context import CodegenContext +from tilefoundry.codegen.context import EmitContext from tilefoundry.codegen.signature import CallableSignature, Signature, TensorSignature from tilefoundry.target import CpuTarget from tilefoundry.target.base import Target from tilefoundry.visitor_registry.registries import codegen_registry -class CpuCodegenContext(CodegenContext): +class CpuCodegenContext(EmitContext): """A compile writing the host unit: the shared context plus what a tensor is here.""" target_kind = CpuTarget @@ -27,8 +27,16 @@ def __init__( *, symbols: Mapping[int, CallableSignature] | None = None, target: Target | None = None, + resolved_launches: Mapping[int, object] | None = None, + codegen_context: object | None = None, ) -> None: - super().__init__(codegen_registry, symbols=symbols, target=target) + super().__init__( + codegen_registry, + symbols=symbols, + target=target, + resolved_launches=resolved_launches, + codegen_context=codegen_context, + ) def local_value(self, signature: Signature) -> str: """The buffer behind a runtime tensor, which is what a device call takes.""" diff --git a/src/tilefoundry/codegen/cpu/module.py b/src/tilefoundry/codegen/cpu/module.py index 547ead958..e7d284d46 100644 --- a/src/tilefoundry/codegen/cpu/module.py +++ b/src/tilefoundry/codegen/cpu/module.py @@ -1,32 +1,31 @@ -"""Emit the host module for the split code-generation pipeline. +"""The host ``.cpp`` translation unit: one entry, and the launches it makes. -Launch entries validate and bind runtime tensors before calling a device shim. -Dispatch entries use a first-match shape predicate and throw on fallback. Output -uses only TVM FFI, DLPack, and standard C++; CUDA syntax and types remain in the -device module and shims. +The entry takes runtime tensors, checks each is on the device its parameter's +storage names, and calls a launch shim per ``Launch`` in its body. Which +arguments that call takes is the callee's answer, reached through the symbol +table, so nothing here spells a device type or a device ABI. Output uses only +TVM FFI, DLPack and standard C++. """ from __future__ import annotations +from dataclasses import replace + +from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cpu.templates import render -from tilefoundry.codegen.cuda.module import shim_symbol -from tilefoundry.codegen.cuda.tir.prim_function import ( - _internal_wrapper_symbol, - _is_hidden_shape_scalar, - _parse_shape_param_name, -) from tilefoundry.codegen.linkable import LinkableFunction, LinkableModule from tilefoundry.codegen.registry import CodeGenerator -from tilefoundry.codegen.signature import LaunchSignature +from tilefoundry.codegen.signature import ( + CallableSignature, + LaunchSignature, + tensor_signature_of, +) from tilefoundry.ir.core import Call, Constant, Var -from tilefoundry.ir.core.module import Module, owning_module -from tilefoundry.ir.core.pattern import DimVarRangePat, locate_dim_var +from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.launch import Launch from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.tir.stmts import Evaluate, Sequential -from tilefoundry.ir.tir.symbol_ref import symbol_call -from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.dim import ( DimAdd, DimFloorDiv, @@ -46,27 +45,6 @@ StorageKind.HOST: "kDLCPU", } - -_LAUNCH_ABI_DECL = ["int", "int", "int", "int", "int", "int", "int", "void*"] - - -def _static_smem(value) -> int: - """Resolve ``dynamic_smem`` to a static int. - - Resolve ``dynamic_smem`` to a static int. Host index-expr codegen - covers the projected grid/block geometry only; a dynamic shared-memory size expression - is not supported. - """ - sv = static_dim_value(value) - if sv is not None: - return sv - raise ValueError( - "emit_host_module: dynamic_smem must be a static int/Constant; a " - f"dynamic shared-memory size expression is not supported, got " - f"{type(value).__name__}" - ) - - _DIM_BINOP_CXX = { DimAdd: "+", DimSub: "-", @@ -106,49 +84,48 @@ def default_visit(self, expr, ctx=None) -> str: raise ValueError(f"emit_host_module: unsupported launch-extent node {type(expr).__name__}") -def _emit_host_int_expr(expr) -> str: - """Lower a launch-extent Expr to a host C++ integer expression. +def _extent(expr) -> str: + """One grid or block extent as a host C++ int expression. - Accepts an ``int`` / integer ``Constant`` (emitted as a literal), a - ``ShapeOf`` (the forwarded tensor's ``shape()`` access), or a - dim-arithmetic ``Call`` over those. Any unsupported node raises — there - is no silent zero/default. + An ``int`` or integer ``Constant`` is written as a literal, a ``ShapeOf`` + reads the forwarded tensor's shape, and dim arithmetic over those composes. + Anything else raises -- there is no silent zero. """ - sv = static_dim_value(expr) - if sv is not None: - return str(sv) - if isinstance(expr, (ShapeOf, Call, Constant)): - return _HostIntExprVisitor().visit(expr) - raise ValueError(f"emit_host_module: unsupported launch-extent node {type(expr).__name__}") - - -def _hidden_names(params) -> set: - return {p.name for p in params if _is_hidden_shape_scalar(p, params)} - - -def _is_user_scalar(p, hidden: set) -> bool: - return p.name not in hidden and isinstance(p.type, TensorType) and not p.type.shape - + value = static_dim_value(expr) + if value is not None: + return str(value) + if not isinstance(expr, (ShapeOf, Call, Constant)): + raise ValueError(f"emit_host_module: unsupported launch-extent node {type(expr).__name__}") + return f"static_cast({_HostIntExprVisitor().visit(expr)})" -def _is_tensor(p, hidden: set) -> bool: - return p.name not in hidden and not _is_user_scalar(p, hidden) +def _static_smem(value) -> int: + """The dynamic shared-memory size, which the host states as a constant. -def _extent(c): - value = static_dim_value(c) - return str(value) if value is not None else f"static_cast({_emit_host_int_expr(c)})" + Host index-expr codegen covers grid and block extents only; a shared-memory + size that is decided at run time has nowhere to be decided. + """ + static = static_dim_value(value) + if static is None: + raise ValueError( + "emit_host_module: dynamic_smem must be a static int/Constant; a " + f"dynamic shared-memory size expression is not supported, got " + f"{type(value).__name__}" + ) + return static -def _call_arg(p, host_names, hidden): - name = host_names[p.name] - if _is_tensor(p, hidden): - return f"{name}.data_ptr()" - if p.name in hidden: - return name - return f"static_cast({name})" +def _reject_unsupported_config(cfg: Launch) -> None: + if cfg.cluster is not None: + raise NotImplementedError("emit_host_module: launch `cluster` is not supported yet") + if cfg.stream is not None: + raise NotImplementedError("emit_host_module: launch `stream` is not supported yet") + if cfg.attrs.entries: + raise NotImplementedError("emit_host_module: launch `attrs` are not supported yet") def _placement_line(name: str, storage) -> str: + """Refuse a tensor that is not where the parameter's storage says it is.""" device_type = _STORAGE_DEVICE_TYPE.get(storage) if device_type is None: raise ValueError( @@ -162,314 +139,142 @@ def _placement_line(name: str, storage) -> str: ) -def _shim_decl(fn: PrimFunction) -> str: - """A types-only ``extern "C"`` forward declaration of *fn*'s launch shim.""" - hidden = _hidden_names(fn.params) - tokens = ["void*" if _is_tensor(p, hidden) else "long long" for p in fn.params] - tokens += _LAUNCH_ABI_DECL - return f'extern "C" void {shim_symbol(fn.name)}({", ".join(tokens)});' +def _launch_statements(entry: PrimFunction) -> tuple[Evaluate, ...]: + """The launches *entry* makes, which is the whole of what a host entry does.""" + body = entry.body + statements = body.body if isinstance(body, Sequential) else () + if not statements or not all( + isinstance(stmt, Evaluate) and isinstance(stmt.callable, Launch) for stmt in statements + ): + raise ValueError( + f"emit_host_module: entry {entry.name!r} body must be one or more Launch statements" + ) + return tuple(statements) -def emit_host_module( - module: Module, functions: tuple[PrimFunction, ...], target: Target -) -> LinkableModule: - """Emit the host ``.cpp`` linkable module for a CPU *entry*. +def _as_this_scope_names_it( + shim: CallableSignature, forwarded: tuple[Var, ...], launch: Evaluate +) -> CallableSignature: + """*shim*'s parameters rewritten as the values this launch passes for them. - *module* is the enclosing ``Module``; the dispatch path resolves each - case's ``SymbolRef`` callee through ``module.lookup`` to read the variant's - parameters. + Order and kind are the callee's and do not change; every name does, because + a caller writes what its own scope calls the thing. The ids the target + states keep their names: the entry was told them under exactly those. """ - if len(functions) != 1: - raise ValueError("emit_host_module: expected exactly one CPU host entry") - entry = functions[0] - body = entry.body - if ( - isinstance(body, Sequential) - and body.body - and all( - isinstance(stmt, Evaluate) and isinstance(stmt.callable, Launch) for stmt in body.body - ) - ): - if len(body.body) > 1: - shim_decls, body_lines, sig = _lower_launches(entry, body.body, module) - else: - shim_decls, body_lines, sig = _lower_launch(entry, body.body[0], module) - elif entry.variants: - pat = entry.variants[0].specializations[0] - loc_idx = locate_dim_var(entry.params, pat.dim_var) - loc = (entry.params[loc_idx[0]], loc_idx[1]) if loc_idx is not None else None - if loc is None: - raise ValueError("emit_host_module: cannot derive specialization subject") - p, axis = loc - subject = ShapeOf(type=TensorType.scalar(DType.i32), param=p, axis=axis) - - def _variant_call(v): - args = [] - for vp in v.params: - parsed = _parse_shape_param_name(vp.name) - if parsed is None: - args.append(next(p for p in entry.params if p.name == vp.name)) - else: - base, ax = parsed - ep = next(p for p in entry.params if p.name == base) - args.append(ShapeOf(type=vp.type, param=ep, axis=ax)) - return symbol_call(v, tuple(args)) - - calls = tuple(_variant_call(v) for v in entry.variants) - shim_decls, body_lines, sig = _lower_dispatch( - entry, - module, - callee_name=entry.name, - subject=subject, - variants=entry.variants, - calls=calls, - ) - else: - raise ValueError(f"emit_host_module: entry {entry.name!r} body must be a single Launch") - source = render( - "cpu_module.cpp.j2", - shim_decls=shim_decls, - internal_host_symbol=_internal_wrapper_symbol(entry.name), - wrapper_params_sig=sig, - body_lines=body_lines, - entry_name=entry.name, + geometry = ( + _extent(launch.args[1]), + _extent(launch.args[4]), + str(_static_smem(launch.callable.dynamic_smem)), ) - return LinkableModule( - target="cpu", - language="cpp", - source=source, - functions=(LinkableFunction(name=entry.name, source=source),), + return replace( + shim, + params=tuple(tensor_signature_of(var) for var in forwarded), + trailing=tuple( + LaunchSignature(name=value, ctype=signature.ctype) + for signature, value in zip(shim.trailing, geometry) + ), ) -CPU_CODE_GENERATOR = CodeGenerator(emit_host_module) - - -def _lower_launch(entry: PrimFunction, evaluate, module): - device_fn = module.lookup(evaluate.args[0].name) - if device_fn.variants: - dim_name = device_fn.variants[0].specializations[0].dim_var - loc_idx = locate_dim_var(device_fn.params, dim_name) - loc = (device_fn.params[loc_idx[0]], loc_idx[1]) if loc_idx is not None else None - if loc is None: - raise ValueError(f"cannot derive specialization subject {dim_name!r}") - p, axis = loc - subject = ShapeOf(type=TensorType.scalar(DType.i32), param=p, axis=axis) - calls = tuple(symbol_call(v, tuple(evaluate.args[1:])) for v in device_fn.variants) - return _lower_dispatch( - entry, - module, - callee_name=device_fn.name, - subject=subject, - variants=device_fn.variants, - calls=calls, +def _forwarded(entry: PrimFunction, launch: Evaluate, callee: PrimFunction) -> tuple[Var, ...]: + """The entry parameters this launch hands the callee, in the callee's order.""" + args = launch.args[7:] + if len(args) != len(callee.params): + raise ValueError( + f"emit_host_module: launch passes {len(args)} args but device function " + f"{callee.name!r} has {len(callee.params)} parameters" ) - return _lower_launches(entry, (evaluate,), module) - - -def _lower_launches(entry: PrimFunction, evaluates, module): - bindings = [] - kinds = {} - used = set() - for evaluate in evaluates: - launch_op = evaluate.callable - device_fn = module.lookup(evaluate.args[0].name) - if device_fn.variants: - raise NotImplementedError( - "emit_host_module: multi-launch entries do not support variants" - ) - dev_params = device_fn.params - _reject_unsupported_config(launch_op) - hidden = _hidden_names(dev_params) - visible_dev = [p for p in dev_params if p.name not in hidden] - args = evaluate.args[1:] - if len(args) != len(visible_dev): - raise ValueError( - f"emit_host_module: launch passes {len(args)} args but device " - f"function {device_fn.name!r} has {len(visible_dev)} host-visible " - "parameters (hidden shape scalars are derived from tensor shapes)" - ) - if not all(isinstance(a, Var) for a in args): + held = {param.name: param for param in entry.params} + for arg in args: + if not isinstance(arg, Var): raise ValueError( "emit_host_module: launch args must be host entry parameters (Var); " "expressions are not accepted" ) - bound = [] - for arg in args: - ep = next((p for p in entry.params if p.name == arg.name), None) - if ep is None: - raise ValueError( - f"emit_host_module: launch arg {arg.name!r} is not a parameter " - f"of entry {entry.name!r}" - ) - bound.append(ep) - local_seen = set() - host_name_of = {} - for vp, ep in zip(visible_dev, bound): - if id(ep) in local_seen: - raise ValueError( - f"emit_host_module: entry parameter {ep.name!r} is bound more " - "than once in one launch" - ) - local_seen.add(id(ep)) - used.add(id(ep)) - kind = _is_tensor(vp, hidden) - if id(ep) in kinds and kinds[id(ep)] != kind: - raise ValueError( - f"emit_host_module: entry parameter {ep.name!r} is bound to " - "incompatible device parameter types across launches" - ) - kinds[id(ep)] = kind - host_name_of[vp.name] = ep.name - bindings.append((evaluate, device_fn, hidden, host_name_of)) - for ep in entry.params: - if id(ep) not in used: + if arg.name not in held: raise ValueError( - f"emit_host_module: entry parameter {ep.name!r} is not used by any " - f"launch in {entry.name!r} -- an unused parameter has no device " - "type to give the wrapper's signature" + f"emit_host_module: launch arg {arg.name!r} is not a parameter " + f"of entry {entry.name!r}" ) - body_lines = [] - shim_decls = [] - - for index, (evaluate, device_fn, hidden, dev_to_host) in enumerate(bindings): - dev_params = device_fn.params - launch_op = evaluate.callable - host_names = { - **dev_to_host, - **{p.name: f"l{index}__{p.name}" for p in dev_params if p.name in hidden}, - } - for p in dev_params: - if _is_tensor(p, hidden): - body_lines.append(_placement_line(host_names[p.name], p.type.storage)) - for p in dev_params: - if p.name in hidden: - base, axis = _parse_shape_param_name(p.name) - host_base = dev_to_host.get(base) - if host_base is None: - raise ValueError( - f"emit_host_module: hidden shape scalar {p.name!r} " - f"references unknown base parameter {base!r}" - ) - body_lines.append( - f"long long {host_names[p.name]} = static_cast(" - f"{host_base}.shape()[{axis}]);" - ) - - owner = owning_module(module, device_fn) or module - launch = LaunchSignature.from_topologies(owner.effective_topologies()) - call_args = [_call_arg(p, host_names, hidden) for p in dev_params] + [ - *(str(d) for d in (*launch.grid, *launch.block)), - str(_static_smem(launch_op.dynamic_smem)), - "nullptr", - ] - body_lines.append(f"{shim_symbol(device_fn.name)}({', '.join(call_args)});") - shim_decls.append(_shim_decl(device_fn)) - sig = ", ".join( - f"tvm::ffi::Tensor {ep.name}" if kinds[id(ep)] else f"int {ep.name}" for ep in entry.params + return tuple(held[arg.name] for arg in args) + + +def _one_launch( + entry: PrimFunction, launch: Evaluate, callee: PrimFunction, ctx: CpuCodegenContext +) -> list[str]: + """The lines one ``Launch`` becomes: the checks it owes, then the call.""" + _reject_unsupported_config(launch.callable) + forwarded = _forwarded(entry, launch, callee) + shim = ctx.signature_of(callee) + bound = _as_this_scope_names_it(shim, forwarded, launch) + lines = [ + _placement_line(var.name, param.type.storage) + for var, param in zip(forwarded, callee.params) + ] + lines.append(f"{shim.name}({ctx.arguments(bound, callee.target)});") + return lines + + +def _declare_callee(callee: PrimFunction, ctx: CpuCodegenContext) -> str: + """The forward declaration of a symbol another unit defines. + + Written by the callee's own target, so this unit and the one that defines + it cannot come to disagree about what the symbol takes. + """ + shim = ctx.signature_of(callee) + return f'extern "C" void {shim.name}({ctx.parameters(shim, callee.target, exported=True)});' + + +def _host_entry( + entry: PrimFunction, + signature: CallableSignature, + body_lines: list[str], + ctx: CpuCodegenContext, +) -> LinkableFunction: + """*entry* in both of its positions, written from the one signature it is called by.""" + return LinkableFunction( + name=entry.name, + declaration=f"void {signature.name}({ctx.parameters(signature)});", + definition=render( + "cpu_entry.cpp.j2", + host_symbol=signature.name, + params=ctx.parameters(signature), + body_lines=body_lines, + entry_name=entry.name, + ), ) - return shim_decls, body_lines, sig - - -def _lower_dispatch(entry: PrimFunction, module, *, callee_name, subject, variants, calls): - if not isinstance(subject, ShapeOf): - raise NotImplementedError( - "emit_host_module: dispatch v1 expects exactly one ShapeOf subject" - ) - for variant in variants: - if len(variant.specializations) != 1 or not isinstance( - variant.specializations[0], DimVarRangePat - ): - raise NotImplementedError( - "emit_host_module: dispatch v1 expects exactly one DimVarRangePat per case" - ) - entry_params = entry.params - entry_names = {p.name for p in entry_params} - hidden = _hidden_names(entry_params) - def _host_name(ref) -> str: - nm = ref.name if isinstance(ref, Var) else ref.param.name - if nm not in entry_names: - raise ValueError( - f"emit_host_module: dispatch arg {nm!r} is not a parameter of entry {entry.name!r}" - ) - return nm - - wrapper_tokens, body_lines = [], [] - for p in entry_params: - if p.name in hidden: - continue - if _is_user_scalar(p, hidden): - wrapper_tokens.append(f"int {p.name}") - else: - wrapper_tokens.append(f"tvm::ffi::Tensor {p.name}") - body_lines.append(_placement_line(p.name, p.type.storage)) - - subj = subject - s = "__tf_dispatch_subject" - body_lines.append( - f"long long {s} = static_cast({_host_name(subj)}.shape()[{subj.axis}]);" - ) +def emit_host_module( + module: Module, + functions: tuple[PrimFunction, ...], + target: Target, + ctx: CpuCodegenContext, +) -> LinkableModule: + """Emit the host ``.cpp`` linkable module for a CPU *entry*. - for variant, call in zip(variants, calls): - if len(call.args) != len(variant.params): - raise ValueError( - f"emit_host_module: dispatch call to {variant.name!r} passes " - f"{len(call.args)} args for {len(variant.params)} parameters" - ) - shim_decls: dict[str, str] = {} - for idx, (variant, call) in enumerate(zip(variants, calls)): - pat = variant.specializations[0] - variant_symbol = variant.name - shim_decls[shim_symbol(variant_symbol)] = _shim_decl(variant) - if len(call.args) != len(variant.params): - raise ValueError( - f"emit_host_module: dispatch call to {variant.name!r} passes " - f"{len(call.args)} args for {len(variant.params)} parameters" - ) - v_hidden = _hidden_names(variant.params) - shim_args = [] - for vp, arg in zip(variant.params, call.args): - if _is_tensor(vp, v_hidden): - shim_args.append(f"{_host_name(arg)}.data_ptr()") - elif vp.name in v_hidden: - if not isinstance(arg, ShapeOf): - raise NotImplementedError( - f"emit_host_module: hidden shape param {vp.name!r} expects a ShapeOf arg" - ) - shim_args.append(f"static_cast({_host_name(arg)}.shape()[{arg.axis}])") - else: - shim_args.append(f"static_cast({_host_name(arg)})") - owner = owning_module(module, variant) or module - launch = LaunchSignature.from_topologies(owner.effective_topologies()) - if launch.grid[0] is None: - raise ValueError( - f"emit_host_module: dispatch variant {variant.name!r} has a " - f"launch-provided (dynamic) CTA extent; the dispatch host path " - f"requires a static grid" - ) - shim_args += [str(d) for d in (*launch.grid, *launch.block, 0)] - shim_args.append("nullptr") - pred = f"(({pat.lo} <= {s}) && ({s} <= {pat.hi}))" - prefix = "if" if idx == 0 else "} else if" - body_lines.append(f"{prefix} ({pred}) {{") - body_lines.append(f" {shim_symbol(variant_symbol)}({', '.join(shim_args)});") - body_lines.append("} else {") - body_lines.append( - f' throw std::runtime_error("tilefoundry: no matching dispatch variant for {entry.name}");' + *module* is the enclosing ``Module``, which resolves each launch's callee; + the FFI export names the symbol it exports, so it is written with the + definition and not in the preamble every declaration follows. + """ + if len(functions) != 1: + raise ValueError("emit_host_module: expected exactly one CPU host entry") + entry = functions[0] + declarations: dict[str, str] = {} + body_lines: list[str] = [] + for launch in _launch_statements(entry): + callee = ctx.callee_of(launch) + declarations[ctx.signature_of(callee).name] = _declare_callee(callee, ctx) + body_lines += _one_launch(entry, launch, callee, ctx) + return LinkableModule( + target="cpu", + language="cpp", + preamble=render("cpu_preamble.cpp.j2", shim_decls=list(declarations.values())), + functions=(_host_entry(entry, ctx.signature_of(entry), body_lines, ctx),), ) - body_lines.append("}") - return list(shim_decls.values()), body_lines, ", ".join(wrapper_tokens) -def _reject_unsupported_config(cfg) -> None: - if cfg.cluster is not None: - raise NotImplementedError("emit_host_module: launch `cluster` is not supported yet") - if cfg.stream is not None: - raise NotImplementedError("emit_host_module: launch `stream` is not supported yet") - if cfg.attrs.entries: - raise NotImplementedError("emit_host_module: launch `attrs` are not supported yet") +CPU_CODE_GENERATOR = CodeGenerator(emit_host_module) __all__ = ["emit_host_module"] diff --git a/src/tilefoundry/codegen/cuda/abi.py b/src/tilefoundry/codegen/cuda/abi.py index ce282307e..cabf6d37c 100644 --- a/src/tilefoundry/codegen/cuda/abi.py +++ b/src/tilefoundry/codegen/cuda/abi.py @@ -8,10 +8,10 @@ from __future__ import annotations -from tilefoundry.codegen.context import CodegenContext +from tilefoundry.codegen.context import EmitContext from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.codegen.cuda.tir.stmts.mesh_scope import mesh_type from tilefoundry.codegen.signature import ( - LAUNCH_ABI, CallableSignature, LaunchSignature, ProgramIdSignature, @@ -21,6 +21,7 @@ tensor_signature_of, ) from tilefoundry.ir.tir.prim_function import PrimFunction +from tilefoundry.ir.types.shard.shard_layout import ShardLayout from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen @@ -85,9 +86,10 @@ def kernel_arguments(kernel: CallableSignature, ctx: CudaCodegenContext) -> str: continue element = ctx.dtype_to_cpp(param.type.dtype.name) written.append(f"static_cast<{element}*>({param.name})") - written += [ - f"static_cast({param.extent_name(axis)})" for axis in param.dynamic_axes - ] + written += [f"static_cast({param.extent_name(axis)})" for axis in param.dynamic_axes] + layout = param.type.layout + if isinstance(layout, ShardLayout): + written.append(f"{mesh_type(layout.mesh)}{{}}") return ", ".join(written) @@ -106,12 +108,11 @@ def _called_as_shim( params=tuple(tensor_signature_of(var) for var in fn.params), output_count=fn.output_count, leading=program_ids, - trailing=LAUNCH_ABI, ) @register_codegen(CudaTarget, Role.CALLEE, TensorSignature) -def _declare_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ...]: +def _declare_tensor(sig: TensorSignature, ctx: EmitContext) -> tuple[str, ...]: """A tensor arrives as a pointer, then one extent per open axis. A static extent is a constant on both sides of the call and costs the @@ -123,10 +124,14 @@ def _declare_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ... pointer, extent = _C_ABI else: pointer, extent = f"{ctx.dtype_to_cpp(sig.type.dtype.name)}*", "int" - return ( + declared = [ f"{pointer} {sig.name}", *(f"{extent} {sig.extent_name(axis)}" for axis in sig.dynamic_axes), - ) + ] + layout = sig.type.layout + if not ctx.exported and isinstance(layout, ShardLayout): + declared.append(f"{mesh_type(layout.mesh)} {sig.name}_mesh") + return tuple(declared) @register_codegen(CudaTarget, Role.CALLEE, LaunchSignature) @@ -148,7 +153,7 @@ def _declare_program_meta(sig: ProgramMetaSignature, ctx: CudaCodegenContext) -> @register_codegen(CudaTarget, Role.CALLER, TensorSignature) -def _pass_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ...]: +def _pass_tensor(sig: TensorSignature, ctx: EmitContext) -> tuple[str, ...]: """The pointer, then the extents, as the calling scope spells them. How many arguments there are and in what order is CUDA's, since CUDA @@ -163,7 +168,7 @@ def _pass_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ...]: @register_codegen(CudaTarget, Role.CALLER, LaunchSignature) @register_codegen(CudaTarget, Role.CALLER, ProgramIdSignature) @register_codegen(CudaTarget, Role.CALLER, ProgramMetaSignature) -def _pass_value(sig: Signature, ctx: CodegenContext) -> tuple[str, ...]: +def _pass_value(sig: Signature, ctx: EmitContext) -> tuple[str, ...]: """One argument the caller already holds, spelled in its own scope.""" return (ctx.local_value(sig),) diff --git a/src/tilefoundry/codegen/cuda/context.py b/src/tilefoundry/codegen/cuda/context.py index fe4e673d9..7ad734de5 100644 --- a/src/tilefoundry/codegen/cuda/context.py +++ b/src/tilefoundry/codegen/cuda/context.py @@ -11,14 +11,15 @@ from collections.abc import Iterable, Mapping -from tilefoundry.codegen.context import CodegenContext +from tilefoundry.codegen.context import EmitContext from tilefoundry.codegen.signature import CallableSignature, TensorSignature -from tilefoundry.codegen.topology import Geometry from tilefoundry.ir.types.dim import DimVar from tilefoundry.target import CudaTarget from tilefoundry.target.base import Target from tilefoundry.visitor_registry.registries import codegen_registry +Geometry = tuple[tuple[object, object, object], tuple[object, object, object]] + _CUDA_CPP: dict[str, str] = { "f32": "float", "bf16": "__nv_bfloat16", @@ -47,7 +48,7 @@ def topology_scope_str(name: str) -> str: ) from None -class CudaCodegenContext(CodegenContext): +class CudaCodegenContext(EmitContext): """A compile writing CUDA: the shared context plus what only CUDA states.""" target_kind = CudaTarget @@ -58,8 +59,14 @@ def __init__( symbols: Mapping[int, CallableSignature] | None = None, target: Target | None = None, launches: Mapping[int, Geometry] | None = None, + codegen_context: object | None = None, ) -> None: - super().__init__(codegen_registry, symbols=symbols, target=target) + super().__init__( + codegen_registry, + symbols=symbols, + target=target, + codegen_context=codegen_context, + ) self._mesh_aliases: dict[int, tuple[str, str]] = {} self.launches: Mapping[int, Geometry] = {} if launches is None else launches """The geometry each device function is launched at, keyed by ``id(fn)``.""" diff --git a/src/tilefoundry/codegen/cuda/module.py b/src/tilefoundry/codegen/cuda/module.py index 0d982308f..a01bb7c19 100644 --- a/src/tilefoundry/codegen/cuda/module.py +++ b/src/tilefoundry/codegen/cuda/module.py @@ -1,125 +1,176 @@ -"""CUDA device linkable-module emitter (split pipeline). - -Emits the device ``.cu`` translation unit: the ``__global__`` kernels (identical -to the single-source path) plus, per kernel, an ``extern "C"`` launch shim that -performs the ``<<>>`` launch. Grid / block / smem / -stream arrive as plain C ABI arguments from the host module, so all CUDA types -stay inside this translation unit. +"""The device ``.cu`` translation unit: one kernel and one shim per function. + +Two passes over each function and nothing between them: one writes the +declaration a caller sees without entering the body, the other writes the +definition by walking it. What is not a function -- the includes, the program +dimensions this unit specializes, the device state something asked for while +it was being written -- settles only once every body has been walked, so it is +assembled last. """ + from __future__ import annotations -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.codegen.cuda.templates import render -from tilefoundry.codegen.cuda.tir.prim_function import ( - _compute_kernel_fields, +import math + +from tilefoundry.codegen.cuda.abi import ( + PROGRAM_META, + PROGRAM_META_CTYPE, + kernel_arguments, + kernel_signature, ) +from tilefoundry.codegen.cuda.context import CudaCodegenContext, topology_scope_str +from tilefoundry.codegen.cuda.templates import render from tilefoundry.codegen.linkable import LinkableFunction, LinkableModule from tilefoundry.codegen.registry import CodeGenerator -from tilefoundry.ir.core.module import Module, owning_module +from tilefoundry.codegen.signature import CallableSignature +from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.prim_function import PrimFunction +from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.target import Target +Geometry = tuple[tuple[object, object, object], tuple[object, object, object]] -def shim_symbol(kernel_name: str) -> str: - """The ``extern "C"`` launch-shim symbol for *kernel_name*. - The ``extern "C"`` launch-shim symbol for *kernel_name* (a plain C - identifier; mangled-variant ``$`` is replaced). +def _emit_function(fn: PrimFunction, ctx: CudaCodegenContext) -> LinkableFunction: + """*fn* in both of its positions, written from the one convention it is called by. + + A caller can name the shim and nothing else, so that is what the + declaration states; the kernel is defined beside it and named only here. """ - return "tilefoundry_" + kernel_name.replace("$", "__") + "_launch" + shim = ctx.signature_of(fn) + kernel = kernel_signature(shim, fn.name) + return LinkableFunction( + name=fn.name, + declaration=f'extern "C" void {shim.name}({ctx.parameters(shim, exported=True)});', + definition=ctx.capture(lambda scope: _define(fn, shim, kernel, scope)), + ) +def _define( + fn: PrimFunction, + shim: CallableSignature, + kernel: CallableSignature, + ctx: CudaCodegenContext, +) -> None: + """The kernel and the shim that launches it, in that order.""" + _define_kernel(fn, kernel, ctx) + ctx.blank() + _define_shim(shim, kernel, ctx) + + +def _define_kernel(fn: PrimFunction, kernel: CallableSignature, ctx: CudaCodegenContext) -> None: + """The ``__global__``: what it was told, then what it was written to do.""" + ctx.reset_barrier_ids() + ctx.emit(f"__global__ void {kernel.name}({ctx.parameters(kernel)}) {{") + ctx.indent() + if kernel.leading: + ctx.emit(f"tilefoundry::program_meta({PROGRAM_META.name});") + ctx.emit_node(fn) + ctx.dedent() + ctx.emit("}") + + +def _define_shim( + shim: CallableSignature, kernel: CallableSignature, ctx: CudaCodegenContext +) -> None: + """The ``extern "C"`` entry: the geometry it was handed, spent on one launch.""" + ctx.emit(f'extern "C" void {shim.name}({ctx.parameters(shim, exported=True)}) {{') + ctx.indent() + _tell_the_block(shim, ctx) + grid_x, block_x, smem = (signature.name for signature in shim.trailing) + ctx.emit(f"dim3 grid({grid_x}, 1, 1);") + ctx.emit(f"dim3 block({block_x}, 1, 1);") + ctx.emit(f"{kernel.name}<<>>({kernel_arguments(kernel, ctx)});") + ctx.dedent() + ctx.emit("}") + + +def _tell_the_block(shim: CallableSignature, ctx: CudaCodegenContext) -> None: + """Gather the ids the shim was told into the block the kernel takes. + + Which levels these are is the signature's answer, so the shim states no + level of its own: one line per id the convention put in front. + """ + if not shim.leading: + return + ctx.emit(f"{PROGRAM_META_CTYPE} {PROGRAM_META.name}{{}};") + for signature in shim.leading: + scope = topology_scope_str(signature.topology_level) + ctx.emit( + f"{PROGRAM_META.name}.program_id[int({scope})] = static_cast({signature.name});" + ) -_LAUNCH_ABI_PARAMS = ( - "int grid_x, int grid_y, int grid_z, " - "int block_x, int block_y, int block_z, " - "int dynamic_smem, void* stream" -) +def _one_geometry(kernels: tuple[PrimFunction, ...], ctx: CudaCodegenContext) -> Geometry: + """The geometry this unit runs at, which every kernel in it states the same. -def _emit_kernel_and_shim(fields) -> str: - shim_params = [] - shim_casts = [] - call_args = [] - for p in fields.params: - if fields.param_kinds[p.name] == "tensor": - cpp = fields.param_cpp_types[p.name] - shim_params.append(f"void* {p.name}") - shim_casts.append(f"{cpp}* {p.name}_p = static_cast<{cpp}*>({p.name});") - call_args.append(f"{p.name}_p") - else: - shim_params.append(f"long long {p.name}") - shim_casts.append(f"int {p.name}_i = static_cast({p.name});") - call_args.append(f"{p.name}_i") - shim_params_sig = ", ".join((*shim_params, _LAUNCH_ABI_PARAMS)) - return render( - "device_kernel.cu.j2", - kernel_name=fields.kernel_name, - kernel_params_sig=fields.kernel_params_sig, - param_wrappers=fields.param_wrappers, - kernel_body=fields.kernel_body, - shim_name=shim_symbol(fields.kernel_name), - shim_params_sig=shim_params_sig, - shim_casts=shim_casts, - kernel_call_args=", ".join(call_args), - ) + A unit specializes ``program_dim`` once, so two kernels launched at + different geometries cannot share one: they are two units. + """ + by_kernel = {fn.name: _launched_at(fn, ctx) for fn in kernels} + distinct = set(by_kernel.values()) + if len(distinct) > 1: + raise ValueError( + f"emit_cuda_module: kernels disagree on launch geometry ({by_kernel}); " + f"one translation unit runs every kernel in it at one geometry" + ) + return distinct.pop() + + +def _launched_at(fn: PrimFunction, ctx: CudaCodegenContext) -> Geometry: + """The geometry *fn*'s call states, which is the only thing that states one.""" + geometry = ctx.launches.get(id(fn)) + if geometry is None: + raise ValueError( + f"emit_cuda_module: nothing launches {fn.name!r}, so no call " + f"states the geometry its unit runs at" + ) + return geometry + + +def _program_dims(geometry: Geometry) -> list[dict[str, str]]: + """One instance count per level: what ``program_dim`` states for this unit. + + ``program_shape`` is derived from these, so a unit specializes the counts + and nothing else. A launch-provided grid has no count here; the runtime + states that one when it is asked. + """ + grid, block = geometry + dims = [] + static_grid = tuple(static_dim_value(value) for value in grid) + if all(value is not None for value in static_grid): + dims.append({"scope": topology_scope_str("cta"), "count": str(math.prod(static_grid))}) + elif static_grid[1:] != (1, 1): + raise ValueError("codegen: only grid.x may be a host-computed dynamic expression") + static_block = tuple(static_dim_value(value) for value in block) + if any(value is None for value in static_block): + raise ValueError("codegen: thread topology extents must be compile-time constants") + dims.append({"scope": topology_scope_str("thread"), "count": str(math.prod(static_block))}) + return dims def emit_cuda_module( - module: Module, cuda_fns: tuple[PrimFunction, ...], target: Target + module: Module, cuda_fns: tuple[PrimFunction, ...], target: Target, ctx: CudaCodegenContext ) -> LinkableModule: - """Emit the device ``.cu`` linkable module for *cuda_fns*. - - Emit the device ``.cu`` linkable module for *cuda_fns* (the CUDA-target - PrimFunctions of a module). - """ - from tilefoundry.codegen.cuda.emit import _topology_dim_specializations # noqa: PLC0415 - - kernel_texts = [] - all_fields = [] - expanded_fns = tuple( - {v.name: v for fn in cuda_fns for v in (fn.variants or (fn,))}.values() - ) - for fn in expanded_fns: - owner = owning_module(module, fn) or module - ctx = CudaCodegenContext(target, owner.effective_topologies()) - fields = _compute_kernel_fields(fn, ctx) - kernel_texts.append(_emit_kernel_and_shim(fields)) - all_fields.append(fields) - if not all_fields: - raise ValueError("emit_cuda_module: no CUDA device kernels") - - - - base = all_fields[0] - for f in all_fields[1:]: - if f.launch != base.launch: - raise ValueError( - "emit_cuda_module: kernels disagree on launch topology — " - f"{base.kernel_name!r} has {base.launch!r} " - f"but {f.kernel_name!r} has {f.launch!r}" - ) - specs = _topology_dim_specializations(base.launch.grid, base.launch.block) - - - uses_grid_barrier = any("tf_grid_bar_state" in text for text in kernel_texts) - source = render( - "cuda_module.cu.j2", - topology_dim_specializations=specs, - kernels="\n".join(kernel_texts), - dynamic_cta=base.launch.grid[0] is None, - uses_grid_barrier=uses_grid_barrier, - ) - functions = tuple( - LinkableFunction(name=f.kernel_name, source=text) - for f, text in zip(all_fields, kernel_texts) - ) - return LinkableModule( - target="cuda", language="cu", source=source, functions=functions + """Emit the device ``.cu`` linkable module for one topology domain.""" + kernels = tuple(fn for fn in cuda_fns if fn.target == target) + if not kernels: + raise ValueError(f"emit_cuda_module: module {module.name!r} has no CUDA device kernels") + functions = tuple(_emit_function(fn, ctx) for fn in kernels) + geometry = _one_geometry(kernels, ctx) + preamble = render( + "cuda_preamble.cu.j2", + program_topology=topology_scope_str( + module.effective_topologies()[0].name if module.effective_topologies() else "cta" + ), + program_dims=_program_dims(geometry), + dynamic_cta=static_dim_value(geometry[0][0]) is None, + grid_barrier_state=ctx.needs_grid_barrier_state, ) + return LinkableModule(target="cuda", language="cu", preamble=preamble, functions=functions) CUDA_CODE_GENERATOR = CodeGenerator(emit_cuda_module) -__all__ = ["CUDA_CODE_GENERATOR", "emit_cuda_module", "shim_symbol"] +__all__ = ["CUDA_CODE_GENERATOR", "emit_cuda_module"] diff --git a/src/tilefoundry/codegen/cuda/tir/prim_function.py b/src/tilefoundry/codegen/cuda/tir/prim_function.py index a623b2dec..8ee019cc6 100644 --- a/src/tilefoundry/codegen/cuda/tir/prim_function.py +++ b/src/tilefoundry/codegen/cuda/tir/prim_function.py @@ -1,233 +1,105 @@ -"""Emit a TIR primitive function's CUDA kernel and host wrapper. +"""Emit the body of one CUDA kernel: its buffers, then its statements. -The wrapper accepts TVM tensors, extracts pointers and hidden integer shape -scalars, and derives launch dimensions from mesh scopes. A specialization entry -containing only dispatch emits no global kernel; its host wrapper selects and -calls variant wrappers because CUDA kernels cannot call host code. -See [codegen §1](docs/spec/codegen.md#1-pipeline). +A kernel is handed pointers and the extents its parameter types leave open, so +the body starts by wrapping each pointer in the cute tensor the ops index. A +specialization prototype has no body of its own: its variants are branches of +the one kernel, chosen on the extent the pattern ranges over, because every +kernel in a translation unit runs at the geometry that unit states once. """ -from __future__ import annotations -from dataclasses import dataclass +from __future__ import annotations -from tilefoundry.codegen.cuda.context import CodegenContext +from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.codegen.cuda.tir.memory.tensor_view import render_shard_layout_value from tilefoundry.codegen.emitter import CudaEmitter -from tilefoundry.codegen.signature import LaunchSignature +from tilefoundry.codegen.signature import TensorSignature, tensor_signature_of +from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.tir.shape import ( - is_hidden_shape_scalar as _is_hidden_shape_scalar, -) -from tilefoundry.ir.tir.shape import ( - parse_shape_var_name as _parse_shape_param_name, -) -from tilefoundry.ir.tir.shape import ( - shape_var_name, -) -from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.shape_helpers import shape_numel_upper_bound from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.target import CudaTarget +from tilefoundry.visitor_registry.registries import Role, register_codegen -def _internal_wrapper_symbol(kernel_name: str) -> str: - """Map a user-facing kernel name to its internal C++ wrapper symbol. +def _wrap_buffer(signature: TensorSignature, ctx: CudaCodegenContext) -> None: + """Wrap one parameter's pointer in the tensor the body indexes it through. - The user-facing name may be ``main`` (collides with ``::main``) or - a mangled variant like ``main$S$1_4`` (``$`` is a GCC extension, - not portable). The internal symbol is always a plain C++ identifier - so the generated source compiles under strict toolchains. + The layout is the envelope the whole dispatch range fits in, so one kernel + covers every shape its type admits; a sharded parameter carries which part + of that envelope this instance holds. """ - return "__tilefoundry_" + kernel_name.replace("$", "__") + "_host" - - -def _param_wrapper(name: str, total: int, cpp_type: str) -> str: + name = signature.name + total = shape_numel_upper_bound(signature.type.shape) layout = f"cute::make_layout(cute::Shape>{{}})" - return ( - f"auto {name}_tensor = cute::make_tensor(" - f"cute::make_gmem_ptr({name}), {layout});" + pointer = f"cute::make_tensor(cute::make_gmem_ptr({name}), {layout})" + shard = signature.type.layout + if not isinstance(shard, ShardLayout): + ctx.emit(f"auto {name}_tensor = {pointer};") + return + preamble, value = render_shard_layout_value( + f"{name}_tensor", shard, ctx.dynamic_extents, signature.type.storage, ctx ) + for line in preamble: + ctx.emit(line) + ctx.emit(f"auto {name}_tensor = tilefoundry::make_shard_tensor({pointer}, {layout}, {value});") -def _param_cpp_types(params: tuple, ctx: CodegenContext) -> dict[str, str]: - """Map each param name → CUDA C++ type from its TensorType dtype.""" - result: dict[str, str] = {} - for p in params: - ty = p.type - if isinstance(ty, TensorType): - result[p.name] = ctx.dtype_to_cpp(ty.dtype.name) - else: - result[p.name] = "float" - return result +def _dispatch(fn: PrimFunction, ctx: CudaCodegenContext) -> None: + """Write the variants of *fn* as branches on the extent they range over. - -def _param_wrapper_shard( - name: str, total: int, shard_layout: ShardLayout, dim_var_runtime=None -) -> str: - """Emit make_shard_tensor wrapping a kernel param with ShardLayout value.""" - global_layout = f"cute::make_layout(cute::Shape>{{}})" - tensor_ref = f"{name}_tensor" - preamble, shard_value = render_shard_layout_value( - tensor_ref, shard_layout, dim_var_runtime, None, ctx - ) - wrapper = ( - f"auto {tensor_ref} = tilefoundry::make_shard_tensor(" - f"cute::make_tensor(cute::make_gmem_ptr({name}), {global_layout}), " - f"{global_layout}, {shard_value});" - ) - return "\n".join([*preamble, wrapper]) - - -@dataclass(frozen=True) -class _KernelFields: - """Represent KernelFields. - - Per-PrimFunction codegen pieces shared by the single-source kernel - emitter and the split device / host fragment emitters. + The extent is already a parameter of this kernel, so the choice costs one + comparison inside the launch the host would otherwise have had to make + before it -- and the host has nothing to gain by making it, since all the + variants run at one geometry anyway. """ - - kernel_name: str - internal_wrapper_name: str - params: tuple - param_cpp_types: dict - param_kinds: dict - kernel_params_sig: str - wrapper_params_sig: str - user_params: tuple - launch_args: str - param_wrappers: list - wrapper_locals: list - kernel_body: str - launch: LaunchSignature - entry_host_only: bool - - -def _compute_kernel_fields(node: PrimFunction, ctx: CodegenContext) -> _KernelFields: - if node.variants: - raise ValueError( - f"PrimFunction {node.name!r} is a specialization prototype; " - "codegen must dispatch variants instead of emitting an empty kernel" - ) - for p in node.params: - ctx.register_kernel_param(p) - - - - - - - ctx._dim_var_runtime = {} - for p in node.params: - ty = p.type - if not isinstance(ty, TensorType): - continue - for axis, dim in enumerate(ty.shape): - if isinstance(dim, DimVar) and dim.name not in ctx._dim_var_runtime: - ctx._dim_var_runtime[dim.name] = shape_var_name(p.name, axis) - - ctx.reset_barrier_ids() - - entry_host_only = bool(node.variants) - - - - body = ctx.capture(lambda inner: CudaEmitter(context=inner).visit(node.body)) - - param_cpp_types = _param_cpp_types(node.params, ctx) - hidden_shape = tuple( - p for p in node.params if _is_hidden_shape_scalar(p, node.params) - ) - hidden_names = {p.name for p in hidden_shape} - user_params = tuple(p for p in node.params if p.name not in hidden_names) - - - - - def _is_user_scalar(p) -> bool: - return ( - p.name not in hidden_names - and isinstance(p.type, TensorType) - and not p.type.shape + ctx.bind_extents(tensor_signature_of(param) for param in fn.params) + for index, variant in enumerate(fn.variants): + pattern = _range_over_one_dimension(fn, variant) + subject = _subject(fn, pattern, ctx) + opening = "if" if index == 0 else "} else if" + ctx.emit(f"{opening} (({pattern.lo} <= {subject}) && ({subject} <= {pattern.hi})) {{") + ctx.indent() + ctx.emit_node(variant) + ctx.dedent() + ctx.emit("} else {") + ctx.indent() + ctx.emit("__trap();") + ctx.dedent() + ctx.emit("}") + + +def _range_over_one_dimension(fn: PrimFunction, variant: PrimFunction) -> DimVarRangePat: + """The pattern selecting *variant*, which a branch can only be a range.""" + pattern = variant.specializations[0] + if not isinstance(pattern, DimVarRangePat): + raise NotImplementedError( + f"CUDA dispatch: {fn.name!r} selects {variant.name!r} by " + f"{type(pattern).__name__}; only a range over one dimension is written" ) + return pattern - def _kind(p) -> str: - if p.name in hidden_names: - return "hidden_scalar" - if _is_user_scalar(p): - return "user_scalar" - return "tensor" - - param_kinds = {p.name: _kind(p) for p in node.params} - - def _kernel_sig_token(p) -> str: - if param_kinds[p.name] != "tensor": - return f"int {p.name}" - return f"{param_cpp_types[p.name]}* {p.name}" - - kernel_params_sig = ", ".join(_kernel_sig_token(p) for p in node.params) - - - - - def _wrapper_param_token(p) -> str: - if param_kinds[p.name] == "user_scalar": - return f"int {p.name}" - return f"tvm::ffi::Tensor {p.name}" - wrapper_params_sig = ", ".join(_wrapper_param_token(p) for p in user_params) - wrapper_locals = [] - for p in hidden_shape: - parsed = _parse_shape_param_name(p.name) - - assert parsed is not None - base, axis = parsed - wrapper_locals.append( - f"int {p.name} = static_cast({base}.shape()[{axis}]);" +def _subject(fn: PrimFunction, pattern: DimVarRangePat, ctx: CudaCodegenContext) -> str: + """Where this kernel reads the dimension *pattern* ranges over.""" + subject = ctx.dynamic_extents.get(pattern.dim_var) + if subject is None: + raise ValueError( + f"CUDA dispatch: {fn.name!r} selects on {pattern.dim_var!r}, which " + f"no parameter's type leaves open, so the kernel is never told it" ) - - - - def _launch_arg(p) -> str: - if param_kinds[p.name] != "tensor": - return p.name - return f"static_cast<{param_cpp_types[p.name]}*>({p.name}.data_ptr())" - - launch_args = ", ".join(_launch_arg(p) for p in node.params) - - - - buffer_params = tuple( - p for p in user_params if param_kinds[p.name] == "tensor" - ) - param_wrappers = [] - for p in buffer_params: - total = shape_numel_upper_bound(p.type.shape) - sl = getattr(p.type, "layout", None) - if isinstance(sl, ShardLayout): - param_wrappers.append( - _param_wrapper_shard(p.name, total, sl, ctx._dim_var_runtime) - ) - else: - param_wrappers.append( - _param_wrapper(p.name, total, param_cpp_types[p.name]) - ) - - launch = LaunchSignature.from_topologies(ctx.topologies) - - codegen_name = node.name - return _KernelFields( - kernel_name=codegen_name, - internal_wrapper_name=_internal_wrapper_symbol(codegen_name), - params=node.params, - param_cpp_types=param_cpp_types, - param_kinds=param_kinds, - kernel_params_sig=kernel_params_sig, - wrapper_params_sig=wrapper_params_sig, - user_params=user_params, - launch_args=launch_args, - param_wrappers=param_wrappers, - wrapper_locals=wrapper_locals, - kernel_body=body, - launch=launch, - entry_host_only=entry_host_only, - ) + return subject + + +@register_codegen(CudaTarget, Role.EMIT, PrimFunction) +def _emit(fn: PrimFunction, ctx: CudaCodegenContext) -> None: + """Write what runs inside one ``__global__``: its buffers, then its statements.""" + if fn.variants: + _dispatch(fn, ctx) + return + params = tuple(tensor_signature_of(param) for param in fn.params) + for param in fn.params: + ctx.register_kernel_param(param) + ctx.bind_extents(params) + for signature in params: + _wrap_buffer(signature, ctx) + CudaEmitter(context=ctx).visit(fn.body) diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/abort.py b/src/tilefoundry/codegen/cuda/tir/stmts/abort.py deleted file mode 100644 index e0e2d42f2..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/abort.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Emit the TIR abort effect.""" - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.abort import Abort -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - - -@register_codegen(CudaTarget, Role.EMIT, Abort) -def _emit(node: Abort, ctx: CudaCodegenContext) -> None: - ctx.emit("__trap();") diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/for_.py b/src/tilefoundry/codegen/cuda/tir/stmts/for_.py deleted file mode 100644 index 56967de56..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/for_.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Emitter for ``tir.For`` — C-style for loop.""" - -from __future__ import annotations - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import For -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - -from .scalar_expr import render_scalar_expr - - -@register_codegen(CudaTarget, Role.EMIT, For) -def _emit(node: For, ctx: CudaCodegenContext) -> None: - iv_name = ctx.name_for(node.induction_var) - start = render_scalar_expr(node.start, ctx) - stop = render_scalar_expr(node.stop, ctx) - step = render_scalar_expr(node.step, ctx) - - ctx.emit(f"for (int {iv_name} = {start}; {iv_name} < {stop}; {iv_name} += {step}) {{") - ctx.indent() - ctx.emit_node(node.body) - ctx.dedent() - ctx.emit("}") diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/if_.py b/src/tilefoundry/codegen/cuda/tir/stmts/if_.py deleted file mode 100644 index 484eed0b5..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/if_.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Emit TIR conditionals with scalar C-style predicates.""" - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import If -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - -from .scalar_expr import render_scalar_expr - - -@register_codegen(CudaTarget, Role.EMIT, If) -def _emit(node: If, ctx: CudaCodegenContext) -> None: - cond = render_scalar_expr(node.cond, ctx) - ctx.emit(f"if ({cond}) {{") - ctx.indent() - ctx.emit_node(node.then_body) - ctx.dedent() - if getattr(node.else_body, "body", None): - ctx.emit("} else {") - ctx.indent() - ctx.emit_node(node.else_body) - ctx.dedent() - ctx.emit("}") - else: - ctx.emit("}") diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/let_stmt.py b/src/tilefoundry/codegen/cuda/tir/stmts/let_stmt.py deleted file mode 100644 index 8ed9e2bae..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/let_stmt.py +++ /dev/null @@ -1,30 +0,0 @@ -"""Emitter for ``tir.LetStmt``. - -LetStmt wraps a value binding: `auto %var = ` then the scoped -body. Dispatch is two-step: - -1. Look up an emitter keyed on the ``Op`` class of ``value.target`` — each - TIR-owned Expr Op (``tir.memory.AllocTensor``, ``tir.memory.PtrOf``, - ``tir.memory.TensorView``, etc.) registers its own emitter with - signature ``(let: LetStmt, ctx) -> None``. -2. Recurse into the Sequential body. -""" -from __future__ import annotations - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.core import Call -from tilefoundry.ir.tir.stmts import LetStmt -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - - -@register_codegen(CudaTarget, Role.EMIT, LetStmt) -def _emit(node: LetStmt, ctx: CudaCodegenContext) -> None: - if not isinstance(node.value, Call): - raise RuntimeError( - f"LetStmt.value must be a Call (TIR-owned Expr Op), " - f"got {type(node.value).__name__}" - ) - op = node.value.target - ctx.handler_for(type(op))(node, ctx) - ctx.emit_node(node.body) diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/return_.py b/src/tilefoundry/codegen/cuda/tir/stmts/return_.py deleted file mode 100644 index ea05e6271..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/return_.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Emitter for `tir.Return`.""" - -from __future__ import annotations - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import Return -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - - -@register_codegen(CudaTarget, Role.EMIT, Return) -def _emit(node: Return, ctx: CudaCodegenContext) -> None: - ctx.emit("return;") diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/sequential.py b/src/tilefoundry/codegen/cuda/tir/stmts/sequential.py deleted file mode 100644 index fadc59595..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/sequential.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Emitter for ``tir.Sequential`` — packs a list of Stmts in order.""" -from __future__ import annotations - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import Sequential -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - - -@register_codegen(CudaTarget, Role.EMIT, Sequential) -def _emit(node: Sequential, ctx: CudaCodegenContext) -> None: - for stmt in node.body: - ctx.emit_node(stmt) diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/while_.py b/src/tilefoundry/codegen/cuda/tir/stmts/while_.py deleted file mode 100644 index df673bc99..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/while_.py +++ /dev/null @@ -1,15 +0,0 @@ -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import While -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - -from .scalar_expr import render_scalar_expr - - -@register_codegen(CudaTarget, Role.EMIT, While) -def _emit(node: While, ctx: CudaCodegenContext) -> None: - ctx.emit(f"while ({render_scalar_expr(node.cond, ctx)}) {{") - ctx.indent() - ctx.emit_node(node.body) - ctx.dedent() - ctx.emit("}") diff --git a/src/tilefoundry/codegen/emitter.py b/src/tilefoundry/codegen/emitter.py index ae40454bb..dcd5ce894 100644 --- a/src/tilefoundry/codegen/emitter.py +++ b/src/tilefoundry/codegen/emitter.py @@ -10,6 +10,7 @@ ) from tilefoundry.codegen.cuda.tir.stmts.scalar_expr import render_scalar_expr from tilefoundry.ir.core import Call +from tilefoundry.ir.tir.abort import Abort from tilefoundry.ir.tir.stmts import ( Evaluate, For, @@ -84,11 +85,13 @@ def visit_While(self, node: While) -> None: def visit_Return(self, node: Return) -> None: self.statement("return") + def visit_Abort(self, node: Abort) -> None: + self.statement("__trap()") + def visit_LetStmt(self, node: LetStmt) -> None: if not isinstance(node.value, Call): raise RuntimeError( - f"LetStmt.value must be a Call (TIR-owned Expr Op), " - f"got {type(node.value).__name__}" + f"LetStmt.value must be a Call (TIR-owned Expr Op), got {type(node.value).__name__}" ) self.emit_let_op(node.value, node) self.visit(node.body) @@ -101,13 +104,13 @@ def visit_MeshScope(self, node: MeshScope) -> None: raise NotImplementedError def emit_evaluate_op(self, call: Call) -> None: - handler = self.context.lookup(type(call.target)) + handler = self.context.handler_for(type(call.target)) if handler is None: raise RuntimeError(f"no codegen handler for Op {type(call.target).__name__}") handler(call, self.context) def emit_let_op(self, call: Call, stmt: LetStmt) -> None: - handler = self.context.lookup(type(call.target)) + handler = self.context.handler_for(type(call.target)) if handler is None: raise RuntimeError(f"no codegen handler for Op {type(call.target).__name__}") handler(stmt, self.context) @@ -140,10 +143,7 @@ def visit_MeshScope(self, node: MeshScope) -> None: ctx.emit(f"using {alias} = {mesh_type_str};") ctx.emit(f"constexpr {alias} {name}_mesh{{}};") if is_slice: - ctx.emit( - f"if (tilefoundry::contains({name}_mesh, " - "tilefoundry::program_ids())) {" - ) + ctx.emit(f"if (tilefoundry::contains({name}_mesh, tilefoundry::program_ids())) {{") ctx.indent() self.visit(node.body) if is_slice: diff --git a/src/tilefoundry/codegen/registry.py b/src/tilefoundry/codegen/registry.py index 5b3048611..3f95cd0d4 100644 --- a/src/tilefoundry/codegen/registry.py +++ b/src/tilefoundry/codegen/registry.py @@ -2,53 +2,6 @@ from __future__ import annotations -from tilefoundry.ir.core.module import Module, module_functions -from tilefoundry.target.base import Target, _target_summary from tilefoundry.target.services import CodeGenerator - -def group_functions_by_target( - module: Module, -) -> dict[Target, tuple["PrimFunction", ...]]: - """Group functions by their equal Target value, preserving source order.""" - from tilefoundry.ir.tir.prim_function import PrimFunction # noqa: PLC0415 - - groups: dict[Target, list[PrimFunction]] = {} - for function in module_functions(module): - if not isinstance(function, PrimFunction): - raise TypeError( - f"tilefoundry: codegen expects PrimFunction values, got " - f"{type(function).__name__} {function.name!r}" - ) - if function.target is None: - raise ValueError( - f"tilefoundry: function {function.name!r} has no resolved Target " - "at codegen grouping" - ) - groups.setdefault(function.target, []).append(function) - """Expose structural variants to their target-specific emitter.""" - for variant in function.variants: - if variant.target is None: - raise ValueError(f"tilefoundry: variant {variant.name!r} has no resolved Target") - groups.setdefault(variant.target, []).append(variant) - - from tilefoundry.target import CudaTarget # noqa: PLC0415 - - device_groups = [ - (target, functions) - for target, functions in groups.items() - if isinstance(target, CudaTarget) - ] - if len(device_groups) > 1: - first_target, first_functions = device_groups[0] - second_target, second_functions = device_groups[1] - raise ValueError( - f"tilefoundry: module {module.name!r} mixes unequal device Targets: " - f"{_target_summary(first_target)} (function {first_functions[0].name!r}) " - f"vs {_target_summary(second_target)} (function {second_functions[0].name!r}); " - "multiple device translation units are not supported" - ) - return {target: tuple(functions) for target, functions in groups.items()} - - -__all__ = ["CodeGenerator", "group_functions_by_target"] +__all__ = ["CodeGenerator"] diff --git a/src/tilefoundry/codegen/signature.py b/src/tilefoundry/codegen/signature.py index 64266be23..bf28c66a1 100644 --- a/src/tilefoundry/codegen/signature.py +++ b/src/tilefoundry/codegen/signature.py @@ -9,10 +9,9 @@ from __future__ import annotations -from collections.abc import Mapping from dataclasses import dataclass -from tilefoundry.ir.core.module import Module, module_functions +from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import DimVar @@ -128,15 +127,14 @@ def all_params(self) -> tuple[Signature, ...]: LAUNCH_ABI = ( LaunchSignature("grid_x", "int"), - LaunchSignature("grid_y", "int"), - LaunchSignature("grid_z", "int"), LaunchSignature("block_x", "int"), - LaunchSignature("block_y", "int"), - LaunchSignature("block_z", "int"), LaunchSignature("dynamic_smem", "int"), - LaunchSignature("stream", "void*"), ) -"""What a launch shim is told beyond the kernel's own arguments.""" +"""Actual cross-translation-unit scalars on a CPU-to-device Launch edge. + +The topology is one-dimensional, so only grid_x and block_x are stated. +Dynamic shared memory remains because its per-Launch value crosses the shim ABI. +""" def tensor_signature_of(var) -> TensorSignature: @@ -187,18 +185,6 @@ def called_as(fn, program_ids: tuple[ProgramIdSignature, ...] = ()) -> CallableS return convention(fn, program_ids) -def symbol_table(module: Module, target: Target) -> Mapping[int, CallableSignature]: - """What every function in the tree is called by, before anything is emitted. - - Keyed by ``id(fn)``: a call site holds the callee itself and never asks - which module it lives in. A specialization variant gets no row, because - nothing calls one: its prototype compiles to the one symbol, and which - variant runs is decided inside it. - """ - program_ids = program_id_params(module, target) - return {id(fn): called_as(fn, program_ids) for fn in module_functions(module)} - - __all__ = [ "LAUNCH_ABI", "CallableSignature", @@ -211,6 +197,5 @@ def symbol_table(module: Module, target: Target) -> Mapping[int, CallableSignatu "UnitSignature", "called_as", "program_id_params", - "symbol_table", "tensor_signature_of", ] diff --git a/src/tilefoundry/compile.py b/src/tilefoundry/compile.py index 306db7bc5..fdf233e9a 100644 --- a/src/tilefoundry/compile.py +++ b/src/tilefoundry/compile.py @@ -10,12 +10,10 @@ import tempfile from dataclasses import dataclass, replace +from tilefoundry.codegen.context_builder import build_codegen_context from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.codegen.linker import link_modules -from tilefoundry.codegen.registry import group_functions_by_target -from tilefoundry.codegen.signature import symbol_table -from tilefoundry.codegen.topology import launch_geometry, topology_domains from tilefoundry.inspection import as_script as _as_script from tilefoundry.ir.core.module import Module from tilefoundry.ir.hir.function import Function as HirFunction @@ -102,7 +100,6 @@ def build( f"conflicts with the Module Target {_target_summary(module_target)}" ) - mod = _build_default_pipeline().run(mod) workdir = os.path.join( tempfile.gettempdir(), f"tilefoundry_build_{mod.entry}_{os.getpid()}_split" ) @@ -120,69 +117,64 @@ def _build_split_runtime_module(mod: Module, *, workdir: str) -> "RuntimeModule" load each other's. Unsupported module shapes raise during codegen -- there is no fallback to a single-source path. """ - device_target = _device_target(mod) + codegen = build_codegen_context(mod) cpu_entry = mod.entry_function() if not isinstance(cpu_entry.target, CpuTarget): raise ValueError( - f"tilefoundry.build: entry {cpu_entry.name!r} is not a CPU host entry " - f"after normalization" + f"tilefoundry.build: entry {cpu_entry.name!r} is not a CPU host entry; " + "build() accepts pass-prepared TIR" ) - symbols = symbol_table(mod, device_target) - launches = launch_geometry(mod) - device_modules = tuple( - device_target.get_code_generator().emit( - domain, - device_fns, - device_target, - CudaCodegenContext(symbols=symbols, target=device_target, launches=launches), - ) - for domain, device_fns in _device_domains(topology_domains(mod), device_target) - ) - host_module = cpu_entry.target.get_code_generator().emit( - mod, - (cpu_entry,), - cpu_entry.target, - CpuCodegenContext(symbols=symbols, target=cpu_entry.target), - ) + device_modules = [] + host_module = None + for group in codegen.groups: + view = codegen.for_group(group) + if isinstance(group.target, CudaTarget): + context = CudaCodegenContext( + symbols=view.symbols, + target=view.target, + launches=view.launches, + codegen_context=view.codegen, + ) + device_modules.append( + group.target.get_code_generator().emit( + group.owner, group.functions, group.target, context + ) + ) + elif isinstance(group.target, CpuTarget): + if cpu_entry not in group.functions: + continue + context = CpuCodegenContext( + symbols=view.symbols, + target=view.target, + resolved_launches=view.resolved_launches, + codegen_context=view.codegen, + ) + host_module = group.target.get_code_generator().emit( + group.owner, (cpu_entry,), group.target, context + ) + if host_module is None: + raise ValueError("tilefoundry.build: root context produced no host entry group") units = (*device_modules, host_module) digest = hashlib.sha256("".join(m.source for m in units).encode("utf-8")).hexdigest()[:16] - loaded_as = replace(symbols[id(cpu_entry)], name=cpu_entry.name) + entry_signature = codegen.symbols.by_function[id(cpu_entry)].callable + loaded_as = replace(entry_signature, name=cpu_entry.name) + device_targets = tuple( + group.target for group in codegen.groups if isinstance(group.target, CudaTarget) + ) + if not device_targets: + raise ValueError(f"tilefoundry.build: module {mod.name!r} has no CUDA device functions") linked_module = link_modules( units, workdir=os.path.join(workdir, digest), lib_name=cpu_entry.name, entry=loaded_as, - cuda_arch=device_target.arch.removeprefix("sm_"), + cuda_arch=device_targets[0].arch.removeprefix("sm_"), ) return load_linked_module(linked_module) -def _device_target(mod: Module) -> Target: - """The one device Target this module's functions run on. - - A second unequal one is refused by the grouping that sees them both, - because one linked artifact holds one device architecture. - """ - targets = [t for t in group_functions_by_target(mod) if isinstance(t, CudaTarget)] - if not targets: - raise ValueError(f"tilefoundry.build: module {mod.name!r} has no CUDA device functions") - return targets[0] - - -def _device_domains(domains, device_target: Target): - """Each *domains* entry that has device functions, and which of them they are. - - A module states its instance counts once, so one that states its own is one - translation unit; a domain holding only host functions is not one at all. - """ - for domain, functions in domains: - device_fns = tuple(fn for fn in functions if fn.target == device_target) - if device_fns: - yield domain, device_fns - - def compile( mod: Module, /, @@ -193,7 +185,8 @@ def compile( *mod* must be a ``Module``. Meshes are derived from the IR body. """ - return build(mod, target=target) + prepared = _build_default_pipeline().run(mod) + return build(prepared, target=target) def _canonical_module_text(mod: Module) -> str: diff --git a/src/tilefoundry/inspection/values.py b/src/tilefoundry/inspection/values.py index 0b8c390cf..26780e3be 100644 --- a/src/tilefoundry/inspection/values.py +++ b/src/tilefoundry/inspection/values.py @@ -6,9 +6,8 @@ from collections.abc import Mapping from dataclasses import dataclass, field, fields, is_dataclass -from tilefoundry.analysis.metadata import TripInterval from tilefoundry.ir.core.metadata import IRMetadata -from tilefoundry.visitor_registry.contexts import TrafficBytes +from tilefoundry.ir.core.values import TripInterval PAIR, PER_UNIT, ENTRY, ENTRIES, FIELD, FIELDS, PARTS, TRIPS = ( "/", @@ -116,34 +115,45 @@ def _single(self, family, value, default=_UNSET): return f"{family}{FIELD}{self.print(value)}" def print_ComputeCostMetadata(self, record, **_): - flops_per_unit = dict(record.flops_per_unit) - ops_per_unit = dict(record.ops_per_unit) - flops = { - dtype: f"{total}{PER_UNIT}{flops_per_unit.get(dtype, 0)}" - for dtype, total in record.flops - } - ops = { - kind: f"{total}{PER_UNIT}{ops_per_unit.get(kind, 0)}" - for kind, total in record.ops - } + def spread_values(spread): + logical = dict(spread.logical) + total = dict(spread.total) + per_unit = [dict(value) for value in spread.per_unit] + names = sorted({*logical, *total, *(name for value in per_unit for name in value)}) + return { + name: ( + f"{logical.get(name, 0)}{PAIR}{total.get(name, 0)}" + + "".join( + f"{PER_UNIT}{level}{ENTRY}{value.get(name, 0)}" + for level, value in zip(record.topologies, per_unit, strict=False) + ) + ) + for name in names + } + return self._record( "compute-cost", ( - ("flops", flops), - ("flops_logical", record.flops_logical), - ("ops", ops), - ("ops_logical", dict(record.ops_logical)), + ("flops", spread_values(record.flops)), + ("other_ops", spread_values(record.other_ops)), ), ) def print_TrafficMetadata(self, record, *, opt_in=frozenset()): - per_unit = dict(record.per_unit) traffic = { level: ( - f"{self.print(moved)}{PER_UNIT}" - f"{self.print(per_unit.get(level, TrafficBytes()))}" + f"{self.print(spread.total)}" + + ( + f"{PER_UNIT}" + + ENTRIES.join( + f"{topology}{ENTRY}{self.print(value)}" + for topology, value in zip(record.topologies, spread.per_unit, strict=False) + ) + if spread.per_unit + else "" + ) ) - for level, moved in record.whole + for level, spread in record.storage.kinds } values = [("traffic", traffic)] if "operands" in opt_in: @@ -159,7 +169,7 @@ def print_MemoryMetadata(self, record, **_): return self._record( "memory", ( - ("peak", {item.level: item.peak_bytes for item in record.footprint}), + ("peak", {item.memory_level: item.peak_bytes for item in record.footprint}), ("persistent", sum(item.persistent_bytes for item in record.footprint), 0), ("advisories", len(record.advisories), 0), ), @@ -167,7 +177,7 @@ def print_MemoryMetadata(self, record, **_): def print_LoopFootprintMetadata(self, record, **_): footprints = { - f"{item.buffer}@{item.level}": PAIR.join( + 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 @@ -240,7 +250,7 @@ def render_comment(record, *, opt_in=frozenset()): def peak_footprint(record): - return {item.level: item.peak_bytes for item in record.footprint} + return {item.memory_level: item.peak_bytes for item in record.footprint} __all__ = [ diff --git a/src/tilefoundry/passes/transforms/host_entry.py b/src/tilefoundry/passes/transforms/host_entry.py index 378b8b7c6..f0d353c55 100644 --- a/src/tilefoundry/passes/transforms/host_entry.py +++ b/src/tilefoundry/passes/transforms/host_entry.py @@ -1,18 +1,22 @@ -"""Synthesize a CPU entry that launches a lone CUDA ``PrimFunction``. +"""Give a device-only module the CPU entry that calls into it. -It copies device parameters and the existing static launch configuration; -dynamic grids remain unsupported. +The entry is synthesized, never retargeted: which target a function records is +what its author wrote, and a pass that rewrites it decides in the wrong place. +Launch geometry is settled here and written into the ``Launch``, so nothing +downstream derives it a second time. A launch-provided (dynamic) grid is not +supported by the synthesized entry. """ + from __future__ import annotations from dataclasses import replace -from tilefoundry.codegen.signature import LaunchSignature from tilefoundry.ir.core import Var from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.launch import launch_call from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.stmts import Sequential +from tilefoundry.passes.pass_base import ModulePass from tilefoundry.target import CpuTarget, CudaTarget _DEFAULT_ENTRY_NAME = "main" @@ -29,35 +33,18 @@ def _fresh_entry_name(taken: set[str]) -> str: return candidate -def insert_default_host_entry(module: Module) -> Module: - """Return a module whose entry is host-callable (CPU-target). - - Return a module whose entry is host-callable (CPU-target): - - a CPU entry is already present → unchanged; - - the entry is a specialization prototype (host-only) → retarget it - to CPU in place of the original (``module.entry`` name unchanged); - - exactly one CUDA device kernel with no CPU entry → synthesize a CPU entry - whose body is a single ``Launch`` of that kernel. - """ - entry_fn = module.entry_function() - if isinstance(entry_fn.target, CpuTarget): - return module - - +def _device_callee(module: Module) -> PrimFunction: + """The one device function a synthesized entry can be written against. + Refuses rather than guesses: a second device function would leave the + entry's parameters undecided, and a CPU function that is not the entry + means the author meant some other function to be called first. + """ if any(isinstance(fn.target, CpuTarget) for fn in module.functions): raise ValueError( "insert_default_host_entry: module has a CPU function that is not " "the entry; refusing to guess the host entry" ) - - - if entry_fn.variants: - cpu_entry = replace(entry_fn, target=CpuTarget()) - new_functions = tuple( - cpu_entry if fn is entry_fn else fn for fn in module.functions - ) - return replace(module, functions=new_functions) device_fns = [ fn for fn in module.functions @@ -66,21 +53,49 @@ def insert_default_host_entry(module: Module) -> Module: if len(device_fns) != 1: raise ValueError( f"insert_default_host_entry: expected exactly one CUDA device " - f"prim_function (or a dispatch entry) and no CPU entry, found " - f"{len(device_fns)} device functions" + f"prim_function and no CPU entry, found {len(device_fns)} device " + f"functions" ) - device_fn = device_fns[0] + return device_fns[0] + + +def _product(values: tuple[object, ...]) -> object: + result: object = 1 + for value in values: + if value is None: + raise ValueError( + "insert_default_host_entry: topology extent cannot be None; " + "declare a DimVar-backed host-computable extent" + ) + result = result * value + return result + + +def _launch_config_of( + module: Module, +) -> tuple[tuple[object, object, object], tuple[object, object, object]]: + """Derive the implicit Launch expressions solely from module topologies.""" + topologies = module.effective_topologies() + grid = _product(tuple(t.size for t in topologies if t.name == "cta")) + block = _product(tuple(t.size for t in topologies if t.name == "thread")) + return (grid, 1, 1), (block, 1, 1) +def insert_default_host_entry(module: Module) -> Module: + """Return a module whose entry is host-callable, synthesizing one if it is not. + + A CPU entry passes through. Otherwise the device entry gains a CPU entry + that launches it with its own parameters. Whether that call reads as one + launch or as a dispatch over variants is the callee's shape to state, so + a prototype and a lone kernel take the same route here. + """ + if isinstance(module.entry_function().target, CpuTarget): + return module + + device_fn = _device_callee(module) entry_params = tuple(Var(type=p.type, name=p.name) for p in device_fn.params) - launch_signature = LaunchSignature.from_topologies(module.effective_topologies()) - if launch_signature.grid[0] is None: - raise ValueError( - f"insert_default_host_entry: device function {device_fn.name!r} has " - f"a launch-provided (dynamic) CTA extent; the implicit host entry " - f"cannot derive its grid — launch it from an explicit host entry" - ) - launch = launch_call(device_fn, entry_params) + grid, block = _launch_config_of(module) + launch = launch_call(device_fn, entry_params, grid, block) name = _fresh_entry_name({fn.name for fn in module.functions}) entry = PrimFunction( name=name, @@ -92,4 +107,18 @@ def insert_default_host_entry(module: Module) -> Module: return replace(module, functions=(*module.functions, entry), entry=name) -__all__ = ["insert_default_host_entry"] +class InsertHostEntryPass(ModulePass): + """Give a device-only module a host-callable entry. + + The only pass a TIR module needs: everything else it is made of was + settled before it became TIR. + """ + + name: str = "insert_host_entry" + requires: tuple[str, ...] = () + + def run(self, module: Module) -> Module: + return insert_default_host_entry(module) + + +__all__ = ["InsertHostEntryPass", "insert_default_host_entry"] diff --git a/src/tilefoundry/visitor_registry/contexts.py b/src/tilefoundry/visitor_registry/contexts.py index c641f259a..f51c5fc8f 100644 --- a/src/tilefoundry/visitor_registry/contexts.py +++ b/src/tilefoundry/visitor_registry/contexts.py @@ -123,14 +123,14 @@ class VerifyContext(TypeInferContext): class CostContext(TypeInferContext): """Cost Evaluator context for one topology window. - ``level=None`` exposes the types as written. A named level projects them to - what one unit of that level holds, letting the same registered evaluator - answer both global and per-unit questions. + ``topology_level=None`` exposes the types as written. A named level projects + them to what one unit of that level holds, letting the same registered + evaluator answer both global and per-unit questions. """ selected_types: Mapping[int, Type] = field(default_factory=dict) selected_output_type: Type | None = None - level: str | None = None + topology_level: str | None = None topologies: tuple[Topology, ...] = () def type_of(self, expr: Expr) -> Type: @@ -141,18 +141,18 @@ def local_type_of(self, expr: Expr) -> Type: """Return ``expr``'s Type in this context's topology window.""" selected = self.selected_types.get(id(expr)) type_ = selected if selected is not None else expr.type - if self.level is None: + if self.topology_level is None: return type_ - return local_type_of(type_, level=self.level, topologies=self.topologies) + return local_type_of(type_, topology_level=self.topology_level, topologies=self.topologies) def local_output_type(self, call: Call) -> Type: """Return the selected candidate output in recursive-local form.""" output = self.selected_output_type if output is None: output = call.type - if self.level is None: + if self.topology_level is None: return output - return local_type_of(output, level=self.level, topologies=self.topologies) + return local_type_of(output, topology_level=self.topology_level, topologies=self.topologies) @dataclass(frozen=True) @@ -174,16 +174,29 @@ class Cost: ``flops`` groups leaf-local logical work by compute ``DType`` so one Op can report mixed work without selecting an ALU/TensorCore implementation. - ``ops`` counts what is not floating point at all -- a comparison, a - select, an integer add -- by the machine operation kind it asks for, because - a dtype is not a kind of work. ``traffic`` carries one entry per operand in call order with + ``service`` counts what is not floating point at all -- a comparison, a + select, an integer add -- by the service it asks for, because a dtype is not + a kind of work. ``traffic`` carries one entry per operand in call order with the result last, so an Op that reads part of an input says so where it knows it. No field names a hardware implementation or a memory level. """ flops: Mapping[DType, int] traffic: tuple[TrafficBytes, ...] - ops: Mapping[str, int] = field(default_factory=dict) + service: Mapping[str, int] = field(default_factory=dict) + sent: tuple[tuple[str, TrafficBytes], ...] = () + """Bytes that left the unit they were on, by the boundary they crossed. + + ``traffic`` says which storage level each operand touched. Data handed + between two units of one topology level touches the same storage at both + ends and has still gone somewhere, so it is stated against that level's + name instead, from the sending unit's own view. + """ + + @property + def ops(self) -> Mapping[str, int]: + """Program-work spelling for non-floating-point operation counts.""" + return self.service @property def bytes(self) -> int: @@ -201,9 +214,9 @@ def __post_init__(self) -> None: raise ValueError("Cost traffic must be non-negative integers") if any( isinstance(value, bool) or not isinstance(value, int) or value < 0 - for value in self.ops.values() + for value in self.service.values() ): - raise ValueError("Cost ops work must be non-negative integers") + raise ValueError("Cost service work must be non-negative integers") __all__ = [ diff --git a/src/tilefoundry/visitor_registry/op_cost.py b/src/tilefoundry/visitor_registry/op_cost.py index eb04ac85f..479774fc5 100644 --- a/src/tilefoundry/visitor_registry/op_cost.py +++ b/src/tilefoundry/visitor_registry/op_cost.py @@ -107,11 +107,11 @@ def _elementwise(call: Call, ctx: CostContext, *, dtype: DType | None = None) -> return Cost({result_dtype: numel(output)}, _traffic(inputs, output)) -def _ops(call: Call, ctx: CostContext, kind: str) -> Cost: - """One machine operation of *kind* per element, and no FLOP work.""" +def _serviced(call: Call, ctx: CostContext, kind: str) -> Cost: + """One result of *kind* per element, and no floating-point work at all.""" inputs = _input_types(call, ctx) output = _output_type(call, ctx) - return Cost({}, _traffic(inputs, output), ops={kind: numel(output)}) + return Cost({}, _traffic(inputs, output), {kind: numel(output)}) @register_cost_evaluator(MatMul) @@ -206,15 +206,15 @@ def _binary(call: Call, ctx: CostContext) -> Cost: A comparison over floats produces booleans, and calling that a bool FLOP asks the target for a rate no machine publishes. What it really asks for is - a predicate, which is a machine operation with a separately stated throughput. + a predicate, which is a service a machine does state a throughput for. ``_PREDICATES`` holds the kinds whose result is one, whatever the operands were. """ kind = call.target.kind if kind in _PREDICATES: - return _ops(call, ctx, "predicate") + return _serviced(call, ctx, "predicate") if _integral(call, ctx): - return _ops(call, ctx, "integer") + return _serviced(call, ctx, "integer") return _elementwise(call, ctx) @@ -236,17 +236,17 @@ def _unary(call: Call, ctx: CostContext) -> Cost: ``_SPECIAL`` holds the kinds the machine answers on its special-function unit, at a rate of its own -- a quarter of the scalar one here. Counting one of those as a single FLOP would put it on the float pipe - at four times the throughput the unit has, so it is a separately priced - operation rather than arithmetic. What is left for ``flops`` is the arithmetic that really is a + at four times the throughput the unit has, so it is a service rather than + arithmetic. What is left for ``flops`` is the arithmetic that really is a multiply or an add. """ kind = call.target.kind if kind in _SPECIAL: - return _ops(call, ctx, "special") + return _serviced(call, ctx, "special") if kind is UnaryKind.NOT: - return _ops(call, ctx, "predicate") + return _serviced(call, ctx, "predicate") if _integral(call, ctx): - return _ops(call, ctx, "integer") + return _serviced(call, ctx, "integer") return _elementwise(call, ctx) @@ -350,7 +350,8 @@ def _softmax(call: Call, ctx: CostContext) -> Cost: @register_cost_evaluator(Where) def _where(call: Call, ctx: CostContext) -> Cost: """Choosing between two values it already has is a select, not arithmetic.""" - return _ops(call, ctx, "select") + return _serviced(call, ctx, "select") + @register_cost_evaluator(LayerNorm) def _layer_norm(call: Call, ctx: CostContext) -> Cost: diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 8cd7805a9..6ec277e6e 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -21,6 +21,7 @@ ) from tilefoundry import func, module from tilefoundry.analysis import ( + Breakdown, ComputeCostMetadata, MemoryHierarchyFacts, MemoryMetadata, @@ -28,6 +29,7 @@ PerformanceServiceFacts, PerformanceSummaryMetadata, RooflineMetadata, + Spread, ThroughputFacts, TrafficMetadata, ) @@ -174,8 +176,7 @@ def test_performance_orders_a_predecessor_materialized_in_a_child_scope() -> Non cross_scope = results["cross_scope"] timed = tuple( - (call, get_metadata(call, PerformanceMetadata)) - for call in _calls(cross_scope.function) + (call, get_metadata(call, PerformanceMetadata)) for call in _calls(cross_scope.function) ) scalar_indices = tuple( (call, record) @@ -233,18 +234,14 @@ def test_a_symbolic_store_stride_preserves_the_literal_control_result() -> None: summary = get_metadata(result.function, PerformanceSummaryMetadata) assert cost is not None and bound is not None and summary is not None observed[name] = ( - dict(cost.ops)["integer"], - dict(cost.ops_per_unit)["integer"], + dict(cost.other_ops.total)["integer"], + dict(cost.other_ops.per_unit[0])["integer"], bound.ideal_ns, summary.timeline.end_ns - summary.timeline.start_ns, ) - literal_service, literal_local, literal_roofline, literal_performance = observed[ - "literal" - ] - symbolic_service, symbolic_local, symbolic_roofline, symbolic_performance = observed[ - "symbolic" - ] + 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 symbolic_local - literal_local == 6 assert symbolic_service - literal_service == 6 * 128 @@ -281,12 +278,8 @@ def last_axis( w: ConstTensor[(_SPLIT_HIDDEN, _SPLIT_OUT), "bf16"], ): with Mesh(("cta",), layout=(_SPLIT_GRID,), names=("unit",)) as mesh: - rows = tf.reshard( - x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem" - ) - strip = tf.reshard( - w[0:_SPLIT_BLOCK, :], (_SPLIT_BLOCK, _SPLIT_OUT @ mesh.unit), "smem" - ) + rows = tf.reshard(x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem") + strip = tf.reshard(w[0:_SPLIT_BLOCK, :], (_SPLIT_BLOCK, _SPLIT_OUT @ mesh.unit), "smem") return tf.matmul(rows, strip) @@ -300,9 +293,7 @@ def strip_major( w: ConstTensor[(_SPLIT_GRID, _SPLIT_HIDDEN, _SPLIT_PER), "bf16"], ): with Mesh(("cta",), layout=(_SPLIT_GRID,), names=("unit",)) as mesh: - rows = tf.reshard( - x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem" - ) + rows = tf.reshard(x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem") strip = tf.reshard( w[:, 0:_SPLIT_BLOCK, :], (_SPLIT_GRID @ mesh.unit, _SPLIT_BLOCK, _SPLIT_PER), @@ -331,8 +322,8 @@ def test_a_matmul_counts_its_rows_once_whichever_axis_the_mesh_split() -> None: cost = get_metadata(product, ComputeCostMetadata) summary = get_metadata(report.function, PerformanceSummaryMetadata) per_layout[name] = ( - dict(cost.flops)["bf16"], - dict(cost.flops_per_unit)["bf16"], + dict(cost.flops.total)["bf16"], + dict(cost.flops.per_unit[0])["bf16"], summary.timeline.end_ns, ) @@ -412,7 +403,7 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N work = next( record for record in records - if record is not None and any(v for _n, v in record.flops_per_unit) + if record is not None and any(v for _n, v in record.flops.per_unit[0]) ) with pytest.raises( @@ -424,14 +415,26 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N with pytest.raises(AnalysisError, match=r"unknown compute dtype 'f9e9m9'"): _local_duration_ns( - replace(work, flops_per_unit=(("f9e9m9", 8),)), + replace(work, flops=replace(work.flops, per_unit=((("f9e9m9", 8),),))), throughput, services, level="cta", ) crossed = TrafficMetadata( - per_unit=((throughput.bandwidth_level, TrafficBytes(read=4096)),) + topologies=("cta",), + storage=Breakdown( + ( + ( + throughput.bandwidth_level, + Spread( + TrafficBytes(read=4096), + TrafficBytes(read=4096), + (TrafficBytes(read=4096),), + ), + ), + ) + ), ) with pytest.raises( AnalysisError, diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index 8989cea65..e849a1af2 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -92,7 +92,7 @@ def assert_performance_contract(result: AnalysisResult) -> None: module_target = result.module.resolve_target() throughput = module_target.get_facts(ThroughputFacts) - services = module_target.get_facts(PerformanceServiceFacts) + services = module_target.get_facts(PerformanceServiceFacts, result.level) scopes = tuple(walk_scopes(build_scopes(result.module, fn))) timed = 0 for expr in collect_exprs(fn.body): @@ -185,17 +185,7 @@ def _every_number_counts_something(result: AnalysisResult) -> None: fn = result.function for expr in (fn, *collect_exprs(fn.body)): for record, rows in ( - ( - ComputeCostMetadata, - ( - "flops", - "flops_logical", - "flops_per_unit", - "ops", - "ops_logical", - "ops_per_unit", - ), - ), + (ComputeCostMetadata, ()), (TrafficMetadata, ()), (MemoryMetadata, ()), (RooflineMetadata, ()), @@ -211,6 +201,12 @@ def _every_number_counts_something(result: AnalysisResult) -> None: continue for name, value in value: assert value >= 0, f"{describe_expr(expr)}: {field}[{name}] = {value}" + if record is ComputeCostMetadata: + for field in ("flops", "other_ops"): + spread = getattr(held, field) + for domain in (spread.logical, spread.total, *spread.per_unit): + for name, value in domain: + 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): diff --git a/tests/analysis/test_analyze_cross_module.py b/tests/analysis/test_analyze_cross_module.py index 0a4a9e676..a9c937457 100644 --- a/tests/analysis/test_analyze_cross_module.py +++ b/tests/analysis/test_analyze_cross_module.py @@ -50,8 +50,8 @@ def _flops(records) -> dict[str, int]: total: dict[str, int] = {} for item in records: record = item[0] if isinstance(item, tuple) else item - for name, spread in record.flops.kinds: - total[name] = total.get(name, 0) + spread.total + for name, value in record.flops.total: + total[name] = total.get(name, 0) + value return total diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 147ea1e4a..d5c8c07a8 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -166,8 +166,8 @@ def test_analyze_help_explains_topology_effects_and_assumptions(capsys) -> None: help_text = capsys.readouterr().out for family in ("compute-cost", "memory", "roofline", "performance"): assert family in help_text - assert "flops_per_unit" in help_text - assert "per-unit traffic" 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 "is an observation, not a bound" in help_text @@ -746,21 +746,24 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> f"# selection requested={','.join(payload['requested'])} " f"executed={','.join(payload['executed'])}", "# compute-cost " - f"flops=f32:{cost['flops']['f32']}@{cost['flops_per_unit']['f32']}" - f" flops-logical={cost['flops_logical']}", + f"flops=f32:{cost['flops']['logical']['f32']}" + f"/{cost['flops']['total']['f32']}" + f"@{payload['topology']}:{cost['flops']['per_unit'][0]['f32']}", "# traffic " - f"traffic=gmem:r{moved['whole']['gmem']['read']}" - f"/w{moved['whole']['gmem']['write']}" - f"@r{moved['per_unit']['gmem']['read']}" - f"/w{moved['per_unit']['gmem']['write']}", + f"traffic=gmem:r{moved['storage']['gmem']['total']['read']}" + f"/w{moved['storage']['gmem']['total']['write']}" + f"@{payload['topology']}:r{moved['storage']['gmem']['per_unit'][0]['read']}" + f"/w{moved['storage']['gmem']['per_unit'][0]['write']}", f"# peak-footprint=gmem:{peak[0]['peak_bytes']}", f"# roofline ideal-ns={bound['ideal_ns']} bound-by={bound['bound_by']}", "# performance root=MoEMegaKernel::experts " f"predicted-ns={summary['timeline']['end_ns']} " f"waves={summary['waves']}", ] - assert payload["totals"]["flops"] == cost["flops"] - assert payload["totals"]["traffic"] == moved["whole"] + assert payload["totals"]["flops"] == cost["flops"]["total"] + assert payload["totals"]["traffic"] == { + name: value["total"] for name, value in moved["storage"].items() + } hoisted = {line.split(" = ", 1)[0] for line in lines if " = Mesh((Topology(" in line} scoped = { diff --git a/tests/codegen/test_codegen_registry.py b/tests/codegen/test_codegen_registry.py deleted file mode 100644 index eca233820..000000000 --- a/tests/codegen/test_codegen_registry.py +++ /dev/null @@ -1,100 +0,0 @@ -"""One codegen registry: who answers, from which side of a call, about what. - -A target is a dimension of the key rather than a word in a registry's name, so -a caller of one target asks the callee's target how a call to it is written -without naming anything of that target's. -""" - -from __future__ import annotations - -import pytest - -from tests.fixtures.tir.square import TirSquare -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.codegen.signature import ( - LAUNCH_ABI, - ProgramIdSignature, - ProgramMetaSignature, - TensorSignature, - UnitSignature, - called_as, -) -from tilefoundry.ir.tir.stmts import Sequential -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.target import CpuTarget, CudaTarget -from tilefoundry.visitor_registry.registries import Role, codegen_registry - -_TARGET = CudaTarget("nvidia.h200_sxm") -_TENSOR = TensorSignature( - name="x", - type=TensorType( - shape=(4, DimVar("S", 1, 8)), - dtype=DType.f32, - layout=None, - storage=StorageKind.GMEM, - ), -) - - -def _ctx() -> CudaCodegenContext: - return CudaCodegenContext(target=_TARGET) - - -def test_a_node_is_written_by_the_target_being_written() -> None: - """Emission keys on the same three things as every other question.""" - assert _ctx().handler_for(Sequential) is codegen_registry.lookup( - (CudaTarget, Role.EMIT, Sequential) - ) - - -def test_a_tensor_declares_as_a_pointer_and_the_axes_its_type_leaves_open() -> None: - """A static extent costs the convention nothing; an open one travels along.""" - assert _ctx().declare(_TENSOR) == ("float* x", "int x_shape_1") - - -def test_what_a_convention_adds_declares_in_the_type_it_states() -> None: - ctx = _ctx() - gpu_id = ProgramIdSignature(name="tilefoundry_gpu_program_id", topology_level="gpu") - assert ctx.declare(LAUNCH_ABI[0]) == ("int grid_x",) - assert ctx.declare(LAUNCH_ABI[-1]) == ("void* stream",) - assert ctx.declare(gpu_id) == ("long long tilefoundry_gpu_program_id",) - assert ctx.declare(ProgramMetaSignature(name="tilefoundry_meta")) == ( - "tilefoundry::ProgramMetaData tilefoundry_meta", - ) - - -def test_the_same_tensor_declares_differently_for_two_targets() -> None: - """The host takes the runtime's tensor, which carries its own extents. - - Both answers come out of one registry and the key says which target gave - it; the host's reads nothing off the context it is handed. - """ - host = codegen_registry.lookup((CpuTarget, Role.CALLEE, TensorSignature)) - assert host(_TENSOR, _ctx()) == ("tvm::ffi::Tensor x",) - - -def test_a_caller_asks_the_callees_target_how_the_call_is_written() -> None: - """The argument list is the callee's answer; the names in it are the caller's.""" - ctx = _ctx() - assert ctx.argument(_TENSOR, _TARGET) == ("x", "x_shape_1") - with pytest.raises(RuntimeError, match=r"nothing registered for \(CpuTarget, Role.CALLER"): - ctx.argument(_TENSOR, CpuTarget()) - - -def test_an_unanswered_question_names_the_whole_key() -> None: - """A miss says which target, which side, and which class went unanswered.""" - with pytest.raises( - RuntimeError, match=r"nothing registered for \(CudaTarget, Role.CALLEE, UnitSignature\)" - ): - _ctx().declare(UnitSignature()) - - -def test_how_a_function_is_called_is_stated_by_its_own_target() -> None: - device = TirSquare.lookup("square_device") - host = TirSquare.lookup("square_host") - assert called_as(device).name == "tilefoundry_square_device_launch" - assert called_as(device).trailing == LAUNCH_ABI - assert called_as(host).name == "tilefoundry_square_host_host" - assert called_as(host).trailing == () diff --git a/tests/codegen/test_device_dispatch.py b/tests/codegen/test_device_dispatch.py index ce924ee8f..a45c52510 100644 --- a/tests/codegen/test_device_dispatch.py +++ b/tests/codegen/test_device_dispatch.py @@ -9,30 +9,40 @@ from __future__ import annotations from tests.fixtures.tir.square import TirSquare +from tilefoundry.codegen.context_builder import build_codegen_context from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cpu.module import emit_host_module from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.codegen.cuda.module import emit_cuda_module -from tilefoundry.codegen.signature import symbol_table -from tilefoundry.codegen.topology import launch_geometry from tilefoundry.target import CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") +def _view(target): + root = build_codegen_context(TirSquare) + group = next(group for group in root.groups if group.target == target) + return root, group, root.for_group(group) + + def _device_source() -> str: + _root, group, view = _view(_CUDA) ctx = CudaCodegenContext( - symbols=symbol_table(TirSquare, _CUDA), + symbols=view.symbols, target=_CUDA, - launches=launch_geometry(TirSquare), + launches=view.launches, ) - kernels = tuple(fn for fn in TirSquare.functions if fn.target == _CUDA) - return emit_cuda_module(TirSquare, kernels, _CUDA, ctx).source + return emit_cuda_module(TirSquare, group.functions, _CUDA, ctx).source def _host_source() -> str: entry = TirSquare.entry_function() - ctx = CpuCodegenContext(symbols=symbol_table(TirSquare, _CUDA), target=entry.target) + _root, group, view = _view(entry.target) + ctx = CpuCodegenContext( + symbols=view.symbols, + target=entry.target, + resolved_launches=view.resolved_launches, + ) return emit_host_module(TirSquare, (entry,), entry.target, ctx).source diff --git a/tests/codegen/test_host_multi_launch.py b/tests/codegen/test_host_multi_launch.py index 7f2fcf9d3..cb01b542e 100644 --- a/tests/codegen/test_host_multi_launch.py +++ b/tests/codegen/test_host_multi_launch.py @@ -7,9 +7,9 @@ import tilefoundry from tilefoundry import module, prim_func +from tilefoundry.codegen.context_builder import build_codegen_context from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cpu.module import emit_host_module -from tilefoundry.codegen.signature import symbol_table from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core import Constant, Var from tilefoundry.ir.core.module import Module @@ -129,8 +129,13 @@ def _module(devices, launches, params) -> tuple[Module, PrimFunction]: def _source(devices, launches, params) -> str: mod, host = _module(devices, launches, params) + root = build_codegen_context(mod) + group = next(group for group in root.groups if host in group.functions) + view = root.for_group(group) ctx = CpuCodegenContext( - symbols=symbol_table(mod, CudaTarget("nvidia.h200_sxm")), target=host.target + symbols=view.symbols, + target=host.target, + resolved_launches=view.resolved_launches, ) return emit_host_module(mod, (host,), host.target, ctx).source @@ -144,8 +149,8 @@ def test_two_launches_keep_order_and_each_block_size() -> None: (a, b), ) assert source.index("tilefoundry_first_launch") < source.index("tilefoundry_second_launch") - assert "1, 32, 1, 1" in source - assert "1, 64, 1, 1" in source + assert "1, 32, 0" in source + assert "1, 64, 0" in source @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/codegen/test_linkable.py b/tests/codegen/test_linkable.py index 41aa72af7..273a4cb54 100644 --- a/tests/codegen/test_linkable.py +++ b/tests/codegen/test_linkable.py @@ -3,11 +3,10 @@ from __future__ import annotations from tests.fixtures.tir.square import TirSquare +from tilefoundry.codegen.context_builder import build_codegen_context from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cpu.module import emit_host_module from tilefoundry.codegen.linkable import LinkableFunction, LinkableModule -from tilefoundry.codegen.signature import symbol_table -from tilefoundry.target import CudaTarget def _linkable_function(name: str) -> LinkableFunction: @@ -27,8 +26,13 @@ def _param_types(parameter_list: str) -> list[str]: def _host_unit() -> LinkableModule: """The host translation unit of a module whose entry launches a shape-dispatched kernel.""" entry = TirSquare.entry_function() + root = build_codegen_context(TirSquare) + group = next(group for group in root.groups if entry in group.functions) + view = root.for_group(group) ctx = CpuCodegenContext( - symbols=symbol_table(TirSquare, CudaTarget("nvidia.h200_sxm")), target=entry.target + symbols=view.symbols, + target=entry.target, + resolved_launches=view.resolved_launches, ) return emit_host_module(TirSquare, (entry,), entry.target, ctx) diff --git a/tests/codegen/test_symbol_table.py b/tests/codegen/test_symbol_table.py deleted file mode 100644 index 397923b4f..000000000 --- a/tests/codegen/test_symbol_table.py +++ /dev/null @@ -1,98 +0,0 @@ -"""The table every emitter reads before it writes a line. - -One row per function a caller can name, settled once for a whole compile, so -a host entry and the shim it calls cannot describe the same call differently. -""" - -from __future__ import annotations - -import pytest - -from tests.fixtures.placed.gpu_placed_rows import GPUS, GpuPlacedRows -from tests.fixtures.tir.square import TirSquare -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.codegen.signature import ( - LAUNCH_ABI, - ProgramIdSignature, - TensorSignature, - program_id_params, - symbol_table, -) -from tilefoundry.ir.core.module import module_functions -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.target import CudaTarget - -_PLACED = CudaTarget("nvidia.h200_sxm", device_count=GPUS) -_GPU_ID = ProgramIdSignature(name="tilefoundry_gpu_program_id", topology_level="gpu") - - -def test_the_ids_a_call_carries_are_the_levels_the_target_states() -> None: - assert program_id_params(GpuPlacedRows, _PLACED) == (_GPU_ID,) - - -def test_a_program_that_names_no_stated_level_carries_no_ids() -> None: - """The same target, a program that never names ``gpu``: nothing to tell it.""" - assert program_id_params(TirSquare, _PLACED) == () - - -def test_a_device_function_is_reached_through_its_shim() -> None: - table = symbol_table(GpuPlacedRows, _PLACED) - shim = table[id(GpuPlacedRows.lookup("copy_rows_device"))] - assert shim.name == "tilefoundry_copy_rows_device_launch" - assert shim.leading == (_GPU_ID,) - assert shim.trailing == LAUNCH_ABI - assert [p.name for p in shim.params] == ["a", "out"] - - -def test_a_host_entry_is_reached_directly_and_told_what_it_places() -> None: - """The id enters the program here: the process is the only thing that knows it.""" - table = symbol_table(GpuPlacedRows, _PLACED) - entry = table[id(GpuPlacedRows.lookup("copy_rows_host"))] - assert entry.name == "tilefoundry_copy_rows_host_host" - assert entry.leading == (_GPU_ID,) - assert entry.trailing == () - - -def test_the_kernels_own_convention_is_not_in_the_table() -> None: - """Only the shim is callable, so only the shim is written down.""" - table = symbol_table(GpuPlacedRows, _PLACED) - assert not any(signature.name.endswith("_kernel") for signature in table.values()) - - -def test_a_prototype_gets_one_row_and_its_variants_none() -> None: - """Nothing calls a variant: the prototype is the one symbol, and it branches.""" - table = symbol_table(TirSquare, _PLACED) - assert sorted(signature.name for signature in table.values()) == [ - "tilefoundry_square_device_launch", - "tilefoundry_square_host_host", - ] - - -def test_a_call_site_looks_its_callee_up_by_identity() -> None: - """Keyed by ``id(fn)``: the call site holds the callee, not its module.""" - assert set(symbol_table(TirSquare, _PLACED)) == { - id(fn) for fn in module_functions(TirSquare) - } - - -def test_the_context_answers_from_the_table_it_was_given() -> None: - device = GpuPlacedRows.lookup("copy_rows_device") - ctx = CudaCodegenContext(symbols=symbol_table(GpuPlacedRows, _PLACED), target=_PLACED) - assert ctx.signature_of(device).name == "tilefoundry_copy_rows_device_launch" - with pytest.raises(KeyError, match="symbol table"): - ctx.signature_of(TirSquare.lookup("square_host")) - - -def test_the_dynamic_axes_are_the_ones_the_type_leaves_open() -> None: - signature = TensorSignature( - name="x", - type=TensorType( - shape=(4, DimVar("S", 1, 8)), - dtype=DType.f32, - layout=None, - storage=StorageKind.GMEM, - ), - ) - assert signature.dynamic_axes == (1,) diff --git a/tests/fixtures/tir/async_sync.py b/tests/fixtures/tir/async_sync.py index 682deefe4..73acbf7f0 100644 --- a/tests/fixtures/tir/async_sync.py +++ b/tests/fixtures/tir/async_sync.py @@ -6,7 +6,7 @@ from tilefoundry.target import CpuTarget, CudaTarget -@module(entry="async_stage_host", target=CudaTarget("nvidia.h200_sxm")) +@module(entry="async_stage_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) class AsyncStage: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def async_stage_device(a: Tensor[(128, 4), "f32"], b: Tensor[(128, 4), "f32"]): diff --git a/tests/fixtures/tir/square.py b/tests/fixtures/tir/square.py index 038e795ed..2d5737358 100644 --- a/tests/fixtures/tir/square.py +++ b/tests/fixtures/tir/square.py @@ -10,7 +10,7 @@ _S = DimVar("S", 1, 256) -@module(entry="square_host", target=CudaTarget("nvidia.h200_sxm")) +@module(entry="square_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) class TirSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def square_device(x: Tensor[(_S,), "f32"]): diff --git a/tests/fixtures/tir/sync.py b/tests/fixtures/tir/sync.py index 7d461e54c..43b95be3c 100644 --- a/tests/fixtures/tir/sync.py +++ b/tests/fixtures/tir/sync.py @@ -7,7 +7,7 @@ from tilefoundry.target import CpuTarget, CudaTarget -@module(entry="sync_square_host", target=CudaTarget("nvidia.h200_sxm")) +@module(entry="sync_square_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) class SyncSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def sync_square_device(a: Tensor[(4, 32), "f32"]): diff --git a/tests/installed/models/smoke_qwen3_1_7b.py b/tests/installed/models/smoke_qwen3_1_7b.py index bc88c2941..51bf9ab81 100644 --- a/tests/installed/models/smoke_qwen3_1_7b.py +++ b/tests/installed/models/smoke_qwen3_1_7b.py @@ -230,10 +230,10 @@ def test_the_mlp_costs_its_three_matrices(tf, shipped_source) -> None: authored = _reported(tf, source, "mlp", None) placed = _reported(tf, source, "placed_mlp", None) authored_flops = authored["totals"]["flops"][DT] - stated = placed["function_records"]["compute-cost"]["flops"][DT] + stated = placed["function_records"]["compute-cost"]["flops"] topologies = placed["function_records"]["compute-cost"]["topologies"] - placed_flops = stated["total"] - per_unit = dict(zip(topologies, stated["per_unit"])) + placed_flops = stated["total"][DT] + per_unit = dict(zip(topologies, (share[DT] for share in stated["per_unit"]))) assert placed_flops == per_unit["cta"] * 128 _holds(authored_flops, _mlp_matmul_flops(), "mlp") diff --git a/tests/ops/ir/test_cache_update.py b/tests/ops/ir/test_cache_update.py index 6350afd50..8e63aa502 100644 --- a/tests/ops/ir/test_cache_update.py +++ b/tests/ops/ir/test_cache_update.py @@ -230,7 +230,7 @@ def test_cache_update_function_analyzes_program_and_cta_cost() -> None: moved = get_metadata(analysed_update, TrafficMetadata) assert result.topology_level == "cta" assert record is not None - assert record.flops.kinds == () + assert record.flops.total == () row_bytes = _WINDOW_BYTES // 4 assert moved.operands == ( TrafficBytes(), diff --git a/tests/ops/ir/test_local.py b/tests/ops/ir/test_local.py index ce2347ea7..28bd4f4ca 100644 --- a/tests/ops/ir/test_local.py +++ b/tests/ops/ir/test_local.py @@ -53,6 +53,6 @@ def test_local_analyzes_as_a_zero_traffic_topology_view() -> None: moved = get_metadata(analysed_local, TrafficMetadata) assert result.topology_level == topology_level assert record is not None - assert record.flops.kinds == () + assert record.flops.total == () assert moved.storage.kinds == () assert moved.operands == (TrafficBytes(), TrafficBytes()) diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 08293a9b6..77d0b14d8 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -35,7 +35,7 @@ def test_handwritten_mma_matches_torch() -> None: torch.testing.assert_close(out, torch.matmul(a.float(), b.float()), rtol=2e-2, atol=2e-2) -@module(entry="tile_host", target=_CUDA) +@module(entry="tile_host", target=_CUDA, topologies=(Topology("thread", 32),)) class MmaTile: """A 16x16 by 16x8 product done as one atom, looped by the tile's shape.""" @@ -68,22 +68,30 @@ def tile_device( ), ) a_tile = T.alloc_tensor( - Tensor[(16, 16), 'bf16', ShardLayout( + Tensor[ + (16, 16), + "bf16", + ShardLayout( layout=Layout(shape=(16, 16), strides=(16, 1)), attrs=(Broadcast(), Broadcast()), mesh=m, - ), 'smem'] + ), + "smem", + ] ) b_tile = T.alloc_tensor( - Tensor[(8, 16), 'bf16', ShardLayout( + Tensor[ + (8, 16), + "bf16", + ShardLayout( layout=Layout(shape=(8, 16), strides=(1, 8)), attrs=(Broadcast(), Broadcast()), mesh=m, - ), 'smem'] - ) - acc = T.alloc_tensor( - Tensor[(16, 8), 'f32', atom.C, 'rmem'] + ), + "smem", + ] ) + acc = T.alloc_tensor(Tensor[(16, 8), "f32", atom.C, "rmem"]) T.copy(a_view, a_tile) T.copy(b_view, b_tile) T.fill(acc, 0.0) diff --git a/tests/ops/tir/cuda/test_tma.py b/tests/ops/tir/cuda/test_tma.py index fa77cab3e..e4a558ab1 100644 --- a/tests/ops/tir/cuda/test_tma.py +++ b/tests/ops/tir/cuda/test_tma.py @@ -73,7 +73,10 @@ def test_refuses_a_shape_change() -> None: verify_prim_function(_pf(_ty(8), _ty(4, storage="smem"))) -@pytest.mark.parametrize(("n", "dtype"), [(5, DType.from_name("f32")), (1, DType.from_name("f32")), (7, DType.from_name("bf16"))]) +@pytest.mark.parametrize( + ("n", "dtype"), + [(5, DType.from_name("f32")), (1, DType.from_name("f32")), (7, DType.from_name("bf16"))], +) def test_admits_a_transfer_off_the_sixteen_byte_grain(n, dtype) -> None: """The grain belongs to one instruction, and the op does not name one. @@ -82,7 +85,11 @@ def test_admits_a_transfer_off_the_sixteen_byte_grain(n, dtype) -> None: verify_prim_function(_pf(_ty(n, dtype), _ty(n, dtype, storage="smem"))) -@module(entry="tma_tiers_host", target=CudaTarget("nvidia.h200_sxm")) +@module( + entry="tma_tiers_host", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("thread", 128),), +) class TmaTiers: """Both staging tiers in one device function, one operand pair each.""" @@ -101,15 +108,18 @@ def tma_tiers_device( ), ) bulk_stage = T.alloc_tensor( - Tensor[(256,), 'f32', ShardLayout( + Tensor[ + (256,), + "f32", + ShardLayout( layout=Layout(shape=(256,), strides=(1,)), attrs=(Broadcast(),), mesh=m, - ), 'smem'] - ) - bulk_bar = T.alloc_tensor( - Tensor[(1,), 'i64', None, 'smem'] + ), + "smem", + ] ) + bulk_bar = T.alloc_tensor(Tensor[(1,), "i64", None, "smem"]) T.mbarrier_init(bulk_bar, arrive_count=1) T.sync(m) T.tma_copy(bulk_view, bulk_stage, bulk_bar) @@ -123,15 +133,18 @@ def tma_tiers_device( ), ) odd_stage = T.alloc_tensor( - Tensor[(5,), 'f32', ShardLayout( + Tensor[ + (5,), + "f32", + ShardLayout( layout=Layout(shape=(5,), strides=(1,)), attrs=(Broadcast(),), mesh=mo, - ), 'smem'] - ) - odd_bar = T.alloc_tensor( - Tensor[(1,), 'i64', None, 'smem'] + ), + "smem", + ] ) + odd_bar = T.alloc_tensor(Tensor[(1,), "i64", None, "smem"]) T.mbarrier_init(odd_bar, arrive_count=1) T.sync(mo) T.tma_copy(odd_view, odd_stage, odd_bar) diff --git a/tests/ops/tir/test_copy.py b/tests/ops/tir/test_copy.py index 02583aa2d..3095c4a75 100644 --- a/tests/ops/tir/test_copy.py +++ b/tests/ops/tir/test_copy.py @@ -19,7 +19,7 @@ _CUDA = CudaTarget("nvidia.h200_sxm") -@module(entry="copy_storage_host", target=_CUDA) +@module(entry="copy_storage_host", target=_CUDA, topologies=(Topology("thread", 128),)) class CopyStorage: """Five operand pairs in one device function, one per thing that can differ.""" @@ -39,44 +39,32 @@ def copy_storage_device( with Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,)), ("t",)) as m: smem_src = T.tensor_view(a_smem, layout=split_rows(m)) smem_dst = T.tensor_view(b_smem, layout=split_rows(m)) - smem_tile = T.alloc_tensor( - Tensor[(128, 4), 'f32', split_rows(m), 'smem'] - ) + smem_tile = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "smem"]) T.copy(smem_src, smem_tile) T.sync(m) T.copy(smem_tile, smem_dst) rmem_src = T.tensor_view(a_rmem, layout=split_rows(m)) rmem_dst = T.tensor_view(b_rmem, layout=split_rows(m)) - rmem_tile = T.alloc_tensor( - Tensor[(128, 4), 'f32', split_rows(m), 'smem'] - ) - rmem_frag = T.alloc_tensor( - Tensor[(128, 4), 'f32', split_rows(m), 'rmem'] - ) + rmem_tile = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "smem"]) + rmem_frag = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "rmem"]) T.copy(rmem_src, rmem_tile) T.sync(m) T.copy(rmem_tile, rmem_frag) T.copy(rmem_frag, rmem_dst) wide_src = T.tensor_view(a_wide, layout=split_rows(m)) wide_dst = T.tensor_view(b_wide, layout=split_rows(m)) - wide_frag = T.alloc_tensor( - Tensor[(128, 4), 'f32', split_rows(m), 'rmem'] - ) + wide_frag = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "rmem"]) T.copy(wide_src, wide_frag) T.copy(wide_frag, wide_dst) narrow_src = T.tensor_view(a_narrow, layout=split_pairs(m)) narrow_dst = T.tensor_view(b_narrow, layout=split_pairs(m)) - narrow_frag = T.alloc_tensor( - Tensor[(128, 2), 'f32', split_pairs(m), 'rmem'] - ) + narrow_frag = T.alloc_tensor(Tensor[(128, 2), "f32", split_pairs(m), "rmem"]) T.copy(narrow_src, narrow_frag) T.copy(narrow_frag, narrow_dst) with Mesh((Topology("thread", 32),), Layout(shape=(32,), strides=(1,)), ("t",)) as mb: bcast_src = T.tensor_view(a_bcast, layout=split_short_rows(mb)) bcast_dst = T.tensor_view(b_bcast, layout=split_short_rows(mb)) - bcast_frag = T.alloc_tensor( - Tensor[(4,), 'f32', broadcast_run(mb), 'rmem'] - ) + bcast_frag = T.alloc_tensor(Tensor[(4,), "f32", broadcast_run(mb), "rmem"]) T.copy(bcast_src, bcast_frag) T.copy(bcast_frag, bcast_dst) diff --git a/tests/ops/tir/test_sync.py b/tests/ops/tir/test_sync.py index 03aba153b..77fffb7fd 100644 --- a/tests/ops/tir/test_sync.py +++ b/tests/ops/tir/test_sync.py @@ -18,7 +18,7 @@ _CUDA = CudaTarget("nvidia.h200_sxm") -@module(entry="slice_host", target=_CUDA) +@module(entry="slice_host", target=_CUDA, topologies=(Topology("thread", 128),)) class SyncSlices: @prim_func(target=_CUDA) def one_warp(a: Tensor[(1, 32), "f32"], o: Tensor[(1, 32), "f32"]): @@ -55,7 +55,11 @@ def slice_host( launch(two_warps, a1, o1, grid=(1, 1, 1), block=(128, 1, 1)) # noqa: F821 -@module(entry="grid_sync_host", target=_CUDA) +@module( + entry="grid_sync_host", + target=_CUDA, + topologies=(Topology("cta", 4), Topology("thread", 128)), +) class GridSync: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def grid_sync_device(a: Tensor[(128,), "f32"]): diff --git a/tests/passes/test_host_entry.py b/tests/passes/test_host_entry.py index a3b9a7fab..ed4452b9f 100644 --- a/tests/passes/test_host_entry.py +++ b/tests/passes/test_host_entry.py @@ -67,29 +67,6 @@ def large(x: Tensor[(_S,), "f32"]): T.sync(thread) -@module(entry="square", topologies=(Topology("thread", 128),)) -class _DisagreeingVariants: - """The same shape, except the two variants mesh over different thread counts.""" - - @prim_func(target=_CUDA) - def square(x: Tensor[(_S,), "f32"]): - pass - - @square.specialize(DimVarRangePat("S", 1, 127)) - def small(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 64),), Layout((64,), (1,))) as thread: - view = T.tensor_view(x, layout=_rows(64)) - T.copy(view, view) - T.sync(thread) - - @square.specialize(DimVarRangePat("S", 128, 255)) - def large(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=_rows(128)) - T.copy(view, view) - T.sync(thread) - - def _launch_of(mod): """The one ``Launch`` the synthesized entry's body is made of.""" (statement,) = mod.entry_function().body.body @@ -128,10 +105,5 @@ def test_a_lone_kernel_reads_its_geometry_off_its_own_body() -> None: assert _geometry(_launch_of(insert_default_host_entry(_OneKernel))) == ((1, 1, 1), (128, 1, 1)) -def test_variants_that_want_different_geometries_are_refused() -> None: - with pytest.raises(ValueError, match="different geometries"): - insert_default_host_entry(_DisagreeingVariants) - - def test_a_module_that_already_has_a_host_entry_is_left_alone() -> None: assert insert_default_host_entry(TirSquare) is TirSquare diff --git a/tests/target/test_target.py b/tests/target/test_target.py index 9736e566d..34ff260f7 100644 --- a/tests/target/test_target.py +++ b/tests/target/test_target.py @@ -20,7 +20,6 @@ from tests.installed.smoke_target.vendor_npu import VendorNpuTarget from tilefoundry import CompilerOptions, DType, build, jit, module from tilefoundry.analysis import AnalysisError, analyze -from tilefoundry.codegen.registry import group_functions_by_target from tilefoundry.dsl import DimVar from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.prim_function import PrimFunction @@ -126,9 +125,7 @@ def test_document_free_target_projects_facts_and_inherits_standard_analysis() -> target = VendorNpuTarget() assert target.get_analyzer("roofline").selector == "roofline" - assert target.get_facts(TopologyLimitFacts, "core") == TopologyLimitFacts( - "core", 256 - ) + assert target.get_facts(TopologyLimitFacts, "core") == TopologyLimitFacts("core", 256) memory = target.get_facts(MemoryHierarchyFacts) assert memory.explicit("gmem").capacity_bytes == 64_000_000_000 assert target.get_facts(ThroughputFacts).peak_for(DType.f32) == 2_000_000_000_000_000 @@ -152,8 +149,7 @@ def test_cuda_projects_one_ctas_share_from_device_rates() -> None: assert services.unit == "cta" assert dict(services.unit_flops) == { - dtype: rate // target.device.sm_count - for dtype, rate in device.peak_flops_per_second + dtype: rate // target.device.sm_count for dtype, rate in device.peak_flops_per_second } assert dict(services.unit_bandwidth) == { device.bandwidth_level: ( @@ -289,36 +285,6 @@ def test_program_topologies_use_target_resource_facts() -> None: assert len(str(error.value)) < 120 -def test_group_functions_by_target_fact_matching() -> None: - """CUDA functions must agree on Target facts before grouping. - - CUDA functions must agree on Target facts before grouping; CPU - functions are exempt from the CUDA fact-matching. - """ - body = Sequential(body=()) - first = PrimFunction(name="first", params=(), body=body, target=CudaTarget("nvidia.h200_sxm")) - second = PrimFunction( - name="second", - params=(), - body=body, - target=CudaTarget( - "nvidia.h200_sxm", - architecture=replace(CudaTarget.hardware.resolve(SM90_ID).value, name="sm_90_alt"), - ), - ) - with pytest.raises(ValueError, match="mixes unequal device Targets") as error: - group_functions_by_target(Module(name="mixed", functions=(first, second), entry="first")) - assert "CudaTarget (cuda)" in str(error.value) - assert "architecture=" not in str(error.value) - assert "device=" not in str(error.value) - assert len(str(error.value)) < 300 - - host = PrimFunction(name="host", params=(), body=body, target=CpuTarget()) - groups = group_functions_by_target(Module(name="mixed", functions=(first, host), entry="host")) - assert tuple(fn.name for fn in groups[first.target]) == ("first",) - assert tuple(fn.name for fn in groups[host.target]) == ("host",) - - def test_build_rejects_unequal_device_targets_before_emitting() -> None: emitted: list[object] = [] @@ -349,12 +315,7 @@ def get_code_generator(self) -> CodeGenerator: with pytest.raises(ValueError) as error: build(mixed) - assert str(error.value) == ( - "tilefoundry: module 'mixed' mixes unequal device Targets: " - "FirstDeviceTarget (tests.target.first_device) (function 'first') vs " - "SecondDeviceTarget (tests.target.second_device) (function 'second'); " - "multiple device translation units are not supported" - ) + assert str(error.value) == "codegen: expected one device Target, found 2" assert emitted == [] From ca65c898d8ade387e56fe24cf93ffc094fcd3a56 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 15 Sep 2026 14:10:01 +0800 Subject: [PATCH 4/7] docs(tutorial): resolve showcase output conflict --- docs/tutorial/showcase.ipynb | 4 ++-- docs/tutorial/showcase.md | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index 62e6201d5..b453daabc 100644 --- a/docs/tutorial/showcase.ipynb +++ b/docs/tutorial/showcase.ipynb @@ -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:337408@cta:42176,f32:51124224@cta:6390528 service=special:262144@cta:32768\n# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144\n# peak-footprint=gmem:10945036,smem:16960\n# roofline ideal-ns=11213 bound-by=memory\n\n v1 = reshard(w_q, layout=ShardLayout(\n layout=Layout((1, 256, 8, 32), None),\n attrs=(S(2),),\n mesh=cta,\n ), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n\n v42 = reshard(w_o, layout=ShardLayout(\n layout=Layout((1, 256, 8, 32), None),\n attrs=(S(2),),\n mesh=cta,\n ), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\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/337408@cta:42176,f32:6390528/51124224@cta:6390528 other-ops=special:32768/262144@cta:32768\n# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144\n# peak-footprint=gmem:10945036,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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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/131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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/131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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:1246400@cta:328672,f32:6410280@cta:801285 service=integer:256@cta:32,special:33040@cta:4130\n# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804\n# peak-footprint=gmem:2950412,rmem:0,smem:33536\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@cta:r32/w0 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@cta:r136/w128 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/1246400@cta:328672,f32:6410280/6410280@cta:801285 other-ops=integer:32/256@cta:32,special:33040/33040@cta:4130\n# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804\n# peak-footprint=gmem:2950412,rmem:0,smem:33536\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@cta:r32/w0 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@cta:r136/w128 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" diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index c9331af5b..471ac58b5 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -871,10 +871,10 @@ for needle in ("reshard(w_q", "reshard(w_o"): # 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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute + 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/131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute + 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/131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute ``` ## 6. Stream the KV cache From 6f67ee254259ba13528c5866bf3dd35995d28628 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 15 Sep 2026 14:20:40 +0800 Subject: [PATCH 5/7] test(analysis): refresh placed metadata golden --- tests/fixtures/inspection/type_printer_sugar.analyzed.txt | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index f3f70372a..ddaaf79a4 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -20,17 +20,17 @@ def composed_mesh_pipeline( ): 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@cta:r512/w0,thread:r64/w0,rmem:r0/w2048@cta:r0/w512,thread:r0/w64 - 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@cta:128,thread:16; traffic traffic=rmem:r2048/w2048@cta:r512/w512,thread:r64/w64 + 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/512@cta:128@thread:16; traffic traffic=rmem:r2048/w2048@cta:r512/w512,thread:r64/w64 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@cta:r512/w0,thread:r64/w0,smem:r0/w2048@cta:r0/w512,thread:r0/w64 - 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@cta:128,thread:128; traffic traffic=smem:r2048/w1024@cta:r512/w256,thread:r512/w256 + 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/512@cta:128@thread:128; traffic traffic=smem:r2048/w1024@cta:r512/w256,thread:r512/w256 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@cta:r256/w256,thread:r256/w256 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@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 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@cta:r64/w0,thread:r8/w0,rmem:r0/w64@cta:r0/w64,thread:r0/w8 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@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 - 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@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 + v10 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost flops=f32:128/128@cta:128@thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 + 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/128@cta:128@thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 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@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 From 36252ea5e42ae5d7e0dac8ce36f9e295b64d50cc Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 15 Sep 2026 15:48:22 +0800 Subject: [PATCH 6/7] refactor(analysis): unify keyed spread metadata --- docs/spec/analysis.md | 18 ++--- docs/tutorial/optimize.ipynb | 4 +- docs/tutorial/optimize.md | 4 +- docs/tutorial/showcase.ipynb | 16 ++--- docs/tutorial/showcase.md | 70 +++++++++---------- src/tilefoundry/analysis/compute_cost.py | 59 ++++++++++------ src/tilefoundry/analysis/metadata.py | 29 ++++---- src/tilefoundry/analysis/roofline.py | 2 +- src/tilefoundry/compile.py | 4 +- src/tilefoundry/inspection/values.py | 55 ++++++--------- tests/analysis/test_analysis_families.py | 20 ++++-- tests/analysis/test_analyze_at_a_size.py | 10 +-- tests/analysis/test_analyze_cross_module.py | 3 +- tests/cli/test_cli.py | 18 ++--- .../type_printer_sugar.analyzed.txt | 26 +++---- tests/ops/ir/test_cache_update.py | 2 +- tests/ops/ir/test_local.py | 2 +- 17 files changed, 179 insertions(+), 163 deletions(-) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 9a4655001..580eda05a 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -169,8 +169,8 @@ class ComputeCostMetadata(IRMetadata): """Typed work in logical, expanded, and topology-unit domains.""" topologies: tuple[str, ...] = () - flops: Spread[tuple[tuple[str, int], ...]] = Spread((), (), ()) - other_ops: Spread[tuple[tuple[str, int], ...]] = Spread((), (), ()) + flops: Breakdown[int] = Breakdown() + other_ops: Breakdown[int] = Breakdown() ``` @@ -185,7 +185,7 @@ record, stated exactly as a Call's is. The whole program's work is not a second record. ```text -compute-cost flops=:/@:[,...] other-ops=:/@:[,...] +compute-cost flops=:@logical,@total,@[,...][;:...] other-ops=:@logical,@total,@[,...][;:...] ``` Every measured Call receives this annotation. Each key pairs the whole quantity @@ -195,10 +195,10 @@ Each reported Call's JSON projection is under its `compute-cost` key: ```text {"topologies": [, ...], - "flops": {"logical": {: }, "total": {: }, - "per_unit": [{: }, ...]}, - "other_ops": {"logical": {: }, "total": {: }, - "per_unit": [{: }, ...]}} + "flops": {: {"logical": , "total": , + "per_unit": [, ...]}}, + "other_ops": {: {"logical": , "total": , + "per_unit": [, ...]}}} ``` - constraints: @@ -540,7 +540,7 @@ Requesting memory adds the Function's own movement, one footprint line, and one line per advisory: ```text -traffic traffic=:r/w@r/w[,...] +traffic traffic=:r/w@total,r/w@[,...][;:...] peak-footprint=:[,:...] advisory="" ``` @@ -559,7 +559,7 @@ is emitted only when asked for ([cli Analyze](./cli.md#analyze)) and is absent from a Function, which has no split: ```text -traffic traffic=:r/w@r/w[,...][ operands=:r/w[,...]] +traffic traffic=:r/w@total,r/w@[,...][;:...] [operands=:r/w[;:...]] ``` Its JSON projection is under the reported value's `traffic` key, with `whole`, diff --git a/docs/tutorial/optimize.ipynb b/docs/tutorial/optimize.ipynb index 6910a0b72..67cb3a31d 100644 --- a/docs/tutorial/optimize.ipynb +++ b/docs/tutorial/optimize.ipynb @@ -54,7 +54,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# traffic traffic=gmem:r71680/w43456@cta:r71680/w43456,thread:r336/w224,rmem:r603496/w488344@cta:r603496/w488344,thread:r2604/w2084\n# roofline ideal-ns=24 bound-by=memory\n" + "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" } ], "source": [ @@ -94,7 +94,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# traffic traffic=gmem:r43008/w14784@cta:r43008/w14784,thread:r224/w112,rmem:r574824/w459672@cta:r574824/w459672,thread:r2492/w1972\n# roofline ideal-ns=13 bound-by=memory\n" + "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" } ], "source": [ diff --git a/docs/tutorial/optimize.md b/docs/tutorial/optimize.md index ddac6c2fe..a9aa7790b 100644 --- a/docs/tutorial/optimize.md +++ b/docs/tutorial/optimize.md @@ -88,7 +88,7 @@ grep -E '^# (traffic|roofline) ' Naive.txt ``` ```text -# traffic traffic=gmem:r71680/w43456@cta:r71680/w43456,thread:r336/w224,rmem:r603496/w488344@cta:r603496/w488344,thread:r2604/w2084 +# traffic traffic=gmem:r71680/w43456@total,r71680/w43456@cta,r336/w224@thread;rmem:r603496/w488344@total,r603496/w488344@cta,r2604/w2084@thread # roofline ideal-ns=24 bound-by=memory ``` @@ -132,7 +132,7 @@ grep -E '^# (traffic|roofline) ' Fused.txt ``` ```text -# traffic traffic=gmem:r43008/w14784@cta:r43008/w14784,thread:r224/w112,rmem:r574824/w459672@cta:r574824/w459672,thread:r2492/w1972 +# traffic traffic=gmem:r43008/w14784@total,r43008/w14784@cta,r224/w112@thread;rmem:r574824/w459672@total,r574824/w459672@cta,r2492/w1972@thread # roofline ideal-ns=13 bound-by=memory ``` diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index b453daabc..ab866d675 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/328896@cta:328896,f32:200448/200448@cta:200448 other-ops=special:1024/1024@cta:1024\n# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592\n# peak-footprint=gmem:1675788\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/131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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@cta:r136/w128 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/131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta\n# peak-footprint=gmem:1675788\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" } ], "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" @@ -137,7 +137,7 @@ { "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/200448@cta:200448` | `gmem:r2225620/w806592@cta:r2225620/w806592` | 1675788 | 632 | memory |\n| 512 | `799488/799488@cta:799488` | `gmem:r4744660/w3202752@cta:r4744660/w3202752` | 2572812 | 1656 | memory |\n| 1024 | `1598208/1598208@cta:1598208` | `gmem:r8103380/w6397632@cta:r8103380/w6397632` | 3768844 | 3022 | memory |\n| 2048 | `3195648/3195648@cta:3195648` | `gmem:r14820820/w12787392@cta:r14820820/w12787392` | 6160908 | 5752 | memory |\n| 4096 | `6390528/6390528@cta:6390528` | `gmem:r28255700/w25566912@cta:r28255700/w25566912` | 10945036 | 11214 | memory |\n| 8192 | `12780288/12780288@cta:12780288` | `gmem:r55125460/w51125952@cta:r55125460/w51125952` | 20513292 | 22136 | memory |\n" + "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` | 1675788 | 632 | memory |\n| 512 | `799488@logical,799488@total,799488@cta` | `gmem:r4744660/w3202752@total,r4744660/w3202752@cta` | 2572812 | 1656 | memory |\n| 1024 | `1598208@logical,1598208@total,1598208@cta` | `gmem:r8103380/w6397632@total,r8103380/w6397632@cta` | 3768844 | 3022 | memory |\n| 2048 | `3195648@logical,3195648@total,3195648@cta` | `gmem:r14820820/w12787392@total,r14820820/w12787392@cta` | 6160908 | 5752 | memory |\n| 4096 | `6390528@logical,6390528@total,6390528@cta` | `gmem:r28255700/w25566912@total,r28255700/w25566912@cta` | 10945036 | 11214 | memory |\n| 8192 | `12780288@logical,12780288@total,12780288@cta` | `gmem:r55125460/w51125952@total,r55125460/w51125952@cta` | 20513292 | 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" @@ -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/2629376@cta:328672,f32:2833728/2833728@cta:354216 other-ops=special:14528/14528@cta:1816\n# traffic traffic=gmem:r5563796/w3721856@cta:r3936212/w3721408,smem:r9597248/w9480000@cta:r1199656/w1185000\n# peak-footprint=gmem:3701260,smem:472160\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# traffic traffic=gmem:r5563796/w3721856@total,r3936212/w3721408@cta;smem:r9597248/w9480000@total,r1199656/w1185000@cta\n# peak-footprint=gmem:3701260;smem:472160\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" @@ -274,7 +274,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# compute-cost flops=bf16:328896/328896@cta:328896,f32:200448/200448@cta:200448 other-ops=special:1024/1024@cta:1024\n# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592\n# peak-footprint=gmem:1675788\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# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta\n# peak-footprint=gmem:1675788\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" @@ -309,7 +309,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# compute-cost flops=bf16:328896/2629376@cta:328672,f32:200448/200448@cta:25056 other-ops=special:1024/1024@cta:128\n# traffic traffic=gmem:r1674644/w264832@cta:r1559508/w264384,smem:r684608/w675392@cta:r85576/w84424\n# peak-footprint=gmem:1540620,smem:33280\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# traffic traffic=gmem:r1674644/w264832@total,r1559508/w264384@cta;smem:r684608/w675392@total,r85576/w84424@cta\n# peak-footprint=gmem:1540620;smem:33280\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" @@ -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/4392896@cta:328672,f32:3239808/6418944@cta:200592 other-ops=integer:9/288@cta:9,special:33056/33152@cta:1036\n# traffic traffic=gmem:r3476884/w4196992@cta:r2558932/w4196544,rmem:r656/w72@cta:r656/w72,smem:r5839296/w5662784@cta:r682464/w672676\n# peak-footprint=gmem:5047436,rmem:16,smem:33408\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@cta:r136/w128 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# 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:5047436;rmem:16;smem:33408\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" } ], "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/337408@cta:42176,f32:6390528/51124224@cta:6390528 other-ops=special:32768/262144@cta:32768\n# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144\n# peak-footprint=gmem:10945036,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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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/131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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/131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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# traffic traffic=gmem:r28254676/w25566912@total,r27967956/w25565792@cta;smem:r331008/w329984@total,r43168/w42144@cta\n# peak-footprint=gmem:10945036;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" } ], "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/1246400@cta:328672,f32:6410280/6410280@cta:801285 other-ops=integer:32/256@cta:32,special:33040/33040@cta:4130\n# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804\n# peak-footprint=gmem:2950412,rmem:0,smem:33536\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@cta:r32/w0 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@cta:r136/w128 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# 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:2950412;rmem:0;smem:33536\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" } ], "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" diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index 471ac58b5..2305af926 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -197,14 +197,14 @@ for needle in ("matmul(hidden, w_q", "cache_update(k_cache", "matmul(v33, w_o"): ```text # 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/328896@cta:328896,f32:200448/200448@cta:200448 other-ops=special:1024/1024@cta:1024 -# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592 +# 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:1675788 # 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/131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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@cta:r136/w128 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/131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 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; 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 ``` The number before `@` is global work or traffic. After it comes one share per @@ -259,12 +259,12 @@ for ctx_len in (128, 512, 1024, 2048, 4096, 8192): | `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound | |---:|---:|---|---:|---:|---| -| 128 | `200448/200448@cta:200448` | `gmem:r2225620/w806592@cta:r2225620/w806592` | 1675788 | 632 | memory | -| 512 | `799488/799488@cta:799488` | `gmem:r4744660/w3202752@cta:r4744660/w3202752` | 2572812 | 1656 | memory | -| 1024 | `1598208/1598208@cta:1598208` | `gmem:r8103380/w6397632@cta:r8103380/w6397632` | 3768844 | 3022 | memory | -| 2048 | `3195648/3195648@cta:3195648` | `gmem:r14820820/w12787392@cta:r14820820/w12787392` | 6160908 | 5752 | memory | -| 4096 | `6390528/6390528@cta:6390528` | `gmem:r28255700/w25566912@cta:r28255700/w25566912` | 10945036 | 11214 | memory | -| 8192 | `12780288/12780288@cta:12780288` | `gmem:r55125460/w51125952@cta:r55125460/w51125952` | 20513292 | 22136 | memory | +| 128 | `200448@logical,200448@total,200448@cta` | `gmem:r2225620/w806592@total,r2225620/w806592@cta` | 1675788 | 632 | memory | +| 512 | `799488@logical,799488@total,799488@cta` | `gmem:r4744660/w3202752@total,r4744660/w3202752@cta` | 2572812 | 1656 | memory | +| 1024 | `1598208@logical,1598208@total,1598208@cta` | `gmem:r8103380/w6397632@total,r8103380/w6397632@cta` | 3768844 | 3022 | memory | +| 2048 | `3195648@logical,3195648@total,3195648@cta` | `gmem:r14820820/w12787392@total,r14820820/w12787392@cta` | 6160908 | 5752 | memory | +| 4096 | `6390528@logical,6390528@total,6390528@cta` | `gmem:r28255700/w25566912@total,r28255700/w25566912@cta` | 10945036 | 11214 | memory | +| 8192 | `12780288@logical,12780288@total,12780288@cta` | `gmem:r55125460/w51125952@total,r55125460/w51125952@cta` | 20513292 | 22136 | memory | The table says: @@ -426,9 +426,9 @@ print(report.partition("\n\n")[0].rstrip()) ```text # 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/2629376@cta:328672,f32:2833728/2833728@cta:354216 other-ops=special:14528/14528@cta:1816 -# traffic traffic=gmem:r5563796/w3721856@cta:r3936212/w3721408,smem:r9597248/w9480000@cta:r1199656/w1185000 -# peak-footprint=gmem:3701260,smem:472160 +# 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:3701260;smem:472160 # roofline ideal-ns=1935 bound-by=memory ``` @@ -554,8 +554,8 @@ for line in report.splitlines(): ``` ```text -# compute-cost flops=bf16:328896/328896@cta:328896,f32:200448/200448@cta:200448 other-ops=special:1024/1024@cta:1024 -# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592 +# 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:1675788 # roofline ideal-ns=632 bound-by=memory ``` @@ -581,9 +581,9 @@ for line in report.splitlines(): ``` ```text -# compute-cost flops=bf16:328896/2629376@cta:328672,f32:200448/200448@cta:25056 other-ops=special:1024/1024@cta:128 -# traffic traffic=gmem:r1674644/w264832@cta:r1559508/w264384,smem:r684608/w675392@cta:r85576/w84424 -# peak-footprint=gmem:1540620,smem:33280 +# 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:1540620;smem:33280 # roofline ideal-ns=405 bound-by=memory ``` @@ -728,12 +728,12 @@ print(next(line.rstrip() for line in annotated.splitlines() if "cache_update(k_c ```text # 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/4392896@cta:328672,f32:3239808/6418944@cta:200592 other-ops=integer:9/288@cta:9,special:33056/33152@cta:1036 -# traffic traffic=gmem:r3476884/w4196992@cta:r2558932/w4196544,rmem:r656/w72@cta:r656/w72,smem:r5839296/w5662784@cta:r682464/w672676 -# peak-footprint=gmem:5047436,rmem:16,smem:33408 +# 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:5047436;rmem:16;smem:33408 # roofline ideal-ns=1599 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@cta:r136/w128 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; 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 ``` The embedded `Stage3_Fused` program is the split-K example for this page. @@ -865,16 +865,16 @@ for needle in ("reshard(w_q", "reshard(w_o"): ```text # 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/337408@cta:42176,f32:6390528/51124224@cta:6390528 other-ops=special:32768/262144@cta:32768 -# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144 -# peak-footprint=gmem:10945036,smem:16960 +# 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:10945036;smem:16960 # 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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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/131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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; 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 - 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@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 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/131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 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 ``` ## 6. Stream the KV cache @@ -1017,13 +1017,13 @@ for needle in ("slice(k_cache", "cache_update(k_cache"): ```text # 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/1246400@cta:328672,f32:6410280/6410280@cta:801285 other-ops=integer:32/256@cta:32,special:33040/33040@cta:4130 -# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804 -# peak-footprint=gmem:2950412,rmem:0,smem:33536 +# 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:2950412;rmem:0;smem:33536 # 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@cta:r32/w0 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@cta:r136/w128 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; 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 ``` ```text diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index 7c760d92f..c8907c334 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, Spread, shares +from .metadata import Breakdown, ComputeCostMetadata, breakdown, shares from .visitor import AnalyzeContext SELECTOR = "compute-cost" @@ -28,9 +28,8 @@ def _counts(values: dict) -> tuple[tuple[str, int], ...]: return tuple(sorted((getattr(kind, "name", kind), value) for kind, value in values.items())) -def _at(spread: Spread, topologies: tuple[str, ...], level: str | None): - index = topologies.index(level) if level in topologies else None - return spread.total if index is None else spread.at(index) +def _at(held: Breakdown[int], topologies: tuple[str, ...], level: str | None): + return tuple(shares(held, topologies, level).items()) def _is_structural_occurrence( @@ -161,8 +160,18 @@ def _call_cost_record( total_other_ops = tuple((kind, value * repeats) for kind, value in unit_other_ops) return ComputeCostMetadata( topologies=tuple(locals_by_unit), - flops=Spread(_counts(logical.flops), total_flops, tuple(flops_by_unit)), - other_ops=Spread(_counts(logical.service), total_other_ops, tuple(other_ops_by_unit)), + flops=breakdown( + dict(total_flops), + tuple(dict(values) for values in flops_by_unit), + 0, + logical=dict(_counts(logical.flops)), + ), + other_ops=breakdown( + dict(total_other_ops), + tuple(dict(values) for values in other_ops_by_unit), + 0, + logical=dict(_counts(logical.service)), + ), ) @@ -201,15 +210,23 @@ def _add(target: dict[str, int], values, trips: int) -> None: target[name] = target.get(name, 0) + value * trips +def _domain_values(held: Breakdown[int], domain: str) -> tuple[tuple[str, int], ...]: + return tuple((name, getattr(spread, domain)) for name, spread in held.kinds) + + def _accumulate(ctx: "ComputeCostContext", record: ComputeCostMetadata, trips: int) -> None: - _add(ctx.flops_logical, record.flops.logical, trips) - _add(ctx.flops, record.flops.total, trips) - _add(ctx.other_ops_logical, record.other_ops.logical, trips) - _add(ctx.other_ops, record.other_ops.total, trips) + _add(ctx.flops_logical, _domain_values(record.flops, "logical"), trips) + _add(ctx.flops, _domain_values(record.flops, "total"), trips) + _add(ctx.other_ops_logical, _domain_values(record.other_ops, "logical"), trips) + _add(ctx.other_ops, _domain_values(record.other_ops, "total"), trips) for index, unit in enumerate(record.topologies): held = ctx.by_unit.setdefault(unit, {"flops": {}, "other_ops": {}}) - _add(held["flops"], record.flops.at(index), trips) - _add(held["other_ops"], record.other_ops.at(index), trips) + _add(held["flops"], ((name, spread.at(index)) for name, spread in record.flops.kinds), trips) + _add( + held["other_ops"], + ((name, spread.at(index)) for name, spread in record.other_ops.kinds), + trips, + ) @dataclass @@ -299,15 +316,17 @@ def analyze_compute_cost(function: Function, context: AnalyzeContext) -> None: function, ComputeCostMetadata( topologies=units, - flops=Spread( - _counts(cost_context.flops_logical), - _counts(cost_context.flops), - tuple(_counts(cost_context.by_unit[u]["flops"]) for u in units), + flops=breakdown( + cost_context.flops, + tuple(cost_context.by_unit[u]["flops"] for u in units), + 0, + logical=cost_context.flops_logical, ), - other_ops=Spread( - _counts(cost_context.other_ops_logical), - _counts(cost_context.other_ops), - tuple(_counts(cost_context.by_unit[u]["other_ops"]) for u in units), + other_ops=breakdown( + cost_context.other_ops, + tuple(cost_context.by_unit[u]["other_ops"] for u in units), + 0, + logical=cost_context.other_ops_logical, ), ), ) diff --git a/src/tilefoundry/analysis/metadata.py b/src/tilefoundry/analysis/metadata.py index 1f85f9b15..9e015eecf 100644 --- a/src/tilefoundry/analysis/metadata.py +++ b/src/tilefoundry/analysis/metadata.py @@ -46,16 +46,21 @@ def names(self) -> tuple[str, ...]: def breakdown[V]( - total: "Mapping[str, V]", per_unit: "Sequence[Mapping[str, V]]", zero: V + total: "Mapping[str, V]", + per_unit: "Sequence[Mapping[str, V]]", + zero: V, + *, + logical: "Mapping[str, V] | None" = None, ) -> Breakdown[V]: - """Gather traffic kinds, filling missing values at each topology level.""" - kinds = sorted({*total, *(kind for share in per_unit for kind in share)}) + """Gather keyed spreads, filling missing values in every counting domain.""" + logical = total if logical is None else logical + kinds = sorted({*logical, *total, *(kind for share in per_unit for kind in share)}) return Breakdown( tuple( ( kind, Spread( - logical=total.get(kind, zero), + logical=logical.get(kind, zero), total=total.get(kind, zero), per_unit=tuple(share.get(kind, zero) for share in per_unit), ), @@ -66,30 +71,24 @@ def breakdown[V]( def shares[V]( - held: Breakdown[V] | Spread[tuple[tuple[str, V], ...]], + held: Breakdown[V], topologies: tuple[str, ...], topology_level: "str | None" = None, ) -> dict[str, V]: - """Each traffic kind's value at one topology level, or its total.""" + """Each kind's value at one topology level, or its total.""" index = topologies.index(topology_level) if topology_level in topologies else None - if isinstance(held, Spread): - value = held.total if index is None else held.at(index) - return dict(value) return { kind: spread.total if index is None else spread.at(index) for kind, spread in held.kinds } -TypedCounts = tuple[tuple[str, int], ...] - - @dataclass(frozen=True) class ComputeCostMetadata(IRMetadata): - """Typed floating-point and other operation work in each counting domain.""" + """Floating-point and other operation kinds, each with one counting spread.""" topologies: tuple[str, ...] = () - flops: Spread[TypedCounts] = Spread((), (), ()) - other_ops: Spread[TypedCounts] = Spread((), (), ()) + flops: Breakdown[int] = Breakdown() + other_ops: Breakdown[int] = Breakdown() @dataclass(frozen=True) diff --git a/src/tilefoundry/analysis/roofline.py b/src/tilefoundry/analysis/roofline.py index 00e88195a..6f24a22fd 100644 --- a/src/tilefoundry/analysis/roofline.py +++ b/src/tilefoundry/analysis/roofline.py @@ -32,7 +32,7 @@ def _totals(held) -> tuple[tuple[str, int], ...]: """One typed work category's whole-program counts.""" - return held.total + return tuple((name, spread.total) for name, spread in held.kinds) def _compute_ns(flops: tuple[tuple[str, int], ...], facts: ThroughputFacts) -> int: diff --git a/src/tilefoundry/compile.py b/src/tilefoundry/compile.py index fdf233e9a..3904a9100 100644 --- a/src/tilefoundry/compile.py +++ b/src/tilefoundry/compile.py @@ -100,9 +100,7 @@ def build( f"conflicts with the Module Target {_target_summary(module_target)}" ) - workdir = os.path.join( - tempfile.gettempdir(), f"tilefoundry_build_{mod.entry}_{os.getpid()}_split" - ) + workdir = tempfile.mkdtemp(prefix=f"tilefoundry_build_{mod.entry}_split_") return _build_split_runtime_module(mod, workdir=workdir) diff --git a/src/tilefoundry/inspection/values.py b/src/tilefoundry/inspection/values.py index 26780e3be..79d8cc8bb 100644 --- a/src/tilefoundry/inspection/values.py +++ b/src/tilefoundry/inspection/values.py @@ -71,7 +71,7 @@ def print(self, value): for item in fields(value) ) if isinstance(value, Mapping): - return ENTRIES.join(f"{k}{ENTRY}{self.print(v)}" for k, v in value.items()) + return ";".join(f"{k}{ENTRY}{self.print(v)}" for k, v in value.items()) if isinstance(value, (tuple, list)): return ENTRIES.join(self.print(v) for v in value) return str(value) @@ -114,47 +114,34 @@ def _single(self, family, value, default=_UNSET): return family return f"{family}{FIELD}{self.print(value)}" - def print_ComputeCostMetadata(self, record, **_): - def spread_values(spread): - logical = dict(spread.logical) - total = dict(spread.total) - per_unit = [dict(value) for value in spread.per_unit] - names = sorted({*logical, *total, *(name for value in per_unit for name in value)}) - return { - name: ( - f"{logical.get(name, 0)}{PAIR}{total.get(name, 0)}" - + "".join( - f"{PER_UNIT}{level}{ENTRY}{value.get(name, 0)}" - for level, value in zip(record.topologies, per_unit, strict=False) - ) - ) - for name in names - } + def _spread(self, spread, topologies, *, logical=True): + values = [] + if logical: + values.append(f"{self.print(spread.logical)}@logical") + values.append(f"{self.print(spread.total)}@total") + values.extend( + f"{self.print(value)}@{topology}" + for topology, value in zip(topologies, spread.per_unit, strict=False) + ) + return self.print(values) + + def _breakdown(self, held, topologies, *, logical=True): + return { + kind: self._spread(spread, topologies, logical=logical) + for kind, spread in held.kinds + } + def print_ComputeCostMetadata(self, record, **_): return self._record( "compute-cost", ( - ("flops", spread_values(record.flops)), - ("other_ops", spread_values(record.other_ops)), + ("flops", self._breakdown(record.flops, record.topologies)), + ("other_ops", self._breakdown(record.other_ops, record.topologies)), ), ) def print_TrafficMetadata(self, record, *, opt_in=frozenset()): - traffic = { - level: ( - f"{self.print(spread.total)}" - + ( - f"{PER_UNIT}" - + ENTRIES.join( - f"{topology}{ENTRY}{self.print(value)}" - for topology, value in zip(record.topologies, spread.per_unit, strict=False) - ) - if spread.per_unit - else "" - ) - ) - for level, spread in record.storage.kinds - } + traffic = self._breakdown(record.storage, record.topologies, logical=False) values = [("traffic", traffic)] if "operands" in opt_in: last = len(record.operands) - 1 diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 6ec277e6e..2109a3e24 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -233,9 +233,11 @@ def test_a_symbolic_store_stride_preserves_the_literal_control_result() -> None: bound = get_metadata(result.function, RooflineMetadata) summary = get_metadata(result.function, PerformanceSummaryMetadata) assert cost is not None and bound is not None and summary is not None + integer_ops = cost.other_ops.of("integer") + assert integer_ops is not None observed[name] = ( - dict(cost.other_ops.total)["integer"], - dict(cost.other_ops.per_unit[0])["integer"], + integer_ops.total, + integer_ops.per_unit[0], bound.ideal_ns, summary.timeline.end_ns - summary.timeline.start_ns, ) @@ -321,9 +323,11 @@ def test_a_matmul_counts_its_rows_once_whichever_axis_the_mesh_split() -> None: ) cost = get_metadata(product, ComputeCostMetadata) summary = get_metadata(report.function, PerformanceSummaryMetadata) + bf16 = cost.flops.of("bf16") + assert bf16 is not None per_layout[name] = ( - dict(cost.flops.total)["bf16"], - dict(cost.flops.per_unit[0])["bf16"], + bf16.total, + bf16.per_unit[0], summary.timeline.end_ns, ) @@ -403,7 +407,8 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N work = next( record for record in records - if record is not None and any(v for _n, v in record.flops.per_unit[0]) + if record is not None + and any(spread.per_unit[0] for _name, spread in record.flops.kinds) ) with pytest.raises( @@ -415,7 +420,10 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N with pytest.raises(AnalysisError, match=r"unknown compute dtype 'f9e9m9'"): _local_duration_ns( - replace(work, flops=replace(work.flops, per_unit=((("f9e9m9", 8),),))), + replace( + work, + flops=Breakdown((*work.flops.kinds, ("f9e9m9", Spread(8, 8, (8,))))), + ), throughput, services, level="cta", diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index e849a1af2..e1a3a106e 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -203,10 +203,12 @@ def _every_number_counts_something(result: AnalysisResult) -> None: assert value >= 0, f"{describe_expr(expr)}: {field}[{name}] = {value}" if record is ComputeCostMetadata: for field in ("flops", "other_ops"): - spread = getattr(held, field) - for domain in (spread.logical, spread.total, *spread.per_unit): - for name, value in domain: - assert value >= 0, f"{describe_expr(expr)}: {field}[{name}] = {value}" + breakdown = getattr(held, field) + 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): diff --git a/tests/analysis/test_analyze_cross_module.py b/tests/analysis/test_analyze_cross_module.py index a9c937457..c9fd8c448 100644 --- a/tests/analysis/test_analyze_cross_module.py +++ b/tests/analysis/test_analyze_cross_module.py @@ -50,7 +50,8 @@ def _flops(records) -> dict[str, int]: total: dict[str, int] = {} for item in records: record = item[0] if isinstance(item, tuple) else item - for name, value in record.flops.total: + for name, spread in record.flops.kinds: + value = spread.total total[name] = total.get(name, 0) + value return total diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index d5c8c07a8..4dbf9d206 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -701,7 +701,7 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> assert cli.main(["analyze", selector, str(operands_path), *flags, "--operands"]) == 0 asked = operands_path.read_text(encoding="utf-8") assert "operands=" not in first - assert "operands=0:r30720/w0,result:r0/w30720" in asked + assert "operands=0:r30720/w0;result:r0/w30720" in asked assert ( cli.main(["analyze", selector, str(operands_json_path), *flags, "--operands", "--json"]) == 0 @@ -746,21 +746,23 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> f"# selection requested={','.join(payload['requested'])} " f"executed={','.join(payload['executed'])}", "# compute-cost " - f"flops=f32:{cost['flops']['logical']['f32']}" - f"/{cost['flops']['total']['f32']}" - f"@{payload['topology']}:{cost['flops']['per_unit'][0]['f32']}", + 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']}" - f"@{payload['topology']}:r{moved['storage']['gmem']['per_unit'][0]['read']}" - f"/w{moved['storage']['gmem']['per_unit'][0]['write']}", + 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']}", f"# roofline ideal-ns={bound['ideal_ns']} bound-by={bound['bound_by']}", "# performance root=MoEMegaKernel::experts " f"predicted-ns={summary['timeline']['end_ns']} " f"waves={summary['waves']}", ] - assert payload["totals"]["flops"] == cost["flops"]["total"] + assert payload["totals"]["flops"] == { + name: spread["total"] for name, spread in cost["flops"].items() + } assert payload["totals"]["traffic"] == { name: value["total"] for name, value in moved["storage"].items() } diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index ddaaf79a4..b6ec3c9eb 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@cta:r512/w0,thread:r64/w0,rmem:r0/w2048@cta:r0/w512,thread:r0/w64 - 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/512@cta:128@thread:16; traffic traffic=rmem:r2048/w2048@cta:r512/w512,thread:r64/w64 - 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@cta:r512/w0,thread:r64/w0,smem:r0/w2048@cta:r0/w512,thread:r0/w64 - 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/512@cta:128@thread:128; traffic traffic=smem:r2048/w1024@cta:r512/w256,thread:r512/w256 - 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@cta:r256/w256,thread:r256/w256 - 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@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 - 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@cta:r64/w0,thread:r8/w0,rmem:r0/w64@cta:r0/w64,thread:r0/w8 + 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/128@cta:128@thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 - 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/128@cta:128@thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 + 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 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@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 - 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@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 - v17 = reshard(v6, layout=Layout((16,), (1,)), storage=gmem) # Tensor[(16,), "f32", Layout((16,), (1,))]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 + 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 return (v5, v13, v15, v17) diff --git a/tests/ops/ir/test_cache_update.py b/tests/ops/ir/test_cache_update.py index 8e63aa502..6350afd50 100644 --- a/tests/ops/ir/test_cache_update.py +++ b/tests/ops/ir/test_cache_update.py @@ -230,7 +230,7 @@ def test_cache_update_function_analyzes_program_and_cta_cost() -> None: moved = get_metadata(analysed_update, TrafficMetadata) assert result.topology_level == "cta" assert record is not None - assert record.flops.total == () + assert record.flops.kinds == () row_bytes = _WINDOW_BYTES // 4 assert moved.operands == ( TrafficBytes(), diff --git a/tests/ops/ir/test_local.py b/tests/ops/ir/test_local.py index 28bd4f4ca..ce2347ea7 100644 --- a/tests/ops/ir/test_local.py +++ b/tests/ops/ir/test_local.py @@ -53,6 +53,6 @@ def test_local_analyzes_as_a_zero_traffic_topology_view() -> None: moved = get_metadata(analysed_local, TrafficMetadata) assert result.topology_level == topology_level assert record is not None - assert record.flops.total == () + assert record.flops.kinds == () assert moved.storage.kinds == () assert moved.operands == (TrafficBytes(), TrafficBytes()) From 71c87b1f4c6b95d8d1c71fc2851624a902f4d9dc Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 15 Sep 2026 15:55:16 +0800 Subject: [PATCH 7/7] test(analysis): read keyed spread JSON --- tests/installed/models/smoke_qwen3_1_7b.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/installed/models/smoke_qwen3_1_7b.py b/tests/installed/models/smoke_qwen3_1_7b.py index 51bf9ab81..bdc0fd5e0 100644 --- a/tests/installed/models/smoke_qwen3_1_7b.py +++ b/tests/installed/models/smoke_qwen3_1_7b.py @@ -232,8 +232,8 @@ def test_the_mlp_costs_its_three_matrices(tf, shipped_source) -> None: authored_flops = authored["totals"]["flops"][DT] stated = placed["function_records"]["compute-cost"]["flops"] topologies = placed["function_records"]["compute-cost"]["topologies"] - placed_flops = stated["total"][DT] - per_unit = dict(zip(topologies, (share[DT] for share in stated["per_unit"]))) + placed_flops = stated[DT]["total"] + per_unit = dict(zip(topologies, stated[DT]["per_unit"])) assert placed_flops == per_unit["cta"] * 128 _holds(authored_flops, _mlp_matmul_flops(), "mlp")