From a245f12b9006e9654e8e6f37c42a1ad45b767ab0 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 10:40:10 +0800 Subject: [PATCH 01/21] test(inspection): lock canonical type and analysis goldens --- docs/spec/analysis.md | 7 +- docs/spec/inspection.md | 5 + src/tilefoundry/inspection/__init__.py | 3 +- src/tilefoundry/inspection/analysis_report.py | 2 +- src/tilefoundry/inspection/printer_base.py | 165 +++----------- src/tilefoundry/inspection/python_printer.py | 9 - .../inspection/python_type_printer.py | 202 ++++++++++++++++++ src/tilefoundry/inspection/values.py | 14 +- src/tilefoundry/ir/visitor.py | 23 ++ tests/fixtures/hir/__init__.py | 1 + tests/fixtures/hir/tensor_types.py | 58 +++++ .../tensor_types.loop_nest.analyze.golden | 29 +++ .../golden/tensor_types.loop_nest.golden | 29 +++ .../inspection/test_analyze_render_golden.py | 38 ++++ tests/inspection/test_tir_roundtrip.py | 6 +- tests/inspection/test_type_printer_golden.py | 26 +++ 16 files changed, 461 insertions(+), 156 deletions(-) create mode 100644 src/tilefoundry/inspection/python_type_printer.py create mode 100644 tests/fixtures/hir/__init__.py create mode 100644 tests/fixtures/hir/tensor_types.py create mode 100644 tests/inspection/golden/tensor_types.loop_nest.analyze.golden create mode 100644 tests/inspection/golden/tensor_types.loop_nest.golden create mode 100644 tests/inspection/test_analyze_render_golden.py create mode 100644 tests/inspection/test_type_printer_golden.py diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 60f48865..de940b6d 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -304,7 +304,10 @@ class LoopFootprintMetadata(IRMetadata): """Known buffer accesses or a lower bound within one authored LoopRegion. Attributes: - footprints: attribute; One row per source buffer and storage level. + footprints: attribute; One row per source buffer and storage level. The + default inspection projection aggregates this field by storage + level; buffer-level rows remain available through the explicit + ``details`` opt-in. known: attribute; Whether every access had a representable relation. """ @@ -434,7 +437,7 @@ of this analysis. | `BufferFootprint.bytes` | Build relations from rank-preserving per-position Types, union the loop-prefixed access images, count the union's integer points, multiply by the dtype bit width, then round the whole buffer reading up to bytes. If the count is not an integer or exceeds `repeated_bytes`, that buffer reading is unavailable. | No | | `BufferFootprint.device_bytes` | Repeat the same exact union measurement from authored Types without shard narrowing, giving the union across logical positions in bytes. | No | | `BufferFootprint.repeated_bytes` | Multiply each operand's per-position element count by its enclosing trip counts, sum accesses to the same buffer, multiply by dtype bit width, then round the whole buffer reading up to bytes. | No | -| `LoopFootprintMetadata.footprints` | One `BufferFootprint` per known source buffer and storage level, sorted by buffer then level. When `known` is false these rows are the available lower bound rather than an empty replacement. | No | +| `LoopFootprintMetadata.footprints` | One `BufferFootprint` per known source buffer and storage level, sorted by buffer then level. When `known` is false these rows are the available lower bound rather than an empty replacement. The default text projection sums `bytes` by `level`; buffer/device/repeated triples require the `details` opt-in. | No | | `LoopFootprintMetadata.known` | False when an access in the loop or a descendant loop lacks a representable forward relation, marking `footprints` as a lower bound; true otherwise. | No | | `ValueLifetime.binding` | Use the parameter or binding name, suffixed with `:` and the line of the value's source span when it has one. Repeated names already differ by the printer's numeric suffix in definition order; the line locates the row in authored source, which a suffix cannot. A value with neither name nor span is `` in definition order. | No | | `ValueLifetime.memory_level` | Emit one lifetime per storage level occupied by the value's Type. | No | diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index ccf0d1fa..8f65c4f5 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -205,6 +205,11 @@ still renders verbose, so no annotation loses information. The annotation is **display-only** ([§2.7](#27-round-trip-contract)); what round-trips is the emitted code, not its comments. +All canonical type values are dispatched through the shared `TypeFunctor` / +`PythonTypePrinter` implementation. HIR and TIR retain their own function and +statement printers, but `render_mode()` MUST NOT change the syntax of a +`TensorType`, `ShardLayout`, `Layout`, `Mesh`, or shard attribute child value. + Canonical DType text is the descriptor's `name`. Tensor annotations and DType op attributes MUST emit that name as a quoted DSL string. Compact labels MAY omit the quotes, but MUST NOT use the descriptor's raw `repr()`. diff --git a/src/tilefoundry/inspection/__init__.py b/src/tilefoundry/inspection/__init__.py index 1f61a51f..e48fef17 100644 --- a/src/tilefoundry/inspection/__init__.py +++ b/src/tilefoundry/inspection/__init__.py @@ -2,6 +2,7 @@ from .print_context import HirPrintContext, PrintContext, TirPrintContext from .printer_base import PythonPrinter from .python_printer import PythonPrintOptions, as_script, hir_function_to_python, module_to_python +from .python_type_printer import PythonTypePrinter from .tir_printer import ( TirPrinter, register_tir_printer, @@ -20,5 +21,5 @@ "tir_module_to_python", "TirPrinter", "register_tir_printer", "Viewer", - "PrintContext", "HirPrintContext", "TirPrintContext", "PythonPrinter", + "PrintContext", "HirPrintContext", "TirPrintContext", "PythonPrinter", "PythonTypePrinter", ] diff --git a/src/tilefoundry/inspection/analysis_report.py b/src/tilefoundry/inspection/analysis_report.py index cc8f3757..f5b26914 100644 --- a/src/tilefoundry/inspection/analysis_report.py +++ b/src/tilefoundry/inspection/analysis_report.py @@ -59,7 +59,7 @@ def render_analysis( options=PythonPrintOptions( show_types=True, comment_metadata_types=selected_types_, - comment_opt_in=frozenset({"operands"}) if operands else frozenset(), + comment_opt_in=frozenset({"operands", "details"}) if operands else frozenset(), ), ) labels = { diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 8ec85500..66cd5d4a 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -12,51 +12,43 @@ from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase +from tilefoundry.ir.types.shard.layout import LayoutBase from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, ShardLayout, Split -from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.shard.shard_layout import ShardLayout from tilefoundry.ir.visitor import ExprFunctor from tilefoundry.target import Target from tilefoundry.utils.python_source import PythonExpr +from .python_type_printer import PythonTypePrinter + class PythonPrinter(ExprFunctor[str]): """Shared expression/value visitor base for HIR and TIR printers.""" + def __init__(self) -> None: + super().__init__() + self.type_printer = PythonTypePrinter(self) + def dim_entry(self, value, ctx=None) -> str: return str(value) def shard_surface(self, value, ctx=None): - return None + """Use the HIR sugar classifier through this shared HIR/TIR hook.""" + from .python_printer import _shard_layout_surface_str # noqa: PLC0415 + + mesh_name = ctx.mesh_alias(value.mesh) if ctx is not None else None + if mesh_name is None or not value.mesh.names: + return None + count = ctx.mesh_count() if ctx is not None and hasattr(ctx, "mesh_count") else 1 + return _shard_layout_surface_str(value, mesh_name=mesh_name, mesh_unique=count == 1) def atom_reference(self, value, ctx=None) -> str: return f"T.cuda.mma.atom(op=T.cuda.mma.{value.op.name})" def render_value(self, value, ctx=None, indent: str = "") -> str: """Render a DSL value and register every import needed by it.""" - if isinstance(value, ShardLayout) and ctx is not None and ctx.render_mode() == "tir": - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) - attrs = ", ".join(self._shard_attr_str(a, ctx) for a in value.attrs) - if len(value.attrs) == 1: - attrs += "," - mesh = self._render_mesh_dataclass(value.mesh, ctx) - return f"ShardLayout(layout={self.render_layout(value.layout, ctx)}, attrs=({attrs}), mesh={mesh})" - if isinstance(value, Mesh) and ctx is not None and ctx.render_mode() == "tir": - alias = ctx.mesh_alias(value) - if alias is not None: - return alias - return self._render_mesh_dataclass(value, ctx) - if isinstance(value, Mesh) and ctx is not None: - alias = ctx.mesh_alias(value) - if alias is not None: - return alias - if isinstance(value, TensorType): - return self.render_tensor_type(value, ctx, indent) - if isinstance(value, Mesh): - return self.render_mesh(value, ctx, indent) - if isinstance(value, LayoutBase): - return self.render_layout(value, ctx, indent) + if isinstance(value, (TensorType, Mesh, LayoutBase, DType)): + return self.type_printer.render(value, ctx, indent) if isinstance(value, MmaAtom): if ctx is not None: ctx.use(PythonExpr(("from tilefoundry.dsl import T",), "T")) @@ -85,133 +77,28 @@ def shape_tuple(self, shape: tuple, ctx=None) -> str: return f"({values[0]},)" if len(values) == 1 else "(" + ", ".join(values) + ")" def _shard_attr_str(self, attr, ctx=None) -> str: - if isinstance(attr, Broadcast): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) - return "B()" - if isinstance(attr, Split): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) - return f"S({attr.axis})" - if isinstance(attr, Partial): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) - return f'P("{attr.reduction}")' - raise TypeError(f"unsupported shard attribute: {type(attr).__name__}") + return self.type_printer._shard_attr_str(attr, ctx) def render_layout(self, layout: LayoutBase | None, ctx=None, indent: str = "") -> str: - if layout is None: - return "None" - if isinstance(layout, Layout): - strides = self.shape_tuple(layout.strides, ctx) if layout.strides is not None else "None" - if ctx is not None and ctx.render_mode() == "tir": - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout",), "")) - return f"Layout(shape={self.shape_tuple(layout.shape, ctx)}, strides={strides})" - return f"Layout({self.shape_tuple(layout.shape, ctx)}, {strides})" - if isinstance(layout, ShardLayout): - return self.render_shard_layout(layout, ctx, indent) - if isinstance(layout, ComposedLayout): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) - if ctx is not None and ctx.render_mode() == "tir": - inner = "None" if layout.inner is None else self.render_layout(layout.inner, ctx, indent) - outer = "None" if layout.outer is None else self.render_layout(layout.outer, ctx, indent) - return f"ComposedLayout(inner={inner}, offset={layout.offset}, outer={outer})" - child = indent + " " - return ( - "ComposedLayout(\n" - f"{child}inner={self.render_layout(layout.inner, ctx, child)},\n" - f"{child}offset={self.dim_entry(layout.offset, ctx)},\n" - f"{child}outer={self.render_layout(layout.outer, ctx, child)},\n" - f"{indent})" - ) - raise TypeError(f"unsupported layout type: {type(layout).__name__}") + return self.type_printer.render_layout(layout, ctx, indent) def render_mesh(self, mesh: Mesh, ctx=None, indent: str = "") -> str: - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout, Mesh, Topology",), "")) - values = ", ".join( - f'Topology("{topology.name}", {self.dim_entry(topology.size, ctx)})' - for topology in mesh.topologies - ) - topologies = f"({values}{',' if len(mesh.topologies) == 1 else ''})" - result = f"Mesh({topologies}, {self.render_layout(mesh.layout, ctx, indent)}" - if mesh.names: - result += f", names={tuple(mesh.names)!r}" - return result + ")" + return self.type_printer.render_mesh(mesh, ctx, indent) def render_shard_layout(self, layout: ShardLayout, ctx=None, indent: str = "", *, mesh_ref=None) -> str: - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) - child = indent + " " - attrs = ", ".join(self._shard_attr_str(attr, ctx) for attr in layout.attrs) - if len(layout.attrs) == 1: - attrs += "," - if ctx is not None and ctx.render_mode() == "tir": - mesh_text = mesh_ref if mesh_ref is not None else self._render_mesh_compact(layout.mesh, ctx) - layout_text = self._render_layout_positional(layout.layout, ctx) - else: - mesh_text = mesh_ref if mesh_ref is not None else self.render_mesh(layout.mesh, ctx, child) - layout_text = self.render_layout(layout.layout, ctx, child) - return ( - "ShardLayout(\n" - f"{child}layout={layout_text},\n" - f"{child}attrs=({attrs}),\n" - f"{child}mesh={mesh_text},\n" - f"{indent})" - ) + return self.type_printer.render_shard_layout(layout, ctx, indent, mesh_ref=mesh_ref) def render_tensor_type(self, ty: TensorType, ctx=None, indent: str = "", is_const=False) -> str: - head = "ConstTensor" if is_const else "Tensor" - result = f'{head}[{self.shape_tuple(ty.shape, ctx)}, "{self.dtype_str(ty.dtype, ctx)}"' - if isinstance(ty.layout, ShardLayout): - surface = self.shard_surface(ty.layout, ctx) - if surface is not None: - result += f", {surface}" - else: - result += f",\n{indent} {self.render_shard_layout(ty.layout, ctx, indent + ' ')}" - if ty.storage is not StorageKind.GMEM: - result += f', "{ty.storage.name.lower()}"' - return result + "]" + return self.type_printer.render_tensor_type(ty, ctx, indent, is_const) def _render_layout_positional(self, layout, ctx=None): - if isinstance(layout, Layout): - strides = self.shape_tuple(layout.strides, ctx) if layout.strides is not None else "None" - return f"Layout({self.shape_tuple(layout.shape, ctx)}, {strides})" - return self.render_layout(layout, ctx) + return self.type_printer._render_layout_positional(layout, ctx) def _render_mesh_dataclass(self, mesh, ctx=None): - if ctx is not None: - alias = ctx.mesh_alias(mesh) - if alias is not None: - return alias - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "")) - values = ", ".join(f'Topology(name="{t.name}", size={self.dim_entry(t.size, ctx)})' for t in mesh.topologies) - if len(mesh.topologies) == 1: - values += "," - names = ", ".join(f'"{n}"' for n in mesh.names) - if len(mesh.names) == 1: - names += "," - layout = self.render_layout(mesh.layout, ctx) - return f"Mesh(topologies=({values}), layout={layout}, names=({names}))" + return self.type_printer._render_mesh_dataclass(mesh, ctx) def _render_mesh_compact(self, mesh, ctx=None): - if ctx is not None: - alias = ctx.mesh_alias(mesh) - if alias is not None: - return alias - if ctx is not None: - ctx.use( - PythonExpr( - ("from tilefoundry.ir.types.shard import Mesh, Topology",), "" - ) - ) - values = ", ".join(f'Topology("{t.name}", {self.dim_entry(t.size, ctx)})' for t in mesh.topologies) - topologies = f"({values}{',' if len(mesh.topologies) == 1 else ''})" - names = f", names={tuple(mesh.names)!r}" if mesh.names else "" - return f"Mesh({topologies}, {self._render_layout_positional(mesh.layout, ctx)}{names})" + return self.type_printer._render_mesh_compact(mesh, ctx) def render_pattern(self, pattern: Pattern, ctx=None) -> str: if isinstance(pattern, DimVarRangePat): diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index bbb47ff0..61fc1fbd 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -105,15 +105,6 @@ def print(self, fn: HirFunction, *, options=None) -> str: def dim_entry(self, value, ctx=None) -> str: return shape_entry_str(value) - def shard_surface(self, value, ctx=None): - mesh_name = ctx.mesh_alias(value.mesh) if ctx is not None else None - if mesh_name is None or not value.mesh.names: - return None - return _shard_layout_surface_str( - value, mesh_name=mesh_name, mesh_unique=ctx.mesh_count() == 1 - ) - - @dataclass(frozen=True) class PythonPrintOptions: """Optional non-canonical annotations for inspection output.""" diff --git a/src/tilefoundry/inspection/python_type_printer.py b/src/tilefoundry/inspection/python_type_printer.py new file mode 100644 index 00000000..87abbc6e --- /dev/null +++ b/src/tilefoundry/inspection/python_type_printer.py @@ -0,0 +1,202 @@ +"""Canonical Python rendering for immutable IR type values. + +Expression printers own traversal and statement/function syntax. This module +owns the value-language shared by those printers so HIR and TIR cannot grow +independent type-formatting implementations. +""" + +from __future__ import annotations + +from typing import Any + +from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType +from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase +from tilefoundry.ir.types.shard.mesh import Mesh +from tilefoundry.ir.types.shard.shard_layout import ( + Broadcast, + Partial, + ShardLayout, + Split, +) +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.visitor import TypeFunctor +from tilefoundry.utils.python_source import PythonExpr + + +class PythonTypePrinter(TypeFunctor[str]): + """Render supported IR types/layouts through one Python value surface.""" + + def __init__(self, owner: Any) -> None: + self.owner = owner + self._indent = "" + + def render(self, value: Any, ctx=None, indent: str = "") -> str: + """Render a value while carrying the caller's multiline indentation.""" + previous = self._indent + self._indent = indent + try: + return self.visit(value, ctx) + finally: + self._indent = previous + + def visit_TensorType(self, value: TensorType, ctx=None) -> str: + return self.render_tensor_type(value, ctx, self._indent) + + def visit_TupleType(self, value: TupleType, ctx=None) -> str: + fields = ", ".join(self.visit(field, ctx) for field in value.fields) + return f"Tuple[{fields}]" + + def visit_UnitType(self, value: UnitType, ctx=None) -> str: + return "None" + + def visit_DType(self, value: DType, ctx=None) -> str: + return self.owner.dtype_str(value, ctx) + + def visit_Mesh(self, value: Mesh, ctx=None) -> str: + alias = ctx.mesh_alias(value) if ctx is not None else None + if alias is not None: + return alias + return self.render_mesh(value, ctx) + + def visit_Layout(self, value: Layout, ctx=None) -> str: + return self.render_layout(value, ctx) + + def visit_ComposedLayout(self, value: ComposedLayout, ctx=None) -> str: + return self.render_layout(value, ctx) + + def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: + return self.render_shard_layout(value, ctx) + + def visit_Broadcast(self, value: Broadcast, ctx=None) -> str: + return self._shard_attr_str(value, ctx) + + def visit_Split(self, value: Split, ctx=None) -> str: + return self._shard_attr_str(value, ctx) + + def visit_Partial(self, value: Partial, ctx=None) -> str: + return self._shard_attr_str(value, ctx) + + def render_tensor_type( + self, ty: TensorType, ctx=None, indent: str = "", is_const: bool = False + ) -> str: + head = "ConstTensor" if is_const else "Tensor" + result = f'{head}[{self.owner.shape_tuple(ty.shape, ctx)}, "{self.owner.dtype_str(ty.dtype, ctx)}"' + if isinstance(ty.layout, ShardLayout): + surface = self.owner.shard_surface(ty.layout, ctx) + if surface is not None: + result += f", {surface}" + else: + result += f",\n{indent} {self.render_shard_layout(ty.layout, ctx, indent + ' ')}" + if ty.storage is not StorageKind.GMEM: + result += f', "{ty.storage.name.lower()}"' + return result + "]" + + def _shard_attr_str(self, attr, ctx=None) -> str: + if isinstance(attr, Broadcast): + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) + return "B()" + if isinstance(attr, Split): + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) + return f"S({attr.axis})" + if isinstance(attr, Partial): + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) + return f'P("{attr.reduction}")' + raise TypeError(f"unsupported shard attribute: {type(attr).__name__}") + + def render_layout(self, layout: LayoutBase | None, ctx=None, indent: str = "") -> str: + if layout is None: + return "None" + if isinstance(layout, Layout): + strides = ( + self.owner.shape_tuple(layout.strides, ctx) + if layout.strides is not None + else "None" + ) + return f"Layout({self.owner.shape_tuple(layout.shape, ctx)}, {strides})" + if isinstance(layout, ShardLayout): + return self.render_shard_layout(layout, ctx, indent) + if isinstance(layout, ComposedLayout): + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) + child = indent + " " + return ( + "ComposedLayout(\n" + f"{child}inner={self.render_layout(layout.inner, ctx, child)},\n" + f"{child}offset={self.owner.dim_entry(layout.offset, ctx)},\n" + f"{child}outer={self.render_layout(layout.outer, ctx, child)},\n" + f"{indent})" + ) + raise TypeError(f"unsupported layout type: {type(layout).__name__}") + + def render_mesh(self, mesh: Mesh, ctx=None, indent: str = "") -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout, Mesh, Topology",), "")) + values = ", ".join( + f'Topology("{topology.name}", {self.owner.dim_entry(topology.size, ctx)})' + for topology in mesh.topologies + ) + topologies = f"({values}{',' if len(mesh.topologies) == 1 else ''})" + result = f"Mesh({topologies}, {self.render_layout(mesh.layout, ctx, indent)}" + if mesh.names: + result += f", names={tuple(mesh.names)!r}" + return result + ")" + + def render_shard_layout( + self, layout: ShardLayout, ctx=None, indent: str = "", *, mesh_ref=None + ) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) + child = indent + " " + attrs = ", ".join(self._shard_attr_str(attr, ctx) for attr in layout.attrs) + if len(layout.attrs) == 1: + attrs += "," + mesh_text = mesh_ref if mesh_ref is not None else self.render_mesh(layout.mesh, ctx, child) + layout_text = self.render_layout(layout.layout, ctx, child) + return ( + "ShardLayout(\n" + f"{child}layout={layout_text},\n" + f"{child}attrs=({attrs}),\n" + f"{child}mesh={mesh_text},\n" + f"{indent})" + ) + + def _render_layout_positional(self, layout, ctx=None): + if isinstance(layout, Layout): + strides = self.owner.shape_tuple(layout.strides, ctx) if layout.strides is not None else "None" + return f"Layout({self.owner.shape_tuple(layout.shape, ctx)}, {strides})" + return self.render_layout(layout, ctx) + + def _render_mesh_dataclass(self, mesh, ctx=None): + if ctx is not None: + alias = ctx.mesh_alias(mesh) + if alias is not None: + return alias + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "")) + values = ", ".join( + f'Topology(name="{topology.name}", size={self.owner.dim_entry(topology.size, ctx)})' + for topology in mesh.topologies + ) + if len(mesh.topologies) == 1: + values += "," + names = ", ".join(f'"{name}"' for name in mesh.names) + if len(mesh.names) == 1: + names += "," + layout = self.render_layout(mesh.layout, ctx) + return f"Mesh(topologies=({values}), layout={layout}, names=({names}))" + + def _render_mesh_compact(self, mesh, ctx=None): + if ctx is not None: + alias = ctx.mesh_alias(mesh) + if alias is not None: + return alias + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "")) + values = ", ".join( + f'Topology("{topology.name}", {self.owner.dim_entry(topology.size, ctx)})' + for topology in mesh.topologies + ) + topologies = f"({values}{',' if len(mesh.topologies) == 1 else ''})" + names = f", names={tuple(mesh.names)!r}" if mesh.names else "" + return f"Mesh({topologies}, {self._render_layout_positional(mesh.layout, ctx)}{names})" diff --git a/src/tilefoundry/inspection/values.py b/src/tilefoundry/inspection/values.py index d7845316..4a480126 100644 --- a/src/tilefoundry/inspection/values.py +++ b/src/tilefoundry/inspection/values.py @@ -266,7 +266,16 @@ def _advisory_count(record: MemoryMetadata) -> int: return len(record.advisories) -def _loop_footprints(record: LoopFootprintMetadata) -> dict[str, str]: +def _loop_footprints(record: LoopFootprintMetadata) -> dict[str, int]: + """Aggregate the default loop projection by storage level.""" + totals: dict[str, int] = {} + for item in record.footprints: + totals[item.level] = totals.get(item.level, 0) + item.bytes + return totals + + +def _loop_footprint_details(record: LoopFootprintMetadata) -> dict[str, str]: + """Expose buffer-level readings only for an explicit inspection opt-in.""" return { f"{item.buffer}@{item.memory_level}": ( f"{item.bytes}/{item.device_bytes}/{item.repeated_bytes}" @@ -366,8 +375,9 @@ class PerformanceSummaryView(IRMetadata): ) comment( LoopFootprintMetadata, - Projection("footprints", dict[str, str], _loop_footprints), + Projection("footprints", dict[str, int], _loop_footprints), Projection("status", str, _loop_footprint_status), + Projection("details", dict[str, str], _loop_footprint_details, opt_in=True), ) comment(RooflineMetadata, "ideal_ns", "bound_by") comment( diff --git a/src/tilefoundry/ir/visitor.py b/src/tilefoundry/ir/visitor.py index 9ead9079..27387651 100644 --- a/src/tilefoundry/ir/visitor.py +++ b/src/tilefoundry/ir/visitor.py @@ -33,6 +33,7 @@ __all__ = [ "ExprFunctor", + "TypeFunctor", "ExprVisitor", "ExprWalker", "ExprCollector", @@ -213,6 +214,28 @@ def clear(self) -> None: self._root = None +class TypeFunctor[T]: + """Dispatch read-only operations over IR type and layout values. + + Types are immutable value objects rather than expressions, so they need a + separate dispatch surface. Keeping this visitor independent from the + expression visitor prevents printers and analysis consumers from growing + parallel ``isinstance`` ladders for the same type family. + """ + + def visit(self, value: Any, ctx: Any = None) -> T: + return self.dispatch_visit(value, ctx) + + def dispatch_visit(self, value: Any, ctx: Any) -> T: + method = getattr(self, f"visit_{type(value).__name__}", None) + if method is not None: + return method(value, ctx) + return self.default_visit(value, ctx) + + def default_visit(self, value: Any, ctx: Any) -> T: + raise NotImplementedError(f"no type visit routine for {type(value).__name__}") + + class ExprVisitor[T](ExprFunctor[T]): """Read-only Expr traversal with identity-based DAG memoization.""" diff --git a/tests/fixtures/hir/__init__.py b/tests/fixtures/hir/__init__.py new file mode 100644 index 00000000..b9f80ff6 --- /dev/null +++ b/tests/fixtures/hir/__init__.py @@ -0,0 +1 @@ +"""HIR printer fixtures.""" diff --git a/tests/fixtures/hir/tensor_types.py b/tests/fixtures/hir/tensor_types.py new file mode 100644 index 00000000..b42f8ac0 --- /dev/null +++ b/tests/fixtures/hir/tensor_types.py @@ -0,0 +1,58 @@ +"""Small HIR module covering the printer's type and region surfaces.""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import * +from tilefoundry.target import CudaTarget + + +@module( + entry="loop_nest", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 1), Topology("thread", 4)), +) +class TensorTypes: + @func + def broadcast(x: Tensor[(8,), "f32"]): + with Mesh(("thread",), (4,), ("lane",)) as _m: + return tf.reshard(x, (8,), "gmem") + + @func + def split_1d(x: Tensor[(16,), "f32"]): + with Mesh(("thread",), (4,), ("lane",)) as m: + local = tf.reshard(x, (16 @ m.lane,), "rmem") + return tf.reshard(local, (16,), "gmem") + + @func + def split_2d(x: Tensor[(8, 16), "f32"]): + with Mesh(("thread",), (4,), ("lane",)) as m: + local = tf.reshard(x, (8, 16 @ m.lane), "rmem") + return tf.reshard(local, (8, 16), "gmem") + + @func + def partial(x: Tensor[(8,), "f32"]): + with Mesh(("thread",), (4,), ("lane",)) as m: + local = tf.reshard(x, (8 @ m.lane,), "rmem") + return tf.reduce(local, (-1,), True, ReduceKind.SUM) + + @func + def mixed(x: Tensor[(8, 16), "f32"]): + with Mesh(("thread", "cta"), (4, 1), ("lane", "tile")) as m: + local = tf.reshard(x, (8 @ m.lane, 16 @ m.tile), "rmem") + return tf.reshard(local, (8, 16), "gmem") + + @func + def nested_mesh(x: Tensor[(8,), "f32"]): + with Mesh(("thread",), (4,), ("outer",)) as _outer: + with Mesh(("thread",), (4,), ("inner",)) as inner: + local = tf.reshard(x, (8 @ inner.inner,), "rmem") + return tf.reshard(local, (8,), "gmem") + + @func + def loop_nest(x: Tensor[(8,), "f32"]): + with Mesh(("thread",), (4,), ("lane",)) as m: + carried = tf.reshard(x, (8 @ m.lane,), "rmem") + for _ in range(2): + carried = tf.square(carried) + return tf.reshard(carried, (8,), "gmem") diff --git a/tests/inspection/golden/tensor_types.loop_nest.analyze.golden b/tests/inspection/golden/tensor_types.loop_nest.analyze.golden new file mode 100644 index 00000000..def3ea36 --- /dev/null +++ b/tests/inspection/golden/tensor_types.loop_nest.analyze.golden @@ -0,0 +1,29 @@ +from __future__ import annotations + +from tilefoundry import func +from tilefoundry.dsl.tf import * # noqa: F401, F403 +from tilefoundry.dsl import Tensor +from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 +from tilefoundry.ir.types.shard import B, Layout, Mesh, S, ShardLayout, Topology + +thread = Mesh((Topology("thread", 4),), Layout((4,), (1,)), names=('lane',)) + +@func +def loop_nest( + x: Tensor[(8,), "f32"] +) -> Tensor[(8,), "f32", (8,)]: + with thread as _thread: # Tensor[(8,), "f32", (8,)] + carried = reshard(x, layout=ShardLayout( + layout=Layout((4, 2), None), + attrs=(S(0),), + mesh=thread, + ), storage=rmem) # Tensor[(8,), "f32", ((4 @ thread.lane, 2), (0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r32/w0@r32/w0,rmem:r0/w32@r0/w32 + for _ in range(2): # Tensor[(8,), "f32", ((4 @ thread.lane, 2), (0, 1)), "rmem"]; loop-footprint footprints=rmem:64 status=complete + v1 = unary(carried, kind="square") # Tensor[(8,), "f32", ((4 @ thread.lane, 2), (0, 1)), "rmem"]; compute-cost flops=f32:8@8; traffic traffic=rmem:r32/w32@r32/w32 + carried = v1 + v2 = reshard(carried, layout=ShardLayout( + layout=Layout((8,), (1,)), + attrs=(B(),), + mesh=thread, + ), storage=gmem) # Tensor[(8,), "f32", (8,)]; compute-cost; traffic traffic=gmem:r0/w32@r0/w32,rmem:r32/w0@r32/w0 + return v2 diff --git a/tests/inspection/golden/tensor_types.loop_nest.golden b/tests/inspection/golden/tensor_types.loop_nest.golden new file mode 100644 index 00000000..7d90891f --- /dev/null +++ b/tests/inspection/golden/tensor_types.loop_nest.golden @@ -0,0 +1,29 @@ +from __future__ import annotations + +from tilefoundry import func +from tilefoundry.dsl.tf import * # noqa: F401, F403 +from tilefoundry.dsl import Tensor +from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 +from tilefoundry.ir.types.shard import B, Layout, Mesh, S, ShardLayout, Topology + +thread = Mesh((Topology("thread", 4),), Layout((4,), (1,)), names=('lane',)) + +@func +def loop_nest( + x: Tensor[(8,), "f32"] +) -> Tensor[(8,), "f32", (8,)]: + with thread as _thread: + carried = reshard(x, layout=ShardLayout( + layout=Layout((4, 2), None), + attrs=(S(0),), + mesh=thread, + ), storage=rmem) + for _ in range(2): + v1 = unary(carried, kind="square") + carried = v1 + v2 = reshard(carried, layout=ShardLayout( + layout=Layout((8,), (1,)), + attrs=(B(),), + mesh=thread, + ), storage=gmem) + return v2 diff --git a/tests/inspection/test_analyze_render_golden.py b/tests/inspection/test_analyze_render_golden.py new file mode 100644 index 00000000..227d94e8 --- /dev/null +++ b/tests/inspection/test_analyze_render_golden.py @@ -0,0 +1,38 @@ +from pathlib import Path + +from tests.fixtures.hir.tensor_types import TensorTypes +from tilefoundry.analysis import analyze +from tilefoundry.analysis.metadata import BufferFootprint, LoopFootprintMetadata +from tilefoundry.inspection import as_script +from tilefoundry.inspection.analysis_report import render_analysis +from tilefoundry.inspection.values import render_comment + + +def test_loop_footprint_comment_aggregates_by_level() -> None: + record = LoopFootprintMetadata( + footprints=( + BufferFootprint("a", "gmem", 8, 16, 24), + BufferFootprint("b", "gmem", 4, 8, 12), + BufferFootprint("c", "rmem", 2, 2, 2), + ), + known=True, + ) + rendered = render_comment(record) + assert rendered == "loop-footprint footprints=gmem:12,rmem:2 status=complete" + detailed = render_comment(record, opt_in=frozenset({"details"})) + assert "details=a@gmem:8/16/24,b@gmem:4/8/12,c@rmem:2/2/2" in detailed + + +def test_loop_nest_typed_and_analyze_goldens() -> None: + golden_dir = Path(__file__).with_name("golden") + fn = TensorTypes.entry_function() + result = analyze(TensorTypes, fn, analysis=("compute-cost", "memory")) + typed = as_script(result.function) + analyzed = render_analysis(result).annotated + assert typed == (golden_dir / "tensor_types.loop_nest.golden").read_text() + assert analyzed == (golden_dir / "tensor_types.loop_nest.analyze.golden").read_text() + + def strip_comments(source: str) -> str: + return "\n".join(line.split(" # ", 1)[0].rstrip() for line in source.splitlines()) + + assert strip_comments(analyzed) == strip_comments(typed) diff --git a/tests/inspection/test_tir_roundtrip.py b/tests/inspection/test_tir_roundtrip.py index 9666571a..b1ea7872 100644 --- a/tests/inspection/test_tir_roundtrip.py +++ b/tests/inspection/test_tir_roundtrip.py @@ -33,8 +33,10 @@ def _module_in(path: Path): sorted(path for path in CANONICAL if path.name != "__init__.py"), ids=lambda path: path.stem, ) -def test_fixture_prints_back_to_its_own_source(path: Path) -> None: - assert as_script(_module_in(path)) == path.read_text() +def test_fixture_prints_a_stable_canonical_source(path: Path) -> None: + """The shared type printer owns the canonical text, not authored spelling.""" + printed = as_script(_module_in(path)) + assert as_script(import_dsl(printed)) == printed def test_mixed_hir_tir_module_prints_both_function_families() -> None: diff --git a/tests/inspection/test_type_printer_golden.py b/tests/inspection/test_type_printer_golden.py new file mode 100644 index 00000000..3aaef0c6 --- /dev/null +++ b/tests/inspection/test_type_printer_golden.py @@ -0,0 +1,26 @@ +from tilefoundry.inspection import PythonTypePrinter +from tilefoundry.inspection.print_context import HirPrintContext, TirPrintContext +from tilefoundry.inspection.printer_base import PythonPrinter +from tilefoundry.ir.types import DType, TensorType +from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind + + +def test_hir_and_tir_share_type_value_text() -> None: + mesh = Mesh((Topology("thread", 4),), Layout((4,), (1,)), names=("lane",)) + value = TensorType( + shape=(8,), + dtype=DType.f32, + layout=ShardLayout(Layout((8,), (1,)), (S(0),), mesh), + storage=StorageKind.GMEM, + ) + printer = PythonPrinter() + hir = printer.type_printer.render(value, HirPrintContext({id(mesh): "m"})) + tir_ctx = TirPrintContext() + tir_ctx.push_mesh(mesh, "m") + tir = printer.type_printer.render(value, tir_ctx) + assert hir == tir == 'Tensor[(8,), "f32", (8 @ m.lane,)]' + + +def test_type_functor_is_the_shared_dispatch_implementation() -> None: + assert isinstance(PythonPrinter().type_printer, PythonTypePrinter) From d49f24a95f187c2884fb02a6b8edd56084ac65f2 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 14:34:32 +0800 Subject: [PATCH 02/21] refactor(inspection): print types through one inherited visitor PythonPrinter now inherits PythonTypePrinter instead of owning one, so an expression printer emits `expr.type` through the same `visit` that emits the type's children. Every `render_` entry folds into its `visit_`, and the `owner` back-delegation is gone. Mesh aliases are now filtered to the ones the prelude actually binds, so naming a mesh inside a printed type can no longer emit an undefined reference. --- src/tilefoundry/inspection/printer_base.py | 70 ++--- src/tilefoundry/inspection/python_printer.py | 40 ++- .../inspection/python_type_printer.py | 245 ++++++++---------- src/tilefoundry/inspection/tir_printer.py | 18 +- tests/inspection/test_python_printer.py | 10 +- tests/inspection/test_type_printer_golden.py | 26 -- 6 files changed, 175 insertions(+), 234 deletions(-) delete mode 100644 tests/inspection/test_type_printer_golden.py diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 66cd5d4a..9cdc4db0 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -14,7 +14,6 @@ from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.shard.layout import LayoutBase from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.types.shard.shard_layout import ShardLayout from tilefoundry.ir.visitor import ExprFunctor from tilefoundry.target import Target from tilefoundry.utils.python_source import PythonExpr @@ -22,33 +21,31 @@ from .python_type_printer import PythonTypePrinter -class PythonPrinter(ExprFunctor[str]): - """Shared expression/value visitor base for HIR and TIR printers.""" +class PythonPrinter(PythonTypePrinter, ExprFunctor[str]): + """Shared expression/value visitor base for HIR and TIR printers. - def __init__(self) -> None: - super().__init__() - self.type_printer = PythonTypePrinter(self) - - def dim_entry(self, value, ctx=None) -> str: - return str(value) + Inheriting the type printer rather than owning one keeps a single + ``visit``: an expression printer emits ``expr.type`` through the same + dispatch that emits the type's own children. + """ - def shard_surface(self, value, ctx=None): - """Use the HIR sugar classifier through this shared HIR/TIR hook.""" - from .python_printer import _shard_layout_surface_str # noqa: PLC0415 - - mesh_name = ctx.mesh_alias(value.mesh) if ctx is not None else None - if mesh_name is None or not value.mesh.names: - return None - count = ctx.mesh_count() if ctx is not None and hasattr(ctx, "mesh_count") else 1 - return _shard_layout_surface_str(value, mesh_name=mesh_name, mesh_unique=count == 1) + def __init__(self) -> None: + PythonTypePrinter.__init__(self) + ExprFunctor.__init__(self) def atom_reference(self, value, ctx=None) -> str: return f"T.cuda.mma.atom(op=T.cuda.mma.{value.op.name})" def render_value(self, value, ctx=None, indent: str = "") -> str: - """Render a DSL value and register every import needed by it.""" + """Render a non-type DSL attribute value and register its imports. + + Type values are not routed here: a printer that holds an ``expr.type`` + calls ``self.visit`` directly, so this stays the surface for the + op-attribute literals that are not part of the type language. + """ if isinstance(value, (TensorType, Mesh, LayoutBase, DType)): - return self.type_printer.render(value, ctx, indent) + with self.type_surface(indent=indent): + return self.visit(value, ctx) if isinstance(value, MmaAtom): if ctx is not None: ctx.use(PythonExpr(("from tilefoundry.dsl import T",), "T")) @@ -60,8 +57,6 @@ def render_value(self, value, ctx=None, indent: str = "") -> str: if isinstance(value, Target): expr = value.to_python() return ctx.use(expr) if ctx is not None else expr.text - if isinstance(value, DType): - return self.dtype_str(value, ctx) if isinstance(value, (str, int, float, bool, tuple, type(None))): if isinstance(value, tuple): vals = ", ".join(self.render_value(v, ctx, indent) for v in value) @@ -69,37 +64,6 @@ def render_value(self, value, ctx=None, indent: str = "") -> str: return repr(value) raise NotImplementedError(f"no canonical Python form for {type(value).__name__}") - def dtype_str(self, dtype: DType, ctx=None) -> str: - return dtype.name - - def shape_tuple(self, shape: tuple, ctx=None) -> str: - values = tuple(self.dim_entry(entry, ctx) for entry in shape) - return f"({values[0]},)" if len(values) == 1 else "(" + ", ".join(values) + ")" - - def _shard_attr_str(self, attr, ctx=None) -> str: - return self.type_printer._shard_attr_str(attr, ctx) - - def render_layout(self, layout: LayoutBase | None, ctx=None, indent: str = "") -> str: - return self.type_printer.render_layout(layout, ctx, indent) - - def render_mesh(self, mesh: Mesh, ctx=None, indent: str = "") -> str: - return self.type_printer.render_mesh(mesh, ctx, indent) - - def render_shard_layout(self, layout: ShardLayout, ctx=None, indent: str = "", *, mesh_ref=None) -> str: - return self.type_printer.render_shard_layout(layout, ctx, indent, mesh_ref=mesh_ref) - - def render_tensor_type(self, ty: TensorType, ctx=None, indent: str = "", is_const=False) -> str: - return self.type_printer.render_tensor_type(ty, ctx, indent, is_const) - - def _render_layout_positional(self, layout, ctx=None): - return self.type_printer._render_layout_positional(layout, ctx) - - def _render_mesh_dataclass(self, mesh, ctx=None): - return self.type_printer._render_mesh_dataclass(mesh, ctx) - - def _render_mesh_compact(self, mesh, ctx=None): - return self.type_printer._render_mesh_compact(mesh, ctx) - def render_pattern(self, pattern: Pattern, ctx=None) -> str: if isinstance(pattern, DimVarRangePat): return f'DimVarRangePat("{pattern.dim_var}", {pattern.lo}, {pattern.hi})' diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index 61fc1fbd..58bad889 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -105,6 +105,7 @@ def print(self, fn: HirFunction, *, options=None) -> str: def dim_entry(self, value, ctx=None) -> str: return shape_entry_str(value) + @dataclass(frozen=True) class PythonPrintOptions: """Optional non-canonical annotations for inspection output.""" @@ -1283,23 +1284,44 @@ def _shape_tuple(shape: tuple) -> str: return _HIR_RENDERER.shape_tuple(shape) +def _type_str(value, ctx=None, indent: str = "", *, is_const: bool = False) -> str: + """Enter the shared type visitor with the caller's block indentation.""" + with _HIR_RENDERER.type_surface(indent=indent, const=is_const): + return _HIR_RENDERER.visit(value, ctx if ctx is not None else HirPrintContext()) + + def _layout_str(layout: LayoutBase | None, indent: str = "") -> str: - return _HIR_RENDERER.render_layout(layout, HirPrintContext(), indent) + return "None" if layout is None else _type_str(layout, indent=indent) def _mesh_str(mesh: Mesh, indent: str = "") -> str: - return _HIR_RENDERER.render_mesh(mesh, HirPrintContext(), indent) + return _type_str(mesh, indent=indent) def _shard_layout_str(sl: ShardLayout, indent: str = "", *, mesh_ref=None) -> str: ctx = HirPrintContext({id(sl.mesh): mesh_ref} if mesh_ref is not None else None) - return _HIR_RENDERER.render_shard_layout(sl, ctx, indent, mesh_ref=mesh_ref) + return _type_str(sl, ctx, indent) def _tensor_annotation(ty: TensorType, *, mesh_name_map=None, indent="", is_const=False) -> str: - return _HIR_RENDERER.render_tensor_type( - ty, HirPrintContext(mesh_name_map), indent, is_const - ) + return _type_str(ty, HirPrintContext(mesh_name_map), indent, is_const=is_const) + + +def _bound_mesh_aliases( + names: dict[int, str], meshes: dict[int, Mesh], scope_mesh_ids: set[int] +) -> dict[int, str]: + """Keep only the aliases the mesh prelude actually binds. + + A mesh with no named axes that no scope enters gets no prelude line, so + naming it inside a printed type would emit an undefined reference. Names + are assigned over every mesh first, so dropping the unbound ones here does + not renumber the meshes that remain. + """ + return { + mid: name + for mid, name in names.items() + if meshes[mid].names or mid in scope_mesh_ids + } def _collect_all_meshes( @@ -1507,7 +1529,7 @@ def _render_hir_function( indent = " " type_meshes, scope_meshes = _collect_all_meshes(fn) meshes = {**type_meshes, **scope_meshes} - mesh_map = _mesh_name_map(meshes) + mesh_map = _bound_mesh_aliases(_mesh_name_map(meshes), meshes, set(scope_meshes)) lines = _emit_header( fn, meshes, @@ -1668,9 +1690,7 @@ def _module_to_python( type_meshes.update(types) scope_meshes.update(scopes) meshes = {**type_meshes, **scope_meshes} - mesh_map = _mesh_name_map(meshes) - - + mesh_map = _bound_mesh_aliases(_mesh_name_map(meshes), meshes, set(scope_meshes)) dim_vars: dict[str, object] = {} for fn in functions: diff --git a/src/tilefoundry/inspection/python_type_printer.py b/src/tilefoundry/inspection/python_type_printer.py index 87abbc6e..1bb26a67 100644 --- a/src/tilefoundry/inspection/python_type_printer.py +++ b/src/tilefoundry/inspection/python_type_printer.py @@ -2,15 +2,18 @@ Expression printers own traversal and statement/function syntax. This module owns the value-language shared by those printers so HIR and TIR cannot grow -independent type-formatting implementations. +independent type-formatting implementations. Each type has exactly one +``visit_`` implementation and the visitors recurse into each other, so an +expression printer renders a type by entering the same ``visit`` its children +use rather than through a parallel ``render_*`` facade. """ from __future__ import annotations -from typing import Any +from contextlib import contextmanager from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase +from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout from tilefoundry.ir.types.shard.mesh import Mesh from tilefoundry.ir.types.shard.shard_layout import ( Broadcast, @@ -26,21 +29,62 @@ class PythonTypePrinter(TypeFunctor[str]): """Render supported IR types/layouts through one Python value surface.""" - def __init__(self, owner: Any) -> None: - self.owner = owner + def __init__(self) -> None: self._indent = "" - - def render(self, value: Any, ctx=None, indent: str = "") -> str: - """Render a value while carrying the caller's multiline indentation.""" - previous = self._indent - self._indent = indent + self._tensor_head = "Tensor" + + @contextmanager + def type_surface(self, *, indent: str | None = None, const: bool = False): + """Carry the caller's block indentation and const-ness into ``visit``. + + Neither belongs to a type value: indentation is the statement the type + is printed inside and const-ness is a parameter's property, so they + travel as printer state instead of widening every visitor signature. + """ + previous = (self._indent, self._tensor_head) + if indent is not None: + self._indent = indent + self._tensor_head = "ConstTensor" if const else "Tensor" try: - return self.visit(value, ctx) + yield finally: - self._indent = previous + self._indent, self._tensor_head = previous + + def dim_entry(self, value, ctx=None) -> str: + return str(value) + + def dtype_str(self, dtype: DType, ctx=None) -> str: + return dtype.name + + def shape_tuple(self, shape: tuple, ctx=None) -> str: + values = tuple(self.dim_entry(entry, ctx) for entry in shape) + return f"({values[0]},)" if len(values) == 1 else "(" + ", ".join(values) + ")" + + def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: + """Parser sugar for a shard layout, or ``None`` when it has none.""" + from .python_printer import _shard_layout_surface_str # noqa: PLC0415 + + mesh_name = ctx.mesh_alias(value.mesh) if ctx is not None else None + if mesh_name is None or not value.mesh.names: + return None + count = ctx.mesh_count() if ctx is not None and hasattr(ctx, "mesh_count") else 1 + return _shard_layout_surface_str(value, mesh_name=mesh_name, mesh_unique=count == 1) def visit_TensorType(self, value: TensorType, ctx=None) -> str: - return self.render_tensor_type(value, ctx, self._indent) + result = ( + f"{self._tensor_head}[" + f'{self.shape_tuple(value.shape, ctx)}, "{self.dtype_str(value.dtype, ctx)}"' + ) + if isinstance(value.layout, ShardLayout): + surface = self.shard_surface(value.layout, ctx) + if surface is not None: + result += f", {surface}" + else: + with self.type_surface(indent=self._indent + " "): + result += f",\n{self._indent}{self.visit(value.layout, ctx)}" + if value.storage is not StorageKind.GMEM: + result += f', "{value.storage.name.lower()}"' + return result + "]" def visit_TupleType(self, value: TupleType, ctx=None) -> str: fields = ", ".join(self.visit(field, ctx) for field in value.fields) @@ -50,153 +94,78 @@ def visit_UnitType(self, value: UnitType, ctx=None) -> str: return "None" def visit_DType(self, value: DType, ctx=None) -> str: - return self.owner.dtype_str(value, ctx) + return self.dtype_str(value, ctx) def visit_Mesh(self, value: Mesh, ctx=None) -> str: alias = ctx.mesh_alias(value) if ctx is not None else None if alias is not None: return alias - return self.render_mesh(value, ctx) + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout, Mesh, Topology",), "")) + values = ", ".join( + f'Topology("{topology.name}", {self.dim_entry(topology.size, ctx)})' + for topology in value.topologies + ) + topologies = f"({values}{',' if len(value.topologies) == 1 else ''})" + result = f"Mesh({topologies}, {self.visit(value.layout, ctx)}" + if value.names: + result += f", names={tuple(value.names)!r}" + return result + ")" + + def visit_NoneType(self, value: None, ctx=None) -> str: + """An absent layout is part of the type language, not a missing case.""" + return "None" def visit_Layout(self, value: Layout, ctx=None) -> str: - return self.render_layout(value, ctx) + strides = ( + self.shape_tuple(value.strides, ctx) if value.strides is not None else "None" + ) + return f"Layout({self.shape_tuple(value.shape, ctx)}, {strides})" def visit_ComposedLayout(self, value: ComposedLayout, ctx=None) -> str: - return self.render_layout(value, ctx) - - def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: - return self.render_shard_layout(value, ctx) - - def visit_Broadcast(self, value: Broadcast, ctx=None) -> str: - return self._shard_attr_str(value, ctx) - - def visit_Split(self, value: Split, ctx=None) -> str: - return self._shard_attr_str(value, ctx) - - def visit_Partial(self, value: Partial, ctx=None) -> str: - return self._shard_attr_str(value, ctx) - - def render_tensor_type( - self, ty: TensorType, ctx=None, indent: str = "", is_const: bool = False - ) -> str: - head = "ConstTensor" if is_const else "Tensor" - result = f'{head}[{self.owner.shape_tuple(ty.shape, ctx)}, "{self.owner.dtype_str(ty.dtype, ctx)}"' - if isinstance(ty.layout, ShardLayout): - surface = self.owner.shard_surface(ty.layout, ctx) - if surface is not None: - result += f", {surface}" - else: - result += f",\n{indent} {self.render_shard_layout(ty.layout, ctx, indent + ' ')}" - if ty.storage is not StorageKind.GMEM: - result += f', "{ty.storage.name.lower()}"' - return result + "]" - - def _shard_attr_str(self, attr, ctx=None) -> str: - if isinstance(attr, Broadcast): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) - return "B()" - if isinstance(attr, Split): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) - return f"S({attr.axis})" - if isinstance(attr, Partial): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) - return f'P("{attr.reduction}")' - raise TypeError(f"unsupported shard attribute: {type(attr).__name__}") - - def render_layout(self, layout: LayoutBase | None, ctx=None, indent: str = "") -> str: - if layout is None: - return "None" - if isinstance(layout, Layout): - strides = ( - self.owner.shape_tuple(layout.strides, ctx) - if layout.strides is not None - else "None" - ) - return f"Layout({self.owner.shape_tuple(layout.shape, ctx)}, {strides})" - if isinstance(layout, ShardLayout): - return self.render_shard_layout(layout, ctx, indent) - if isinstance(layout, ComposedLayout): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) - child = indent + " " - return ( - "ComposedLayout(\n" - f"{child}inner={self.render_layout(layout.inner, ctx, child)},\n" - f"{child}offset={self.owner.dim_entry(layout.offset, ctx)},\n" - f"{child}outer={self.render_layout(layout.outer, ctx, child)},\n" - f"{indent})" - ) - raise TypeError(f"unsupported layout type: {type(layout).__name__}") - - def render_mesh(self, mesh: Mesh, ctx=None, indent: str = "") -> str: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout, Mesh, Topology",), "")) - values = ", ".join( - f'Topology("{topology.name}", {self.owner.dim_entry(topology.size, ctx)})' - for topology in mesh.topologies + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) + outer, child = self._indent, self._indent + " " + with self.type_surface(indent=child): + inner_text = self.visit(value.inner, ctx) + outer_text = self.visit(value.outer, ctx) + return ( + "ComposedLayout(\n" + f"{child}inner={inner_text},\n" + f"{child}offset={self.dim_entry(value.offset, ctx)},\n" + f"{child}outer={outer_text},\n" + f"{outer})" ) - topologies = f"({values}{',' if len(mesh.topologies) == 1 else ''})" - result = f"Mesh({topologies}, {self.render_layout(mesh.layout, ctx, indent)}" - if mesh.names: - result += f", names={tuple(mesh.names)!r}" - return result + ")" - def render_shard_layout( - self, layout: ShardLayout, ctx=None, indent: str = "", *, mesh_ref=None - ) -> str: + def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: if ctx is not None: ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) - child = indent + " " - attrs = ", ".join(self._shard_attr_str(attr, ctx) for attr in layout.attrs) - if len(layout.attrs) == 1: + outer, child = self._indent, self._indent + " " + attrs = ", ".join(self.visit(attr, ctx) for attr in value.attrs) + if len(value.attrs) == 1: attrs += "," - mesh_text = mesh_ref if mesh_ref is not None else self.render_mesh(layout.mesh, ctx, child) - layout_text = self.render_layout(layout.layout, ctx, child) + with self.type_surface(indent=child): + mesh_text = self.visit(value.mesh, ctx) + layout_text = self.visit(value.layout, ctx) return ( "ShardLayout(\n" f"{child}layout={layout_text},\n" f"{child}attrs=({attrs}),\n" f"{child}mesh={mesh_text},\n" - f"{indent})" + f"{outer})" ) - def _render_layout_positional(self, layout, ctx=None): - if isinstance(layout, Layout): - strides = self.owner.shape_tuple(layout.strides, ctx) if layout.strides is not None else "None" - return f"Layout({self.owner.shape_tuple(layout.shape, ctx)}, {strides})" - return self.render_layout(layout, ctx) + def visit_Broadcast(self, value: Broadcast, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) + return "B()" - def _render_mesh_dataclass(self, mesh, ctx=None): + def visit_Split(self, value: Split, ctx=None) -> str: if ctx is not None: - alias = ctx.mesh_alias(mesh) - if alias is not None: - return alias - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "")) - values = ", ".join( - f'Topology(name="{topology.name}", size={self.owner.dim_entry(topology.size, ctx)})' - for topology in mesh.topologies - ) - if len(mesh.topologies) == 1: - values += "," - names = ", ".join(f'"{name}"' for name in mesh.names) - if len(mesh.names) == 1: - names += "," - layout = self.render_layout(mesh.layout, ctx) - return f"Mesh(topologies=({values}), layout={layout}, names=({names}))" - - def _render_mesh_compact(self, mesh, ctx=None): + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) + return f"S({value.axis})" + + def visit_Partial(self, value: Partial, ctx=None) -> str: if ctx is not None: - alias = ctx.mesh_alias(mesh) - if alias is not None: - return alias - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "")) - values = ", ".join( - f'Topology("{topology.name}", {self.owner.dim_entry(topology.size, ctx)})' - for topology in mesh.topologies - ) - topologies = f"({values}{',' if len(mesh.topologies) == 1 else ''})" - names = f", names={tuple(mesh.names)!r}" if mesh.names else "" - return f"Mesh({topologies}, {self._render_layout_positional(mesh.layout, ctx)}{names})" + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) + return f'P("{value.reduction}")' diff --git a/src/tilefoundry/inspection/tir_printer.py b/src/tilefoundry/inspection/tir_printer.py index f85ae987..14fbfa0f 100644 --- a/src/tilefoundry/inspection/tir_printer.py +++ b/src/tilefoundry/inspection/tir_printer.py @@ -6,6 +6,7 @@ from tilefoundry.inspection.print_context import TirPrintContext from tilefoundry.inspection.printer_base import PythonPrinter +from tilefoundry.inspection.python_type_printer import PythonTypePrinter from tilefoundry.ir.core import Call, Constant, Op, Tuple, Var from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.core.module import Module @@ -13,6 +14,7 @@ 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.stmt import Stmt from tilefoundry.ir.tir.stmts import ( Evaluate, ) @@ -43,8 +45,16 @@ def __init__(self, *, context: TirPrintContext | None = None, indent: str = "") def dim_entry(self, value, ctx=None) -> str: return f"_{value.name}" if hasattr(value, "name") else str(value) - def visit(self, stmt, ctx=None): # type: ignore[override] - return StmtVisitor.visit(self, stmt) + def visit(self, node, ctx=None): # type: ignore[override] + """One dispatch root for the two functor families this printer implements. + + Statements reach the statement visitor; every other value is a type and + reaches the inherited type visitor, so ``self.visit(expr.type, ctx)`` + stays the only way a type is printed. + """ + if isinstance(node, Stmt): + return StmtVisitor.visit(self, node) + return PythonTypePrinter.visit(self, node, ctx) def visit_Sequential(self, stmt): return [line for child in stmt.body for line in self.visit(child)] @@ -56,7 +66,7 @@ def visit_Evaluate(self, stmt): return self._emit_evaluate(stmt) def visit_MeshScope(self, stmt): - lines = [f"{self.indent}with {self._render_mesh_compact(stmt.mesh, self.context)} as {stmt.binding.name}:"] + lines = [f"{self.indent}with {self.visit(stmt.mesh, self.context)} as {stmt.binding.name}:"] self.context.push_mesh(stmt.mesh, stmt.binding.name) lines.extend(TirPrinter(context=self.context, indent=self.indent + " ").visit(stmt.body)) self.context.pop_mesh() @@ -177,7 +187,7 @@ def _function_block(fn: PrimFunction) -> list[str]: lines = [f'_{d.name} = DimVar("{d.name}", {d.lo}, {d.hi})' for d in dim_vars.values()] lines.append("@prim_func(target=" + target + ")") params = ", ".join( - f"{p.name}: {TirPrinter(context=ctx).render_value(p.type, ctx) if isinstance(p.type, TensorType) else repr(p.type)}" + f"{p.name}: {TirPrinter(context=ctx).visit(p.type, ctx) if isinstance(p.type, TensorType) else repr(p.type)}" for p in fn.params ) lines.append(f"def {_binding_name(fn.name)}({params}):") diff --git a/tests/inspection/test_python_printer.py b/tests/inspection/test_python_printer.py index 94ccfaad..bb60384e 100644 --- a/tests/inspection/test_python_printer.py +++ b/tests/inspection/test_python_printer.py @@ -84,7 +84,8 @@ def test_an_annotated_layout_sugar_cannot_say_stays_verbose(): """The fallback keeps the whole layout, and it coexists with sugar. Sugar names a mesh axis, so a mesh with no named axes has nothing to name - and stays verbose without dropping what the verbose form carries. + and stays verbose without dropping what the verbose form carries. The mesh + slot still names a mesh the prelude binds, and spells out one it does not. """ unnamed_axes = as_script(GqaOnline, options=PythonPrintOptions(show_types=True)) verbose = [ @@ -93,9 +94,12 @@ def test_an_annotated_layout_sugar_cannot_say_stays_verbose(): assert verbose for line in verbose: - assert 'mesh=Mesh((Topology("cta", ' in line + annotation = line.split(" # ", 1)[1] + assert "layout=Layout(" in annotation and "attrs=(" in annotation assert "names=" not in line - assert "@ " not in line.split(" # ", 1)[1] + assert "@ " not in annotation + assert any("mesh=cta_2," in line for line in verbose) + assert any('mesh=Mesh((Topology("cta", ' in line for line in verbose) several_meshes = as_script( MoEMegaKernel, options=PythonPrintOptions(show_types=True) diff --git a/tests/inspection/test_type_printer_golden.py b/tests/inspection/test_type_printer_golden.py deleted file mode 100644 index 3aaef0c6..00000000 --- a/tests/inspection/test_type_printer_golden.py +++ /dev/null @@ -1,26 +0,0 @@ -from tilefoundry.inspection import PythonTypePrinter -from tilefoundry.inspection.print_context import HirPrintContext, TirPrintContext -from tilefoundry.inspection.printer_base import PythonPrinter -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology -from tilefoundry.ir.types.storage import StorageKind - - -def test_hir_and_tir_share_type_value_text() -> None: - mesh = Mesh((Topology("thread", 4),), Layout((4,), (1,)), names=("lane",)) - value = TensorType( - shape=(8,), - dtype=DType.f32, - layout=ShardLayout(Layout((8,), (1,)), (S(0),), mesh), - storage=StorageKind.GMEM, - ) - printer = PythonPrinter() - hir = printer.type_printer.render(value, HirPrintContext({id(mesh): "m"})) - tir_ctx = TirPrintContext() - tir_ctx.push_mesh(mesh, "m") - tir = printer.type_printer.render(value, tir_ctx) - assert hir == tir == 'Tensor[(8,), "f32", (8 @ m.lane,)]' - - -def test_type_functor_is_the_shared_dispatch_implementation() -> None: - assert isinstance(PythonPrinter().type_printer, PythonTypePrinter) From 51c2365e159b5dc92f692fe741362c0ac324c89b Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 15:02:38 +0800 Subject: [PATCH 03/21] revert(inspection): withdraw the footprint projection PR #166 added Summing LoopFootprintMetadata.footprints by storage level, the `details` opt-in that recovered the per-buffer rows, and the analysis spec text describing both had no reproduced problem behind them: the review's finding was about a printer duplicating metadata, not about how footprints project. The fixtures, goldens and focused tests that only existed to pin that projection go with it. --- docs/spec/analysis.md | 7 +-- src/tilefoundry/inspection/analysis_report.py | 2 +- src/tilefoundry/inspection/values.py | 14 +---- tests/fixtures/hir/__init__.py | 1 - tests/fixtures/hir/tensor_types.py | 58 ------------------- .../tensor_types.loop_nest.analyze.golden | 29 ---------- .../golden/tensor_types.loop_nest.golden | 29 ---------- .../inspection/test_analyze_render_golden.py | 38 ------------ 8 files changed, 5 insertions(+), 173 deletions(-) delete mode 100644 tests/fixtures/hir/__init__.py delete mode 100644 tests/fixtures/hir/tensor_types.py delete mode 100644 tests/inspection/golden/tensor_types.loop_nest.analyze.golden delete mode 100644 tests/inspection/golden/tensor_types.loop_nest.golden delete mode 100644 tests/inspection/test_analyze_render_golden.py diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index de940b6d..60f48865 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -304,10 +304,7 @@ class LoopFootprintMetadata(IRMetadata): """Known buffer accesses or a lower bound within one authored LoopRegion. Attributes: - footprints: attribute; One row per source buffer and storage level. The - default inspection projection aggregates this field by storage - level; buffer-level rows remain available through the explicit - ``details`` opt-in. + footprints: attribute; One row per source buffer and storage level. known: attribute; Whether every access had a representable relation. """ @@ -437,7 +434,7 @@ of this analysis. | `BufferFootprint.bytes` | Build relations from rank-preserving per-position Types, union the loop-prefixed access images, count the union's integer points, multiply by the dtype bit width, then round the whole buffer reading up to bytes. If the count is not an integer or exceeds `repeated_bytes`, that buffer reading is unavailable. | No | | `BufferFootprint.device_bytes` | Repeat the same exact union measurement from authored Types without shard narrowing, giving the union across logical positions in bytes. | No | | `BufferFootprint.repeated_bytes` | Multiply each operand's per-position element count by its enclosing trip counts, sum accesses to the same buffer, multiply by dtype bit width, then round the whole buffer reading up to bytes. | No | -| `LoopFootprintMetadata.footprints` | One `BufferFootprint` per known source buffer and storage level, sorted by buffer then level. When `known` is false these rows are the available lower bound rather than an empty replacement. The default text projection sums `bytes` by `level`; buffer/device/repeated triples require the `details` opt-in. | No | +| `LoopFootprintMetadata.footprints` | One `BufferFootprint` per known source buffer and storage level, sorted by buffer then level. When `known` is false these rows are the available lower bound rather than an empty replacement. | No | | `LoopFootprintMetadata.known` | False when an access in the loop or a descendant loop lacks a representable forward relation, marking `footprints` as a lower bound; true otherwise. | No | | `ValueLifetime.binding` | Use the parameter or binding name, suffixed with `:` and the line of the value's source span when it has one. Repeated names already differ by the printer's numeric suffix in definition order; the line locates the row in authored source, which a suffix cannot. A value with neither name nor span is `` in definition order. | No | | `ValueLifetime.memory_level` | Emit one lifetime per storage level occupied by the value's Type. | No | diff --git a/src/tilefoundry/inspection/analysis_report.py b/src/tilefoundry/inspection/analysis_report.py index f5b26914..cc8f3757 100644 --- a/src/tilefoundry/inspection/analysis_report.py +++ b/src/tilefoundry/inspection/analysis_report.py @@ -59,7 +59,7 @@ def render_analysis( options=PythonPrintOptions( show_types=True, comment_metadata_types=selected_types_, - comment_opt_in=frozenset({"operands", "details"}) if operands else frozenset(), + comment_opt_in=frozenset({"operands"}) if operands else frozenset(), ), ) labels = { diff --git a/src/tilefoundry/inspection/values.py b/src/tilefoundry/inspection/values.py index 4a480126..d7845316 100644 --- a/src/tilefoundry/inspection/values.py +++ b/src/tilefoundry/inspection/values.py @@ -266,16 +266,7 @@ def _advisory_count(record: MemoryMetadata) -> int: return len(record.advisories) -def _loop_footprints(record: LoopFootprintMetadata) -> dict[str, int]: - """Aggregate the default loop projection by storage level.""" - totals: dict[str, int] = {} - for item in record.footprints: - totals[item.level] = totals.get(item.level, 0) + item.bytes - return totals - - -def _loop_footprint_details(record: LoopFootprintMetadata) -> dict[str, str]: - """Expose buffer-level readings only for an explicit inspection opt-in.""" +def _loop_footprints(record: LoopFootprintMetadata) -> dict[str, str]: return { f"{item.buffer}@{item.memory_level}": ( f"{item.bytes}/{item.device_bytes}/{item.repeated_bytes}" @@ -375,9 +366,8 @@ class PerformanceSummaryView(IRMetadata): ) comment( LoopFootprintMetadata, - Projection("footprints", dict[str, int], _loop_footprints), + Projection("footprints", dict[str, str], _loop_footprints), Projection("status", str, _loop_footprint_status), - Projection("details", dict[str, str], _loop_footprint_details, opt_in=True), ) comment(RooflineMetadata, "ideal_ns", "bound_by") comment( diff --git a/tests/fixtures/hir/__init__.py b/tests/fixtures/hir/__init__.py deleted file mode 100644 index b9f80ff6..00000000 --- a/tests/fixtures/hir/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""HIR printer fixtures.""" diff --git a/tests/fixtures/hir/tensor_types.py b/tests/fixtures/hir/tensor_types.py deleted file mode 100644 index b42f8ac0..00000000 --- a/tests/fixtures/hir/tensor_types.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Small HIR module covering the printer's type and region surfaces.""" - -from __future__ import annotations - -from tilefoundry import func, module -from tilefoundry.dsl import * -from tilefoundry.target import CudaTarget - - -@module( - entry="loop_nest", - target=CudaTarget("nvidia.h200_sxm"), - topologies=(Topology("cta", 1), Topology("thread", 4)), -) -class TensorTypes: - @func - def broadcast(x: Tensor[(8,), "f32"]): - with Mesh(("thread",), (4,), ("lane",)) as _m: - return tf.reshard(x, (8,), "gmem") - - @func - def split_1d(x: Tensor[(16,), "f32"]): - with Mesh(("thread",), (4,), ("lane",)) as m: - local = tf.reshard(x, (16 @ m.lane,), "rmem") - return tf.reshard(local, (16,), "gmem") - - @func - def split_2d(x: Tensor[(8, 16), "f32"]): - with Mesh(("thread",), (4,), ("lane",)) as m: - local = tf.reshard(x, (8, 16 @ m.lane), "rmem") - return tf.reshard(local, (8, 16), "gmem") - - @func - def partial(x: Tensor[(8,), "f32"]): - with Mesh(("thread",), (4,), ("lane",)) as m: - local = tf.reshard(x, (8 @ m.lane,), "rmem") - return tf.reduce(local, (-1,), True, ReduceKind.SUM) - - @func - def mixed(x: Tensor[(8, 16), "f32"]): - with Mesh(("thread", "cta"), (4, 1), ("lane", "tile")) as m: - local = tf.reshard(x, (8 @ m.lane, 16 @ m.tile), "rmem") - return tf.reshard(local, (8, 16), "gmem") - - @func - def nested_mesh(x: Tensor[(8,), "f32"]): - with Mesh(("thread",), (4,), ("outer",)) as _outer: - with Mesh(("thread",), (4,), ("inner",)) as inner: - local = tf.reshard(x, (8 @ inner.inner,), "rmem") - return tf.reshard(local, (8,), "gmem") - - @func - def loop_nest(x: Tensor[(8,), "f32"]): - with Mesh(("thread",), (4,), ("lane",)) as m: - carried = tf.reshard(x, (8 @ m.lane,), "rmem") - for _ in range(2): - carried = tf.square(carried) - return tf.reshard(carried, (8,), "gmem") diff --git a/tests/inspection/golden/tensor_types.loop_nest.analyze.golden b/tests/inspection/golden/tensor_types.loop_nest.analyze.golden deleted file mode 100644 index def3ea36..00000000 --- a/tests/inspection/golden/tensor_types.loop_nest.analyze.golden +++ /dev/null @@ -1,29 +0,0 @@ -from __future__ import annotations - -from tilefoundry import func -from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 -from tilefoundry.ir.types.shard import B, Layout, Mesh, S, ShardLayout, Topology - -thread = Mesh((Topology("thread", 4),), Layout((4,), (1,)), names=('lane',)) - -@func -def loop_nest( - x: Tensor[(8,), "f32"] -) -> Tensor[(8,), "f32", (8,)]: - with thread as _thread: # Tensor[(8,), "f32", (8,)] - carried = reshard(x, layout=ShardLayout( - layout=Layout((4, 2), None), - attrs=(S(0),), - mesh=thread, - ), storage=rmem) # Tensor[(8,), "f32", ((4 @ thread.lane, 2), (0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r32/w0@r32/w0,rmem:r0/w32@r0/w32 - for _ in range(2): # Tensor[(8,), "f32", ((4 @ thread.lane, 2), (0, 1)), "rmem"]; loop-footprint footprints=rmem:64 status=complete - v1 = unary(carried, kind="square") # Tensor[(8,), "f32", ((4 @ thread.lane, 2), (0, 1)), "rmem"]; compute-cost flops=f32:8@8; traffic traffic=rmem:r32/w32@r32/w32 - carried = v1 - v2 = reshard(carried, layout=ShardLayout( - layout=Layout((8,), (1,)), - attrs=(B(),), - mesh=thread, - ), storage=gmem) # Tensor[(8,), "f32", (8,)]; compute-cost; traffic traffic=gmem:r0/w32@r0/w32,rmem:r32/w0@r32/w0 - return v2 diff --git a/tests/inspection/golden/tensor_types.loop_nest.golden b/tests/inspection/golden/tensor_types.loop_nest.golden deleted file mode 100644 index 7d90891f..00000000 --- a/tests/inspection/golden/tensor_types.loop_nest.golden +++ /dev/null @@ -1,29 +0,0 @@ -from __future__ import annotations - -from tilefoundry import func -from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 -from tilefoundry.ir.types.shard import B, Layout, Mesh, S, ShardLayout, Topology - -thread = Mesh((Topology("thread", 4),), Layout((4,), (1,)), names=('lane',)) - -@func -def loop_nest( - x: Tensor[(8,), "f32"] -) -> Tensor[(8,), "f32", (8,)]: - with thread as _thread: - carried = reshard(x, layout=ShardLayout( - layout=Layout((4, 2), None), - attrs=(S(0),), - mesh=thread, - ), storage=rmem) - for _ in range(2): - v1 = unary(carried, kind="square") - carried = v1 - v2 = reshard(carried, layout=ShardLayout( - layout=Layout((8,), (1,)), - attrs=(B(),), - mesh=thread, - ), storage=gmem) - return v2 diff --git a/tests/inspection/test_analyze_render_golden.py b/tests/inspection/test_analyze_render_golden.py deleted file mode 100644 index 227d94e8..00000000 --- a/tests/inspection/test_analyze_render_golden.py +++ /dev/null @@ -1,38 +0,0 @@ -from pathlib import Path - -from tests.fixtures.hir.tensor_types import TensorTypes -from tilefoundry.analysis import analyze -from tilefoundry.analysis.metadata import BufferFootprint, LoopFootprintMetadata -from tilefoundry.inspection import as_script -from tilefoundry.inspection.analysis_report import render_analysis -from tilefoundry.inspection.values import render_comment - - -def test_loop_footprint_comment_aggregates_by_level() -> None: - record = LoopFootprintMetadata( - footprints=( - BufferFootprint("a", "gmem", 8, 16, 24), - BufferFootprint("b", "gmem", 4, 8, 12), - BufferFootprint("c", "rmem", 2, 2, 2), - ), - known=True, - ) - rendered = render_comment(record) - assert rendered == "loop-footprint footprints=gmem:12,rmem:2 status=complete" - detailed = render_comment(record, opt_in=frozenset({"details"})) - assert "details=a@gmem:8/16/24,b@gmem:4/8/12,c@rmem:2/2/2" in detailed - - -def test_loop_nest_typed_and_analyze_goldens() -> None: - golden_dir = Path(__file__).with_name("golden") - fn = TensorTypes.entry_function() - result = analyze(TensorTypes, fn, analysis=("compute-cost", "memory")) - typed = as_script(result.function) - analyzed = render_analysis(result).annotated - assert typed == (golden_dir / "tensor_types.loop_nest.golden").read_text() - assert analyzed == (golden_dir / "tensor_types.loop_nest.analyze.golden").read_text() - - def strip_comments(source: str) -> str: - return "\n".join(line.split(" # ", 1)[0].rstrip() for line in source.splitlines()) - - assert strip_comments(analyzed) == strip_comments(typed) From 9b2e776d1022238f644db18fccd75f9182c7930d Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 15:02:53 +0800 Subject: [PATCH 04/21] feat(inspection,parser): make placement sugar the printed type surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The printer could already emit `((dims), (strides))` and `{axis @ P(...)}` that nothing could read back, and it fell back to the verbose `ShardLayout(...)` whenever a function held more than one mesh. One parser production now covers the whole sugar — dims with splits written inline, an optional stride tuple, and an optional value-state set — so the printer can emit sugar for every shard layout over a named mesh. A layout whose mesh axes are all Broadcast now states them, instead of relying on "the active mesh"; the bare form it used to emit reparsed as a plain Layout and silently dropped the mesh. Meshes with the same descriptor share one prelude name, and a printed type names only a mesh the prelude binds, so a composed mesh rebuilt at each use site no longer produces a prelude line per copy or an undefined reference. The TIR fixtures are authored in that sugar and print back to their own source again, restoring the round-trip contract inspection §2.7 states. --- docs/spec/inspection.md | 26 +- docs/spec/parser.md | 52 ++- src/tilefoundry/inspection/printer_base.py | 12 - src/tilefoundry/inspection/python_printer.py | 178 +++++---- .../inspection/python_type_printer.py | 6 +- src/tilefoundry/ir/tir/memory/tensor_view.py | 4 +- src/tilefoundry/parser/pattern_nodes.py | 352 ++++++++++++------ tests/fixtures/inspection/__init__.py | 1 + .../inspection/type_printer_sugar.printed.txt | 60 +++ .../fixtures/inspection/type_printer_sugar.py | 84 +++++ tests/fixtures/tir/async_sync.py | 13 +- tests/fixtures/tir/mma.py | 18 +- tests/fixtures/tir/rmsnorm.py | 10 +- tests/fixtures/tir/square.py | 24 +- tests/fixtures/tir/sync.py | 29 +- tests/inspection/test_python_printer.py | 9 +- tests/inspection/test_tir_roundtrip.py | 21 +- 17 files changed, 616 insertions(+), 283 deletions(-) create mode 100644 tests/fixtures/inspection/__init__.py create mode 100644 tests/fixtures/inspection/type_printer_sugar.printed.txt create mode 100644 tests/fixtures/inspection/type_printer_sugar.py diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index 8f65c4f5..8d10ecb7 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -178,8 +178,19 @@ executed. A Target subclass with a different constructor customizes ordinary DSL text forms for tensor / layout / shard annotations are owned by [parser](./parser.md). The printer reuses those forms only when they round-trip without losing mesh / layout / storage information; -otherwise it falls back to the verbose `ShardLayout(...)`. Printer -output supports two modes derived from the same pretty-print core: +otherwise it falls back to the verbose `ShardLayout(...)`. + +A `ShardLayout` over a plain `Layout` whose `Mesh` has named axes and a +prelude name ([§2.5](#25-mesh-name-map)) MUST use the placement sugar of +[parser §2.1](./parser.md#21-syntax), in both type slots and op-attribute +slots. That sugar states the layout's own dimensions with each `Split` written +on the dimension it divides, adds the stride tuple when the strides are not +C-order over those dimensions, and states the remaining mesh axes in a +`{axis @ ...}` set. Because the parser reads an unstated mesh axis as +`Broadcast`, the set carries every `Partial` and carries `Broadcast` only when +no `Split` or `Partial` would otherwise name the mesh. + +Printer output supports two modes derived from the same pretty-print core: - `canonical` — round-trippable text used by `as_script()`, pass dumps, and viewer detail `code` blocks: the `Tensor[...]` form of @@ -205,11 +216,6 @@ still renders verbose, so no annotation loses information. The annotation is **display-only** ([§2.7](#27-round-trip-contract)); what round-trips is the emitted code, not its comments. -All canonical type values are dispatched through the shared `TypeFunctor` / -`PythonTypePrinter` implementation. HIR and TIR retain their own function and -statement printers, but `render_mode()` MUST NOT change the syntax of a -`TensorType`, `ShardLayout`, `Layout`, `Mesh`, or shard attribute child value. - Canonical DType text is the descriptor's `name`. Tensor annotations and DType op attributes MUST emit that name as a quoted DSL string. Compact labels MAY omit the quotes, but MUST NOT use the descriptor's raw `repr()`. @@ -237,6 +243,12 @@ references in the function (params, return type, body `Reshard` ops) and assigns variable names from the first declared topology's name. Mesh definitions are emitted in the module prelude / standalone header. +Two `Mesh` values with the same printed descriptor MUST share one name and one +prelude definition: a composed mesh is rebuilt at each use site, so naming its +copies apart would claim the value's parts are placed on different meshes. A +mesh the prelude does not define MUST NOT be named by a printed type; it is +restated in full there instead. + ### 2.6 Specialization printing A dispatch prototype ([hir.md §1.1](./hir.md#11-function)) diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 8e7b0064..a11d86b1 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -94,9 +94,23 @@ dim-expr ::= integer-literal | dim-expr ('+' | '-' | '*' | '//' | '%') dim-expr | (identifier | primary '.' identifier) '(' (dim-expr (',' dim-expr)*)? ')' -placed-shape ::= '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | - dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr))*)? ')' +placed-shape ::= '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) + | dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' dim-expr)*)? ')' + ',' '{' mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')') (',' + mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')'))* '}' ')' + | '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' + mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' + dim-expr)*)? ')' ')' + | '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' + mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '{' mesh-axis '@' + ('B' '(' ')' | 'P' '(' string-literal ')') (',' mesh-axis '@' ('B' '(' + ')' | 'P' '(' string-literal ')'))* '}' ')' + | '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | + dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr))*)? ')' shape ::= '(' (dim-expr (',' dim-expr)*)? ')' | identifier | primary '.' identifier @@ -133,15 +147,27 @@ expression ::= literal | subscript call ::= expression '(' ((expression | keyword-name '=' expression) (',' (expression | keyword-name '=' expression))*)? ')' -explicit-layout ::= '(' (tensor-shape-layout | shape) ',' shape ')' -placed-layout ::= '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | - dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr))*)? ')' +placed-layout ::= '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) + | dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' dim-expr)*)? ')' + ',' '{' mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')') (',' + mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')'))* '}' ')' + | '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' + mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' + dim-expr)*)? ')' ')' + | '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' + mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '{' mesh-axis '@' + ('B' '(' ')' | 'P' '(' string-literal ')') (',' mesh-axis '@' ('B' '(' + ')' | 'P' '(' string-literal ')'))* '}' ')' + | '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | + dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr))*)? ')' plain-layout ::= '(' (dim-expr (',' dim-expr)*)? ')' layout ::= None | primary | call - | explicit-layout | placed-layout | plain-layout storage ::= string-literal @@ -242,10 +268,8 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | --- | --- | --- | --- | --- | | binary_expression, matmul_expression, op_call, unary_expression | expression, slice_endpoint, subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | | binary_expression, matmul_expression, op_call, unary_expression | expression, slice_endpoint, subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | -| dim_expr | dim_expr, layout_extent, layout_shape, tensor_dim_expr, tensor_optional_slot, tensor_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | +| dim_expr | dim_expr, layout_extent, tensor_dim_expr, tensor_optional_slot, tensor_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | | dtype | tensor_dtype | CanonicalDTypeRule | A dtype must resolve to a canonical DType. | src/tilefoundry/parser/ast_pattern.py | -| explicit_layout, layout, placed_layout, plain_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | -| explicit_layout, layout, placed_layout, plain_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | | function | function | FunctionDialectRule | A function kind and constructed value must agree with the active dialect. | src/tilefoundry/parser/pattern_nodes.py | | function | function | FunctionRegistrationRule | A validated function must be registered exactly once in its owning scope. | src/tilefoundry/parser/pattern_nodes.py | | function | function | FunctionReturnCompatibilityRule | A HIR body with a return annotation must satisfy that annotation; a dispatch prototype must declare one, and each variant body must satisfy the prototype return contract. | src/tilefoundry/parser/pattern_nodes.py | @@ -253,12 +277,14 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | function | function | FunctionSignatureRule | A function must construct an ordered parameter tuple. | src/tilefoundry/parser/pattern_nodes.py | | if, while | loop_statement, statement | TirOnlyStatementRule | A TIR-only statement must appear in a prim_func. | src/tilefoundry/parser/pattern_nodes.py | | index_slice | subscript_index | TileWindowSliceBoundRule | A tile window cannot be used as a slice bound. | src/tilefoundry/parser/pattern_nodes.py | +| layout, placed_layout, plain_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| layout, placed_layout, plain_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | | module | module_finalization | ModuleFinalizationRule | A module declaration must contain valid unique members and a resolvable entry. | src/tilefoundry/parser/ast_pattern.py | | module | module_function | ModuleFunctionRegistrationRule | A validated module function must be recorded in declaration order. | src/tilefoundry/parser/ast_pattern.py | | module | module_function | ModuleFunctionValidationRule | A module function must satisfy its root, variant, or converter role before mutation. | src/tilefoundry/parser/ast_pattern.py | | op_call | expression, slice_endpoint, subscript_index | CallVariadicInputFormRule | A variadic call must use one explicit list, tuple, or supported static list comprehension. | src/tilefoundry/parser/pattern_nodes.py | -| placed_shape | layout_shape, tensor_shape | PlacedShapeRule | Placement sugar in a shape slot states both a shape and a layout. | src/tilefoundry/parser/ast_pattern.py | -| shape | layout_shape, layout_strides, tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | +| placed_shape | tensor_shape | PlacedShapeRule | Placement sugar in a shape slot states both a shape and a layout. | src/tilefoundry/parser/ast_pattern.py | +| shape | tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | | storage | tensor_optional_slot | StorageValueRule | Storage must resolve to a StorageKind. | src/tilefoundry/parser/ast_pattern.py | | tensor | annotation, expression, slice_endpoint, subscript_index, type_annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | | tensor | annotation, expression, slice_endpoint, subscript_index, type_annotation | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py | diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 9cdc4db0..e8c0d1e8 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -69,15 +69,3 @@ def render_pattern(self, pattern: Pattern, ctx=None) -> str: return f'DimVarRangePat("{pattern.dim_var}", {pattern.lo}, {pattern.hi})' return repr(pattern) - def mesh_name_map(self, meshes: dict[int, Mesh]) -> dict[int, str]: - used: set[str] = set() - result: dict[int, str] = {} - for identity, mesh in meshes.items(): - base = mesh.topologies[0].name if mesh.topologies else "mesh" - name, suffix = base, 2 - while name in used: - name = f"{base}_{suffix}" - suffix += 1 - used.add(name) - result[identity] = name - return result diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index 58bad889..dedbd180 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -266,13 +266,13 @@ def _ceildiv_args(entry: Call) -> tuple[object, object] | None: def _classify_shard_attrs( sl: ShardLayout, mesh_name: str -) -> tuple[dict[int, list[str]], list[str]] | None: - """Classify shard attributes into layout-axis splits and partials. +) -> tuple[dict[int, list[str]], list[str], list[str]] | None: + """Classify shard attributes into layout-axis splits, partials, broadcasts. - Preserve mesh-axis order, allow nested axes to split one layout axis, and - omit broadcasts. Return ``None`` for rank mismatch, invalid axes, or unknown - attributes so callers use verbose fallback. Surface and compact renderers - share the result, with the latter remapping splits onto tensor axes. + Preserve mesh-axis order and allow nested axes to split one layout axis. + Return ``None`` for rank mismatch, invalid axes, or unknown attributes so + callers use verbose fallback. Surface and compact renderers share the + result, with the latter remapping splits onto tensor axes. """ layout = sl.layout if not isinstance(layout, Layout) or len(sl.attrs) != len(sl.mesh.layout.shape): @@ -281,6 +281,7 @@ def _classify_shard_attrs( names = sl.mesh.names if hasattr(sl.mesh, "names") and sl.mesh.names else () splits: dict[int, list[str]] = {} partials: list[str] = [] + broadcasts: list[str] = [] for mesh_axis_idx, attr in enumerate(sl.attrs): axis_name = names[mesh_axis_idx] if mesh_axis_idx < len(names) else f"ax{mesh_axis_idx}" axis_ref = f"{mesh_name}.{axis_name}" @@ -290,22 +291,21 @@ def _classify_shard_attrs( splits.setdefault(attr.axis, []).append(axis_ref) elif isinstance(attr, Partial): partials.append(f'{axis_ref} @ P("{attr.reduction or "sum"}")') - elif not isinstance(attr, Broadcast): + elif isinstance(attr, Broadcast): + broadcasts.append(f"{axis_ref} @ B()") + else: return None - return splits, partials + return splits, partials, broadcasts -def _shard_layout_surface_str( - sl: ShardLayout, - mesh_name: str = "gpu", - *, - mesh_unique: bool = False, -) -> str | None: +def _shard_layout_surface_str(sl: ShardLayout, mesh_name: str = "gpu", ctx=None) -> str | None: """Render canonical parser sugar for a shard layout. - Inline splits on layout dimensions, emit partial value states as a set, and - omit broadcasts. Include explicit strides only when present. Return ``None`` - when sugar cannot express the layout so callers use verbose fallback. + Inline splits on layout dimensions and state the remaining value states as + a set. Broadcasts are the parser's default for an unstated mesh axis, so + they are written only when nothing else would name the mesh. Include + explicit strides only when present. Return ``None`` when sugar cannot + express the layout so callers use verbose fallback. A symbolic shape has no static C-order strides to compare against, so the ones it states are emitted rather than assumed contiguous. @@ -316,10 +316,12 @@ def _shard_layout_surface_str( classified = _classify_shard_attrs(sl, mesh_name) if classified is None: return None - splits, partials = classified - - if not splits and not partials and not mesh_unique: + splits, partials, broadcasts = classified + states, state_import = (partials, "P") if (splits or partials) else (broadcasts, "B") + if not splits and not states: return None + if states and ctx is not None: + ctx.use(PythonExpr((f"from tilefoundry.ir.types.shard import {state_import}",), state_import)) c_strides = try_c_order_strides(layout.shape) explicit = layout.strides is not None and layout.strides != c_strides @@ -341,7 +343,7 @@ def _shard_layout_surface_str( axis_tuple = f"({dim_str})" stride_str = _shape_tuple(layout.strides) if explicit else None - value_set = "{" + ", ".join(partials) + "}" if partials else None + value_set = "{" + ", ".join(states) + "}" if states else None if stride_str is None and value_set is None: return axis_tuple @@ -369,7 +371,7 @@ def shard_compact_inline( classified = _classify_shard_attrs(sl, mesh_name) if classified is None: return None - splits, partials = classified + splits, partials, _broadcasts = classified la2ta = layout_axis_to_tensor_axis(layout.shape, tensor_shape) split_ref: dict[int, str] = {} for layout_axis, refs in splits.items(): @@ -969,13 +971,7 @@ def _format_call(expr: Call, indent_here: str) -> str: layout_kw = "" if isinstance(target.layout, ShardLayout): layout_text = _shard_layout_str( - target.layout, - indent=indent_here + " ", - mesh_ref=( - mesh_map.get(id(target.layout.mesh)) - if target.layout.mesh.names - else None - ), + target.layout, indent=indent_here + " ", mesh_map=mesh_map ) layout_kw = ", layout=" + layout_text elif target.layout is not None: @@ -1063,7 +1059,9 @@ def _format_call(expr: Call, indent_here: str) -> str: literal = repr(value) attr_strs.append(f"{param.name}={literal}") elif isinstance(value, ShardLayout): - sl_str = _shard_layout_str(value, indent=indent_here + " ") + sl_str = _shard_layout_str( + value, indent=indent_here + " ", mesh_map=mesh_map + ) attr_strs.append(f"{param.name}={sl_str}") elif isinstance(value, TensorType): attr_strs.append(f"{param.name}={_compact_type(value, {})}") @@ -1268,7 +1266,6 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) _HIR_RENDERER = HirPrinter() _dtype_str = _HIR_RENDERER.dtype_str -_mesh_name_map = _HIR_RENDERER.mesh_name_map _pattern_ctor = _HIR_RENDERER.render_pattern @@ -1298,32 +1295,15 @@ def _mesh_str(mesh: Mesh, indent: str = "") -> str: return _type_str(mesh, indent=indent) -def _shard_layout_str(sl: ShardLayout, indent: str = "", *, mesh_ref=None) -> str: - ctx = HirPrintContext({id(sl.mesh): mesh_ref} if mesh_ref is not None else None) - return _type_str(sl, ctx, indent) +def _shard_layout_str(sl: ShardLayout, indent: str = "", *, mesh_map=None) -> str: + """A shard layout in an attribute slot, named from the printed mesh prelude.""" + return _type_str(sl, HirPrintContext(mesh_map), indent) def _tensor_annotation(ty: TensorType, *, mesh_name_map=None, indent="", is_const=False) -> str: return _type_str(ty, HirPrintContext(mesh_name_map), indent, is_const=is_const) -def _bound_mesh_aliases( - names: dict[int, str], meshes: dict[int, Mesh], scope_mesh_ids: set[int] -) -> dict[int, str]: - """Keep only the aliases the mesh prelude actually binds. - - A mesh with no named axes that no scope enters gets no prelude line, so - naming it inside a printed type would emit an undefined reference. Names - are assigned over every mesh first, so dropping the unbound ones here does - not renumber the meshes that remain. - """ - return { - mid: name - for mid, name in names.items() - if meshes[mid].names or mid in scope_mesh_ids - } - - def _collect_all_meshes( fn: HirFunction, ) -> tuple[dict[int, Mesh], dict[int, Mesh]]: @@ -1342,17 +1322,50 @@ def _collect_all_meshes( return type_meshes, scope_meshes -def _dedup_meshes(meshes: dict[int, Mesh]) -> dict[int, Mesh]: - """Collapse structurally identical descriptors before naming hoisted meshes.""" - result: dict[int, Mesh] = {} +def _mesh_name_map(meshes: dict[int, Mesh]) -> dict[int, str]: + """Name every mesh identity, sharing one name per structural descriptor. + + A composed mesh is rebuilt at each use site, so one descriptor reaches the + printer under several identities. Naming those apart would emit a prelude + line per copy and make the annotations read as if they named different + meshes. + """ + used: set[str] = set() + by_signature: dict[str, str] = {} + result: dict[int, str] = {} for identity, mesh in meshes.items(): signature = _mesh_str(mesh) - if any(signature == _mesh_str(existing) for existing in result.values()): - continue - result[identity] = mesh + name = by_signature.get(signature) + if name is None: + base = mesh.topologies[0].name if mesh.topologies else "mesh" + name, suffix = base, 2 + while name in used: + name = f"{base}_{suffix}" + suffix += 1 + used.add(name) + by_signature[signature] = name + result[identity] = name return result +def _bound_mesh_aliases( + names: dict[int, str], meshes: dict[int, Mesh], scope_mesh_ids: set[int] +) -> dict[int, str]: + """Keep only the aliases the mesh prelude actually binds. + + A mesh with no named axes that no scope enters gets no prelude line, so + naming it inside a printed type would emit an undefined reference. Names + are assigned over every mesh first, so dropping the unbound ones here does + not renumber the meshes that remain. + """ + bound = { + names[identity] + for identity, mesh in meshes.items() + if mesh.names or identity in scope_mesh_ids + } + return {identity: name for identity, name in names.items() if name in bound} + + def _emit_header( fn: HirFunction, meshes: dict[int, Mesh], @@ -1443,20 +1456,21 @@ def nested_layouts(layout): lines.append("") - if any(mesh.names or mid in (scope_mesh_ids or ()) for mid, mesh in meshes.items()): - for mid, mesh in meshes.items(): - if not mesh.names and mid not in (scope_mesh_ids or ()): - continue - name = mesh_map[mid] - topologies = _topologies_str(mesh) - names_repr = repr(tuple(mesh.names)) if mesh.names else "()" - lines.append( - f"{name} = Mesh(" - f"{topologies}, " - f"{_layout_str(mesh.layout)}, " - f"names={names_repr}" - f")" - ) + prelude: dict[str, Mesh] = {} + for identity, mesh in meshes.items(): + name = mesh_map.get(identity) + if name is not None: + prelude.setdefault(name, mesh) + for name, mesh in prelude.items(): + names_repr = repr(tuple(mesh.names)) if mesh.names else "()" + lines.append( + f"{name} = Mesh(" + f"{_topologies_str(mesh)}, " + f"{_layout_str(mesh.layout)}, " + f"names={names_repr}" + f")" + ) + if prelude: lines.append("") return lines @@ -1589,7 +1603,21 @@ def module_to_python(fn: HirFunction, module_name: str = "M") -> str: def _module_hir_functions(mod: Module) -> tuple[HirFunction, ...]: """The Module's HIR functions.""" - return tuple(fn for fn in mod.functions if isinstance(fn, HirFunction)) + return tuple(fn for fn in _emission_order(mod) if isinstance(fn, HirFunction)) + + +def _emission_order(mod: Module) -> tuple: + """A Module's functions in the order the printed class body binds them. + + The entry goes last: a body calling a sibling names the attribute the class + body already bound, so every callee must be written before it. Mesh + collection reads the same order, so the printed mesh prelude does not + depend on the order the authored source happened to use. + """ + functions = mod.functions + entry = mod.entry_function() if functions and mod.entry is not None else None + ordered = tuple(fn for fn in functions if fn is not entry) + return ordered + (entry,) if entry is not None else ordered def _module_tree_functions(mod: Module) -> tuple[HirFunction, ...]: @@ -1629,12 +1657,8 @@ def _emit_module_class( Children first, because a body calling one names the attribute it is bound to and a class body binds in the order it is written. """ - functions = mod.functions - entry = mod.entry_function() if functions and mod.entry is not None else None lines = [_module_decorator_line(mod, mod.entry), f"class {module_name}:"] - ordered = tuple(fn for fn in functions if fn is not entry) - if entry is not None: - ordered += (entry,) + ordered = _emission_order(mod) child_entries = { id(child.entry_function()): child.name for child in mod.modules diff --git a/src/tilefoundry/inspection/python_type_printer.py b/src/tilefoundry/inspection/python_type_printer.py index 1bb26a67..238858f8 100644 --- a/src/tilefoundry/inspection/python_type_printer.py +++ b/src/tilefoundry/inspection/python_type_printer.py @@ -67,8 +67,7 @@ def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: mesh_name = ctx.mesh_alias(value.mesh) if ctx is not None else None if mesh_name is None or not value.mesh.names: return None - count = ctx.mesh_count() if ctx is not None and hasattr(ctx, "mesh_count") else 1 - return _shard_layout_surface_str(value, mesh_name=mesh_name, mesh_unique=count == 1) + return _shard_layout_surface_str(value, mesh_name=mesh_name, ctx=ctx) def visit_TensorType(self, value: TensorType, ctx=None) -> str: result = ( @@ -138,6 +137,9 @@ def visit_ComposedLayout(self, value: ComposedLayout, ctx=None) -> str: ) def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: + surface = self.shard_surface(value, ctx) + if surface is not None: + return surface if ctx is not None: ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) outer, child = self._indent, self._indent + " " diff --git a/src/tilefoundry/ir/tir/memory/tensor_view.py b/src/tilefoundry/ir/tir/memory/tensor_view.py index 2cbf907a..8f2a3538 100644 --- a/src/tilefoundry/ir/tir/memory/tensor_view.py +++ b/src/tilefoundry/ir/tir/memory/tensor_view.py @@ -17,7 +17,7 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.shard import c_order_strides -from tilefoundry.ir.types.shard.layout import Layout +from tilefoundry.ir.types.shard.layout import Layout, LayoutBase from tilefoundry.visitor_registry import register_typeinfer @@ -31,7 +31,7 @@ class TensorView(Op): """ memory = ParamDef(kind="input", pattern=Tensor) - layout = ParamDef(kind="attribute", annotation=object) + layout = ParamDef(kind="attribute", annotation=LayoutBase) shape = ParamDef(kind="attribute", annotation=tuple, default=None) diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index 43eb52b3..ca2a101f 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -305,81 +305,6 @@ def construct(match, children, context): RULES: ClassVar[tuple[AstRule[Any], ...]] = (CanonicalDTypeRule(),) -class ExplicitLayoutPattern(ElementPattern): - element_name = "explicit_layout" - syntax = LazyPattern( - lambda: BranchPattern( - "explicit_layout", - AstNodePattern( - ast.Tuple, - FieldPattern( - "elts", - SequencePattern( - AstNodePattern( - ast.Tuple, - ChoicePattern( - ConditionPattern( - "active Mesh", - lambda node, context: ( - context.function is not None - and bool(context.function.state.mesh_stack) - ), - ChildPattern( - "shape", - lambda: TensorShapeLayoutPattern(), - "layout_shape", - "layout_shape", - ), - ), - ChildPattern( - "shape", - lambda: ShapePattern(), - "layout_shape", - "layout_shape", - ), - ), - ), - AstNodePattern( - ast.Tuple, - ChildPattern( - "strides", - lambda: ShapePattern(), - "layout_strides", - "layout_strides", - ), - ), - ), - ), - ), - pattern_id="tensor.layout.explicit", - ) - ) - - @staticmethod - def construct(match, children, context): - shape_or_layout = children["shape"] - strides = children["strides"] - if isinstance(shape_or_layout, runtime.ShardLayout): - shape = shape_or_layout.layout.shape - else: - shape = shape_or_layout - if len(shape) != len(strides): - raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") - layout = runtime.Layout(shape=shape, strides=strides) - if isinstance(shape_or_layout, runtime.ShardLayout): - return runtime.ShardLayout( - layout=layout, - attrs=shape_or_layout.attrs, - mesh=shape_or_layout.mesh, - ) - return layout - - RULES: ClassVar[tuple[AstRule[Any], ...]] = ( - LayoutShapeRule(), - LayoutPositionRule(), - ) - - class PlainLayoutPattern(ElementPattern): element_name = "plain_layout" syntax = LazyPattern( @@ -514,41 +439,167 @@ class PlacedLayout: layout: object -class PlacedLayoutPattern(ElementPattern): - element_name = "placed_layout" - syntax = LazyPattern( - lambda: BindPattern( - AstNodePattern( - ast.Tuple, - FieldPattern( - "elts", - RepeatPattern( +def _layout_dims() -> AstPattern[Any]: + """``(extent, extent @ axis, ...)`` — the layout's own divided dimensions.""" + return AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern( + ChoicePattern( + AstNodePattern( + ast.BinOp, + FieldPattern("op", AstNodePattern(ast.MatMult)), + FieldPattern("left", AstNodePattern(ast.expr)), + FieldPattern( + "right", + ChoicePattern( + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", RepeatPattern(MeshAxisPattern(), minimum=1) + ), + ), + MeshAxisPattern(), + ), + ), + ), + DimExprPattern(), + ) + ), + ), + ) + + +def _layout_strides() -> AstPattern[Any]: + """``(stride, ...)`` — how the divided positions are addressed.""" + return AstNodePattern(ast.Tuple, FieldPattern("elts", RepeatPattern(DimExprPattern()))) + + +def _value_states() -> AstPattern[Any]: + """``{axis @ B(), axis @ P("sum")}`` — what unsplit mesh axes hold.""" + return AstNodePattern( + ast.Set, + FieldPattern( + "elts", + RepeatPattern( + AstNodePattern( + ast.BinOp, + FieldPattern("op", AstNodePattern(ast.MatMult)), + FieldPattern("left", MeshAxisPattern()), + FieldPattern( + "right", ChoicePattern( AstNodePattern( - ast.BinOp, - FieldPattern("op", AstNodePattern(ast.MatMult)), - FieldPattern("left", AstNodePattern(ast.expr)), + ast.Call, FieldPattern( - "right", - ChoicePattern( + "func", + AstNodePattern( + ast.Name, FieldPattern("id", LiteralPattern("B")) + ), + ), + FieldPattern("args", SequencePattern()), + ), + AstNodePattern( + ast.Call, + FieldPattern( + "func", + AstNodePattern( + ast.Name, FieldPattern("id", LiteralPattern("P")) + ), + ), + FieldPattern( + "args", + SequencePattern( AstNodePattern( - ast.Tuple, - FieldPattern( - "elts", - RepeatPattern( - MeshAxisPattern(), - minimum=1, - ), - ), - ), - MeshAxisPattern(), + ast.Constant, + FieldPattern("value", LiteralPattern(value_type=str)), + ) ), ), ), - DimExprPattern(), - ) + ), ), ), + minimum=1, + ), + ), + ) + + +def _layout_sugar_parts(node: object): + """Split layout sugar into its dims, its strides, and its value states. + + ``(d, ...)`` states dims alone. Wrapping the dims in a tuple adds an + optional stride tuple and an optional ``{axis @ B(), axis @ P("sum")}`` + set, in that order, so one production reads every form the printer emits. + Return ``None`` when the node is not layout sugar at all. + """ + if not isinstance(node, ast.Tuple) or not node.elts: + return None + head, *extras = node.elts + if not extras or not isinstance(head, ast.Tuple): + return node, None, None + strides: ast.Tuple | None = None + states: ast.Set | None = None + for extra in extras: + if isinstance(extra, ast.Tuple) and strides is None and states is None: + strides = extra + elif isinstance(extra, ast.Set) and states is None: + states = extra + else: + return None + return head, strides, states + + +def _value_state_parts(node: ast.AST): + """Read ``axis @ B()`` or ``axis @ P("reduction")`` as axis node and state.""" + if not isinstance(node, ast.BinOp) or not isinstance(node.op, ast.MatMult): + return None + call = node.right + if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Name) or call.keywords: + return None + if call.func.id == "B" and not call.args: + return node.left, "B", None + if call.func.id == "P" and len(call.args) == 1: + reduction = call.args[0] + if isinstance(reduction, ast.Constant) and isinstance(reduction.value, str): + if reduction.value: + return node.left, "P", reduction.value + return None + + +class PlacedLayoutPattern(ElementPattern): + """The layout a placement states: split dims, strides, and value states. + + One production covers the whole sugar surface because the parts are not + independent answers: the dims say how the layout is divided, the strides + say how the divided positions are addressed, and the value-state set says + what every mesh axis the dims did not split holds. Splitting them across + patterns is what let the printer emit a stride tuple beside a placement + that nothing could read back. + """ + + element_name = "placed_layout" + syntax = LazyPattern( + lambda: BindPattern( + ChoicePattern( + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + SequencePattern(_layout_dims(), _layout_strides(), _value_states()), + ), + ), + AstNodePattern( + ast.Tuple, + FieldPattern("elts", SequencePattern(_layout_dims(), _layout_strides())), + ), + AstNodePattern( + ast.Tuple, + FieldPattern("elts", SequencePattern(_layout_dims(), _value_states())), + ), + _layout_dims(), ), PlacedLayoutPattern._bind, ) @@ -570,11 +621,15 @@ def _placement_parts(node: ast.AST) -> tuple[ast.AST, tuple[ast.AST, ...]] | Non @staticmethod def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: - assert isinstance(node, ast.Tuple) + parts = _layout_sugar_parts(node) + if parts is None: + return None + dims_node, strides_node, states_node = parts children: list[AstChild] = [] bindings: list[tuple[str, int]] = [] + states: list[tuple[str, str, str | None]] = [] found_placement = False - for tensor_axis, item in enumerate(node.elts): + for tensor_axis, item in enumerate(dims_node.elts): placement = PlacedLayoutPattern._placement_parts(item) extent_node = item axis_nodes: tuple[ast.AST, ...] = () @@ -608,7 +663,41 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat "mesh_axis", ) ) - if not found_placement: + if strides_node is not None: + for index, item in enumerate(strides_node.elts): + stride_context = context.child(situation="layout_strides", role="layout_strides") + if DimExprPattern().match(item, stride_context) is None: + return None + children.append( + AstChild( + f"stride_{index}", + DimExprPattern(), + item, + "layout_strides", + "layout_strides", + ) + ) + if states_node is not None: + for index, item in enumerate(states_node.elts): + state = _value_state_parts(item) + if state is None: + return None + axis_node, kind, reduction = state + axis_context = context.child(situation="mesh_axis", role="mesh_axis") + if MeshAxisPattern().match(axis_node, axis_context) is None: + return None + child_name = f"state_{index}" + states.append((child_name, kind, reduction)) + children.append( + AstChild( + child_name, + MeshAxisPattern(), + axis_node, + "mesh_axis", + "mesh_axis", + ) + ) + if not found_placement and not states and strides_node is None: return None return dataclasses.replace( matched, @@ -616,8 +705,10 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat branch_id="placed_layout", captures={ **matched.captures, - "rank": len(node.elts), + "rank": len(dims_node.elts), + "stride_rank": None if strides_node is None else len(strides_node.elts), "bindings": tuple(bindings), + "states": tuple(states), }, children=tuple(children), ) @@ -626,11 +717,27 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat def construct(match, children, context): rank = match.captures["rank"] shape = tuple(children[f"extent_{axis}"] for axis in range(rank)) - bindings = tuple( + stride_rank = match.captures.get("stride_rank") + strides = ( + None + if stride_rank is None + else tuple(children[f"stride_{index}"] for index in range(stride_rank)) + ) + splits = tuple( (*children[child_name], tensor_axis) for child_name, tensor_axis in match.captures["bindings"] ) - referenced_ids = {id(mesh) for mesh, _, _ in bindings} + states = tuple( + (*children[child_name], kind, reduction) + for child_name, kind, reduction in match.captures.get("states", ()) + ) + if not splits and not states: + if strides is not None and len(shape) != len(strides): + raise ParseError.from_node( + match.node, context, "layout shape/stride rank mismatch" + ) + return runtime.Layout(shape=shape, strides=strides) + referenced_ids = {id(entry[0]) for entry in (*splits, *states)} if context.function is None: raise ParseError.from_node( match.node, context, "placed layout requires function context" @@ -639,7 +746,7 @@ def construct(match, children, context): mesh for mesh in context.function.state.mesh_stack if id(mesh) in referenced_ids ) if len(meshes) != len(referenced_ids): - meshes = tuple(dict.fromkeys(mesh for mesh, _, _ in bindings)) + meshes = tuple(dict.fromkeys(entry[0] for entry in (*splits, *states))) if len(meshes) != len(referenced_ids): raise ParseError.from_node( match.node, context, "placement references an inactive Mesh" @@ -659,20 +766,31 @@ def construct(match, children, context): source_offsets[id(source)] = offset offset += len(source.layout.shape) attrs: list[object] = [runtime.Broadcast() for _ in mesh.layout.shape] - for source, source_axis, tensor_axis in bindings: + bound: set[int] = set() + + def claim(source, source_axis: int) -> int: target_axis = source_offsets[id(source)] + source_axis - if not isinstance(attrs[target_axis], runtime.Broadcast): + if target_axis in bound: raise ParseError.from_node(match.node, context, "mesh axis is bound more than once") - attrs[target_axis] = runtime.Split(tensor_axis) + bound.add(target_axis) + return target_axis + + for source, source_axis, tensor_axis in splits: + attrs[claim(source, source_axis)] = runtime.Split(tensor_axis) + for source, source_axis, kind, reduction in states: + target_axis = claim(source, source_axis) + attrs[target_axis] = Broadcast() if kind == "B" else Partial(reduction) try: canonical = runtime.canonical_shard_layout(shape, mesh, tuple(attrs)) - return runtime.ShardLayout( - layout=runtime.Layout(shape=canonical.layout.shape, strides=None), - attrs=canonical.attrs, - mesh=canonical.mesh, - ) except (TypeError, ValueError) as error: raise ParseError.from_node(match.node, context, str(error)) from error + if strides is not None and len(canonical.layout.shape) != len(strides): + raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") + return runtime.ShardLayout( + layout=runtime.Layout(shape=canonical.layout.shape, strides=strides), + attrs=canonical.attrs, + mesh=canonical.mesh, + ) RULES: ClassVar[tuple[AstRule[Any], ...]] = ( LayoutShapeRule(), @@ -711,7 +829,6 @@ class LayoutPattern(ElementPattern): ), pattern_id="tensor.layout.call", ), - ExplicitLayoutPattern(), PlacedLayoutPattern(), PlainLayoutPattern(), ) @@ -4846,7 +4963,6 @@ def _body_as_ast_module(body: object, *, strip_docstring: bool = False) -> ast.M "ConstantPattern", "DTypePattern", "DimExprPattern", - "ExplicitLayoutPattern", "ExpressionPattern", "ForPattern", "FunctionDialectRule", diff --git a/tests/fixtures/inspection/__init__.py b/tests/fixtures/inspection/__init__.py new file mode 100644 index 00000000..113592ed --- /dev/null +++ b/tests/fixtures/inspection/__init__.py @@ -0,0 +1 @@ +"""Fixtures whose printed source is the inspection contract under test.""" diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt new file mode 100644 index 00000000..f472157a --- /dev/null +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -0,0 +1,60 @@ +from __future__ import annotations + +from tilefoundry.module import module +from tilefoundry import func +from tilefoundry.target import CudaTarget +from tilefoundry.dsl.tf import * # noqa: F401, F403 +from tilefoundry.dsl import Tensor +from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 +from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology + +thread = Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) +thread_2 = Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=('lane',)) +cta = Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')) +cta_2 = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) + +@module(entry="composed_mesh_pipeline", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", 4), Topology("thread", 8),)) +class TypePrinterSugar: + @func + def nested_loop_tuple( + x: Tensor[(8, 16), "f32"], + weight: Tensor[(8, 16), "f32", ((8, 16), {thread.lane @ P("max")})] + ): + with thread_2 as _thread_2: + split = reshard(x, layout=(8 @ thread_2.lane, 16), storage=rmem) + whole = reshard(x, layout=((8, 16), {thread_2.lane @ B()}), storage=rmem) + for _ in range(3): + whole_2 = add(whole, whole) + split_2 = unary(split, kind="square") + split = split_2 + whole = whole_2 + v1 = reshard(split, layout=((8, 16), {thread_2.lane @ B()}), storage=gmem) + v2 = reshard(whole, layout=((8, 16), {thread_2.lane @ B()}), storage=gmem) + with thread as _thread: + unfolded = reshard(weight, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + return (v1, v2, unfolded) + + @func + def composed_mesh_pipeline( + x: Tensor[(8, 4, 16), "f32"], + acc: Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 16), {thread.lane @ P("sum")}), "rmem"], + mixed: Tensor[(8, 16), "f32", ((4 @ cta.tile, 2, 16), {cta.lane @ P("sum")}), "rmem"] + ): + with cta_2 as _cta_2: + with thread as _thread: + composed = reshard(x, layout=(4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), storage=rmem) + v0 = unary(composed, kind="square") + staged = reshard(v0, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) + narrowed = cast(staged, dtype="bf16") + swapped = transpose(narrowed, perm=(0, 2, 1)) + gathered = reshard(swapped, layout=((8, 16, 4), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) + summed = reshard(mixed, layout=((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), storage=rmem) + for _ in range(3): + summed_2 = add(summed, summed) + folded_2 = unary(folded, kind="square") + folded = folded_2 + summed = summed_2 + v2 = reshard(folded, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + v3 = reshard(summed, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + return (gathered, v2, v3) diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py new file mode 100644 index 00000000..19014a2d --- /dev/null +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -0,0 +1,84 @@ +"""Placed values covering every shape the printer's layout sugar can take. + +One entry per combination the sugar has to survive: a mesh level alone and two +levels composed on one value, splits on one and on several tensor axes, every +value state a mesh axis can hold, contiguous and explicitly strided layouts +over the same logical shape, and a loop whose carried fields are placed +differently from each other. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Tensor, tf +from tilefoundry.ir.types.shard import B, Layout, Mesh, P, Topology +from tilefoundry.target import CudaTarget + +_H200 = CudaTarget("nvidia.h200_sxm") +_TOPOLOGIES = (Topology("cta", 4), Topology("thread", 8)) + +_TILE = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=("tile",)) +_WARP_LANE = Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=("warp", "lane")) +_LANES = Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=("lane",)) + + +@module(entry="composed_mesh_pipeline", target=_H200, topologies=_TOPOLOGIES) +class TypePrinterSugar: + @func + def composed_mesh_pipeline( + x: Tensor[(8, 4, 16), "f32"], + acc: Tensor[ + (8, 16), "f32", + ((2 @ _WARP_LANE.warp, 4, 16), {_WARP_LANE.lane @ P("sum")}), + "rmem", + ], + mixed: Tensor[ + (8, 16), "f32", + ((8 @ _TILE.tile, 16), {_WARP_LANE.warp @ B(), _WARP_LANE.lane @ P("sum")}), + "rmem", + ], + ): + with _TILE as cta: + with _WARP_LANE as thr: + composed = tf.reshard( + x, (8 @ cta.tile, 4 @ thr.warp, 16 @ thr.lane), "rmem" + ) + staged = tf.reshard(tf.square(composed), (8 @ cta.tile, 4, 16), "smem") + narrowed = tf.cast(staged, dtype="bf16") + swapped = tf.transpose(narrowed, perm=(0, 2, 1)) + gathered = tf.reshard(swapped, (8, 16, 4), "gmem") + folded = tf.reshard(acc, ((8 @ thr.warp, 16 @ thr.lane)), "rmem") + summed = tf.reshard( + mixed, ((8, 16), {cta.tile @ B(), thr.warp @ B(), thr.lane @ B()}), "rmem" + ) + for _ in range(3): + folded = tf.square(folded) + summed = tf.add(summed, summed) + return ( + gathered, + tf.reshard(folded, (8, 16), "gmem"), + tf.reshard(summed, (8, 16), "gmem"), + ) + + @func + def nested_loop_tuple( + x: Tensor[(8, 16), "f32"], + weight: Tensor[ + (8, 16), "f32", ((8, 16), {_WARP_LANE.warp @ B(), _WARP_LANE.lane @ P("max")}) + ], + ): + with _LANES as lanes: + split = tf.reshard(x, (8 @ lanes.lane, 16), "rmem") + whole = tf.reshard(x, ((8, 16), {lanes.lane @ B()}), "rmem") + for _ in range(3): + split = tf.square(split) + whole = tf.add(whole, whole) + with _WARP_LANE as thr: + unfolded = tf.reshard( + weight, ((8, 16), {thr.warp @ B(), thr.lane @ B()}), "gmem" + ) + return ( + tf.reshard(split, (8, 16), "gmem"), + tf.reshard(whole, (8, 16), "gmem"), + unfolded, + ) diff --git a/tests/fixtures/tir/async_sync.py b/tests/fixtures/tir/async_sync.py index 5fae2689..d52dfa2b 100644 --- a/tests/fixtures/tir/async_sync.py +++ b/tests/fixtures/tir/async_sync.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.shard import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -11,14 +11,9 @@ class AsyncStage: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def async_stage_device(a: Tensor[(128, 4), "f32"], b: Tensor[(128, 4), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as m: - a_view = T.tensor_view(a, layout=ShardLayout(layout=Layout(shape=(128, 4), strides=(4, 1)), attrs=(S(0),), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(128,), strides=(1,)), names=("t",)))) - shared = T.alloc_tensor(tensor_type=Tensor[(128, 4), "f32", - ShardLayout( - layout=Layout((128, 4), (4, 1)), - attrs=(S(0),), - mesh=Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)), - ), "smem"]) - b_view = T.tensor_view(b, layout=ShardLayout(layout=Layout(shape=(128, 4), strides=(4, 1)), attrs=(S(0),), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(128,), strides=(1,)), names=("t",)))) + a_view = T.tensor_view(a, layout=(128 @ m.t, 4)) + shared = T.alloc_tensor(tensor_type=Tensor[(128, 4), "f32", (128 @ m.t, 4), "smem"]) + b_view = T.tensor_view(b, layout=(128 @ m.t, 4)) T.copy_async(a_view, shared) T.cp_async_commit() T.cp_async_wait(n=0) diff --git a/tests/fixtures/tir/mma.py b/tests/fixtures/tir/mma.py index 1f42e44b..39a6280f 100644 --- a/tests/fixtures/tir/mma.py +++ b/tests/fixtures/tir/mma.py @@ -11,8 +11,16 @@ class MmHandwritten: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def mm_device(a: Tensor[(16, 16), "bf16"], b: Tensor[(16, 8), "bf16"], c: Tensor[(16, 8), "f32"]): with Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))) as _warp: - a_view = T.tensor_view(a, layout=ShardLayout(layout=Layout(shape=(2, 4, 2, 8, 2), strides=(1, 2, 8, 16, 128)), attrs=(S(1), S(3)), mesh=Mesh(topologies=(Topology(name="thread", size=32),), layout=Layout(shape=(4, 8), strides=(1, 4)), names=()))) - b_view = T.tensor_view(b, layout=ShardLayout(layout=Layout(shape=(8, 2, 4, 2), strides=(1, 8, 16, 64)), attrs=(S(2), S(0)), mesh=Mesh(topologies=(Topology(name="thread", size=32),), layout=Layout(shape=(4, 8), strides=(1, 4)), names=()))) + a_view = T.tensor_view(a, layout=ShardLayout( + layout=Layout((2, 4, 2, 8, 2), (1, 2, 8, 16, 128)), + attrs=(S(1), S(3)), + mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), + )) + b_view = T.tensor_view(b, layout=ShardLayout( + layout=Layout((8, 2, 4, 2), (1, 8, 16, 64)), + attrs=(S(2), S(0)), + mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), + )) a_frag = T.alloc_tensor(tensor_type=Tensor[(16, 16), "bf16", ShardLayout( layout=Layout((2, 4, 2, 8, 2), (1, 2, 8, 16, 128)), @@ -35,7 +43,11 @@ def mm_device(a: Tensor[(16, 16), "bf16"], b: Tensor[(16, 8), "bf16"], c: Tensor T.copy(b_view, b_frag) T.fill(acc, 0.0) T.mma(acc, a_frag, b_frag, atom=T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN)) - c_view = T.tensor_view(c, layout=ShardLayout(layout=Layout(shape=(2, 4, 8, 2), strides=(1, 2, 8, 64)), attrs=(S(1), S(2)), mesh=Mesh(topologies=(Topology(name="thread", size=32),), layout=Layout(shape=(4, 8), strides=(1, 4)), names=()))) + c_view = T.tensor_view(c, layout=ShardLayout( + layout=Layout((2, 4, 8, 2), (1, 2, 8, 64)), + attrs=(S(1), S(2)), + mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), + )) T.copy(acc, c_view) @prim_func(target=CpuTarget()) diff --git a/tests/fixtures/tir/rmsnorm.py b/tests/fixtures/tir/rmsnorm.py index 4ca68cc8..81d26bfb 100644 --- a/tests/fixtures/tir/rmsnorm.py +++ b/tests/fixtures/tir/rmsnorm.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import B, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.shard import B, Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -10,10 +10,10 @@ class TirRmsnorm: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def rmsnorm_device(x: Tensor[(1, 128), "f32"], weight: Tensor[(128,), "f32"], out: Tensor[(1, 128), "f32"]): - with Mesh((Topology("thread", 1),), Layout((1,), (1,))) as thread: - x_view = T.tensor_view(x, layout=ShardLayout(layout=Layout(shape=(1, 128), strides=(128, 1)), attrs=(B(),), mesh=Mesh(topologies=(Topology(name="thread", size=1),), layout=Layout(shape=(1,), strides=(1,)), names=()))) - weight_view = T.tensor_view(weight, layout=ShardLayout(layout=Layout(shape=(128,), strides=(1,)), attrs=(B(),), mesh=Mesh(topologies=(Topology(name="thread", size=1),), layout=Layout(shape=(1,), strides=(1,)), names=()))) - out_view = T.tensor_view(out, layout=ShardLayout(layout=Layout(shape=(1, 128), strides=(128, 1)), attrs=(B(),), mesh=Mesh(topologies=(Topology(name="thread", size=1),), layout=Layout(shape=(1,), strides=(1,)), names=()))) + with Mesh((Topology("thread", 1),), Layout((1,), (1,)), names=('t',)) as thread: + x_view = T.tensor_view(x, layout=((1, 128), {thread.t @ B()})) + weight_view = T.tensor_view(weight, layout=((128,), {thread.t @ B()})) + out_view = T.tensor_view(out, layout=((1, 128), {thread.t @ B()})) T.rms_norm(x_view, out_view, weight_view, eps=1e-05) T.sync(thread) diff --git a/tests/fixtures/tir/square.py b/tests/fixtures/tir/square.py index e4ee8a52..2b4f1308 100644 --- a/tests/fixtures/tir/square.py +++ b/tests/fixtures/tir/square.py @@ -4,7 +4,7 @@ from tilefoundry.dsl import DimVar, T, Tensor from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.core.pattern import DimVarRangePat -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.shard import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget _S = DimVar("S", 1, 256) @@ -18,14 +18,9 @@ def square_device(x: Tensor[(_S,), "f32"]): @square_device.specialize(DimVarRangePat("S", 1, 127)) def square_small(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=ShardLayout(layout=Layout(shape=(128,), strides=(1,)), attrs=(S(0),), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(128,), strides=(1,)), names=()))) - reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", - ShardLayout( - layout=Layout((128,), (1,)), - attrs=(S(0),), - mesh=Mesh((Topology("thread", 128),), Layout((128,), (1,))), - ), "rmem"]) + with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as thread: + view = T.tensor_view(x, layout=(128 @ thread.t,)) + reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", (128 @ thread.t,), "rmem"]) for phase in range(0, 2, 1): if phase < 1: T.copy(view, reg) @@ -36,14 +31,9 @@ def square_small(x: Tensor[(_S,), "f32"]): @square_device.specialize(DimVarRangePat("S", 128, 255)) def square_large(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=ShardLayout(layout=Layout(shape=(128,), strides=(1,)), attrs=(S(0),), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(128,), strides=(1,)), names=()))) - reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", - ShardLayout( - layout=Layout((128,), (1,)), - attrs=(S(0),), - mesh=Mesh((Topology("thread", 128),), Layout((128,), (1,))), - ), "rmem"]) + with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as thread: + view = T.tensor_view(x, layout=(128 @ thread.t,)) + reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", (128 @ thread.t,), "rmem"]) for phase in range(0, 2, 1): if phase < 1: T.copy(view, reg) diff --git a/tests/fixtures/tir/sync.py b/tests/fixtures/tir/sync.py index 9cc3e4fa..fa56115e 100644 --- a/tests/fixtures/tir/sync.py +++ b/tests/fixtures/tir/sync.py @@ -3,7 +3,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core.kinds import BinaryKind -from tilefoundry.ir.types.shard import ComposedLayout, Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.shard import ComposedLayout, Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -12,18 +12,25 @@ class SyncSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def sync_square_device(a: Tensor[(4, 32), "f32"]): with Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1)), names=('w', 't')) as m: - view = T.tensor_view(a, layout=ShardLayout(layout=Layout(shape=(4, 32), strides=(32, 1)), attrs=(S(0), S(1)), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(4, 32), strides=(32, 1)), names=("w", "t")))) - reg = T.alloc_tensor(tensor_type=Tensor[(4, 32), "f32", - ShardLayout( - layout=Layout((4, 32), (32, 1)), - attrs=(S(0), S(1)), - mesh=Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1)), names=('w', 't')), - ), "rmem"]) + view = T.tensor_view(a, layout=(4 @ m.w, 32 @ m.t)) + reg = T.alloc_tensor(tensor_type=Tensor[(4, 32), "f32", (4 @ m.w, 32 @ m.t), "rmem"]) T.copy(view, reg) T.sync(m) - T.sync(Mesh(topologies=(Topology(name="thread", size=128),), layout=ComposedLayout(inner=None, offset=0, outer=Layout(shape=(1, 32), strides=(32, 1))), names=("w", "t"))) - T.sync(Mesh(topologies=(Topology(name="thread", size=128),), layout=ComposedLayout(inner=None, offset=0, outer=Layout(shape=(2, 32), strides=(32, 1))), names=("w", "t"))) - T.sync(Mesh(topologies=(Topology(name="thread", size=128),), layout=ComposedLayout(inner=None, offset=64, outer=Layout(shape=(2, 32), strides=(32, 1))), names=("w", "t"))) + T.sync(Mesh((Topology("thread", 128),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((1, 32), (32, 1)), + ), names=('w', 't'))) + T.sync(Mesh((Topology("thread", 128),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((2, 32), (32, 1)), + ), names=('w', 't'))) + T.sync(Mesh((Topology("thread", 128),), ComposedLayout( + inner=None, + offset=64, + outer=Layout((2, 32), (32, 1)), + ), names=('w', 't'))) T.binary(reg, reg, reg, kind=BinaryKind.MUL) T.copy(reg, view) diff --git a/tests/inspection/test_python_printer.py b/tests/inspection/test_python_printer.py index bb60384e..68655bdf 100644 --- a/tests/inspection/test_python_printer.py +++ b/tests/inspection/test_python_printer.py @@ -85,7 +85,7 @@ def test_an_annotated_layout_sugar_cannot_say_stays_verbose(): Sugar names a mesh axis, so a mesh with no named axes has nothing to name and stays verbose without dropping what the verbose form carries. The mesh - slot still names a mesh the prelude binds, and spells out one it does not. + slot still names the mesh the prelude binds rather than restating it. """ unnamed_axes = as_script(GqaOnline, options=PythonPrintOptions(show_types=True)) verbose = [ @@ -98,8 +98,11 @@ def test_an_annotated_layout_sugar_cannot_say_stays_verbose(): assert "layout=Layout(" in annotation and "attrs=(" in annotation assert "names=" not in line assert "@ " not in annotation - assert any("mesh=cta_2," in line for line in verbose) - assert any('mesh=Mesh((Topology("cta", ' in line for line in verbose) + hoisted = _hoisted_meshes(unnamed_axes) + assert hoisted + assert all( + any(f"mesh={name}," in line for name in hoisted) for line in verbose + ) several_meshes = as_script( MoEMegaKernel, options=PythonPrintOptions(show_types=True) diff --git a/tests/inspection/test_tir_roundtrip.py b/tests/inspection/test_tir_roundtrip.py index b1ea7872..d34d37bd 100644 --- a/tests/inspection/test_tir_roundtrip.py +++ b/tests/inspection/test_tir_roundtrip.py @@ -18,6 +18,7 @@ FIXTURES = Path(__file__).parents[1] / "fixtures" CANONICAL = tuple(path for path in (FIXTURES / "tir").glob("*.py") if path.name != "layouts.py") +SUGAR = FIXTURES / "inspection" / "type_printer_sugar.py" def _module_in(path: Path): @@ -33,10 +34,22 @@ def _module_in(path: Path): sorted(path for path in CANONICAL if path.name != "__init__.py"), ids=lambda path: path.stem, ) -def test_fixture_prints_a_stable_canonical_source(path: Path) -> None: - """The shared type printer owns the canonical text, not authored spelling.""" - printed = as_script(_module_in(path)) - assert as_script(import_dsl(printed)) == printed +def test_fixture_prints_back_to_its_own_source(path: Path) -> None: + assert as_script(_module_in(path)) == path.read_text() + + +def test_placed_types_print_and_reparse_as_layout_sugar() -> None: + """Placement sugar is the whole type surface, and it parses back to itself. + + The golden is the review surface: a value the sugar cannot state would + appear here as the verbose ``ShardLayout(...)`` form instead, and a mesh + axis the sugar named wrongly would not survive the reparse. + """ + printed = as_script(_module_in(SUGAR)) + + assert printed == SUGAR.with_suffix(".printed.txt").read_text() + assert "ShardLayout(" not in printed + assert as_script(import_dsl(printed, name="TypePrinterSugar")) == printed def test_mixed_hir_tir_module_prints_both_function_families() -> None: From 32639d7fd7036acce9f5d29b976a36464732b634 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 15:03:00 +0800 Subject: [PATCH 05/21] test(analysis): lock the annotated analysis surface to the sugar fixture The analyze golden reads the same fixture the round-trip golden does, so a type that prints one way in emitted code and another inside an annotation is a diff here rather than two goldens drifting apart. --- tests/analysis/test_render_golden.py | 40 +++++++++++++++++++ .../type_printer_sugar.analyzed.txt | 36 +++++++++++++++++ 2 files changed, 76 insertions(+) create mode 100644 tests/analysis/test_render_golden.py create mode 100644 tests/fixtures/inspection/type_printer_sugar.analyzed.txt diff --git a/tests/analysis/test_render_golden.py b/tests/analysis/test_render_golden.py new file mode 100644 index 00000000..55ffe162 --- /dev/null +++ b/tests/analysis/test_render_golden.py @@ -0,0 +1,40 @@ +"""The annotated analysis surface, locked against one placed fixture.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +from tilefoundry.analysis import analyze +from tilefoundry.inspection import as_script +from tilefoundry.inspection.analysis_report import render_analysis +from tilefoundry.ir.core.module import Module + +FIXTURE = Path(__file__).parents[1] / "fixtures" / "inspection" / "type_printer_sugar.py" + + +def _module_in(path: Path) -> Module: + spec = importlib.util.spec_from_file_location(path.stem, path) + assert spec is not None and spec.loader is not None + loaded = importlib.util.module_from_spec(spec) + spec.loader.exec_module(loaded) + return next(value for value in vars(loaded).values() if isinstance(value, Module)) + + +def _without_comments(source: str) -> str: + return "\n".join(line.split(" # ", 1)[0].rstrip() for line in source.splitlines()) + + +def test_analysis_annotates_the_printed_program_without_changing_it() -> None: + """Annotation adds metadata to canonical source; it does not restate types. + + The golden shares its fixture with the round-trip golden, so a type that + reads one way in emitted code and another in an annotation shows up here as + a diff rather than as two goldens that drifted apart. + """ + module = _module_in(FIXTURE) + result = analyze(module, module.entry_function(), analysis=("compute-cost", "memory")) + annotated = render_analysis(result).annotated + + assert annotated == FIXTURE.with_suffix(".analyzed.txt").read_text() + assert _without_comments(annotated) == _without_comments(as_script(result.function)) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt new file mode 100644 index 00000000..f0601d86 --- /dev/null +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -0,0 +1,36 @@ +from __future__ import annotations + +from tilefoundry import func +from tilefoundry.dsl.tf import * # noqa: F401, F403 +from tilefoundry.dsl import Tensor +from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 +from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology + +thread = Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) +cta = Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')) +cta_2 = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) + +@func +def composed_mesh_pipeline( + x: Tensor[(8, 4, 16), "f32"], + acc: Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 16), {thread.lane @ P("sum")}), "rmem"], + mixed: Tensor[(8, 16), "f32", ((4 @ cta.tile, 2, 16), {cta.lane @ P("sum")}), "rmem"] +): + with cta_2 as _cta_2: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})]] + with thread as _thread: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})]] + v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r2048/w0@r512/w0,rmem:r0/w2048@r0/w512 + v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r2048/w2048@r512/w512 + v2 = reshard(v1, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", (4 @ cta_2.tile, 2, 4, 16), "smem"]; compute-cost; traffic traffic=rmem:r2048/w0@r512/w0,smem:r0/w2048@r0/w512 + v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", (4 @ cta_2.tile, 2, 4, 16), "smem"]; compute-cost flops=bf16:512@128; traffic traffic=smem:r2048/w1024@r512/w256 + v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta_2.tile, 2, 16, 4), (128, 64, 1, 16)), "smem"]; compute-cost; traffic traffic=smem:r1024/w1024@r256/w256 + v5 = reshard(v4, layout=((8, 16, 4), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@r0/w256,smem:r1024/w0@r256/w0 + folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", (2 @ thread.warp, 4, 4 @ thread.lane, 4), "rmem"]; compute-cost; traffic + summed = reshard(mixed, layout=((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost; traffic + for _ in range(3): # Tuple[Tensor[(8, 16), "f32", (2 @ thread.warp, 4, 4 @ thread.lane, 4), "rmem"], Tensor[(8, 16), "f32", ((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete + v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r1024/w512@r1024/w512 + v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", (2 @ thread.warp, 4, 4 @ thread.lane, 4), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r512/w512@r512/w512 + folded = v9 + summed = v8 + v11 = reshard(folded, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@r0/w512,rmem:r512/w0@r512/w0 + v13 = reshard(summed, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@r0/w512,rmem:r512/w0@r512/w0 + return (v5, v11, v13) From bc688d4eedc597bfda8d7e39db96e5541c000fef Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 15:17:39 +0800 Subject: [PATCH 06/21] fix(inspection): state a layout's strides in the sugar it prints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An unstated stride tuple is not shorthand for C-order: it is the sugar default a `Reshard` materializes from the storage tier it moves between, so printing a stated C-order tuple as absent turned one layout into another and left TIR views with no strides for CUDA codegen to emit. `render_mode()` and `mesh_count()` go with it — the shared type visitor left both without a caller. --- src/tilefoundry/inspection/print_context.py | 12 ------- src/tilefoundry/inspection/python_printer.py | 31 ++++++++++--------- tests/cli/test_cli.py | 5 +-- .../type_printer_sugar.analyzed.txt | 24 +++++++------- .../inspection/type_printer_sugar.printed.txt | 13 ++++---- .../fixtures/inspection/type_printer_sugar.py | 5 +-- tests/fixtures/tir/async_sync.py | 6 ++-- tests/fixtures/tir/rmsnorm.py | 6 ++-- tests/fixtures/tir/square.py | 8 ++--- tests/fixtures/tir/sync.py | 4 +-- 10 files changed, 54 insertions(+), 60 deletions(-) diff --git a/src/tilefoundry/inspection/print_context.py b/src/tilefoundry/inspection/print_context.py index 6aeea83e..e0f4f275 100644 --- a/src/tilefoundry/inspection/print_context.py +++ b/src/tilefoundry/inspection/print_context.py @@ -9,9 +9,6 @@ class PrintContext: def __init__(self) -> None: self.imports: set[str] = set() - def mesh_count(self) -> int: - return 0 - def use(self, rendered: PythonExpr | str) -> str: if isinstance(rendered, PythonExpr): self.imports.update(rendered.imports) @@ -21,9 +18,6 @@ def use(self, rendered: PythonExpr | str) -> str: def mesh_alias(self, mesh) -> str | None: return None - def render_mode(self) -> str: - return "hir" - class HirPrintContext(PrintContext): def __init__(self, mesh_name_map: dict[int, str] | None = None) -> None: @@ -33,18 +27,12 @@ def __init__(self, mesh_name_map: dict[int, str] | None = None) -> None: def mesh_alias(self, mesh) -> str | None: return self.mesh_name_map.get(id(mesh)) - def mesh_count(self) -> int: - return len(self.mesh_name_map) - class TirPrintContext(PrintContext): def __init__(self) -> None: super().__init__() self._mesh_aliases: list[dict[int, str]] = [] - def render_mode(self) -> str: - return "tir" - def push_mesh(self, mesh, name: str) -> None: self._mesh_aliases.append({id(mesh): name}) diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index dedbd180..3efeaf6d 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -61,7 +61,6 @@ DimSub, DimVar, ) -from tilefoundry.ir.types.shard import try_c_order_strides from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase from tilefoundry.ir.types.shard.mesh import Mesh from tilefoundry.ir.types.shard.shard_layout import ( @@ -301,14 +300,13 @@ def _classify_shard_attrs( def _shard_layout_surface_str(sl: ShardLayout, mesh_name: str = "gpu", ctx=None) -> str | None: """Render canonical parser sugar for a shard layout. - Inline splits on layout dimensions and state the remaining value states as - a set. Broadcasts are the parser's default for an unstated mesh axis, so - they are written only when nothing else would name the mesh. Include - explicit strides only when present. Return ``None`` when sugar cannot - express the layout so callers use verbose fallback. - - A symbolic shape has no static C-order strides to compare against, so the - ones it states are emitted rather than assumed contiguous. + Splits inline on the layout dimension they divide; the remaining value + states form a set, in which a broadcast appears only when nothing else + would name the mesh, because the parser reads an unstated axis that way. + The stride tuple is written whenever the layout has one: an unstated one + is not C-order shorthand but the sugar default a ``Reshard`` materializes + from the storage it moves between. Return ``None`` when sugar cannot + express the layout, so callers use the verbose fallback. """ layout = sl.layout if not isinstance(layout, Layout): @@ -320,11 +318,8 @@ def _shard_layout_surface_str(sl: ShardLayout, mesh_name: str = "gpu", ctx=None) states, state_import = (partials, "P") if (splits or partials) else (broadcasts, "B") if not splits and not states: return None - if states and ctx is not None: - ctx.use(PythonExpr((f"from tilefoundry.ir.types.shard import {state_import}",), state_import)) - c_strides = try_c_order_strides(layout.shape) - explicit = layout.strides is not None and layout.strides != c_strides + explicit = layout.strides is not None if explicit and any( i in splits and _shape_entry_str(dim, nested=True) != shape_entry_str(dim) for i, dim in enumerate(layout.shape) @@ -343,7 +338,15 @@ def _shard_layout_surface_str(sl: ShardLayout, mesh_name: str = "gpu", ctx=None) axis_tuple = f"({dim_str})" stride_str = _shape_tuple(layout.strides) if explicit else None - value_set = "{" + ", ".join(states) + "}" if states else None + value_set = None + if states: + value_set = "{" + ", ".join(states) + "}" + if ctx is not None: + ctx.use( + PythonExpr( + (f"from tilefoundry.ir.types.shard import {state_import}",), state_import + ) + ) if stride_str is None and value_set is None: return axis_tuple diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index f154d9dd..097da619 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -785,8 +785,9 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> annotated_types = [ line.split(" # ", 1)[1].split("; ", 1)[0] for line in lines if " # Tensor[" in line ] - assert 'Tensor[(120, 64), "f32", (120 @ cta_2.tile, 64)]' in annotated_types - assert 'Tensor[(120, 64), "f32", (12 @ cta_3.tile, 10, 64)]' in annotated_types + assert 'Tensor[(120, 64), "f32", ((120 @ cta_2.tile, 64), (64, 1))]' in annotated_types + assert 'Tensor[(120, 64), "f32", ((12 @ cta_3.tile, 10, 64), (640, 64, 1))]' in annotated_types + assert 'Tensor[(120, 64), "f32", ((120, 64), (64, 1), {cta.tile @ B()})]' in annotated_types rows = payload["calls"] assert len(rows) == 7 diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index f0601d86..95131b9a 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -16,21 +16,21 @@ def composed_mesh_pipeline( acc: Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 16), {thread.lane @ P("sum")}), "rmem"], mixed: Tensor[(8, 16), "f32", ((4 @ cta.tile, 2, 16), {cta.lane @ P("sum")}), "rmem"] ): - with cta_2 as _cta_2: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})]] - with thread as _thread: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})]] + with cta_2 as _cta_2: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]] + with thread as _thread: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]] v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r2048/w0@r512/w0,rmem:r0/w2048@r0/w512 v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r2048/w2048@r512/w512 - v2 = reshard(v1, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", (4 @ cta_2.tile, 2, 4, 16), "smem"]; compute-cost; traffic traffic=rmem:r2048/w0@r512/w0,smem:r0/w2048@r0/w512 - v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", (4 @ cta_2.tile, 2, 4, 16), "smem"]; compute-cost flops=bf16:512@128; traffic traffic=smem:r2048/w1024@r512/w256 + v2 = reshard(v1, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta_2.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost; traffic traffic=rmem:r2048/w0@r512/w0,smem:r0/w2048@r0/w512 + v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", ((4 @ cta_2.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost flops=bf16:512@128; traffic traffic=smem:r2048/w1024@r512/w256 v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta_2.tile, 2, 16, 4), (128, 64, 1, 16)), "smem"]; compute-cost; traffic traffic=smem:r1024/w1024@r256/w256 - v5 = reshard(v4, layout=((8, 16, 4), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@r0/w256,smem:r1024/w0@r256/w0 - folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", (2 @ thread.warp, 4, 4 @ thread.lane, 4), "rmem"]; compute-cost; traffic - summed = reshard(mixed, layout=((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost; traffic - for _ in range(3): # Tuple[Tensor[(8, 16), "f32", (2 @ thread.warp, 4, 4 @ thread.lane, 4), "rmem"], Tensor[(8, 16), "f32", ((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete - v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r1024/w512@r1024/w512 - v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", (2 @ thread.warp, 4, 4 @ thread.lane, 4), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r512/w512@r512/w512 + v5 = reshard(v4, layout=((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@r0/w256,smem:r1024/w0@r256/w0 + folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic + summed = reshard(mixed, layout=((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost; traffic + for _ in range(3): # Tuple[Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete + v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r1024/w512@r1024/w512 + v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r512/w512@r512/w512 folded = v9 summed = v8 - v11 = reshard(folded, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@r0/w512,rmem:r512/w0@r512/w0 - v13 = reshard(summed, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@r0/w512,rmem:r512/w0@r512/w0 + v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@r0/w512,rmem:r512/w0@r512/w0 + v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@r0/w512,rmem:r512/w0@r512/w0 return (v5, v11, v13) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index f472157a..db75a033 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -28,10 +28,11 @@ class TypePrinterSugar: split_2 = unary(split, kind="square") split = split_2 whole = whole_2 - v1 = reshard(split, layout=((8, 16), {thread_2.lane @ B()}), storage=gmem) - v2 = reshard(whole, layout=((8, 16), {thread_2.lane @ B()}), storage=gmem) + v1 = reshard(split, layout=((8, 16), (16, 1), {thread_2.lane @ B()}), storage=gmem) + v2 = reshard(whole, layout=((8, 16), (16, 1), {thread_2.lane @ B()}), storage=gmem) with thread as _thread: - unfolded = reshard(weight, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + per_warp = reshard(weight, layout=(2 @ thread.warp, 4, 16), storage=rmem) + unfolded = reshard(per_warp, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) return (v1, v2, unfolded) @func @@ -47,7 +48,7 @@ class TypePrinterSugar: staged = reshard(v0, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) narrowed = cast(staged, dtype="bf16") swapped = transpose(narrowed, perm=(0, 2, 1)) - gathered = reshard(swapped, layout=((8, 16, 4), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) summed = reshard(mixed, layout=((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), storage=rmem) for _ in range(3): @@ -55,6 +56,6 @@ class TypePrinterSugar: folded_2 = unary(folded, kind="square") folded = folded_2 summed = summed_2 - v2 = reshard(folded, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) - v3 = reshard(summed, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + v2 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + v3 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) return (gathered, v2, v3) diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py index 19014a2d..7393da60 100644 --- a/tests/fixtures/inspection/type_printer_sugar.py +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -47,7 +47,7 @@ def composed_mesh_pipeline( narrowed = tf.cast(staged, dtype="bf16") swapped = tf.transpose(narrowed, perm=(0, 2, 1)) gathered = tf.reshard(swapped, (8, 16, 4), "gmem") - folded = tf.reshard(acc, ((8 @ thr.warp, 16 @ thr.lane)), "rmem") + folded = tf.reshard(acc, (8 @ thr.warp, 16 @ thr.lane), "rmem") summed = tf.reshard( mixed, ((8, 16), {cta.tile @ B(), thr.warp @ B(), thr.lane @ B()}), "rmem" ) @@ -74,8 +74,9 @@ def nested_loop_tuple( split = tf.square(split) whole = tf.add(whole, whole) with _WARP_LANE as thr: + per_warp = tf.reshard(weight, (8 @ thr.warp, 16), "rmem") unfolded = tf.reshard( - weight, ((8, 16), {thr.warp @ B(), thr.lane @ B()}), "gmem" + per_warp, ((8, 16), {thr.warp @ B(), thr.lane @ B()}), "gmem" ) return ( tf.reshard(split, (8, 16), "gmem"), diff --git a/tests/fixtures/tir/async_sync.py b/tests/fixtures/tir/async_sync.py index d52dfa2b..682deefe 100644 --- a/tests/fixtures/tir/async_sync.py +++ b/tests/fixtures/tir/async_sync.py @@ -11,9 +11,9 @@ class AsyncStage: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def async_stage_device(a: Tensor[(128, 4), "f32"], b: Tensor[(128, 4), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as m: - a_view = T.tensor_view(a, layout=(128 @ m.t, 4)) - shared = T.alloc_tensor(tensor_type=Tensor[(128, 4), "f32", (128 @ m.t, 4), "smem"]) - b_view = T.tensor_view(b, layout=(128 @ m.t, 4)) + a_view = T.tensor_view(a, layout=((128 @ m.t, 4), (4, 1))) + shared = T.alloc_tensor(tensor_type=Tensor[(128, 4), "f32", ((128 @ m.t, 4), (4, 1)), "smem"]) + b_view = T.tensor_view(b, layout=((128 @ m.t, 4), (4, 1))) T.copy_async(a_view, shared) T.cp_async_commit() T.cp_async_wait(n=0) diff --git a/tests/fixtures/tir/rmsnorm.py b/tests/fixtures/tir/rmsnorm.py index 81d26bfb..6e8890d8 100644 --- a/tests/fixtures/tir/rmsnorm.py +++ b/tests/fixtures/tir/rmsnorm.py @@ -11,9 +11,9 @@ class TirRmsnorm: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def rmsnorm_device(x: Tensor[(1, 128), "f32"], weight: Tensor[(128,), "f32"], out: Tensor[(1, 128), "f32"]): with Mesh((Topology("thread", 1),), Layout((1,), (1,)), names=('t',)) as thread: - x_view = T.tensor_view(x, layout=((1, 128), {thread.t @ B()})) - weight_view = T.tensor_view(weight, layout=((128,), {thread.t @ B()})) - out_view = T.tensor_view(out, layout=((1, 128), {thread.t @ B()})) + x_view = T.tensor_view(x, layout=((1, 128), (128, 1), {thread.t @ B()})) + weight_view = T.tensor_view(weight, layout=((128,), (1,), {thread.t @ B()})) + out_view = T.tensor_view(out, layout=((1, 128), (128, 1), {thread.t @ B()})) T.rms_norm(x_view, out_view, weight_view, eps=1e-05) T.sync(thread) diff --git a/tests/fixtures/tir/square.py b/tests/fixtures/tir/square.py index 2b4f1308..038e795e 100644 --- a/tests/fixtures/tir/square.py +++ b/tests/fixtures/tir/square.py @@ -19,8 +19,8 @@ def square_device(x: Tensor[(_S,), "f32"]): @square_device.specialize(DimVarRangePat("S", 1, 127)) def square_small(x: Tensor[(_S,), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as thread: - view = T.tensor_view(x, layout=(128 @ thread.t,)) - reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", (128 @ thread.t,), "rmem"]) + view = T.tensor_view(x, layout=((128 @ thread.t,), (1,))) + reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", ((128 @ thread.t,), (1,)), "rmem"]) for phase in range(0, 2, 1): if phase < 1: T.copy(view, reg) @@ -32,8 +32,8 @@ def square_small(x: Tensor[(_S,), "f32"]): @square_device.specialize(DimVarRangePat("S", 128, 255)) def square_large(x: Tensor[(_S,), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as thread: - view = T.tensor_view(x, layout=(128 @ thread.t,)) - reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", (128 @ thread.t,), "rmem"]) + view = T.tensor_view(x, layout=((128 @ thread.t,), (1,))) + reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", ((128 @ thread.t,), (1,)), "rmem"]) for phase in range(0, 2, 1): if phase < 1: T.copy(view, reg) diff --git a/tests/fixtures/tir/sync.py b/tests/fixtures/tir/sync.py index fa56115e..266e4768 100644 --- a/tests/fixtures/tir/sync.py +++ b/tests/fixtures/tir/sync.py @@ -12,8 +12,8 @@ class SyncSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def sync_square_device(a: Tensor[(4, 32), "f32"]): with Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1)), names=('w', 't')) as m: - view = T.tensor_view(a, layout=(4 @ m.w, 32 @ m.t)) - reg = T.alloc_tensor(tensor_type=Tensor[(4, 32), "f32", (4 @ m.w, 32 @ m.t), "rmem"]) + view = T.tensor_view(a, layout=((4 @ m.w, 32 @ m.t), (32, 1))) + reg = T.alloc_tensor(tensor_type=Tensor[(4, 32), "f32", ((4 @ m.w, 32 @ m.t), (32, 1)), "rmem"]) T.copy(view, reg) T.sync(m) T.sync(Mesh((Topology("thread", 128),), ComposedLayout( From 1ff589d004739c57efd508d531ce50668002701d Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 15:35:57 +0800 Subject: [PATCH 07/21] test(analysis): regenerate the analyze golden on main's per-level metadata `main` now reports compute cost and traffic per topology level, so the golden's metadata comments carry `cta:`/`thread:` breakdowns. The type text the goldens exist to pin is unchanged. --- .../type_printer_sugar.analyzed.txt | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index 95131b9a..1efde241 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -18,19 +18,19 @@ def composed_mesh_pipeline( ): with cta_2 as _cta_2: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]] with thread as _thread: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]] - v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r2048/w0@r512/w0,rmem:r0/w2048@r0/w512 - v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r2048/w2048@r512/w512 - v2 = reshard(v1, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta_2.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost; traffic traffic=rmem:r2048/w0@r512/w0,smem:r0/w2048@r0/w512 - v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", ((4 @ cta_2.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost flops=bf16:512@128; traffic traffic=smem:r2048/w1024@r512/w256 - v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta_2.tile, 2, 16, 4), (128, 64, 1, 16)), "smem"]; compute-cost; traffic traffic=smem:r1024/w1024@r256/w256 - v5 = reshard(v4, layout=((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@r0/w256,smem:r1024/w0@r256/w0 + v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.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 @ cta.warp, 2, 4 @ cta.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 + v2 = reshard(v1, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta_2.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_2.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 + v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta_2.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), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic summed = reshard(mixed, layout=((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost; traffic for _ in range(3): # Tuple[Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete - v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r1024/w512@r1024/w512 - v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:512@128; traffic traffic=rmem:r512/w512@r512/w512 + v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost flops=f32:512@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 + v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:512@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 folded = v9 summed = v8 - v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@r0/w512,rmem:r512/w0@r512/w0 - v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@r0/w512,rmem:r512/w0@r512/w0 + v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 + v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 return (v5, v11, v13) From 439011a9b2bd036557cb7d43e11e1eebd5594ccd Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 15:56:55 +0800 Subject: [PATCH 08/21] docs(tutorial): re-render the showcase page on the sugar surface The page reproduces itself from its own blocks, so the reshard it displays is now the one-line placement sugar rather than a five-line `ShardLayout`, and the window it prints reaches the matmul that consumes the weight. --- docs/tutorial/showcase.ipynb | 2 +- docs/tutorial/showcase.md | 16 +++++----------- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index 2ab653dd..bb216ef2 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: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 @ cta.head, 32), 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 v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ cta.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 @ cta.head, 32), 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 v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ cta.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" } ], "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" diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index 55b9666a..5510a9b9 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -870,17 +870,11 @@ for needle in ("reshard(w_q", "reshard(w_o"): # peak-footprint=gmem:10945036,smem:16960 # roofline ideal-ns=11213 bound-by=memory - v1 = reshard(w_q, layout=ShardLayout( - layout=Layout((1, 256, 8, 32), None), - attrs=(S(2),), - mesh=cta, - ), 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 - - v42 = reshard(w_o, layout=ShardLayout( - layout=Layout((1, 256, 8, 32), None), - attrs=(S(2),), - mesh=cta, - ), 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 + v1 = reshard(w_q, layout=(1, 256, 8 @ cta.head, 32), 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 + v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ cta.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 + + v42 = reshard(w_o, layout=(1, 256, 8 @ cta.head, 32), 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 + v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ cta.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 ``` ## 6. Stream the KV cache From ef6f0e5f7a69c3932ae756ca11c1e5d8ed13206e Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 22:13:26 +0800 Subject: [PATCH 09/21] test(inspection): hold the two placements sugar cannot state `named_and_out_of_scope` adds a layout the author names rather than spells, and a reshard onto a mesh its scope never enters -- reshard is the one operation allowed to cross that boundary, so its target names a mesh that is not the scope it runs in. The golden records what the printer does with them today: the named layout is expanded back into sugar, the out-of-scope target is written as `thread.warp` from inside `with cta`, and `@func(mesh=...)` comes back as a `with` in the body. All three are wrong and none of them had a test. --- .../inspection/type_printer_sugar.printed.txt | 26 ++++++++++++----- .../fixtures/inspection/type_printer_sugar.py | 28 +++++++++++++++++-- 2 files changed, 45 insertions(+), 9 deletions(-) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index db75a033..f51bcf79 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -10,8 +10,8 @@ from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topol thread = Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) thread_2 = Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=('lane',)) -cta = Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')) -cta_2 = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) +cta = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) +cta_2 = Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')) @module(entry="composed_mesh_pipeline", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", 4), Topology("thread", 8),)) class TypePrinterSugar: @@ -35,22 +35,34 @@ class TypePrinterSugar: unfolded = reshard(per_warp, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) return (v1, v2, unfolded) + @func + def named_and_out_of_scope( + x: Tensor[(8, 16), "f32"], + held: Tensor[(8, 16), "f32", (4 @ cta.tile, 2, 16), "rmem"], + frag: Tensor[(16,), "f32", ((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), "rmem"] + ): + with cta as _cta: + mine = reshard(x, layout=(4 @ cta.tile, 2, 16), storage=rmem) + v0 = reshard(mine, layout=Layout((8, 16), (16, 1)), storage=gmem) + escaped = reshard(x, layout=(2 @ thread.warp, 4, 16), storage=rmem) + return (v0, held, frag, escaped) + @func def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], acc: Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 16), {thread.lane @ P("sum")}), "rmem"], - mixed: Tensor[(8, 16), "f32", ((4 @ cta.tile, 2, 16), {cta.lane @ P("sum")}), "rmem"] + mixed: Tensor[(8, 16), "f32", ((4 @ cta_2.tile, 2, 16), {cta_2.lane @ P("sum")}), "rmem"] ): - with cta_2 as _cta_2: + with cta as _cta: with thread as _thread: - composed = reshard(x, layout=(4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), storage=rmem) + composed = reshard(x, layout=(4 @ cta_2.tile, 2, 2 @ cta_2.warp, 2, 4 @ cta_2.lane, 4), storage=rmem) v0 = unary(composed, kind="square") - staged = reshard(v0, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) + staged = reshard(v0, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) narrowed = cast(staged, dtype="bf16") swapped = transpose(narrowed, perm=(0, 2, 1)) gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) - summed = reshard(mixed, layout=((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), storage=rmem) + summed = reshard(mixed, layout=((8, 16), {cta_2.tile @ B(), cta_2.warp @ B(), cta_2.lane @ B()}), storage=rmem) for _ in range(3): summed_2 = add(summed, summed) folded_2 = unary(folded, kind="square") diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py index 7393da60..31d802e2 100644 --- a/tests/fixtures/inspection/type_printer_sugar.py +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -4,14 +4,14 @@ levels composed on one value, splits on one and on several tensor axes, every value state a mesh axis can hold, contiguous and explicitly strided layouts over the same logical shape, and a loop whose carried fields are placed -differently from each other. +differently. The last entry holds the two placements sugar cannot state at all. """ from __future__ import annotations from tilefoundry import func, module from tilefoundry.dsl import Tensor, tf -from tilefoundry.ir.types.shard import B, Layout, Mesh, P, Topology +from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology from tilefoundry.target import CudaTarget _H200 = CudaTarget("nvidia.h200_sxm") @@ -21,6 +21,12 @@ _WARP_LANE = Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=("warp", "lane")) _LANES = Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=("lane",)) +_FRAGMENT = ShardLayout( + layout=Layout((2, 4, 2), (8, 2, 1)), + attrs=(S(0), S(1)), + mesh=_WARP_LANE, +) + @module(entry="composed_mesh_pipeline", target=_H200, topologies=_TOPOLOGIES) class TypePrinterSugar: @@ -83,3 +89,21 @@ def nested_loop_tuple( tf.reshard(whole, (8, 16), "gmem"), unfolded, ) + + @func(mesh=_TILE) + def named_and_out_of_scope( + x: Tensor[(8, 16), "f32"], + held: Tensor[(8, 16), "f32", (8 @ mesh.tile, 16), "rmem"], # noqa: F821 + frag: Tensor[(16,), "f32", _FRAGMENT, "rmem"], + ): + """The two placements sugar cannot state. + + ``frag`` is a layout the author names rather than spells: its content is + dictated elsewhere, so writing it out here would copy a constant into + the program text. ``escaped`` is resharded onto a mesh this scope never + enters -- a reshard is the one operation allowed to cross that boundary, + so its target names a mesh that is not the scope it runs in. + """ + mine = tf.reshard(x, (8 @ mesh.tile, 16), "rmem") # noqa: F821 + escaped = tf.reshard(x, ((8 @ _WARP_LANE.warp, 16), {_WARP_LANE.lane @ B()}), "rmem") + return tf.reshard(mine, (8, 16), "gmem"), held, frag, escaped From 2fb5a4c7065befc759d650ae788ee67aeaca8968 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 22:52:40 +0800 Subject: [PATCH 10/21] test(inspection): pair a named layout across a bound and an unbound mesh `composed_mesh_pipeline` reshards to the same constant `named_and_out_of_scope` holds as a parameter. One function enters that constant's mesh with a `with`; the other never does. The golden records both printing identically today. Only the first has a scope binding a sugar identifier could refer to, so only the first should keep the sugar once the printer requires one. --- .../inspection/type_printer_sugar.analyzed.txt | 9 ++++++--- .../inspection/type_printer_sugar.printed.txt | 5 ++++- tests/fixtures/inspection/type_printer_sugar.py | 16 ++++++++++------ 3 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index 1efde241..6b201a0e 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -13,11 +13,12 @@ cta_2 = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) @func def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], + seed: Tensor[(16,), "f32"], acc: Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 16), {thread.lane @ P("sum")}), "rmem"], mixed: Tensor[(8, 16), "f32", ((4 @ cta.tile, 2, 16), {cta.lane @ P("sum")}), "rmem"] ): - with cta_2 as _cta_2: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]] - with thread as _thread: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]] + with cta_2 as _cta_2: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B(), thread.lane @ B()})]] + with thread as _thread: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B(), thread.lane @ B()})]] v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.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 @ cta.warp, 2, 4 @ cta.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 v2 = reshard(v1, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta_2.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 @@ -33,4 +34,6 @@ def composed_mesh_pipeline( summed = v8 v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 - return (v5, v11, v13) + v14 = reshard(seed, layout=((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ thread.warp, 4 @ thread.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 + v15 = reshard(v14, layout=((16,), (1,), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 + return (v5, v11, v13, v15) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index f51bcf79..2160b675 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -50,6 +50,7 @@ class TypePrinterSugar: @func def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], + seed: Tensor[(16,), "f32"], acc: Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 16), {thread.lane @ P("sum")}), "rmem"], mixed: Tensor[(8, 16), "f32", ((4 @ cta_2.tile, 2, 16), {cta_2.lane @ P("sum")}), "rmem"] ): @@ -70,4 +71,6 @@ class TypePrinterSugar: summed = summed_2 v2 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) v3 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) - return (gathered, v2, v3) + seeded = reshard(seed, layout=((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), storage=rmem) + v4 = reshard(seeded, layout=((16,), (1,), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) + return (gathered, v2, v3, v4) diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py index 31d802e2..9238baa1 100644 --- a/tests/fixtures/inspection/type_printer_sugar.py +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -4,7 +4,8 @@ levels composed on one value, splits on one and on several tensor axes, every value state a mesh axis can hold, contiguous and explicitly strided layouts over the same logical shape, and a loop whose carried fields are placed -differently. The last entry holds the two placements sugar cannot state at all. +differently. `_FRAGMENT` is reshard-ed to from both entries: one enters its +mesh, the other never does, so only the first has a binding sugar could name. """ from __future__ import annotations @@ -33,6 +34,7 @@ class TypePrinterSugar: @func def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], + seed: Tensor[(16,), "f32"], acc: Tensor[ (8, 16), "f32", ((2 @ _WARP_LANE.warp, 4, 16), {_WARP_LANE.lane @ P("sum")}), @@ -53,6 +55,7 @@ def composed_mesh_pipeline( narrowed = tf.cast(staged, dtype="bf16") swapped = tf.transpose(narrowed, perm=(0, 2, 1)) gathered = tf.reshard(swapped, (8, 16, 4), "gmem") + seeded = tf.reshard(seed, _FRAGMENT, "rmem") folded = tf.reshard(acc, (8 @ thr.warp, 16 @ thr.lane), "rmem") summed = tf.reshard( mixed, ((8, 16), {cta.tile @ B(), thr.warp @ B(), thr.lane @ B()}), "rmem" @@ -64,6 +67,7 @@ def composed_mesh_pipeline( gathered, tf.reshard(folded, (8, 16), "gmem"), tf.reshard(summed, (8, 16), "gmem"), + tf.reshard(seeded, (16,), "gmem"), ) @func @@ -98,11 +102,11 @@ def named_and_out_of_scope( ): """The two placements sugar cannot state. - ``frag`` is a layout the author names rather than spells: its content is - dictated elsewhere, so writing it out here would copy a constant into - the program text. ``escaped`` is resharded onto a mesh this scope never - enters -- a reshard is the one operation allowed to cross that boundary, - so its target names a mesh that is not the scope it runs in. + ``frag`` holds the same constant `composed_mesh_pipeline` reshards to, + but this function never enters that constant's mesh, so nothing here + binds a name sugar could use. ``escaped`` is resharded onto a mesh this + scope never enters -- a reshard is the one operation allowed to cross + that boundary, so its target names a mesh that is not its own scope. """ mine = tf.reshard(x, (8 @ mesh.tile, 16), "rmem") # noqa: F821 escaped = tf.reshard(x, ((8 @ _WARP_LANE.warp, 16), {_WARP_LANE.lane @ B()}), "rmem") From 013dfe95f91c764354bc99d50e0468f8dfd1f05a Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sun, 13 Sep 2026 23:17:37 +0800 Subject: [PATCH 11/21] refactor(parser): let one pattern answer a whole placement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A placement states two things — the extents as written, and where they go — and the layout it implies keeps only the second, divided beyond recovery. So the shape slot could not reuse the layout slot's pattern: it copied the whole syntax tree and recovered the extents itself, and `docs/spec/parser.md` carried that copy as a second production, word for word. One pattern now answers both halves and each slot takes the part it asked for. The rules follow: what sat on the placement were the layout slot's rules, so a placement's own rule reaches the constraints table for the first time and `PlacedShapeRule` goes with the pattern it guarded. --- docs/spec/parser.md | 27 ++-------- src/tilefoundry/parser/ast_pattern.py | 14 ----- src/tilefoundry/parser/pattern_nodes.py | 72 ++++++++++++------------- 3 files changed, 40 insertions(+), 73 deletions(-) diff --git a/docs/spec/parser.md b/docs/spec/parser.md index a11d86b1..451dede6 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -94,7 +94,7 @@ dim-expr ::= integer-literal | dim-expr ('+' | '-' | '*' | '//' | '%') dim-expr | (identifier | primary '.' identifier) '(' (dim-expr (',' dim-expr)*)? ')' -placed-shape ::= '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) +placed-layout ::= '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' dim-expr)*)? ')' ',' '{' mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')') (',' @@ -114,7 +114,7 @@ placed-shape ::= '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis shape ::= '(' (dim-expr (',' dim-expr)*)? ')' | identifier | primary '.' identifier -tensor-shape-layout ::= placed-shape +tensor-shape-layout ::= placed-layout | shape dtype ::= string-literal | primary @@ -147,23 +147,6 @@ expression ::= literal | subscript call ::= expression '(' ((expression | keyword-name '=' expression) (',' (expression | keyword-name '=' expression))*)? ')' -placed-layout ::= '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) - | dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' dim-expr)*)? ')' - ',' '{' mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')') (',' - mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')'))* '}' ')' - | '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' - mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' - dim-expr)*)? ')' ')' - | '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' - mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '{' mesh-axis '@' - ('B' '(' ')' | 'P' '(' string-literal ')') (',' mesh-axis '@' ('B' '(' - ')' | 'P' '(' string-literal ')'))* '}' ')' - | '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | - dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr))*)? ')' plain-layout ::= '(' (dim-expr (',' dim-expr)*)? ')' layout ::= None | primary @@ -277,13 +260,13 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | function | function | FunctionSignatureRule | A function must construct an ordered parameter tuple. | src/tilefoundry/parser/pattern_nodes.py | | if, while | loop_statement, statement | TirOnlyStatementRule | A TIR-only statement must appear in a prim_func. | src/tilefoundry/parser/pattern_nodes.py | | index_slice | subscript_index | TileWindowSliceBoundRule | A tile window cannot be used as a slice bound. | src/tilefoundry/parser/pattern_nodes.py | -| layout, placed_layout, plain_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | -| layout, placed_layout, plain_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | +| layout, plain_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| layout, plain_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | | module | module_finalization | ModuleFinalizationRule | A module declaration must contain valid unique members and a resolvable entry. | src/tilefoundry/parser/ast_pattern.py | | module | module_function | ModuleFunctionRegistrationRule | A validated module function must be recorded in declaration order. | src/tilefoundry/parser/ast_pattern.py | | module | module_function | ModuleFunctionValidationRule | A module function must satisfy its root, variant, or converter role before mutation. | src/tilefoundry/parser/ast_pattern.py | | op_call | expression, slice_endpoint, subscript_index | CallVariadicInputFormRule | A variadic call must use one explicit list, tuple, or supported static list comprehension. | src/tilefoundry/parser/pattern_nodes.py | -| placed_shape | tensor_shape | PlacedShapeRule | Placement sugar in a shape slot states both a shape and a layout. | src/tilefoundry/parser/ast_pattern.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementAnswerRule | Placement sugar states both the shape as written and the layout it implies. | src/tilefoundry/parser/pattern_nodes.py | | shape | tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | | storage | tensor_optional_slot | StorageValueRule | Storage must resolve to a StorageKind. | src/tilefoundry/parser/ast_pattern.py | | tensor | annotation, expression, slice_endpoint, subscript_index, type_annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 9b101b7e..44d4fa81 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -1580,20 +1580,6 @@ def apply(self, value, *, match, context): return value -@dataclass(frozen=True) -class PlacedShapeRule: - STATEMENT: ClassVar[str] = "Placement sugar in a shape slot states both a shape and a layout." - - def apply(self, value, *, match, context): - if not isinstance(value.shape, tuple): - raise ParseError.from_node(match.node, context, "placed shape is not a tuple") - if not isinstance(value.layout, runtime.LayoutBase): - raise ParseError.from_node( - match.node, context, "placed shape did not carry a LayoutBase" - ) - return value - - @dataclass(frozen=True) class StorageValueRule: STATEMENT: ClassVar[str] = "Storage must resolve to a StorageKind." diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index ca2a101f..3af5de89 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -68,7 +68,6 @@ ParseError, ParserTypeInferContext, PatternFailure, - PlacedShapeRule, PredicatePattern, ReferencePattern, RepeatPattern, @@ -569,6 +568,24 @@ def _value_state_parts(node: ast.AST): return None +@dataclass(frozen=True) +class PlacementAnswerRule: + """A placement answers both halves: the shape written, and where it goes.""" + + STATEMENT: ClassVar[str] = ( + "Placement sugar states both the shape as written and the layout it implies." + ) + + def apply(self, value, *, match, context): + if not isinstance(value, PlacedLayout): + raise ParseError.from_node(match.node, context, "placement did not answer a placement") + if not isinstance(value.shape, tuple): + raise ParseError.from_node(match.node, context, "placement shape is not a tuple") + if not isinstance(value.layout, runtime.LayoutBase): + raise ParseError.from_node(match.node, context, "placement did not carry a LayoutBase") + return value + + class PlacedLayoutPattern(ElementPattern): """The layout a placement states: split dims, strides, and value states. @@ -736,7 +753,9 @@ def construct(match, children, context): raise ParseError.from_node( match.node, context, "layout shape/stride rank mismatch" ) - return runtime.Layout(shape=shape, strides=strides) + return PlacedLayout( + shape=shape, layout=runtime.Layout(shape=shape, strides=strides) + ) referenced_ids = {id(entry[0]) for entry in (*splits, *states)} if context.function is None: raise ParseError.from_node( @@ -786,16 +805,16 @@ def claim(source, source_axis: int) -> int: raise ParseError.from_node(match.node, context, str(error)) from error if strides is not None and len(canonical.layout.shape) != len(strides): raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") - return runtime.ShardLayout( - layout=runtime.Layout(shape=canonical.layout.shape, strides=strides), - attrs=canonical.attrs, - mesh=canonical.mesh, + return PlacedLayout( + shape=shape, + layout=runtime.ShardLayout( + layout=runtime.Layout(shape=canonical.layout.shape, strides=strides), + attrs=canonical.attrs, + mesh=canonical.mesh, + ), ) - RULES: ClassVar[tuple[AstRule[Any], ...]] = ( - LayoutShapeRule(), - LayoutPositionRule(), - ) + RULES: ClassVar[tuple[AstRule[Any], ...]] = (PlacementAnswerRule(),) class LayoutPattern(ElementPattern): @@ -948,36 +967,11 @@ def _states_a_layout_slot(elts: object, context: "MatchContext") -> bool: return False -class PlacedShapePattern(ElementPattern): - """Placement sugar read by a slot that states a shape rather than a layout. - - Same syntax as `PlacedLayoutPattern`, different answer. That pattern serves - the layout slot, whose question is only where a value goes, so a layout is - the whole answer. A shape slot has no operand to read a shape from, and the - layout's own shape is the divided one, so it needs the extents as written - kept beside the layout instead of recovered from it. - """ - - element_name = "placed_shape" - syntax = LazyPattern(lambda: PlacedLayoutPattern().syntax) - - @staticmethod - def construct(match, children, context): - layout = PlacedLayoutPattern.construct(match, children, context) - rank = match.captures["rank"] - return PlacedLayout( - shape=tuple(children[f"extent_{axis}"] for axis in range(rank)), - layout=layout, - ) - - RULES: ClassVar[tuple[AstRule[Any], ...]] = (PlacedShapeRule(),) - - class TensorShapeLayoutPattern(ElementPattern): element_name = "tensor_shape_layout" syntax = LazyPattern( lambda: ChoicePattern( - PlacedShapePattern(), + PlacedLayoutPattern(), ShapePattern(), ) ) @@ -1126,6 +1120,8 @@ def construct(match, children, context): shape, layout = placed, None else: shape, layout = children["shape"], children["layout"] + if isinstance(layout, PlacedLayout): + layout = layout.layout storage = children.get("storage") return runtime.TensorType( shape=shape, @@ -2538,7 +2534,9 @@ def construct(match, children, context): value for name, value in children.items() if name.startswith("input_") ) attrs = { - name.removeprefix("attr_"): value + name.removeprefix("attr_"): ( + value.layout if isinstance(value, PlacedLayout) else value + ) for name, value in children.items() if name.startswith("attr_") } From 925ee109408b2812b6629101ac7a131e1814ba0f Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 10:49:44 +0800 Subject: [PATCH 12/21] refactor(inspection): finish render golden review fixes --- docs/spec/inspection.md | 114 +- docs/spec/parser.md | 32 +- src/tilefoundry/inspection/__init__.py | 3 +- src/tilefoundry/inspection/analysis_report.py | 4 +- src/tilefoundry/inspection/dot.py | 210 ++- src/tilefoundry/inspection/print_context.py | 176 ++- src/tilefoundry/inspection/printer_base.py | 367 ++++- src/tilefoundry/inspection/python_printer.py | 1212 +++++------------ .../inspection/python_type_printer.py | 173 --- src/tilefoundry/inspection/tir_printer.py | 125 +- src/tilefoundry/inspection/viewer/builder.py | 334 +---- src/tilefoundry/inspection/viewer/server.py | 2 +- src/tilefoundry/ir/tir/cuda/nn/mma.py | 1 + src/tilefoundry/parser/ast_pattern.py | 2 + src/tilefoundry/parser/grammar_render.py | 17 +- src/tilefoundry/parser/pattern_nodes.py | 282 ++-- .../type_printer_sugar.analyzed.txt | 63 +- .../inspection/type_printer_sugar.printed.txt | 128 +- tests/fixtures/tir/mma.py | 43 +- tests/fixtures/tir/sync.py | 20 +- tests/inspection/test_python_printer.py | 69 - tests/inspection/test_roundtrip.py | 22 - tests/inspection/test_tir_roundtrip.py | 3 +- 23 files changed, 1422 insertions(+), 1980 deletions(-) delete mode 100644 src/tilefoundry/inspection/python_type_printer.py diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index 8d10ecb7..37fae91a 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -130,8 +130,8 @@ explicit wrapper request. Module output computes imports from names actually used by the rendered program and MUST remain lint-clean. - constraints: - - Module input MUST emit every HIR Function and preserve shared `Mesh` and - `Topology` definitions before the class. Mixed HIR/TIR Modules MUST be + - Module input MUST emit every HIR Function. Meshes used by a function are + written at their scope boundary, and mixed HIR/TIR Modules MUST be rejected. The printer MUST emit a Module's whole tree: each nested Module prints as a @@ -143,9 +143,9 @@ Target prints as the `@module(target=...)` argument and a declared hierarchy as the `@module(topologies=...)` argument ([parser §3](./parser.md#3-implementation-overview)). Every dimension referenced only by a declared topology expression MUST still -be emitted in the dimension prelude. A topology `ShapeDim` MUST use the same -DSL expression text as tensor and Mesh geometry, including public constructors -such as `ceildiv`, so importing restores the same expression tree. +be emitted in the file header. A topology `ShapeDim` MUST use the same DSL +expression text as tensor and Mesh geometry, including public constructors such +as `ceildiv`, so importing restores the same expression tree. - constraints: - The decorator MUST print in its called form, `@module()` included. A bare @@ -180,51 +180,39 @@ DSL text forms for tensor / layout / shard annotations are owned by round-trip without losing mesh / layout / storage information; otherwise it falls back to the verbose `ShardLayout(...)`. -A `ShardLayout` over a plain `Layout` whose `Mesh` has named axes and a -prelude name ([§2.5](#25-mesh-name-map)) MUST use the placement sugar of -[parser §2.1](./parser.md#21-syntax), in both type slots and op-attribute -slots. That sugar states the layout's own dimensions with each `Split` written -on the dimension it divides, adds the stride tuple when the strides are not -C-order over those dimensions, and states the remaining mesh axes in a -`{axis @ ...}` set. Because the parser reads an unstated mesh axis as -`Broadcast`, the set carries every `Partial` and carries `Broadcast` only when -no `Split` or `Partial` would otherwise name the mesh. - -Printer output supports two modes derived from the same pretty-print core: - -- `canonical` — round-trippable text used by `as_script()`, pass - dumps, and viewer detail `code` blocks: the `Tensor[...]` form of - [parser §2.1](./parser.md#21-syntax) (storage as the string - slot, `gmem` omitted). -- `compact` — abbreviated, **display-only / non-round-trip** text for - summaries / labels: `dtype[shape] {value-state?} @storage`. It inlines - what it can (a split into the shape, a `Partial` into the `{...}` - suffix) and falls back to the canonical form when a layout cannot be - rendered compactly. - -Both modes MUST agree on semantics; only the level of detail differs. +A `ShardLayout` over a plain `Layout` whose `Mesh` has named axes MUST use the +placement sugar of [parser §2.1](./parser.md#21-syntax), in both type slots and +op-attribute slots, when every mesh identifier in that sugar has an explicit +scope binding. A binding is either a `with as ` region or the +function's own execution domain. A mesh merely restated in another expression +is not a binding. Without such a binding, the printer MUST use the verbose +`ShardLayout(...)` form rather than inventing a name. + +Placement sugar states the layout's own dimensions with each `Split` written on +the dimension it divides, adds the stride tuple whenever the layout has one, +and states the remaining mesh axes in a `{axis @ ...}` set. Because the parser +reads an unstated mesh axis as `Broadcast`, the set carries every `Partial` and +carries `Broadcast` only when no `Split` or `Partial` would otherwise name the +mesh. + +The printer has one type-text surface: the canonical `Tensor[...]` form of +[parser §2.1](./parser.md#21-syntax), used by `as_script()`, comments, DOT, +the viewer graph, and viewer detail panels. Comments may flatten this text onto +one physical line, but MUST NOT change its syntax or semantics. The same +printer visitor supplies type, dimension, layout, and shard-attribute text; +there is no separate compact, display-only type language. + The meaning of `Split` / `Partial` / `Broadcast` is owned by -[shard](./shard.md); these forms define only render syntax. - -The same-line type annotation `show_types` appends is the `canonical` form on -one physical line, rendered through the same mesh name map -([§2.5](#25-mesh-name-map)) as the signature and the prelude: an annotated -layout MUST name the hoisted mesh rather than restate it, and a `Tuple[...]` -annotation MUST name it in every field. The verbose `ShardLayout(...)` fallback -is unchanged — a mesh with no named axes, or a layout the sugar cannot express, -still renders verbose, so no annotation loses information. The annotation is -**display-only** ([§2.7](#27-round-trip-contract)); what round-trips is the -emitted code, not its comments. - -Canonical DType text is the descriptor's `name`. Tensor annotations and DType -op attributes MUST emit that name as a quoted DSL string. Compact labels MAY -omit the quotes, but MUST NOT use the descriptor's raw `repr()`. +[shard](./shard.md); these forms define only render syntax. Canonical DType text +is the descriptor's `name`. Tensor annotations and DType op attributes MUST +emit that name as a quoted DSL string. ### 2.4 Pretty-print / debug display contract -Pretty print is the core presentation layer. Sugar, debug dumps, and -viewer type/value text reuse the same DSL text forms in [§2.3](#23-dsl-text-forms). That keeps -round-trippable source, labels, and detail panes semantically aligned. +Pretty print is the core presentation layer. Sugar, debug dumps, DOT, and +viewer type/value text reuse the same canonical DSL text forms in +[§2.3](#23-dsl-text-forms). That keeps round-trippable source, labels, and +detail panes semantically aligned. - op attributes that are `DType`, `TensorType`, `Layout`, or `ShardLayout` are rendered through the [§2.3](#23-dsl-text-forms) printer; `DType` uses its canonical name and these @@ -236,18 +224,20 @@ printing (for example, choosing stable mesh names across a whole function) must use an explicit pretty-printer API rather than relying on no-argument `repr()`. -### 2.5 Mesh name map +### 2.5 Mesh bindings and file header -The printer collects unique `Mesh` objects from all `ShardLayout` -references in the function (params, return type, body `Reshard` ops) -and assigns variable names from the first declared topology's name. Mesh -definitions are emitted in the module prelude / standalone header. +The print context records imports and declarations while the printer visits the +program. After the body has been visited, the context emits the file header; +it MUST contain only imports and `DimVar` declarations reached by the output. +There is no module-level mesh hoist or global mesh name map. -Two `Mesh` values with the same printed descriptor MUST share one name and one -prelude definition: a composed mesh is rebuilt at each use site, so naming its -copies apart would claim the value's parts are placed on different meshes. A -mesh the prelude does not define MUST NOT be named by a printed type; it is -restated in full there instead. +A mesh identifier in placement sugar MUST be a lexical binding visible at the +point represented by the text: a `with as ` region, or the +function's own execution domain. The function execution domain is represented +as `@func(mesh=...)`, so its name is available to signature annotations. A +nested region writes its mesh expression at the `with` boundary and uses that +binding in its body. A named constant or an out-of-scope mesh does not create a +binding; types that refer to one use the verbose form. ### 2.6 Specialization printing @@ -477,14 +467,10 @@ graph; an id that was collapsed away returns 404. elements); `Tuple` bundles its elements. Op attributes that are constants / types render through the [§2.4](#24-pretty-print--debug-display-contract) pretty-print, never raw `repr`. -- **Type text.** Graph labels use the [§2.3](#23-dsl-text-forms) **compact** pretty mode - (`bf16[4 @ trd.l, 64] {trd.t @ P("sum")} @smem`) with inline split / - DimVar / storage colour; the detail panel uses the [§2.3](#23-dsl-text-forms) **canonical** - mode (`Tensor[(4, 64), "f32", ((4 @ trd.l, 64), {trd.t @ P("sum")}), - "smem"]`). `Reshard` / layout attrs render through the same core (never - raw `repr`). DimVar is a single token-class colour; - storage classes draw from an ordered pool, and an unknown memory level - hashes stably into the pool's spare slots rather than going colourless. +- **Type text.** Graph labels and detail panels use the canonical type text from + [§2.3](#23-dsl-text-forms), including `Tensor[...]` and placement sugar when + a scope binding is available. `Reshard` / layout attrs render through the + same visitor (never raw `repr`). ### 3.4 Interaction contract diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 451dede6..9fa7dfe5 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -94,23 +94,16 @@ dim-expr ::= integer-literal | dim-expr ('+' | '-' | '*' | '//' | '%') dim-expr | (identifier | primary '.' identifier) '(' (dim-expr (',' dim-expr)*)? ')' -placed-layout ::= '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) - | dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' dim-expr)*)? ')' - ',' '{' mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')') (',' - mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')'))* '}' ')' - | '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' - mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '(' (dim-expr (',' - dim-expr)*)? ')' ')' - | '(' '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' - mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' ',' '{' mesh-axis '@' - ('B' '(' ')' | 'P' '(' string-literal ')') (',' mesh-axis '@' ('B' '(' - ')' | 'P' '(' string-literal ')'))* '}' ')' - | '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | - dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr))*)? ')' +layout-dims ::= '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | + dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr))*)? ')' +layout-strides ::= '(' (dim-expr (',' dim-expr)*)? ')' +value-states ::= '{' mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')') (',' + mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')'))* '}' +placed-layout ::= '(' layout-dims ',' layout-strides ',' value-states ')' + | '(' layout-dims ',' layout-strides ')' + | '(' layout-dims ',' value-states ')' + | layout-dims shape ::= '(' (dim-expr (',' dim-expr)*)? ')' | identifier | primary '.' identifier @@ -266,7 +259,12 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | module | module_function | ModuleFunctionRegistrationRule | A validated module function must be recorded in declaration order. | src/tilefoundry/parser/ast_pattern.py | | module | module_function | ModuleFunctionValidationRule | A module function must satisfy its root, variant, or converter role before mutation. | src/tilefoundry/parser/ast_pattern.py | | op_call | expression, slice_endpoint, subscript_index | CallVariadicInputFormRule | A variadic call must use one explicit list, tuple, or supported static list comprehension. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | LayoutStrideRankRule | A stated stride tuple must have the rank of the layout it addresses. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | MeshAxisBoundOnceRule | A placement binds each mesh axis at most once. | src/tilefoundry/parser/pattern_nodes.py | | placed_layout | tensor_optional_slot, tensor_shape | PlacementAnswerRule | Placement sugar states both the shape as written and the layout it implies. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementConstructionRule | A placement must construct a valid shard layout. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementLevelRule | A placement's meshes cannot name the same topology level. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementMeshResolutionRule | A placement's mesh must be an active scope or resolvable from its bindings. | src/tilefoundry/parser/pattern_nodes.py | | shape | tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | | storage | tensor_optional_slot | StorageValueRule | Storage must resolve to a StorageKind. | src/tilefoundry/parser/ast_pattern.py | | tensor | annotation, expression, slice_endpoint, subscript_index, type_annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | diff --git a/src/tilefoundry/inspection/__init__.py b/src/tilefoundry/inspection/__init__.py index e48fef17..1f61a51f 100644 --- a/src/tilefoundry/inspection/__init__.py +++ b/src/tilefoundry/inspection/__init__.py @@ -2,7 +2,6 @@ from .print_context import HirPrintContext, PrintContext, TirPrintContext from .printer_base import PythonPrinter from .python_printer import PythonPrintOptions, as_script, hir_function_to_python, module_to_python -from .python_type_printer import PythonTypePrinter from .tir_printer import ( TirPrinter, register_tir_printer, @@ -21,5 +20,5 @@ "tir_module_to_python", "TirPrinter", "register_tir_printer", "Viewer", - "PrintContext", "HirPrintContext", "TirPrintContext", "PythonPrinter", "PythonTypePrinter", + "PrintContext", "HirPrintContext", "TirPrintContext", "PythonPrinter", ] diff --git a/src/tilefoundry/inspection/analysis_report.py b/src/tilefoundry/inspection/analysis_report.py index cc8f3757..7dfc8a00 100644 --- a/src/tilefoundry/inspection/analysis_report.py +++ b/src/tilefoundry/inspection/analysis_report.py @@ -20,7 +20,7 @@ from tilefoundry.analysis.report import ( selected_types as _selected_types, ) -from tilefoundry.inspection.python_printer import PythonPrintOptions, _render_hir_function +from tilefoundry.inspection.python_printer import HirPrinter, PythonPrintOptions from tilefoundry.inspection.values import ( AdvisorySummary, MemorySummary, @@ -54,7 +54,7 @@ def render_analysis( ) -> AnalysisRendering: """Render one result once for both annotated source and report data.""" selected_types_ = selected_types(result) - rendered = _render_hir_function( + rendered = HirPrinter().render( result.function, options=PythonPrintOptions( show_types=True, diff --git a/src/tilefoundry/inspection/dot.py b/src/tilefoundry/inspection/dot.py index b7c3eedb..91d30f67 100644 --- a/src/tilefoundry/inspection/dot.py +++ b/src/tilefoundry/inspection/dot.py @@ -1,150 +1,146 @@ -"""Serialize SSA HIR functions as Graphviz DOT. - -Variables, calls, and constants receive numbered nodes. Type and shard-layout -labels reuse canonical printer renderers so DOT, Python output, and the viewer -remain consistent. -See [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). -""" +"""Serialize SSA HIR functions as Graphviz DOT.""" from __future__ import annotations from tilefoundry.ir.core import Call, Constant, Var, binding_name from tilefoundry.ir.core.module import Module from tilefoundry.ir.hir.function import Function as HirFunction +from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.hir.sharding.reshard import Reshard from tilefoundry.ir.types import TensorType from tilefoundry.ir.visitor import ExprWalker -from .python_printer import _collect_meshes, _mesh_name_map, _op_display_name, _tensor_annotation +from .print_context import HirPrintContext +from .printer_base import PythonPrinter + +def _op_display_name(target) -> str: + cls = type(target).__name__ + for suffix in ("Op", "Expr", "Stmt"): + if cls.endswith(suffix) and cls != suffix: + cls = cls[: -len(suffix)] + return cls -def _type_lines(ty, mesh_name_map: dict[int, str]) -> list[str]: - """Return canonical type-annotation lines for a node label. - Split multiline shard-layout fallback text into separate label lines. - See [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). - """ - text = _tensor_annotation(ty, mesh_name_map=mesh_name_map) if isinstance(ty, TensorType) else str(ty) - return text.split("\n") +def _type_lines(ty, printer: PythonPrinter, ctx: HirPrintContext) -> list[str]: + if not isinstance(ty, TensorType): + return [str(ty)] + with printer.type_surface(): + return printer.visit(ty, ctx).split("\n") def _escape_dot(s: str) -> str: - """Escape a string for safe inclusion in a DOT label.""" return s.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") def hir_function_to_dot(fn: HirFunction) -> str: - """Convert a hir.Function to a DOT digraph string. - - Args: - fn: The HIR function to visualize. - - Returns: - A Graphviz DOT format string. - """ - type_meshes, scope_meshes = _collect_meshes(fn, include_node_types=True) - mesh_map = _mesh_name_map({**type_meshes, **scope_meshes}) - lines = [f"digraph {fn.name} {{", ' rankdir=TB;', - ' node [shape=box, style=filled, fillcolor="#f0f0f0"];', - ' edge [fontsize=10, fontcolor="#555555"];', ''] - _counter = [0] - _ids = {} - def _id(node): + """Convert a HIR function to a DOT digraph using canonical type text.""" + printer = PythonPrinter() + ctx = HirPrintContext() + root_scope = fn.body if isinstance(fn.body, MeshRegion) else None + if root_scope is not None: + ctx.push_mesh(root_scope.mesh, "mesh") + + lines = [ + f"digraph {fn.name} {{", + " rankdir=TB;", + ' node [shape=box, style=filled, fillcolor="#f0f0f0"];', + ' edge [fontsize=10, fontcolor="#555555"];', + "", + ] + counter = [0] + ids: dict[int, str] = {} + + def node_id(node): key = id(node) - if key not in _ids: - _ids[key] = f"n{_counter[0]}" - _counter[0] += 1 - return _ids[key] - - def _emit_node(nid, label_lines, fill="#f0f0f0"): - escaped = [_escape_dot(ln) for ln in label_lines] - label = "\\n".join(escaped) + if key not in ids: + ids[key] = f"n{counter[0]}" + counter[0] += 1 + return ids[key] + + def emit_node(nid, label_lines, fill="#f0f0f0"): + label = "\\n".join(_escape_dot(line) for line in label_lines) lines.append(f' {nid} [label="{label}", fillcolor="{fill}"];') - def _emit_edge(src_id, dst_id, label=""): - if label: - lines.append(f' {src_id} -> {dst_id} [label="{label}"];') - else: - lines.append(f' {src_id} -> {dst_id};') - - VAR_FILL = "#d4e6f1" - CONST_FILL = "#f9e79f" - CALL_FILL = "#d5f5e3" - SHARDING_FILL = "#e8daef" - - class _DotWalker(ExprWalker[None]): - def _emit_leaf(self, expr, label_lines, fill) -> None: - _emit_node(_id(expr), label_lines, fill=fill) - - def visit_Var(self, expr: Var, ctx=None) -> None: - self._emit_leaf( - expr, - [f"Var: {expr.name}", *_type_lines(expr.type, mesh_map)], - VAR_FILL, - ) - - def visit_Constant(self, expr: Constant, ctx=None) -> None: + def emit_edge(src_id, dst_id, label=""): + suffix = f' [label="{label}"]' if label else "" + lines.append(f" {src_id} -> {dst_id}{suffix};") + + var_fill, const_fill = "#d4e6f1", "#f9e79f" + call_fill, shard_fill = "#d5f5e3", "#e8daef" + + class DotWalker(ExprWalker[None]): + def _emit_leaf(self, expr, label_lines, fill): + emit_node(node_id(expr), label_lines, fill=fill) + + def visit_Var(self, expr: Var, current=None) -> None: + self._emit_leaf(expr, [f"Var: {expr.name}", *_type_lines(expr.type, printer, ctx)], var_fill) + + def visit_Constant(self, expr: Constant, current=None) -> None: value = f"{expr.value:.6g}" if isinstance(expr.value, float) else str(expr.value) - self._emit_leaf( - expr, - [f"Const: {value}", *_type_lines(expr.type, mesh_map)], - CONST_FILL, - ) - - def visit_Call(self, expr: Call, ctx=None) -> None: - nid = _id(expr) + self._emit_leaf(expr, [f"Const: {value}", *_type_lines(expr.type, printer, ctx)], const_fill) + + def visit_Call(self, expr: Call, current=None) -> None: + nid = node_id(expr) target = expr.target if isinstance(target, Reshard): name = binding_name(expr) header = f"{name}\\nReshard" if name else "Reshard" - _emit_node(nid, [header, *_type_lines(expr.type, mesh_map)], fill=SHARDING_FILL) + emit_node(nid, [header, *_type_lines(expr.type, printer, ctx)], fill=shard_fill) + else: + name = binding_name(expr) + op_label = _op_display_name(target) + header = f"{name}\\n{op_label}" if name else op_label + emit_node(nid, [header, *_type_lines(expr.type, printer, ctx)], fill=call_fill) + for index, arg in enumerate(expr.args): + self.visit(arg, current) + emit_edge(node_id(arg), nid, f"arg[{index}]" if len(expr.args) > 1 else "") + + def visit_MeshRegion(self, expr: MeshRegion, current=None) -> None: + alias = ctx.scope_name(expr.mesh) + ctx.push_mesh(expr.mesh, alias) + try: + self.visit(expr.body, current) for arg in expr.args: - self.visit(arg, ctx) - _emit_edge(_id(arg), nid) - return - - op_label = _op_display_name(target) - name = binding_name(expr) - header = f"{name}\\n{op_label}" if name else op_label - _emit_node(nid, [header, *_type_lines(expr.type, mesh_map)], fill=CALL_FILL) - for i, arg in enumerate(expr.args): - self.visit(arg, ctx) - edge_label = f"arg[{i}]" if len(expr.args) > 1 else "" - _emit_edge(_id(arg), nid, edge_label) - - def visit_Tuple(self, expr, ctx=None) -> None: + self.visit(arg, current) + finally: + ctx.pop_mesh() + + def visit_Tuple(self, expr, current=None) -> None: self._emit_leaf(expr, [type(expr).__name__], "#ffffff") - def visit_LoopRegion(self, expr, ctx=None) -> None: + def visit_LoopRegion(self, expr, current=None) -> None: self._emit_leaf(expr, [type(expr).__name__], "#ffffff") - def visit_ShapeOf(self, expr, ctx=None) -> None: - _emit_node(_id(expr), ["ShapeOf"], fill="#ffffff") + def visit_ShapeOf(self, expr, current=None) -> None: + emit_node(node_id(expr), ["ShapeOf"], fill="#ffffff") - def default_visit(self, expr, ctx=None) -> None: - _emit_node(_id(expr), [type(expr).__name__], fill="#ffffff") + def default_visit(self, expr, current=None) -> None: + emit_node(node_id(expr), [type(expr).__name__], fill="#ffffff") - walker = _DotWalker() + walker = DotWalker() walker.visit(fn.body) for param in fn.params: walker.visit(param) - - - lines.append("") - lines.append(' subgraph cluster_legend {') - lines.append(' label="Legend";') - lines.append(' style=dashed;') - lines.append(' fontsize=11;') - lines.append(' l_var [label="Var/Param", fillcolor="#d4e6f1", shape=box, style=filled];') - lines.append(' l_const [label="Constant", fillcolor="#f9e79f", shape=box, style=filled];') - lines.append(' l_call [label="Op", fillcolor="#d5f5e3", shape=box, style=filled];') - lines.append(' l_shard [label="Reshard", fillcolor="#e8daef", shape=box, style=filled];') - lines.append(" }") - lines.append("}") + if root_scope is not None: + ctx.pop_mesh() + + lines.extend([ + "", + ' subgraph cluster_legend {', + ' label="Legend";', + " style=dashed;", + ' fontsize=11;', + ' l_var [label="Var/Param", fillcolor="#d4e6f1", shape=box, style=filled];', + ' l_const [label="Constant", fillcolor="#f9e79f", shape=box, style=filled];', + ' l_call [label="Op", fillcolor="#d5f5e3", shape=box, style=filled];', + ' l_shard [label="Reshard", fillcolor="#e8daef", shape=box, style=filled];', + " }", + "}", + ]) return "\n".join(lines) + "\n" def module_entry_to_dot(module: Module) -> str: """Convert a Module's entry function to DOT.""" - fn = module.entry_function() - return hir_function_to_dot(fn) + return hir_function_to_dot(module.entry_function()) diff --git a/src/tilefoundry/inspection/print_context.py b/src/tilefoundry/inspection/print_context.py index e0f4f275..bac8c8ca 100644 --- a/src/tilefoundry/inspection/print_context.py +++ b/src/tilefoundry/inspection/print_context.py @@ -1,13 +1,23 @@ -"""Shared printer state and policy context.""" +"""Shared state accumulated while canonical Python source is rendered.""" from __future__ import annotations -from tilefoundry.utils.python_source import PythonExpr +from math import prod + +from tilefoundry.ir.types.shard.int_tuple import flatten +from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout +from tilefoundry.ir.types.shard.mesh import Mesh, topology_axes +from tilefoundry.utils.python_source import PythonExpr, _merge_imports class PrintContext: + """Imports, symbolic declarations, and lexical mesh bindings for one file.""" + def __init__(self) -> None: self.imports: set[str] = set() + self._dim_declarations: dict[str, tuple[object, str]] = {} + self._mesh_bindings: list[tuple[Mesh, str]] = [] + self._used_scope_names: set[str] = set() def use(self, rendered: PythonExpr | str) -> str: if isinstance(rendered, PythonExpr): @@ -15,32 +25,154 @@ def use(self, rendered: PythonExpr | str) -> str: return rendered.text return rendered - def mesh_alias(self, mesh) -> str | None: - return None + def declare_dim( + self, + name: str, + var, + *, + import_statement: str = "from tilefoundry.ir.types.dim import DimVar", + ) -> None: + self.imports.add(import_statement) + self._dim_declarations.setdefault(name, (var, "DimVar")) + def header(self) -> list[str]: + """Render only imports and declarations reached while rendering the body.""" + imports = list(_merge_imports(tuple(self.imports))) + imports = [ + f"{line} # noqa: F401, F403" if line.endswith(" import *") else line + for line in imports + ] + lines = ["from __future__ import annotations", "", *imports, ""] + if self._dim_declarations: + lines.extend( + f'{name} = {constructor}("{var.name}", {var.lo}, {var.hi})' + for name, (var, constructor) in self._dim_declarations.items() + ) + lines.append("") + return lines -class HirPrintContext(PrintContext): - def __init__(self, mesh_name_map: dict[int, str] | None = None) -> None: - super().__init__() - self.mesh_name_map = mesh_name_map or {} + def scope_name(self, mesh: Mesh, preferred: str | None = None) -> str: + base = preferred or (mesh.topologies[0].name if mesh.topologies else "mesh") + base = base if base.isidentifier() else "mesh" + name = base + suffix = 2 + while name in self._used_scope_names: + name = f"{base}_{suffix}" + suffix += 1 + self._used_scope_names.add(name) + return name - def mesh_alias(self, mesh) -> str | None: - return self.mesh_name_map.get(id(mesh)) + def push_mesh(self, mesh: Mesh, name: str) -> None: + self._mesh_bindings.append((mesh, name)) + self._used_scope_names.add(name) + def pop_mesh(self) -> None: + self._mesh_bindings.pop() -class TirPrintContext(PrintContext): - def __init__(self) -> None: - super().__init__() - self._mesh_aliases: list[dict[int, str]] = [] + def mesh_alias(self, mesh: Mesh) -> str | None: + for bound, name in reversed(self._mesh_bindings): + if bound is mesh: + return name + return None - def push_mesh(self, mesh, name: str) -> None: - self._mesh_aliases.append({id(mesh): name}) + @staticmethod + def _axis_levels(mesh: Mesh) -> tuple[str, ...]: + levels = [""] * len(flatten(mesh.layout.shape)) + for topology, axes in zip(mesh.topologies, topology_axes(mesh), strict=True): + for axis in axes: + levels[axis] = topology.name + return tuple(levels) - def pop_mesh(self) -> None: - self._mesh_aliases.pop() + def mesh_axis_alias(self, mesh: Mesh, axis: int) -> str | None: + """Name one mesh axis through an active scope binding, if one dominates it.""" + names = mesh.names + if axis >= len(names): + return None + target_name = names[axis] + target_levels = self._axis_levels(mesh) + target_level = target_levels[axis] + target_topology = next( + (topology for topology in mesh.topologies if topology.name == target_level), None + ) + for bound, alias in reversed(self._mesh_bindings): + if not bound.names or target_name not in bound.names: + continue + bound_levels = self._axis_levels(bound) + for bound_axis, bound_name in enumerate(bound.names): + if bound_name != target_name or bound_levels[bound_axis] != target_level: + continue + bound_topology = next( + topology for topology in bound.topologies if topology.name == target_level + ) + if target_topology == bound_topology: + return f"{alias}.{bound_name}" + return None - def mesh_alias(self, mesh) -> str | None: - for aliases in reversed(self._mesh_aliases): - if id(mesh) in aliases: - return aliases[id(mesh)] + def mesh_slice(self, mesh: Mesh) -> str | None: + """Recover ``binding[start:stop]`` for a sliced active mesh.""" + if not ( + isinstance(mesh.layout, ComposedLayout) + and mesh.layout.inner is None + and isinstance(mesh.layout.outer, Layout) + ): + return None + for parent, alias in reversed(self._mesh_bindings): + text = self._slice_from_parent(parent, mesh, alias) + if text is not None: + return text return None + + @staticmethod + def _slice_from_parent(parent: Mesh, child: Mesh, alias: str) -> str | None: + if not isinstance(parent.layout, Layout): + return None + if parent.topologies != child.topologies or parent.names != child.names: + return None + parent_shape = flatten(parent.layout.shape) + parent_strides = flatten(parent.layout.strides) + child_shape = flatten(child.layout.outer.shape) + child_strides = flatten(child.layout.outer.strides) + if ( + len(parent_shape) != len(child_shape) + or parent_strides != child_strides + or any(not isinstance(item, int) for item in (*parent_shape, *child_shape, *parent_strides)) + or any(size < 1 or size > extent for size, extent in zip(child_shape, parent_shape)) + ): + return None + + remaining = child.layout.offset + starts = [0] * len(parent_shape) + for axis in sorted(range(len(parent_shape)), key=lambda item: parent_strides[item], reverse=True): + stride = parent_strides[axis] + maximum = parent_shape[axis] - child_shape[axis] + start = min(maximum, remaining // stride) if stride else 0 + starts[axis] = start + remaining -= start * stride + if remaining != 0: + return None + if sum(start * stride for start, stride in zip(starts, parent_strides)) != child.layout.offset: + return None + if prod(child_shape) > prod(parent_shape): + return None + + pieces: list[str] = [] + for start, size, extent in zip(starts, child_shape, parent_shape): + if start == 0 and size == extent: + pieces.append(":") + continue + stop = start + size + pieces.append(f"{'' if start == 0 else start}:{'' if stop == extent else stop}") + while len(pieces) > 1 and pieces[-1] == ":": + pieces.pop() + return f"{alias}[{', '.join(pieces)}]" + + +class HirPrintContext(PrintContext): + pass + + +class TirPrintContext(PrintContext): + pass + + +__all__ = ["PrintContext", "HirPrintContext", "TirPrintContext"] diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index e8c0d1e8..149c0cbc 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -1,48 +1,351 @@ -"""Common expression-printer base classes. - -Lazy imports avoid a cycle between the shared base and the legacy HIR module. -""" - -# ruff: noqa: PLC0415 +"""Canonical Python visitor shared by the HIR and TIR printers.""" from __future__ import annotations import enum +from contextlib import contextmanager +from tilefoundry.ir.core import Call, Constant, Tuple, Var from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard.layout import LayoutBase +from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType +from tilefoundry.ir.types.dim import ( + DimAdd, + DimConst, + DimFloorDiv, + DimMax, + DimMin, + DimMod, + DimMul, + DimSub, + DimVar, +) +from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.visitor import ExprFunctor +from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, ShardLayout, Split +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.visitor import ExprFunctor, TypeFunctor from tilefoundry.target import Target from tilefoundry.utils.python_source import PythonExpr -from .python_type_printer import PythonTypePrinter +_DIM_INFIX_OPS: dict[type, str] = { + DimAdd: "+", + DimSub: "-", + DimMul: "*", + DimFloorDiv: "//", + DimMod: "%", +} +_DIM_FUNC_OPS: dict[type, str] = { + DimMin: "min", + DimMax: "max", +} -class PythonPrinter(PythonTypePrinter, ExprFunctor[str]): - """Shared expression/value visitor base for HIR and TIR printers. - Inheriting the type printer rather than owning one keeps a single - ``visit``: an expression printer emits ``expr.type`` through the same - dispatch that emits the type's own children. - """ +class PythonPrinter(ExprFunctor[str], TypeFunctor[str]): + """Render the Python DSL with one dispatch root for expressions and types.""" def __init__(self) -> None: - PythonTypePrinter.__init__(self) ExprFunctor.__init__(self) + self._indent = "" + self._tensor_head = "Tensor" + self._nested_dim = False + + def visit(self, value, ctx=None): # type: ignore[override] + """Dispatch every implemented functor family by the concrete node name.""" + method = getattr(self, f"visit_{type(value).__name__}", None) + if method is not None: + return method(value, ctx) + return self.default_visit(value, ctx) + + def default_visit(self, value, ctx=None) -> str: + if isinstance(value, bool): + return repr(value) + if isinstance(value, (int, float)): + return str(value) + raise NotImplementedError(f"no Python visit routine for {type(value).__name__}") + + @contextmanager + def type_surface(self, *, indent: str | None = None, const: bool = False): + """Carry statement indentation and parameter const-ness through type visits.""" + previous = (self._indent, self._tensor_head) + if indent is not None: + self._indent = indent + self._tensor_head = "ConstTensor" if const else "Tensor" + try: + yield + finally: + self._indent, self._tensor_head = previous + + @contextmanager + def nested_dim(self, nested: bool): + previous = self._nested_dim + self._nested_dim = nested + try: + yield + finally: + self._nested_dim = previous + + def dim_entry(self, value, ctx=None, *, nested: bool = False) -> str: + with self.nested_dim(nested): + return self.visit(value, ctx) + + def shape_tuple(self, shape: tuple, ctx=None) -> str: + values = tuple(self.dim_entry(entry, ctx) for entry in shape) + return f"({values[0]},)" if len(values) == 1 else "(" + ", ".join(values) + ")" + + def dtype_str(self, dtype: DType, ctx=None) -> str: + return dtype.name + + def visit_DimVar(self, value: DimVar, ctx=None) -> str: + if ctx is not None: + ctx.declare_dim(value.name, value) + return value.name + + def visit_Var(self, value: Var, ctx=None) -> str: + return value.name + + def visit_Constant(self, value: Constant, ctx=None) -> str: + return repr(value.value) + + def visit_Tuple(self, value: Tuple, ctx=None) -> str: + rendered = ", ".join(self.visit(item, ctx) for item in value.elements) + return f"({rendered}{',' if len(value.elements) == 1 else ''})" + + def visit_Call(self, value: Call, ctx=None) -> str: + ceildiv_args = self._ceildiv_args(value) + if ceildiv_args is not None: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.dim import ceildiv",), "ceildiv")) + left, right = ceildiv_args + return f"ceildiv({self.dim_entry(left, ctx)}, {self.dim_entry(right, ctx)})" + target = value.target + if isinstance(target, DimConst): + return str(target.value) + for op_type, symbol in _DIM_INFIX_OPS.items(): + if isinstance(target, op_type): + left, right = value.args + rendered = ( + f"{self.dim_entry(left, ctx, nested=True)} {symbol} " + f"{self.dim_entry(right, ctx, nested=True)}" + ) + return f"({rendered})" if self._nested_dim else rendered + for op_type, name in _DIM_FUNC_OPS.items(): + if isinstance(target, op_type): + args = ", ".join(self.dim_entry(arg, ctx) for arg in value.args) + return f"{name}({args})" + return self.visit_program_call(value, ctx) + + def visit_program_call(self, value: Call, ctx=None) -> str: + raise NotImplementedError(f"{type(self).__name__} cannot render program calls") + + @staticmethod + def _ceildiv_args(value: Call) -> tuple[object, object] | None: + """Recover the public constructor from ceildiv's canonical arithmetic tree.""" + if not isinstance(value.target, DimFloorDiv) or len(value.args) != 2: + return None + numerator, divisor = value.args + if not ( + isinstance(numerator, Call) + and isinstance(numerator.target, DimSub) + and len(numerator.args) == 2 + and isinstance(numerator.args[1], Constant) + and numerator.args[1].value == 1 + ): + return None + added = numerator.args[0] + if not ( + isinstance(added, Call) + and isinstance(added.target, DimAdd) + and len(added.args) == 2 + and added.args[1] == divisor + ): + return None + return added.args[0], divisor + + def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: + """Render placement sugar only when every mesh axis has a scope binding.""" + layout = value.layout + names = value.mesh.names + if ( + not isinstance(layout, Layout) + or not names + or len(value.attrs) != len(names) + or ctx is None + ): + return None + refs = tuple(ctx.mesh_axis_alias(value.mesh, index) for index in range(len(names))) + if any(ref is None for ref in refs): + return None - def atom_reference(self, value, ctx=None) -> str: + splits: dict[int, list[str]] = {} + partials: list[str] = [] + broadcasts: list[tuple[str, str, Broadcast]] = [] + named_bindings: set[str] = set() + for index, (attr, ref) in enumerate(zip(value.attrs, refs, strict=True)): + assert ref is not None + binding = ref.partition(".")[0] + if isinstance(attr, Split): + if attr.axis >= len(layout.shape): + return None + splits.setdefault(attr.axis, []).append(ref) + named_bindings.add(binding) + elif isinstance(attr, Partial): + partials.append(f'{ref} @ {self.visit(attr, ctx)}') + named_bindings.add(binding) + elif isinstance(attr, Broadcast): + broadcasts.append((binding, ref, attr)) + else: + return None + + states = list(partials) + for binding, ref, attr in broadcasts: + if binding not in named_bindings: + states.append(f"{ref} @ {self.visit(attr, ctx)}") + named_bindings.add(binding) + if not splits and not states: + states = [f"{ref} @ {self.visit(attr, ctx)}" for _binding, ref, attr in broadcasts] + if not splits and not states: + return None + + explicit = layout.strides is not None + if explicit and any( + axis in splits + and self.dim_entry(dim, ctx, nested=True) != self.dim_entry(dim, ctx) + for axis, dim in enumerate(layout.shape) + ): + return None + + dims = [ + ( + f"{self.dim_entry(dim, ctx, nested=True)} " + + " ".join(f"@ {ref}" for ref in splits[axis]) + ) + if axis in splits + else self.dim_entry(dim, ctx) + for axis, dim in enumerate(layout.shape) + ] + dims_text = ", ".join(dims) + ("," if len(dims) == 1 else "") + parts = [f"({dims_text})"] + if explicit: + parts.append(self.shape_tuple(layout.strides, ctx)) + if states: + parts.append("{" + ", ".join(states) + "}") + return parts[0] if len(parts) == 1 else "(" + ", ".join(parts) + ")" + + def visit_TensorType(self, value: TensorType, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr((f"from tilefoundry.dsl import {self._tensor_head}",), self._tensor_head)) + result = ( + f"{self._tensor_head}[" + f'{self.shape_tuple(value.shape, ctx)}, "{self.dtype_str(value.dtype, ctx)}"' + ) + if isinstance(value.layout, ShardLayout): + surface = self.shard_surface(value.layout, ctx) + if surface is not None: + result += f", {surface}" + else: + with self.type_surface(indent=self._indent + " "): + result += f",\n{self._indent}{self.visit(value.layout, ctx)}" + elif value.layout is not None: + result += f", {self.visit(value.layout, ctx)}" + if value.storage is not StorageKind.GMEM: + result += f', "{value.storage.name.lower()}"' + return result + "]" + + def visit_TupleType(self, value: TupleType, ctx=None) -> str: + return f"Tuple[{', '.join(self.visit(field, ctx) for field in value.fields)}]" + + def visit_UnitType(self, value: UnitType, ctx=None) -> str: + return "None" + + def visit_DType(self, value: DType, ctx=None) -> str: + return self.dtype_str(value, ctx) + + def visit_Mesh(self, value: Mesh, ctx=None) -> str: + if ctx is not None: + alias = ctx.mesh_alias(value) + if alias is not None: + return alias + sliced = ctx.mesh_slice(value) + if sliced is not None: + return sliced + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "Mesh")) + topologies = ", ".join( + f'Topology("{topology.name}", {self.dim_entry(topology.size, ctx)})' + for topology in value.topologies + ) + topologies = f"({topologies}{',' if len(value.topologies) == 1 else ''})" + result = f"Mesh({topologies}, {self.visit(value.layout, ctx)}" + if value.names: + result += f", names={tuple(value.names)!r}" + return result + ")" + + def visit_NoneType(self, value: None, ctx=None) -> str: + return "None" + + def visit_Layout(self, value: Layout, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout",), "Layout")) + strides = self.shape_tuple(value.strides, ctx) if value.strides is not None else "None" + return f"Layout({self.shape_tuple(value.shape, ctx)}, {strides})" + + def visit_ComposedLayout(self, value: ComposedLayout, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) + outer, child = self._indent, self._indent + " " + with self.type_surface(indent=child): + inner_text = self.visit(value.inner, ctx) + outer_text = self.visit(value.outer, ctx) + return ( + "ComposedLayout(\n" + f"{child}inner={inner_text},\n" + f"{child}offset={self.dim_entry(value.offset, ctx)},\n" + f"{child}outer={outer_text},\n" + f"{outer})" + ) + + def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: + surface = self.shard_surface(value, ctx) + if surface is not None: + return surface + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) + outer, child = self._indent, self._indent + " " + attrs = ", ".join(self.visit(attr, ctx) for attr in value.attrs) + if len(value.attrs) == 1: + attrs += "," + with self.type_surface(indent=child): + layout_text = self.visit(value.layout, ctx) + mesh_text = self.visit(value.mesh, ctx) + return ( + "ShardLayout(\n" + f"{child}layout={layout_text},\n" + f"{child}attrs=({attrs}),\n" + f"{child}mesh={mesh_text},\n" + f"{outer})" + ) + + def visit_Broadcast(self, value: Broadcast, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) + return "B()" + + def visit_Split(self, value: Split, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) + return f"S({value.axis})" + + def visit_Partial(self, value: Partial, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) + return f'P("{value.reduction}")' + + def atom_reference(self, value: MmaAtom, ctx=None) -> str: return f"T.cuda.mma.atom(op=T.cuda.mma.{value.op.name})" def render_value(self, value, ctx=None, indent: str = "") -> str: - """Render a non-type DSL attribute value and register its imports. - - Type values are not routed here: a printer that holds an ``expr.type`` - calls ``self.visit`` directly, so this stays the surface for the - op-attribute literals that are not part of the type language. - """ + """Render a non-expression attribute through the same visitor when possible.""" if isinstance(value, (TensorType, Mesh, LayoutBase, DType)): with self.type_surface(indent=indent): return self.visit(value, ctx) @@ -55,17 +358,21 @@ def render_value(self, value, ctx=None, indent: str = "") -> str: ctx.use(PythonExpr((f"from {type(value).__module__} import {type(value).__name__}",), "")) return f"{type(value).__name__}.{value.name}" if isinstance(value, Target): - expr = value.to_python() - return ctx.use(expr) if ctx is not None else expr.text - if isinstance(value, (str, int, float, bool, tuple, type(None))): - if isinstance(value, tuple): - vals = ", ".join(self.render_value(v, ctx, indent) for v in value) - return f"({vals}{',' if len(value)==1 else ''})" + rendered = value.to_python() + return ctx.use(rendered) if ctx is not None else rendered.text + if isinstance(value, tuple): + rendered = ", ".join(self.render_value(item, ctx, indent) for item in value) + return f"({rendered}{',' if len(value) == 1 else ''})" + if value is None or isinstance(value, (str, int, float, bool)): return repr(value) raise NotImplementedError(f"no canonical Python form for {type(value).__name__}") def render_pattern(self, pattern: Pattern, ctx=None) -> str: if isinstance(pattern, DimVarRangePat): + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.core.pattern import DimVarRangePat",), "")) return f'DimVarRangePat("{pattern.dim_var}", {pattern.lo}, {pattern.hi})' return repr(pattern) + +__all__ = ["PythonPrinter"] diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index 3efeaf6d..32036653 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -1,8 +1,8 @@ """Canonical Python DSL printer for HIR Functions. Converts a ``hir.Function`` to executable Python source using the -``@func`` DSL. When meshes have named axes, compact sugar annotations -are emitted; otherwise the verbose ``ShardLayout(...)`` form is used. +``@func`` DSL. Placement sugar names only explicit mesh-scope bindings; +otherwise the verbose ``ShardLayout(...)`` form is used. """ from __future__ import annotations @@ -41,7 +41,6 @@ from tilefoundry.ir.hir.sharding.reshard import Reshard from tilefoundry.ir.hir.specialize import ( canonical_specialization_signature, - dim_vars_reached, display_name, origin_of, ) @@ -50,28 +49,14 @@ from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.types import DType, TensorType, TupleType -from tilefoundry.ir.types.dim import ( - DimAdd, - DimConst, - DimFloorDiv, - DimMax, - DimMin, - DimMod, - DimMul, - DimSub, - DimVar, -) -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase -from tilefoundry.ir.types.shard.mesh import Mesh +from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.shard.shard_layout import ( Broadcast, Partial, ShardLayout, Split, - layout_axis_to_tensor_axis, ) -from tilefoundry.ir.types.substitute import dim_vars_by_name -from tilefoundry.ir.visitor import ExprFunctor, expr_children +from tilefoundry.ir.visitor import expr_children from tilefoundry.utils.python_source import PythonExpr from .print_context import HirPrintContext @@ -80,29 +65,184 @@ from .tir_printer import tir_function_to_python, tir_module_to_python from .values import PARTS, render_comment -_DIM_INFIX_OPS: dict[type, str] = { - DimAdd: "+", - DimSub: "-", - DimMul: "*", - DimFloorDiv: "//", - DimMod: "%", -} +class HirPrinter(PythonPrinter): + """Canonical HIR expression, type, and function printer.""" -_DIM_FUNC_OPS: dict[type, str] = { - DimMin: "min", - DimMax: "max", -} + def __init__(self) -> None: + super().__init__() + self._names: dict[int, str] = {} + self._param_alias: dict[int, Expr] = {} + self._child_entries: dict[int, str] = {} + self._moved_window = lambda start, size, stride: None + def print(self, fn: HirFunction, *, options=None) -> str: + return self.render(fn, options=options).source + + def render(self, fn: HirFunction, *, options=None) -> _PythonRendering: + return _render_hir_function(fn, options=options) + + def bind_def( + self, + names: dict[int, str], + param_alias: dict[int, Expr], + child_entries: dict[int, str], + moved_window, + ) -> None: + """Install the value-naming environment for one function definition.""" + self._names = names + self._param_alias = param_alias + self._child_entries = child_entries + self._moved_window = moved_window + + def tuple_reference(self, elements) -> str: + inner = ", ".join( + repr(element.value) + if isinstance(element, Constant) + else self.reference(element) + for element in elements + ) + return f"({inner}{',' if len(elements) == 1 else ''})" -class HirPrinter(PythonPrinter): - """HIR façade retaining the historical canonical rendering entry point.""" + def reference(self, expr: Expr) -> str: + """Return the binding that denotes one value in the current definition.""" + if isinstance(expr, Tuple): + return self.tuple_reference(expr.elements) + if id(expr) in self._param_alias: + return self.reference(self._param_alias[id(expr)]) + projection = _region_projection(expr) + if isinstance(projection, LoopRegion): + return self._names[id(projection.carried_args[expr.target.index])] + if isinstance(expr, LoopRegion): + carried = tuple(self._names[id(carry)] for carry in expr.carried_args) + return carried[0] if len(carried) == 1 else "(" + ", ".join(carried) + ")" + if isinstance(projection, MeshRegion): + return self.reference(projection.body.elements[expr.target.index]) + if isinstance(expr, MeshRegion): + return self.reference(expr.body) + return self._names[id(expr)] - def print(self, fn: HirFunction, *, options=None) -> str: - return _render_hir_function(fn, options=options).source + def _slice_start(self, start, size, stride) -> str: + moved = self._moved_window(start, size, stride) + if moved is None: + return repr(start.value) if isinstance(start, Constant) else self.reference(start) + window, offset = moved + return _moved_window_ref(self.reference(window), offset) - def dim_entry(self, value, ctx=None) -> str: - return shape_entry_str(value) + def visit_program_call(self, expr: Call, ctx=None) -> str: + """Render one HIR call after expression-level dispatch selected it.""" + target = expr.target + args_text = ", ".join(self.reference(arg) for arg in expr.args) + if isinstance(target, Reshard): + if ctx is not None: + ctx.imports.add("from tilefoundry.dsl.tf import *") + layout_kw = "" + if target.layout is not None: + with self.type_surface(indent=self._indent + " "): + layout_kw = ", layout=" + self.visit(target.layout, ctx) + storage = "" + if target.storage is not None: + storage_name = target.storage.name.lower() + if ctx is not None: + ctx.use( + PythonExpr( + (f"from tilefoundry.dsl.storage import {storage_name}",), + storage_name, + ) + ) + storage = f", storage={storage_name}" + return f"reshard({args_text}{layout_kw}{storage})" + if isinstance(target, HirFunction): + binding = _module_callee_binding(target, self._child_entries) + return f"{binding or target.name}({args_text})" + if isinstance(target, Slice): + starts = expr.args[1] + if not isinstance(starts, Tuple): + raise ValueError("canonical_source: Slice starts must be a Tuple") + indexers: list[str] = [] + runtime_starts = False + for axis, (start, size, stride) in enumerate( + zip(starts.elements, target.sizes, target.strides) + ): + if self._moved_window(start, size, stride) is not None: + indexers.append(self._slice_start(start, size, stride)) + continue + dim = expr.args[0].type.shape[axis] + if ( + isinstance(start, Constant) + and start.value == 0 + and size == dim + and stride == 1 + ): + indexers.append(":") + continue + if not ( + isinstance(start, Constant) + and isinstance(start.value, int) + and isinstance(size, int) + and isinstance(stride, int) + ): + runtime_starts = True + break + begin = int(start.value) + stop = begin + size * stride + indexers.append( + f"{begin}:{stop}" if stride == 1 else f"{begin}:{stop}:{stride}" + ) + if runtime_starts: + if ctx is not None: + ctx.imports.add("from tilefoundry.dsl.tf import *") + start_refs = ", ".join( + self._slice_start(start, size, stride) + for start, size, stride in zip( + starts.elements, target.sizes, target.strides + ) + ) + if len(starts.elements) == 1: + start_refs += "," + return ( + f"slice({self.reference(expr.args[0])}, ({start_refs}), " + f"sizes={_attr_tuple_str(target.sizes, self, ctx)}, " + f"strides={_attr_tuple_str(target.strides, self, ctx)})" + ) + return f"{self.reference(expr.args[0])}[{', '.join(indexers)}]" + + if ctx is not None: + ctx.imports.add("from tilefoundry.dsl.tf import *") + alias_name = _kinded_alias_name(target) + suppressed = {"kind"} if alias_name is not None else set() + attrs: list[str] = [] + for param in type(target).params(): + if param.kind != "attribute": + continue + value = getattr(target, param.name, None) + if value is None or param.name in suppressed or param.name == "layout": + continue + if isinstance(value, str): + attrs.append(f'{param.name}="{value}"') + elif isinstance(value, DType): + attrs.append(f'{param.name}="{value.name}"') + elif isinstance(value, enum.Enum) and isinstance(value.value, str): + attrs.append(f'{param.name}="{value.value}"') + elif isinstance(value, float): + if math.isinf(value): + literal = "-1e999" if value < 0 else "1e999" + elif math.isnan(value): + literal = "(1e999 - 1e999)" + else: + literal = repr(value) + attrs.append(f"{param.name}={literal}") + elif isinstance(value, (ShardLayout, TensorType)): + with self.type_surface(indent=self._indent + " "): + rendered = self.visit(value, ctx) + if isinstance(value, TensorType): + rendered = " ".join(rendered.split()) + attrs.append(f"{param.name}={rendered}") + elif isinstance(value, tuple): + attrs.append(f"{param.name}={_attr_tuple_str(value, self, ctx)}") + else: + attrs.append(f"{param.name}={value}") + return f"{_op_name(target)}({', '.join([*(self.reference(arg) for arg in expr.args), *attrs])})" @dataclass(frozen=True) @@ -134,22 +274,16 @@ def _physical_line_count(lines: list[str]) -> int: return sum(line.count("\n") + 1 for line in lines) -def _compact_type(ty: object, mesh_name_map: dict[int, str]) -> str: - """One physical-line, DSL-shaped type annotation for inspection output. - - Takes the same mesh-name map the signature and mesh prelude are rendered - from, so an annotated layout names the hoisted mesh instead of restating it. - """ - if isinstance(ty, TensorType): - rendered = _tensor_annotation(ty, mesh_name_map=mesh_name_map) +def _compact_type(ty: object, printer: PythonPrinter, ctx) -> str: + """One physical-line, DSL-shaped type annotation for inspection output.""" + if isinstance(ty, (TensorType, TupleType)): + with printer.type_surface(): + rendered = printer.visit(ty, ctx) return " ".join(rendered.split()) - if isinstance(ty, TupleType): - fields = ", ".join(_compact_type(field, mesh_name_map) for field in ty.fields) - return f"Tuple[{fields}]" return repr(ty) -def _comments(expr: Expr, options: PythonPrintOptions, mesh_name_map: dict[int, str]) -> str: +def _comments(expr: Expr, options: PythonPrintOptions, printer: PythonPrinter, ctx) -> str: """Return same-line annotations for one printed statement. Omit the binding because the left-hand side already carries its emitted name @@ -163,7 +297,7 @@ def _comments(expr: Expr, options: PythonPrintOptions, mesh_name_map: dict[int, """ comments: list[str] = [] if options.show_types: - comments.append(_compact_type(expr.type, mesh_name_map)) + comments.append(_compact_type(expr.type, printer, ctx)) for metadata_type in options.comment_metadata_types: metadata = get_metadata(expr, metadata_type) if metadata is None: @@ -174,219 +308,6 @@ def _comments(expr: Expr, options: PythonPrintOptions, mesh_name_map: dict[int, return f" # {PARTS.join(comments)}" if comments else "" -def shape_entry_str(entry: object) -> str: - """Render one tensor shape entry in canonical human-readable form. - - Static integers remain literals, dimension variables use their names, and - arithmetic expression trees use infix or function syntax. The printer and - viewer share this rendering instead of exposing dataclass representations. - See [types §4](docs/spec/types.md#4-dim--symbolic-shape-dimensions) and - [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). - """ - return _shape_entry_str(entry, nested=False) - - -class _ShapeEntryVisitor(ExprFunctor[str]): - def __init__(self, nested: bool) -> None: - super().__init__() - self.nested = nested - - def visit_DimVar(self, entry: DimVar, ctx=None) -> str: - return entry.name - - def visit_Var(self, entry: Var, ctx=None) -> str: - return entry.name - - def visit_Constant(self, entry: Constant, ctx=None) -> str: - return str(entry.value) - - def visit_Call(self, entry: Call, ctx=None) -> str: - ceildiv_args = _ceildiv_args(entry) - if ceildiv_args is not None: - a, b = ceildiv_args - return f"ceildiv({self._render(a, False)}, {self._render(b, False)})" - target = entry.target - if isinstance(target, DimConst): - return str(target.value) - for op_cls, sym in _DIM_INFIX_OPS.items(): - if isinstance(target, op_cls): - a, b = entry.args - rendered = f"{self._render(a, True)} {sym} {self._render(b, True)}" - return f"({rendered})" if self.nested else rendered - for op_cls, fname in _DIM_FUNC_OPS.items(): - if isinstance(target, op_cls): - rendered = ", ".join(self._render(arg, False) for arg in entry.args) - return f"{fname}({rendered})" - return repr(entry) - - def default_visit(self, entry, ctx=None) -> str: - if isinstance(entry, bool): - return repr(entry) - if isinstance(entry, int): - return str(entry) - return repr(entry) - - def _render(self, entry, nested: bool) -> str: - previous = self.nested - self.nested = nested - try: - return self.visit(entry, None) - finally: - self.nested = previous - - -def _shape_entry_str(entry: object, *, nested: bool) -> str: - return _ShapeEntryVisitor(nested).visit(entry) - - -def _ceildiv_args(entry: Call) -> tuple[object, object] | None: - """Recover the public constructor from ceildiv's canonical arithmetic tree.""" - if not isinstance(entry.target, DimFloorDiv) or len(entry.args) != 2: - return None - numerator, divisor = entry.args - if not ( - isinstance(numerator, Call) - and isinstance(numerator.target, DimSub) - and len(numerator.args) == 2 - and isinstance(numerator.args[1], Constant) - and numerator.args[1].value == 1 - ): - return None - added = numerator.args[0] - if not ( - isinstance(added, Call) - and isinstance(added.target, DimAdd) - and len(added.args) == 2 - and added.args[1] == divisor - ): - return None - return added.args[0], divisor - - -def _classify_shard_attrs( - sl: ShardLayout, mesh_name: str -) -> tuple[dict[int, list[str]], list[str], list[str]] | None: - """Classify shard attributes into layout-axis splits, partials, broadcasts. - - Preserve mesh-axis order and allow nested axes to split one layout axis. - Return ``None`` for rank mismatch, invalid axes, or unknown attributes so - callers use verbose fallback. Surface and compact renderers share the - result, with the latter remapping splits onto tensor axes. - """ - layout = sl.layout - if not isinstance(layout, Layout) or len(sl.attrs) != len(sl.mesh.layout.shape): - return None - layout_rank = len(layout.shape) - names = sl.mesh.names if hasattr(sl.mesh, "names") and sl.mesh.names else () - splits: dict[int, list[str]] = {} - partials: list[str] = [] - broadcasts: list[str] = [] - for mesh_axis_idx, attr in enumerate(sl.attrs): - axis_name = names[mesh_axis_idx] if mesh_axis_idx < len(names) else f"ax{mesh_axis_idx}" - axis_ref = f"{mesh_name}.{axis_name}" - if isinstance(attr, Split): - if attr.axis >= layout_rank: - return None - splits.setdefault(attr.axis, []).append(axis_ref) - elif isinstance(attr, Partial): - partials.append(f'{axis_ref} @ P("{attr.reduction or "sum"}")') - elif isinstance(attr, Broadcast): - broadcasts.append(f"{axis_ref} @ B()") - else: - return None - return splits, partials, broadcasts - - -def _shard_layout_surface_str(sl: ShardLayout, mesh_name: str = "gpu", ctx=None) -> str | None: - """Render canonical parser sugar for a shard layout. - - Splits inline on the layout dimension they divide; the remaining value - states form a set, in which a broadcast appears only when nothing else - would name the mesh, because the parser reads an unstated axis that way. - The stride tuple is written whenever the layout has one: an unstated one - is not C-order shorthand but the sugar default a ``Reshard`` materializes - from the storage it moves between. Return ``None`` when sugar cannot - express the layout, so callers use the verbose fallback. - """ - layout = sl.layout - if not isinstance(layout, Layout): - return None - classified = _classify_shard_attrs(sl, mesh_name) - if classified is None: - return None - splits, partials, broadcasts = classified - states, state_import = (partials, "P") if (splits or partials) else (broadcasts, "B") - if not splits and not states: - return None - - explicit = layout.strides is not None - if explicit and any( - i in splits and _shape_entry_str(dim, nested=True) != shape_entry_str(dim) - for i, dim in enumerate(layout.shape) - ): - return None - - dims = [ - f"{_shape_entry_str(d, nested=True)} {' '.join(f'@ {r}' for r in splits[i])}" - if i in splits - else shape_entry_str(d) - for i, d in enumerate(layout.shape) - ] - dim_str = ", ".join(dims) - if len(dims) == 1: - dim_str += "," - axis_tuple = f"({dim_str})" - - stride_str = _shape_tuple(layout.strides) if explicit else None - value_set = None - if states: - value_set = "{" + ", ".join(states) + "}" - if ctx is not None: - ctx.use( - PythonExpr( - (f"from tilefoundry.ir.types.shard import {state_import}",), state_import - ) - ) - - if stride_str is None and value_set is None: - return axis_tuple - parts = [axis_tuple] - if stride_str is not None: - parts.append(stride_str) - if value_set is not None: - parts.append(value_set) - return "(" + ", ".join(parts) + ")" - - -def shard_compact_inline( - sl: ShardLayout, mesh_name: str, tensor_shape: tuple -) -> tuple[dict[int, str], list[str]] | None: - """Decompose a shard layout for compact tensor-axis display. - - Map splits to tensor axes, collect ordered partial states, and omit - broadcasts. Return ``None`` for ambiguous split positions, invalid axes, - unknown attributes, or rank mismatch so callers fall back to canonical - rendering. Attribute classification is shared with surface rendering. - """ - layout = sl.layout - if not isinstance(layout, Layout): - return None - classified = _classify_shard_attrs(sl, mesh_name) - if classified is None: - return None - splits, partials, _broadcasts = classified - la2ta = layout_axis_to_tensor_axis(layout.shape, tensor_shape) - split_ref: dict[int, str] = {} - for layout_axis, refs in splits.items(): - if len(refs) != 1: - return None - t_axis = la2ta[layout_axis] - if t_axis in split_ref: - return None - split_ref[t_axis] = refs[0] - return split_ref, partials - - def _moved_window_ref(name: str, offset: int) -> str: """A tile-window indexer, carrying the compile-time offset that moves it.""" if offset == 0: @@ -394,7 +315,7 @@ def _moved_window_ref(name: str, offset: int) -> str: return f"{name} + {offset}" if offset > 0 else f"{name} - {-offset}" -def _attr_tuple_str(value: tuple) -> str: +def _attr_tuple_str(value: tuple, printer: PythonPrinter, ctx) -> str: """Render an attribute's tuple value as a Python tuple literal. A shape-valued attribute -- `new_shape`, a tile's extents -- can hold a @@ -404,7 +325,7 @@ def _attr_tuple_str(value: tuple) -> str: the header emits binds the name, not the repr. """ rendered = tuple( - shape_entry_str(entry) if _is_dim_entry(entry) else repr(entry) + printer.visit(entry, ctx) if _is_dim_entry(entry) else repr(entry) for entry in value ) if len(rendered) == 1: @@ -414,10 +335,7 @@ def _attr_tuple_str(value: tuple) -> str: def _is_dim_entry(entry: object) -> bool: """Whether *entry* is a dimension rather than a plain attribute value.""" - return isinstance(entry, (DimVar, Var, Constant)) or ( - isinstance(entry, Call) - and isinstance(entry.target, (DimConst, *_DIM_INFIX_OPS, *_DIM_FUNC_OPS)) - ) + return isinstance(entry, (DimVar, Var, Constant, Call)) def _tensor_import_names(fn: HirFunction) -> str: @@ -577,7 +495,7 @@ def _where_str(metadata: ScheduleConstraintMetadata) -> str: fields.append(f"layout={_layout_constraint_str(layout)}") for item in metadata.constraints: if isinstance(item, MeshConstraint): - fields.append(f"mesh={_mesh_str(item.mesh)}") + fields.append(f"mesh={PythonPrinter().visit(item.mesh, HirPrintContext())}") elif isinstance(item, StorageConstraint): fields.append(f'storage="{item.storage.name.lower()}"') return "where(" + ", ".join(fields) + ")" @@ -612,52 +530,6 @@ def iter_exprs(root: Expr | None, seen: set[int] | None = None) -> Iterator[Expr yield root -def _collect_meshes( - fn: HirFunction, - *, - include_node_types: bool = False, -) -> tuple[dict[int, Mesh], dict[int, Mesh]]: - """Collect meshes needed before emitting a function header. - - The function header declares every mesh and builds the name map used by - parameter annotations and region bodies, so collection must precede body - emission. It walks params, return type, and body references once; the - optional node-type scan is used by the viewer for intermediate annotations. - - With ``include_node_types=True`` (the viewer's wider scan, via - ``viewer.builder._collect_view_meshes``) every node's own result type is - also walked, since the viewer renders shard sugar on intermediate types too. - """ - type_meshes: dict[int, Mesh] = {} - scope_meshes: dict[int, Mesh] = {} - - def _add_layout(layout) -> None: - if isinstance(layout, ShardLayout): - type_meshes.setdefault(id(layout.mesh), layout.mesh) - - def _add_type(ty) -> None: - if isinstance(ty, TensorType): - _add_layout(ty.layout) - elif isinstance(ty, TupleType): - for f in ty.fields: - _add_type(f) - - for p in fn.params: - _add_type(p.type) - _add_type(fn.return_type) - - expressions = tuple(iter_exprs(fn.body)) - for expr in expressions: - if include_node_types: - _add_type(getattr(expr, "type", None)) - if isinstance(expr, Call) and isinstance(expr.target, Reshard): - _add_layout(expr.target.layout) - if isinstance(expr, MeshRegion): - scope_meshes.setdefault(id(expr.mesh), expr.mesh) - - return type_meshes, scope_meshes - - def _region_projection(expr: Expr) -> LoopRegion | MeshRegion | None: """Return the region projected by a one-argument ``TupleGetItem``.""" if not ( @@ -693,7 +565,7 @@ def _module_callee_binding(target: HirFunction, child_entries: dict[int, str]) - def _emit_def( - fn: HirFunction, def_name: str, mesh_map: dict[int, str], indent: str, + fn: HirFunction, def_name: str, ctx: HirPrintContext, indent: str, options: PythonPrintOptions, child_entries: dict[int, str] | None = None, *, line_offset: int = 0, @@ -708,7 +580,14 @@ def _emit_def( """ child_entries = {} if child_entries is None else child_entries lines: list[str] = [] - + printer = HirPrinter() + root_mesh = ( + fn.body.mesh + if isinstance(fn.body, MeshRegion) and not fn.specializations + else None + ) + if root_mesh is not None: + ctx.push_mesh(root_mesh, "mesh") _counter = [0] _names: dict[int, str] = {} @@ -867,81 +746,28 @@ def _assign_name(expr: Expr) -> str: if isinstance(expr, LoopRegion): for carry in expr.carried_args: _assign_name(carry) + printer.bind_def(_names, _param_alias, child_entries, _moved_window) - def _tuple_literal(elements) -> str: - inner = ", ".join( - repr(el.value) if isinstance(el, Constant) else _expr_ref(el) - for el in elements - ) - if len(elements) == 1: - inner += "," - return f"({inner})" - - def _expr_ref(expr: Expr) -> str: - if id(expr) in _param_alias: - return _expr_ref(_param_alias[id(expr)]) - projection = _region_projection(expr) - if isinstance(projection, LoopRegion): - return _names[id(projection.carried_args[expr.target.index])] - if isinstance(expr, LoopRegion): - - - - carried = tuple(_names[id(carry)] for carry in expr.carried_args) - if len(carried) == 1: - return carried[0] - return "(" + ", ".join(carried) + ")" - if isinstance(projection, MeshRegion): - return _expr_ref(projection.body.elements[expr.target.index]) - if isinstance(expr, MeshRegion): - return _arg_ref(expr.body) - return _names[id(expr)] - - def _arg_ref(a) -> str: - - - - return _tuple_literal(a.elements) if isinstance(a, Tuple) else _expr_ref(a) - - def _start_ref(start, size, stride) -> str: - """One Slice start as source. - - A moved tile window prints as the move itself -- the offset is a - compile-time constant, so it belongs in the indexer rather than in a - statement of its own. - """ - moved = _moved_window(start, size, stride) - if moved is None: - return repr(start.value) if isinstance(start, Constant) else _expr_ref(start) - window, offset = moved - return _moved_window_ref(_expr_ref(window), offset) - - - - return_ty = fn.return_type + return_type = fn.return_type arrow = "" - if isinstance(return_ty, TensorType): - arrow = " -> " + _tensor_annotation( - return_ty, mesh_name_map=mesh_map, indent=indent - ) - elif not isinstance(return_ty, TupleType): + if isinstance(return_type, TensorType): + with printer.type_surface(indent=indent): + arrow = " -> " + printer.visit(return_type, ctx) + elif not isinstance(return_type, TupleType): arrow = " -> None" lines.append(f"def {def_name}(") - param_strs = [] - for p in fn.params: - name = _names[id(p)] - if isinstance(p.type, TensorType): - ann = _tensor_annotation( - p.type, mesh_name_map=mesh_map, indent=indent, is_const=p.is_const, - ) - param_strs.append(f"{indent}{name}: {ann}") + params: list[str] = [] + for param in fn.params: + name = _names[id(param)] + if isinstance(param.type, TensorType): + with printer.type_surface(indent=indent, const=param.is_const): + annotation = printer.visit(param.type, ctx) + params.append(f"{indent}{name}: {annotation}") else: - param_strs.append(f"{indent}{name}") - - - for index, text in enumerate(param_strs): - suffix = "," if index < len(param_strs) - 1 else "" + params.append(f"{indent}{name}") + for index, text in enumerate(params): + suffix = "," if index < len(params) - 1 else "" lines.extend((text + suffix).split("\n")) lines.append(f"){arrow}:") @@ -949,135 +775,14 @@ def _start_ref(start, size, stride) -> str: line = _constraint_line(param, indent, _names[id(param)]) if line is not None: lines.append(line) - - if fn.body is None: lines.append(f"{indent}pass") + if root_mesh is not None: + ctx.pop_mesh() return lines printed: set[int] = {id(param) for param in fn.params} - def _format_call(expr: Call, indent_here: str) -> str: - """Render a Call's RHS expression text. - - Render a Call's RHS expression text: the ``reshard(...)`` / - ``(...)`` special forms, else ``op_name(args, attr=val, - ...)``. Shared by the inline (tile-loop body) emitter and the - top-level emit loop so an attribute-rendering rule (``ShardLayout``, - ``DType``, ...) only needs one edit. A reshard that gathers back to the - whole names no mesh, so its target is a plain ``Layout`` and there is no - mesh reference to abbreviate. - """ - target = expr.target - args_str = ", ".join(_arg_ref(arg) for arg in expr.args) - if isinstance(target, Reshard): - layout_kw = "" - if isinstance(target.layout, ShardLayout): - layout_text = _shard_layout_str( - target.layout, indent=indent_here + " ", mesh_map=mesh_map - ) - layout_kw = ", layout=" + layout_text - elif target.layout is not None: - layout_kw = ", layout=" + _layout_str(target.layout, indent_here + " ") - storage = ( - f", storage={target.storage.name.lower()}" - if target.storage is not None - else "" - ) - return f"reshard({args_str}{layout_kw}{storage})" - if isinstance(target, HirFunction): - binding = _module_callee_binding(target, child_entries) - return f"{binding or target.name}({args_str})" - if isinstance(target, Slice): - indexers = [] - starts = expr.args[1] - if not isinstance(starts, Tuple): - raise ValueError("canonical_source: Slice starts must be a Tuple") - runtime_starts = False - for axis, (start, size, stride) in enumerate( - zip(starts.elements, target.sizes, target.strides) - ): - if _moved_window(start, size, stride) is not None: - indexers.append(_start_ref(start, size, stride)) - continue - dim = expr.args[0].type.shape[axis] - if ( - isinstance(start, Constant) - and start.value == 0 - and size == dim - and stride == 1 - ): - indexers.append(":") - continue - if not ( - isinstance(start, Constant) - and isinstance(start.value, int) - and isinstance(size, int) - and isinstance(stride, int) - ): - runtime_starts = True - break - begin = int(start.value) - stop = begin + size * stride - indexers.append( - f"{begin}:{stop}" if stride == 1 else f"{begin}:{stop}:{stride}" - ) - if runtime_starts: - start_refs = ", ".join( - _start_ref(start, size, stride) - for start, size, stride in zip( - starts.elements, target.sizes, target.strides - ) - ) - if len(starts.elements) == 1: - start_refs += "," - return ( - f"slice({_arg_ref(expr.args[0])}, ({start_refs}), " - f"sizes={_attr_tuple_str(target.sizes)}, " - f"strides={_attr_tuple_str(target.strides)})" - ) - return f"{_arg_ref(expr.args[0])}[{', '.join(indexers)}]" - - alias_name = _kinded_alias_name(target) - suppress_attrs = {"kind"} if alias_name is not None else set() - attr_strs = [] - for param in type(target).params(): - if param.kind != "attribute": - continue - value = getattr(target, param.name, None) - if value is None or param.name in suppress_attrs or param.name == "layout": - continue - if isinstance(value, str): - attr_strs.append(f'{param.name}="{value}"') - elif isinstance(value, DType): - attr_strs.append(f'{param.name}="{value.name}"') - elif isinstance(value, enum.Enum) and isinstance(value.value, str): - attr_strs.append(f'{param.name}="{value.value}"') - elif isinstance(value, float): - if math.isinf(value): - literal = "-1e999" if value < 0 else "1e999" - elif math.isnan(value): - literal = "(1e999 - 1e999)" - else: - literal = repr(value) - attr_strs.append(f"{param.name}={literal}") - elif isinstance(value, ShardLayout): - sl_str = _shard_layout_str( - value, indent=indent_here + " ", mesh_map=mesh_map - ) - attr_strs.append(f"{param.name}={sl_str}") - elif isinstance(value, TensorType): - attr_strs.append(f"{param.name}={_compact_type(value, {})}") - elif isinstance(value, tuple): - attr_strs.append(f"{param.name}={_attr_tuple_str(value)}") - else: - attr_strs.append(f"{param.name}={value}") - - - - arguments = [_arg_ref(arg) for arg in expr.args] + attr_strs - return f"{_op_name(target)}({', '.join(arguments)})" - def _emit_inline_call(expr: Call, level: str) -> None: name = _names[id(expr)] if statements is not None: @@ -1085,77 +790,58 @@ def _emit_inline_call(expr: Call, level: str) -> None: value=name, line=line_offset + _physical_line_count(lines) + 1, ) + with printer.type_surface(indent=level): + rendered = printer.visit(expr, ctx) lines.append( - f"{level}{name} = {_format_call(expr, level)}" - f"{_comments(expr, options, mesh_map)}" + f"{level}{name} = {rendered}" + f"{_comments(expr, options, printer, ctx)}" ) printed.add(id(expr)) - class _ExprEmitter(ExprFunctor[None]): - def __init__(self, level: str) -> None: - super().__init__() - self.level = level - - def emit(self, expr: Expr) -> None: - self.visit(expr) - - def visit(self, expr, ctx=None): - key = id(expr) - if key in printed: - return None - if key in _inlined_start_ids: - printed.add(key) - return None - return super().visit(expr, ctx) - - def visit_Var(self, expr: Var, ctx=None) -> None: - printed.add(id(expr)) - - def visit_Constant(self, expr: Constant, ctx=None) -> None: + def _emit_expr(expr: Expr, level: str) -> None: + key = id(expr) + if key in printed: + return + if key in _inlined_start_ids: + printed.add(key) + return + if isinstance(expr, Var): + printed.add(key) + return + if isinstance(expr, Constant): lines.append( - f"{self.level}{_names[id(expr)]} = {repr(expr.value)}" - f"{_comments(expr, options, mesh_map)}" + f"{level}{_names[id(expr)]} = {repr(expr.value)}" + f"{_comments(expr, options, printer, ctx)}" ) - printed.add(id(expr)) - - def visit_Tuple(self, expr: Tuple, ctx=None) -> None: + printed.add(key) + return + if isinstance(expr, Tuple): for element in expr.elements: if not isinstance(element, Constant): - self.visit(element, ctx) - printed.add(id(expr)) - - def visit_LoopRegion(self, expr: LoopRegion, ctx=None) -> None: - _emit_loop_region(expr, self.level) - - def visit_MeshRegion(self, expr: MeshRegion, ctx=None) -> None: - _emit_mesh_region(expr, self.level) - - def visit_Call(self, expr: Call, ctx=None) -> None: + _emit_expr(element, level) + printed.add(key) + return + if isinstance(expr, LoopRegion): + _emit_loop_region(expr, level) + return + if isinstance(expr, MeshRegion): + _emit_mesh_region(expr, level) + return + if isinstance(expr, Call): projection = _region_projection(expr) if isinstance(projection, LoopRegion): - _emit_loop_region(projection, self.level) - printed.add(id(expr)) + _emit_loop_region(projection, level) + printed.add(key) return if isinstance(projection, MeshRegion): - _emit_mesh_region(projection, self.level) - printed.add(id(expr)) + _emit_mesh_region(projection, level) + printed.add(key) return for arg in expr.args: - self.visit(arg, ctx) - _emit_inline_call(expr, self.level) - - def default_visit(self, expr, ctx=None) -> None: - return None - - _expr_emitter = _ExprEmitter("") - - def _emit_expr(expr: Expr, level: str) -> None: - previous = _expr_emitter.level - _expr_emitter.level = level - try: - _expr_emitter.emit(expr) - finally: - _expr_emitter.level = previous + _emit_expr(arg, level) + _emit_inline_call(expr, level) + return + printed.add(key) def _emit_loop_region(region: LoopRegion, level: str) -> None: key = id(region) @@ -1165,23 +851,24 @@ def _emit_loop_region(region: LoopRegion, level: str) -> None: _emit_expr(init, level) for carry in region.carried_args: printed.add(id(carry)) - extent = shape_entry_str(region.extent) - step = shape_entry_str(region.step) - start = shape_entry_str(region.start) + extent = printer.visit(region.extent, ctx) + step = printer.visit(region.step, ctx) + start = printer.visit(region.start, ctx) if id(region.induction_var) in _tile_window_steps: + ctx.imports.add("from tilefoundry.dsl.tf import *") loop = f"tile({extent}, {step})" elif region.start == 0 and region.step == 1: loop = f"range({extent})" else: loop = f"range({start}, {extent}, {step})" - lines.append(f"{level}for {region.induction_var.name} in {loop}:{_comments(region, options, mesh_map)}") + lines.append(f"{level}for {region.induction_var.name} in {loop}:{_comments(region, options, printer, ctx)}") printed.add(key) inner = level + " " _emit_expr(region.body, inner) for value in region.yield_values: _emit_expr(value, inner) for carry, value in zip(region.carried_args, region.yield_values): - lines.append(f"{inner}{_names[id(carry)]} = {_expr_ref(value)}") + lines.append(f"{inner}{_names[id(carry)]} = {printer.reference(value)}") def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) -> None: key = id(region) @@ -1189,19 +876,30 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) return for arg in region.args: _emit_expr(arg, level) - mesh_name = mesh_map[id(region.mesh)] + if root_mesh is not None and region is fn.body: + printed.add(key) + _emit_expr(region.body, level) + if terminal: + lines.append(f"{level}return {printer.reference(region.body)}") + return + mesh_text = printer.visit(region.mesh, ctx) + mesh_name = ctx.scope_name(region.mesh) lines.append( - f"{level}with {mesh_name} as _{mesh_name}:" - f"{_comments(region, options, mesh_map)}" + f"{level}with {mesh_text} as {mesh_name}:" + f"{_comments(region, options, printer, ctx)}" ) printed.add(key) inner = level + " " - if terminal and isinstance(region.body, MeshRegion): - _emit_mesh_region(region.body, inner, terminal=True) - return - _emit_expr(region.body, inner) - if terminal: - lines.append(f"{inner}return {_arg_ref(region.body)}") + ctx.push_mesh(region.mesh, mesh_name) + try: + if terminal and isinstance(region.body, MeshRegion): + _emit_mesh_region(region.body, inner, terminal=True) + return + _emit_expr(region.body, inner) + if terminal: + lines.append(f"{inner}return {printer.reference(region.body)}") + finally: + ctx.pop_mesh() for expr in _order: if ( @@ -1224,7 +922,7 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) continue if isinstance(expr, Constant): name = _names[id(expr)] - lines.append(f"{indent}{name} = {repr(expr.value)}{_comments(expr, options, mesh_map)}") + lines.append(f"{indent}{name} = {repr(expr.value)}{_comments(expr, options, printer, ctx)}") line = _constraint_line(expr, indent, name) if line is not None: lines.append(line) @@ -1243,9 +941,11 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) value=name, line=line_offset + _physical_line_count(lines) + 1, ) + with printer.type_surface(indent=indent): + rendered = printer.visit(expr, ctx) lines.append( - f"{indent}{name} = {_format_call(expr, indent)}" - f"{_comments(expr, options, mesh_map)}" + f"{indent}{name} = {rendered}" + f"{_comments(expr, options, printer, ctx)}" ) line = _constraint_line(expr, indent, name) if line is not None: @@ -1256,227 +956,30 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) if not isinstance(fn.body, MeshRegion): if isinstance(fn.body, Tuple): - lines.append(f"{indent}return {_tuple_literal(fn.body.elements)}") + lines.append(f"{indent}return {printer.tuple_reference(fn.body.elements)}") elif isinstance(fn.body, LoopRegion): values = tuple(_names[id(carry)] for carry in fn.body.carried_args) result = values[0] if len(values) == 1 else "(" + ", ".join(values) + ")" lines.append(f"{indent}return {result}") else: - body_name = _expr_ref(fn.body) + body_name = printer.reference(fn.body) lines.append(f"{indent}return {body_name}") + if root_mesh is not None: + ctx.pop_mesh() return lines -_HIR_RENDERER = HirPrinter() -_dtype_str = _HIR_RENDERER.dtype_str -_pattern_ctor = _HIR_RENDERER.render_pattern - - -def _topologies_str(mesh: Mesh) -> str: - values = ", ".join( - f'Topology("{topology.name}", {shape_entry_str(topology.size)})' - for topology in mesh.topologies - ) - return f"({values}{',' if len(mesh.topologies) == 1 else ''})" - - -def _shape_tuple(shape: tuple) -> str: - return _HIR_RENDERER.shape_tuple(shape) - - -def _type_str(value, ctx=None, indent: str = "", *, is_const: bool = False) -> str: - """Enter the shared type visitor with the caller's block indentation.""" - with _HIR_RENDERER.type_surface(indent=indent, const=is_const): - return _HIR_RENDERER.visit(value, ctx if ctx is not None else HirPrintContext()) - -def _layout_str(layout: LayoutBase | None, indent: str = "") -> str: - return "None" if layout is None else _type_str(layout, indent=indent) - - -def _mesh_str(mesh: Mesh, indent: str = "") -> str: - return _type_str(mesh, indent=indent) - - -def _shard_layout_str(sl: ShardLayout, indent: str = "", *, mesh_map=None) -> str: - """A shard layout in an attribute slot, named from the printed mesh prelude.""" - return _type_str(sl, HirPrintContext(mesh_map), indent) - - -def _tensor_annotation(ty: TensorType, *, mesh_name_map=None, indent="", is_const=False) -> str: - return _type_str(ty, HirPrintContext(mesh_name_map), indent, is_const=is_const) - - -def _collect_all_meshes( - fn: HirFunction, -) -> tuple[dict[int, Mesh], dict[int, Mesh]]: - """Meshes referenced by *fn* and every specialization variant. - - Meshes referenced by *fn* and every specialization variant — the - printer's mesh-name map must stay stable across the base prototype and - each ``.specialize`` block. - """ - type_meshes: dict[int, Mesh] = {} - scope_meshes: dict[int, Mesh] = {} - for f in (fn, *fn.variants): - types, scopes = _collect_meshes(f) - type_meshes.update(types) - scope_meshes.update(scopes) - return type_meshes, scope_meshes - - -def _mesh_name_map(meshes: dict[int, Mesh]) -> dict[int, str]: - """Name every mesh identity, sharing one name per structural descriptor. - - A composed mesh is rebuilt at each use site, so one descriptor reaches the - printer under several identities. Naming those apart would emit a prelude - line per copy and make the annotations read as if they named different - meshes. - """ - used: set[str] = set() - by_signature: dict[str, str] = {} - result: dict[int, str] = {} - for identity, mesh in meshes.items(): - signature = _mesh_str(mesh) - name = by_signature.get(signature) - if name is None: - base = mesh.topologies[0].name if mesh.topologies else "mesh" - name, suffix = base, 2 - while name in used: - name = f"{base}_{suffix}" - suffix += 1 - used.add(name) - by_signature[signature] = name - result[identity] = name - return result - - -def _bound_mesh_aliases( - names: dict[int, str], meshes: dict[int, Mesh], scope_mesh_ids: set[int] -) -> dict[int, str]: - """Keep only the aliases the mesh prelude actually binds. - - A mesh with no named axes that no scope enters gets no prelude line, so - naming it inside a printed type would emit an undefined reference. Names - are assigned over every mesh first, so dropping the unbound ones here does - not renumber the meshes that remain. - """ - bound = { - names[identity] - for identity, mesh in meshes.items() - if mesh.names or identity in scope_mesh_ids - } - return {identity: name for identity, name in names.items() if name in bound} - - -def _emit_header( - fn: HirFunction, - meshes: dict[int, Mesh], - mesh_map: dict[int, str], - indent: str, - *, - for_module: bool = False, - target: object | None = None, - dim_vars: "dict[str, object] | None" = None, - scope_mesh_ids: set[int] | None = None, -) -> list[str]: - """Import header + mesh-prelude shared by ``hir_function_to_python`` and ``_module_to_python``. - - Import header + mesh-prelude shared by ``hir_function_to_python`` and - ``_module_to_python`` — the only source for the imports/mesh-defs a - dispatch prototype needs (the conditional ``DimVarRangePat`` import for - ``fn.variants``, the ``ConstTensor``/``Tensor`` selection), so standalone - and module-wrapped output cannot drift out of sync. - """ - lines: list[str] = ["from __future__ import annotations", ""] +def _new_hir_context(*, for_module: bool = False, target=None) -> HirPrintContext: + """Create a HIR context with imports owned by the surrounding file.""" + ctx = HirPrintContext() if for_module: - lines.append("from tilefoundry.module import module") - lines.append("from tilefoundry import func") + ctx.imports.add("from tilefoundry.module import module") + ctx.imports.add("from tilefoundry import func") if target is not None: - rendered: PythonExpr = target.to_python() - lines.extend(rendered.imports) - lines.append("from tilefoundry.dsl.tf import * # noqa: F401, F403") - lines.append(f"from tilefoundry.dsl import {_tensor_import_names(fn)}") - lines.append("from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401") - shard_names = {"Layout", "Mesh", "Topology"} if meshes else set() - if for_module: - shard_names.add("Topology") - layouts = [] - functions = (fn, *fn.variants) - for current in functions: - for param in current.params: - if isinstance(param.type, TensorType): - layouts.append(param.type.layout) - if isinstance(current.return_type, TensorType): - layouts.append(current.return_type.layout) - for current in functions: - for expr in iter_exprs(current.body): - if isinstance(expr.type, TensorType): - layouts.append(expr.type.layout) - if isinstance(expr, Call): - for candidate in vars(expr.target).values() if hasattr(expr.target, "__dict__") else (): - if isinstance(candidate, LayoutBase): - layouts.append(candidate) - elif isinstance(candidate, TensorType): - layouts.append(candidate.layout) - def nested_layouts(layout): - if isinstance(layout, ComposedLayout): - yield layout - if layout.inner is not None: - yield from nested_layouts(layout.inner) - if layout.outer is not None: - yield from nested_layouts(layout.outer) - elif isinstance(layout, ShardLayout): - yield layout - yield from nested_layouts(layout.layout) - - all_layouts = [nested for layout in layouts for nested in nested_layouts(layout)] - if any(isinstance(layout, ShardLayout) for layout in all_layouts): - shard_names.add("ShardLayout") - if any(isinstance(layout, ComposedLayout) for layout in all_layouts): - shard_names.add("ComposedLayout") - for layout in all_layouts: - if isinstance(layout, ShardLayout): - for attr in layout.attrs: - shard_names.add({Broadcast: "B", Split: "S", Partial: "P"}.get(type(attr), "")) - shard_names.discard("") - shard_import = ", ".join(sorted(shard_names)) - if shard_import: - lines.append(f"from tilefoundry.ir.types.shard import {shard_import}") - if fn.variants: - lines.append("from tilefoundry.ir.core.pattern import DimVarRangePat") - if dim_vars: - lines.append("from tilefoundry.ir.types.dim import DimVar, ceildiv") - lines.append("") - - - - - - if dim_vars: - for name, var in dim_vars.items(): - lines.append(f'{name} = DimVar("{var.name}", {var.lo}, {var.hi})') - lines.append("") - - - prelude: dict[str, Mesh] = {} - for identity, mesh in meshes.items(): - name = mesh_map.get(identity) - if name is not None: - prelude.setdefault(name, mesh) - for name, mesh in prelude.items(): - names_repr = repr(tuple(mesh.names)) if mesh.names else "()" - lines.append( - f"{name} = Mesh(" - f"{_topologies_str(mesh)}, " - f"{_layout_str(mesh.layout)}, " - f"names={names_repr}" - f")" - ) - if prelude: - lines.append("") - return lines - + rendered = target.to_python() + ctx.imports.update(rendered.imports) + return ctx def _variant_binding_name(variant: HirFunction) -> str: """Return a valid source binding for a variant without display metadata.""" @@ -1488,7 +991,7 @@ def _variant_binding_name(variant: HirFunction) -> str: def _emit_decorated_defs( - fn: HirFunction, mesh_map: dict[int, str], indent: str, options: PythonPrintOptions, + fn: HirFunction, ctx: HirPrintContext, indent: str, options: PythonPrintOptions, child_entries: dict[int, str] | None = None, *, line_offset: int = 0, @@ -1501,12 +1004,20 @@ def _emit_decorated_defs( render identically. See [inspection §2.6](docs/spec/inspection.md#26-specialization-printing). """ - lines: list[str] = ["@func"] + printer = HirPrinter() + decorator = "@func" + if isinstance(fn.body, MeshRegion) and not fn.specializations: + root_ctx = HirPrintContext() + with printer.type_surface(indent=indent): + mesh_text = printer.visit(fn.body.mesh, root_ctx) + ctx.imports.update(root_ctx.imports) + decorator = f"@func(mesh={mesh_text})" + lines: list[str] = [decorator] lines.extend( _emit_def( fn, fn.name, - mesh_map, + ctx, indent, options, child_entries, @@ -1519,11 +1030,11 @@ def _emit_decorated_defs( for variant in fn.variants: lines.append("") lines.append( - f"@{fn.name}.specialize({_pattern_ctor(variant.specializations[0])})" + f"@{fn.name}.specialize({HirPrinter().render_pattern(variant.specializations[0], ctx)})" ) lines.extend( _emit_def( - variant, _variant_binding_name(variant), mesh_map, indent, options, + variant, _variant_binding_name(variant), ctx, indent, options, child_entries, line_offset=line_offset + _physical_line_count(lines), statements=statements, @@ -1539,34 +1050,27 @@ def _render_hir_function( A normal function prints as a single ``@func``. A dispatch prototype (``variants != ()``) prints as a ``pass``-bodied ``@func`` base followed by - one ``@.specialize(pattern)`` block per variant. When the function - uses meshes with named axes, compact sugar form is emitted; otherwise the - verbose ``ShardLayout(...)`` form is used. + one ``@.specialize(pattern)`` block per variant. Placement sugar is + emitted only where an explicit mesh-scope binding dominates its use. """ indent = " " - type_meshes, scope_meshes = _collect_all_meshes(fn) - meshes = {**type_meshes, **scope_meshes} - mesh_map = _bound_mesh_aliases(_mesh_name_map(meshes), meshes, set(scope_meshes)) - lines = _emit_header( + ctx = _new_hir_context() + statements: dict[int, _PrintedStatement] = {} + lines = _emit_decorated_defs( fn, - meshes, - mesh_map, + ctx, indent, - dim_vars=dim_vars_reached(fn), - scope_mesh_ids=set(scope_meshes), + options or PythonPrintOptions(), + line_offset=0, + statements=statements, ) - statements: dict[int, _PrintedStatement] = {} - lines.extend( - _emit_decorated_defs( - fn, - mesh_map, - indent, - options or PythonPrintOptions(), - line_offset=_physical_line_count(lines), - statements=statements, - ) - ) - return _PythonRendering("\n".join(lines) + "\n", statements) + header = ctx.header() + header_lines = _physical_line_count(header) + statements = { + identity: _PrintedStatement(value=item.value, line=item.line + header_lines) + for identity, item in statements.items() + } + return _PythonRendering("\n".join(header + lines) + "\n", statements) def hir_function_to_python( @@ -1614,8 +1118,8 @@ def _emission_order(mod: Module) -> tuple: The entry goes last: a body calling a sibling names the attribute the class body already bound, so every callee must be written before it. Mesh - collection reads the same order, so the printed mesh prelude does not - depend on the order the authored source happened to use. + traversal reads the same order, so output does not depend on the order the + authored source happened to use. """ functions = mod.functions entry = mod.entry_function() if functions and mod.entry is not None else None @@ -1631,19 +1135,22 @@ def _module_tree_functions(mod: Module) -> tuple[HirFunction, ...]: return tuple(functions) -def _module_decorator_line(mod: Module, entry_name: str | None) -> str: +def _module_decorator_line(mod: Module, entry_name: str | None, ctx: HirPrintContext) -> str: """Render the context this Module declares as an ``@module(...)`` line. Always the called form. A bare decorator has not run while the class body is evaluated, so a body naming a child call could not resolve it. """ kwargs: list[str] = [] if entry_name is None else [f'entry="{entry_name}"'] + printer = HirPrinter() if mod.target is not None: rendered: PythonExpr = mod.target.to_python() + ctx.imports.update(rendered.imports) kwargs.append(f"target={rendered.text}") if mod.topologies is not None: + ctx.imports.add("from tilefoundry.ir.types.shard import Topology") topo_strs = [ - f'Topology("{t.name}", {shape_entry_str(t.size)})' + f'Topology("{t.name}", {printer.visit(t.size, ctx)})' for t in mod.topologies ] rendered_topologies = f'({", ".join(topo_strs)},)' if topo_strs else "()" @@ -1652,7 +1159,7 @@ def _module_decorator_line(mod: Module, entry_name: str | None) -> str: def _emit_module_class( - mod: Module, module_name: str, mesh_map: dict[int, str], indent: str, + mod: Module, module_name: str, ctx: HirPrintContext, indent: str, options: PythonPrintOptions, ) -> list[str]: """One ``@module`` class block: its nested Modules, then its functions. @@ -1660,7 +1167,7 @@ def _emit_module_class( Children first, because a body calling one names the attribute it is bound to and a class body binds in the order it is written. """ - lines = [_module_decorator_line(mod, mod.entry), f"class {module_name}:"] + lines = [_module_decorator_line(mod, mod.entry, ctx), f"class {module_name}:"] ordered = _emission_order(mod) child_entries = { id(child.entry_function()): child.name @@ -1668,14 +1175,16 @@ def _emit_module_class( if child.entry is not None and isinstance(child.entry_function(), HirFunction) } blocks: list[list[str]] = [ - _emit_module_class(child, child.name, mesh_map, indent, options) + _emit_module_class(child, child.name, ctx, indent, options) for child in mod.modules ] for fn in ordered: if isinstance(fn, HirFunction): - blocks.append(_emit_decorated_defs(fn, mesh_map, indent, options, child_entries)) + blocks.append(_emit_decorated_defs(fn, ctx, indent, options, child_entries)) elif isinstance(fn, PrimFunction): - blocks.append(_tir_function_block(fn)) + tir_block = _tir_function_block(fn) + ctx.imports.update(tir_block.imports) + blocks.append(tir_block) else: raise TypeError(f"Python printer cannot serialize {type(fn).__name__}") for index, block in enumerate(blocks): @@ -1708,48 +1217,13 @@ def _module_to_python( raise TypeError("Module printer requires a function entry") - header_of = entry if entry is not None else functions[0] indent4 = " " - type_meshes: dict[int, Mesh] = {} - scope_meshes: dict[int, Mesh] = {} - for fn in functions: - types, scopes = _collect_all_meshes(fn) - type_meshes.update(types) - scope_meshes.update(scopes) - meshes = {**type_meshes, **scope_meshes} - mesh_map = _bound_mesh_aliases(_mesh_name_map(meshes), meshes, set(scope_meshes)) - - dim_vars: dict[str, object] = {} - for fn in functions: - dim_vars.update(dim_vars_reached(fn)) - for node in _module_tree(root): - dim_vars.update(dim_vars_by_name(node.topologies or ())) - lines = _emit_header( - header_of, meshes, mesh_map, indent4, for_module=True, target=root.target, - dim_vars=dim_vars, scope_mesh_ids=set(scope_meshes), - ) - if any(isinstance(fn, PrimFunction) for node in _module_tree(root) for fn in node.functions): - lines = [ - line.replace("from tilefoundry import func", "from tilefoundry import func, prim_func") - for line in lines - ] - tensor_line = next(i for i, line in enumerate(lines) if line.startswith("from tilefoundry.dsl import ")) - lines[tensor_line] = lines[tensor_line].replace("import ", "import T, ") - target_imports = sorted({fn.target.to_python().imports[0] for node in _module_tree(root) for fn in node.functions if isinstance(fn, PrimFunction)}) - lines[2:2] = target_imports - tensor_names = "ConstTensor, Tensor" if any( - param.is_const for fn in functions for param in fn.params - ) else "Tensor" - lines = [ - f"from tilefoundry.dsl import {tensor_names}" if line.startswith("from tilefoundry.dsl import Tensor") else line - for line in lines - ] - lines.extend( - _emit_module_class( - root, module_name, mesh_map, indent4, options or PythonPrintOptions(), - ) + ctx = _new_hir_context(for_module=True, target=root.target) + lines = _emit_module_class( + root, module_name, ctx, indent4, options or PythonPrintOptions(), ) - return "\n".join(lines) + "\n" + header = ctx.header() + return "\n".join(header + lines) + "\n" def _module_tree(root: Module) -> Iterator[Module]: diff --git a/src/tilefoundry/inspection/python_type_printer.py b/src/tilefoundry/inspection/python_type_printer.py deleted file mode 100644 index 238858f8..00000000 --- a/src/tilefoundry/inspection/python_type_printer.py +++ /dev/null @@ -1,173 +0,0 @@ -"""Canonical Python rendering for immutable IR type values. - -Expression printers own traversal and statement/function syntax. This module -owns the value-language shared by those printers so HIR and TIR cannot grow -independent type-formatting implementations. Each type has exactly one -``visit_`` implementation and the visitors recurse into each other, so an -expression printer renders a type by entering the same ``visit`` its children -use rather than through a parallel ``render_*`` facade. -""" - -from __future__ import annotations - -from contextlib import contextmanager - -from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout -from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.types.shard.shard_layout import ( - Broadcast, - Partial, - ShardLayout, - Split, -) -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.ir.visitor import TypeFunctor -from tilefoundry.utils.python_source import PythonExpr - - -class PythonTypePrinter(TypeFunctor[str]): - """Render supported IR types/layouts through one Python value surface.""" - - def __init__(self) -> None: - self._indent = "" - self._tensor_head = "Tensor" - - @contextmanager - def type_surface(self, *, indent: str | None = None, const: bool = False): - """Carry the caller's block indentation and const-ness into ``visit``. - - Neither belongs to a type value: indentation is the statement the type - is printed inside and const-ness is a parameter's property, so they - travel as printer state instead of widening every visitor signature. - """ - previous = (self._indent, self._tensor_head) - if indent is not None: - self._indent = indent - self._tensor_head = "ConstTensor" if const else "Tensor" - try: - yield - finally: - self._indent, self._tensor_head = previous - - def dim_entry(self, value, ctx=None) -> str: - return str(value) - - def dtype_str(self, dtype: DType, ctx=None) -> str: - return dtype.name - - def shape_tuple(self, shape: tuple, ctx=None) -> str: - values = tuple(self.dim_entry(entry, ctx) for entry in shape) - return f"({values[0]},)" if len(values) == 1 else "(" + ", ".join(values) + ")" - - def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: - """Parser sugar for a shard layout, or ``None`` when it has none.""" - from .python_printer import _shard_layout_surface_str # noqa: PLC0415 - - mesh_name = ctx.mesh_alias(value.mesh) if ctx is not None else None - if mesh_name is None or not value.mesh.names: - return None - return _shard_layout_surface_str(value, mesh_name=mesh_name, ctx=ctx) - - def visit_TensorType(self, value: TensorType, ctx=None) -> str: - result = ( - f"{self._tensor_head}[" - f'{self.shape_tuple(value.shape, ctx)}, "{self.dtype_str(value.dtype, ctx)}"' - ) - if isinstance(value.layout, ShardLayout): - surface = self.shard_surface(value.layout, ctx) - if surface is not None: - result += f", {surface}" - else: - with self.type_surface(indent=self._indent + " "): - result += f",\n{self._indent}{self.visit(value.layout, ctx)}" - if value.storage is not StorageKind.GMEM: - result += f', "{value.storage.name.lower()}"' - return result + "]" - - def visit_TupleType(self, value: TupleType, ctx=None) -> str: - fields = ", ".join(self.visit(field, ctx) for field in value.fields) - return f"Tuple[{fields}]" - - def visit_UnitType(self, value: UnitType, ctx=None) -> str: - return "None" - - def visit_DType(self, value: DType, ctx=None) -> str: - return self.dtype_str(value, ctx) - - def visit_Mesh(self, value: Mesh, ctx=None) -> str: - alias = ctx.mesh_alias(value) if ctx is not None else None - if alias is not None: - return alias - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout, Mesh, Topology",), "")) - values = ", ".join( - f'Topology("{topology.name}", {self.dim_entry(topology.size, ctx)})' - for topology in value.topologies - ) - topologies = f"({values}{',' if len(value.topologies) == 1 else ''})" - result = f"Mesh({topologies}, {self.visit(value.layout, ctx)}" - if value.names: - result += f", names={tuple(value.names)!r}" - return result + ")" - - def visit_NoneType(self, value: None, ctx=None) -> str: - """An absent layout is part of the type language, not a missing case.""" - return "None" - - def visit_Layout(self, value: Layout, ctx=None) -> str: - strides = ( - self.shape_tuple(value.strides, ctx) if value.strides is not None else "None" - ) - return f"Layout({self.shape_tuple(value.shape, ctx)}, {strides})" - - def visit_ComposedLayout(self, value: ComposedLayout, ctx=None) -> str: - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) - outer, child = self._indent, self._indent + " " - with self.type_surface(indent=child): - inner_text = self.visit(value.inner, ctx) - outer_text = self.visit(value.outer, ctx) - return ( - "ComposedLayout(\n" - f"{child}inner={inner_text},\n" - f"{child}offset={self.dim_entry(value.offset, ctx)},\n" - f"{child}outer={outer_text},\n" - f"{outer})" - ) - - def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: - surface = self.shard_surface(value, ctx) - if surface is not None: - return surface - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) - outer, child = self._indent, self._indent + " " - attrs = ", ".join(self.visit(attr, ctx) for attr in value.attrs) - if len(value.attrs) == 1: - attrs += "," - with self.type_surface(indent=child): - mesh_text = self.visit(value.mesh, ctx) - layout_text = self.visit(value.layout, ctx) - return ( - "ShardLayout(\n" - f"{child}layout={layout_text},\n" - f"{child}attrs=({attrs}),\n" - f"{child}mesh={mesh_text},\n" - f"{outer})" - ) - - def visit_Broadcast(self, value: Broadcast, ctx=None) -> str: - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) - return "B()" - - def visit_Split(self, value: Split, ctx=None) -> str: - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) - return f"S({value.axis})" - - def visit_Partial(self, value: Partial, ctx=None) -> str: - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) - return f'P("{value.reduction}")' diff --git a/src/tilefoundry/inspection/tir_printer.py b/src/tilefoundry/inspection/tir_printer.py index 14fbfa0f..5ca8a4df 100644 --- a/src/tilefoundry/inspection/tir_printer.py +++ b/src/tilefoundry/inspection/tir_printer.py @@ -6,7 +6,6 @@ from tilefoundry.inspection.print_context import TirPrintContext from tilefoundry.inspection.printer_base import PythonPrinter -from tilefoundry.inspection.python_type_printer import PythonTypePrinter from tilefoundry.ir.core import Call, Constant, Op, Tuple, Var from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.core.module import Module @@ -14,7 +13,6 @@ 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.stmt import Stmt from tilefoundry.ir.tir.stmts import ( Evaluate, ) @@ -42,89 +40,90 @@ def __init__(self, *, context: TirPrintContext | None = None, indent: str = "") self.context = context or TirPrintContext() self.indent = indent - def dim_entry(self, value, ctx=None) -> str: - return f"_{value.name}" if hasattr(value, "name") else str(value) + def visit_DimVar(self, value, ctx=None) -> str: + return f"_{value.name}" - def visit(self, node, ctx=None): # type: ignore[override] - """One dispatch root for the two functor families this printer implements. - - Statements reach the statement visitor; every other value is a type and - reaches the inherited type visitor, so ``self.visit(expr.type, ctx)`` - stays the only way a type is printed. - """ - if isinstance(node, Stmt): - return StmtVisitor.visit(self, node) - return PythonTypePrinter.visit(self, node, ctx) + def dim_entry(self, value, ctx=None, *, nested: bool = False) -> str: + return super().dim_entry(value, ctx, nested=nested) - def visit_Sequential(self, stmt): + def visit(self, node, ctx=None): # type: ignore[override] + """Dispatch statements, expressions, and types by their concrete name.""" + return PythonPrinter.visit(self, node, ctx) + + def visit_Var(self, expr: Var, ctx=None) -> str: + return expr.name + + def visit_Constant(self, expr: Constant, ctx=None) -> str: + return repr(expr.value) + + def visit_SymbolRef(self, expr: SymbolRef, ctx=None) -> str: + return _binding_name(expr.name) + + def visit_ShapeOf(self, expr: ShapeOf, ctx=None) -> str: + return f"shape_of({expr.param.name}, axis={expr.axis})" + + def visit_Op(self, expr: Op, ctx=None) -> str: + name = getattr(getattr(expr, "_op_schema", None), "name", type(expr).__name__.lower()) + self.context.use(PythonExpr(("from tilefoundry.dsl import T",), "T")) + return f"T.{name}" + + def visit_program_call(self, expr: Call, ctx=None) -> str: + target = expr.target + scalar_binary = { + BinaryKind.EQ: "==", BinaryKind.NE: "!=", BinaryKind.LT: "<", + BinaryKind.LE: "<=", BinaryKind.GT: ">", BinaryKind.GE: ">=", BinaryKind.AND: "and", + } + kind = getattr(target, "kind", None) + if kind in scalar_binary and len(expr.args) == 2 and expr.type.dtype is DType.bool: + return f"{self.visit(expr.args[0])} {scalar_binary[kind]} {self.visit(expr.args[1])}" + name = getattr(getattr(target, "_op_schema", None), "name", None) or re.sub( + r"(?", BinaryKind.GE:">=", BinaryKind.AND:"and"} - kind = getattr(target, "kind", None) - if kind in scalar_binary and len(expr.args)==2 and expr.type.dtype is DType.bool: - return f"{self._expr(expr.args[0])} {scalar_binary[kind]} {self._expr(expr.args[1])}" - name = getattr(getattr(target, "_op_schema", None), "name", None) or re.sub(r"(? list[str]: callee, grid = stmt.args[0], stmt.args[1:4] block = stmt.args[4:7] forwarded = stmt.args[7:] - return [f"{indent}launch({printer._expr(callee)}, {printer._join_args(forwarded)}, grid={printer._expr(Tuple(type=grid[0].type, elements=tuple(grid)))}, block={printer._expr(Tuple(type=block[0].type, elements=tuple(block)))}) # noqa: F821"] + return [f"{indent}launch({printer.visit(callee)}, {printer._join_args(forwarded)}, grid={printer.visit(Tuple(type=grid[0].type, elements=tuple(grid)))}, block={printer.visit(Tuple(type=block[0].type, elements=tuple(block)))}) # noqa: F821"] @register_tir_printer(Op) @@ -172,8 +171,8 @@ def _print_op_evaluate(stmt: Evaluate, printer: TirPrinter) -> list[str]: continue rendered = printer.render_value(value, printer.context, printer.indent + " ") attrs.append(rendered if op_name == "sync" and p.name == "mesh" else f"{p.name}={rendered}") - rendered_args = [printer._expr(arg) for arg in args] - return [f"{indent}{printer._expr(target)}({', '.join(rendered_args + attrs)})"] + rendered_args = [printer.visit(arg) for arg in args] + return [f"{indent}{printer.visit_Op(target)}({', '.join(rendered_args + attrs)})"] def _function_block(fn: PrimFunction) -> list[str]: diff --git a/src/tilefoundry/inspection/viewer/builder.py b/src/tilefoundry/inspection/viewer/builder.py index d8dd7e51..ea354944 100644 --- a/src/tilefoundry/inspection/viewer/builder.py +++ b/src/tilefoundry/inspection/viewer/builder.py @@ -6,31 +6,18 @@ import graphviz -from tilefoundry.inspection.python_printer import ( - _DIM_FUNC_OPS, - _DIM_INFIX_OPS, - _collect_meshes, - _mesh_name_map, - _op_display_name, - _shard_layout_str, - _shard_layout_surface_str, - _tensor_annotation, - shape_entry_str, - shard_compact_inline, -) +from tilefoundry.inspection.print_context import HirPrintContext +from tilefoundry.inspection.printer_base import PythonPrinter from tilefoundry.ir.core import Tuple as HirTuple from tilefoundry.ir.core.expr import Call, Constant, Var from tilefoundry.ir.hir.function import Function as HirFunction +from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.types import DType -from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard.mesh import Mesh from tilefoundry.ir.types.shard.shard_layout import ShardLayout from tilefoundry.ir.types.tensor_type import TensorType, TupleType -from tilefoundry.ir.visitor import ExprVisitor from .htmltable import Cell, Span, Table from .palette import ( - DIMVAR_COLOR, HAIR, INK, MUTED, @@ -38,7 +25,6 @@ depth_border, depth_fill, exprkind_color, - storage_color, ) @@ -68,25 +54,6 @@ def _renderable_functions(root) -> list[tuple[str, "HirFunction"]]: return out -def _collect_view_meshes(root) -> dict[int, "Mesh"]: - """Collect unique ``Mesh`` objects referenced anywhere in *root*. - - Collect unique ``Mesh`` objects referenced anywhere in *root* — params, - return type, every node's result type, and ``Reshard`` layout attrs. - - The viewer renders shard sugar on intermediate node result types too, so it - needs a wider mesh-name map than the printer's param/return/Reshard scan: - built on the shared ``python_printer._collect_meshes`` with - ``include_node_types=True``. - """ - meshes: dict[int, Mesh] = {} - for _label, fn in _renderable_functions(root): - type_meshes, scope_meshes = _collect_meshes(fn, include_node_types=True) - meshes.update(type_meshes) - meshes.update(scope_meshes) - return meshes - - @dataclass class DetailRef: """Minimal click-lookup reference. NOT a graph/IR model. @@ -105,11 +72,11 @@ class DetailRef: class DetailIndex: """Detail lookup index only. ``dict[visual_id, DetailRef]``. - Carries the per-view ``mesh_name_map`` so the on-demand detail endpoint can - render shard sugar with the same stable mesh names as the graph. + Carries the per-view print context so detail text uses the same bindings + as the graph. """ entries: dict[str, DetailRef] = field(default_factory=dict) - mesh_name_map: dict[int, str] = field(default_factory=dict) + context: HirPrintContext = field(default_factory=HirPrintContext) def add(self, visual_id: str, ref: DetailRef) -> None: @@ -128,286 +95,124 @@ def get(self, visual_id: str) -> DetailRef | None: def _format_constant(c: Constant) -> str: - """Compact constant rendering (ported from the old viewer's pretty view). - - Compact constant rendering (ported from the old viewer's pretty - view). Scalars: ``const(1)`` / ``const(1.0f)``; sequences: - ``const([1.0f, 2.0f, ...])`` truncated to the first 8 elements. - """ ty = getattr(c, "type", None) suffix = ( _CONST_DTYPE_SUFFIX.get(ty.dtype.name, "") - if isinstance(ty, TensorType) and isinstance(ty.dtype, DType) else "" + if isinstance(ty, TensorType) and isinstance(ty.dtype, DType) + else "" ) - def _fmt(v) -> str: - if isinstance(v, bool): - return repr(v) - if isinstance(v, float): - return f"{v}{suffix}" - return repr(v) + def format_value(value): + if isinstance(value, bool): + return repr(value) + if isinstance(value, float): + return f"{value}{suffix}" + return repr(value) - val = c.value + value = c.value if isinstance(ty, TensorType) and ty.shape == (): - return f"const({_fmt(val)})" + return f"const({format_value(value)})" try: - items = list(val) + items = list(value) except TypeError: - return f"const({_fmt(val)})" - head = ", ".join(_fmt(v) for v in items[:8]) + return f"const({format_value(value)})" + head = ", ".join(format_value(item) for item in items[:8]) tail = ", ..." if len(items) > 8 else "" return f"const([{head}{tail}])" -def _shard_layout_text(sl: ShardLayout, mesh_name_map: dict[int, str] | None) -> str: - """Shard layout text. - - Render a bare ``ShardLayout`` attr (e.g. a ``Reshard`` layout) through the - canonical sugar core, falling back to the verbose ``ShardLayout(...)`` form - when the mesh is unnamed or the layout is not sugar-expressible. - """ - mesh_name = mesh_name_map.get(id(sl.mesh)) if mesh_name_map else None - if mesh_name and getattr(sl.mesh, "names", None): - mesh_unique = mesh_name_map is not None and len(mesh_name_map) == 1 - sugar = _shard_layout_surface_str(sl, mesh_name=mesh_name, mesh_unique=mesh_unique) - if sugar is not None: - return sugar - return _shard_layout_str(sl) - +def _type_text(ty, context=None) -> str: + printer = PythonPrinter() + if context is None: + context = HirPrintContext() + with printer.type_surface(): + return printer.visit(ty, context) -def _pretty_attr_value(value, *, full: bool = False, mesh_name_map: dict[int, str] | None = None) -> str: - """Readable rendering of an Op attribute value. - Readable rendering of an Op attribute value — pretty constants / - tuples / lists / types instead of raw ``repr``. ``full`` selects the - canonical type form (detail panel) vs the compact one (graph label). - """ +def _pretty_attr_value(value, *, context=None) -> str: if isinstance(value, Constant): return _format_constant(value) if isinstance(value, DType): return value.name - if isinstance(value, (TensorType, TupleType)): - return ( - type_to_canonical_pretty(value, mesh_name_map=mesh_name_map) - if full - else type_to_compact_pretty(value, mesh_name_map=mesh_name_map) - ) - if isinstance(value, ShardLayout): - return _shard_layout_text(value, mesh_name_map) + if isinstance(value, (TensorType, TupleType, ShardLayout)): + return _type_text(value, context) if isinstance(value, tuple): - inner = ", ".join(_pretty_attr_value(v, full=full, mesh_name_map=mesh_name_map) for v in value) + inner = ", ".join(_pretty_attr_value(v, context=context) for v in value) return f"({inner}{',' if len(value) == 1 else ''})" if isinstance(value, list): - return "[" + ", ".join(_pretty_attr_value(v, full=full, mesh_name_map=mesh_name_map) for v in value) + "]" + return "[" + ", ".join(_pretty_attr_value(v, context=context) for v in value) + "]" if isinstance(value, str): return value return repr(value) -def _op_attributes( - target, *, full: bool = False, mesh_name_map: dict[int, str] | None = None -) -> list[tuple[str, str]]: - """Op attributes. - - The Op's non-input (attribute) params as ``(name, pretty-value)`` - pairs — e.g. ``("axis", "2")`` / ``("begin", "(const(0), ...)")``. - ``full`` selects canonical (detail) vs compact (graph) type text. - Empty when the Op has no attributes or doesn't expose ``params()``. - """ +def _op_attributes(target, *, context=None) -> list[tuple[str, str]]: try: pdefs = type(target).params() except (AttributeError, TypeError): return [] - out = [] - for p in pdefs: - if getattr(p, "kind", None) == "attribute": - out.append(( - p.name, - _pretty_attr_value(getattr(target, p.name, None), full=full, mesh_name_map=mesh_name_map), - )) - return out + return [ + (p.name, _pretty_attr_value(getattr(target, p.name, None), context=context)) + for p in pdefs + if getattr(p, "kind", None) == "attribute" + ] -class _DimSpanVisitor(ExprVisitor[list[Span]]): - def visit_DimVar(self, dim: DimVar, ctx=None) -> list[Span]: - return [Span(text=dim.name, color=DIMVAR_COLOR, bold=True)] - - def visit_Constant(self, dim: Constant, ctx=None) -> list[Span]: - return [Span(text=str(dim.value))] - - def visit_Call(self, dim: Call, ctx=None) -> list[Span]: - target = dim.target - for op_cls, sym in _DIM_INFIX_OPS.items(): - if isinstance(target, op_cls): - spans = list(self.visit(dim.args[0], ctx)) - spans.append(Span(text=f" {sym} ")) - spans.extend(self.visit(dim.args[1], ctx)) - return spans - for op_cls, fname in _DIM_FUNC_OPS.items(): - if isinstance(target, op_cls): - spans = [Span(text=f"{fname}(")] - for i, arg in enumerate(dim.args): - if i: - spans.append(Span(text=", ")) - spans.extend(self.visit(arg, ctx)) - spans.append(Span(text=")")) - return spans - return [Span(text=shape_entry_str(dim))] - - def default_visit(self, dim, ctx=None) -> list[Span]: - return [Span(text=shape_entry_str(dim))] - - -def _format_dim(dim) -> list[Span]: - """Render a shape dimension as inline viewer spans. - - Share arithmetic syntax with the canonical Python printer. Dimension - variables use one token-class color while their text and detail identify the - symbol; integers remain plain. - See [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). - """ - return _DimSpanVisitor().visit(dim) - - -def _shard_inline(ty: TensorType, mesh_name_map: dict[int, str] | None): - """Compact shard decomposition for *ty*. - - Compact shard decomposition for *ty*: ``(split_ref_by_tensor_axis, - partials)`` or ``None`` when there is no named sugar-expressible shard - layout (caller renders the plain shape). - """ - layout = getattr(ty, "layout", None) - if not isinstance(layout, ShardLayout): - return None - mesh_name = mesh_name_map.get(id(layout.mesh)) if mesh_name_map else None - if not (mesh_name and getattr(layout.mesh, "names", None)): - return None - return shard_compact_inline(layout, mesh_name, ty.shape) - - -def _compact_type_spans(ty, mesh_name_map: dict[int, str] | None = None) -> list[Span]: - """Render compact graph-label types as colored inline spans. - - Tint dimension variables and storage, inline shard splits on tensor axes, - and append partial states. Layouts that cannot be inlined use canonical - annotation text. ``type_to_compact_pretty`` joins spans as plain text. - See [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). - """ - if isinstance(ty, TensorType): - inline = _shard_inline(ty, mesh_name_map) - if isinstance(ty.layout, ShardLayout) and inline is None: - - return [Span(text=type_to_canonical_pretty(ty, mesh_name_map=mesh_name_map))] - split_ref, partials = inline if inline is not None else ({}, []) - dtype = ty.dtype.name if hasattr(ty.dtype, "name") else str(ty.dtype) - spans: list[Span] = [Span(text=f"{dtype}[")] - for i, d in enumerate(ty.shape): - if i: - spans.append(Span(text=", ")) - spans.extend(_format_dim(d)) - if i in split_ref: - spans.append(Span(text=f" @ {split_ref[i]}")) - spans.append(Span(text="]")) - if partials: - spans.append(Span(text=" {" + ", ".join(partials) + "}")) - storage = getattr(ty, "storage", None) - if storage: - spans.append(Span(text=" @")) - spans.append(Span(text=str(storage), color=storage_color(str(storage)), bold=True)) - return spans - if isinstance(ty, TupleType): - - - spans: list[Span] = [Span(text="⟨")] - for i, sub in enumerate(ty.fields): - if i: - spans.append(Span(text=", ")) - spans.extend(_compact_type_spans(sub, mesh_name_map)) - spans.append(Span(text="⟩")) - return spans - return [Span(text=str(ty))] - - -def type_to_compact_pretty(ty, mesh_name_map: dict[int, str] | None = None) -> str: - """Render [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms) compact text.""" - return "".join(s.text for s in _compact_type_spans(ty, mesh_name_map)) - - -def type_to_canonical_pretty(ty, mesh_name_map: dict[int, str] | None = None) -> str: - """Render [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms) canonical text.""" - if isinstance(ty, TensorType): - return _tensor_annotation(ty, mesh_name_map=mesh_name_map) - if isinstance(ty, TupleType): - return "(" + ", ".join(type_to_canonical_pretty(f, mesh_name_map) for f in ty.fields) + ")" - return str(ty) - - -def _returns_of(ty, mesh_name_map: dict[int, str] | None = None) -> list[dict]: +def _returns_of(ty, context=None) -> list[dict]: if isinstance(ty, TupleType): return [ - {"idx": i, "type": type_to_canonical_pretty(f, mesh_name_map)} - for i, f in enumerate(ty.fields) + {"idx": i, "type": _type_text(field, context)} + for i, field in enumerate(ty.fields) ] if ty is None: return [] - return [{"idx": 0, "type": type_to_canonical_pretty(ty, mesh_name_map)}] + return [{"idx": 0, "type": _type_text(ty, context)}] def format_detail( - visual_id: str, ref: "DetailRef", mesh_name_map: dict[int, str] | None = None + visual_id: str, ref: "DetailRef", context: HirPrintContext | None = None ) -> dict: - """Format a detail-panel payload from a ``DetailRef`` on demand. - - Format a detail-panel payload from a ``DetailRef`` on demand (no - pre-baked JSON in the index). Shape: - ``{id, kind, name, params:[{name,type}], returns:[{idx,type}], attrs:[{key,value}]}``. - - ``mesh_name_map`` (the per-view map carried on ``DetailIndex``) lets the - canonical type text render shard sugar with the same stable mesh names as - the graph. - """ + """Format a detail-panel payload from a live HIR reference.""" + context = context or HirPrintContext() expr = ref.hir_expr name = ref.kind params: list[dict] = [] attrs: list[dict] = [] returns: list[dict] = [] - mm = mesh_name_map if isinstance(expr, HirFunction): name = expr.name - params = [{"name": p.name, "type": type_to_canonical_pretty(p.type, mm)} for p in expr.params] - returns = _returns_of(expr.return_type, mm) + params = [{"name": p.name, "type": _type_text(p.type, context)} for p in expr.params] + returns = _returns_of(expr.return_type, context) elif isinstance(expr, Var): name = expr.name - returns = _returns_of(expr.type, mm) + returns = _returns_of(expr.type, context) elif isinstance(expr, Constant): name = _format_constant(expr) - returns = _returns_of(expr.type, mm) + returns = _returns_of(expr.type, context) elif isinstance(expr, HirTuple): name = "Tuple" - params = [{"name": f"e{i}", "type": type_to_canonical_pretty(el.type, mm)} - for i, el in enumerate(expr.elements)] - returns = _returns_of(expr.type, mm) + params = [{"name": f"e{i}", "type": _type_text(item.type, context)} for i, item in enumerate(expr.elements)] + returns = _returns_of(expr.type, context) elif isinstance(expr, Call): - tgt = expr.target - if isinstance(tgt, HirFunction): - name = tgt.name - pnames = [p.name for p in tgt.params] + target = expr.target + if isinstance(target, HirFunction): + name = target.name + pnames = [p.name for p in target.params] else: - name = _op_display_name(tgt) + name = _op_display_name(target) try: - pnames = [p.name for p in type(tgt).params() if p.kind == "input"] + pnames = [p.name for p in type(target).params() if p.kind == "input"] except (AttributeError, TypeError): pnames = [] - attrs = [{"key": k, "value": v} for k, v in _op_attributes(tgt, full=True, mesh_name_map=mm)] + attrs = [{"key": key, "value": value} for key, value in _op_attributes(target, context=context)] params = [ - {"name": pnames[i] if i < len(pnames) else f"in{i}", - "type": type_to_canonical_pretty(a.type, mm)} - for i, a in enumerate(expr.args) + {"name": pnames[i] if i < len(pnames) else f"in{i}", "type": _type_text(arg.type, context)} + for i, arg in enumerate(expr.args) ] - returns = _returns_of(expr.type, mm) + returns = _returns_of(expr.type, context) else: - returns = _returns_of(getattr(expr, "type", None), mm) + returns = _returns_of(getattr(expr, "type", None), context) return {"id": visual_id, "kind": ref.kind, "name": name, "params": params, "returns": returns, "attrs": attrs} @@ -427,8 +232,11 @@ def __init__(self, root, collapsed: set[str] | None = None) -> None: self.collapsed = set(collapsed or ()) - self.mesh_name_map = _mesh_name_map(_collect_view_meshes(root)) - self.index = DetailIndex(mesh_name_map=self.mesh_name_map) + self.printer = PythonPrinter() + self.context = HirPrintContext() + if isinstance(root, HirFunction) and isinstance(root.body, MeshRegion): + self.context.push_mesh(root.body.mesh, "mesh") + self.index = DetailIndex(context=self.context) @@ -660,7 +468,7 @@ def _emit_function_node( width = 2 + n_params title.add_row( Cell( - spans=tuple(_compact_type_spans(fn.return_type, self.mesh_name_map)), + spans=(Span(text=_type_text(fn.return_type, self.context)),), colspan=width, color=INK, bold=True, font_size=12, ) ) @@ -716,7 +524,7 @@ def _walk_expr( Cell(text=f"Var {expr.name}", href="javascript:void(0)", title=f"expr:{vid}", bgcolor=exprkind_color("Var"), color="#ffffff", bold=True) ) - tbl.add_row(Cell(spans=tuple(_compact_type_spans(expr.type, self.mesh_name_map)), color=MUTED, font_size=11)) + tbl.add_row(Cell(spans=(Span(text=_type_text(expr.type, self.context)),), color=MUTED, font_size=11)) if output_slot is not None: tbl.add_row(self._output_marker_row(output_slot, 1)) g.node(vid, label=tbl.to_html()) @@ -733,7 +541,7 @@ def _walk_expr( Cell(text=_format_constant(expr), href="javascript:void(0)", title=f"expr:{vid}", bgcolor=exprkind_color("Constant"), color="#ffffff", bold=True) ) - tbl.add_row(Cell(spans=tuple(_compact_type_spans(expr.type, self.mesh_name_map)), color=MUTED, font_size=11)) + tbl.add_row(Cell(spans=(Span(text=_type_text(expr.type, self.context)),), color=MUTED, font_size=11)) if output_slot is not None: tbl.add_row(self._output_marker_row(output_slot, 1)) g.node(vid, label=tbl.to_html()) @@ -847,14 +655,14 @@ def _emit_call( - for key, val in _op_attributes(call.target, mesh_name_map=self.mesh_name_map): + for key, val in _op_attributes(call.target, context=self.context): tbl.add_row( Cell(text=f"{key}: {val}", colspan=width, color=MUTED, font_size=11, align="LEFT") ) tbl.add_row( Cell( - spans=tuple(_compact_type_spans(call.type, self.mesh_name_map)), + spans=(Span(text=_type_text(call.type, self.context)),), colspan=width, color=INK, bold=True, @@ -963,5 +771,5 @@ def _emit_tuple( __all__ = [ "DetailIndex", "DetailRef", "ViewerBuilder", - "format_detail", "type_to_compact_pretty", "type_to_canonical_pretty", + "format_detail", ] diff --git a/src/tilefoundry/inspection/viewer/server.py b/src/tilefoundry/inspection/viewer/server.py index afb651b0..9bae9ab1 100644 --- a/src/tilefoundry/inspection/viewer/server.py +++ b/src/tilefoundry/inspection/viewer/server.py @@ -102,7 +102,7 @@ def _send_expr(self, visual_id: str) -> None: if ref is None: self._send_json(404, {"error": "unknown visual_id", "id": visual_id}) return - self._send_json(200, format_detail(visual_id, ref, index.mesh_name_map)) + self._send_json(200, format_detail(visual_id, ref, index.context)) @staticmethod def _parse_collapsed(query: dict[str, list[str]]) -> set[str]: diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma.py b/src/tilefoundry/ir/tir/cuda/nn/mma.py index d7f20cef..491b7328 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma.py @@ -109,6 +109,7 @@ def _(call: "Call", ctx: "VerifyContext") -> None: _SM80_THREAD_MESH = Mesh( topologies=(Topology("thread", 32),), layout=Layout(shape=(4, 8), strides=(1, 4)), + names=("warp", "lane"), ) diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 44d4fa81..cf769406 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -78,6 +78,7 @@ Broadcast, Layout, Mesh, + Partial, ShardLayout, Split, Topology, @@ -239,6 +240,7 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex CallableType=CallableType, BindingSubstitutionCloner=BindingSubstitutionCloner, Broadcast=Broadcast, + Partial=Partial, Binary=Binary, BinaryKind=BinaryKind, Constant=Constant, diff --git a/src/tilefoundry/parser/grammar_render.py b/src/tilefoundry/parser/grammar_render.py index 844ecad1..22a0b6b4 100644 --- a/src/tilefoundry/parser/grammar_render.py +++ b/src/tilefoundry/parser/grammar_render.py @@ -96,8 +96,19 @@ class RenderVisitor: def __init__(self, *, line_width: int = 100): self.line_width = line_width self._seen: set[str] = set() + self._seen_named: set[str] = set() self._productions: list[tuple[str, _Expr]] = [] + def _named_pattern(self, pattern: Any) -> _Expr | None: + name = getattr(pattern, "grammar_name", None) + if not isinstance(name, str): + return None + grammar_name = _grammar_name(name) + if grammar_name not in self._seen_named: + self._seen_named.add(grammar_name) + self._productions.append((grammar_name, self._ast_node(pattern, allow_named=False))) + return _text(grammar_name) + def _element(self, pattern: ElementPattern[Any]) -> _Expr: name = pattern.element_name if not name: @@ -136,7 +147,11 @@ def _optional_field(self, fields: dict[str, object], name: str, fallback: str) - pattern = pattern.pattern return _text(fallback) if pattern is None else self.visit(pattern) - def _ast_node(self, pattern: AstNodePattern) -> _Expr: + def _ast_node(self, pattern: AstNodePattern, *, allow_named: bool = True) -> _Expr: + if allow_named: + named = self._named_pattern(pattern) + if named is not None: + return named node_type = pattern.node_type for part in pattern.parts: self.visit(part) diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index 3af5de89..89bd0b56 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -440,7 +440,7 @@ class PlacedLayout: def _layout_dims() -> AstPattern[Any]: """``(extent, extent @ axis, ...)`` — the layout's own divided dimensions.""" - return AstNodePattern( + pattern = AstNodePattern( ast.Tuple, FieldPattern( "elts", @@ -468,16 +468,20 @@ def _layout_dims() -> AstPattern[Any]: ), ), ) + pattern.grammar_name = "layout_dims" + return pattern def _layout_strides() -> AstPattern[Any]: """``(stride, ...)`` — how the divided positions are addressed.""" - return AstNodePattern(ast.Tuple, FieldPattern("elts", RepeatPattern(DimExprPattern()))) + pattern = AstNodePattern(ast.Tuple, FieldPattern("elts", RepeatPattern(DimExprPattern()))) + pattern.grammar_name = "layout_strides" + return pattern def _value_states() -> AstPattern[Any]: """``{axis @ B(), axis @ P("sum")}`` — what unsplit mesh axes hold.""" - return AstNodePattern( + pattern = AstNodePattern( ast.Set, FieldPattern( "elts", @@ -524,6 +528,8 @@ def _value_states() -> AstPattern[Any]: ), ), ) + pattern.grammar_name = "value_states" + return pattern def _layout_sugar_parts(node: object): @@ -568,6 +574,125 @@ def _value_state_parts(node: ast.AST): return None +@dataclass(frozen=True) +class _PlacementCandidate: + shape: tuple + strides: tuple | None + splits: tuple[tuple[object, int, int], ...] + states: tuple[tuple[object, int, str, str | None], ...] + + +def _placement_meshes(value: _PlacementCandidate, context: MatchContext, match): + referenced_ids = {id(entry[0]) for entry in (*value.splits, *value.states)} + if not referenced_ids: + return () + if context.function is None: + raise ParseError.from_node( + match.node, context, "placed layout requires function context" + ) + meshes = tuple( + mesh for mesh in context.function.state.mesh_stack if id(mesh) in referenced_ids + ) + if len(meshes) != len(referenced_ids): + meshes = tuple(dict.fromkeys(entry[0] for entry in (*value.splits, *value.states))) + if len(meshes) != len(referenced_ids): + raise ParseError.from_node(match.node, context, "placement references an inactive Mesh") + return meshes + + +@dataclass(frozen=True) +class LayoutStrideRankRule: + STATEMENT: ClassVar[str] = ( + "A stated stride tuple must have the rank of the layout it addresses." + ) + + def apply(self, value, *, match, context): + if value.strides is not None and len(value.shape) != len(value.strides): + raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") + return value + + +@dataclass(frozen=True) +class PlacementMeshResolutionRule: + STATEMENT: ClassVar[str] = ( + "A placement's mesh must be an active scope or resolvable from its bindings." + ) + + def apply(self, value, *, match, context): + _placement_meshes(value, context, match) + return value + + +@dataclass(frozen=True) +class PlacementLevelRule: + STATEMENT: ClassVar[str] = "A placement's meshes cannot name the same topology level." + + def apply(self, value, *, match, context): + meshes = _placement_meshes(value, context, match) + levels = [topology.name for mesh in meshes for topology in mesh.topologies] + if len(levels) != len(set(levels)): + duplicates = sorted(name for name in set(levels) if levels.count(name) > 1) + raise ParseError.from_node( + match.node, context, + f"a layout can split one level once; two of these meshes name {duplicates}", + ) + return value + + +@dataclass(frozen=True) +class MeshAxisBoundOnceRule: + STATEMENT: ClassVar[str] = "A placement binds each mesh axis at most once." + + def apply(self, value, *, match, context): + seen: set[tuple[int, int]] = set() + for mesh, axis, *_ in (*value.splits, *value.states): + key = (id(mesh), axis) + if key in seen: + raise ParseError.from_node(match.node, context, "mesh axis is bound more than once") + seen.add(key) + return value + + +@dataclass(frozen=True) +class PlacementConstructionRule: + """Materialize the candidate only after placement invariants have run.""" + + STATEMENT: ClassVar[str] = "A placement must construct a valid shard layout." + + def apply(self, value, *, match, context): + if not value.splits and not value.states: + return PlacedLayout( + shape=value.shape, layout=runtime.Layout(shape=value.shape, strides=value.strides) + ) + meshes = _placement_meshes(value, context, match) + mesh = meshes[0] if len(meshes) == 1 else runtime.composed(meshes) + source_offsets: dict[int, int] = {} + offset = 0 + for source in meshes: + source_offsets[id(source)] = offset + offset += len(source.layout.shape) + attrs: list[object] = [runtime.Broadcast() for _ in mesh.layout.shape] + for source, source_axis, tensor_axis in value.splits: + attrs[source_offsets[id(source)] + source_axis] = runtime.Split(tensor_axis) + for source, source_axis, kind, reduction in value.states: + target_axis = source_offsets[id(source)] + source_axis + attrs[target_axis] = runtime.Broadcast() if kind == "B" else runtime.Partial(reduction) + try: + canonical = runtime.canonical_shard_layout(value.shape, mesh, tuple(attrs)) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + if value.strides is not None and len(canonical.layout.shape) != len(value.strides): + raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") + return PlacedLayout( + shape=value.shape, + layout=runtime.ShardLayout( + layout=runtime.Layout(shape=canonical.layout.shape, strides=value.strides), + attrs=canonical.attrs, + mesh=canonical.mesh, + ), + ) + + @dataclass(frozen=True) class PlacementAnswerRule: """A placement answers both halves: the shape written, and where it goes.""" @@ -587,15 +712,7 @@ def apply(self, value, *, match, context): class PlacedLayoutPattern(ElementPattern): - """The layout a placement states: split dims, strides, and value states. - - One production covers the whole sugar surface because the parts are not - independent answers: the dims say how the layout is divided, the strides - say how the divided positions are addressed, and the value-state set says - what every mesh axis the dims did not split holds. Splitting them across - patterns is what let the printer emit a stride tuple beside a placement - that nothing could read back. - """ + """The one placement production, including dimensions, strides, and states.""" element_name = "placed_layout" syntax = LazyPattern( @@ -624,7 +741,6 @@ class PlacedLayoutPattern(ElementPattern): @staticmethod def _placement_parts(node: ast.AST) -> tuple[ast.AST, tuple[ast.AST, ...]] | None: - """Flatten ``extent @ axis @ axis`` into one extent and its axes.""" if not isinstance(node, ast.BinOp) or not isinstance(node.op, ast.MatMult): return None left = PlacedLayoutPattern._placement_parts(node.left) @@ -656,44 +772,20 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat extent_context = context.child(situation="layout_extent", role="layout_extent") if DimExprPattern().match(extent_node, extent_context) is None: return None - children.append( - AstChild( - f"extent_{tensor_axis}", - DimExprPattern(), - extent_node, - "layout_extent", - "layout_extent", - ) - ) + children.append(AstChild(f"extent_{tensor_axis}", DimExprPattern(), extent_node, "layout_extent", "layout_extent")) for mesh_axis, axis_node in enumerate(axis_nodes): axis_context = context.child(situation="mesh_axis", role="mesh_axis") if MeshAxisPattern().match(axis_node, axis_context) is None: return None child_name = f"binding_{tensor_axis}_{mesh_axis}" bindings.append((child_name, tensor_axis)) - children.append( - AstChild( - child_name, - MeshAxisPattern(), - axis_node, - "mesh_axis", - "mesh_axis", - ) - ) + children.append(AstChild(child_name, MeshAxisPattern(), axis_node, "mesh_axis", "mesh_axis")) if strides_node is not None: for index, item in enumerate(strides_node.elts): stride_context = context.child(situation="layout_strides", role="layout_strides") if DimExprPattern().match(item, stride_context) is None: return None - children.append( - AstChild( - f"stride_{index}", - DimExprPattern(), - item, - "layout_strides", - "layout_strides", - ) - ) + children.append(AstChild(f"stride_{index}", DimExprPattern(), item, "layout_strides", "layout_strides")) if states_node is not None: for index, item in enumerate(states_node.elts): state = _value_state_parts(item) @@ -705,29 +797,16 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat return None child_name = f"state_{index}" states.append((child_name, kind, reduction)) - children.append( - AstChild( - child_name, - MeshAxisPattern(), - axis_node, - "mesh_axis", - "mesh_axis", - ) - ) + children.append(AstChild(child_name, MeshAxisPattern(), axis_node, "mesh_axis", "mesh_axis")) if not found_placement and not states and strides_node is None: return None return dataclasses.replace( - matched, - pattern_id="tensor.layout.placed", - branch_id="placed_layout", + matched, pattern_id="tensor.layout.placed", branch_id="placed_layout", captures={ - **matched.captures, - "rank": len(dims_node.elts), + **matched.captures, "rank": len(dims_node.elts), "stride_rank": None if strides_node is None else len(strides_node.elts), - "bindings": tuple(bindings), - "states": tuple(states), - }, - children=tuple(children), + "bindings": tuple(bindings), "states": tuple(states), + }, children=tuple(children), ) @staticmethod @@ -735,86 +814,15 @@ def construct(match, children, context): rank = match.captures["rank"] shape = tuple(children[f"extent_{axis}"] for axis in range(rank)) stride_rank = match.captures.get("stride_rank") - strides = ( - None - if stride_rank is None - else tuple(children[f"stride_{index}"] for index in range(stride_rank)) - ) - splits = tuple( - (*children[child_name], tensor_axis) - for child_name, tensor_axis in match.captures["bindings"] - ) - states = tuple( - (*children[child_name], kind, reduction) - for child_name, kind, reduction in match.captures.get("states", ()) - ) - if not splits and not states: - if strides is not None and len(shape) != len(strides): - raise ParseError.from_node( - match.node, context, "layout shape/stride rank mismatch" - ) - return PlacedLayout( - shape=shape, layout=runtime.Layout(shape=shape, strides=strides) - ) - referenced_ids = {id(entry[0]) for entry in (*splits, *states)} - if context.function is None: - raise ParseError.from_node( - match.node, context, "placed layout requires function context" - ) - meshes = tuple( - mesh for mesh in context.function.state.mesh_stack if id(mesh) in referenced_ids - ) - if len(meshes) != len(referenced_ids): - meshes = tuple(dict.fromkeys(entry[0] for entry in (*splits, *states))) - if len(meshes) != len(referenced_ids): - raise ParseError.from_node( - match.node, context, "placement references an inactive Mesh" - ) - levels = [topology.name for mesh in meshes for topology in mesh.topologies] - if len(levels) != len(set(levels)): - duplicates = sorted(name for name in set(levels) if levels.count(name) > 1) - raise ParseError.from_node( - match.node, - context, - f"a layout can split one level once; two of these meshes name {duplicates}", - ) - mesh = meshes[0] if len(meshes) == 1 else runtime.composed(meshes) - source_offsets: dict[int, int] = {} - offset = 0 - for source in meshes: - source_offsets[id(source)] = offset - offset += len(source.layout.shape) - attrs: list[object] = [runtime.Broadcast() for _ in mesh.layout.shape] - bound: set[int] = set() + strides = None if stride_rank is None else tuple(children[f"stride_{index}"] for index in range(stride_rank)) + splits = tuple((*children[child_name], tensor_axis) for child_name, tensor_axis in match.captures["bindings"]) + states = tuple((*children[child_name], kind, reduction) for child_name, kind, reduction in match.captures.get("states", ())) + return _PlacementCandidate(shape, strides, splits, states) - def claim(source, source_axis: int) -> int: - target_axis = source_offsets[id(source)] + source_axis - if target_axis in bound: - raise ParseError.from_node(match.node, context, "mesh axis is bound more than once") - bound.add(target_axis) - return target_axis - - for source, source_axis, tensor_axis in splits: - attrs[claim(source, source_axis)] = runtime.Split(tensor_axis) - for source, source_axis, kind, reduction in states: - target_axis = claim(source, source_axis) - attrs[target_axis] = Broadcast() if kind == "B" else Partial(reduction) - try: - canonical = runtime.canonical_shard_layout(shape, mesh, tuple(attrs)) - except (TypeError, ValueError) as error: - raise ParseError.from_node(match.node, context, str(error)) from error - if strides is not None and len(canonical.layout.shape) != len(strides): - raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") - return PlacedLayout( - shape=shape, - layout=runtime.ShardLayout( - layout=runtime.Layout(shape=canonical.layout.shape, strides=strides), - attrs=canonical.attrs, - mesh=canonical.mesh, - ), - ) - - RULES: ClassVar[tuple[AstRule[Any], ...]] = (PlacementAnswerRule(),) + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + LayoutStrideRankRule(), MeshAxisBoundOnceRule(), PlacementMeshResolutionRule(), + PlacementLevelRule(), PlacementConstructionRule(), PlacementAnswerRule(), + ) class LayoutPattern(ElementPattern): diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index 6b201a0e..c52ea140 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -1,39 +1,44 @@ from __future__ import annotations from tilefoundry import func -from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.dsl import Tensor -from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 +from tilefoundry.dsl.storage import gmem, rmem, smem +from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology -thread = Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) -cta = Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')) -cta_2 = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) - -@func +@func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], - acc: Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 16), {thread.lane @ P("sum")}), "rmem"], - mixed: Tensor[(8, 16), "f32", ((4 @ cta.tile, 2, 16), {cta.lane @ P("sum")}), "rmem"] + acc: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((2, 4, 16), None), + attrs=(S(0), P("sum")), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + ), "rmem"], + mixed: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((4, 2, 16), None), + attrs=(S(0), B(), P("sum")), + mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), + ), "rmem"] ): - with cta_2 as _cta_2: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B(), thread.lane @ B()})]] - with thread as _thread: # Tuple[Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})], Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B(), thread.lane @ B()})]] - v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ cta.warp, 2, 4 @ cta.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 @ cta.warp, 2, 4 @ cta.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 - v2 = reshard(v1, layout=(4 @ cta_2.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta_2.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_2.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 - v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta_2.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), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 - folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic - summed = reshard(mixed, layout=((8, 16), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost; traffic - for _ in range(3): # Tuple[Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete - v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {cta.tile @ B(), cta.warp @ B(), cta.lane @ B()}), "rmem"]; compute-cost flops=f32:512@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 - v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:512@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 - folded = v9 - summed = v8 - v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 - v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 - v14 = reshard(seed, layout=((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ thread.warp, 4 @ thread.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 - v15 = reshard(v14, layout=((16,), (1,), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) # Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B(), thread.lane @ B()})]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 - return (v5, v11, v13, v15) + with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread: # Tuple[Tensor[(8, 16, 4), "bf16", ShardLayout( layout=Layout((8, 16, 4), (64, 4, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(8, 16), "f32", ShardLayout( layout=Layout((8, 16), (16, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(8, 16), "f32", ShardLayout( layout=Layout((8, 16), (16, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(16,), "f32", ShardLayout( layout=Layout((16,), (1,)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )]] + v0 = reshard(x, layout=(4 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.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 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.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 + v2 = reshard(v1, layout=(4 @ mesh.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ mesh.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 @ mesh.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 + v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ mesh.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), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 + folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic + summed = reshard(mixed, layout=((8, 16), {mesh.tile @ B(), thread.warp @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]; compute-cost; traffic + for _ in range(3): # Tuple[Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete + v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]; compute-cost flops=f32:512@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 + v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:512@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 + folded = v9 + summed = v8 + v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 + v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 + v14 = reshard(seed, layout=((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ thread.warp, 4 @ thread.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 + v15 = reshard(v14, layout=((16,), (1,), {thread.warp @ B()}), storage=gmem) # Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 + return (v5, v11, v13, v15) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index 2160b675..882d8bcf 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -1,76 +1,92 @@ from __future__ import annotations -from tilefoundry.module import module from tilefoundry import func -from tilefoundry.target import CudaTarget -from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.dsl import Tensor -from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 +from tilefoundry.dsl.storage import gmem, rmem, smem +from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology - -thread = Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) -thread_2 = Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=('lane',)) -cta = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) -cta_2 = Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')) +from tilefoundry.module import module +from tilefoundry.target import CudaTarget @module(entry="composed_mesh_pipeline", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", 4), Topology("thread", 8),)) class TypePrinterSugar: - @func + @func(mesh=Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=('lane',))) def nested_loop_tuple( x: Tensor[(8, 16), "f32"], - weight: Tensor[(8, 16), "f32", ((8, 16), {thread.lane @ P("max")})] + weight: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((8, 16), None), + attrs=(B(), P("max")), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + )] ): - with thread_2 as _thread_2: - split = reshard(x, layout=(8 @ thread_2.lane, 16), storage=rmem) - whole = reshard(x, layout=((8, 16), {thread_2.lane @ B()}), storage=rmem) - for _ in range(3): - whole_2 = add(whole, whole) - split_2 = unary(split, kind="square") - split = split_2 - whole = whole_2 - v1 = reshard(split, layout=((8, 16), (16, 1), {thread_2.lane @ B()}), storage=gmem) - v2 = reshard(whole, layout=((8, 16), (16, 1), {thread_2.lane @ B()}), storage=gmem) - with thread as _thread: - per_warp = reshard(weight, layout=(2 @ thread.warp, 4, 16), storage=rmem) - unfolded = reshard(per_warp, layout=((8, 16), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) - return (v1, v2, unfolded) + split = reshard(x, layout=(8 @ mesh.lane, 16), storage=rmem) + whole = reshard(x, layout=((8, 16), {mesh.lane @ B()}), storage=rmem) + for _ in range(3): + whole_2 = add(whole, whole) + split_2 = unary(split, kind="square") + split = split_2 + whole = whole_2 + v1 = reshard(split, layout=((8, 16), (16, 1), {mesh.lane @ B()}), storage=gmem) + v2 = reshard(whole, layout=((8, 16), (16, 1), {mesh.lane @ B()}), storage=gmem) + with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread: + per_warp = reshard(weight, layout=(2 @ thread.warp, 4, 16), storage=rmem) + unfolded = reshard(per_warp, layout=((8, 16), {thread.warp @ B()}), storage=gmem) + return (v1, v2, unfolded) - @func + @func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) def named_and_out_of_scope( x: Tensor[(8, 16), "f32"], - held: Tensor[(8, 16), "f32", (4 @ cta.tile, 2, 16), "rmem"], - frag: Tensor[(16,), "f32", ((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), "rmem"] + held: Tensor[(8, 16), "f32", (4 @ mesh.tile, 2, 16), "rmem"], + frag: Tensor[(16,), "f32", + ShardLayout( + layout=Layout((2, 4, 2), (8, 2, 1)), + attrs=(S(0), S(1)), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + ), "rmem"] ): - with cta as _cta: - mine = reshard(x, layout=(4 @ cta.tile, 2, 16), storage=rmem) - v0 = reshard(mine, layout=Layout((8, 16), (16, 1)), storage=gmem) - escaped = reshard(x, layout=(2 @ thread.warp, 4, 16), storage=rmem) - return (v0, held, frag, escaped) + mine = reshard(x, layout=(4 @ mesh.tile, 2, 16), storage=rmem) + v0 = reshard(mine, layout=Layout((8, 16), (16, 1)), storage=gmem) + escaped = reshard(x, layout=ShardLayout( + layout=Layout((2, 4, 16), None), + attrs=(S(0), B()), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + ), storage=rmem) + return (v0, held, frag, escaped) - @func + @func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], - acc: Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 16), {thread.lane @ P("sum")}), "rmem"], - mixed: Tensor[(8, 16), "f32", ((4 @ cta_2.tile, 2, 16), {cta_2.lane @ P("sum")}), "rmem"] + acc: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((2, 4, 16), None), + attrs=(S(0), P("sum")), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + ), "rmem"], + mixed: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((4, 2, 16), None), + attrs=(S(0), B(), P("sum")), + mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), + ), "rmem"] ): - with cta as _cta: - with thread as _thread: - composed = reshard(x, layout=(4 @ cta_2.tile, 2, 2 @ cta_2.warp, 2, 4 @ cta_2.lane, 4), storage=rmem) - v0 = unary(composed, kind="square") - staged = reshard(v0, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) - narrowed = cast(staged, dtype="bf16") - swapped = transpose(narrowed, perm=(0, 2, 1)) - gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) - folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) - summed = reshard(mixed, layout=((8, 16), {cta_2.tile @ B(), cta_2.warp @ B(), cta_2.lane @ B()}), storage=rmem) - for _ in range(3): - summed_2 = add(summed, summed) - folded_2 = unary(folded, kind="square") - folded = folded_2 - summed = summed_2 - v2 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) - v3 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) - seeded = reshard(seed, layout=((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), storage=rmem) - v4 = reshard(seeded, layout=((16,), (1,), {thread.warp @ B(), thread.lane @ B()}), storage=gmem) - return (gathered, v2, v3, v4) + with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread_2: + composed = reshard(x, layout=(4 @ mesh.tile, 2, 2 @ thread_2.warp, 2, 4 @ thread_2.lane, 4), storage=rmem) + v0 = unary(composed, kind="square") + staged = reshard(v0, layout=(4 @ mesh.tile, 2, 4, 16), storage=smem) + narrowed = cast(staged, dtype="bf16") + swapped = transpose(narrowed, perm=(0, 2, 1)) + gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {thread_2.warp @ B()}), storage=gmem) + folded = reshard(acc, layout=(2 @ thread_2.warp, 4, 4 @ thread_2.lane, 4), storage=rmem) + summed = reshard(mixed, layout=((8, 16), {mesh.tile @ B(), thread_2.warp @ B()}), storage=rmem) + for _ in range(3): + summed_2 = add(summed, summed) + folded_2 = unary(folded, kind="square") + folded = folded_2 + summed = summed_2 + v2 = reshard(folded, layout=((8, 16), (16, 1), {thread_2.warp @ B()}), storage=gmem) + v3 = reshard(summed, layout=((8, 16), (16, 1), {thread_2.warp @ B()}), storage=gmem) + seeded = reshard(seed, layout=((2 @ thread_2.warp, 4 @ thread_2.lane, 2), (8, 2, 1)), storage=rmem) + v4 = reshard(seeded, layout=((16,), (1,), {thread_2.warp @ B()}), storage=gmem) + return (gathered, v2, v3, v4) diff --git a/tests/fixtures/tir/mma.py b/tests/fixtures/tir/mma.py index 39a6280f..4ef4038c 100644 --- a/tests/fixtures/tir/mma.py +++ b/tests/fixtures/tir/mma.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.shard import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -10,44 +10,17 @@ class MmHandwritten: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def mm_device(a: Tensor[(16, 16), "bf16"], b: Tensor[(16, 8), "bf16"], c: Tensor[(16, 8), "f32"]): - with Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))) as _warp: - a_view = T.tensor_view(a, layout=ShardLayout( - layout=Layout((2, 4, 2, 8, 2), (1, 2, 8, 16, 128)), - attrs=(S(1), S(3)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - )) - b_view = T.tensor_view(b, layout=ShardLayout( - layout=Layout((8, 2, 4, 2), (1, 8, 16, 64)), - attrs=(S(2), S(0)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - )) - a_frag = T.alloc_tensor(tensor_type=Tensor[(16, 16), "bf16", - ShardLayout( - layout=Layout((2, 4, 2, 8, 2), (1, 2, 8, 16, 128)), - attrs=(S(1), S(3)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - ), "rmem"]) - b_frag = T.alloc_tensor(tensor_type=Tensor[(16, 8), "bf16", - ShardLayout( - layout=Layout((8, 2, 4, 2), (1, 8, 16, 64)), - attrs=(S(2), S(0)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - ), "rmem"]) - acc = T.alloc_tensor(tensor_type=Tensor[(16, 8), "f32", - ShardLayout( - layout=Layout((2, 4, 8, 2), (1, 2, 8, 64)), - attrs=(S(1), S(2)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - ), "rmem"]) + with Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4)), names=('warp', 'lane')) as _warp: + a_view = T.tensor_view(a, layout=((2, 4 @ _warp.warp, 2, 8 @ _warp.lane, 2), (1, 2, 8, 16, 128))) + b_view = T.tensor_view(b, layout=((8 @ _warp.lane, 2, 4 @ _warp.warp, 2), (1, 8, 16, 64))) + a_frag = T.alloc_tensor(tensor_type=Tensor[(16, 16), "bf16", ((2, 4 @ _warp.warp, 2, 8 @ _warp.lane, 2), (1, 2, 8, 16, 128)), "rmem"]) + b_frag = T.alloc_tensor(tensor_type=Tensor[(16, 8), "bf16", ((8 @ _warp.lane, 2, 4 @ _warp.warp, 2), (1, 8, 16, 64)), "rmem"]) + acc = T.alloc_tensor(tensor_type=Tensor[(16, 8), "f32", ((2, 4 @ _warp.warp, 8 @ _warp.lane, 2), (1, 2, 8, 64)), "rmem"]) T.copy(a_view, a_frag) T.copy(b_view, b_frag) T.fill(acc, 0.0) T.mma(acc, a_frag, b_frag, atom=T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN)) - c_view = T.tensor_view(c, layout=ShardLayout( - layout=Layout((2, 4, 8, 2), (1, 2, 8, 64)), - attrs=(S(1), S(2)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - )) + c_view = T.tensor_view(c, layout=((2, 4 @ _warp.warp, 8 @ _warp.lane, 2), (1, 2, 8, 64))) T.copy(acc, c_view) @prim_func(target=CpuTarget()) diff --git a/tests/fixtures/tir/sync.py b/tests/fixtures/tir/sync.py index 266e4768..7d461e54 100644 --- a/tests/fixtures/tir/sync.py +++ b/tests/fixtures/tir/sync.py @@ -3,7 +3,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core.kinds import BinaryKind -from tilefoundry.ir.types.shard import ComposedLayout, Layout, Mesh, Topology +from tilefoundry.ir.types.shard import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -16,21 +16,9 @@ def sync_square_device(a: Tensor[(4, 32), "f32"]): reg = T.alloc_tensor(tensor_type=Tensor[(4, 32), "f32", ((4 @ m.w, 32 @ m.t), (32, 1)), "rmem"]) T.copy(view, reg) T.sync(m) - T.sync(Mesh((Topology("thread", 128),), ComposedLayout( - inner=None, - offset=0, - outer=Layout((1, 32), (32, 1)), - ), names=('w', 't'))) - T.sync(Mesh((Topology("thread", 128),), ComposedLayout( - inner=None, - offset=0, - outer=Layout((2, 32), (32, 1)), - ), names=('w', 't'))) - T.sync(Mesh((Topology("thread", 128),), ComposedLayout( - inner=None, - offset=64, - outer=Layout((2, 32), (32, 1)), - ), names=('w', 't'))) + T.sync(m[:1]) + T.sync(m[:2]) + T.sync(m[2:]) T.binary(reg, reg, reg, kind=BinaryKind.MUL) T.copy(reg, view) diff --git a/tests/inspection/test_python_printer.py b/tests/inspection/test_python_printer.py index 68655bdf..df80c77d 100644 --- a/tests/inspection/test_python_printer.py +++ b/tests/inspection/test_python_printer.py @@ -6,13 +6,9 @@ *same* Target when the emitted source is executed. """ -import re from dataclasses import dataclass, fields, replace from tests._source import import_dsl -from tests.fixtures.placed.gqa_decode import GqaOnline -from tests.fixtures.placed.moe_mega_kernel import MoEMegaKernel -from tests.fixtures.placed.prefill_decode_attention import PrefillDecodeAttention from tests.fixtures.placed.rmsnorm import RmsnormModule from tilefoundry.inspection import PythonPrintOptions, as_script from tilefoundry.inspection.analysis_report import _type_text @@ -51,71 +47,6 @@ def test_inspection_types_are_opt_in_same_line_comments(): ) -def _hoisted_meshes(source: str) -> set[str]: - return { - line.split(" = ", 1)[0] for line in source.splitlines() if " = Mesh((Topology(" in line - } - - -def test_annotated_tuple_fields_each_name_the_hoisted_mesh(): - """Every field of a tuple-typed annotation reaches the prelude's mesh name. - - A loop carrying several tensors annotates a ``Tuple[...]``, so the field - walk has to carry the same mesh name map the signature is rendered from. A - field that lost it would restate the whole mesh as ``ShardLayout(...)``. - """ - annotated = as_script( - PrefillDecodeAttention, options=PythonPrintOptions(show_types=True) - ) - hoisted = _hoisted_meshes(annotated) - tuples = [ - line.split(" # ", 1)[1] for line in annotated.splitlines() if " # Tuple[" in line - ] - - assert len(tuples) == 2 - for annotation in tuples: - assert "ShardLayout(" not in annotation - named = set(re.findall(r"@ (\w+)\.", annotation)) - assert named - assert named <= hoisted - - -def test_an_annotated_layout_sugar_cannot_say_stays_verbose(): - """The fallback keeps the whole layout, and it coexists with sugar. - - Sugar names a mesh axis, so a mesh with no named axes has nothing to name - and stays verbose without dropping what the verbose form carries. The mesh - slot still names the mesh the prelude binds rather than restating it. - """ - unnamed_axes = as_script(GqaOnline, options=PythonPrintOptions(show_types=True)) - verbose = [ - line for line in unnamed_axes.splitlines() if " # " in line and "ShardLayout(" in line - ] - - assert verbose - for line in verbose: - annotation = line.split(" # ", 1)[1] - assert "layout=Layout(" in annotation and "attrs=(" in annotation - assert "names=" not in line - assert "@ " not in annotation - hoisted = _hoisted_meshes(unnamed_axes) - assert hoisted - assert all( - any(f"mesh={name}," in line for name in hoisted) for line in verbose - ) - - several_meshes = as_script( - MoEMegaKernel, options=PythonPrintOptions(show_types=True) - ) - annotations = [ - line.split(" # ", 1)[1].split("; ", 1)[0] - for line in several_meshes.splitlines() - if " # Tensor[" in line - ] - - assert any("@ cta_2.tile" in annotation for annotation in annotations) - - def test_binding_metadata_names_the_emitted_binding(): tensor_type = TensorType.scalar(DType.f32) source = Var(name="source", type=tensor_type) diff --git a/tests/inspection/test_roundtrip.py b/tests/inspection/test_roundtrip.py index 02dd6428..92e75fe6 100644 --- a/tests/inspection/test_roundtrip.py +++ b/tests/inspection/test_roundtrip.py @@ -217,28 +217,6 @@ def test_nested_composed_shard_layout_roundtrips_without_flattening() -> None: assert as_script(import_dsl(printed)) == printed -def test_compound_split_without_explicit_strides_stays_compact() -> None: - fn = import_dsl( - "from __future__ import annotations\n" - "from tilefoundry import func\n" - "from tilefoundry.dsl import DimVar, Tensor\n" - "from tilefoundry.ir.types.shard import Layout, Mesh, Topology\n" - "\n" - 'n = DimVar("n", 1, 65)\n' - "tiles = ((n - 1) // 8) + 1\n" - 'cta = Mesh((Topology("cta", tiles),), Layout((tiles,), (1,)), names=("tile",))\n' - "\n" - "@func\n" - 'def f(x: Tensor[(tiles, 8), "f32", (tiles @ cta.tile, 8)]):\n' - " return x\n" - ) - - printed = as_script(fn) - - assert "(((n - 1) // 8) + 1) @ cta.tile" in printed - assert as_script(import_dsl(printed)) == printed - - def test_carry_updates_print_last_without_shadowing_the_old_value() -> None: fn = import_dsl( _HEADER + "\n@func\n" diff --git a/tests/inspection/test_tir_roundtrip.py b/tests/inspection/test_tir_roundtrip.py index d34d37bd..9104df18 100644 --- a/tests/inspection/test_tir_roundtrip.py +++ b/tests/inspection/test_tir_roundtrip.py @@ -48,7 +48,6 @@ def test_placed_types_print_and_reparse_as_layout_sugar() -> None: printed = as_script(_module_in(SUGAR)) assert printed == SUGAR.with_suffix(".printed.txt").read_text() - assert "ShardLayout(" not in printed assert as_script(import_dsl(printed, name="TypePrinterSugar")) == printed @@ -104,7 +103,7 @@ def test_tir_for_if_and_sync_mesh_forms_roundtrip() -> None: ) printed = as_script(function) assert "T.sync(thread)" in printed - assert "T.sync(Mesh(" in printed + assert "T.sync(thread[:])" in printed assert as_script(import_dsl(printed, name="device")) == printed From 0a506dd2912d8ddf9179ab8c3cef08ea49a8ce03 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 11:39:14 +0800 Subject: [PATCH 13/21] fix(ci): align render and mma review contracts --- tests/cli/test_cli.py | 17 +++++++++++------ tests/ops/tir/cuda/test_mma.py | 6 +++++- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 097da619..7345d588 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -743,8 +743,8 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> assert annotated.count("reshard(tokens") == 2 assert "v0 = reshard(tokens" in annotated assert "v3 = reshard(tokens" in annotated - assert "offset=0" in annotated - assert "offset=120" in annotated + assert "with cta[:120] as cta_2:" in annotated + assert "with cta_3[120:] as cta_4:" in annotated summary = payload["function_records"]["performance"] cost = payload["function_records"]["compute-cost"] @@ -780,8 +780,13 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> } hoisted = {line.split(" = ", 1)[0] for line in lines if " = Mesh((Topology(" in line} - scoped = {line.lstrip().split()[1] for line in lines if line.lstrip().startswith("with ")} - assert hoisted == {"cta", "cta_2"} | scoped + scoped = { + line.rsplit(" as ", 1)[1].split(":", 1)[0] + for line in lines + if line.lstrip().startswith("with ") + } + assert hoisted == set() + assert scoped == {"cta", "cta_2", "cta_3", "cta_4"} annotated_types = [ line.split(" # ", 1)[1].split("; ", 1)[0] for line in lines if " # Tensor[" in line ] @@ -813,9 +818,9 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> else: assert "; performance=" not in statement annotated_meshes = set(re.findall(r"@ (\w+)\.", statement.split(" # ", 1)[1])) - assert annotated_meshes <= hoisted + assert annotated_meshes <= scoped placed_meshes = set(re.findall(r"mesh=(\w+),", statement)) - assert placed_meshes <= hoisted + assert placed_meshes <= scoped assert len([line for line in lines if "; performance=" in line]) == len(timed) assert len({row["value"] for row in rows}) == len(rows) assert "units=" not in annotated diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 09fa693a..08293a9b 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -46,7 +46,11 @@ def tile_device( c: Tensor[(128,), "f32"], ): atom = T.cuda.mma.atom(op=_OP) - with Mesh((Topology("thread", 32),), _MESH_LAYOUT) as m: + with Mesh( + (Topology("thread", 32),), + _MESH_LAYOUT, + names=("warp", "lane"), + ) as m: a_view = T.tensor_view( a, layout=ShardLayout( From 17ee2581625723a977d8b73388fc9aafd2b52a51 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 12:21:17 +0800 Subject: [PATCH 14/21] docs(tutorial): refresh showcase after printer refactor --- docs/tutorial/showcase.ipynb | 4 ++-- docs/tutorial/showcase.md | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index bb216ef2..a7aa768f 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=(1, 256, 8 @ cta.head, 32), 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 v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ cta.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 @ cta.head, 32), 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 v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ cta.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=(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" } ], "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 5510a9b9..45bb5856 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -870,11 +870,11 @@ for needle in ("reshard(w_q", "reshard(w_o"): # peak-footprint=gmem:10945036,smem:16960 # roofline ideal-ns=11213 bound-by=memory - v1 = reshard(w_q, layout=(1, 256, 8 @ cta.head, 32), 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 - v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ cta.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 + 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 - v42 = reshard(w_o, layout=(1, 256, 8 @ cta.head, 32), 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 - v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ cta.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 + 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 ``` ## 6. Stream the KV cache @@ -1022,7 +1022,7 @@ for needle in ("slice(k_cache", "cache_update(k_cache"): # 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 + 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 ``` From e1e2fdeec72c4e9b7cf1aff4bbddc4c241e47371 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 14:12:27 +0800 Subject: [PATCH 15/21] test(inspection): align mesh sugar fixture with review --- .../type_printer_sugar.analyzed.txt | 47 +++++++++---------- .../inspection/type_printer_sugar.printed.txt | 41 +++++++--------- .../fixtures/inspection/type_printer_sugar.py | 45 +++++++++--------- 3 files changed, 61 insertions(+), 72 deletions(-) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index c52ea140..ca8ba22e 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -6,16 +6,11 @@ from tilefoundry.dsl.storage import gmem, rmem, smem from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology -@func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) +@func(mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane'))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], - acc: Tensor[(8, 16), "f32", - ShardLayout( - layout=Layout((2, 4, 16), None), - attrs=(S(0), P("sum")), - mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), - ), "rmem"], + acc: Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem"], mixed: Tensor[(8, 16), "f32", ShardLayout( layout=Layout((4, 2, 16), None), @@ -23,22 +18,22 @@ def composed_mesh_pipeline( mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), ), "rmem"] ): - with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread: # Tuple[Tensor[(8, 16, 4), "bf16", ShardLayout( layout=Layout((8, 16, 4), (64, 4, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(8, 16), "f32", ShardLayout( layout=Layout((8, 16), (16, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(8, 16), "f32", ShardLayout( layout=Layout((8, 16), (16, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(16,), "f32", ShardLayout( layout=Layout((16,), (1,)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )]] - v0 = reshard(x, layout=(4 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.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 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.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 - v2 = reshard(v1, layout=(4 @ mesh.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ mesh.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 @ mesh.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 - v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ mesh.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), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 - folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic - summed = reshard(mixed, layout=((8, 16), {mesh.tile @ B(), thread.warp @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]; compute-cost; traffic - for _ in range(3): # Tuple[Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete - v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]; compute-cost flops=f32:512@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 - v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:512@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 - folded = v9 - summed = v8 - v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 - v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 - v14 = reshard(seed, layout=((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ thread.warp, 4 @ thread.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 - v15 = reshard(v14, layout=((16,), (1,), {thread.warp @ B()}), storage=gmem) # Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 - return (v5, v11, v13, v15) + with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: # Tuple[Tensor[(8, 16, 4), "bf16", ShardLayout( layout=Layout((8, 16, 4), (64, 4, 1)), attrs=(B(),), mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)), )], Tensor[(16,), "f32", ((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), "rmem"]] + 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 + 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 + 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): # Tuple[Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]]; 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 + 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 + return (v5, v13, v15, v17) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index 882d8bcf..fb23fb86 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -54,16 +54,11 @@ class TypePrinterSugar: ), storage=rmem) return (v0, held, frag, escaped) - @func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) + @func(mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane'))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], - acc: Tensor[(8, 16), "f32", - ShardLayout( - layout=Layout((2, 4, 16), None), - attrs=(S(0), P("sum")), - mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), - ), "rmem"], + acc: Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem"], mixed: Tensor[(8, 16), "f32", ShardLayout( layout=Layout((4, 2, 16), None), @@ -71,22 +66,22 @@ class TypePrinterSugar: mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), ), "rmem"] ): - with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread_2: - composed = reshard(x, layout=(4 @ mesh.tile, 2, 2 @ thread_2.warp, 2, 4 @ thread_2.lane, 4), storage=rmem) + with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: + composed = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) v0 = unary(composed, kind="square") - staged = reshard(v0, layout=(4 @ mesh.tile, 2, 4, 16), storage=smem) + staged = reshard(v0, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) narrowed = cast(staged, dtype="bf16") swapped = transpose(narrowed, perm=(0, 2, 1)) - gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {thread_2.warp @ B()}), storage=gmem) - folded = reshard(acc, layout=(2 @ thread_2.warp, 4, 4 @ thread_2.lane, 4), storage=rmem) - summed = reshard(mixed, layout=((8, 16), {mesh.tile @ B(), thread_2.warp @ B()}), storage=rmem) - for _ in range(3): - summed_2 = add(summed, summed) - folded_2 = unary(folded, kind="square") - folded = folded_2 - summed = summed_2 - v2 = reshard(folded, layout=((8, 16), (16, 1), {thread_2.warp @ B()}), storage=gmem) - v3 = reshard(summed, layout=((8, 16), (16, 1), {thread_2.warp @ B()}), storage=gmem) - seeded = reshard(seed, layout=((2 @ thread_2.warp, 4 @ thread_2.lane, 2), (8, 2, 1)), storage=rmem) - v4 = reshard(seeded, layout=((16,), (1,), {thread_2.warp @ B()}), storage=gmem) - return (gathered, v2, v3, v4) + gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) + seeded = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) + folded = reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage=rmem) + summed = reshard(mixed, layout=((8, 16), {mesh.warp @ B()}), storage=rmem) + for _ in range(3): + summed_2 = add(summed, summed) + folded_2 = unary(folded, kind="square") + folded = folded_2 + summed = summed_2 + v4 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) + v5 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) + v6 = reshard(seeded, layout=Layout((16,), (1,)), storage=gmem) + return (gathered, v4, v5, v6) diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py index 9238baa1..34959130 100644 --- a/tests/fixtures/inspection/type_printer_sugar.py +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -31,7 +31,7 @@ @module(entry="composed_mesh_pipeline", target=_H200, topologies=_TOPOLOGIES) class TypePrinterSugar: - @func + @func(mesh=_WARP_LANE) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], @@ -47,28 +47,27 @@ def composed_mesh_pipeline( ], ): with _TILE as cta: - with _WARP_LANE as thr: - composed = tf.reshard( - x, (8 @ cta.tile, 4 @ thr.warp, 16 @ thr.lane), "rmem" - ) - staged = tf.reshard(tf.square(composed), (8 @ cta.tile, 4, 16), "smem") - narrowed = tf.cast(staged, dtype="bf16") - swapped = tf.transpose(narrowed, perm=(0, 2, 1)) - gathered = tf.reshard(swapped, (8, 16, 4), "gmem") - seeded = tf.reshard(seed, _FRAGMENT, "rmem") - folded = tf.reshard(acc, (8 @ thr.warp, 16 @ thr.lane), "rmem") - summed = tf.reshard( - mixed, ((8, 16), {cta.tile @ B(), thr.warp @ B(), thr.lane @ B()}), "rmem" - ) - for _ in range(3): - folded = tf.square(folded) - summed = tf.add(summed, summed) - return ( - gathered, - tf.reshard(folded, (8, 16), "gmem"), - tf.reshard(summed, (8, 16), "gmem"), - tf.reshard(seeded, (16,), "gmem"), - ) + composed = tf.reshard( + x, (8 @ cta.tile, 4 @ mesh.warp, 16 @ mesh.lane), "rmem" + ) + staged = tf.reshard(tf.square(composed), (8 @ cta.tile, 4, 16), "smem") + narrowed = tf.cast(staged, dtype="bf16") + swapped = tf.transpose(narrowed, perm=(0, 2, 1)) + gathered = tf.reshard(swapped, (8, 16, 4), "gmem") + seeded = tf.reshard(seed, _FRAGMENT, "rmem") + folded = tf.reshard(acc, (8 @ mesh.warp, 16 @ mesh.lane), "rmem") + summed = tf.reshard( + mixed, ((8, 16), {mesh.warp @ B(), mesh.lane @ B()}), "rmem" + ) + for _ in range(3): + folded = tf.square(folded) + summed = tf.add(summed, summed) + return ( + gathered, + tf.reshard(folded, (8, 16), "gmem"), + tf.reshard(summed, (8, 16), "gmem"), + tf.reshard(seeded, (16,), "gmem"), + ) @func def nested_loop_tuple( From 5c5f396d1746a95a0aca2f51e718084e634bbc30 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 14:44:49 +0800 Subject: [PATCH 16/21] Revert "test(inspection): align mesh sugar fixture with review" This reverts commit e1e2fdeec72c4e9b7cf1aff4bbddc4c241e47371. --- .../type_printer_sugar.analyzed.txt | 47 ++++++++++--------- .../inspection/type_printer_sugar.printed.txt | 41 +++++++++------- .../fixtures/inspection/type_printer_sugar.py | 45 +++++++++--------- 3 files changed, 72 insertions(+), 61 deletions(-) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index ca8ba22e..c52ea140 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -6,11 +6,16 @@ from tilefoundry.dsl.storage import gmem, rmem, smem from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology -@func(mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane'))) +@func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], - acc: Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem"], + acc: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((2, 4, 16), None), + attrs=(S(0), P("sum")), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + ), "rmem"], mixed: Tensor[(8, 16), "f32", ShardLayout( layout=Layout((4, 2, 16), None), @@ -18,22 +23,22 @@ def composed_mesh_pipeline( mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), ), "rmem"] ): - with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: # Tuple[Tensor[(8, 16, 4), "bf16", ShardLayout( layout=Layout((8, 16, 4), (64, 4, 1)), attrs=(B(),), mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)), )], Tensor[(16,), "f32", ((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), "rmem"]] - 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 - 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 - 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): # Tuple[Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]]; 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 - 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 - return (v5, v13, v15, v17) + with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread: # Tuple[Tensor[(8, 16, 4), "bf16", ShardLayout( layout=Layout((8, 16, 4), (64, 4, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(8, 16), "f32", ShardLayout( layout=Layout((8, 16), (16, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(8, 16), "f32", ShardLayout( layout=Layout((8, 16), (16, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(16,), "f32", ShardLayout( layout=Layout((16,), (1,)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )]] + v0 = reshard(x, layout=(4 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.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 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.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 + v2 = reshard(v1, layout=(4 @ mesh.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ mesh.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 @ mesh.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 + v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ mesh.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), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 + folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic + summed = reshard(mixed, layout=((8, 16), {mesh.tile @ B(), thread.warp @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]; compute-cost; traffic + for _ in range(3): # Tuple[Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete + v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]; compute-cost flops=f32:512@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 + v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:512@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 + folded = v9 + summed = v8 + v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 + v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 + v14 = reshard(seed, layout=((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ thread.warp, 4 @ thread.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 + v15 = reshard(v14, layout=((16,), (1,), {thread.warp @ B()}), storage=gmem) # Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 + return (v5, v11, v13, v15) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index fb23fb86..882d8bcf 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -54,11 +54,16 @@ class TypePrinterSugar: ), storage=rmem) return (v0, held, frag, escaped) - @func(mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane'))) + @func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], - acc: Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem"], + acc: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((2, 4, 16), None), + attrs=(S(0), P("sum")), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + ), "rmem"], mixed: Tensor[(8, 16), "f32", ShardLayout( layout=Layout((4, 2, 16), None), @@ -66,22 +71,22 @@ class TypePrinterSugar: mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), ), "rmem"] ): - with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: - composed = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) + with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread_2: + composed = reshard(x, layout=(4 @ mesh.tile, 2, 2 @ thread_2.warp, 2, 4 @ thread_2.lane, 4), storage=rmem) v0 = unary(composed, kind="square") - staged = reshard(v0, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) + staged = reshard(v0, layout=(4 @ mesh.tile, 2, 4, 16), storage=smem) narrowed = cast(staged, dtype="bf16") swapped = transpose(narrowed, perm=(0, 2, 1)) - gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) - seeded = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) - folded = reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage=rmem) - summed = reshard(mixed, layout=((8, 16), {mesh.warp @ B()}), storage=rmem) - for _ in range(3): - summed_2 = add(summed, summed) - folded_2 = unary(folded, kind="square") - folded = folded_2 - summed = summed_2 - v4 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) - v5 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) - v6 = reshard(seeded, layout=Layout((16,), (1,)), storage=gmem) - return (gathered, v4, v5, v6) + gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {thread_2.warp @ B()}), storage=gmem) + folded = reshard(acc, layout=(2 @ thread_2.warp, 4, 4 @ thread_2.lane, 4), storage=rmem) + summed = reshard(mixed, layout=((8, 16), {mesh.tile @ B(), thread_2.warp @ B()}), storage=rmem) + for _ in range(3): + summed_2 = add(summed, summed) + folded_2 = unary(folded, kind="square") + folded = folded_2 + summed = summed_2 + v2 = reshard(folded, layout=((8, 16), (16, 1), {thread_2.warp @ B()}), storage=gmem) + v3 = reshard(summed, layout=((8, 16), (16, 1), {thread_2.warp @ B()}), storage=gmem) + seeded = reshard(seed, layout=((2 @ thread_2.warp, 4 @ thread_2.lane, 2), (8, 2, 1)), storage=rmem) + v4 = reshard(seeded, layout=((16,), (1,), {thread_2.warp @ B()}), storage=gmem) + return (gathered, v2, v3, v4) diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py index 34959130..9238baa1 100644 --- a/tests/fixtures/inspection/type_printer_sugar.py +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -31,7 +31,7 @@ @module(entry="composed_mesh_pipeline", target=_H200, topologies=_TOPOLOGIES) class TypePrinterSugar: - @func(mesh=_WARP_LANE) + @func def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], @@ -47,27 +47,28 @@ def composed_mesh_pipeline( ], ): with _TILE as cta: - composed = tf.reshard( - x, (8 @ cta.tile, 4 @ mesh.warp, 16 @ mesh.lane), "rmem" - ) - staged = tf.reshard(tf.square(composed), (8 @ cta.tile, 4, 16), "smem") - narrowed = tf.cast(staged, dtype="bf16") - swapped = tf.transpose(narrowed, perm=(0, 2, 1)) - gathered = tf.reshard(swapped, (8, 16, 4), "gmem") - seeded = tf.reshard(seed, _FRAGMENT, "rmem") - folded = tf.reshard(acc, (8 @ mesh.warp, 16 @ mesh.lane), "rmem") - summed = tf.reshard( - mixed, ((8, 16), {mesh.warp @ B(), mesh.lane @ B()}), "rmem" - ) - for _ in range(3): - folded = tf.square(folded) - summed = tf.add(summed, summed) - return ( - gathered, - tf.reshard(folded, (8, 16), "gmem"), - tf.reshard(summed, (8, 16), "gmem"), - tf.reshard(seeded, (16,), "gmem"), - ) + with _WARP_LANE as thr: + composed = tf.reshard( + x, (8 @ cta.tile, 4 @ thr.warp, 16 @ thr.lane), "rmem" + ) + staged = tf.reshard(tf.square(composed), (8 @ cta.tile, 4, 16), "smem") + narrowed = tf.cast(staged, dtype="bf16") + swapped = tf.transpose(narrowed, perm=(0, 2, 1)) + gathered = tf.reshard(swapped, (8, 16, 4), "gmem") + seeded = tf.reshard(seed, _FRAGMENT, "rmem") + folded = tf.reshard(acc, (8 @ thr.warp, 16 @ thr.lane), "rmem") + summed = tf.reshard( + mixed, ((8, 16), {cta.tile @ B(), thr.warp @ B(), thr.lane @ B()}), "rmem" + ) + for _ in range(3): + folded = tf.square(folded) + summed = tf.add(summed, summed) + return ( + gathered, + tf.reshard(folded, (8, 16), "gmem"), + tf.reshard(summed, (8, 16), "gmem"), + tf.reshard(seeded, (16,), "gmem"), + ) @func def nested_loop_tuple( From 56465bd16c40da8818444007247da0f80278c7fa Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 14:44:49 +0800 Subject: [PATCH 17/21] Revert "docs(tutorial): refresh showcase after printer refactor" This reverts commit 17ee2581625723a977d8b73388fc9aafd2b52a51. --- docs/tutorial/showcase.ipynb | 4 ++-- docs/tutorial/showcase.md | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index a7aa768f..bb216ef2 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=(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=(1, 256, 8 @ cta.head, 32), 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 v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ cta.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 @ cta.head, 32), 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 v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ cta.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" } ], "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 45bb5856..5510a9b9 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -870,11 +870,11 @@ for needle in ("reshard(w_q", "reshard(w_o"): # 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@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 @ cta.head, 32), 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 + v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ cta.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 - 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 + v42 = reshard(w_o, layout=(1, 256, 8 @ cta.head, 32), 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 + v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ cta.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 ``` ## 6. Stream the KV cache @@ -1022,7 +1022,7 @@ for needle in ("slice(k_cache", "cache_update(k_cache"): # 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 + 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 ``` From bcf8179796792f5d8f09d698fc194f7249bac6f5 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 22:36:01 +0800 Subject: [PATCH 18/21] fix(inspection/parser): enforce lexical mesh placement bindings [M0-M2] --- docs/spec/inspection.md | 4 + docs/spec/parser.md | 9 ++- src/tilefoundry/inspection/python_printer.py | 2 +- src/tilefoundry/parser/ast_pattern.py | 31 +++++++- src/tilefoundry/parser/pattern_nodes.py | 28 +++---- .../type_printer_sugar.analyzed.txt | 47 +++++------ .../inspection/type_printer_sugar.printed.txt | 41 +++++----- .../fixtures/inspection/type_printer_sugar.py | 78 +++++++++++-------- tests/parser/test_calls.py | 20 ++++- 9 files changed, 153 insertions(+), 107 deletions(-) diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index 37fae91a..d6e29d01 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -239,6 +239,10 @@ nested region writes its mesh expression at the `with` boundary and uses that binding in its body. A named constant or an out-of-scope mesh does not create a binding; types that refer to one use the verbose form. +Analysis annotations use the same type text for value-producing statements, but +a structural `MeshRegion` line MUST NOT restate the aggregate type of its body. +The region may still carry selected analysis metadata. + ### 2.6 Specialization printing A dispatch prototype ([hir.md §1.1](./hir.md#11-function)) diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 9fa7dfe5..bcdbee62 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -264,13 +264,20 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | placed_layout | tensor_optional_slot, tensor_shape | PlacementAnswerRule | Placement sugar states both the shape as written and the layout it implies. | src/tilefoundry/parser/pattern_nodes.py | | placed_layout | tensor_optional_slot, tensor_shape | PlacementConstructionRule | A placement must construct a valid shard layout. | src/tilefoundry/parser/pattern_nodes.py | | placed_layout | tensor_optional_slot, tensor_shape | PlacementLevelRule | A placement's meshes cannot name the same topology level. | src/tilefoundry/parser/pattern_nodes.py | -| placed_layout | tensor_optional_slot, tensor_shape | PlacementMeshResolutionRule | A placement's mesh must be an active scope or resolvable from its bindings. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementMeshResolutionRule | A placement's mesh must be a lexical mesh binding. | src/tilefoundry/parser/pattern_nodes.py | | shape | tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | | storage | tensor_optional_slot | StorageValueRule | Storage must resolve to a StorageKind. | src/tilefoundry/parser/ast_pattern.py | | tensor | annotation, expression, slice_endpoint, subscript_index, type_annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | | tensor | annotation, expression, slice_endpoint, subscript_index, type_annotation | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py | +A `mesh-axis` used by placement sugar MUST resolve to a mesh binding in the +current lexical scope. A module or closure name that resolves to a `Mesh` does +not become a placement binding. Such an external value remains valid as the +context expression of `with ... as ...` or as the value supplied to +`@func(mesh=...)`; the resulting lexical binding is the name placement sugar +may use. + ## 3. Implementation Overview | Component | Responsibility | diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index 32036653..8e444a14 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -296,7 +296,7 @@ def _comments(expr: Expr, options: PythonPrintOptions, printer: PythonPrinter, c the boundary between those two languages. """ comments: list[str] = [] - if options.show_types: + if options.show_types and not isinstance(expr, MeshRegion): comments.append(_compact_type(expr.type, printer, ctx)) for metadata_type in options.comment_metadata_types: metadata = get_metadata(expr, metadata_type) diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index cf769406..4925ef44 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -747,12 +747,26 @@ def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFai class LexicalScope: """Parser-local lexical frames shared by sequential child construction.""" - def __init__(self, frames: tuple[Mapping[str, object], ...] | None = None): + def __init__( + self, + frames: tuple[Mapping[str, object], ...] | None = None, + mesh_bindings: tuple[set[str], ...] | None = None, + ): source = frames or ({},) self._frames = [dict(frame) for frame in source] + self._mesh_bindings = [ + set(names) for names in (mesh_bindings or tuple(set() for _ in source)) + ] + if len(self._frames) != len(self._mesh_bindings): + raise ValueError("lexical frames and mesh bindings must have the same length") def define(self, name: str, value: object) -> None: self._frames[-1][name] = value + self._mesh_bindings[-1].discard(name) + + def define_mesh(self, name: str, value: object) -> None: + self._frames[-1][name] = value + self._mesh_bindings[-1].add(name) def lookup(self, name: str) -> object | None: for frame in reversed(self._frames): @@ -760,15 +774,26 @@ def lookup(self, name: str) -> object | None: return frame[name] return None + def lookup_mesh(self, name: str) -> object | None: + for frame, mesh_names in reversed(tuple(zip(self._frames, self._mesh_bindings))): + if name in frame: + return frame[name] if name in mesh_names else None + return None + def fork(self) -> LexicalScope: - return LexicalScope(tuple(self._frames) + ({},)) + return LexicalScope( + tuple(self._frames) + ({},), + tuple(self._mesh_bindings) + (set(),), + ) def push_frame(self) -> None: self._frames.append({}) + self._mesh_bindings.append(set()) def pop_frame(self) -> dict[str, object]: if len(self._frames) == 1: raise RuntimeError("cannot pop the root lexical frame") + self._mesh_bindings.pop() return self._frames.pop() def items(self): @@ -1264,7 +1289,7 @@ def from_function(cls, function: FuncParserContext) -> MatchContext: context, ) resolved_mesh = dataclasses.replace(function.mesh, topologies=topologies) - scope.define("mesh", resolved_mesh) + scope.define_mesh("mesh", resolved_mesh) scope.define( _TYPE_INFER_CONTEXT, ParserTypeInferContext(child_resolver=provider, current_mesh=resolved_mesh), diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index 89bd0b56..9cdec03f 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -385,15 +385,11 @@ def construct(match, children, context): else: binding = node.value.id axis_name = node.attr - mesh = context.lexical_scope.lookup(binding) - if mesh is None: - try: - reference = node if isinstance(node, ast.Name) else node.value - mesh = _resolve_reference(reference, context) - except ParseError: - mesh = None + mesh = context.lexical_scope.lookup_mesh(binding) if not isinstance(mesh, runtime.Mesh): - raise ParseError.from_node(node, context, f"{binding!r} is not an active Mesh") + raise ParseError.from_node( + node, context, f"{binding!r} is not a lexical Mesh binding" + ) if axis_name is None: if len(mesh.layout.shape) != 1: raise ParseError.from_node( @@ -590,13 +586,11 @@ def _placement_meshes(value: _PlacementCandidate, context: MatchContext, match): raise ParseError.from_node( match.node, context, "placed layout requires function context" ) - meshes = tuple( - mesh for mesh in context.function.state.mesh_stack if id(mesh) in referenced_ids - ) + meshes = tuple(dict.fromkeys(entry[0] for entry in (*value.splits, *value.states))) if len(meshes) != len(referenced_ids): - meshes = tuple(dict.fromkeys(entry[0] for entry in (*value.splits, *value.states))) - if len(meshes) != len(referenced_ids): - raise ParseError.from_node(match.node, context, "placement references an inactive Mesh") + raise ParseError.from_node( + match.node, context, "placement references a non-lexical Mesh binding" + ) return meshes @@ -614,9 +608,7 @@ def apply(self, value, *, match, context): @dataclass(frozen=True) class PlacementMeshResolutionRule: - STATEMENT: ClassVar[str] = ( - "A placement's mesh must be an active scope or resolvable from its bindings." - ) + STATEMENT: ClassVar[str] = "A placement's mesh must be a lexical mesh binding." def apply(self, value, *, match, context): _placement_meshes(value, context, match) @@ -3362,7 +3354,7 @@ def construct(match, children, context): binding = context.values.get("mesh_binding") _enter_mesh_scope(context, mesh, match) if isinstance(binding, str): - context.lexical_scope.define(binding, mesh) + context.lexical_scope.define_mesh(binding, mesh) context.function.state.mesh_stack.append(mesh) return mesh diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index c52ea140..36585389 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -6,16 +6,11 @@ from tilefoundry.dsl.storage import gmem, rmem, smem from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology -@func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) +@func(mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane'))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], - acc: Tensor[(8, 16), "f32", - ShardLayout( - layout=Layout((2, 4, 16), None), - attrs=(S(0), P("sum")), - mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), - ), "rmem"], + acc: Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem"], mixed: Tensor[(8, 16), "f32", ShardLayout( layout=Layout((4, 2, 16), None), @@ -23,22 +18,22 @@ def composed_mesh_pipeline( mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), ), "rmem"] ): - with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread: # Tuple[Tensor[(8, 16, 4), "bf16", ShardLayout( layout=Layout((8, 16, 4), (64, 4, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(8, 16), "f32", ShardLayout( layout=Layout((8, 16), (16, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(8, 16), "f32", ShardLayout( layout=Layout((8, 16), (16, 1)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )], Tensor[(16,), "f32", ShardLayout( layout=Layout((16,), (1,)), attrs=(B(), B()), mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), )]] - v0 = reshard(x, layout=(4 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.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 @ mesh.tile, 2, 2 @ thread.warp, 2, 4 @ thread.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 - v2 = reshard(v1, layout=(4 @ mesh.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ mesh.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 @ mesh.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 - v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ mesh.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), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 - folded = reshard(acc, layout=(2 @ thread.warp, 4, 4 @ thread.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic - summed = reshard(mixed, layout=((8, 16), {mesh.tile @ B(), thread.warp @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]; compute-cost; traffic - for _ in range(3): # Tuple[Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete - v8 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.tile @ B(), thread.warp @ B()}), "rmem"]; compute-cost flops=f32:512@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 - v9 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ thread.warp, 4, 4 @ thread.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:512@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 - folded = v9 - summed = v8 - v11 = reshard(folded, layout=((8, 16), (16, 1), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 - v13 = reshard(summed, layout=((8, 16), (16, 1), {thread.warp @ B()}), storage=gmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 - v14 = reshard(seed, layout=((2 @ thread.warp, 4 @ thread.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ thread.warp, 4 @ thread.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 - v15 = reshard(v14, layout=((16,), (1,), {thread.warp @ B()}), storage=gmem) # Tensor[(16,), "f32", ((16,), (1,), {thread.warp @ B()})]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 - return (v5, v11, v13, v15) + 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 + 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 + 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): # Tuple[Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]]; 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 + 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 + return (v5, v13, v15, v17) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index 882d8bcf..fb23fb86 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -54,16 +54,11 @@ class TypePrinterSugar: ), storage=rmem) return (v0, held, frag, escaped) - @func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) + @func(mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane'))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], - acc: Tensor[(8, 16), "f32", - ShardLayout( - layout=Layout((2, 4, 16), None), - attrs=(S(0), P("sum")), - mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), - ), "rmem"], + acc: Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem"], mixed: Tensor[(8, 16), "f32", ShardLayout( layout=Layout((4, 2, 16), None), @@ -71,22 +66,22 @@ class TypePrinterSugar: mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), ), "rmem"] ): - with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread_2: - composed = reshard(x, layout=(4 @ mesh.tile, 2, 2 @ thread_2.warp, 2, 4 @ thread_2.lane, 4), storage=rmem) + with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: + composed = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) v0 = unary(composed, kind="square") - staged = reshard(v0, layout=(4 @ mesh.tile, 2, 4, 16), storage=smem) + staged = reshard(v0, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) narrowed = cast(staged, dtype="bf16") swapped = transpose(narrowed, perm=(0, 2, 1)) - gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {thread_2.warp @ B()}), storage=gmem) - folded = reshard(acc, layout=(2 @ thread_2.warp, 4, 4 @ thread_2.lane, 4), storage=rmem) - summed = reshard(mixed, layout=((8, 16), {mesh.tile @ B(), thread_2.warp @ B()}), storage=rmem) - for _ in range(3): - summed_2 = add(summed, summed) - folded_2 = unary(folded, kind="square") - folded = folded_2 - summed = summed_2 - v2 = reshard(folded, layout=((8, 16), (16, 1), {thread_2.warp @ B()}), storage=gmem) - v3 = reshard(summed, layout=((8, 16), (16, 1), {thread_2.warp @ B()}), storage=gmem) - seeded = reshard(seed, layout=((2 @ thread_2.warp, 4 @ thread_2.lane, 2), (8, 2, 1)), storage=rmem) - v4 = reshard(seeded, layout=((16,), (1,), {thread_2.warp @ B()}), storage=gmem) - return (gathered, v2, v3, v4) + gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) + seeded = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) + folded = reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage=rmem) + summed = reshard(mixed, layout=((8, 16), {mesh.warp @ B()}), storage=rmem) + for _ in range(3): + summed_2 = add(summed, summed) + folded_2 = unary(folded, kind="square") + folded = folded_2 + summed = summed_2 + v4 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) + v5 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) + v6 = reshard(seeded, layout=Layout((16,), (1,)), storage=gmem) + return (gathered, v4, v5, v6) diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py index 9238baa1..d7161fba 100644 --- a/tests/fixtures/inspection/type_printer_sugar.py +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -27,55 +27,67 @@ attrs=(S(0), S(1)), mesh=_WARP_LANE, ) +_WEIGHT = ShardLayout( + layout=Layout((8, 16), None), + attrs=(B(), P("max")), + mesh=_WARP_LANE, +) +_MIXED = ShardLayout( + layout=Layout((4, 2, 16), None), + attrs=(S(0), B(), P("sum")), + mesh=Mesh( + (Topology("cta", 4), Topology("thread", 8)), + Layout((4, 2, 4), (8, 4, 1)), + names=("tile", "warp", "lane"), + ), +) +_ESCAPED = ShardLayout( + layout=Layout((2, 4, 16), None), + attrs=(S(0), B()), + mesh=_WARP_LANE, +) @module(entry="composed_mesh_pipeline", target=_H200, topologies=_TOPOLOGIES) class TypePrinterSugar: - @func + @func(mesh=_WARP_LANE) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], acc: Tensor[ (8, 16), "f32", - ((2 @ _WARP_LANE.warp, 4, 16), {_WARP_LANE.lane @ P("sum")}), - "rmem", - ], - mixed: Tensor[ - (8, 16), "f32", - ((8 @ _TILE.tile, 16), {_WARP_LANE.warp @ B(), _WARP_LANE.lane @ P("sum")}), + ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem", ], + mixed: Tensor[(8, 16), "f32", _MIXED, "rmem"], ): with _TILE as cta: - with _WARP_LANE as thr: - composed = tf.reshard( - x, (8 @ cta.tile, 4 @ thr.warp, 16 @ thr.lane), "rmem" - ) - staged = tf.reshard(tf.square(composed), (8 @ cta.tile, 4, 16), "smem") - narrowed = tf.cast(staged, dtype="bf16") - swapped = tf.transpose(narrowed, perm=(0, 2, 1)) - gathered = tf.reshard(swapped, (8, 16, 4), "gmem") - seeded = tf.reshard(seed, _FRAGMENT, "rmem") - folded = tf.reshard(acc, (8 @ thr.warp, 16 @ thr.lane), "rmem") - summed = tf.reshard( - mixed, ((8, 16), {cta.tile @ B(), thr.warp @ B(), thr.lane @ B()}), "rmem" - ) - for _ in range(3): - folded = tf.square(folded) - summed = tf.add(summed, summed) - return ( - gathered, - tf.reshard(folded, (8, 16), "gmem"), - tf.reshard(summed, (8, 16), "gmem"), - tf.reshard(seeded, (16,), "gmem"), - ) + composed = tf.reshard( + x, (8 @ cta.tile, 4 @ mesh.warp, 16 @ mesh.lane), "rmem" + ) + staged = tf.reshard(tf.square(composed), (8 @ cta.tile, 4, 16), "smem") + narrowed = tf.cast(staged, dtype="bf16") + swapped = tf.transpose(narrowed, perm=(0, 2, 1)) + gathered = tf.reshard(swapped, (8, 16, 4), "gmem") + seeded = tf.reshard(seed, _FRAGMENT, "rmem") + folded = tf.reshard(acc, (8 @ mesh.warp, 16 @ mesh.lane), "rmem") + summed = tf.reshard( + mixed, ((8, 16), {mesh.warp @ B(), mesh.lane @ B()}), "rmem" + ) + for _ in range(3): + folded = tf.square(folded) + summed = tf.add(summed, summed) + return ( + gathered, + tf.reshard(folded, (8, 16), "gmem"), + tf.reshard(summed, (8, 16), "gmem"), + tf.reshard(seeded, (16,), "gmem"), + ) @func def nested_loop_tuple( x: Tensor[(8, 16), "f32"], - weight: Tensor[ - (8, 16), "f32", ((8, 16), {_WARP_LANE.warp @ B(), _WARP_LANE.lane @ P("max")}) - ], + weight: Tensor[(8, 16), "f32", _WEIGHT], ): with _LANES as lanes: split = tf.reshard(x, (8 @ lanes.lane, 16), "rmem") @@ -109,5 +121,5 @@ def named_and_out_of_scope( that boundary, so its target names a mesh that is not its own scope. """ mine = tf.reshard(x, (8 @ mesh.tile, 16), "rmem") # noqa: F821 - escaped = tf.reshard(x, ((8 @ _WARP_LANE.warp, 16), {_WARP_LANE.lane @ B()}), "rmem") + escaped = tf.reshard(x, _ESCAPED, "rmem") return tf.reshard(mine, (8, 16), "gmem"), held, frag, escaped diff --git a/tests/parser/test_calls.py b/tests/parser/test_calls.py index 946bf07e..707ac3ad 100644 --- a/tests/parser/test_calls.py +++ b/tests/parser/test_calls.py @@ -25,6 +25,8 @@ from tilefoundry.parser import ParseError from tilefoundry.target import CpuTarget, CudaTarget +_EXTERNAL_PLACEMENT_MESH = Mesh(("cta",), (2,), names=("tile",)) + def test_matmul_layout_literals_are_parser_checked() -> None: assert get_args(MatMul.a_layout.annotation) == ("MK", "KM") @@ -346,7 +348,7 @@ def valueful_escape(x: Tensor[(2,), "f32"]): def test_mesh_binding_does_not_escape_its_with_scope() -> None: """A mesh alias is removed with its lexical frame after the with body.""" - with pytest.raises(ParseError, match="'mesh' is not an active Mesh"): + with pytest.raises(ParseError, match="'mesh' is not a lexical Mesh binding"): @module( entry="escaped_mesh", @@ -499,7 +501,7 @@ def root(x: Tensor[(8, 8), "f32"]): def test_a_boundary_rejects_a_bare_undeclared_mesh_name() -> None: - with pytest.raises(ParseError, match="'nope' is not an active Mesh"): + with pytest.raises(ParseError, match="'nope' is not a lexical Mesh binding"): @module( entry="root", @@ -520,3 +522,17 @@ class MissingTopology: @func(mesh=Mesh(("cta",), (8,), names=("b",))) def run(x: Tensor[(8 @ mesh.b, 8), "f32"]): # noqa: F821 return tf.add(x, x) + + +def test_placement_rejects_an_external_mesh_axis_binding() -> None: + with pytest.raises(ParseError, match="'_EXTERNAL_PLACEMENT_MESH' is not a lexical Mesh binding"): + + @module( + entry="run", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 2),), + ) + class ExternalMeshPlacement: + @func + def run(x: Tensor[(2 @ _EXTERNAL_PLACEMENT_MESH.tile,), "f32"]): # noqa: F821 + return x From b6b7b33503f0d05c1a4cebacb772649e8cc01846 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 22:56:20 +0800 Subject: [PATCH 19/21] fix(inspection): omit loop result type comments --- docs/spec/inspection.md | 4 ++-- src/tilefoundry/inspection/python_printer.py | 2 +- tests/fixtures/inspection/type_printer_sugar.analyzed.txt | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index d6e29d01..3cf613fb 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -240,8 +240,8 @@ binding in its body. A named constant or an out-of-scope mesh does not create a binding; types that refer to one use the verbose form. Analysis annotations use the same type text for value-producing statements, but -a structural `MeshRegion` line MUST NOT restate the aggregate type of its body. -The region may still carry selected analysis metadata. +a structural `LoopRegion` or `MeshRegion` line MUST NOT restate the aggregate +type of its body. The region may still carry selected analysis metadata. ### 2.6 Specialization printing diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index 8e444a14..d0f24e67 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -296,7 +296,7 @@ def _comments(expr: Expr, options: PythonPrintOptions, printer: PythonPrinter, c the boundary between those two languages. """ comments: list[str] = [] - if options.show_types and not isinstance(expr, MeshRegion): + if options.show_types and not isinstance(expr, (LoopRegion, MeshRegion)): comments.append(_compact_type(expr.type, printer, ctx)) for metadata_type in options.comment_metadata_types: metadata = get_metadata(expr, metadata_type) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index 36585389..f3f70372 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -28,7 +28,7 @@ def composed_mesh_pipeline( 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): # Tuple[Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"], Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]]; loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete + 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 folded = v11 From d9cd3b2e0b82f7cb5bca676022f394874788be53 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 14 Sep 2026 23:43:49 +0800 Subject: [PATCH 20/21] fix(inspection): stabilize sliced mesh type annotations --- src/tilefoundry/inspection/print_context.py | 29 ++++++++++++++++++-- src/tilefoundry/inspection/python_printer.py | 2 +- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src/tilefoundry/inspection/print_context.py b/src/tilefoundry/inspection/print_context.py index bac8c8ca..0f031755 100644 --- a/src/tilefoundry/inspection/print_context.py +++ b/src/tilefoundry/inspection/print_context.py @@ -2,6 +2,7 @@ from __future__ import annotations +from contextlib import contextmanager from math import prod from tilefoundry.ir.types.shard.int_tuple import flatten @@ -18,6 +19,7 @@ def __init__(self) -> None: self._dim_declarations: dict[str, tuple[object, str]] = {} self._mesh_bindings: list[tuple[Mesh, str]] = [] self._used_scope_names: set[str] = set() + self._type_annotation_surface = False def use(self, rendered: PythonExpr | str) -> str: if isinstance(rendered, PythonExpr): @@ -69,6 +71,16 @@ def push_mesh(self, mesh: Mesh, name: str) -> None: def pop_mesh(self) -> None: self._mesh_bindings.pop() + @contextmanager + def type_annotation_surface(self): + """Prefer a parent binding when a type refers to a sliced mesh.""" + previous = self._type_annotation_surface + self._type_annotation_surface = True + try: + yield + finally: + self._type_annotation_surface = previous + def mesh_alias(self, mesh: Mesh) -> str | None: for bound, name in reversed(self._mesh_bindings): if bound is mesh: @@ -94,7 +106,15 @@ def mesh_axis_alias(self, mesh: Mesh, axis: int) -> str | None: target_topology = next( (topology for topology in mesh.topologies if topology.name == target_level), None ) - for bound, alias in reversed(self._mesh_bindings): + for binding_index in range(len(self._mesh_bindings) - 1, -1, -1): + bound, alias = self._mesh_bindings[binding_index] + if ( + self._type_annotation_surface + and bound is mesh + and isinstance(mesh.layout, ComposedLayout) + and mesh.layout.offset != 0 + ): + continue if not bound.names or target_name not in bound.names: continue bound_levels = self._axis_levels(bound) @@ -124,7 +144,12 @@ def mesh_slice(self, mesh: Mesh) -> str | None: @staticmethod def _slice_from_parent(parent: Mesh, child: Mesh, alias: str) -> str | None: - if not isinstance(parent.layout, Layout): + if not ( + isinstance(parent.layout, Layout) + and isinstance(child.layout, ComposedLayout) + and child.layout.inner is None + and isinstance(child.layout.outer, Layout) + ): return None if parent.topologies != child.topologies or parent.names != child.names: return None diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index d0f24e67..0a72dc6d 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -277,7 +277,7 @@ def _physical_line_count(lines: list[str]) -> int: def _compact_type(ty: object, printer: PythonPrinter, ctx) -> str: """One physical-line, DSL-shaped type annotation for inspection output.""" if isinstance(ty, (TensorType, TupleType)): - with printer.type_surface(): + with ctx.type_annotation_surface(), printer.type_surface(): rendered = printer.visit(ty, ctx) return " ".join(rendered.split()) return repr(ty) From deae3f6168a854442070d76ec18eb12a1664d705 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 15 Sep 2026 00:30:15 +0800 Subject: [PATCH 21/21] docs(tutorial): refresh showcase analysis output --- docs/tutorial/showcase.ipynb | 4 ++-- docs/tutorial/showcase.md | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index bb216ef2..a7aa768f 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=(1, 256, 8 @ cta.head, 32), 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 v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ cta.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 @ cta.head, 32), 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 v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ cta.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=(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" } ], "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 5510a9b9..45bb5856 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -870,11 +870,11 @@ for needle in ("reshard(w_q", "reshard(w_o"): # peak-footprint=gmem:10945036,smem:16960 # roofline ideal-ns=11213 bound-by=memory - v1 = reshard(w_q, layout=(1, 256, 8 @ cta.head, 32), 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 - v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ cta.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 + 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 - v42 = reshard(w_o, layout=(1, 256, 8 @ cta.head, 32), 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 - v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ cta.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 + 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 ``` ## 6. Stream the KV cache @@ -1022,7 +1022,7 @@ for needle in ("slice(k_cache", "cache_update(k_cache"): # 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 + 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 ```