From a7f79d3ab2aa452e825a5a15af626746f711c293 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Fri, 25 Sep 2026 22:06:42 +0800 Subject: [PATCH 1/8] refactor(mesh): support layered mesh composition --- docs/spec/code-organization.md | 11 +- docs/spec/shard.md | 36 ++- src/tilefoundry/analysis/compute_cost.py | 4 +- src/tilefoundry/ir/hir/specialize.py | 10 +- src/tilefoundry/ir/mesh_scope.py | 91 +------ src/tilefoundry/ir/types/mesh.py | 237 +++++++++++++----- src/tilefoundry/parser/ast_pattern.py | 4 +- src/tilefoundry/parser/pattern_nodes.py | 8 +- src/tilefoundry/visitor_registry/typeinfer.py | 5 +- tests/analysis/test_analysis_invariants.py | 5 +- tests/fixtures/meshes.py | 9 + tests/ir/test_function_call_typeinfer.py | 8 +- tests/ir/test_visitor.py | 5 +- tests/ir/types/test_local_layout.py | 21 +- tests/ir/types/test_mesh.py | 60 ++++- tests/ir/types/test_tensor_type.py | 7 +- tests/ops/ir/test_arange.py | 4 +- tests/ops/ir/test_argmax.py | 23 +- tests/ops/ir/test_binary.py | 14 +- tests/ops/ir/test_cache_update.py | 14 +- tests/ops/ir/test_cast.py | 5 +- tests/ops/ir/test_clamp.py | 4 +- tests/ops/ir/test_concat.py | 11 +- tests/ops/ir/test_conv2d.py | 12 +- tests/ops/ir/test_layer_norm.py | 6 +- tests/ops/ir/test_matmul.py | 7 +- tests/ops/ir/test_quant.py | 39 ++- tests/ops/ir/test_relu.py | 4 +- tests/ops/ir/test_repeat_interleave.py | 11 +- tests/ops/ir/test_reshape.py | 19 +- tests/ops/ir/test_reshard.py | 4 +- tests/ops/ir/test_rms_norm.py | 11 +- tests/ops/ir/test_rope.py | 6 +- tests/ops/ir/test_sigmoid.py | 4 +- tests/ops/ir/test_silu.py | 4 +- tests/ops/ir/test_slice.py | 9 +- tests/ops/ir/test_softmax.py | 4 +- tests/ops/ir/test_softplus.py | 4 +- tests/ops/ir/test_split.py | 9 +- tests/ops/ir/test_stack.py | 29 ++- tests/ops/ir/test_tanh.py | 4 +- tests/ops/ir/test_topk.py | 25 +- tests/ops/ir/test_transpose.py | 15 +- tests/ops/ir/test_unary.py | 11 +- tests/ops/ir/test_where.py | 11 +- 45 files changed, 538 insertions(+), 306 deletions(-) create mode 100644 tests/fixtures/meshes.py diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 31123dc6..bc21e6aa 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -19,8 +19,8 @@ truth for the directory's structure and invariants. |---|---|---| | `ir/core/` | [core-ir](./core-ir.md) | Shared node algebra: `Module` / `Expr` / `Var` / `Constant` / `Tuple` / `Op` / `Call` / `Stmt` (base class) / `OpSchema` / `ParamDef` / call-graph and ownership queries / typed metadata attach-detach and diagnostics / `@register_op` / `@register_alias` / `op_registry` / `errors`. | | `ir/types/` | [types](./types.md) | Type-system root: `Type` / `TensorType` / `TupleType` / `UnitType` / `CallableType` / `DType` / `StorageKind` / `resolve_storage` / local projections (`local_type_of`) / tensor-leaf, byte-by-storage, and topology-extent queries / `dim.*` (with their typeinfer). | -| `ir/types/{int_tuple,stride,layout,layout_algebra,shard_layout,mesh}.py` | [shard](./shard.md) | `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`), filed as CuTe files them: int tuples (`flatten` / `unflatten` / `repeat_like` / `product`), strides (`compact_major` / `idx2crd` / `crd2idx`), layouts and the algebra over them each in their own module. | -| `ir/mesh_scope.py` | [shard](./shard.md) | Which scope a statement stands inside and what it admits: `merge_mesh`, `device_layout`, `covered_by_scope`, `check_topology`. Neither a type nor a visitor, so it sits beside `ir/isl_interop.py` rather than in either. | +| `ir/types/{int_tuple,stride,layout,layout_algebra,shard_layout,mesh}.py` | [shard](./shard.md) | `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`), filed as CuTe files them: int tuples (`flatten` / `unflatten` / `repeat_like` / `product`), strides (`compact_major` / `idx2crd` / `crd2idx`), layouts and the algebra over them each in their own module; mesh construction and separation stay with `Mesh`. | +| `ir/mesh_scope.py` | [shard](./shard.md) | Which scope a statement stands inside and what it admits: `device_layout`, `covered_by_scope`, `check_topology`. Neither a type nor a visitor, so it sits beside `ir/isl_interop.py` rather than in either. | | `ir/constraints/` | [parser](./parser.md) | Authored `where(layout=..., mesh=..., storage=...)` constraint records: the shared base plus layout, mesh, and storage constraints, attached by the parser and read back by the Python printer. | | `ir/visitor.py` | [visitor-mutator](./visitor-mutator.md) | `ExprFunctor` / `ExprVisitor` / `ExprWalker` / `ExprCollector` / `ExprCloner` / `BindingSubstitutionCloner` / `StmtVisitor` / `StmtMutator` / `StmtExprMutator`, plus `collect_exprs`, value-operand/function-value queries, and the canonical `PrimFunction` walk and rewrite entries. | | `ir/isl_interop.py` | [types](./types.md) | Interoperation between dimension and shape IR values and isl: expression rendering and decoding, normalization, value ranges, and shape-domain construction. Pure isl operations remain in `utils/isl_utils.py`. | @@ -181,6 +181,13 @@ go through [parser §2](./parser.md#2-syntax-and-rules). **Rule 7 — what template files contain.** `codegen//templates/*.j2` carry boilerplate assembly only; emitters live in Python walkers. +### 2.1 Package export rule + +`tilefoundry.ir.types` re-exports type classes and `make_*` constructors. Other +functions are imported from the module that owns them; for example, +`make_mesh` is available at the package surface, while `separate` is imported +from `tilefoundry.ir.types.mesh`. + ## 3. Multi-agent parallelism guarantee The lock granularity is a single `(node, target)` pair. The naming rules in diff --git a/docs/spec/shard.md b/docs/spec/shard.md index cf45d465..c64a2741 100644 --- a/docs/spec/shard.md +++ b/docs/spec/shard.md @@ -327,18 +327,26 @@ that numbering. Mesh composition uses the following rules: -- `merge_mesh(meshes)` MUST `append` a mesh whose level names are disjoint from +- `make_mesh(*meshes)` MUST append a mesh whose level names are disjoint from those in force, MUST replace them entirely when the inner mesh names every - one, and MUST `replace` the trailing levels when the inner mesh names a - suffix of them, keeping the levels above and their names unchanged. Level - names overlapping in any other way MUST be rejected rather than decomposed. -- `append` and `replace` concatenate the per-level arrangements; no stride or - offset is rescaled, because each level already states its own numbering. -- `merge_mesh(meshes)` invokes `check_topology` on its result. For each named + one, and MUST replace the trailing levels when the inner mesh names a suffix + of them, keeping the levels above and their names unchanged. Level names + overlapping in any other way MUST be rejected rather than decomposed. +- Append concatenates the per-level arrangements. When an appended mesh is + sliced, its per-level start is re-encoded in the combined device numbering + and retained in the result's `ComposedLayout.offset`. +- Replacing a suffix with a sliced mesh MUST be rejected. The slice belongs on + the combined multi-level mesh, where every axis contributing to its device + offset is stated together. Replacing an unsliced suffix and replacing the + whole mesh retain their existing behavior. +- `make_mesh(*meshes)` invokes `check_topology` on its result. For each named level with a concrete declared extent, its position count MUST NOT exceed that extent; symbolic extents are deferred until dimensions are bound. A level stating a run is already bounded by `Mesh.__getitem__` and is not checked again. +- `separate(mesh)` is the inverse construction: it returns one single-level + `Mesh` per topology, with that level's arrangement, slice start, and share of + the axis names. It is imported from `tilefoundry.ir.types.mesh`. HIR `MeshRegion` applies this composition only at its body boundary. Its `args` are evaluated in the enclosing scope and are not recomposed merely because the @@ -347,9 +355,11 @@ value is consumed by a region. - constraints: - `Topology` construction rejects a `None` size. `Mesh` construction rejects a `None` entry in its layout shape and normalizes the layout into one mode - per level; beyond that it performs no position-consistency check. Its `topologies` field is a - `tuple[Topology, ...]`; helpers such as `make_mesh` construct that tuple for - handwritten Python. + per level; beyond that it performs no position-consistency check. Its + `topologies` field is a `tuple[Topology | str, ...]`: a single-level mesh + may temporarily name its topology by string until the parser resolves it + from the enclosing module declaration, while a multi-level mesh MUST carry + `Topology` values. Topology names in one mesh MUST be unique. - The author surface is `with Mesh(("cta",), layout=(128,)) as cta:`. The parser resolves the non-empty tuple of declared topology names to the `Topology` tuple before it constructs the record. A bare string and the @@ -369,7 +379,9 @@ value is consumed by a region. of the levels below it; a stride that division does not divide exactly MUST be refused. A Mesh naming one level takes every axis into its single mode, which needs no division and therefore no declared extent. A Mesh naming - several MUST NOT also be sliced. + several may be sliced: the key addresses all level axes in order, every + level retains its own arrangement, and the slice offset is the sum of each + selected start times that axis's step in device numbering. - Nested single-level Mesh scopes compose to exactly that shape: the axes join outermost first and each outer stride is scaled by the positions below it. A value distributed at two levels at once may therefore be written either @@ -395,7 +407,7 @@ value is consumed by a region. The placed-layout constructor has one additional guard: a single layout may split a named level only once. If two distinct meshes used by one placed layout name the same topology level, parsing MUST reject that layout at its source -node. This is a layout-construction rule, independent of `merge_mesh()`'s +node. This is a layout-construction rule, independent of `make_mesh()`'s scope-composition rules. ### 5.1 `Placement` diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index 91b5a3d6..6b1fc2bc 100644 --- a/src/tilefoundry/analysis/compute_cost.py +++ b/src/tilefoundry/analysis/compute_cost.py @@ -9,9 +9,9 @@ from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion -from tilefoundry.ir.mesh_scope import merge_mesh from tilefoundry.ir.types import DType, Mesh from tilefoundry.ir.types.layout import ComposedLayout, get, size +from tilefoundry.ir.types.mesh import make_mesh from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.visitor_registry.contexts import CostContext, FunctionScope, TrafficBytes from tilefoundry.visitor_registry.visitors import CostEvaluator @@ -257,7 +257,7 @@ def visit_MeshRegion(self, expr: MeshRegion, ctx: ComputeCostContext) -> None: child = next(item for item in ctx.current.children if item.owner is expr) for arg in expr.args: self.visit(arg, ctx) - mesh = merge_mesh((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh + mesh = make_mesh(ctx.current_mesh, expr.mesh) if ctx.current_mesh else expr.mesh topologies = ctx.module.effective_topologies() positions = { unit: _scope_position_count(mesh, unit, topologies) for unit in ctx.locals_by_unit diff --git a/src/tilefoundry/ir/hir/specialize.py b/src/tilefoundry/ir/hir/specialize.py index 2ef70f1c..2b81ec1c 100644 --- a/src/tilefoundry/ir/hir/specialize.py +++ b/src/tilefoundry/ir/hir/specialize.py @@ -17,8 +17,8 @@ from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion -from tilefoundry.ir.mesh_scope import merge_mesh from tilefoundry.ir.types.dim import is_dim_expr +from tilefoundry.ir.types.mesh import make_mesh from tilefoundry.ir.types.substitute import ( dim_vars_by_name, has_symbolic_dims, @@ -223,18 +223,14 @@ def visit_MeshRegion(self, expr: MeshRegion, ctx: InstantiateContext) -> Expr: mesh = substitute_mesh_dims(expr.mesh, ctx.dims) new_args = tuple(self.visit(arg, ctx) for arg in expr.args) new_params = tuple( - param - if new_arg.type == param.type - else Var(type=new_arg.type, name=param.name) + param if new_arg.type == param.type else Var(type=new_arg.type, name=param.name) for param, new_arg in zip(expr.params, new_args, strict=True) ) for old, new in zip(expr.params, new_params, strict=True): if old is not new: ctx.subst[id(old)] = new current_mesh = ( - merge_mesh((ctx.type_ctx.current_mesh, mesh)) - if ctx.type_ctx.current_mesh - else mesh + make_mesh(ctx.type_ctx.current_mesh, mesh) if ctx.type_ctx.current_mesh else mesh ) body_ctx = dataclasses.replace( ctx, diff --git a/src/tilefoundry/ir/mesh_scope.py b/src/tilefoundry/ir/mesh_scope.py index ce2f870a..38166f92 100644 --- a/src/tilefoundry/ir/mesh_scope.py +++ b/src/tilefoundry/ir/mesh_scope.py @@ -10,36 +10,11 @@ from __future__ import annotations from tilefoundry.ir.types.int_tuple import flatten, product -from tilefoundry.ir.types.layout import ComposedLayout, Layout, get, rank, size +from tilefoundry.ir.types.layout import ComposedLayout, Layout, size from tilefoundry.ir.types.layout_algebra import is_inverse_projectable -from tilefoundry.ir.types.mesh import Mesh +from tilefoundry.ir.types.mesh import Mesh, _levels, _starts from tilefoundry.ir.types.storage import StorageKind, resolve_storage -from tilefoundry.ir.types.stride import compact_major, idx2crd - - -def _levels(mesh: Mesh) -> tuple[Layout, ...]: - """Each level's own arrangement: mode ``i`` of the mesh is level ``i``.""" - stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout - if stated is None: - raise ValueError( - "a mesh whose slice states an identity box states no arrangement of " - "its own, so its levels select nothing" - ) - return tuple(get(stated, index) for index in range(rank(stated))) - - -def _starts(mesh: Mesh) -> tuple[int, ...]: - """Where each level's run begins, read out of the offset the mesh states. - - The offset is one index in the numbering the device gives every position, - and the levels are its shape, so the coordinate it stands for is what each - level starts at. - """ - offset = mesh.layout.offset if isinstance(mesh.layout, ComposedLayout) else 0 - sizes = tuple(getattr(topology, "size", 1) for topology in mesh.topologies) - if not isinstance(offset, int) or any(not isinstance(one, int) for one in sizes): - return (0,) * len(sizes) - return tuple(idx2crd(offset, sizes, compact_major(sizes))) +from tilefoundry.ir.types.stride import compact_major def device_layout(mesh: Mesh) -> Layout: @@ -180,70 +155,10 @@ def check_topology(mesh: Mesh) -> None: ) -def _joined(topologies: tuple, levels: tuple, names: tuple) -> Mesh: - """One mesh out of the levels it names, each stating its own arrangement.""" - if len(levels) == 1: - return Mesh(topologies=topologies, layout=levels[0], names=names) - return Mesh( - topologies=topologies, - layout=Layout( - shape=tuple(tuple(flatten(one.shape)) for one in levels), - strides=tuple(tuple(flatten(one.strides)) for one in levels), - ), - names=names, - ) - - -def _named(mesh: Mesh) -> tuple[str, ...]: - return tuple(getattr(topology, "name", topology) for topology in mesh.topologies) - - -def merge_mesh(meshes: "tuple[Mesh, ...]") -> Mesh: - """The scope in force once each of *meshes* has been entered in turn. - - A scope naming levels none of those in force name is appended below them. - One naming every level in force replaces them. One naming a suffix of them - replaces that suffix and keeps what is above. Any other overlap is refused - rather than decomposed: which positions the half-named levels would then - state is nobody's statement. No stride is rescaled, because every level - already states its own numbering. - """ - result = meshes[0] - for inner in meshes[1:]: - here, there = _named(result), _named(inner) - if set(here).isdisjoint(there): - result = _joined( - (*result.topologies, *inner.topologies), - (*_levels(result), *_levels(inner)), - (*result.names, *inner.names), - ) - elif set(here) <= set(there): - result = inner - elif len(there) < len(here) and here[-len(there) :] == there: - kept = len(here) - len(there) - above = _levels(result)[:kept] - named = sum(len(flatten(one.shape)) for one in above) - result = _joined( - (*result.topologies[:kept], *inner.topologies), - (*above, *_levels(inner)), - (*result.names[:named], *inner.names), - ) - else: - shared = sorted(set(here) & set(there)) - unnamed = sorted(set(here) - set(there)) - raise ValueError( - f"{shared} named again while {unnamed} is not; a scope either " - "replaces the levels in force or adds levels below them" - ) - check_topology(result) - return result - - __all__ = [ "check_topology", "device_layout", "covered_by_scope", - "merge_mesh", "mesh_scope_matches_required_scope", "states_consistent_positions", "storage_reaches", diff --git a/src/tilefoundry/ir/types/mesh.py b/src/tilefoundry/ir/types/mesh.py index 8aed90c0..c571a65b 100644 --- a/src/tilefoundry/ir/types/mesh.py +++ b/src/tilefoundry/ir/types/mesh.py @@ -1,11 +1,10 @@ from __future__ import annotations -import math from dataclasses import dataclass from tilefoundry.ir.types.layout import ComposedLayout, Layout, LayoutBase, flatten, get from tilefoundry.ir.types.layout import rank as _rank -from tilefoundry.ir.types.stride import compact_row_major, try_compact_major +from tilefoundry.ir.types.stride import compact_major, compact_row_major, crd2idx, idx2crd from tilefoundry.ir.types.tensor_type import ShapeDim @@ -44,6 +43,15 @@ class Mesh: names: tuple[str, ...] = () def __post_init__(self) -> None: + if len(self.topologies) > 1 and any( + not isinstance(topology, Topology) for topology in self.topologies + ): + raise ValueError("a multi-level Mesh requires Topology values") + topology_names = tuple( + getattr(topology, "name", topology) for topology in self.topologies + ) + if len(set(topology_names)) != len(topology_names): + raise ValueError(f"Mesh topology names must be unique, got {topology_names!r}") object.__setattr__(self, "layout", _nested(self.layout, tuple(self.topologies))) for axis, extent in enumerate(flatten(self.layout.shape)): if extent is None: @@ -57,60 +65,75 @@ def __getitem__(self, key) -> "Mesh": Missing axes are full slices; integers select extent one. The result preserves topology and names while recording the sub-box as a - ``ComposedLayout``. Only a mesh naming one level is sliced: a slice and - a level boundary would otherwise both decide which positions these are. + ``ComposedLayout``. Each level retains its own arrangement, while the + slice offset uses the device's numbering across all levels. See [shard §5](docs/spec/shard.md#5-mesh). """ - if len(self.topologies) != 1: - raise ValueError("cannot slice a mesh that names several levels") if isinstance(self.layout, ComposedLayout): raise ValueError("cannot slice an already-sliced mesh (nested slice unsupported)") - level = get(self.layout, 0) - shape = level.shape - strides = level.strides - rank = len(shape) + levels = _levels(self) + rank = sum(len(flatten(level.shape)) for level in levels) keys = key if isinstance(key, tuple) else (key,) if len(keys) > rank: raise ValueError(f"mesh slice has {len(keys)} indices but the mesh has {rank} axes") keys = keys + (slice(None),) * (rank - len(keys)) - sub_shape: list[int] = [] + sub_levels: list[Layout] = [] offset = 0 - for axis, (k, extent, stride) in enumerate(zip(keys, shape, strides)): - if not isinstance(extent, int) or not isinstance(stride, int): - raise ValueError(f"cannot slice mesh axis {axis} with a dynamic extent/stride") - if isinstance(k, int): - start = k + extent if k < 0 else k - if not (0 <= start < extent): - raise ValueError( - f"mesh slice index {k} out of range for axis {axis} (extent {extent})" - ) - sel = 1 - elif isinstance(k, slice): - if k.step not in (None, 1): - raise ValueError(f"mesh slice step must be 1 (axis {axis})") - start = 0 if k.start is None else (k.start + extent if k.start < 0 else k.start) - stop = extent if k.stop is None else (k.stop + extent if k.stop < 0 else k.stop) - if not (0 <= start <= stop <= extent): + axis = 0 + units = ( + (1,) + if len(self.topologies) == 1 + else compact_major(tuple(topology.size for topology in self.topologies)) + ) + for level, unit in zip(levels, units): + level_shape = tuple(flatten(level.shape)) + stated = level.strides + level_strides = ( + tuple(flatten(stated)) if stated is not None else compact_row_major(level_shape) + ) + sub_shape: list[int] = [] + for k, extent, stride in zip( + keys[axis : axis + len(level_shape)], level_shape, level_strides + ): + if not isinstance(extent, int) or not isinstance(stride, int): + raise ValueError(f"cannot slice mesh axis {axis} with a dynamic extent/stride") + if isinstance(k, int): + start = k + extent if k < 0 else k + if not (0 <= start < extent): + raise ValueError( + f"mesh slice index {k} out of range for axis {axis} (extent {extent})" + ) + selected = 1 + elif isinstance(k, slice): + if k.step not in (None, 1): + raise ValueError(f"mesh slice step must be 1 (axis {axis})") + start = 0 if k.start is None else (k.start + extent if k.start < 0 else k.start) + stop = extent if k.stop is None else (k.stop + extent if k.stop < 0 else k.stop) + if not (0 <= start <= stop <= extent): + raise ValueError( + f"mesh slice {k.start}:{k.stop} out of range for axis " + f"{axis} (extent {extent})" + ) + selected = stop - start + if selected == 0: + raise ValueError(f"mesh slice selects an empty range on axis {axis}") + else: raise ValueError( - f"mesh slice {k.start}:{k.stop} out of range for axis " - f"{axis} (extent {extent})" + f"mesh slice index must be int or slice, got {type(k).__name__}" ) - sel = stop - start - if sel == 0: - raise ValueError(f"mesh slice selects an empty range on axis {axis}") - else: - raise ValueError(f"mesh slice index must be int or slice, got {type(k).__name__}") - offset += start * stride - sub_shape.append(sel) + offset += start * stride * unit + sub_shape.append(selected) + axis += 1 + sub_levels.append(Layout(tuple(sub_shape), level_strides)) return Mesh( topologies=self.topologies, layout=ComposedLayout( inner=None, offset=offset, - outer=Layout(shape=(tuple(sub_shape),), strides=(tuple(strides),)), + outer=_joined_layout(tuple(sub_levels)), ), names=self.names, ) @@ -138,11 +161,6 @@ def _nested(layout, topologies: tuple) -> "Layout | ComposedLayout": extents = tuple(flatten(layout)) layout = Layout(shape=extents, strides=compact_row_major(extents)) if isinstance(layout, ComposedLayout): - if len(topologies) != 1: - raise ValueError( - "a mesh naming several levels states one arrangement per level; a " - "slice and a level boundary cannot both decide which positions these are" - ) if layout.outer is None or _levelled(layout.outer, topologies): return layout return ComposedLayout( @@ -162,10 +180,10 @@ def _nested(layout, topologies: tuple) -> "Layout | ComposedLayout": below = 1 for topology in reversed(topologies): units.insert(0, below) - size = getattr(topology, "size", None) + size = topology.size if not isinstance(size, int) or isinstance(size, bool) or size < 1: raise ValueError( - f"mesh level {getattr(topology, 'name', topology)!r} states extent " + f"mesh level {topology.name!r} states extent " f"{size!r}; cutting one arrangement at the level boundaries needs " "each of their position counts" ) @@ -212,33 +230,116 @@ def _nested(layout, topologies: tuple) -> "Layout | ComposedLayout": return Layout(shape=tuple(shape), strides=tuple(strides)) -def make_mesh( +def _levels(mesh: Mesh) -> tuple[Layout, ...]: + """Each level's arrangement, in that level's own numbering.""" + stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout + if stated is None: + raise ValueError( + "a mesh whose slice states an identity box states no arrangement of " + "its own, so its levels select nothing" + ) + return tuple(get(stated, index) for index in range(_rank(stated))) + + +def _starts(mesh: Mesh) -> tuple[int, ...]: + """Where each level's run starts, decoded from the device-numbered offset.""" + offset = mesh.layout.offset if isinstance(mesh.layout, ComposedLayout) else 0 + if len(mesh.topologies) == 1 and isinstance(offset, int): + return (offset,) + sizes = tuple(topology.size for topology in mesh.topologies) + if not isinstance(offset, int) or any(not isinstance(one, int) for one in sizes): + return (0,) * len(sizes) + return tuple(idx2crd(offset, sizes, compact_major(sizes))) + + +def _joined_layout(levels: tuple[Layout, ...]) -> Layout: + if len(levels) == 1: + return levels[0] + return Layout( + shape=tuple(tuple(flatten(level.shape)) for level in levels), + strides=tuple(tuple(flatten(level.strides)) for level in levels), + ) + - layout_shape: tuple, - names: "tuple[str, ...] | None" = None, - topology: "str | Topology" = "gpu", +def _joined( + topologies: tuple[Topology, ...], + levels: tuple[Layout, ...], + starts: tuple[int, ...], + names: tuple[str, ...], + *, + sliced: bool, ) -> Mesh: - """Convenience constructor for a ``Mesh`` with the given axis extents and C-order strides. + layout: Layout | ComposedLayout = _joined_layout(levels) + if sliced: + sizes = tuple(getattr(topology, "size", None) for topology in topologies) + if not all(isinstance(size, int) for size in sizes): + raise ValueError("joining sliced meshes needs static topology extents") + layout = ComposedLayout(None, crd2idx(starts, sizes, compact_major(sizes)), layout) + return Mesh(topologies, layout, names) + + +def _named(mesh: Mesh) -> tuple[str, ...]: + return tuple(getattr(topology, "name", topology) for topology in mesh.topologies) + + +def make_mesh(*meshes: Mesh) -> Mesh: + """Compose nested mesh scopes, preserving slices in device numbering.""" + if not meshes: + raise ValueError("make_mesh requires at least one mesh") + result = meshes[0] + for inner in meshes[1:]: + here, there = _named(result), _named(inner) + if set(here).isdisjoint(there): + result = _joined( + (*result.topologies, *inner.topologies), + (*_levels(result), *_levels(inner)), + (*_starts(result), *_starts(inner)), + (*result.names, *inner.names), + sliced=isinstance(result.layout, ComposedLayout) + or isinstance(inner.layout, ComposedLayout), + ) + elif set(here) <= set(there): + result = inner + elif len(there) < len(here) and here[-len(there) :] == there: + if isinstance(inner.layout, ComposedLayout): + raise ValueError("cannot replace a mesh suffix with a sliced mesh") + kept = len(here) - len(there) + above = _levels(result)[:kept] + named = sum(len(flatten(level.shape)) for level in above) + result = _joined( + (*result.topologies[:kept], *inner.topologies), + (*above, *_levels(inner)), + (*_starts(result)[:kept], *_starts(inner)), + (*result.names[:named], *inner.names), + sliced=isinstance(result.layout, ComposedLayout), + ) + else: + shared = sorted(set(here) & set(there)) + unnamed = sorted(set(here) - set(there)) + raise ValueError( + f"{shared} named again while {unnamed} is not; a scope either " + "replaces the levels in force or adds levels below them" + ) + from tilefoundry.ir.mesh_scope import check_topology # noqa: PLC0415 - Convenience constructor for a ``Mesh`` with the given (logical) axis - extents and C-order strides. ``names`` defaults to ``a, b, c, ...`` (or - ``g`` for a single axis) so a caller states only the extents instead of - hand-building a ``Mesh``. + check_topology(result) + return result - ``topology`` accepts an explicit ``Topology`` or the ``"gpu"``-shorthand - default; a raw string is resolved here into a real ``Topology`` sized to - the domain. - """ - if names is None: - names = ("g",) if len(layout_shape) == 1 else tuple("abcdef"[: len(layout_shape)]) - if isinstance(topology, str): - topology = Topology(topology, math.prod(layout_shape)) - layout_shape = tuple(layout_shape) - return Mesh( - topologies=(topology,), - layout=Layout(shape=layout_shape, strides=try_compact_major(layout_shape)), - names=tuple(names), - ) + +def separate(mesh: Mesh) -> tuple[Mesh, ...]: + """Split a mesh into one mesh per topology, retaining per-level slices.""" + sliced = isinstance(mesh.layout, ComposedLayout) + names_at = 0 + separated: list[Mesh] = [] + for topology, level, start in zip(mesh.topologies, _levels(mesh), _starts(mesh)): + axis_count = len(flatten(level.shape)) + names = mesh.names[names_at : names_at + axis_count] + layout: Layout | ComposedLayout = level + if sliced: + layout = ComposedLayout(None, start, level) + separated.append(Mesh((topology,), layout, names)) + names_at += axis_count + return tuple(separated) -__all__ = ["Mesh", "Topology", "make_mesh"] +__all__ = ["Mesh", "Topology", "make_mesh", "separate"] diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index ce2fcc43..af4cd911 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -49,7 +49,6 @@ from tilefoundry.ir.hir.tensor.slice import Slice, slice_size from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.isl_interop import normalize_dim -from tilefoundry.ir.mesh_scope import merge_mesh from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.tir.stmts import ( @@ -88,6 +87,7 @@ simplify_dim, ) from tilefoundry.ir.types.layout import LayoutBase +from tilefoundry.ir.types.mesh import make_mesh from tilefoundry.ir.types.shard_layout import canonical_shard_layout from tilefoundry.ir.types.storage import StorageKind, resolve_storage from tilefoundry.ir.types.stride import compact_row_major @@ -299,7 +299,7 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex DISPLAY_NAME=DISPLAY_NAME, compact_row_major=compact_row_major, canonical_shard_layout=canonical_shard_layout, - merge_mesh=merge_mesh, + make_mesh=make_mesh, dim_expr=dim_expr, normalize_dim=normalize_dim, static_dim_value=static_dim_value, diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index c3b9ecfd..6e6ed486 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -658,7 +658,7 @@ def apply(self, value, *, match, context): 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.merge_mesh(meshes) + mesh = meshes[0] if len(meshes) == 1 else runtime.make_mesh(*meshes) source_offsets: dict[int, int] = {} offset = 0 for source in meshes: @@ -3449,11 +3449,7 @@ def _enter_mesh_scope(context, mesh, match): """ infer = _parser_infer_context(context) try: - entered_mesh = ( - runtime.merge_mesh((infer.current_mesh, mesh)) - if infer.current_mesh - else mesh - ) + entered_mesh = runtime.make_mesh(infer.current_mesh, mesh) if infer.current_mesh else mesh except ValueError as error: raise ParseError.from_node(match.node, context, str(error)) from error context.lexical_scope.push_frame() diff --git a/src/tilefoundry/visitor_registry/typeinfer.py b/src/tilefoundry/visitor_registry/typeinfer.py index bb2275e2..138d9086 100644 --- a/src/tilefoundry/visitor_registry/typeinfer.py +++ b/src/tilefoundry/visitor_registry/typeinfer.py @@ -14,9 +14,10 @@ from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.hir.sharding.reshard import Reshard as HirReshard -from tilefoundry.ir.mesh_scope import covered_by_scope, merge_mesh, storage_reaches +from tilefoundry.ir.mesh_scope import covered_by_scope, storage_reaches from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.types.callable_type import callable_type_for +from tilefoundry.ir.types.mesh import make_mesh from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.ir.types.substitute import canonicalize_dims from tilefoundry.ir.types.tensor_type import TupleType, Type @@ -220,7 +221,7 @@ def visit_MeshRegion(self, expr: MeshRegion, ctx: TypeInferContext) -> Type: from tilefoundry.ir.hir.verify import _verify_isolated # noqa: PLC0415 _verify_isolated(expr, ctx) - mesh = merge_mesh((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh + mesh = make_mesh(ctx.current_mesh, expr.mesh) if ctx.current_mesh else expr.mesh return self.visit(expr.body, replace(ctx, current_mesh=mesh, memo=memo)) def visit_Function(self, fn: Function, ctx: TypeInferContext) -> Type: diff --git a/tests/analysis/test_analysis_invariants.py b/tests/analysis/test_analysis_invariants.py index 28170286..3f8b10df 100644 --- a/tests/analysis/test_analysis_invariants.py +++ b/tests/analysis/test_analysis_invariants.py @@ -48,10 +48,11 @@ from tilefoundry.ir.isl_interop import index_set from tilefoundry.ir.types import ( DType, + Layout, + Mesh, TensorType, Topology, TupleType, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -129,7 +130,7 @@ def test_a_boundary_reaching_past_its_operand_is_held_to_what_it_was_handed() -> that rather than from a read nobody could perform. """ cta = Topology("cta", 2) - mesh = make_mesh((2,), ("c",), topology=cta) + mesh = Mesh((cta,), Layout((2,), (1,)), ("c",)) destination = make_shard_tensor_type((8,), mesh=mesh, attrs=(ShardSplit(0),), dtype=DType.f32) update = make_shard_tensor_type((4,), mesh=mesh, attrs=(ShardSplit(0),), dtype=DType.f32) call = Call( diff --git a/tests/fixtures/meshes.py b/tests/fixtures/meshes.py new file mode 100644 index 00000000..156705f2 --- /dev/null +++ b/tests/fixtures/meshes.py @@ -0,0 +1,9 @@ +from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, Topology + +CTA = Mesh((Topology("cta", 4),), Layout((4,), (1,))) +THR = Mesh((Topology("thread", 384),), Layout((384,), (1,))) +CT = Mesh( + (Topology("cta", 4), Topology("thread", 384)), + Layout(((4,), (384,)), ((1,), (1,))), +) +RUN = ComposedLayout(None, 128, Layout(((4,), (128,)), ((1,), (1,)))) diff --git a/tests/ir/test_function_call_typeinfer.py b/tests/ir/test_function_call_typeinfer.py index 3b56b063..dfcffed5 100644 --- a/tests/ir/test_function_call_typeinfer.py +++ b/tests/ir/test_function_call_typeinfer.py @@ -21,8 +21,10 @@ from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.types import ( DType, + Layout, + Mesh, + Topology, TupleType, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -32,7 +34,7 @@ from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor _F = DType.f32 -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _PLAIN = make_tensor_type((4, 8), _F) _SPLIT0 = make_shard_tensor_type((4, 8), mesh=_M, attrs=(Split(0),)) @@ -152,7 +154,7 @@ def test_plain_formal_rejects_shape_or_dtype_mismatch(): def test_function_call_preserves_partial_in_tuple_return(): - mesh_ab = make_mesh((2, 4), ("a", "b")) + mesh_ab = Mesh((Topology("gpu", 8),), Layout((2, 4), (4, 1)), ("a", "b")) partial = make_shard_tensor_type((4, 8), mesh=mesh_ab, attrs=(Broadcast(), Partial("max"))) param = Var(type=_PLAIN, name="x") return_type = TupleType(fields=(_PLAIN, _PLAIN)) diff --git a/tests/ir/test_visitor.py b/tests/ir/test_visitor.py index 2547b241..3b0f5f33 100644 --- a/tests/ir/test_visitor.py +++ b/tests/ir/test_visitor.py @@ -33,8 +33,7 @@ While, ) from tilefoundry.ir.tir.symbol_ref import SymbolRef -from tilefoundry.ir.types import CallableType, DType, TensorType, UnitType, make_mesh -from tilefoundry.ir.types.mesh import Topology +from tilefoundry.ir.types import CallableType, DType, Layout, Mesh, TensorType, Topology, UnitType from tilefoundry.ir.types.storage import StorageKind from tilefoundry.ir.visitor import ( ExprCloner, @@ -355,7 +354,7 @@ def _seq(*items) -> Sequential: While(cond=_var("c"), body=_seq(_eval_call(Copy(), _var("s2"), _var("d2")))), If(cond=_var("c2"), then_body=_seq(), else_body=_seq()), MeshScope( - mesh=make_mesh((2,), topology=Topology(name="chip", size=2)), + mesh=Mesh((Topology(name="chip", size=2),), Layout((2,), (1,)), ("g",)), binding=binding, body=_seq(_eval_call(Copy(), _var("s3"), _var("d3"))), ), diff --git a/tests/ir/types/test_local_layout.py b/tests/ir/types/test_local_layout.py index 994c5773..b58e5469 100644 --- a/tests/ir/types/test_local_layout.py +++ b/tests/ir/types/test_local_layout.py @@ -15,7 +15,7 @@ import pytest -from tilefoundry.ir.types import DType, Mesh, Topology, make_mesh, make_shard_tensor_type +from tilefoundry.ir.types import DType, Layout, Mesh, Topology, make_shard_tensor_type from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.shard_layout import ( Broadcast, @@ -24,13 +24,14 @@ local_layout_and_offset, shard_layout_of, ) +from tilefoundry.ir.types.stride import try_compact_major from tilefoundry.ir.types.utils import local_type_of _GPU, _THREAD = Topology("gpu", 2), Topology("thread", 32) def _mesh(topology: Topology, extents: tuple[int, ...], names: tuple[str, ...]) -> Mesh: - return make_mesh(extents, names, topology=topology) + return Mesh((topology,), Layout(extents, try_compact_major(extents)), names) _CASES = { @@ -120,7 +121,7 @@ def test_every_instance_together_tile_the_tensor(case: str) -> None: def test_the_layout_keeps_the_whole_tensors_strides() -> None: """A slice is the same rows the same distance apart, begun further in.""" - mesh = make_mesh((2,), ("g",), topology=_GPU) + mesh = Mesh((_GPU,), Layout((2,), (1,)), ("g",)) held = make_shard_tensor_type((4, 4), mesh=mesh, attrs=(Split(1),), dtype=DType.f32) shard = shard_layout_of(held.layout) @@ -134,10 +135,8 @@ def test_the_outer_axis_steps_over_what_the_inner_one_holds() -> None: One step of the outer axis clears the eight the inner one holds, rather than the whole the outer axis was already narrowed out of. """ - mesh = make_mesh((4, 8), ("w", "t"), topology=_THREAD) - held = make_shard_tensor_type( - (32,), mesh=mesh, attrs=(Split(0), Split(0)), dtype=DType.f32 - ) + mesh = Mesh((_THREAD,), Layout((4, 8), (8, 1)), ("w", "t")) + held = make_shard_tensor_type((32,), mesh=mesh, attrs=(Split(0), Split(0)), dtype=DType.f32) shard = shard_layout_of(held.layout) offsets = [local_layout_and_offset(shard, (32,), (pid,))[1] for pid in range(32)] @@ -146,10 +145,8 @@ def test_the_outer_axis_steps_over_what_the_inner_one_holds() -> None: def test_a_level_with_no_id_is_left_whole() -> None: """A level the host did not place divides nothing; the device does that.""" - mesh = make_mesh((2, 4), ("g", "c"), topology=_GPU) - held = make_shard_tensor_type( - (8, 4), mesh=mesh, attrs=(Split(0), Split(1)), dtype=DType.f32 - ) + mesh = Mesh((_GPU,), Layout((2, 4), (4, 1)), ("g", "c")) + held = make_shard_tensor_type((8, 4), mesh=mesh, attrs=(Split(0), Split(1)), dtype=DType.f32) shard = shard_layout_of(held.layout) layout, offset = local_layout_and_offset(shard, (8, 4), (None,)) @@ -163,7 +160,7 @@ def test_an_extent_its_mesh_axis_does_not_divide_is_refused() -> None: shape asked about here is one the layout was not built for -- which is the only way the two can disagree, and worth naming rather than slicing. """ - mesh = make_mesh((4,), ("g",), topology=_GPU) + mesh = Mesh((_GPU,), Layout((4,), (1,)), ("g",)) held = make_shard_tensor_type((8,), mesh=mesh, attrs=(Split(0),), dtype=DType.f32) shard = shard_layout_of(held.layout) diff --git a/tests/ir/types/test_mesh.py b/tests/ir/types/test_mesh.py index ea4a4e0a..f32c7fc6 100644 --- a/tests/ir/types/test_mesh.py +++ b/tests/ir/types/test_mesh.py @@ -2,20 +2,23 @@ import pytest +from tests.fixtures.meshes import CT, CTA, RUN, THR from tilefoundry.ir.mesh_scope import ( check_topology, + covered_by_scope, mesh_scope_matches_required_scope, states_consistent_positions, ) -from tilefoundry.ir.types import Layout, Mesh, Topology, make_mesh +from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, Topology, make_mesh from tilefoundry.ir.types.int_tuple import product from tilefoundry.ir.types.layout_algebra import size +from tilefoundry.ir.types.mesh import separate def test_mesh_position_consistency_is_an_explicit_predicate() -> None: - matching = make_mesh((32,), topology="thread") - mismatching = make_mesh((32,), topology=Topology("thread", 64)) - explicit = make_mesh((8,), topology=Topology("cta", 8)) + matching = Mesh((Topology("thread", 32),), Layout((32,), (1,)), ("g",)) + mismatching = Mesh((Topology("thread", 64),), Layout((32,), (1,)), ("g",)) + explicit = Mesh((Topology("cta", 8),), Layout((8,), (1,)), ("g",)) assert product(matching.topologies) == 32 assert states_consistent_positions(matching) @@ -45,13 +48,13 @@ def test_mesh_is_a_frozen_record_without_axis_attributes() -> None: assert not hasattr(mesh, "topology") assert not hasattr(mesh, "axes") - normalized = make_mesh((4, 8), topology="cta") + normalized = Mesh((Topology("cta", 32),), Layout((4, 8), (8, 1)), ("a", "b")) assert normalized.topologies == (Topology("cta", 32),) assert normalized.layout == Layout(shape=((4, 8),), strides=((8, 1),)) def test_mesh_slice_keeps_the_parent_topologies() -> None: - mesh = make_mesh((4, 32), topology="thread") + mesh = Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1)), ("a", "b")) sliced = mesh[0, :] @@ -70,8 +73,49 @@ def test_check_topology_rejects_positions_beyond_a_declared_extent() -> None: def test_mesh_value_equality_is_by_value() -> None: - left = make_mesh((8,), topology="thread") - right = make_mesh((8,), topology="thread") + left = Mesh((Topology("thread", 8),), Layout((8,), (1,)), ("g",)) + right = Mesh((Topology("thread", 8),), Layout((8,), (1,)), ("g",)) assert left == right assert hash(left) == hash(right) + + +def test_make_mesh_appends_a_sliced_scope_with_its_offset() -> None: + assert make_mesh(CTA, THR[128:256]).layout == RUN + assert not covered_by_scope(make_mesh(CTA, THR[0:128]), make_mesh(CTA, THR[128:256])) + + +def test_make_mesh_refuses_a_sliced_suffix_replacement() -> None: + with pytest.raises(ValueError): + make_mesh(CT, THR[128:256]) + + +def test_mesh_with_several_levels_slices_in_device_numbering() -> None: + assert CT[:, 128:256].layout == RUN + assert CT[:, 128:256] == make_mesh(CTA, THR[128:256]) + assert CT[1:3, 128:256].layout == ComposedLayout( + None, 1 * 384 + 128, Layout(((2,), (128,)), ((1,), (1,))) + ) + + +def test_separate_undoes_make_mesh() -> None: + mesh = CT[1:3, 128:256] + assert separate(mesh) == ( + Mesh( + (Topology("cta", 4),), + ComposedLayout(None, 1, Layout((2,), (1,))), + ), + Mesh( + (Topology("thread", 384),), + ComposedLayout(None, 128, Layout((128,), (1,))), + ), + ) + assert make_mesh(*separate(mesh)) == mesh + + +def test_mesh_refuses_a_repeated_topology_name() -> None: + with pytest.raises(ValueError): + Mesh( + (Topology("thread", 4), Topology("thread", 32)), + Layout(((4,), (32,)), ((1,), (1,))), + ) diff --git a/tests/ir/types/test_tensor_type.py b/tests/ir/types/test_tensor_type.py index abedbabe..9ec51d8e 100644 --- a/tests/ir/types/test_tensor_type.py +++ b/tests/ir/types/test_tensor_type.py @@ -15,7 +15,6 @@ Split, TensorType, Topology, - make_mesh, make_shard_tensor_type, ) from tilefoundry.ir.types.dim import DimVar, ceildiv @@ -62,7 +61,9 @@ def test_zero_extent_has_zero_logical_and_local_size() -> None: assert numel(type) == 0 assert tensor_bytes(type) == 0 - sharded = make_shard_tensor_type((0,), mesh=make_mesh((2,)), attrs=(Split(0),)) + sharded = make_shard_tensor_type( + (0,), mesh=Mesh((Topology("gpu", 2),), Layout((2,), (1,)), ("g",)), attrs=(Split(0),) + ) assert local_type_of(sharded, topology_level="gpu", topologies=(Topology("gpu", 2),)).shape == ( 1, 0, @@ -84,7 +85,7 @@ def test_size_rejects_symbolic_and_negative_extents() -> None: def test_local_type_rejects_a_zero_mesh_extent() -> None: - mesh = make_mesh((0,)) + mesh = Mesh((Topology("gpu", 0),), Layout((0,), (1,)), ("g",)) layout = ShardLayout(Layout(shape=(0,), strides=(1,)), (Split(0),), mesh) type = TensorType(shape=(0,), dtype=DType.f32, layout=layout, storage="gmem") diff --git a/tests/ops/ir/test_arange.py b/tests/ops/ir/test_arange.py index 05c663c5..190c11b5 100644 --- a/tests/ops/ir/test_arange.py +++ b/tests/ops/ir/test_arange.py @@ -18,9 +18,9 @@ from tilefoundry.ir.hir.specialize import residual_dims, specialize_concretely from tilefoundry.ir.hir.tensor.arange import Arange from tilefoundry.ir.isl_interop import normalize_dim -from tilefoundry.ir.mesh_scope import merge_mesh from tilefoundry.ir.types import DType, Layout, Mesh, TensorType, Topology from tilefoundry.ir.types.dim import ceildiv +from tilefoundry.ir.types.mesh import make_mesh from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.contexts import TrafficBytes, TypeInferContext from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor @@ -131,7 +131,7 @@ def test_an_unbound_mesh_coordinate_is_rejected() -> None: def test_an_inner_mesh_coordinate_is_bound_by_a_multilevel_scope() -> None: cta = Mesh((Topology("cta", 2),), Layout((2,), (1,)), ("c",)) - current = merge_mesh((cta, _COORD_MESH)) + current = make_mesh(cta, _COORD_MESH) assert ( TypeInferVisitor().visit(_coord(), TypeInferContext(current_mesh=current)) == _COORD_INDEX ) diff --git a/tests/ops/ir/test_argmax.py b/tests/ops/ir/test_argmax.py index 1391eed6..f7d49963 100644 --- a/tests/ops/ir/test_argmax.py +++ b/tests/ops/ir/test_argmax.py @@ -16,8 +16,9 @@ from tilefoundry.ir.types import ( DType, Layout, + Mesh, ShardLayout, - make_mesh, + Topology, make_shard_tensor_type, make_tensor_type, ) @@ -45,13 +46,25 @@ TypeInferCase( "partial_input_rejected", ArgMax(), - (make_shard_tensor_type((4, 256), mesh=make_mesh((4,)), attrs=(Partial("max"),)),), + ( + make_shard_tensor_type( + (4, 256), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Partial("max"),), + ), + ), ExpectedError(match="x carries Partial"), ), TypeInferCase( "reduction_axis_split_rejected", ArgMax(axis=-1), - (make_shard_tensor_type((4, 256), mesh=make_mesh((4,)), attrs=(Split(1),)),), + ( + make_shard_tensor_type( + (4, 256), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Split(1),), + ), + ), ExpectedError(match=r"reduction axis 1.*Split-sharded.*Reshard"), ), ] @@ -66,7 +79,9 @@ def test_argmax_layout_describes_result_and_preserves_surviving_split(): plain = make_tensor_type((4, 256), DType.f32, layout=Layout(shape=(4, 256), strides=(256, 1))) assert infer_call(ArgMax(axis=-1), plain).layout == Layout(shape=(4,), strides=(1,)) - sharded = make_shard_tensor_type((4, 256), mesh=make_mesh((4,)), attrs=(Split(0),)) + sharded = make_shard_tensor_type( + (4, 256), mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), attrs=(Split(0),) + ) result = infer_call(ArgMax(axis=-1), sharded) assert isinstance(result.layout, ShardLayout) assert result.layout.attrs == (Split(0),) diff --git a/tests/ops/ir/test_binary.py b/tests/ops/ir/test_binary.py index 4cb9b8a8..23308c6c 100644 --- a/tests/ops/ir/test_binary.py +++ b/tests/ops/ir/test_binary.py @@ -19,8 +19,14 @@ from tilefoundry.ir.core.errors import VerifyError from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.hir.math.binary import Binary -from tilefoundry.ir.types import DType, make_mesh, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.layout import Layout +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + Topology, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.shard_layout import Broadcast, Partial, Split from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.contexts import TrafficBytes @@ -31,8 +37,8 @@ _F = DType.f32 -_M = make_mesh((4,)) -_MAB = make_mesh((2, 4), ("a", "b")) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) +_MAB = Mesh((Topology("gpu", 8),), Layout((2, 4), (4, 1)), ("a", "b")) _PSUM = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("sum"),)) _PMAX = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("max"),)) _BCAST = make_tensor_type((16, 8), _F) diff --git a/tests/ops/ir/test_cache_update.py b/tests/ops/ir/test_cache_update.py index be0bce3a..258be776 100644 --- a/tests/ops/ir/test_cache_update.py +++ b/tests/ops/ir/test_cache_update.py @@ -31,8 +31,8 @@ from tilefoundry.ir.hir.tensor.cache_update import CacheUpdate from tilefoundry.ir.types import ( DType, + Layout, Topology, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -76,7 +76,10 @@ def test_cache_update_evaluate(): CacheUpdate(), ( make_shard_tensor_type( - (1, 16, 4, 8), DType.bf16, mesh=make_mesh((4,)), attrs=(Partial("sum"),) + (1, 16, 4, 8), + DType.bf16, + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Partial("sum"),), ), make_tensor_type((1,), DType.i32), make_tensor_type((1,), DType.i32), @@ -92,7 +95,10 @@ def test_cache_update_evaluate(): make_tensor_type((1,), DType.i32), make_tensor_type((1,), DType.i32), make_shard_tensor_type( - (1, 4, 4, 8), DType.bf16, mesh=make_mesh((4,)), attrs=(Partial("sum"),) + (1, 4, 4, 8), + DType.bf16, + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Partial("sum"),), ), ), ExpectedError(match="new carries Partial"), @@ -166,7 +172,7 @@ def test_cache_update_cost_is_independent_of_cache_length(cache_len) -> None: _CTA = Topology("cta", 2) -_CTA_MESH = make_mesh((2,), topology=_CTA) +_CTA_MESH = Mesh((_CTA,), Layout((2,), (1,)), ("g",)) @module( diff --git a/tests/ops/ir/test_cast.py b/tests/ops/ir/test_cast.py index f8d70f7c..e035bb61 100644 --- a/tests/ops/ir/test_cast.py +++ b/tests/ops/ir/test_cast.py @@ -17,11 +17,10 @@ from tilefoundry.evaluator import evaluate from tilefoundry.evaluator.value import EvalError from tilefoundry.ir.hir.tensor.cast import Cast -from tilefoundry.ir.types import DType, make_mesh, make_tensor_type -from tilefoundry.ir.types.layout import Layout +from tilefoundry.ir.types import DType, Layout, Mesh, Topology, make_tensor_type from tilefoundry.ir.types.shard_layout import ShardLayout, Split -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) def test_cast_carries_sharded_layout(): diff --git a/tests/ops/ir/test_clamp.py b/tests/ops/ir/test_clamp.py index 95e624df..65db6aae 100644 --- a/tests/ops/ir/test_clamp.py +++ b/tests/ops/ir/test_clamp.py @@ -12,11 +12,11 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.math.clamp import Clamp -from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types import Layout, Mesh, Topology, make_shard_tensor_type from tilefoundry.ir.types.shard_layout import Partial _OP = Clamp(min_val=-1.0, max_val=1.0) -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _PSUM = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("sum"),)) _PMAX = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("max"),)) diff --git a/tests/ops/ir/test_concat.py b/tests/ops/ir/test_concat.py index bee5b0cf..aa154017 100644 --- a/tests/ops/ir/test_concat.py +++ b/tests/ops/ir/test_concat.py @@ -8,11 +8,18 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.tensor.concat import Concat -from tilefoundry.ir.types import DType, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + Topology, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.shard_layout import Split _F = DType.f32 -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) def test_concat_propagates_a_split_outside_the_concat_axis(): diff --git a/tests/ops/ir/test_conv2d.py b/tests/ops/ir/test_conv2d.py index 8db6b3e0..d1086fac 100644 --- a/tests/ops/ir/test_conv2d.py +++ b/tests/ops/ir/test_conv2d.py @@ -20,9 +20,9 @@ from tilefoundry.ir.types import ( DType, Layout, + Mesh, ShardLayout, Topology, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -38,7 +38,7 @@ _X = make_tensor_type((2, 4, 7, 7), _F) _W = make_tensor_type((6, 2, 3, 3), _F) _BIAS = make_tensor_type((6,), _F) -_MESH = make_mesh((2,)) +_MESH = Mesh((Topology("gpu", 2),), Layout((2,), (1,)), ("g",)) VALIDATION_CASES = [ @@ -245,7 +245,11 @@ def test_conv2d_relation_accepts_exact_contraction_partial() -> None: "translated_spatial_split_is_underivable", Conv2D(stride=(1, 1), padding=(4, 0), dilation=(1, 1), groups=1), ( - make_shard_tensor_type((1, 4, 8, 8), mesh=make_mesh((4,)), attrs=(Split(2),)), + make_shard_tensor_type( + (1, 4, 8, 8), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Split(2),), + ), make_tensor_type((4, 4, 1, 1)), make_tensor_type((4,)), ), @@ -260,7 +264,7 @@ def test_conv2d_rejects_underivable_ownership(case) -> None: _CTA = Topology("cta", 2) -_CTA_MESH = make_mesh((2,), topology=_CTA) +_CTA_MESH = Mesh((_CTA,), Layout((2,), (1,)), ("g",)) _INPUT_BYTES = 2 * 4 * 7 * 7 * 4 _WEIGHT_BYTES = 6 * 2 * 3 * 3 * 4 _BIAS_BYTES = 6 * 4 diff --git a/tests/ops/ir/test_layer_norm.py b/tests/ops/ir/test_layer_norm.py index ce87d012..281631dc 100644 --- a/tests/ops/ir/test_layer_norm.py +++ b/tests/ops/ir/test_layer_norm.py @@ -16,8 +16,10 @@ from tilefoundry.ir.hir.nn.layer_norm import LayerNorm from tilefoundry.ir.types import ( DType, + Layout, + Mesh, ShardLayout, - make_mesh, + Topology, make_shard_tensor_type, make_tensor_type, ) @@ -25,7 +27,7 @@ _OP = LayerNorm(axis=-1, eps=1e-5) _F = DType.f32 -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _X = make_tensor_type((4, 8), _F) _W = make_tensor_type((8,), _F) _B = make_tensor_type((8,), _F) diff --git a/tests/ops/ir/test_matmul.py b/tests/ops/ir/test_matmul.py index a872ce78..07cc3e2a 100644 --- a/tests/ops/ir/test_matmul.py +++ b/tests/ops/ir/test_matmul.py @@ -25,8 +25,9 @@ from tilefoundry.ir.hir.nn.matmul import MatMul from tilefoundry.ir.types import ( DType, + Layout, + Mesh, Topology, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -39,11 +40,11 @@ _MM = MatMul() -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _CTA = Topology("cta", 5) -_CTA_MESH = make_mesh((5,), topology=_CTA) +_CTA_MESH = Mesh((_CTA,), Layout((5,), (1,)), ("g",)) COST_CASES = [ diff --git a/tests/ops/ir/test_quant.py b/tests/ops/ir/test_quant.py index cf025fc8..42297bd0 100644 --- a/tests/ops/ir/test_quant.py +++ b/tests/ops/ir/test_quant.py @@ -23,9 +23,10 @@ from tilefoundry.ir.types import ( DType, Layout, + Mesh, ShardLayout, + Topology, TupleType, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -99,13 +100,25 @@ TypeInferCase( "partial_input_rejected", Quant(), - (make_shard_tensor_type((1, 2048), mesh=make_mesh((4,)), attrs=(Partial("max"),)),), + ( + make_shard_tensor_type( + (1, 2048), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Partial("max"),), + ), + ), ExpectedError(match="x carries Partial"), ), TypeInferCase( "last_split_through_group_rejected", Quant(group=128), - (make_shard_tensor_type((2, 256), mesh=make_mesh((4,)), attrs=(Split(1),)),), + ( + make_shard_tensor_type( + (2, 256), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Split(1),), + ), + ), ExpectedError(match=r"last axis 1 Split cuts through group=128.*Reshard"), ), TypeInferCase( @@ -117,7 +130,7 @@ (2, 4, 256), (1024, 256, 2), (Split(1),), - make_mesh((4,)), + Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), ), ), ExpectedError(match=r"last axis 1 Split cuts through group=128.*Reshard"), @@ -131,7 +144,7 @@ (2, 4, 256), (1024, 128, 1), (Split(1),), - make_mesh((4,)), + Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), ), ), ExpectedError(match=r"last axis 1 Split cuts through group=128.*Reshard"), @@ -218,7 +231,12 @@ def test_quant_plain_layouts_describe_each_result() -> None: ids=("outer_axis", "whole_group_last_axis"), ) def test_quant_propagates_representable_sharding(shape, split_axis, expected_scale_shape) -> None: - source = make_shard_tensor_type(shape, _BF, mesh=make_mesh((4,)), attrs=(Split(split_axis),)) + source = make_shard_tensor_type( + shape, + _BF, + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Split(split_axis),), + ) quantized, scale = infer_call(Quant(group=128), source).fields assert quantized.shape == shape @@ -236,7 +254,7 @@ def test_quant_accepts_factorized_contiguous_whole_groups() -> None: (2, 256, 4), (1024, 4, 1), (Split(1),), - make_mesh((4,)), + Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), dtype=_BF, ) quantized, scale = infer_call(Quant(group=128), source).fields @@ -247,7 +265,12 @@ def test_quant_accepts_factorized_contiguous_whole_groups() -> None: def test_quant_drops_fully_broadcast_mesh_ownership() -> None: - source = make_shard_tensor_type((2, 256), _BF, mesh=make_mesh((4,)), attrs=(Broadcast(),)) + source = make_shard_tensor_type( + (2, 256), + _BF, + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Broadcast(),), + ) quantized, scale = infer_call(Quant(group=128), source).fields assert quantized.layout is None diff --git a/tests/ops/ir/test_relu.py b/tests/ops/ir/test_relu.py index d968556f..44a9241a 100644 --- a/tests/ops/ir/test_relu.py +++ b/tests/ops/ir/test_relu.py @@ -10,11 +10,11 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.relu import ReLU -from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types import Layout, Mesh, Topology, make_shard_tensor_type from tilefoundry.ir.types.shard_layout import Partial _OP = ReLU() -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _PSUM = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("sum"),)) _PMAX = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("max"),)) diff --git a/tests/ops/ir/test_repeat_interleave.py b/tests/ops/ir/test_repeat_interleave.py index 68232f4d..e67d6fad 100644 --- a/tests/ops/ir/test_repeat_interleave.py +++ b/tests/ops/ir/test_repeat_interleave.py @@ -15,11 +15,18 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.tensor.repeat_interleave import RepeatInterleave -from tilefoundry.ir.types import DType, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + Topology, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.shard_layout import Split _F = DType.f32 -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) CASES = [ TypeInferCase( diff --git a/tests/ops/ir/test_reshape.py b/tests/ops/ir/test_reshape.py index b718720b..79f7841b 100644 --- a/tests/ops/ir/test_reshape.py +++ b/tests/ops/ir/test_reshape.py @@ -25,8 +25,9 @@ from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.types import ( Layout, + Mesh, ShardLayout, - make_mesh, + Topology, make_shard_tensor_type, make_tensor_type, ) @@ -39,7 +40,7 @@ ) from tilefoundry.ir.types.storage import StorageKind -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) def _reshape(new_shape): @@ -90,7 +91,13 @@ def test_straddling_split_fails_closed(): TypeInferCase( "straddle_fails_closed", _reshape((3, 8)), - (make_shard_tensor_type((6, 4), mesh=make_mesh((2,)), attrs=(Split(0),)),), + ( + make_shard_tensor_type( + (6, 4), + mesh=Mesh((Topology("gpu", 2),), Layout((2,), (1,)), ("g",)), + attrs=(Split(0),), + ), + ), ExpectedError(match="align"), ) ) @@ -125,7 +132,11 @@ def test_split_remaps_partial_carries(): """ ty = infer_call( _reshape((1, 32, 128)), - make_shard_tensor_type((32, 128), mesh=make_mesh((2, 2)), attrs=(Split(0), Partial("sum"))), + make_shard_tensor_type( + (32, 128), + mesh=Mesh((Topology("gpu", 4),), Layout((2, 2), (2, 1)), ("a", "b")), + attrs=(Split(0), Partial("sum")), + ), ) assert tuple(ty.shape) == (1, 32, 128) assert _split_mesh_axes(ty) == {0} diff --git a/tests/ops/ir/test_reshard.py b/tests/ops/ir/test_reshard.py index d898c8b5..1d9ef7d9 100644 --- a/tests/ops/ir/test_reshard.py +++ b/tests/ops/ir/test_reshard.py @@ -18,7 +18,7 @@ ) from tilefoundry.dsl.storage import gmem, rmem from tilefoundry.ir.hir.sharding.reshard import Reshard -from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Topology, make_mesh, make_tensor_type +from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Topology, make_tensor_type from tilefoundry.ir.types.dim import DimMul, DimVar, simplify_dim from tilefoundry.ir.types.shard_layout import Split from tilefoundry.ir.types.storage import StorageKind @@ -28,7 +28,7 @@ def _shard_layout(shape) -> ShardLayout: return ShardLayout( layout=Layout(shape=shape, strides=tuple([1] * len(shape))), attrs=(), - mesh=make_mesh((128,), topology=Topology("cta", 128)), + mesh=Mesh((Topology("cta", 128),), Layout((128,), (1,)), ("g",)), ) diff --git a/tests/ops/ir/test_rms_norm.py b/tests/ops/ir/test_rms_norm.py index 849afa47..8fda44d9 100644 --- a/tests/ops/ir/test_rms_norm.py +++ b/tests/ops/ir/test_rms_norm.py @@ -19,11 +19,18 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.rms_norm import RMSNorm -from tilefoundry.ir.types import DType, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + Topology, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.shard_layout import Partial _RMS = RMSNorm(eps=1e-6) -_PARTIAL_MESH = make_mesh((4,)) +_PARTIAL_MESH = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) CASES = [ TypeInferCase( diff --git a/tests/ops/ir/test_rope.py b/tests/ops/ir/test_rope.py index 511aa09c..00cb358b 100644 --- a/tests/ops/ir/test_rope.py +++ b/tests/ops/ir/test_rope.py @@ -19,15 +19,17 @@ from tilefoundry.ir.hir.nn.rope import RoPE from tilefoundry.ir.types import ( DType, + Layout, + Mesh, + Topology, TupleType, - make_mesh, make_shard_tensor_type, make_tensor_type, ) from tilefoundry.ir.types.shard_layout import Partial _BF = DType.bf16 -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) def _rope_inputs(q_shape, k_shape, *, q=None, k=None, cos=None, sin=None, pos=None): diff --git a/tests/ops/ir/test_sigmoid.py b/tests/ops/ir/test_sigmoid.py index 41458f05..cb593b39 100644 --- a/tests/ops/ir/test_sigmoid.py +++ b/tests/ops/ir/test_sigmoid.py @@ -8,10 +8,10 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.sigmoid import Sigmoid -from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types import Layout, Mesh, Topology, make_shard_tensor_type from tilefoundry.ir.types.shard_layout import Partial -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) def test_sigmoid_rejects_partial_sum_input(): diff --git a/tests/ops/ir/test_silu.py b/tests/ops/ir/test_silu.py index b084943f..c7a0a55f 100644 --- a/tests/ops/ir/test_silu.py +++ b/tests/ops/ir/test_silu.py @@ -12,11 +12,11 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.silu import Silu -from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types import Layout, Mesh, Topology, make_shard_tensor_type from tilefoundry.ir.types.shard_layout import Partial _OP = Silu() -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _PSUM = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("sum"),)) _PMAX = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("max"),)) diff --git a/tests/ops/ir/test_slice.py b/tests/ops/ir/test_slice.py index a51adc4d..79f28467 100644 --- a/tests/ops/ir/test_slice.py +++ b/tests/ops/ir/test_slice.py @@ -21,8 +21,9 @@ ComposedLayout, DType, Layout, + Mesh, + Topology, TupleType, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -33,7 +34,7 @@ from tilefoundry.visitor_registry.visitors import CostEvaluator _F = DType.f32 -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) def _slice_call(source, starts, sizes, strides, *, source_expr=None): @@ -150,7 +151,7 @@ def test_runtime_window_preserves_distribution_without_claiming_an_offset(): def test_runtime_window_before_a_split_axis_preserves_the_split_target(): seq = DimVar("slice_seq", 1, 4097) - mesh = make_mesh((16,)) + mesh = Mesh((Topology("gpu", 16),), Layout((16,), (1,)), ("g",)) source = make_shard_tensor_type((1, seq, 16, 128), mesh=mesh, attrs=(Split(2),)) start = Var(type=make_tensor_type((), DType.i64), name="start") @@ -165,7 +166,7 @@ def test_runtime_window_before_a_split_axis_preserves_the_split_target(): def test_fused_gqa_qkv_slices_keep_distribution_visible_to_consumers(): """32 Q / 8 KV heads use group slices 4/1/1 and retain HKV sharding.""" - mesh = make_mesh((8,)) + mesh = Mesh((Topology("gpu", 8),), Layout((8,), (1,)), ("g",)) source = make_tensor_type( (64, 8, 6, 16), _F, diff --git a/tests/ops/ir/test_softmax.py b/tests/ops/ir/test_softmax.py index 7c91562a..aa965a90 100644 --- a/tests/ops/ir/test_softmax.py +++ b/tests/ops/ir/test_softmax.py @@ -8,13 +8,13 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.softmax import SoftMax -from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types import Layout, Mesh, Topology, make_shard_tensor_type from tilefoundry.ir.types.shard_layout import Partial def test_softmax_typeinfer_partial_input_errors(): - m = make_mesh((4,)) + m = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) run_typeinfer_case( TypeInferCase( "partial_sum_errors", diff --git a/tests/ops/ir/test_softplus.py b/tests/ops/ir/test_softplus.py index e792bf8f..412c39b5 100644 --- a/tests/ops/ir/test_softplus.py +++ b/tests/ops/ir/test_softplus.py @@ -16,11 +16,11 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.math.softplus import Softplus -from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types import Layout, Mesh, Topology, make_shard_tensor_type from tilefoundry.ir.types.shard_layout import Partial _OP = Softplus() -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _PSUM = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("sum"),)) diff --git a/tests/ops/ir/test_split.py b/tests/ops/ir/test_split.py index 85777742..db8afc5d 100644 --- a/tests/ops/ir/test_split.py +++ b/tests/ops/ir/test_split.py @@ -20,9 +20,10 @@ from tilefoundry.ir.types import ( DType, Layout, + Mesh, ShardLayout, + Topology, TupleType, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -68,7 +69,9 @@ def test_split_rebuilds_plain_and_sharded_result_layouts(): plain_parts = infer_call(Split(axis=0, num_splits=4), plain).fields assert all(part.layout == Layout(shape=(4, 8), strides=(8, 1)) for part in plain_parts) - sharded = make_shard_tensor_type((16, 8), mesh=make_mesh((4,)), attrs=(SplitAttr(0),)) + sharded = make_shard_tensor_type( + (16, 8), mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), attrs=(SplitAttr(0),) + ) sharded_parts = infer_call(Split(axis=1, num_splits=2), sharded).fields assert all(isinstance(part.layout, ShardLayout) for part in sharded_parts) assert all(part.layout.attrs == (SplitAttr(0),) for part in sharded_parts) @@ -91,7 +94,7 @@ def test_split_values_carry_their_exact_inferred_field_types() -> None: (8, 4), dtype=DType.f32, storage="rmem", - mesh=make_mesh((4,)), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), attrs=(SplitAttr(0),), ) op = Split(axis=1, num_splits=2) diff --git a/tests/ops/ir/test_stack.py b/tests/ops/ir/test_stack.py index 0b4e2b5a..a51d312c 100644 --- a/tests/ops/ir/test_stack.py +++ b/tests/ops/ir/test_stack.py @@ -19,9 +19,10 @@ from tilefoundry.ir.types import ( DType, Layout, + Mesh, Partial, ShardLayout, - make_mesh, + Topology, make_shard_tensor_type, make_tensor_type, ) @@ -48,7 +49,7 @@ def test_stack_plain_result_has_a_fresh_layout_and_value() -> None: def test_stack_relation_carries_a_single_split_past_the_inserted_axis() -> None: - mesh = make_mesh((4,)) + mesh = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) plain = make_tensor_type((2, 8)) sharded = make_shard_tensor_type((2, 8), mesh=mesh, attrs=(Split(1),)) @@ -61,7 +62,11 @@ def test_stack_relation_carries_a_single_split_past_the_inserted_axis() -> None: def test_stack_relation_carries_uniform_partial_slices() -> None: - partial = make_shard_tensor_type((2, 8), mesh=make_mesh((4,)), attrs=(Partial("sum"),)) + partial = make_shard_tensor_type( + (2, 8), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Partial("sum"),), + ) result = infer_call(Stack(axis=1), partial, partial) @@ -75,8 +80,16 @@ def test_stack_relation_carries_uniform_partial_slices() -> None: "incompatible_split_axes", Stack(axis=0), ( - make_shard_tensor_type((8, 8), mesh=make_mesh((4,)), attrs=(Split(0),)), - make_shard_tensor_type((8, 8), mesh=make_mesh((4,)), attrs=(Split(1),)), + make_shard_tensor_type( + (8, 8), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Split(0),), + ), + make_shard_tensor_type( + (8, 8), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Split(1),), + ), ), ExpectedError(match=r"input 1 .*incompatible.*Reshard"), ), @@ -84,7 +97,11 @@ def test_stack_relation_carries_uniform_partial_slices() -> None: "partial_and_plain_slices", Stack(axis=0), ( - make_shard_tensor_type((8, 8), mesh=make_mesh((4,)), attrs=(Partial("sum"),)), + make_shard_tensor_type( + (8, 8), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Partial("sum"),), + ), make_tensor_type((8, 8)), ), ExpectedError(match=r"input 1 does not carry Partial.*Reshard"), diff --git a/tests/ops/ir/test_tanh.py b/tests/ops/ir/test_tanh.py index cec7be7d..f93c83d1 100644 --- a/tests/ops/ir/test_tanh.py +++ b/tests/ops/ir/test_tanh.py @@ -12,11 +12,11 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.tanh import Tanh -from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types import Layout, Mesh, Topology, make_shard_tensor_type from tilefoundry.ir.types.shard_layout import Partial _OP = Tanh() -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _PSUM = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("sum"),)) _PMAX = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("max"),)) diff --git a/tests/ops/ir/test_topk.py b/tests/ops/ir/test_topk.py index b6a84f5c..014a5913 100644 --- a/tests/ops/ir/test_topk.py +++ b/tests/ops/ir/test_topk.py @@ -33,8 +33,9 @@ from tilefoundry.ir.types import ( DType, Layout, + Mesh, + Topology, TupleType, - make_mesh, make_shard_tensor_type, make_tensor_type, ) @@ -57,13 +58,25 @@ TypeInferCase( "split_on_selected_axis_rejected", TopK(k=2, axis=-1), - (make_shard_tensor_type((4, 256), mesh=make_mesh((4,)), attrs=(Split(1),)),), + ( + make_shard_tensor_type( + (4, 256), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Split(1),), + ), + ), ExpectedError(match="must not be Split-sharded"), ), TypeInferCase( "partial_input_rejected", TopK(k=2, axis=-1), - (make_shard_tensor_type((4, 256), mesh=make_mesh((4,)), attrs=(Partial("max"),)),), + ( + make_shard_tensor_type( + (4, 256), + mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), + attrs=(Partial("max"),), + ), + ), ExpectedError(match="x carries Partial"), ), ] @@ -110,7 +123,9 @@ def test_topk_output_layout_shrinks_selected_axis_preserving_split(): A Split on a non-selected axis must be preserved, and the output shard layout's selected axis must shrink to k so size(layout)==size(shape). """ - x_ty = make_shard_tensor_type((4, 256), mesh=make_mesh((4,)), attrs=(Split(0),)) + x_ty = make_shard_tensor_type( + (4, 256), mesh=Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), attrs=(Split(0),) + ) out = infer_call(TopK(k=6, axis=-1), x_ty) values_ty, indices_ty = out.fields assert values_ty.shape == (4, 6) and indices_ty.shape == (4, 6) @@ -149,7 +164,7 @@ def test_topk_all_broadcast_layout_with_dynamic_dim(): (256, s), None, (Broadcast(),), - make_mesh((4,)), + Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)), dtype=_F32, ) values_ty, indices_ty = infer_call(TopK(k=6, axis=0), x_ty).fields diff --git a/tests/ops/ir/test_transpose.py b/tests/ops/ir/test_transpose.py index f7c9cac6..529b9175 100644 --- a/tests/ops/ir/test_transpose.py +++ b/tests/ops/ir/test_transpose.py @@ -15,7 +15,14 @@ raw_shard_tensor_type, ) from tilefoundry.ir.hir.tensor.transpose import Transpose -from tilefoundry.ir.types import DType, Layout, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + Topology, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.shard_layout import ( Broadcast, ShardLayout, @@ -23,7 +30,11 @@ shard_layout_local_shape, ) -_M = make_mesh((1, 128, 8, 32), ("cluster", "cta", "warp", "lane")) +_M = Mesh( + (Topology("gpu", 32768),), + Layout((1, 128, 8, 32), (32768, 256, 32, 1)), + ("cluster", "cta", "warp", "lane"), +) _B4 = (Broadcast(), Broadcast(), Broadcast(), Broadcast()) _T10 = Transpose(perm=(1, 0)) diff --git a/tests/ops/ir/test_unary.py b/tests/ops/ir/test_unary.py index dbbbb77e..0b2cf1c3 100644 --- a/tests/ops/ir/test_unary.py +++ b/tests/ops/ir/test_unary.py @@ -20,7 +20,14 @@ ) from tilefoundry.ir.core.kinds import UnaryKind from tilefoundry.ir.hir.math.unary import Unary -from tilefoundry.ir.types import DType, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + Topology, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.shard_layout import Partial from tilefoundry.visitor_registry.contexts import TrafficBytes @@ -28,7 +35,7 @@ _EXP = Unary(kind=UnaryKind.EXP) _ABS = Unary(kind=UnaryKind.ABS) _RSQRT = Unary(kind=UnaryKind.RSQRT) -_M = make_mesh((4,)) +_M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) _PSUM = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("sum"),)) _PMAX = make_shard_tensor_type((16, 8), mesh=_M, attrs=(Partial("max"),)) diff --git a/tests/ops/ir/test_where.py b/tests/ops/ir/test_where.py index 38ba1bdc..601f32ad 100644 --- a/tests/ops/ir/test_where.py +++ b/tests/ops/ir/test_where.py @@ -15,12 +15,19 @@ from tests.ops.ir.typeinfer_utils import ExpectedError, TypeInferCase, run_typeinfer_case from tilefoundry.evaluator import evaluate from tilefoundry.ir.hir.tensor.where import Where -from tilefoundry.ir.types import DType, Layout, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + Topology, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.shard_layout import Split from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.contexts import TrafficBytes -_MESH = make_mesh((4,)) +_MESH = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) def test_where_evaluates_right_aligned_broadcast(): From a2dd332cfc1467180991666d75fdd80f23c57f2d Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Fri, 25 Sep 2026 22:13:19 +0800 Subject: [PATCH 2/8] fix(mesh): keep topology checks with mesh types --- docs/spec/code-organization.md | 4 ++-- src/tilefoundry/ir/mesh_scope.py | 22 +-------------------- src/tilefoundry/ir/types/mesh.py | 33 ++++++++++++++++++++++++++------ tests/ir/types/test_mesh.py | 3 +-- 4 files changed, 31 insertions(+), 31 deletions(-) diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index bc21e6aa..b5e71e8c 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -19,8 +19,8 @@ truth for the directory's structure and invariants. |---|---|---| | `ir/core/` | [core-ir](./core-ir.md) | Shared node algebra: `Module` / `Expr` / `Var` / `Constant` / `Tuple` / `Op` / `Call` / `Stmt` (base class) / `OpSchema` / `ParamDef` / call-graph and ownership queries / typed metadata attach-detach and diagnostics / `@register_op` / `@register_alias` / `op_registry` / `errors`. | | `ir/types/` | [types](./types.md) | Type-system root: `Type` / `TensorType` / `TupleType` / `UnitType` / `CallableType` / `DType` / `StorageKind` / `resolve_storage` / local projections (`local_type_of`) / tensor-leaf, byte-by-storage, and topology-extent queries / `dim.*` (with their typeinfer). | -| `ir/types/{int_tuple,stride,layout,layout_algebra,shard_layout,mesh}.py` | [shard](./shard.md) | `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`), filed as CuTe files them: int tuples (`flatten` / `unflatten` / `repeat_like` / `product`), strides (`compact_major` / `idx2crd` / `crd2idx`), layouts and the algebra over them each in their own module; mesh construction and separation stay with `Mesh`. | -| `ir/mesh_scope.py` | [shard](./shard.md) | Which scope a statement stands inside and what it admits: `device_layout`, `covered_by_scope`, `check_topology`. Neither a type nor a visitor, so it sits beside `ir/isl_interop.py` rather than in either. | +| `ir/types/{int_tuple,stride,layout,layout_algebra,shard_layout,mesh}.py` | [shard](./shard.md) | `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`), filed as CuTe files them: int tuples (`flatten` / `unflatten` / `repeat_like` / `product`), strides (`compact_major` / `idx2crd` / `crd2idx`), layouts and the algebra over them each in their own module; mesh construction, separation, and topology-bound checking stay with `Mesh`. | +| `ir/mesh_scope.py` | [shard](./shard.md) | Which scope a statement stands inside and what it admits: `device_layout`, `covered_by_scope`. Neither a type nor a visitor, so it sits beside `ir/isl_interop.py` rather than in either. | | `ir/constraints/` | [parser](./parser.md) | Authored `where(layout=..., mesh=..., storage=...)` constraint records: the shared base plus layout, mesh, and storage constraints, attached by the parser and read back by the Python printer. | | `ir/visitor.py` | [visitor-mutator](./visitor-mutator.md) | `ExprFunctor` / `ExprVisitor` / `ExprWalker` / `ExprCollector` / `ExprCloner` / `BindingSubstitutionCloner` / `StmtVisitor` / `StmtMutator` / `StmtExprMutator`, plus `collect_exprs`, value-operand/function-value queries, and the canonical `PrimFunction` walk and rewrite entries. | | `ir/isl_interop.py` | [types](./types.md) | Interoperation between dimension and shape IR values and isl: expression rendering and decoding, normalization, value ranges, and shape-domain construction. Pure isl operations remain in `utils/isl_utils.py`. | diff --git a/src/tilefoundry/ir/mesh_scope.py b/src/tilefoundry/ir/mesh_scope.py index 38166f92..f674867f 100644 --- a/src/tilefoundry/ir/mesh_scope.py +++ b/src/tilefoundry/ir/mesh_scope.py @@ -12,7 +12,7 @@ from tilefoundry.ir.types.int_tuple import flatten, product from tilefoundry.ir.types.layout import ComposedLayout, Layout, size from tilefoundry.ir.types.layout_algebra import is_inverse_projectable -from tilefoundry.ir.types.mesh import Mesh, _levels, _starts +from tilefoundry.ir.types.mesh import Mesh, _levels, _starts, check_topology from tilefoundry.ir.types.storage import StorageKind, resolve_storage from tilefoundry.ir.types.stride import compact_major @@ -135,26 +135,6 @@ def _flat(mesh: Mesh) -> Layout: ) -def check_topology(mesh: Mesh) -> None: - """Reject static mesh positions beyond their declared topology extents. - - A constant slice is already bounded by ``Mesh.__getitem__``; its shortened - axes no longer land on full topology boundaries and are therefore accepted. - """ - if isinstance(mesh.layout, ComposedLayout): - return - for topology, arrangement in zip(mesh.topologies, _levels(mesh)): - declared = getattr(topology, "size", None) - if not isinstance(declared, int) or isinstance(declared, bool): - continue - count = product(tuple(flatten(arrangement.shape))) - if isinstance(count, int) and count > declared: - raise ValueError( - f"mesh level {getattr(topology, 'name', topology)!r} has {count} " - f"positions, exceeding declared extent {declared}" - ) - - __all__ = [ "check_topology", "device_layout", diff --git a/src/tilefoundry/ir/types/mesh.py b/src/tilefoundry/ir/types/mesh.py index c571a65b..0038d02a 100644 --- a/src/tilefoundry/ir/types/mesh.py +++ b/src/tilefoundry/ir/types/mesh.py @@ -2,6 +2,7 @@ from dataclasses import dataclass +from tilefoundry.ir.types.int_tuple import product from tilefoundry.ir.types.layout import ComposedLayout, Layout, LayoutBase, flatten, get from tilefoundry.ir.types.layout import rank as _rank from tilefoundry.ir.types.stride import compact_major, compact_row_major, crd2idx, idx2crd @@ -244,14 +245,36 @@ def _levels(mesh: Mesh) -> tuple[Layout, ...]: def _starts(mesh: Mesh) -> tuple[int, ...]: """Where each level's run starts, decoded from the device-numbered offset.""" offset = mesh.layout.offset if isinstance(mesh.layout, ComposedLayout) else 0 - if len(mesh.topologies) == 1 and isinstance(offset, int): + if not isinstance(offset, int): + return (0,) * len(mesh.topologies) + if len(mesh.topologies) == 1: return (offset,) sizes = tuple(topology.size for topology in mesh.topologies) - if not isinstance(offset, int) or any(not isinstance(one, int) for one in sizes): + if any(not isinstance(one, int) for one in sizes): return (0,) * len(sizes) return tuple(idx2crd(offset, sizes, compact_major(sizes))) +def check_topology(mesh: Mesh) -> None: + """Reject static mesh positions beyond their declared topology extents. + + A constant slice is already bounded by ``Mesh.__getitem__``; its shortened + axes no longer land on full topology boundaries and are therefore accepted. + """ + if isinstance(mesh.layout, ComposedLayout): + return + for topology, arrangement in zip(mesh.topologies, _levels(mesh)): + declared = getattr(topology, "size", None) + if not isinstance(declared, int) or isinstance(declared, bool): + continue + count = product(tuple(flatten(arrangement.shape))) + if isinstance(count, int) and count > declared: + raise ValueError( + f"mesh level {getattr(topology, 'name', topology)!r} has {count} " + f"positions, exceeding declared extent {declared}" + ) + + def _joined_layout(levels: tuple[Layout, ...]) -> Layout: if len(levels) == 1: return levels[0] @@ -271,7 +294,7 @@ def _joined( ) -> Mesh: layout: Layout | ComposedLayout = _joined_layout(levels) if sliced: - sizes = tuple(getattr(topology, "size", None) for topology in topologies) + sizes = tuple(topology.size for topology in topologies) if not all(isinstance(size, int) for size in sizes): raise ValueError("joining sliced meshes needs static topology extents") layout = ComposedLayout(None, crd2idx(starts, sizes, compact_major(sizes)), layout) @@ -320,8 +343,6 @@ def make_mesh(*meshes: Mesh) -> Mesh: f"{shared} named again while {unnamed} is not; a scope either " "replaces the levels in force or adds levels below them" ) - from tilefoundry.ir.mesh_scope import check_topology # noqa: PLC0415 - check_topology(result) return result @@ -342,4 +363,4 @@ def separate(mesh: Mesh) -> tuple[Mesh, ...]: return tuple(separated) -__all__ = ["Mesh", "Topology", "make_mesh", "separate"] +__all__ = ["Mesh", "Topology", "check_topology", "make_mesh", "separate"] diff --git a/tests/ir/types/test_mesh.py b/tests/ir/types/test_mesh.py index f32c7fc6..51b4b55e 100644 --- a/tests/ir/types/test_mesh.py +++ b/tests/ir/types/test_mesh.py @@ -4,7 +4,6 @@ from tests.fixtures.meshes import CT, CTA, RUN, THR from tilefoundry.ir.mesh_scope import ( - check_topology, covered_by_scope, mesh_scope_matches_required_scope, states_consistent_positions, @@ -12,7 +11,7 @@ from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, Topology, make_mesh from tilefoundry.ir.types.int_tuple import product from tilefoundry.ir.types.layout_algebra import size -from tilefoundry.ir.types.mesh import separate +from tilefoundry.ir.types.mesh import check_topology, separate def test_mesh_position_consistency_is_an_explicit_predicate() -> None: From 692fa9cb95ba733dcb71349c1a126fc40368c818 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Fri, 25 Sep 2026 22:37:45 +0800 Subject: [PATCH 3/8] refactor(pattern): introduce operation pattern language --- docs/spec/code-organization.md | 9 +- docs/spec/core-ir.md | 81 +- docs/spec/hir.md | 10 +- docs/spec/inspection.md | 4 +- docs/tutorial/showcase.ipynb | 8 +- docs/tutorial/showcase.md | 10 +- src/tilefoundry/__init__.py | 94 ++- src/tilefoundry/analysis/check.py | 33 +- .../codegen/cuda/tir/prim_function.py | 8 +- src/tilefoundry/dsl/__init__.py | 4 +- src/tilefoundry/evaluator/interpreter.py | 51 +- src/tilefoundry/inspection/printer_base.py | 28 +- src/tilefoundry/inspection/python_printer.py | 248 +++--- src/tilefoundry/inspection/tir_printer.py | 80 +- src/tilefoundry/ir/clause/__init__.py | 24 + .../ir/{constraints => clause}/base.py | 37 +- .../ir/{constraints => clause}/layout.py | 15 +- .../ir/{constraints => clause}/mesh.py | 10 +- .../ir/{constraints => clause}/storage.py | 6 +- src/tilefoundry/ir/constraints/__init__.py | 24 - src/tilefoundry/ir/core/overload.py | 13 +- src/tilefoundry/ir/core/param_def.py | 27 +- src/tilefoundry/ir/core/pattern.py | 153 ---- src/tilefoundry/ir/hir/function.py | 6 +- src/tilefoundry/ir/hir/math/binary.py | 2 +- src/tilefoundry/ir/hir/math/clamp.py | 4 +- src/tilefoundry/ir/hir/math/softplus.py | 4 +- src/tilefoundry/ir/hir/math/unary.py | 2 +- src/tilefoundry/ir/hir/nn/conv2d.py | 2 +- src/tilefoundry/ir/hir/nn/gelu.py | 4 +- src/tilefoundry/ir/hir/nn/layer_norm.py | 12 +- src/tilefoundry/ir/hir/nn/matmul.py | 2 +- src/tilefoundry/ir/hir/nn/relu.py | 4 +- src/tilefoundry/ir/hir/nn/rms_norm.py | 2 +- src/tilefoundry/ir/hir/nn/rope.py | 14 +- src/tilefoundry/ir/hir/nn/sigmoid.py | 4 +- src/tilefoundry/ir/hir/nn/silu.py | 4 +- src/tilefoundry/ir/hir/nn/softmax.py | 2 +- src/tilefoundry/ir/hir/nn/tanh.py | 4 +- src/tilefoundry/ir/hir/sharding/local.py | 2 +- src/tilefoundry/ir/hir/sharding/mesh_coord.py | 2 +- src/tilefoundry/ir/hir/sharding/reshard.py | 4 +- src/tilefoundry/ir/hir/specialize.py | 35 +- src/tilefoundry/ir/hir/tensor/argmax.py | 2 +- src/tilefoundry/ir/hir/tensor/cache_update.py | 12 +- src/tilefoundry/ir/hir/tensor/cast.py | 2 +- src/tilefoundry/ir/hir/tensor/concat.py | 2 +- src/tilefoundry/ir/hir/tensor/full_like.py | 2 +- src/tilefoundry/ir/hir/tensor/index_add.py | 8 +- src/tilefoundry/ir/hir/tensor/index_copy.py | 8 +- src/tilefoundry/ir/hir/tensor/index_select.py | 14 +- src/tilefoundry/ir/hir/tensor/insert_slice.py | 29 +- src/tilefoundry/ir/hir/tensor/quant.py | 13 +- src/tilefoundry/ir/hir/tensor/rank.py | 2 +- src/tilefoundry/ir/hir/tensor/reduce.py | 2 +- .../ir/hir/tensor/repeat_interleave.py | 8 +- src/tilefoundry/ir/hir/tensor/reshape.py | 4 +- src/tilefoundry/ir/hir/tensor/shape_of.py | 2 +- src/tilefoundry/ir/hir/tensor/slice.py | 36 +- src/tilefoundry/ir/hir/tensor/split.py | 2 +- src/tilefoundry/ir/hir/tensor/stack.py | 2 +- src/tilefoundry/ir/hir/tensor/topk.py | 2 +- src/tilefoundry/ir/hir/tensor/transpose.py | 2 +- .../ir/hir/tensor/tuple_get_item.py | 10 +- src/tilefoundry/ir/hir/verify.py | 20 +- src/tilefoundry/ir/pattern/__init__.py | 103 +++ src/tilefoundry/ir/pattern/constraint.py | 171 ++++ src/tilefoundry/ir/pattern/match.py | 218 +++++ src/tilefoundry/ir/pattern/pattern.py | 788 ++++++++++++++++++ src/tilefoundry/ir/pattern/utils.py | 68 ++ src/tilefoundry/ir/tir/arith.py | 2 +- src/tilefoundry/ir/tir/async_copy.py | 2 +- src/tilefoundry/ir/tir/clamp.py | 2 +- src/tilefoundry/ir/tir/cuda/memory/tma.py | 2 +- src/tilefoundry/ir/tir/cuda/nn/mma.py | 5 +- src/tilefoundry/ir/tir/cuda/sync/mbarrier.py | 2 +- src/tilefoundry/ir/tir/dot.py | 2 +- src/tilefoundry/ir/tir/memory/copy.py | 2 +- src/tilefoundry/ir/tir/memory/fill.py | 2 +- src/tilefoundry/ir/tir/memory/memory_span.py | 2 +- src/tilefoundry/ir/tir/memory/ptr_of.py | 2 +- src/tilefoundry/ir/tir/memory/tensor_view.py | 2 +- src/tilefoundry/ir/tir/nn/relu.py | 2 +- src/tilefoundry/ir/tir/nn/rms_norm.py | 2 +- src/tilefoundry/ir/tir/prim_function.py | 6 +- src/tilefoundry/ir/tir/reduce.py | 2 +- src/tilefoundry/ir/tir/verify.py | 15 +- src/tilefoundry/parser/ast_pattern.py | 8 +- src/tilefoundry/parser/pattern_nodes.py | 240 ++++-- src/tilefoundry/script.py | 25 +- tests/analysis/test_analysis_invariants.py | 2 +- tests/dsl/test_dsl_surface.py | 16 +- tests/evaluator/test_eval_core.py | 6 +- tests/fixtures/placed/gqa_decode.py | 6 +- .../placed/prefill_decode_attention.py | 6 +- tests/fixtures/placed/qwen3_1_7b_pd.py | 6 +- .../placed/specialize_through_call.py | 26 +- tests/fixtures/tir/square.py | 22 +- .../inspection/test_module_tree_roundtrip.py | 14 +- tests/inspection/test_specialization_print.py | 14 +- .../models/deepseek_v4_flash/test_moe.py | 10 +- .../test_dynamic_shape_dispatch.py | 6 +- tests/ir/core/test_overload.py | 19 +- tests/ir/core/test_param_def.py | 15 +- tests/ir/core/test_pattern.py | 83 +- tests/ir/core/test_register_alias.py | 6 +- tests/ir/core/test_specialize.py | 14 +- tests/ir/pattern/test_mesh_pattern.py | 41 + tests/ir/test_dim_var_envelope.py | 10 +- tests/parser/test_calls.py | 6 +- tests/parser/test_functions.py | 9 +- tests/passes/test_host_entry.py | 6 +- 112 files changed, 2289 insertions(+), 1005 deletions(-) create mode 100644 src/tilefoundry/ir/clause/__init__.py rename src/tilefoundry/ir/{constraints => clause}/base.py (59%) rename src/tilefoundry/ir/{constraints => clause}/layout.py (79%) rename src/tilefoundry/ir/{constraints => clause}/mesh.py (59%) rename src/tilefoundry/ir/{constraints => clause}/storage.py (82%) delete mode 100644 src/tilefoundry/ir/constraints/__init__.py delete mode 100644 src/tilefoundry/ir/core/pattern.py create mode 100644 src/tilefoundry/ir/pattern/__init__.py create mode 100644 src/tilefoundry/ir/pattern/constraint.py create mode 100644 src/tilefoundry/ir/pattern/match.py create mode 100644 src/tilefoundry/ir/pattern/pattern.py create mode 100644 src/tilefoundry/ir/pattern/utils.py create mode 100644 tests/ir/pattern/test_mesh_pattern.py diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index b5e71e8c..997b35bb 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -18,10 +18,11 @@ truth for the directory's structure and invariants. | Directory | Owning spec | Contents | |---|---|---| | `ir/core/` | [core-ir](./core-ir.md) | Shared node algebra: `Module` / `Expr` / `Var` / `Constant` / `Tuple` / `Op` / `Call` / `Stmt` (base class) / `OpSchema` / `ParamDef` / call-graph and ownership queries / typed metadata attach-detach and diagnostics / `@register_op` / `@register_alias` / `op_registry` / `errors`. | +| `ir/pattern/` | [core-ir](./core-ir.md) | Operation-declaration predicates: composable pattern values in `pattern.py`, match/binding and rendering mechanics in `match.py`, cross-operand relations in `constraint.py`, and pattern construction/specialization helpers in `utils.py`. | | `ir/types/` | [types](./types.md) | Type-system root: `Type` / `TensorType` / `TupleType` / `UnitType` / `CallableType` / `DType` / `StorageKind` / `resolve_storage` / local projections (`local_type_of`) / tensor-leaf, byte-by-storage, and topology-extent queries / `dim.*` (with their typeinfer). | | `ir/types/{int_tuple,stride,layout,layout_algebra,shard_layout,mesh}.py` | [shard](./shard.md) | `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`), filed as CuTe files them: int tuples (`flatten` / `unflatten` / `repeat_like` / `product`), strides (`compact_major` / `idx2crd` / `crd2idx`), layouts and the algebra over them each in their own module; mesh construction, separation, and topology-bound checking stay with `Mesh`. | | `ir/mesh_scope.py` | [shard](./shard.md) | Which scope a statement stands inside and what it admits: `device_layout`, `covered_by_scope`. Neither a type nor a visitor, so it sits beside `ir/isl_interop.py` rather than in either. | -| `ir/constraints/` | [parser](./parser.md) | Authored `where(layout=..., mesh=..., storage=...)` constraint records: the shared base plus layout, mesh, and storage constraints, attached by the parser and read back by the Python printer. | +| `ir/clause/` | [parser](./parser.md) | Authored `where(layout=..., mesh=..., storage=...)` constraint records: the shared base plus layout, mesh, and storage constraints, attached by the parser and read back by the Python printer. | | `ir/visitor.py` | [visitor-mutator](./visitor-mutator.md) | `ExprFunctor` / `ExprVisitor` / `ExprWalker` / `ExprCollector` / `ExprCloner` / `BindingSubstitutionCloner` / `StmtVisitor` / `StmtMutator` / `StmtExprMutator`, plus `collect_exprs`, value-operand/function-value queries, and the canonical `PrimFunction` walk and rewrite entries. | | `ir/isl_interop.py` | [types](./types.md) | Interoperation between dimension and shape IR values and isl: expression rendering and decoding, normalization, value ranges, and shape-domain construction. Pure isl operations remain in `utils/isl_utils.py`. | | `ir/hir/` | [hir](./hir.md) | HIR Op layer; one subdirectory per category (`math/` / `tensor/` / `nn/` / `shape/` / `sharding/`). One real Op per `.py` ([§2](#2-file-naming-and-content-rules) rule 1); surface-alias schemas have no per-name file and live in each category's `aliases.py` ([§2](#2-file-naming-and-content-rules) rule 5). | @@ -101,9 +102,9 @@ physical directory layout reflects that boundary directly. contracts are distinct even though both are consumed across the codegen boundary. -`ir/constraints/`, `visitor_registry/`, and `dump/` are cross-cutting packages; -their stable responsibilities are owned by [parser](./parser.md), -[visitor-registry](./visitor-registry.md), and [inspection](./inspection.md), +`ir/pattern/`, `ir/clause/`, `visitor_registry/`, and `dump/` are cross-cutting packages; +their stable responsibilities are owned by [core-ir](./core-ir.md), +[parser](./parser.md), [visitor-registry](./visitor-registry.md), and [inspection](./inspection.md), respectively. Their internal file layout is not a per-Op contract. ## 2. File naming and content rules diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index df256ac9..288de64f 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -482,6 +482,11 @@ class Op: (see [tir §2.3](./tir.md#23-tir-ops)). ```python +class MemoryEffect(Flag): + READ = auto() + WRITE = auto() + + class ParamDef: """Declare one Op input or attribute. @@ -491,6 +496,7 @@ class ParamDef: pattern: attribute; Optional input-type predicate. optional: attribute; Whether None is accepted. default: attribute; Call-site default or the required-value sentinel. + effect: attribute; Declared storage effect, or None when undeclared. """ kind: Literal["input", "attribute"] @@ -498,10 +504,14 @@ class ParamDef: pattern: Pattern | None = None optional: bool = False default: Any = MISSING + effect: MemoryEffect | None = None ``` - constraints: - a single Op parameter descriptor; the order of input-kind ParamDefs fixes `Call.args` position. + - `MemoryEffect` is a `Flag` with `READ` and `WRITE`; an input may declare + either or both. `None` means undeclared, while the zero flag is invalid. + Attributes cannot carry a memory effect. Example: @@ -592,62 +602,85 @@ and specialization dispatch. class Pattern: """Carry a reusable dispatch predicate.""" - def match(self, subject) -> bool: ... + def match(self, subject, captures=None) -> Match | None: ... ``` - constraints: - shared by parser dispatch (`ParamDef.pattern`) and specialization dispatch (`Function.specializations` / `PrimFunction.specializations`). + - a successful match returns a truthy `Match` carrying named captures; + failure is `None`. + +The implementation is split by responsibility under `ir/pattern/`: + +- `pattern.py` defines `Pattern` and the composable classes + `OrPattern`, `AndPattern`, `SequencePattern`, `CapturePattern`, + `ConstraintPattern`, `GuardPattern`, `SwitchPattern`, `RangePattern`, + `MultipleOfPattern`, `OneOfPattern`, `AttrPattern`, `BitsPattern`, + `LayoutPattern`, `SwizzlePattern`, `ComposedLayoutPattern`, `MeshPattern`, + `ShardLayoutPattern`, `ScalarPattern`, `TensorPattern`, and + `WildcardPattern`. It also owns the `Scalar` and `Tensor` singletons. +- `match.py` owns matches, captures, symbolic resolution, layout-frame reading, + and the shared description helpers. +- `constraint.py` owns cross-operand `Constraint`, `DistinctConstraint`, + `SameConstraint`, and `SameModesConstraint` values. +- `utils.py` owns exact-layout construction plus specialization naming and + dimension lookup. + +`LayoutPattern` checks `forward` and `injective` over the whole flattened +arrangement by default. With `per_mode=True`, it checks each top-level mode +independently; `MeshPattern` requires this explicit form because each mesh +level uses its own numbering space. `MeshPattern` never changes the supplied +pattern implicitly. Two consumer surfaces: - **Parser dispatch** — `ParamDef.pattern` ([§2.3](#23-op)) is matched against an argument's `Expr.type` during overload resolution. Subclasses used: - `ScalarPat` (rank-0), `TensorPat(rank?, dtype?)` (non-scalar), and - `AndPat(parts)` (conjunction). Two singletons are exported as - convenience: `Scalar = ScalarPat()` and `Tensor = TensorPat()`. + `ScalarPattern` (rank-0), `TensorPattern(rank?, dtype?)` (non-scalar), and + `AndPattern(parts)` (conjunction). Two singletons are exported as + convenience: `Scalar = ScalarPattern()` and `Tensor = TensorPattern()`. - **Specialization dispatch** — patterns appearing in `hir.Function.specializations` ([hir.md §1.1](./hir.md#11-function)) and `tir.PrimFunction.specializations` describe which runtime shape range a variant covers. The HIR→TIR lowering inspects each pattern's fields directly; it does not call `match`. -### 3.1 `DimVarRangePat` +### 3.1 `RangePattern` ```python -class DimVarRangePat(Pattern): - """Match one sub-range of a named dimension. +class RangePattern(Pattern): + """Match a closed integer range, optionally naming a dimension. Attributes: - dim_var: attribute; Name of the dimension. - lo: attribute; Inclusive lower bound. - hi: attribute; Exclusive upper bound. + dim_var: attribute; Specialization dimension name, or empty otherwise. + lo: attribute; Optional inclusive lower bound. + hi: attribute; Optional inclusive upper bound. """ dim_var: str = "" - lo: int = 0 - hi: int = 0 + lo: int | None = None + hi: int | None = None ``` - constraints: - - This is the per-variant sub-range for a named `DimVar`; `match(v)` is - `lo <= v <= hi` and ignores `dim_var`. - - `dim_var` MUST be a non-empty `str` — the name of the `DimVar` the - range applies to. The lowering resolves it to a runtime + - At least one of `lo` and `hi` MUST be stated. Each stated bound MUST be a + plain `int` (`bool` is rejected), and two stated bounds MUST satisfy + `lo <= hi`. `RangePattern(lo=k)` and `RangePattern(hi=k)` express the + one-sided relations formerly represented by separate pattern classes. + - A specialization states both bounds and a non-empty `dim_var`: the name + of the `DimVar` the range applies to. The lowering resolves it to a runtime `ShapeOf(param, axis)` by walking the enclosing function signature. - - `lo` and `hi` MUST be plain `int`s (`bool` is rejected). - - The interval is closed `[lo, hi]`; construction MUST satisfy `lo <= hi`. A single-point - range is `[k, k+1)`. - - `match(value)` returns `True` for an `int` value `v` iff - `lo <= v <= hi`. The `dim_var` field does not participate in - `match`. + - A two-sided interval is closed `[lo, hi]`; a single-point range is `[k, k]`. + - `match(value)` succeeds for an `int` value `v` iff every stated bound + admits it. The `dim_var` field does not participate in `match`. - The pattern references a `DimVar` by name only. The envelope of the named dim lives on the `DimVar(name, lo, hi)` itself (see [types.md §4](./types.md#4-dim--symbolic-shape-dimensions)); the - `DimVarRangePat` carries the per-variant sub-range. Envelope + `RangePattern` carries the per-variant sub-range. Envelope containment (`pattern ⊆ DimVar envelope`) is checked in signature context — by the `@tilefoundry.func` validator and the - HIR→TIR lowering — not by `DimVarRangePat.__post_init__`. + HIR→TIR lowering — not by `RangePattern.__post_init__`. ## 4. Shared operation kinds diff --git a/docs/spec/hir.md b/docs/spec/hir.md index 2626fac6..f05000f0 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -236,7 +236,7 @@ the Op typeinfer registry. `TypeInferVisitor` handles it directly as a - Within a `Function` signature, every occurrence of a same-name `DimVar` across `params` and `return_type` MUST agree on its `(lo, hi)` bounds; a disagreement is a verify error. A - `DimVarRangePat` specialization MUST anchor to a `DimVar` reachable + `RangePattern` specialization MUST anchor to a `DimVar` reachable from an input parameter and lie within that `DimVar`'s envelope (see **Shape dispatch and specializations** below). @@ -290,15 +290,15 @@ freeze** below). `return_type`: a variant specializes the body, not the signature. A variant runs in the same execution domain as its base because both are owned by the same `Module`. -- A variant carries exactly one `DimVarRangePat` in `specializations`. +- A variant carries exactly one `RangePattern` in `specializations`. The canonical signature is `";".join(f"{p.dim_var}${p.lo}_{p.hi}" for p in specializations)` - (v0 allows only `DimVarRangePat`). Two variants of one base MUST have + (v0 allows only `RangePattern`). Two variants of one base MUST have distinct canonical signatures. *Envelope coverage.* A dispatched function's parameter `TensorType.shape` carries a `DimVar(name, lo, hi)` whose `(lo, hi)` is -the dispatch envelope; `DimVarRangePat` references that `DimVar` by name. +the dispatch envelope; `RangePattern` references that `DimVar` by name. The variants' closed ranges MUST **partition** the envelope — pairwise **disjoint** and jointly **complete** (their union is exactly the DimVar's half-open `[lo, hi)` envelope). Adjacent closed ranges are written @@ -310,7 +310,7 @@ typeinferred, lowered, or evaluated as a body. Only its variants carry executable bodies. There is no base body to fall back to. *Dispatch resolution.* A `Call` whose target is a dispatch prototype -(`variants != ()`) is a dispatch call: the variant whose `DimVarRangePat` +(`variants != ()`) is a dispatch call: the variant whose `RangePattern` matches is selected and is the call's result. Evaluation selects from the call's concrete argument shapes; specialization selects from the caller's stated dimension bindings. Both use the same variant table. A shape outside diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index 21222f33..edbc5a35 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -256,12 +256,12 @@ variant as an `@.specialize(pattern)` block in declared order: def f(x: Tensor[(S,), "f32"]) -> Tensor[(S,), "f32"]: pass -@f.specialize(DimVarRangePat("S", 1, 3)) +@f.specialize(RangePattern("S", 1, 3)) def small_sequence(x: Tensor[(S,), "f32"]) -> Tensor[(S,), "f32"]: ... ``` -The pattern prints in its constructor form (`DimVarRangePat("S", 1, 3)`; +The pattern prints in its constructor form (`RangePattern("S", 1, 3)`; other `Pattern` subclasses fall back to `repr(pattern)`). The emitted binding mirrors the authoring surface ([parser.md §2.1](./parser.md#21-syntax)); when an IR variant has no display label, the printer synthesizes a valid binding diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index 34c5b0b4..e8887434 100644 --- a/docs/tutorial/showcase.ipynb +++ b/docs/tutorial/showcase.ipynb @@ -42,7 +42,7 @@ } }, "outputs": [], - "source": "from __future__ import annotations\n\nimport ast\nimport math\nimport re\nimport sys\nfrom pathlib import Path\n\nfrom tilefoundry import func, module\nfrom tilefoundry.analysis import analyze as run_analysis\nfrom tilefoundry.dsl import ConstTensor, DimVar, DimVarRangePat, Mesh, Tensor, tf\nfrom tilefoundry.dsl.tf import * # noqa: F401, F403 - bare tile() in the fused body\nfrom tilefoundry.inspection.analysis_report import render_analysis, render_text\nfrom tilefoundry.ir.types import Topology\nfrom tilefoundry.target import CudaTarget\n\nHIDDEN = 256\nQUERY_HEADS = 8\nKV_HEADS = 2\nHEAD_DIM = 32\nKV_DIM = KV_HEADS * HEAD_DIM\nGQA_GROUP = QUERY_HEADS // KV_HEADS\nROPE_CONTEXT = 8192\nCTX = DimVar(\"ctx_len\", 1, ROPE_CONTEXT + 1)\nSCALE = 1.0 / math.sqrt(HEAD_DIM)\nWORKERS = 4\nBLOCK = 128\n\n_H200 = CudaTarget(\"nvidia.h200_sxm\")\n_CTA = Topology(\"cta\", 132)\n" + "source": "from __future__ import annotations\n\nimport ast\nimport math\nimport re\nimport sys\nfrom pathlib import Path\n\nfrom tilefoundry import func, module\nfrom tilefoundry.analysis import analyze as run_analysis\nfrom tilefoundry.dsl import ConstTensor, DimVar, RangePattern, Mesh, Tensor, tf\nfrom tilefoundry.dsl.tf import * # noqa: F401, F403 - bare tile() in the fused body\nfrom tilefoundry.inspection.analysis_report import render_analysis, render_text\nfrom tilefoundry.ir.types import Topology\nfrom tilefoundry.target import CudaTarget\n\nHIDDEN = 256\nQUERY_HEADS = 8\nKV_HEADS = 2\nHEAD_DIM = 32\nKV_DIM = KV_HEADS * HEAD_DIM\nGQA_GROUP = QUERY_HEADS // KV_HEADS\nROPE_CONTEXT = 8192\nCTX = DimVar(\"ctx_len\", 1, ROPE_CONTEXT + 1)\nSCALE = 1.0 / math.sqrt(HEAD_DIM)\nWORKERS = 4\nBLOCK = 128\n\n_H200 = CudaTarget(\"nvidia.h200_sxm\")\n_CTA = Topology(\"cta\", 132)\n" }, { "cell_type": "markdown", @@ -145,7 +145,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "The table says:\n\n```text\nper-CTA work = global work -> no authored split yet\nweight bytes = fixed -> projection weights are a staging target\ncache scan = grows with ctx_len -> a full-cache residency decision will fail first\n```\n\n## 2. Specialize at the cache-only capacity estimate\n\nThe full-cache sharded program in `Stage2_Sharded` places one local query head's K and V cache in smem. A cache-only estimate follows from the target capacity:\n\n```text\nbytes per ctx per CTA = K + V\n = 2 * HEAD_DIM * sizeof(bf16)\n = 2 * 32 * 2\n = 128 B\n\nT = floor(232448 B / 128 B)\n = 1816\n```\n\nThis is not the complete placement peak: other simultaneously resident values also occupy smem. `Stage1_Specialized` still uses the estimate to express two closed `DimVarRangePat` variants, `[1, 1816]` and `[1817, 8192]`; the CP placement report below decides whether either variant actually fits. The Stage1 body is deliberately the unsplit baseline, so the dispatch contract can be read independently from the later implementations.\n" + "source": "The table says:\n\n```text\nper-CTA work = global work -> no authored split yet\nweight bytes = fixed -> projection weights are a staging target\ncache scan = grows with ctx_len -> a full-cache residency decision will fail first\n```\n\n## 2. Specialize at the cache-only capacity estimate\n\nThe full-cache sharded program in `Stage2_Sharded` places one local query head's K and V cache in smem. A cache-only estimate follows from the target capacity:\n\n```text\nbytes per ctx per CTA = K + V\n = 2 * HEAD_DIM * sizeof(bf16)\n = 2 * 32 * 2\n = 128 B\n\nT = floor(232448 B / 128 B)\n = 1816\n```\n\nThis is not the complete placement peak: other simultaneously resident values also occupy smem. `Stage1_Specialized` still uses the estimate to express two closed `RangePattern` variants, `[1, 1816]` and `[1817, 8192]`; the CP placement report below decides whether either variant actually fits. The Stage1 body is deliberately the unsplit baseline, so the dispatch contract can be read independently from the later implementations.\n" }, { "cell_type": "code", @@ -156,7 +156,7 @@ } }, "outputs": [], - "source": "SMEM_BUDGET = 232448\nCACHE_BYTES_PER_CONTEXT_PER_CTA = 2 * HEAD_DIM * 2\nSPECIALIZE_T = SMEM_BUDGET // CACHE_BYTES_PER_CONTEXT_PER_CTA\n\n\n@module(entry=\"gqa_decode\", target=_H200, topologies=(_CTA,))\nclass Stage1_Specialized:\n \"\"\"Dispatch the same unsplit kernel at the measured context boundary.\"\"\"\n\n @func\n def _decode_core(\n hidden: Tensor[(1, 1, HIDDEN), \"bf16\"],\n w_q: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n w_k: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_v: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_o: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n cur_pos: Tensor[(1,), \"i32\"],\n write_len: Tensor[(1,), \"i32\"],\n pos_ids: Tensor[(1,), \"i32\"],\n cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n ) -> Tensor[(1, 1, HIDDEN), \"bf16\"]:\n q = tf.reshape(tf.matmul(hidden, w_q), new_shape=(1, 1, QUERY_HEADS, HEAD_DIM))\n k = tf.reshape(tf.matmul(hidden, w_k), new_shape=(1, 1, KV_HEADS, HEAD_DIM))\n v = tf.reshape(tf.matmul(hidden, w_v), new_shape=(1, 1, KV_HEADS, HEAD_DIM))\n q_rope, k_rope = tf.rope(q, k, cos_cache, sin_cache, pos_ids)\n k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope)\n v_all = tf.cache_update(v_cache, cur_pos, write_len, v)\n k_heads = tf.transpose(\n tf.repeat_interleave(k_all, repeats=GQA_GROUP, axis=2), perm=(0, 2, 1, 3)\n )\n v_heads = tf.transpose(\n tf.repeat_interleave(v_all, repeats=GQA_GROUP, axis=2), perm=(0, 2, 1, 3)\n )\n q_f32 = tf.cast(q_rope, dtype=\"f32\")\n k_f32 = tf.cast(k_heads, dtype=\"f32\")\n v_f32 = tf.cast(v_heads, dtype=\"f32\")\n scaled_q = q_f32 * tf.full_like(q_f32, value=SCALE)\n q_e = tf.reshape(scaled_q, new_shape=(1, 1, QUERY_HEADS, 1, HEAD_DIM))\n k_e = tf.reshape(k_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM))\n v_e = tf.reshape(v_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM))\n scores = tf.reduce(q_e * k_e, axes=(-1,), keepdim=True, kind=\"sum\")\n peak = tf.reduce(scores, axes=(-2,), keepdim=True, kind=\"max\")\n weights = tf.exp(scores - peak)\n normalizer = tf.reduce(weights, axes=(-2,), keepdim=False, kind=\"sum\")\n weighted = tf.reduce(weights * v_e, axes=(-2,), keepdim=False, kind=\"sum\")\n attended = tf.cast(weighted / normalizer, dtype=\"bf16\")\n return tf.matmul(tf.reshape(attended, new_shape=(1, 1, HIDDEN)), w_o)\n\n @func\n def gqa_decode(\n hidden: Tensor[(1, 1, HIDDEN), \"bf16\"],\n w_q: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n w_k: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_v: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_o: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n cur_pos: Tensor[(1,), \"i32\"],\n write_len: Tensor[(1,), \"i32\"],\n pos_ids: Tensor[(1,), \"i32\"],\n cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n ) -> Tensor[(1, 1, HIDDEN), \"bf16\"]:\n pass\n\n @gqa_decode.specialize(DimVarRangePat(\"ctx_len\", 1, SPECIALIZE_T))\n def short_context(\n hidden: Tensor[(1, 1, HIDDEN), \"bf16\"],\n w_q: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n w_k: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_v: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_o: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n cur_pos: Tensor[(1,), \"i32\"],\n write_len: Tensor[(1,), \"i32\"],\n pos_ids: Tensor[(1,), \"i32\"],\n cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n ) -> Tensor[(1, 1, HIDDEN), \"bf16\"]:\n return _decode_core(\n hidden, w_q, w_k, w_v, w_o, k_cache, v_cache,\n cur_pos, write_len, pos_ids, cos_cache, sin_cache,\n )\n\n @gqa_decode.specialize(DimVarRangePat(\"ctx_len\", SPECIALIZE_T + 1, ROPE_CONTEXT))\n def long_context(\n hidden: Tensor[(1, 1, HIDDEN), \"bf16\"],\n w_q: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n w_k: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_v: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_o: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n cur_pos: Tensor[(1,), \"i32\"],\n write_len: Tensor[(1,), \"i32\"],\n pos_ids: Tensor[(1,), \"i32\"],\n cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n ) -> Tensor[(1, 1, HIDDEN), \"bf16\"]:\n return _decode_core(\n hidden, w_q, w_k, w_v, w_o, k_cache, v_cache,\n cur_pos, write_len, pos_ids, cos_cache, sin_cache,\n )\n\n\ngqa_decode_specialized = Stage1_Specialized.entry_function()\n" + "source": "SMEM_BUDGET = 232448\nCACHE_BYTES_PER_CONTEXT_PER_CTA = 2 * HEAD_DIM * 2\nSPECIALIZE_T = SMEM_BUDGET // CACHE_BYTES_PER_CONTEXT_PER_CTA\n\n\n@module(entry=\"gqa_decode\", target=_H200, topologies=(_CTA,))\nclass Stage1_Specialized:\n \"\"\"Dispatch the same unsplit kernel at the measured context boundary.\"\"\"\n\n @func\n def _decode_core(\n hidden: Tensor[(1, 1, HIDDEN), \"bf16\"],\n w_q: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n w_k: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_v: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_o: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n cur_pos: Tensor[(1,), \"i32\"],\n write_len: Tensor[(1,), \"i32\"],\n pos_ids: Tensor[(1,), \"i32\"],\n cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n ) -> Tensor[(1, 1, HIDDEN), \"bf16\"]:\n q = tf.reshape(tf.matmul(hidden, w_q), new_shape=(1, 1, QUERY_HEADS, HEAD_DIM))\n k = tf.reshape(tf.matmul(hidden, w_k), new_shape=(1, 1, KV_HEADS, HEAD_DIM))\n v = tf.reshape(tf.matmul(hidden, w_v), new_shape=(1, 1, KV_HEADS, HEAD_DIM))\n q_rope, k_rope = tf.rope(q, k, cos_cache, sin_cache, pos_ids)\n k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope)\n v_all = tf.cache_update(v_cache, cur_pos, write_len, v)\n k_heads = tf.transpose(\n tf.repeat_interleave(k_all, repeats=GQA_GROUP, axis=2), perm=(0, 2, 1, 3)\n )\n v_heads = tf.transpose(\n tf.repeat_interleave(v_all, repeats=GQA_GROUP, axis=2), perm=(0, 2, 1, 3)\n )\n q_f32 = tf.cast(q_rope, dtype=\"f32\")\n k_f32 = tf.cast(k_heads, dtype=\"f32\")\n v_f32 = tf.cast(v_heads, dtype=\"f32\")\n scaled_q = q_f32 * tf.full_like(q_f32, value=SCALE)\n q_e = tf.reshape(scaled_q, new_shape=(1, 1, QUERY_HEADS, 1, HEAD_DIM))\n k_e = tf.reshape(k_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM))\n v_e = tf.reshape(v_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM))\n scores = tf.reduce(q_e * k_e, axes=(-1,), keepdim=True, kind=\"sum\")\n peak = tf.reduce(scores, axes=(-2,), keepdim=True, kind=\"max\")\n weights = tf.exp(scores - peak)\n normalizer = tf.reduce(weights, axes=(-2,), keepdim=False, kind=\"sum\")\n weighted = tf.reduce(weights * v_e, axes=(-2,), keepdim=False, kind=\"sum\")\n attended = tf.cast(weighted / normalizer, dtype=\"bf16\")\n return tf.matmul(tf.reshape(attended, new_shape=(1, 1, HIDDEN)), w_o)\n\n @func\n def gqa_decode(\n hidden: Tensor[(1, 1, HIDDEN), \"bf16\"],\n w_q: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n w_k: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_v: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_o: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n cur_pos: Tensor[(1,), \"i32\"],\n write_len: Tensor[(1,), \"i32\"],\n pos_ids: Tensor[(1,), \"i32\"],\n cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n ) -> Tensor[(1, 1, HIDDEN), \"bf16\"]:\n pass\n\n @gqa_decode.specialize(RangePattern(\"ctx_len\", 1, SPECIALIZE_T))\n def short_context(\n hidden: Tensor[(1, 1, HIDDEN), \"bf16\"],\n w_q: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n w_k: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_v: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_o: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n cur_pos: Tensor[(1,), \"i32\"],\n write_len: Tensor[(1,), \"i32\"],\n pos_ids: Tensor[(1,), \"i32\"],\n cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n ) -> Tensor[(1, 1, HIDDEN), \"bf16\"]:\n return _decode_core(\n hidden, w_q, w_k, w_v, w_o, k_cache, v_cache,\n cur_pos, write_len, pos_ids, cos_cache, sin_cache,\n )\n\n @gqa_decode.specialize(RangePattern(\"ctx_len\", SPECIALIZE_T + 1, ROPE_CONTEXT))\n def long_context(\n hidden: Tensor[(1, 1, HIDDEN), \"bf16\"],\n w_q: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n w_k: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_v: ConstTensor[(1, HIDDEN, KV_DIM), \"bf16\"],\n w_o: ConstTensor[(1, HIDDEN, HIDDEN), \"bf16\"],\n k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), \"bf16\"],\n cur_pos: Tensor[(1,), \"i32\"],\n write_len: Tensor[(1,), \"i32\"],\n pos_ids: Tensor[(1,), \"i32\"],\n cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), \"bf16\"],\n ) -> Tensor[(1, 1, HIDDEN), \"bf16\"]:\n return _decode_core(\n hidden, w_q, w_k, w_v, w_o, k_cache, v_cache,\n cur_pos, write_len, pos_ids, cos_cache, sin_cache,\n )\n\n\ngqa_decode_specialized = Stage1_Specialized.entry_function()\n" }, { "cell_type": "markdown", @@ -470,7 +470,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "```text\nweight staging: static tensor -> one output slice -> reusable for the step\ncache staging: growing context -> one block -> compute -> next block\n```\n\n## Feature ledger\n\nThe page uses this embedded ladder for features orthogonal to GQA:\n\n| feature | live program |\n|---|---|\n| `@module(entry/target/topologies)` | `Stage0_Naive` |\n| `@func`, `Tensor`, `ConstTensor`, `DimVar` | `Stage0_Naive` |\n| `pass` prototype and `@f.specialize(DimVarRangePat)` | `Stage1_Specialized` |\n| single `Mesh`, shard sugar `X @ m.axis`, `reshard` to smem/gmem | `Stage2_Sharded` |\n| split-K worker mesh and online softmax state | `Stage3_Fused` |\n| weight staging and output gather | `Stage4_WeightPrepared` |\n| cache update, block scan, `matmul`, `rope`, `reduce`, `cast` | `Stage5_CachePrepared` |\n| nested Mesh, `rmem`, rank-changing `reshard`, multi-level `Topology` | not in this embedded ladder |\n| runtime weight converter | [migrate](migrate.md) |\n\nThe command surface used by this page is:\n\n```text\n--compute-cost logical work and executed totals\n--memory movement, residency, and placement peaks\n--roofline ideal bound and limiting resource\n--performance per-level execution projection\n--operands operand split in annotated call lines\n--dim bind ctx_len for one static analysis run\n--json write the same report data as JSON\n```\n\nFor example, the JSON form writes to a path just like the text form. The Bash cell runs the\nCLI with `--json`; the following Python cell loads the JSON report and prints a stable summary.\n" + "source": "```text\nweight staging: static tensor -> one output slice -> reusable for the step\ncache staging: growing context -> one block -> compute -> next block\n```\n\n## Feature ledger\n\nThe page uses this embedded ladder for features orthogonal to GQA:\n\n| feature | live program |\n|---|---|\n| `@module(entry/target/topologies)` | `Stage0_Naive` |\n| `@func`, `Tensor`, `ConstTensor`, `DimVar` | `Stage0_Naive` |\n| `pass` prototype and `@f.specialize(RangePattern)` | `Stage1_Specialized` |\n| single `Mesh`, shard sugar `X @ m.axis`, `reshard` to smem/gmem | `Stage2_Sharded` |\n| split-K worker mesh and online softmax state | `Stage3_Fused` |\n| weight staging and output gather | `Stage4_WeightPrepared` |\n| cache update, block scan, `matmul`, `rope`, `reduce`, `cast` | `Stage5_CachePrepared` |\n| nested Mesh, `rmem`, rank-changing `reshard`, multi-level `Topology` | not in this embedded ladder |\n| runtime weight converter | [migrate](migrate.md) |\n\nThe command surface used by this page is:\n\n```text\n--compute-cost logical work and executed totals\n--memory movement, residency, and placement peaks\n--roofline ideal bound and limiting resource\n--performance per-level execution projection\n--operands operand split in annotated call lines\n--dim bind ctx_len for one static analysis run\n--json write the same report data as JSON\n```\n\nFor example, the JSON form writes to a path just like the text form. The Bash cell runs the\nCLI with `--json`; the following Python cell loads the JSON report and prints a stable summary.\n" }, { "cell_type": "code", diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index 8361e841..f1fafca5 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -80,7 +80,7 @@ from pathlib import Path from tilefoundry import func, module from tilefoundry.analysis import analyze as run_analysis -from tilefoundry.dsl import ConstTensor, DimVar, DimVarRangePat, Mesh, Tensor, tf +from tilefoundry.dsl import ConstTensor, DimVar, RangePattern, Mesh, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 - bare tile() in the fused body from tilefoundry.inspection.analysis_report import render_analysis, render_text from tilefoundry.ir.types import Topology @@ -286,7 +286,7 @@ T = floor(232448 B / 128 B) = 1816 ``` -This is not the complete placement peak: other simultaneously resident values also occupy smem. `Stage1_Specialized` still uses the estimate to express two closed `DimVarRangePat` variants, `[1, 1816]` and `[1817, 8192]`; the CP placement report below decides whether either variant actually fits. The Stage1 body is deliberately the unsplit baseline, so the dispatch contract can be read independently from the later implementations. +This is not the complete placement peak: other simultaneously resident values also occupy smem. `Stage1_Specialized` still uses the estimate to express two closed `RangePattern` variants, `[1, 1816]` and `[1817, 8192]`; the CP placement report below decides whether either variant actually fits. The Stage1 body is deliberately the unsplit baseline, so the dispatch contract can be read independently from the later implementations. @@ -359,7 +359,7 @@ class Stage1_Specialized: ) -> Tensor[(1, 1, HIDDEN), "bf16"]: pass - @gqa_decode.specialize(DimVarRangePat("ctx_len", 1, SPECIALIZE_T)) + @gqa_decode.specialize(RangePattern("ctx_len", 1, SPECIALIZE_T)) def short_context( hidden: Tensor[(1, 1, HIDDEN), "bf16"], w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"], @@ -379,7 +379,7 @@ class Stage1_Specialized: cur_pos, write_len, pos_ids, cos_cache, sin_cache, ) - @gqa_decode.specialize(DimVarRangePat("ctx_len", SPECIALIZE_T + 1, ROPE_CONTEXT)) + @gqa_decode.specialize(RangePattern("ctx_len", SPECIALIZE_T + 1, ROPE_CONTEXT)) def long_context( hidden: Tensor[(1, 1, HIDDEN), "bf16"], w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"], @@ -1114,7 +1114,7 @@ The page uses this embedded ladder for features orthogonal to GQA: |---|---| | `@module(entry/target/topologies)` | `Stage0_Naive` | | `@func`, `Tensor`, `ConstTensor`, `DimVar` | `Stage0_Naive` | -| `pass` prototype and `@f.specialize(DimVarRangePat)` | `Stage1_Specialized` | +| `pass` prototype and `@f.specialize(RangePattern)` | `Stage1_Specialized` | | single `Mesh`, shard sugar `X @ m.axis`, `reshard` to smem/gmem | `Stage2_Sharded` | | split-K worker mesh and online softmax state | `Stage3_Fused` | | weight staging and output gather | `Stage4_WeightPrepared` | diff --git a/src/tilefoundry/__init__.py b/src/tilefoundry/__init__.py index caecfadc..c5251840 100644 --- a/src/tilefoundry/__init__.py +++ b/src/tilefoundry/__init__.py @@ -34,17 +34,32 @@ ) -from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern +from tilefoundry.ir.pattern import RangePattern, Pattern from tilefoundry.ir.types import DType, TensorType, TupleType, Type from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types import B, Broadcast, ComposedLayout, Dynamic, IntTuple, Layout, LayoutBase, Mesh, P, Partial, S, ShardAttr, ShardLayout, Split, Swizzle, Topology +from tilefoundry.ir.types import ( + B, + Broadcast, + ComposedLayout, + Dynamic, + IntTuple, + Layout, + LayoutBase, + Mesh, + P, + Partial, + S, + ShardAttr, + ShardLayout, + Split, + Swizzle, + Topology, +) from tilefoundry.ir.tir.stmt import Stmt - - from tilefoundry.ir.types import _register_dim_typeinfer @@ -52,8 +67,6 @@ from tilefoundry.ir import tir as _tir # noqa: F401 - - from tilefoundry.visitor_registry import op_cost as _op_cost # noqa: F401 @@ -70,6 +83,7 @@ _register_dim_typeinfer() + def view(root, *, port: int = 0, open_browser: bool = True) -> int: """Start the interactive HIR viewer for *root* (Function or Module). @@ -77,28 +91,58 @@ def view(root, *, port: int = 0, open_browser: bool = True) -> int: """ return _Viewer(root).serve(port=port, open_browser=open_browser) + __all__ = [ "__version__", - - "Expr", "Var", "Constant", "Call", "Stmt", "TupleGetItem", - "Op", "ParameterInfo", + "Expr", + "Var", + "Constant", + "Call", + "Stmt", + "TupleGetItem", + "Op", + "ParameterInfo", "DispatchRegistry", - "typeinfer_registry", "verify_stmt_registry", "cost_evaluator_registry", - "register_typeinfer", "register_verify_stmt", "register_cost_evaluator", - "TypeInferContext", "FunctionScope", + "typeinfer_registry", + "verify_stmt_registry", + "cost_evaluator_registry", + "register_typeinfer", + "register_verify_stmt", + "register_cost_evaluator", + "TypeInferContext", + "FunctionScope", "VerifyError", - - "DType", "TensorType", "TupleType", "Type", - "Pattern", "DimVarRangePat", "DimVar", - - "IntTuple", "LayoutBase", "Layout", "Swizzle", "ComposedLayout", - "Topology", "Mesh", - "ShardAttr", "Split", "Partial", "Broadcast", "Dynamic", "ShardLayout", - "S", "P", "B", - - "func", "prim_func", "intrinsic", "module", - - "build", "compile", "jit", - "normalize_to_module", "CompilerOptions", + "DType", + "TensorType", + "TupleType", + "Type", + "Pattern", + "RangePattern", + "DimVar", + "IntTuple", + "LayoutBase", + "Layout", + "Swizzle", + "ComposedLayout", + "Topology", + "Mesh", + "ShardAttr", + "Split", + "Partial", + "Broadcast", + "Dynamic", + "ShardLayout", + "S", + "P", + "B", + "func", + "prim_func", + "intrinsic", + "module", + "build", + "compile", + "jit", + "normalize_to_module", + "CompilerOptions", "view", ] diff --git a/src/tilefoundry/analysis/check.py b/src/tilefoundry/analysis/check.py index a872f237..75cde1e7 100644 --- a/src/tilefoundry/analysis/check.py +++ b/src/tilefoundry/analysis/check.py @@ -25,7 +25,6 @@ reachable_functions, subtree, ) -from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.hir.specialize import ( @@ -35,6 +34,7 @@ is_concrete, specialize_concretely, ) +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.types import Topology, callable_type_for from tilefoundry.ir.types.substitute import ( DimSubstitutionError, @@ -92,8 +92,7 @@ def selected_topology(self) -> Topology: """The topology level the analyses were asked about.""" if self.topology_level is None: raise AnalysisError( - "no topology level was selected, so results cannot carry an " - "execution placement" + "no topology level was selected, so results cannot carry an execution placement" ) return self.module.resolve_topology(self.topology_level) @@ -209,9 +208,7 @@ def resolve_program_geometry( execution_module = _substitute_module_tree(module, dims) except DimSubstitutionError as error: raise SpecializationError(str(error)) from None - _require_concrete_geometry( - execution_module, function, error_type=SpecializationError - ) + _require_concrete_geometry(execution_module, function, error_type=SpecializationError) return execution_module, function @@ -228,7 +225,7 @@ def visit(fn: Function) -> None: return seen.add(id(fn)) for pattern in fn.specializations: - if isinstance(pattern, DimVarRangePat): + if isinstance(pattern, RangePattern): found.add(pattern.dim_var) for variant in fn.variants: visit(variant) @@ -264,8 +261,7 @@ def _required_owner(module: Module, function: Function) -> Module: def _substitute_module_tree(module: Module, dims: Mapping[str, int]) -> Module: effective = tuple( - substitute_topology_dims(topology, dims) - for topology in module.effective_topologies() + substitute_topology_dims(topology, dims) for topology in module.effective_topologies() ) def declared(node: Module) -> Module: @@ -436,9 +432,7 @@ def visit_Call(self, expr: Call, ctx: Mapping[int, Expr]) -> Expr: def visit_MeshRegion(self, expr: MeshRegion, ctx: Mapping[int, Expr]) -> MeshRegion: """Inline through a HIR execution region while preserving its boundary.""" - params = tuple( - replace(param, metadata=_view_metadata(param)) for param in expr.params - ) + params = tuple(replace(param, metadata=_view_metadata(param)) for param in expr.params) body_ctx = {**ctx, **{id(old): new for old, new in zip(expr.params, params)}} rebuilt = replace( expr, @@ -519,9 +513,7 @@ def expr( def _inline_view(module: Module, function: Function, budget: int) -> Function: declared_resources = _resource_parameters(module, function) paths = _module_paths(module) - params = tuple( - replace(param, metadata=_view_metadata(param)) for param in function.params - ) + params = tuple(replace(param, metadata=_view_metadata(param)) for param in function.params) resources: dict[_ResourceKey, Var] = {} appended: list[Var] = [] for key, declaration in declared_resources: @@ -537,9 +529,7 @@ def _inline_view(module: Module, function: Function, budget: int) -> Function: view_params = (*params, *appended) env = {id(old): new for old, new in zip(function.params, params)} - inliner = _Inliner( - module, resources, paths, {param.name for param in view_params} - ) + inliner = _Inliner(module, resources, paths, {param.name for param in view_params}) body = inliner.function_body(function, env, (function.name,), frozenset()) size = len(collect_exprs(body)) if size > budget: @@ -564,9 +554,7 @@ def _inline_view(module: Module, function: Function, budget: int) -> Function: class InlineCloner: """Clone and inline one authored Function into an analysis view.""" - def __init__( - self, module: Module, function: Function, budget: int = _INLINE_NODES - ) -> None: + def __init__(self, module: Module, function: Function, budget: int = _INLINE_NODES) -> None: self.module = module self.function = function self.budget = budget @@ -595,8 +583,7 @@ def check_program( """ if isinstance(budget, bool) or not isinstance(budget, int) or budget < 0: raise AnalysisError( - f"inlining {function.name!r} needs a non-negative integer node budget, " - f"got {budget!r}" + f"inlining {function.name!r} needs a non-negative integer node budget, got {budget!r}" ) derived = InlineCloner(module, function, budget).clone() inference_type( diff --git a/src/tilefoundry/codegen/cuda/tir/prim_function.py b/src/tilefoundry/codegen/cuda/tir/prim_function.py index 54d967a0..a49053c7 100644 --- a/src/tilefoundry/codegen/cuda/tir/prim_function.py +++ b/src/tilefoundry/codegen/cuda/tir/prim_function.py @@ -13,7 +13,7 @@ from tilefoundry.codegen.cuda.tir.memory.tensor_view import render_shard_layout_value from tilefoundry.codegen.emitter import CudaEmitter from tilefoundry.codegen.signature import TensorSignature, tensor_signature_of -from tilefoundry.ir.core.pattern import DimVarRangePat +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.ir.types.utils import shape_numel_upper_bound @@ -68,10 +68,10 @@ def _dispatch(fn: PrimFunction, ctx: CudaCodegenContext) -> None: ctx.emit("}") -def _range_over_one_dimension(fn: PrimFunction, variant: PrimFunction) -> DimVarRangePat: +def _range_over_one_dimension(fn: PrimFunction, variant: PrimFunction) -> RangePattern: """The pattern selecting *variant*, which a branch can only be a range.""" pattern = variant.specializations[0] - if not isinstance(pattern, DimVarRangePat): + if not isinstance(pattern, RangePattern): raise NotImplementedError( f"CUDA dispatch: {fn.name!r} selects {variant.name!r} by " f"{type(pattern).__name__}; only a range over one dimension is written" @@ -79,7 +79,7 @@ def _range_over_one_dimension(fn: PrimFunction, variant: PrimFunction) -> DimVar return pattern -def _subject(fn: PrimFunction, pattern: DimVarRangePat, ctx: CudaCodegenContext) -> str: +def _subject(fn: PrimFunction, pattern: RangePattern, ctx: CudaCodegenContext) -> str: """Where this kernel reads the dimension *pattern* ranges over.""" subject = ctx.dynamic_extents.get(pattern.dim_var) if subject is None: diff --git a/src/tilefoundry/dsl/__init__.py b/src/tilefoundry/dsl/__init__.py index 468a8fef..13adc4a8 100644 --- a/src/tilefoundry/dsl/__init__.py +++ b/src/tilefoundry/dsl/__init__.py @@ -16,7 +16,7 @@ from tilefoundry.script import func -from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern +from tilefoundry.ir.pattern import RangePattern, Pattern from tilefoundry.ir.types.dim import DimVar, ceildiv from tilefoundry.ir.types import B, Broadcast, Mesh, P, Partial, S, Split, Topology from tilefoundry.ir.core.kinds import ReduceKind, UnaryKind, BinaryKind @@ -28,7 +28,7 @@ "Tensor", "func", "Pattern", - "DimVarRangePat", + "RangePattern", "DimVar", "ceildiv", "Mesh", diff --git a/src/tilefoundry/evaluator/interpreter.py b/src/tilefoundry/evaluator/interpreter.py index 5e36be8c..784eb00c 100644 --- a/src/tilefoundry/evaluator/interpreter.py +++ b/src/tilefoundry/evaluator/interpreter.py @@ -2,6 +2,7 @@ Walks a HIR ``Function`` body and returns concrete torch values. """ + from __future__ import annotations from typing import Any @@ -20,18 +21,16 @@ to_torch_dtype, ) from tilefoundry.ir.core import Call, Constant, Tuple, Var, describe_expr -from tilefoundry.ir.core.pattern import locate_dim_var from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion +from tilefoundry.ir.pattern import locate_dim_var from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.utils import types_compatible from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.utils.spec_ref import spec_ref_render -_MISSING_PREPARED_WEIGHT = ( - "[runtime §1.3](docs/spec/runtime.md#13-runtimedecoratorpy)" -) +_MISSING_PREPARED_WEIGHT = "[runtime §1.3](docs/spec/runtime.md#13-runtimedecoratorpy)" def _device_of(values) -> str | None: @@ -81,8 +80,7 @@ def _bind_dim_vars(params, values) -> dict[str, int]: prev = binding.get(dim.name) if prev is not None and prev != size: raise EvalError( - f"evaluator: inconsistent binding for DimVar " - f"{dim.name!r}: {prev} vs {size}" + f"evaluator: inconsistent binding for DimVar {dim.name!r}: {prev} vs {size}" ) binding[dim.name] = size return binding @@ -127,9 +125,7 @@ def visit_MeshRegion(self, region: MeshRegion, ctx: EvaluateContext) -> Value: def visit_leaf_Var(self, var: Var, _operands, ctx: EvaluateContext) -> Value: raise EvalError(f"evaluator: unbound variable {var.name!r}") - def visit_leaf_Constant( - self, const: Constant, _operands, ctx: EvaluateContext - ) -> TensorValue: + def visit_leaf_Constant(self, const: Constant, _operands, ctx: EvaluateContext) -> TensorValue: data = torch.as_tensor( const.value, dtype=to_torch_dtype(const.type.dtype), device=ctx.device ) @@ -150,13 +146,9 @@ def visit_leaf_Call(self, call: Call, args, ctx: EvaluateContext) -> Value: except Exception as error: raise EvalError(f"evaluator: {describe_expr(call)}: {error}") from error - def _call_function( - self, callee: Function, arg_values, ctx: EvaluateContext - ) -> Value: + def _call_function(self, callee: Function, arg_values, ctx: EvaluateContext) -> Value: child = child_module_instance(ctx.loaded_module, callee) - supplied = [ - param for param in callee.params if not (child is not None and param.is_const) - ] + supplied = [param for param in callee.params if not (child is not None and param.is_const)] if len(arg_values) != len(supplied): kind = "activation(s)" if child is not None else "args" raise EvalError( @@ -208,17 +200,11 @@ def visit_LoopRegion(self, region: LoopRegion, ctx: EvaluateContext) -> Value: extent = self._resolve_loop_field(region.extent, "extent", ctx) step = self._resolve_loop_field(region.step, "step", ctx) if start < 0: - raise EvalError( - f"evaluator: LoopRegion start must be non-negative, got {start}" - ) + raise EvalError(f"evaluator: LoopRegion start must be non-negative, got {start}") if extent < 0: - raise EvalError( - f"evaluator: LoopRegion extent must be non-negative, got {extent}" - ) + raise EvalError(f"evaluator: LoopRegion extent must be non-negative, got {extent}") if step <= 0: - raise EvalError( - f"evaluator: LoopRegion step must be positive, got {step}" - ) + raise EvalError(f"evaluator: LoopRegion step must be positive, got {step}") indices = range(start, extent, step) def iter_memo(i: int, carried) -> dict: @@ -227,9 +213,9 @@ def iter_memo(i: int, carried) -> dict: id(iv): ( iv, TensorValue( - data=torch.as_tensor(i, dtype=iv_dtype, device=ctx.device), - type=iv.type, - ), + data=torch.as_tensor(i, dtype=iv_dtype, device=ctx.device), + type=iv.type, + ), ), } for phi, value in zip(region.carried_args, carried): @@ -237,14 +223,11 @@ def iter_memo(i: int, carried) -> dict: return memo if not region.carried_args: - last = None for i in indices: last = EvaluatorVisitor(memo=iter_memo(i, ())).visit(region.body, ctx) if last is None: - raise EvalError( - "evaluator: LoopRegion has an empty iteration domain" - ) + raise EvalError("evaluator: LoopRegion has an empty iteration domain") return last carried = list(init_values) @@ -318,9 +301,7 @@ def _run_selected(loaded_module, fn: Function, *activations, device: str | None) ) supplied = iter(activations) args = [ - _read_weight(loaded_module, param, device) - if param.is_const - else next(supplied) + _read_weight(loaded_module, param, device) if param.is_const else next(supplied) for param in fn.params ] return _run_bound(fn, args, device=device, reading=loaded_module) @@ -335,7 +316,7 @@ def _unwrap(value: Value) -> Any: def _select_variant(callee: Function, arg_values) -> Function: - """Pick the variant whose ``DimVarRangePat`` matches the runtime arg shapes. + """Pick the variant whose ``RangePattern`` matches the runtime arg shapes. Errors unless exactly one matches — dispatch never falls back to the prototype body. diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 3bf4e006..1bdb7c6a 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -6,9 +6,9 @@ from contextlib import contextmanager from tilefoundry.ir.core import Call, Constant, Tuple, Var -from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord from tilefoundry.ir.mesh_scope import device_layout +from tilefoundry.ir.pattern import Pattern, RangePattern from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType from tilefoundry.ir.types.dim import ( @@ -214,10 +214,7 @@ def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: or ctx is None ): return None - if any( - isinstance(entry, tuple) - for entry in (*layout.shape, *(layout.strides or ())) - ): + if any(isinstance(entry, tuple) for entry in (*layout.shape, *(layout.strides or ()))): 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): @@ -236,7 +233,7 @@ def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: splits.setdefault(attr.axis, []).append(ref) named_bindings.add(binding) elif isinstance(attr, Partial): - partials.append(f'{ref} @ {self.visit(attr, ctx)}') + partials.append(f"{ref} @ {self.visit(attr, ctx)}") named_bindings.add(binding) elif isinstance(attr, Broadcast): broadcasts.append((binding, ref, attr)) @@ -255,8 +252,7 @@ def shard_surface(self, value: ShardLayout, ctx=None) -> str | 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) + 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 @@ -280,7 +276,9 @@ def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: 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)) + 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)}"' @@ -410,7 +408,11 @@ def render_value(self, value, ctx=None, indent: str = "") -> str: return self.atom_reference(value, ctx) if isinstance(value, enum.Enum): if ctx is not None: - ctx.use(PythonExpr((f"from {type(value).__module__} import {type(value).__name__}",), "")) + ctx.use( + PythonExpr( + (f"from {type(value).__module__} import {type(value).__name__}",), "" + ) + ) return f"{type(value).__name__}.{value.name}" if isinstance(value, Target): rendered = value.to_python() @@ -423,10 +425,10 @@ def render_value(self, value, ctx=None, indent: str = "") -> str: 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 isinstance(pattern, RangePattern): 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})' + ctx.use(PythonExpr(("from tilefoundry.ir.pattern import RangePattern",), "")) + return f'RangePattern("{pattern.dim_var}", {pattern.lo}, {pattern.hi})' return repr(pattern) diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index 848be7da..90be24cf 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -13,14 +13,14 @@ from collections.abc import Iterator from dataclasses import dataclass -from tilefoundry.ir.constraints import ( - LayoutConstraint, - MeshConstraint, - ScheduleConstraintMetadata, - StorageConstraint, - constraint_metadata, +from tilefoundry.ir.clause import ( + LayoutClause, + MeshClause, + StorageClause, + WhereClauseMetadata, + clause_metadata, ) -from tilefoundry.ir.constraints.layout import is_layout_wildcard +from tilefoundry.ir.clause.layout import is_layout_wildcard from tilefoundry.ir.core import ( Call, Constant, @@ -97,9 +97,7 @@ def bind_def( def tuple_reference(self, elements) -> str: inner = ", ".join( - repr(element.value) - if isinstance(element, Constant) - else self.reference(element) + repr(element.value) if isinstance(element, Constant) else self.reference(element) for element in elements ) return f"({inner}{',' if len(elements) == 1 else ''})" @@ -168,12 +166,7 @@ def visit_program_call(self, expr: Call, ctx=None) -> str: 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 - ): + if isinstance(start, Constant) and start.value == 0 and size == dim and stride == 1: indexers.append(":") continue if not ( @@ -186,17 +179,13 @@ def visit_program_call(self, expr: Call, ctx=None) -> str: break begin = int(start.value) stop = begin + size * stride - indexers.append( - f"{begin}:{stop}" if stride == 1 else f"{begin}:{stop}:{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 - ) + for start, size, stride in zip(starts.elements, target.sizes, target.strides) ) if len(starts.elements) == 1: start_refs += "," @@ -325,8 +314,7 @@ def _attr_tuple_str(value: tuple, printer: PythonPrinter, ctx) -> str: the header emits binds the name, not the repr. """ rendered = tuple( - printer.visit(entry, ctx) if _is_dim_entry(entry) else repr(entry) - for entry in value + printer.visit(entry, ctx) if _is_dim_entry(entry) else repr(entry) for entry in value ) if len(rendered) == 1: return f"({rendered[0]},)" @@ -396,19 +384,33 @@ def _kinded_alias_name(target) -> str | None: def _build_kinded_alias_maps(): return ( { - BinaryKind.ADD: "add", BinaryKind.SUB: "sub", BinaryKind.MUL: "mul", - BinaryKind.DIV: "div", BinaryKind.FLOOR_DIV: "floor_div", - BinaryKind.MOD: "mod", BinaryKind.MIN: "min", BinaryKind.MAX: "max", - BinaryKind.EQ: "cmp_eq", BinaryKind.NE: "cmp_ne", - BinaryKind.LT: "cmp_lt", BinaryKind.LE: "cmp_le", - BinaryKind.GT: "cmp_gt", BinaryKind.GE: "cmp_ge", - BinaryKind.AND: "logical_and", BinaryKind.OR: "logical_or", + BinaryKind.ADD: "add", + BinaryKind.SUB: "sub", + BinaryKind.MUL: "mul", + BinaryKind.DIV: "div", + BinaryKind.FLOOR_DIV: "floor_div", + BinaryKind.MOD: "mod", + BinaryKind.MIN: "min", + BinaryKind.MAX: "max", + BinaryKind.EQ: "cmp_eq", + BinaryKind.NE: "cmp_ne", + BinaryKind.LT: "cmp_lt", + BinaryKind.LE: "cmp_le", + BinaryKind.GT: "cmp_gt", + BinaryKind.GE: "cmp_ge", + BinaryKind.AND: "logical_and", + BinaryKind.OR: "logical_or", }, { - UnaryKind.NEG: "neg", UnaryKind.ABS: "abs", UnaryKind.NOT: "logical_not", - UnaryKind.EXP: "exp", UnaryKind.LOG: "log", - UnaryKind.CEIL: "ceil", UnaryKind.ROUND: "round", - UnaryKind.EXP2: "exp2", UnaryKind.LOG2: "log2", + UnaryKind.NEG: "neg", + UnaryKind.ABS: "abs", + UnaryKind.NOT: "logical_not", + UnaryKind.EXP: "exp", + UnaryKind.LOG: "log", + UnaryKind.CEIL: "ceil", + UnaryKind.ROUND: "round", + UnaryKind.EXP2: "exp2", + UnaryKind.LOG2: "log2", }, ) @@ -448,7 +450,7 @@ def _constraint_value_str(value: object) -> str: return repr(value) -def _layout_constraint_str(constraint: LayoutConstraint) -> str: +def _layout_constraint_str(constraint: LayoutClause) -> str: split_bindings = { attr.axis: (topology, attr) for topology, attr in constraint.bindings @@ -458,17 +460,12 @@ def _layout_constraint_str(constraint: LayoutConstraint) -> str: for index, extent in enumerate(constraint.layout.shape): if index in split_bindings: topology, _ = split_bindings[index] - dims.append( - f"{_constraint_value_str(extent)} @ " - f"{_constraint_value_str(topology)}" - ) + dims.append(f"{_constraint_value_str(extent)} @ {_constraint_value_str(topology)}") else: dims.append(_constraint_value_str(extent)) dims_str = "(" + ", ".join(dims) + ("," if len(dims) == 1 else "") + ")" bindings = [ - (topology, attr) - for topology, attr in constraint.bindings - if not isinstance(attr, Split) + (topology, attr) for topology, attr in constraint.bindings if not isinstance(attr, Split) ] if not bindings: return dims_str @@ -477,32 +474,30 @@ def _layout_constraint_str(constraint: LayoutConstraint) -> str: if isinstance(attr, Broadcast): binding_str.append(f"{_constraint_value_str(topology)} @ B()") elif isinstance(attr, Partial): - binding_str.append( - f'{_constraint_value_str(topology)} @ P("{attr.reduction}")' - ) - else: # pragma: no cover - LayoutConstraint validates this type + binding_str.append(f'{_constraint_value_str(topology)} @ P("{attr.reduction}")') + else: # pragma: no cover - LayoutClause validates this type raise TypeError(f"unsupported layout binding {type(attr).__name__}") return f"({dims_str}, {{{', '.join(binding_str)}}})" -def _where_str(metadata: ScheduleConstraintMetadata) -> str: +def _where_str(metadata: WhereClauseMetadata) -> str: layout = next( - (item for item in metadata.constraints if isinstance(item, LayoutConstraint)), + (item for item in metadata.constraints if isinstance(item, LayoutClause)), None, ) fields: list[str] = [] if layout is not None: fields.append(f"layout={_layout_constraint_str(layout)}") for item in metadata.constraints: - if isinstance(item, MeshConstraint): + if isinstance(item, MeshClause): fields.append(f"mesh={PythonPrinter().visit(item.mesh, HirPrintContext())}") - elif isinstance(item, StorageConstraint): + elif isinstance(item, StorageClause): fields.append(f'storage="{item.storage.name.lower()}"') return "where(" + ", ".join(fields) + ")" def _constraint_line(expr: Expr, indent: str, name: str) -> str | None: - metadata = constraint_metadata(expr) + metadata = clause_metadata(expr) if metadata is None: return None return f"{indent}{name}: {_where_str(metadata)}" @@ -533,9 +528,7 @@ def iter_exprs(root: Expr | None, seen: set[int] | None = None) -> Iterator[Expr def _region_projection(expr: Expr) -> LoopRegion | MeshRegion | None: """Return the region projected by a one-argument ``TupleGetItem``.""" if not ( - isinstance(expr, Call) - and isinstance(expr.target, TupleGetItem) - and len(expr.args) == 1 + isinstance(expr, Call) and isinstance(expr.target, TupleGetItem) and len(expr.args) == 1 ): return None region = expr.args[0] @@ -565,8 +558,12 @@ def _module_callee_binding(target: HirFunction, child_entries: dict[int, str]) - def _emit_def( - fn: HirFunction, def_name: str, ctx: HirPrintContext, indent: str, - options: PythonPrintOptions, child_entries: dict[int, str] | None = None, + fn: HirFunction, + def_name: str, + ctx: HirPrintContext, + indent: str, + options: PythonPrintOptions, + child_entries: dict[int, str] | None = None, *, line_offset: int = 0, statements: dict[int, _PrintedStatement] | None = None, @@ -581,19 +578,13 @@ 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 - ) + 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] = {} - - _seen: set[int] = set() _order: list[Expr] = list(iter_exprs(fn.body, _seen)) for p in fn.params: @@ -607,7 +598,6 @@ def _emit_def( for param, arg in zip(scope.params, scope.args, strict=True) } - _op_names_set: set[str] = set() for expr in _order: if isinstance(expr, Call): @@ -634,30 +624,24 @@ def _emit_def( for expr in _order: if not isinstance(expr, LoopRegion): continue - if ( - any( - isinstance(candidate, Call) - and isinstance(candidate.target, Slice) - and id(candidate) not in collapsed_slice_ids - and len(candidate.args) == 2 - and isinstance(candidate.args[1], Tuple) - and any( - window_base(start)[0] is expr.induction_var - and size == expr.step - and stride == 1 - for start, size, stride in zip( - candidate.args[1].elements, - candidate.target.sizes, - candidate.target.strides, - ) + if any( + isinstance(candidate, Call) + and isinstance(candidate.target, Slice) + and id(candidate) not in collapsed_slice_ids + and len(candidate.args) == 2 + and isinstance(candidate.args[1], Tuple) + and any( + window_base(start)[0] is expr.induction_var and size == expr.step and stride == 1 + for start, size, stride in zip( + candidate.args[1].elements, + candidate.target.sizes, + candidate.target.strides, ) - for candidate in _order ) + for candidate in _order ): _tile_window_steps[id(expr.induction_var)] = expr.step - for carry, init, value in zip( - expr.carried_args, expr.init_args, expr.yield_values - ): + for carry, init, value in zip(expr.carried_args, expr.init_args, expr.yield_values): _forced_names[id(carry)] = _sanitize_name(carry.name) _forced_names[id(init)] = _sanitize_name(carry.name) for _ in iter_exprs(expr.body, _grid_internal_ids): @@ -686,18 +670,14 @@ def _emit_def( if isinstance(expr, MeshRegion): for _ in iter_exprs(expr.body, _mesh_region_internal_ids): pass + def _moved_window(start, size, stride): """The tile window and offset *start* moves it by, else ``None``.""" window, offset = window_base(start) - if ( - isinstance(window, Var) - and stride == 1 - and _tile_window_steps.get(id(window)) == size - ): + if isinstance(window, Var) and stride == 1 and _tile_window_steps.get(id(window)) == size: return window, offset return None - _inlined_start_ids = { id(start) for expr in _order @@ -738,7 +718,6 @@ def _assign_name(expr: Expr) -> str: _names[key] = name return name - for expr in _order: _assign_name(expr) for expr in _order: @@ -791,10 +770,7 @@ def _emit_inline_call(expr: Call, level: str) -> None: ) with printer.type_surface(indent=level): rendered = printer.visit(expr, ctx) - lines.append( - f"{level}{name} = {rendered}" - f"{_comments(expr, options, printer, ctx)}" - ) + lines.append(f"{level}{name} = {rendered}{_comments(expr, options, printer, ctx)}") printed.add(id(expr)) def _emit_expr(expr: Expr, level: str) -> None: @@ -864,7 +840,9 @@ def _emit_loop_region(region: LoopRegion, level: str) -> None: 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, printer, ctx)}") + 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) @@ -888,8 +866,7 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) mesh_text = printer.visit(region.mesh, ctx) mesh_name = ctx.scope_name(region.mesh) lines.append( - f"{level}with {mesh_text} as {mesh_name}:" - f"{_comments(region, options, printer, ctx)}" + f"{level}with {mesh_text} as {mesh_name}:{_comments(region, options, printer, ctx)}" ) printed.add(key) inner = level + " " @@ -925,17 +902,15 @@ 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, printer, ctx)}") + 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) printed.add(id(expr)) continue if isinstance(expr, Tuple): - - - - continue if isinstance(expr, Call): name = _names[id(expr)] @@ -946,17 +921,12 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) ) with printer.type_surface(indent=indent): rendered = printer.visit(expr, ctx) - lines.append( - f"{indent}{name} = {rendered}" - f"{_comments(expr, options, printer, ctx)}" - ) + lines.append(f"{indent}{name} = {rendered}{_comments(expr, options, printer, ctx)}") line = _constraint_line(expr, indent, name) if line is not None: lines.append(line) printed.add(id(expr)) - - if not isinstance(fn.body, MeshRegion): if isinstance(fn.body, Tuple): lines.append(f"{indent}return {printer.tuple_reference(fn.body.elements)}") @@ -972,7 +942,6 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) return lines - def _new_hir_context(*, for_module: bool = False, target=None) -> HirPrintContext: """Create a HIR context with imports owned by the surrounding file.""" ctx = HirPrintContext() @@ -984,6 +953,7 @@ def _new_hir_context(*, for_module: bool = False, target=None) -> HirPrintContex 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.""" label = display_name(variant) @@ -994,7 +964,10 @@ def _variant_binding_name(variant: HirFunction) -> str: def _emit_decorated_defs( - fn: HirFunction, ctx: HirPrintContext, indent: str, options: PythonPrintOptions, + fn: HirFunction, + ctx: HirPrintContext, + indent: str, + options: PythonPrintOptions, child_entries: dict[int, str] | None = None, *, line_offset: int = 0, @@ -1029,7 +1002,6 @@ def _emit_decorated_defs( ) ) - for variant in fn.variants: lines.append("") lines.append( @@ -1037,7 +1009,11 @@ def _emit_decorated_defs( ) lines.extend( _emit_def( - variant, _variant_binding_name(variant), ctx, indent, options, + variant, + _variant_binding_name(variant), + ctx, + indent, + options, child_entries, line_offset=line_offset + _physical_line_count(lines), statements=statements, @@ -1047,7 +1023,9 @@ def _emit_decorated_defs( def _render_hir_function( - fn: HirFunction, *, options: PythonPrintOptions | None = None, + fn: HirFunction, + *, + options: PythonPrintOptions | None = None, ) -> _PythonRendering: """Render a HIR Function and locate every Call equation in the same pass. @@ -1077,14 +1055,18 @@ def _render_hir_function( def hir_function_to_python( - fn: HirFunction, *, options: PythonPrintOptions | None = None, + fn: HirFunction, + *, + options: PythonPrintOptions | None = None, ) -> str: """Convert a HIR Function to canonical Python DSL source.""" return HirPrinter().print(fn, options=options) def as_script( - fn: HirFunction | PrimFunction | Module, *, module: str | None = None, + fn: HirFunction | PrimFunction | Module, + *, + module: str | None = None, options: PythonPrintOptions | None = None, ) -> str: """Convert an HIR function or module to Python DSL source. @@ -1101,7 +1083,9 @@ def as_script( return tir_function_to_python(fn, options=options) if module is not None: if isinstance(fn, PrimFunction): - return tir_module_to_python(Module(name=module, functions=(fn,), entry=fn.name), options=options) + return tir_module_to_python( + Module(name=module, functions=(fn,), entry=fn.name), options=options + ) return _module_to_python(fn, module, options=options) return hir_function_to_python(fn, options=options) @@ -1152,17 +1136,17 @@ def _module_decorator_line(mod: Module, entry_name: str | None, ctx: HirPrintCon kwargs.append(f"target={rendered.text}") if mod.topologies is not None: ctx.imports.add("from tilefoundry.ir.types import Topology") - topo_strs = [ - f'Topology("{t.name}", {printer.visit(t.size, ctx)})' - for t in mod.topologies - ] - rendered_topologies = f'({", ".join(topo_strs)},)' if topo_strs else "()" + topo_strs = [f'Topology("{t.name}", {printer.visit(t.size, ctx)})' for t in mod.topologies] + rendered_topologies = f"({', '.join(topo_strs)},)" if topo_strs else "()" kwargs.append(f"topologies={rendered_topologies}") return f"@module({', '.join(kwargs)})" def _emit_module_class( - mod: Module, module_name: str, ctx: HirPrintContext, 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. @@ -1178,8 +1162,7 @@ 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, ctx, indent, options) - for child in mod.modules + _emit_module_class(child, child.name, ctx, indent, options) for child in mod.modules ] for fn in ordered: if isinstance(fn, HirFunction): @@ -1198,8 +1181,10 @@ def _emit_module_class( def _module_to_python( - fn_or_module: HirFunction | Module, module_name: str | None = None, - *, options: PythonPrintOptions | None = None, + fn_or_module: HirFunction | Module, + module_name: str | None = None, + *, + options: PythonPrintOptions | None = None, ) -> str: """Render a function or a whole Module tree as ``@module`` source.""" if isinstance(fn_or_module, Module): @@ -1219,11 +1204,14 @@ def _module_to_python( if entry is not None and not isinstance(entry, (HirFunction, PrimFunction)): raise TypeError("Module printer requires a function entry") - indent4 = " " ctx = _new_hir_context(for_module=True, target=root.target) lines = _emit_module_class( - root, module_name, ctx, indent4, options or PythonPrintOptions(), + root, + module_name, + ctx, + indent4, + options or PythonPrintOptions(), ) header = ctx.header() return "\n".join(header + lines) + "\n" diff --git a/src/tilefoundry/inspection/tir_printer.py b/src/tilefoundry/inspection/tir_printer.py index ed0f5162..afced9cf 100644 --- a/src/tilefoundry/inspection/tir_printer.py +++ b/src/tilefoundry/inspection/tir_printer.py @@ -74,21 +74,29 @@ def visit_program_call(self, expr: Call, ctx=None) -> str: if isinstance(target, Slice): return self._window_subscript(expr, ctx) scalar_binary = { - BinaryKind.EQ: "==", BinaryKind.NE: "!=", BinaryKind.LT: "<", - BinaryKind.LE: "<=", BinaryKind.GT: ">", BinaryKind.GE: ">=", BinaryKind.AND: "and", + 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"(? str: spans = [] starts = expr.args[1].elements for start, size, stride in zip(starts, expr.target.sizes, expr.target.strides): - low = ( - self.dim_entry(start, ctx) - if is_dim_op_call(start) - else self.visit(start, ctx) - ) + low = self.dim_entry(start, ctx) if is_dim_op_call(start) else self.visit(start, ctx) high = f"{low} + {size * stride}" spans.append(f"{low}:{high}" if stride == 1 else f"{low}:{high}:{stride}") return f"{self.visit(expr.args[0], ctx)}[{', '.join(spans)}]" @@ -135,8 +139,7 @@ def visit_For(self, stmt, ctx=None): print context holds, so the bounds are rendered through it. """ bounds = ", ".join( - self.visit(bound, self.context) - for bound in (stmt.start, stmt.stop, stmt.step) + self.visit(bound, self.context) for bound in (stmt.start, stmt.stop, stmt.step) ) lines = [f"{self.indent}for {stmt.induction_var.name} in range({bounds}):"] lines.extend(TirPrinter(context=self.context, indent=self.indent + " ").visit(stmt.body)) @@ -144,24 +147,35 @@ def visit_For(self, stmt, ctx=None): def visit_If(self, stmt, ctx=None): lines = [f"{self.indent}if {self.visit(stmt.cond)}:"] - lines.extend(TirPrinter(context=self.context, indent=self.indent + " ").visit(stmt.then_body)) + lines.extend( + TirPrinter(context=self.context, indent=self.indent + " ").visit(stmt.then_body) + ) if stmt.else_body.body: lines.append(f"{self.indent}else:") - lines.extend(TirPrinter(context=self.context, indent=self.indent + " ").visit(stmt.else_body)) + lines.extend( + TirPrinter(context=self.context, indent=self.indent + " ").visit(stmt.else_body) + ) return lines def visit_While(self, stmt, ctx=None): - return [f"{self.indent}while {self.visit(stmt.cond)}:"] + TirPrinter(context=self.context, indent=self.indent + " ").visit(stmt.body) + return [f"{self.indent}while {self.visit(stmt.cond)}:"] + TirPrinter( + context=self.context, indent=self.indent + " " + ).visit(stmt.body) def visit_Return(self, stmt, ctx=None): return [f"{self.indent}return"] - def _join_args(self, args): return ", ".join(self.visit(arg) for arg in args) + def _join_args(self, args): + return ", ".join(self.visit(arg) for arg in args) def _emit_evaluate(self, stmt): - handler = _STMT_PRINTERS.get(type(stmt.callable)) or (_STMT_PRINTERS.get(Op) if isinstance(stmt.callable, Op) else None) + handler = _STMT_PRINTERS.get(type(stmt.callable)) or ( + _STMT_PRINTERS.get(Op) if isinstance(stmt.callable, Op) else None + ) if handler is None: - raise NotImplementedError(f"TIR printer has no emitter for {type(stmt.callable).__name__}") + raise NotImplementedError( + f"TIR printer has no emitter for {type(stmt.callable).__name__}" + ) return handler(stmt, self) @@ -170,9 +184,11 @@ def _emit_evaluate(self, stmt): def register_tir_printer(node_type: type): """Register the source emitter for one TIR callable/statement type.""" + def decorate(fn): _STMT_PRINTERS[node_type] = fn return fn + return decorate @@ -186,7 +202,9 @@ def _print_launch(stmt: Evaluate, printer: TirPrinter) -> 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.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"] + 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) @@ -213,7 +231,13 @@ def _function_block(fn: PrimFunction) -> list[str]: target = ctx.use(fn.target.to_python()) ctx.use(PythonExpr(("from tilefoundry import prim_func",), "prim_func")) ctx.use(PythonExpr(("from tilefoundry.dsl import Tensor",), "Tensor")) - dim_vars = {d.name: d for p in fn.params if isinstance(p.type, TensorType) for d in p.type.shape if hasattr(d, "name")} + dim_vars = { + d.name: d + for p in fn.params + if isinstance(p.type, TensorType) + for d in p.type.shape + if hasattr(d, "name") + } if dim_vars: ctx.use(PythonExpr(("from tilefoundry.dsl import DimVar",), "DimVar")) lines = [f'_{d.name} = DimVar("{d.name}", {d.lo}, {d.hi})' for d in dim_vars.values()] @@ -226,12 +250,16 @@ def _function_block(fn: PrimFunction) -> list[str]: body = TirPrinter(context=ctx, indent=" ").visit(fn.body) lines.extend(body or [" pass"]) if fn.variants: - ctx.use(PythonExpr(("from tilefoundry.ir.core.pattern import DimVarRangePat",), "DimVarRangePat")) + ctx.use(PythonExpr(("from tilefoundry.ir.pattern import RangePattern",), "RangePattern")) for variant in fn.variants: pat = variant.specializations[0] lines.append("") - lines.append(f"@{_binding_name(fn.name)}.specialize({TirPrinter(context=ctx).render_pattern(pat, ctx)})") - lines.append(f"def {_binding_name(getattr(variant, '_display_name', variant.name))}({params}):") + lines.append( + f"@{_binding_name(fn.name)}.specialize({TirPrinter(context=ctx).render_pattern(pat, ctx)})" + ) + lines.append( + f"def {_binding_name(getattr(variant, '_display_name', variant.name))}({params}):" + ) vbody = TirPrinter(context=ctx, indent=" ").visit(variant.body) lines.extend(vbody or [" pass"]) return _RenderedLines(lines, ctx.imports) @@ -281,7 +309,9 @@ def tir_module_to_python(mod: Module, module_name: str | None = None, *, options if index: lines.append("") lines.extend((" " + line) if line else "" for line in block if " = DimVar(" not in line) - declarations = list(dict.fromkeys(line for block in blocks for line in block if " = DimVar(" in line)) + declarations = list( + dict.fromkeys(line for block in blocks for line in block if " = DimVar(" in line) + ) if declarations: remaining = [line for line in lines if line not in declarations] while remaining and not remaining[0]: diff --git a/src/tilefoundry/ir/clause/__init__.py b/src/tilefoundry/ir/clause/__init__.py new file mode 100644 index 00000000..f6d0178d --- /dev/null +++ b/src/tilefoundry/ir/clause/__init__.py @@ -0,0 +1,24 @@ +"""Typed, stage-neutral where-clause values.""" + +from .base import ( + ClauseProvenance, + SourceLocation, + WhereClause, + WhereClauseMetadata, + clause_metadata, +) +from .layout import LayoutClause, is_layout_wildcard +from .mesh import MeshClause +from .storage import StorageClause + +__all__ = [ + "ClauseProvenance", + "LayoutClause", + "is_layout_wildcard", + "MeshClause", + "WhereClause", + "WhereClauseMetadata", + "SourceLocation", + "StorageClause", + "clause_metadata", +] diff --git a/src/tilefoundry/ir/constraints/base.py b/src/tilefoundry/ir/clause/base.py similarity index 59% rename from src/tilefoundry/ir/constraints/base.py rename to src/tilefoundry/ir/clause/base.py index 554e70e3..d778469c 100644 --- a/src/tilefoundry/ir/constraints/base.py +++ b/src/tilefoundry/ir/clause/base.py @@ -1,4 +1,4 @@ -"""Stage-neutral schedule constraint metadata and source locations.""" +"""Stage-neutral where-clause metadata and source locations.""" from __future__ import annotations @@ -24,54 +24,51 @@ def describe(self) -> str: return f"{self.filename}:{self.line}:{self.column}" -class ConstraintProvenance(Enum): - """Source category for a schedule constraint.""" +class ClauseProvenance(Enum): + """Source category for a where clause.""" AUTHOR = "author" @dataclass(frozen=True) -class ScheduleConstraint: +class WhereClause: """Base value for one stage-neutral hard constraint.""" source_loc: SourceLocation = field(default_factory=SourceLocation) - provenance: ConstraintProvenance = ConstraintProvenance.AUTHOR + provenance: ClauseProvenance = ClauseProvenance.AUTHOR @dataclass(frozen=True) -class ScheduleConstraintMetadata(IRMetadata): +class WhereClauseMetadata(IRMetadata): """Aggregate hard constraints attached to one concrete tensor Expr.""" - constraints: tuple[ScheduleConstraint, ...] = () + constraints: tuple[WhereClause, ...] = () source_loc: SourceLocation = field(default_factory=SourceLocation) def __post_init__(self) -> None: constraints = tuple(self.constraints) if not constraints: - raise ValueError( - f"schedule constraints at {self.source_loc.describe()} cannot be empty" - ) - if any(not isinstance(item, ScheduleConstraint) for item in constraints): - bad = next(item for item in constraints if not isinstance(item, ScheduleConstraint)) + raise ValueError(f"where clauses at {self.source_loc.describe()} cannot be empty") + if any(not isinstance(item, WhereClause) for item in constraints): + bad = next(item for item in constraints if not isinstance(item, WhereClause)) raise TypeError( - f"schedule constraint metadata expects ScheduleConstraint values, " - f"got {type(bad).__name__}" + f"where-clause metadata expects WhereClause values, got {type(bad).__name__}" ) object.__setattr__(self, "constraints", constraints) -def constraint_metadata(expr: Any) -> ScheduleConstraintMetadata | None: +def clause_metadata(expr: Any) -> WhereClauseMetadata | None: """Return schedule metadata attached to ``expr``, if present.""" for item in getattr(expr, "metadata", ()): - if type(item) is ScheduleConstraintMetadata: + if type(item) is WhereClauseMetadata: return item return None __all__ = [ - "ConstraintProvenance", - "ScheduleConstraint", - "ScheduleConstraintMetadata", + "ClauseProvenance", + "WhereClause", + "WhereClauseMetadata", "SourceLocation", - "constraint_metadata", + "clause_metadata", ] diff --git a/src/tilefoundry/ir/constraints/layout.py b/src/tilefoundry/ir/clause/layout.py similarity index 79% rename from src/tilefoundry/ir/constraints/layout.py rename to src/tilefoundry/ir/clause/layout.py index 811bbf8b..425edec2 100644 --- a/src/tilefoundry/ir/constraints/layout.py +++ b/src/tilefoundry/ir/clause/layout.py @@ -6,7 +6,7 @@ from tilefoundry.ir.types import Layout, ShardAttr -from .base import ScheduleConstraint +from .base import WhereClause @dataclass(frozen=True) @@ -26,7 +26,7 @@ def is_layout_wildcard(value: object) -> bool: @dataclass(frozen=True) -class LayoutConstraint(ScheduleConstraint): +class LayoutClause(WhereClause): """Fix a Layout pattern and its authored ShardAttr bindings.""" layout: Layout = Layout(shape=()) @@ -34,18 +34,13 @@ class LayoutConstraint(ScheduleConstraint): def __post_init__(self) -> None: if not isinstance(self.layout, Layout): - raise TypeError( - f"layout constraint requires Layout, got " - f"{type(self.layout).__name__}" - ) + raise TypeError(f"layout constraint requires Layout, got {type(self.layout).__name__}") bindings = tuple(self.bindings) for topology, attr in bindings: if not isinstance(topology, str) or not topology: raise ValueError("layout binding topology must be non-empty") if not isinstance(attr, ShardAttr): - raise TypeError( - f"layout binding requires ShardAttr, got {type(attr).__name__}" - ) + raise TypeError(f"layout binding requires ShardAttr, got {type(attr).__name__}") if len({topology for topology, _ in bindings}) != len(bindings): raise ValueError("layout constraint cannot bind one topology more than once") object.__setattr__(self, "bindings", bindings) @@ -57,6 +52,6 @@ def physical_shape(self) -> tuple: __all__ = [ - "LayoutConstraint", + "LayoutClause", "is_layout_wildcard", ] diff --git a/src/tilefoundry/ir/constraints/mesh.py b/src/tilefoundry/ir/clause/mesh.py similarity index 59% rename from src/tilefoundry/ir/constraints/mesh.py rename to src/tilefoundry/ir/clause/mesh.py index 3ee090b7..473b2610 100644 --- a/src/tilefoundry/ir/constraints/mesh.py +++ b/src/tilefoundry/ir/clause/mesh.py @@ -6,20 +6,18 @@ from tilefoundry.ir.types import Mesh -from .base import ScheduleConstraint +from .base import WhereClause @dataclass(frozen=True) -class MeshConstraint(ScheduleConstraint): +class MeshClause(WhereClause): """Filter an eventual ShardLayout by one existing Mesh value.""" mesh: Mesh | None = None def __post_init__(self) -> None: if not isinstance(self.mesh, Mesh): - raise TypeError( - f"mesh constraint requires a Mesh, got {type(self.mesh).__name__}" - ) + raise TypeError(f"mesh constraint requires a Mesh, got {type(self.mesh).__name__}") -__all__ = ["MeshConstraint"] +__all__ = ["MeshClause"] diff --git a/src/tilefoundry/ir/constraints/storage.py b/src/tilefoundry/ir/clause/storage.py similarity index 82% rename from src/tilefoundry/ir/constraints/storage.py rename to src/tilefoundry/ir/clause/storage.py index b40df58a..7e35f4b1 100644 --- a/src/tilefoundry/ir/constraints/storage.py +++ b/src/tilefoundry/ir/clause/storage.py @@ -6,11 +6,11 @@ from tilefoundry.ir.types.storage import StorageKind, resolve_storage -from .base import ScheduleConstraint +from .base import WhereClause @dataclass(frozen=True) -class StorageConstraint(ScheduleConstraint): +class StorageClause(WhereClause): """Filter a value by one current IR StorageKind.""" storage: StorageKind | None = None @@ -22,4 +22,4 @@ def __post_init__(self) -> None: object.__setattr__(self, "storage", value) -__all__ = ["StorageConstraint"] +__all__ = ["StorageClause"] diff --git a/src/tilefoundry/ir/constraints/__init__.py b/src/tilefoundry/ir/constraints/__init__.py deleted file mode 100644 index a1eb41f2..00000000 --- a/src/tilefoundry/ir/constraints/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Typed, stage-neutral scheduling constraint values.""" - -from .base import ( - ConstraintProvenance, - ScheduleConstraint, - ScheduleConstraintMetadata, - SourceLocation, - constraint_metadata, -) -from .layout import LayoutConstraint, is_layout_wildcard -from .mesh import MeshConstraint -from .storage import StorageConstraint - -__all__ = [ - "ConstraintProvenance", - "LayoutConstraint", - "is_layout_wildcard", - "MeshConstraint", - "ScheduleConstraint", - "ScheduleConstraintMetadata", - "SourceLocation", - "StorageConstraint", - "constraint_metadata", -] diff --git a/src/tilefoundry/ir/core/overload.py b/src/tilefoundry/ir/core/overload.py index f36a6c3c..ae728dff 100644 --- a/src/tilefoundry/ir/core/overload.py +++ b/src/tilefoundry/ir/core/overload.py @@ -40,19 +40,14 @@ def _pattern_matches(pd: ParamDef, arg_type: Any) -> bool: """True iff ``pd.pattern`` accepts ``arg_type`` (or no pattern given).""" if pd.pattern is None: return True - return pd.pattern.match(arg_type) - - - + return pd.pattern.match(arg_type) is not None class OverloadError(LookupError): """No OpSchema candidate matched the given arg types.""" -def filter_candidates( - candidates: Iterable[OpSchema], arg_types: Sequence[Any] -) -> list[OpSchema]: +def filter_candidates(candidates: Iterable[OpSchema], arg_types: Sequence[Any]) -> list[OpSchema]: """Return candidates whose arity + every input pattern matches. Order is preserved; this is the raw filter without first-match @@ -75,9 +70,7 @@ def filter_candidates( return out -def resolve( - candidates: Iterable[OpSchema], arg_types: Sequence[Any] -) -> OpSchema: +def resolve(candidates: Iterable[OpSchema], arg_types: Sequence[Any]) -> OpSchema: """Return the first matching candidate (F3 first-match lock). Raises :class:`OverloadError` if no candidate matches. diff --git a/src/tilefoundry/ir/core/param_def.py b/src/tilefoundry/ir/core/param_def.py index a2c6796a..3d366a9d 100644 --- a/src/tilefoundry/ir/core/param_def.py +++ b/src/tilefoundry/ir/core/param_def.py @@ -11,6 +11,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from enum import Flag, auto from typing import Any, Literal @@ -40,6 +41,13 @@ def __bool__(self) -> bool: _ParamKind = Literal["input", "attribute"] +class MemoryEffect(Flag): + """Storage effects declared by one tensor operand.""" + + READ = auto() + WRITE = auto() + + @dataclass class ParamDef: """Class-body descriptor for an Op parameter. @@ -47,7 +55,8 @@ class ParamDef: Use as: ``src = ParamDef(kind="input", pattern=Tensor)``. The ``__set_name__`` hook records the attribute name on the - descriptor instance for later reflection. + descriptor instance for later reflection. ``effect=None`` means the + operand's memory effect is undeclared, not effect-free. """ kind: _ParamKind @@ -55,13 +64,19 @@ class ParamDef: pattern: "Pattern | None" = None optional: bool = False default: Any = MISSING + effect: MemoryEffect | None = None _attr_name: str = field(default="", init=False, repr=False) def __post_init__(self) -> None: - if self.kind not in ("input", "attribute"): raise ValueError(f"ParamDef.kind must be 'input' or 'attribute', got {self.kind!r}") + if self.effect is not None and not isinstance(self.effect, MemoryEffect): + raise TypeError("effect must be a MemoryEffect flag or None (undeclared)") + if self.kind == "attribute" and self.effect: + raise ValueError("attributes do not read or write tensor storage") + if self.kind == "input" and self.effect == MemoryEffect(0): + raise ValueError("tensor operands must declare READ, WRITE, or both") def __set_name__(self, owner: type, name: str) -> None: @@ -107,4 +122,10 @@ def collect_param_defs(cls: type) -> tuple["ParamDef", ...]: return tuple(seen[name] for name in order) -__all__ = ["ParamDef", "MISSING", "_MissingType", "collect_param_defs"] +__all__ = [ + "MISSING", + "MemoryEffect", + "ParamDef", + "_MissingType", + "collect_param_defs", +] diff --git a/src/tilefoundry/ir/core/pattern.py b/src/tilefoundry/ir/core/pattern.py deleted file mode 100644 index d87a016d..00000000 --- a/src/tilefoundry/ir/core/pattern.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Declarative predicates for overload and specialization dispatch. - -``ParamDef.pattern`` filters parser overloads; ``DimVarRangePat`` selects HIR -specializations. Patterns do not participate in static type checking. - -See [core-ir §3](docs/spec/core-ir.md#3-pattern) and -[hir §2](docs/spec/hir.md#2-function-specialization-api). -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - - -@dataclass(frozen=True) -class Pattern: - """Base. A reusable predicate. - - Subclasses override :meth:`match` to express their constraint. - The ``subject`` of :meth:`match` depends on the consumer: parser - dispatch passes an IR ``Type``; specialization-dispatch lowering - inspects the pattern's own fields and does not call :meth:`match`. - """ - - def match(self, subject: Any) -> bool: - raise NotImplementedError - - -@dataclass(frozen=True) -class ScalarPat(Pattern): - """Matches rank-0 tensor (``shape == ()``).""" - - def match(self, subject: Any) -> bool: - shape = getattr(subject, "shape", None) - return shape == () - - -@dataclass(frozen=True) -class TensorPat(Pattern): - """Matches non-scalar tensor (``shape != ()``). - - Optional ``rank`` and ``dtype`` further constrain the shape length - and dtype. Default (no constraints) matches any non-scalar tensor. - """ - - rank: int | None = None - dtype: Any = None - - def match(self, subject: Any) -> bool: - shape = getattr(subject, "shape", None) - if shape is None: - return False - if shape == (): - return False - if self.rank is not None and len(shape) != self.rank: - return False - if self.dtype is not None: - ty_dtype = getattr(subject, "dtype", None) - if ty_dtype != self.dtype: - return False - return True - - -@dataclass(frozen=True) -class AndPat(Pattern): - """All children must match.""" - - parts: tuple[Pattern, ...] = field(default_factory=tuple) - - def match(self, subject: Any) -> bool: - return all(p.match(subject) for p in self.parts) - - -@dataclass(frozen=True) -class DimVarRangePat(Pattern): - """Match ``lo <= value <= hi`` for a named specialization dimension. - - ``dim_var`` identifies the runtime shape source but is not inspected by - :meth:`match`, which receives only the scalar value. - - See [core-ir §3.1](docs/spec/core-ir.md#31-dimvarrangepat). - """ - - dim_var: str = "" - lo: int = 0 - hi: int = 0 - - def __post_init__(self) -> None: - if not isinstance(self.dim_var, str) or not self.dim_var: - raise ValueError( - f"DimVarRangePat: dim_var must be a non-empty str, got {self.dim_var!r}" - ) - if not isinstance(self.lo, int) or isinstance(self.lo, bool): - raise TypeError(f"DimVarRangePat: lo must be int, got {type(self.lo).__name__}") - if not isinstance(self.hi, int) or isinstance(self.hi, bool): - raise TypeError(f"DimVarRangePat: hi must be int, got {type(self.hi).__name__}") - if self.lo > self.hi: - raise ValueError( - f"DimVarRangePat({self.dim_var!r}, {self.lo}, {self.hi}): " - f"requires lo <= hi (closed [lo, hi])" - ) - - def match(self, subject: Any) -> bool: - if isinstance(subject, bool) or not isinstance(subject, int): - return False - return self.lo <= subject <= self.hi - - -def locate_dim_var(params: tuple, name: str) -> tuple[int, int] | None: - """First ``(param_index, axis)`` where a ``DimVar`` named *name* appears in *params*' shapes. - - First ``(param_index, axis)`` where a ``DimVar`` named *name* appears in - *params*' shapes. - - Canonical scan order is ``(param_index ascending, axis ascending)`` — the - single dispatch-subject rule shared by HIR→TIR lowering and the reference - evaluator's variant selection. - """ - for i, p in enumerate(params): - shape = getattr(p.type, "shape", None) - if shape is None: - continue - for axis, dim in enumerate(shape): - if getattr(dim, "name", None) == name: - return (i, axis) - return None - - -def _mangle_variant_name(name: str, specializations: tuple[Pattern, ...]) -> str: - if len(specializations) != 1 or not isinstance(specializations[0], DimVarRangePat): - raise TypeError("variant requires exactly one DimVarRangePat") - pat = specializations[0] - return f"{name}${pat.dim_var}${pat.lo}_{pat.hi}" - - -Scalar: ScalarPat = ScalarPat() - - -Tensor: TensorPat = TensorPat() - - -__all__ = [ - "Pattern", - "ScalarPat", - "TensorPat", - "AndPat", - "DimVarRangePat", - "Scalar", - "Tensor", - "locate_dim_var", - "_mangle_variant_name", -] diff --git a/src/tilefoundry/ir/hir/function.py b/src/tilefoundry/ir/hir/function.py index c0560817..2700abf6 100644 --- a/src/tilefoundry/ir/hir/function.py +++ b/src/tilefoundry/ir/hir/function.py @@ -4,7 +4,7 @@ from tilefoundry import evaluator from tilefoundry.ir.core import Expr, Var -from tilefoundry.ir.core.pattern import Pattern +from tilefoundry.ir.pattern import Pattern from tilefoundry.ir.types import Type, callable_type_for from tilefoundry.ir.types.substitute import canonicalize_dims @@ -34,9 +34,7 @@ class Function(Expr): _specialized_dims: tuple[tuple[str, int], ...] | None = field( default=None, compare=False, hash=False, repr=False ) - _display_name: str | None = field( - default=None, compare=False, hash=False, repr=False - ) + _display_name: str | None = field(default=None, compare=False, hash=False, repr=False) @classmethod def build( diff --git a/src/tilefoundry/ir/hir/math/binary.py b/src/tilefoundry/ir/hir/math/binary.py index 2e840773..ee185f98 100644 --- a/src/tilefoundry/ir/hir/math/binary.py +++ b/src/tilefoundry/ir/hir/math/binary.py @@ -16,10 +16,10 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._helpers import broadcast_shapes, resolve_anchor_storage from tilefoundry.ir.hir._shard_checks import check_multilinear_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, Layout, TensorType from tilefoundry.ir.types.shard_layout import ( Broadcast, diff --git a/src/tilefoundry/ir/hir/math/clamp.py b/src/tilefoundry/ir/hir/math/clamp.py index 74e756eb..7ecf19d6 100644 --- a/src/tilefoundry/ir/hir/math/clamp.py +++ b/src/tilefoundry/ir/hir/math/clamp.py @@ -11,9 +11,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -33,8 +33,6 @@ class Clamp(Op): max_val = ParamDef(kind="attribute", annotation=float) - - @register_typeinfer(Clamp) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: x_ty = ctx.type_of(call.args[0]) diff --git a/src/tilefoundry/ir/hir/math/softplus.py b/src/tilefoundry/ir/hir/math/softplus.py index 7ddd9571..eb7bb7dc 100644 --- a/src/tilefoundry/ir/hir/math/softplus.py +++ b/src/tilefoundry/ir/hir/math/softplus.py @@ -6,9 +6,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -26,8 +26,6 @@ class Softplus(Op): x = ParamDef(kind="input", pattern=Tensor) - - @register_typeinfer(Softplus) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: x_ty = ctx.type_of(call.args[0]) diff --git a/src/tilefoundry/ir/hir/math/unary.py b/src/tilefoundry/ir/hir/math/unary.py index 9a70d7af..9c07eddc 100644 --- a/src/tilefoundry/ir/hir/math/unary.py +++ b/src/tilefoundry/ir/hir/math/unary.py @@ -13,9 +13,9 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.kinds import UnaryKind from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( diff --git a/src/tilefoundry/ir/hir/nn/conv2d.py b/src/tilefoundry/ir/hir/nn/conv2d.py index 3b0939c8..9f2f4887 100644 --- a/src/tilefoundry/ir/hir/nn/conv2d.py +++ b/src/tilefoundry/ir/hir/nn/conv2d.py @@ -8,10 +8,10 @@ from tilefoundry.ir.core import Expr, Op from tilefoundry.ir.core.expr import Call, Constant from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import check_multilinear_partials from tilefoundry.ir.isl_interop import normalize_dim +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import Layout, TensorType from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimSub, simplify_dim from tilefoundry.ir.types.shard_layout import Split, shard_layout_of, split_target_axes diff --git a/src/tilefoundry/ir/hir/nn/gelu.py b/src/tilefoundry/ir/hir/nn/gelu.py index a25c78ff..7ef48b1e 100644 --- a/src/tilefoundry/ir/hir/nn/gelu.py +++ b/src/tilefoundry/ir/hir/nn/gelu.py @@ -6,9 +6,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -31,8 +31,6 @@ class Gelu(Op): approximate = ParamDef(kind="attribute", annotation=str, default="tanh") - - @register_typeinfer(Gelu) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: x_ty = ctx.type_of(call.args[0]) diff --git a/src/tilefoundry/ir/hir/nn/layer_norm.py b/src/tilefoundry/ir/hir/nn/layer_norm.py index 202d0d7a..a681fabc 100644 --- a/src/tilefoundry/ir/hir/nn/layer_norm.py +++ b/src/tilefoundry/ir/hir/nn/layer_norm.py @@ -7,9 +7,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.shard_layout import ShardLayout, split_target_axes from tilefoundry.visitor_registry import register_typeinfer @@ -44,9 +44,7 @@ def _normalized_axis(call: "Call", ctx: "TypeInferContext", rank: int) -> int: def _reject_normalized_splits(call, ctx, name, type_, first_axis: int) -> None: if not isinstance(type_.layout, ShardLayout): return - for mesh_axis, logical_axis in enumerate( - split_target_axes(type_.layout, type_.shape) - ): + for mesh_axis, logical_axis in enumerate(split_target_axes(type_.layout, type_.shape)): if logical_axis is not None and logical_axis >= first_axis: ctx.error( call, @@ -123,9 +121,9 @@ def _layer_norm_access(call: "Call", ctx) -> AccessRelations: where = f" : {' and '.join(guards)}" if guards else "" row = AffineAccess(isl.map(f"{{ [{domain}] -> [{', '.join(names)}]{where} }}")) belongs = logical_axes_of(x, x) - suffix = ", ".join( - names[position] for position, owner in enumerate(belongs) if owner >= axis - ) or "0" + suffix = ( + ", ".join(names[position] for position, owner in enumerate(belongs) if owner >= axis) or "0" + ) across = AffineAccess(isl.map(f"{{ [{domain}] -> [{suffix}]{where} }}")) return iterating( rows, diff --git a/src/tilefoundry/ir/hir/nn/matmul.py b/src/tilefoundry/ir/hir/nn/matmul.py index 289e26ec..d1de46f1 100644 --- a/src/tilefoundry/ir/hir/nn/matmul.py +++ b/src/tilefoundry/ir/hir/nn/matmul.py @@ -9,10 +9,10 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._helpers import broadcast_shapes, is_one, resolve_anchor_storage from tilefoundry.ir.hir._shard_checks import check_multilinear_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.shard_layout import shard_layout_of, split_target_axes from tilefoundry.visitor_registry import register_typeinfer diff --git a/src/tilefoundry/ir/hir/nn/relu.py b/src/tilefoundry/ir/hir/nn/relu.py index 5fb7f620..b6685518 100644 --- a/src/tilefoundry/ir/hir/nn/relu.py +++ b/src/tilefoundry/ir/hir/nn/relu.py @@ -6,9 +6,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -24,8 +24,6 @@ class ReLU(Op): x = ParamDef(kind="input", pattern=Tensor) - - @register_typeinfer(ReLU) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: x_ty = ctx.type_of(call.args[0]) diff --git a/src/tilefoundry/ir/hir/nn/rms_norm.py b/src/tilefoundry/ir/hir/nn/rms_norm.py index 3474ab5a..eb5b2b8c 100644 --- a/src/tilefoundry/ir/hir/nn/rms_norm.py +++ b/src/tilefoundry/ir/hir/nn/rms_norm.py @@ -14,9 +14,9 @@ from tilefoundry.evaluator.value import TensorValue, to_torch_dtype from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( diff --git a/src/tilefoundry/ir/hir/nn/rope.py b/src/tilefoundry/ir/hir/nn/rope.py index 484f4e09..87843009 100644 --- a/src/tilefoundry/ir/hir/nn/rope.py +++ b/src/tilefoundry/ir/hir/nn/rope.py @@ -18,10 +18,10 @@ from tilefoundry.evaluator.value import TensorValue, TupleValue, to_torch_dtype from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import check_multilinear_partials, reject_partials from tilefoundry.ir.isl_interop import index_set +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TupleType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -113,13 +113,15 @@ def _rope_access_relation(call: "Call", ctx: "TypeInferContext") -> AccessRelati logical_table = ctx.type_of(call.args[operand]) rows = len(logical_table.shape) - 1 tables.append( - BoundaryRelation(reached_at( + BoundaryRelation( + reached_at( rank + 1, table, logical_table, {rows: carried.get(head_dim, "0")}, free=tuple(range(rows)), - )) + ) + ) ) return iterating( (*q_ty.shape, 2), @@ -128,13 +130,15 @@ def _rope_access_relation(call: "Call", ctx: "TypeInferContext") -> AccessRelati BoundaryRelation(value), BoundaryRelation(grouped), *tables, - BoundaryRelation(reached_at( + BoundaryRelation( + reached_at( rank + 1, positions, ctx.type_of(call.args[4]), {}, free=tuple(range(len(ctx.type_of(call.args[4]).shape))), - )), + ) + ), ), outputs=( BoundaryRelation(value), diff --git a/src/tilefoundry/ir/hir/nn/sigmoid.py b/src/tilefoundry/ir/hir/nn/sigmoid.py index 70c575d9..471df553 100644 --- a/src/tilefoundry/ir/hir/nn/sigmoid.py +++ b/src/tilefoundry/ir/hir/nn/sigmoid.py @@ -6,9 +6,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -24,8 +24,6 @@ class Sigmoid(Op): x = ParamDef(kind="input", pattern=Tensor) - - @register_typeinfer(Sigmoid) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: x_ty = ctx.type_of(call.args[0]) diff --git a/src/tilefoundry/ir/hir/nn/silu.py b/src/tilefoundry/ir/hir/nn/silu.py index 828b130c..1f60104f 100644 --- a/src/tilefoundry/ir/hir/nn/silu.py +++ b/src/tilefoundry/ir/hir/nn/silu.py @@ -8,9 +8,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -26,8 +26,6 @@ class Silu(Op): x = ParamDef(kind="input", pattern=Tensor) - - @register_typeinfer(Silu) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: x_ty = ctx.type_of(call.args[0]) diff --git a/src/tilefoundry/ir/hir/nn/softmax.py b/src/tilefoundry/ir/hir/nn/softmax.py index 785a9947..05affe5c 100644 --- a/src/tilefoundry/ir/hir/nn/softmax.py +++ b/src/tilefoundry/ir/hir/nn/softmax.py @@ -7,9 +7,9 @@ from tilefoundry.evaluator.value import TensorValue, to_torch_dtype from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( diff --git a/src/tilefoundry/ir/hir/nn/tanh.py b/src/tilefoundry/ir/hir/nn/tanh.py index 137d422b..565515d6 100644 --- a/src/tilefoundry/ir/hir/nn/tanh.py +++ b/src/tilefoundry/ir/hir/nn/tanh.py @@ -6,9 +6,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -24,8 +24,6 @@ class Tanh(Op): x = ParamDef(kind="input", pattern=Tensor) - - @register_typeinfer(Tanh) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: x_ty = ctx.type_of(call.args[0]) diff --git a/src/tilefoundry/ir/hir/sharding/local.py b/src/tilefoundry/ir/hir/sharding/local.py index 67dd8449..6ff8d056 100644 --- a/src/tilefoundry/ir/hir/sharding/local.py +++ b/src/tilefoundry/ir/hir/sharding/local.py @@ -4,8 +4,8 @@ from tilefoundry.evaluator.value import EvalError from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.shard_layout import ShardLayout, Split diff --git a/src/tilefoundry/ir/hir/sharding/mesh_coord.py b/src/tilefoundry/ir/hir/sharding/mesh_coord.py index ea571240..600c8dc1 100644 --- a/src/tilefoundry/ir/hir/sharding/mesh_coord.py +++ b/src/tilefoundry/ir/hir/sharding/mesh_coord.py @@ -6,9 +6,9 @@ from tilefoundry.evaluator.value import EvalError from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Scalar from tilefoundry.ir.core.register import register_op from tilefoundry.ir.mesh_scope import covered_by_scope +from tilefoundry.ir.pattern import Scalar from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.mesh import Mesh diff --git a/src/tilefoundry/ir/hir/sharding/reshard.py b/src/tilefoundry/ir/hir/sharding/reshard.py index 73dd473b..37988787 100644 --- a/src/tilefoundry/ir/hir/sharding/reshard.py +++ b/src/tilefoundry/ir/hir/sharding/reshard.py @@ -4,8 +4,8 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import DimMul, simplify_dim from tilefoundry.ir.types.layout import Layout @@ -166,8 +166,6 @@ class Reshard(Op): storage = ParamDef(kind="attribute", default=None) - - @register_access_relation(Reshard) def _reshard_access(call: "Call", ctx) -> AccessRelations: """Every logical index reads itself. Where those bytes go is a separate fact. diff --git a/src/tilefoundry/ir/hir/specialize.py b/src/tilefoundry/ir/hir/specialize.py index 2b81ec1c..76279752 100644 --- a/src/tilefoundry/ir/hir/specialize.py +++ b/src/tilefoundry/ir/hir/specialize.py @@ -14,9 +14,9 @@ from collections.abc import Mapping from tilefoundry.ir.core import Call, Constant, Expr, Op, Tuple, Var -from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion +from tilefoundry.ir.pattern import Pattern, RangePattern from tilefoundry.ir.types.dim import is_dim_expr from tilefoundry.ir.types.mesh import make_mesh from tilefoundry.ir.types.substitute import ( @@ -40,7 +40,7 @@ def canonical_specialization_signature( """Deterministic identity string for a Function's specialization tuple.""" parts: list[str] = [] for pat in specializations: - if isinstance(pat, DimVarRangePat): + if isinstance(pat, RangePattern): parts.append(f"{pat.dim_var}${pat.lo}_{pat.hi}") else: parts.append(repr(pat)) @@ -63,9 +63,7 @@ def display_name(fn: Function) -> str | None: return getattr(fn, DISPLAY_NAME, None) -def _record_provenance( - derived: Function, origin: Function, dims: Mapping[str, int] | None -) -> None: +def _record_provenance(derived: Function, origin: Function, dims: Mapping[str, int] | None) -> None: """Note that *derived* is *origin*, at *dims* when a size was chosen. These fields are declared on Function with ``compare=False`` because they @@ -81,9 +79,7 @@ def _record_provenance( derived._specialized_dims = tuple(sorted(dims.items())) -def _record_complete_bindings( - function: Function, dims: Mapping[str, int] -) -> Function: +def _record_complete_bindings(function: Function, dims: Mapping[str, int]) -> Function: """Record a public call's complete program bindings on a derived Function.""" if bound_dims_of(function) is None: derived = dataclasses.replace(function) @@ -137,7 +133,7 @@ def _covers(fn: Function, variant: Function, dims: Mapping[str, int]) -> bool: the caller does not yet know which implementation they are asking for. """ for pattern in variant.specializations: - if not isinstance(pattern, DimVarRangePat): + if not isinstance(pattern, RangePattern): continue if pattern.dim_var not in dims: raise SpecializationError( @@ -154,7 +150,7 @@ def _coverage(fn: Function) -> str: ", ".join( f"{pattern.dim_var} in [{pattern.lo}, {pattern.hi}]" for pattern in variant.specializations - if isinstance(pattern, DimVarRangePat) + if isinstance(pattern, RangePattern) ) or "everything" for variant in fn.variants @@ -185,7 +181,7 @@ def specialize_function( present = set(residual_dims(chosen)) for pattern in chosen.specializations: - if isinstance(pattern, DimVarRangePat): + if isinstance(pattern, RangePattern): present.add(pattern.dim_var) unknown = sorted(set(dims) - present) if unknown: @@ -256,14 +252,9 @@ def visit_Call(self, call: Call, ctx: InstantiateContext) -> Expr: new_args = tuple(self.visit(arg, ctx) for arg in call.args) new_target = call.target if isinstance(new_target, Function): - new_target = _specialize_callee( - new_target, ctx.dims, ctx.type_ctx - ) + new_target = _specialize_callee(new_target, ctx.dims, ctx.type_ctx) new_target = _substitute_op_dims(new_target, ctx.dims) - if ( - all(new is old for new, old in zip(new_args, call.args)) - and new_target is call.target - ): + if all(new is old for new, old in zip(new_args, call.args)) and new_target is call.target: return call rebuilt = dataclasses.replace(call, args=new_args, target=new_target) return self._retyped(rebuilt, ctx) @@ -272,9 +263,7 @@ def visit_LoopRegion(self, region: LoopRegion, ctx: InstantiateContext) -> Expr: """Rebuild loop bindings and shape fields excluded by generic cloning.""" new_inits = tuple(self.visit(arg, ctx) for arg in region.init_args) new_phis = tuple( - old_phi - if new_init.type == old_phi.type - else Var(type=new_init.type, name=old_phi.name) + old_phi if new_init.type == old_phi.type else Var(type=new_init.type, name=old_phi.name) for old_phi, new_init in zip(region.carried_args, new_inits) ) for old_phi, new_phi in zip(region.carried_args, new_phis): @@ -309,9 +298,7 @@ def default_visit(self, expr: Expr, ctx: InstantiateContext) -> Expr: return rebuilt if rebuilt is expr else self._retyped(rebuilt, ctx) def _retyped(self, rebuilt: Expr, ctx: InstantiateContext) -> Expr: - return dataclasses.replace( - rebuilt, type=ctx.type_visitor.visit(rebuilt, ctx.type_ctx) - ) + return dataclasses.replace(rebuilt, type=ctx.type_visitor.visit(rebuilt, ctx.type_ctx)) @dataclasses.dataclass diff --git a/src/tilefoundry/ir/hir/tensor/argmax.py b/src/tilefoundry/ir/hir/tensor/argmax.py index 45ff7bd2..e4073880 100644 --- a/src/tilefoundry/ir/hir/tensor/argmax.py +++ b/src/tilefoundry/ir/hir/tensor/argmax.py @@ -14,9 +14,9 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, Layout, TensorType from tilefoundry.ir.types.shard_layout import ( Split, diff --git a/src/tilefoundry/ir/hir/tensor/cache_update.py b/src/tilefoundry/ir/hir/tensor/cache_update.py index 71c10e57..d2472ea2 100644 --- a/src/tilefoundry/ir/hir/tensor/cache_update.py +++ b/src/tilefoundry/ir/hir/tensor/cache_update.py @@ -9,9 +9,9 @@ from tilefoundry.evaluator.value import EvalError, TensorValue from tilefoundry.ir.core import Constant, Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import require_matching_partial_state +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.shard_layout import Split, shard_layout_of, split_target_axes @@ -38,8 +38,6 @@ class CacheUpdate(Op): new = ParamDef(kind="input", pattern=Tensor) - - def _limit(cache: tuple, supplied: tuple) -> int | None: """The most rows one call may write: the fewer of what each side states.""" stated = [ @@ -71,8 +69,6 @@ def _rows(expr) -> object: return expr - - def _row_limit(offsets: tuple, extents: tuple, limit: int | None) -> tuple: """The most the row window may extend, at the position holding the rows. @@ -121,7 +117,8 @@ def _cache_update_access(call: "Call", ctx) -> AccessRelations: BoundaryRelation(complement), BoundaryRelation(control_read(rank, ctx, call.args[1])), BoundaryRelation(control_read(rank, ctx, call.args[2])), - BoundaryRelation(window_source( + BoundaryRelation( + window_source( offsets, rank, logical_new, @@ -129,7 +126,8 @@ def _cache_update_access(call: "Call", ctx) -> AccessRelations: {axis: f"d{axis}" for axis in range(rank)}, (None, rows), ceilings, - )), + ) + ), ), outputs=(BoundaryRelation(reached),), ), diff --git a/src/tilefoundry/ir/hir/tensor/cast.py b/src/tilefoundry/ir/hir/tensor/cast.py index a8063dab..0dc1c81d 100644 --- a/src/tilefoundry/ir/hir/tensor/cast.py +++ b/src/tilefoundry/ir/hir/tensor/cast.py @@ -4,8 +4,8 @@ from tilefoundry.evaluator.value import TensorValue, to_torch_dtype from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.shard_layout import shard_layout_of from tilefoundry.visitor_registry import register_typeinfer diff --git a/src/tilefoundry/ir/hir/tensor/concat.py b/src/tilefoundry/ir/hir/tensor/concat.py index 1eebfcaa..e25a8f5b 100644 --- a/src/tilefoundry/ir/hir/tensor/concat.py +++ b/src/tilefoundry/ir/hir/tensor/concat.py @@ -10,7 +10,6 @@ from tilefoundry.ir.core import Expr, Op from tilefoundry.ir.core.expr import Call from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._helpers import resolve_anchor_storage from tilefoundry.ir.hir._shard_checks import ( @@ -18,6 +17,7 @@ require_uniform_partial_slices, ) from tilefoundry.ir.isl_interop import normalize_dim_entries +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import Layout, Split, TensorType from tilefoundry.ir.types.dim import DimAdd, simplify_dim from tilefoundry.ir.types.shard_layout import shard_layout_of, split_target_axes diff --git a/src/tilefoundry/ir/hir/tensor/full_like.py b/src/tilefoundry/ir/hir/tensor/full_like.py index d84786d4..490054e1 100644 --- a/src/tilefoundry/ir/hir/tensor/full_like.py +++ b/src/tilefoundry/ir/hir/tensor/full_like.py @@ -14,8 +14,8 @@ from tilefoundry.evaluator.value import TensorValue, to_torch_dtype from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( diff --git a/src/tilefoundry/ir/hir/tensor/index_add.py b/src/tilefoundry/ir/hir/tensor/index_add.py index ad7e6f0b..e4b922c4 100644 --- a/src/tilefoundry/ir/hir/tensor/index_add.py +++ b/src/tilefoundry/ir/hir/tensor/index_add.py @@ -6,10 +6,10 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials from tilefoundry.ir.hir.tensor.index_select import _norm_dim +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.shard_layout import Split, shard_layout_of from tilefoundry.visitor_registry import register_typeinfer @@ -151,14 +151,12 @@ def _index_add_access(call: "Call", ctx) -> AccessRelations: payload = reached_at(rank, src, ctx.type_of(call.args[2]), carried, free=(dim,)) return iterating( dst.shape, - AccessRelations( + AccessRelations( inputs=( BoundaryRelation(rows), BoundaryRelation(named), BoundaryRelation(payload), ), - outputs=( - BoundaryRelation(rows), - ), + outputs=(BoundaryRelation(rows),), ), ) diff --git a/src/tilefoundry/ir/hir/tensor/index_copy.py b/src/tilefoundry/ir/hir/tensor/index_copy.py index a305cea5..03dcfc14 100644 --- a/src/tilefoundry/ir/hir/tensor/index_copy.py +++ b/src/tilefoundry/ir/hir/tensor/index_copy.py @@ -6,10 +6,10 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir.tensor.index_add import _infer_index_write from tilefoundry.ir.hir.tensor.index_select import _norm_dim +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -81,14 +81,12 @@ def _index_copy_access(call: "Call", ctx) -> AccessRelations: payload = reached_at(rank, src, ctx.type_of(call.args[2]), carried, free=(dim,)) return iterating( dst.shape, - AccessRelations( + AccessRelations( inputs=( BoundaryRelation(identity), BoundaryRelation(named), BoundaryRelation(payload), ), - outputs=( - BoundaryRelation(rows), - ), + outputs=(BoundaryRelation(rows),), ), ) diff --git a/src/tilefoundry/ir/hir/tensor/index_select.py b/src/tilefoundry/ir/hir/tensor/index_select.py index ce67e657..e0fa5aa5 100644 --- a/src/tilefoundry/ir/hir/tensor/index_select.py +++ b/src/tilefoundry/ir/hir/tensor/index_select.py @@ -6,8 +6,8 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.layout import Layout from tilefoundry.ir.types.shard_layout import ( @@ -128,12 +128,14 @@ def _index_select_access_relation(call: "Call", ctx) -> AccessRelations: out_shape, AccessRelations( inputs=( - BoundaryRelation(reached_at(rank, source_ty, logical_source, carried, free=(axis,))), - BoundaryRelation(reached_at(rank, index_ty, logical_index, {0: carried.get(axis, "0")})), - ), - outputs=( - BoundaryRelation(identity_access(rank)), + BoundaryRelation( + reached_at(rank, source_ty, logical_source, carried, free=(axis,)) + ), + BoundaryRelation( + reached_at(rank, index_ty, logical_index, {0: carried.get(axis, "0")}) + ), ), + outputs=(BoundaryRelation(identity_access(rank)),), ), ) diff --git a/src/tilefoundry/ir/hir/tensor/insert_slice.py b/src/tilefoundry/ir/hir/tensor/insert_slice.py index 820723ec..c7e46d6d 100644 --- a/src/tilefoundry/ir/hir/tensor/insert_slice.py +++ b/src/tilefoundry/ir/hir/tensor/insert_slice.py @@ -6,9 +6,9 @@ from tilefoundry.evaluator.value import TensorValue, TupleValue from tilefoundry.ir.core import Constant, Op, Tuple from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Scalar, Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import require_matching_partial_state +from tilefoundry.ir.pattern import Scalar, Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.visitor_registry import register_typeinfer @@ -34,8 +34,6 @@ class InsertSlice(Op): offsets = ParamDef(kind="input", pattern=Scalar) - - def _offset_axes(call: "Call", rank: int) -> tuple: """Where the window starts on each axis, as a number or as the value it is. @@ -47,15 +45,11 @@ def _offset_axes(call: "Call", rank: int) -> tuple: given = call.args[2] if isinstance(given, Tuple): return tuple( - int(item.value) - if isinstance(item, Constant) and isinstance(item.value, int) - else item + int(item.value) if isinstance(item, Constant) and isinstance(item.value, int) else item for item in given.elements ) start = ( - int(given.value) - if isinstance(given, Constant) and isinstance(given.value, int) - else given + int(given.value) if isinstance(given, Constant) and isinstance(given.value, int) else given ) return (start, *(0 for _ in range(rank - 1))) @@ -78,23 +72,16 @@ def _insert_slice_access(call: "Call", ctx) -> AccessRelations: complement, written = placed_window( offsets, tuple(update.shape), rank, within=tuple(result.shape) ) - read_update = window_source( - offsets, rank, update, update, logical_coordinates(result, result) - ) + read_update = window_source(offsets, rank, update, update, logical_coordinates(result, result)) return iterating( result.shape, - AccessRelations( + AccessRelations( inputs=( BoundaryRelation(complement), BoundaryRelation(read_update), - *( - BoundaryRelation(control_read(rank, ctx, arg)) - for arg in call.args[2:] - ), - ), - outputs=( - BoundaryRelation(written), + *(BoundaryRelation(control_read(rank, ctx, arg)) for arg in call.args[2:]), ), + outputs=(BoundaryRelation(written),), ), ) @@ -170,8 +157,6 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: return dst_ty - - @register_eval(InsertSlice) def _eval_insert_slice(ctx): dst = ctx.args[0].data diff --git a/src/tilefoundry/ir/hir/tensor/quant.py b/src/tilefoundry/ir/hir/tensor/quant.py index abce8917..af1a81bc 100644 --- a/src/tilefoundry/ir/hir/tensor/quant.py +++ b/src/tilefoundry/ir/hir/tensor/quant.py @@ -13,9 +13,9 @@ from tilefoundry.evaluator.value import EvalError, TensorValue, TupleValue, to_torch_dtype from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, Layout, ShardLayout, TensorType, TupleType from tilefoundry.ir.types.dim import DimFloorDiv, simplify_dim from tilefoundry.ir.types.layout import flatten @@ -127,8 +127,7 @@ def _result_layouts(call, ctx, x_ty, scale_shape, group: int): except ValueError as error: ctx.error( call, - f"cannot derive result sharding: {error}; use an explicit " - "Reshard before Quant", + f"cannot derive result sharding: {error}; use an explicit Reshard before Quant", ) if x_ty.layout is None: return None, None @@ -193,9 +192,7 @@ def _eval_quant(ctx): group = ctx.op.group last = x.shape[-1] if last % group: - raise EvalError( - f"Quant: runtime last dim {last} not divisible by group={group}" - ) + raise EvalError(f"Quant: runtime last dim {last} not divisible by group={group}") grouped = x.reshape(*x.shape[:-1], last // group, group) absmax = grouped.abs().amax(dim=-1) scale = torch.where(absmax == 0, torch.ones_like(absmax), absmax / 448.0) @@ -236,9 +233,7 @@ def _quant_access_relation(call: "Call", ctx: "TypeInferContext") -> AccessRelat outer = ", ".join(f"i{k}" for k in range(rank - 1)) last = f"i{rank - 1}" out_dims = (outer + ", ") if outer else "" - scale_rel = AffineAccess( - isl.map(f"{{ [{dims}] -> [{out_dims}floor({last}/{group})] }}") - ) + scale_rel = AffineAccess(isl.map(f"{{ [{dims}] -> [{out_dims}floor({last}/{group})] }}")) return iterating( x_ty.shape, diff --git a/src/tilefoundry/ir/hir/tensor/rank.py b/src/tilefoundry/ir/hir/tensor/rank.py index e054f501..78db23b3 100644 --- a/src/tilefoundry/ir/hir/tensor/rank.py +++ b/src/tilefoundry/ir/hir/tensor/rank.py @@ -6,8 +6,8 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( diff --git a/src/tilefoundry/ir/hir/tensor/reduce.py b/src/tilefoundry/ir/hir/tensor/reduce.py index 76ff3a0b..3f711b4d 100644 --- a/src/tilefoundry/ir/hir/tensor/reduce.py +++ b/src/tilefoundry/ir/hir/tensor/reduce.py @@ -10,9 +10,9 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.kinds import ReduceKind from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import Layout, TensorType from tilefoundry.ir.types.shard_layout import canonical_shard_layout, shard_layout_of from tilefoundry.ir.types.stride import try_compact_major diff --git a/src/tilefoundry/ir/hir/tensor/repeat_interleave.py b/src/tilefoundry/ir/hir/tensor/repeat_interleave.py index 0cfe9bc0..cbd999d8 100644 --- a/src/tilefoundry/ir/hir/tensor/repeat_interleave.py +++ b/src/tilefoundry/ir/hir/tensor/repeat_interleave.py @@ -7,8 +7,8 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.shard_layout import Broadcast, ShardLayout from tilefoundry.visitor_registry import register_typeinfer @@ -100,9 +100,11 @@ def _repeat_interleave_access(call: "Call", ctx) -> AccessRelations: produced *= extent if isinstance(extent, int) else 1 return iterating( out_shape, - AccessRelations( + AccessRelations( inputs=( - BoundaryRelation(AffineAccess(isl.multi_aff(f"{{ [{domain}] -> [{', '.join(reads)}] }}"))), + BoundaryRelation( + AffineAccess(isl.multi_aff(f"{{ [{domain}] -> [{', '.join(reads)}] }}")) + ), ), outputs=(BoundaryRelation(identity_access(rank)),), ), diff --git a/src/tilefoundry/ir/hir/tensor/reshape.py b/src/tilefoundry/ir/hir/tensor/reshape.py index 53eec571..009cb442 100644 --- a/src/tilefoundry/ir/hir/tensor/reshape.py +++ b/src/tilefoundry/ir/hir/tensor/reshape.py @@ -6,8 +6,8 @@ from tilefoundry.evaluator.value import EvalError, TensorValue from tilefoundry.ir.core import Call, Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import ComposedLayout, TensorType from tilefoundry.ir.types.layout import Layout, flatten from tilefoundry.ir.types.shard_layout import ( @@ -32,8 +32,6 @@ class Reshape(Op): new_shape = ParamDef(kind="attribute", annotation=tuple) - - def _reshape_view(call: "Call", ctx) -> tuple: """Where a result coordinate sits in the source it was renamed from.""" out_shape = tuple(call.target.new_shape) diff --git a/src/tilefoundry/ir/hir/tensor/shape_of.py b/src/tilefoundry/ir/hir/tensor/shape_of.py index b46d89ca..6f876521 100644 --- a/src/tilefoundry/ir/hir/tensor/shape_of.py +++ b/src/tilefoundry/ir/hir/tensor/shape_of.py @@ -7,8 +7,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.expr import Constant from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, TensorType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( diff --git a/src/tilefoundry/ir/hir/tensor/slice.py b/src/tilefoundry/ir/hir/tensor/slice.py index 6d0106b0..47ed6d1a 100644 --- a/src/tilefoundry/ir/hir/tensor/slice.py +++ b/src/tilefoundry/ir/hir/tensor/slice.py @@ -8,9 +8,9 @@ from tilefoundry.ir.core import Expr, Op, Tuple from tilefoundry.ir.core.expr import Call, Constant from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.isl_interop import dim_range +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import ComposedLayout, Layout, ShardLayout, Swizzle, TensorType from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimMul, DimSub, simplify_dim from tilefoundry.ir.types.int_tuple import flatten @@ -41,12 +41,6 @@ def __init__(self, **attrs): super().__init__(**attrs) - - - - - - class _Unbounded(ValueError): """A relation would have a parameter nothing can bound.""" @@ -145,11 +139,7 @@ def start(self, value, axis: int) -> str: if literal is None: self.params.append((name, value)) self.guards.append(f"0 <= {name}") - reach = ( - name - if self.size == "1" - else f"{name} + ({self.size} - 1) * {self.stride}" - ) + reach = name if self.size == "1" else f"{name} + ({self.size} - 1) * {self.stride}" self.guards.append(f"{reach} <= {self.extent} - 1") return name @@ -332,9 +322,7 @@ def _slice_shard_layout(call, ctx, x_ty, source, starts, inherited_offset): narrow_positions[tensor_axis] = position new_shape[position] = op.sizes[tensor_axis] if new_strides is not None: - new_strides[position] = _dim_mul( - new_strides[position], op.strides[tensor_axis] - ) + new_strides[position] = _dim_mul(new_strides[position], op.strides[tensor_axis]) sharded = ShardLayout( layout=Layout( @@ -345,8 +333,7 @@ def _slice_shard_layout(call, ctx, x_ty, source, starts, inherited_offset): mesh=source.mesh, ) if any(start is None for start in static_starts) or not all( - isinstance(stride, int) and not isinstance(stride, bool) - for stride in op.strides + isinstance(stride, int) and not isinstance(stride, bool) for stride in op.strides ): return sharded @@ -457,9 +444,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: source = x_ty.layout inherited_offset = 0 inherited_inner = None - if isinstance(source, ComposedLayout) and isinstance( - source.outer, (Layout, ShardLayout) - ): + if isinstance(source, ComposedLayout) and isinstance(source.outer, (Layout, ShardLayout)): if isinstance(source.inner, Swizzle) and isinstance(source.outer, Layout): inherited_inner = source.inner elif source.inner is not None: @@ -482,9 +467,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: new_layout = None if isinstance(source, ShardLayout): - new_layout = _slice_shard_layout( - call, ctx, x_ty, source, starts, inherited_offset - ) + new_layout = _slice_shard_layout(call, ctx, x_ty, source, starts, inherited_offset) elif isinstance(source, Layout) and source.strides is not None: static_starts = [] steps = [] @@ -499,9 +482,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: static_starts.append(int(start.value)) steps.append(stride) else: - found = window_image( - source, tuple(static_starts), tuple(layout_shape), tuple(steps) - ) + found = window_image(source, tuple(static_starts), tuple(layout_shape), tuple(steps)) if found is not None: moved, window = found new_layout = ComposedLayout( @@ -527,8 +508,7 @@ def _eval_slice(ctx): for axis, (start, size, stride) in enumerate(zip(start_values, sizes, strides)): if start < 0 or size < 0 or stride <= 0: raise EvalError( - f"Slice: invalid window on axis {axis}: start={start}, size={size}, " - f"stride={stride}" + f"Slice: invalid window on axis {axis}: start={start}, size={size}, stride={stride}" ) last = start if size == 0 else start + (size - 1) * stride if size and last >= ctx.args[0].data.shape[axis]: diff --git a/src/tilefoundry/ir/hir/tensor/split.py b/src/tilefoundry/ir/hir/tensor/split.py index 1569064c..cc900fa8 100644 --- a/src/tilefoundry/ir/hir/tensor/split.py +++ b/src/tilefoundry/ir/hir/tensor/split.py @@ -7,8 +7,8 @@ from tilefoundry.evaluator.value import TensorValue, TupleValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import Broadcast, Layout, Partial, ShardLayout, TensorType, TupleType from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.shard_layout import Split as ShardSplit diff --git a/src/tilefoundry/ir/hir/tensor/stack.py b/src/tilefoundry/ir/hir/tensor/stack.py index 067f1ae3..615ba20b 100644 --- a/src/tilefoundry/ir/hir/tensor/stack.py +++ b/src/tilefoundry/ir/hir/tensor/stack.py @@ -9,13 +9,13 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._helpers import resolve_anchor_storage from tilefoundry.ir.hir._shard_checks import ( reject_dynamic_shards, require_uniform_partial_slices, ) +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import Layout, TensorType from tilefoundry.ir.types.shard_layout import shard_layout_of from tilefoundry.ir.types.stride import try_compact_major diff --git a/src/tilefoundry/ir/hir/tensor/topk.py b/src/tilefoundry/ir/hir/tensor/topk.py index aa63cd2f..bf9a2d35 100644 --- a/src/tilefoundry/ir/hir/tensor/topk.py +++ b/src/tilefoundry/ir/hir/tensor/topk.py @@ -16,9 +16,9 @@ from tilefoundry.evaluator.value import TensorValue, TupleValue, to_torch_dtype from tilefoundry.ir.core import Call, Constant, Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, Layout, TensorType, TupleType from tilefoundry.ir.types.dim import ( DimAdd, diff --git a/src/tilefoundry/ir/hir/tensor/transpose.py b/src/tilefoundry/ir/hir/tensor/transpose.py index 671f5099..a1249b6c 100644 --- a/src/tilefoundry/ir/hir/tensor/transpose.py +++ b/src/tilefoundry/ir/hir/tensor/transpose.py @@ -7,8 +7,8 @@ from tilefoundry.evaluator.value import TensorValue from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import ComposedLayout, Layout, TensorType from tilefoundry.ir.types.shard_layout import shard_layout_of from tilefoundry.ir.types.stride import try_compact_major diff --git a/src/tilefoundry/ir/hir/tensor/tuple_get_item.py b/src/tilefoundry/ir/hir/tensor/tuple_get_item.py index a2b4b456..33f6a641 100644 --- a/src/tilefoundry/ir/hir/tensor/tuple_get_item.py +++ b/src/tilefoundry/ir/hir/tensor/tuple_get_item.py @@ -3,8 +3,8 @@ from tilefoundry.evaluator.registry import register_eval from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TupleType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -25,13 +25,7 @@ class TupleGetItem(Op): index = ParamDef(kind="attribute", annotation=int) - - - - -register_access_relation(TupleGetItem)( - view_relations(0, field=lambda call, ctx: call.target.index) -) +register_access_relation(TupleGetItem)(view_relations(0, field=lambda call, ctx: call.target.index)) @register_typeinfer(TupleGetItem) diff --git a/src/tilefoundry/ir/hir/verify.py b/src/tilefoundry/ir/hir/verify.py index 685b596a..67f7a451 100644 --- a/src/tilefoundry/ir/hir/verify.py +++ b/src/tilefoundry/ir/hir/verify.py @@ -2,8 +2,8 @@ from tilefoundry.ir.core import Expr, VerifyError from tilefoundry.ir.core.expr import Call, Var -from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.hir.mesh_region import MeshRegion +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.tir.stmt import Stmt from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import DimVar @@ -17,9 +17,7 @@ def _verify_isolated(region: MeshRegion, ctx=None) -> None: """Ensure a region body reaches captured values only through its params.""" if len(region.params) != len(region.args): - message = ( - f"region has {len(region.params)} params but {len(region.args)} args" - ) + message = f"region has {len(region.params)} params but {len(region.args)} args" if ctx is not None: ctx.error(region, message) raise VerifyError(f"MeshRegion: {message}") @@ -28,8 +26,7 @@ def _verify_isolated(region: MeshRegion, ctx=None) -> None: if not leaked: return message = ( - "region is not isolated: body reads an args value directly; " - "reference its param instead" + "region is not isolated: body reads an args value directly; reference its param instead" ) if ctx is not None: ctx.error(region, message) @@ -59,7 +56,6 @@ def verify_function(fn: Function, *, module=None) -> None: _reject_stmt_nodes(fn.body) - def _verify_variants(base: Function, *, module=None) -> None: """Verify a dispatch prototype's variants and their envelope partition.""" base_param_types = tuple(p.type for p in base.params) @@ -106,9 +102,9 @@ def _verify_partition(base: Function) -> None: ranges: list[tuple[int, int]] = [] for v in base.variants: for pat in v.specializations: - if not isinstance(pat, DimVarRangePat): + if not isinstance(pat, RangePattern): raise VerifyError( - f"hir Function {base.name!r}: only DimVarRangePat is " + f"hir Function {base.name!r}: only RangePattern is " f"supported for dispatch (got {type(pat).__name__})" ) dim_vars.add(pat.dim_var) @@ -219,19 +215,19 @@ def _verify_signature_dim_vars(fn: Function) -> None: _check_signature_dim_var_consistency(fn) param_bounds = _collect_param_dim_vars(fn) for pat in fn.specializations: - if not isinstance(pat, DimVarRangePat): + if not isinstance(pat, RangePattern): continue dv_bounds = param_bounds.get(pat.dim_var) if dv_bounds is None: raise VerifyError( - f"specialization DimVarRangePat({pat.dim_var!r}, {pat.lo}, " + f"specialization RangePattern({pat.dim_var!r}, {pat.lo}, " f"{pat.hi}) references unknown DimVar (specializations must " f"anchor to a DimVar reachable from an input parameter)" ) lo, hi = dv_bounds if not (lo <= pat.lo and pat.hi <= hi): raise VerifyError( - f"DimVarRangePat ({pat.dim_var!r}, {pat.lo}, {pat.hi}) is not " + f"RangePattern ({pat.dim_var!r}, {pat.lo}, {pat.hi}) is not " f"contained in DimVar envelope [{lo}, {hi})" ) diff --git a/src/tilefoundry/ir/pattern/__init__.py b/src/tilefoundry/ir/pattern/__init__.py new file mode 100644 index 00000000..b59e2694 --- /dev/null +++ b/src/tilefoundry/ir/pattern/__init__.py @@ -0,0 +1,103 @@ +"""Public operation-declaration pattern language.""" + +from .constraint import ( + Constraint, + DistinctConstraint, + SameConstraint, + SameModesConstraint, + affine_part, +) +from .match import ( + ABSENT, + ARRANGEMENT, + OPAQUE, + UNNAMED_PLACE, + Match, + affine_frame, + alternatives_of, + between_rules, + evaluated, + fits, + grouping, + is_symbolic, + matched, + refusals_between, + relations_of, + resolved, +) +from .pattern import ( + AndPattern, + AttrPattern, + BitsPattern, + CapturePattern, + ComposedLayoutPattern, + ConstraintPattern, + GuardPattern, + LayoutPattern, + MeshPattern, + MultipleOfPattern, + OneOfPattern, + OrPattern, + Pattern, + RangePattern, + Scalar, + ScalarPattern, + SequencePattern, + ShardLayoutPattern, + SwitchPattern, + SwizzlePattern, + Tensor, + TensorPattern, + WildcardPattern, +) +from .utils import _mangle_variant_name, arrangement_pattern, locate_dim_var + +__all__ = [ + "ABSENT", + "ARRANGEMENT", + "AndPattern", + "AttrPattern", + "BitsPattern", + "CapturePattern", + "ComposedLayoutPattern", + "Constraint", + "ConstraintPattern", + "DistinctConstraint", + "GuardPattern", + "LayoutPattern", + "Match", + "MeshPattern", + "MultipleOfPattern", + "OPAQUE", + "OneOfPattern", + "OrPattern", + "Pattern", + "RangePattern", + "SameConstraint", + "SameModesConstraint", + "Scalar", + "ScalarPattern", + "SequencePattern", + "ShardLayoutPattern", + "SwizzlePattern", + "SwitchPattern", + "Tensor", + "TensorPattern", + "UNNAMED_PLACE", + "WildcardPattern", + "_mangle_variant_name", + "affine_frame", + "affine_part", + "alternatives_of", + "arrangement_pattern", + "between_rules", + "evaluated", + "fits", + "grouping", + "is_symbolic", + "locate_dim_var", + "matched", + "refusals_between", + "relations_of", + "resolved", +] diff --git a/src/tilefoundry/ir/pattern/constraint.py b/src/tilefoundry/ir/pattern/constraint.py new file mode 100644 index 00000000..1d41797b --- /dev/null +++ b/src/tilefoundry/ir/pattern/constraint.py @@ -0,0 +1,171 @@ +"""Relations between operands of one operation declaration.""" + +from __future__ import annotations + +from dataclasses import dataclass +from math import prod + +from tilefoundry.ir.types import Broadcast, ComposedLayout, Layout, ShardLayout, Swizzle +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.tensor_type import TensorType + +from .match import _named + + +@dataclass(frozen=True) +class Constraint: + """One relation between two named operands.""" + + field: str + left: str + right: str + + def pair(self, operands: dict): + held = tuple(operands.get(name) for name in (self.left, self.right)) + if any(value is None for value in held): + return None + found = tuple(getattr(value, self.field, None) for value in held) + return None if any(value is None for value in found) else found + + def holds(self, operands: dict) -> bool: + raise NotImplementedError + + def written(self) -> str: + raise NotImplementedError + + def refused(self, operands: dict) -> str: + pair = self.pair(operands) + if pair is None: + return self.written() + left, right = pair + return ( + f"{self.left}.{self.field}={_named(left)} " + f"{self.right}.{self.field}={_named(right)}, reads {self.written()}" + ) + + +@dataclass(frozen=True) +class DistinctConstraint(Constraint): + def holds(self, operands: dict) -> bool: + pair = self.pair(operands) + return pair is None or pair[0] != pair[1] + + def written(self) -> str: + return f"{self.left}.{self.field} != {self.right}.{self.field}" + + +@dataclass(frozen=True) +class SameConstraint(Constraint): + def holds(self, operands: dict) -> bool: + pair = self.pair(operands) + return pair is None or pair[0] == pair[1] + + def written(self) -> str: + return f"{self.left}.{self.field} = {self.right}.{self.field}" + + +def affine_part(layout): + """Return the strided affine part beneath shard frames and swizzles.""" + if isinstance(layout, ShardLayout): + if not all(isinstance(attr, Broadcast) for attr in layout.attrs): + return None + layout = layout.layout + if isinstance(layout, ComposedLayout): + if layout.inner is not None and not isinstance(layout.inner, Swizzle): + return None + layout = layout.outer + return layout if isinstance(layout, Layout) and layout.strides is not None else None + + +@dataclass(frozen=True, init=False) +class SameModesConstraint(Constraint): + """Require two tensor layouts to run along the same final mode of one axis.""" + + def __init__(self, left: str, right: str): + object.__setattr__(self, "field", "layout") + object.__setattr__(self, "left", left) + object.__setattr__(self, "right", right) + + def pair(self, operands: dict): + held = tuple(operands.get(name) for name in (self.left, self.right)) + return None if any(not isinstance(value, TensorType) for value in held) else held + + @staticmethod + def reading(tensor: TensorType): + layout = affine_part(tensor.layout) + if layout is None: + return None, f"{tensor.layout!r} is no strided arrangement" + shape, strides = tuple(layout.shape), tuple(layout.strides) + extents = tuple(tensor.shape) + if len(shape) != len(extents) or any( + prod(flatten(group)) != extent for group, extent in zip(shape, extents) + ): + return None, ( + f"{layout!r} is not one group of modes per axis of the {extents} tile it arranges" + ) + unit = [] + for axis, (group, steps) in enumerate(zip(shape, strides)): + walked = tuple( + step for extent, step in zip(flatten(group), flatten(steps)) if extent > 1 + ) + unit += [ + (axis, position, position == len(walked) - 1) + for position, step in enumerate(walked) + if step == 1 + ] + if len(unit) != 1: + return None, f"{layout!r} has {len(unit)} modes at step 1, not one" + return unit[0], None + + def holds(self, operands: dict) -> bool: + pair = self.pair(operands) + if pair is None: + return True + (left, _), (right, _) = (self.reading(value) for value in pair) + return ( + left is not None and right is not None and left[0] == right[0] and left[2] and right[2] + ) + + def written(self) -> str: + return ( + f"{self.left} and {self.right} step 1 along the last mode of one tile " + "axis, each grouped by tile axis" + ) + + def refused(self, operands: dict) -> str: + pair = self.pair(operands) + if pair is None: + return self.written() + readings = tuple( + (name, *self.reading(value)) for name, value in zip((self.left, self.right), pair) + ) + for name, _, why in readings: + if why is not None: + return f"{name}.layout: {why}; reads {self.written()}" + ( + (left_name, (axis, _, left_fastest), _), + ( + right_name, + (other, _, right_fastest), + _, + ), + ) = readings + if axis != other: + return ( + f"{left_name} steps 1 along tile axis {axis} and {right_name} along " + f"tile axis {other}: that is a transpose; reads {self.written()}" + ) + name = left_name if not left_fastest else right_name + return ( + f"{name} does not step 1 along the last mode of tile axis {axis}; " + f"reads {self.written()}" + ) + + +__all__ = [ + "Constraint", + "DistinctConstraint", + "SameConstraint", + "SameModesConstraint", + "affine_part", +] diff --git a/src/tilefoundry/ir/pattern/match.py b/src/tilefoundry/ir/pattern/match.py new file mode 100644 index 00000000..27c18f83 --- /dev/null +++ b/src/tilefoundry/ir/pattern/match.py @@ -0,0 +1,218 @@ +"""Shared matching, resolution, and rendering helpers for IR patterns.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from tilefoundry.ir.clause.layout import is_layout_wildcard +from tilefoundry.ir.types import ComposedLayout +from tilefoundry.ir.types.dim import DimFloorDiv, DimMul, DimVar, is_dim_op_call +from tilefoundry.ir.types.substitute import DimSubstitutionError, substitute_shape_dim + +UNNAMED_PLACE = "_" +ARRANGEMENT = "every arrangement" +ABSENT = type("Absent", (), {"__repr__": lambda self: "ABSENT"})() +OPAQUE = object() + + +@dataclass(frozen=True) +class Match: + """The bindings produced by a successful match.""" + + captures: dict = field(default_factory=dict) + + +def _pattern_type(): + from .pattern import Pattern # noqa: PLC0415 - pattern protocol cycle + + return Pattern + + +def _named(value): + """One enumerated field as an author writes it, or None when unstated.""" + return None if value is None else getattr(value, "name", str(value)).lower() + + +def _extents(bindings) -> dict: + return {name: value for name, value in dict(bindings or {}).items() if type(value) is int} + + +def resolved(value, bindings): + """Resolve symbols and nested patterns under *bindings*.""" + if isinstance(value, _pattern_type()): + return value.resolve(bindings) + if isinstance(value, DimVar) or is_dim_op_call(value): + return substitute_shape_dim(value, _extents(bindings)) + if isinstance(value, tuple): + return tuple(resolved(item, bindings) for item in value) + return value + + +def evaluated(value, captures): + """Evaluate one symbolic dimension, or return None while it is unresolved.""" + try: + held = substitute_shape_dim(value, _extents(captures)) + except DimSubstitutionError: + return None + return held if type(held) is int else None + + +def is_symbolic(value) -> bool: + return isinstance(value, DimVar) or is_dim_op_call(value) + + +def written_dim(value) -> str: + if isinstance(value, DimVar): + return value.name + if is_dim_op_call(value): + left, right = (written_dim(arg) for arg in value.args) + if isinstance(value.target, DimFloorDiv): + return f"{left}/{right}" + if isinstance(value.target, DimMul): + return f"{left}*{right}" + return f"({left} {type(value.target).__name__} {right})" + return str(getattr(value, "value", value)) + + +def written_binding(value) -> str: + return getattr(value, "name", str(value)) + + +def written_bindings(bindings) -> str: + return ", ".join(f"{name}={written_binding(value)}" for name, value in bindings) + + +def written_tuple(items) -> str: + written = ", ".join(items) + return f"({written},)" if len(items) == 1 else f"({written})" + + +def written_grouped(modes) -> str: + if isinstance(modes, tuple): + return written_tuple(tuple(written_grouped(mode) for mode in modes)) + return written_place(modes) + + +def written_place(value, name: str = UNNAMED_PLACE) -> str: + if isinstance(value, _pattern_type()): + return value.describe(name) + if is_symbolic(value): + return written_dim(value) + return UNNAMED_PLACE if is_layout_wildcard(value) else str(value) + + +def written_field(value) -> str | None: + if isinstance(value, _pattern_type()): + return value.describe() + return _named(value) + + +def written_alternatives(items, name: str = UNNAMED_PLACE) -> str: + held = tuple( + (written_bindings(bindings), written_place(alternative, name)) + for bindings, alternative in items + ) + width = max((len(label) for label, _ in held), default=0) + lines = [] + for label, written in held: + first, *rest = written.splitlines() or ("",) + lines.append(first if not width else f"{label.ljust(width)} {first}") + lines.extend(line if not width else f"{' ' * (width + 2)}{line}" for line in rest) + return "\n".join(lines) + + +def alternatives_of(pattern, bindings=()) -> tuple: + if isinstance(pattern, _pattern_type()): + return pattern.alternatives(bindings) + return ((tuple(bindings), pattern),) + + +def matched(pattern, subject, captures=None) -> Match | None: + """Match a nested pattern, symbolic dimension, wildcard, or fixed value.""" + held = Match(dict(captures or {})) + if isinstance(pattern, DimVar): + if pattern.name in held.captures: + return held if held.captures[pattern.name] == subject else None + if type(subject) is not int or not pattern.lo <= subject < pattern.hi: + return None + return Match({**held.captures, pattern.name: subject}) + if is_dim_op_call(pattern): + found = evaluated(pattern, held.captures) + return held if found is not None and found == subject else None + if isinstance(pattern, _pattern_type()): + return pattern.match(subject, held.captures) + return held if fits(pattern, subject) else None + + +def fits(wanted, held) -> bool: + """Whether one pattern position admits one concrete value.""" + if wanted is None: + return True + if isinstance(wanted, _pattern_type()) or is_symbolic(wanted): + return matched(wanted, held) is not None + if is_layout_wildcard(wanted): + return type(held) is int + return wanted == held + + +def grouping(modes): + """The tuple nesting of *modes*, with the concrete modes erased.""" + if isinstance(modes, tuple): + return tuple(grouping(mode) for mode in modes) + return None + + +def relations_of(values) -> tuple[str, ...]: + lines: dict[str, None] = {} + for value in values: + if isinstance(value, _pattern_type()): + lines.update(dict.fromkeys(value.relations())) + return tuple(lines) + + +def between_rules(op_type) -> tuple: + return tuple(getattr(op_type, "between", ())) + + +def refusals_between(op_type, operands: dict) -> tuple[str, ...]: + return tuple( + rule.refused(operands) for rule in between_rules(op_type) if not rule.holds(operands) + ) + + +def affine_frame(layout) -> tuple | None: + """Read a bare layout or a composition through no transform as offset + layout.""" + if not isinstance(layout, ComposedLayout): + return 0, layout + if layout.inner is not None: + return None + return layout.offset, layout.outer + + +__all__ = [ + "ABSENT", + "ARRANGEMENT", + "Match", + "OPAQUE", + "UNNAMED_PLACE", + "_named", + "affine_frame", + "alternatives_of", + "between_rules", + "evaluated", + "fits", + "grouping", + "is_symbolic", + "matched", + "refusals_between", + "relations_of", + "resolved", + "written_alternatives", + "written_binding", + "written_bindings", + "written_dim", + "written_field", + "written_grouped", + "written_place", + "written_tuple", +] diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py new file mode 100644 index 00000000..bb98d4f8 --- /dev/null +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -0,0 +1,788 @@ +"""Composable predicates for operation declarations and specialization dispatch.""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from typing import Any + +from tilefoundry.ir.types import ( + Broadcast, + ComposedLayout, + Layout, + Mesh, + ShardLayout, + Swizzle, + TensorType, + make_mesh, +) +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.layout_algebra import is_inverse_projectable +from tilefoundry.ir.types.mesh import separate + +from .match import ( + ABSENT, + ARRANGEMENT, + UNNAMED_PLACE, + Match, + _named, + affine_frame, + alternatives_of, + evaluated, + grouping, + matched, + relations_of, + resolved, + written_alternatives, + written_binding, + written_grouped, + written_place, + written_tuple, +) + + +@dataclass(frozen=True) +class Pattern: + """Base class for a predicate that returns bindings or ``None``.""" + + def match(self, subject, captures=None) -> Match | None: + raise NotImplementedError + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return type(self).__name__ + + def relations(self) -> tuple[str, ...]: + return () + + def alternatives(self, bindings=()) -> tuple: + return ((tuple(bindings), self),) + + def rules(self, arrangements=None) -> tuple[str, ...]: + return self.relations() + + def resolve(self, bindings): + return replace( + self, + **{name: resolved(value, bindings) for name, value in vars(self).items()}, + ) + + +@dataclass(frozen=True) +class WildcardPattern(Pattern): + """Match any value without binding it.""" + + def match(self, subject, captures=None) -> Match: + return Match(dict(captures or {})) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return name + + +@dataclass(frozen=True, init=False) +class OrPattern(Pattern): + patterns: tuple + + def __init__(self, *patterns): + object.__setattr__(self, "patterns", tuple(patterns)) + + def match(self, subject, captures=None): + for pattern in self.patterns: + held = matched(pattern, subject, captures) + if held is not None: + return held + return None + + def describe(self, name: str = UNNAMED_PLACE) -> str: + if not any(isinstance(pattern, Pattern) for pattern in self.patterns): + values = "{" + ", ".join(written_place(p) for p in self.patterns) + "}" + return values if name == UNNAMED_PLACE else f"{name} in {values}" + return written_alternatives(self.alternatives(), name) + + def relations(self) -> tuple[str, ...]: + return relations_of(self.patterns) + + def alternatives(self, bindings=()) -> tuple: + return tuple( + held for pattern in self.patterns for held in alternatives_of(pattern, bindings) + ) + + def resolve(self, bindings): + held = tuple( + pattern + for pattern in (resolved(one, bindings) for one in self.patterns) + if pattern is not ABSENT + ) + return OrPattern(*held) if held else ABSENT + + +@dataclass(frozen=True) +class AndPattern(Pattern): + parts: tuple = field(default_factory=tuple) + + def match(self, subject, captures=None): + held = Match(dict(captures or {})) + for pattern in self.parts: + held = matched(pattern, subject, held.captures) + if held is None: + return None + return held + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return " and ".join(written_place(pattern, name) for pattern in self.parts) + + def relations(self) -> tuple[str, ...]: + return relations_of(self.parts) + + +@dataclass(frozen=True, init=False) +class SequencePattern(Pattern): + patterns: tuple + + def __init__(self, *patterns): + object.__setattr__(self, "patterns", tuple(patterns)) + + def match(self, subject, captures=None): + if not isinstance(subject, (tuple, list)) or len(subject) != len(self.patterns): + return None + held = Match(dict(captures or {})) + for pattern, value in zip(self.patterns, subject): + held = matched(pattern, value, held.captures) + if held is None: + return None + return held + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return written_tuple(tuple(written_place(p, name) for p in self.patterns)) + + def relations(self) -> tuple[str, ...]: + return relations_of(self.patterns) + + def resolve(self, bindings): + return SequencePattern(*(resolved(pattern, bindings) for pattern in self.patterns)) + + +@dataclass(frozen=True) +class CapturePattern(Pattern): + name: str + pattern: object = None + + def match(self, subject, captures=None): + held = dict(captures or {}) + if self.name in held: + return Match(held) if held[self.name] == subject else None + found = matched(self.pattern, subject, held) + if found is None: + return None + return Match({**found.captures, self.name: subject}) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return self.name + + def relations(self) -> tuple[str, ...]: + if self.pattern is None: + return () + return (written_place(self.pattern, self.name), *relations_of((self.pattern,))) + + +@dataclass(frozen=True, init=False) +class ConstraintPattern(Pattern): + patterns: tuple + + def __init__(self, *patterns): + if not patterns: + raise ValueError("a constraint pattern must state at least one constraint") + object.__setattr__(self, "patterns", tuple(patterns)) + + def match(self, subject, captures=None): + held = Match(dict(captures or {})) + for pattern in self.patterns: + held = matched(pattern, subject, held.captures) + if held is None: + return None + return held + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return " and ".join(written_place(pattern, name) for pattern in self.patterns) + + def resolve(self, bindings): + return self + + +@dataclass(frozen=True) +class MultipleOfPattern(Pattern): + unit: int + + def __post_init__(self): + if type(self.unit) is not int or self.unit <= 0: + raise ValueError("MultipleOfPattern unit must be a positive int") + + def match(self, subject, captures=None): + return ( + Match(dict(captures or {})) + if type(subject) is int and subject % self.unit == 0 + else None + ) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return f"{name} % {self.unit} = 0" + + +@dataclass(frozen=True) +class RangePattern(Pattern): + """A closed integer range, optionally naming a specialization dimension.""" + + dim_var: str = "" + lo: int | None = None + hi: int | None = None + + def __post_init__(self): + if not isinstance(self.dim_var, str): + raise TypeError("RangePattern dim_var must be a str") + for name in ("lo", "hi"): + value = getattr(self, name) + if value is not None and (not isinstance(value, int) or isinstance(value, bool)): + raise TypeError( + f"RangePattern: {name} must be int or None, got {type(value).__name__}" + ) + if self.lo is not None and self.hi is not None and self.lo > self.hi: + raise ValueError("RangePattern requires lo <= hi (closed [lo, hi])") + if self.lo is None and self.hi is None: + raise ValueError("RangePattern must state a lower or upper bound") + if self.dim_var and (self.lo is None or self.hi is None): + raise ValueError("a named RangePattern must state both lo and hi") + + def match(self, subject, captures=None): + if isinstance(subject, bool) or not isinstance(subject, int): + return None + if self.lo is not None and subject < self.lo: + return None + if self.hi is not None and subject > self.hi: + return None + return Match(dict(captures or {})) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + if self.lo is None: + return f"{name} <= {self.hi}" + if self.hi is None: + return f"{self.lo} <= {name}" + return f"{self.lo} <= {name} <= {self.hi}" + + +@dataclass(frozen=True) +class OneOfPattern(Pattern): + values: tuple + + def __post_init__(self): + if len(self.values) < 2: + raise ValueError("OneOfPattern requires at least two values") + + def match(self, subject, captures=None): + return ( + Match(dict(captures or {})) if any(subject == value for value in self.values) else None + ) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return f"{name} in {{{', '.join(_named(value) for value in self.values)}}}" + + +@dataclass(frozen=True) +class AttrPattern(Pattern): + attr: str + pattern: object + + def match(self, subject, captures=None): + if not hasattr(subject, self.attr): + return None + return matched(self.pattern, getattr(subject, self.attr), captures) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return written_place(self.pattern, f"{name}.{self.attr}") + + +@dataclass(frozen=True) +class BitsPattern(Pattern): + dtype: str + pattern: object + + def match(self, subject, captures=None): + width = getattr(dict(captures or {}).get(self.dtype), "bit_width", None) + if type(subject) is not int or type(width) is not int: + return None + return matched(self.pattern, subject * width, captures) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return written_place(self.pattern, f"{name} * {self.dtype}.bit_width") + + +@dataclass(frozen=True, init=False) +class SwitchPattern(Pattern): + param: str + branches: tuple + + def __init__(self, param, branches): + object.__setattr__(self, "param", param) + object.__setattr__(self, "branches", tuple(dict(branches).items())) + + def match(self, subject, captures=None): + held = dict(captures or {}) + if self.param in held: + wanted = held[self.param] + pattern = next((p for value, p in self.branches if value == wanted), None) + return None if pattern is None else matched(pattern, subject, held) + for value, pattern in self.branches: + found = matched(pattern, subject, {**held, self.param: value}) + if found is not None: + return found + return None + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return written_alternatives(self.alternatives(), name) + + def relations(self) -> tuple[str, ...]: + return relations_of(tuple(pattern for _, pattern in self.branches)) + + def refusal(self, subject, captures=None) -> str | None: + held = dict(captures or {}) + if self.param not in held: + return ( + None + if self.match(subject, held) is not None + else f"{subject!r} is none of {len(self.branches)} branches" + ) + pattern = next((p for value, p in self.branches if value == held[self.param]), None) + if pattern is None: + return f"{self.param}={written_binding(held[self.param])} selects no branch" + explained = getattr(pattern, "refusal", None) + if explained is not None: + return explained(subject, held) + return ( + None + if matched(pattern, subject, held) is not None + else f"{subject!r} is not {written_place(pattern)}" + ) + + def alternatives(self, bindings=()) -> tuple: + return tuple( + held + for value, pattern in self.branches + for held in alternatives_of(pattern, (*bindings, (self.param, value))) + ) + + def resolve(self, bindings): + held = dict(bindings or {}) + if self.param in held: + pattern = next((p for value, p in self.branches if value == held[self.param]), ABSENT) + return ABSENT if pattern is ABSENT else resolved(pattern, held) + branches = {value: resolved(pattern, held) for value, pattern in self.branches} + branches = {value: p for value, p in branches.items() if p is not ABSENT} + return SwitchPattern(self.param, branches) if branches else ABSENT + + +@dataclass(frozen=True) +class GuardPattern(Pattern): + symbol: object + condition: Pattern + pattern: object + + def match(self, subject, captures=None): + value = evaluated(self.symbol, captures) + if value is None or matched(self.condition, value, captures) is None: + return None + return matched(self.pattern, subject, captures) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return written_place(self.pattern, name) + + def relations(self) -> tuple[str, ...]: + return ( + *relations_of((self.pattern,)), + self.condition.describe(str(self.symbol)), + ) + + def resolve(self, bindings): + value = evaluated(self.symbol, bindings) + if value is None: + return GuardPattern(self.symbol, self.condition, resolved(self.pattern, bindings)) + return ( + resolved(self.pattern, bindings) + if matched(self.condition, value) is not None + else ABSENT + ) + + def fixed(self): + return None + + +@dataclass(frozen=True) +class LayoutPattern(Pattern): + """An affine layout, preserving grouping and optional whole-layout rules.""" + + shape: tuple + strides: tuple + forward: bool = True + injective: bool = True + per_mode: bool = False + + def positions(self) -> tuple: + return (*flatten(self.shape), *flatten(self.strides)) + + def match(self, subject, captures=None): + framed = affine_frame(subject) + layout = None if framed is None else framed[1] + if not isinstance(layout, Layout) or layout.strides is None: + return None + if grouping(tuple(layout.shape)) != grouping(tuple(self.shape)) or grouping( + tuple(layout.strides) + ) != grouping(tuple(self.strides)): + return None + extents = tuple(flatten(layout.shape)) + strides = tuple(flatten(layout.strides)) + if any(type(number) is not int for number in (*extents, *strides)): + return None + if any(number <= 0 for number in extents): + return None + held = matched(SequencePattern(*self.positions()), (*extents, *strides), captures) + if held is None: + return None + arrangements = ( + tuple( + Layout(tuple(flatten(shape)), tuple(flatten(steps))) + for shape, steps in zip(self.shape, self.strides) + ) + if self.per_mode + else (Layout(extents, strides),) + ) + for arrangement in arrangements: + held_strides = tuple(flatten(arrangement.strides)) + if self.forward and any(stride < 0 for stride in held_strides): + return None + if self.injective and not is_inverse_projectable(arrangement): + return None + return held + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return ( + f"Layout({written_grouped(tuple(self.shape))}, {written_grouped(tuple(self.strides))})" + ) + + def relations(self) -> tuple[str, ...]: + lines = list(relations_of(self.positions())) + subject = "each top-level mode" if self.per_mode else ARRANGEMENT + if self.forward: + lines.append(f"{subject} has no backward step") + if self.injective: + lines.append(f"{subject} reaches each of its own slots exactly once") + return tuple(lines) + + def fixed(self): + if any(isinstance(value, Pattern) for value in self.positions()): + return None + return Layout(tuple(self.shape), tuple(self.strides)) + + +@dataclass(frozen=True) +class SwizzlePattern(Pattern): + bits: int + base: int + shift: int + + def match(self, subject, captures=None): + return ( + Match(dict(captures or {})) + if subject == Swizzle(self.bits, self.base, self.shift) + else None + ) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return f"Swizzle({self.bits}, {self.base}, {self.shift})" + + def fixed(self): + return Swizzle(self.bits, self.base, self.shift) + + +@dataclass(frozen=True) +class ComposedLayoutPattern(Pattern): + """Match the three fields of a concrete ``ComposedLayout`` only.""" + + inner: object = None + offset: object = None + outer: object = None + + def match(self, subject, captures=None): + if not isinstance(subject, ComposedLayout): + return None + held = Match(dict(captures or {})) + for pattern, value in ( + (self.inner, subject.inner), + (self.offset, subject.offset), + (self.outer, subject.outer), + ): + held = matched(pattern, value, held.captures) + if held is None: + return None + return held + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return ( + f"ComposedLayout({written_place(self.inner)}, " + f"{written_place(self.offset)}, {written_place(self.outer)})" + ) + + def relations(self) -> tuple[str, ...]: + return relations_of((self.inner, self.offset, self.outer)) + + def fixed(self): + held = tuple( + value.fixed() if isinstance(value, Pattern) and hasattr(value, "fixed") else value + for value in (self.inner, self.offset, self.outer) + ) + return None if any(value is None for value in held[1:]) else ComposedLayout(*held) + + +@dataclass(frozen=True) +class MeshPattern(Pattern): + """Match named mesh levels after separating and recomposing them. + + A ``ComposedLayoutPattern`` here deliberately matches only a sliced mesh, + because an unsliced mesh carries a bare ``Layout``. To accept both forms, + use ``OrPattern(ComposedLayoutPattern(offset=..., outer=L), L)``. + """ + + topologies: tuple[str, ...] + layout: object + + def __post_init__(self): + if ( + not self.topologies + or any(not isinstance(name, str) or not name for name in self.topologies) + or len(set(self.topologies)) != len(self.topologies) + ): + raise ValueError("MeshPattern topologies must be unique non-empty names") + + def require_per_mode(pattern) -> None: + if isinstance(pattern, OrPattern): + for alternative in pattern.patterns: + require_per_mode(alternative) + return + if isinstance(pattern, ComposedLayoutPattern): + pattern = pattern.outer + if not isinstance(pattern, LayoutPattern) or not pattern.per_mode: + raise ValueError( + "MeshPattern layout must be a LayoutPattern(per_mode=True), " + "or a ComposedLayoutPattern whose outer uses per_mode=True" + ) + + require_per_mode(self.layout) + + def match(self, subject, captures=None): + if not isinstance(subject, Mesh): + return None + picked = tuple( + level + for level in separate(subject) + if getattr(level.topologies[0], "name", level.topologies[0]) in self.topologies + ) + found = tuple(getattr(level.topologies[0], "name", level.topologies[0]) for level in picked) + if len(picked) != len(self.topologies) or set(found) != set(self.topologies): + return None + return matched(self.layout, make_mesh(*picked).layout, captures) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return f"Mesh({self.topologies!r}, {written_place(self.layout)})" + + def relations(self) -> tuple[str, ...]: + return relations_of((self.layout,)) + + +@dataclass(frozen=True) +class ScalarPattern(Pattern): + def match(self, subject, captures=None): + return ( + Match(dict(captures or {})) + if isinstance(subject, TensorType) and subject.shape == () + else None + ) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return "scalar" + + +@dataclass(frozen=True) +class TensorPattern(Pattern): + """Match a non-scalar ``TensorType`` field by field.""" + + rank: int | None = None + dtype: Any = None + shape: tuple | None = None + storage: Any = None + layout: Pattern | None = None + + def match(self, subject, captures=None): + if not isinstance(subject, TensorType) or subject.shape == (): + return None + wanted = ( + (self.rank, len(subject.shape)), + (self.dtype, subject.dtype), + (self.storage, subject.storage), + ) + if self.shape is not None: + wanted = ((SequencePattern(*self.shape), tuple(subject.shape)), *wanted) + if self.layout is not None: + wanted = (*wanted, (self.layout, subject.layout)) + held = Match(dict(captures or {})) + for pattern, value in wanted: + held = matched(pattern, value, held.captures) + if held is None: + return None + return held + + def refusal(self, subject, captures=None) -> str | None: + if not isinstance(subject, TensorType) or subject.shape == (): + return f"{subject!r} is no tensor" + wanted = ( + ( + "shape", + None if self.shape is None else SequencePattern(*self.shape), + tuple(subject.shape), + ), + ("rank", self.rank, len(subject.shape)), + ("dtype", self.dtype, subject.dtype), + ("storage", self.storage, subject.storage), + ("layout", self.layout, subject.layout), + ) + held = Match(dict(captures or {})) + for name, pattern, value in wanted: + if pattern is None: + continue + found = matched(pattern, value, held.captures) + if found is not None: + held = found + continue + explained = getattr(pattern, "refusal", None) + if name == "layout" and explained is not None: + return explained(value, held.captures) + return f"its {name} is {_named(value) if name != 'shape' else value}" + return None + + def describe(self, name: str = UNNAMED_PLACE, arrangements=None) -> str: + stated = [] + if self.shape is not None: + stated.append("shape=" + written_tuple(tuple(written_place(x) for x in self.shape))) + if self.rank is not None: + stated.append(f"rank={self.rank}") + if self.dtype is not None: + stated.append(f"dtype={_named(self.dtype)}") + if self.storage is not None: + stated.append(f"storage={_named(self.storage)}") + head = " ".join(stated) if stated else "any tensor" + return ( + f"{head}, in any arrangement" + if self.layout is None + else (f"{head}, held in {self.layout.describe()}") + ) + + def relations(self) -> tuple[str, ...]: + return relations_of( + ( + self.rank, + self.dtype, + self.storage, + *(self.shape or ()), + *((self.layout,) if self.layout is not None else ()), + ) + ) + + +@dataclass(frozen=True) +class ShardLayoutPattern(Pattern): + """Match a sharded layout's arrangement, shard attrs, and mesh frame.""" + + arrangement: object + attrs: tuple + mesh: Mesh + + def match(self, subject, captures=None): + if not isinstance(subject, ShardLayout): + return None + extra = len(subject.attrs) - len(self.attrs) + if ( + extra < 0 + or subject.attrs[extra:] != self.attrs + or any(not isinstance(attr, Broadcast) for attr in subject.attrs[:extra]) + ): + return None + framed = affine_frame(subject.mesh.layout) + if framed is None: + return None + frame = framed[1] + if extra: + if not isinstance(frame, Layout): + return None + frame = Layout(frame.shape[extra:], frame.strides[extra:]) + subject_names = tuple( + getattr(topology, "name", topology) for topology in subject.mesh.topologies + ) + wanted_names = tuple( + getattr(topology, "name", topology) for topology in self.mesh.topologies + ) + if frame != self.mesh.layout or subject_names != wanted_names: + return None + return self.reads(subject.layout, captures) + + def reads(self, layout, captures=None): + return matched(self.arrangement, layout, captures) + + def accepts_layout(self, layout) -> bool: + return self.reads(layout) is not None + + def alternatives(self, bindings=()) -> tuple: + return alternatives_of(self.arrangement, bindings) + + def relations(self) -> tuple[str, ...]: + return relations_of((self.arrangement,)) + + def rules(self, arrangements=None) -> tuple[str, ...]: + items = self.alternatives() if arrangements is None else tuple(arrangements) + return relations_of(tuple(pattern for _, pattern in items)) + + def describe(self, name: str = UNNAMED_PLACE, arrangements=None) -> str: + items = self.alternatives() if arrangements is None else tuple(arrangements) + head = f"{len(items)} arrangement{'' if len(items) == 1 else 's'}:" + written = written_alternatives(items).splitlines() + return "\n".join( + ( + head, + *(f" {line}" for line in written), + *(f" {rule}" for rule in self.rules(items)), + ) + ) + + +Scalar: ScalarPattern = ScalarPattern() +Tensor: TensorPattern = TensorPattern() + + +__all__ = [ + "AndPattern", + "AttrPattern", + "BitsPattern", + "CapturePattern", + "ComposedLayoutPattern", + "ConstraintPattern", + "GuardPattern", + "LayoutPattern", + "MeshPattern", + "MultipleOfPattern", + "OneOfPattern", + "OrPattern", + "Pattern", + "RangePattern", + "Scalar", + "ScalarPattern", + "SequencePattern", + "ShardLayoutPattern", + "SwizzlePattern", + "SwitchPattern", + "Tensor", + "TensorPattern", + "WildcardPattern", +] diff --git a/src/tilefoundry/ir/pattern/utils.py b/src/tilefoundry/ir/pattern/utils.py new file mode 100644 index 00000000..30eec000 --- /dev/null +++ b/src/tilefoundry/ir/pattern/utils.py @@ -0,0 +1,68 @@ +"""Construction and specialization helpers for IR patterns.""" + +from __future__ import annotations + +from tilefoundry.ir.types import ComposedLayout, Swizzle + +from .pattern import ( + ComposedLayoutPattern, + LayoutPattern, + Pattern, + RangePattern, + SwizzlePattern, +) + + +def arrangement_pattern( + layout, + *, + forward: bool = True, + injective: bool = True, + per_mode: bool = False, +): + """Build the exact pattern for one authored arrangement.""" + if isinstance(layout, ComposedLayout): + inner = layout.inner + return ComposedLayoutPattern( + SwizzlePattern(inner.bits, inner.base, inner.shift) + if isinstance(inner, Swizzle) + else inner, + layout.offset, + arrangement_pattern( + layout.outer, + forward=forward, + injective=injective, + per_mode=per_mode, + ), + ) + return LayoutPattern( + tuple(layout.shape), + tuple(layout.strides), + forward=forward, + injective=injective, + per_mode=per_mode, + ) + + +def locate_dim_var(params: tuple, name: str) -> tuple[int, int] | None: + """Return the first parameter/axis carrying a ``DimVar`` named *name*.""" + for index, param in enumerate(params): + shape = getattr(param.type, "shape", None) + if shape is None: + continue + for axis, dim in enumerate(shape): + if getattr(dim, "name", None) == name: + return index, axis + return None + + +def _mangle_variant_name(name: str, specializations: tuple[Pattern, ...]) -> str: + if len(specializations) != 1 or not isinstance(specializations[0], RangePattern): + raise TypeError("variant requires exactly one RangePattern") + pattern = specializations[0] + if not pattern.dim_var or pattern.lo is None or pattern.hi is None: + raise TypeError("variant RangePattern requires dim_var, lo, and hi") + return f"{name}${pattern.dim_var}${pattern.lo}_{pattern.hi}" + + +__all__ = ["_mangle_variant_name", "arrangement_pattern", "locate_dim_var"] diff --git a/src/tilefoundry/ir/tir/arith.py b/src/tilefoundry/ir/tir/arith.py index 6dbeb1ee..35bf8814 100644 --- a/src/tilefoundry/ir/tir/arith.py +++ b/src/tilefoundry/ir/tir/arith.py @@ -15,8 +15,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.kinds import BinaryKind, UnaryKind from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/async_copy.py b/src/tilefoundry/ir/tir/async_copy.py index 3826ebf3..d4b204e8 100644 --- a/src/tilefoundry/ir/tir/async_copy.py +++ b/src/tilefoundry/ir/tir/async_copy.py @@ -4,8 +4,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/clamp.py b/src/tilefoundry/ir/tir/clamp.py index 3a5e8773..2edaa107 100644 --- a/src/tilefoundry/ir/tir/clamp.py +++ b/src/tilefoundry/ir/tir/clamp.py @@ -4,8 +4,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/cuda/memory/tma.py b/src/tilefoundry/ir/tir/cuda/memory/tma.py index 42e967f1..107c1e23 100644 --- a/src/tilefoundry/ir/tir/cuda/memory/tma.py +++ b/src/tilefoundry/ir/tir/cuda/memory/tma.py @@ -8,8 +8,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma.py b/src/tilefoundry/ir/tir/cuda/nn/mma.py index bc851736..fb9775e2 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma.py @@ -11,8 +11,8 @@ from tilefoundry.ir.core import Op, VerifyError from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology, UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -89,8 +89,7 @@ def _(call: "Call", ctx: "VerifyContext") -> None: ) if not any( - mesh_scope_matches_required_scope(s, atom.required_scope) - for s in ctx.mesh_scope + mesh_scope_matches_required_scope(s, atom.required_scope) for s in ctx.mesh_scope ): raise VerifyError( "T.mma: no enclosing mesh scope hosts the atom's required thread " diff --git a/src/tilefoundry/ir/tir/cuda/sync/mbarrier.py b/src/tilefoundry/ir/tir/cuda/sync/mbarrier.py index 594efff5..26230d6e 100644 --- a/src/tilefoundry/ir/tir/cuda/sync/mbarrier.py +++ b/src/tilefoundry/ir/tir/cuda/sync/mbarrier.py @@ -7,8 +7,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Scalar, Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Scalar, Tensor from tilefoundry.ir.types import UnitType from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/dot.py b/src/tilefoundry/ir/tir/dot.py index 5f18c4e3..f712b2c6 100644 --- a/src/tilefoundry/ir/tir/dot.py +++ b/src/tilefoundry/ir/tir/dot.py @@ -7,8 +7,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.ir.types.shard_layout import ShardLayout, shard_layout_local_shape from tilefoundry.ir.types.storage import StorageKind diff --git a/src/tilefoundry/ir/tir/memory/copy.py b/src/tilefoundry/ir/tir/memory/copy.py index 97252dff..0c01724d 100644 --- a/src/tilefoundry/ir/tir/memory/copy.py +++ b/src/tilefoundry/ir/tir/memory/copy.py @@ -10,8 +10,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/memory/fill.py b/src/tilefoundry/ir/tir/memory/fill.py index 9134ffe7..3e009730 100644 --- a/src/tilefoundry/ir/tir/memory/fill.py +++ b/src/tilefoundry/ir/tir/memory/fill.py @@ -8,8 +8,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/memory/memory_span.py b/src/tilefoundry/ir/tir/memory/memory_span.py index 4b63d858..47d2dd5e 100644 --- a/src/tilefoundry/ir/tir/memory/memory_span.py +++ b/src/tilefoundry/ir/tir/memory/memory_span.py @@ -8,8 +8,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer diff --git a/src/tilefoundry/ir/tir/memory/ptr_of.py b/src/tilefoundry/ir/tir/memory/ptr_of.py index d12993b4..13ed75b1 100644 --- a/src/tilefoundry/ir/tir/memory/ptr_of.py +++ b/src/tilefoundry/ir/tir/memory/ptr_of.py @@ -9,8 +9,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry import register_typeinfer diff --git a/src/tilefoundry/ir/tir/memory/tensor_view.py b/src/tilefoundry/ir/tir/memory/tensor_view.py index 0ab48a3f..1655a4e7 100644 --- a/src/tilefoundry/ir/tir/memory/tensor_view.py +++ b/src/tilefoundry/ir/tir/memory/tensor_view.py @@ -13,8 +13,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.layout import Layout, LayoutBase from tilefoundry.ir.types.stride import compact_row_major diff --git a/src/tilefoundry/ir/tir/nn/relu.py b/src/tilefoundry/ir/tir/nn/relu.py index 53d89549..573c9bc3 100644 --- a/src/tilefoundry/ir/tir/nn/relu.py +++ b/src/tilefoundry/ir/tir/nn/relu.py @@ -9,8 +9,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/nn/rms_norm.py b/src/tilefoundry/ir/tir/nn/rms_norm.py index a519b6af..e866a9fd 100644 --- a/src/tilefoundry/ir/tir/nn/rms_norm.py +++ b/src/tilefoundry/ir/tir/nn/rms_norm.py @@ -10,8 +10,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/prim_function.py b/src/tilefoundry/ir/tir/prim_function.py index 348d8f86..1f7fa122 100644 --- a/src/tilefoundry/ir/tir/prim_function.py +++ b/src/tilefoundry/ir/tir/prim_function.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from tilefoundry.ir.core import Var -from tilefoundry.ir.core.pattern import Pattern +from tilefoundry.ir.pattern import Pattern from tilefoundry.ir.tir.stmt import Stmt from tilefoundry.ir.tir.stmts import Sequential from tilefoundry.target.base import Target, target_instance @@ -41,7 +41,9 @@ def __post_init__(self) -> None: def add_variant(self, variant: "PrimFunction") -> None: if getattr(self, "_sealed", False): - raise RuntimeError(f"tir PrimFunction {self.name!r}: cannot add a specialization variant after sealing") + raise RuntimeError( + f"tir PrimFunction {self.name!r}: cannot add a specialization variant after sealing" + ) self.variants = (*self.variants, variant) diff --git a/src/tilefoundry/ir/tir/reduce.py b/src/tilefoundry/ir/tir/reduce.py index 15cb2391..e36521a6 100644 --- a/src/tilefoundry/ir/tir/reduce.py +++ b/src/tilefoundry/ir/tir/reduce.py @@ -5,8 +5,8 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.kinds import ReduceKind from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor from tilefoundry.ir.types import UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/verify.py b/src/tilefoundry/ir/tir/verify.py index 923434bb..a8ae1a90 100644 --- a/src/tilefoundry/ir/tir/verify.py +++ b/src/tilefoundry/ir/tir/verify.py @@ -11,12 +11,12 @@ from tilefoundry.ir.core import Expr, Var, VerifyError from tilefoundry.ir.core.expr import Call, Constant -from tilefoundry.ir.core.pattern import DimVarRangePat, locate_dim_var from tilefoundry.ir.hir.function import ( Function as HirFunction, ) from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord from tilefoundry.ir.hir.verify import verify_function +from tilefoundry.ir.pattern import RangePattern, locate_dim_var from tilefoundry.ir.types import DType, TensorType, UnitType from tilefoundry.ir.types.callable_type import callable_type_for_prim_function from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimMax, DimMin, DimMod, DimMul, DimSub @@ -57,9 +57,11 @@ def verify_prim_function( _check_param_homogeneity(fn) if fn.variants: for variant in fn.variants: - if len(variant.specializations) != 1 or not isinstance(variant.specializations[0], DimVarRangePat): + if len(variant.specializations) != 1 or not isinstance( + variant.specializations[0], RangePattern + ): raise VerifyError( - f"PrimFunction {fn.name!r}: each variant must have one DimVarRangePat" + f"PrimFunction {fn.name!r}: each variant must have one RangePattern" ) pat = variant.specializations[0] if locate_dim_var(fn.params, pat.dim_var) is None: @@ -272,8 +274,7 @@ def _check_bound_coordinates(field: str, bound, scope) -> None: ) if not any(held is mesh or held == mesh for held in scope): raise VerifyError( - f"For.{field} reads a coordinate of {mesh!r}, which no enclosing " - "MeshScope binds" + f"For.{field} reads a coordinate of {mesh!r}, which no enclosing MeshScope binds" ) @@ -537,9 +538,7 @@ def verify_module(fns) -> None: if isinstance(fns, Module): fns = module_functions(fns) prim_fns = [f for f in fns if isinstance(f, PrimFunction)] - prim_fns_with_variants = [ - variant for f in prim_fns for variant in (f, *f.variants) - ] + prim_fns_with_variants = [variant for f in prim_fns for variant in (f, *f.variants)] for f in fns: if isinstance(f, HirFunction): verify_function(f) diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index af4cd911..34df2619 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -34,7 +34,6 @@ from tilefoundry.ir.core.kinds import BinaryKind, UnaryKind from tilefoundry.ir.core.module import Module from tilefoundry.ir.core.op_schema import OpSchema -from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.math.binary import Binary @@ -49,6 +48,7 @@ from tilefoundry.ir.hir.tensor.slice import Slice, slice_size from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.isl_interop import normalize_dim +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.tir.stmts import ( @@ -255,7 +255,7 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex DimMul=DimMul, DimSub=DimSub, DimVar=DimVar, - DimVarRangePat=DimVarRangePat, + RangePattern=RangePattern, Evaluate=Evaluate, For=For, Expr=Expr, @@ -1199,9 +1199,7 @@ def _finalize(self, cls: type) -> object: verify_function(function, module=result) inference_type( function, - runtime.TypeInferContext( - scope=runtime.FunctionScope(result, function) - ), + runtime.TypeInferContext(scope=runtime.FunctionScope(result, function)), ranges=True, ) elif isinstance(function, runtime.PrimFunction): diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index 6e6ed486..b60a9bed 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -14,23 +14,23 @@ from dataclasses import dataclass from typing import Any, ClassVar, Literal, get_args, get_origin -from tilefoundry.ir.constraints import ( - ConstraintProvenance, - LayoutConstraint, - MeshConstraint, - ScheduleConstraintMetadata, +from tilefoundry.ir.clause import ( + ClauseProvenance, + LayoutClause, + MeshClause, SourceLocation, - StorageConstraint, + StorageClause, + WhereClauseMetadata, ) -from tilefoundry.ir.constraints.layout import _LAYOUT_WILDCARD +from tilefoundry.ir.clause.layout import _LAYOUT_WILDCARD from tilefoundry.ir.core import ( BindingMetadata, RangeMetadata, attach_metadata, get_metadata, ) -from tilefoundry.ir.core.pattern import _mangle_variant_name from tilefoundry.ir.hir.nn.matmul import MatMul +from tilefoundry.ir.pattern import _mangle_variant_name from tilefoundry.ir.tir.launch import launch_call from tilefoundry.ir.types import Broadcast, Layout, Partial, Split, TensorType from tilefoundry.ir.types.dim import DimVar @@ -388,9 +388,7 @@ def construct(match, children, context): axis_name = node.attr mesh = context.lexical_scope.lookup_mesh(binding) if not isinstance(mesh, runtime.Mesh): - raise ParseError.from_node( - node, context, f"{binding!r} is not a lexical Mesh binding" - ) + raise ParseError.from_node(node, context, f"{binding!r} is not a lexical Mesh binding") if axis_name is None: if len(flatten(mesh.layout).shape) != 1: raise ParseError.from_node( @@ -584,9 +582,7 @@ def _placement_meshes(value: _PlacementCandidate, context: MatchContext, match): if not referenced_ids: return () if context.function is None: - raise ParseError.from_node( - match.node, context, "placed layout requires function context" - ) + raise ParseError.from_node(match.node, context, "placed layout requires function context") meshes = tuple(dict.fromkeys(entry[0] for entry in (*value.splits, *value.states))) if len(meshes) != len(referenced_ids): raise ParseError.from_node( @@ -626,7 +622,8 @@ def apply(self, value, *, match, context): 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, + match.node, + context, f"a layout can split one level once; two of these meshes name {duplicates}", ) return value @@ -765,20 +762,38 @@ 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) @@ -790,16 +805,23 @@ 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 @@ -807,14 +829,28 @@ 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", ())) + 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) RULES: ClassVar[tuple[AstRule[Any], ...]] = ( - LayoutStrideRankRule(), MeshAxisBoundOnceRule(), PlacementMeshResolutionRule(), - PlacementLevelRule(), PlacementConstructionRule(), PlacementAnswerRule(), + LayoutStrideRankRule(), + MeshAxisBoundOnceRule(), + PlacementMeshResolutionRule(), + PlacementLevelRule(), + PlacementConstructionRule(), + PlacementAnswerRule(), ) @@ -1233,8 +1269,7 @@ class TupleTypePattern(ElementPattern): @staticmethod def construct(match, children, context): fields = tuple( - children[f"field_{index}"] - for index in range(match.captures.get("field_count", 1)) + children[f"field_{index}"] for index in range(match.captures.get("field_count", 1)) ) if not all(isinstance(field, (runtime.TensorType, runtime.TupleType)) for field in fields): raise ParseError.from_node( @@ -1353,7 +1388,7 @@ def resolve_extent(item: ast.AST): bindings.extend(_parse_constraint_bindings(extras[0])) if len({topology for topology, _ in bindings}) != len(bindings): raise ValueError("layout constraint cannot bind one topology more than once") - return LayoutConstraint(layout=Layout(shape=tuple(shape)), bindings=tuple(bindings)) + return LayoutClause(layout=Layout(shape=tuple(shape)), bindings=tuple(bindings)) class WhereAnnotationPattern(ElementPattern): @@ -1396,23 +1431,23 @@ def construct(match, children, context): dataclasses.replace( _parse_layout_constraint(keyword.value, context), source_loc=location, - provenance=ConstraintProvenance.AUTHOR, + provenance=ClauseProvenance.AUTHOR, ) ) elif keyword.arg == "mesh": constraints.append( - MeshConstraint( + MeshClause( mesh=_where_static(keyword.value, context), source_loc=location, - provenance=ConstraintProvenance.AUTHOR, + provenance=ClauseProvenance.AUTHOR, ) ) elif keyword.arg == "storage": constraints.append( - StorageConstraint( + StorageClause( storage=_where_static(keyword.value, context), source_loc=location, - provenance=ConstraintProvenance.AUTHOR, + provenance=ClauseProvenance.AUTHOR, ) ) else: @@ -1422,7 +1457,7 @@ def construct(match, children, context): ) except (TypeError, ValueError) as error: raise ParseError.from_node(node, context, str(error)) from error - return ScheduleConstraintMetadata(constraints=tuple(constraints), source_loc=location) + return WhereClauseMetadata(constraints=tuple(constraints), source_loc=location) RULES: ClassVar[tuple[AstRule[Any], ...]] = () @@ -3472,6 +3507,7 @@ def _scoped_region(mesh, body, params=(), args=()): def _mesh_scope_captures(node, context): """Capture outer expression bindings for one region boundary.""" + def free_names(statements, outer_bound=frozenset()): local = _directly_bound_names(statements) visible = outer_bound | local @@ -3574,9 +3610,7 @@ def _rebind_through_region(context, mesh, names, frame, node, params=(), args=() _bind_region_results(context, scoped, [values[0][0]], node) return values[0][0] tuple_type = runtime.TupleType(fields=tuple(value.type for _name, value in values)) - tuple_body = runtime.IrTuple( - type=tuple_type, elements=tuple(value for _name, value in values) - ) + tuple_body = runtime.IrTuple(type=tuple_type, elements=tuple(value for _name, value in values)) scoped = _scoped_region(mesh, tuple_body, params, args) _bind_region_results(context, scoped, [name for name, _value in values], node) return values[0][0] @@ -3834,9 +3868,7 @@ def _statement_names(node: object, context: MatchContext) -> tuple[str, ...]: target = child.targets[0] targets = target.elts if isinstance(target, ast.Tuple) else (target,) found.extend( - item.id - for item in targets - if isinstance(item, ast.Name) and item.id not in found + item.id for item in targets if isinstance(item, ast.Name) and item.id not in found ) return tuple(found) @@ -3959,19 +3991,14 @@ def _mentions_mesh_coordinate(node: ast.AST, context: MatchContext) -> bool: return False -def _iterator_arity_failure( - kind: str, count: int, node: ast.Call -) -> PatternFailure | None: +def _iterator_arity_failure(kind: str, count: int, node: ast.Call) -> PatternFailure | None: """Describe an invalid tile/range arity, if this call has one.""" if count in ({2, 3} if kind == "tile" else {1, 2, 3}): return None if kind == "tile" and count == 1: detail = "tile(extent) is not supported; use range(extent)" elif kind == "tile": - detail = ( - "tile() takes 2 or 3 arguments, (stop, step) or " - f"(start, stop, step), got {count}" - ) + detail = f"tile() takes 2 or 3 arguments, (stop, step) or (start, stop, step), got {count}" else: detail = f"range() takes 1 to 3 arguments, got {count}" return PatternFailure("loop_header", node, detail) @@ -4105,13 +4132,16 @@ def _bind( @staticmethod def construct(match, children, context): if context.function is not None and context.function.dialect == "tir": - iv = runtime.Var(type=runtime.TensorType.scalar(runtime.DType.i64), name=match.captures["target"]) + iv = runtime.Var( + type=runtime.TensorType.scalar(runtime.DType.i64), name=match.captures["target"] + ) values = dict(match.captures["defaults"]) values.update({name: value for name, value in children.items() if name != "carry"}) bounds = [values[name] for name in ("start", "extent", "step")] bounds = [_constant(v) if isinstance(v, (bool, int, float)) else v for v in bounds] bounds = [ - bound if runtime.static_dim_value(bound) is not None + bound + if runtime.static_dim_value(bound) is not None else runtime.normalize_dim(bound) for bound in bounds ] @@ -4245,14 +4275,39 @@ def construct(match, children, context): class ForPattern(ElementPattern): element_name = "for" syntax = LazyPattern( - lambda: BranchPattern("loop", AstNodePattern( - ast.For, - ChildPattern("header", LoopHeaderPattern(), "loop_header"), - FieldPattern("body", ChildPattern("body", ChoicePattern( - ConditionPattern("tir loop body", lambda node, context: context.function is not None and context.function.dialect == "tir", BlockPattern()), - ConditionPattern("hir loop body", lambda node, context: context.function is None or context.function.dialect == "hir", LoopBodyPattern()), - ), "loop_body", transform=_body_as_ast_module)), - ), pattern_id="statement.for") + lambda: BranchPattern( + "loop", + AstNodePattern( + ast.For, + ChildPattern("header", LoopHeaderPattern(), "loop_header"), + FieldPattern( + "body", + ChildPattern( + "body", + ChoicePattern( + ConditionPattern( + "tir loop body", + lambda node, context: ( + context.function is not None + and context.function.dialect == "tir" + ), + BlockPattern(), + ), + ConditionPattern( + "hir loop body", + lambda node, context: ( + context.function is None or context.function.dialect == "hir" + ), + LoopBodyPattern(), + ), + ), + "loop_body", + transform=_body_as_ast_module, + ), + ), + ), + pattern_id="statement.for", + ) ) @staticmethod @@ -4297,7 +4352,8 @@ class TirOnlyStatementRule: def apply(self, value, *, match, context): if context.function is None or context.function.dialect != "tir": raise ParseError.from_node( - match.node, context, + match.node, + context, f"{match.element_name} is a TIR statement; HIR does not support it", ) return value @@ -4311,10 +4367,15 @@ class IfPattern(ElementPattern): AstNodePattern( ast.If, CapturePattern("cond_node", lambda node, context: node.test), - FieldPattern("body", ChildPattern("then", BlockPattern(), "block", transform=_body_as_ast_module)), + FieldPattern( + "body", + ChildPattern("then", BlockPattern(), "block", transform=_body_as_ast_module), + ), FieldPattern( "orelse", - OptionalPattern(ChildPattern("else", BlockPattern(), "block", transform=_body_as_ast_module)), + OptionalPattern( + ChildPattern("else", BlockPattern(), "block", transform=_body_as_ast_module) + ), ), ), pattern_id="statement.if", @@ -4331,17 +4392,29 @@ def construct(match, children, context): RULES: ClassVar[tuple[AstRule[Any], ...]] = (TirOnlyStatementRule(),) + class WhilePattern(ElementPattern): element_name = "while" - syntax = LazyPattern(lambda: BranchPattern("while", AstNodePattern( - ast.While, - CapturePattern("cond_node", lambda node, context: node.test), - FieldPattern("body", ChildPattern("body", BlockPattern(), "block", transform=_body_as_ast_module)), - ), pattern_id="statement.while")) + syntax = LazyPattern( + lambda: BranchPattern( + "while", + AstNodePattern( + ast.While, + CapturePattern("cond_node", lambda node, context: node.test), + FieldPattern( + "body", + ChildPattern("body", BlockPattern(), "block", transform=_body_as_ast_module), + ), + ), + pattern_id="statement.while", + ) + ) @staticmethod def construct(match, children, context): - return runtime.While(_tir_scalar_expr(match.captures["cond_node"], context), children["body"]) + return runtime.While( + _tir_scalar_expr(match.captures["cond_node"], context), children["body"] + ) RULES: ClassVar[tuple[AstRule[Any], ...]] = (TirOnlyStatementRule(),) @@ -4588,14 +4661,14 @@ def construct(match, children, context): if context.function.dialect == "tir" and not isinstance(value, runtime.Expr): context.lexical_scope.define(match.captures["name"], value) return None - if isinstance(annotation, ScheduleConstraintMetadata): + if isinstance(annotation, WhereClauseMetadata): if context.function.dialect != "hir" or not isinstance(value.type, TensorType): raise ParseError.from_node( match.node, context, "where annotation requires a tensor-valued HIR Expr", ) - previous = get_metadata(value, ScheduleConstraintMetadata) + previous = get_metadata(value, WhereClauseMetadata) if previous is not None: binding = get_metadata(value, BindingMetadata) label = binding.name if binding is not None else "" @@ -4691,9 +4764,7 @@ class BlockPattern(ElementPattern): def _bind(node, _context, matched): assert isinstance(node, ast.Module) escaping = _block_escaping_names(node.body) - child_values = { - f"statement_{index}": values for index, values in escaping.items() - } + child_values = {f"statement_{index}": values for index, values in escaping.items()} return dataclasses.replace( matched, children=tuple( @@ -4761,7 +4832,11 @@ def apply(self, value, *, match, context): kind = context.function.function_kind if context.function.dialect == "hir" and kind == "prim_func": raise ParseError.from_node(match.node, context, "prim_func requires tir dialect") - if context.function.dialect == "tir" and kind != "prim_func" and context.function.role is not FunctionRole.VARIANT: + if ( + context.function.dialect == "tir" + and kind != "prim_func" + and context.function.role is not FunctionRole.VARIANT + ): raise ParseError.from_node(match.node, context, f"{kind} requires hir dialect") expected = runtime.Function if context.function.dialect == "hir" else runtime.PrimFunction if not isinstance(value, expected): @@ -4843,7 +4918,9 @@ def _validate_standalone( if function_context.role is FunctionRole.ROOT: return base = function_context.base - expected_base = runtime.Function if function_context.dialect == "hir" else runtime.PrimFunction + expected_base = ( + runtime.Function if function_context.dialect == "hir" else runtime.PrimFunction + ) if not isinstance(base, expected_base): raise ParseError.from_node(node, match_context, "standalone role lacks a matching base") if getattr(base, "_sealed", False): @@ -4969,16 +5046,13 @@ def construct(match, children, context): match.node, context, "function mesh could not be resolved" ) region_params = tuple( - runtime.Var(type=param.type, name=param.name) - for param in params + runtime.Var(type=param.type, name=param.name) for param in params ) body = runtime.BindingSubstitutionCloner().visit( body, {id(old): new for old, new in zip(params, region_params, strict=True)}, ) - body = _scoped_region( - outer_mesh, body, params=region_params, args=params - ) + body = _scoped_region(outer_mesh, body, params=region_params, args=params) declared_return = ( None if declared_return is None else canonicalize_dims(declared_return) ) @@ -5008,7 +5082,7 @@ def construct(match, children, context): setattr(function, runtime.DISPLAY_NAME, match.captures["name"]) if getattr(context.function, "dialect", None) == "tir" and specializations: pat = specializations[0] - if isinstance(pat, runtime.DimVarRangePat): + if isinstance(pat, runtime.RangePattern): function.name = _mangle_variant_name(function_name, (pat,)) else: function.name = function_name @@ -5034,7 +5108,7 @@ def construct(match, children, context): **kwargs, specializations=specializations, ) - if specializations and isinstance(specializations[0], runtime.DimVarRangePat): + if specializations and isinstance(specializations[0], runtime.RangePattern): function.name = _mangle_variant_name(function.name, (specializations[0],)) function._display_name = match.captures["name"] define = getattr(context.function.module_scope, "define", None) diff --git a/src/tilefoundry/script.py b/src/tilefoundry/script.py index 3d8d0971..1a0a4693 100644 --- a/src/tilefoundry/script.py +++ b/src/tilefoundry/script.py @@ -13,9 +13,9 @@ from typing import Any, Callable, ClassVar, Literal, Mapping from tilefoundry.ir.core.module import Module -from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern, _mangle_variant_name from tilefoundry.ir.hir.function import Function as HirFunction from tilefoundry.ir.hir.verify import verify_function +from tilefoundry.ir.pattern import Pattern, RangePattern, _mangle_variant_name from tilefoundry.ir.tir.intrinsic import intrinsic as _intrinsic from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.verify import verify_prim_function @@ -107,9 +107,7 @@ class ParsedFuncRules: ParsedFuncKind.VARIANT: HandleRule( lambda fn, key: _mangle_variant_name(fn.name, (key,)), "base" ), - ParsedFuncKind.CONVERTER: HandleRule( - lambda fn, key: f"{fn.name}.converter[{key}]", "base" - ), + ParsedFuncKind.CONVERTER: HandleRule(lambda fn, key: f"{fn.name}.converter[{key}]", "base"), } @classmethod @@ -132,9 +130,7 @@ def check( ) -def _binding_scope( - kind: ParsedFuncKind, entry: _Entry | None, base_name: str | None -) -> str: +def _binding_scope(kind: ParsedFuncKind, entry: _Entry | None, base_name: str | None) -> str: if entry is not None: module_name = entry.owner_name or "" if kind is ParsedFuncKind.VARIANT: @@ -163,9 +159,9 @@ def _validate_one_pattern(pattern: Any) -> Pattern: f"tilefoundry.specialize: pattern must be a Pattern instance, got " f"{type(pattern).__name__}" ) - if not isinstance(pattern, DimVarRangePat): + if not isinstance(pattern, RangePattern): raise TypeError( - f"tilefoundry.specialize: only DimVarRangePat is supported for v0, " + f"tilefoundry.specialize: only RangePattern is supported for v0, " f"got {type(pattern).__name__}" ) return pattern @@ -182,8 +178,7 @@ def _validate_converter_weight_name(base: HirFunction, weight_name: str) -> None ) return raise TypeError( - f"tilefoundry.converter: {base.name!r} has no ConstTensor param named " - f"{weight_name!r}" + f"tilefoundry.converter: {base.name!r} has no ConstTensor param named {weight_name!r}" ) @@ -426,7 +421,7 @@ def _wrap(fn_inner): def _specialize(self: HirFunction, pattern: Any): - """``@base.specialize(DimVarRangePat(...))`` — register a shape variant. + """``@base.specialize(RangePattern(...))`` — register a shape variant. Parses the decorated ``def`` into a variant ``hir.Function`` and appends it to ``base.variants``. The identifier becomes the variant's display label and the @@ -467,7 +462,6 @@ def _wrap_variant(fn_inner): return _wrap_variant - HirFunction.specialize = _specialize PrimFunction.specialize = _specialize @@ -492,10 +486,7 @@ def _wrap_converter(fn_inner): topologies=_enclosing_topologies(), ) if ir.body is None: - raise TypeError( - "tilefoundry.converter: a converter must have a real body, " - "not `pass`" - ) + raise TypeError("tilefoundry.converter: a converter must have a real body, not `pass`") ir.name = f"{self.name}.converter[{weight_name}]" verify_function(ir) diff --git a/tests/analysis/test_analysis_invariants.py b/tests/analysis/test_analysis_invariants.py index 3f8b10df..6d7ad327 100644 --- a/tests/analysis/test_analysis_invariants.py +++ b/tests/analysis/test_analysis_invariants.py @@ -42,10 +42,10 @@ from tilefoundry.ir.core.op import Op from tilefoundry.ir.core.op_registry import iter_schemas from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor as TensorPattern from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice from tilefoundry.ir.hir.tensor.slice import Slice as SliceOp from tilefoundry.ir.isl_interop import index_set +from tilefoundry.ir.pattern import Tensor as TensorPattern from tilefoundry.ir.types import ( DType, Layout, diff --git a/tests/dsl/test_dsl_surface.py b/tests/dsl/test_dsl_surface.py index 3a52f39e..705c5e48 100644 --- a/tests/dsl/test_dsl_surface.py +++ b/tests/dsl/test_dsl_surface.py @@ -5,13 +5,13 @@ import pytest from tilefoundry import func -from tilefoundry.dsl import DimVar, DimVarRangePat, T, Tensor, tf +from tilefoundry.dsl import DimVar, RangePattern, T, Tensor, tf from tilefoundry.inspection import as_script from tilefoundry.ir.core.op_registry import _schemas_by_dialect_name from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor as TensorPat from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir.specialize import display_name +from tilefoundry.ir.pattern import Tensor as TensorPattern from tilefoundry.ir.types.dim import DimVar as IrDimVar @@ -35,8 +35,8 @@ def test_a_dialect_namespace_resolves_only_its_own_ops() -> None: @register_op(dialect="tf", category="math", name="my_add") class _MyAdd: - a = ParamDef(kind="input", pattern=TensorPat) - b = ParamDef(kind="input", pattern=TensorPat) + a = ParamDef(kind="input", pattern=TensorPattern) + b = ParamDef(kind="input", pattern=TensorPattern) def __init__(self, **kw): self.kw = kw @@ -61,12 +61,12 @@ def sub(x: Tensor[(_S,), "f32"]) -> Tensor[(_S,), "f32"]: pass -@sub.specialize(DimVarRangePat("S", 1, 2)) +@sub.specialize(RangePattern("S", 1, 2)) def narrow_s(x: Tensor[(_S,), "f32"]) -> Tensor[(_S,), "f32"]: return x -@sub.specialize(DimVarRangePat("S", 4, 6)) +@sub.specialize(RangePattern("S", 4, 6)) def wide_s(x: Tensor[(_S,), "f32"]) -> Tensor[(_S,), "f32"]: return x @@ -77,8 +77,8 @@ def test_func_specializations_parse_to_variants() -> None: variants = sub.variants assert len(variants) == 2 assert [v.name for v in variants] == ["sub", "sub"] - assert variants[0].specializations == (DimVarRangePat("S", 1, 2),) - assert variants[1].specializations == (DimVarRangePat("S", 4, 6),) + assert variants[0].specializations == (RangePattern("S", 1, 2),) + assert variants[1].specializations == (RangePattern("S", 4, 6),) assert display_name(variants[0]) == "narrow_s" assert display_name(variants[1]) == "wide_s" diff --git a/tests/evaluator/test_eval_core.py b/tests/evaluator/test_eval_core.py index 9dd1a099..b787cbc4 100644 --- a/tests/evaluator/test_eval_core.py +++ b/tests/evaluator/test_eval_core.py @@ -23,7 +23,7 @@ ScaledChild, ) from tilefoundry import func, module -from tilefoundry.dsl import DimVarRangePat, Tensor +from tilefoundry.dsl import RangePattern, Tensor from tilefoundry.dsl.tf import * # noqa: F401, F403 — bare op bindings for @func bodies from tilefoundry.evaluator import evaluate from tilefoundry.ir.core import Var @@ -118,7 +118,7 @@ def load(self, name: str): def subtree(self, name: str) -> "_Weights": prefix = f"{name}." return _Weights( - {k[len(prefix):]: v for k, v in self.values.items() if k.startswith(prefix)} + {k[len(prefix) :]: v for k, v in self.values.items() if k.startswith(prefix)} ) @@ -176,7 +176,7 @@ class _Dispatch: def dispatch(x: Tensor[(_N_EVAL,), "f32"]) -> Tensor[(_N_EVAL,), "f32"]: pass - @dispatch.specialize(DimVarRangePat("N_eval", 1, 7)) + @dispatch.specialize(RangePattern("N_eval", 1, 7)) def dynamic_variant(x: Tensor[(_N_EVAL,), "f32"]) -> Tensor[(_N_EVAL,), "f32"]: return scaled(x) # noqa: F821 diff --git a/tests/fixtures/placed/gqa_decode.py b/tests/fixtures/placed/gqa_decode.py index 5d1a110d..2f035b27 100644 --- a/tests/fixtures/placed/gqa_decode.py +++ b/tests/fixtures/placed/gqa_decode.py @@ -16,7 +16,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Tensor, tf # noqa: F401 — tf used by the @func body from tilefoundry.dsl.tf import * # noqa: F401, F403 — bare op names for the @func body -from tilefoundry.ir.core.pattern import DimVarRangePat +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.types import Broadcast, Layout, Mesh, ShardLayout, Topology from tilefoundry.ir.types.dim import DimVar @@ -81,7 +81,7 @@ def gqa_online_attend( pass - @gqa_online_attend.specialize(DimVarRangePat("ctx_len", 0, SMALL_CONTEXT_T)) + @gqa_online_attend.specialize(RangePattern("ctx_len", 0, SMALL_CONTEXT_T)) def head_on_cta( q: Tensor[(1, S, _HQ, _D), "bf16"], k_cache: Tensor[(1, C, _HKV, _D), "bf16", _CACHE_LAYOUT], @@ -213,7 +213,7 @@ def _ctx_combine( corr_n = tf.exp(score_n - m_all) return tf.cast((o * corr + corr_n * v_n) / (l_blk * corr + corr_n), dtype="bf16") - @gqa_online_attend.specialize(DimVarRangePat("ctx_len", SMALL_CONTEXT_T + 1, MAX_CTX)) + @gqa_online_attend.specialize(RangePattern("ctx_len", SMALL_CONTEXT_T + 1, MAX_CTX)) def ctx_split_kv( q: Tensor[(1, S, _HQ, _D), "bf16"], k_cache: Tensor[(1, C, _HKV, _D), "bf16", _CACHE_LAYOUT], diff --git a/tests/fixtures/placed/prefill_decode_attention.py b/tests/fixtures/placed/prefill_decode_attention.py index 4df99e90..c676d775 100644 --- a/tests/fixtures/placed/prefill_decode_attention.py +++ b/tests/fixtures/placed/prefill_decode_attention.py @@ -10,7 +10,7 @@ import math from tilefoundry import func, module -from tilefoundry.dsl import DimVar, DimVarRangePat, Mesh, Tensor, ceildiv, tf +from tilefoundry.dsl import DimVar, Mesh, RangePattern, Tensor, ceildiv, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 -- bare tile() in authored bodies from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget @@ -38,7 +38,7 @@ def attend( ) -> Tensor[(1, SEQ, HEADS, HEAD_DIM), "bf16"]: pass - @attend.specialize(DimVarRangePat("seq", 1, 1)) + @attend.specialize(RangePattern("seq", 1, 1)) def decode( q: Tensor[(1, SEQ, HEADS, HEAD_DIM), "bf16"], k_cache: Tensor[(1, CTX, HEADS, HEAD_DIM), "bf16"], @@ -112,7 +112,7 @@ def decode( "gmem", ) - @attend.specialize(DimVarRangePat("seq", 2, 4096)) + @attend.specialize(RangePattern("seq", 2, 4096)) def prefill( q: Tensor[(1, SEQ, HEADS, HEAD_DIM), "bf16"], k_cache: Tensor[(1, CTX, HEADS, HEAD_DIM), "bf16"], diff --git a/tests/fixtures/placed/qwen3_1_7b_pd.py b/tests/fixtures/placed/qwen3_1_7b_pd.py index cc40f818..93610412 100644 --- a/tests/fixtures/placed/qwen3_1_7b_pd.py +++ b/tests/fixtures/placed/qwen3_1_7b_pd.py @@ -14,7 +14,7 @@ from pathlib import Path from tilefoundry import func, module -from tilefoundry.dsl import ConstTensor, DimVar, DimVarRangePat, Mesh, Tensor, tf +from tilefoundry.dsl import ConstTensor, DimVar, Mesh, RangePattern, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401,F403 from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget @@ -419,7 +419,7 @@ def model( ) -> Tensor[(SEQ, V), "f32"]: pass - @model.specialize(DimVarRangePat("seq", 2, 8192)) # noqa: F821 + @model.specialize(RangePattern("seq", 2, 8192)) # noqa: F821 def prefill( ids: Tensor[(SEQ,), "i32"], w_embed: ConstTensor[(V, HID), "bf16"], @@ -507,7 +507,7 @@ def prefill( lg = tf.insert_slice(lg, tf.reshard(acc, (ROWS, BN), "gmem"), (m, n)) return lg - @model.specialize(DimVarRangePat("seq", 1, 1)) # noqa: F821 + @model.specialize(RangePattern("seq", 1, 1)) # noqa: F821 def decode( ids: Tensor[(SEQ,), "i32"], w_embed: ConstTensor[(V, HID), "bf16"], diff --git a/tests/fixtures/placed/specialize_through_call.py b/tests/fixtures/placed/specialize_through_call.py index afd30fc3..40511f83 100644 --- a/tests/fixtures/placed/specialize_through_call.py +++ b/tests/fixtures/placed/specialize_through_call.py @@ -1,7 +1,7 @@ """A dispatch on a callee: check and analyze both select its implementation.""" from tilefoundry import func, module -from tilefoundry.dsl import DimVar, DimVarRangePat, Mesh, Tensor, tf +from tilefoundry.dsl import DimVar, Mesh, RangePattern, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget @@ -21,18 +21,14 @@ class ToCallee: def pick(x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"]) -> Tensor[(1, D), "f32"]: pass - @pick.specialize(DimVarRangePat("n", 1, BOUND - 1)) - def pick_small( - x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"] - ) -> Tensor[(1, D), "f32"]: + @pick.specialize(RangePattern("n", 1, BOUND - 1)) + def pick_small(x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"]) -> Tensor[(1, D), "f32"]: with Mesh(("cta",), layout=(W,), names=("w",)) as m: xs = tf.reshard(x, (1, D @ m.w), "smem") return tf.reshard(xs + xs, (1, D), "gmem") - @pick.specialize(DimVarRangePat("n", BOUND, N_MAX)) - def pick_big( - x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"] - ) -> Tensor[(1, D), "f32"]: + @pick.specialize(RangePattern("n", BOUND, N_MAX)) + def pick_big(x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"]) -> Tensor[(1, D), "f32"]: with Mesh(("cta",), layout=(W,), names=("w",)) as m: xs = tf.reshard(x, (1, D @ m.w), "smem") return tf.reshard(xs + xs + xs, (1, D), "gmem") @@ -50,18 +46,14 @@ class Direct: def pick(x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"]) -> Tensor[(1, D), "f32"]: pass - @pick.specialize(DimVarRangePat("n", 1, BOUND - 1)) - def pick_small( - x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"] - ) -> Tensor[(1, D), "f32"]: + @pick.specialize(RangePattern("n", 1, BOUND - 1)) + def pick_small(x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"]) -> Tensor[(1, D), "f32"]: with Mesh(("cta",), layout=(W,), names=("w",)) as m: xs = tf.reshard(x, (1, D @ m.w), "smem") return tf.reshard(xs + xs, (1, D), "gmem") - @pick.specialize(DimVarRangePat("n", BOUND, N_MAX)) - def pick_big( - x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"] - ) -> Tensor[(1, D), "f32"]: + @pick.specialize(RangePattern("n", BOUND, N_MAX)) + def pick_big(x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"]) -> Tensor[(1, D), "f32"]: with Mesh(("cta",), layout=(W,), names=("w",)) as m: xs = tf.reshard(x, (1, D @ m.w), "smem") return tf.reshard(xs + xs + xs, (1, D), "gmem") diff --git a/tests/fixtures/tir/square.py b/tests/fixtures/tir/square.py index db78553c..b7ab2844 100644 --- a/tests/fixtures/tir/square.py +++ b/tests/fixtures/tir/square.py @@ -3,24 +3,28 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import DimVar, T, Tensor from tilefoundry.ir.core.kinds import BinaryKind -from tilefoundry.ir.core.pattern import DimVarRangePat +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget _S = DimVar("S", 1, 256) -@module(entry="square_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) +@module( + entry="square_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),) +) class TirSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def square_device(x: Tensor[(_S,), "f32"]): pass - @square_device.specialize(DimVarRangePat("S", 1, 127)) + @square_device.specialize(RangePattern("S", 1, 127)) def square_small(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as thread: + with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=("t",)) as thread: view = T.tensor_view(x, layout=((128 @ thread.t,), (1,))) - reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", ((128 @ thread.t,), (1,)), "rmem"]) + 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) @@ -29,11 +33,13 @@ def square_small(x: Tensor[(_S,), "f32"]): T.copy(reg, view) T.sync(thread) - @square_device.specialize(DimVarRangePat("S", 128, 255)) + @square_device.specialize(RangePattern("S", 128, 255)) def square_large(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as thread: + with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=("t",)) as thread: view = T.tensor_view(x, layout=((128 @ thread.t,), (1,))) - reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", ((128 @ thread.t,), (1,)), "rmem"]) + 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/inspection/test_module_tree_roundtrip.py b/tests/inspection/test_module_tree_roundtrip.py index 148ba4f4..ede221bf 100644 --- a/tests/inspection/test_module_tree_roundtrip.py +++ b/tests/inspection/test_module_tree_roundtrip.py @@ -20,8 +20,8 @@ from tilefoundry.dsl import ( # noqa: F401 ConstTensor, DimVar, - DimVarRangePat, Mesh, + RangePattern, Tensor, tf, ) @@ -73,9 +73,7 @@ def test_prefill_decode_specializations_survive_the_round_trip() -> None: assert len(variants) == 2 for variant in variants: targets = { - type(expr.target) - for expr in collect_exprs(variant.body) - if isinstance(expr, Call) + type(expr.target) for expr in collect_exprs(variant.body) if isinstance(expr, Call) } assert Arange in targets assert Where in targets @@ -265,9 +263,7 @@ def test_a_child_before_the_functions_naming_it() -> None: @module(entry="run") class _WeightedAtAnySize: @func - def run( - x: Tensor[(_N, 8), "f32"], w: ConstTensor[(8, 8), "f32"] - ) -> Tensor[(_N, 8), "f32"]: + def run(x: Tensor[(_N, 8), "f32"], w: ConstTensor[(8, 8), "f32"]) -> Tensor[(_N, 8), "f32"]: return tf.matmul(x, w) @@ -279,7 +275,7 @@ class _Dispatching: def dispatch(x: Tensor[(_N, 8), "f32"]) -> Tensor[(_N, 8), "f32"]: pass - @dispatch.specialize(DimVarRangePat("n_print", 1, 8)) + @dispatch.specialize(RangePattern("n_print", 1, 8)) def child_dispatch(x: Tensor[(_N, 8), "f32"]) -> Tensor[(_N, 8), "f32"]: return leaf(x) # noqa: F821 @@ -298,7 +294,7 @@ def test_a_child_call_in_a_specialization_body_survives_the_round_trip() -> None (imported_child,) = imported.modules (variant,) = imported.entry_function().variants assert imported.entry_function().body is None - assert variant.specializations == (DimVarRangePat("n_print", 1, 8),) + assert variant.specializations == (RangePattern("n_print", 1, 8),) assert variant.body.target is imported_child.entry_function() assert len(variant.body.args) == 1 assert [param.is_const for param in variant.body.target.params] == [False, True] diff --git a/tests/inspection/test_specialization_print.py b/tests/inspection/test_specialization_print.py index 5f8fc16a..468e2d16 100644 --- a/tests/inspection/test_specialization_print.py +++ b/tests/inspection/test_specialization_print.py @@ -9,8 +9,8 @@ from tilefoundry.inspection import as_script from tilefoundry.ir.core import Var -from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.hir.function import Function as HirFunction +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.types import make_tensor_type from tilefoundry.ir.types.dim import DimVar @@ -27,7 +27,7 @@ def _fn(*, body_is_self: bool, lo: int = 0, hi: int = 0) -> HirFunction: params=(x,), body=x if body_is_self else None, return_type=ty, - specializations=(DimVarRangePat("S", lo, hi),) if lo else (), + specializations=(RangePattern("S", lo, hi),) if lo else (), ) @@ -43,7 +43,7 @@ def test_prototype_prints_pass_base_and_specialize_blocks() -> None: The base is a pass-bodied prototype and each variant a ``.specialize`` block over a generated binding. The ``@module``-wrapped form must emit the same - ``DimVarRangePat`` import as the standalone form — module and standalone output + ``RangePattern`` import as the standalone form — module and standalone output share one header emitter, so a construct requiring an extra import in one mode requires it in both. """ @@ -54,10 +54,10 @@ def test_prototype_prints_pass_base_and_specialize_blocks() -> None: assert " pass" in src assert "def variant_S_1_2(" in src assert "def variant_S_4_6(" in src - assert '@main.specialize(DimVarRangePat("S", 1, 2))' in src - assert '@main.specialize(DimVarRangePat("S", 4, 6))' in src + assert '@main.specialize(RangePattern("S", 1, 2))' in src + assert '@main.specialize(RangePattern("S", 4, 6))' in src - assert "from tilefoundry.ir.core.pattern import DimVarRangePat" in src + assert "from tilefoundry.ir.pattern import RangePattern" in src compile(src, "", "exec") assert "@func\ndef main(" in standalone @@ -67,5 +67,5 @@ def test_prototype_prints_pass_base_and_specialize_blocks() -> None: def test_normal_function_omits_specialize() -> None: src = as_script(_fn(body_is_self=True)) assert ".specialize(" not in src - assert "DimVarRangePat" not in src + assert "RangePattern" not in src assert " pass" not in src diff --git a/tests/integration/models/deepseek_v4_flash/test_moe.py b/tests/integration/models/deepseek_v4_flash/test_moe.py index 74645532..406f6f21 100644 --- a/tests/integration/models/deepseek_v4_flash/test_moe.py +++ b/tests/integration/models/deepseek_v4_flash/test_moe.py @@ -2,7 +2,7 @@ from tests.models.deepseek_v4_flash.model import REAL, deepseek_v4_flash_module from tilefoundry.inspection import as_script -from tilefoundry.ir.constraints import LayoutConstraint, constraint_metadata +from tilefoundry.ir.clause import LayoutClause, clause_metadata from tilefoundry.ir.core import Call, Tuple from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.nn.matmul import MatMul @@ -53,8 +53,8 @@ def test_root_helpers_and_constraints_keep_real_model_contract() -> None: for call in _calls(deepseek_v4_flash_moe) if isinstance(call.target, Function) and call.target.name == "moe_topk" ) - routed = constraint_metadata(routed_call).constraints[0] - assert isinstance(routed, LayoutConstraint) + routed = clause_metadata(routed_call).constraints[0] + assert isinstance(routed, LayoutClause) assert repr(routed.layout.shape[0]) == "_" assert routed.layout.shape[1:] == (N_ACT, DIM) assert routed.bindings == (("cta", Split(1)),) @@ -65,8 +65,8 @@ def test_root_helpers_and_constraints_keep_real_model_contract() -> None: for call in _calls(deepseek_v4_flash_moe) if isinstance(call.target, Function) and call.target.name == "combine_expert_outputs" ) - combined = constraint_metadata(combined_call).constraints[0] - assert isinstance(combined, LayoutConstraint) + combined = clause_metadata(combined_call).constraints[0] + assert isinstance(combined, LayoutClause) assert combined.bindings == (("cta", Broadcast()),) diff --git a/tests/integration/test_dynamic_shape_dispatch.py b/tests/integration/test_dynamic_shape_dispatch.py index e86e04bf..af970f98 100644 --- a/tests/integration/test_dynamic_shape_dispatch.py +++ b/tests/integration/test_dynamic_shape_dispatch.py @@ -13,7 +13,7 @@ import tilefoundry from tilefoundry import func, module -from tilefoundry.dsl import DimVar, DimVarRangePat, Tensor +from tilefoundry.dsl import DimVar, RangePattern, Tensor from tilefoundry.dsl.tf import * # noqa: F401, F403 — binds bare ``mul`` / ``add`` from tilefoundry.target import CudaTarget @@ -26,11 +26,11 @@ class Dispatch: def main(x: Tensor[(_S,), "f32"]) -> Tensor[(_S,), "f32"]: pass - @main.specialize(DimVarRangePat("S", 1, 3)) + @main.specialize(RangePattern("S", 1, 3)) def small_shape(x: Tensor[(_S,), "f32"]) -> Tensor[(_S,), "f32"]: return mul(x, x) # noqa: F821 (bound via ``from tilefoundry.dsl.tf import *``) - @main.specialize(DimVarRangePat("S", 4, 7)) + @main.specialize(RangePattern("S", 4, 7)) def large_shape(x: Tensor[(_S,), "f32"]) -> Tensor[(_S,), "f32"]: return add(x, x) # noqa: F821 diff --git a/tests/ir/core/test_overload.py b/tests/ir/core/test_overload.py index 183400c5..8ee31be8 100644 --- a/tests/ir/core/test_overload.py +++ b/tests/ir/core/test_overload.py @@ -2,7 +2,6 @@ from __future__ import annotations -from dataclasses import dataclass from typing import Any import pytest @@ -10,18 +9,12 @@ from tilefoundry.ir.core.op_schema import OpSchema from tilefoundry.ir.core.overload import OverloadError, filter_candidates, resolve from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Scalar, Tensor, TensorPat +from tilefoundry.ir.pattern import Scalar, Tensor, TensorPattern +from tilefoundry.ir.types import TensorType - -@dataclass(frozen=True) -class _FakeType: - shape: tuple[int, ...] - dtype: str = "f32" - - -_S = _FakeType(shape=()) -_T1 = _FakeType(shape=(8,)) -_T2 = _FakeType(shape=(4, 8)) +_S = TensorType.umat_scalar() +_T1 = TensorType.umat_tensor((8,)) +_T2 = TensorType.umat_tensor((4, 8)) def _schema(name: str, *patterns: Any, defaults: tuple = ()) -> OpSchema: @@ -49,7 +42,7 @@ def __init__(self, **kw: Any) -> None: ... def test_resolve_picks_first_matching_candidate() -> None: """Arity + pattern filter; first-match wins; raises when no match.""" - rank2 = _schema("matmul", TensorPat(rank=2), TensorPat(rank=2)) + rank2 = _schema("matmul", TensorPattern(rank=2), TensorPattern(rank=2)) any_t = _schema("matmul", Tensor, Tensor) assert resolve([rank2, any_t], [_T2, _T2]) is rank2 diff --git a/tests/ir/core/test_param_def.py b/tests/ir/core/test_param_def.py index 63b90065..36fe6dc0 100644 --- a/tests/ir/core/test_param_def.py +++ b/tests/ir/core/test_param_def.py @@ -9,7 +9,7 @@ import pytest -from tilefoundry.ir.core.param_def import ParamDef +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef def test_paramdef_rejects_an_unknown_kind_and_keeps_required_independent() -> None: @@ -30,3 +30,16 @@ def test_paramdef_rejects_an_unknown_kind_and_keeps_required_independent() -> No omittable = ParamDef(kind="attribute", default=0) assert not omittable.is_required and omittable.has_default + + +def test_paramdef_effect_is_an_explicit_input_only_flag() -> None: + read_write = ParamDef(kind="input", effect=MemoryEffect.READ | MemoryEffect.WRITE) + assert read_write.effect == MemoryEffect.READ | MemoryEffect.WRITE + assert ParamDef(kind="input").effect is None + + with pytest.raises(TypeError, match="MemoryEffect"): + ParamDef(kind="input", effect="read") # type: ignore[arg-type] + with pytest.raises(ValueError, match="attributes"): + ParamDef(kind="attribute", effect=MemoryEffect.READ) + with pytest.raises(ValueError, match="READ, WRITE"): + ParamDef(kind="input", effect=MemoryEffect(0)) diff --git a/tests/ir/core/test_pattern.py b/tests/ir/core/test_pattern.py index 183526ef..0366f655 100644 --- a/tests/ir/core/test_pattern.py +++ b/tests/ir/core/test_pattern.py @@ -1,62 +1,59 @@ -"""Pattern — minimal contract.""" - -from __future__ import annotations - -from dataclasses import dataclass +"""Pattern — core scalar, tensor, composition, and range contracts.""" import pytest -from tilefoundry.ir.core.pattern import ( - AndPat, - DimVarRangePat, +from tilefoundry.ir.pattern import ( + AndPattern, + RangePattern, Scalar, Tensor, - TensorPat, + TensorPattern, ) +from tilefoundry.ir.types import DType, TensorType -@dataclass(frozen=True) -class FakeTy: - shape: tuple[int, ...] - dtype: str = "f32" +def _tensor(shape: tuple[int, ...], dtype: DType = DType.f32) -> TensorType: + return TensorType.umat_tensor(shape, dtype) def test_pattern_match_contract() -> None: """Singletons + parametric patterns + And combinator share one contract.""" - assert Scalar.match(FakeTy(shape=())) - assert not Scalar.match(FakeTy(shape=(3,))) - assert Tensor.match(FakeTy(shape=(3, 4))) - assert not Tensor.match(FakeTy(shape=())) - - rank2_bf16 = TensorPat(rank=2, dtype="bf16") - assert rank2_bf16.match(FakeTy(shape=(3, 4), dtype="bf16")) - assert not rank2_bf16.match(FakeTy(shape=(3,), dtype="bf16")) - assert not rank2_bf16.match(FakeTy(shape=(3, 4), dtype="f32")) - - combined = AndPat(parts=(TensorPat(rank=2), TensorPat(dtype="f16"))) - assert combined.match(FakeTy(shape=(3, 4), dtype="f16")) - assert not combined.match(FakeTy(shape=(3,), dtype="f16")) - assert AndPat(parts=()).match(FakeTy(shape=())) - - -def test_dim_var_range_pat_contract() -> None: - """Half-open ``[lo, hi)`` match semantics and the ``lo < hi`` rule. - - A single point is spelled ``[k, k+1)``; ``lo >= hi`` is an empty range - and rejected at construction. Non-int values (incl. ``bool``, which - subclasses ``int`` but is not a shape value) never match. - """ - p = DimVarRangePat("S", 1, 4) + assert Scalar.match(TensorType.umat_scalar()) + assert not Scalar.match(_tensor((3,))) + assert Tensor.match(_tensor((3, 4))) + assert not Tensor.match(TensorType.umat_scalar()) + assert not Tensor.match(type("FakeTy", (), {"shape": (3, 4)})()) + + rank2_bf16 = TensorPattern(rank=2, dtype=DType.bf16) + assert rank2_bf16.match(_tensor((3, 4), DType.bf16)) + assert not rank2_bf16.match(_tensor((3,), DType.bf16)) + assert not rank2_bf16.match(_tensor((3, 4), DType.f32)) + + combined = AndPattern(parts=(TensorPattern(rank=2), TensorPattern(dtype=DType.f16))) + assert combined.match(_tensor((3, 4), DType.f16)) + assert not combined.match(_tensor((3,), DType.f16)) + assert AndPattern(parts=()).match(TensorType.umat_scalar()) + + +def test_range_pattern_contract() -> None: + """Specialization ranges are closed and reject non-integer values.""" + p = RangePattern("S", 1, 4) assert p.match(1) and p.match(3) assert p.match(4) assert not p.match(0) assert not p.match(2.0) assert not p.match(True) - single = DimVarRangePat("S", 3, 4) - assert single.match(3) - assert not single.match(2) and single.match(4) - - assert DimVarRangePat("S", 4, 4).match(4) + closed = RangePattern("S", 3, 4) + assert closed.match(3) + assert not closed.match(2) and closed.match(4) + + assert RangePattern("S", 4, 4).match(4) + assert RangePattern(lo=3).match(3) and RangePattern(lo=3).match(30) + assert RangePattern(hi=3).match(-3) and RangePattern(hi=3).match(3) + with pytest.raises(ValueError, match="lower or upper"): + RangePattern() + with pytest.raises(ValueError, match="both lo and hi"): + RangePattern("S", lo=1) with pytest.raises(ValueError, match="lo <= hi"): - DimVarRangePat("S", 5, 4) + RangePattern("S", 5, 4) diff --git a/tests/ir/core/test_register_alias.py b/tests/ir/core/test_register_alias.py index b12b3c4f..957c987a 100644 --- a/tests/ir/core/test_register_alias.py +++ b/tests/ir/core/test_register_alias.py @@ -27,10 +27,10 @@ iter_schema_names, ) from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor as TensorPat from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.math.unary import Unary +from tilefoundry.ir.pattern import Tensor as TensorPattern from tilefoundry.ir.types import DType @@ -84,8 +84,8 @@ class _A(_DummyBase): @register_op(dialect="T", category="nn", name="testdup_relu") class _B(_DummyBase): - src = ParamDef(kind="input", pattern=TensorPat) - dst = ParamDef(kind="input", pattern=TensorPat) + src = ParamDef(kind="input", pattern=TensorPattern) + dst = ParamDef(kind="input", pattern=TensorPattern) bucket = get_schemas("T", "testdup_relu") assert [s.op_class for s in bucket] == [_A, _B] diff --git a/tests/ir/core/test_specialize.py b/tests/ir/core/test_specialize.py index eeceaf0d..49e08bff 100644 --- a/tests/ir/core/test_specialize.py +++ b/tests/ir/core/test_specialize.py @@ -13,7 +13,7 @@ ) from tests.fixtures.placed.specialize_through_call import ToCallee from tilefoundry import func, module -from tilefoundry.dsl import DimVarRangePat, Tensor, Topology, tf +from tilefoundry.dsl import RangePattern, Tensor, Topology, tf from tilefoundry.dsl.tf import * # noqa: F401,F403 -- names resolved dynamically from tilefoundry.evaluator import evaluate from tilefoundry.ir.hir.specialize import ( @@ -47,11 +47,11 @@ class _MissingCalleeDimension: def pick(x: Tensor[(_CALL_N,), "f32"]) -> Tensor[(_CALL_N,), "f32"]: pass - @pick.specialize(DimVarRangePat("call_n", 1, _DISPATCH_BOUND - 1)) + @pick.specialize(RangePattern("call_n", 1, _DISPATCH_BOUND - 1)) def pick_small(x: Tensor[(_CALL_N,), "f32"]) -> Tensor[(_CALL_N,), "f32"]: return tf.add(x, x) - @pick.specialize(DimVarRangePat("call_n", _DISPATCH_BOUND, _N_MAX)) + @pick.specialize(RangePattern("call_n", _DISPATCH_BOUND, _N_MAX)) def pick_big(x: Tensor[(_CALL_N,), "f32"]) -> Tensor[(_CALL_N,), "f32"]: return tf.add(tf.add(x, x), x) @@ -70,11 +70,11 @@ class _NestedDispatch: def inner(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: pass - @inner.specialize(DimVarRangePat("nested_n", 1, _DISPATCH_BOUND - 1)) + @inner.specialize(RangePattern("nested_n", 1, _DISPATCH_BOUND - 1)) def inner_small(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: return tf.add(x, x) - @inner.specialize(DimVarRangePat("nested_n", _DISPATCH_BOUND, _N_MAX)) + @inner.specialize(RangePattern("nested_n", _DISPATCH_BOUND, _N_MAX)) def inner_big(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: return tf.add(tf.add(x, x), x) @@ -82,11 +82,11 @@ def inner_big(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: def mid(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: pass - @mid.specialize(DimVarRangePat("nested_n", 1, _DISPATCH_BOUND - 1)) + @mid.specialize(RangePattern("nested_n", 1, _DISPATCH_BOUND - 1)) def mid_small(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: return inner(x) # noqa: F821 - @mid.specialize(DimVarRangePat("nested_n", _DISPATCH_BOUND, _N_MAX)) + @mid.specialize(RangePattern("nested_n", _DISPATCH_BOUND, _N_MAX)) def mid_big(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: return inner(x) # noqa: F821 diff --git a/tests/ir/pattern/test_mesh_pattern.py b/tests/ir/pattern/test_mesh_pattern.py new file mode 100644 index 00000000..23cc95be --- /dev/null +++ b/tests/ir/pattern/test_mesh_pattern.py @@ -0,0 +1,41 @@ +import pytest + +from tests.fixtures.meshes import CT, CTA +from tilefoundry.ir.pattern import ( + CapturePattern, + ComposedLayoutPattern, + LayoutPattern, + MeshPattern, + MultipleOfPattern, + WildcardPattern, +) + + +def test_mesh_pattern_matches_the_levels_it_names(): + with pytest.raises(ValueError, match="per_mode=True"): + MeshPattern( + ("thread",), + ComposedLayoutPattern( + outer=LayoutPattern(((128,),), ((1,),)) + ), + ) + + warpgroup = MeshPattern( + ("thread",), + ComposedLayoutPattern( + offset=CapturePattern("p0", MultipleOfPattern(128)), + outer=LayoutPattern(((128,),), ((1,),), per_mode=True), + ), + ) + assert warpgroup.match(CT[1:3, 128:256]).captures["p0"] == 128 + assert warpgroup.match(CT[1:3, 64:192]) is None + both = MeshPattern( + ("cta", "thread"), + ComposedLayoutPattern( + offset=WildcardPattern(), + outer=LayoutPattern(((2,), (128,)), ((1,), (1,)), per_mode=True), + ), + ) + assert both.match(CT[1:3, 128:256]) is not None + assert both.match(CT[0:1, 128:256]) is None + assert warpgroup.match(CTA) is None diff --git a/tests/ir/test_dim_var_envelope.py b/tests/ir/test_dim_var_envelope.py index d457a9d5..92075bdc 100644 --- a/tests/ir/test_dim_var_envelope.py +++ b/tests/ir/test_dim_var_envelope.py @@ -12,9 +12,9 @@ import pytest from tilefoundry.ir.core import Var, VerifyError -from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.hir.function import Function as HirFunction from tilefoundry.ir.hir.verify import verify_function +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.tensor_type import TupleType @@ -51,7 +51,7 @@ def test_a_signature_may_not_forge_its_dim_var_envelope() -> None: s = DimVar(name="S_env", lo=1, hi=8) forged = _identity_fn( params=(Var(type=_tensor((s,)), name="x"),), - specializations=(DimVarRangePat("S_env", 0, 99),), + specializations=(RangePattern("S_env", 0, 99),), ) with pytest.raises(VerifyError, match="not contained in DimVar envelope"): verify_function(forged) @@ -59,7 +59,7 @@ def test_a_signature_may_not_forge_its_dim_var_envelope() -> None: known = Var(type=_tensor((DimVar(name="S_known", lo=1, hi=8),)), name="x") with pytest.raises(VerifyError, match="references unknown DimVar"): verify_function( - _identity_fn(params=(known,), specializations=(DimVarRangePat("OTHER", 1, 3),)) + _identity_fn(params=(known,), specializations=(RangePattern("OTHER", 1, 3),)) ) r = DimVar(name="R_ret_only", lo=1, hi=8) @@ -68,7 +68,7 @@ def test_a_signature_may_not_forge_its_dim_var_envelope() -> None: _identity_fn( params=(Var(type=_tensor((4,)), name="x"),), return_type=_tensor((r,)), - specializations=(DimVarRangePat("R_ret_only", 1, 3),), + specializations=(RangePattern("R_ret_only", 1, 3),), ) ) @@ -103,7 +103,7 @@ def _dispatch_proto(name: str, env, ranges): params=(x,), body=x, return_type=ty, - specializations=(DimVarRangePat(name, lo, hi),), + specializations=(RangePattern(name, lo, hi),), ) ) return base diff --git a/tests/parser/test_calls.py b/tests/parser/test_calls.py index d78f8414..f1530585 100644 --- a/tests/parser/test_calls.py +++ b/tests/parser/test_calls.py @@ -10,7 +10,6 @@ from tilefoundry import func, module, prim_func from tilefoundry.dsl import Mesh, Tensor, tf from tilefoundry.ir.core import Call, Constant, Tuple, VerifyError -from tilefoundry.ir.core.pattern import Tensor as TensorPattern from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.hir.nn.matmul import MatMul @@ -20,6 +19,7 @@ from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.hir.tensor.stack import Stack from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem +from tilefoundry.ir.pattern import Tensor as TensorPattern from tilefoundry.ir.types import DType, Topology from tilefoundry.parser import ParseError from tilefoundry.target import CpuTarget, CudaTarget @@ -524,7 +524,9 @@ def run(x: Tensor[(8 @ mesh.b, 8), "f32"]): # noqa: F821 def test_placement_rejects_an_external_mesh_axis_binding() -> None: - with pytest.raises(ParseError, match="'_EXTERNAL_PLACEMENT_MESH' is not a lexical Mesh binding"): + with pytest.raises( + ParseError, match="'_EXTERNAL_PLACEMENT_MESH' is not a lexical Mesh binding" + ): @module( entry="run", diff --git a/tests/parser/test_functions.py b/tests/parser/test_functions.py index 0bc1cbef..c6f10f84 100644 --- a/tests/parser/test_functions.py +++ b/tests/parser/test_functions.py @@ -13,8 +13,8 @@ from tilefoundry.inspection import as_script from tilefoundry.ir.core import Call, SourceSpanMetadata, get_metadata from tilefoundry.ir.core.module import Module, subtree -from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.hir.function import Function +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.types import TupleType from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.storage import StorageKind @@ -110,7 +110,7 @@ def root( ) -> tuple[Tensor[(size,), "f32"], Tensor[(size,), "f32"]]: pass - @root.specialize(DimVarRangePat("prototype_size", 1, 8)) + @root.specialize(RangePattern("prototype_size", 1, 8)) def both( x: Tensor[(size,), "f32"], ) -> tuple[Tensor[(size,), "f32"], Tensor[(size,), "f32"]]: @@ -139,7 +139,7 @@ def root( ) -> tuple[Tensor[(size,), "f32"], Tensor[(size,), "f32"]]: pass - @root.specialize(DimVarRangePat("prototype_mismatch_size", 1, 8)) + @root.specialize(RangePattern("prototype_mismatch_size", 1, 8)) def scalar(x: Tensor[(size,), "f32"]): return x @@ -216,8 +216,7 @@ def test_every_parsed_call_knows_where_it_came_from(source: Path) -> None: ("tile(10)", "tile(extent) is not supported; use range(extent)"), ( "tile(1, 2, 3, 4)", - "tile() takes 2 or 3 arguments, (stop, step) or " - "(start, stop, step), got 4", + "tile() takes 2 or 3 arguments, (stop, step) or (start, stop, step), got 4", ), ("range(1, 2, 3, 4)", "range() takes 1 to 3 arguments, got 4"), ("steps(1, 2)", "loop iterator must be tile(...) or range(...)"), diff --git a/tests/passes/test_host_entry.py b/tests/passes/test_host_entry.py index 91257ae6..ceeb45b3 100644 --- a/tests/passes/test_host_entry.py +++ b/tests/passes/test_host_entry.py @@ -12,7 +12,7 @@ from tests.fixtures.tir.square import TirSquare from tilefoundry import module, prim_func from tilefoundry.dsl import DimVar, T, Tensor -from tilefoundry.ir.core.pattern import DimVarRangePat +from tilefoundry.ir.pattern import RangePattern from tilefoundry.ir.tir.launch import Launch from tilefoundry.ir.tir.stmts import Evaluate from tilefoundry.ir.types import Layout, Mesh, S, ShardLayout, Topology @@ -52,14 +52,14 @@ class _Prototype: def square(x: Tensor[(_S,), "f32"]): pass - @square.specialize(DimVarRangePat("S", 1, 127)) + @square.specialize(RangePattern("S", 1, 127)) def small(x: Tensor[(_S,), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: view = T.tensor_view(x, layout=_rows(128)) T.copy(view, view) T.sync(thread) - @square.specialize(DimVarRangePat("S", 128, 255)) + @square.specialize(RangePattern("S", 128, 255)) def large(x: Tensor[(_S,), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: view = T.tensor_view(x, layout=_rows(128)) From 4302ad9fd00addf05e2b380c851d15f4cccc1d4d Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Fri, 25 Sep 2026 22:42:18 +0800 Subject: [PATCH 4/8] fix(pattern): validate matched layout modes --- src/tilefoundry/ir/pattern/pattern.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index bb98d4f8..2bf4aa27 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -446,7 +446,7 @@ def match(self, subject, captures=None): arrangements = ( tuple( Layout(tuple(flatten(shape)), tuple(flatten(steps))) - for shape, steps in zip(self.shape, self.strides) + for shape, steps in zip(layout.shape, layout.strides) ) if self.per_mode else (Layout(extents, strides),) From 0b75688e65bbaeaaa7b9b5ac4598cf01769543dd Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Fri, 25 Sep 2026 23:14:58 +0800 Subject: [PATCH 5/8] refactor(tir): declare pointer-based memory operations --- docs/spec/codegen.md | 6 +- docs/spec/runtime.md | 8 + docs/spec/tir.md | 70 ++++++--- docs/spec/types.md | 23 ++- include/tilefoundry/runtime/cuda/ops/copy.cuh | 6 + src/tilefoundry/codegen/cuda/context.py | 12 ++ .../codegen/cuda/tir/memory/ptr_of.py | 7 +- .../codegen/cuda/tir/memory/tensor_view.py | 141 ++++++++++-------- .../codegen/cuda/tir/prim_function.py | 1 + src/tilefoundry/inspection/printer_base.py | 18 ++- src/tilefoundry/inspection/tir_printer.py | 135 +++++++++++++++-- src/tilefoundry/ir/pattern/__init__.py | 18 ++- src/tilefoundry/ir/pattern/utils.py | 78 +++++++++- src/tilefoundry/ir/tir/__init__.py | 2 + src/tilefoundry/ir/tir/cast.py | 42 ++++++ src/tilefoundry/ir/tir/memory/copy.py | 50 ++++++- src/tilefoundry/ir/tir/memory/fill.py | 15 +- src/tilefoundry/ir/tir/memory/ptr_of.py | 13 +- src/tilefoundry/ir/tir/memory/tensor_view.py | 77 +++++++--- src/tilefoundry/ir/types/__init__.py | 2 + src/tilefoundry/ir/types/pointer.py | 26 ++++ src/tilefoundry/ir/types/substitute.py | 14 +- src/tilefoundry/ir/types/tensor_type.py | 6 +- src/tilefoundry/parser/ast_pattern.py | 2 + src/tilefoundry/parser/pattern_nodes.py | 30 +++- tests/codegen/test_host_multi_launch.py | 8 +- tests/codegen/test_submodule_topology.py | 8 +- .../inspection/type_printer_sugar.printed.txt | 18 +-- tests/fixtures/placed/gpu_placed_rows.py | 4 +- tests/fixtures/tir/async_sync.py | 16 +- tests/fixtures/tir/mma.py | 51 +++++-- tests/fixtures/tir/rmsnorm.py | 16 +- tests/fixtures/tir/square.py | 4 +- tests/fixtures/tir/sync.py | 14 +- .../test_dynamic_cta_tir_handwritten.py | 3 +- tests/ops/tir/cuda/test_mma.py | 6 +- tests/ops/tir/cuda/test_swizzle.py | 17 +-- tests/ops/tir/cuda/test_tma.py | 4 +- tests/ops/tir/test_copy.py | 20 +-- tests/ops/tir/test_dot.py | 27 +++- tests/ops/tir/test_elementwise.py | 33 ++-- tests/ops/tir/test_reduce.py | 56 +++++-- tests/ops/tir/test_sync.py | 8 +- tests/passes/test_host_entry.py | 6 +- 44 files changed, 851 insertions(+), 270 deletions(-) create mode 100644 src/tilefoundry/ir/tir/cast.py create mode 100644 src/tilefoundry/ir/types/pointer.py diff --git a/docs/spec/codegen.md b/docs/spec/codegen.md index 2eceb7bb..deaf194b 100644 --- a/docs/spec/codegen.md +++ b/docs/spec/codegen.md @@ -144,6 +144,8 @@ class CudaCodegenContext(CodegenContext): def bind_extents(self, params) -> None: ... def reset_barrier_ids(self) -> None: ... def alloc_barrier_id(self) -> int: ... + def reset_smem_base(self) -> None: ... + def smem_base(self) -> str: ... def dtype_to_cpp(self, dtype_name: str) -> str: ... @@ -181,6 +183,9 @@ class CpuCodegenContext(CodegenContext): types alone. `launches` is the geometry each device function is called at, keyed by `id(fn)`, settled where the `Launch` was written ([passes §7.3](./passes.md#73-insert_default_host_entry)). + - `smem_base` declares and returns one byte-addressed dynamic shared-memory + base per kernel. A numeric `TensorView` address adds its byte offset before + converting the result to a CuTe shared-memory pointer. - A target subclass owns the type strings and hardware counters only it can state; a handler MUST reach them through the context rather than reading the IR for them. Other helpers MAY be added per target. @@ -440,4 +445,3 @@ variant runs is decided on the device. the call already carries -- the parameter the open axis expands into. A shape outside every variant's range is a call-contract violation and the kernel traps. - diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index e46738dc..20346a54 100644 --- a/docs/spec/runtime.md +++ b/docs/spec/runtime.md @@ -1259,6 +1259,9 @@ helpers genuinely capable of host compilation (e.g. `local_tensor`, ```cpp // include/tilefoundry/runtime/cuda/ops/copy.cuh +template +CUTE_HOST_DEVICE auto tensor_view(TPointer pointer, TLayout layout); + template __device__ void copy(TSrc const &src, TDst &dst); @@ -1267,6 +1270,11 @@ __device__ void copy_async(TSrc const &src, TDst &dst); ``` +**`tensor_view`.** Rebuilds the CuTe tensor engine that a TIR `TensorView` +describes from its typed pointer and emitted layout. Residency stays on the +pointer engine; a later `make_shard_tensor` adds the logical global and shard +layouts without recovering the source tensor. + **`copy`.** One entry, taking tensors, with the transfer shape, the strides, the element diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 1565cf45..7509fc63 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -8,7 +8,7 @@ the work, structural Stmts carry control flow. `body` is a `Sequential`; the function returns no value. - **Stmt tree**: function bodies are nested Stmts only. Exprs appear inside Stmt fields (e.g. `LetStmt.value`, `For.start`). -- **Effect Ops** (`Copy`, `Fill`, `Mma`, `ReLU`, `RMSNorm`, `Reduce`) +- **Effect Ops** (`Copy`, `Fill`, `Cast`, `Mma`, `ReLU`, `RMSNorm`, `Reduce`) are value-class Ops registered with `@register_op`; in Stmt position they are invoked as `Evaluate(op, args)` ([§1.4](#14-evaluate)). @@ -497,40 +497,53 @@ class PtrOf(Op): """Value form; take the device address of a tensor. Attributes: - x: input; the tensor whose device address is taken. + tensor: input; the tensor whose device address is taken. """ - x: Tensor + tensor: Tensor ``` -- constraints: [] +- constraints: + - returns `PointerType(tensor.dtype, tensor.storage)`; it does not preserve + the tensor's shape or layout in the pointer type. ##### TensorView ```python class TensorView(Op): - """Value form; derive a sub-view of a tensor. + """Value form; construct a logical tensor over a typed pointer. Attributes: - memory: input; the base tensor (may be a ``PtrOf`` result). + pointer: input; a ``PointerType`` value or an smem byte offset. coordinates: optional trailing inputs; one absolute element start per logical window axis (or one absolute flat start for a rank-1 view). - layout: attribute; the sub-view descriptor — a plain ``Layout`` or a - ``ShardLayout`` placed over ``memory``. - shape: attribute; optional logical-shape override (reshape). + dtype: optional attribute; element type stated for a numeric address. + storage: optional attribute; storage stated for a numeric address. + layout: attribute; the view descriptor. + shape: attribute; logical shape, optionally inherited from ``PtrOf``. """ - memory: Tensor + pointer: object + dtype: str | None = None + storage: StorageKind | None = None layout: object shape: tuple | None = None ``` - constraints: + - An explicit `shape` is authoritative. If omitted, it is inherited only + when the input syntax is exactly `T.ptr_of(tensor)`, from that tensor's + `TensorType.shape`. No other pointer provenance and no layout is inspected + to guess a shape. + - An integer input is a byte offset from the kernel's dynamic shared-memory + base. It MUST state `dtype`, `storage="smem"`, and `shape`; booleans and + non-integer numeric addresses are invalid. + - A `PointerType` input MAY restate `dtype` or `storage`, but any stated value + MUST equal the pointer descriptor. - With trailing coordinates, codegen derives the view at those absolute element starts. A coordinate is not a tile ordinal and MUST NOT be multiplied by the view extent. - The coordinate count MUST match the logical window rank before any shard-owned layout axes are removed locally. - - `memory` MAY be an allocated `ShardTensor`; in that case the view - reprojects the same storage with a new shard layout, using the tensor's - engine rather than its existing shard layout. + - `T.ptr_of` MAY point at an allocated `ShardTensor`; the view rebuilds over + the engine pointer rather than reusing the existing shard layout. ##### Copy ```python @@ -545,7 +558,12 @@ class Copy(Op): src: Tensor dst: Tensor ``` -- constraints: [] +- constraints: + - `src` declares `READ`; `dst` declares `WRITE`. + - both operands are whole-byte tensors in gmem, smem, or rmem, with equal + dtype. Their storages MAY be equal; same-storage copy is still a byte move. + - `scope` optionally states any non-empty run of threads. `rmem_layout` and + `smem_layout` optionally state the author's landing arrangements. ##### Fill ```python @@ -560,7 +578,25 @@ class Fill(Op): tensor: Tensor value: Tensor ``` -- constraints: [] +- constraints: + - `tensor` declares `WRITE`; `value` declares `READ` and MUST be scalar. + - a nonconstant value's dtype MUST equal the destination dtype. Constant zero + is convertible and MAY use the parser's default scalar dtype. + - `scope` optionally states any non-empty run of threads. + +##### Cast +```python +class Cast(Op): + """Effect form; convert a register tile to another dtype.""" + + src: Tensor + dst: Tensor + scope: Mesh | None = None +``` +- constraints: + - `src` declares `READ`; `dst` declares `WRITE`; both are rmem tensors. + - the operands have equal shapes and distinct dtypes. + - `scope` optionally states any non-empty run of threads. #### NN Ops (`tir.nn.*`) @@ -964,10 +1000,10 @@ atom = T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN) with Mesh((Topology("thread", 32),), Layout(shape=(4, 8), strides=(1, 4))) as warp: a_frag = T.alloc_tensor(TensorType(..., layout=atom.A, storage=rmem)) acc = T.alloc_tensor(TensorType(..., layout=atom.C, storage=rmem)) - T.copy(T.tensor_view(a, layout=atom.A), a_frag) # load + T.copy(T.tensor_view(T.ptr_of(a), layout=atom.A), a_frag) # load T.fill(acc, 0.0) T.mma(acc, a_frag, b_frag, atom=atom) # compute - T.copy(acc, T.tensor_view(c, layout=atom.C)) # store + T.copy(acc, T.tensor_view(T.ptr_of(c), layout=atom.C)) # store ``` - The author allocates each register fragment with the matching diff --git a/docs/spec/types.md b/docs/spec/types.md index 6b7d8c13..31961872 100644 --- a/docs/spec/types.md +++ b/docs/spec/types.md @@ -7,6 +7,7 @@ flowchart TB TupleType["TupleType"] UnitType["UnitType"] CallableType["CallableType"] + PointerType["PointerType"] DType["DType"] dim["dim ops"] @@ -16,10 +17,12 @@ flowchart TB TupleType -. member of .-> Type UnitType -. member of .-> Type CallableType -. member of .-> Type + PointerType -. member of .-> Type DType -. dtype .-> TensorType dim -. shape elements .-> TensorType Layout -. layout .-> TensorType + DType -. dtype .-> PointerType TensorType -. element of .-> TupleType Type -. return type and parameter types .-> CallableType @@ -28,11 +31,29 @@ flowchart TB ## 1. `Type` ```python -Type = TensorType | TupleType | UnitType | CallableType +Type = TensorType | TupleType | UnitType | CallableType | PointerType ``` --- +### 1.1 `PointerType` + +```python +class PointerType: + """Typed physical-memory engine consumed by ``T.tensor_view``.""" + + dtype: DType + storage: StorageKind +``` + +- constraints: + - `dtype` is the element type read through the pointer. + - `storage` is concrete and normalized to `StorageKind`; `None` is invalid. + - A pointer carries no logical shape or layout. `TensorView` supplies those + facts when it reconstructs a tensor value. + +--- + ## 2. `TensorType` ```python diff --git a/include/tilefoundry/runtime/cuda/ops/copy.cuh b/include/tilefoundry/runtime/cuda/ops/copy.cuh index 5903af75..bc598c24 100644 --- a/include/tilefoundry/runtime/cuda/ops/copy.cuh +++ b/include/tilefoundry/runtime/cuda/ops/copy.cuh @@ -4,6 +4,12 @@ #include "copy/copy_impl.h" +/// Rebuild a CuTe tensor view from the typed pointer carried by TIR. +template +CUTE_HOST_DEVICE auto tensor_view(TPointer pointer, TLayout layout) { + return cute::make_tensor(pointer, layout); +} + /// Copy one projected slice to another. template __device__ void copy(TSrc const &src, TDst &dst) { diff --git a/src/tilefoundry/codegen/cuda/context.py b/src/tilefoundry/codegen/cuda/context.py index 7ad734de..52f27641 100644 --- a/src/tilefoundry/codegen/cuda/context.py +++ b/src/tilefoundry/codegen/cuda/context.py @@ -75,6 +75,18 @@ def __init__( self._next_barrier_id = 1 self.needs_grid_barrier_state = False """Set while emitting a grid barrier, which the module declares state for.""" + self._has_smem_base = False + + def reset_smem_base(self) -> None: + """Start a kernel with no dynamic shared-memory base declaration.""" + self._has_smem_base = False + + def smem_base(self) -> str: + """Return the byte-addressed dynamic shared-memory base, declaring it once.""" + if not self._has_smem_base: + self.emit("extern __shared__ unsigned char smem_base[];") + self._has_smem_base = True + return "smem_base" def bind_extents(self, params: Iterable[TensorSignature]) -> None: """Say where the kernel reads the extent of every dimension its types leave open. diff --git a/src/tilefoundry/codegen/cuda/tir/memory/ptr_of.py b/src/tilefoundry/codegen/cuda/tir/memory/ptr_of.py index 3384b265..b92a45be 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/ptr_of.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/ptr_of.py @@ -18,8 +18,7 @@ def _emit(let_stmt, ctx: CudaCodegenContext) -> None: call = let_stmt.value src = call.args[0] src_name = ctx.name_for(src) + if ctx.is_kernel_param(src): + src_name += "_tensor" var_name = ctx.name_for(let_stmt.var) - - - - ctx.emit(f"auto {var_name} = {src_name}.engine.data();") + ctx.emit(f"auto {var_name} = {src_name}.data();") diff --git a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py index 8691ce96..f92816b2 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py @@ -20,6 +20,7 @@ ) from tilefoundry.ir.core import Call, Constant from tilefoundry.ir.core.kinds import ReduceKind +from tilefoundry.ir.tir.memory.ptr_of import PtrOf from tilefoundry.ir.tir.memory.tensor_view import TensorView from tilefoundry.ir.tir.stmts import LetStmt from tilefoundry.ir.tir.sync import participation @@ -66,9 +67,7 @@ def _render_layout_type(layout: LayoutBase) -> str: shape_args = ", ".join(f"cute::Int<{s}>" for s in layout.shape) stride_args = ", ".join(f"cute::Int<{s}>" for s in layout.strides) return f"cute::Layout, cute::Stride<{stride_args}>>" - raise NotImplementedError( - f"tensor_view: no CuTe layout type for {type(layout).__name__}" - ) + raise NotImplementedError(f"tensor_view: no CuTe layout type for {type(layout).__name__}") def _render_layout_value(layout: LayoutBase, dim, stride) -> str: @@ -89,12 +88,9 @@ def _render_layout_value(layout: LayoutBase, dim, stride) -> str: shape_args = ", ".join(dim(d) for d in layout.shape) stride_args = ", ".join(stride(s) for s in layout.strides) return ( - f"cute::make_layout(cute::make_shape({shape_args}), " - f"cute::make_stride({stride_args}))" + f"cute::make_layout(cute::make_shape({shape_args}), cute::make_stride({stride_args}))" ) - raise NotImplementedError( - f"tensor_view: no CuTe layout value for {type(layout).__name__}" - ) + raise NotImplementedError(f"tensor_view: no CuTe layout value for {type(layout).__name__}") def _scope_mesh_value(mesh, ctx) -> "str | None": @@ -104,9 +100,7 @@ def _scope_mesh_value(mesh, ctx) -> "str | None": entry = ctx._mesh_aliases.get(id(mesh)) if entry is None: inline = mesh_type(mesh) - entry = next( - (e for e in ctx._mesh_aliases.values() if e[1] == inline), None - ) + entry = next((e for e in ctx._mesh_aliases.values() if e[1] == inline), None) if entry is None: return None alias = entry[0] @@ -197,9 +191,7 @@ def register_strides(sl: SL) -> tuple[int, ...]: return tuple(strides) -def render_shard_layout_value( - var_name: str, sl: SL, dynamic_extents=None, storage=None, ctx=None -): +def render_shard_layout_value(var_name: str, sl: SL, dynamic_extents=None, storage=None, ctx=None): """Render a shard layout as runtime C++ preamble and value expression. Static values retain the type produced by the type renderer. Runtime @@ -353,14 +345,49 @@ def _coord_ref(index_var, ctx: CudaCodegenContext) -> str: return _CoordinateVisitor().visit(index_var, ctx) +def _tensor_ref(var, ctx: CudaCodegenContext) -> str: + name = ctx.name_for(var) + return f"{name}_tensor" if ctx.is_kernel_param(var) else name + + +def _pointer_ref(pointer, ctx: CudaCodegenContext) -> tuple[str, object | None]: + """Render a TensorView pointer and return its syntactic source tensor, if any.""" + if isinstance(pointer, Call) and isinstance(pointer.target, PtrOf): + source = pointer.args[0] + return f"{_tensor_ref(source, ctx)}.data()", source + if isinstance(pointer, Constant): + cpp_type = ctx.dtype_to_cpp(pointer.type.dtype.name) + base = ctx.smem_base() + return ( + f"cute::make_smem_ptr(reinterpret_cast<{cpp_type} *>({base} + {pointer.value}))", + None, + ) + return ctx.name_for(pointer), None + + +def _plain_layout_value(layout: LayoutBase) -> str: + def dim(value): + return f"cute::Int<{int(upper_bound(value))}>{{}}" + + def stride(value): + return f"cute::Int<{int(value)}>{{}}" + + return _render_layout_value(layout, dim, stride) + + @register_codegen(CudaTarget, Role.EMIT, TensorView) def _emit(let: LetStmt, ctx: CudaCodegenContext) -> None: call = let.value - memory_var = call.args[0] + pointer = call.args[0] + pointer_ref, memory_var = _pointer_ref(pointer, ctx) var_name = ctx.name_for(let.var) layout = call.target.layout if len(call.args) > 1: + if memory_var is None: + raise NotImplementedError( + "tensor_view coordinates require a syntactic T.ptr_of(tensor) source" + ) mem_name = ctx.name_for(memory_var) if len(call.args) > 2: @@ -456,54 +483,40 @@ def _emit(let: LetStmt, ctx: CudaCodegenContext) -> None: return if isinstance(layout, SL): - mem_name = ctx.name_for(memory_var) - - if ctx.is_kernel_param(memory_var): - tensor_ref = f"{mem_name}_tensor" - global_total = shape_numel_upper_bound(memory_var.type.shape) - global_layout = f"cute::make_layout(cute::Shape>{{}})" - preamble, shard_value = render_shard_layout_value( - var_name, - layout, - ctx.dynamic_extents, - getattr(let.var.type, "storage", None), - ctx, - ) - for line in preamble: - ctx.emit(line) - ctx.emit( - f"auto {var_name} = tilefoundry::make_shard_tensor(" - f"{tensor_ref}, {global_layout}, {shard_value});" - ) - else: - source_layout = getattr(memory_var.type, "layout", None) - if isinstance(source_layout, SL): - ctx.emit(f"auto {var_name}_tensor = {mem_name}.engine;") + target_total = shape_numel_upper_bound(let.var.type.shape) + target_global = f"cute::make_layout(cute::Shape>{{}})" + if memory_var is not None and not ctx.is_kernel_param(memory_var): + local_shape = shard_layout_local_shape(layout) + local_shape = tuple(s for s in local_shape if s != 1) or (1,) + if len(local_shape) > 1: + shape_args = ", ".join(f"cute::Int<{int(s)}>" for s in local_shape) + engine_layout = f"cute::make_layout(cute::Shape<{shape_args}>{{}})" else: - local_shape = shard_layout_local_shape(layout) - local_shape = tuple(s for s in local_shape if s != 1) or (1,) - if len(local_shape) > 1: - shape_args = ", ".join(f"cute::Int<{int(s)}>" for s in local_shape) - tensor_layout = f"cute::make_layout(cute::Shape<{shape_args}>{{}})" - else: - tensor_layout = ( - f"cute::make_layout(cute::Shape>{{}})" - ) - ctx.emit( - f"auto {var_name}_tensor = cute::make_tensor({mem_name}, {tensor_layout});" + engine_layout = ( + f"cute::make_layout(cute::Shape>{{}})" ) - target_total = shape_numel_upper_bound(let.var.type.shape) - target_global = f"cute::make_layout(cute::Shape>{{}})" - preamble, shard_value = render_shard_layout_value( - var_name, - layout, - ctx.dynamic_extents, - getattr(let.var.type, "storage", None), - ctx, - ) - for line in preamble: - ctx.emit(line) - ctx.emit( - f"auto {var_name} = tilefoundry::make_shard_tensor(" - f"{var_name}_tensor, {target_global}, {shard_value});" - ) + else: + engine_layout = target_global + ctx.emit( + f"auto {var_name}_tensor = " + f"tilefoundry::ops::tensor_view({pointer_ref}, {engine_layout});" + ) + preamble, shard_value = render_shard_layout_value( + var_name, + layout, + ctx.dynamic_extents, + getattr(let.var.type, "storage", None), + ctx, + ) + for line in preamble: + ctx.emit(line) + ctx.emit( + f"auto {var_name} = tilefoundry::make_shard_tensor(" + f"{var_name}_tensor, {target_global}, {shard_value});" + ) + return + + ctx.emit( + f"auto {var_name} = tilefoundry::ops::tensor_view(" + f"{pointer_ref}, {_plain_layout_value(layout)});" + ) diff --git a/src/tilefoundry/codegen/cuda/tir/prim_function.py b/src/tilefoundry/codegen/cuda/tir/prim_function.py index a49053c7..b8cbc317 100644 --- a/src/tilefoundry/codegen/cuda/tir/prim_function.py +++ b/src/tilefoundry/codegen/cuda/tir/prim_function.py @@ -93,6 +93,7 @@ def _subject(fn: PrimFunction, pattern: RangePattern, ctx: CudaCodegenContext) - @register_codegen(CudaTarget, Role.EMIT, PrimFunction) def _emit(fn: PrimFunction, ctx: CudaCodegenContext) -> None: """Write what runs inside one ``__global__``: its buffers, then its statements.""" + ctx.reset_smem_base() if fn.variants: _dispatch(fn, ctx) return diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 1bdb7c6a..adb0458c 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -3,6 +3,7 @@ from __future__ import annotations import enum +import json from contextlib import contextmanager from tilefoundry.ir.core import Call, Constant, Tuple, Var @@ -10,7 +11,7 @@ from tilefoundry.ir.mesh_scope import device_layout from tilefoundry.ir.pattern import Pattern, RangePattern from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom -from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType +from tilefoundry.ir.types import DType, PointerType, TensorType, TupleType, UnitType from tilefoundry.ir.types.dim import ( DimAdd, DimConst, @@ -299,6 +300,16 @@ def visit_TensorType(self, value: TensorType, ctx=None) -> str: def visit_TupleType(self, value: TupleType, ctx=None) -> str: return f"Tuple[{', '.join(self.visit(field, ctx) for field in value.fields)}]" + def visit_PointerType(self, value: PointerType, ctx=None) -> str: + if ctx is not None: + ctx.use( + PythonExpr( + ("from tilefoundry.ir.types import DType, PointerType, StorageKind",), + "", + ) + ) + return f"PointerType(DType.{value.dtype.name}, StorageKind.{value.storage.name})" + def visit_UnitType(self, value: UnitType, ctx=None) -> str: return "None" @@ -326,7 +337,8 @@ def visit_Mesh(self, value: Mesh, ctx=None) -> str: ) result = f"Mesh({topologies}, {self.visit(written, ctx)}" if value.names: - result += f", names={tuple(value.names)!r}" + names = ", ".join(json.dumps(name) for name in value.names) + result += f", names=({names}{',' if len(value.names) == 1 else ''})" return result + ")" def visit_NoneType(self, value: None, ctx=None) -> str: @@ -399,7 +411,7 @@ def atom_reference(self, value: MmaAtom, ctx=None) -> str: def render_value(self, value, ctx=None, indent: str = "") -> str: """Render a non-expression attribute through the same visitor when possible.""" - if isinstance(value, (TensorType, Mesh, LayoutBase, DType)): + if isinstance(value, (TensorType, PointerType, Mesh, LayoutBase, DType)): with self.type_surface(indent=indent): return self.visit(value, ctx) if isinstance(value, MmaAtom): diff --git a/src/tilefoundry/inspection/tir_printer.py b/src/tilefoundry/inspection/tir_printer.py index afced9cf..c3f5ea9a 100644 --- a/src/tilefoundry/inspection/tir_printer.py +++ b/src/tilefoundry/inspection/tir_printer.py @@ -12,17 +12,51 @@ from tilefoundry.ir.hir.function import Function as HirFunction from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.tir.launch import Launch +from tilefoundry.ir.tir.memory import AllocTensor from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.tir.stmts import ( Evaluate, ) from tilefoundry.ir.tir.symbol_ref import SymbolRef -from tilefoundry.ir.types import DType, TensorType +from tilefoundry.ir.types import DType, PointerType, TensorType from tilefoundry.ir.types.dim import is_dim_op_call from tilefoundry.ir.visitor import StmtVisitor from tilefoundry.utils.python_source import PythonExpr, _merge_imports +_LINE_LENGTH = 100 + + +def _split_top_level(text: str) -> list[str]: + """Split comma-separated Python fragments without cutting nested forms.""" + parts: list[str] = [] + start = 0 + depth = 0 + quote = None + escaped = False + for index, char in enumerate(text): + if quote is not None: + if escaped: + escaped = False + elif char == "\\": + escaped = True + elif char == quote: + quote = None + continue + if char in {'"', "'"}: + quote = char + elif char in "([{": + depth += 1 + elif char in ")]}": + depth -= 1 + elif char == "," and depth == 0: + parts.append(text[start:index].strip()) + start = index + 1 + tail = text[start:].strip() + if tail: + parts.append(tail) + return parts + class _RenderedLines(list): def __init__(self, lines, imports): @@ -119,13 +153,70 @@ def visit_Sequential(self, stmt, ctx=None): return [line for child in stmt.body for line in self.visit(child)] def visit_LetStmt(self, stmt, ctx=None): - return [f"{self.indent}{stmt.var.name} = {self.visit(stmt.value)}"] + self.visit(stmt.body) + rendered = self.visit(stmt.value) + line = f"{self.indent}{stmt.var.name} = {rendered}" + if len(line) + 4 <= _LINE_LENGTH or not isinstance(stmt.value, Call): + return [line] + self.visit(stmt.body) + if isinstance(stmt.value.target, AllocTensor): + lines = self._wrapped_alloc(stmt) + else: + lines = self._wrapped_call(f"{stmt.var.name} = ", rendered) + return lines + self.visit(stmt.body) + + def _wrapped_call(self, prefix: str, rendered: str) -> list[str]: + head, separator, arguments = rendered.partition("(") + if not separator or not arguments.endswith(")"): + return [f"{self.indent}{prefix}{rendered}"] + content = arguments[:-1] + continuation = f"{self.indent} {content}" + if len(continuation) + 4 <= _LINE_LENGTH: + middle = [continuation] + else: + middle = [f"{self.indent} {part}," for part in _split_top_level(content)] + return [f"{self.indent}{prefix}{head}(", *middle, f"{self.indent})"] + + def _wrapped_alloc(self, stmt) -> list[str]: + target = stmt.value.target + tensor = self.render_value(target.tensor_type, self.context) + argument = f"tensor_type={tensor}" + continuation = f"{self.indent} {argument}" + if len(continuation) + 4 <= _LINE_LENGTH: + middle = [continuation] + else: + assert tensor.startswith("Tensor[") and tensor.endswith("]") + fields = _split_top_level(tensor[len("Tensor[") : -1]) + packed = ", ".join(fields) + field_indent = f"{self.indent} " + if len(field_indent + packed) + 4 <= _LINE_LENGTH: + field_lines = [field_indent + packed] + else: + field_lines = [field_indent + field + "," for field in fields] + middle = [ + f"{self.indent} tensor_type=Tensor[", + *field_lines, + f"{self.indent} ]", + ] + return [ + f"{self.indent}{stmt.var.name} = T.alloc_tensor(", + *middle, + f"{self.indent})", + ] def visit_Evaluate(self, stmt, ctx=None): return self._emit_evaluate(stmt) def visit_MeshScope(self, stmt, ctx=None): - lines = [f"{self.indent}with {self.visit(stmt.mesh, self.context)} as {stmt.binding.name}:"] + rendered = self.visit(stmt.mesh, self.context) + line = f"{self.indent}with {rendered} as {stmt.binding.name}:" + if len(line) + 4 <= _LINE_LENGTH: + lines = [line] + else: + content = rendered.removeprefix("Mesh(").removesuffix(")") + lines = [ + f"{self.indent}with Mesh(", + f"{self.indent} {content}", + f"{self.indent}) 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() @@ -223,7 +314,16 @@ def _print_op_evaluate(stmt: Evaluate, printer: TirPrinter) -> list[str]: 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.visit(arg) for arg in args] - return [f"{indent}{printer.visit_Op(target)}({', '.join(rendered_args + attrs)})"] + head = printer.visit_Op(target) + arguments = rendered_args + attrs + line = f"{indent}{head}({', '.join(arguments)})" + if len(line) + 4 <= _LINE_LENGTH: + return [line] + return [ + f"{indent}{head}(", + *(f"{indent} {argument}," for argument in arguments), + f"{indent})", + ] def _function_block(fn: PrimFunction) -> list[str]: @@ -243,10 +343,14 @@ 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).visit(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, PointerType)) else repr(p.type)}" for p in fn.params ) - lines.append(f"def {_binding_name(fn.name)}({params}):") + definition = f"def {_binding_name(fn.name)}({params}):" + if len(definition) + 4 <= _LINE_LENGTH: + lines.append(definition) + else: + lines.extend((f"def {_binding_name(fn.name)}(", f" {params}", "):")) body = TirPrinter(context=ctx, indent=" ").visit(fn.body) lines.extend(body or [" pass"]) if fn.variants: @@ -257,9 +361,12 @@ def _function_block(fn: PrimFunction) -> list[str]: lines.append( f"@{_binding_name(fn.name)}.specialize({TirPrinter(context=ctx).render_pattern(pat, ctx)})" ) - lines.append( - f"def {_binding_name(getattr(variant, '_display_name', variant.name))}({params}):" - ) + variant_name = _binding_name(getattr(variant, "_display_name", variant.name)) + definition = f"def {variant_name}({params}):" + if len(definition) + 4 <= _LINE_LENGTH: + lines.append(definition) + else: + lines.extend((f"def {variant_name}(", f" {params}", "):")) vbody = TirPrinter(context=ctx, indent=" ").visit(variant.body) lines.extend(vbody or [" pass"]) return _RenderedLines(lines, ctx.imports) @@ -290,7 +397,15 @@ def tir_module_to_python(mod: Module, module_name: str | None = None, *, options imports.add("from tilefoundry.ir.types import Topology") rendered = ", ".join(f'Topology("{t.name}", {t.size!r})' for t in mod.topologies) kwargs.append(f"topologies=({rendered},)" if rendered else "topologies=()") - lines.append(f"@module({', '.join(kwargs)})") + decorator = f"@module({', '.join(kwargs)})" + if len(decorator) > _LINE_LENGTH: + packed = f" {', '.join(kwargs)}" + if len(packed) <= _LINE_LENGTH: + lines.extend(("@module(", packed, ")")) + else: + lines.extend(("@module(", *(f" {item}," for item in kwargs), ")")) + else: + lines.append(decorator) lines.append(f"class {name}:") blocks: list[list[str]] = [] for child in mod.modules: diff --git a/src/tilefoundry/ir/pattern/__init__.py b/src/tilefoundry/ir/pattern/__init__.py index b59e2694..a3344174 100644 --- a/src/tilefoundry/ir/pattern/__init__.py +++ b/src/tilefoundry/ir/pattern/__init__.py @@ -50,7 +50,17 @@ TensorPattern, WildcardPattern, ) -from .utils import _mangle_variant_name, arrangement_pattern, locate_dim_var +from .utils import ( + MOVED_STORAGES, + WHOLE_BYTES, + _mangle_variant_name, + any_threads, + arrangement_pattern, + dtype_place, + locate_dim_var, + moved_tile, + storage_place, +) __all__ = [ "ABSENT", @@ -67,6 +77,7 @@ "LayoutPattern", "Match", "MeshPattern", + "MOVED_STORAGES", "MultipleOfPattern", "OPAQUE", "OneOfPattern", @@ -85,19 +96,24 @@ "TensorPattern", "UNNAMED_PLACE", "WildcardPattern", + "WHOLE_BYTES", "_mangle_variant_name", "affine_frame", "affine_part", "alternatives_of", + "any_threads", "arrangement_pattern", "between_rules", "evaluated", "fits", "grouping", + "dtype_place", "is_symbolic", "locate_dim_var", "matched", + "moved_tile", "refusals_between", "relations_of", "resolved", + "storage_place", ] diff --git a/src/tilefoundry/ir/pattern/utils.py b/src/tilefoundry/ir/pattern/utils.py index 30eec000..3e361e9e 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -2,16 +2,80 @@ from __future__ import annotations -from tilefoundry.ir.types import ComposedLayout, Swizzle +from tilefoundry.ir.core.param_def import ParamDef +from tilefoundry.ir.types import ComposedLayout, Mesh, StorageKind, Swizzle from .pattern import ( + AttrPattern, + CapturePattern, ComposedLayoutPattern, LayoutPattern, + MeshPattern, + MultipleOfPattern, + OneOfPattern, + OrPattern, Pattern, RangePattern, SwizzlePattern, + TensorPattern, + WildcardPattern, ) +MOVED_STORAGES = (StorageKind.GMEM, StorageKind.SMEM, StorageKind.RMEM) +WHOLE_BYTES = AttrPattern("bit_width", MultipleOfPattern(8)) + + +def moved_tile(index: int, storage=None, layout=None) -> TensorPattern: + """A moved tensor tile, with dtype and storage captures named by end.""" + storages = MOVED_STORAGES if storage is None else storage + return TensorPattern( + dtype=CapturePattern(dtype_place(index), WHOLE_BYTES), + storage=( + CapturePattern(storage_place(index), OneOfPattern(tuple(storages))) + if isinstance(storages, tuple) + else storages + ), + layout=layout, + ) + + +def dtype_place(index: int) -> str: + """The capture name for transfer end *index*'s dtype.""" + return f"dtype{index}" + + +def storage_place(index: int) -> str: + """The capture name for transfer end *index*'s storage.""" + return f"storage{index}" + + +_ANY_THREADS = OrPattern( + ComposedLayoutPattern( + offset=WildcardPattern(), + outer=LayoutPattern( + ((CapturePattern("n", RangePattern(lo=1)),),), + ((1,),), + per_mode=True, + ), + ), + LayoutPattern( + ((CapturePattern("n", RangePattern(lo=1)),),), + ((1,),), + per_mode=True, + ), +) + + +def any_threads() -> ParamDef: + """Declare an optional scope spanning one or more threads.""" + return ParamDef( + kind="attribute", + annotation=Mesh, + pattern=MeshPattern(("thread",), _ANY_THREADS), + optional=True, + default=None, + ) + def arrangement_pattern( layout, @@ -65,4 +129,14 @@ def _mangle_variant_name(name: str, specializations: tuple[Pattern, ...]) -> str return f"{name}${pattern.dim_var}${pattern.lo}_{pattern.hi}" -__all__ = ["_mangle_variant_name", "arrangement_pattern", "locate_dim_var"] +__all__ = [ + "MOVED_STORAGES", + "WHOLE_BYTES", + "_mangle_variant_name", + "any_threads", + "arrangement_pattern", + "dtype_place", + "locate_dim_var", + "moved_tile", + "storage_place", +] diff --git a/src/tilefoundry/ir/tir/__init__.py b/src/tilefoundry/ir/tir/__init__.py index 502e8903..63f37e92 100644 --- a/src/tilefoundry/ir/tir/__init__.py +++ b/src/tilefoundry/ir/tir/__init__.py @@ -5,6 +5,7 @@ from .abort import Abort from .async_copy import CopyAsync, CpAsyncCommit, CpAsyncWait +from .cast import Cast from .launch import Launch from .prim_function import PrimFunction from .shape import ShapeOf @@ -47,6 +48,7 @@ def _auto_import(pkg_name: str) -> None: "classify", "participation", "CopyAsync", + "Cast", "CpAsyncCommit", "CpAsyncWait", ] diff --git a/src/tilefoundry/ir/tir/cast.py b/src/tilefoundry/ir/tir/cast.py new file mode 100644 index 00000000..a329eb53 --- /dev/null +++ b/src/tilefoundry/ir/tir/cast.py @@ -0,0 +1,42 @@ +"""Effect-ful TIR dtype conversion operation.""" + +from __future__ import annotations + +from tilefoundry.ir.core import Op +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef +from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import ( + DistinctConstraint, + SameConstraint, + TensorPattern, + any_threads, +) +from tilefoundry.ir.tir.memory.copy import verify_between +from tilefoundry.ir.types import StorageKind, UnitType +from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt + +_REGISTER = TensorPattern(storage=StorageKind.RMEM) + + +@register_op(dialect="T", category="arith", name="cast") +class Cast(Op): + """Convert a register tile to another dtype without changing its shape.""" + + src = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=_REGISTER) + dst = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=_REGISTER) + between = ( + SameConstraint("shape", "src", "dst"), + DistinctConstraint("dtype", "src", "dst"), + ) + scope = any_threads() + + +@register_typeinfer(Cast) +def _(call: "Call", ctx: "TypeInferContext") -> UnitType: + return UnitType() + + +register_verify_stmt(Cast)(verify_between) + + +__all__ = ["Cast"] diff --git a/src/tilefoundry/ir/tir/memory/copy.py b/src/tilefoundry/ir/tir/memory/copy.py index 0c01724d..f0ec45fa 100644 --- a/src/tilefoundry/ir/tir/memory/copy.py +++ b/src/tilefoundry/ir/tir/memory/copy.py @@ -9,10 +9,14 @@ from __future__ import annotations from tilefoundry.ir.core import Op -from tilefoundry.ir.core.param_def import ParamDef +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import Tensor -from tilefoundry.ir.types import UnitType +from tilefoundry.ir.pattern import ( + any_threads, + between_rules, + moved_tile, +) +from tilefoundry.ir.types import LayoutBase, UnitType from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -21,8 +25,12 @@ class Copy(Op): """Copies ``src`` into ``dst`` (in-place memory write).""" - src = ParamDef(kind="input", pattern=Tensor) - dst = ParamDef(kind="input", pattern=Tensor) + src = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=moved_tile(0)) + dst = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=moved_tile(1)) + rmem_layout = ParamDef(kind="attribute", annotation=LayoutBase, optional=True, default=None) + smem_layout = ParamDef(kind="attribute", annotation=LayoutBase, optional=True, default=None) + + scope = any_threads() @register_typeinfer(Copy) @@ -41,6 +49,38 @@ def _(call: "Call", ctx: "VerifyContext") -> None: ctx.error(call, f"Copy dtype mismatch: {src.dtype} vs {dst.dtype}") +def _input_params(op_type: type) -> tuple: + return tuple(param for param in op_type._op_schema.signature if param.kind == "input") + + +def verify_between(call, ctx, lead: str = "") -> None: + """Hold one call to the relations declared between its operands.""" + op_type = type(call.target) + rules = between_rules(op_type) + if not rules: + return + names = tuple(param.name for param in _input_params(op_type)) + operands = dict(zip(names, (ctx.type_of(arg) for arg in call.args))) + for rule in rules: + if not rule.holds(operands): + ctx.error(call, lead + rule.refused(operands)) + + +def verify_operands(call, ctx, label: str) -> None: + """Hold each operand to the pattern declared for its parameter.""" + for param, arg in zip(_input_params(type(call.target)), call.args): + if param.pattern is None: + continue + value = ctx.type_of(arg) + if param.pattern.match(value) is None: + ctx.error( + call, + f"{label} {param.name} is {tuple(value.shape)} " + f"{value.dtype.name} storage={value.storage}: " + f"{param.pattern.refusal(value)}", + ) + + def _is_copyable_shard(src_ty, dst_ty) -> bool: """Both sides carry a ShardLayout describing the same per-thread buffer.""" src_sl = getattr(src_ty, "layout", None) diff --git a/src/tilefoundry/ir/tir/memory/fill.py b/src/tilefoundry/ir/tir/memory/fill.py index 3e009730..3f3cb61a 100644 --- a/src/tilefoundry/ir/tir/memory/fill.py +++ b/src/tilefoundry/ir/tir/memory/fill.py @@ -6,10 +6,10 @@ from __future__ import annotations -from tilefoundry.ir.core import Op -from tilefoundry.ir.core.param_def import ParamDef +from tilefoundry.ir.core import Constant, Op +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import Tensor +from tilefoundry.ir.pattern import Scalar, any_threads, moved_tile from tilefoundry.ir.types import UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -18,8 +18,9 @@ class Fill(Op): """Fills ``tensor`` element-wise with ``value`` (rank-0 scalar).""" - tensor = ParamDef(kind="input", pattern=Tensor) - value = ParamDef(kind="input", pattern=Tensor) + tensor = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=moved_tile(0)) + value = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=Scalar) + scope = any_threads() @register_typeinfer(Fill) @@ -33,5 +34,7 @@ def _(call: "Call", ctx: "VerifyContext") -> None: v_ty = ctx.type_of(call.args[1]) if v_ty.shape != (): ctx.error(call, "Fill value must be rank-0 scalar") - if v_ty.dtype != t_ty.dtype: + if v_ty.dtype != t_ty.dtype and not ( + isinstance(call.args[1], Constant) and call.args[1].value == 0 + ): ctx.error(call, f"Fill dtype mismatch: {v_ty.dtype} vs {t_ty.dtype}") diff --git a/src/tilefoundry/ir/tir/memory/ptr_of.py b/src/tilefoundry/ir/tir/memory/ptr_of.py index 13ed75b1..125bbbab 100644 --- a/src/tilefoundry/ir/tir/memory/ptr_of.py +++ b/src/tilefoundry/ir/tir/memory/ptr_of.py @@ -1,8 +1,6 @@ """TIR view Expr Op: `tir.view.PtrOf`. -Takes an Expr of tensor/scalar type, returns a raw -pointer descriptor. Placeholder: typeinfer returns the input type until a -dedicated PointerType lands. +Takes a tensor Expr and returns its typed physical-memory pointer descriptor. """ from __future__ import annotations @@ -11,7 +9,7 @@ from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.core.register import register_op from tilefoundry.ir.pattern import Tensor -from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types import PointerType from tilefoundry.visitor_registry import register_typeinfer @@ -19,9 +17,10 @@ class PtrOf(Op): """Take the device address of a tensor for downstream view ops (value form).""" - x = ParamDef(kind="input", pattern=Tensor) + tensor = ParamDef(kind="input", pattern=Tensor) @register_typeinfer(PtrOf) -def _(call: "Call", ctx: "TypeInferContext") -> TensorType: - return ctx.type_of(call.args[0]) +def _(call: "Call", ctx: "TypeInferContext") -> PointerType: + tensor = ctx.type_of(call.args[0]) + return PointerType(tensor.dtype, tensor.storage) diff --git a/src/tilefoundry/ir/tir/memory/tensor_view.py b/src/tilefoundry/ir/tir/memory/tensor_view.py index 1655a4e7..5095e07f 100644 --- a/src/tilefoundry/ir/tir/memory/tensor_view.py +++ b/src/tilefoundry/ir/tir/memory/tensor_view.py @@ -1,7 +1,6 @@ """TIR view Expr Op: `tir.memory.TensorView`. -Constructs a logical tensor view over a memory source -(tensor / ptr / span). ``layout`` can be a plain ``Layout`` (→ plain view) +Constructs a logical tensor view over a typed pointer. ``layout`` can be a plain ``Layout`` or a ``ShardLayout`` (→ shard tensor view, no allocation). A slice view carries an absolute element-start coordinate per axis after the @@ -11,12 +10,12 @@ from __future__ import annotations -from tilefoundry.ir.core import Op +from tilefoundry.ir.core import Call, Constant, Op from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import Tensor -from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types import DType, PointerType, StorageKind, TensorType from tilefoundry.ir.types.layout import Layout, LayoutBase +from tilefoundry.ir.types.storage import resolve_storage from tilefoundry.ir.types.stride import compact_row_major from tilefoundry.visitor_registry import register_typeinfer @@ -26,29 +25,73 @@ class TensorView(Op): """Derive a sub-view of a tensor (value form). ``layout`` updates the ``ShardLayout`` / cute ``Layout``; the optional - ``shape`` overrides the logical shape (reshape). ``memory`` MAY be a - ``PtrOf`` result (ptr + offset). + ``shape`` overrides the logical shape. When omitted, it is inherited only + from a syntactic ``PtrOf(tensor)`` input. Integer inputs are byte offsets + from the dynamic shared-memory base and require dtype, storage, and shape. """ - memory = ParamDef(kind="input", pattern=Tensor) + pointer = ParamDef(kind="input") + dtype = ParamDef(kind="attribute", annotation=str, optional=True, default=None) + storage = ParamDef(kind="attribute", annotation=StorageKind, optional=True, default=None) layout = ParamDef(kind="attribute", annotation=LayoutBase) shape = ParamDef(kind="attribute", annotation=tuple, default=None) @register_typeinfer(TensorView) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: - src_ty = ctx.type_of(call.args[0]) + pointer = ctx.type_of(call.args[0]) op = call.target - new_layout = op.layout - new_shape = op.shape if op.shape is not None else src_ty.shape - - if len(call.args) > 1: - new_shape = call.type.shape + origin = call.args[0] + inherited_shape = None + if isinstance(origin, Call): + from .ptr_of import PtrOf # noqa: PLC0415 + + if isinstance(origin.target, PtrOf): + tensor = ctx.type_of(origin.args[0]) + if not isinstance(tensor, TensorType): + ctx.error(call, "PtrOf input must have TensorType") + inherited_shape = tensor.shape + + if not isinstance(pointer, PointerType): + address = call.args[0] + if ( + not isinstance(address, Constant) + or isinstance(address.value, bool) + or not isinstance(address.value, int) + ): + ctx.error(call, "tensor_view input must be a pointer") + if op.dtype is None or op.storage is None or op.shape is None: + ctx.error( + call, + "shared-memory byte offset requires dtype, storage, and shape", + ) + try: + storage = resolve_storage(op.storage) + pointer = PointerType(DType.from_name(op.dtype), storage) + except (TypeError, ValueError) as error: + ctx.error(call, str(error)) + if pointer.storage is not StorageKind.SMEM: + ctx.error(call, "integer tensor_view address requires storage='smem'") + address.type = pointer + elif op.dtype is not None or op.storage is not None: + try: + stated = PointerType( + pointer.dtype if op.dtype is None else DType.from_name(op.dtype), + pointer.storage if op.storage is None else resolve_storage(op.storage), + ) + except (TypeError, ValueError) as error: + ctx.error(call, str(error)) + if stated != pointer: + ctx.error(call, "tensor_view pointer dtype/storage mismatch") + + new_shape = op.shape if op.shape is not None else inherited_shape + if new_shape is None: + ctx.error(call, "tensor_view shape is required unless pointer is T.ptr_of(tensor)") return TensorType( shape=new_shape, - dtype=src_ty.dtype, - layout=new_layout, - storage=src_ty.storage, + dtype=pointer.dtype, + layout=op.layout, + storage=pointer.storage, ) diff --git a/src/tilefoundry/ir/types/__init__.py b/src/tilefoundry/ir/types/__init__.py index c63273b6..f82479df 100644 --- a/src/tilefoundry/ir/types/__init__.py +++ b/src/tilefoundry/ir/types/__init__.py @@ -19,6 +19,7 @@ ) from .mesh import Mesh, Topology, make_mesh from .placement import Placement +from .pointer import PointerType from .tensor_type import TensorType, TupleType, Type, UnitType from .utils import make_shard_tensor_type, make_tensor_type from .callable_type import CallableType, callable_type_for @@ -41,6 +42,7 @@ "P", "Partial", "Placement", + "PointerType", "S", "ShardAttr", "ShardLayout", diff --git a/src/tilefoundry/ir/types/pointer.py b/src/tilefoundry/ir/types/pointer.py new file mode 100644 index 00000000..2d587d30 --- /dev/null +++ b/src/tilefoundry/ir/types/pointer.py @@ -0,0 +1,26 @@ +"""Typed physical-memory pointer descriptors.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .dtype import DType +from .storage import StorageKind, resolve_storage + + +@dataclass(frozen=True) +class PointerType: + """A typed physical-memory engine consumed by ``T.tensor_view``.""" + + dtype: DType + storage: StorageKind + + def __post_init__(self) -> None: + normalized = resolve_storage(self.storage) + if normalized is None: + raise TypeError("PointerType.storage must be a StorageKind, not None") + if normalized is not self.storage: + object.__setattr__(self, "storage", normalized) + + +__all__ = ["PointerType"] diff --git a/src/tilefoundry/ir/types/substitute.py b/src/tilefoundry/ir/types/substitute.py index b81d5628..ebf458c7 100644 --- a/src/tilefoundry/ir/types/substitute.py +++ b/src/tilefoundry/ir/types/substitute.py @@ -14,6 +14,7 @@ from tilefoundry.ir.isl_interop import normalize_dim from .dim import _DIM_OP_TYPES, DimVar, simplify_dim +from .pointer import PointerType from .tensor_type import TensorType, TupleType, Type @@ -131,6 +132,8 @@ def _collect(value: object, found: dict[str, "DimVar"]) -> None: for field in value.fields: _collect(field, found) return + if isinstance(value, PointerType): + return if isinstance(value, DimVar): found[value.name] = value return @@ -195,6 +198,8 @@ def substitute_dims(value: Type, bindings: Mapping[str, int]) -> Type: if fields == value.fields: return value return TupleType(fields=fields) + if isinstance(value, PointerType): + return value return value @@ -218,6 +223,8 @@ def canonicalize_dims(value: Type) -> Type: if fields == value.fields: return value return TupleType(fields=fields) + if isinstance(value, PointerType): + return value return value @@ -400,6 +407,8 @@ def has_symbolic_dims(value: object) -> bool: return has_symbolic_dims(value.shape) or has_symbolic_dims(value.layout) if isinstance(value, TupleType): return any(has_symbolic_dims(field) for field in value.fields) + if isinstance(value, PointerType): + return False if isinstance(value, Topology): return has_symbolic_dims(value.size) if isinstance(value, Mesh): @@ -407,10 +416,7 @@ def has_symbolic_dims(value: object) -> bool: if isinstance(value, ShardLayout): return has_symbolic_dims(value.layout) or has_symbolic_dims(value.mesh) if isinstance(value, ComposedLayout): - return any( - has_symbolic_dims(entry) - for entry in (value.inner, value.offset, value.outer) - ) + return any(has_symbolic_dims(entry) for entry in (value.inner, value.offset, value.outer)) if isinstance(value, Layout): return has_symbolic_dims(value.shape) or has_symbolic_dims(value.strides) if isinstance(value, tuple): diff --git a/src/tilefoundry/ir/types/tensor_type.py b/src/tilefoundry/ir/types/tensor_type.py index c27848e3..da532ad0 100644 --- a/src/tilefoundry/ir/types/tensor_type.py +++ b/src/tilefoundry/ir/types/tensor_type.py @@ -76,9 +76,7 @@ def umat_scalar(dtype: DType = DType.i64) -> "TensorType": """ from .layout import EMPTY_LAYOUT # noqa: PLC0415 - cycle guard - return TensorType( - shape=(), dtype=dtype, layout=EMPTY_LAYOUT, storage=StorageKind.UMAT - ) + return TensorType(shape=(), dtype=dtype, layout=EMPTY_LAYOUT, storage=StorageKind.UMAT) @staticmethod def umat_tensor(shape: tuple, dtype: DType = DType.i64) -> "TensorType": @@ -109,6 +107,6 @@ class UnitType: """ -Type = Union[TensorType, TupleType, UnitType, "CallableType"] +Type = Union[TensorType, TupleType, UnitType, "CallableType", "PointerType"] __all__ = ["DType", "TensorType", "TupleType", "UnitType", "Type"] diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 34df2619..e61d1e38 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -68,6 +68,7 @@ Layout, Mesh, Partial, + PointerType, ShardLayout, Split, TensorType, @@ -245,6 +246,7 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex BindingSubstitutionCloner=BindingSubstitutionCloner, Broadcast=Broadcast, Partial=Partial, + PointerType=PointerType, Binary=Binary, BinaryKind=BinaryKind, Constant=Constant, diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index b60a9bed..c92f3f03 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -1184,17 +1184,35 @@ def _head(node: ast.expr) -> str | None: class ScalarTypePattern(ElementPattern): element_name = "scalar_type" syntax = LazyPattern( - lambda: BranchPattern( - "type_reference", - ReferencePattern(), - pattern_id="type.reference", + lambda: ChoicePattern( + BranchPattern( + "type_reference", + ReferencePattern(), + pattern_id="type.reference", + ), + BranchPattern( + "type_call", + ChildPattern( + "value", + lambda: StaticCallPattern(), + "type_constructor", + ), + pattern_id="type.call", + ), ) ) @staticmethod def construct(match, children, context): - value = _resolve_reference(match.node, context) - if not isinstance(value, (runtime.TensorType, runtime.TupleType, runtime.UnitType)): + value = ( + children["value"] + if match.branch_id == "type_call" + else _resolve_reference(match.node, context) + ) + if not isinstance( + value, + (runtime.TensorType, runtime.TupleType, runtime.UnitType, runtime.PointerType), + ): raise ParseError.from_node(match.node, context, "annotation did not resolve to IR Type") return value diff --git a/tests/codegen/test_host_multi_launch.py b/tests/codegen/test_host_multi_launch.py index 671459ba..9302ba60 100644 --- a/tests/codegen/test_host_multi_launch.py +++ b/tests/codegen/test_host_multi_launch.py @@ -40,7 +40,7 @@ class _TwoCopyLaunches: def first(x: Tensor[(1, 128), "f32"], out: Tensor[(1, 128), "f32"]): with Mesh((Topology("thread", 1),), Layout((1,), (1,))) as thread: view = T.tensor_view( - x, + T.ptr_of(x), layout=ShardLayout( layout=Layout((1, 128), (128, 1)), attrs=(B(),), @@ -48,7 +48,7 @@ def first(x: Tensor[(1, 128), "f32"], out: Tensor[(1, 128), "f32"]): ), ) out_view = T.tensor_view( - out, + T.ptr_of(out), layout=ShardLayout( layout=Layout((1, 128), (128, 1)), attrs=(B(),), @@ -62,7 +62,7 @@ def first(x: Tensor[(1, 128), "f32"], out: Tensor[(1, 128), "f32"]): def second(x: Tensor[(1, 128), "f32"], out: Tensor[(1, 128), "f32"]): with Mesh((Topology("thread", 1),), Layout((1,), (1,))) as thread: view = T.tensor_view( - x, + T.ptr_of(x), layout=ShardLayout( layout=Layout((1, 128), (128, 1)), attrs=(B(),), @@ -70,7 +70,7 @@ def second(x: Tensor[(1, 128), "f32"], out: Tensor[(1, 128), "f32"]): ), ) out_view = T.tensor_view( - out, + T.ptr_of(out), layout=ShardLayout( layout=Layout((1, 128), (128, 1)), attrs=(B(),), diff --git a/tests/codegen/test_submodule_topology.py b/tests/codegen/test_submodule_topology.py index e7a3ff1d..1b3bb97f 100644 --- a/tests/codegen/test_submodule_topology.py +++ b/tests/codegen/test_submodule_topology.py @@ -21,8 +21,8 @@ class _Left: def left_device(x: Tensor[(32,), "f32"], out: Tensor[(32,), "f32"]): with Mesh((Topology("thread", 32),), Layout((32,), (1,))) as thread: layout = ShardLayout(Layout((32,), (1,)), (Split(0),), thread) - src = T.tensor_view(x, layout=layout) - dst = T.tensor_view(out, layout=layout) + src = T.tensor_view(T.ptr_of(x), layout=layout) + dst = T.tensor_view(T.ptr_of(out), layout=layout) T.copy(src, dst) T.sync(thread) @@ -33,8 +33,8 @@ class _Right: def right_device(x: Tensor[(128,), "f32"], out: Tensor[(128,), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: layout = ShardLayout(Layout((128,), (1,)), (Split(0),), thread) - src = T.tensor_view(x, layout=layout) - dst = T.tensor_view(out, layout=layout) + src = T.tensor_view(T.ptr_of(x), layout=layout) + dst = T.tensor_view(T.ptr_of(out), layout=layout) T.copy(src, dst) T.sync(thread) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index 02ece31a..4c6cea29 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -10,14 +10,14 @@ 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(mesh=Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=('lane',))) + @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", ShardLayout( layout=Layout((8, 16), None), attrs=(B(), P("max")), - mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=("warp", "lane")), )] ): split = reshard(x, layout=(8 @ mesh.lane, 16), storage=rmem) @@ -29,12 +29,12 @@ class TypePrinterSugar: 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: + 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(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) + @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 @ mesh.tile, 2, 16), "rmem"], @@ -42,7 +42,7 @@ class TypePrinterSugar: 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')), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=("warp", "lane")), ), "rmem"] ): mine = reshard(x, layout=(4 @ mesh.tile, 2, 16), storage=rmem) @@ -50,11 +50,11 @@ class TypePrinterSugar: 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')), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=("warp", "lane")), ), 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("thread", 8),), Layout((2, 4), (4, 1)), names=("warp", "lane"))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], @@ -63,10 +63,10 @@ class TypePrinterSugar: 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')), + 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: + 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 @ cta.tile, 2, 4, 16), storage=smem) diff --git a/tests/fixtures/placed/gpu_placed_rows.py b/tests/fixtures/placed/gpu_placed_rows.py index 51c62aef..0647523a 100644 --- a/tests/fixtures/placed/gpu_placed_rows.py +++ b/tests/fixtures/placed/gpu_placed_rows.py @@ -46,8 +46,8 @@ def copy_rows_device(a: Tensor[(ROWS, COLS), "f32"], out: Tensor[(ROWS, COLS), " ("g", "c", "t"), ) as m: rows = _split_rows(m) - source = T.tensor_view(a, layout=rows) - written = T.tensor_view(out, layout=rows) + source = T.tensor_view(T.ptr_of(a), layout=rows) + written = T.tensor_view(T.ptr_of(out), layout=rows) T.copy(source, written) @prim_func(target=CpuTarget()) diff --git a/tests/fixtures/tir/async_sync.py b/tests/fixtures/tir/async_sync.py index 44820b28..43c6efb7 100644 --- a/tests/fixtures/tir/async_sync.py +++ b/tests/fixtures/tir/async_sync.py @@ -6,14 +6,20 @@ from tilefoundry.target import CpuTarget, CudaTarget -@module(entry="async_stage_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) +@module( + entry="async_stage_host", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("thread", 128),), +) class AsyncStage: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def async_stage_device(a: Tensor[(128, 4), "f32"], b: Tensor[(128, 4), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as m: - 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))) + with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=("t",)) as m: + a_view = T.tensor_view(T.ptr_of(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(T.ptr_of(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/mma.py b/tests/fixtures/tir/mma.py index 60a67446..63a0213c 100644 --- a/tests/fixtures/tir/mma.py +++ b/tests/fixtures/tir/mma.py @@ -9,18 +9,51 @@ @module(entry="mm_host", target=CudaTarget("nvidia.h200_sxm")) 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)), 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"]) + 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)), names=("warp", "lane") + ) as _warp: + a_view = T.tensor_view( + T.ptr_of(a), layout=((2, 4 @ _warp.warp, 2, 8 @ _warp.lane, 2), (1, 2, 8, 16, 128)) + ) + b_view = T.tensor_view( + T.ptr_of(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=((2, 4 @ _warp.warp, 8 @ _warp.lane, 2), (1, 2, 8, 64))) + 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( + T.ptr_of(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/rmsnorm.py b/tests/fixtures/tir/rmsnorm.py index 1aee5c21..ad4fc5b2 100644 --- a/tests/fixtures/tir/rmsnorm.py +++ b/tests/fixtures/tir/rmsnorm.py @@ -9,14 +9,18 @@ @module(entry="rmsnorm_host", target=CudaTarget("nvidia.h200_sxm")) 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), (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()})) + 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(T.ptr_of(x), layout=((1, 128), (128, 1), {thread.t @ B()})) + weight_view = T.tensor_view(T.ptr_of(weight), layout=((128,), (1,), {thread.t @ B()})) + out_view = T.tensor_view(T.ptr_of(out), layout=((1, 128), (128, 1), {thread.t @ B()})) T.rms_norm(x_view, out_view, weight_view, eps=1e-05) T.sync(thread) @prim_func(target=CpuTarget()) - def rmsnorm_host(x: Tensor[(1, 128), "f32"], weight: Tensor[(128,), "f32"], out: Tensor[(1, 128), "f32"]): + def rmsnorm_host( + x: Tensor[(1, 128), "f32"], weight: Tensor[(128,), "f32"], out: Tensor[(1, 128), "f32"] + ): launch(rmsnorm_device, x, weight, out, grid=(1, 1, 1), block=(1, 1, 1)) # noqa: F821 diff --git a/tests/fixtures/tir/square.py b/tests/fixtures/tir/square.py index b7ab2844..11620487 100644 --- a/tests/fixtures/tir/square.py +++ b/tests/fixtures/tir/square.py @@ -21,7 +21,7 @@ def square_device(x: Tensor[(_S,), "f32"]): @square_device.specialize(RangePattern("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,), (1,))) + view = T.tensor_view(T.ptr_of(x), layout=((128 @ thread.t,), (1,))) reg = T.alloc_tensor( tensor_type=Tensor[(128,), "f32", ((128 @ thread.t,), (1,)), "rmem"] ) @@ -36,7 +36,7 @@ def square_small(x: Tensor[(_S,), "f32"]): @square_device.specialize(RangePattern("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,), (1,))) + view = T.tensor_view(T.ptr_of(x), layout=((128 @ thread.t,), (1,))) reg = T.alloc_tensor( tensor_type=Tensor[(128,), "f32", ((128 @ thread.t,), (1,)), "rmem"] ) diff --git a/tests/fixtures/tir/sync.py b/tests/fixtures/tir/sync.py index 3d4f7984..40472ba3 100644 --- a/tests/fixtures/tir/sync.py +++ b/tests/fixtures/tir/sync.py @@ -7,13 +7,19 @@ from tilefoundry.target import CpuTarget, CudaTarget -@module(entry="sync_square_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) +@module( + entry="sync_square_host", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("thread", 128),), +) class SyncSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def sync_square_device(a: Tensor[(4, 32), "f32"]): - 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), (32, 1))) - reg = T.alloc_tensor(tensor_type=Tensor[(4, 32), "f32", ((4 @ m.w, 32 @ m.t), (32, 1)), "rmem"]) + with Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1)), names=("w", "t")) as m: + view = T.tensor_view(T.ptr_of(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(m[:1]) diff --git a/tests/integration/test_dynamic_cta_tir_handwritten.py b/tests/integration/test_dynamic_cta_tir_handwritten.py index 44a1846e..47d1fc7e 100644 --- a/tests/integration/test_dynamic_cta_tir_handwritten.py +++ b/tests/integration/test_dynamic_cta_tir_handwritten.py @@ -34,13 +34,14 @@ def test_handwritten_tir_dynamic_cta_matches_torch_at_several_shapes() -> None: One compiled artifact squares the tensor at three ``Ntile`` shapes via the host-computed grid; all match torch with no recompile. """ + @module(entry="dyn_square_host") class DynSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def dyn_square(a: Tensor[(_NT, _TILE), "f32"]): with Mesh((Topology("cta", _NT),), Layout(shape=(_NT,), strides=(1,))) as cta: a_view = T.tensor_view( - a, + T.ptr_of(a), layout=ShardLayout( layout=Layout(shape=(_NT, _TILE), strides=(_TILE, 1)), attrs=(Split(0),), diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 431630e7..531ea6e7 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -52,7 +52,7 @@ def tile_device( names=("warp", "lane"), ) as m: a_view = T.tensor_view( - a, + T.ptr_of(a), layout=ShardLayout( layout=Layout(shape=(256,), strides=(1,)), attrs=(Broadcast(), Broadcast()), @@ -60,7 +60,7 @@ def tile_device( ), ) b_view = T.tensor_view( - b, + T.ptr_of(b), layout=ShardLayout( layout=Layout(shape=(128,), strides=(1,)), attrs=(Broadcast(), Broadcast()), @@ -97,7 +97,7 @@ def tile_device( T.fill(acc, 0.0) T.sync(m) T.mma(acc, a_tile, b_tile) - c_view = T.tensor_view(c, layout=atom.C) + c_view = T.tensor_view(T.ptr_of(c), layout=atom.C) T.copy(acc, c_view) @prim_func(target=CpuTarget()) diff --git a/tests/ops/tir/cuda/test_swizzle.py b/tests/ops/tir/cuda/test_swizzle.py index f6d8bb4b..e7aa6fc2 100644 --- a/tests/ops/tir/cuda/test_swizzle.py +++ b/tests/ops/tir/cuda/test_swizzle.py @@ -38,9 +38,7 @@ def _swizzled_rows(mesh: Mesh) -> ShardLayout: address really moves, and the four elements of a row stay together. """ return ShardLayout( - ComposedLayout( - inner=_SWIZZLE, offset=0, outer=Layout((_ROWS, _COLS), (_COLS, 1)) - ), + ComposedLayout(inner=_SWIZZLE, offset=0, outer=Layout((_ROWS, _COLS), (_COLS, 1))), (Split(0),), mesh, ) @@ -56,21 +54,16 @@ def swizzled_square_device( dst: Tensor[(_ROWS, _COLS), "f32"], ): with Mesh((Topology("thread", _ROWS),), Layout((_ROWS,), (1,)), ("t",)) as threads: - src_view = T.tensor_view(src, layout=_rows(threads)) - dst_view = T.tensor_view(dst, layout=_rows(threads)) - tile = T.alloc_tensor( - Tensor[(_ROWS, _COLS), "f32", _swizzled_rows(threads), "smem"] - ) - fragment = T.alloc_tensor( - Tensor[(_ROWS, _COLS), "f32", _rows(threads), "rmem"] - ) + src_view = T.tensor_view(T.ptr_of(src), layout=_rows(threads)) + dst_view = T.tensor_view(T.ptr_of(dst), layout=_rows(threads)) + tile = T.alloc_tensor(Tensor[(_ROWS, _COLS), "f32", _swizzled_rows(threads), "smem"]) + fragment = T.alloc_tensor(Tensor[(_ROWS, _COLS), "f32", _rows(threads), "rmem"]) T.copy(src_view, tile) T.sync(threads) T.copy(tile, fragment) T.binary(fragment, fragment, fragment, kind=BinaryKind.MUL) T.copy(fragment, dst_view) - @prim_func(target=CpuTarget()) def swizzled_square_host( src: Tensor[(_ROWS, _COLS), "f32"], diff --git a/tests/ops/tir/cuda/test_tma.py b/tests/ops/tir/cuda/test_tma.py index f27c299c..cc19f15f 100644 --- a/tests/ops/tir/cuda/test_tma.py +++ b/tests/ops/tir/cuda/test_tma.py @@ -101,7 +101,7 @@ def tma_tiers_device( ): with Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,)), ("t",)) as m: bulk_view = T.tensor_view( - bulk_a, + T.ptr_of(bulk_a), layout=ShardLayout( layout=Layout(shape=(256,), strides=(1,)), attrs=(Broadcast(),), mesh=m ), @@ -126,7 +126,7 @@ def tma_tiers_device( T.copy(bulk_stage, bulk_b) with Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,)), ("t",)) as mo: odd_view = T.tensor_view( - odd_a, + T.ptr_of(odd_a), layout=ShardLayout( layout=Layout(shape=(5,), strides=(1,)), attrs=(Broadcast(),), mesh=mo ), diff --git a/tests/ops/tir/test_copy.py b/tests/ops/tir/test_copy.py index 2215d437..32520d9b 100644 --- a/tests/ops/tir/test_copy.py +++ b/tests/ops/tir/test_copy.py @@ -37,33 +37,33 @@ def copy_storage_device( b_narrow: Tensor[(128, 2), "f32"], ): with Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,)), ("t",)) as m: - smem_src = T.tensor_view(a_smem, layout=split_rows(m)) - smem_dst = T.tensor_view(b_smem, layout=split_rows(m)) + smem_src = T.tensor_view(T.ptr_of(a_smem), layout=split_rows(m)) + smem_dst = T.tensor_view(T.ptr_of(b_smem), layout=split_rows(m)) smem_tile = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "smem"]) T.copy(smem_src, smem_tile) T.sync(m) T.copy(smem_tile, smem_dst) - rmem_src = T.tensor_view(a_rmem, layout=split_rows(m)) - rmem_dst = T.tensor_view(b_rmem, layout=split_rows(m)) + rmem_src = T.tensor_view(T.ptr_of(a_rmem), layout=split_rows(m)) + rmem_dst = T.tensor_view(T.ptr_of(b_rmem), layout=split_rows(m)) rmem_tile = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "smem"]) rmem_frag = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "rmem"]) T.copy(rmem_src, rmem_tile) T.sync(m) T.copy(rmem_tile, rmem_frag) T.copy(rmem_frag, rmem_dst) - wide_src = T.tensor_view(a_wide, layout=split_rows(m)) - wide_dst = T.tensor_view(b_wide, layout=split_rows(m)) + wide_src = T.tensor_view(T.ptr_of(a_wide), layout=split_rows(m)) + wide_dst = T.tensor_view(T.ptr_of(b_wide), layout=split_rows(m)) wide_frag = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "rmem"]) T.copy(wide_src, wide_frag) T.copy(wide_frag, wide_dst) - narrow_src = T.tensor_view(a_narrow, layout=split_pairs(m)) - narrow_dst = T.tensor_view(b_narrow, layout=split_pairs(m)) + narrow_src = T.tensor_view(T.ptr_of(a_narrow), layout=split_pairs(m)) + narrow_dst = T.tensor_view(T.ptr_of(b_narrow), layout=split_pairs(m)) narrow_frag = T.alloc_tensor(Tensor[(128, 2), "f32", split_pairs(m), "rmem"]) T.copy(narrow_src, narrow_frag) T.copy(narrow_frag, narrow_dst) with Mesh((Topology("thread", 32),), Layout(shape=(32,), strides=(1,)), ("t",)) as mb: - bcast_src = T.tensor_view(a_bcast, layout=split_short_rows(mb)) - bcast_dst = T.tensor_view(b_bcast, layout=split_short_rows(mb)) + bcast_src = T.tensor_view(T.ptr_of(a_bcast), layout=split_short_rows(mb)) + bcast_dst = T.tensor_view(T.ptr_of(b_bcast), layout=split_short_rows(mb)) bcast_frag = T.alloc_tensor(Tensor[(4,), "f32", broadcast_run(mb), "rmem"]) T.copy(bcast_src, bcast_frag) T.copy(bcast_frag, bcast_dst) diff --git a/tests/ops/tir/test_dot.py b/tests/ops/tir/test_dot.py index 234125d5..65dff20e 100644 --- a/tests/ops/tir/test_dot.py +++ b/tests/ops/tir/test_dot.py @@ -78,7 +78,7 @@ def dot_warp( ): with Mesh((Topology("thread", 32),), Layout(shape=(32,), strides=(1,)), ("t",)) as mw: warp_a_view = T.tensor_view( - warp_a, + T.ptr_of(warp_a), layout=ShardLayout( layout=Layout(shape=(32, 32), strides=(32, 1)), attrs=(Split(0),), @@ -86,13 +86,13 @@ def dot_warp( ), ) warp_b_view = T.tensor_view( - warp_b, + T.ptr_of(warp_b), layout=ShardLayout( layout=Layout(shape=(32,), strides=(1,)), attrs=(Broadcast(),), mesh=mw ), ) warp_c_view = T.tensor_view( - warp_c, + T.ptr_of(warp_c), layout=ShardLayout( layout=Layout(shape=(32,), strides=(1,)), attrs=(Split(0),), mesh=mw ), @@ -107,9 +107,15 @@ def dot_cta( cta_a: Tensor[(128,), "f32"], cta_b: Tensor[(128,), "f32"], cta_c: Tensor[(1,), "f32"] ): with Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,)), ("t",)) as mc: - cta_a_view = T.tensor_view(cta_a, layout=ShardLayout(Layout((128,), (1,)), (Split(0),), mc)) - cta_b_view = T.tensor_view(cta_b, layout=ShardLayout(Layout((128,), (1,)), (Split(0),), mc)) - cta_c_view = T.tensor_view(cta_c, layout=ShardLayout(Layout((1,), (1,)), (Broadcast(),), mc)) + cta_a_view = T.tensor_view( + T.ptr_of(cta_a), layout=ShardLayout(Layout((128,), (1,)), (Split(0),), mc) + ) + cta_b_view = T.tensor_view( + T.ptr_of(cta_b), layout=ShardLayout(Layout((128,), (1,)), (Split(0),), mc) + ) + cta_c_view = T.tensor_view( + T.ptr_of(cta_c), layout=ShardLayout(Layout((1,), (1,)), (Broadcast(),), mc) + ) cta_ws = T.alloc_tensor(Tensor[(4,), "f32", None, "smem"]) T.dot(cta_a_view, cta_b_view, cta_c_view, cta_ws) @@ -120,7 +126,14 @@ class DotTiers: cta = DotCta @prim_func(target=CpuTarget()) - def dot_tiers_host(warp_a: Tensor[(32, 32), "f32"], warp_b: Tensor[(32,), "f32"], warp_c: Tensor[(32,), "f32"], cta_a: Tensor[(128,), "f32"], cta_b: Tensor[(128,), "f32"], cta_c: Tensor[(1,), "f32"]): + def dot_tiers_host( + warp_a: Tensor[(32, 32), "f32"], + warp_b: Tensor[(32,), "f32"], + warp_c: Tensor[(32,), "f32"], + cta_a: Tensor[(128,), "f32"], + cta_b: Tensor[(128,), "f32"], + cta_c: Tensor[(1,), "f32"], + ): launch(warp.dot_warp, warp_a, warp_b, warp_c, grid=(1, 1, 1), block=(32, 1, 1)) launch(cta.dot_cta, cta_a, cta_b, cta_c, grid=(1, 1, 1), block=(128, 1, 1)) diff --git a/tests/ops/tir/test_elementwise.py b/tests/ops/tir/test_elementwise.py index e2a6d0c1..f3e6601f 100644 --- a/tests/ops/tir/test_elementwise.py +++ b/tests/ops/tir/test_elementwise.py @@ -36,18 +36,12 @@ def col_bcast_device( a: Tensor[(32,), "f32"], r: Tensor[(4,), "f32"], out: Tensor[(32,), "f32"] ): with Mesh((Topology("thread", 32),), Layout(shape=(32,), strides=(1,)), ("t",)) as m: - a_view = T.tensor_view(a, layout=bcast((32,), (1,), m)) - r_view = T.tensor_view(r, layout=bcast((4,), (1,), m)) - out_view = T.tensor_view(out, layout=bcast((32,), (1,), m)) - lhs = T.alloc_tensor( - Tensor[(4, 8), 'f32', bcast((4, 8), (8, 1), m), 'smem'] - ) - col = T.alloc_tensor( - Tensor[(4, 1), 'f32', bcast((4, 1), (1, 1), m), 'smem'] - ) - dst = T.alloc_tensor( - Tensor[(4, 8), 'f32', bcast((4, 8), (8, 1), m), 'smem'] - ) + a_view = T.tensor_view(T.ptr_of(a), layout=bcast((32,), (1,), m)) + r_view = T.tensor_view(T.ptr_of(r), layout=bcast((4,), (1,), m)) + out_view = T.tensor_view(T.ptr_of(out), layout=bcast((32,), (1,), m)) + lhs = T.alloc_tensor(Tensor[(4, 8), "f32", bcast((4, 8), (8, 1), m), "smem"]) + col = T.alloc_tensor(Tensor[(4, 1), "f32", bcast((4, 1), (1, 1), m), "smem"]) + dst = T.alloc_tensor(Tensor[(4, 8), "f32", bcast((4, 8), (8, 1), m), "smem"]) T.copy(a_view, lhs) T.copy(r_view, col) T.sync(m) @@ -67,8 +61,8 @@ def two_levels_device(a: Tensor[(_ROWS,), "f32"], out: Tensor[(_ROWS,), "f32"]): ("c", "t"), ) as m: row = ShardLayout(Layout((_ROWS,), (1,)), (Split(0), Split(0)), m) - source = T.tensor_view(a, layout=row) - written = T.tensor_view(out, layout=row) + source = T.tensor_view(T.ptr_of(a), layout=row) + written = T.tensor_view(T.ptr_of(out), layout=row) held = T.alloc_tensor(Tensor[(_ROWS,), "f32", row, "rmem"]) T.copy(source, held) T.binary(held, held, held, kind=BinaryKind.MUL) @@ -87,8 +81,8 @@ def two_axes_device(a: Tensor[(_ROWS,), "f32"], out: Tensor[(_ROWS,), "f32"]): ("w", "t"), ) as m: row = ShardLayout(Layout((_ROWS,), (1,)), (Split(0), Split(0)), m) - source = T.tensor_view(a, layout=row) - written = T.tensor_view(out, layout=row) + source = T.tensor_view(T.ptr_of(a), layout=row) + written = T.tensor_view(T.ptr_of(out), layout=row) held = T.alloc_tensor(Tensor[(_ROWS,), "f32", row, "rmem"]) T.copy(source, held) T.binary(held, held, held, kind=BinaryKind.MUL) @@ -114,7 +108,12 @@ def elementwise_host( axis_out: Tensor[(_ROWS,), "f32"], ): launch( # noqa: F821 - bcast_tier.col_bcast_device, a, r, out, grid=(1, 1, 1), block=(32, 1, 1) # noqa: F821 + bcast_tier.col_bcast_device, + a, + r, + out, + grid=(1, 1, 1), + block=(32, 1, 1), # noqa: F821 ) launch( # noqa: F821 two_levels.two_levels_device, # noqa: F821 diff --git a/tests/ops/tir/test_reduce.py b/tests/ops/tir/test_reduce.py index 76e00063..d28fa6af 100644 --- a/tests/ops/tir/test_reduce.py +++ b/tests/ops/tir/test_reduce.py @@ -25,8 +25,12 @@ class _Plain: @prim_func(target=_CUDA) def plain(a_plain: Tensor[(128, 8), "f32"], out_plain: Tensor[(128,), "f32"]): with Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,)), ("t",)) as mp: - plain_src = T.tensor_view(a_plain, layout=ShardLayout(Layout((128, 8), (8, 1)), (Split(0),), mp)) - plain_dst = T.tensor_view(out_plain, layout=ShardLayout(Layout((128,), (1,)), (Split(0),), mp)) + plain_src = T.tensor_view( + T.ptr_of(a_plain), layout=ShardLayout(Layout((128, 8), (8, 1)), (Split(0),), mp) + ) + plain_dst = T.tensor_view( + T.ptr_of(out_plain), layout=ShardLayout(Layout((128,), (1,)), (Split(0),), mp) + ) T.reduce(plain_src, plain_dst, axes=(1,), kind=ReduceKind.MEAN) @@ -35,8 +39,12 @@ class _Warp: @prim_func(target=_CUDA) def intra_warp(a_warp: Tensor[(32, 4), "f32"], out_warp: Tensor[(1,), "f32"]): with Mesh((Topology("thread", 32),), Layout(shape=(32,), strides=(1,)), ("t",)) as mw: - warp_src = T.tensor_view(a_warp, layout=ShardLayout(Layout((32, 4), (4, 1)), (Split(0),), mw)) - warp_dst = T.tensor_view(out_warp, layout=ShardLayout(Layout((1,), (1,)), (Broadcast(),), mw)) + warp_src = T.tensor_view( + T.ptr_of(a_warp), layout=ShardLayout(Layout((32, 4), (4, 1)), (Split(0),), mw) + ) + warp_dst = T.tensor_view( + T.ptr_of(out_warp), layout=ShardLayout(Layout((1,), (1,)), (Broadcast(),), mw) + ) T.reduce(warp_src, warp_dst, axes=(1,), kind=ReduceKind.ABS_MAX) @@ -44,9 +52,17 @@ def intra_warp(a_warp: Tensor[(32, 4), "f32"], out_warp: Tensor[(1,), "f32"]): class _Cta: @prim_func(target=_CUDA) def intra_cta(a_cta: Tensor[(4, 32, 8), "f32"], out_cta: Tensor[(1,), "f32"]): - with Mesh((Topology("thread", 128),), Layout(shape=(4, 32), strides=(32, 1)), ("w", "t")) as mc: - cta_src = T.tensor_view(a_cta, layout=ShardLayout(Layout((4, 32, 8), (256, 8, 1)), (Split(0), Split(1)), mc)) - cta_dst = T.tensor_view(out_cta, layout=ShardLayout(Layout((1,), (1,)), (Broadcast(), Broadcast()), mc)) + with Mesh( + (Topology("thread", 128),), Layout(shape=(4, 32), strides=(32, 1)), ("w", "t") + ) as mc: + cta_src = T.tensor_view( + T.ptr_of(a_cta), + layout=ShardLayout(Layout((4, 32, 8), (256, 8, 1)), (Split(0), Split(1)), mc), + ) + cta_dst = T.tensor_view( + T.ptr_of(out_cta), + layout=ShardLayout(Layout((1,), (1,)), (Broadcast(), Broadcast()), mc), + ) cta_ws = T.alloc_tensor(Tensor[(4,), "f32", None, "smem"]) T.reduce(cta_src, cta_dst, cta_ws, axes=(2,), kind=ReduceKind.MEAN) @@ -55,9 +71,17 @@ def intra_cta(a_cta: Tensor[(4, 32, 8), "f32"], out_cta: Tensor[(1,), "f32"]): class _Cross: @prim_func(target=_CUDA) def cross_warp(a_cross: Tensor[(4, 32), "f32"], out_cross: Tensor[(1, 32), "f32"]): - with Mesh((Topology("thread", 128),), Layout(shape=(4, 32), strides=(32, 1)), ("w", "t")) as mx: - cross_src = T.tensor_view(a_cross, layout=ShardLayout(Layout((4, 32), (32, 1)), (Split(0), Split(1)), mx)) - cross_dst = T.tensor_view(out_cross, layout=ShardLayout(Layout((1, 32), (32, 1)), (Broadcast(), Split(1)), mx)) + with Mesh( + (Topology("thread", 128),), Layout(shape=(4, 32), strides=(32, 1)), ("w", "t") + ) as mx: + cross_src = T.tensor_view( + T.ptr_of(a_cross), + layout=ShardLayout(Layout((4, 32), (32, 1)), (Split(0), Split(1)), mx), + ) + cross_dst = T.tensor_view( + T.ptr_of(out_cross), + layout=ShardLayout(Layout((1, 32), (32, 1)), (Broadcast(), Split(1)), mx), + ) cross_ws = T.alloc_tensor(Tensor[(128,), "f32", None, "smem"]) T.reduce(cross_src, cross_dst, cross_ws, axes=(0,), kind=ReduceKind.ABS_MAX) @@ -71,10 +95,14 @@ class ReduceTiers: @prim_func(target=CpuTarget()) def reduce_tiers_host( - a_plain: Tensor[(128, 8), "f32"], out_plain: Tensor[(128,), "f32"], - a_warp: Tensor[(32, 4), "f32"], out_warp: Tensor[(1,), "f32"], - a_cta: Tensor[(4, 32, 8), "f32"], out_cta: Tensor[(1,), "f32"], - a_cross: Tensor[(4, 32), "f32"], out_cross: Tensor[(1, 32), "f32"], + a_plain: Tensor[(128, 8), "f32"], + out_plain: Tensor[(128,), "f32"], + a_warp: Tensor[(32, 4), "f32"], + out_warp: Tensor[(1,), "f32"], + a_cta: Tensor[(4, 32, 8), "f32"], + out_cta: Tensor[(1,), "f32"], + a_cross: Tensor[(4, 32), "f32"], + out_cross: Tensor[(1, 32), "f32"], ): launch(plain.plain, a_plain, out_plain, grid=(1, 1, 1), block=(128, 1, 1)) launch(warp.intra_warp, a_warp, out_warp, grid=(1, 1, 1), block=(32, 1, 1)) diff --git a/tests/ops/tir/test_sync.py b/tests/ops/tir/test_sync.py index 4bee2d96..53961b10 100644 --- a/tests/ops/tir/test_sync.py +++ b/tests/ops/tir/test_sync.py @@ -25,8 +25,8 @@ def one_warp(a: Tensor[(1, 32), "f32"], o: Tensor[(1, 32), "f32"]): with Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1))) as full: with full[0:1] as m: sl = ShardLayout(Layout((1, 32), (32, 1)), (S(0), S(1)), m) - src = T.tensor_view(a, layout=sl) - dst = T.tensor_view(o, layout=sl) + src = T.tensor_view(T.ptr_of(a), layout=sl) + dst = T.tensor_view(T.ptr_of(o), layout=sl) reg = T.alloc_tensor(Tensor[(1, 32), "f32", sl, "rmem"]) T.copy(src, reg) T.sync(m) @@ -37,8 +37,8 @@ def two_warps(a: Tensor[(2, 32), "f32"], o: Tensor[(2, 32), "f32"]): with Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1))) as full: with full[2:4] as m: sl = ShardLayout(Layout((2, 32), (32, 1)), (S(0), S(1)), m) - src = T.tensor_view(a, layout=sl) - dst = T.tensor_view(o, layout=sl) + src = T.tensor_view(T.ptr_of(a), layout=sl) + dst = T.tensor_view(T.ptr_of(o), layout=sl) reg = T.alloc_tensor(Tensor[(2, 32), "f32", sl, "rmem"]) T.copy(src, reg) T.sync(m) diff --git a/tests/passes/test_host_entry.py b/tests/passes/test_host_entry.py index ceeb45b3..08ef9b69 100644 --- a/tests/passes/test_host_entry.py +++ b/tests/passes/test_host_entry.py @@ -39,7 +39,7 @@ class _OneKernel: @prim_func(target=_CUDA) def copy_one(x: Tensor[(128,), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=_rows(128)) + view = T.tensor_view(T.ptr_of(x), layout=_rows(128)) T.copy(view, view) T.sync(thread) @@ -55,14 +55,14 @@ def square(x: Tensor[(_S,), "f32"]): @square.specialize(RangePattern("S", 1, 127)) def small(x: Tensor[(_S,), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=_rows(128)) + view = T.tensor_view(T.ptr_of(x), layout=_rows(128)) T.copy(view, view) T.sync(thread) @square.specialize(RangePattern("S", 128, 255)) def large(x: Tensor[(_S,), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=_rows(128)) + view = T.tensor_view(T.ptr_of(x), layout=_rows(128)) T.copy(view, view) T.sync(thread) From 7aca95f0cfdad4d10d507ffcb37aa771cdb6c660 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Fri, 25 Sep 2026 23:22:18 +0800 Subject: [PATCH 6/8] fix(tir): reject pointer view coordinates --- docs/spec/tir.md | 8 +---- src/tilefoundry/ir/tir/cast.py | 2 +- src/tilefoundry/ir/tir/memory/copy.py | 38 +------------------- src/tilefoundry/ir/tir/memory/tensor_view.py | 6 ++-- src/tilefoundry/ir/tir/verify.py | 36 +++++++++++++++++-- 5 files changed, 39 insertions(+), 51 deletions(-) diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 7509fc63..2e62fe9e 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -513,8 +513,6 @@ class TensorView(Op): Attributes: pointer: input; a ``PointerType`` value or an smem byte offset. - coordinates: optional trailing inputs; one absolute element start per - logical window axis (or one absolute flat start for a rank-1 view). dtype: optional attribute; element type stated for a numeric address. storage: optional attribute; storage stated for a numeric address. layout: attribute; the view descriptor. @@ -537,11 +535,7 @@ class TensorView(Op): non-integer numeric addresses are invalid. - A `PointerType` input MAY restate `dtype` or `storage`, but any stated value MUST equal the pointer descriptor. - - With trailing coordinates, codegen derives the view at those absolute - element starts. A coordinate is not a tile ordinal and MUST NOT be - multiplied by the view extent. - - The coordinate count MUST match the logical window rank before any - shard-owned layout axes are removed locally. + - Trailing coordinate inputs are not supported on pointer views. - `T.ptr_of` MAY point at an allocated `ShardTensor`; the view rebuilds over the engine pointer rather than reusing the existing shard layout. diff --git a/src/tilefoundry/ir/tir/cast.py b/src/tilefoundry/ir/tir/cast.py index a329eb53..622986c1 100644 --- a/src/tilefoundry/ir/tir/cast.py +++ b/src/tilefoundry/ir/tir/cast.py @@ -11,7 +11,7 @@ TensorPattern, any_threads, ) -from tilefoundry.ir.tir.memory.copy import verify_between +from tilefoundry.ir.tir.verify import verify_between from tilefoundry.ir.types import StorageKind, UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/memory/copy.py b/src/tilefoundry/ir/tir/memory/copy.py index f0ec45fa..5a50f7af 100644 --- a/src/tilefoundry/ir/tir/memory/copy.py +++ b/src/tilefoundry/ir/tir/memory/copy.py @@ -11,11 +11,7 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import ( - any_threads, - between_rules, - moved_tile, -) +from tilefoundry.ir.pattern import any_threads, moved_tile from tilefoundry.ir.types import LayoutBase, UnitType from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -49,38 +45,6 @@ def _(call: "Call", ctx: "VerifyContext") -> None: ctx.error(call, f"Copy dtype mismatch: {src.dtype} vs {dst.dtype}") -def _input_params(op_type: type) -> tuple: - return tuple(param for param in op_type._op_schema.signature if param.kind == "input") - - -def verify_between(call, ctx, lead: str = "") -> None: - """Hold one call to the relations declared between its operands.""" - op_type = type(call.target) - rules = between_rules(op_type) - if not rules: - return - names = tuple(param.name for param in _input_params(op_type)) - operands = dict(zip(names, (ctx.type_of(arg) for arg in call.args))) - for rule in rules: - if not rule.holds(operands): - ctx.error(call, lead + rule.refused(operands)) - - -def verify_operands(call, ctx, label: str) -> None: - """Hold each operand to the pattern declared for its parameter.""" - for param, arg in zip(_input_params(type(call.target)), call.args): - if param.pattern is None: - continue - value = ctx.type_of(arg) - if param.pattern.match(value) is None: - ctx.error( - call, - f"{label} {param.name} is {tuple(value.shape)} " - f"{value.dtype.name} storage={value.storage}: " - f"{param.pattern.refusal(value)}", - ) - - def _is_copyable_shard(src_ty, dst_ty) -> bool: """Both sides carry a ShardLayout describing the same per-thread buffer.""" src_sl = getattr(src_ty, "layout", None) diff --git a/src/tilefoundry/ir/tir/memory/tensor_view.py b/src/tilefoundry/ir/tir/memory/tensor_view.py index 5095e07f..71090257 100644 --- a/src/tilefoundry/ir/tir/memory/tensor_view.py +++ b/src/tilefoundry/ir/tir/memory/tensor_view.py @@ -2,10 +2,6 @@ Constructs a logical tensor view over a typed pointer. ``layout`` can be a plain ``Layout`` or a ``ShardLayout`` (→ shard tensor view, no allocation). - -A slice view carries an absolute element-start coordinate per axis after the -memory source: a single coordinate is a flat rank-1 window, while multiple -coordinates are a per-axis N-D window. """ from __future__ import annotations @@ -39,6 +35,8 @@ class TensorView(Op): @register_typeinfer(TensorView) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: + if len(call.args) > 1: + ctx.error(call, "coordinates are not supported on a pointer view") pointer = ctx.type_of(call.args[0]) op = call.target origin = call.args[0] diff --git a/src/tilefoundry/ir/tir/verify.py b/src/tilefoundry/ir/tir/verify.py index a8ae1a90..2595188f 100644 --- a/src/tilefoundry/ir/tir/verify.py +++ b/src/tilefoundry/ir/tir/verify.py @@ -16,7 +16,7 @@ ) from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord from tilefoundry.ir.hir.verify import verify_function -from tilefoundry.ir.pattern import RangePattern, locate_dim_var +from tilefoundry.ir.pattern import RangePattern, between_rules, locate_dim_var from tilefoundry.ir.types import DType, TensorType, UnitType from tilefoundry.ir.types.callable_type import callable_type_for_prim_function from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimMax, DimMin, DimMod, DimMul, DimSub @@ -50,6 +50,38 @@ _PRIM_FUNCTION = "[tir §1.3](docs/spec/tir.md#13-primfunction)" +def _input_params(op_type: type) -> tuple: + return tuple(param for param in op_type._op_schema.signature if param.kind == "input") + + +def verify_between(call, ctx, lead: str = "") -> None: + """Hold one call to the relations declared between its operands.""" + op_type = type(call.target) + rules = between_rules(op_type) + if not rules: + return + names = tuple(param.name for param in _input_params(op_type)) + operands = dict(zip(names, (ctx.type_of(arg) for arg in call.args))) + for rule in rules: + if not rule.holds(operands): + ctx.error(call, lead + rule.refused(operands)) + + +def verify_operands(call, ctx, label: str) -> None: + """Hold each operand to the pattern declared for its parameter.""" + for param, arg in zip(_input_params(type(call.target)), call.args): + if param.pattern is None: + continue + value = ctx.type_of(arg) + if param.pattern.match(value) is None: + ctx.error( + call, + f"{label} {param.name} is {tuple(value.shape)} " + f"{value.dtype.name} storage={value.storage}: " + f"{param.pattern.refusal(value)}", + ) + + def verify_prim_function( fn: PrimFunction, *, module_fns: Iterable[PrimFunction] | Module = () ) -> None: @@ -570,4 +602,4 @@ def verify_module(fns) -> None: ) -__all__ = ["verify_prim_function", "verify_module"] +__all__ = ["verify_between", "verify_module", "verify_operands", "verify_prim_function"] From 6271797bc8f37841d55584150cdbd7c83ad684c1 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Fri, 25 Sep 2026 23:44:22 +0800 Subject: [PATCH 7/8] feat(ir): support dynamic tuple indexing --- docs/spec/codegen.md | 15 ++++ docs/spec/hir.md | 15 ++-- docs/spec/parser.md | 13 ++- src/tilefoundry/codegen/context.py | 15 ++++ .../codegen/cuda/tir/tuple_get_item.py | 61 +++++++++++++ src/tilefoundry/codegen/emitter.py | 12 ++- src/tilefoundry/inspection/printer_base.py | 7 ++ src/tilefoundry/inspection/python_printer.py | 33 +++++-- src/tilefoundry/inspection/viewer/builder.py | 11 ++- .../ir/hir/tensor/tuple_get_item.py | 85 ++++++++++++++++--- src/tilefoundry/parser/pattern_nodes.py | 14 +-- tests/ops/ir/test_topk.py | 6 +- tests/parser/test_calls.py | 4 +- 13 files changed, 252 insertions(+), 39 deletions(-) create mode 100644 src/tilefoundry/codegen/cuda/tir/tuple_get_item.py diff --git a/docs/spec/codegen.md b/docs/spec/codegen.md index deaf194b..1b2b3a09 100644 --- a/docs/spec/codegen.md +++ b/docs/spec/codegen.md @@ -192,6 +192,21 @@ class CpuCodegenContext(CodegenContext): ### 2.4 Effect Op dispatch +CUDA value emission lowers `TupleGetItem` according to its index form. A constant +index names the selected tuple element directly. A dynamic index requires the +homogeneous tuple established by type inference and emits +`cute::array{a, b, c}[index]`; dimension arithmetic in the index remains runtime +C++ arithmetic. + +An IR `Tuple` is structural and has no target-side storage. A `LetStmt` binding +one emits no C++ variable and continues with its body; its authored name (for +example `lhs_stages`) therefore does not appear in generated C++. Each dynamic +`TupleGetItem` use constructs its own `cute::array`, so indexing the same tuple +N times constructs N arrays. Constant selection remains valid for heterogeneous +tuples because it does not materialize an aggregate. The codegen context records +the structural tuple by its fresh SSA `Var` identity so consumers can recover +its elements without target-side storage. + Effect Ops (`Copy`, `Fill`, `Mma`, `tir.nn.*`, ...) appear in Stmt position as `Evaluate(op, args)` rather than as Stmt subclasses. The walker matches `Evaluate` and dispatches on `type(callable)` through diff --git a/docs/spec/hir.md b/docs/spec/hir.md index f05000f0..b3f19f6b 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -943,20 +943,25 @@ class Stack(Op): ```python class TupleGetItem(Op): - """Extract one field from a tuple-typed expression. + """Extract one field from a tuple-typed expression by scalar index. Attributes: tuple_value: input; Tuple-typed expression. - index: attribute; Static field index. + index: input; Scalar field index. """ tuple_value: Expr - index: int + index: Expr ``` - constraints: - - `tuple_value.type` MUST be `TupleType` and `index` MUST be in range. - - The result type MUST be exactly the selected field type. + - `tuple_value.type` MUST be `TupleType`. A constant `index` MUST be in range, + and the result type is exactly the selected field type. + - A dynamic `index` requires a non-empty tuple whose field types are all equal; + the result has that common field type. + - Access through a constant index reads the selected field's leaf span. Dynamic + access conservatively reads every leaf because the selected field is known + only at runtime. ##### Reshape ```python diff --git a/docs/spec/parser.md b/docs/spec/parser.md index c8005e90..b67d20ff 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -48,10 +48,10 @@ physical source-file coordinates with a one-based start column. ### 1.3 Tuple Binding Metadata -For `a, b = producer(...)`, detached `TupleGetItem(index=0)` and -`TupleGetItem(index=1)` lexical values carry the respective target Name spans (`a` and `b`) and -matching `BindingMetadata`; later reads do not replace that identity. A multi-carry loop's -derived projections carry the `for` statement span and their carry binding name. +For `a, b = producer(...)`, detached `TupleGetItem` values with scalar inputs `0` and `1` +carry the respective target Name spans (`a` and `b`) and matching `BindingMetadata`; later +reads do not replace that identity. A multi-carry loop's derived projections carry the `for` +statement span and their carry binding name. ### 1.4 Context and Diagnostics @@ -83,6 +83,11 @@ lowered from, so a bound naming a mesh coordinate reads back as it was printed. ## 2. Syntax and Rules +Tuple subscripting lowers `stages[index]` to `TupleGetItem(stages, index)`. +Literal negative indices are normalized against the tuple arity before lowering. +A non-literal scalar index is accepted only when type inference can establish one +common field type for every possible result. + ### 2.1 Syntax diff --git a/src/tilefoundry/codegen/context.py b/src/tilefoundry/codegen/context.py index 7f06cd04..12fd7fc3 100644 --- a/src/tilefoundry/codegen/context.py +++ b/src/tilefoundry/codegen/context.py @@ -12,6 +12,7 @@ from collections.abc import Callable, Mapping from tilefoundry.codegen.signature import CallableSignature, Signature, TensorSignature +from tilefoundry.ir.core import Tuple from tilefoundry.ir.tir.stmts import Evaluate from tilefoundry.target.base import Target from tilefoundry.visitor_registry.registries import DispatchRegistry, Role, spelled @@ -48,6 +49,7 @@ def __init__( self._lines: list[str] = [] self._indent = 0 self._var_names: dict[int, str] = {} + self._tuple_values: dict[int, Tuple] = {} self._counter = 0 self._kernel_param_ids: set[int] = set() @@ -74,6 +76,19 @@ def register_kernel_param(self, var) -> None: self._var_names[key] = var.name self._kernel_param_ids.add(key) + def register_tuple(self, var, value: Tuple) -> None: + """Record the structural Tuple bound to one fresh SSA variable.""" + self._tuple_values[id(var)] = value + + def tuple_for(self, var) -> Tuple: + """Return the structural Tuple bound to *var*, refusing an unknown value.""" + try: + return self._tuple_values[id(var)] + except KeyError: + raise KeyError( + f"codegen: {getattr(var, 'name', var)!r} has no structural tuple binding" + ) from None + def is_kernel_param(self, var) -> bool: return id(var) in self._kernel_param_ids diff --git a/src/tilefoundry/codegen/cuda/tir/tuple_get_item.py b/src/tilefoundry/codegen/cuda/tir/tuple_get_item.py new file mode 100644 index 00000000..453753af --- /dev/null +++ b/src/tilefoundry/codegen/cuda/tir/tuple_get_item.py @@ -0,0 +1,61 @@ +"""CUDA value emission for tuple indexing.""" + +from __future__ import annotations + +from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.ir.core import Call, Constant, Tuple, Var +from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem +from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimMod, DimMul, DimSub +from tilefoundry.target import CudaTarget +from tilefoundry.visitor_registry.registries import Role, register_codegen + +_DIM_SYMBOLS = { + DimAdd: "+", + DimSub: "-", + DimMul: "*", + DimFloorDiv: "/", + DimMod: "%", +} + + +def _value(expr, ctx: CudaCodegenContext) -> str: + if isinstance(expr, Constant): + if isinstance(expr.value, bool): + return "true" if expr.value else "false" + return str(expr.value) + name = ctx.name_for(expr) + return f"{name}_tensor" if ctx.is_kernel_param(expr) else name + + +def _index(expr, ctx: CudaCodegenContext) -> str: + if isinstance(expr, Constant) and isinstance(expr.value, int): + return str(expr.value) + if isinstance(expr, Var): + return ctx.name_for(expr) + if isinstance(expr, Call): + for op_type, symbol in _DIM_SYMBOLS.items(): + if isinstance(expr.target, op_type): + lhs, rhs = expr.args + return f"({_index(lhs, ctx)} {symbol} {_index(rhs, ctx)})" + raise NotImplementedError( + f"CUDA TupleGetItem index {type(expr).__name__} is not scalar dimension arithmetic" + ) + + +@register_codegen(CudaTarget, Role.EMIT, TupleGetItem) +def _emit(let_stmt, ctx: CudaCodegenContext) -> None: + call = let_stmt.value + held, index = call.args + if isinstance(held, Var): + held = ctx.tuple_for(held) + if not isinstance(held, Tuple): + raise ValueError("CUDA TupleGetItem requires a tuple value") + result = ctx.name_for(let_stmt.var) + if isinstance(index, Constant): + selected = index.value + if not isinstance(selected, int) or isinstance(selected, bool): + raise ValueError("CUDA TupleGetItem constant index must be an integer") + ctx.emit(f"auto {result} = {_value(held.elements[selected], ctx)};") + return + elements = ", ".join(_value(element, ctx) for element in held.elements) + ctx.emit(f"auto {result} = cute::array{{{elements}}}[{_index(index, ctx)}];") diff --git a/src/tilefoundry/codegen/emitter.py b/src/tilefoundry/codegen/emitter.py index 0c86048b..119d3a17 100644 --- a/src/tilefoundry/codegen/emitter.py +++ b/src/tilefoundry/codegen/emitter.py @@ -9,7 +9,7 @@ program_topologies, ) from tilefoundry.codegen.cuda.tir.stmts.scalar_expr import render_scalar_expr -from tilefoundry.ir.core import Call +from tilefoundry.ir.core import Call, Tuple from tilefoundry.ir.tir.abort import Abort from tilefoundry.ir.tir.stmts import ( Evaluate, @@ -88,6 +88,16 @@ def visit_Abort(self, node: Abort) -> None: self.statement("__trap()") def visit_LetStmt(self, node: LetStmt) -> None: + """Emit a stored call result, or pass through a structural tuple binding. + + A Tuple has no storage of its own. Its consumers lower the elements + directly, so its authored name is absent from generated C++ and each + TupleGetItem use may materialize its own target-side aggregate. + """ + if isinstance(node.value, Tuple): + self.context.register_tuple(node.var, node.value) + self.visit(node.body) + return if not isinstance(node.value, Call): raise RuntimeError( f"LetStmt.value must be a Call (TIR-owned Expr Op), got {type(node.value).__name__}" diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index adb0458c..591b8c35 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -8,6 +8,7 @@ from tilefoundry.ir.core import Call, Constant, Tuple, Var from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord +from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.mesh_scope import device_layout from tilefoundry.ir.pattern import Pattern, RangePattern from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom @@ -133,6 +134,8 @@ def visit_Call(self, value: Call, ctx=None) -> str: left, right = ceildiv_args return f"ceildiv({self.dim_entry(left, ctx)}, {self.dim_entry(right, ctx)})" target = value.target + if isinstance(target, TupleGetItem): + return self._tuple_get_item_text(value, ctx) if isinstance(target, MeshCoord): return self._mesh_coordinate_text(value, target, ctx) if isinstance(target, DimConst): @@ -151,6 +154,10 @@ def visit_Call(self, value: Call, ctx=None) -> str: return f"{name}({args})" return self.visit_program_call(value, ctx) + def _tuple_get_item_text(self, value: Call, ctx=None) -> str: + held, index = value.args + return f"{self.visit(held, ctx)}[{self.visit(index, ctx)}]" + def _mesh_coordinate_text(self, value: Call, target: MeshCoord, ctx) -> str: """Render one coordinate through the active binding of its mesh.""" axis = static_dim_value(value.args[0]) if value.args else None diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index 90be24cf..49684367 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -110,16 +110,20 @@ def reference(self, expr: Expr) -> str: 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])] + return self._names[id(projection.carried_args[_projection_index(expr)])] 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]) + return self.reference(projection.body.elements[_projection_index(expr)]) if isinstance(expr, MeshRegion): return self.reference(expr.body) return self._names[id(expr)] + def _tuple_get_item_text(self, expr: Call, ctx=None) -> str: + held, index = expr.args + return f"{self.reference(held)}[{self.visit(index, ctx)}]" + def _slice_start(self, start, size, stride) -> str: moved = self._moved_window(start, size, stride) if moved is None: @@ -526,9 +530,13 @@ def iter_exprs(root: Expr | None, seen: set[int] | None = None) -> Iterator[Expr def _region_projection(expr: Expr) -> LoopRegion | MeshRegion | None: - """Return the region projected by a one-argument ``TupleGetItem``.""" + """Return the region projected by a constant-index ``TupleGetItem``.""" if not ( - isinstance(expr, Call) and isinstance(expr.target, TupleGetItem) and len(expr.args) == 1 + isinstance(expr, Call) + and isinstance(expr.target, TupleGetItem) + and len(expr.args) == 2 + and isinstance(expr.args[1], Constant) + and isinstance(expr.args[1].value, int) ): return None region = expr.args[0] @@ -539,6 +547,13 @@ def _region_projection(expr: Expr) -> LoopRegion | MeshRegion | None: return None +def _projection_index(expr: Call) -> int: + """The literal field selected by a region projection.""" + index = expr.args[1] + assert isinstance(index, Constant) and isinstance(index.value, int) + return index.value + + def _module_callee_binding(target: HirFunction, child_entries: dict[int, str]) -> str | None: """The attribute a call on *target* was written through, if a child's entry. @@ -690,6 +705,13 @@ def _moved_window(start, size, stride): ) if _moved_window(start, size, stride) is not None } + _tuple_index_ids: set[int] = set() + for expr in _order: + if isinstance(expr, Call) and isinstance(expr.target, TupleGetItem): + for nested in iter_exprs(expr.args[1], set()): + if not isinstance(nested, Var): + _tuple_index_ids.add(id(nested)) + _inlined_start_ids.update(_tuple_index_ids) def _assign_name(expr: Expr) -> str: key = id(expr) @@ -719,7 +741,8 @@ def _assign_name(expr: Expr) -> str: return name for expr in _order: - _assign_name(expr) + if id(expr) not in _tuple_index_ids: + _assign_name(expr) for expr in _order: if isinstance(expr, LoopRegion): for carry in expr.carried_args: diff --git a/src/tilefoundry/inspection/viewer/builder.py b/src/tilefoundry/inspection/viewer/builder.py index bfa7ce49..4b81a046 100644 --- a/src/tilefoundry/inspection/viewer/builder.py +++ b/src/tilefoundry/inspection/viewer/builder.py @@ -686,9 +686,14 @@ def _emit_call( - tuple_index = ( - call.target.index if _op_display_name(call.target) == "TupleGetItem" else None - ) + tuple_index = None + if ( + _op_display_name(call.target) == "TupleGetItem" + and len(call.args) == 2 + and isinstance(call.args[1], Constant) + and isinstance(call.args[1].value, int) + ): + tuple_index = call.args[1].value for i, arg in enumerate(call.args): src = self._walk_expr( g, arg, call_path=call_path, visited=visited, local_counter=local_counter diff --git a/src/tilefoundry/ir/hir/tensor/tuple_get_item.py b/src/tilefoundry/ir/hir/tensor/tuple_get_item.py index 33f6a641..1bc0b0a1 100644 --- a/src/tilefoundry/ir/hir/tensor/tuple_get_item.py +++ b/src/tilefoundry/ir/hir/tensor/tuple_get_item.py @@ -1,31 +1,75 @@ from __future__ import annotations +import isl + from tilefoundry.evaluator.registry import register_eval -from tilefoundry.ir.core import Op +from tilefoundry.evaluator.value import EvalError, TensorValue +from tilefoundry.ir.core import Constant, Op from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import Tensor +from tilefoundry.ir.pattern import Scalar, Tensor from tilefoundry.ir.types import TupleType from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( + AccessRelations, + AffineAccess, + BoundaryRelation, + control_read, + identity_access, + iterating, + leaf_span, + leaves_of, register_access_relation, - view_relations, ) @register_op(name="tuple_get_item") class TupleGetItem(Op): - """Extract a field of a tuple-typed expression by static index. + """Extract a field of a tuple-typed expression by scalar index. Representing extraction as a Call keeps multi-output consumers in the HIR SSA expression model. See [hir §1](docs/spec/hir.md#1-hir-expr-constructs). """ tuple_value = ParamDef(kind="input", pattern=Tensor) - index = ParamDef(kind="attribute", annotation=int) + index = ParamDef(kind="input", pattern=Scalar) -register_access_relation(TupleGetItem)(view_relations(0, field=lambda call, ctx: call.target.index)) +@register_access_relation(TupleGetItem) +def _access_relations(call: "Call", ctx: "AccessContext") -> AccessRelations: + held = ctx.type_of(call.args[0]) + if not isinstance(held, TupleType): + raise ValueError("TupleGetItem access requires a TupleType operand") + index = call.args[1] + if isinstance(index, Constant): + taken = index.value + if ( + not isinstance(taken, int) + or isinstance(taken, bool) + or not 0 <= taken < len(held.fields) + ): + raise ValueError(f"TupleGetItem index {taken!r} out of range") + result = held.fields[taken] + begin, count = leaf_span(held, taken) + else: + if not held.fields: + raise ValueError("TupleGetItem dynamic access requires a non-empty TupleType") + result = held.fields[0] + begin, count = 0, len(leaves_of(held)) + walks = getattr(result, "shape", ()) or () + rank = len(walks) + coordinates = ", ".join(f"d{axis}" for axis in range(rank)) + reads = AffineAccess(isl.map(f"{{ [{coordinates}] -> [l] : {begin} <= l < {begin + count} }}")) + return iterating( + walks, + AccessRelations( + inputs=( + BoundaryRelation(reads), + BoundaryRelation(control_read(rank, ctx, index)), + ), + outputs=(BoundaryRelation(identity_access(rank)),), + ), + ) @register_typeinfer(TupleGetItem) @@ -33,15 +77,34 @@ def _(call: "Call", ctx: "TypeInferContext"): tup_ty = ctx.type_of(call.args[0]) if not isinstance(tup_ty, TupleType): ctx.error(call, "TupleGetItem on non-TupleType") - idx = call.target.index - if idx < 0 or idx >= len(tup_ty.fields): - ctx.error(call, f"TupleGetItem index {idx} out of range") - return tup_ty.fields[idx] + index = call.args[1] + if isinstance(index, Constant): + taken = index.value + if ( + not isinstance(taken, int) + or isinstance(taken, bool) + or taken < 0 + or taken >= len(tup_ty.fields) + ): + ctx.error(call, f"TupleGetItem index {taken!r} out of range") + return tup_ty.fields[taken] + if not tup_ty.fields: + ctx.error(call, "TupleGetItem dynamic index requires a non-empty TupleType") + first = tup_ty.fields[0] + if any(field != first for field in tup_ty.fields[1:]): + ctx.error(call, "TupleGetItem dynamic index requires homogeneous fields") + return first @register_eval(TupleGetItem) def _eval_tuple_get_item(ctx): - return ctx.args[0].elements[ctx.op.index] + index = ctx.args[1] + if not isinstance(index, TensorValue) or index.data.numel() != 1: + raise EvalError("evaluator: TupleGetItem index is a single integer") + taken = int(index.data.reshape(-1)[0].item()) + if not 0 <= taken < len(ctx.args[0].elements): + raise EvalError(f"evaluator: TupleGetItem index {taken} out of range") + return ctx.args[0].elements[taken] __all__ = ["TupleGetItem"] diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index c92f3f03..b9de19c9 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -3097,12 +3097,14 @@ def construct(match, children, context): value = children["value"] index = children["index"] if isinstance(value.type, runtime.TupleType): - if isinstance(index, bool) or not isinstance(index, int): + if isinstance(index, bool) or not isinstance(index, (int, runtime.Expr)): raise ParseError.from_node( - match.node, context, "Tuple subscript requires an integer literal" + match.node, context, "Tuple subscript requires an integer expression" ) - normalized = index + len(value.type.fields) if index < 0 else index - return _infer_call(runtime.TupleGetItem(index=normalized), (value,), context) + if isinstance(index, int): + normalized = index + len(value.type.fields) if index < 0 else index + index = _constant(normalized) + return _infer_call(runtime.TupleGetItem(), (value, index), context) if not isinstance(value.type, runtime.TensorType): raise ParseError.from_node( match.node, context, "subscript requires TensorType or TupleType" @@ -3592,7 +3594,7 @@ def _bind_region_results(context, region, names, node): context.lexical_scope.define(names[0], region) return for index, name in enumerate(names): - projection = _infer_call(runtime.TupleGetItem(index=index), (region,), context) + projection = _infer_call(runtime.TupleGetItem(), (region, _constant(index)), context) attach_authored_metadata( projection, node, @@ -4524,7 +4526,7 @@ def construct(match, children, context): attach_metadata(value, BindingMetadata(parent_name)) for index, (name, target_node) in enumerate(zip(names, target_nodes)): assert isinstance(target_node, ast.Name) - projection = _infer_call(runtime.TupleGetItem(index=index), (value,), context) + projection = _infer_call(runtime.TupleGetItem(), (value, _constant(index)), context) attach_authored_metadata( projection, target_node, diff --git a/tests/ops/ir/test_topk.py b/tests/ops/ir/test_topk.py index 014a5913..08d52fae 100644 --- a/tests/ops/ir/test_topk.py +++ b/tests/ops/ir/test_topk.py @@ -24,7 +24,7 @@ run_typeinfer_case, ) from tilefoundry.evaluator import evaluate -from tilefoundry.ir.core import Call, Var +from tilefoundry.ir.core import Call, Constant, Var from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.tensor.index_select import IndexSelect from tilefoundry.ir.hir.tensor.reshape import Reshape @@ -34,6 +34,7 @@ DType, Layout, Mesh, + TensorType, Topology, TupleType, make_shard_tensor_type, @@ -280,7 +281,8 @@ def test_topk_dynamic_k_downstream_index_select_shape_consistent(): topk_ty = TypeInferVisitor().visit(topk_call, TypeInferContext()) topk_call = replace(topk_call, type=topk_ty) - idx_call = Call(type=topk_ty.fields[1], target=TupleGetItem(index=1), args=(topk_call,)) + index = Constant(type=TensorType.umat_scalar(), value=1) + idx_call = Call(type=topk_ty.fields[1], target=TupleGetItem(), args=(topk_call, index)) idx_ty = TypeInferVisitor().visit(idx_call, TypeInferContext()) idx_call = replace(idx_call, type=idx_ty) diff --git a/tests/parser/test_calls.py b/tests/parser/test_calls.py index f1530585..b5f077d9 100644 --- a/tests/parser/test_calls.py +++ b/tests/parser/test_calls.py @@ -316,8 +316,8 @@ def multi_escape(x: Tensor[(2,), "f32"]): assert isinstance(scope, MeshRegion) assert isinstance(scope.body, Tuple) assert len(scope.body.elements) == 2 - assert first.target.index == 0 - assert second.target.index == 1 + assert isinstance(first.args[1], Constant) and first.args[1].value == 0 + assert isinstance(second.args[1], Constant) and second.args[1].value == 1 def test_valueful_mesh_region_also_wraps_escaping_bindings() -> None: From a0b8ae7151a0418363a88371e2c4523195ed590c Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Fri, 25 Sep 2026 23:51:28 +0800 Subject: [PATCH 8/8] test(inspection): update tuple indexing snapshots --- tests/fixtures/inspection/type_printer_sugar.analyzed.txt | 6 +++--- tests/inspection/test_module_tree_roundtrip.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index d42a40dc..c0bec297 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -6,7 +6,7 @@ from tilefoundry.dsl.storage import gmem, rmem, smem from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.ir.types 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("thread", 8),), Layout((2, 4), (4, 1)), names=("warp", "lane"))) def composed_mesh_pipeline( x: Tensor[(8, 4, 16), "f32"], seed: Tensor[(16,), "f32"], @@ -15,10 +15,10 @@ def composed_mesh_pipeline( 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')), + 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: + 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; memory traffic=gmem:r2.00KB/w0@logical,r2.00KB/w0@total,r512B/w0@cta,r64B/w0@thread;rmem:r0/w2.00KB@logical,r0/w2.00KB@total,r0/w512B@cta,r0/w64B@thread footprint=x:2.00KB v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@logical,512@total,128@cta,16@thread; memory traffic=rmem:r2.00KB/w2.00KB@logical,r2.00KB/w2.00KB@total,r512B/w512B@cta,r64B/w64B@thread v2 = reshard(v1, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost; memory traffic=rmem:r2.00KB/w0@logical,r2.00KB/w0@total,r512B/w0@cta,r64B/w0@thread;smem:r0/w2.00KB@logical,r0/w2.00KB@total,r0/w512B@cta,r0/w64B@thread diff --git a/tests/inspection/test_module_tree_roundtrip.py b/tests/inspection/test_module_tree_roundtrip.py index ede221bf..4e59e15c 100644 --- a/tests/inspection/test_module_tree_roundtrip.py +++ b/tests/inspection/test_module_tree_roundtrip.py @@ -304,6 +304,6 @@ def test_hir_function_dot_keeps_loop_regions_as_opaque_leaves() -> None: """The public DOT form keeps structured regions as white leaf boxes.""" dot = hir_function_to_dot(static_online_attend.entry_function()) - assert len(dot.splitlines()) == 29 + assert len(dot.splitlines()) == 33 assert 'label="LoopRegion", fillcolor="#ffffff"' in dot assert "TupleType(" not in dot