diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index f9c06acf..d5de9def 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -1323,8 +1323,10 @@ widening; the two modules do not define a second affine graph representation. parameter is already fixed to one integer point. An empty parameter context MUST count as zero; a non-empty context with an unbounded parameter MUST report unknown. - - `IterationScope.trips()` MUST fix child and parent domains to the same parameter - point before dividing, and take the maximum of those ratios. + - `IterationScope.trips()` MUST count a loop off its bounds, not its domain: + `max(1, ceil(max(stop - start) / step))`, refusing an unbounded span. + `trips_precision` is `EXACT` where that span reaches one value and + `WIDENED` where it reaches several, the count then being an upper bound. ### 2.2 Target-selected Analyzers diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index e1444dca..31123dc6 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -19,7 +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/shard/` | [shard](./shard.md) | Shard / layout sublayer: `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`). The physical nesting reflects the spec's "sublayer" relationship. | +| `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/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`. | @@ -74,14 +75,11 @@ physical directory layout reflects that boundary directly. **Reading notes:** - `ir/` holds the IR proper and its sublayers only. `ir/types/` is the - root of the type system; `ir/types/shard/` is its shard / layout - sublayer ([architecture §3](./architecture.md#3-type-system)). The - physical nesting reflects the spec's conceptual "sublayer". -- The placement of `shard/` under `types/` is a filing decision, not a - consumer restriction: `Topology` / `Mesh` / `Layout` / `ShardLayout` - are consumed directly by `parser`, `tir`, and `codegen`. The - hierarchy expresses "role in the type system", not "who may import - it". + root of the type system, and the shard / layout types sit in it + directly ([architecture §3](./architecture.md#3-type-system)): they + are types, and a folder naming them a sublayer said nothing a reader + could use. Their filing follows CuTe's, one module per thing the + algebra is written against. - `codegen/` and `parser/` sit outside `ir/`. By the [architecture §1](./architecture.md#1-spec-relationship-map) pipeline they are the front-end producer and back-end consumer of diff --git a/docs/spec/hir.md b/docs/spec/hir.md index 462dc927..2626fac6 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -663,12 +663,14 @@ Tensor structural operations; consensus ops (`Transpose` / `Slice` / `Concat` ([torch tensor manipulation ops](https://pytorch.org/docs/stable/torch.html#indexing-slicing-joining-mutating-ops)). `Transpose`, statically positioned `Slice`, and `Reshape` derive a view layout from -their input when it states one. An input with `layout=None` produces a view with -`layout=None`. Neither case says that the view materialized. +their input when it states one. For `Slice` and `Reshape`, an input with +`layout=None` produces a view with `layout=None`. Neither case says that the +view materialized. - `Transpose` MUST permute the layout shape and strides by the same permutation as the tensor shape. A `ShardLayout` MUST remap its split positions through - the registered relation. + the registered relation. An input with `layout=None` MUST permute the C order + it stands for and state the result as a `Layout`. - `Slice` is normalized as `Slice(x, starts, sizes=..., strides=...)`. `starts` is a tuple of rank-0 integer operands; `sizes` and `strides` are `ShapeDim` attributes stored in the same IR normal form as every other dim. @@ -681,6 +683,11 @@ their input when it states one. An input with `layout=None` produces a view with offset is the source offset plus the starts multiplied by the source strides, and its outer layout carries the sliced shape and retained strides (multiplied by any slice step). +- A source axis written as a group of modes has no single stride: its start + MUST be read through the group as a mixed-radix coordinate, and the window + MUST be the group's least-stepping modes, as many as its size takes, in the + order the group wrote them. A start that is not a multiple of the size, a + size landing on no mode boundary, and a step MUST each be refused. - A `ShardLayout` slice MUST preserve its mesh attributes when every narrowed logical axis is unsplit. The corresponding primitive layout position takes the window size and stepped stride. Static starts wrap that shard layout in a diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index 1b7b53e9..21222f33 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -185,12 +185,13 @@ placement sugar of [parser §2.1](./parser.md#21-syntax), in both type slots and op-attribute slots, when every mesh identifier in that sugar has an explicit scope binding. A binding is either a `with as ` region or the function's own execution domain. A mesh merely restated in another expression -is not a binding. Without such a binding, the printer MUST use the verbose -`ShardLayout(...)` form rather than inventing a name. +is not a binding, and the printer never invents a name for one. -Placement sugar states the layout's own dimensions with each `Split` written on -the dimension it divides, adds the stride tuple whenever the layout has one, -and states the remaining mesh axes in a `{axis @ ...}` set. Because the parser +Placement sugar states one dimension per tensor axis, with each `Split` written +on the dimension it divides, adds the stride tuple whenever the layout has one, +and states the remaining mesh axes in a `{axis @ ...}` set. A layout that groups +an axis's modes has no sugar; elsewhere a group is written as the shape tuple +it is. Because the parser reads an unstated mesh axis as `Broadcast`, the set carries every `Partial` and carries `Broadcast` only when no `Split` or `Partial` would otherwise name the mesh. diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 1dae50af..c8005e90 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -75,6 +75,12 @@ inside a region but bound outside it are captured as `MeshRegion.args`, with a fresh `MeshRegion.params` binding used by the body. Capture is performed one region boundary at a time, so nested regions pass a value through each door. +A loop body holds `with Mesh(...)` statements, and the loop carries what one +binds. Because the body repeats, a name the `with` reads on its way to binding +it escapes as a name read after it does, and the loop carries both. A TIR loop +bound that is not a literal is read as the dimension arithmetic the loop was +lowered from, so a bound naming a mesh coordinate reads back as it was printed. + ## 2. Syntax and Rules ### 2.1 Syntax diff --git a/docs/spec/shard.md b/docs/spec/shard.md index 8ab3d31b..cf45d465 100644 --- a/docs/spec/shard.md +++ b/docs/spec/shard.md @@ -285,7 +285,7 @@ class Mesh: Attributes: topologies: attribute; Ordered topology sequence. - layout: attribute; Mesh layout or constant sliced layout. + layout: attribute; One nested arrangement, mode i being level i. names: attribute; Human-readable layout-axis names. """ @@ -303,42 +303,51 @@ class Mesh: Field meanings: - `topologies` — the ordered device-domain descriptions (`name` + `size`) -- `layout` — the mesh's own shape / strides (a `Layout`); a constant slice - (`m[...]`) replaces it with a `ComposedLayout` recording the sub-box - ([tir §1.5](./tir.md#15-sync)) +- `layout` — a nested `Layout` whose mode `i` is level `i`'s own arrangement; a + constant slice (`m[...]`) replaces it with a `ComposedLayout` recording the + sub-box ([tir §1.5](./tir.md#15-sync)) - `names` — optional human-readable names (`cta.x`, `cta.y`, …) -`Mesh` describes the parallel device domain; it is not a tensor layout -object. +Every mesh MUST state one arrangement per level it names, each in that level's +own numbering, as one nested `Layout` whose mode `i` is level `i` -- a mesh +naming one level included. `Mesh` normalizes what was written into that form at +construction: one arrangement over all the levels' axes is cut at the level +boundaries, each level taking axes left to right until their extents multiply to +its own size and its steps divided by what the levels under it hold. An axis of +extent one adds no positions and joins the level being filled. A boundary no +prefix of axes lands on MUST be refused. + +Because the layout is a nested `Layout`, the layout operations read it: `get(mesh.layout, i)` +is level `i`'s own arrangement, `flatten(mesh.layout)` is every level's axes end +to end, and `repeat_like` against `mesh.layout.shape` states per axis whatever is +stated per level. `device_layout(mesh)` is those axes as the device numbers its +positions, each stepping by what its level states times what the levels under it +hold, and a slice's `ComposedLayout.offset` is where the whole mesh starts in +that numbering. Mesh composition uses the following rules: -- `composed(meshes)` MUST replace all levels in force when the inner mesh names - every current level, and MUST concatenate meshes whose level-name sets are - disjoint. For concatenation, each outer stride and offset is scaled by the - product of positions in the levels below it: - `offset = outer_offset * below + inner_offset`. -- A sliced mesh is composable; its slice offset remains a - `ComposedLayout` offset after composition. A mesh whose level names partially - overlap the levels in force MUST be rejected rather than decomposed. -- `composed(meshes)` invokes `check_topology` on its result. For each named - level with a concrete declared extent, its projected position count MUST NOT - exceed that extent; symbolic extents are deferred until dimensions are bound. - This check does not reject same-level nesting, because replacement determines - the final position count. +- `merge_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 + 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. 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 value is consumed by a region. -`Mesh` MAY carry more than one `Topology` (e.g. `warp(4) x thread(32)`); the -full sequence is always `topologies`. - - constraints: - `Topology` construction rejects a `None` size. `Mesh` construction rejects - a `None` entry in its layout shape. Beyond that explicit-extent check, - `Mesh` is a frozen record: it performs no construction-time normalization - or position-consistency check. Its `topologies` field is a + 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. - The author surface is `with Mesh(("cta",), layout=(128,)) as cta:`. The @@ -358,8 +367,9 @@ full sequence is always `topologies`. - The layout of the positions one level has is the axes up to and including that level's segment, with their strides divided by the product of the sizes of the levels below it; a stride that division does not divide exactly MUST - be refused. A Mesh naming one level states its own layout and is not - projected. A Mesh naming several MUST NOT also be sliced. + 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. - 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 @@ -385,7 +395,7 @@ full sequence is always `topologies`. 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 `composed()`'s +node. This is a layout-construction rule, independent of `merge_mesh()`'s scope-composition rules. ### 5.1 `Placement` @@ -762,6 +772,10 @@ copy that displacement to a materialized consumer. ## 9. Layout construction and mesh-scope projection +Making the steps of a compact arrangement and reading an index back into the +coordinate that reaches it are not operations on layouts, so they are filed +apart from the algebra, as CuTe files `stride.hpp` apart from `layout.hpp`. + ```python class NotProjectable(ValueError): """Report that a layout cannot serve as a mesh execution scope.""" diff --git a/docs/spec/tir.md b/docs/spec/tir.md index a06e5309..1565cf45 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -179,6 +179,9 @@ class PrimFunction(Stmt): - **`MeshScope` mesh in scope**. Any embedded `ShardLayout` MUST reference a mesh on the active TIR `MeshScope` traversal cache or a parameter's `ShardLayout.mesh`. +- **`For` bound coordinates**. A `MeshCoord` read by `For.start` / `.stop` / + `.step` MUST name a literal in-range axis of a mesh bound by an enclosing + `MeshScope`. - **`Evaluate.callable`**. When `callable` is a `SymbolRef` ([§2.1](#21-symbolref)), module-level resolution MUST find exactly one `PrimFunction` of that name in the enclosing `Module`, `args` length diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index 259fde54..34c5b0b4 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.shard 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, 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" }, { "cell_type": "markdown", diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index c58524b7..8361e841 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -83,7 +83,7 @@ from tilefoundry.analysis import analyze as run_analysis from tilefoundry.dsl import ConstTensor, DimVar, DimVarRangePat, 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.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget HIDDEN = 256 diff --git a/src/tilefoundry/__init__.py b/src/tilefoundry/__init__.py index 285bb204..caecfadc 100644 --- a/src/tilefoundry/__init__.py +++ b/src/tilefoundry/__init__.py @@ -37,24 +37,7 @@ from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern from tilefoundry.ir.types import DType, TensorType, TupleType, Type from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import ( - B, - Broadcast, - ComposedLayout, - Swizzle, - Dynamic, - IntTuple, - Layout, - LayoutBase, - Mesh, - P, - Partial, - S, - ShardAttr, - ShardLayout, - Split, - 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 diff --git a/src/tilefoundry/analysis/access.py b/src/tilefoundry/analysis/access.py index f88d44cc..f09fb250 100644 --- a/src/tilefoundry/analysis/access.py +++ b/src/tilefoundry/analysis/access.py @@ -14,8 +14,7 @@ from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.isl_interop import index_set from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.utils import local_type_of +from tilefoundry.ir.types.utils import local_type_of, static_dim_value from tilefoundry.utils.isl_utils import cardinality, has_unbounded_param from tilefoundry.visitor_registry.access_relation import ( BoundaryRelation, diff --git a/src/tilefoundry/analysis/check.py b/src/tilefoundry/analysis/check.py index 4cc47b7d..a872f237 100644 --- a/src/tilefoundry/analysis/check.py +++ b/src/tilefoundry/analysis/check.py @@ -35,17 +35,14 @@ is_concrete, specialize_concretely, ) -from tilefoundry.ir.types import callable_type_for -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard import ( - Topology, -) +from tilefoundry.ir.types import Topology, callable_type_for from tilefoundry.ir.types.substitute import ( DimSubstitutionError, dim_vars_by_name, substitute_shape_dim, substitute_topology_dims, ) +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.ir.visitor import BindingSubstitutionCloner, collect_exprs from tilefoundry.target import UnsupportedCapabilityError from tilefoundry.target.facts import TopologyFacts diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index 801123d0..91b5a3d6 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.types import DType -from tilefoundry.ir.types.shard import Mesh, composed, topology_axes -from tilefoundry.ir.types.shard.mesh import _positions_layout +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.visitor import ExprVisitor from tilefoundry.visitor_registry.contexts import CostContext, FunctionScope, TrafficBytes from tilefoundry.visitor_registry.visitors import CostEvaluator @@ -187,18 +187,18 @@ def _scope_position_count(mesh: Mesh, topology_level: str | None, topologies: tu return 1 declared = {topology.name: index for index, topology in enumerate(topologies)} selected = declared[topology_level] - shape, _strides, _offset = _positions_layout(mesh) + stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout positions = 1 - for topology, axes in zip(mesh.topologies, topology_axes(mesh)): + for index, topology in enumerate(mesh.topologies): if declared[topology.name] > selected: continue - for axis in axes: - extent = shape[axis] - if not isinstance(extent, int) or isinstance(extent, bool) or extent < 1: - raise AnalysisError( - f"compute-cost: mesh axis {axis} needs a positive static extent, got {extent!r}" - ) - positions *= extent + count = size(get(stated, index)) + if not isinstance(count, int) or isinstance(count, bool) or count < 1: + raise AnalysisError( + f"compute-cost: mesh level {topology.name!r} needs positive static " + f"extents, got {count!r}" + ) + positions *= count return positions @@ -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 = composed((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh + mesh = merge_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/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index 6d948c51..473cf647 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -16,16 +16,10 @@ from tilefoundry.ir.core.module import Module from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord -from tilefoundry.ir.types import DType, TensorType, TupleType, Type -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard import ( - ComposedLayout, - Layout, - Mesh, - flatten, - topology_axes, - try_c_order_strides, -) +from tilefoundry.ir.types import DType, Mesh, TensorType, TupleType, Type +from tilefoundry.ir.types.int_tuple import repeat_like +from tilefoundry.ir.types.layout import ComposedLayout, flatten +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.target.base import Target, UnsupportedCapabilityError from tilefoundry.target.facts import TopologyFacts from tilefoundry.utils.isl_utils import cardinality @@ -87,45 +81,18 @@ def window(self) -> int | None: return -1 if self.space else None -def _layout_image(mesh: Mesh) -> tuple[int, tuple[int, ...], tuple[int, ...]] | None: - """Return a statically stated ``(offset, shape, strides)`` image.""" - layout = mesh.layout - if isinstance(layout, Layout): - stated = layout - offset = 0 - elif ( - isinstance(layout, ComposedLayout) - and layout.inner is None - and isinstance(layout.outer, Layout) - ): - stated = layout.outer - offset = layout.offset - else: +def _static_axes(mesh: Mesh) -> "tuple[int, tuple[int, ...], tuple[int, ...]] | None": + """A mesh's device-numbered positions, when every one of them is a number.""" + try: + axes = flatten(mesh.layout) + offset = (mesh.layout.offset if isinstance(mesh.layout, ComposedLayout) else 0) + except ValueError: return None - - shape = flatten(stated.shape) - if stated.strides is None: - strides = try_c_order_strides(shape) - if strides is None: - return None - else: - strides = tuple(static_dim_value(value) for value in flatten(stated.strides)) - if any(value is None for value in strides): - return None - - static_shape = tuple(static_dim_value(value) for value in shape) - static_offset = static_dim_value(offset) - if ( - static_offset is None - or any(value is None for value in static_shape) - or len(static_shape) != len(strides) - ): + shape = tuple(static_dim_value(extent) for extent in flatten(axes.shape)) + strides = tuple(static_dim_value(stride) for stride in flatten(axes.strides)) + if any(value is None for value in (*shape, *strides, static_dim_value(offset))): return None - return ( - static_offset, - tuple(value for value in static_shape if value is not None), - tuple(value for value in strides if value is not None), - ) + return offset, shape, strides def _mesh_parameters(scope: IterationScope) -> tuple[tuple[str, Call], ...]: @@ -144,7 +111,7 @@ def _linear_position( if not parameters: return 0, () mesh = parameters[0][1].target.mesh - image = _layout_image(mesh) + image = _static_axes(mesh) if image is None: return None offset, shape, strides = image @@ -340,7 +307,7 @@ def space_axes( wave_reached = boundary.space_wave_reached(window) shared = [] for axis, parameter_name in enumerate(boundary.axis_parameters): - shape = flatten(boundary.mesh.layout.shape) + shape = flatten(flatten(boundary.mesh.layout).shape) extent = static_dim_value(shape[axis]) if axis < len(shape) else None unit_reached = boundary.unit_reached(window, parameter_name) if extent is not None and extent > 1 and unit_reached is not None: @@ -351,11 +318,14 @@ def space_axes( def axis_label(mesh: Mesh, axis: int) -> str: """Name a mesh axis, falling back to its factual numeric position.""" - level_name = "" - for level, axes in zip(mesh.topologies, topology_axes(mesh), strict=True): - if axis in axes: - level_name = getattr(level, "name", str(level)) - break + stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout + levels = flatten( + tuple( + repeat_like(mode, getattr(level, "name", str(level))) + for mode, level in zip(stated.shape, mesh.topologies, strict=True) + ) + ) + level_name = levels[axis] if axis < len(levels) else "" if axis < len(mesh.names): name = mesh.names[axis] return f"{level_name}.{name}" if level_name else name @@ -364,7 +334,7 @@ def axis_label(mesh: Mesh, axis: int) -> str: def shared_units(mesh: Mesh, axes: tuple[int, ...], wave: tuple[int, int]) -> int: """Count wave positions after projecting out axes that do not share data.""" - image = _layout_image(mesh) + image = _static_axes(mesh) if image is None: return 1 offset, shape, strides = image @@ -491,7 +461,7 @@ def _axis_parameters( and (axis := static_dim_value(coordinate.args[0])) is not None } return tuple( - names_by_axis.get(axis) for axis in range(len(flatten(mesh.layout.shape))) + names_by_axis.get(axis) for axis in range(len(flatten(flatten(mesh.layout).shape))) ) diff --git a/src/tilefoundry/analysis/iteration_scope.py b/src/tilefoundry/analysis/iteration_scope.py index 69998ba9..52a08e51 100644 --- a/src/tilefoundry/analysis/iteration_scope.py +++ b/src/tilefoundry/analysis/iteration_scope.py @@ -12,15 +12,11 @@ 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.types.shard import Mesh +from tilefoundry.ir.isl_interop import dim_range +from tilefoundry.ir.types import Mesh +from tilefoundry.ir.types.dim import DimSub, simplify_dim +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.ir.visitor import expr_children -from tilefoundry.utils.isl_utils import ( - PARAM_POINT_LIMIT, - ParameterBoxTooLarge, - UnboundedParameterBox, - cardinality, - param_points, -) from tilefoundry.visitor_registry.access_relation import ( AccessRelations, access_relation_registry, @@ -29,7 +25,7 @@ ) from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext -from .access import Access, resolve_access +from .access import Access, AccessPrecision, resolve_access from .errors import AnalysisError from .loop_domain import induction_name, iteration_domain @@ -89,47 +85,64 @@ def enclosing_mesh(self) -> Mesh | None: cursor = cursor.parent return None + def _counted(self) -> tuple[int, AccessPrecision]: + """This scope's iteration count relative to its parent, and how exact it is. + + A loop runs its own span, not the grid it is entered from, so the count + is read off ``stop - start``: exactly where that span reaches one value + -- however it is written, so bounds naming an enclosing coordinate that + cancel are exact -- and at its widest where it reaches several. The + widest span is an upper bound on every coordinate's count, which is + what every consumer of a trip count already reads it as. + """ + if isinstance(self.owner, MeshRegion) or self.parent is None: + return 1, AccessPrecision.EXACT + owner = self.owner + span = simplify_dim(DimSub, (owner.extent, owner.start)) + step = static_dim_value(owner.step) + if step is None: + raise AnalysisError( + f"loop {induction_name(owner)!r} steps by {owner.step!r}, which is not " + "a number; how often it runs is not stated" + ) + widest = static_dim_value(span) + precision = AccessPrecision.EXACT + if widest is None: + try: + bounds = dim_range(span) + except (ArithmeticError, TypeError, ValueError): + bounds = None + if bounds is None: + raise AnalysisError( + f"loop {induction_name(owner)!r} runs from {owner.start!r} to " + f"{owner.extent!r}, which state no bounded span; its trip count " + "cannot be determined" + ) + widest = bounds[1] - 1 + if bounds[1] - bounds[0] > 1: + precision = AccessPrecision.WIDENED + if step <= 0 or widest <= 0: + return 1, precision + return max(1, -(-widest // step)), precision + def trips(self) -> int: """Return this scope's iteration count relative to its parent.""" - if isinstance(self.owner, MeshRegion): - return 1 cached = getattr(self, "_trips_cache", None) - if cached is not None: - return cached - if self.parent is None: - return 1 - if isinstance(self.owner, LoopRegion): - start, extent, step = self.owner.start, self.owner.extent, self.owner.step - if all(isinstance(value, int) for value in (start, extent, step)): - result = 1 if step <= 0 or extent <= start else -(-(extent - start) // step) - self._trips_cache = result - return result - domain = self.domain - parent = self.parent.domain.align_params(domain.get_space()) - domain = domain.align_params(parent.get_space()) - try: - points = param_points(domain.params().intersect(parent.params())) - except UnboundedParameterBox as error: - raise AnalysisError( - f"loop {induction_name(self.owner)!r} has unbounded parameter " - f"{error.parameter!r}, so its trip count cannot be determined" - ) from error - except ParameterBoxTooLarge as error: - raise AnalysisError( - f"loop {induction_name(self.owner)!r} has a parameter box exceeding " - f"the {PARAM_POINT_LIMIT}-point analysis limit, so its trip count cannot be " - "determined" - ) from error - ratios = [] - for point in points: - amount = cardinality(domain.intersect_params(point)) - parent_count = cardinality(parent.intersect_params(point)) - if amount is None or not parent_count: - continue - ratios.append(max(1, amount // parent_count)) - result = max(ratios, default=1) - self._trips_cache = result - return result + if cached is None: + cached = self._counted() + self._trips_cache = cached + return cached[0] + + @property + def trips_precision(self) -> AccessPrecision: + """How exactly :meth:`trips` counts this scope's iterations. + + ``EXACT`` where the span between the bounds is a number, ``WIDENED`` + where it is read at its widest because a bound names something no one + here fixes -- a coordinate of the scope the loop is entered from. + """ + self.trips() + return self._trips_cache[1] class ScopeBuilder: diff --git a/src/tilefoundry/analysis/loop_domain.py b/src/tilefoundry/analysis/loop_domain.py index e515f833..53982b24 100644 --- a/src/tilefoundry/analysis/loop_domain.py +++ b/src/tilefoundry/analysis/loop_domain.py @@ -8,7 +8,7 @@ from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.isl_interop import dim_to_isl_expr -from tilefoundry.ir.types.shape_helpers import static_dim_value +from tilefoundry.ir.types.utils import static_dim_value from .errors import AnalysisError diff --git a/src/tilefoundry/analysis/memory.py b/src/tilefoundry/analysis/memory.py index c619ec2e..3194eceb 100644 --- a/src/tilefoundry/analysis/memory.py +++ b/src/tilefoundry/analysis/memory.py @@ -18,8 +18,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.types import TensorType, TupleType, Type, bytes_by_storage +from tilefoundry.ir.types import TensorType, TupleType, Type from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.utils import bytes_by_storage from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.utils.units import format_bytes from tilefoundry.visitor_registry.access_relation import ( diff --git a/src/tilefoundry/analysis/performance.py b/src/tilefoundry/analysis/performance.py index cf0fcd6d..587dc3f2 100644 --- a/src/tilefoundry/analysis/performance.py +++ b/src/tilefoundry/analysis/performance.py @@ -11,7 +11,7 @@ from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.mesh_region import MeshRegion -from tilefoundry.ir.types.shape_helpers import static_dim_value +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.target import UnsupportedCapabilityError from tilefoundry.target.facts import TopologyFacts diff --git a/src/tilefoundry/analysis/report.py b/src/tilefoundry/analysis/report.py index 88cdc926..7c3ac06f 100644 --- a/src/tilefoundry/analysis/report.py +++ b/src/tilefoundry/analysis/report.py @@ -24,7 +24,7 @@ from tilefoundry.ir.core.module import Module from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion -from tilefoundry.ir.types import tensor_types +from tilefoundry.ir.types.utils import tensor_types from tilefoundry.ir.visitor import collect_exprs _FAMILIES: dict[type[IRMetadata], str] = {} diff --git a/src/tilefoundry/codegen/cpu/module.py b/src/tilefoundry/codegen/cpu/module.py index e7d284d4..7b9511cc 100644 --- a/src/tilefoundry/codegen/cpu/module.py +++ b/src/tilefoundry/codegen/cpu/module.py @@ -26,17 +26,9 @@ from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.tir.stmts import Evaluate, Sequential -from tilefoundry.ir.types.dim import ( - DimAdd, - DimFloorDiv, - DimMax, - DimMin, - DimMod, - DimMul, - DimSub, -) -from tilefoundry.ir.types.shape_helpers import static_dim_value +from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimMax, DimMin, DimMod, DimMul, DimSub from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.target import Target diff --git a/src/tilefoundry/codegen/cuda/abi.py b/src/tilefoundry/codegen/cuda/abi.py index cabf6d37..1b566fe0 100644 --- a/src/tilefoundry/codegen/cuda/abi.py +++ b/src/tilefoundry/codegen/cuda/abi.py @@ -21,7 +21,7 @@ tensor_signature_of, ) from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen diff --git a/src/tilefoundry/codegen/cuda/module.py b/src/tilefoundry/codegen/cuda/module.py index a01bb7c1..c736de08 100644 --- a/src/tilefoundry/codegen/cuda/module.py +++ b/src/tilefoundry/codegen/cuda/module.py @@ -25,7 +25,7 @@ from tilefoundry.codegen.signature import CallableSignature from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.types.shape_helpers import static_dim_value +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.target import Target Geometry = tuple[tuple[object, object, object], tuple[object, object, object]] diff --git a/src/tilefoundry/codegen/cuda/tir/arith.py b/src/tilefoundry/codegen/cuda/tir/arith.py index 4bdfbb57..fca17315 100644 --- a/src/tilefoundry/codegen/cuda/tir/arith.py +++ b/src/tilefoundry/codegen/cuda/tir/arith.py @@ -4,13 +4,13 @@ from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.ir.tir.arith import Binary, BinaryKind, Unary, UnaryKind -from tilefoundry.ir.types.shape_helpers import ( +from tilefoundry.ir.types.shard_layout import ShardLayout, shard_layout_local_shape +from tilefoundry.ir.types.utils import ( shape_has_dim_var, shape_numel_upper_bound, shape_runtime_total, shape_upper_bound, ) -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, shard_layout_local_shape from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen diff --git a/src/tilefoundry/codegen/cuda/tir/mbarrier.py b/src/tilefoundry/codegen/cuda/tir/mbarrier.py index 4172a28f..3400fc38 100644 --- a/src/tilefoundry/codegen/cuda/tir/mbarrier.py +++ b/src/tilefoundry/codegen/cuda/tir/mbarrier.py @@ -10,7 +10,7 @@ MBarrierInvalidate, MBarrierWaitParity, ) -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen diff --git a/src/tilefoundry/codegen/cuda/tir/memory/alloc_tensor.py b/src/tilefoundry/codegen/cuda/tir/memory/alloc_tensor.py index a21bedff..ef25ac5e 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/alloc_tensor.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/alloc_tensor.py @@ -16,13 +16,9 @@ from tilefoundry.codegen.cuda.tir.stmts.mesh_scope import program_topologies from tilefoundry.ir.tir.memory import AllocTensor from tilefoundry.ir.tir.stmts import LetStmt -from tilefoundry.ir.types.shape_helpers import ( - shape_numel_upper_bound, - shape_upper_bound, - upper_bound, -) -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, shard_layout_local_shape +from tilefoundry.ir.types.shard_layout import ShardLayout, shard_layout_local_shape from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.utils import shape_numel_upper_bound, shape_upper_bound, upper_bound from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen diff --git a/src/tilefoundry/codegen/cuda/tir/memory/copy.py b/src/tilefoundry/codegen/cuda/tir/memory/copy.py index 354f3230..11f02d21 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/copy.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/copy.py @@ -12,8 +12,8 @@ from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.ir.tir.memory.copy import Copy -from tilefoundry.ir.types.shape_helpers import shape_has_dim_var, shape_runtime_total -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.ir.types.shard_layout import ShardLayout +from tilefoundry.ir.types.utils import shape_has_dim_var, shape_runtime_total from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen diff --git a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py index 6f01a17a..8691ce96 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py @@ -23,19 +23,21 @@ from tilefoundry.ir.tir.memory.tensor_view import TensorView from tilefoundry.ir.tir.stmts import LetStmt from tilefoundry.ir.tir.sync import participation +from tilefoundry.ir.types import ComposedLayout from tilefoundry.ir.types.dim import DimAdd, DimMul, DimSub, DimVar -from tilefoundry.ir.types.shape_helpers import shape_numel_upper_bound, upper_bound -from tilefoundry.ir.types.shard import c_order_strides, swizzle_of -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.layout import Layout, LayoutBase, flatten +from tilefoundry.ir.types.layout_algebra import swizzle_of +from tilefoundry.ir.types.shard_layout import ( Broadcast, Dynamic, Partial, Split, shard_layout_local_shape, ) -from tilefoundry.ir.types.shard.shard_layout import ShardLayout as SL +from tilefoundry.ir.types.shard_layout import ShardLayout as SL from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.stride import compact_row_major +from tilefoundry.ir.types.utils import shape_numel_upper_bound, upper_bound from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen @@ -217,24 +219,17 @@ def render_shard_layout_value( "and rebuilding this layout from register strides would drop it" ) sll = Layout(shape=sll.shape, strides=register_strides(sl)) - mesh_layout = sl.mesh.layout - if isinstance(mesh_layout, ComposedLayout): - mesh_outer = mesh_layout.outer - if not isinstance(mesh_outer, Layout) or mesh_outer.strides is None: + mesh_base = 0 + if isinstance(sl.mesh.layout, ComposedLayout): + if sl.mesh.layout.outer is None: raise NotImplementedError( - "render_shard_layout_value: a sliced mesh needs its participating " - "box as a strided Layout; this states an identity box, with no sub-box" + "render_shard_layout_value: a sliced mesh needs its participating box " + "as a strided Layout; this states an identity box, with no sub-box" ) + mesh_base = sl.mesh.layout.offset participation(sl.mesh) - ml_shape, ml_strides, ml_base = ( - mesh_outer.shape, - mesh_outer.strides, - int(mesh_layout.offset), - ) - else: - if mesh_layout.strides is None: - raise NotImplementedError("render_shard_layout_value: mesh layout needs strides") - ml_shape, ml_strides, ml_base = mesh_layout.shape, mesh_layout.strides, 0 + mesh_axes = flatten(sl.mesh.layout) + ml_shape, ml_strides, ml_base = mesh_axes.shape, mesh_axes.strides, int(mesh_base) topo = program_topologies(sl.mesh)[0] def _static_dim(value, what): @@ -399,7 +394,7 @@ def _emit(let: LetStmt, ctx: CudaCodegenContext) -> None: ) source_stride_args = ", ".join( f"cute::Int<{stride}>{{}}" - for stride in c_order_strides( + for stride in compact_row_major( tuple(int(upper_bound(dim)) for dim in source_shape) ) ) diff --git a/src/tilefoundry/codegen/cuda/tir/prim_function.py b/src/tilefoundry/codegen/cuda/tir/prim_function.py index 8ee019cc..54d967a0 100644 --- a/src/tilefoundry/codegen/cuda/tir/prim_function.py +++ b/src/tilefoundry/codegen/cuda/tir/prim_function.py @@ -15,8 +15,8 @@ from tilefoundry.codegen.signature import TensorSignature, tensor_signature_of from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.types.shape_helpers import shape_numel_upper_bound -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.ir.types.shard_layout import ShardLayout +from tilefoundry.ir.types.utils import shape_numel_upper_bound from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py index d7f17863..3eb47ba1 100644 --- a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py +++ b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py @@ -7,8 +7,8 @@ from tilefoundry.codegen.cuda.context import topology_scope_str from tilefoundry.ir.tir.sync import participation -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout -from tilefoundry.ir.types.shard.mesh import Mesh, Topology, positions_at +from tilefoundry.ir.types.layout import ComposedLayout, Layout, flatten, get +from tilefoundry.ir.types.mesh import Mesh, Topology from tilefoundry.target import validate_cuda_topology_levels @@ -40,17 +40,19 @@ def _validate_topology(mesh: Mesh, target) -> None: def _levelwise_layout(mesh: Mesh, topos) -> str: """One nest per topology level, each in that level's own numbering. - A flat shape says nothing about which axes are whose, so the runtime's - ``get`` needs the boundary stated. ``level_axes`` hands the axes to - the levels left to right and ``positions_at`` divides each level's strides - by what the levels under it contribute, so a nest reads as the layout that - level would have alone. The composite is a grouping, not a map: only the + A mesh states one arrangement per level already, each in that level's own + numbering, so a nest is that arrangement written out: nothing here decides + where a boundary falls. The composite is a grouping, not a map: only the nests are evaluated. A level every instance shares keeps a mode of one, so that ``get`` still has something to pick. """ shapes, strides = [], [] for topology in topos: - level_shape, level_strides = positions_at(mesh, topology.name) + stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout + index = [t.name for t in mesh.topologies].index(topology.name) + arrangement = get(stated, index) + level_shape = tuple(flatten(arrangement.shape)) + level_strides = tuple(flatten(arrangement.strides)) if not level_shape: level_shape, level_strides = (1,), (0,) shapes.append( @@ -80,14 +82,15 @@ def mesh_type(mesh: Mesh) -> str: "CUDA mesh emission: a sliced mesh needs its participating box " "as a strided Layout; this states an identity box, with no sub-box" ) - shape, strides, base = outer.shape, outer.strides, participation(mesh).base + stated, base = outer, participation(mesh).base else: if layout_value.strides is None: raise NotImplementedError("CUDA mesh emission: mesh layout needs strides") - shape, strides, base = layout_value.shape, layout_value.strides, 0 + stated, base = layout_value, 0 if len(topos) == 1: - shape_types = ", ".join(f"cute::Int<{s}>" for s in shape) - stride_types = ", ".join(f"cute::Int<{s}>" for s in strides) + level = get(stated, 0) + shape_types = ", ".join(f"cute::Int<{s}>" for s in flatten(level.shape)) + stride_types = ", ".join(f"cute::Int<{s}>" for s in flatten(level.strides)) layout = ( f"cute::Layout, " f"cute::Stride<{stride_types}>>" @@ -113,7 +116,7 @@ def _is_dynamic_mesh(mesh: Mesh) -> bool: """ if any(t.size is None for t in program_topologies(mesh)): return True - return any(s is None for s in mesh.layout.shape) + return any(s is None for s in flatten(mesh.layout).shape) __all__ = ["mesh_type", "program_topologies"] diff --git a/src/tilefoundry/codegen/emitter.py b/src/tilefoundry/codegen/emitter.py index dcd5ce89..0c86048b 100644 --- a/src/tilefoundry/codegen/emitter.py +++ b/src/tilefoundry/codegen/emitter.py @@ -21,8 +21,7 @@ Sequential, While, ) -from tilefoundry.ir.types import UnitType -from tilefoundry.ir.types.shard.layout import ComposedLayout +from tilefoundry.ir.types import ComposedLayout, UnitType from tilefoundry.ir.visitor import StmtVisitor diff --git a/src/tilefoundry/dsl/__init__.py b/src/tilefoundry/dsl/__init__.py index b11f9da9..468a8fef 100644 --- a/src/tilefoundry/dsl/__init__.py +++ b/src/tilefoundry/dsl/__init__.py @@ -18,16 +18,7 @@ from tilefoundry.script import func from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern from tilefoundry.ir.types.dim import DimVar, ceildiv -from tilefoundry.ir.types.shard import ( - Mesh, - Topology, - Split, - Partial, - Broadcast, - S, - P, - B, -) +from tilefoundry.ir.types import B, Broadcast, Mesh, P, Partial, S, Split, Topology from tilefoundry.ir.core.kinds import ReduceKind, UnaryKind, BinaryKind __all__ = [ diff --git a/src/tilefoundry/evaluator/value.py b/src/tilefoundry/evaluator/value.py index e309dbca..fb9b6a0b 100644 --- a/src/tilefoundry/evaluator/value.py +++ b/src/tilefoundry/evaluator/value.py @@ -6,7 +6,7 @@ import torch from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.ir.types.shard_layout import ShardLayout class EvalError(Exception): diff --git a/src/tilefoundry/inspection/print_context.py b/src/tilefoundry/inspection/print_context.py index 0f031755..2f3ae346 100644 --- a/src/tilefoundry/inspection/print_context.py +++ b/src/tilefoundry/inspection/print_context.py @@ -5,9 +5,9 @@ from contextlib import contextmanager from math import prod -from tilefoundry.ir.types.shard.int_tuple import flatten -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout -from tilefoundry.ir.types.shard.mesh import Mesh, topology_axes +from tilefoundry.ir.types.int_tuple import repeat_like +from tilefoundry.ir.types.layout import ComposedLayout, Layout, flatten +from tilefoundry.ir.types.mesh import Mesh from tilefoundry.utils.python_source import PythonExpr, _merge_imports @@ -89,11 +89,14 @@ def mesh_alias(self, mesh: Mesh) -> str | None: @staticmethod def _axis_levels(mesh: Mesh) -> tuple[str, ...]: - levels = [""] * len(flatten(mesh.layout.shape)) - for topology, axes in zip(mesh.topologies, topology_axes(mesh), strict=True): - for axis in axes: - levels[axis] = topology.name - return tuple(levels) + """The level each of the mesh's axes stands in, one name per axis.""" + stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout + return flatten( + tuple( + repeat_like(mode, topology.name) + for mode, topology in zip(stated.shape, mesh.topologies, strict=True) + ) + ) def mesh_axis_alias(self, mesh: Mesh, axis: int) -> str | None: """Name one mesh axis through an active scope binding, if one dominates it.""" @@ -111,8 +114,7 @@ def mesh_axis_alias(self, mesh: Mesh, axis: int) -> str | None: if ( self._type_annotation_surface and bound is mesh - and isinstance(mesh.layout, ComposedLayout) - and mesh.layout.offset != 0 + and (mesh.layout.offset if isinstance(mesh.layout, ComposedLayout) else 0) != 0 ): continue if not bound.names or target_name not in bound.names: @@ -130,11 +132,7 @@ def mesh_axis_alias(self, mesh: Mesh, axis: int) -> str | None: def mesh_slice(self, mesh: Mesh) -> str | None: """Recover ``binding[start:stop]`` for a sliced active mesh.""" - if not ( - isinstance(mesh.layout, ComposedLayout) - and mesh.layout.inner is None - and isinstance(mesh.layout.outer, Layout) - ): + if not isinstance(mesh.layout, ComposedLayout): return None for parent, alias in reversed(self._mesh_bindings): text = self._slice_from_parent(parent, mesh, alias) diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 63e39816..3bf4e006 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.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.tir.cuda.nn.mma_atom import MmaAtom from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType from tilefoundry.ir.types.dim import ( @@ -21,11 +22,11 @@ DimSub, DimVar, ) -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase, Swizzle -from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, ShardLayout, Split +from tilefoundry.ir.types.layout import ComposedLayout, Layout, LayoutBase, Swizzle, flatten +from tilefoundry.ir.types.mesh import Mesh +from tilefoundry.ir.types.shard_layout import Broadcast, Partial, ShardLayout, Split from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.ir.visitor import ExprFunctor, TypeFunctor from tilefoundry.target import Target from tilefoundry.utils.python_source import PythonExpr @@ -89,6 +90,15 @@ def nested_dim(self, nested: bool): self._nested_dim = previous def dim_entry(self, value, ctx=None, *, nested: bool = False) -> str: + """One entry of a shape or stride tuple, which may itself be a group. + + A layout groups the modes of one tensor axis by writing them as a + tuple in that axis's place, so an entry is read as the shape tuple it + is rather than as a single dimension. + """ + if isinstance(value, tuple): + entries = ", ".join(self.dim_entry(item, ctx) for item in value) + return f"({entries}{',' if len(value) == 1 else ''})" with self.nested_dim(nested): return self.visit(value, ctx) @@ -143,7 +153,7 @@ def visit_Call(self, value: Call, ctx=None) -> str: 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 - if axis is None or axis < 0 or axis >= len(target.mesh.layout.shape): + if axis is None or axis < 0 or axis >= len(flatten(target.mesh.layout).shape): raise ValueError("MeshCoord requires a literal in-range axis to print") if ctx is None: raise ValueError("MeshCoord requires an active mesh binding to print") @@ -189,7 +199,12 @@ def _ceildiv_args(value: Call) -> tuple[object, object] | None: return added.args[0], divisor def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: - """Render placement sugar only when every mesh axis has a scope binding.""" + """Render placement sugar only when every mesh axis has a scope binding. + + The sugar states one extent per tensor axis, each a dimension + expression, so it declines a layout whose modes are grouped by tile + axis: writing the groups in would emit a line the parser refuses. + """ layout = value.layout names = value.mesh.names if ( @@ -199,6 +214,11 @@ 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 ())) + ): + return None refs = tuple(ctx.mesh_axis_alias(value.mesh, index) for index in range(len(names))) if any(ref is None for ref in refs): return None @@ -295,13 +315,18 @@ def visit_Mesh(self, value: Mesh, ctx=None) -> str: sliced = ctx.mesh_slice(value) if sliced is not None: return sliced - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "Mesh")) + ctx.use(PythonExpr(("from tilefoundry.ir.types import Mesh, Topology",), "Mesh")) topologies = ", ".join( f'Topology("{topology.name}", {self.dim_entry(topology.size, ctx)})' for topology in value.topologies ) topologies = f"({topologies}{',' if len(value.topologies) == 1 else ''})" - result = f"Mesh({topologies}, {self.visit(value.layout, ctx)}" + written = device_layout(value) + if isinstance(value.layout, ComposedLayout): + written = ComposedLayout( + inner=value.layout.inner, offset=value.layout.offset, outer=written + ) + result = f"Mesh({topologies}, {self.visit(written, ctx)}" if value.names: result += f", names={tuple(value.names)!r}" return result + ")" @@ -311,18 +336,18 @@ def visit_NoneType(self, value: None, ctx=None) -> str: def visit_Layout(self, value: Layout, ctx=None) -> str: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout",), "Layout")) + ctx.use(PythonExpr(("from tilefoundry.ir.types import Layout",), "Layout")) strides = self.shape_tuple(value.strides, ctx) if value.strides is not None else "None" return f"Layout({self.shape_tuple(value.shape, ctx)}, {strides})" def visit_Swizzle(self, value: Swizzle, ctx=None) -> str: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Swizzle",), "")) + ctx.use(PythonExpr(("from tilefoundry.ir.types import Swizzle",), "")) return f"Swizzle({value.bits}, {value.base}, {value.shift})" def visit_ComposedLayout(self, value: ComposedLayout, ctx=None) -> str: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) + ctx.use(PythonExpr(("from tilefoundry.ir.types import ComposedLayout",), "")) outer, child = self._indent, self._indent + " " with self.type_surface(indent=child): inner_text = self.visit(value.inner, ctx) @@ -340,7 +365,7 @@ def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: if surface is not None: return surface if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) + ctx.use(PythonExpr(("from tilefoundry.ir.types import ShardLayout",), "")) outer, child = self._indent, self._indent + " " attrs = ", ".join(self.visit(attr, ctx) for attr in value.attrs) if len(value.attrs) == 1: @@ -358,17 +383,17 @@ def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: def visit_Broadcast(self, value: Broadcast, ctx=None) -> str: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) + ctx.use(PythonExpr(("from tilefoundry.ir.types import B",), "B")) return "B()" def visit_Split(self, value: Split, ctx=None) -> str: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) + ctx.use(PythonExpr(("from tilefoundry.ir.types import S",), "S")) return f"S({value.axis})" def visit_Partial(self, value: Partial, ctx=None) -> str: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) + ctx.use(PythonExpr(("from tilefoundry.ir.types import P",), "P")) return f'P("{value.reduction}")' def atom_reference(self, value: MmaAtom, ctx=None) -> str: diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index eba48dcf..848be7da 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -50,7 +50,7 @@ from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.types import DType, TensorType, TupleType from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.shard_layout import ( Broadcast, Partial, ShardLayout, @@ -1151,7 +1151,7 @@ def _module_decorator_line(mod: Module, entry_name: str | None, ctx: HirPrintCon ctx.imports.update(rendered.imports) kwargs.append(f"target={rendered.text}") if mod.topologies is not None: - ctx.imports.add("from tilefoundry.ir.types.shard import Topology") + 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 diff --git a/src/tilefoundry/inspection/tir_printer.py b/src/tilefoundry/inspection/tir_printer.py index 5ca8a4df..ed0f5162 100644 --- a/src/tilefoundry/inspection/tir_printer.py +++ b/src/tilefoundry/inspection/tir_printer.py @@ -10,6 +10,7 @@ from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.core.module import Module 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.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf @@ -18,6 +19,7 @@ ) from tilefoundry.ir.tir.symbol_ref import SymbolRef from tilefoundry.ir.types import DType, 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 @@ -69,6 +71,8 @@ def visit_Op(self, expr: Op, ctx=None) -> str: def visit_program_call(self, expr: Call, ctx=None) -> str: target = expr.target + 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", @@ -88,6 +92,25 @@ def visit_program_call(self, expr: Call, ctx=None) -> str: self.context.use(PythonExpr(("from tilefoundry.dsl import T",), "T")) return f"T.{name}({', '.join(args)})" + def _window_subscript(self, expr: Call, ctx=None) -> str: + """A window as the subscript it was authored as, which reads back. + + The generic call form names an Op the parser does not bind, so a window + is written as the slice it is: each axis runs from its start to that + start plus what the window reaches, with a step where it has one. + """ + 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) + ) + 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)}]" + def visit_Sequential(self, stmt, ctx=None): return [line for child in stmt.body for line in self.visit(child)] @@ -105,7 +128,17 @@ def visit_MeshScope(self, stmt, ctx=None): return lines def visit_For(self, stmt, ctx=None): - lines = [f"{self.indent}for {stmt.induction_var.name} in range({self.visit(stmt.start)}, {self.visit(stmt.stop)}, {self.visit(stmt.step)}):"] + """A loop, its bounds read through the scopes the loop stands in. + + A bound may name a coordinate of an enclosing mesh -- a CTA grid deals + each CTA its own tiles -- and which scope binds that mesh is what the + 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) + ) + lines = [f"{self.indent}for {stmt.induction_var.name} in range({bounds}):"] lines.extend(TirPrinter(context=self.context, indent=self.indent + " ").visit(stmt.body)) return lines @@ -226,7 +259,7 @@ def tir_module_to_python(mod: Module, module_name: str | None = None, *, options imports.update(target.imports) kwargs.append(f"target={target.text}") if mod.topologies is not None: - imports.add("from tilefoundry.ir.types.shard import Topology") + 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)})") diff --git a/src/tilefoundry/inspection/viewer/builder.py b/src/tilefoundry/inspection/viewer/builder.py index ea354944..bfa7ce49 100644 --- a/src/tilefoundry/inspection/viewer/builder.py +++ b/src/tilefoundry/inspection/viewer/builder.py @@ -13,7 +13,7 @@ from tilefoundry.ir.hir.function import Function as HirFunction from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.types import DType -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.ir.types.tensor_type import TensorType, TupleType from .htmltable import Cell, Span, Table diff --git a/src/tilefoundry/ir/constraints/layout.py b/src/tilefoundry/ir/constraints/layout.py index 76cadf82..811bbf8b 100644 --- a/src/tilefoundry/ir/constraints/layout.py +++ b/src/tilefoundry/ir/constraints/layout.py @@ -4,7 +4,7 @@ from dataclasses import dataclass -from tilefoundry.ir.types.shard import Layout, ShardAttr +from tilefoundry.ir.types import Layout, ShardAttr from .base import ScheduleConstraint diff --git a/src/tilefoundry/ir/constraints/mesh.py b/src/tilefoundry/ir/constraints/mesh.py index 56267146..3ee090b7 100644 --- a/src/tilefoundry/ir/constraints/mesh.py +++ b/src/tilefoundry/ir/constraints/mesh.py @@ -4,7 +4,7 @@ from dataclasses import dataclass -from tilefoundry.ir.types.shard import Mesh +from tilefoundry.ir.types import Mesh from .base import ScheduleConstraint diff --git a/src/tilefoundry/ir/core/module.py b/src/tilefoundry/ir/core/module.py index 1475eaf0..b8b58fe6 100644 --- a/src/tilefoundry/ir/core/module.py +++ b/src/tilefoundry/ir/core/module.py @@ -17,7 +17,7 @@ from tilefoundry.evaluator.value import tensor_type_of from tilefoundry.ir.hir.function import Function as HirFunction from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.types.shard.mesh import Topology +from tilefoundry.ir.types.mesh import Topology from tilefoundry.ir.types.substitute import canonicalize_topology_dims from tilefoundry.ir.types.tensor_type import TensorType from tilefoundry.ir.types.utils import types_compatible diff --git a/src/tilefoundry/ir/hir/_shard_checks.py b/src/tilefoundry/ir/hir/_shard_checks.py index 3ad7220e..4edcc24c 100644 --- a/src/tilefoundry/ir/hir/_shard_checks.py +++ b/src/tilefoundry/ir/hir/_shard_checks.py @@ -8,7 +8,7 @@ from __future__ import annotations -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.shard_layout import ( Broadcast, Dynamic, Partial, diff --git a/src/tilefoundry/ir/hir/math/binary.py b/src/tilefoundry/ir/hir/math/binary.py index 97ecd3fc..2e840773 100644 --- a/src/tilefoundry/ir/hir/math/binary.py +++ b/src/tilefoundry/ir/hir/math/binary.py @@ -20,10 +20,15 @@ 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.types import DType, TensorType -from tilefoundry.ir.types.shard import Layout, canonical_shard_layout, try_c_order_strides -from tilefoundry.ir.types.shard.shard_layout import Broadcast, ShardLayout, shard_layout_of +from tilefoundry.ir.types import DType, Layout, TensorType +from tilefoundry.ir.types.shard_layout import ( + Broadcast, + ShardLayout, + canonical_shard_layout, + shard_layout_of, +) from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.stride import try_compact_major from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -178,7 +183,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: ctx.error(call, f"Binary {op.kind.name}: {e}") storage = resolve_anchor_storage(ctx, call, lhs_ty.storage, rhs_ty.storage) if layout is None and storage in (StorageKind.RMEM, StorageKind.SMEM) and out_shape: - layout = Layout(shape=out_shape, strides=try_c_order_strides(out_shape)) + layout = Layout(shape=out_shape, strides=try_compact_major(out_shape)) return TensorType( shape=out_shape, dtype=out_dtype, diff --git a/src/tilefoundry/ir/hir/mesh_region.py b/src/tilefoundry/ir/hir/mesh_region.py index 1aa1f652..53f89540 100644 --- a/src/tilefoundry/ir/hir/mesh_region.py +++ b/src/tilefoundry/ir/hir/mesh_region.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from tilefoundry.ir.core import Expr, Var -from tilefoundry.ir.types.shard.mesh import Mesh +from tilefoundry.ir.types.mesh import Mesh @dataclass(unsafe_hash=True) diff --git a/src/tilefoundry/ir/hir/nn/conv2d.py b/src/tilefoundry/ir/hir/nn/conv2d.py index b9526e73..3b0939c8 100644 --- a/src/tilefoundry/ir/hir/nn/conv2d.py +++ b/src/tilefoundry/ir/hir/nn/conv2d.py @@ -12,11 +12,11 @@ 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.types import TensorType +from tilefoundry.ir.types import Layout, TensorType from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimSub, simplify_dim -from tilefoundry.ir.types.shape_helpers import i64_const, static_dim_value -from tilefoundry.ir.types.shard import Layout, try_c_order_strides -from tilefoundry.ir.types.shard.shard_layout import Split, shard_layout_of, split_target_axes +from tilefoundry.ir.types.shard_layout import Split, shard_layout_of, split_target_axes +from tilefoundry.ir.types.stride import try_compact_major +from tilefoundry.ir.types.utils import i64_const, static_dim_value from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -292,7 +292,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: call, f"cannot derive input ownership: {error}; use an explicit Reshard before Conv2D", ) - layout = shard or Layout(shape=out_shape, strides=try_c_order_strides(out_shape)) + layout = shard or Layout(shape=out_shape, strides=try_compact_major(out_shape)) return TensorType( shape=out_shape, dtype=x.dtype, diff --git a/src/tilefoundry/ir/hir/nn/layer_norm.py b/src/tilefoundry/ir/hir/nn/layer_norm.py index 96d24dd0..202d0d7a 100644 --- a/src/tilefoundry/ir/hir/nn/layer_norm.py +++ b/src/tilefoundry/ir/hir/nn/layer_norm.py @@ -11,7 +11,7 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import reject_partials from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, split_target_axes +from tilefoundry.ir.types.shard_layout import ShardLayout, split_target_axes from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, diff --git a/src/tilefoundry/ir/hir/nn/matmul.py b/src/tilefoundry/ir/hir/nn/matmul.py index d74f2e7a..289e26ec 100644 --- a/src/tilefoundry/ir/hir/nn/matmul.py +++ b/src/tilefoundry/ir/hir/nn/matmul.py @@ -14,7 +14,7 @@ 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.types import TensorType -from tilefoundry.ir.types.shard.shard_layout import shard_layout_of, split_target_axes +from tilefoundry.ir.types.shard_layout import shard_layout_of, split_target_axes from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, diff --git a/src/tilefoundry/ir/hir/sharding/local.py b/src/tilefoundry/ir/hir/sharding/local.py index 781a40b4..67dd8449 100644 --- a/src/tilefoundry/ir/hir/sharding/local.py +++ b/src/tilefoundry/ir/hir/sharding/local.py @@ -7,8 +7,9 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, Split +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.shard_layout import ShardLayout, Split +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( measures_without_reading, @@ -33,7 +34,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: new_shape = list(x_ty.shape) for mesh_axis, attr in enumerate(sl.attrs): if isinstance(attr, Split): - mesh_extent = sl.mesh.layout.shape[mesh_axis] + mesh_extent = flatten(sl.mesh.layout).shape[mesh_axis] dim = new_shape[attr.axis] v = static_dim_value(dim) if v is not None: diff --git a/src/tilefoundry/ir/hir/sharding/mesh_coord.py b/src/tilefoundry/ir/hir/sharding/mesh_coord.py index d3340971..ea571240 100644 --- a/src/tilefoundry/ir/hir/sharding/mesh_coord.py +++ b/src/tilefoundry/ir/hir/sharding/mesh_coord.py @@ -8,11 +8,12 @@ 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.types import DType, TensorType -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.types.shard.scope_match import covered_by_scope +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.mesh import Mesh from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( measures_without_reading, @@ -47,7 +48,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TypeInferResults: ctx.error(call, "MeshCoord.mesh must be bound by the current mesh scope") if not call.args: ctx.error(call, "missing required input 'axis'") - shape = call.target.mesh.layout.shape + shape = flatten(call.target.mesh.layout).shape axis = static_dim_value(call.args[0]) if axis is not None and not 0 <= axis < len(shape): ctx.error(call, f"axis {axis} is out of range for rank-{len(shape)} mesh") diff --git a/src/tilefoundry/ir/hir/sharding/reshard.py b/src/tilefoundry/ir/hir/sharding/reshard.py index b6a76bc6..73dd473b 100644 --- a/src/tilefoundry/ir/hir/sharding/reshard.py +++ b/src/tilefoundry/ir/hir/sharding/reshard.py @@ -8,14 +8,14 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import DimMul, simplify_dim -from tilefoundry.ir.types.shard import c_order_strides -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.layout import Layout +from tilefoundry.ir.types.shard_layout import ( ShardLayout, Split, shard_layout_local_shape, ) from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.stride import compact_row_major from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -46,7 +46,7 @@ def _c_order_strides(shape: tuple) -> tuple: symbolic stride for the axes above it via ``_dim_mul``; static inner strides stay plain ints. """ - return c_order_strides(shape, mul=_dim_mul) + return compact_row_major(shape, mul=_dim_mul) def _shared_engine_strides(sl: ShardLayout) -> tuple: diff --git a/src/tilefoundry/ir/hir/specialize.py b/src/tilefoundry/ir/hir/specialize.py index f7be3a84..2ef70f1c 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.shard.mesh import composed from tilefoundry.ir.types.substitute import ( dim_vars_by_name, has_symbolic_dims, @@ -232,7 +232,7 @@ def visit_MeshRegion(self, expr: MeshRegion, ctx: InstantiateContext) -> Expr: if old is not new: ctx.subst[id(old)] = new current_mesh = ( - composed((ctx.type_ctx.current_mesh, mesh)) + merge_mesh((ctx.type_ctx.current_mesh, mesh)) if ctx.type_ctx.current_mesh else mesh ) @@ -392,8 +392,8 @@ def _substitute_op_dims(target: object, dims: Mapping[str, int]) -> object: """Substitute bindings in an operation's shape-valued attributes.""" if isinstance(target, Function) or not isinstance(target, Op): return target - from tilefoundry.ir.types.shard.layout import LayoutBase # noqa: PLC0415 - from tilefoundry.ir.types.shard.mesh import Mesh # noqa: PLC0415 + from tilefoundry.ir.types.layout import LayoutBase # noqa: PLC0415 + from tilefoundry.ir.types.mesh import Mesh # noqa: PLC0415 from tilefoundry.ir.types.substitute import ( # noqa: PLC0415 substitute_layout_dims, substitute_mesh_dims, diff --git a/src/tilefoundry/ir/hir/tensor/argmax.py b/src/tilefoundry/ir/hir/tensor/argmax.py index 69c4b6d3..45ff7bd2 100644 --- a/src/tilefoundry/ir/hir/tensor/argmax.py +++ b/src/tilefoundry/ir/hir/tensor/argmax.py @@ -17,13 +17,14 @@ 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.types import DType, TensorType -from tilefoundry.ir.types.shard import ( - Layout, +from tilefoundry.ir.types import DType, Layout, TensorType +from tilefoundry.ir.types.shard_layout import ( + Split, canonical_shard_layout, - try_c_order_strides, + shard_layout_of, + split_target_axes, ) -from tilefoundry.ir.types.shard.shard_layout import Split, shard_layout_of, split_target_axes +from tilefoundry.ir.types.stride import try_compact_major from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -72,7 +73,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: new_layout = ( None if x_ty.layout is None - else Layout(shape=out_shape, strides=try_c_order_strides(out_shape)) + else Layout(shape=out_shape, strides=try_compact_major(out_shape)) ) if source_shard is not None: relation = relations_of(call, ctx) diff --git a/src/tilefoundry/ir/hir/tensor/cache_update.py b/src/tilefoundry/ir/hir/tensor/cache_update.py index 198bf762..71c10e57 100644 --- a/src/tilefoundry/ir/hir/tensor/cache_update.py +++ b/src/tilefoundry/ir/hir/tensor/cache_update.py @@ -13,8 +13,8 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import require_matching_partial_state from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import shard_layout_of -from tilefoundry.ir.types.shard.shard_layout import Split, split_target_axes +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.shard_layout import Split, shard_layout_of, split_target_axes from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -182,7 +182,7 @@ def _divided_axes(type_) -> set[int]: if layout is None: return set() targets = split_target_axes(layout, type_.shape) - mesh = layout.mesh.layout.shape if layout.mesh is not None else () + mesh = flatten(layout.mesh.layout).shape if layout.mesh is not None else () return { targets[mesh_axis] for mesh_axis, attr in enumerate(layout.attrs) diff --git a/src/tilefoundry/ir/hir/tensor/cast.py b/src/tilefoundry/ir/hir/tensor/cast.py index 753fe72f..a8063dab 100644 --- a/src/tilefoundry/ir/hir/tensor/cast.py +++ b/src/tilefoundry/ir/hir/tensor/cast.py @@ -7,7 +7,7 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard.shard_layout import shard_layout_of +from tilefoundry.ir.types.shard_layout import shard_layout_of from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( identity_relations, diff --git a/src/tilefoundry/ir/hir/tensor/concat.py b/src/tilefoundry/ir/hir/tensor/concat.py index ab4695e6..1eebfcaa 100644 --- a/src/tilefoundry/ir/hir/tensor/concat.py +++ b/src/tilefoundry/ir/hir/tensor/concat.py @@ -18,15 +18,10 @@ require_uniform_partial_slices, ) from tilefoundry.ir.isl_interop import normalize_dim_entries -from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types import Layout, Split, TensorType from tilefoundry.ir.types.dim import DimAdd, simplify_dim -from tilefoundry.ir.types.shard import ( - Layout, - Split, - shard_layout_of, - try_c_order_strides, -) -from tilefoundry.ir.types.shard.shard_layout import split_target_axes +from tilefoundry.ir.types.shard_layout import shard_layout_of, split_target_axes +from tilefoundry.ir.types.stride import try_compact_major from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -178,7 +173,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: "layout; use an explicit Reshard before Concat", ) else: - layout = Layout(shape=new_shape, strides=try_c_order_strides(new_shape)) + layout = Layout(shape=new_shape, strides=try_compact_major(new_shape)) storage = resolve_anchor_storage(ctx, call, *(t.storage for t in types)) return TensorType(shape=new_shape, dtype=base.dtype, layout=layout, storage=storage) diff --git a/src/tilefoundry/ir/hir/tensor/index_add.py b/src/tilefoundry/ir/hir/tensor/index_add.py index fc8d212f..ad7e6f0b 100644 --- a/src/tilefoundry/ir/hir/tensor/index_add.py +++ b/src/tilefoundry/ir/hir/tensor/index_add.py @@ -11,8 +11,7 @@ from tilefoundry.ir.hir._shard_checks import reject_partials from tilefoundry.ir.hir.tensor.index_select import _norm_dim from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import shard_layout_of -from tilefoundry.ir.types.shard.shard_layout import Split +from tilefoundry.ir.types.shard_layout import Split, shard_layout_of from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, diff --git a/src/tilefoundry/ir/hir/tensor/index_select.py b/src/tilefoundry/ir/hir/tensor/index_select.py index 0132273b..ce67e657 100644 --- a/src/tilefoundry/ir/hir/tensor/index_select.py +++ b/src/tilefoundry/ir/hir/tensor/index_select.py @@ -9,14 +9,14 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.layout_algebra import prefix_product -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.layout import Layout +from tilefoundry.ir.types.shard_layout import ( Partial, ShardLayout, Split, split_target_axes, ) +from tilefoundry.ir.types.stride import compact_col_major from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -71,7 +71,7 @@ def _index_select_shard_layout(call, ctx, x_ty, dim: int, out_shape: tuple): f"dim {dim} index_select over a shard layout with multiple Split " "axes including the selected dim; cannot derive an output layout", ) - natural = Layout(shape=out_shape, strides=prefix_product(out_shape)) + natural = Layout(shape=out_shape, strides=compact_col_major(out_shape)) if on_dim: mesh_idx = on_dim[0] new_attrs = tuple( diff --git a/src/tilefoundry/ir/hir/tensor/insert_slice.py b/src/tilefoundry/ir/hir/tensor/insert_slice.py index f223a0cf..820723ec 100644 --- a/src/tilefoundry/ir/hir/tensor/insert_slice.py +++ b/src/tilefoundry/ir/hir/tensor/insert_slice.py @@ -10,7 +10,7 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import require_matching_partial_state from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shape_helpers import static_dim_value +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, diff --git a/src/tilefoundry/ir/hir/tensor/quant.py b/src/tilefoundry/ir/hir/tensor/quant.py index 8df733a9..abce8917 100644 --- a/src/tilefoundry/ir/hir/tensor/quant.py +++ b/src/tilefoundry/ir/hir/tensor/quant.py @@ -16,22 +16,19 @@ 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.types import DType, TensorType, TupleType +from tilefoundry.ir.types import DType, Layout, ShardLayout, TensorType, TupleType from tilefoundry.ir.types.dim import DimFloorDiv, simplify_dim -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard import ( - Layout, - ShardLayout, - canonical_shard_layout, - try_c_order_strides, -) -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.shard_layout import ( Broadcast, Split, + canonical_shard_layout, layout_axis_to_tensor_axis, shard_layout_local_shape, split_target_axes, ) +from tilefoundry.ir.types.stride import try_compact_major +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -95,7 +92,7 @@ def _logical_shard_attrs(call, ctx, x_ty, group: int): if not isinstance(attr, Split) or target != last_axis: continue split_extent = static_dim_value(layout.layout.shape[attr.axis]) - mesh_extent = static_dim_value(layout.mesh.layout.shape[mesh_axis]) + mesh_extent = static_dim_value(flatten(layout.mesh.layout).shape[mesh_axis]) aligned = ( aligned and split_extent is not None @@ -136,8 +133,8 @@ def _result_layouts(call, ctx, x_ty, scale_shape, group: int): if x_ty.layout is None: return None, None return ( - Layout(shape=x_ty.shape, strides=try_c_order_strides(x_ty.shape)), - Layout(shape=scale_shape, strides=try_c_order_strides(scale_shape)), + Layout(shape=x_ty.shape, strides=try_compact_major(x_ty.shape)), + Layout(shape=scale_shape, strides=try_compact_major(scale_shape)), ) diff --git a/src/tilefoundry/ir/hir/tensor/reduce.py b/src/tilefoundry/ir/hir/tensor/reduce.py index 9339c766..76ff3a0b 100644 --- a/src/tilefoundry/ir/hir/tensor/reduce.py +++ b/src/tilefoundry/ir/hir/tensor/reduce.py @@ -13,13 +13,9 @@ 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.types import TensorType -from tilefoundry.ir.types.shard import ( - Layout, - canonical_shard_layout, - try_c_order_strides, -) -from tilefoundry.ir.types.shard.shard_layout import shard_layout_of +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 from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -79,7 +75,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: new_layout = ( None if x_ty.layout is None - else Layout(shape=out_shape, strides=try_c_order_strides(out_shape)) + else Layout(shape=out_shape, strides=try_compact_major(out_shape)) ) source_shard = shard_layout_of(x_ty.layout) if source_shard is not None: diff --git a/src/tilefoundry/ir/hir/tensor/repeat_interleave.py b/src/tilefoundry/ir/hir/tensor/repeat_interleave.py index 33097fa8..0cfe9bc0 100644 --- a/src/tilefoundry/ir/hir/tensor/repeat_interleave.py +++ b/src/tilefoundry/ir/hir/tensor/repeat_interleave.py @@ -10,7 +10,7 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.shard.shard_layout import Broadcast, ShardLayout +from tilefoundry.ir.types.shard_layout import Broadcast, ShardLayout from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, diff --git a/src/tilefoundry/ir/hir/tensor/reshape.py b/src/tilefoundry/ir/hir/tensor/reshape.py index 97c57680..53eec571 100644 --- a/src/tilefoundry/ir/hir/tensor/reshape.py +++ b/src/tilefoundry/ir/hir/tensor/reshape.py @@ -8,15 +8,15 @@ 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.types import TensorType -from tilefoundry.ir.types.shard import ComposedLayout, try_c_order_strides -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types import ComposedLayout, TensorType +from tilefoundry.ir.types.layout import Layout, flatten +from tilefoundry.ir.types.shard_layout import ( Broadcast, ShardLayout, Split, ) from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.stride import try_compact_major from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( identity_access, @@ -87,7 +87,7 @@ def _carry_sharded_reshape(layout: ShardLayout, new_shape: tuple): axis_strides = axis_layout.strides n_axis = len(axis_shape) - mesh_shape = layout.mesh.layout.shape + mesh_shape = flatten(layout.mesh.layout).shape split_mesh_extent: dict[int, int] = {} for mesh_axis_idx, attr in enumerate(layout.attrs): if isinstance(attr, Split) and mesh_axis_idx < len(mesh_shape): @@ -193,18 +193,18 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: source = x_ty.layout if isinstance(source, Layout): source_strides = source.strides - expected_strides = try_c_order_strides(source.shape) + expected_strides = try_compact_major(source.shape) if source_strides is None or source_strides == expected_strides: new_layout = Layout( shape=new_shape, - strides=try_c_order_strides(new_shape), + strides=try_compact_major(new_shape), ) elif ( isinstance(source, ComposedLayout) and isinstance(source.outer, Layout) and ( source.outer.strides is None - or source.outer.strides == try_c_order_strides(source.outer.shape) + or source.outer.strides == try_compact_major(source.outer.shape) ) ): new_layout = ComposedLayout( @@ -212,7 +212,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: offset=source.offset, outer=Layout( shape=new_shape, - strides=try_c_order_strides(new_shape), + strides=try_compact_major(new_shape), ), ) return TensorType( diff --git a/src/tilefoundry/ir/hir/tensor/slice.py b/src/tilefoundry/ir/hir/tensor/slice.py index 7220b0de..6d0106b0 100644 --- a/src/tilefoundry/ir/hir/tensor/slice.py +++ b/src/tilefoundry/ir/hir/tensor/slice.py @@ -11,26 +11,15 @@ 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.types import TensorType -from tilefoundry.ir.types.dim import ( - DimAdd, - DimFloorDiv, - DimMul, - DimSub, - simplify_dim, -) -from tilefoundry.ir.types.shape_helpers import i64_const -from tilefoundry.ir.types.shard import ( - ComposedLayout, - Layout, - ShardLayout, - Swizzle, -) -from tilefoundry.ir.types.shard.shard_layout import ( +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 +from tilefoundry.ir.types.shard_layout import ( layout_axis_to_tensor_axis, split_target_axes, ) from tilefoundry.ir.types.substitute import dim_vars_by_name +from tilefoundry.ir.types.utils import i64_const from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( @@ -377,6 +366,62 @@ def _slice_shard_layout(call, ctx, x_ty, source, starts, inherited_offset): return ComposedLayout(inner=None, offset=offset, outer=sharded) +def window_image(layout, starts: tuple, sizes: tuple, steps: tuple) -> "tuple[int, Layout] | None": + """Where a window of *layout* starts and how it lies, at static starts. + + One axis at a time, in the tensor's own order. An axis written as a single + mode is a run: its start walks its stride into the offset and the window + keeps that stride, stepped. An axis written as a group has no single stride + to walk, so its start is read through the group as the mixed-radix + coordinate it is and the window is the modes the group steps the least by, + as many as its size takes, in the order the group wrote them. A grouped + axis straddling modes is refused; ``None`` is no window at all. + """ + if not isinstance(layout, Layout) or layout.strides is None: + return None + shape, strides = tuple(layout.shape), tuple(layout.strides) + if not len(shape) == len(strides) == len(starts) == len(sizes) == len(steps): + return None + + offset = 0 + window_shape: list = [] + window_strides: list = [] + for group, walk, start, size, step in zip(shape, strides, starts, sizes, steps): + if not isinstance(group, tuple) and not isinstance(walk, tuple): + if any(_literal(value) is None for value in (walk, start, step)): + return None + offset += _literal(start) * _literal(walk) + window_shape.append(size) + window_strides.append(_literal(walk) * _literal(step)) + continue + + extents, walked = flatten(group), flatten(walk) + numbers = (*extents, *walked, start, size, step) + if len(extents) != len(walked) or any(_literal(value) is None for value in numbers): + return None + extents = tuple(_literal(value) for value in extents) + walked = tuple(_literal(value) for value in walked) + start, size, step = _literal(start), _literal(size), _literal(step) + if step != 1 or size < 1 or start % size: + return None + + below, reach = 1, 1 + kept: list[int] = [] + for index in sorted(range(len(extents)), key=lambda axis: walked[axis]): + offset += ((start // below) % extents[index]) * walked[index] + if below < size: + kept.append(index) + reach *= extents[index] + below *= extents[index] + if reach != size: + return None + kept.sort() + window_shape.append(tuple(extents[index] for index in kept)) + window_strides.append(tuple(walked[index] for index in kept)) + + return offset, Layout(shape=tuple(window_shape), strides=tuple(window_strides)) + + @register_typeinfer(Slice) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: """The window's type, its source layout carried where the window keeps it. @@ -454,18 +499,16 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: static_starts.append(int(start.value)) steps.append(stride) else: - new_layout = ComposedLayout( - inner=inherited_inner, - offset=inherited_offset - + sum( - start * stride - for start, stride in zip(static_starts, source.strides) - ), - outer=Layout( - shape=layout_shape, - strides=tuple(stride * step for stride, step in zip(source.strides, steps)), - ), + found = window_image( + source, tuple(static_starts), tuple(layout_shape), tuple(steps) ) + if found is not None: + moved, window = found + new_layout = ComposedLayout( + inner=inherited_inner, + offset=inherited_offset + moved, + outer=window, + ) return TensorType(shape=shape, dtype=x_ty.dtype, layout=new_layout, storage=x_ty.storage) @@ -497,4 +540,4 @@ def _eval_slice(ctx): return TensorValue(data=ctx.args[0].data[tuple(key)], type=ctx.result_type) -__all__ = ["Slice", "slice_size", "window_base"] +__all__ = ["Slice", "slice_size", "window_base", "window_image"] diff --git a/src/tilefoundry/ir/hir/tensor/split.py b/src/tilefoundry/ir/hir/tensor/split.py index b5f7412f..1569064c 100644 --- a/src/tilefoundry/ir/hir/tensor/split.py +++ b/src/tilefoundry/ir/hir/tensor/split.py @@ -9,20 +9,18 @@ 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.types import TensorType, TupleType -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard import ( - Broadcast, - Layout, - Partial, - ShardLayout, +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 +from tilefoundry.ir.types.shard_layout import Split as SplitAttr +from tilefoundry.ir.types.shard_layout import ( canonical_shard_layout, + layout_axis_to_tensor_axis, shard_layout_of, - try_c_order_strides, + split_target_axes, ) -from tilefoundry.ir.types.shard.shard_layout import Split as ShardSplit -from tilefoundry.ir.types.shard.shard_layout import Split as SplitAttr -from tilefoundry.ir.types.shard.shard_layout import layout_axis_to_tensor_axis, split_target_axes +from tilefoundry.ir.types.stride import try_compact_major +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -55,7 +53,7 @@ def _reject_redistribution(ctx, call: "Call", x_ty, axis: int, parts: int) -> No if layout is None: return targets = split_target_axes(layout, x_ty.shape) - mesh = layout.mesh.layout.shape if layout.mesh is not None else () + mesh = flatten(layout.mesh.layout).shape if layout.mesh is not None else () for mesh_axis, attr in enumerate(layout.attrs): divides = mesh_axis < len(mesh) and mesh[mesh_axis] > 1 if isinstance(attr, ShardSplit) and divides and targets[mesh_axis] == axis: @@ -105,7 +103,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TupleType: elif x_ty.layout is None: part_layout = None else: - part_layout = Layout(shape=part_shape, strides=try_c_order_strides(part_shape)) + part_layout = Layout(shape=part_shape, strides=try_compact_major(part_shape)) part_ty = TensorType( shape=part_shape, dtype=x_ty.dtype, layout=part_layout, storage=x_ty.storage ) diff --git a/src/tilefoundry/ir/hir/tensor/stack.py b/src/tilefoundry/ir/hir/tensor/stack.py index 648672a0..067f1ae3 100644 --- a/src/tilefoundry/ir/hir/tensor/stack.py +++ b/src/tilefoundry/ir/hir/tensor/stack.py @@ -16,12 +16,9 @@ reject_dynamic_shards, require_uniform_partial_slices, ) -from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.shard import ( - Layout, - shard_layout_of, - try_c_order_strides, -) +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 from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -89,7 +86,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: "layout; use an explicit Reshard before Stack", ) else: - layout = Layout(shape=new_shape, strides=try_c_order_strides(new_shape)) + layout = Layout(shape=new_shape, strides=try_compact_major(new_shape)) storage = resolve_anchor_storage(ctx, call, *(t.storage for t in types)) return TensorType(shape=new_shape, dtype=base.dtype, layout=layout, storage=storage) diff --git a/src/tilefoundry/ir/hir/tensor/topk.py b/src/tilefoundry/ir/hir/tensor/topk.py index 51c91bce..aa63cd2f 100644 --- a/src/tilefoundry/ir/hir/tensor/topk.py +++ b/src/tilefoundry/ir/hir/tensor/topk.py @@ -19,7 +19,7 @@ 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.types import DType, TensorType, TupleType +from tilefoundry.ir.types import DType, Layout, TensorType, TupleType from tilefoundry.ir.types.dim import ( DimAdd, DimFloorDiv, @@ -30,13 +30,13 @@ DimVar, is_dim_expr, ) -from tilefoundry.ir.types.shard import Layout, try_c_order_strides -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.shard_layout import ( ShardLayout, Split, layout_axis_to_tensor_axis, shard_layout_of, ) +from tilefoundry.ir.types.stride import try_compact_major from tilefoundry.ir.types.tensor_type import ShapeDim from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.visitor_registry import register_typeinfer @@ -151,7 +151,7 @@ def _canonical_shard(sl: "ShardLayout", out_shape) -> "ShardLayout": when the shape is non-static; ``attrs`` and ``mesh`` pass through. """ out_shape = tuple(out_shape) - strides = try_c_order_strides(out_shape) or tuple(1 for _ in out_shape) + strides = try_compact_major(out_shape) or tuple(1 for _ in out_shape) return ShardLayout( layout=Layout(shape=out_shape, strides=strides), attrs=sl.attrs, @@ -202,7 +202,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TupleType: new_layout = ( None if x_ty.layout is None - else Layout(shape=out_shape, strides=try_c_order_strides(out_shape)) + else Layout(shape=out_shape, strides=try_compact_major(out_shape)) ) if source_shard is not None: relation = relations_of(call, ctx) diff --git a/src/tilefoundry/ir/hir/tensor/transpose.py b/src/tilefoundry/ir/hir/tensor/transpose.py index 67d9def8..671f5099 100644 --- a/src/tilefoundry/ir/hir/tensor/transpose.py +++ b/src/tilefoundry/ir/hir/tensor/transpose.py @@ -9,9 +9,9 @@ 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.types import TensorType -from tilefoundry.ir.types.shard import ComposedLayout, Layout, try_c_order_strides -from tilefoundry.ir.types.shard.shard_layout import shard_layout_of +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 from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AffineAccess, @@ -39,7 +39,7 @@ def _strides(type_: TensorType) -> tuple | None: return None if layout.strides is not None: return tuple(layout.strides) - return try_c_order_strides(tuple(layout.shape)) + return try_compact_major(tuple(layout.shape)) def _transpose_view(call: "Call", ctx) -> tuple: @@ -71,6 +71,14 @@ def _transpose_view(call: "Call", ctx) -> tuple: @register_typeinfer(Transpose) def _(call: "Call", ctx: "TypeInferContext") -> TensorType: + """The same bytes with the axes in another order, and the strides to match. + + A tensor that states no layout is the C order every layer of this IR reads + it as, and permuting that is not the C order of the result: a (K, M) read + with strides (M, 1) transposes to an (M, K) view whose strides are (1, M), + not (K, 1). So an unstated layout is written out as the strides it stands + for and those are permuted with the shape. + """ x_ty = ctx.type_of(call.args[0]) perm = call.target.perm if len(perm) != len(x_ty.shape): @@ -86,6 +94,10 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: new_layout = derived else: source = x_ty.layout + if source is None: + source = Layout(shape=tuple(x_ty.shape), strides=try_compact_major(tuple(x_ty.shape))) + if source.strides is None: + source = None if isinstance(source, Layout): new_layout = Layout( shape=tuple(source.shape[p] for p in perm), diff --git a/src/tilefoundry/ir/hir/tensor/where.py b/src/tilefoundry/ir/hir/tensor/where.py index cfe40499..d2c7bd0a 100644 --- a/src/tilefoundry/ir/hir/tensor/where.py +++ b/src/tilefoundry/ir/hir/tensor/where.py @@ -13,10 +13,10 @@ from tilefoundry.ir.hir._helpers import broadcast_shapes, is_one, resolve_anchor_storage from tilefoundry.ir.hir._shard_checks import reject_partials from tilefoundry.ir.hir.math.binary import _merge_layout -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import Layout, try_c_order_strides -from tilefoundry.ir.types.shard.shard_layout import Broadcast, shard_layout_of +from tilefoundry.ir.types import DType, Layout, TensorType +from tilefoundry.ir.types.shard_layout import Broadcast, shard_layout_of from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.stride import try_compact_major from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, @@ -131,7 +131,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: storage = resolve_anchor_storage(ctx, call, input_.storage, other.storage) if layout is None and storage in (StorageKind.RMEM, StorageKind.SMEM) and out_shape: - layout = Layout(shape=out_shape, strides=try_c_order_strides(out_shape)) + layout = Layout(shape=out_shape, strides=try_compact_major(out_shape)) return TensorType( shape=out_shape, dtype=input_.dtype, diff --git a/src/tilefoundry/ir/mesh_scope.py b/src/tilefoundry/ir/mesh_scope.py new file mode 100644 index 00000000..ce2f870a --- /dev/null +++ b/src/tilefoundry/ir/mesh_scope.py @@ -0,0 +1,250 @@ +"""Which scope a program runs in, and what that scope admits. + +Entering a `with Mesh(...)` composes the scope in force with the one it names; +a value is read where the positions it is laid out over are the positions the +scope runs on. Both questions are about the mesh a statement stands inside, +which is neither a type nor a visitor, so they live here rather than with +either. +""" + +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_algebra import is_inverse_projectable +from tilefoundry.ir.types.mesh import Mesh +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))) + + +def device_layout(mesh: Mesh) -> Layout: + """The whole mesh's axes as the device numbers its positions. + + Each level states its own numbering, and a position of one level is a + position within its parent, so an axis steps by what it states times what + the levels under it hold. That product is the numbering a coordinate of + the whole mesh is read in, and the one a mesh is written down in. + """ + shape: list = [] + strides: list = [] + sizes = tuple(getattr(topology, "size", 1) for topology in mesh.topologies) + units = compact_major(sizes, major="row") if all( + isinstance(one, int) for one in sizes + ) else (1,) * len(sizes) + for arrangement, unit in zip(_levels(mesh), units): + stated = arrangement.strides + shape.extend(flatten(arrangement.shape)) + strides.extend( + step * unit + for step in ( + flatten(stated) + if stated is not None + else compact_major(tuple(flatten(arrangement.shape))) + ) + ) + return Layout(shape=tuple(shape), strides=tuple(strides)) + + +def _selected(arrangement: Layout, start: int) -> tuple[tuple, tuple, int]: + """One level's positions as a set of them reads: its modes, and where it starts. + + An axis of one position names no instance, and modes written in another + order state the same positions, so the modes come back sorted by step with + the adjacent ones joined and the ones of a single position left out. + """ + strides = arrangement.strides + if strides is None: + return tuple(flatten(arrangement.shape)), (), start + modes = [ + (extent, stride) + for extent, stride in zip(flatten(arrangement.shape), flatten(strides)) + if extent != 1 + ] + joined: list[list] = [] + for extent, stride in sorted(modes, key=lambda mode: (mode[1], mode[0])): + if joined and joined[-1][0] * joined[-1][1] == stride: + joined[-1][0] *= extent + else: + joined.append([extent, stride]) + return ( + tuple(extent for extent, _ in joined), + tuple(stride for _, stride in joined), + start, + ) + + +def covered_by_scope(mesh: Mesh, current: Mesh) -> bool: + """Whether *mesh* selects exactly the positions the enclosing scope does. + + Level by level, on the positions each level states rather than on the axes + standing where it does: a scope that is part of a level -- one warp of a + CTA's threads -- says which positions it is by where its run starts and how + its modes step, and a value laid out over that same run is inside it + however either of them wrote the axes down. + """ + scope = { + getattr(topology, "name", topology): _selected(arrangement, start) + for topology, arrangement, start in zip( + current.topologies, _levels(current), _starts(current) + ) + } + return all( + getattr(topology, "name", topology) in scope + and _selected(arrangement, start) + == scope[getattr(topology, "name", topology)] + for topology, arrangement, start in zip( + mesh.topologies, _levels(mesh), _starts(mesh) + ) + ) + + +def storage_reaches(storage, mesh: Mesh, current: Mesh) -> bool: + """Whether *storage* reaches across a coarser value-to-scope boundary.""" + if current.topologies[-1].name in {topology.name for topology in mesh.topologies}: + return True + try: + storage = resolve_storage(storage) + except (TypeError, ValueError): + return False + return storage in {StorageKind.GMEM, StorageKind.SMEM} + + +def states_consistent_positions(mesh: Mesh) -> bool: + """Whether the positions a mesh states are the ones its levels declare.""" + return product(mesh.topologies) == size(mesh.layout) + + +def mesh_scope_matches_required_scope(current: Mesh, required: Mesh) -> bool: + """True iff ``current`` provides the thread participation ``required`` needs.""" + if current.topologies[0].name != required.topologies[0].name: + return False + if not states_consistent_positions(current) or not states_consistent_positions(required): + return False + here, there = _flat(current), _flat(required) + if not is_inverse_projectable(here): + return False + return here.shape == there.shape and here.strides == there.strides + + +def _flat(mesh: Mesh) -> Layout: + stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout + strides = stated.strides + return Layout( + shape=tuple(flatten(stated.shape)), + strides=None if strides is None else tuple(flatten(strides)), + ) + + +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(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/tir/cuda/nn/mma.py b/src/tilefoundry/ir/tir/cuda/nn/mma.py index 491b7328..bc851736 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma.py @@ -13,14 +13,7 @@ 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.types import DType, UnitType -from tilefoundry.ir.types.shard import ( - Layout, - Mesh, - ShardLayout, - Split, - Topology, -) +from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology, UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt from .mma_atom import MmaAtom, MmaOpSpec @@ -91,7 +84,7 @@ def _(call: "Call", ctx: "VerifyContext") -> None: call, f"Mma {role} fragment layout does not match atom {_ATOM_ROLE[role]}", ) - from tilefoundry.ir.types.shard.scope_match import ( # noqa: PLC0415 + from tilefoundry.ir.mesh_scope import ( # noqa: PLC0415 mesh_scope_matches_required_scope, ) diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py index 81f21c53..86aca13f 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py @@ -4,8 +4,7 @@ from dataclasses import dataclass -from tilefoundry.ir.types import DType -from tilefoundry.ir.types.shard import Mesh, ShardLayout +from tilefoundry.ir.types import DType, Mesh, ShardLayout @dataclass(frozen=True) diff --git a/src/tilefoundry/ir/tir/dot.py b/src/tilefoundry/ir/tir/dot.py index a0cf9a74..5f18c4e3 100644 --- a/src/tilefoundry/ir/tir/dot.py +++ b/src/tilefoundry/ir/tir/dot.py @@ -10,7 +10,7 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import UnitType -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, shard_layout_local_shape +from tilefoundry.ir.types.shard_layout import ShardLayout, shard_layout_local_shape from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/tir/launch.py b/src/tilefoundry/ir/tir/launch.py index 8a288303..87548253 100644 --- a/src/tilefoundry/ir/tir/launch.py +++ b/src/tilefoundry/ir/tir/launch.py @@ -121,6 +121,8 @@ def launch_call( CallableType, DType, TensorType, + ) + from tilefoundry.ir.types.callable_type import ( # noqa: PLC0415 callable_type_for_prim_function, ) from tilefoundry.ir.types.dim import ( # noqa: PLC0415 diff --git a/src/tilefoundry/ir/tir/memory/copy.py b/src/tilefoundry/ir/tir/memory/copy.py index a4a78700..97252dff 100644 --- a/src/tilefoundry/ir/tir/memory/copy.py +++ b/src/tilefoundry/ir/tir/memory/copy.py @@ -13,7 +13,7 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import UnitType -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +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/tensor_view.py b/src/tilefoundry/ir/tir/memory/tensor_view.py index 8f2a3538..0ab48a3f 100644 --- a/src/tilefoundry/ir/tir/memory/tensor_view.py +++ b/src/tilefoundry/ir/tir/memory/tensor_view.py @@ -16,8 +16,8 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.shard import c_order_strides -from tilefoundry.ir.types.shard.layout import Layout, LayoutBase +from tilefoundry.ir.types.layout import Layout, LayoutBase +from tilefoundry.ir.types.stride import compact_row_major from tilefoundry.visitor_registry import register_typeinfer @@ -54,7 +54,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: def _c_order_strides(src_shape: tuple) -> list: """C-order contiguous strides of the source buffer a slice view reads.""" - return list(c_order_strides(tuple(src_shape))) + return list(compact_row_major(tuple(src_shape))) def layout_for_slice(src_shape: tuple, axis: int, sliced_shape: tuple) -> Layout: diff --git a/src/tilefoundry/ir/tir/stmts.py b/src/tilefoundry/ir/tir/stmts.py index b0a60816..a6f3270b 100644 --- a/src/tilefoundry/ir/tir/stmts.py +++ b/src/tilefoundry/ir/tir/stmts.py @@ -13,7 +13,7 @@ from tilefoundry.ir.core import Expr, Var from tilefoundry.ir.tir.stmt import Stmt -from tilefoundry.ir.types.shard.mesh import Mesh +from tilefoundry.ir.types.mesh import Mesh @dataclass(unsafe_hash=True) diff --git a/src/tilefoundry/ir/tir/symbol_ref.py b/src/tilefoundry/ir/tir/symbol_ref.py index e81fc8b7..4189f382 100644 --- a/src/tilefoundry/ir/tir/symbol_ref.py +++ b/src/tilefoundry/ir/tir/symbol_ref.py @@ -22,7 +22,7 @@ def symbol_call(callee, args) -> "Evaluate": # noqa: F821 -- lazy Evaluate callee ``PrimFunction`` by symbol. """ from tilefoundry.ir.tir.stmts import Evaluate # noqa: PLC0415 - from tilefoundry.ir.types import callable_type_for_prim_function # noqa: PLC0415 + from tilefoundry.ir.types.callable_type import callable_type_for_prim_function # noqa: PLC0415 ref = SymbolRef(name=callee.name, type=callable_type_for_prim_function(callee)) return Evaluate(callable=ref, args=tuple(args)) diff --git a/src/tilefoundry/ir/tir/sync.py b/src/tilefoundry/ir/tir/sync.py index 435875af..6337f01c 100644 --- a/src/tilefoundry/ir/tir/sync.py +++ b/src/tilefoundry/ir/tir/sync.py @@ -8,12 +8,13 @@ from tilefoundry.ir.core import Op, VerifyError from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.mesh_scope import device_layout from tilefoundry.ir.types import UnitType -from tilefoundry.ir.types.shard import product -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout -from tilefoundry.ir.types.shard.layout_algebra import apply as _apply -from tilefoundry.ir.types.shard.layout_algebra import size as _size -from tilefoundry.ir.types.shard.mesh import Mesh +from tilefoundry.ir.types.int_tuple import flatten, product +from tilefoundry.ir.types.layout import ComposedLayout, Layout, get +from tilefoundry.ir.types.layout_algebra import apply as _apply +from tilefoundry.ir.types.layout_algebra import size as _size +from tilefoundry.ir.types.mesh import Mesh from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt _WARP_SIZE = 32 @@ -44,9 +45,10 @@ def _legal_slice_of(m: Mesh, e: Mesh) -> bool: region = m.layout if not isinstance(region, ComposedLayout): return False - outer = region.outer - p = e.layout - if not isinstance(outer, Layout) or outer.strides != p.strides: + if not isinstance(region.outer, Layout): + return False + outer, p = get(region.outer, 0), get(e.layout, 0) + if outer.strides != p.strides: return False sub, pshape = outer.shape, p.shape if len(sub) != len(pshape): @@ -78,7 +80,7 @@ def _legal_slice_of(m: Mesh, e: Mesh) -> bool: def _mesh_str(mesh: Mesh) -> str: """A mesh as its topologies and the shape they are viewed through.""" topologies = ", ".join(f"{t.name}({t.size})" for t in mesh.topologies) - return f"({topologies})[{', '.join(str(d) for d in mesh.layout.shape)}]" + return f"({topologies})[{', '.join(str(d) for d in flatten(mesh.layout.shape))}]" def _no_enclosing_mesh_error(m: Mesh, scope: "tuple[Mesh, ...]") -> VerifyError: @@ -155,17 +157,15 @@ class Participation: def _participant_layout(mesh: Mesh) -> "tuple[Layout, int]": """The (outer layout, offset) describing which threads participate. - For a sliced mesh ``layout`` is a ``ComposedLayout`` whose ``outer`` is the - participating sub-box and ``offset`` the slice origin; for an un-sliced mesh - the whole plain-``Layout`` ``layout`` participates at offset 0. + Every level states its own run, and which threads participate is those + runs read as the device numbers them: the mesh's axes, at its offset. """ - ly = mesh.layout - if isinstance(ly, ComposedLayout): - outer = ly.outer - if not isinstance(outer, Layout): - raise VerifyError("T.sync: mesh slice must be a plain-Layout affine scope") - return outer, ly.offset - return ly, 0 + try: + return device_layout(mesh), ( + mesh.layout.offset if isinstance(mesh.layout, ComposedLayout) else 0 + ) + except ValueError as error: + raise VerifyError(f"T.sync: {error}") from error def participation(mesh: Mesh) -> Participation: diff --git a/src/tilefoundry/ir/tir/verify.py b/src/tilefoundry/ir/tir/verify.py index 95d2984b..923434bb 100644 --- a/src/tilefoundry/ir/tir/verify.py +++ b/src/tilefoundry/ir/tir/verify.py @@ -15,26 +15,17 @@ 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.types import ( - DType, - TensorType, - UnitType, - callable_type_for_prim_function, -) -from tilefoundry.ir.types.dim import ( - DimAdd, - DimFloorDiv, - DimMax, - DimMin, - DimMod, - DimMul, - DimSub, -) -from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +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 +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.mesh import Mesh +from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.ir.visitor import ExprVisitor +from tilefoundry.ir.types.utils import static_dim_value +from tilefoundry.ir.visitor import ExprVisitor, collect_exprs from tilefoundry.target import CudaTarget from tilefoundry.utils.spec_ref import spec_ref_render from tilefoundry.visitor_registry import verify_stmt_registry @@ -145,6 +136,12 @@ def _walk_stmt(stmt, ctx, scope, fn, module_fn_map, bound_var_ids: set[int]): _check_rank0_int(ctx, stmt, stmt.start, "For.start") _check_rank0_int(ctx, stmt, stmt.stop, "For.stop") _check_rank0_int(ctx, stmt, stmt.step, "For.step") + for field, bound in ( + ("start", stmt.start), + ("stop", stmt.stop), + ("step", stmt.step), + ): + _check_bound_coordinates(field, bound, scope) if isinstance(stmt.step, Constant) and stmt.step.value == 0: raise VerifyError("For.step must not be 0") iv_ty = stmt.induction_var.type @@ -256,6 +253,30 @@ def _reject_nested_alloc_tensor(expr: Expr, *, at_letstmt_value: bool) -> None: visitor.visit(expr) +def _check_bound_coordinates(field: str, bound, scope) -> None: + """Hold a loop bound's mesh coordinates to the scopes around the loop. + + A coordinate is a coordinate of one scope, read through that scope's + binding, so a bound naming one has to stand inside a ``MeshScope`` of the + mesh it names -- the enclosing scope is what says which unit it is about. + """ + for expr in collect_exprs(bound): + if not (isinstance(expr, Call) and isinstance(expr.target, MeshCoord)): + continue + mesh = expr.target.mesh + axis = static_dim_value(expr.args[0]) if expr.args else None + if axis is None or not 0 <= axis < len(flatten(mesh.layout).shape): + raise VerifyError( + f"For.{field} reads axis {axis!r} of a mesh of rank " + f"{len(flatten(mesh.layout).shape)}" + ) + 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" + ) + + def _check_embedded_sharding(expr: Expr, scope, fn): """Require embedded shard layouts to use a scoped or parameter mesh. diff --git a/src/tilefoundry/ir/types/__init__.py b/src/tilefoundry/ir/types/__init__.py index 7e4fa538..c63273b6 100644 --- a/src/tilefoundry/ir/types/__init__.py +++ b/src/tilefoundry/ir/types/__init__.py @@ -3,44 +3,59 @@ # ruff: noqa: I001 -- curated re-export order; alphabetical sort breaks staged imports. from .dtype import BoolDType, DType, FloatDType, IntegerDType -from .tensor_type import TensorType, TupleType, Type, UnitType -from .utils import ( - bytes_by_storage, - local_type_of, - make_shard_tensor_type, - make_tensor_type, - numel, - tensor_types, - tensor_bytes, - topology_extent, -) -from .callable_type import ( - CallableType, - callable_type_for, - callable_type_for_prim_function, +from .storage import StorageKind +from .int_tuple import IntTuple +from .layout import ComposedLayout, Layout, LayoutBase, Swizzle +from .shard_layout import ( + B, + Broadcast, + Dynamic, + P, + Partial, + S, + ShardAttr, + ShardLayout, + Split, ) +from .mesh import Mesh, Topology, make_mesh +from .placement import Placement +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 __all__ = [ + "B", "BoolDType", + "Broadcast", "CallableType", + "ComposedLayout", "DType", + "Dynamic", "FloatDType", + "IntTuple", "IntegerDType", + "Layout", + "LayoutBase", + "Mesh", + "P", + "Partial", + "Placement", + "S", + "ShardAttr", + "ShardLayout", + "Split", + "StorageKind", + "Swizzle", "TensorType", + "Topology", "TupleType", "Type", "UnitType", "callable_type_for", - "callable_type_for_prim_function", - "bytes_by_storage", - "local_type_of", + "make_mesh", "make_shard_tensor_type", "make_tensor_type", - "numel", - "tensor_types", - "tensor_bytes", - "topology_extent", ] diff --git a/src/tilefoundry/ir/types/int_tuple.py b/src/tilefoundry/ir/types/int_tuple.py new file mode 100644 index 00000000..ff81f27f --- /dev/null +++ b/src/tilefoundry/ir/types/int_tuple.py @@ -0,0 +1,77 @@ +"""IntTuple alias + helpers.""" + +from __future__ import annotations + +from typing import Union, overload + +from tilefoundry.ir.types.tensor_type import ShapeDim + +IntTuple = Union[int, tuple["IntTuple", ...]] + + +@overload +def flatten(t: IntTuple) -> tuple[int, ...]: ... + + +@overload +def flatten(t: object) -> tuple[object, ...]: ... + + +def flatten(t: object) -> tuple[object, ...]: + if not isinstance(t, tuple): + return (t,) + return tuple(value for item in t for value in flatten(item)) + + +def product(t) -> "ShapeDim": + from .mesh import Topology # noqa: PLC0415 + + result = 1 + for v in flatten(t): + if isinstance(v, Topology): + v = v.size + result *= v + return result + + +def repeat_like(profile, value) -> object: + """CuTe ``repeat_like``: *value* at every leaf, nested like *profile*. + + What a flat tuple has to say about a nested one is said by building it + against the nesting and flattening that, rather than by working out which + flat positions each mode covers. + """ + if not isinstance(profile, tuple): + return value + return tuple(repeat_like(item, value) for item in profile) + + +def _unflatten(flat: tuple, profile) -> tuple: + """Take *profile*'s worth of *flat*, returning it nested and what is left.""" + if not isinstance(profile, tuple): + if not flat: + raise ValueError("unflatten: the profile asks for more modes than the tuple has") + return flat[0], flat[1:] + nested: list = [] + for item in profile: + value, flat = _unflatten(flat, item) + nested.append(value) + return tuple(nested), flat + + +def unflatten(flat: tuple, profile) -> tuple: + """CuTe ``unflatten``: nest a flat tuple to *profile*'s structure. + + Only *profile*'s nesting is read, never its leaves, so the profile may be + the grouping itself. ``flatten(unflatten(t, p)) == t``. + """ + nested, rest = _unflatten(flat, profile) + if rest: + raise ValueError( + f"unflatten: the profile accounts for {len(flat) - len(rest)} of the " + f"tuple's {len(flat)} modes" + ) + return nested + + +__all__ = ["IntTuple", "flatten", "product", "repeat_like", "unflatten"] diff --git a/src/tilefoundry/ir/types/shard/layout.py b/src/tilefoundry/ir/types/layout.py similarity index 61% rename from src/tilefoundry/ir/types/shard/layout.py rename to src/tilefoundry/ir/types/layout.py index 665ba70f..89504d1b 100644 --- a/src/tilefoundry/ir/types/shard/layout.py +++ b/src/tilefoundry/ir/types/layout.py @@ -3,7 +3,8 @@ from dataclasses import dataclass from typing import Optional -from .int_tuple import flatten +from .int_tuple import flatten as _flat +from .int_tuple import product class LayoutBase: @@ -111,4 +112,73 @@ def shape(self) -> tuple: EMPTY_LAYOUT = Layout(shape=(), strides=()) -__all__ = ["LayoutBase", "Layout", "Swizzle", "ComposedLayout", "EMPTY_LAYOUT"] +def size(layout: Layout) -> int: + return product(layout.shape) + + +def flatten(layout): + """CuTe ``flatten``: every mode at the top level, of an arrangement or a tuple. + + CuTe spells this once for each (``layout.hpp`` and the tuple algorithms); + here one name reads both, because which was handed over is plain from what + comes back. + """ + if not isinstance(layout, LayoutBase): + return _flat(layout) + if isinstance(layout, ComposedLayout): + return flatten(layout.outer) if layout.outer is not None else EMPTY_LAYOUT + strides = getattr(layout, "strides", None) + return Layout( + shape=_flat(layout.shape), + strides=None if strides is None else _flat(strides), + ) + + +def unflatten(layout: LayoutBase, profile) -> "Layout": + """CuTe ``unflatten``: a flat arrangement nested to *profile*'s shape.""" + from .int_tuple import unflatten as unflatten_tuple # noqa: PLC0415 - cycle guard + + strides = getattr(layout, "strides", None) + return Layout( + shape=unflatten_tuple(tuple(layout.shape), profile), + strides=None if strides is None else unflatten_tuple(tuple(strides), profile), + ) + + +def rank(layout: LayoutBase) -> int: + """CuTe ``rank``: how many modes a layout states at its top level.""" + return len(layout.shape) + + +def get(layout: LayoutBase, index: int) -> "Layout": + """CuTe ``get``: one mode of a layout, as a layout of its own.""" + shape = layout.shape[index] + strides = layout.strides[index] if getattr(layout, "strides", None) is not None else None + return Layout( + shape=shape if isinstance(shape, tuple) else (shape,), + strides=None if strides is None else (strides if isinstance(strides, tuple) else (strides,)), + ) + + +def take(layout: LayoutBase, begin: int, end: int) -> "Layout": + """CuTe ``take``: the modes in ``[begin, end)``, as a layout.""" + strides = getattr(layout, "strides", None) + return Layout( + shape=tuple(layout.shape[begin:end]), + strides=None if strides is None else tuple(strides[begin:end]), + ) + + +__all__ = [ + "LayoutBase", + "flatten", + "size", + "unflatten", + "Layout", + "Swizzle", + "ComposedLayout", + "EMPTY_LAYOUT", + "get", + "rank", + "take", +] diff --git a/src/tilefoundry/ir/types/shard/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py similarity index 83% rename from src/tilefoundry/ir/types/shard/layout_algebra.py rename to src/tilefoundry/ir/types/layout_algebra.py index 43e7985a..82303e35 100644 --- a/src/tilefoundry/ir/types/shard/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -12,8 +12,10 @@ from typing import Optional, Union -from .int_tuple import flatten, product -from .layout import ComposedLayout, Layout, Swizzle +from tilefoundry.ir.types.layout import flatten + +from .layout import ComposedLayout, Layout, Swizzle, size +from .stride import compact_col_major, idx2crd class NotProjectable(ValueError): @@ -27,52 +29,7 @@ def _shape(layout: Layout) -> tuple[int, ...]: def _stride(layout: Layout) -> tuple[int, ...]: if layout.strides is not None: return layout.strides - return prefix_product(_shape(layout)) - - -def prefix_product(shape: tuple[int, ...]) -> tuple[int, ...]: - """Exclusive prefix product (column-major natural strides).""" - out: list[int] = [] - acc = 1 - for s in shape: - out.append(acc) - acc *= s - return tuple(out) - - -def c_order_strides(shape: tuple, *, mul=None) -> tuple: - """Row-major (C-order) contiguous strides. - - Row-major (C-order) contiguous strides: ``strides[-1] == 1``, - ``strides[i] == strides[i+1] * shape[i+1]``. - - The single home for this computation. *mul* defaults to ``int`` - multiplication; pass a dim-expression fold (e.g. wrapping - ``simplify_dim(DimMul, ...)``) for shapes with symbolic entries. - """ - if not shape: - return () - if mul is None: - mul = lambda a, b: a * b # noqa: E731 - strides = [1] * len(shape) - for i in range(len(shape) - 2, -1, -1): - strides[i] = mul(strides[i + 1], shape[i + 1]) - return tuple(strides) - - -def try_c_order_strides(shape: tuple) -> tuple[int, ...] | None: - """``c_order_strides`` when every entry is a static non-bool ``int``, else ``None``. - - ``c_order_strides`` when every entry is a static non-bool ``int``, - else ``None`` (symbolic / dynamic shapes have no static strides). - """ - if not all(isinstance(s, int) and not isinstance(s, bool) for s in shape): - return None - return c_order_strides(shape) - - -def size(layout: Layout) -> int: - return product(layout.shape) + return compact_col_major(_shape(layout)) def swizzle_of(layout: object) -> Optional[Swizzle]: @@ -140,39 +97,6 @@ def cosize(layout: Union[Layout, ComposedLayout]) -> int: return apply(layout, size(layout) - 1) + 1 -def idx2crd(idx: int, shape: tuple[int, ...], stride: tuple[int, ...]) -> tuple[int, ...]: - """Per-mode ``(idx // stride_i) % shape_i`` (CuTe ``idx2crd``).""" - return tuple((idx // d) % s for s, d in zip(shape, stride)) - - -def _unflatten(flat: tuple, profile) -> tuple: - """Take *profile*'s worth of *flat*, returning it nested and what is left.""" - if not isinstance(profile, tuple): - if not flat: - raise ValueError("unflatten: the profile asks for more modes than the tuple has") - return flat[0], flat[1:] - nested: list = [] - for item in profile: - value, flat = _unflatten(flat, item) - nested.append(value) - return tuple(nested), flat - - -def unflatten(flat: tuple, profile) -> tuple: - """CuTe ``unflatten``: nest a flat tuple to *profile*'s structure. - - Only *profile*'s nesting is read, never its leaves, so the profile may be - the grouping itself. ``flatten(unflatten(t, p)) == t``. - """ - nested, rest = _unflatten(flat, profile) - if rest: - raise ValueError( - f"unflatten: the profile accounts for {len(flat) - len(rest)} of the " - f"tuple's {len(flat)} modes" - ) - return nested - - def coalesce(layout: Union[Layout, ComposedLayout]): """Flatten + merge contiguous modes, drop shape-1 modes (CuTe ``coalesce``). @@ -252,7 +176,7 @@ def _right_inverse_layout(layout: Layout) -> Layout: current_idx = 1 shape = _shape(layout) stride = _stride(layout) - triples = sorted(zip(stride, shape, prefix_product(shape))) + triples = sorted(zip(stride, shape, compact_col_major(shape))) for st, sh, rstride in triples: if sh == 1: continue @@ -278,7 +202,7 @@ def _is_identity_inner(inner: object) -> bool: if inner is None: return True if isinstance(inner, Layout): - return _stride(inner) == prefix_product(_shape(inner)) + return _stride(inner) == compact_col_major(_shape(inner)) return False @@ -465,7 +389,7 @@ def project(scope: ComposedLayout, t: int) -> Optional[tuple[int, ...]]: return None shape = _shape(outer) - return idx2crd(coord_1d, shape, prefix_product(shape)) + return idx2crd(coord_1d, shape, compact_col_major(shape)) def contains(scope: ComposedLayout, t: int) -> bool: @@ -475,14 +399,10 @@ def contains(scope: ComposedLayout, t: int) -> bool: __all__ = [ "NotProjectable", - "prefix_product", - "size", "swizzle_of", "composition", "cosize", "apply", - "idx2crd", - "unflatten", "coalesce", "complement", "is_inverse_projectable", diff --git a/src/tilefoundry/ir/types/mesh.py b/src/tilefoundry/ir/types/mesh.py new file mode 100644 index 00000000..8aed90c0 --- /dev/null +++ b/src/tilefoundry/ir/types/mesh.py @@ -0,0 +1,244 @@ +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.tensor_type import ShapeDim + + +@dataclass(frozen=True) +class Topology: + """Name one hardware level and its explicit static or symbolic size.""" + + name: str + + size: "ShapeDim" + + def __post_init__(self) -> None: + if self.size is None: + raise ValueError( + f"Topology {self.name!r}: extent must be explicit; None is not " + "a ShapeDim. The rule: " + "tilefoundry spec target topology-levels" + ) + + +@dataclass(frozen=True) +class Mesh: + """Describe hardware levels, logical positions, and axis names. + + ``layout`` runs parallel to ``topologies``: one arrangement per level, in + that level's own numbering, so a level states which of its positions it + selects without any level above it entering the answer. A constant slice + replaces one level's arrangement with a ``ComposedLayout`` whose ``offset`` + is where that level's run starts. + + See [shard §5](docs/spec/shard.md#5-mesh). + """ + + topologies: tuple[Topology | str, ...] + layout: "Layout | ComposedLayout" + names: tuple[str, ...] = () + + def __post_init__(self) -> None: + object.__setattr__(self, "layout", _nested(self.layout, tuple(self.topologies))) + for axis, extent in enumerate(flatten(self.layout.shape)): + if extent is None: + raise ValueError( + f"Mesh: layout axis {axis} must have an explicit extent; " + "None is not a ShapeDim. The rule: tilefoundry spec shard mesh" + ) + + def __getitem__(self, key) -> "Mesh": + """Return a constant sub-mesh selected by integers or unit-step slices. + + 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. + + 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) + 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] = [] + 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): + raise ValueError( + f"mesh slice {k.start}:{k.stop} out of range for axis " + f"{axis} (extent {extent})" + ) + 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) + + return Mesh( + topologies=self.topologies, + layout=ComposedLayout( + inner=None, + offset=offset, + outer=Layout(shape=(tuple(sub_shape),), strides=(tuple(strides),)), + ), + names=self.names, + ) + + +def _levelled(layout, topologies: tuple) -> bool: + """Whether *layout* already states one mode per level the mesh names.""" + return _rank(layout) == len(topologies) and all( + isinstance(mode, tuple) for mode in layout.shape + ) + + +def _nested(layout, topologies: tuple) -> "Layout | ComposedLayout": + """What a mesh was written as, with every level's axes under its own mode. + + A mesh naming one level states that level's arrangement directly. One + naming several may state a tuple of extents or one arrangement over all of + their axes; the axes are handed to the levels left to right, each taking + them until their extents multiply to its own size, and each level's steps + are divided by what the levels under it hold. So mode ``i`` of what comes + back is level ``i``'s own arrangement, and a boundary no prefix of axes + lands on is refused rather than guessed. + """ + if not isinstance(layout, LayoutBase): + 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( + inner=layout.inner, + offset=layout.offset, + outer=_nested(layout.outer, topologies), + ) + if _levelled(layout, topologies): + return layout + + extents = tuple(flatten(layout.shape)) + stated = layout.strides if layout.strides is not None else compact_row_major(extents) + steps = tuple(flatten(stated)) + if len(topologies) == 1: + return Layout(shape=(extents,), strides=(steps,)) + units: list[int] = [] + below = 1 + for topology in reversed(topologies): + units.insert(0, below) + size = getattr(topology, "size", None) + 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"{size!r}; cutting one arrangement at the level boundaries needs " + "each of their position counts" + ) + below *= size + + shape: list = [] + strides: list = [] + axis = 0 + for topology, unit in zip(topologies, units): + size = topology.size + taken_extents: list = [] + taken_steps: list = [] + reach = 1 + while axis < len(extents) and (reach < size or extents[axis] == 1): + extent, step = extents[axis], steps[axis] + if not isinstance(extent, int) or isinstance(extent, bool): + raise ValueError( + f"mesh layout axis {axis} states extent {extent!r}; cutting one " + "arrangement at the level boundaries needs concrete axis extents" + ) + if not isinstance(step, int) or isinstance(step, bool) or step % unit: + raise ValueError( + f"mesh axis {axis} steps by {step!r}, which the {unit} positions " + f"below {topology.name!r} do not divide; its positions are not " + "that level's" + ) + reach *= extent + taken_extents.append(extent) + taken_steps.append(step // unit) + axis += 1 + if reach != size: + raise ValueError( + f"mesh axes {extents} do not land on the boundary of level " + f"{topology.name!r} at {size}: the axes up to there multiply to " + f"{reach}. Write the axis that straddles it as the two axes it is" + ) + shape.append(tuple(taken_extents)) + strides.append(tuple(taken_steps)) + if axis != len(extents): + raise ValueError( + f"mesh layout has {len(extents)} axes but the levels it names account " + f"for {axis}; every axis belongs to one of them" + ) + return Layout(shape=tuple(shape), strides=tuple(strides)) + + +def make_mesh( + + layout_shape: tuple, + names: "tuple[str, ...] | None" = None, + topology: "str | Topology" = "gpu", +) -> Mesh: + """Convenience constructor for a ``Mesh`` with the given axis extents and C-order strides. + + 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``. + + ``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), + ) + + +__all__ = ["Mesh", "Topology", "make_mesh"] diff --git a/src/tilefoundry/ir/types/shard/placement.py b/src/tilefoundry/ir/types/placement.py similarity index 97% rename from src/tilefoundry/ir/types/shard/placement.py rename to src/tilefoundry/ir/types/placement.py index 8c160d86..ca913946 100644 --- a/src/tilefoundry/ir/types/shard/placement.py +++ b/src/tilefoundry/ir/types/placement.py @@ -16,7 +16,7 @@ from collections.abc import Mapping from dataclasses import dataclass -from tilefoundry.ir.types.shard.mesh import Topology +from tilefoundry.ir.types.mesh import Topology @dataclass(frozen=True) diff --git a/src/tilefoundry/ir/types/shape_helpers.py b/src/tilefoundry/ir/types/shape_helpers.py deleted file mode 100644 index c26614ae..00000000 --- a/src/tilefoundry/ir/types/shape_helpers.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Handle static and symbolic entries in tensor shapes. - -Upper-bound helpers produce compile-time allocation sizes. Runtime-total -helpers combine static factors with expressions for actual symbolic extents. - -See [types §4](docs/spec/types.md#4-dim--symbolic-shape-dimensions). -""" - -from __future__ import annotations - -from .dim import DimVar - - -def static_dim_value(dim): - """Return the compile-time ``int`` value of a *static* shape dim, else ``None``. - - A static dim is a plain ``int`` or an integer-valued ``Constant`` (the latter - only appears transiently before ``TensorType`` canonicalizes it to ``int``). - ``DimVar`` / dynamic dim ``Call`` exprs are not static → ``None``. The - detection is exact (real ``Constant`` with an ``int`` value), never "anything - with a ``.value``"; the ``Constant`` import is deferred to dodge the - ``ir.core.expr`` ↔ ``ir.types`` cycle and fails closed (returns ``None``). - """ - if isinstance(dim, int) and not isinstance(dim, bool): - return dim - try: - from tilefoundry.ir.core.expr import Constant # noqa: PLC0415 - cycle guard - except ImportError: # pragma: no cover - import-cycle guard - return None - if isinstance(dim, Constant) and isinstance(dim.value, int) and not isinstance(dim.value, bool): - return int(dim.value) - return None - - -def i64_const(value: int) -> "Constant": - """The canonical i64 shape-scalar ``Constant`` (meta-scalar typed).""" - from tilefoundry.ir.core.expr import Constant # noqa: PLC0415 - cycle guard - - from .tensor_type import TensorType # noqa: PLC0415 - cycle guard - - return Constant(type=TensorType.umat_scalar(), value=int(value)) - - -def upper_bound(dim) -> int: - """Return a concrete int upper-bound element count for ``dim``.""" - if isinstance(dim, DimVar): - return int(dim.hi) - 1 - static = static_dim_value(dim) - if static is not None: - return static - return int(dim) - - -def shape_numel_upper_bound(shape) -> int: - """Product of per-dim upper bounds. - - Product of per-dim upper bounds: the static element count a buffer or - layout must hold across every runtime shape in the dispatch envelope. - """ - n = 1 - for s in shape: - n *= upper_bound(s) - return n - - -def shape_upper_bound(shape) -> tuple[int, ...]: - """Map ``upper_bound`` over every entry of *shape*.""" - return tuple(upper_bound(s) for s in shape) - - -def shape_has_dim_var(shape) -> bool: - """True iff *shape* contains at least one ``DimVar`` entry.""" - return any(isinstance(s, DimVar) for s in shape) - - -def shape_runtime_total(shape, dim_var_expr: dict[str, str]) -> object: - """Return the runtime element count of *shape*. - - All-static shape → an ``int``. Any ``DimVar`` axis pulls its - runtime extent from ``dim_var_expr[name]``; the result is a C++ - expression string ``"(a * b * ...)"`` that the codegen splices - verbatim into the generated source. Static dims fold into a single - leading constant factor when present, otherwise the constant is - elided. - """ - if not shape: - return 1 - static_prod = 1 - dyn_terms: list[str] = [] - for s in shape: - if isinstance(s, DimVar): - expr = dim_var_expr.get(s.name) - if expr is None: - static_prod *= upper_bound(s) - else: - dyn_terms.append(expr) - else: - static_prod *= upper_bound(s) - if not dyn_terms: - return static_prod - if static_prod == 1: - if len(dyn_terms) == 1: - return dyn_terms[0] - return "(" + " * ".join(dyn_terms) + ")" - return "(" + " * ".join([str(static_prod), *dyn_terms]) + ")" - - -__all__ = [ - "static_dim_value", - "i64_const", - "upper_bound", - "shape_numel_upper_bound", - "shape_upper_bound", - "shape_has_dim_var", - "shape_runtime_total", -] diff --git a/src/tilefoundry/ir/types/shard/__init__.py b/src/tilefoundry/ir/types/shard/__init__.py deleted file mode 100644 index 9ed5d14a..00000000 --- a/src/tilefoundry/ir/types/shard/__init__.py +++ /dev/null @@ -1,75 +0,0 @@ -from __future__ import annotations - -# ruff: noqa: I001 -- curated re-export order; alphabetical sort breaks staged imports. - -from .int_tuple import IntTuple, flatten, product -from .layout import ComposedLayout, Layout, LayoutBase, Swizzle -from .layout_algebra import ( - c_order_strides, - composition, - prefix_product, - swizzle_of, - try_c_order_strides, -) -from .mesh import ( - Mesh, - Topology, - check_topology, - composed, - positions_at, - topology_axes, - topology_projection, -) -from .placement import Placement -from .scope_match import covered_by_scope, storage_reaches -from .shard_layout import ( - B, - Broadcast, - Dynamic, - P, - Partial, - S, - ShardAttr, - ShardLayout, - Split, - canonical_shard_layout, - shard_layout_of, -) -from .utils import make_mesh - -__all__ = [ - "IntTuple", - "flatten", - "product", - "c_order_strides", - "try_c_order_strides", - "prefix_product", - "LayoutBase", - "Layout", - "Swizzle", - "ComposedLayout", - "composition", - "swizzle_of", - "Topology", - "check_topology", - "composed", - "positions_at", - "topology_axes", - "topology_projection", - "Placement", - "covered_by_scope", - "storage_reaches", - "Mesh", - "make_mesh", - "ShardAttr", - "Split", - "Partial", - "Broadcast", - "Dynamic", - "ShardLayout", - "S", - "P", - "B", - "canonical_shard_layout", - "shard_layout_of", -] diff --git a/src/tilefoundry/ir/types/shard/int_tuple.py b/src/tilefoundry/ir/types/shard/int_tuple.py deleted file mode 100644 index 05b726f1..00000000 --- a/src/tilefoundry/ir/types/shard/int_tuple.py +++ /dev/null @@ -1,37 +0,0 @@ -"""IntTuple alias + helpers.""" - -from __future__ import annotations - -from typing import Union, overload - -from tilefoundry.ir.types.tensor_type import ShapeDim - -IntTuple = Union[int, tuple["IntTuple", ...]] - - -@overload -def flatten(t: IntTuple) -> tuple[int, ...]: ... - - -@overload -def flatten(t: object) -> tuple[object, ...]: ... - - -def flatten(t: object) -> tuple[object, ...]: - if not isinstance(t, tuple): - return (t,) - return tuple(value for item in t for value in flatten(item)) - - -def product(t) -> "ShapeDim": - from .mesh import Topology # noqa: PLC0415 - - result = 1 - for v in flatten(t): - if isinstance(v, Topology): - v = v.size - result *= v - return result - - -__all__ = ["IntTuple", "flatten", "product"] diff --git a/src/tilefoundry/ir/types/shard/local.py b/src/tilefoundry/ir/types/shard/local.py deleted file mode 100644 index 790281d1..00000000 --- a/src/tilefoundry/ir/types/shard/local.py +++ /dev/null @@ -1,171 +0,0 @@ -"""Which part of a sharded tensor one mesh instance holds. - -The same two answers under the same two names as the device side -(layout/shard_layout.cuh): ``local_layout`` is what an instance holds and -``local_layout_and_offset`` adds where its part begins. The host states the -tensor shape as well, because ``canonical_shard_layout`` factors a split axis -away from the tensor's own rank, and the ids as well, because a level the -placement leaves unfixed is the device's to divide rather than this side's. - -See [shard §7.6](docs/spec/shard.md#76-local_layout). -""" - -from __future__ import annotations - -from collections.abc import Iterable - -from .int_tuple import flatten -from .layout import Layout -from .layout_algebra import c_order_strides, idx2crd -from .mesh import positions_at, topology_axes -from .shard_layout import ShardLayout, layout_axis_to_tensor_axis, split_target_axes - - -def _extents(shard: ShardLayout) -> tuple[int, ...]: - """Each mesh axis's extent, flat and in the order the attrs index them.""" - values = tuple(flatten(shard.mesh.layout.shape)) - for value in values: - if not isinstance(value, int) or isinstance(value, bool): - raise ValueError( - f"local: this mesh states a symbolic extent {value!r}; which part is " - f"one instance's is a number, so bind the mesh first" - ) - return values - - -def _cutting(shard: ShardLayout, tensor_shape: tuple) -> tuple[int | None, ...]: - """Per mesh axis, the tensor axis its attr cuts, or ``None`` when none.""" - return split_target_axes(shard, tuple(tensor_shape)) - - -def _strides(shard: ShardLayout, tensor_shape: tuple) -> tuple[int, ...]: - """Each tensor axis's own stride, read off the factored layout. - - A split axis is stated as mesh-sized positions plus a residual, so the - step of the tensor axis itself is the last of them; a layout that - materialized no strides is the C order its canonical form would have. - """ - fallback = c_order_strides(tuple(tensor_shape)) - stated = shard.layout.strides - if stated is None: - return fallback - axes = layout_axis_to_tensor_axis(shard.layout.shape, tuple(tensor_shape)) - found = dict(zip(axes, stated, strict=False)) - return tuple(found.get(axis, fallback[axis]) for axis in range(len(tensor_shape))) - - -def _narrowed( - shard: ShardLayout, tensor_shape: tuple, axis: int, dividing: Iterable[int] -) -> int: - """Tensor *axis* divided by each mesh axis in *dividing* that cuts it. - - Several may. Two axes of one level cut a tensor axis into a grid, and so - do two levels, one taking a block of what the other left; each division is - of what the previous one left. An extent its mesh axis does not divide is - refused by name: a shard would then not be one slice. - """ - extents = _extents(shard) - allowed = set(dividing) - held = tensor_shape[axis] - for mesh_axis, target in enumerate(_cutting(shard, tensor_shape)): - if target != axis or mesh_axis >= len(extents) or mesh_axis not in allowed: - continue - count = extents[mesh_axis] - if held % count: - raise ValueError( - f"local: axis {axis} has extent {held}, which its mesh axis of " - f"{count} positions does not divide; a shard would not be one slice" - ) - held //= count - return held - - -def _inner(shard: ShardLayout, tensor_shape: tuple, mesh_axis: int, axis: int) -> int: - """How much of tensor *axis* one step of *mesh_axis* steps over. - - The axes cutting one tensor axis are ordered outermost first, so a step of - one of them clears everything the axes inside it hold. - """ - extents = _extents(shard) - product = 1 - for other, target in enumerate(_cutting(shard, tensor_shape)): - if other > mesh_axis and target == axis and other < len(extents): - product *= extents[other] - return product - - -def _stride(shard: ShardLayout, tensor_shape: tuple, mesh_axis: int, axis: int) -> int: - """What one step of *mesh_axis* costs the slice's origin, in elements. - - As the device's ``stride(sl)`` has it: the extent one instance is left - of the tensor axis, times what the axes inside this one hold of it, times - the tensor axis's own stride. Every axis cutting it divides here, unfixed - or not, because a step must clear what the device will divide too. - """ - return ( - _narrowed(shard, tensor_shape, axis, range(len(_extents(shard)))) - * _inner(shard, tensor_shape, mesh_axis, axis) - * _strides(shard, tensor_shape)[axis] - ) - - -def _positions(shard: ShardLayout, ids: tuple[int | None, ...]) -> dict[int, int]: - """Each mesh axis's coordinate, for the levels an id was given for. - - A level with no id is left unfixed and names no coordinate, so the axes it - owns divide nothing and the whole of them stays. - """ - names = tuple(topology.name for topology in shard.mesh.topologies) - found: dict[int, int] = {} - for index, mesh_axes in enumerate(topology_axes(shard.mesh)): - program_id = ids[index] if index < len(ids) else None - if program_id is None: - continue - coord = idx2crd(program_id, *positions_at(shard.mesh, names[index])) - for mesh_axis, position in zip(mesh_axes, coord, strict=True): - found[mesh_axis] = position - return found - - -def local_layout( - shard: ShardLayout, tensor_shape: tuple, ids: tuple[int | None, ...] -) -> Layout: - """What one instance holds: each tensor axis narrowed by what cuts it. - - Narrowed by the axes *ids* names, and no others: an axis nothing cuts is - held whole, and so is one cut only by a level the placement left unfixed -- - ``cta`` and ``thread`` are the device's to divide. The strides stay the - whole tensor's, as they do on the device: a slice of it is the same rows, - the same distance apart. - """ - dividing = _positions(shard, ids).keys() - return Layout( - shape=tuple( - _narrowed(shard, tensor_shape, axis, dividing) - for axis in range(len(tensor_shape)) - ), - strides=_strides(shard, tensor_shape), - ) - - -def local_layout_and_offset( - shard: ShardLayout, tensor_shape: tuple, ids: tuple[int | None, ...] -) -> tuple[Layout, int]: - """That layout, and how far into the tensor this instance's part begins. - - What ``cute::slice_and_offset`` is to a Layout, as on the device: the - offset is the instance's mesh coordinate dotted with one stride per mesh - axis. It is counted in the strides ``local_layout`` reports, so a reader - laid out otherwise would read the wrong elements. - """ - cutting = _cutting(shard, tensor_shape) - offset = 0 - for mesh_axis, position in _positions(shard, ids).items(): - axis = cutting[mesh_axis] if mesh_axis < len(cutting) else None - if axis is None: - continue - offset += position * _stride(shard, tensor_shape, mesh_axis, axis) - return local_layout(shard, tensor_shape, ids), offset - - -__all__ = ["local_layout", "local_layout_and_offset"] diff --git a/src/tilefoundry/ir/types/shard/mesh.py b/src/tilefoundry/ir/types/shard/mesh.py deleted file mode 100644 index 6c6e5071..00000000 --- a/src/tilefoundry/ir/types/shard/mesh.py +++ /dev/null @@ -1,364 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from functools import lru_cache - -from tilefoundry.ir.types.shard.int_tuple import flatten -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout -from tilefoundry.ir.types.shard.layout_algebra import c_order_strides, unflatten -from tilefoundry.ir.types.tensor_type import ShapeDim - - -@dataclass(frozen=True) -class Topology: - """Name one hardware level and its explicit static or symbolic size.""" - - name: str - - size: "ShapeDim" - - def __post_init__(self) -> None: - if self.size is None: - raise ValueError( - f"Topology {self.name!r}: extent must be explicit; None is not " - "a ShapeDim. The rule: " - "tilefoundry spec target topology-levels" - ) - - -@dataclass(frozen=True) -class Mesh: - """Describe hardware levels, logical positions, and axis names. - - A constant slice replaces ``layout`` with a ``ComposedLayout`` whose - ``offset`` and ``outer`` describe the selected sub-box. It remains a - compile-time descriptor outside the IR/SSA graph. - - See [shard §5](docs/spec/shard.md#5-mesh). - """ - - topologies: tuple[Topology | str, ...] - layout: "Layout | ComposedLayout" - names: tuple[str, ...] = () - - def __post_init__(self) -> None: - layout = self.layout - if isinstance(layout, tuple): - layout = Layout(shape=layout, strides=c_order_strides(layout)) - object.__setattr__(self, "layout", layout) - - for axis, extent in enumerate(flatten(self.layout.shape)): - if extent is None: - raise ValueError( - f"Mesh: layout axis {axis} must have an explicit extent; " - "None is not a ShapeDim. The rule: tilefoundry spec shard mesh" - ) - - def __getitem__(self, key) -> "Mesh": - """Return a constant sub-mesh selected by integers or unit-step slices. - - Missing axes are full slices; integers select extent one. The result - preserves topology and names while recording the sub-box as a - ``ComposedLayout``. Dynamic layouts and nested slices are rejected. - - See [shard §5](docs/spec/shard.md#5-mesh). - """ - if isinstance(self.layout, ComposedLayout): - raise ValueError("cannot slice an already-sliced mesh (nested slice unsupported)") - shape = self.layout.shape - strides = self.layout.strides - rank = len(shape) - 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] = [] - 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): - raise ValueError( - f"mesh slice {k.start}:{k.stop} out of range for axis {axis} (extent {extent})" - ) - 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) - - sliced = ComposedLayout( - inner=None, - offset=offset, - outer=Layout(shape=tuple(sub_shape), strides=strides), - ) - return Mesh( - topologies=self.topologies, - layout=sliced, - names=self.names, - ) - - -def topology_axes(mesh: "Mesh") -> tuple[tuple[int, ...], ...]: - """Which of *mesh*'s layout axes belong to each level it names, in order. - - The axes are handed to the levels left to right, and a level takes them - until their extents multiply to exactly its own size. A boundary that no - prefix of axes lands on is refused rather than guessed: one axis of four - positions across a two-CTA and two-thread boundary could be either half of - it, and picking one would place work somewhere nobody said. - - A mesh naming one level owns all of them, which is the answer this gives - without any of the arithmetic. - """ - if len(mesh.topologies) == 1: - return (tuple(range(len(flatten(mesh.layout.shape)))),) - extents = flatten(mesh.layout.shape) - found: list[tuple[int, ...]] = [] - axis = 0 - for topology in mesh.topologies: - size = topology.size - if not isinstance(size, int) or isinstance(size, bool) or size < 1: - raise ValueError( - f"mesh level {topology.name!r} states extent {size!r}; segmenting a " - "mesh that names several levels needs each of their position counts" - ) - taken: list[int] = [] - product = 1 - while product < size and axis < len(extents): - extent = extents[axis] - if not isinstance(extent, int) or isinstance(extent, bool): - raise ValueError( - f"mesh layout axis {axis} states extent {extent!r}; segmenting a " - "mesh that names several levels needs concrete axis extents" - ) - product *= extent - taken.append(axis) - axis += 1 - if product != size: - raise ValueError( - f"mesh axes {tuple(extents)} do not land on the boundary of level " - f"{topology.name!r} at {size}: the axes up to there multiply to " - f"{product}. Write the axis that straddles it as the two axes it is" - ) - found.append(tuple(taken)) - if axis != len(extents): - raise ValueError( - f"mesh layout has {len(extents)} axes but the levels it names account " - f"for {axis}; every axis belongs to one of them" - ) - return tuple(found) - - -def _positions_layout(mesh: Mesh) -> tuple[tuple, tuple, int]: - """Return flattened shape, strides, and offset for a supported mesh layout.""" - if isinstance(mesh.layout, Layout): - return flatten(mesh.layout.shape), flatten(mesh.layout.strides), 0 - if mesh.layout.inner is None and isinstance(mesh.layout.outer, Layout): - return ( - flatten(mesh.layout.outer.shape), - flatten(mesh.layout.outer.strides), - mesh.layout.offset, - ) - raise ValueError(f"mesh levels {mesh.topologies!r} have an unsupported layout") - - -@lru_cache(maxsize=None) -def grouped_layout(mesh: Mesh) -> Layout: - """*mesh*'s positions with one mode per level it names, in mesh numbering. - - The grouping only says which axes are whose; the strides still count - positions of the whole mesh. What one level alone would call them is that - mode divided by :func:`positions_below`. - """ - shape, strides, _offset = _positions_layout(mesh) - if any(stride is None for stride in strides): - strides = c_order_strides(shape) - profile = topology_axes(mesh) - return Layout(shape=unflatten(shape, profile), strides=unflatten(strides, profile)) - - -def positions_below(mesh: Mesh, index: int) -> int: - """How many positions the levels under the one at *index* contribute.""" - below = 1 - for topology in mesh.topologies[index + 1 :]: - if not isinstance(topology.size, int) or isinstance(topology.size, bool): - raise ValueError( - f"mesh level {mesh.topologies[index].name!r} has a symbolic level below it" - ) - below *= topology.size - return below - - -def _index_of(mesh: Mesh, topology_level: str) -> int: - names = tuple(topology.name for topology in mesh.topologies) - if topology_level not in names: - raise ValueError(f"mesh names levels {names}, not {topology_level!r}") - return names.index(topology_level) - - -def _divided(strides: tuple, below: int, topology_level: str) -> tuple[int, ...]: - """*strides* read as the level's own, or a refusal that they are not.""" - divided: list[int] = [] - for axis, stride in enumerate(strides): - if not isinstance(stride, int) or isinstance(stride, bool) or stride % below: - raise ValueError( - f"mesh axis {axis} has stride {stride!r}, which the {below} positions " - f"below {topology_level!r} do not divide; its positions are not that level's" - ) - divided.append(stride // below) - return tuple(divided) - - -def positions_at(mesh: Mesh, topology_level: str) -> tuple[tuple, tuple]: - """Return one named level's shape and normalized strides, every axis of it. - - An axis of one position is still that level's axis: dropping it here would - leave the level's own layout narrower than the attrs written against those - axes, and nothing downstream could say which attr went with which mode. - """ - index = _index_of(mesh, topology_level) - grouped = grouped_layout(mesh) - below = positions_below(mesh, index) - shape = flatten(grouped.shape[index]) - return shape, _divided(flatten(grouped.strides[index]), below, topology_level) - - -def composed(meshes: "tuple[Mesh, ...]") -> "Mesh": - """Compose scopes, replacing an existing level when the inner names it.""" - if len(meshes) == 1: - check_topology(meshes[0]) - return meshes[0] - - def positions(mesh: Mesh) -> int: - count = 1 - for topology in mesh.topologies: - size = topology.size - if not isinstance(size, int) or isinstance(size, bool) or size < 1: - raise ValueError( - f"mesh level {topology.name!r} states extent {size!r}; " - "composing scopes needs each level's position count" - ) - count *= size - return count - - def concatenate(outer: Mesh, inner: Mesh) -> Mesh: - outer_shape, outer_strides, outer_offset = _positions_layout(outer) - inner_shape, inner_strides, inner_offset = _positions_layout(inner) - for mesh, strides in ((outer, outer_strides), (inner, inner_strides)): - if any(not isinstance(stride, int) or isinstance(stride, bool) for stride in strides): - raise ValueError( - f"mesh levels {mesh.topologies!r} need concrete strides to compose" - ) - below = positions(inner) - layout = Layout( - shape=(*outer_shape, *inner_shape), - strides=(*(stride * below for stride in outer_strides), *inner_strides), - ) - offset = outer_offset * below + inner_offset - sliced = isinstance(outer.layout, ComposedLayout) or isinstance( - inner.layout, ComposedLayout - ) - return Mesh( - topologies=(*outer.topologies, *inner.topologies), - layout=( - layout if not sliced else ComposedLayout(inner=None, offset=offset, outer=layout) - ), - names=(*outer.names, *inner.names), - ) - - result = meshes[0] - for inner in meshes[1:]: - current_names = {topology.name for topology in result.topologies} - inner_names = {topology.name for topology in inner.topologies} - if current_names.isdisjoint(inner_names): - result = concatenate(result, inner) - continue - if current_names <= inner_names: - result = inner - continue - raise ValueError( - f"{sorted(current_names & inner_names)} named again while " - f"{sorted(current_names - inner_names)} is not; a scope either " - "replaces the levels in force or adds levels below them" - ) - check_topology(result) - return result - - -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 - shape, _strides, _offset = _positions_layout(mesh) - for topology, axes in zip(mesh.topologies, topology_axes(mesh)): - if not isinstance(topology.size, int) or isinstance(topology.size, bool): - continue - count = 1 - for axis in axes: - extent = shape[axis] - if not isinstance(extent, int) or isinstance(extent, bool): - count = None - break - count *= extent - if count is not None and count > topology.size: - raise ValueError( - f"mesh level {topology.name!r} has {count} positions, exceeding declared extent {topology.size}" - ) - - -def topology_projection(mesh: "Mesh", topology_level: str) -> Layout: - """The layout of the positions *level* has, out of a mesh that names more. - - A position at one level is a position within its parent, so the projection - keeps every axis up to and including that level's own and divides their - strides by what the deeper levels contribute. Asking a single-level mesh - returns its own layout untouched. - """ - index = _index_of(mesh, topology_level) - if len(mesh.topologies) == 1: - if isinstance(mesh.layout, Layout): - return mesh.layout - raise ValueError("a sliced mesh states its own layout; it is not projected") - if not isinstance(mesh.layout, Layout): - raise ValueError( - "a mesh naming several levels cannot also be sliced; the slice and the " - "level boundary would both be deciding which positions these are" - ) - grouped = grouped_layout(mesh) - below = positions_below(mesh, index) - shape = flatten(grouped.shape[: index + 1]) - strides = _divided(flatten(grouped.strides[: index + 1]), below, topology_level) - return Layout(shape=shape, strides=strides) - - -__all__ = [ - "Mesh", - "Topology", - "composed", - "grouped_layout", - "positions_at", - "positions_below", - "topology_axes", - "topology_projection", - "check_topology", -] diff --git a/src/tilefoundry/ir/types/shard/scope_match.py b/src/tilefoundry/ir/types/shard/scope_match.py deleted file mode 100644 index 130c06f3..00000000 --- a/src/tilefoundry/ir/types/shard/scope_match.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Match TIR scope declarations against required hardware participation. - -The TIR verifier keeps its own tuple of enclosing mesh values. CUDA MMA uses -these structural predicates to decide whether an atom's required thread scope -is hosted by one of them; HIR execution-domain visibility is checked separately -by ``covered_by_scope`` and ``storage_reaches``. -""" - -from __future__ import annotations - -from ..storage import StorageKind, resolve_storage -from .int_tuple import product -from .layout import Layout -from .layout_algebra import is_inverse_projectable, size -from .mesh import Mesh, positions_at - - -def _as_layout(mesh: Mesh) -> Layout: - return Layout(shape=tuple(mesh.layout.shape), strides=tuple(mesh.layout.strides)) - - -def states_consistent_positions(mesh: Mesh) -> bool: - return product(mesh.topologies) == size(mesh.layout) - - -def mesh_scope_matches_required_scope(current: Mesh, required: Mesh) -> bool: - """True iff ``current`` provides the thread participation ``required`` needs.""" - if current.topologies[0].name != required.topologies[0].name: - return False - - cur_layout = _as_layout(current) - req_layout = _as_layout(required) - - if not states_consistent_positions(current) or not states_consistent_positions(required): - return False - - if not is_inverse_projectable(cur_layout): - return False - - return cur_layout.shape == req_layout.shape and cur_layout.strides == req_layout.strides - - -def _positions(mesh: Mesh, topology_level: str) -> tuple[tuple, tuple]: - """One level's positions, with the axes of one position left out. - - An axis of one position names no instance, so two scopes state the same - positions whether or not either of them wrote such an axis down. - """ - shape, strides = positions_at(mesh, topology_level) - kept = tuple(axis for axis, extent in enumerate(shape) if extent != 1) - return tuple(shape[axis] for axis in kept), tuple(strides[axis] for axis in kept) - - -def covered_by_scope(mesh: Mesh, current: Mesh) -> bool: - """Whether *mesh* names no finer positions than the enclosing scope.""" - scope = {topology.name: _positions(current, topology.name) for topology in current.topologies} - return all( - topology.name in scope and _positions(mesh, topology.name) == scope[topology.name] - for topology in mesh.topologies - ) - - -def storage_reaches(storage, mesh: Mesh, current: Mesh) -> bool: - """Whether *storage* reaches across a coarser value-to-scope boundary.""" - if current.topologies[-1].name in {topology.name for topology in mesh.topologies}: - return True - try: - storage = resolve_storage(storage) - except (TypeError, ValueError): - return False - return storage in {StorageKind.GMEM, StorageKind.SMEM} - - -__all__ = [ - "covered_by_scope", - "mesh_scope_matches_required_scope", - "states_consistent_positions", - "storage_reaches", -] diff --git a/src/tilefoundry/ir/types/shard/utils.py b/src/tilefoundry/ir/types/shard/utils.py deleted file mode 100644 index 0de6b9cc..00000000 --- a/src/tilefoundry/ir/types/shard/utils.py +++ /dev/null @@ -1,35 +0,0 @@ -from __future__ import annotations - -import math - -from .layout import Layout -from .layout_algebra import try_c_order_strides -from .mesh import Mesh, Topology - - -def make_mesh( - layout_shape: tuple, - names: "tuple[str, ...] | None" = None, - topology: "str | Topology" = "gpu", -) -> Mesh: - """Convenience constructor for a ``Mesh`` with the given axis extents and C-order strides. - - 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``. - - ``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_c_order_strides(layout_shape)), - names=tuple(names), - ) diff --git a/src/tilefoundry/ir/types/shard/shard_layout.py b/src/tilefoundry/ir/types/shard_layout.py similarity index 53% rename from src/tilefoundry/ir/types/shard/shard_layout.py rename to src/tilefoundry/ir/types/shard_layout.py index 67d01fb9..ba7a8fb1 100644 --- a/src/tilefoundry/ir/types/shard/shard_layout.py +++ b/src/tilefoundry/ir/types/shard_layout.py @@ -2,9 +2,13 @@ from dataclasses import dataclass -from .layout import ComposedLayout, Layout, LayoutBase -from .layout_algebra import try_c_order_strides +from tilefoundry.ir.types.int_tuple import repeat_like +from tilefoundry.ir.types.layout import flatten, rank +from tilefoundry.ir.types.stride import idx2crd + +from .layout import ComposedLayout, Layout, LayoutBase, get from .mesh import Mesh +from .stride import compact_row_major, try_compact_major class ShardAttr: @@ -78,7 +82,7 @@ def canonical_shard_layout(logical_shape: tuple, mesh: Mesh, attrs: tuple) -> "S See [shard §7.1.1](docs/spec/shard.md#711-layoutshape). """ - mesh_shape = mesh.layout.shape + mesh_shape = flatten(mesh.layout).shape bindings: dict[int, list[int]] = {} for mesh_axis, attr in enumerate(attrs): if isinstance(attr, Split): @@ -131,7 +135,7 @@ def canonical_shard_layout(logical_shape: tuple, mesh: Mesh, attrs: tuple) -> "S ) layout_shape = tuple(layout_shape) return ShardLayout( - layout=Layout(shape=layout_shape, strides=try_c_order_strides(layout_shape)), + layout=Layout(shape=layout_shape, strides=try_compact_major(layout_shape)), attrs=remapped_attrs, mesh=mesh, ) @@ -151,7 +155,7 @@ def shard_layout_local_shape( See [shard §7](docs/spec/shard.md#7-shardlayout). """ - mesh_shape = sl.mesh.layout.shape + mesh_shape = flatten(sl.mesh.layout).shape local = list(sl.layout.shape) for mesh_axis_idx, attr in enumerate(sl.attrs): if mesh_axis_idx >= len(mesh_shape): @@ -194,7 +198,7 @@ def layout_axis_to_tensor_axis(layout_shape: tuple, tensor_shape: tuple) -> list See [shard §7.1.1](docs/spec/shard.md#711-layoutshape). """ - from ..shape_helpers import static_dim_value # noqa: PLC0415 - cycle guard + from .utils import static_dim_value # noqa: PLC0415 - cycle guard result: list[int] = [] layout_idx = 0 @@ -236,7 +240,169 @@ def split_target_axes(sl: "ShardLayout", tensor_shape: tuple) -> tuple: return tuple(la2ta[a.axis] if isinstance(a, Split) else None for a in sl.attrs) +def _extents(shard: ShardLayout) -> tuple[int, ...]: + """Each mesh axis's extent, flat and in the order the attrs index them.""" + values = tuple(flatten(flatten(shard.mesh.layout).shape)) + for value in values: + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError( + f"local: this mesh states a symbolic extent {value!r}; which part is " + f"one instance's is a number, so bind the mesh first" + ) + return values + + +def _cutting(shard: ShardLayout, tensor_shape: tuple) -> tuple[int | None, ...]: + """Per mesh axis, the tensor axis its attr cuts, or ``None`` when none.""" + return split_target_axes(shard, tuple(tensor_shape)) + + +def _strides(shard: ShardLayout, tensor_shape: tuple) -> tuple[int, ...]: + """Each tensor axis's own stride, read off the factored layout. + + A split axis is stated as mesh-sized positions plus a residual, so the + step of the tensor axis itself is the last of them; a layout that + materialized no strides is the C order its canonical form would have. + """ + fallback = compact_row_major(tuple(tensor_shape)) + stated = shard.layout.strides + if stated is None: + return fallback + axes = layout_axis_to_tensor_axis(shard.layout.shape, tuple(tensor_shape)) + found = dict(zip(axes, stated, strict=False)) + return tuple(found.get(axis, fallback[axis]) for axis in range(len(tensor_shape))) + + +def _narrowed( + shard: ShardLayout, tensor_shape: tuple, axis: int, dividing: Iterable[int] +) -> int: + """Tensor *axis* divided by each mesh axis in *dividing* that cuts it. + + Several may. Two axes of one level cut a tensor axis into a grid, and so + do two levels, one taking a block of what the other left; each division is + of what the previous one left. An extent its mesh axis does not divide is + refused by name: a shard would then not be one slice. + """ + extents = _extents(shard) + allowed = set(dividing) + held = tensor_shape[axis] + for mesh_axis, target in enumerate(_cutting(shard, tensor_shape)): + if target != axis or mesh_axis >= len(extents) or mesh_axis not in allowed: + continue + count = extents[mesh_axis] + if held % count: + raise ValueError( + f"local: axis {axis} has extent {held}, which its mesh axis of " + f"{count} positions does not divide; a shard would not be one slice" + ) + held //= count + return held + + +def _inner(shard: ShardLayout, tensor_shape: tuple, mesh_axis: int, axis: int) -> int: + """How much of tensor *axis* one step of *mesh_axis* steps over. + + The axes cutting one tensor axis are ordered outermost first, so a step of + one of them clears everything the axes inside it hold. + """ + extents = _extents(shard) + product = 1 + for other, target in enumerate(_cutting(shard, tensor_shape)): + if other > mesh_axis and target == axis and other < len(extents): + product *= extents[other] + return product + + +def _stride(shard: ShardLayout, tensor_shape: tuple, mesh_axis: int, axis: int) -> int: + """What one step of *mesh_axis* costs the slice's origin, in elements. + + As the device's ``stride(sl)`` has it: the extent one instance is left + of the tensor axis, times what the axes inside this one hold of it, times + the tensor axis's own stride. Every axis cutting it divides here, unfixed + or not, because a step must clear what the device will divide too. + """ + return ( + _narrowed(shard, tensor_shape, axis, range(len(_extents(shard)))) + * _inner(shard, tensor_shape, mesh_axis, axis) + * _strides(shard, tensor_shape)[axis] + ) + + +def _positions(shard: ShardLayout, ids: tuple[int | None, ...]) -> dict[int, int]: + """Each mesh axis's coordinate, for the levels an id was given for. + + A level with no id is left unfixed and names no coordinate, so the axes it + owns divide nothing and the whole of them stays. Each level reads its own + id into a coordinate shaped like its own modes, so flattening what the + levels read together gives one entry per mesh axis. + """ + mesh_layout = shard.mesh.layout + stated = mesh_layout.outer if isinstance(mesh_layout, ComposedLayout) else mesh_layout + read: list = [] + for index in range(rank(stated)): + arrangement = get(stated, index) + program_id = ids[index] if index < len(ids) else None + if program_id is None: + read.append(repeat_like(arrangement.shape, None)) + continue + read.append( + idx2crd( + program_id, + tuple(flatten(arrangement.shape)), + tuple(flatten(arrangement.strides)), + ) + ) + return { + axis: position + for axis, position in enumerate(flatten(tuple(read))) + if position is not None + } + + +def local_layout( + shard: ShardLayout, tensor_shape: tuple, ids: tuple[int | None, ...] +) -> Layout: + """What one instance holds: each tensor axis narrowed by what cuts it. + + Narrowed by the axes *ids* names, and no others: an axis nothing cuts is + held whole, and so is one cut only by a level the placement left unfixed -- + ``cta`` and ``thread`` are the device's to divide. The strides stay the + whole tensor's, as they do on the device: a slice of it is the same rows, + the same distance apart. + """ + dividing = _positions(shard, ids).keys() + return Layout( + shape=tuple( + _narrowed(shard, tensor_shape, axis, dividing) + for axis in range(len(tensor_shape)) + ), + strides=_strides(shard, tensor_shape), + ) + + +def local_layout_and_offset( + shard: ShardLayout, tensor_shape: tuple, ids: tuple[int | None, ...] +) -> tuple[Layout, int]: + """That layout, and how far into the tensor this instance's part begins. + + What ``cute::slice_and_offset`` is to a Layout, as on the device: the + offset is the instance's mesh coordinate dotted with one stride per mesh + axis. It is counted in the strides ``local_layout`` reports, so a reader + laid out otherwise would read the wrong elements. + """ + cutting = _cutting(shard, tensor_shape) + offset = 0 + for mesh_axis, position in _positions(shard, ids).items(): + axis = cutting[mesh_axis] if mesh_axis < len(cutting) else None + if axis is None: + continue + offset += position * _stride(shard, tensor_shape, mesh_axis, axis) + return local_layout(shard, tensor_shape, ids), offset + + __all__ = [ + "local_layout", + "local_layout_and_offset", "ShardAttr", "Split", "Partial", diff --git a/src/tilefoundry/ir/types/stride.py b/src/tilefoundry/ir/types/stride.py new file mode 100644 index 00000000..cc54ba1e --- /dev/null +++ b/src/tilefoundry/ir/types/stride.py @@ -0,0 +1,100 @@ +"""Strides: the steps a compact arrangement walks, and reading one back. + +CuTe keeps this apart from the layout algebra (``stride.hpp``): making the +steps of a compact arrangement, and turning an index back into the coordinate +that reaches it, are not operations on layouts. +""" + +from __future__ import annotations + + +def compact_major(shape: tuple, *, major: str = "row", mul=None, current=1) -> tuple: + """CuTe ``compact_major``: the steps a compact arrangement of *shape* walks. + + ``"col"`` makes mode zero the fastest and ``"row"`` the last one, which is + the one difference between CuTe's two spellings of it. A mode that is + itself a group of them is walked the same way inside what the modes beside + it leave it, so the steps come back shaped like the shape they were made + for. *mul* defaults to integer multiplication; pass a dim-expression fold + for a shape whose entries are symbolic. + """ + if not isinstance(shape, tuple): + return current + if not shape: + return () + if mul is None: + mul = lambda a, b: a * b # noqa: E731 + order = range(len(shape)) if major == "col" else range(len(shape) - 1, -1, -1) + strides: list = [1] * len(shape) + acc = current + for index in order: + strides[index] = compact_major(shape[index], major=major, mul=mul, current=acc) + acc = mul(acc, _product(shape[index], mul)) + return tuple(strides) + + +def _product(shape, mul): + if not isinstance(shape, tuple): + return shape + total = 1 + for one in shape: + total = mul(total, _product(one, mul)) + return total + + +def compact_col_major(shape: tuple, *, mul=None) -> tuple: + """CuTe ``compact_col_major``: mode zero walks fastest.""" + return compact_major(shape, major="col", mul=mul) + + +def compact_row_major(shape: tuple, *, mul=None) -> tuple: + """CuTe ``compact_row_major``: the last mode walks fastest.""" + return compact_major(shape, major="row", mul=mul) + + +def try_compact_major(shape: tuple, *, major: str = "row") -> "tuple | None": + """:func:`compact_major`, or ``None`` where an extent is not a static int.""" + if not all( + isinstance(one, int) and not isinstance(one, bool) + for one in _flat(shape) + ): + return None + return compact_major(shape, major=major) + + +def _flat(shape): + if not isinstance(shape, tuple): + return (shape,) + return tuple(one for item in shape for one in _flat(item)) + + +def idx2crd(idx: int, shape: tuple, stride: tuple) -> tuple: + """CuTe ``idx2crd``: the coordinate in ```` an index reaches. + + Per mode ``(idx // stride) % shape``, recursing wherever a mode is itself a + group of them, so what comes back is shaped like the arrangement it was + read against. Which mode walks fastest is the strides' to say, not this + function's. + """ + return tuple( + idx2crd(idx, one, step) if isinstance(one, tuple) else (idx // step) % one + for one, step in zip(shape, stride) + ) + + +def crd2idx(crd: tuple, shape: tuple, stride: tuple) -> int: + """CuTe ``crd2idx``: the index a coordinate in ```` reaches.""" + total = 0 + for value, one, step in zip(crd, shape, stride): + total += crd2idx(value, one, step) if isinstance(one, tuple) else value * step + return total + + +__all__ = [ + "compact_col_major", + "compact_major", + "compact_row_major", + "crd2idx", + "idx2crd", + "try_compact_major", +] diff --git a/src/tilefoundry/ir/types/substitute.py b/src/tilefoundry/ir/types/substitute.py index 595a6fc2..b81d5628 100644 --- a/src/tilefoundry/ir/types/substitute.py +++ b/src/tilefoundry/ir/types/substitute.py @@ -112,9 +112,9 @@ def dim_vars_by_name(value: object) -> dict[str, "DimVar"]: def _shard_types() -> tuple[type, ...]: """Shard geometry descriptors, imported at call time to avoid a cycle.""" - from .shard.layout import ComposedLayout, Layout # noqa: PLC0415 - from .shard.mesh import Mesh, Topology # noqa: PLC0415 - from .shard.shard_layout import ShardLayout # noqa: PLC0415 + from .layout import ComposedLayout, Layout # noqa: PLC0415 + from .mesh import Mesh, Topology # noqa: PLC0415 + from .shard_layout import ShardLayout # noqa: PLC0415 return (Layout, ComposedLayout, ShardLayout, Mesh, Topology) @@ -257,13 +257,20 @@ def canonicalize_topology_dims(topology: object) -> object: return Topology(topology.name, size) +def _map_mesh_layout(layout, rebuild): + """Rebuild what a mesh stated, whether that is one arrangement or one per level.""" + if isinstance(layout, tuple): + return tuple(rebuild(one) for one in layout) + return rebuild(layout) + + def _canonicalize_mesh_dims(mesh: object) -> object: _, _, _, Mesh, _ = _shard_types() if not isinstance(mesh, Mesh): return mesh topologies = tuple(canonicalize_topology_dims(item) for item in mesh.topologies) - layout = _canonicalize_layout_dims(mesh.layout) - if topologies == mesh.topologies and layout is mesh.layout: + layout = _map_mesh_layout(mesh.layout, _canonicalize_layout_dims) + if topologies == mesh.topologies and layout == mesh.layout: return mesh return Mesh(topologies=topologies, layout=layout, names=mesh.names) @@ -328,8 +335,8 @@ def substitute_mesh_dims(mesh: object, bindings: Mapping[str, int]) -> object: if not isinstance(mesh, Mesh): return mesh topologies = tuple(substitute_topology_dims(item, bindings) for item in mesh.topologies) - layout = substitute_layout_dims(mesh.layout, bindings) - if topologies == mesh.topologies and layout is mesh.layout: + layout = _map_mesh_layout(mesh.layout, lambda one: substitute_layout_dims(one, bindings)) + if topologies == mesh.topologies and layout == mesh.layout: return mesh return Mesh(topologies=topologies, layout=layout, names=mesh.names) diff --git a/src/tilefoundry/ir/types/tensor_type.py b/src/tilefoundry/ir/types/tensor_type.py index cd5c92e4..c27848e3 100644 --- a/src/tilefoundry/ir/types/tensor_type.py +++ b/src/tilefoundry/ir/types/tensor_type.py @@ -74,7 +74,7 @@ def umat_scalar(dtype: DType = DType.i64) -> "TensorType": Shape elements and dimension arithmetic results use this single form so structural type equality holds across construction sites. """ - from .shard.layout import EMPTY_LAYOUT # noqa: PLC0415 - cycle guard + from .layout import EMPTY_LAYOUT # noqa: PLC0415 - cycle guard return TensorType( shape=(), dtype=dtype, layout=EMPTY_LAYOUT, storage=StorageKind.UMAT @@ -83,7 +83,7 @@ def umat_scalar(dtype: DType = DType.i64) -> "TensorType": @staticmethod def umat_tensor(shape: tuple, dtype: DType = DType.i64) -> "TensorType": """Canonical ranked value whose residency is not decided.""" - from .shard.layout import EMPTY_LAYOUT # noqa: PLC0415 - cycle guard + from .layout import EMPTY_LAYOUT # noqa: PLC0415 - cycle guard return TensorType( shape=tuple(shape), diff --git a/src/tilefoundry/ir/types/utils.py b/src/tilefoundry/ir/types/utils.py index 5aa6494d..7836a7f0 100644 --- a/src/tilefoundry/ir/types/utils.py +++ b/src/tilefoundry/ir/types/utils.py @@ -3,22 +3,21 @@ import math from typing import Optional +from tilefoundry.ir.types.int_tuple import repeat_like +from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.storage import StorageKind from .dtype import DType -from .shard import ( - ComposedLayout, - Layout, - Mesh, +from .layout import ComposedLayout, Layout +from .layout_algebra import size +from .mesh import Mesh, Topology +from .shard_layout import ( ShardLayout, Split, - Topology, canonical_shard_layout, shard_layout_of, - topology_axes, + split_target_axes, ) -from .shard.layout_algebra import size -from .shard.shard_layout import split_target_axes from .tensor_type import TensorType, TupleType, Type @@ -146,7 +145,7 @@ def topology_extent(type: Type, name: str) -> int | None: names = tuple(topology.name for topology in layout.mesh.topologies) if len(names) != 1 or names[0] != name: continue - count = size(layout.mesh.layout) + count = size(flatten(layout.mesh.layout)) if not isinstance(count, int) or isinstance(count, bool) or count <= 0: raise ValueError( f"topology_extent: {name!r} needs a positive static layout size" @@ -213,7 +212,7 @@ def local_type_of( for mesh_axis, tensor_axis in enumerate(split_target_axes(layout, type.shape)): if tensor_axis is None: continue - extent = layout.mesh.layout.shape[mesh_axis] + extent = flatten(layout.mesh.layout).shape[mesh_axis] if extent is None: local[tensor_axis] = 1 continue @@ -300,20 +299,29 @@ def _local_layout_shape( ) ) declared = {topology.name: index for index, topology in enumerate(topologies)} - axis_topology_level: dict[int, int] = {} - for topology, axes in zip(layout.mesh.topologies, topology_axes(layout.mesh)): - position = declared.get(topology.name) - if position is None: + for topology in layout.mesh.topologies: + if topology.name not in declared: raise ValueError( f"local_type_of: shard uses undeclared topology level {topology.name!r}" ) - for mesh_axis in axes: - axis_topology_level[mesh_axis] = position - mesh_shape = layout.mesh.layout.shape + mesh_layout = layout.mesh.layout + stated = mesh_layout.outer if isinstance(mesh_layout, ComposedLayout) else mesh_layout + axis_topology_level = flatten( + tuple( + repeat_like(mode, declared[topology.name]) + for mode, topology in zip(stated.shape, layout.mesh.topologies, strict=True) + ) + ) + mesh_shape = flatten(layout.mesh.layout).shape for mesh_axis, attr in enumerate(layout.attrs): if not isinstance(attr, Split): continue - if axis_topology_level.get(mesh_axis, selected_topology_level) > selected_topology_level: + here = ( + axis_topology_level[mesh_axis] + if mesh_axis < len(axis_topology_level) + else selected_topology_level + ) + if here > selected_topology_level: continue if mesh_axis >= len(mesh_shape): raise ValueError("local_type_of: shard attribute exceeds mesh layout rank") @@ -341,3 +349,99 @@ def _local_layout_shape( shape[axis] //= extent _require_concrete(shape) return tuple(shape) + + +def static_dim_value(dim): + """Return the compile-time ``int`` value of a *static* shape dim, else ``None``. + + A static dim is a plain ``int`` or an integer-valued ``Constant`` (the latter + only appears transiently before ``TensorType`` canonicalizes it to ``int``). + ``DimVar`` / dynamic dim ``Call`` exprs are not static → ``None``. The + detection is exact (real ``Constant`` with an ``int`` value), never "anything + with a ``.value``". + """ + from .dim import Constant # noqa: PLC0415 - cycle guard + + if isinstance(dim, int) and not isinstance(dim, bool): + return dim + if isinstance(dim, Constant) and isinstance(dim.value, int) and not isinstance(dim.value, bool): + return int(dim.value) + return None + + +def i64_const(value: int) -> "Constant": + """The canonical i64 shape-scalar ``Constant`` (meta-scalar typed).""" + from .dim import Constant # noqa: PLC0415 - cycle guard + from .tensor_type import TensorType # noqa: PLC0415 - cycle guard + + return Constant(type=TensorType.umat_scalar(), value=int(value)) + + +def upper_bound(dim) -> int: + """Return a concrete int upper-bound element count for ``dim``.""" + from .dim import DimVar # noqa: PLC0415 - cycle guard + + if isinstance(dim, DimVar): + return int(dim.hi) - 1 + static = static_dim_value(dim) + if static is not None: + return static + return int(dim) + + +def shape_numel_upper_bound(shape) -> int: + """Product of per-dim upper bounds. + + Product of per-dim upper bounds: the static element count a buffer or + layout must hold across every runtime shape in the dispatch envelope. + """ + n = 1 + for s in shape: + n *= upper_bound(s) + return n + + +def shape_upper_bound(shape) -> tuple[int, ...]: + """Map ``upper_bound`` over every entry of *shape*.""" + return tuple(upper_bound(s) for s in shape) + + +def shape_has_dim_var(shape) -> bool: + """True iff *shape* contains at least one ``DimVar`` entry.""" + from .dim import DimVar # noqa: PLC0415 - cycle guard + + return any(isinstance(s, DimVar) for s in shape) + + +def shape_runtime_total(shape, dim_var_expr: dict[str, str]) -> object: + """Return the runtime element count of *shape*. + + All-static shape → an ``int``. Any ``DimVar`` axis pulls its + runtime extent from ``dim_var_expr[name]``; the result is a C++ + expression string ``"(a * b * ...)"`` that the codegen splices + verbatim into the generated source. Static dims fold into a single + leading constant factor when present, otherwise the constant is + elided. + """ + from .dim import DimVar # noqa: PLC0415 - cycle guard + + if not shape: + return 1 + static_prod = 1 + dyn_terms: list[str] = [] + for s in shape: + if isinstance(s, DimVar): + expr = dim_var_expr.get(s.name) + if expr is None: + static_prod *= upper_bound(s) + else: + dyn_terms.append(expr) + else: + static_prod *= upper_bound(s) + if not dyn_terms: + return static_prod + if static_prod == 1: + if len(dyn_terms) == 1: + return dyn_terms[0] + return "(" + " * ".join(dyn_terms) + ")" + return "(" + " * ".join([str(static_prod), *dyn_terms]) + ")" diff --git a/src/tilefoundry/module.py b/src/tilefoundry/module.py index e3721f0d..b0cca068 100644 --- a/src/tilefoundry/module.py +++ b/src/tilefoundry/module.py @@ -57,7 +57,7 @@ def enclosing_declaration(frame: FrameType | None) -> _Entry | None: def _validate(topologies) -> tuple: - from tilefoundry.ir.types.shard.mesh import Topology # noqa: PLC0415 + from tilefoundry.ir.types.mesh import Topology # noqa: PLC0415 if not isinstance(topologies, tuple) or not all( isinstance(topology, Topology) for topology in topologies diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 08715ff2..ce2fcc43 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -49,6 +49,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.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 ( @@ -62,7 +63,19 @@ While, ) from tilefoundry.ir.tir.symbol_ref import SymbolRef -from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType +from tilefoundry.ir.types import ( + Broadcast, + DType, + Layout, + Mesh, + Partial, + ShardLayout, + Split, + TensorType, + Topology, + TupleType, + UnitType, +) from tilefoundry.ir.types.callable_type import CallableType from tilefoundry.ir.types.dim import ( DimAdd, @@ -74,20 +87,11 @@ dim_expr, simplify_dim, ) -from tilefoundry.ir.types.shard import ( - Broadcast, - Layout, - Mesh, - Partial, - ShardLayout, - Split, - Topology, - c_order_strides, - canonical_shard_layout, - composed, -) -from tilefoundry.ir.types.shard.layout import LayoutBase +from tilefoundry.ir.types.layout import LayoutBase +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 +from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.ir.visitor import BindingSubstitutionCloner from tilefoundry.target import MemoryHierarchyFacts, Target, UnsupportedCapabilityError from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext @@ -293,11 +297,12 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex UnitType=UnitType, Var=Var, DISPLAY_NAME=DISPLAY_NAME, - c_order_strides=c_order_strides, + compact_row_major=compact_row_major, canonical_shard_layout=canonical_shard_layout, - composed=composed, + merge_mesh=merge_mesh, dim_expr=dim_expr, normalize_dim=normalize_dim, + static_dim_value=static_dim_value, slice_size=slice_size, simplify_dim=simplify_dim, resolve_storage=resolve_storage, diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index f3737cf4..c3b9ecfd 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -32,9 +32,9 @@ from tilefoundry.ir.core.pattern import _mangle_variant_name from tilefoundry.ir.hir.nn.matmul import MatMul from tilefoundry.ir.tir.launch import launch_call -from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types import Broadcast, Layout, Partial, Split, TensorType from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Broadcast, Layout, Partial, Split +from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.substitute import canonicalize_dims from tilefoundry.ir.types.utils import types_compatible @@ -340,7 +340,7 @@ def construct(match, children, context): shape = tuple(children.values()) layout = runtime.Layout( shape=shape, - strides=runtime.c_order_strides(shape, mul=operator.mul), + strides=runtime.compact_row_major(shape, mul=operator.mul), ) if ( context.situation != "mesh_layout" @@ -350,7 +350,7 @@ def construct(match, children, context): mesh = context.function.state.mesh_stack[-1] return runtime.ShardLayout( layout=layout, - attrs=tuple(runtime.Broadcast() for _ in mesh.layout.shape), + attrs=tuple(runtime.Broadcast() for _ in flatten(mesh.layout).shape), mesh=mesh, ) return layout @@ -392,7 +392,7 @@ def construct(match, children, context): node, context, f"{binding!r} is not a lexical Mesh binding" ) if axis_name is None: - if len(mesh.layout.shape) != 1: + if len(flatten(mesh.layout).shape) != 1: raise ParseError.from_node( node, context, "bare Mesh placement requires a one-axis mesh" ) @@ -658,13 +658,13 @@ 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.composed(meshes) + mesh = meshes[0] if len(meshes) == 1 else runtime.merge_mesh(meshes) source_offsets: dict[int, int] = {} offset = 0 for source in meshes: source_offsets[id(source)] = offset - offset += len(source.layout.shape) - attrs: list[object] = [runtime.Broadcast() for _ in mesh.layout.shape] + offset += len(flatten(source.layout).shape) + attrs: list[object] = [runtime.Broadcast() for _ in flatten(mesh.layout).shape] for source, source_axis, tensor_axis in value.splits: attrs[source_offsets[id(source)] + source_axis] = runtime.Split(tensor_axis) for source, source_axis, kind, reduction in value.states: @@ -3130,7 +3130,7 @@ def _bind( ) if axis is None and node.attr in {"x", "y", "z"}: candidate = ("x", "y", "z").index(node.attr) - if candidate < len(mesh.layout.shape): + if candidate < len(flatten(mesh.layout).shape): axis = candidate if axis is None: named = ", ".join(mesh.names) @@ -3152,7 +3152,7 @@ def construct(match, children, context): raise ParseError.from_node(match.node, context, "mesh coordinate lacks context") mesh = match.captures["mesh"] axis = match.captures["axis"] - extent = mesh.layout.shape[axis] + extent = flatten(mesh.layout).shape[axis] if isinstance(extent, bool) or not isinstance(extent, int): raise ParseError.from_node( match.node, context, "mesh coordinate requires a concrete axis extent" @@ -3338,6 +3338,10 @@ def construct(match, children, context): "Mesh topologies must be a tuple", ) names = children.get("names", ()) + if any(isinstance(topology, str) for topology in topology_names): + topology_names = _resolve_mesh_topologies_at( + topology_names, context.function.topologies, match.node, context + ) try: mesh = runtime.Mesh( topologies=topology_names, @@ -3346,7 +3350,6 @@ def construct(match, children, context): ) except (TypeError, ValueError) as error: raise ParseError.from_node(match.node, context, str(error)) from error - mesh = _resolved_mesh(mesh, match, context) elif match.branch_id == "mesh_reference": mesh = children["value"] if not isinstance(mesh, runtime.Mesh): @@ -3401,18 +3404,37 @@ def _loaded_names(statements): ) -def _block_escaping_names(statements): - """Return each with child's escaping bindings in one reverse block scan.""" +def _read_before_bound(statements): + """Return names a block reads before it binds them, in statement order. + + A name a block reads on its way to binding it came from outside the block: + an accumulator reads what the last round left before it writes this one. + """ + bound: set[str] = set() + live: set[str] = set() + for statement in statements: + live.update(_loaded_names((statement,)) - bound) + bound.update(_directly_bound_names((statement,))) + return frozenset(live) + + +def _block_escaping_names(statements, *, repeated: bool = False): + """Return each with child's escaping bindings in one reverse block scan. + + A block that repeats carries what it reads on its way to binding it, so a + `with` that both reads and binds a name states that name's next value and + the region it is bound through is what the round after reads. + """ read_after: set[str] = set() escaping: dict[int, Mapping[str, object]] = {} for index in range(len(statements) - 1, -1, -1): statement = statements[index] if isinstance(statement, ast.With): - escaping[index] = { - "escaping_names": frozenset( - _directly_bound_names(statement.body) & read_after - ) - } + bound = _directly_bound_names(statement.body) + reached = set(read_after) + if repeated: + reached.update(_read_before_bound(statement.body)) + escaping[index] = {"escaping_names": frozenset(bound & reached)} read_after.update(_loaded_names((statement,))) return escaping @@ -3428,7 +3450,7 @@ def _enter_mesh_scope(context, mesh, match): infer = _parser_infer_context(context) try: entered_mesh = ( - runtime.composed((infer.current_mesh, mesh)) + runtime.merge_mesh((infer.current_mesh, mesh)) if infer.current_mesh else mesh ) @@ -3534,7 +3556,9 @@ def _rebind_through_region(context, mesh, names, frame, node, params=(), args=() A region's computed values are named, and those names are how code after it reads them. Rebinding each through the region is what keeps who-ran-this on - the graph instead of leaving it to be guessed from a layout. + the graph instead of leaving it to be guessed from a layout. The first name + rebound comes back, so a caller with no body value of its own still has one + that reaches the region. """ values = [ (name, frame[name]) @@ -3552,13 +3576,14 @@ def _rebind_through_region(context, mesh, names, frame, node, params=(), args=() if len(values) == 1: scoped = _scoped_region(mesh, values[0][1], params, args) _bind_region_results(context, scoped, [values[0][0]], node) - return + 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) ) scoped = _scoped_region(mesh, tuple_body, params, args) _bind_region_results(context, scoped, [name for name, _value in values], node) + return values[0][0] class WithPattern(ElementPattern): @@ -3658,13 +3683,14 @@ def construct(match, children, context): escaping = context.values.get("escaping_names", frozenset()) params = match.captures.get("region_params", ()) args = match.captures.get("region_args", ()) + rebound = None if escaping: - _rebind_through_region( + rebound = _rebind_through_region( context, mesh, escaping, frame, match.node, params, args ) if body is not None: return _scoped_region(mesh, body, params, args) - return None + return None if rebound is None else context.lexical_scope.lookup(rebound) binding = runtime.Var( type=runtime.TensorType.scalar(runtime.DType.i64, storage=runtime.StorageKind.RMEM), name=match.captures["binding"], @@ -3788,13 +3814,36 @@ class LoopCarryStatementPattern(ElementPattern): "loop_carry_statement", AstNodePattern( ast.stmt, - CapturePattern("names", lambda node, context: ()), + CapturePattern("names", LoopCarryStatementPattern._statement_names), ), pattern_id="loop.carry_statement", ), ) ) + @staticmethod + def _statement_names(node: object, context: MatchContext) -> tuple[str, ...]: + """Names a statement binds that the loop carries rather than the loop itself. + + An assignment states its own; a `with` states what its body binds, + because a region delimits who runs the work and not who owns the value. + Every other statement binds nothing the loop carries. + """ + if not isinstance(node, ast.With): + return () + found: list[str] = [] + for child in ast.walk(node): + if not isinstance(child, ast.Assign): + continue + 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 + ) + return tuple(found) + @staticmethod def _target_names(node: object, context: MatchContext) -> tuple[str, ...]: assert isinstance(node, ast.Assign) @@ -4065,6 +4114,11 @@ def construct(match, children, context): 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 + else runtime.normalize_dim(bound) + for bound in bounds + ] context.lexical_scope.push_frame() context.lexical_scope.define(match.captures["target"], iv) return (iv, *bounds) @@ -4118,40 +4172,73 @@ def construct(match, children, context): class LoopBodyPattern(ElementPattern): element_name = "loop_body" syntax = LazyPattern( - lambda: BranchPattern( - "loop_body", - AstNodePattern( - ast.Module, - PredicatePattern( - "assignment-suite", - lambda node, context: ( - not any( - isinstance(statement, (ast.Return, ast.With, ast.Expr, ast.Pass)) - for statement in node.body - ) + lambda: BindPattern( + BranchPattern( + "loop_body", + AstNodePattern( + ast.Module, + PredicatePattern( + "assignment-suite", + lambda node, context: ( + not any( + isinstance(statement, (ast.Return, ast.Expr, ast.Pass)) + for statement in node.body + ) + ), ), - ), - FieldPattern( - "body", - RepeatPattern( - ChildPattern( - "statement_{index}", - StatementPattern(), - "loop_statement", - "loop_statement", - ) + FieldPattern( + "body", + RepeatPattern( + ChildPattern( + "statement_{index}", + StatementPattern(), + "loop_statement", + "loop_statement", + ) + ), ), ), + pattern_id="loop.body", ), - pattern_id="loop.body", + LoopBodyPattern._bind, ) ) + @staticmethod + def _bind(node, _context, matched): + """Tell each `with` in the body which of its names the loop reads later. + + A `with Mesh(...)` states who runs the statements it holds, so a value + assigned inside one and read after it leaves through the region rather + than around it. Which names those are is the same reverse scan a + function block makes, asked of the loop's own statements. + """ + assert isinstance(node, ast.Module) + escaping = _block_escaping_names(node.body, repeated=True) + child_values = {f"statement_{index}": values for index, values in escaping.items()} + return dataclasses.replace( + matched, + children=tuple( + dataclasses.replace(child, values={**child.values, **child_values[child.name]}) + if child.name in child_values + else child + for child in matched.children + ), + ) + @staticmethod def construct(match, children, context): + """The loop body's value: the last statement that has one. + + A `with` that only rebinds the names escaping it has no value of its + own -- the names it rebound reach its region, and the loop reads them + as its carried values -- so the body's value is the statement before. + """ if not children: raise ParseError.from_node(match.node, context, "loop body cannot be empty") - value = tuple(children.values())[-1] + value = next( + (item for item in reversed(tuple(children.values())) if item is not None), None + ) if not isinstance(value, runtime.Expr): raise ParseError.from_node(match.node, context, "loop body must yield an Expr") return value diff --git a/src/tilefoundry/runtime/decorator.py b/src/tilefoundry/runtime/decorator.py index 26775fcf..be149117 100644 --- a/src/tilefoundry/runtime/decorator.py +++ b/src/tilefoundry/runtime/decorator.py @@ -11,7 +11,8 @@ from typing import Callable from tilefoundry.ir.core.module import Module, _refuse_bare_call -from tilefoundry.ir.types.shard import Placement, shard_layout_of +from tilefoundry.ir.types import Placement +from tilefoundry.ir.types.shard_layout import shard_layout_of from tilefoundry.ir.types.tensor_type import TensorType from tilefoundry.runtime.function import RuntimeFunction from tilefoundry.runtime.module import RuntimeModule diff --git a/src/tilefoundry/runtime/measure.py b/src/tilefoundry/runtime/measure.py index 53f006f4..0b005beb 100644 --- a/src/tilefoundry/runtime/measure.py +++ b/src/tilefoundry/runtime/measure.py @@ -11,7 +11,7 @@ import torch from tilefoundry.evaluator.value import from_torch_dtype -from tilefoundry.ir.types.shard import shard_layout_of +from tilefoundry.ir.types.shard_layout import shard_layout_of from tilefoundry.runtime.tensor import ShardTensor _NEAR_ZERO = 1e-12 diff --git a/src/tilefoundry/runtime/module.py b/src/tilefoundry/runtime/module.py index b09eea45..b6c963eb 100644 --- a/src/tilefoundry/runtime/module.py +++ b/src/tilefoundry/runtime/module.py @@ -9,7 +9,7 @@ from typing import Callable from tilefoundry.codegen.signature import CallableSignature -from tilefoundry.ir.types.shard import Placement +from tilefoundry.ir.types import Placement from tilefoundry.runtime.resource import RuntimeResource __all__ = ["CompiledModule", "RuntimeModule"] diff --git a/src/tilefoundry/runtime/tensor.py b/src/tilefoundry/runtime/tensor.py index ce04a562..ea8d28cf 100644 --- a/src/tilefoundry/runtime/tensor.py +++ b/src/tilefoundry/runtime/tensor.py @@ -14,12 +14,8 @@ import torch -from tilefoundry.ir.types.shard import ( - Placement, - Topology, - shard_layout_of, -) -from tilefoundry.ir.types.shard.local import local_layout_and_offset +from tilefoundry.ir.types import Placement, Topology +from tilefoundry.ir.types.shard_layout import local_layout_and_offset, shard_layout_of from tilefoundry.ir.types.tensor_type import TensorType diff --git a/src/tilefoundry/script.py b/src/tilefoundry/script.py index 025cd491..3d8d0971 100644 --- a/src/tilefoundry/script.py +++ b/src/tilefoundry/script.py @@ -19,7 +19,7 @@ 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 -from tilefoundry.ir.types.shard import Mesh +from tilefoundry.ir.types import Mesh from tilefoundry.module import UNDECLARED, _Entry from tilefoundry.parser import FuncParserContext, FunctionRole, parse_function from tilefoundry.parser.ast_pattern import ( diff --git a/src/tilefoundry/target/base.py b/src/tilefoundry/target/base.py index 37c09399..1b5e77df 100644 --- a/src/tilefoundry/target/base.py +++ b/src/tilefoundry/target/base.py @@ -10,8 +10,7 @@ from types import MappingProxyType from typing import Any, ClassVar, Mapping, TypeVar -from tilefoundry.ir.types import DType -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import DType, Topology from tilefoundry.target.facts import TopologyFacts, facts_result from tilefoundry.target.hardware.envelope import ( DuplicateRegistrationError, diff --git a/src/tilefoundry/visitor_registry/access_relation.py b/src/tilefoundry/visitor_registry/access_relation.py index aac1c5b0..c2a26144 100644 --- a/src/tilefoundry/visitor_registry/access_relation.py +++ b/src/tilefoundry/visitor_registry/access_relation.py @@ -16,9 +16,9 @@ from tilefoundry.ir.hir._helpers import is_one from tilefoundry.ir.isl_interop import index_set, isl_to_dim, shape_to_isl_domain -from tilefoundry.ir.types import TensorType, TupleType, Type, tensor_bytes -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard.shard_layout import layout_axis_to_tensor_axis +from tilefoundry.ir.types import TensorType, TupleType, Type +from tilefoundry.ir.types.shard_layout import layout_axis_to_tensor_axis +from tilefoundry.ir.types.utils import static_dim_value, tensor_bytes from tilefoundry.utils.isl_utils import cardinality from .registries import DispatchRegistry diff --git a/src/tilefoundry/visitor_registry/contexts.py b/src/tilefoundry/visitor_registry/contexts.py index 8d737110..7c467d05 100644 --- a/src/tilefoundry/visitor_registry/contexts.py +++ b/src/tilefoundry/visitor_registry/contexts.py @@ -14,7 +14,7 @@ get_metadata, ) from tilefoundry.ir.core.stmt import Stmt -from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.ir.types import Mesh, Topology from tilefoundry.ir.types.tensor_type import DType, Type from tilefoundry.ir.types.utils import local_type_of diff --git a/src/tilefoundry/visitor_registry/op_cost.py b/src/tilefoundry/visitor_registry/op_cost.py index 479774fc..a5fd3e50 100644 --- a/src/tilefoundry/visitor_registry/op_cost.py +++ b/src/tilefoundry/visitor_registry/op_cost.py @@ -55,9 +55,15 @@ from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.hir.tensor.where import Where from tilefoundry.ir.hir.tensor.zeros import Zeros -from tilefoundry.ir.types import DType, IntegerDType, TensorType, Type, numel, tensor_bytes -from tilefoundry.ir.types.shard import ShardLayout, flatten, shard_layout_of, topology_axes -from tilefoundry.ir.types.shard.shard_layout import layout_axis_to_tensor_axis, split_target_axes +from tilefoundry.ir.types import DType, IntegerDType, ShardLayout, TensorType, Type +from tilefoundry.ir.types.int_tuple import repeat_like +from tilefoundry.ir.types.layout import ComposedLayout, flatten +from tilefoundry.ir.types.shard_layout import ( + layout_axis_to_tensor_axis, + shard_layout_of, + split_target_axes, +) +from tilefoundry.ir.types.utils import numel, tensor_bytes from tilefoundry.visitor_registry.access_relation import logical_axes_of from .contexts import Cost, CostContext, TrafficBytes @@ -550,12 +556,21 @@ def _sent(source, destination) -> tuple[tuple[str, TrafficBytes], ...]: if before is None or after is None or before == after: return () mesh = shard_layout_of(source.layout).mesh - extents = flatten(mesh.layout.shape) + extents = flatten(flatten(mesh.layout).shape) + stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout + levels = flatten( + tuple( + repeat_like(mode, topology.name) + for mode, topology in zip(stated.shape, mesh.topologies, strict=True) + ) + ) held = tensor_bytes(source) moved: dict[str, int] = {} - for topology, axes in zip(mesh.topologies, topology_axes(mesh)): + for topology in mesh.topologies: units = 1 - for mesh_axis in axes: + for mesh_axis, level in enumerate(levels): + if level != topology.name: + continue if before.get(mesh_axis) != after.get(mesh_axis): units *= extents[mesh_axis] if units > 1: diff --git a/src/tilefoundry/visitor_registry/shard_propagate.py b/src/tilefoundry/visitor_registry/shard_propagate.py index d8d8f585..533be5e5 100644 --- a/src/tilefoundry/visitor_registry/shard_propagate.py +++ b/src/tilefoundry/visitor_registry/shard_propagate.py @@ -12,19 +12,17 @@ import isl -from tilefoundry.ir.types.shard import ( - Layout, - ShardLayout, - canonical_shard_layout, - try_c_order_strides, -) -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types import Layout, ShardLayout +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.shard_layout import ( Broadcast, Partial, Split, + canonical_shard_layout, layout_axis_to_tensor_axis, shard_layout_of, ) +from tilefoundry.ir.types.stride import try_compact_major from tilefoundry.utils.isl_utils import as_multi_aff, equates, involved_dims from tilefoundry.visitor_registry.access_relation import boundary_maps @@ -168,7 +166,7 @@ def _carrier_layout( if fresh_strides: - c = try_c_order_strides(tuple(new_shape)) or tuple(1 for _ in new_shape) + c = try_compact_major(tuple(new_shape)) or tuple(1 for _ in new_shape) new_strides = [ 0 if (isinstance(sz, int) and sz == 1) else cc for sz, cc in zip(new_shape, c) @@ -242,7 +240,7 @@ def derive_output_shard_layout( if not sharded: return None mesh = sharded[0][1].mesh - mesh_rank = len(mesh.layout.shape) + mesh_rank = len(flatten(mesh.layout).shape) *input_maps, output_map = boundary_maps(relations) out_access = _result_access(output_map, folded=True) diff --git a/src/tilefoundry/visitor_registry/typeinfer.py b/src/tilefoundry/visitor_registry/typeinfer.py index 68761213..bb2275e2 100644 --- a/src/tilefoundry/visitor_registry/typeinfer.py +++ b/src/tilefoundry/visitor_registry/typeinfer.py @@ -14,11 +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.tir.shape import ShapeOf from tilefoundry.ir.types.callable_type import callable_type_for -from tilefoundry.ir.types.shard.mesh import composed -from tilefoundry.ir.types.shard.scope_match import covered_by_scope, storage_reaches -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +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 from tilefoundry.ir.types.utils import types_compatible @@ -221,7 +220,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 = composed((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh + mesh = merge_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 e13833a2..28170286 100644 --- a/tests/analysis/test_analysis_invariants.py +++ b/tests/analysis/test_analysis_invariants.py @@ -49,14 +49,15 @@ from tilefoundry.ir.types import ( DType, TensorType, + Topology, TupleType, + make_mesh, make_shard_tensor_type, make_tensor_type, - tensor_bytes, ) -from tilefoundry.ir.types.shard import Topology, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Split as ShardSplit +from tilefoundry.ir.types.shard_layout import Split as ShardSplit from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.types.utils import tensor_bytes from tilefoundry.visitor_registry.access_relation import ( AccessRelations, AffineAccess, diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index ac88c25d..d32346b3 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -44,9 +44,7 @@ variant_for, ) from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice -from tilefoundry.ir.types.shard import ( - Topology, -) +from tilefoundry.ir.types import Topology from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget, PerformanceServiceFacts, ThroughputFacts diff --git a/tests/analysis/test_analyze_composition.py b/tests/analysis/test_analyze_composition.py index 1a00fd56..b135c66a 100644 --- a/tests/analysis/test_analyze_composition.py +++ b/tests/analysis/test_analyze_composition.py @@ -12,7 +12,7 @@ from tilefoundry.analysis.errors import AnalysisError from tilefoundry.analysis.registry import Analyzer from tilefoundry.ir.core.module import Module -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget, Target diff --git a/tests/analysis/test_analyze_cross_module.py b/tests/analysis/test_analyze_cross_module.py index 6eeef3bc..10011c45 100644 --- a/tests/analysis/test_analyze_cross_module.py +++ b/tests/analysis/test_analyze_cross_module.py @@ -23,9 +23,9 @@ from tilefoundry.ir.core.module import reachable_functions from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.nn.matmul import MatMul -from tilefoundry.ir.types import tensor_types -from tilefoundry.ir.types.shard.layout import ComposedLayout -from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.ir.types.layout import ComposedLayout, flatten +from tilefoundry.ir.types.shard_layout import ShardLayout +from tilefoundry.ir.types.utils import tensor_types from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget @@ -216,5 +216,5 @@ def test_each_placed_branch_keeps_its_slice_on_its_primitive_results() -> None: assert placed, branch assert {op for op, _ in placed} - {"Reshard"}, branch for _op, layout in placed: - assert layout.outer.shape == shape + assert flatten(layout.outer).shape == shape assert layout.offset == offset diff --git a/tests/analysis/test_mesh_region_cost.py b/tests/analysis/test_mesh_region_cost.py index c2469b2c..dd920d60 100644 --- a/tests/analysis/test_mesh_region_cost.py +++ b/tests/analysis/test_mesh_region_cost.py @@ -12,8 +12,7 @@ from tilefoundry.ir.core import Call, Var, VerifyError, get_metadata from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.mesh_region import MeshRegion -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import Layout +from tilefoundry.ir.types import DType, Layout, TensorType from tilefoundry.ir.types.storage import StorageKind from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget diff --git a/tests/analysis/test_shard_propagate.py b/tests/analysis/test_shard_propagate.py index 17ab30d3..6ff754ee 100644 --- a/tests/analysis/test_shard_propagate.py +++ b/tests/analysis/test_shard_propagate.py @@ -14,9 +14,15 @@ import isl import pytest -from tilefoundry.ir.types import make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Topology -from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, Split +from tilefoundry.ir.types import ( + Layout, + Mesh, + ShardLayout, + Topology, + make_shard_tensor_type, + make_tensor_type, +) +from tilefoundry.ir.types.shard_layout import Broadcast, Partial, Split from tilefoundry.visitor_registry.access_relation import ( AccessRelations, AffineAccess, diff --git a/tests/codegen/test_host_multi_launch.py b/tests/codegen/test_host_multi_launch.py index cb01b542..671459ba 100644 --- a/tests/codegen/test_host_multi_launch.py +++ b/tests/codegen/test_host_multi_launch.py @@ -17,8 +17,17 @@ from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.stmts import Evaluate, Sequential from tilefoundry.ir.tir.symbol_ref import SymbolRef -from tilefoundry.ir.types import CallableType, DType, TensorType, UnitType -from tilefoundry.ir.types.shard import B, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types import ( + B, + CallableType, + DType, + Layout, + Mesh, + ShardLayout, + TensorType, + Topology, + UnitType, +) from tilefoundry.ir.types.storage import StorageKind from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/codegen/test_reduce_tag_table.py b/tests/codegen/test_reduce_tag_table.py index 84a1f882..03f202d7 100644 --- a/tests/codegen/test_reduce_tag_table.py +++ b/tests/codegen/test_reduce_tag_table.py @@ -10,7 +10,7 @@ from tilefoundry.codegen.cuda.tir.memory.tensor_view import _render_attr from tilefoundry.codegen.cuda.tir.reduce import REDUCE_TAG from tilefoundry.ir.core.kinds import ReduceKind -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types.shard_layout import Partial def test_every_reduce_kind_has_a_runtime_tag() -> None: diff --git a/tests/codegen/test_shard_layout_value.py b/tests/codegen/test_shard_layout_value.py index fc65e973..4e0de216 100644 --- a/tests/codegen/test_shard_layout_value.py +++ b/tests/codegen/test_shard_layout_value.py @@ -12,8 +12,8 @@ from tilefoundry.codegen.cuda.tir.memory.tensor_view import render_shard_layout_value from tilefoundry.codegen.cuda.tir.stmts.mesh_scope import mesh_type from tilefoundry.ir.core.errors import VerifyError -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology -from tilefoundry.ir.types.shard.layout import ComposedLayout +from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Split, Topology +from tilefoundry.ir.types.layout import ComposedLayout _BLOCK = Mesh( (Topology("thread", 128),), Layout(shape=(4, 32), strides=(32, 1)), ("w", "t") diff --git a/tests/codegen/test_submodule_topology.py b/tests/codegen/test_submodule_topology.py index 7d162b03..e7a3ff1d 100644 --- a/tests/codegen/test_submodule_topology.py +++ b/tests/codegen/test_submodule_topology.py @@ -9,7 +9,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor from tilefoundry.dsl.tf import * # noqa: F401,F403 -- bind bare launch -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology +from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Split, Topology from tilefoundry.target import CpuTarget, CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") diff --git a/tests/dsl/test_module_context.py b/tests/dsl/test_module_context.py index 57a9addc..888da7ff 100644 --- a/tests/dsl/test_module_context.py +++ b/tests/dsl/test_module_context.py @@ -21,7 +21,7 @@ from tilefoundry.ir.core import VerifyError from tilefoundry.ir.core.module import Module from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CpuTarget, CudaTarget _CTA = context_fixture.CONTEXT_CTA @@ -190,7 +190,7 @@ def plain(x: Tensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: _MEMBER_CONTEXT_SOURCE = """ import tilefoundry from tilefoundry.dsl import Mesh, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget @tilefoundry.module( diff --git a/tests/dsl/test_module_decorator.py b/tests/dsl/test_module_decorator.py index 1ca138d3..cd83645b 100644 --- a/tests/dsl/test_module_decorator.py +++ b/tests/dsl/test_module_decorator.py @@ -23,7 +23,7 @@ from tilefoundry.evaluator.value import EvalError from tilefoundry.ir.core.errors import VerifyError from tilefoundry.ir.core.module import Module -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.runtime.resource import DictResource from tilefoundry.target import CpuTarget, CudaTarget from tilefoundry.utils.spec_ref import spec_ref_render diff --git a/tests/fixtures/diagnostics/dynamic_trip_count.py b/tests/fixtures/diagnostics/dynamic_trip_count.py index 2eb15017..cae421cb 100644 --- a/tests/fixtures/diagnostics/dynamic_trip_count.py +++ b/tests/fixtures/diagnostics/dynamic_trip_count.py @@ -4,7 +4,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget CTAS = 132 diff --git a/tests/fixtures/diagnostics/tuple_projection_diagnostic.py b/tests/fixtures/diagnostics/tuple_projection_diagnostic.py index f6643663..cf46727f 100644 --- a/tests/fixtures/diagnostics/tuple_projection_diagnostic.py +++ b/tests/fixtures/diagnostics/tuple_projection_diagnostic.py @@ -4,7 +4,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index 24bdb207..d42a40dc 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -4,7 +4,7 @@ from tilefoundry import func from tilefoundry.dsl import Tensor from tilefoundry.dsl.storage import gmem, rmem, smem from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology +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'))) def composed_mesh_pipeline( diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt index fb23fb86..02ece31a 100644 --- a/tests/fixtures/inspection/type_printer_sugar.printed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -4,7 +4,7 @@ from tilefoundry import func from tilefoundry.dsl import Tensor from tilefoundry.dsl.storage import gmem, rmem, smem from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology +from tilefoundry.ir.types import B, Layout, Mesh, P, S, ShardLayout, Topology from tilefoundry.module import module from tilefoundry.target import CudaTarget diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py index d7161fba..ac99eb37 100644 --- a/tests/fixtures/inspection/type_printer_sugar.py +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -12,7 +12,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Tensor, tf -from tilefoundry.ir.types.shard import B, Layout, Mesh, P, S, ShardLayout, Topology +from tilefoundry.ir.types import B, Layout, Mesh, P, S, ShardLayout, Topology from tilefoundry.target import CudaTarget _H200 = CudaTarget("nvidia.h200_sxm") diff --git a/tests/fixtures/logical/authored_constraint.py b/tests/fixtures/logical/authored_constraint.py index 9dde1b76..33a15319 100644 --- a/tests/fixtures/logical/authored_constraint.py +++ b/tests/fixtures/logical/authored_constraint.py @@ -4,7 +4,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Tensor, tf -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.target import CudaTarget _CTA_MESH = Mesh((Topology("cta", 8),), Layout((8,), (1,))) diff --git a/tests/fixtures/logical/gqa_static.py b/tests/fixtures/logical/gqa_static.py index c395cc58..cc8ce155 100644 --- a/tests/fixtures/logical/gqa_static.py +++ b/tests/fixtures/logical/gqa_static.py @@ -14,7 +14,7 @@ from tilefoundry import func from tilefoundry.dsl import Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.target import CudaTarget S = 1 diff --git a/tests/fixtures/logical/hir_composition.py b/tests/fixtures/logical/hir_composition.py index 8197f84b..759239e3 100644 --- a/tests/fixtures/logical/hir_composition.py +++ b/tests/fixtures/logical/hir_composition.py @@ -10,7 +10,7 @@ from tilefoundry import func, module from tilefoundry.dsl import ConstTensor, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget _H200 = CudaTarget("nvidia.h200_sxm") diff --git a/tests/fixtures/logical/matmul_rms_norm.py b/tests/fixtures/logical/matmul_rms_norm.py index fa31f257..264a9951 100644 --- a/tests/fixtures/logical/matmul_rms_norm.py +++ b/tests/fixtures/logical/matmul_rms_norm.py @@ -7,7 +7,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Tensor from tilefoundry.dsl.tf import matmul, rms_norm -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget diff --git a/tests/fixtures/logical/module_context.py b/tests/fixtures/logical/module_context.py index 4adf77aa..d18a0e7b 100644 --- a/tests/fixtures/logical/module_context.py +++ b/tests/fixtures/logical/module_context.py @@ -4,7 +4,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget CONTEXT_CTA = Topology("cta", 132) diff --git a/tests/fixtures/placed/derived_prefill.py b/tests/fixtures/placed/derived_prefill.py index acfb6e66..8860c474 100644 --- a/tests/fixtures/placed/derived_prefill.py +++ b/tests/fixtures/placed/derived_prefill.py @@ -2,7 +2,7 @@ from tilefoundry import func, module from tilefoundry.dsl import DimVar, Mesh, Tensor, ceildiv, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget PREFILL_N = DimVar("prefill_n", 1, 65) diff --git a/tests/fixtures/placed/flash_split_k_decode.py b/tests/fixtures/placed/flash_split_k_decode.py index 4be05542..f5b6b9ff 100644 --- a/tests/fixtures/placed/flash_split_k_decode.py +++ b/tests/fixtures/placed/flash_split_k_decode.py @@ -15,7 +15,7 @@ from tilefoundry import func, module from tilefoundry.dsl import DimVar, Mesh, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget HEADS = 16 diff --git a/tests/fixtures/placed/gemm_schedules.py b/tests/fixtures/placed/gemm_schedules.py index 266439cf..51fb80e5 100644 --- a/tests/fixtures/placed/gemm_schedules.py +++ b/tests/fixtures/placed/gemm_schedules.py @@ -5,7 +5,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget _H200 = CudaTarget("nvidia.h200_sxm") diff --git a/tests/fixtures/placed/gpu_placed_rows.py b/tests/fixtures/placed/gpu_placed_rows.py index c7d0c6ce..51c62aef 100644 --- a/tests/fixtures/placed/gpu_placed_rows.py +++ b/tests/fixtures/placed/gpu_placed_rows.py @@ -13,7 +13,7 @@ import tilefoundry.codegen.cuda # noqa: F401 -- trigger emitter autodiscovery from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology +from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Split, Topology from tilefoundry.target import CpuTarget, CudaTarget GPUS, CTAS, THREADS = 2, 4, 16 diff --git a/tests/fixtures/placed/gqa_decode.py b/tests/fixtures/placed/gqa_decode.py index dc4b9e37..5d1a110d 100644 --- a/tests/fixtures/placed/gqa_decode.py +++ b/tests/fixtures/placed/gqa_decode.py @@ -17,8 +17,8 @@ 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.types import Broadcast, Layout, Mesh, ShardLayout, Topology from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Broadcast, Layout, Mesh, ShardLayout, Topology HEAD_DIM = 128 NUM_Q_HEADS = 32 diff --git a/tests/fixtures/placed/hand_checked.py b/tests/fixtures/placed/hand_checked.py index 498a447b..426bb2fc 100644 --- a/tests/fixtures/placed/hand_checked.py +++ b/tests/fixtures/placed/hand_checked.py @@ -6,7 +6,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget S, K, N = 8, 4, 6 diff --git a/tests/fixtures/placed/mesh_slice_start.py b/tests/fixtures/placed/mesh_slice_start.py index 06554b08..7c715e21 100644 --- a/tests/fixtures/placed/mesh_slice_start.py +++ b/tests/fixtures/placed/mesh_slice_start.py @@ -4,7 +4,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget D, BLK, W, N = 64, 8, 4, 32 diff --git a/tests/fixtures/placed/moe_mega_kernel.py b/tests/fixtures/placed/moe_mega_kernel.py index d22b4d61..ca44b9cf 100644 --- a/tests/fixtures/placed/moe_mega_kernel.py +++ b/tests/fixtures/placed/moe_mega_kernel.py @@ -4,7 +4,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget _H200 = CudaTarget("nvidia.h200_sxm") diff --git a/tests/fixtures/placed/persistent_gemm_flat.py b/tests/fixtures/placed/persistent_gemm_flat.py index 8a84b1f7..c276cbbb 100644 --- a/tests/fixtures/placed/persistent_gemm_flat.py +++ b/tests/fixtures/placed/persistent_gemm_flat.py @@ -5,7 +5,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget M = 3840 diff --git a/tests/fixtures/placed/persistent_gemm_tiled.py b/tests/fixtures/placed/persistent_gemm_tiled.py index 2f12e6e8..d6a7e228 100644 --- a/tests/fixtures/placed/persistent_gemm_tiled.py +++ b/tests/fixtures/placed/persistent_gemm_tiled.py @@ -5,7 +5,7 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget M = 3840 diff --git a/tests/fixtures/placed/prefill_decode_attention.py b/tests/fixtures/placed/prefill_decode_attention.py index ffffa6ab..4df99e90 100644 --- a/tests/fixtures/placed/prefill_decode_attention.py +++ b/tests/fixtures/placed/prefill_decode_attention.py @@ -12,7 +12,7 @@ from tilefoundry import func, module from tilefoundry.dsl import DimVar, DimVarRangePat, Mesh, Tensor, ceildiv, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 -- bare tile() in authored bodies -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget SEQ = DimVar("seq", 1, 4097) diff --git a/tests/fixtures/placed/qwen3_1_7b_pd.py b/tests/fixtures/placed/qwen3_1_7b_pd.py index f9089ec5..cc40f818 100644 --- a/tests/fixtures/placed/qwen3_1_7b_pd.py +++ b/tests/fixtures/placed/qwen3_1_7b_pd.py @@ -16,7 +16,7 @@ from tilefoundry import func, module from tilefoundry.dsl import ConstTensor, DimVar, DimVarRangePat, Mesh, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401,F403 -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget _CONFIG = Path(__file__).parents[2] / "models" / "qwen3_1_7b" / "config.json" diff --git a/tests/fixtures/placed/specialize_through_call.py b/tests/fixtures/placed/specialize_through_call.py index 84ced88e..afd30fc3 100644 --- a/tests/fixtures/placed/specialize_through_call.py +++ b/tests/fixtures/placed/specialize_through_call.py @@ -3,7 +3,7 @@ from tilefoundry import func, module from tilefoundry.dsl import DimVar, DimVarRangePat, Mesh, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget D, W, BOUND = 64, 4, 128 diff --git a/tests/fixtures/placed/symbolic_offset.py b/tests/fixtures/placed/symbolic_offset.py index 657dfb15..d34e8f56 100644 --- a/tests/fixtures/placed/symbolic_offset.py +++ b/tests/fixtures/placed/symbolic_offset.py @@ -5,7 +5,7 @@ from tilefoundry import func, module from tilefoundry.dsl import ConstTensor, DimVar, Mesh, Tensor, tf from tilefoundry.dsl.tf import * # noqa: F401, F403 -- bare tile in authored bodies -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget _HIDDEN = 2048 diff --git a/tests/fixtures/placed/tiny_tp_decoder.py b/tests/fixtures/placed/tiny_tp_decoder.py index a42706c2..ed2dcda5 100644 --- a/tests/fixtures/placed/tiny_tp_decoder.py +++ b/tests/fixtures/placed/tiny_tp_decoder.py @@ -12,8 +12,9 @@ from tilefoundry import func, module from tilefoundry.dsl import ConstTensor, Mesh, Tensor, Topology, tf -from tilefoundry.ir.types.shard import Layout, Split, canonical_shard_layout -from tilefoundry.ir.types.shard import Mesh as ShardMesh +from tilefoundry.ir.types import Layout, Split +from tilefoundry.ir.types import Mesh as ShardMesh +from tilefoundry.ir.types.shard_layout import canonical_shard_layout from tilefoundry.runtime import runtime_func, runtime_module from tilefoundry.target import CudaTarget diff --git a/tests/fixtures/placed/tp_all_to_all.py b/tests/fixtures/placed/tp_all_to_all.py index 093805d3..5e455d3a 100644 --- a/tests/fixtures/placed/tp_all_to_all.py +++ b/tests/fixtures/placed/tp_all_to_all.py @@ -12,8 +12,9 @@ from tilefoundry import func, module from tilefoundry.dsl import Mesh, Tensor, Topology, tf -from tilefoundry.ir.types.shard import Layout, Split, canonical_shard_layout -from tilefoundry.ir.types.shard import Mesh as ShardMesh +from tilefoundry.ir.types import Layout, Split +from tilefoundry.ir.types import Mesh as ShardMesh +from tilefoundry.ir.types.shard_layout import canonical_shard_layout from tilefoundry.target import CudaTarget GPUS, CTAS, R, C = 2, 4, 8, 8 diff --git a/tests/fixtures/schedule/__init__.py b/tests/fixtures/schedule/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/schedule/plain/__init__.py b/tests/fixtures/schedule/plain/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/schedule/plain/gemm_8192x17408x5120_cta_grid.py b/tests/fixtures/schedule/plain/gemm_8192x17408x5120_cta_grid.py new file mode 100644 index 00000000..a32f6374 --- /dev/null +++ b/tests/fixtures/schedule/plain/gemm_8192x17408x5120_cta_grid.py @@ -0,0 +1,38 @@ +"""The real matmul as plain HIR: a CTA grid, and no instruction chosen yet. + +M = 8192, K = 5120, N = 17408, split over a 64 x 68 grid of CTAs, one +128 x 256 tile of the output each (``M @ cta.bm``, ``N @ cta.bn``), with K +streamed through shared memory 64 at a time. Nothing here is scheduled and +nothing names a thread: this is the program an author asks ``schedule +candidates`` about before writing the schedule, so what it states is each +CTA's tiles and the storage they move between. +""" +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, Topology, tf +from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops +from tilefoundry.target import CudaTarget + +M = 8192 +K = 5120 +N = 17408 +BM = 128 +BN = 256 +BK = 64 +GM = M // BM +GN = N // BN + + +@module(entry="gemm", target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", GM * GN), Topology("thread", 384))) +class GEMM_8192X17408X5120_CTA_GRID: + @func + def gemm(a: Tensor[(M, K), "bf16"], + b: Tensor[(K, N), "bf16"]) -> Tensor[(M, N), "bf16"]: + with Mesh(("cta",), layout=(GM, GN), names=("bm", "bn")) as cta: + acc = tf.zeros(Tensor[(M @ cta.bm, N @ cta.bn), "f32", "rmem"]) + for k in tile(K, BK): + at = tf.reshard(a[:, k], (M @ cta.bm, BK), "smem") + bt = tf.reshard(b[k, :], (BK, N @ cta.bn), "smem") + part = tf.matmul(at, bt) + acc = acc + tf.reshard(tf.cast(part, "f32"), (M @ cta.bm, N @ cta.bn), "rmem") + return tf.reshard(tf.cast(acc, "bf16"), (M, N), "gmem") diff --git a/tests/fixtures/schedule/plain/gemm_relu_gemm_smem_staged.py b/tests/fixtures/schedule/plain/gemm_relu_gemm_smem_staged.py new file mode 100644 index 00000000..f4191393 --- /dev/null +++ b/tests/fixtures/schedule/plain/gemm_relu_gemm_smem_staged.py @@ -0,0 +1,60 @@ +"""The matmul-ReLU-matmul of `gemm_relu_gemm_tiled`, with both operand windows staged. + +Same computation, same tiling, same accumulator: the one difference is that +each k step names the shared-memory tile it contracts out of instead of letting +the contraction read the global window straight. Staging is a `reshard` and the +instruction that performs it is an offer --- `T.copy_async`, `T.tma` or a plain +`T.copy` --- so what used to be a second implementation of the contraction is a +decision of its own, and ReLU stays the first matmul's tile-level epilogue. +""" +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, Topology, tf +from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops +from tilefoundry.target import CudaTarget + +M = 1024 +N = 2048 +K = 2048 +BM = 32 +BN = 16 +BK = 32 + + +@module(entry="gemm", target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 1), Topology("thread", 512))) +class GEMM_RELU_GEMM_SMEM_STAGED: + @func + def gemm(a: Tensor[(M, K), "bf16"], + b: Tensor[(K, N), "bf16"], + c: Tensor[(N, N), "bf16"]) -> Tensor[(M, N), "bf16"]: + with Mesh(("cta",), layout=(1,), names=("g",)) as _cta: + first = tf.zeros(Tensor[(M, N), "bf16"]) + for m in tile(M, BM): + for n in tile(N, BN): + acc = tf.zeros(Tensor[(BM, BN), "f32", (BM, BN), "rmem"]) + for k in tile(K, BK): + a_s = tf.reshard(a[m, k], (BM, BK), "smem") + b_s = tf.reshard(b[k, n], (BK, BN), "smem") + lhs = tf.cast(a_s, dtype="f32") + rhs = tf.cast(b_s, dtype="f32") + partial = tf.reshard( + tf.matmul(lhs, rhs), (BM, BN), "rmem" + ) + acc = acc + partial + first = tf.insert_slice(first, tf.relu(tf.cast(acc, dtype="bf16")), (m, n)) + + result = tf.zeros(Tensor[(M, N), "bf16"]) + for m2 in tile(M, BM): + for n2 in tile(N, BN): + acc2 = tf.zeros(Tensor[(BM, BN), "f32", (BM, BN), "rmem"]) + for k2 in tile(N, BK): + first_s = tf.reshard(first[m2, k2], (BM, BK), "smem") + c_s = tf.reshard(c[k2, n2], (BK, BN), "smem") + lhs2 = tf.cast(first_s, dtype="f32") + rhs2 = tf.cast(c_s, dtype="f32") + partial2 = tf.reshard( + tf.matmul(lhs2, rhs2), (BM, BN), "rmem" + ) + acc2 = acc2 + partial2 + result = tf.insert_slice(result, tf.cast(acc2, dtype="bf16"), (m2, n2)) + return result diff --git a/tests/fixtures/schedule/plain/gemm_relu_gemm_tiled.py b/tests/fixtures/schedule/plain/gemm_relu_gemm_tiled.py new file mode 100644 index 00000000..0f0d4888 --- /dev/null +++ b/tests/fixtures/schedule/plain/gemm_relu_gemm_tiled.py @@ -0,0 +1,51 @@ +"""Tiled matmul-ReLU-matmul HIR input for the scheduling walkthrough. + +ReLU is the first matmul's epilogue, so every op stays at tile level. +""" +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, Topology, tf +from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops +from tilefoundry.target import CudaTarget + +M = 1024 +N = 2048 +K = 2048 +BM = 32 +BN = 16 +BK = 32 + + +@module(entry="gemm", target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 1), Topology("thread", 512))) +class GEMM_RELU_GEMM_TILED: + @func + def gemm(a: Tensor[(M, K), "bf16"], + b: Tensor[(K, N), "bf16"], + c: Tensor[(N, N), "bf16"]) -> Tensor[(M, N), "bf16"]: + with Mesh(("cta",), layout=(1,), names=("g",)) as _cta: + first = tf.zeros(Tensor[(M, N), "bf16"]) + for m in tile(M, BM): + for n in tile(N, BN): + acc = tf.zeros(Tensor[(BM, BN), "f32", (BM, BN), "rmem"]) + for k in tile(K, BK): + lhs = tf.cast(a[m, k], dtype="f32") + rhs = tf.cast(b[k, n], dtype="f32") + partial = tf.reshard( + tf.matmul(lhs, rhs), (BM, BN), "rmem" + ) + acc = acc + partial + first = tf.insert_slice(first, tf.relu(tf.cast(acc, dtype="bf16")), (m, n)) + + result = tf.zeros(Tensor[(M, N), "bf16"]) + for m2 in tile(M, BM): + for n2 in tile(N, BN): + acc2 = tf.zeros(Tensor[(BM, BN), "f32", (BM, BN), "rmem"]) + for k2 in tile(N, BK): + lhs2 = tf.cast(first[m2, k2], dtype="f32") + rhs2 = tf.cast(c[k2, n2], dtype="f32") + partial2 = tf.reshard( + tf.matmul(lhs2, rhs2), (BM, BN), "rmem" + ) + acc2 = acc2 + partial2 + result = tf.insert_slice(result, tf.cast(acc2, dtype="bf16"), (m2, n2)) + return result diff --git a/tests/fixtures/schedule/plain/gemm_relu_gemm_untiled.py b/tests/fixtures/schedule/plain/gemm_relu_gemm_untiled.py new file mode 100644 index 00000000..b52ce9fa --- /dev/null +++ b/tests/fixtures/schedule/plain/gemm_relu_gemm_untiled.py @@ -0,0 +1,18 @@ +"""Untiled baseline retained for the original placement regression workflow.""" +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, Topology, tf +from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops +from tilefoundry.target import CudaTarget + + +@module(entry="gemm", target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 1), Topology("thread", 512))) +class GEMM_RELU_GEMM_UNTILED: + @func + def gemm(a: Tensor[(1024, 2048), "bf16"], + b: Tensor[(2048, 2048), "bf16"], + c: Tensor[(2048, 2048), "bf16"]) -> Tensor[(1024, 2048), "bf16"]: + with Mesh(("cta",), layout=(1,), names=("g",)) as _cta: + x = tf.matmul(a, b) + y = tf.relu(x) + return tf.matmul(y, c) diff --git a/tests/fixtures/shapes/window_programs.py b/tests/fixtures/shapes/window_programs.py index 676b0faa..9828e904 100644 --- a/tests/fixtures/shapes/window_programs.py +++ b/tests/fixtures/shapes/window_programs.py @@ -4,7 +4,7 @@ from tilefoundry import func, module from tilefoundry.dsl import DimVar, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget WINDOW_SEQ = DimVar("seq", 4, 64) diff --git a/tests/fixtures/tir/async_sync.py b/tests/fixtures/tir/async_sync.py index 73acbf7f..44820b28 100644 --- a/tests/fixtures/tir/async_sync.py +++ b/tests/fixtures/tir/async_sync.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/fixtures/tir/layouts.py b/tests/fixtures/tir/layouts.py index 4c8449dc..29de3eea 100644 --- a/tests/fixtures/tir/layouts.py +++ b/tests/fixtures/tir/layouts.py @@ -2,15 +2,16 @@ from __future__ import annotations -from tilefoundry.ir.types.shard import Layout, ShardLayout, Split -from tilefoundry.ir.types.shard.shard_layout import Broadcast +from tilefoundry.ir.types import Layout, ShardLayout, Split +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.shard_layout import Broadcast def bcast(shape, strides, mesh) -> ShardLayout: """A tile held whole by every instance, so ``local()`` returns that tile.""" return ShardLayout( layout=Layout(shape=shape, strides=strides), - attrs=tuple(Broadcast() for _ in mesh.layout.shape), + attrs=tuple(Broadcast() for _ in flatten(mesh.layout).shape), mesh=mesh, ) diff --git a/tests/fixtures/tir/mma.py b/tests/fixtures/tir/mma.py index 4ef4038c..60a67446 100644 --- a/tests/fixtures/tir/mma.py +++ b/tests/fixtures/tir/mma.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/fixtures/tir/rmsnorm.py b/tests/fixtures/tir/rmsnorm.py index 6e8890d8..1aee5c21 100644 --- a/tests/fixtures/tir/rmsnorm.py +++ b/tests/fixtures/tir/rmsnorm.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import B, Layout, Mesh, Topology +from tilefoundry.ir.types import B, Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/fixtures/tir/square.py b/tests/fixtures/tir/square.py index 2d573735..db78553c 100644 --- a/tests/fixtures/tir/square.py +++ b/tests/fixtures/tir/square.py @@ -4,7 +4,7 @@ from tilefoundry.dsl import DimVar, T, Tensor from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.core.pattern import DimVarRangePat -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget _S = DimVar("S", 1, 256) diff --git a/tests/fixtures/tir/sync.py b/tests/fixtures/tir/sync.py index 43b95be3..3d4f7984 100644 --- a/tests/fixtures/tir/sync.py +++ b/tests/fixtures/tir/sync.py @@ -3,7 +3,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core.kinds import BinaryKind -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/inspection/test_roundtrip.py b/tests/inspection/test_roundtrip.py index 92e75fe6..01098d26 100644 --- a/tests/inspection/test_roundtrip.py +++ b/tests/inspection/test_roundtrip.py @@ -18,7 +18,7 @@ ) _SHARD_IMPORT = ( - "from tilefoundry.ir.types.shard import (\n" + "from tilefoundry.ir.types import (\n" " B, S, P, ComposedLayout, Layout, Mesh, ShardLayout, Topology,\n" ")\n" ) diff --git a/tests/inspection/test_tir_roundtrip.py b/tests/inspection/test_tir_roundtrip.py index 9104df18..efaa0f6a 100644 --- a/tests/inspection/test_tir_roundtrip.py +++ b/tests/inspection/test_tir_roundtrip.py @@ -89,7 +89,7 @@ def test_tir_for_if_and_sync_mesh_forms_roundtrip() -> None: function = import_dsl( "from tilefoundry import prim_func\n" "from tilefoundry.dsl import T, Tensor\n" - "from tilefoundry.ir.types.shard import Layout, Mesh, Topology\n" + "from tilefoundry.ir.types import Layout, Mesh, Topology\n" "from tilefoundry.target import CudaTarget\n\n" "@prim_func(target=CudaTarget('nvidia.h200_sxm'))\n" "def device(a: Tensor[(64,), 'f32'], out: Tensor[(64,), 'f32']):\n" diff --git a/tests/integration/models/deepseek_v4_flash/test_moe.py b/tests/integration/models/deepseek_v4_flash/test_moe.py index 561715be..74645532 100644 --- a/tests/integration/models/deepseek_v4_flash/test_moe.py +++ b/tests/integration/models/deepseek_v4_flash/test_moe.py @@ -10,7 +10,7 @@ from tilefoundry.ir.hir.tensor.reduce import Reduce from tilefoundry.ir.hir.tensor.topk import TopK from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem -from tilefoundry.ir.types.shard import Broadcast, Split +from tilefoundry.ir.types import Broadcast, Split def _walk(expr, seen=None): diff --git a/tests/integration/test_dynamic_cta_tir_handwritten.py b/tests/integration/test_dynamic_cta_tir_handwritten.py index 1b188194..44a1846e 100644 --- a/tests/integration/test_dynamic_cta_tir_handwritten.py +++ b/tests/integration/test_dynamic_cta_tir_handwritten.py @@ -15,8 +15,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import DimVar, T, Tensor from tilefoundry.ir.core.kinds import BinaryKind -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology +from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, TensorType, Topology from tilefoundry.ir.types.storage import StorageKind from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/integration/test_host_launch.py b/tests/integration/test_host_launch.py index 018f30c8..a2a680f5 100644 --- a/tests/integration/test_host_launch.py +++ b/tests/integration/test_host_launch.py @@ -22,8 +22,17 @@ from tilefoundry.ir.tir.stmts import Evaluate, Sequential from tilefoundry.ir.tir.symbol_ref import SymbolRef from tilefoundry.ir.tir.verify import verify_prim_function -from tilefoundry.ir.types import CallableType, DType, TensorType, UnitType -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types import ( + CallableType, + DType, + Layout, + Mesh, + S, + ShardLayout, + TensorType, + Topology, + UnitType, +) from tilefoundry.ir.types.storage import StorageKind from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/ir/test_function_call_typeinfer.py b/tests/ir/test_function_call_typeinfer.py index be78fa59..3b56b063 100644 --- a/tests/ir/test_function_call_typeinfer.py +++ b/tests/ir/test_function_call_typeinfer.py @@ -19,10 +19,15 @@ from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice -from tilefoundry.ir.types import DType, TupleType, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + DType, + TupleType, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.dim import DimMul, DimVar, simplify_dim -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, Split +from tilefoundry.ir.types.shard_layout import Broadcast, Partial, Split from tilefoundry.visitor_registry.contexts import TypeInferContext from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor diff --git a/tests/ir/test_shard_layout_local_shape.py b/tests/ir/test_shard_layout_local_shape.py index 157a918b..300a0ac8 100644 --- a/tests/ir/test_shard_layout_local_shape.py +++ b/tests/ir/test_shard_layout_local_shape.py @@ -12,9 +12,9 @@ import pytest from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.mesh import Mesh, Topology -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.layout import Layout +from tilefoundry.ir.types.mesh import Mesh, Topology +from tilefoundry.ir.types.shard_layout import ( Broadcast, Partial, ShardLayout, diff --git a/tests/ir/test_simplify_dim.py b/tests/ir/test_simplify_dim.py index 939aee0a..a4ca2fa6 100644 --- a/tests/ir/test_simplify_dim.py +++ b/tests/ir/test_simplify_dim.py @@ -15,7 +15,16 @@ from tilefoundry.ir.hir.math.unary import Unary from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice -from tilefoundry.ir.types import DType, TensorType, TupleType +from tilefoundry.ir.types import ( + ComposedLayout, + DType, + Layout, + Mesh, + ShardLayout, + TensorType, + Topology, + TupleType, +) from tilefoundry.ir.types.dim import ( DimAdd, DimFloorDiv, @@ -27,8 +36,7 @@ DimVar, simplify_dim, ) -from tilefoundry.ir.types.shard import ComposedLayout, Layout, Mesh, ShardLayout, Topology -from tilefoundry.ir.types.shard.shard_layout import Broadcast +from tilefoundry.ir.types.shard_layout import Broadcast from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor diff --git a/tests/ir/test_visitor.py b/tests/ir/test_visitor.py index 118a92cc..2547b241 100644 --- a/tests/ir/test_visitor.py +++ b/tests/ir/test_visitor.py @@ -33,9 +33,8 @@ While, ) from tilefoundry.ir.tir.symbol_ref import SymbolRef -from tilefoundry.ir.types import CallableType, DType, TensorType, UnitType -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.mesh import Topology +from tilefoundry.ir.types import CallableType, DType, TensorType, UnitType, make_mesh +from tilefoundry.ir.types.mesh import Topology from tilefoundry.ir.types.storage import StorageKind from tilefoundry.ir.visitor import ( ExprCloner, diff --git a/tests/ir/types/test_local_layout.py b/tests/ir/types/test_local_layout.py index e04124c9..994c5773 100644 --- a/tests/ir/types/test_local_layout.py +++ b/tests/ir/types/test_local_layout.py @@ -15,10 +15,15 @@ import pytest -from tilefoundry.ir.types import DType, make_shard_tensor_type -from tilefoundry.ir.types.shard import Mesh, Topology, make_mesh, shard_layout_of -from tilefoundry.ir.types.shard.local import local_layout, local_layout_and_offset -from tilefoundry.ir.types.shard.shard_layout import Broadcast, Split +from tilefoundry.ir.types import DType, Mesh, Topology, make_mesh, make_shard_tensor_type +from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.shard_layout import ( + Broadcast, + Split, + local_layout, + local_layout_and_offset, + shard_layout_of, +) from tilefoundry.ir.types.utils import local_type_of _GPU, _THREAD = Topology("gpu", 2), Topology("thread", 32) @@ -53,7 +58,7 @@ def _mesh(topology: Topology, extents: tuple[int, ...], names: tuple[str, ...]) def _instances(mesh: Mesh) -> int: """How many programs the mesh spreads a tensor over.""" count = 1 - for extent in mesh.layout.shape: + for extent in flatten(mesh.layout).shape: count *= extent return count diff --git a/tests/ir/types/test_mesh.py b/tests/ir/types/test_mesh.py index e9daaf42..ea4a4e0a 100644 --- a/tests/ir/types/test_mesh.py +++ b/tests/ir/types/test_mesh.py @@ -2,19 +2,14 @@ import pytest -from tilefoundry.ir.types.shard import ( - Layout, - Mesh, - Topology, +from tilefoundry.ir.mesh_scope import ( check_topology, - make_mesh, - product, -) -from tilefoundry.ir.types.shard.layout_algebra import size -from tilefoundry.ir.types.shard.scope_match import ( mesh_scope_matches_required_scope, states_consistent_positions, ) +from tilefoundry.ir.types import Layout, Mesh, Topology, make_mesh +from tilefoundry.ir.types.int_tuple import product +from tilefoundry.ir.types.layout_algebra import size def test_mesh_position_consistency_is_an_explicit_predicate() -> None: @@ -45,14 +40,14 @@ def test_mesh_is_a_frozen_record_without_axis_attributes() -> None: mesh = Mesh(topologies, layout, ("warp", "lane")) assert mesh.topologies is topologies - assert mesh.layout is layout + assert mesh.layout == Layout(shape=((4, 8),), strides=((8, 1),)) assert mesh.names == ("warp", "lane") assert not hasattr(mesh, "topology") assert not hasattr(mesh, "axes") normalized = make_mesh((4, 8), topology="cta") assert normalized.topologies == (Topology("cta", 32),) - assert normalized.layout == Layout(shape=(4, 8), strides=(8, 1)) + assert normalized.layout == Layout(shape=((4, 8),), strides=((8, 1),)) def test_mesh_slice_keeps_the_parent_topologies() -> None: @@ -61,7 +56,7 @@ def test_mesh_slice_keeps_the_parent_topologies() -> None: sliced = mesh[0, :] assert sliced.topologies is mesh.topologies - assert sliced.layout.shape == (1, 32) + assert sliced.layout.shape == ((1, 32),) def test_check_topology_rejects_positions_beyond_a_declared_extent() -> None: diff --git a/tests/ir/types/test_mma_fragment_layouts.py b/tests/ir/types/test_mma_fragment_layouts.py index 5934c3b5..96d2c1c8 100644 --- a/tests/ir/types/test_mma_fragment_layouts.py +++ b/tests/ir/types/test_mma_fragment_layouts.py @@ -12,8 +12,8 @@ from tilefoundry.ir.core import Call, Var from tilefoundry.ir.hir.sharding.reshard import Reshard from tilefoundry.ir.tir.cuda.nn.mma import SM80_16x8x16_F32BF16BF16F32_TN, make_atom -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import ShardLayout, Split, product +from tilefoundry.ir.types import DType, ShardLayout, Split, TensorType +from tilefoundry.ir.types.int_tuple import flatten, product from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.typeinfer import inference_type @@ -47,12 +47,12 @@ def _product(shape: tuple[int, ...]) -> int: def _per_thread_size(sl: ShardLayout) -> int: """Layout product divided by mesh size (= per-thread element count).""" - mesh_size = _product(sl.mesh.layout.shape) + mesh_size = _product(flatten(sl.mesh.layout.shape)) return _product(sl.layout.shape) // mesh_size def _check_split_extents_match_mesh(sl: ShardLayout) -> None: - mesh_shape = sl.mesh.layout.shape + mesh_shape = flatten(sl.mesh.layout.shape) assert len(sl.attrs) == len(mesh_shape), ( f"attrs len {len(sl.attrs)} != mesh rank {len(mesh_shape)}" ) diff --git a/tests/ir/types/test_shard_layout.py b/tests/ir/types/test_shard_layout.py index 67cbdc0c..67c8cd16 100644 --- a/tests/ir/types/test_shard_layout.py +++ b/tests/ir/types/test_shard_layout.py @@ -4,8 +4,8 @@ import pytest -from tilefoundry.ir.types.shard import Layout -from tilefoundry.ir.types.shard import layout_algebra as la +from tilefoundry.ir.types import Layout +from tilefoundry.ir.types import layout_algebra as la @pytest.mark.parametrize( diff --git a/tests/ir/types/test_tensor_type.py b/tests/ir/types/test_tensor_type.py index 1375aa04..abedbabe 100644 --- a/tests/ir/types/test_tensor_type.py +++ b/tests/ir/types/test_tensor_type.py @@ -5,25 +5,21 @@ import pytest from tilefoundry.ir.types import ( - DType, - TensorType, - local_type_of, - make_shard_tensor_type, - numel, - tensor_bytes, -) -from tilefoundry.ir.types.dim import DimVar, ceildiv -from tilefoundry.ir.types.shard import ( Broadcast, ComposedLayout, + DType, Layout, Mesh, Partial, ShardLayout, Split, + TensorType, Topology, make_mesh, + make_shard_tensor_type, ) +from tilefoundry.ir.types.dim import DimVar, ceildiv +from tilefoundry.ir.types.utils import local_type_of, numel, tensor_bytes def test_tensor_type_equality_over_a_dim_var_shape_entry() -> None: diff --git a/tests/models/access_footprint/model.py b/tests/models/access_footprint/model.py index 3c50ad5f..09c7cd49 100644 --- a/tests/models/access_footprint/model.py +++ b/tests/models/access_footprint/model.py @@ -5,7 +5,7 @@ from tests.fixtures.placed.flash_split_k_decode import FlashSplitKDecode from tilefoundry import func, module from tilefoundry.dsl import ConstTensor, Mesh, Tensor, tf -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget _H200 = CudaTarget("nvidia.h200_sxm") diff --git a/tests/models/corpus.py b/tests/models/corpus.py index 1ebc62d8..6041fc29 100644 --- a/tests/models/corpus.py +++ b/tests/models/corpus.py @@ -31,7 +31,7 @@ from tilefoundry.ir.core.module import Module, function_selectors, select from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target.base import Target MODELS_ROOT = Path(__file__).parent diff --git a/tests/models/deepseek_v4_flash/model.py b/tests/models/deepseek_v4_flash/model.py index 20cb14a7..bfe7b97b 100644 --- a/tests/models/deepseek_v4_flash/model.py +++ b/tests/models/deepseek_v4_flash/model.py @@ -56,7 +56,7 @@ from tilefoundry import func, module from tilefoundry.dsl import ConstTensor, ReduceKind, Tensor, tf from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget _MAIN_ROPE: "tuple | None" = None diff --git a/tests/models/fixtures.py b/tests/models/fixtures.py index d590a1ca..7d3b8ac8 100644 --- a/tests/models/fixtures.py +++ b/tests/models/fixtures.py @@ -17,7 +17,7 @@ from __future__ import annotations from tests.models.corpus import TargetFixture -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget, TopologyFacts, TopologyLevelFacts from tilefoundry.target.amx.target import AmxTarget from tilefoundry.target.base import Target diff --git a/tests/models/gemma2_2b/model.py b/tests/models/gemma2_2b/model.py index eec3c4f5..4da43d53 100644 --- a/tests/models/gemma2_2b/model.py +++ b/tests/models/gemma2_2b/model.py @@ -104,7 +104,7 @@ from tilefoundry.dsl import ConstTensor, Tensor, tf # noqa: F401 — tf used by @func bodies from tilefoundry.dsl.tf import * # noqa: F401, F403 from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget diff --git a/tests/models/kimi_linear_48b_a3b/model.py b/tests/models/kimi_linear_48b_a3b/model.py index 065e49f5..9a8b74b5 100644 --- a/tests/models/kimi_linear_48b_a3b/model.py +++ b/tests/models/kimi_linear_48b_a3b/model.py @@ -61,7 +61,7 @@ from tilefoundry.dsl import Tensor, tf # noqa: F401 — tf used by @func bodies from tilefoundry.dsl.tf import * # noqa: F401, F403 — bare op bindings for @func bodies from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget # ── the checkpoint's own configuration class ───────────────────────────────── diff --git a/tests/models/minicpm3_4b/model.py b/tests/models/minicpm3_4b/model.py index 6eb649e0..3863f94b 100644 --- a/tests/models/minicpm3_4b/model.py +++ b/tests/models/minicpm3_4b/model.py @@ -101,7 +101,7 @@ from tilefoundry.dsl import ConstTensor, Tensor, tf # noqa: F401 — tf used by @func bodies from tilefoundry.dsl.tf import * # noqa: F401, F403 — bare op bindings for @func bodies from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget diff --git a/tests/models/qwen2_5_1_5b/model.py b/tests/models/qwen2_5_1_5b/model.py index a0a1a448..18428dcd 100644 --- a/tests/models/qwen2_5_1_5b/model.py +++ b/tests/models/qwen2_5_1_5b/model.py @@ -65,7 +65,7 @@ from tilefoundry.dsl import ConstTensor, Tensor, tf # noqa: F401 — tf used by @func bodies from tilefoundry.dsl.tf import * # noqa: F401, F403 — bare op bindings for @func bodies from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget diff --git a/tests/models/qwen3_1_7b/model.py b/tests/models/qwen3_1_7b/model.py index 874d0830..cf72ac21 100644 --- a/tests/models/qwen3_1_7b/model.py +++ b/tests/models/qwen3_1_7b/model.py @@ -60,7 +60,7 @@ from tilefoundry.dsl import ConstTensor, Mesh, Tensor, tf # noqa: F401 — used by @func bodies from tilefoundry.dsl.tf import * # noqa: F401, F403 — bare op bindings for @func bodies from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget diff --git a/tests/models/qwen3_5_35b_a3b/model.py b/tests/models/qwen3_5_35b_a3b/model.py index e917fa29..491f39e2 100644 --- a/tests/models/qwen3_5_35b_a3b/model.py +++ b/tests/models/qwen3_5_35b_a3b/model.py @@ -21,7 +21,7 @@ from tilefoundry.dsl.tf import * # noqa: F401, F403 -- bare op bindings from tilefoundry.evaluator import to_torch_dtype from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import CudaTarget diff --git a/tests/ops/ir/cost_utils.py b/tests/ops/ir/cost_utils.py index 1001e96e..7807f956 100644 --- a/tests/ops/ir/cost_utils.py +++ b/tests/ops/ir/cost_utils.py @@ -6,8 +6,7 @@ from dataclasses import dataclass, field, replace from tilefoundry.ir.core import Call, Var -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import DType, TensorType, Topology from tilefoundry.visitor_registry.contexts import CostContext, TrafficBytes, TypeInferContext from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor from tilefoundry.visitor_registry.visitors import CostEvaluator diff --git a/tests/ops/ir/test_arange.py b/tests/ops/ir/test_arange.py index 4ad8c6dc..05c663c5 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.types import DType, TensorType +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.shard import Layout, Mesh, Topology, composed 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 = composed((cta, _COORD_MESH)) + current = merge_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 3231a6cf..1391eed6 100644 --- a/tests/ops/ir/test_argmax.py +++ b/tests/ops/ir/test_argmax.py @@ -13,9 +13,15 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.tensor.argmax import ArgMax -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import Layout, ShardLayout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types import ( + DType, + Layout, + ShardLayout, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) +from tilefoundry.ir.types.shard_layout import ( Partial, Split, layout_axis_to_tensor_axis, diff --git a/tests/ops/ir/test_binary.py b/tests/ops/ir/test_binary.py index 6c19623b..4cb9b8a8 100644 --- a/tests/ops/ir/test_binary.py +++ b/tests/ops/ir/test_binary.py @@ -19,10 +19,9 @@ 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_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, Split +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.shard_layout import Broadcast, Partial, Split from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.contexts import TrafficBytes diff --git a/tests/ops/ir/test_cache_update.py b/tests/ops/ir/test_cache_update.py index 4362b4a7..be0bce3a 100644 --- a/tests/ops/ir/test_cache_update.py +++ b/tests/ops/ir/test_cache_update.py @@ -31,11 +31,12 @@ from tilefoundry.ir.hir.tensor.cache_update import CacheUpdate from tilefoundry.ir.types import ( DType, + Topology, + make_mesh, make_shard_tensor_type, make_tensor_type, ) -from tilefoundry.ir.types.shard import Topology, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types.shard_layout import Partial from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.contexts import CostContext, TrafficBytes, TypeInferContext diff --git a/tests/ops/ir/test_cast.py b/tests/ops/ir/test_cast.py index 0715a3ca..f8d70f7c 100644 --- a/tests/ops/ir/test_cast.py +++ b/tests/ops/ir/test_cast.py @@ -17,10 +17,9 @@ 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_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, Split +from tilefoundry.ir.types import DType, make_mesh, make_tensor_type +from tilefoundry.ir.types.layout import Layout +from tilefoundry.ir.types.shard_layout import ShardLayout, Split _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_clamp.py b/tests/ops/ir/test_clamp.py index 637b9e37..95e624df 100644 --- a/tests/ops/ir/test_clamp.py +++ b/tests/ops/ir/test_clamp.py @@ -12,9 +12,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.math.clamp import Clamp -from tilefoundry.ir.types import make_shard_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import make_mesh, 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,)) diff --git a/tests/ops/ir/test_concat.py b/tests/ops/ir/test_concat.py index a63316f0..bee5b0cf 100644 --- a/tests/ops/ir/test_concat.py +++ b/tests/ops/ir/test_concat.py @@ -8,9 +8,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.tensor.concat import Concat -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Split +from tilefoundry.ir.types import DType, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types.shard_layout import Split _F = DType.f32 _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_conv2d.py b/tests/ops/ir/test_conv2d.py index 8eab27c2..8db6b3e0 100644 --- a/tests/ops/ir/test_conv2d.py +++ b/tests/ops/ir/test_conv2d.py @@ -17,9 +17,16 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.conv2d import Conv2D -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import Layout, ShardLayout, Topology, make_mesh -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types import ( + DType, + Layout, + ShardLayout, + Topology, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) +from tilefoundry.ir.types.shard_layout import ( Partial, Split, split_target_axes, diff --git a/tests/ops/ir/test_layer_norm.py b/tests/ops/ir/test_layer_norm.py index 652d914e..ce87d012 100644 --- a/tests/ops/ir/test_layer_norm.py +++ b/tests/ops/ir/test_layer_norm.py @@ -14,9 +14,14 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.layer_norm import LayerNorm -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import ShardLayout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial, Split, split_target_axes +from tilefoundry.ir.types import ( + DType, + ShardLayout, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) +from tilefoundry.ir.types.shard_layout import Partial, Split, split_target_axes _OP = LayerNorm(axis=-1, eps=1e-5) _F = DType.f32 diff --git a/tests/ops/ir/test_local.py b/tests/ops/ir/test_local.py index 847e045d..63d30ea1 100644 --- a/tests/ops/ir/test_local.py +++ b/tests/ops/ir/test_local.py @@ -8,7 +8,7 @@ from tilefoundry.dsl import Mesh, Tensor, Topology, tf from tilefoundry.ir.core import Call, get_metadata from tilefoundry.ir.hir.sharding.local import Local -from tilefoundry.ir.types.shard import Layout +from tilefoundry.ir.types import Layout from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.contexts import TrafficBytes diff --git a/tests/ops/ir/test_matmul.py b/tests/ops/ir/test_matmul.py index a9207970..a872ce78 100644 --- a/tests/ops/ir/test_matmul.py +++ b/tests/ops/ir/test_matmul.py @@ -23,9 +23,14 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.matmul import MatMul -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import Topology, make_mesh -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types import ( + DType, + Topology, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) +from tilefoundry.ir.types.shard_layout import ( Partial, Split, ) diff --git a/tests/ops/ir/test_quant.py b/tests/ops/ir/test_quant.py index 0aabdb9c..cf025fc8 100644 --- a/tests/ops/ir/test_quant.py +++ b/tests/ops/ir/test_quant.py @@ -22,13 +22,15 @@ from tilefoundry.ir.hir.tensor.quant import Quant from tilefoundry.ir.types import ( DType, + Layout, + ShardLayout, TupleType, + make_mesh, make_shard_tensor_type, make_tensor_type, ) from tilefoundry.ir.types.dim import DimFloorDiv, DimVar -from tilefoundry.ir.types.shard import Layout, ShardLayout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.shard_layout import ( Broadcast, Partial, Split, diff --git a/tests/ops/ir/test_relu.py b/tests/ops/ir/test_relu.py index 280f7ee0..d968556f 100644 --- a/tests/ops/ir/test_relu.py +++ b/tests/ops/ir/test_relu.py @@ -10,9 +10,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.relu import ReLU -from tilefoundry.ir.types import make_shard_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types.shard_layout import Partial _OP = ReLU() _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_repeat_interleave.py b/tests/ops/ir/test_repeat_interleave.py index f7fefa23..68232f4d 100644 --- a/tests/ops/ir/test_repeat_interleave.py +++ b/tests/ops/ir/test_repeat_interleave.py @@ -15,9 +15,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.tensor.repeat_interleave import RepeatInterleave -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Split +from tilefoundry.ir.types import DType, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types.shard_layout import Split _F = DType.f32 _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_reshape.py b/tests/ops/ir/test_reshape.py index 0a825988..b718720b 100644 --- a/tests/ops/ir/test_reshape.py +++ b/tests/ops/ir/test_reshape.py @@ -23,10 +23,15 @@ ) from tilefoundry.ir.hir.sharding.reshard import Reshard from tilefoundry.ir.hir.tensor.reshape import Reshape -from tilefoundry.ir.types import make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + Layout, + ShardLayout, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Layout, ShardLayout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types.shard_layout import ( Broadcast, Partial, Split, diff --git a/tests/ops/ir/test_reshard.py b/tests/ops/ir/test_reshard.py index 8e62a7bb..d898c8b5 100644 --- a/tests/ops/ir/test_reshard.py +++ b/tests/ops/ir/test_reshard.py @@ -18,10 +18,9 @@ ) from tilefoundry.dsl.storage import gmem, rmem from tilefoundry.ir.hir.sharding.reshard import Reshard -from tilefoundry.ir.types import make_tensor_type +from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Topology, make_mesh, make_tensor_type from tilefoundry.ir.types.dim import DimMul, DimVar, simplify_dim -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Topology, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Split +from tilefoundry.ir.types.shard_layout import Split from tilefoundry.ir.types.storage import StorageKind diff --git a/tests/ops/ir/test_rms_norm.py b/tests/ops/ir/test_rms_norm.py index a880c3df..849afa47 100644 --- a/tests/ops/ir/test_rms_norm.py +++ b/tests/ops/ir/test_rms_norm.py @@ -19,9 +19,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.rms_norm import RMSNorm -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import DType, make_mesh, 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,)) diff --git a/tests/ops/ir/test_rope.py b/tests/ops/ir/test_rope.py index 4c909120..511aa09c 100644 --- a/tests/ops/ir/test_rope.py +++ b/tests/ops/ir/test_rope.py @@ -17,9 +17,14 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.rope import RoPE -from tilefoundry.ir.types import DType, TupleType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import ( + DType, + 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,)) diff --git a/tests/ops/ir/test_shape_metadata.py b/tests/ops/ir/test_shape_metadata.py index a7eb9329..09bcfd02 100644 --- a/tests/ops/ir/test_shape_metadata.py +++ b/tests/ops/ir/test_shape_metadata.py @@ -13,7 +13,7 @@ from tilefoundry.ir.hir.tensor.rank import Rank from tilefoundry.ir.hir.tensor.shape_of import ShapeOf from tilefoundry.ir.types import TensorType, make_tensor_type -from tilefoundry.ir.types.shard.layout import EMPTY_LAYOUT +from tilefoundry.ir.types.layout import EMPTY_LAYOUT from tilefoundry.ir.types.storage import StorageKind from tilefoundry.ir.visitor import collect_exprs from tilefoundry.visitor_registry.contexts import TrafficBytes diff --git a/tests/ops/ir/test_sigmoid.py b/tests/ops/ir/test_sigmoid.py index 0a4da3cd..41458f05 100644 --- a/tests/ops/ir/test_sigmoid.py +++ b/tests/ops/ir/test_sigmoid.py @@ -8,9 +8,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.sigmoid import Sigmoid -from tilefoundry.ir.types import make_shard_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types.shard_layout import Partial _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_silu.py b/tests/ops/ir/test_silu.py index 766242eb..b084943f 100644 --- a/tests/ops/ir/test_silu.py +++ b/tests/ops/ir/test_silu.py @@ -12,9 +12,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.silu import Silu -from tilefoundry.ir.types import make_shard_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types.shard_layout import Partial _OP = Silu() _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_slice.py b/tests/ops/ir/test_slice.py index 58a8f0e8..a51adc4d 100644 --- a/tests/ops/ir/test_slice.py +++ b/tests/ops/ir/test_slice.py @@ -17,10 +17,17 @@ from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.tensor.slice import Slice, slice_size from tilefoundry.ir.isl_interop import normalize_dim -from tilefoundry.ir.types import DType, TupleType, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types import ( + ComposedLayout, + DType, + Layout, + TupleType, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) from tilefoundry.ir.types.dim import DimMul, DimVar, simplify_dim -from tilefoundry.ir.types.shard import ComposedLayout, Layout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, Split, shard_layout_of +from tilefoundry.ir.types.shard_layout import ShardLayout, Split, shard_layout_of from tilefoundry.visitor_registry.contexts import CostContext, TrafficBytes from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor from tilefoundry.visitor_registry.visitors import CostEvaluator diff --git a/tests/ops/ir/test_softmax.py b/tests/ops/ir/test_softmax.py index 5881d760..7c91562a 100644 --- a/tests/ops/ir/test_softmax.py +++ b/tests/ops/ir/test_softmax.py @@ -8,9 +8,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.softmax import SoftMax -from tilefoundry.ir.types import make_shard_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types.shard_layout import Partial def test_softmax_typeinfer_partial_input_errors(): diff --git a/tests/ops/ir/test_softplus.py b/tests/ops/ir/test_softplus.py index a5afe1d4..e792bf8f 100644 --- a/tests/ops/ir/test_softplus.py +++ b/tests/ops/ir/test_softplus.py @@ -16,9 +16,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.math.softplus import Softplus -from tilefoundry.ir.types import make_shard_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types.shard_layout import Partial _OP = Softplus() _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_split.py b/tests/ops/ir/test_split.py index 288a9449..85777742 100644 --- a/tests/ops/ir/test_split.py +++ b/tests/ops/ir/test_split.py @@ -17,10 +17,17 @@ from tilefoundry.evaluator.registry import eval_registry from tilefoundry.evaluator.value import TensorValue, TupleValue from tilefoundry.ir.hir.tensor.split import Split -from tilefoundry.ir.types import DType, TupleType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import Layout, ShardLayout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Split as SplitAttr -from tilefoundry.ir.types.shard.shard_layout import shard_layout_local_shape +from tilefoundry.ir.types import ( + DType, + Layout, + ShardLayout, + TupleType, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) +from tilefoundry.ir.types.shard_layout import Split as SplitAttr +from tilefoundry.ir.types.shard_layout import shard_layout_local_shape from tilefoundry.visitor_registry.contexts import TrafficBytes CASES = [ diff --git a/tests/ops/ir/test_stack.py b/tests/ops/ir/test_stack.py index 27fc4016..0b4e2b5a 100644 --- a/tests/ops/ir/test_stack.py +++ b/tests/ops/ir/test_stack.py @@ -16,9 +16,16 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.tensor.stack import Stack -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import Layout, Partial, ShardLayout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Split, split_target_axes +from tilefoundry.ir.types import ( + DType, + Layout, + Partial, + ShardLayout, + make_mesh, + make_shard_tensor_type, + make_tensor_type, +) +from tilefoundry.ir.types.shard_layout import Split, split_target_axes from tilefoundry.visitor_registry.contexts import TrafficBytes diff --git a/tests/ops/ir/test_tanh.py b/tests/ops/ir/test_tanh.py index 6bdff3f5..cec7be7d 100644 --- a/tests/ops/ir/test_tanh.py +++ b/tests/ops/ir/test_tanh.py @@ -12,9 +12,8 @@ run_typeinfer_case, ) from tilefoundry.ir.hir.nn.tanh import Tanh -from tilefoundry.ir.types import make_shard_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import make_mesh, make_shard_tensor_type +from tilefoundry.ir.types.shard_layout import Partial _OP = Tanh() _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_topk.py b/tests/ops/ir/test_topk.py index f7f23e54..b6a84f5c 100644 --- a/tests/ops/ir/test_topk.py +++ b/tests/ops/ir/test_topk.py @@ -32,13 +32,14 @@ from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.types import ( DType, + Layout, TupleType, + make_mesh, make_shard_tensor_type, make_tensor_type, ) from tilefoundry.ir.types.dim import DimVar, dim_min -from tilefoundry.ir.types.shard import Layout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, ShardLayout, Split +from tilefoundry.ir.types.shard_layout import Broadcast, Partial, ShardLayout, Split from tilefoundry.visitor_registry.contexts import TypeInferContext from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor diff --git a/tests/ops/ir/test_transpose.py b/tests/ops/ir/test_transpose.py index 791c1be6..f7c9cac6 100644 --- a/tests/ops/ir/test_transpose.py +++ b/tests/ops/ir/test_transpose.py @@ -15,9 +15,8 @@ raw_shard_tensor_type, ) from tilefoundry.ir.hir.tensor.transpose import Transpose -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import Layout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import ( +from tilefoundry.ir.types import DType, Layout, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types.shard_layout import ( Broadcast, ShardLayout, Split, @@ -29,12 +28,13 @@ _T10 = Transpose(perm=(1, 0)) -def test_plain_input_permutes_its_layout_when_one_is_stated(): +def test_plain_input_permutes_its_layout(): + """An unstated layout is the C order it stands for, and is permuted as one.""" source = make_tensor_type((16, 8), DType.bf16, layout=Layout(shape=(16, 8), strides=(8, 1))) ty = infer_call(_T10, source) assert ty.layout == Layout(shape=(8, 16), strides=(1, 8)) - assert infer_call(_T10, make_tensor_type((16, 8), DType.bf16)).layout is None + assert infer_call(_T10, make_tensor_type((16, 8), DType.bf16)).layout == ty.layout def test_factorized_split_reorders_subaxes(): diff --git a/tests/ops/ir/test_unary.py b/tests/ops/ir/test_unary.py index 8013e803..dbbbb77e 100644 --- a/tests/ops/ir/test_unary.py +++ b/tests/ops/ir/test_unary.py @@ -20,9 +20,8 @@ ) from tilefoundry.ir.core.kinds import UnaryKind from tilefoundry.ir.hir.math.unary import Unary -from tilefoundry.ir.types import DType, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import make_mesh -from tilefoundry.ir.types.shard.shard_layout import Partial +from tilefoundry.ir.types import DType, make_mesh, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types.shard_layout import Partial from tilefoundry.visitor_registry.contexts import TrafficBytes _NEG = Unary(kind=UnaryKind.NEG) diff --git a/tests/ops/ir/test_where.py b/tests/ops/ir/test_where.py index facda266..38ba1bdc 100644 --- a/tests/ops/ir/test_where.py +++ b/tests/ops/ir/test_where.py @@ -15,9 +15,8 @@ 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, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.shard import Layout, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Split +from tilefoundry.ir.types import DType, Layout, make_mesh, 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 diff --git a/tests/ops/ir/typeinfer_utils.py b/tests/ops/ir/typeinfer_utils.py index 0c6b2146..06bfe775 100644 --- a/tests/ops/ir/typeinfer_utils.py +++ b/tests/ops/ir/typeinfer_utils.py @@ -15,8 +15,8 @@ from tilefoundry.ir.core import Call, Var from tilefoundry.ir.core.errors import VerifyError from tilefoundry.ir.types import DType, TensorType, TupleType -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, Split, shard_layout_local_shape +from tilefoundry.ir.types.layout import Layout +from tilefoundry.ir.types.shard_layout import ShardLayout, Split, shard_layout_local_shape from tilefoundry.visitor_registry.contexts import TypeInferContext from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor diff --git a/tests/ops/tir/cuda/test_mbarrier.py b/tests/ops/tir/cuda/test_mbarrier.py index 318a1efb..006e77d9 100644 --- a/tests/ops/tir/cuda/test_mbarrier.py +++ b/tests/ops/tir/cuda/test_mbarrier.py @@ -21,8 +21,7 @@ from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.stmts import Evaluate, Return, Sequential from tilefoundry.ir.tir.verify import verify_prim_function -from tilefoundry.ir.types import DType, make_tensor_type -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import DType, Layout, Mesh, Topology, make_tensor_type from tilefoundry.target import CpuTarget, CudaTarget _SMEM_BAR = make_tensor_type((1,), DType.from_name("i64"), storage="smem") diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 77d0b14d..431630e7 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -13,8 +13,8 @@ from tests.fixtures.tir.mma import MmHandwritten from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Topology -from tilefoundry.ir.types.shard.shard_layout import Broadcast +from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.shard_layout import Broadcast from tilefoundry.target import CpuTarget, CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") diff --git a/tests/ops/tir/cuda/test_swizzle.py b/tests/ops/tir/cuda/test_swizzle.py index ca42d239..f6d8bb4b 100644 --- a/tests/ops/tir/cuda/test_swizzle.py +++ b/tests/ops/tir/cuda/test_swizzle.py @@ -16,15 +16,7 @@ from tilefoundry.dsl import T, Tensor from tilefoundry.inspection import as_script from tilefoundry.ir.core.kinds import BinaryKind -from tilefoundry.ir.types.shard import ( - ComposedLayout, - Layout, - Mesh, - ShardLayout, - Split, - Swizzle, - Topology, -) +from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, ShardLayout, Split, Swizzle, Topology from tilefoundry.target import CpuTarget, CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") diff --git a/tests/ops/tir/cuda/test_tma.py b/tests/ops/tir/cuda/test_tma.py index e4a558ab..f27c299c 100644 --- a/tests/ops/tir/cuda/test_tma.py +++ b/tests/ops/tir/cuda/test_tma.py @@ -17,9 +17,8 @@ from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.stmts import Evaluate, Return, Sequential from tilefoundry.ir.tir.verify import verify_prim_function -from tilefoundry.ir.types import DType, make_tensor_type -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Topology -from tilefoundry.ir.types.shard.shard_layout import Broadcast +from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Topology, make_tensor_type +from tilefoundry.ir.types.shard_layout import Broadcast from tilefoundry.target import CpuTarget, CudaTarget _BAR = make_tensor_type((1,), DType.from_name("i64"), storage="smem") diff --git a/tests/ops/tir/test_copy.py b/tests/ops/tir/test_copy.py index 3095c4a7..2215d437 100644 --- a/tests/ops/tir/test_copy.py +++ b/tests/ops/tir/test_copy.py @@ -13,7 +13,7 @@ from tests.fixtures.tir.layouts import broadcast_run, split_pairs, split_rows, split_short_rows from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, Topology +from tilefoundry.ir.types import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") diff --git a/tests/ops/tir/test_dot.py b/tests/ops/tir/test_dot.py index f6d7a279..234125d5 100644 --- a/tests/ops/tir/test_dot.py +++ b/tests/ops/tir/test_dot.py @@ -17,9 +17,8 @@ from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.stmts import Evaluate, Return, Sequential from tilefoundry.ir.tir.verify import verify_prim_function -from tilefoundry.ir.types import DType, make_tensor_type -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology -from tilefoundry.ir.types.shard.shard_layout import Broadcast +from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology, make_tensor_type +from tilefoundry.ir.types.shard_layout import Broadcast from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/ops/tir/test_elementwise.py b/tests/ops/tir/test_elementwise.py index 14ab4002..e2a6d0c1 100644 --- a/tests/ops/tir/test_elementwise.py +++ b/tests/ops/tir/test_elementwise.py @@ -18,7 +18,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core.kinds import BinaryKind -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology +from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Split, Topology from tilefoundry.target import CpuTarget, CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") diff --git a/tests/ops/tir/test_launch.py b/tests/ops/tir/test_launch.py index 467f711a..0ce415ca 100644 --- a/tests/ops/tir/test_launch.py +++ b/tests/ops/tir/test_launch.py @@ -11,10 +11,8 @@ from tilefoundry.ir.tir.stmts import Evaluate, Sequential from tilefoundry.ir.tir.symbol_ref import SymbolRef from tilefoundry.ir.tir.verify import verify_module -from tilefoundry.ir.types import ( - callable_type_for_prim_function, -) -from tilefoundry.ir.types.shard import Layout +from tilefoundry.ir.types import Layout +from tilefoundry.ir.types.callable_type import callable_type_for_prim_function from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/ops/tir/test_reduce.py b/tests/ops/tir/test_reduce.py index c399b66a..76e00063 100644 --- a/tests/ops/tir/test_reduce.py +++ b/tests/ops/tir/test_reduce.py @@ -13,8 +13,8 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core.kinds import ReduceKind -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology -from tilefoundry.ir.types.shard.shard_layout import Broadcast +from tilefoundry.ir.types import Layout, Mesh, ShardLayout, Split, Topology +from tilefoundry.ir.types.shard_layout import Broadcast from tilefoundry.target import CpuTarget, CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") diff --git a/tests/ops/tir/test_sync.py b/tests/ops/tir/test_sync.py index 77fffb7f..4bee2d96 100644 --- a/tests/ops/tir/test_sync.py +++ b/tests/ops/tir/test_sync.py @@ -12,7 +12,7 @@ from tests.fixtures.tir.sync import SyncSquare from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types import Layout, Mesh, S, ShardLayout, Topology from tilefoundry.target import CpuTarget, CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") diff --git a/tests/parser/test_calls.py b/tests/parser/test_calls.py index 707ac3ad..d78f8414 100644 --- a/tests/parser/test_calls.py +++ b/tests/parser/test_calls.py @@ -20,8 +20,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.types import DType -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import DType, Topology from tilefoundry.parser import ParseError from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/passes/test_host_entry.py b/tests/passes/test_host_entry.py index ed4452b9..91257ae6 100644 --- a/tests/passes/test_host_entry.py +++ b/tests/passes/test_host_entry.py @@ -15,7 +15,7 @@ from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.tir.launch import Launch from tilefoundry.ir.tir.stmts import Evaluate -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types import Layout, Mesh, S, ShardLayout, Topology from tilefoundry.passes.transforms import insert_default_host_entry from tilefoundry.target import CpuTarget, CudaTarget diff --git a/tests/runtime/resource/test_placed_launch.py b/tests/runtime/resource/test_placed_launch.py index e4afa7c7..c885d622 100644 --- a/tests/runtime/resource/test_placed_launch.py +++ b/tests/runtime/resource/test_placed_launch.py @@ -23,7 +23,7 @@ ROWS_PER_CARD, GpuPlacedRows, ) -from tilefoundry.ir.types.shard import Placement +from tilefoundry.ir.types import Placement from tilefoundry.runtime.resource import DictResource diff --git a/tests/runtime/resource/test_weight_loading.py b/tests/runtime/resource/test_weight_loading.py index 2e388938..53d691f5 100644 --- a/tests/runtime/resource/test_weight_loading.py +++ b/tests/runtime/resource/test_weight_loading.py @@ -23,7 +23,7 @@ TinyTPDecoderLM, TinyTPDecoderLMTwin, ) -from tilefoundry.ir.types.shard import Placement +from tilefoundry.ir.types import Placement from tilefoundry.runtime.resource import DictResource, SafetensorsResource diff --git a/tests/schedule/__init__.py b/tests/schedule/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/schedule/test_fixtures.py b/tests/schedule/test_fixtures.py new file mode 100644 index 00000000..441242c6 --- /dev/null +++ b/tests/schedule/test_fixtures.py @@ -0,0 +1,36 @@ +"""The scheduling corpus's plain programs, read the way a scheduler reads them. + +These are the programs an author states before choosing an instruction: a tiled +matmul chain, the same chain with both operand windows staged, its untiled +baseline, and the real matmul over a CTA grid. What each one is here to witness +is that it parses, holds together, and can be measured -- the three questions +anything downstream asks before it offers a schedule at all. +""" + +from __future__ import annotations + +import importlib + +import pytest + +from tilefoundry.analysis.api import analyze +from tilefoundry.analysis.check import check_program + +PLAIN = ( + "gemm_8192x17408x5120_cta_grid", + "gemm_relu_gemm_smem_staged", + "gemm_relu_gemm_tiled", + "gemm_relu_gemm_untiled", +) + + +@pytest.mark.parametrize("name", PLAIN) +def test_plain_program_is_analyzable(name: str) -> None: + module = importlib.import_module(f"tests.fixtures.schedule.plain.{name}") + program = next( + value for value in vars(module).values() if type(value).__name__ == "Module" + ) + entry = next(function for function in program.functions if function.name == "gemm") + check_program(program, entry) + result = analyze(program, entry, analysis=("memory", "performance")) + assert result.metadata_types diff --git a/tests/scripts/test_no_machine_paths_lint.py b/tests/scripts/test_no_machine_paths_lint.py index 508f3f39..6034a3de 100644 --- a/tests/scripts/test_no_machine_paths_lint.py +++ b/tests/scripts/test_no_machine_paths_lint.py @@ -53,7 +53,7 @@ def lint(): 'path = Path(__file__).parent / "model" / "decoder_layer.py"', 'CKPT_DIR = os.environ["TILEFOUNDRY_QWEN35_CKPT"]', 'shard = "model-00001-of-00001.safetensors"', - "from tilefoundry.ir.types.shard import Layout", + "from tilefoundry.ir.types import Layout", 'doc = "/usr/share/doc"', f'note = "put it under {_HOME}//checkouts"', f"# see {_HOME}/someone/notes.md # no-machine-path: allow", diff --git a/tests/target/test_amx_target.py b/tests/target/test_amx_target.py index b066db62..6b0b933a 100644 --- a/tests/target/test_amx_target.py +++ b/tests/target/test_amx_target.py @@ -9,9 +9,8 @@ import pytest -from tilefoundry.ir.types import DType +from tilefoundry.ir.types import DType, Topology from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Topology from tilefoundry.target import ( AmxTarget, TopologyFacts, diff --git a/tests/target/test_target.py b/tests/target/test_target.py index 4cca2776..1db1a64d 100644 --- a/tests/target/test_target.py +++ b/tests/target/test_target.py @@ -24,7 +24,7 @@ from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.stmts import Sequential -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import Topology from tilefoundry.target import ( CpuTarget, CudaTarget, diff --git a/tests/target/test_target_facts.py b/tests/target/test_target_facts.py index bfb70bb6..810be1c2 100644 --- a/tests/target/test_target_facts.py +++ b/tests/target/test_target_facts.py @@ -11,8 +11,7 @@ MemoryHierarchyFacts, ThroughputFacts, ) -from tilefoundry.ir.types import DType -from tilefoundry.ir.types.shard import Topology +from tilefoundry.ir.types import DType, Topology from tilefoundry.target import ( AmxTarget, CudaTarget,