diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index ccf0d1fa..3cf613fb 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -130,8 +130,8 @@ explicit wrapper request. Module output computes imports from names actually used by the rendered program and MUST remain lint-clean. - constraints: - - Module input MUST emit every HIR Function and preserve shared `Mesh` and - `Topology` definitions before the class. Mixed HIR/TIR Modules MUST be + - Module input MUST emit every HIR Function. Meshes used by a function are + written at their scope boundary, and mixed HIR/TIR Modules MUST be rejected. The printer MUST emit a Module's whole tree: each nested Module prints as a @@ -143,9 +143,9 @@ Target prints as the `@module(target=...)` argument and a declared hierarchy as the `@module(topologies=...)` argument ([parser §3](./parser.md#3-implementation-overview)). Every dimension referenced only by a declared topology expression MUST still -be emitted in the dimension prelude. A topology `ShapeDim` MUST use the same -DSL expression text as tensor and Mesh geometry, including public constructors -such as `ceildiv`, so importing restores the same expression tree. +be emitted in the file header. A topology `ShapeDim` MUST use the same DSL +expression text as tensor and Mesh geometry, including public constructors such +as `ceildiv`, so importing restores the same expression tree. - constraints: - The decorator MUST print in its called form, `@module()` included. A bare @@ -178,42 +178,41 @@ executed. A Target subclass with a different constructor customizes ordinary DSL text forms for tensor / layout / shard annotations are owned by [parser](./parser.md). The printer reuses those forms only when they round-trip without losing mesh / layout / storage information; -otherwise it falls back to the verbose `ShardLayout(...)`. Printer -output supports two modes derived from the same pretty-print core: - -- `canonical` — round-trippable text used by `as_script()`, pass - dumps, and viewer detail `code` blocks: the `Tensor[...]` form of - [parser §2.1](./parser.md#21-syntax) (storage as the string - slot, `gmem` omitted). -- `compact` — abbreviated, **display-only / non-round-trip** text for - summaries / labels: `dtype[shape] {value-state?} @storage`. It inlines - what it can (a split into the shape, a `Partial` into the `{...}` - suffix) and falls back to the canonical form when a layout cannot be - rendered compactly. - -Both modes MUST agree on semantics; only the level of detail differs. +otherwise it falls back to the verbose `ShardLayout(...)`. + +A `ShardLayout` over a plain `Layout` whose `Mesh` has named axes MUST use the +placement sugar of [parser §2.1](./parser.md#21-syntax), in both type slots and +op-attribute slots, when every mesh identifier in that sugar has an explicit +scope binding. A binding is either a `with as ` region or the +function's own execution domain. A mesh merely restated in another expression +is not a binding. Without such a binding, the printer MUST use the verbose +`ShardLayout(...)` form rather than inventing a name. + +Placement sugar states the layout's own dimensions with each `Split` written on +the dimension it divides, adds the stride tuple whenever the layout has one, +and states the remaining mesh axes in a `{axis @ ...}` set. Because the parser +reads an unstated mesh axis as `Broadcast`, the set carries every `Partial` and +carries `Broadcast` only when no `Split` or `Partial` would otherwise name the +mesh. + +The printer has one type-text surface: the canonical `Tensor[...]` form of +[parser §2.1](./parser.md#21-syntax), used by `as_script()`, comments, DOT, +the viewer graph, and viewer detail panels. Comments may flatten this text onto +one physical line, but MUST NOT change its syntax or semantics. The same +printer visitor supplies type, dimension, layout, and shard-attribute text; +there is no separate compact, display-only type language. + The meaning of `Split` / `Partial` / `Broadcast` is owned by -[shard](./shard.md); these forms define only render syntax. - -The same-line type annotation `show_types` appends is the `canonical` form on -one physical line, rendered through the same mesh name map -([§2.5](#25-mesh-name-map)) as the signature and the prelude: an annotated -layout MUST name the hoisted mesh rather than restate it, and a `Tuple[...]` -annotation MUST name it in every field. The verbose `ShardLayout(...)` fallback -is unchanged — a mesh with no named axes, or a layout the sugar cannot express, -still renders verbose, so no annotation loses information. The annotation is -**display-only** ([§2.7](#27-round-trip-contract)); what round-trips is the -emitted code, not its comments. - -Canonical DType text is the descriptor's `name`. Tensor annotations and DType -op attributes MUST emit that name as a quoted DSL string. Compact labels MAY -omit the quotes, but MUST NOT use the descriptor's raw `repr()`. +[shard](./shard.md); these forms define only render syntax. Canonical DType text +is the descriptor's `name`. Tensor annotations and DType op attributes MUST +emit that name as a quoted DSL string. ### 2.4 Pretty-print / debug display contract -Pretty print is the core presentation layer. Sugar, debug dumps, and -viewer type/value text reuse the same DSL text forms in [§2.3](#23-dsl-text-forms). That keeps -round-trippable source, labels, and detail panes semantically aligned. +Pretty print is the core presentation layer. Sugar, debug dumps, DOT, and +viewer type/value text reuse the same canonical DSL text forms in +[§2.3](#23-dsl-text-forms). That keeps round-trippable source, labels, and +detail panes semantically aligned. - op attributes that are `DType`, `TensorType`, `Layout`, or `ShardLayout` are rendered through the [§2.3](#23-dsl-text-forms) printer; `DType` uses its canonical name and these @@ -225,12 +224,24 @@ printing (for example, choosing stable mesh names across a whole function) must use an explicit pretty-printer API rather than relying on no-argument `repr()`. -### 2.5 Mesh name map +### 2.5 Mesh bindings and file header + +The print context records imports and declarations while the printer visits the +program. After the body has been visited, the context emits the file header; +it MUST contain only imports and `DimVar` declarations reached by the output. +There is no module-level mesh hoist or global mesh name map. + +A mesh identifier in placement sugar MUST be a lexical binding visible at the +point represented by the text: a `with as ` region, or the +function's own execution domain. The function execution domain is represented +as `@func(mesh=...)`, so its name is available to signature annotations. A +nested region writes its mesh expression at the `with` boundary and uses that +binding in its body. A named constant or an out-of-scope mesh does not create a +binding; types that refer to one use the verbose form. -The printer collects unique `Mesh` objects from all `ShardLayout` -references in the function (params, return type, body `Reshard` ops) -and assigns variable names from the first declared topology's name. Mesh -definitions are emitted in the module prelude / standalone header. +Analysis annotations use the same type text for value-producing statements, but +a structural `LoopRegion` or `MeshRegion` line MUST NOT restate the aggregate +type of its body. The region may still carry selected analysis metadata. ### 2.6 Specialization printing @@ -460,14 +471,10 @@ graph; an id that was collapsed away returns 404. elements); `Tuple` bundles its elements. Op attributes that are constants / types render through the [§2.4](#24-pretty-print--debug-display-contract) pretty-print, never raw `repr`. -- **Type text.** Graph labels use the [§2.3](#23-dsl-text-forms) **compact** pretty mode - (`bf16[4 @ trd.l, 64] {trd.t @ P("sum")} @smem`) with inline split / - DimVar / storage colour; the detail panel uses the [§2.3](#23-dsl-text-forms) **canonical** - mode (`Tensor[(4, 64), "f32", ((4 @ trd.l, 64), {trd.t @ P("sum")}), - "smem"]`). `Reshard` / layout attrs render through the same core (never - raw `repr`). DimVar is a single token-class colour; - storage classes draw from an ordered pool, and an unknown memory level - hashes stably into the pool's spare slots rather than going colourless. +- **Type text.** Graph labels and detail panels use the canonical type text from + [§2.3](#23-dsl-text-forms), including `Tensor[...]` and placement sugar when + a scope binding is available. `Reshard` / layout attrs render through the + same visitor (never raw `repr`). ### 3.4 Interaction contract diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 8e7b0064..bcdbee62 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -94,13 +94,20 @@ dim-expr ::= integer-literal | dim-expr ('+' | '-' | '*' | '//' | '%') dim-expr | (identifier | primary '.' identifier) '(' (dim-expr (',' dim-expr)*)? ')' -placed-shape ::= '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | +layout-dims ::= '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | dim-expr))*)? ')' +layout-strides ::= '(' (dim-expr (',' dim-expr)*)? ')' +value-states ::= '{' mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')') (',' + mesh-axis '@' ('B' '(' ')' | 'P' '(' string-literal ')'))* '}' +placed-layout ::= '(' layout-dims ',' layout-strides ',' value-states ')' + | '(' layout-dims ',' layout-strides ')' + | '(' layout-dims ',' value-states ')' + | layout-dims shape ::= '(' (dim-expr (',' dim-expr)*)? ')' | identifier | primary '.' identifier -tensor-shape-layout ::= placed-shape +tensor-shape-layout ::= placed-layout | shape dtype ::= string-literal | primary @@ -133,15 +140,10 @@ expression ::= literal | subscript call ::= expression '(' ((expression | keyword-name '=' expression) (',' (expression | keyword-name '=' expression))*)? ')' -explicit-layout ::= '(' (tensor-shape-layout | shape) ',' shape ')' -placed-layout ::= '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | - dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | - mesh-axis) | dim-expr))*)? ')' plain-layout ::= '(' (dim-expr (',' dim-expr)*)? ')' layout ::= None | primary | call - | explicit-layout | placed-layout | plain-layout storage ::= string-literal @@ -242,10 +244,8 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | --- | --- | --- | --- | --- | | binary_expression, matmul_expression, op_call, unary_expression | expression, slice_endpoint, subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | | binary_expression, matmul_expression, op_call, unary_expression | expression, slice_endpoint, subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | -| dim_expr | dim_expr, layout_extent, layout_shape, tensor_dim_expr, tensor_optional_slot, tensor_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | +| dim_expr | dim_expr, layout_extent, tensor_dim_expr, tensor_optional_slot, tensor_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | | dtype | tensor_dtype | CanonicalDTypeRule | A dtype must resolve to a canonical DType. | src/tilefoundry/parser/ast_pattern.py | -| explicit_layout, layout, placed_layout, plain_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | -| explicit_layout, layout, placed_layout, plain_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | | function | function | FunctionDialectRule | A function kind and constructed value must agree with the active dialect. | src/tilefoundry/parser/pattern_nodes.py | | function | function | FunctionRegistrationRule | A validated function must be registered exactly once in its owning scope. | src/tilefoundry/parser/pattern_nodes.py | | function | function | FunctionReturnCompatibilityRule | A HIR body with a return annotation must satisfy that annotation; a dispatch prototype must declare one, and each variant body must satisfy the prototype return contract. | src/tilefoundry/parser/pattern_nodes.py | @@ -253,17 +253,31 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | function | function | FunctionSignatureRule | A function must construct an ordered parameter tuple. | src/tilefoundry/parser/pattern_nodes.py | | if, while | loop_statement, statement | TirOnlyStatementRule | A TIR-only statement must appear in a prim_func. | src/tilefoundry/parser/pattern_nodes.py | | index_slice | subscript_index | TileWindowSliceBoundRule | A tile window cannot be used as a slice bound. | src/tilefoundry/parser/pattern_nodes.py | +| layout, plain_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| layout, plain_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | | module | module_finalization | ModuleFinalizationRule | A module declaration must contain valid unique members and a resolvable entry. | src/tilefoundry/parser/ast_pattern.py | | module | module_function | ModuleFunctionRegistrationRule | A validated module function must be recorded in declaration order. | src/tilefoundry/parser/ast_pattern.py | | module | module_function | ModuleFunctionValidationRule | A module function must satisfy its root, variant, or converter role before mutation. | src/tilefoundry/parser/ast_pattern.py | | op_call | expression, slice_endpoint, subscript_index | CallVariadicInputFormRule | A variadic call must use one explicit list, tuple, or supported static list comprehension. | src/tilefoundry/parser/pattern_nodes.py | -| placed_shape | layout_shape, tensor_shape | PlacedShapeRule | Placement sugar in a shape slot states both a shape and a layout. | src/tilefoundry/parser/ast_pattern.py | -| shape | layout_shape, layout_strides, tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | +| placed_layout | tensor_optional_slot, tensor_shape | LayoutStrideRankRule | A stated stride tuple must have the rank of the layout it addresses. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | MeshAxisBoundOnceRule | A placement binds each mesh axis at most once. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementAnswerRule | Placement sugar states both the shape as written and the layout it implies. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementConstructionRule | A placement must construct a valid shard layout. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementLevelRule | A placement's meshes cannot name the same topology level. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | tensor_optional_slot, tensor_shape | PlacementMeshResolutionRule | A placement's mesh must be a lexical mesh binding. | src/tilefoundry/parser/pattern_nodes.py | +| shape | tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | | storage | tensor_optional_slot | StorageValueRule | Storage must resolve to a StorageKind. | src/tilefoundry/parser/ast_pattern.py | | tensor | annotation, expression, slice_endpoint, subscript_index, type_annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | | tensor | annotation, expression, slice_endpoint, subscript_index, type_annotation | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py | +A `mesh-axis` used by placement sugar MUST resolve to a mesh binding in the +current lexical scope. A module or closure name that resolves to a `Mesh` does +not become a placement binding. Such an external value remains valid as the +context expression of `with ... as ...` or as the value supplied to +`@func(mesh=...)`; the resulting lexical binding is the name placement sugar +may use. + ## 3. Implementation Overview | Component | Responsibility | diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index 2ab653dd..a7aa768f 100644 --- a/docs/tutorial/showcase.ipynb +++ b/docs/tutorial/showcase.ipynb @@ -411,7 +411,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:337408@cta:42176,f32:51124224@cta:6390528 service=special:262144@cta:32768\n# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144\n# peak-footprint=gmem:10945036,smem:16960\n# roofline ideal-ns=11213 bound-by=memory\n\n v1 = reshard(w_q, layout=ShardLayout(\n layout=Layout((1, 256, 8, 32), None),\n attrs=(S(2),),\n mesh=cta,\n ), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n\n v42 = reshard(w_o, layout=ShardLayout(\n layout=Layout((1, 256, 8, 32), None),\n attrs=(S(2),),\n mesh=cta,\n ), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:337408@cta:42176,f32:51124224@cta:6390528 service=special:262144@cta:32768\n# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144\n# peak-footprint=gmem:10945036,smem:16960\n# roofline ideal-ns=11213 bound-by=memory\n\n v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute\n\n v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage4-4096.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nlines = annotated.splitlines()\nfor needle in (\"reshard(w_q\", \"reshard(w_o\"):\n start = next(index for index, line in enumerate(lines) if needle in line)\n end = start\n while end + 1 < len(lines):\n end += 1\n if end > start and \" # \" in lines[end]:\n break\n print()\n print(\"\\n\".join(line.rstrip() for line in lines[start : end + 1]))\n" @@ -462,7 +462,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:1246400@cta:328672,f32:6410280@cta:801285 service=integer:256@cta:32,special:33040@cta:4130\n# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804\n# peak-footprint=gmem:2950412,rmem:0,smem:33536\n# roofline ideal-ns=2474 bound-by=memory\n\n v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; traffic traffic=rmem:r32/w0@cta:r32/w0 operands=0:r0/w0,1:r32/w0,result:r0/w0; roofline\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:1246400@cta:328672,f32:6410280@cta:801285 service=integer:256@cta:32,special:33040@cta:4130\n# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804\n# peak-footprint=gmem:2950412,rmem:0,smem:33536\n# roofline ideal-ns=2474 bound-by=memory\n\n v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; traffic traffic=rmem:r32/w0@cta:r32/w0 operands=0:r0/w0,1:r32/w0,result:r0/w0; roofline\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage5-4096.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nprint()\nfor needle in (\"slice(k_cache\", \"cache_update(k_cache\"):\n print(next(line.rstrip() for line in annotated.splitlines() if needle in line))\n" diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index 55b9666a..45bb5856 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -870,17 +870,11 @@ for needle in ("reshard(w_q", "reshard(w_o"): # peak-footprint=gmem:10945036,smem:16960 # roofline ideal-ns=11213 bound-by=memory - v1 = reshard(w_q, layout=ShardLayout( - layout=Layout((1, 256, 8, 32), None), - attrs=(S(2),), - mesh=cta, - ), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory - - v42 = reshard(w_o, layout=ShardLayout( - layout=Layout((1, 256, 8, 32), None), - attrs=(S(2),), - mesh=cta, - ), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory + v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory + v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute + + v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory + v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute ``` ## 6. Stream the KV cache @@ -1028,7 +1022,7 @@ for needle in ("slice(k_cache", "cache_update(k_cache"): # peak-footprint=gmem:2950412,rmem:0,smem:33536 # roofline ideal-ns=2474 bound-by=memory - v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=rmem:r32/w0@cta:r32/w0 operands=0:r0/w0,1:r32/w0,result:r0/w0; roofline + v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=rmem:r32/w0@cta:r32/w0 operands=0:r0/w0,1:r32/w0,result:r0/w0; roofline v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory ``` diff --git a/src/tilefoundry/inspection/analysis_report.py b/src/tilefoundry/inspection/analysis_report.py index cc8f3757..7dfc8a00 100644 --- a/src/tilefoundry/inspection/analysis_report.py +++ b/src/tilefoundry/inspection/analysis_report.py @@ -20,7 +20,7 @@ from tilefoundry.analysis.report import ( selected_types as _selected_types, ) -from tilefoundry.inspection.python_printer import PythonPrintOptions, _render_hir_function +from tilefoundry.inspection.python_printer import HirPrinter, PythonPrintOptions from tilefoundry.inspection.values import ( AdvisorySummary, MemorySummary, @@ -54,7 +54,7 @@ def render_analysis( ) -> AnalysisRendering: """Render one result once for both annotated source and report data.""" selected_types_ = selected_types(result) - rendered = _render_hir_function( + rendered = HirPrinter().render( result.function, options=PythonPrintOptions( show_types=True, diff --git a/src/tilefoundry/inspection/dot.py b/src/tilefoundry/inspection/dot.py index b7c3eedb..91d30f67 100644 --- a/src/tilefoundry/inspection/dot.py +++ b/src/tilefoundry/inspection/dot.py @@ -1,150 +1,146 @@ -"""Serialize SSA HIR functions as Graphviz DOT. - -Variables, calls, and constants receive numbered nodes. Type and shard-layout -labels reuse canonical printer renderers so DOT, Python output, and the viewer -remain consistent. -See [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). -""" +"""Serialize SSA HIR functions as Graphviz DOT.""" from __future__ import annotations from tilefoundry.ir.core import Call, Constant, Var, binding_name from tilefoundry.ir.core.module import Module from tilefoundry.ir.hir.function import Function as HirFunction +from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.hir.sharding.reshard import Reshard from tilefoundry.ir.types import TensorType from tilefoundry.ir.visitor import ExprWalker -from .python_printer import _collect_meshes, _mesh_name_map, _op_display_name, _tensor_annotation +from .print_context import HirPrintContext +from .printer_base import PythonPrinter + +def _op_display_name(target) -> str: + cls = type(target).__name__ + for suffix in ("Op", "Expr", "Stmt"): + if cls.endswith(suffix) and cls != suffix: + cls = cls[: -len(suffix)] + return cls -def _type_lines(ty, mesh_name_map: dict[int, str]) -> list[str]: - """Return canonical type-annotation lines for a node label. - Split multiline shard-layout fallback text into separate label lines. - See [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). - """ - text = _tensor_annotation(ty, mesh_name_map=mesh_name_map) if isinstance(ty, TensorType) else str(ty) - return text.split("\n") +def _type_lines(ty, printer: PythonPrinter, ctx: HirPrintContext) -> list[str]: + if not isinstance(ty, TensorType): + return [str(ty)] + with printer.type_surface(): + return printer.visit(ty, ctx).split("\n") def _escape_dot(s: str) -> str: - """Escape a string for safe inclusion in a DOT label.""" return s.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") def hir_function_to_dot(fn: HirFunction) -> str: - """Convert a hir.Function to a DOT digraph string. - - Args: - fn: The HIR function to visualize. - - Returns: - A Graphviz DOT format string. - """ - type_meshes, scope_meshes = _collect_meshes(fn, include_node_types=True) - mesh_map = _mesh_name_map({**type_meshes, **scope_meshes}) - lines = [f"digraph {fn.name} {{", ' rankdir=TB;', - ' node [shape=box, style=filled, fillcolor="#f0f0f0"];', - ' edge [fontsize=10, fontcolor="#555555"];', ''] - _counter = [0] - _ids = {} - def _id(node): + """Convert a HIR function to a DOT digraph using canonical type text.""" + printer = PythonPrinter() + ctx = HirPrintContext() + root_scope = fn.body if isinstance(fn.body, MeshRegion) else None + if root_scope is not None: + ctx.push_mesh(root_scope.mesh, "mesh") + + lines = [ + f"digraph {fn.name} {{", + " rankdir=TB;", + ' node [shape=box, style=filled, fillcolor="#f0f0f0"];', + ' edge [fontsize=10, fontcolor="#555555"];', + "", + ] + counter = [0] + ids: dict[int, str] = {} + + def node_id(node): key = id(node) - if key not in _ids: - _ids[key] = f"n{_counter[0]}" - _counter[0] += 1 - return _ids[key] - - def _emit_node(nid, label_lines, fill="#f0f0f0"): - escaped = [_escape_dot(ln) for ln in label_lines] - label = "\\n".join(escaped) + if key not in ids: + ids[key] = f"n{counter[0]}" + counter[0] += 1 + return ids[key] + + def emit_node(nid, label_lines, fill="#f0f0f0"): + label = "\\n".join(_escape_dot(line) for line in label_lines) lines.append(f' {nid} [label="{label}", fillcolor="{fill}"];') - def _emit_edge(src_id, dst_id, label=""): - if label: - lines.append(f' {src_id} -> {dst_id} [label="{label}"];') - else: - lines.append(f' {src_id} -> {dst_id};') - - VAR_FILL = "#d4e6f1" - CONST_FILL = "#f9e79f" - CALL_FILL = "#d5f5e3" - SHARDING_FILL = "#e8daef" - - class _DotWalker(ExprWalker[None]): - def _emit_leaf(self, expr, label_lines, fill) -> None: - _emit_node(_id(expr), label_lines, fill=fill) - - def visit_Var(self, expr: Var, ctx=None) -> None: - self._emit_leaf( - expr, - [f"Var: {expr.name}", *_type_lines(expr.type, mesh_map)], - VAR_FILL, - ) - - def visit_Constant(self, expr: Constant, ctx=None) -> None: + def emit_edge(src_id, dst_id, label=""): + suffix = f' [label="{label}"]' if label else "" + lines.append(f" {src_id} -> {dst_id}{suffix};") + + var_fill, const_fill = "#d4e6f1", "#f9e79f" + call_fill, shard_fill = "#d5f5e3", "#e8daef" + + class DotWalker(ExprWalker[None]): + def _emit_leaf(self, expr, label_lines, fill): + emit_node(node_id(expr), label_lines, fill=fill) + + def visit_Var(self, expr: Var, current=None) -> None: + self._emit_leaf(expr, [f"Var: {expr.name}", *_type_lines(expr.type, printer, ctx)], var_fill) + + def visit_Constant(self, expr: Constant, current=None) -> None: value = f"{expr.value:.6g}" if isinstance(expr.value, float) else str(expr.value) - self._emit_leaf( - expr, - [f"Const: {value}", *_type_lines(expr.type, mesh_map)], - CONST_FILL, - ) - - def visit_Call(self, expr: Call, ctx=None) -> None: - nid = _id(expr) + self._emit_leaf(expr, [f"Const: {value}", *_type_lines(expr.type, printer, ctx)], const_fill) + + def visit_Call(self, expr: Call, current=None) -> None: + nid = node_id(expr) target = expr.target if isinstance(target, Reshard): name = binding_name(expr) header = f"{name}\\nReshard" if name else "Reshard" - _emit_node(nid, [header, *_type_lines(expr.type, mesh_map)], fill=SHARDING_FILL) + emit_node(nid, [header, *_type_lines(expr.type, printer, ctx)], fill=shard_fill) + else: + name = binding_name(expr) + op_label = _op_display_name(target) + header = f"{name}\\n{op_label}" if name else op_label + emit_node(nid, [header, *_type_lines(expr.type, printer, ctx)], fill=call_fill) + for index, arg in enumerate(expr.args): + self.visit(arg, current) + emit_edge(node_id(arg), nid, f"arg[{index}]" if len(expr.args) > 1 else "") + + def visit_MeshRegion(self, expr: MeshRegion, current=None) -> None: + alias = ctx.scope_name(expr.mesh) + ctx.push_mesh(expr.mesh, alias) + try: + self.visit(expr.body, current) for arg in expr.args: - self.visit(arg, ctx) - _emit_edge(_id(arg), nid) - return - - op_label = _op_display_name(target) - name = binding_name(expr) - header = f"{name}\\n{op_label}" if name else op_label - _emit_node(nid, [header, *_type_lines(expr.type, mesh_map)], fill=CALL_FILL) - for i, arg in enumerate(expr.args): - self.visit(arg, ctx) - edge_label = f"arg[{i}]" if len(expr.args) > 1 else "" - _emit_edge(_id(arg), nid, edge_label) - - def visit_Tuple(self, expr, ctx=None) -> None: + self.visit(arg, current) + finally: + ctx.pop_mesh() + + def visit_Tuple(self, expr, current=None) -> None: self._emit_leaf(expr, [type(expr).__name__], "#ffffff") - def visit_LoopRegion(self, expr, ctx=None) -> None: + def visit_LoopRegion(self, expr, current=None) -> None: self._emit_leaf(expr, [type(expr).__name__], "#ffffff") - def visit_ShapeOf(self, expr, ctx=None) -> None: - _emit_node(_id(expr), ["ShapeOf"], fill="#ffffff") + def visit_ShapeOf(self, expr, current=None) -> None: + emit_node(node_id(expr), ["ShapeOf"], fill="#ffffff") - def default_visit(self, expr, ctx=None) -> None: - _emit_node(_id(expr), [type(expr).__name__], fill="#ffffff") + def default_visit(self, expr, current=None) -> None: + emit_node(node_id(expr), [type(expr).__name__], fill="#ffffff") - walker = _DotWalker() + walker = DotWalker() walker.visit(fn.body) for param in fn.params: walker.visit(param) - - - lines.append("") - lines.append(' subgraph cluster_legend {') - lines.append(' label="Legend";') - lines.append(' style=dashed;') - lines.append(' fontsize=11;') - lines.append(' l_var [label="Var/Param", fillcolor="#d4e6f1", shape=box, style=filled];') - lines.append(' l_const [label="Constant", fillcolor="#f9e79f", shape=box, style=filled];') - lines.append(' l_call [label="Op", fillcolor="#d5f5e3", shape=box, style=filled];') - lines.append(' l_shard [label="Reshard", fillcolor="#e8daef", shape=box, style=filled];') - lines.append(" }") - lines.append("}") + if root_scope is not None: + ctx.pop_mesh() + + lines.extend([ + "", + ' subgraph cluster_legend {', + ' label="Legend";', + " style=dashed;", + ' fontsize=11;', + ' l_var [label="Var/Param", fillcolor="#d4e6f1", shape=box, style=filled];', + ' l_const [label="Constant", fillcolor="#f9e79f", shape=box, style=filled];', + ' l_call [label="Op", fillcolor="#d5f5e3", shape=box, style=filled];', + ' l_shard [label="Reshard", fillcolor="#e8daef", shape=box, style=filled];', + " }", + "}", + ]) return "\n".join(lines) + "\n" def module_entry_to_dot(module: Module) -> str: """Convert a Module's entry function to DOT.""" - fn = module.entry_function() - return hir_function_to_dot(fn) + return hir_function_to_dot(module.entry_function()) diff --git a/src/tilefoundry/inspection/print_context.py b/src/tilefoundry/inspection/print_context.py index 6aeea83e..0f031755 100644 --- a/src/tilefoundry/inspection/print_context.py +++ b/src/tilefoundry/inspection/print_context.py @@ -1,16 +1,25 @@ -"""Shared printer state and policy context.""" +"""Shared state accumulated while canonical Python source is rendered.""" from __future__ import annotations -from tilefoundry.utils.python_source import PythonExpr +from 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.utils.python_source import PythonExpr, _merge_imports class PrintContext: + """Imports, symbolic declarations, and lexical mesh bindings for one file.""" + def __init__(self) -> None: self.imports: set[str] = set() - - def mesh_count(self) -> int: - return 0 + self._dim_declarations: dict[str, tuple[object, str]] = {} + self._mesh_bindings: list[tuple[Mesh, str]] = [] + self._used_scope_names: set[str] = set() + self._type_annotation_surface = False def use(self, rendered: PythonExpr | str) -> str: if isinstance(rendered, PythonExpr): @@ -18,41 +27,177 @@ def use(self, rendered: PythonExpr | str) -> str: return rendered.text return rendered - def mesh_alias(self, mesh) -> str | None: + def declare_dim( + self, + name: str, + var, + *, + import_statement: str = "from tilefoundry.ir.types.dim import DimVar", + ) -> None: + self.imports.add(import_statement) + self._dim_declarations.setdefault(name, (var, "DimVar")) + + def header(self) -> list[str]: + """Render only imports and declarations reached while rendering the body.""" + imports = list(_merge_imports(tuple(self.imports))) + imports = [ + f"{line} # noqa: F401, F403" if line.endswith(" import *") else line + for line in imports + ] + lines = ["from __future__ import annotations", "", *imports, ""] + if self._dim_declarations: + lines.extend( + f'{name} = {constructor}("{var.name}", {var.lo}, {var.hi})' + for name, (var, constructor) in self._dim_declarations.items() + ) + lines.append("") + return lines + + def scope_name(self, mesh: Mesh, preferred: str | None = None) -> str: + base = preferred or (mesh.topologies[0].name if mesh.topologies else "mesh") + base = base if base.isidentifier() else "mesh" + name = base + suffix = 2 + while name in self._used_scope_names: + name = f"{base}_{suffix}" + suffix += 1 + self._used_scope_names.add(name) + return name + + def push_mesh(self, mesh: Mesh, name: str) -> None: + self._mesh_bindings.append((mesh, name)) + self._used_scope_names.add(name) + + def pop_mesh(self) -> None: + self._mesh_bindings.pop() + + @contextmanager + def type_annotation_surface(self): + """Prefer a parent binding when a type refers to a sliced mesh.""" + previous = self._type_annotation_surface + self._type_annotation_surface = True + try: + yield + finally: + self._type_annotation_surface = previous + + def mesh_alias(self, mesh: Mesh) -> str | None: + for bound, name in reversed(self._mesh_bindings): + if bound is mesh: + return name return None - def render_mode(self) -> str: - return "hir" + @staticmethod + def _axis_levels(mesh: Mesh) -> tuple[str, ...]: + levels = [""] * len(flatten(mesh.layout.shape)) + for topology, axes in zip(mesh.topologies, topology_axes(mesh), strict=True): + for axis in axes: + levels[axis] = topology.name + return tuple(levels) + + def mesh_axis_alias(self, mesh: Mesh, axis: int) -> str | None: + """Name one mesh axis through an active scope binding, if one dominates it.""" + names = mesh.names + if axis >= len(names): + return None + target_name = names[axis] + target_levels = self._axis_levels(mesh) + target_level = target_levels[axis] + target_topology = next( + (topology for topology in mesh.topologies if topology.name == target_level), None + ) + for binding_index in range(len(self._mesh_bindings) - 1, -1, -1): + bound, alias = self._mesh_bindings[binding_index] + if ( + self._type_annotation_surface + and bound is mesh + and isinstance(mesh.layout, ComposedLayout) + and mesh.layout.offset != 0 + ): + continue + if not bound.names or target_name not in bound.names: + continue + bound_levels = self._axis_levels(bound) + for bound_axis, bound_name in enumerate(bound.names): + if bound_name != target_name or bound_levels[bound_axis] != target_level: + continue + bound_topology = next( + topology for topology in bound.topologies if topology.name == target_level + ) + if target_topology == bound_topology: + return f"{alias}.{bound_name}" + return None + def mesh_slice(self, mesh: Mesh) -> str | None: + """Recover ``binding[start:stop]`` for a sliced active mesh.""" + if not ( + isinstance(mesh.layout, ComposedLayout) + and mesh.layout.inner is None + and isinstance(mesh.layout.outer, Layout) + ): + return None + for parent, alias in reversed(self._mesh_bindings): + text = self._slice_from_parent(parent, mesh, alias) + if text is not None: + return text + return None -class HirPrintContext(PrintContext): - def __init__(self, mesh_name_map: dict[int, str] | None = None) -> None: - super().__init__() - self.mesh_name_map = mesh_name_map or {} + @staticmethod + def _slice_from_parent(parent: Mesh, child: Mesh, alias: str) -> str | None: + if not ( + isinstance(parent.layout, Layout) + and isinstance(child.layout, ComposedLayout) + and child.layout.inner is None + and isinstance(child.layout.outer, Layout) + ): + return None + if parent.topologies != child.topologies or parent.names != child.names: + return None + parent_shape = flatten(parent.layout.shape) + parent_strides = flatten(parent.layout.strides) + child_shape = flatten(child.layout.outer.shape) + child_strides = flatten(child.layout.outer.strides) + if ( + len(parent_shape) != len(child_shape) + or parent_strides != child_strides + or any(not isinstance(item, int) for item in (*parent_shape, *child_shape, *parent_strides)) + or any(size < 1 or size > extent for size, extent in zip(child_shape, parent_shape)) + ): + return None + + remaining = child.layout.offset + starts = [0] * len(parent_shape) + for axis in sorted(range(len(parent_shape)), key=lambda item: parent_strides[item], reverse=True): + stride = parent_strides[axis] + maximum = parent_shape[axis] - child_shape[axis] + start = min(maximum, remaining // stride) if stride else 0 + starts[axis] = start + remaining -= start * stride + if remaining != 0: + return None + if sum(start * stride for start, stride in zip(starts, parent_strides)) != child.layout.offset: + return None + if prod(child_shape) > prod(parent_shape): + return None + + pieces: list[str] = [] + for start, size, extent in zip(starts, child_shape, parent_shape): + if start == 0 and size == extent: + pieces.append(":") + continue + stop = start + size + pieces.append(f"{'' if start == 0 else start}:{'' if stop == extent else stop}") + while len(pieces) > 1 and pieces[-1] == ":": + pieces.pop() + return f"{alias}[{', '.join(pieces)}]" - def mesh_alias(self, mesh) -> str | None: - return self.mesh_name_map.get(id(mesh)) - def mesh_count(self) -> int: - return len(self.mesh_name_map) +class HirPrintContext(PrintContext): + pass class TirPrintContext(PrintContext): - def __init__(self) -> None: - super().__init__() - self._mesh_aliases: list[dict[int, str]] = [] - - def render_mode(self) -> str: - return "tir" + pass - def push_mesh(self, mesh, name: str) -> None: - self._mesh_aliases.append({id(mesh): name}) - def pop_mesh(self) -> None: - self._mesh_aliases.pop() - - def mesh_alias(self, mesh) -> str | None: - for aliases in reversed(self._mesh_aliases): - if id(mesh) in aliases: - return aliases[id(mesh)] - return None +__all__ = ["PrintContext", "HirPrintContext", "TirPrintContext"] diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 8ec85500..149c0cbc 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -1,232 +1,378 @@ -"""Common expression-printer base classes. - -Lazy imports avoid a cycle between the shared base and the legacy HIR module. -""" - -# ruff: noqa: PLC0415 +"""Canonical Python visitor shared by the HIR and TIR printers.""" from __future__ import annotations import enum +from contextlib import contextmanager +from tilefoundry.ir.core import Call, Constant, Tuple, Var from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom -from tilefoundry.ir.types import DType, TensorType +from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType +from tilefoundry.ir.types.dim import ( + DimAdd, + DimConst, + DimFloorDiv, + DimMax, + DimMin, + DimMod, + DimMul, + DimSub, + DimVar, +) from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase from tilefoundry.ir.types.shard.mesh import Mesh from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, ShardLayout, Split from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.ir.visitor import ExprFunctor +from tilefoundry.ir.visitor import ExprFunctor, TypeFunctor from tilefoundry.target import Target from tilefoundry.utils.python_source import PythonExpr +_DIM_INFIX_OPS: dict[type, str] = { + DimAdd: "+", + DimSub: "-", + DimMul: "*", + DimFloorDiv: "//", + DimMod: "%", +} -class PythonPrinter(ExprFunctor[str]): - """Shared expression/value visitor base for HIR and TIR printers.""" +_DIM_FUNC_OPS: dict[type, str] = { + DimMin: "min", + DimMax: "max", +} - def dim_entry(self, value, ctx=None) -> str: - return str(value) - def shard_surface(self, value, ctx=None): - return None +class PythonPrinter(ExprFunctor[str], TypeFunctor[str]): + """Render the Python DSL with one dispatch root for expressions and types.""" - def atom_reference(self, value, ctx=None) -> str: - return f"T.cuda.mma.atom(op=T.cuda.mma.{value.op.name})" + def __init__(self) -> None: + ExprFunctor.__init__(self) + self._indent = "" + self._tensor_head = "Tensor" + self._nested_dim = False - def render_value(self, value, ctx=None, indent: str = "") -> str: - """Render a DSL value and register every import needed by it.""" - if isinstance(value, ShardLayout) and ctx is not None and ctx.render_mode() == "tir": - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) - attrs = ", ".join(self._shard_attr_str(a, ctx) for a in value.attrs) - if len(value.attrs) == 1: - attrs += "," - mesh = self._render_mesh_dataclass(value.mesh, ctx) - return f"ShardLayout(layout={self.render_layout(value.layout, ctx)}, attrs=({attrs}), mesh={mesh})" - if isinstance(value, Mesh) and ctx is not None and ctx.render_mode() == "tir": - alias = ctx.mesh_alias(value) - if alias is not None: - return alias - return self._render_mesh_dataclass(value, ctx) - if isinstance(value, Mesh) and ctx is not None: - alias = ctx.mesh_alias(value) - if alias is not None: - return alias - if isinstance(value, TensorType): - return self.render_tensor_type(value, ctx, indent) - if isinstance(value, Mesh): - return self.render_mesh(value, ctx, indent) - if isinstance(value, LayoutBase): - return self.render_layout(value, ctx, indent) - if isinstance(value, MmaAtom): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.dsl import T",), "T")) - return self.atom_reference(value, ctx) - if isinstance(value, enum.Enum): - if ctx is not None: - ctx.use(PythonExpr((f"from {type(value).__module__} import {type(value).__name__}",), "")) - return f"{type(value).__name__}.{value.name}" - if isinstance(value, Target): - expr = value.to_python() - return ctx.use(expr) if ctx is not None else expr.text - if isinstance(value, DType): - return self.dtype_str(value, ctx) - if isinstance(value, (str, int, float, bool, tuple, type(None))): - if isinstance(value, tuple): - vals = ", ".join(self.render_value(v, ctx, indent) for v in value) - return f"({vals}{',' if len(value)==1 else ''})" + def visit(self, value, ctx=None): # type: ignore[override] + """Dispatch every implemented functor family by the concrete node name.""" + method = getattr(self, f"visit_{type(value).__name__}", None) + if method is not None: + return method(value, ctx) + return self.default_visit(value, ctx) + + def default_visit(self, value, ctx=None) -> str: + if isinstance(value, bool): return repr(value) - raise NotImplementedError(f"no canonical Python form for {type(value).__name__}") + if isinstance(value, (int, float)): + return str(value) + raise NotImplementedError(f"no Python visit routine for {type(value).__name__}") - def dtype_str(self, dtype: DType, ctx=None) -> str: - return dtype.name + @contextmanager + def type_surface(self, *, indent: str | None = None, const: bool = False): + """Carry statement indentation and parameter const-ness through type visits.""" + previous = (self._indent, self._tensor_head) + if indent is not None: + self._indent = indent + self._tensor_head = "ConstTensor" if const else "Tensor" + try: + yield + finally: + self._indent, self._tensor_head = previous + + @contextmanager + def nested_dim(self, nested: bool): + previous = self._nested_dim + self._nested_dim = nested + try: + yield + finally: + self._nested_dim = previous + + def dim_entry(self, value, ctx=None, *, nested: bool = False) -> str: + with self.nested_dim(nested): + return self.visit(value, ctx) def shape_tuple(self, shape: tuple, ctx=None) -> str: values = tuple(self.dim_entry(entry, ctx) for entry in shape) return f"({values[0]},)" if len(values) == 1 else "(" + ", ".join(values) + ")" - def _shard_attr_str(self, attr, ctx=None) -> str: - if isinstance(attr, Broadcast): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) - return "B()" - if isinstance(attr, Split): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) - return f"S({attr.axis})" - if isinstance(attr, Partial): - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) - return f'P("{attr.reduction}")' - raise TypeError(f"unsupported shard attribute: {type(attr).__name__}") - - def render_layout(self, layout: LayoutBase | None, ctx=None, indent: str = "") -> str: - if layout is None: - return "None" - if isinstance(layout, Layout): - strides = self.shape_tuple(layout.strides, ctx) if layout.strides is not None else "None" - if ctx is not None and ctx.render_mode() == "tir": - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout",), "")) - return f"Layout(shape={self.shape_tuple(layout.shape, ctx)}, strides={strides})" - return f"Layout({self.shape_tuple(layout.shape, ctx)}, {strides})" - if isinstance(layout, ShardLayout): - return self.render_shard_layout(layout, ctx, indent) - if isinstance(layout, ComposedLayout): + def dtype_str(self, dtype: DType, ctx=None) -> str: + return dtype.name + + def visit_DimVar(self, value: DimVar, ctx=None) -> str: + if ctx is not None: + ctx.declare_dim(value.name, value) + return value.name + + def visit_Var(self, value: Var, ctx=None) -> str: + return value.name + + def visit_Constant(self, value: Constant, ctx=None) -> str: + return repr(value.value) + + def visit_Tuple(self, value: Tuple, ctx=None) -> str: + rendered = ", ".join(self.visit(item, ctx) for item in value.elements) + return f"({rendered}{',' if len(value.elements) == 1 else ''})" + + def visit_Call(self, value: Call, ctx=None) -> str: + ceildiv_args = self._ceildiv_args(value) + if ceildiv_args is not None: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) - if ctx is not None and ctx.render_mode() == "tir": - inner = "None" if layout.inner is None else self.render_layout(layout.inner, ctx, indent) - outer = "None" if layout.outer is None else self.render_layout(layout.outer, ctx, indent) - return f"ComposedLayout(inner={inner}, offset={layout.offset}, outer={outer})" - child = indent + " " - return ( - "ComposedLayout(\n" - f"{child}inner={self.render_layout(layout.inner, ctx, child)},\n" - f"{child}offset={self.dim_entry(layout.offset, ctx)},\n" - f"{child}outer={self.render_layout(layout.outer, ctx, child)},\n" - f"{indent})" + ctx.use(PythonExpr(("from tilefoundry.ir.types.dim import ceildiv",), "ceildiv")) + left, right = ceildiv_args + return f"ceildiv({self.dim_entry(left, ctx)}, {self.dim_entry(right, ctx)})" + target = value.target + if isinstance(target, DimConst): + return str(target.value) + for op_type, symbol in _DIM_INFIX_OPS.items(): + if isinstance(target, op_type): + left, right = value.args + rendered = ( + f"{self.dim_entry(left, ctx, nested=True)} {symbol} " + f"{self.dim_entry(right, ctx, nested=True)}" + ) + return f"({rendered})" if self._nested_dim else rendered + for op_type, name in _DIM_FUNC_OPS.items(): + if isinstance(target, op_type): + args = ", ".join(self.dim_entry(arg, ctx) for arg in value.args) + return f"{name}({args})" + return self.visit_program_call(value, ctx) + + def visit_program_call(self, value: Call, ctx=None) -> str: + raise NotImplementedError(f"{type(self).__name__} cannot render program calls") + + @staticmethod + def _ceildiv_args(value: Call) -> tuple[object, object] | None: + """Recover the public constructor from ceildiv's canonical arithmetic tree.""" + if not isinstance(value.target, DimFloorDiv) or len(value.args) != 2: + return None + numerator, divisor = value.args + if not ( + isinstance(numerator, Call) + and isinstance(numerator.target, DimSub) + and len(numerator.args) == 2 + and isinstance(numerator.args[1], Constant) + and numerator.args[1].value == 1 + ): + return None + added = numerator.args[0] + if not ( + isinstance(added, Call) + and isinstance(added.target, DimAdd) + and len(added.args) == 2 + and added.args[1] == divisor + ): + return None + return added.args[0], divisor + + def shard_surface(self, value: ShardLayout, ctx=None) -> str | None: + """Render placement sugar only when every mesh axis has a scope binding.""" + layout = value.layout + names = value.mesh.names + if ( + not isinstance(layout, Layout) + or not names + or len(value.attrs) != len(names) + or ctx is None + ): + return None + refs = tuple(ctx.mesh_axis_alias(value.mesh, index) for index in range(len(names))) + if any(ref is None for ref in refs): + return None + + splits: dict[int, list[str]] = {} + partials: list[str] = [] + broadcasts: list[tuple[str, str, Broadcast]] = [] + named_bindings: set[str] = set() + for index, (attr, ref) in enumerate(zip(value.attrs, refs, strict=True)): + assert ref is not None + binding = ref.partition(".")[0] + if isinstance(attr, Split): + if attr.axis >= len(layout.shape): + return None + splits.setdefault(attr.axis, []).append(ref) + named_bindings.add(binding) + elif isinstance(attr, Partial): + partials.append(f'{ref} @ {self.visit(attr, ctx)}') + named_bindings.add(binding) + elif isinstance(attr, Broadcast): + broadcasts.append((binding, ref, attr)) + else: + return None + + states = list(partials) + for binding, ref, attr in broadcasts: + if binding not in named_bindings: + states.append(f"{ref} @ {self.visit(attr, ctx)}") + named_bindings.add(binding) + if not splits and not states: + states = [f"{ref} @ {self.visit(attr, ctx)}" for _binding, ref, attr in broadcasts] + if not splits and not states: + return None + + explicit = layout.strides is not None + if explicit and any( + axis in splits + and self.dim_entry(dim, ctx, nested=True) != self.dim_entry(dim, ctx) + for axis, dim in enumerate(layout.shape) + ): + return None + + dims = [ + ( + f"{self.dim_entry(dim, ctx, nested=True)} " + + " ".join(f"@ {ref}" for ref in splits[axis]) ) - raise TypeError(f"unsupported layout type: {type(layout).__name__}") + if axis in splits + else self.dim_entry(dim, ctx) + for axis, dim in enumerate(layout.shape) + ] + dims_text = ", ".join(dims) + ("," if len(dims) == 1 else "") + parts = [f"({dims_text})"] + if explicit: + parts.append(self.shape_tuple(layout.strides, ctx)) + if states: + parts.append("{" + ", ".join(states) + "}") + return parts[0] if len(parts) == 1 else "(" + ", ".join(parts) + ")" + + def visit_TensorType(self, value: TensorType, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr((f"from tilefoundry.dsl import {self._tensor_head}",), self._tensor_head)) + result = ( + f"{self._tensor_head}[" + f'{self.shape_tuple(value.shape, ctx)}, "{self.dtype_str(value.dtype, ctx)}"' + ) + if isinstance(value.layout, ShardLayout): + surface = self.shard_surface(value.layout, ctx) + if surface is not None: + result += f", {surface}" + else: + with self.type_surface(indent=self._indent + " "): + result += f",\n{self._indent}{self.visit(value.layout, ctx)}" + elif value.layout is not None: + result += f", {self.visit(value.layout, ctx)}" + if value.storage is not StorageKind.GMEM: + result += f', "{value.storage.name.lower()}"' + return result + "]" + + def visit_TupleType(self, value: TupleType, ctx=None) -> str: + return f"Tuple[{', '.join(self.visit(field, ctx) for field in value.fields)}]" + + def visit_UnitType(self, value: UnitType, ctx=None) -> str: + return "None" + + def visit_DType(self, value: DType, ctx=None) -> str: + return self.dtype_str(value, ctx) - def render_mesh(self, mesh: Mesh, ctx=None, indent: str = "") -> str: + def visit_Mesh(self, value: Mesh, ctx=None) -> str: if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout, Mesh, Topology",), "")) - values = ", ".join( + alias = ctx.mesh_alias(value) + if alias is not None: + return alias + sliced = ctx.mesh_slice(value) + if sliced is not None: + return sliced + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "Mesh")) + topologies = ", ".join( f'Topology("{topology.name}", {self.dim_entry(topology.size, ctx)})' - for topology in mesh.topologies + for topology in value.topologies ) - topologies = f"({values}{',' if len(mesh.topologies) == 1 else ''})" - result = f"Mesh({topologies}, {self.render_layout(mesh.layout, ctx, indent)}" - if mesh.names: - result += f", names={tuple(mesh.names)!r}" + topologies = f"({topologies}{',' if len(value.topologies) == 1 else ''})" + result = f"Mesh({topologies}, {self.visit(value.layout, ctx)}" + if value.names: + result += f", names={tuple(value.names)!r}" return result + ")" - def render_shard_layout(self, layout: ShardLayout, ctx=None, indent: str = "", *, mesh_ref=None) -> str: + def visit_NoneType(self, value: None, ctx=None) -> str: + return "None" + + def visit_Layout(self, value: Layout, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Layout",), "Layout")) + strides = self.shape_tuple(value.strides, ctx) if value.strides is not None else "None" + return f"Layout({self.shape_tuple(value.shape, ctx)}, {strides})" + + def visit_ComposedLayout(self, value: ComposedLayout, ctx=None) -> str: + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ComposedLayout",), "")) + outer, child = self._indent, self._indent + " " + with self.type_surface(indent=child): + inner_text = self.visit(value.inner, ctx) + outer_text = self.visit(value.outer, ctx) + return ( + "ComposedLayout(\n" + f"{child}inner={inner_text},\n" + f"{child}offset={self.dim_entry(value.offset, ctx)},\n" + f"{child}outer={outer_text},\n" + f"{outer})" + ) + + def visit_ShardLayout(self, value: ShardLayout, ctx=None) -> str: + surface = self.shard_surface(value, ctx) + if surface is not None: + return surface if ctx is not None: ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import ShardLayout",), "")) - child = indent + " " - attrs = ", ".join(self._shard_attr_str(attr, ctx) for attr in layout.attrs) - if len(layout.attrs) == 1: + outer, child = self._indent, self._indent + " " + attrs = ", ".join(self.visit(attr, ctx) for attr in value.attrs) + if len(value.attrs) == 1: attrs += "," - if ctx is not None and ctx.render_mode() == "tir": - mesh_text = mesh_ref if mesh_ref is not None else self._render_mesh_compact(layout.mesh, ctx) - layout_text = self._render_layout_positional(layout.layout, ctx) - else: - mesh_text = mesh_ref if mesh_ref is not None else self.render_mesh(layout.mesh, ctx, child) - layout_text = self.render_layout(layout.layout, ctx, child) + with self.type_surface(indent=child): + layout_text = self.visit(value.layout, ctx) + mesh_text = self.visit(value.mesh, ctx) return ( "ShardLayout(\n" f"{child}layout={layout_text},\n" f"{child}attrs=({attrs}),\n" f"{child}mesh={mesh_text},\n" - f"{indent})" + f"{outer})" ) - def render_tensor_type(self, ty: TensorType, ctx=None, indent: str = "", is_const=False) -> str: - head = "ConstTensor" if is_const else "Tensor" - result = f'{head}[{self.shape_tuple(ty.shape, ctx)}, "{self.dtype_str(ty.dtype, ctx)}"' - if isinstance(ty.layout, ShardLayout): - surface = self.shard_surface(ty.layout, ctx) - if surface is not None: - result += f", {surface}" - else: - result += f",\n{indent} {self.render_shard_layout(ty.layout, ctx, indent + ' ')}" - if ty.storage is not StorageKind.GMEM: - result += f', "{ty.storage.name.lower()}"' - return result + "]" - - def _render_layout_positional(self, layout, ctx=None): - if isinstance(layout, Layout): - strides = self.shape_tuple(layout.strides, ctx) if layout.strides is not None else "None" - return f"Layout({self.shape_tuple(layout.shape, ctx)}, {strides})" - return self.render_layout(layout, ctx) - - def _render_mesh_dataclass(self, mesh, ctx=None): + def visit_Broadcast(self, value: Broadcast, ctx=None) -> str: if ctx is not None: - alias = ctx.mesh_alias(mesh) - if alias is not None: - return alias - if ctx is not None: - ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import Mesh, Topology",), "")) - values = ", ".join(f'Topology(name="{t.name}", size={self.dim_entry(t.size, ctx)})' for t in mesh.topologies) - if len(mesh.topologies) == 1: - values += "," - names = ", ".join(f'"{n}"' for n in mesh.names) - if len(mesh.names) == 1: - names += "," - layout = self.render_layout(mesh.layout, ctx) - return f"Mesh(topologies=({values}), layout={layout}, names=({names}))" - - def _render_mesh_compact(self, mesh, ctx=None): + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import B",), "B")) + return "B()" + + def visit_Split(self, value: Split, ctx=None) -> str: if ctx is not None: - alias = ctx.mesh_alias(mesh) - if alias is not None: - return alias + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import S",), "S")) + return f"S({value.axis})" + + def visit_Partial(self, value: Partial, ctx=None) -> str: if ctx is not None: - ctx.use( - PythonExpr( - ("from tilefoundry.ir.types.shard import Mesh, Topology",), "" - ) - ) - values = ", ".join(f'Topology("{t.name}", {self.dim_entry(t.size, ctx)})' for t in mesh.topologies) - topologies = f"({values}{',' if len(mesh.topologies) == 1 else ''})" - names = f", names={tuple(mesh.names)!r}" if mesh.names else "" - return f"Mesh({topologies}, {self._render_layout_positional(mesh.layout, ctx)}{names})" + ctx.use(PythonExpr(("from tilefoundry.ir.types.shard import P",), "P")) + return f'P("{value.reduction}")' + + def atom_reference(self, value: MmaAtom, ctx=None) -> str: + return f"T.cuda.mma.atom(op=T.cuda.mma.{value.op.name})" + + def render_value(self, value, ctx=None, indent: str = "") -> str: + """Render a non-expression attribute through the same visitor when possible.""" + if isinstance(value, (TensorType, Mesh, LayoutBase, DType)): + with self.type_surface(indent=indent): + return self.visit(value, ctx) + if isinstance(value, MmaAtom): + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.dsl import T",), "T")) + return self.atom_reference(value, ctx) + if isinstance(value, enum.Enum): + if ctx is not None: + ctx.use(PythonExpr((f"from {type(value).__module__} import {type(value).__name__}",), "")) + return f"{type(value).__name__}.{value.name}" + if isinstance(value, Target): + rendered = value.to_python() + return ctx.use(rendered) if ctx is not None else rendered.text + if isinstance(value, tuple): + rendered = ", ".join(self.render_value(item, ctx, indent) for item in value) + return f"({rendered}{',' if len(value) == 1 else ''})" + if value is None or isinstance(value, (str, int, float, bool)): + return repr(value) + raise NotImplementedError(f"no canonical Python form for {type(value).__name__}") def render_pattern(self, pattern: Pattern, ctx=None) -> str: if isinstance(pattern, DimVarRangePat): + if ctx is not None: + ctx.use(PythonExpr(("from tilefoundry.ir.core.pattern import DimVarRangePat",), "")) return f'DimVarRangePat("{pattern.dim_var}", {pattern.lo}, {pattern.hi})' return repr(pattern) - def mesh_name_map(self, meshes: dict[int, Mesh]) -> dict[int, str]: - used: set[str] = set() - result: dict[int, str] = {} - for identity, mesh in meshes.items(): - base = mesh.topologies[0].name if mesh.topologies else "mesh" - name, suffix = base, 2 - while name in used: - name = f"{base}_{suffix}" - suffix += 1 - used.add(name) - result[identity] = name - return result + +__all__ = ["PythonPrinter"] diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index bbb47ff0..0a72dc6d 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -1,8 +1,8 @@ """Canonical Python DSL printer for HIR Functions. Converts a ``hir.Function`` to executable Python source using the -``@func`` DSL. When meshes have named axes, compact sugar annotations -are emitted; otherwise the verbose ``ShardLayout(...)`` form is used. +``@func`` DSL. Placement sugar names only explicit mesh-scope bindings; +otherwise the verbose ``ShardLayout(...)`` form is used. """ from __future__ import annotations @@ -41,7 +41,6 @@ from tilefoundry.ir.hir.sharding.reshard import Reshard from tilefoundry.ir.hir.specialize import ( canonical_specialization_signature, - dim_vars_reached, display_name, origin_of, ) @@ -50,29 +49,14 @@ from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.types import DType, TensorType, TupleType -from tilefoundry.ir.types.dim import ( - DimAdd, - DimConst, - DimFloorDiv, - DimMax, - DimMin, - DimMod, - DimMul, - DimSub, - DimVar, -) -from tilefoundry.ir.types.shard import try_c_order_strides -from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase -from tilefoundry.ir.types.shard.mesh import Mesh +from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.shard.shard_layout import ( Broadcast, Partial, ShardLayout, Split, - layout_axis_to_tensor_axis, ) -from tilefoundry.ir.types.substitute import dim_vars_by_name -from tilefoundry.ir.visitor import ExprFunctor, expr_children +from tilefoundry.ir.visitor import expr_children from tilefoundry.utils.python_source import PythonExpr from .print_context import HirPrintContext @@ -81,37 +65,184 @@ from .tir_printer import tir_function_to_python, tir_module_to_python from .values import PARTS, render_comment -_DIM_INFIX_OPS: dict[type, str] = { - DimAdd: "+", - DimSub: "-", - DimMul: "*", - DimFloorDiv: "//", - DimMod: "%", -} +class HirPrinter(PythonPrinter): + """Canonical HIR expression, type, and function printer.""" -_DIM_FUNC_OPS: dict[type, str] = { - DimMin: "min", - DimMax: "max", -} + def __init__(self) -> None: + super().__init__() + self._names: dict[int, str] = {} + self._param_alias: dict[int, Expr] = {} + self._child_entries: dict[int, str] = {} + self._moved_window = lambda start, size, stride: None + def print(self, fn: HirFunction, *, options=None) -> str: + return self.render(fn, options=options).source + + def render(self, fn: HirFunction, *, options=None) -> _PythonRendering: + return _render_hir_function(fn, options=options) + + def bind_def( + self, + names: dict[int, str], + param_alias: dict[int, Expr], + child_entries: dict[int, str], + moved_window, + ) -> None: + """Install the value-naming environment for one function definition.""" + self._names = names + self._param_alias = param_alias + self._child_entries = child_entries + self._moved_window = moved_window + + def tuple_reference(self, elements) -> str: + inner = ", ".join( + repr(element.value) + if isinstance(element, Constant) + else self.reference(element) + for element in elements + ) + return f"({inner}{',' if len(elements) == 1 else ''})" -class HirPrinter(PythonPrinter): - """HIR façade retaining the historical canonical rendering entry point.""" + def reference(self, expr: Expr) -> str: + """Return the binding that denotes one value in the current definition.""" + if isinstance(expr, Tuple): + return self.tuple_reference(expr.elements) + if id(expr) in self._param_alias: + return self.reference(self._param_alias[id(expr)]) + projection = _region_projection(expr) + if isinstance(projection, LoopRegion): + return self._names[id(projection.carried_args[expr.target.index])] + if isinstance(expr, LoopRegion): + carried = tuple(self._names[id(carry)] for carry in expr.carried_args) + return carried[0] if len(carried) == 1 else "(" + ", ".join(carried) + ")" + if isinstance(projection, MeshRegion): + return self.reference(projection.body.elements[expr.target.index]) + if isinstance(expr, MeshRegion): + return self.reference(expr.body) + return self._names[id(expr)] - def print(self, fn: HirFunction, *, options=None) -> str: - return _render_hir_function(fn, options=options).source + def _slice_start(self, start, size, stride) -> str: + moved = self._moved_window(start, size, stride) + if moved is None: + return repr(start.value) if isinstance(start, Constant) else self.reference(start) + window, offset = moved + return _moved_window_ref(self.reference(window), offset) - def dim_entry(self, value, ctx=None) -> str: - return shape_entry_str(value) + def visit_program_call(self, expr: Call, ctx=None) -> str: + """Render one HIR call after expression-level dispatch selected it.""" + target = expr.target + args_text = ", ".join(self.reference(arg) for arg in expr.args) + if isinstance(target, Reshard): + if ctx is not None: + ctx.imports.add("from tilefoundry.dsl.tf import *") + layout_kw = "" + if target.layout is not None: + with self.type_surface(indent=self._indent + " "): + layout_kw = ", layout=" + self.visit(target.layout, ctx) + storage = "" + if target.storage is not None: + storage_name = target.storage.name.lower() + if ctx is not None: + ctx.use( + PythonExpr( + (f"from tilefoundry.dsl.storage import {storage_name}",), + storage_name, + ) + ) + storage = f", storage={storage_name}" + return f"reshard({args_text}{layout_kw}{storage})" + if isinstance(target, HirFunction): + binding = _module_callee_binding(target, self._child_entries) + return f"{binding or target.name}({args_text})" + if isinstance(target, Slice): + starts = expr.args[1] + if not isinstance(starts, Tuple): + raise ValueError("canonical_source: Slice starts must be a Tuple") + indexers: list[str] = [] + runtime_starts = False + for axis, (start, size, stride) in enumerate( + zip(starts.elements, target.sizes, target.strides) + ): + if self._moved_window(start, size, stride) is not None: + indexers.append(self._slice_start(start, size, stride)) + continue + dim = expr.args[0].type.shape[axis] + if ( + isinstance(start, Constant) + and start.value == 0 + and size == dim + and stride == 1 + ): + indexers.append(":") + continue + if not ( + isinstance(start, Constant) + and isinstance(start.value, int) + and isinstance(size, int) + and isinstance(stride, int) + ): + runtime_starts = True + break + begin = int(start.value) + stop = begin + size * stride + indexers.append( + f"{begin}:{stop}" if stride == 1 else f"{begin}:{stop}:{stride}" + ) + if runtime_starts: + if ctx is not None: + ctx.imports.add("from tilefoundry.dsl.tf import *") + start_refs = ", ".join( + self._slice_start(start, size, stride) + for start, size, stride in zip( + starts.elements, target.sizes, target.strides + ) + ) + if len(starts.elements) == 1: + start_refs += "," + return ( + f"slice({self.reference(expr.args[0])}, ({start_refs}), " + f"sizes={_attr_tuple_str(target.sizes, self, ctx)}, " + f"strides={_attr_tuple_str(target.strides, self, ctx)})" + ) + return f"{self.reference(expr.args[0])}[{', '.join(indexers)}]" - def shard_surface(self, value, ctx=None): - mesh_name = ctx.mesh_alias(value.mesh) if ctx is not None else None - if mesh_name is None or not value.mesh.names: - return None - return _shard_layout_surface_str( - value, mesh_name=mesh_name, mesh_unique=ctx.mesh_count() == 1 - ) + if ctx is not None: + ctx.imports.add("from tilefoundry.dsl.tf import *") + alias_name = _kinded_alias_name(target) + suppressed = {"kind"} if alias_name is not None else set() + attrs: list[str] = [] + for param in type(target).params(): + if param.kind != "attribute": + continue + value = getattr(target, param.name, None) + if value is None or param.name in suppressed or param.name == "layout": + continue + if isinstance(value, str): + attrs.append(f'{param.name}="{value}"') + elif isinstance(value, DType): + attrs.append(f'{param.name}="{value.name}"') + elif isinstance(value, enum.Enum) and isinstance(value.value, str): + attrs.append(f'{param.name}="{value.value}"') + elif isinstance(value, float): + if math.isinf(value): + literal = "-1e999" if value < 0 else "1e999" + elif math.isnan(value): + literal = "(1e999 - 1e999)" + else: + literal = repr(value) + attrs.append(f"{param.name}={literal}") + elif isinstance(value, (ShardLayout, TensorType)): + with self.type_surface(indent=self._indent + " "): + rendered = self.visit(value, ctx) + if isinstance(value, TensorType): + rendered = " ".join(rendered.split()) + attrs.append(f"{param.name}={rendered}") + elif isinstance(value, tuple): + attrs.append(f"{param.name}={_attr_tuple_str(value, self, ctx)}") + else: + attrs.append(f"{param.name}={value}") + return f"{_op_name(target)}({', '.join([*(self.reference(arg) for arg in expr.args), *attrs])})" @dataclass(frozen=True) @@ -143,22 +274,16 @@ def _physical_line_count(lines: list[str]) -> int: return sum(line.count("\n") + 1 for line in lines) -def _compact_type(ty: object, mesh_name_map: dict[int, str]) -> str: - """One physical-line, DSL-shaped type annotation for inspection output. - - Takes the same mesh-name map the signature and mesh prelude are rendered - from, so an annotated layout names the hoisted mesh instead of restating it. - """ - if isinstance(ty, TensorType): - rendered = _tensor_annotation(ty, mesh_name_map=mesh_name_map) +def _compact_type(ty: object, printer: PythonPrinter, ctx) -> str: + """One physical-line, DSL-shaped type annotation for inspection output.""" + if isinstance(ty, (TensorType, TupleType)): + with ctx.type_annotation_surface(), printer.type_surface(): + rendered = printer.visit(ty, ctx) return " ".join(rendered.split()) - if isinstance(ty, TupleType): - fields = ", ".join(_compact_type(field, mesh_name_map) for field in ty.fields) - return f"Tuple[{fields}]" return repr(ty) -def _comments(expr: Expr, options: PythonPrintOptions, mesh_name_map: dict[int, str]) -> str: +def _comments(expr: Expr, options: PythonPrintOptions, printer: PythonPrinter, ctx) -> str: """Return same-line annotations for one printed statement. Omit the binding because the left-hand side already carries its emitted name @@ -171,8 +296,8 @@ def _comments(expr: Expr, options: PythonPrintOptions, mesh_name_map: dict[int, the boundary between those two languages. """ comments: list[str] = [] - if options.show_types: - comments.append(_compact_type(expr.type, mesh_name_map)) + if options.show_types and not isinstance(expr, (LoopRegion, MeshRegion)): + comments.append(_compact_type(expr.type, printer, ctx)) for metadata_type in options.comment_metadata_types: metadata = get_metadata(expr, metadata_type) if metadata is None: @@ -183,213 +308,6 @@ def _comments(expr: Expr, options: PythonPrintOptions, mesh_name_map: dict[int, return f" # {PARTS.join(comments)}" if comments else "" -def shape_entry_str(entry: object) -> str: - """Render one tensor shape entry in canonical human-readable form. - - Static integers remain literals, dimension variables use their names, and - arithmetic expression trees use infix or function syntax. The printer and - viewer share this rendering instead of exposing dataclass representations. - See [types §4](docs/spec/types.md#4-dim--symbolic-shape-dimensions) and - [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). - """ - return _shape_entry_str(entry, nested=False) - - -class _ShapeEntryVisitor(ExprFunctor[str]): - def __init__(self, nested: bool) -> None: - super().__init__() - self.nested = nested - - def visit_DimVar(self, entry: DimVar, ctx=None) -> str: - return entry.name - - def visit_Var(self, entry: Var, ctx=None) -> str: - return entry.name - - def visit_Constant(self, entry: Constant, ctx=None) -> str: - return str(entry.value) - - def visit_Call(self, entry: Call, ctx=None) -> str: - ceildiv_args = _ceildiv_args(entry) - if ceildiv_args is not None: - a, b = ceildiv_args - return f"ceildiv({self._render(a, False)}, {self._render(b, False)})" - target = entry.target - if isinstance(target, DimConst): - return str(target.value) - for op_cls, sym in _DIM_INFIX_OPS.items(): - if isinstance(target, op_cls): - a, b = entry.args - rendered = f"{self._render(a, True)} {sym} {self._render(b, True)}" - return f"({rendered})" if self.nested else rendered - for op_cls, fname in _DIM_FUNC_OPS.items(): - if isinstance(target, op_cls): - rendered = ", ".join(self._render(arg, False) for arg in entry.args) - return f"{fname}({rendered})" - return repr(entry) - - def default_visit(self, entry, ctx=None) -> str: - if isinstance(entry, bool): - return repr(entry) - if isinstance(entry, int): - return str(entry) - return repr(entry) - - def _render(self, entry, nested: bool) -> str: - previous = self.nested - self.nested = nested - try: - return self.visit(entry, None) - finally: - self.nested = previous - - -def _shape_entry_str(entry: object, *, nested: bool) -> str: - return _ShapeEntryVisitor(nested).visit(entry) - - -def _ceildiv_args(entry: Call) -> tuple[object, object] | None: - """Recover the public constructor from ceildiv's canonical arithmetic tree.""" - if not isinstance(entry.target, DimFloorDiv) or len(entry.args) != 2: - return None - numerator, divisor = entry.args - if not ( - isinstance(numerator, Call) - and isinstance(numerator.target, DimSub) - and len(numerator.args) == 2 - and isinstance(numerator.args[1], Constant) - and numerator.args[1].value == 1 - ): - return None - added = numerator.args[0] - if not ( - isinstance(added, Call) - and isinstance(added.target, DimAdd) - and len(added.args) == 2 - and added.args[1] == divisor - ): - return None - return added.args[0], divisor - - -def _classify_shard_attrs( - sl: ShardLayout, mesh_name: str -) -> tuple[dict[int, list[str]], list[str]] | None: - """Classify shard attributes into layout-axis splits and partials. - - Preserve mesh-axis order, allow nested axes to split one layout axis, and - omit broadcasts. Return ``None`` for rank mismatch, invalid axes, or unknown - attributes so callers use verbose fallback. Surface and compact renderers - share the result, with the latter remapping splits onto tensor axes. - """ - layout = sl.layout - if not isinstance(layout, Layout) or len(sl.attrs) != len(sl.mesh.layout.shape): - return None - layout_rank = len(layout.shape) - names = sl.mesh.names if hasattr(sl.mesh, "names") and sl.mesh.names else () - splits: dict[int, list[str]] = {} - partials: list[str] = [] - for mesh_axis_idx, attr in enumerate(sl.attrs): - axis_name = names[mesh_axis_idx] if mesh_axis_idx < len(names) else f"ax{mesh_axis_idx}" - axis_ref = f"{mesh_name}.{axis_name}" - if isinstance(attr, Split): - if attr.axis >= layout_rank: - return None - splits.setdefault(attr.axis, []).append(axis_ref) - elif isinstance(attr, Partial): - partials.append(f'{axis_ref} @ P("{attr.reduction or "sum"}")') - elif not isinstance(attr, Broadcast): - return None - return splits, partials - - -def _shard_layout_surface_str( - sl: ShardLayout, - mesh_name: str = "gpu", - *, - mesh_unique: bool = False, -) -> str | None: - """Render canonical parser sugar for a shard layout. - - Inline splits on layout dimensions, emit partial value states as a set, and - omit broadcasts. Include explicit strides only when present. Return ``None`` - when sugar cannot express the layout so callers use verbose fallback. - - A symbolic shape has no static C-order strides to compare against, so the - ones it states are emitted rather than assumed contiguous. - """ - layout = sl.layout - if not isinstance(layout, Layout): - return None - classified = _classify_shard_attrs(sl, mesh_name) - if classified is None: - return None - splits, partials = classified - - if not splits and not partials and not mesh_unique: - return None - - c_strides = try_c_order_strides(layout.shape) - explicit = layout.strides is not None and layout.strides != c_strides - if explicit and any( - i in splits and _shape_entry_str(dim, nested=True) != shape_entry_str(dim) - for i, dim in enumerate(layout.shape) - ): - return None - - dims = [ - f"{_shape_entry_str(d, nested=True)} {' '.join(f'@ {r}' for r in splits[i])}" - if i in splits - else shape_entry_str(d) - for i, d in enumerate(layout.shape) - ] - dim_str = ", ".join(dims) - if len(dims) == 1: - dim_str += "," - axis_tuple = f"({dim_str})" - - stride_str = _shape_tuple(layout.strides) if explicit else None - value_set = "{" + ", ".join(partials) + "}" if partials else None - - if stride_str is None and value_set is None: - return axis_tuple - parts = [axis_tuple] - if stride_str is not None: - parts.append(stride_str) - if value_set is not None: - parts.append(value_set) - return "(" + ", ".join(parts) + ")" - - -def shard_compact_inline( - sl: ShardLayout, mesh_name: str, tensor_shape: tuple -) -> tuple[dict[int, str], list[str]] | None: - """Decompose a shard layout for compact tensor-axis display. - - Map splits to tensor axes, collect ordered partial states, and omit - broadcasts. Return ``None`` for ambiguous split positions, invalid axes, - unknown attributes, or rank mismatch so callers fall back to canonical - rendering. Attribute classification is shared with surface rendering. - """ - layout = sl.layout - if not isinstance(layout, Layout): - return None - classified = _classify_shard_attrs(sl, mesh_name) - if classified is None: - return None - splits, partials = classified - la2ta = layout_axis_to_tensor_axis(layout.shape, tensor_shape) - split_ref: dict[int, str] = {} - for layout_axis, refs in splits.items(): - if len(refs) != 1: - return None - t_axis = la2ta[layout_axis] - if t_axis in split_ref: - return None - split_ref[t_axis] = refs[0] - return split_ref, partials - - def _moved_window_ref(name: str, offset: int) -> str: """A tile-window indexer, carrying the compile-time offset that moves it.""" if offset == 0: @@ -397,7 +315,7 @@ def _moved_window_ref(name: str, offset: int) -> str: return f"{name} + {offset}" if offset > 0 else f"{name} - {-offset}" -def _attr_tuple_str(value: tuple) -> str: +def _attr_tuple_str(value: tuple, printer: PythonPrinter, ctx) -> str: """Render an attribute's tuple value as a Python tuple literal. A shape-valued attribute -- `new_shape`, a tile's extents -- can hold a @@ -407,7 +325,7 @@ def _attr_tuple_str(value: tuple) -> str: the header emits binds the name, not the repr. """ rendered = tuple( - shape_entry_str(entry) if _is_dim_entry(entry) else repr(entry) + printer.visit(entry, ctx) if _is_dim_entry(entry) else repr(entry) for entry in value ) if len(rendered) == 1: @@ -417,10 +335,7 @@ def _attr_tuple_str(value: tuple) -> str: def _is_dim_entry(entry: object) -> bool: """Whether *entry* is a dimension rather than a plain attribute value.""" - return isinstance(entry, (DimVar, Var, Constant)) or ( - isinstance(entry, Call) - and isinstance(entry.target, (DimConst, *_DIM_INFIX_OPS, *_DIM_FUNC_OPS)) - ) + return isinstance(entry, (DimVar, Var, Constant, Call)) def _tensor_import_names(fn: HirFunction) -> str: @@ -580,7 +495,7 @@ def _where_str(metadata: ScheduleConstraintMetadata) -> str: fields.append(f"layout={_layout_constraint_str(layout)}") for item in metadata.constraints: if isinstance(item, MeshConstraint): - fields.append(f"mesh={_mesh_str(item.mesh)}") + fields.append(f"mesh={PythonPrinter().visit(item.mesh, HirPrintContext())}") elif isinstance(item, StorageConstraint): fields.append(f'storage="{item.storage.name.lower()}"') return "where(" + ", ".join(fields) + ")" @@ -615,52 +530,6 @@ def iter_exprs(root: Expr | None, seen: set[int] | None = None) -> Iterator[Expr yield root -def _collect_meshes( - fn: HirFunction, - *, - include_node_types: bool = False, -) -> tuple[dict[int, Mesh], dict[int, Mesh]]: - """Collect meshes needed before emitting a function header. - - The function header declares every mesh and builds the name map used by - parameter annotations and region bodies, so collection must precede body - emission. It walks params, return type, and body references once; the - optional node-type scan is used by the viewer for intermediate annotations. - - With ``include_node_types=True`` (the viewer's wider scan, via - ``viewer.builder._collect_view_meshes``) every node's own result type is - also walked, since the viewer renders shard sugar on intermediate types too. - """ - type_meshes: dict[int, Mesh] = {} - scope_meshes: dict[int, Mesh] = {} - - def _add_layout(layout) -> None: - if isinstance(layout, ShardLayout): - type_meshes.setdefault(id(layout.mesh), layout.mesh) - - def _add_type(ty) -> None: - if isinstance(ty, TensorType): - _add_layout(ty.layout) - elif isinstance(ty, TupleType): - for f in ty.fields: - _add_type(f) - - for p in fn.params: - _add_type(p.type) - _add_type(fn.return_type) - - expressions = tuple(iter_exprs(fn.body)) - for expr in expressions: - if include_node_types: - _add_type(getattr(expr, "type", None)) - if isinstance(expr, Call) and isinstance(expr.target, Reshard): - _add_layout(expr.target.layout) - if isinstance(expr, MeshRegion): - scope_meshes.setdefault(id(expr.mesh), expr.mesh) - - return type_meshes, scope_meshes - - def _region_projection(expr: Expr) -> LoopRegion | MeshRegion | None: """Return the region projected by a one-argument ``TupleGetItem``.""" if not ( @@ -696,7 +565,7 @@ def _module_callee_binding(target: HirFunction, child_entries: dict[int, str]) - def _emit_def( - fn: HirFunction, def_name: str, mesh_map: dict[int, str], indent: str, + fn: HirFunction, def_name: str, ctx: HirPrintContext, indent: str, options: PythonPrintOptions, child_entries: dict[int, str] | None = None, *, line_offset: int = 0, @@ -711,7 +580,14 @@ def _emit_def( """ child_entries = {} if child_entries is None else child_entries lines: list[str] = [] - + printer = HirPrinter() + root_mesh = ( + fn.body.mesh + if isinstance(fn.body, MeshRegion) and not fn.specializations + else None + ) + if root_mesh is not None: + ctx.push_mesh(root_mesh, "mesh") _counter = [0] _names: dict[int, str] = {} @@ -870,81 +746,28 @@ def _assign_name(expr: Expr) -> str: if isinstance(expr, LoopRegion): for carry in expr.carried_args: _assign_name(carry) + printer.bind_def(_names, _param_alias, child_entries, _moved_window) - def _tuple_literal(elements) -> str: - inner = ", ".join( - repr(el.value) if isinstance(el, Constant) else _expr_ref(el) - for el in elements - ) - if len(elements) == 1: - inner += "," - return f"({inner})" - - def _expr_ref(expr: Expr) -> str: - if id(expr) in _param_alias: - return _expr_ref(_param_alias[id(expr)]) - projection = _region_projection(expr) - if isinstance(projection, LoopRegion): - return _names[id(projection.carried_args[expr.target.index])] - if isinstance(expr, LoopRegion): - - - - carried = tuple(_names[id(carry)] for carry in expr.carried_args) - if len(carried) == 1: - return carried[0] - return "(" + ", ".join(carried) + ")" - if isinstance(projection, MeshRegion): - return _expr_ref(projection.body.elements[expr.target.index]) - if isinstance(expr, MeshRegion): - return _arg_ref(expr.body) - return _names[id(expr)] - - def _arg_ref(a) -> str: - - - - return _tuple_literal(a.elements) if isinstance(a, Tuple) else _expr_ref(a) - - def _start_ref(start, size, stride) -> str: - """One Slice start as source. - - A moved tile window prints as the move itself -- the offset is a - compile-time constant, so it belongs in the indexer rather than in a - statement of its own. - """ - moved = _moved_window(start, size, stride) - if moved is None: - return repr(start.value) if isinstance(start, Constant) else _expr_ref(start) - window, offset = moved - return _moved_window_ref(_expr_ref(window), offset) - - - - return_ty = fn.return_type + return_type = fn.return_type arrow = "" - if isinstance(return_ty, TensorType): - arrow = " -> " + _tensor_annotation( - return_ty, mesh_name_map=mesh_map, indent=indent - ) - elif not isinstance(return_ty, TupleType): + if isinstance(return_type, TensorType): + with printer.type_surface(indent=indent): + arrow = " -> " + printer.visit(return_type, ctx) + elif not isinstance(return_type, TupleType): arrow = " -> None" lines.append(f"def {def_name}(") - param_strs = [] - for p in fn.params: - name = _names[id(p)] - if isinstance(p.type, TensorType): - ann = _tensor_annotation( - p.type, mesh_name_map=mesh_map, indent=indent, is_const=p.is_const, - ) - param_strs.append(f"{indent}{name}: {ann}") + params: list[str] = [] + for param in fn.params: + name = _names[id(param)] + if isinstance(param.type, TensorType): + with printer.type_surface(indent=indent, const=param.is_const): + annotation = printer.visit(param.type, ctx) + params.append(f"{indent}{name}: {annotation}") else: - param_strs.append(f"{indent}{name}") - - - for index, text in enumerate(param_strs): - suffix = "," if index < len(param_strs) - 1 else "" + params.append(f"{indent}{name}") + for index, text in enumerate(params): + suffix = "," if index < len(params) - 1 else "" lines.extend((text + suffix).split("\n")) lines.append(f"){arrow}:") @@ -952,139 +775,14 @@ def _start_ref(start, size, stride) -> str: line = _constraint_line(param, indent, _names[id(param)]) if line is not None: lines.append(line) - - if fn.body is None: lines.append(f"{indent}pass") + if root_mesh is not None: + ctx.pop_mesh() return lines printed: set[int] = {id(param) for param in fn.params} - def _format_call(expr: Call, indent_here: str) -> str: - """Render a Call's RHS expression text. - - Render a Call's RHS expression text: the ``reshard(...)`` / - ``(...)`` special forms, else ``op_name(args, attr=val, - ...)``. Shared by the inline (tile-loop body) emitter and the - top-level emit loop so an attribute-rendering rule (``ShardLayout``, - ``DType``, ...) only needs one edit. A reshard that gathers back to the - whole names no mesh, so its target is a plain ``Layout`` and there is no - mesh reference to abbreviate. - """ - target = expr.target - args_str = ", ".join(_arg_ref(arg) for arg in expr.args) - if isinstance(target, Reshard): - layout_kw = "" - if isinstance(target.layout, ShardLayout): - layout_text = _shard_layout_str( - target.layout, - indent=indent_here + " ", - mesh_ref=( - mesh_map.get(id(target.layout.mesh)) - if target.layout.mesh.names - else None - ), - ) - layout_kw = ", layout=" + layout_text - elif target.layout is not None: - layout_kw = ", layout=" + _layout_str(target.layout, indent_here + " ") - storage = ( - f", storage={target.storage.name.lower()}" - if target.storage is not None - else "" - ) - return f"reshard({args_str}{layout_kw}{storage})" - if isinstance(target, HirFunction): - binding = _module_callee_binding(target, child_entries) - return f"{binding or target.name}({args_str})" - if isinstance(target, Slice): - indexers = [] - starts = expr.args[1] - if not isinstance(starts, Tuple): - raise ValueError("canonical_source: Slice starts must be a Tuple") - runtime_starts = False - for axis, (start, size, stride) in enumerate( - zip(starts.elements, target.sizes, target.strides) - ): - if _moved_window(start, size, stride) is not None: - indexers.append(_start_ref(start, size, stride)) - continue - dim = expr.args[0].type.shape[axis] - if ( - isinstance(start, Constant) - and start.value == 0 - and size == dim - and stride == 1 - ): - indexers.append(":") - continue - if not ( - isinstance(start, Constant) - and isinstance(start.value, int) - and isinstance(size, int) - and isinstance(stride, int) - ): - runtime_starts = True - break - begin = int(start.value) - stop = begin + size * stride - indexers.append( - f"{begin}:{stop}" if stride == 1 else f"{begin}:{stop}:{stride}" - ) - if runtime_starts: - start_refs = ", ".join( - _start_ref(start, size, stride) - for start, size, stride in zip( - starts.elements, target.sizes, target.strides - ) - ) - if len(starts.elements) == 1: - start_refs += "," - return ( - f"slice({_arg_ref(expr.args[0])}, ({start_refs}), " - f"sizes={_attr_tuple_str(target.sizes)}, " - f"strides={_attr_tuple_str(target.strides)})" - ) - return f"{_arg_ref(expr.args[0])}[{', '.join(indexers)}]" - - alias_name = _kinded_alias_name(target) - suppress_attrs = {"kind"} if alias_name is not None else set() - attr_strs = [] - for param in type(target).params(): - if param.kind != "attribute": - continue - value = getattr(target, param.name, None) - if value is None or param.name in suppress_attrs or param.name == "layout": - continue - if isinstance(value, str): - attr_strs.append(f'{param.name}="{value}"') - elif isinstance(value, DType): - attr_strs.append(f'{param.name}="{value.name}"') - elif isinstance(value, enum.Enum) and isinstance(value.value, str): - attr_strs.append(f'{param.name}="{value.value}"') - elif isinstance(value, float): - if math.isinf(value): - literal = "-1e999" if value < 0 else "1e999" - elif math.isnan(value): - literal = "(1e999 - 1e999)" - else: - literal = repr(value) - attr_strs.append(f"{param.name}={literal}") - elif isinstance(value, ShardLayout): - sl_str = _shard_layout_str(value, indent=indent_here + " ") - attr_strs.append(f"{param.name}={sl_str}") - elif isinstance(value, TensorType): - attr_strs.append(f"{param.name}={_compact_type(value, {})}") - elif isinstance(value, tuple): - attr_strs.append(f"{param.name}={_attr_tuple_str(value)}") - else: - attr_strs.append(f"{param.name}={value}") - - - - arguments = [_arg_ref(arg) for arg in expr.args] + attr_strs - return f"{_op_name(target)}({', '.join(arguments)})" - def _emit_inline_call(expr: Call, level: str) -> None: name = _names[id(expr)] if statements is not None: @@ -1092,77 +790,58 @@ def _emit_inline_call(expr: Call, level: str) -> None: value=name, line=line_offset + _physical_line_count(lines) + 1, ) + with printer.type_surface(indent=level): + rendered = printer.visit(expr, ctx) lines.append( - f"{level}{name} = {_format_call(expr, level)}" - f"{_comments(expr, options, mesh_map)}" + f"{level}{name} = {rendered}" + f"{_comments(expr, options, printer, ctx)}" ) printed.add(id(expr)) - class _ExprEmitter(ExprFunctor[None]): - def __init__(self, level: str) -> None: - super().__init__() - self.level = level - - def emit(self, expr: Expr) -> None: - self.visit(expr) - - def visit(self, expr, ctx=None): - key = id(expr) - if key in printed: - return None - if key in _inlined_start_ids: - printed.add(key) - return None - return super().visit(expr, ctx) - - def visit_Var(self, expr: Var, ctx=None) -> None: - printed.add(id(expr)) - - def visit_Constant(self, expr: Constant, ctx=None) -> None: + def _emit_expr(expr: Expr, level: str) -> None: + key = id(expr) + if key in printed: + return + if key in _inlined_start_ids: + printed.add(key) + return + if isinstance(expr, Var): + printed.add(key) + return + if isinstance(expr, Constant): lines.append( - f"{self.level}{_names[id(expr)]} = {repr(expr.value)}" - f"{_comments(expr, options, mesh_map)}" + f"{level}{_names[id(expr)]} = {repr(expr.value)}" + f"{_comments(expr, options, printer, ctx)}" ) - printed.add(id(expr)) - - def visit_Tuple(self, expr: Tuple, ctx=None) -> None: + printed.add(key) + return + if isinstance(expr, Tuple): for element in expr.elements: if not isinstance(element, Constant): - self.visit(element, ctx) - printed.add(id(expr)) - - def visit_LoopRegion(self, expr: LoopRegion, ctx=None) -> None: - _emit_loop_region(expr, self.level) - - def visit_MeshRegion(self, expr: MeshRegion, ctx=None) -> None: - _emit_mesh_region(expr, self.level) - - def visit_Call(self, expr: Call, ctx=None) -> None: + _emit_expr(element, level) + printed.add(key) + return + if isinstance(expr, LoopRegion): + _emit_loop_region(expr, level) + return + if isinstance(expr, MeshRegion): + _emit_mesh_region(expr, level) + return + if isinstance(expr, Call): projection = _region_projection(expr) if isinstance(projection, LoopRegion): - _emit_loop_region(projection, self.level) - printed.add(id(expr)) + _emit_loop_region(projection, level) + printed.add(key) return if isinstance(projection, MeshRegion): - _emit_mesh_region(projection, self.level) - printed.add(id(expr)) + _emit_mesh_region(projection, level) + printed.add(key) return for arg in expr.args: - self.visit(arg, ctx) - _emit_inline_call(expr, self.level) - - def default_visit(self, expr, ctx=None) -> None: - return None - - _expr_emitter = _ExprEmitter("") - - def _emit_expr(expr: Expr, level: str) -> None: - previous = _expr_emitter.level - _expr_emitter.level = level - try: - _expr_emitter.emit(expr) - finally: - _expr_emitter.level = previous + _emit_expr(arg, level) + _emit_inline_call(expr, level) + return + printed.add(key) def _emit_loop_region(region: LoopRegion, level: str) -> None: key = id(region) @@ -1172,23 +851,24 @@ def _emit_loop_region(region: LoopRegion, level: str) -> None: _emit_expr(init, level) for carry in region.carried_args: printed.add(id(carry)) - extent = shape_entry_str(region.extent) - step = shape_entry_str(region.step) - start = shape_entry_str(region.start) + extent = printer.visit(region.extent, ctx) + step = printer.visit(region.step, ctx) + start = printer.visit(region.start, ctx) if id(region.induction_var) in _tile_window_steps: + ctx.imports.add("from tilefoundry.dsl.tf import *") loop = f"tile({extent}, {step})" elif region.start == 0 and region.step == 1: loop = f"range({extent})" else: loop = f"range({start}, {extent}, {step})" - lines.append(f"{level}for {region.induction_var.name} in {loop}:{_comments(region, options, mesh_map)}") + lines.append(f"{level}for {region.induction_var.name} in {loop}:{_comments(region, options, printer, ctx)}") printed.add(key) inner = level + " " _emit_expr(region.body, inner) for value in region.yield_values: _emit_expr(value, inner) for carry, value in zip(region.carried_args, region.yield_values): - lines.append(f"{inner}{_names[id(carry)]} = {_expr_ref(value)}") + lines.append(f"{inner}{_names[id(carry)]} = {printer.reference(value)}") def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) -> None: key = id(region) @@ -1196,19 +876,30 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) return for arg in region.args: _emit_expr(arg, level) - mesh_name = mesh_map[id(region.mesh)] + if root_mesh is not None and region is fn.body: + printed.add(key) + _emit_expr(region.body, level) + if terminal: + lines.append(f"{level}return {printer.reference(region.body)}") + return + mesh_text = printer.visit(region.mesh, ctx) + mesh_name = ctx.scope_name(region.mesh) lines.append( - f"{level}with {mesh_name} as _{mesh_name}:" - f"{_comments(region, options, mesh_map)}" + f"{level}with {mesh_text} as {mesh_name}:" + f"{_comments(region, options, printer, ctx)}" ) printed.add(key) inner = level + " " - if terminal and isinstance(region.body, MeshRegion): - _emit_mesh_region(region.body, inner, terminal=True) - return - _emit_expr(region.body, inner) - if terminal: - lines.append(f"{inner}return {_arg_ref(region.body)}") + ctx.push_mesh(region.mesh, mesh_name) + try: + if terminal and isinstance(region.body, MeshRegion): + _emit_mesh_region(region.body, inner, terminal=True) + return + _emit_expr(region.body, inner) + if terminal: + lines.append(f"{inner}return {printer.reference(region.body)}") + finally: + ctx.pop_mesh() for expr in _order: if ( @@ -1231,7 +922,7 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) continue if isinstance(expr, Constant): name = _names[id(expr)] - lines.append(f"{indent}{name} = {repr(expr.value)}{_comments(expr, options, mesh_map)}") + lines.append(f"{indent}{name} = {repr(expr.value)}{_comments(expr, options, printer, ctx)}") line = _constraint_line(expr, indent, name) if line is not None: lines.append(line) @@ -1250,9 +941,11 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) value=name, line=line_offset + _physical_line_count(lines) + 1, ) + with printer.type_surface(indent=indent): + rendered = printer.visit(expr, ctx) lines.append( - f"{indent}{name} = {_format_call(expr, indent)}" - f"{_comments(expr, options, mesh_map)}" + f"{indent}{name} = {rendered}" + f"{_comments(expr, options, printer, ctx)}" ) line = _constraint_line(expr, indent, name) if line is not None: @@ -1263,190 +956,30 @@ def _emit_mesh_region(region: MeshRegion, level: str, *, terminal: bool = False) if not isinstance(fn.body, MeshRegion): if isinstance(fn.body, Tuple): - lines.append(f"{indent}return {_tuple_literal(fn.body.elements)}") + lines.append(f"{indent}return {printer.tuple_reference(fn.body.elements)}") elif isinstance(fn.body, LoopRegion): values = tuple(_names[id(carry)] for carry in fn.body.carried_args) result = values[0] if len(values) == 1 else "(" + ", ".join(values) + ")" lines.append(f"{indent}return {result}") else: - body_name = _expr_ref(fn.body) + body_name = printer.reference(fn.body) lines.append(f"{indent}return {body_name}") + if root_mesh is not None: + ctx.pop_mesh() return lines -_HIR_RENDERER = HirPrinter() -_dtype_str = _HIR_RENDERER.dtype_str -_mesh_name_map = _HIR_RENDERER.mesh_name_map -_pattern_ctor = _HIR_RENDERER.render_pattern - - -def _topologies_str(mesh: Mesh) -> str: - values = ", ".join( - f'Topology("{topology.name}", {shape_entry_str(topology.size)})' - for topology in mesh.topologies - ) - return f"({values}{',' if len(mesh.topologies) == 1 else ''})" - - -def _shape_tuple(shape: tuple) -> str: - return _HIR_RENDERER.shape_tuple(shape) - - -def _layout_str(layout: LayoutBase | None, indent: str = "") -> str: - return _HIR_RENDERER.render_layout(layout, HirPrintContext(), indent) - - -def _mesh_str(mesh: Mesh, indent: str = "") -> str: - return _HIR_RENDERER.render_mesh(mesh, HirPrintContext(), indent) - - -def _shard_layout_str(sl: ShardLayout, indent: str = "", *, mesh_ref=None) -> str: - ctx = HirPrintContext({id(sl.mesh): mesh_ref} if mesh_ref is not None else None) - return _HIR_RENDERER.render_shard_layout(sl, ctx, indent, mesh_ref=mesh_ref) - - -def _tensor_annotation(ty: TensorType, *, mesh_name_map=None, indent="", is_const=False) -> str: - return _HIR_RENDERER.render_tensor_type( - ty, HirPrintContext(mesh_name_map), indent, is_const - ) - - -def _collect_all_meshes( - fn: HirFunction, -) -> tuple[dict[int, Mesh], dict[int, Mesh]]: - """Meshes referenced by *fn* and every specialization variant. - Meshes referenced by *fn* and every specialization variant — the - printer's mesh-name map must stay stable across the base prototype and - each ``.specialize`` block. - """ - type_meshes: dict[int, Mesh] = {} - scope_meshes: dict[int, Mesh] = {} - for f in (fn, *fn.variants): - types, scopes = _collect_meshes(f) - type_meshes.update(types) - scope_meshes.update(scopes) - return type_meshes, scope_meshes - - -def _dedup_meshes(meshes: dict[int, Mesh]) -> dict[int, Mesh]: - """Collapse structurally identical descriptors before naming hoisted meshes.""" - result: dict[int, Mesh] = {} - for identity, mesh in meshes.items(): - signature = _mesh_str(mesh) - if any(signature == _mesh_str(existing) for existing in result.values()): - continue - result[identity] = mesh - return result - - -def _emit_header( - fn: HirFunction, - meshes: dict[int, Mesh], - mesh_map: dict[int, str], - indent: str, - *, - for_module: bool = False, - target: object | None = None, - dim_vars: "dict[str, object] | None" = None, - scope_mesh_ids: set[int] | None = None, -) -> list[str]: - """Import header + mesh-prelude shared by ``hir_function_to_python`` and ``_module_to_python``. - - Import header + mesh-prelude shared by ``hir_function_to_python`` and - ``_module_to_python`` — the only source for the imports/mesh-defs a - dispatch prototype needs (the conditional ``DimVarRangePat`` import for - ``fn.variants``, the ``ConstTensor``/``Tensor`` selection), so standalone - and module-wrapped output cannot drift out of sync. - """ - lines: list[str] = ["from __future__ import annotations", ""] +def _new_hir_context(*, for_module: bool = False, target=None) -> HirPrintContext: + """Create a HIR context with imports owned by the surrounding file.""" + ctx = HirPrintContext() if for_module: - lines.append("from tilefoundry.module import module") - lines.append("from tilefoundry import func") + ctx.imports.add("from tilefoundry.module import module") + ctx.imports.add("from tilefoundry import func") if target is not None: - rendered: PythonExpr = target.to_python() - lines.extend(rendered.imports) - lines.append("from tilefoundry.dsl.tf import * # noqa: F401, F403") - lines.append(f"from tilefoundry.dsl import {_tensor_import_names(fn)}") - lines.append("from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401") - shard_names = {"Layout", "Mesh", "Topology"} if meshes else set() - if for_module: - shard_names.add("Topology") - layouts = [] - functions = (fn, *fn.variants) - for current in functions: - for param in current.params: - if isinstance(param.type, TensorType): - layouts.append(param.type.layout) - if isinstance(current.return_type, TensorType): - layouts.append(current.return_type.layout) - for current in functions: - for expr in iter_exprs(current.body): - if isinstance(expr.type, TensorType): - layouts.append(expr.type.layout) - if isinstance(expr, Call): - for candidate in vars(expr.target).values() if hasattr(expr.target, "__dict__") else (): - if isinstance(candidate, LayoutBase): - layouts.append(candidate) - elif isinstance(candidate, TensorType): - layouts.append(candidate.layout) - def nested_layouts(layout): - if isinstance(layout, ComposedLayout): - yield layout - if layout.inner is not None: - yield from nested_layouts(layout.inner) - if layout.outer is not None: - yield from nested_layouts(layout.outer) - elif isinstance(layout, ShardLayout): - yield layout - yield from nested_layouts(layout.layout) - - all_layouts = [nested for layout in layouts for nested in nested_layouts(layout)] - if any(isinstance(layout, ShardLayout) for layout in all_layouts): - shard_names.add("ShardLayout") - if any(isinstance(layout, ComposedLayout) for layout in all_layouts): - shard_names.add("ComposedLayout") - for layout in all_layouts: - if isinstance(layout, ShardLayout): - for attr in layout.attrs: - shard_names.add({Broadcast: "B", Split: "S", Partial: "P"}.get(type(attr), "")) - shard_names.discard("") - shard_import = ", ".join(sorted(shard_names)) - if shard_import: - lines.append(f"from tilefoundry.ir.types.shard import {shard_import}") - if fn.variants: - lines.append("from tilefoundry.ir.core.pattern import DimVarRangePat") - if dim_vars: - lines.append("from tilefoundry.ir.types.dim import DimVar, ceildiv") - lines.append("") - - - - - - if dim_vars: - for name, var in dim_vars.items(): - lines.append(f'{name} = DimVar("{var.name}", {var.lo}, {var.hi})') - lines.append("") - - - if any(mesh.names or mid in (scope_mesh_ids or ()) for mid, mesh in meshes.items()): - for mid, mesh in meshes.items(): - if not mesh.names and mid not in (scope_mesh_ids or ()): - continue - name = mesh_map[mid] - topologies = _topologies_str(mesh) - names_repr = repr(tuple(mesh.names)) if mesh.names else "()" - lines.append( - f"{name} = Mesh(" - f"{topologies}, " - f"{_layout_str(mesh.layout)}, " - f"names={names_repr}" - f")" - ) - lines.append("") - return lines - + rendered = target.to_python() + ctx.imports.update(rendered.imports) + return ctx def _variant_binding_name(variant: HirFunction) -> str: """Return a valid source binding for a variant without display metadata.""" @@ -1458,7 +991,7 @@ def _variant_binding_name(variant: HirFunction) -> str: def _emit_decorated_defs( - fn: HirFunction, mesh_map: dict[int, str], indent: str, options: PythonPrintOptions, + fn: HirFunction, ctx: HirPrintContext, indent: str, options: PythonPrintOptions, child_entries: dict[int, str] | None = None, *, line_offset: int = 0, @@ -1471,12 +1004,20 @@ def _emit_decorated_defs( render identically. See [inspection §2.6](docs/spec/inspection.md#26-specialization-printing). """ - lines: list[str] = ["@func"] + printer = HirPrinter() + decorator = "@func" + if isinstance(fn.body, MeshRegion) and not fn.specializations: + root_ctx = HirPrintContext() + with printer.type_surface(indent=indent): + mesh_text = printer.visit(fn.body.mesh, root_ctx) + ctx.imports.update(root_ctx.imports) + decorator = f"@func(mesh={mesh_text})" + lines: list[str] = [decorator] lines.extend( _emit_def( fn, fn.name, - mesh_map, + ctx, indent, options, child_entries, @@ -1489,11 +1030,11 @@ def _emit_decorated_defs( for variant in fn.variants: lines.append("") lines.append( - f"@{fn.name}.specialize({_pattern_ctor(variant.specializations[0])})" + f"@{fn.name}.specialize({HirPrinter().render_pattern(variant.specializations[0], ctx)})" ) lines.extend( _emit_def( - variant, _variant_binding_name(variant), mesh_map, indent, options, + variant, _variant_binding_name(variant), ctx, indent, options, child_entries, line_offset=line_offset + _physical_line_count(lines), statements=statements, @@ -1509,34 +1050,27 @@ def _render_hir_function( A normal function prints as a single ``@func``. A dispatch prototype (``variants != ()``) prints as a ``pass``-bodied ``@func`` base followed by - one ``@.specialize(pattern)`` block per variant. When the function - uses meshes with named axes, compact sugar form is emitted; otherwise the - verbose ``ShardLayout(...)`` form is used. + one ``@.specialize(pattern)`` block per variant. Placement sugar is + emitted only where an explicit mesh-scope binding dominates its use. """ indent = " " - type_meshes, scope_meshes = _collect_all_meshes(fn) - meshes = {**type_meshes, **scope_meshes} - mesh_map = _mesh_name_map(meshes) - lines = _emit_header( + ctx = _new_hir_context() + statements: dict[int, _PrintedStatement] = {} + lines = _emit_decorated_defs( fn, - meshes, - mesh_map, + ctx, indent, - dim_vars=dim_vars_reached(fn), - scope_mesh_ids=set(scope_meshes), - ) - statements: dict[int, _PrintedStatement] = {} - lines.extend( - _emit_decorated_defs( - fn, - mesh_map, - indent, - options or PythonPrintOptions(), - line_offset=_physical_line_count(lines), - statements=statements, - ) + options or PythonPrintOptions(), + line_offset=0, + statements=statements, ) - return _PythonRendering("\n".join(lines) + "\n", statements) + header = ctx.header() + header_lines = _physical_line_count(header) + statements = { + identity: _PrintedStatement(value=item.value, line=item.line + header_lines) + for identity, item in statements.items() + } + return _PythonRendering("\n".join(header + lines) + "\n", statements) def hir_function_to_python( @@ -1576,7 +1110,21 @@ def module_to_python(fn: HirFunction, module_name: str = "M") -> str: def _module_hir_functions(mod: Module) -> tuple[HirFunction, ...]: """The Module's HIR functions.""" - return tuple(fn for fn in mod.functions if isinstance(fn, HirFunction)) + return tuple(fn for fn in _emission_order(mod) if isinstance(fn, HirFunction)) + + +def _emission_order(mod: Module) -> tuple: + """A Module's functions in the order the printed class body binds them. + + The entry goes last: a body calling a sibling names the attribute the class + body already bound, so every callee must be written before it. Mesh + traversal reads the same order, so output does not depend on the order the + authored source happened to use. + """ + functions = mod.functions + entry = mod.entry_function() if functions and mod.entry is not None else None + ordered = tuple(fn for fn in functions if fn is not entry) + return ordered + (entry,) if entry is not None else ordered def _module_tree_functions(mod: Module) -> tuple[HirFunction, ...]: @@ -1587,19 +1135,22 @@ def _module_tree_functions(mod: Module) -> tuple[HirFunction, ...]: return tuple(functions) -def _module_decorator_line(mod: Module, entry_name: str | None) -> str: +def _module_decorator_line(mod: Module, entry_name: str | None, ctx: HirPrintContext) -> str: """Render the context this Module declares as an ``@module(...)`` line. Always the called form. A bare decorator has not run while the class body is evaluated, so a body naming a child call could not resolve it. """ kwargs: list[str] = [] if entry_name is None else [f'entry="{entry_name}"'] + printer = HirPrinter() if mod.target is not None: rendered: PythonExpr = mod.target.to_python() + ctx.imports.update(rendered.imports) kwargs.append(f"target={rendered.text}") if mod.topologies is not None: + ctx.imports.add("from tilefoundry.ir.types.shard import Topology") topo_strs = [ - f'Topology("{t.name}", {shape_entry_str(t.size)})' + f'Topology("{t.name}", {printer.visit(t.size, ctx)})' for t in mod.topologies ] rendered_topologies = f'({", ".join(topo_strs)},)' if topo_strs else "()" @@ -1608,7 +1159,7 @@ def _module_decorator_line(mod: Module, entry_name: str | None) -> str: def _emit_module_class( - mod: Module, module_name: str, mesh_map: dict[int, str], indent: str, + mod: Module, module_name: str, ctx: HirPrintContext, indent: str, options: PythonPrintOptions, ) -> list[str]: """One ``@module`` class block: its nested Modules, then its functions. @@ -1616,26 +1167,24 @@ def _emit_module_class( Children first, because a body calling one names the attribute it is bound to and a class body binds in the order it is written. """ - functions = mod.functions - entry = mod.entry_function() if functions and mod.entry is not None else None - lines = [_module_decorator_line(mod, mod.entry), f"class {module_name}:"] - ordered = tuple(fn for fn in functions if fn is not entry) - if entry is not None: - ordered += (entry,) + lines = [_module_decorator_line(mod, mod.entry, ctx), f"class {module_name}:"] + ordered = _emission_order(mod) child_entries = { id(child.entry_function()): child.name for child in mod.modules if child.entry is not None and isinstance(child.entry_function(), HirFunction) } blocks: list[list[str]] = [ - _emit_module_class(child, child.name, mesh_map, indent, options) + _emit_module_class(child, child.name, ctx, indent, options) for child in mod.modules ] for fn in ordered: if isinstance(fn, HirFunction): - blocks.append(_emit_decorated_defs(fn, mesh_map, indent, options, child_entries)) + blocks.append(_emit_decorated_defs(fn, ctx, indent, options, child_entries)) elif isinstance(fn, PrimFunction): - blocks.append(_tir_function_block(fn)) + tir_block = _tir_function_block(fn) + ctx.imports.update(tir_block.imports) + blocks.append(tir_block) else: raise TypeError(f"Python printer cannot serialize {type(fn).__name__}") for index, block in enumerate(blocks): @@ -1668,50 +1217,13 @@ def _module_to_python( raise TypeError("Module printer requires a function entry") - header_of = entry if entry is not None else functions[0] indent4 = " " - type_meshes: dict[int, Mesh] = {} - scope_meshes: dict[int, Mesh] = {} - for fn in functions: - types, scopes = _collect_all_meshes(fn) - type_meshes.update(types) - scope_meshes.update(scopes) - meshes = {**type_meshes, **scope_meshes} - mesh_map = _mesh_name_map(meshes) - - - - dim_vars: dict[str, object] = {} - for fn in functions: - dim_vars.update(dim_vars_reached(fn)) - for node in _module_tree(root): - dim_vars.update(dim_vars_by_name(node.topologies or ())) - lines = _emit_header( - header_of, meshes, mesh_map, indent4, for_module=True, target=root.target, - dim_vars=dim_vars, scope_mesh_ids=set(scope_meshes), - ) - if any(isinstance(fn, PrimFunction) for node in _module_tree(root) for fn in node.functions): - lines = [ - line.replace("from tilefoundry import func", "from tilefoundry import func, prim_func") - for line in lines - ] - tensor_line = next(i for i, line in enumerate(lines) if line.startswith("from tilefoundry.dsl import ")) - lines[tensor_line] = lines[tensor_line].replace("import ", "import T, ") - target_imports = sorted({fn.target.to_python().imports[0] for node in _module_tree(root) for fn in node.functions if isinstance(fn, PrimFunction)}) - lines[2:2] = target_imports - tensor_names = "ConstTensor, Tensor" if any( - param.is_const for fn in functions for param in fn.params - ) else "Tensor" - lines = [ - f"from tilefoundry.dsl import {tensor_names}" if line.startswith("from tilefoundry.dsl import Tensor") else line - for line in lines - ] - lines.extend( - _emit_module_class( - root, module_name, mesh_map, indent4, options or PythonPrintOptions(), - ) + ctx = _new_hir_context(for_module=True, target=root.target) + lines = _emit_module_class( + root, module_name, ctx, indent4, options or PythonPrintOptions(), ) - return "\n".join(lines) + "\n" + header = ctx.header() + return "\n".join(header + lines) + "\n" def _module_tree(root: Module) -> Iterator[Module]: diff --git a/src/tilefoundry/inspection/tir_printer.py b/src/tilefoundry/inspection/tir_printer.py index f85ae987..5ca8a4df 100644 --- a/src/tilefoundry/inspection/tir_printer.py +++ b/src/tilefoundry/inspection/tir_printer.py @@ -40,81 +40,90 @@ def __init__(self, *, context: TirPrintContext | None = None, indent: str = "") self.context = context or TirPrintContext() self.indent = indent - def dim_entry(self, value, ctx=None) -> str: - return f"_{value.name}" if hasattr(value, "name") else str(value) - - def visit(self, stmt, ctx=None): # type: ignore[override] - return StmtVisitor.visit(self, stmt) - - def visit_Sequential(self, stmt): + def visit_DimVar(self, value, ctx=None) -> str: + return f"_{value.name}" + + def dim_entry(self, value, ctx=None, *, nested: bool = False) -> str: + return super().dim_entry(value, ctx, nested=nested) + + def visit(self, node, ctx=None): # type: ignore[override] + """Dispatch statements, expressions, and types by their concrete name.""" + return PythonPrinter.visit(self, node, ctx) + + def visit_Var(self, expr: Var, ctx=None) -> str: + return expr.name + + def visit_Constant(self, expr: Constant, ctx=None) -> str: + return repr(expr.value) + + def visit_SymbolRef(self, expr: SymbolRef, ctx=None) -> str: + return _binding_name(expr.name) + + def visit_ShapeOf(self, expr: ShapeOf, ctx=None) -> str: + return f"shape_of({expr.param.name}, axis={expr.axis})" + + def visit_Op(self, expr: Op, ctx=None) -> str: + name = getattr(getattr(expr, "_op_schema", None), "name", type(expr).__name__.lower()) + self.context.use(PythonExpr(("from tilefoundry.dsl import T",), "T")) + return f"T.{name}" + + def visit_program_call(self, expr: Call, ctx=None) -> str: + target = expr.target + scalar_binary = { + BinaryKind.EQ: "==", BinaryKind.NE: "!=", BinaryKind.LT: "<", + BinaryKind.LE: "<=", BinaryKind.GT: ">", BinaryKind.GE: ">=", BinaryKind.AND: "and", + } + kind = getattr(target, "kind", None) + if kind in scalar_binary and len(expr.args) == 2 and expr.type.dtype is DType.bool: + return f"{self.visit(expr.args[0])} {scalar_binary[kind]} {self.visit(expr.args[1])}" + name = getattr(getattr(target, "_op_schema", None), "name", None) or re.sub( + r"(?", BinaryKind.GE:">=", BinaryKind.AND:"and"} - kind = getattr(target, "kind", None) - if kind in scalar_binary and len(expr.args)==2 and expr.type.dtype is DType.bool: - return f"{self._expr(expr.args[0])} {scalar_binary[kind]} {self._expr(expr.args[1])}" - name = getattr(getattr(target, "_op_schema", None), "name", None) or re.sub(r"(? list[str]: callee, grid = stmt.args[0], stmt.args[1:4] block = stmt.args[4:7] forwarded = stmt.args[7:] - return [f"{indent}launch({printer._expr(callee)}, {printer._join_args(forwarded)}, grid={printer._expr(Tuple(type=grid[0].type, elements=tuple(grid)))}, block={printer._expr(Tuple(type=block[0].type, elements=tuple(block)))}) # noqa: F821"] + return [f"{indent}launch({printer.visit(callee)}, {printer._join_args(forwarded)}, grid={printer.visit(Tuple(type=grid[0].type, elements=tuple(grid)))}, block={printer.visit(Tuple(type=block[0].type, elements=tuple(block)))}) # noqa: F821"] @register_tir_printer(Op) @@ -162,8 +171,8 @@ def _print_op_evaluate(stmt: Evaluate, printer: TirPrinter) -> list[str]: continue rendered = printer.render_value(value, printer.context, printer.indent + " ") attrs.append(rendered if op_name == "sync" and p.name == "mesh" else f"{p.name}={rendered}") - rendered_args = [printer._expr(arg) for arg in args] - return [f"{indent}{printer._expr(target)}({', '.join(rendered_args + attrs)})"] + rendered_args = [printer.visit(arg) for arg in args] + return [f"{indent}{printer.visit_Op(target)}({', '.join(rendered_args + attrs)})"] def _function_block(fn: PrimFunction) -> list[str]: @@ -177,7 +186,7 @@ def _function_block(fn: PrimFunction) -> list[str]: lines = [f'_{d.name} = DimVar("{d.name}", {d.lo}, {d.hi})' for d in dim_vars.values()] lines.append("@prim_func(target=" + target + ")") params = ", ".join( - f"{p.name}: {TirPrinter(context=ctx).render_value(p.type, ctx) if isinstance(p.type, TensorType) else repr(p.type)}" + f"{p.name}: {TirPrinter(context=ctx).visit(p.type, ctx) if isinstance(p.type, TensorType) else repr(p.type)}" for p in fn.params ) lines.append(f"def {_binding_name(fn.name)}({params}):") diff --git a/src/tilefoundry/inspection/viewer/builder.py b/src/tilefoundry/inspection/viewer/builder.py index d8dd7e51..ea354944 100644 --- a/src/tilefoundry/inspection/viewer/builder.py +++ b/src/tilefoundry/inspection/viewer/builder.py @@ -6,31 +6,18 @@ import graphviz -from tilefoundry.inspection.python_printer import ( - _DIM_FUNC_OPS, - _DIM_INFIX_OPS, - _collect_meshes, - _mesh_name_map, - _op_display_name, - _shard_layout_str, - _shard_layout_surface_str, - _tensor_annotation, - shape_entry_str, - shard_compact_inline, -) +from tilefoundry.inspection.print_context import HirPrintContext +from tilefoundry.inspection.printer_base import PythonPrinter from tilefoundry.ir.core import Tuple as HirTuple from tilefoundry.ir.core.expr import Call, Constant, Var from tilefoundry.ir.hir.function import Function as HirFunction +from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.types import DType -from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard.mesh import Mesh from tilefoundry.ir.types.shard.shard_layout import ShardLayout from tilefoundry.ir.types.tensor_type import TensorType, TupleType -from tilefoundry.ir.visitor import ExprVisitor from .htmltable import Cell, Span, Table from .palette import ( - DIMVAR_COLOR, HAIR, INK, MUTED, @@ -38,7 +25,6 @@ depth_border, depth_fill, exprkind_color, - storage_color, ) @@ -68,25 +54,6 @@ def _renderable_functions(root) -> list[tuple[str, "HirFunction"]]: return out -def _collect_view_meshes(root) -> dict[int, "Mesh"]: - """Collect unique ``Mesh`` objects referenced anywhere in *root*. - - Collect unique ``Mesh`` objects referenced anywhere in *root* — params, - return type, every node's result type, and ``Reshard`` layout attrs. - - The viewer renders shard sugar on intermediate node result types too, so it - needs a wider mesh-name map than the printer's param/return/Reshard scan: - built on the shared ``python_printer._collect_meshes`` with - ``include_node_types=True``. - """ - meshes: dict[int, Mesh] = {} - for _label, fn in _renderable_functions(root): - type_meshes, scope_meshes = _collect_meshes(fn, include_node_types=True) - meshes.update(type_meshes) - meshes.update(scope_meshes) - return meshes - - @dataclass class DetailRef: """Minimal click-lookup reference. NOT a graph/IR model. @@ -105,11 +72,11 @@ class DetailRef: class DetailIndex: """Detail lookup index only. ``dict[visual_id, DetailRef]``. - Carries the per-view ``mesh_name_map`` so the on-demand detail endpoint can - render shard sugar with the same stable mesh names as the graph. + Carries the per-view print context so detail text uses the same bindings + as the graph. """ entries: dict[str, DetailRef] = field(default_factory=dict) - mesh_name_map: dict[int, str] = field(default_factory=dict) + context: HirPrintContext = field(default_factory=HirPrintContext) def add(self, visual_id: str, ref: DetailRef) -> None: @@ -128,286 +95,124 @@ def get(self, visual_id: str) -> DetailRef | None: def _format_constant(c: Constant) -> str: - """Compact constant rendering (ported from the old viewer's pretty view). - - Compact constant rendering (ported from the old viewer's pretty - view). Scalars: ``const(1)`` / ``const(1.0f)``; sequences: - ``const([1.0f, 2.0f, ...])`` truncated to the first 8 elements. - """ ty = getattr(c, "type", None) suffix = ( _CONST_DTYPE_SUFFIX.get(ty.dtype.name, "") - if isinstance(ty, TensorType) and isinstance(ty.dtype, DType) else "" + if isinstance(ty, TensorType) and isinstance(ty.dtype, DType) + else "" ) - def _fmt(v) -> str: - if isinstance(v, bool): - return repr(v) - if isinstance(v, float): - return f"{v}{suffix}" - return repr(v) + def format_value(value): + if isinstance(value, bool): + return repr(value) + if isinstance(value, float): + return f"{value}{suffix}" + return repr(value) - val = c.value + value = c.value if isinstance(ty, TensorType) and ty.shape == (): - return f"const({_fmt(val)})" + return f"const({format_value(value)})" try: - items = list(val) + items = list(value) except TypeError: - return f"const({_fmt(val)})" - head = ", ".join(_fmt(v) for v in items[:8]) + return f"const({format_value(value)})" + head = ", ".join(format_value(item) for item in items[:8]) tail = ", ..." if len(items) > 8 else "" return f"const([{head}{tail}])" -def _shard_layout_text(sl: ShardLayout, mesh_name_map: dict[int, str] | None) -> str: - """Shard layout text. - - Render a bare ``ShardLayout`` attr (e.g. a ``Reshard`` layout) through the - canonical sugar core, falling back to the verbose ``ShardLayout(...)`` form - when the mesh is unnamed or the layout is not sugar-expressible. - """ - mesh_name = mesh_name_map.get(id(sl.mesh)) if mesh_name_map else None - if mesh_name and getattr(sl.mesh, "names", None): - mesh_unique = mesh_name_map is not None and len(mesh_name_map) == 1 - sugar = _shard_layout_surface_str(sl, mesh_name=mesh_name, mesh_unique=mesh_unique) - if sugar is not None: - return sugar - return _shard_layout_str(sl) - +def _type_text(ty, context=None) -> str: + printer = PythonPrinter() + if context is None: + context = HirPrintContext() + with printer.type_surface(): + return printer.visit(ty, context) -def _pretty_attr_value(value, *, full: bool = False, mesh_name_map: dict[int, str] | None = None) -> str: - """Readable rendering of an Op attribute value. - Readable rendering of an Op attribute value — pretty constants / - tuples / lists / types instead of raw ``repr``. ``full`` selects the - canonical type form (detail panel) vs the compact one (graph label). - """ +def _pretty_attr_value(value, *, context=None) -> str: if isinstance(value, Constant): return _format_constant(value) if isinstance(value, DType): return value.name - if isinstance(value, (TensorType, TupleType)): - return ( - type_to_canonical_pretty(value, mesh_name_map=mesh_name_map) - if full - else type_to_compact_pretty(value, mesh_name_map=mesh_name_map) - ) - if isinstance(value, ShardLayout): - return _shard_layout_text(value, mesh_name_map) + if isinstance(value, (TensorType, TupleType, ShardLayout)): + return _type_text(value, context) if isinstance(value, tuple): - inner = ", ".join(_pretty_attr_value(v, full=full, mesh_name_map=mesh_name_map) for v in value) + inner = ", ".join(_pretty_attr_value(v, context=context) for v in value) return f"({inner}{',' if len(value) == 1 else ''})" if isinstance(value, list): - return "[" + ", ".join(_pretty_attr_value(v, full=full, mesh_name_map=mesh_name_map) for v in value) + "]" + return "[" + ", ".join(_pretty_attr_value(v, context=context) for v in value) + "]" if isinstance(value, str): return value return repr(value) -def _op_attributes( - target, *, full: bool = False, mesh_name_map: dict[int, str] | None = None -) -> list[tuple[str, str]]: - """Op attributes. - - The Op's non-input (attribute) params as ``(name, pretty-value)`` - pairs — e.g. ``("axis", "2")`` / ``("begin", "(const(0), ...)")``. - ``full`` selects canonical (detail) vs compact (graph) type text. - Empty when the Op has no attributes or doesn't expose ``params()``. - """ +def _op_attributes(target, *, context=None) -> list[tuple[str, str]]: try: pdefs = type(target).params() except (AttributeError, TypeError): return [] - out = [] - for p in pdefs: - if getattr(p, "kind", None) == "attribute": - out.append(( - p.name, - _pretty_attr_value(getattr(target, p.name, None), full=full, mesh_name_map=mesh_name_map), - )) - return out + return [ + (p.name, _pretty_attr_value(getattr(target, p.name, None), context=context)) + for p in pdefs + if getattr(p, "kind", None) == "attribute" + ] -class _DimSpanVisitor(ExprVisitor[list[Span]]): - def visit_DimVar(self, dim: DimVar, ctx=None) -> list[Span]: - return [Span(text=dim.name, color=DIMVAR_COLOR, bold=True)] - - def visit_Constant(self, dim: Constant, ctx=None) -> list[Span]: - return [Span(text=str(dim.value))] - - def visit_Call(self, dim: Call, ctx=None) -> list[Span]: - target = dim.target - for op_cls, sym in _DIM_INFIX_OPS.items(): - if isinstance(target, op_cls): - spans = list(self.visit(dim.args[0], ctx)) - spans.append(Span(text=f" {sym} ")) - spans.extend(self.visit(dim.args[1], ctx)) - return spans - for op_cls, fname in _DIM_FUNC_OPS.items(): - if isinstance(target, op_cls): - spans = [Span(text=f"{fname}(")] - for i, arg in enumerate(dim.args): - if i: - spans.append(Span(text=", ")) - spans.extend(self.visit(arg, ctx)) - spans.append(Span(text=")")) - return spans - return [Span(text=shape_entry_str(dim))] - - def default_visit(self, dim, ctx=None) -> list[Span]: - return [Span(text=shape_entry_str(dim))] - - -def _format_dim(dim) -> list[Span]: - """Render a shape dimension as inline viewer spans. - - Share arithmetic syntax with the canonical Python printer. Dimension - variables use one token-class color while their text and detail identify the - symbol; integers remain plain. - See [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). - """ - return _DimSpanVisitor().visit(dim) - - -def _shard_inline(ty: TensorType, mesh_name_map: dict[int, str] | None): - """Compact shard decomposition for *ty*. - - Compact shard decomposition for *ty*: ``(split_ref_by_tensor_axis, - partials)`` or ``None`` when there is no named sugar-expressible shard - layout (caller renders the plain shape). - """ - layout = getattr(ty, "layout", None) - if not isinstance(layout, ShardLayout): - return None - mesh_name = mesh_name_map.get(id(layout.mesh)) if mesh_name_map else None - if not (mesh_name and getattr(layout.mesh, "names", None)): - return None - return shard_compact_inline(layout, mesh_name, ty.shape) - - -def _compact_type_spans(ty, mesh_name_map: dict[int, str] | None = None) -> list[Span]: - """Render compact graph-label types as colored inline spans. - - Tint dimension variables and storage, inline shard splits on tensor axes, - and append partial states. Layouts that cannot be inlined use canonical - annotation text. ``type_to_compact_pretty`` joins spans as plain text. - See [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms). - """ - if isinstance(ty, TensorType): - inline = _shard_inline(ty, mesh_name_map) - if isinstance(ty.layout, ShardLayout) and inline is None: - - return [Span(text=type_to_canonical_pretty(ty, mesh_name_map=mesh_name_map))] - split_ref, partials = inline if inline is not None else ({}, []) - dtype = ty.dtype.name if hasattr(ty.dtype, "name") else str(ty.dtype) - spans: list[Span] = [Span(text=f"{dtype}[")] - for i, d in enumerate(ty.shape): - if i: - spans.append(Span(text=", ")) - spans.extend(_format_dim(d)) - if i in split_ref: - spans.append(Span(text=f" @ {split_ref[i]}")) - spans.append(Span(text="]")) - if partials: - spans.append(Span(text=" {" + ", ".join(partials) + "}")) - storage = getattr(ty, "storage", None) - if storage: - spans.append(Span(text=" @")) - spans.append(Span(text=str(storage), color=storage_color(str(storage)), bold=True)) - return spans - if isinstance(ty, TupleType): - - - spans: list[Span] = [Span(text="⟨")] - for i, sub in enumerate(ty.fields): - if i: - spans.append(Span(text=", ")) - spans.extend(_compact_type_spans(sub, mesh_name_map)) - spans.append(Span(text="⟩")) - return spans - return [Span(text=str(ty))] - - -def type_to_compact_pretty(ty, mesh_name_map: dict[int, str] | None = None) -> str: - """Render [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms) compact text.""" - return "".join(s.text for s in _compact_type_spans(ty, mesh_name_map)) - - -def type_to_canonical_pretty(ty, mesh_name_map: dict[int, str] | None = None) -> str: - """Render [inspection §2.3](docs/spec/inspection.md#23-dsl-text-forms) canonical text.""" - if isinstance(ty, TensorType): - return _tensor_annotation(ty, mesh_name_map=mesh_name_map) - if isinstance(ty, TupleType): - return "(" + ", ".join(type_to_canonical_pretty(f, mesh_name_map) for f in ty.fields) + ")" - return str(ty) - - -def _returns_of(ty, mesh_name_map: dict[int, str] | None = None) -> list[dict]: +def _returns_of(ty, context=None) -> list[dict]: if isinstance(ty, TupleType): return [ - {"idx": i, "type": type_to_canonical_pretty(f, mesh_name_map)} - for i, f in enumerate(ty.fields) + {"idx": i, "type": _type_text(field, context)} + for i, field in enumerate(ty.fields) ] if ty is None: return [] - return [{"idx": 0, "type": type_to_canonical_pretty(ty, mesh_name_map)}] + return [{"idx": 0, "type": _type_text(ty, context)}] def format_detail( - visual_id: str, ref: "DetailRef", mesh_name_map: dict[int, str] | None = None + visual_id: str, ref: "DetailRef", context: HirPrintContext | None = None ) -> dict: - """Format a detail-panel payload from a ``DetailRef`` on demand. - - Format a detail-panel payload from a ``DetailRef`` on demand (no - pre-baked JSON in the index). Shape: - ``{id, kind, name, params:[{name,type}], returns:[{idx,type}], attrs:[{key,value}]}``. - - ``mesh_name_map`` (the per-view map carried on ``DetailIndex``) lets the - canonical type text render shard sugar with the same stable mesh names as - the graph. - """ + """Format a detail-panel payload from a live HIR reference.""" + context = context or HirPrintContext() expr = ref.hir_expr name = ref.kind params: list[dict] = [] attrs: list[dict] = [] returns: list[dict] = [] - mm = mesh_name_map if isinstance(expr, HirFunction): name = expr.name - params = [{"name": p.name, "type": type_to_canonical_pretty(p.type, mm)} for p in expr.params] - returns = _returns_of(expr.return_type, mm) + params = [{"name": p.name, "type": _type_text(p.type, context)} for p in expr.params] + returns = _returns_of(expr.return_type, context) elif isinstance(expr, Var): name = expr.name - returns = _returns_of(expr.type, mm) + returns = _returns_of(expr.type, context) elif isinstance(expr, Constant): name = _format_constant(expr) - returns = _returns_of(expr.type, mm) + returns = _returns_of(expr.type, context) elif isinstance(expr, HirTuple): name = "Tuple" - params = [{"name": f"e{i}", "type": type_to_canonical_pretty(el.type, mm)} - for i, el in enumerate(expr.elements)] - returns = _returns_of(expr.type, mm) + params = [{"name": f"e{i}", "type": _type_text(item.type, context)} for i, item in enumerate(expr.elements)] + returns = _returns_of(expr.type, context) elif isinstance(expr, Call): - tgt = expr.target - if isinstance(tgt, HirFunction): - name = tgt.name - pnames = [p.name for p in tgt.params] + target = expr.target + if isinstance(target, HirFunction): + name = target.name + pnames = [p.name for p in target.params] else: - name = _op_display_name(tgt) + name = _op_display_name(target) try: - pnames = [p.name for p in type(tgt).params() if p.kind == "input"] + pnames = [p.name for p in type(target).params() if p.kind == "input"] except (AttributeError, TypeError): pnames = [] - attrs = [{"key": k, "value": v} for k, v in _op_attributes(tgt, full=True, mesh_name_map=mm)] + attrs = [{"key": key, "value": value} for key, value in _op_attributes(target, context=context)] params = [ - {"name": pnames[i] if i < len(pnames) else f"in{i}", - "type": type_to_canonical_pretty(a.type, mm)} - for i, a in enumerate(expr.args) + {"name": pnames[i] if i < len(pnames) else f"in{i}", "type": _type_text(arg.type, context)} + for i, arg in enumerate(expr.args) ] - returns = _returns_of(expr.type, mm) + returns = _returns_of(expr.type, context) else: - returns = _returns_of(getattr(expr, "type", None), mm) + returns = _returns_of(getattr(expr, "type", None), context) return {"id": visual_id, "kind": ref.kind, "name": name, "params": params, "returns": returns, "attrs": attrs} @@ -427,8 +232,11 @@ def __init__(self, root, collapsed: set[str] | None = None) -> None: self.collapsed = set(collapsed or ()) - self.mesh_name_map = _mesh_name_map(_collect_view_meshes(root)) - self.index = DetailIndex(mesh_name_map=self.mesh_name_map) + self.printer = PythonPrinter() + self.context = HirPrintContext() + if isinstance(root, HirFunction) and isinstance(root.body, MeshRegion): + self.context.push_mesh(root.body.mesh, "mesh") + self.index = DetailIndex(context=self.context) @@ -660,7 +468,7 @@ def _emit_function_node( width = 2 + n_params title.add_row( Cell( - spans=tuple(_compact_type_spans(fn.return_type, self.mesh_name_map)), + spans=(Span(text=_type_text(fn.return_type, self.context)),), colspan=width, color=INK, bold=True, font_size=12, ) ) @@ -716,7 +524,7 @@ def _walk_expr( Cell(text=f"Var {expr.name}", href="javascript:void(0)", title=f"expr:{vid}", bgcolor=exprkind_color("Var"), color="#ffffff", bold=True) ) - tbl.add_row(Cell(spans=tuple(_compact_type_spans(expr.type, self.mesh_name_map)), color=MUTED, font_size=11)) + tbl.add_row(Cell(spans=(Span(text=_type_text(expr.type, self.context)),), color=MUTED, font_size=11)) if output_slot is not None: tbl.add_row(self._output_marker_row(output_slot, 1)) g.node(vid, label=tbl.to_html()) @@ -733,7 +541,7 @@ def _walk_expr( Cell(text=_format_constant(expr), href="javascript:void(0)", title=f"expr:{vid}", bgcolor=exprkind_color("Constant"), color="#ffffff", bold=True) ) - tbl.add_row(Cell(spans=tuple(_compact_type_spans(expr.type, self.mesh_name_map)), color=MUTED, font_size=11)) + tbl.add_row(Cell(spans=(Span(text=_type_text(expr.type, self.context)),), color=MUTED, font_size=11)) if output_slot is not None: tbl.add_row(self._output_marker_row(output_slot, 1)) g.node(vid, label=tbl.to_html()) @@ -847,14 +655,14 @@ def _emit_call( - for key, val in _op_attributes(call.target, mesh_name_map=self.mesh_name_map): + for key, val in _op_attributes(call.target, context=self.context): tbl.add_row( Cell(text=f"{key}: {val}", colspan=width, color=MUTED, font_size=11, align="LEFT") ) tbl.add_row( Cell( - spans=tuple(_compact_type_spans(call.type, self.mesh_name_map)), + spans=(Span(text=_type_text(call.type, self.context)),), colspan=width, color=INK, bold=True, @@ -963,5 +771,5 @@ def _emit_tuple( __all__ = [ "DetailIndex", "DetailRef", "ViewerBuilder", - "format_detail", "type_to_compact_pretty", "type_to_canonical_pretty", + "format_detail", ] diff --git a/src/tilefoundry/inspection/viewer/server.py b/src/tilefoundry/inspection/viewer/server.py index afb651b0..9bae9ab1 100644 --- a/src/tilefoundry/inspection/viewer/server.py +++ b/src/tilefoundry/inspection/viewer/server.py @@ -102,7 +102,7 @@ def _send_expr(self, visual_id: str) -> None: if ref is None: self._send_json(404, {"error": "unknown visual_id", "id": visual_id}) return - self._send_json(200, format_detail(visual_id, ref, index.mesh_name_map)) + self._send_json(200, format_detail(visual_id, ref, index.context)) @staticmethod def _parse_collapsed(query: dict[str, list[str]]) -> set[str]: diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma.py b/src/tilefoundry/ir/tir/cuda/nn/mma.py index d7f20cef..491b7328 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma.py @@ -109,6 +109,7 @@ def _(call: "Call", ctx: "VerifyContext") -> None: _SM80_THREAD_MESH = Mesh( topologies=(Topology("thread", 32),), layout=Layout(shape=(4, 8), strides=(1, 4)), + names=("warp", "lane"), ) diff --git a/src/tilefoundry/ir/tir/memory/tensor_view.py b/src/tilefoundry/ir/tir/memory/tensor_view.py index 2cbf907a..8f2a3538 100644 --- a/src/tilefoundry/ir/tir/memory/tensor_view.py +++ b/src/tilefoundry/ir/tir/memory/tensor_view.py @@ -17,7 +17,7 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.shard import c_order_strides -from tilefoundry.ir.types.shard.layout import Layout +from tilefoundry.ir.types.shard.layout import Layout, LayoutBase from tilefoundry.visitor_registry import register_typeinfer @@ -31,7 +31,7 @@ class TensorView(Op): """ memory = ParamDef(kind="input", pattern=Tensor) - layout = ParamDef(kind="attribute", annotation=object) + layout = ParamDef(kind="attribute", annotation=LayoutBase) shape = ParamDef(kind="attribute", annotation=tuple, default=None) diff --git a/src/tilefoundry/ir/visitor.py b/src/tilefoundry/ir/visitor.py index 9ead9079..27387651 100644 --- a/src/tilefoundry/ir/visitor.py +++ b/src/tilefoundry/ir/visitor.py @@ -33,6 +33,7 @@ __all__ = [ "ExprFunctor", + "TypeFunctor", "ExprVisitor", "ExprWalker", "ExprCollector", @@ -213,6 +214,28 @@ def clear(self) -> None: self._root = None +class TypeFunctor[T]: + """Dispatch read-only operations over IR type and layout values. + + Types are immutable value objects rather than expressions, so they need a + separate dispatch surface. Keeping this visitor independent from the + expression visitor prevents printers and analysis consumers from growing + parallel ``isinstance`` ladders for the same type family. + """ + + def visit(self, value: Any, ctx: Any = None) -> T: + return self.dispatch_visit(value, ctx) + + def dispatch_visit(self, value: Any, ctx: Any) -> T: + method = getattr(self, f"visit_{type(value).__name__}", None) + if method is not None: + return method(value, ctx) + return self.default_visit(value, ctx) + + def default_visit(self, value: Any, ctx: Any) -> T: + raise NotImplementedError(f"no type visit routine for {type(value).__name__}") + + class ExprVisitor[T](ExprFunctor[T]): """Read-only Expr traversal with identity-based DAG memoization.""" diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 9b101b7e..4925ef44 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -78,6 +78,7 @@ Broadcast, Layout, Mesh, + Partial, ShardLayout, Split, Topology, @@ -239,6 +240,7 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex CallableType=CallableType, BindingSubstitutionCloner=BindingSubstitutionCloner, Broadcast=Broadcast, + Partial=Partial, Binary=Binary, BinaryKind=BinaryKind, Constant=Constant, @@ -745,12 +747,26 @@ def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFai class LexicalScope: """Parser-local lexical frames shared by sequential child construction.""" - def __init__(self, frames: tuple[Mapping[str, object], ...] | None = None): + def __init__( + self, + frames: tuple[Mapping[str, object], ...] | None = None, + mesh_bindings: tuple[set[str], ...] | None = None, + ): source = frames or ({},) self._frames = [dict(frame) for frame in source] + self._mesh_bindings = [ + set(names) for names in (mesh_bindings or tuple(set() for _ in source)) + ] + if len(self._frames) != len(self._mesh_bindings): + raise ValueError("lexical frames and mesh bindings must have the same length") def define(self, name: str, value: object) -> None: self._frames[-1][name] = value + self._mesh_bindings[-1].discard(name) + + def define_mesh(self, name: str, value: object) -> None: + self._frames[-1][name] = value + self._mesh_bindings[-1].add(name) def lookup(self, name: str) -> object | None: for frame in reversed(self._frames): @@ -758,15 +774,26 @@ def lookup(self, name: str) -> object | None: return frame[name] return None + def lookup_mesh(self, name: str) -> object | None: + for frame, mesh_names in reversed(tuple(zip(self._frames, self._mesh_bindings))): + if name in frame: + return frame[name] if name in mesh_names else None + return None + def fork(self) -> LexicalScope: - return LexicalScope(tuple(self._frames) + ({},)) + return LexicalScope( + tuple(self._frames) + ({},), + tuple(self._mesh_bindings) + (set(),), + ) def push_frame(self) -> None: self._frames.append({}) + self._mesh_bindings.append(set()) def pop_frame(self) -> dict[str, object]: if len(self._frames) == 1: raise RuntimeError("cannot pop the root lexical frame") + self._mesh_bindings.pop() return self._frames.pop() def items(self): @@ -1262,7 +1289,7 @@ def from_function(cls, function: FuncParserContext) -> MatchContext: context, ) resolved_mesh = dataclasses.replace(function.mesh, topologies=topologies) - scope.define("mesh", resolved_mesh) + scope.define_mesh("mesh", resolved_mesh) scope.define( _TYPE_INFER_CONTEXT, ParserTypeInferContext(child_resolver=provider, current_mesh=resolved_mesh), @@ -1580,20 +1607,6 @@ def apply(self, value, *, match, context): return value -@dataclass(frozen=True) -class PlacedShapeRule: - STATEMENT: ClassVar[str] = "Placement sugar in a shape slot states both a shape and a layout." - - def apply(self, value, *, match, context): - if not isinstance(value.shape, tuple): - raise ParseError.from_node(match.node, context, "placed shape is not a tuple") - if not isinstance(value.layout, runtime.LayoutBase): - raise ParseError.from_node( - match.node, context, "placed shape did not carry a LayoutBase" - ) - return value - - @dataclass(frozen=True) class StorageValueRule: STATEMENT: ClassVar[str] = "Storage must resolve to a StorageKind." diff --git a/src/tilefoundry/parser/grammar_render.py b/src/tilefoundry/parser/grammar_render.py index 844ecad1..22a0b6b4 100644 --- a/src/tilefoundry/parser/grammar_render.py +++ b/src/tilefoundry/parser/grammar_render.py @@ -96,8 +96,19 @@ class RenderVisitor: def __init__(self, *, line_width: int = 100): self.line_width = line_width self._seen: set[str] = set() + self._seen_named: set[str] = set() self._productions: list[tuple[str, _Expr]] = [] + def _named_pattern(self, pattern: Any) -> _Expr | None: + name = getattr(pattern, "grammar_name", None) + if not isinstance(name, str): + return None + grammar_name = _grammar_name(name) + if grammar_name not in self._seen_named: + self._seen_named.add(grammar_name) + self._productions.append((grammar_name, self._ast_node(pattern, allow_named=False))) + return _text(grammar_name) + def _element(self, pattern: ElementPattern[Any]) -> _Expr: name = pattern.element_name if not name: @@ -136,7 +147,11 @@ def _optional_field(self, fields: dict[str, object], name: str, fallback: str) - pattern = pattern.pattern return _text(fallback) if pattern is None else self.visit(pattern) - def _ast_node(self, pattern: AstNodePattern) -> _Expr: + def _ast_node(self, pattern: AstNodePattern, *, allow_named: bool = True) -> _Expr: + if allow_named: + named = self._named_pattern(pattern) + if named is not None: + return named node_type = pattern.node_type for part in pattern.parts: self.visit(part) diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index 43eb52b3..9cdec03f 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -68,7 +68,6 @@ ParseError, ParserTypeInferContext, PatternFailure, - PlacedShapeRule, PredicatePattern, ReferencePattern, RepeatPattern, @@ -305,81 +304,6 @@ def construct(match, children, context): RULES: ClassVar[tuple[AstRule[Any], ...]] = (CanonicalDTypeRule(),) -class ExplicitLayoutPattern(ElementPattern): - element_name = "explicit_layout" - syntax = LazyPattern( - lambda: BranchPattern( - "explicit_layout", - AstNodePattern( - ast.Tuple, - FieldPattern( - "elts", - SequencePattern( - AstNodePattern( - ast.Tuple, - ChoicePattern( - ConditionPattern( - "active Mesh", - lambda node, context: ( - context.function is not None - and bool(context.function.state.mesh_stack) - ), - ChildPattern( - "shape", - lambda: TensorShapeLayoutPattern(), - "layout_shape", - "layout_shape", - ), - ), - ChildPattern( - "shape", - lambda: ShapePattern(), - "layout_shape", - "layout_shape", - ), - ), - ), - AstNodePattern( - ast.Tuple, - ChildPattern( - "strides", - lambda: ShapePattern(), - "layout_strides", - "layout_strides", - ), - ), - ), - ), - ), - pattern_id="tensor.layout.explicit", - ) - ) - - @staticmethod - def construct(match, children, context): - shape_or_layout = children["shape"] - strides = children["strides"] - if isinstance(shape_or_layout, runtime.ShardLayout): - shape = shape_or_layout.layout.shape - else: - shape = shape_or_layout - if len(shape) != len(strides): - raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") - layout = runtime.Layout(shape=shape, strides=strides) - if isinstance(shape_or_layout, runtime.ShardLayout): - return runtime.ShardLayout( - layout=layout, - attrs=shape_or_layout.attrs, - mesh=shape_or_layout.mesh, - ) - return layout - - RULES: ClassVar[tuple[AstRule[Any], ...]] = ( - LayoutShapeRule(), - LayoutPositionRule(), - ) - - class PlainLayoutPattern(ElementPattern): element_name = "plain_layout" syntax = LazyPattern( @@ -461,15 +385,11 @@ def construct(match, children, context): else: binding = node.value.id axis_name = node.attr - mesh = context.lexical_scope.lookup(binding) - if mesh is None: - try: - reference = node if isinstance(node, ast.Name) else node.value - mesh = _resolve_reference(reference, context) - except ParseError: - mesh = None + mesh = context.lexical_scope.lookup_mesh(binding) if not isinstance(mesh, runtime.Mesh): - raise ParseError.from_node(node, context, f"{binding!r} is not an active Mesh") + raise ParseError.from_node( + node, context, f"{binding!r} is not a lexical Mesh binding" + ) if axis_name is None: if len(mesh.layout.shape) != 1: raise ParseError.from_node( @@ -514,41 +434,298 @@ class PlacedLayout: layout: object -class PlacedLayoutPattern(ElementPattern): - element_name = "placed_layout" - syntax = LazyPattern( - lambda: BindPattern( - AstNodePattern( - ast.Tuple, - FieldPattern( - "elts", - RepeatPattern( +def _layout_dims() -> AstPattern[Any]: + """``(extent, extent @ axis, ...)`` — the layout's own divided dimensions.""" + pattern = AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern( + ChoicePattern( + AstNodePattern( + ast.BinOp, + FieldPattern("op", AstNodePattern(ast.MatMult)), + FieldPattern("left", AstNodePattern(ast.expr)), + FieldPattern( + "right", + ChoicePattern( + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", RepeatPattern(MeshAxisPattern(), minimum=1) + ), + ), + MeshAxisPattern(), + ), + ), + ), + DimExprPattern(), + ) + ), + ), + ) + pattern.grammar_name = "layout_dims" + return pattern + + +def _layout_strides() -> AstPattern[Any]: + """``(stride, ...)`` — how the divided positions are addressed.""" + pattern = AstNodePattern(ast.Tuple, FieldPattern("elts", RepeatPattern(DimExprPattern()))) + pattern.grammar_name = "layout_strides" + return pattern + + +def _value_states() -> AstPattern[Any]: + """``{axis @ B(), axis @ P("sum")}`` — what unsplit mesh axes hold.""" + pattern = AstNodePattern( + ast.Set, + FieldPattern( + "elts", + RepeatPattern( + AstNodePattern( + ast.BinOp, + FieldPattern("op", AstNodePattern(ast.MatMult)), + FieldPattern("left", MeshAxisPattern()), + FieldPattern( + "right", ChoicePattern( AstNodePattern( - ast.BinOp, - FieldPattern("op", AstNodePattern(ast.MatMult)), - FieldPattern("left", AstNodePattern(ast.expr)), + ast.Call, FieldPattern( - "right", - ChoicePattern( + "func", + AstNodePattern( + ast.Name, FieldPattern("id", LiteralPattern("B")) + ), + ), + FieldPattern("args", SequencePattern()), + ), + AstNodePattern( + ast.Call, + FieldPattern( + "func", + AstNodePattern( + ast.Name, FieldPattern("id", LiteralPattern("P")) + ), + ), + FieldPattern( + "args", + SequencePattern( AstNodePattern( - ast.Tuple, - FieldPattern( - "elts", - RepeatPattern( - MeshAxisPattern(), - minimum=1, - ), - ), - ), - MeshAxisPattern(), + ast.Constant, + FieldPattern("value", LiteralPattern(value_type=str)), + ) ), ), ), - DimExprPattern(), - ) + ), ), ), + minimum=1, + ), + ), + ) + pattern.grammar_name = "value_states" + return pattern + + +def _layout_sugar_parts(node: object): + """Split layout sugar into its dims, its strides, and its value states. + + ``(d, ...)`` states dims alone. Wrapping the dims in a tuple adds an + optional stride tuple and an optional ``{axis @ B(), axis @ P("sum")}`` + set, in that order, so one production reads every form the printer emits. + Return ``None`` when the node is not layout sugar at all. + """ + if not isinstance(node, ast.Tuple) or not node.elts: + return None + head, *extras = node.elts + if not extras or not isinstance(head, ast.Tuple): + return node, None, None + strides: ast.Tuple | None = None + states: ast.Set | None = None + for extra in extras: + if isinstance(extra, ast.Tuple) and strides is None and states is None: + strides = extra + elif isinstance(extra, ast.Set) and states is None: + states = extra + else: + return None + return head, strides, states + + +def _value_state_parts(node: ast.AST): + """Read ``axis @ B()`` or ``axis @ P("reduction")`` as axis node and state.""" + if not isinstance(node, ast.BinOp) or not isinstance(node.op, ast.MatMult): + return None + call = node.right + if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Name) or call.keywords: + return None + if call.func.id == "B" and not call.args: + return node.left, "B", None + if call.func.id == "P" and len(call.args) == 1: + reduction = call.args[0] + if isinstance(reduction, ast.Constant) and isinstance(reduction.value, str): + if reduction.value: + return node.left, "P", reduction.value + return None + + +@dataclass(frozen=True) +class _PlacementCandidate: + shape: tuple + strides: tuple | None + splits: tuple[tuple[object, int, int], ...] + states: tuple[tuple[object, int, str, str | None], ...] + + +def _placement_meshes(value: _PlacementCandidate, context: MatchContext, match): + referenced_ids = {id(entry[0]) for entry in (*value.splits, *value.states)} + if not referenced_ids: + return () + if context.function is None: + raise ParseError.from_node( + match.node, context, "placed layout requires function context" + ) + meshes = tuple(dict.fromkeys(entry[0] for entry in (*value.splits, *value.states))) + if len(meshes) != len(referenced_ids): + raise ParseError.from_node( + match.node, context, "placement references a non-lexical Mesh binding" + ) + return meshes + + +@dataclass(frozen=True) +class LayoutStrideRankRule: + STATEMENT: ClassVar[str] = ( + "A stated stride tuple must have the rank of the layout it addresses." + ) + + def apply(self, value, *, match, context): + if value.strides is not None and len(value.shape) != len(value.strides): + raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") + return value + + +@dataclass(frozen=True) +class PlacementMeshResolutionRule: + STATEMENT: ClassVar[str] = "A placement's mesh must be a lexical mesh binding." + + def apply(self, value, *, match, context): + _placement_meshes(value, context, match) + return value + + +@dataclass(frozen=True) +class PlacementLevelRule: + STATEMENT: ClassVar[str] = "A placement's meshes cannot name the same topology level." + + def apply(self, value, *, match, context): + meshes = _placement_meshes(value, context, match) + levels = [topology.name for mesh in meshes for topology in mesh.topologies] + if len(levels) != len(set(levels)): + duplicates = sorted(name for name in set(levels) if levels.count(name) > 1) + raise ParseError.from_node( + match.node, context, + f"a layout can split one level once; two of these meshes name {duplicates}", + ) + return value + + +@dataclass(frozen=True) +class MeshAxisBoundOnceRule: + STATEMENT: ClassVar[str] = "A placement binds each mesh axis at most once." + + def apply(self, value, *, match, context): + seen: set[tuple[int, int]] = set() + for mesh, axis, *_ in (*value.splits, *value.states): + key = (id(mesh), axis) + if key in seen: + raise ParseError.from_node(match.node, context, "mesh axis is bound more than once") + seen.add(key) + return value + + +@dataclass(frozen=True) +class PlacementConstructionRule: + """Materialize the candidate only after placement invariants have run.""" + + STATEMENT: ClassVar[str] = "A placement must construct a valid shard layout." + + def apply(self, value, *, match, context): + if not value.splits and not value.states: + return PlacedLayout( + shape=value.shape, layout=runtime.Layout(shape=value.shape, strides=value.strides) + ) + meshes = _placement_meshes(value, context, match) + mesh = meshes[0] if len(meshes) == 1 else runtime.composed(meshes) + source_offsets: dict[int, int] = {} + offset = 0 + for source in meshes: + source_offsets[id(source)] = offset + offset += len(source.layout.shape) + attrs: list[object] = [runtime.Broadcast() for _ in mesh.layout.shape] + for source, source_axis, tensor_axis in value.splits: + attrs[source_offsets[id(source)] + source_axis] = runtime.Split(tensor_axis) + for source, source_axis, kind, reduction in value.states: + target_axis = source_offsets[id(source)] + source_axis + attrs[target_axis] = runtime.Broadcast() if kind == "B" else runtime.Partial(reduction) + try: + canonical = runtime.canonical_shard_layout(value.shape, mesh, tuple(attrs)) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + if value.strides is not None and len(canonical.layout.shape) != len(value.strides): + raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") + return PlacedLayout( + shape=value.shape, + layout=runtime.ShardLayout( + layout=runtime.Layout(shape=canonical.layout.shape, strides=value.strides), + attrs=canonical.attrs, + mesh=canonical.mesh, + ), + ) + + +@dataclass(frozen=True) +class PlacementAnswerRule: + """A placement answers both halves: the shape written, and where it goes.""" + + STATEMENT: ClassVar[str] = ( + "Placement sugar states both the shape as written and the layout it implies." + ) + + def apply(self, value, *, match, context): + if not isinstance(value, PlacedLayout): + raise ParseError.from_node(match.node, context, "placement did not answer a placement") + if not isinstance(value.shape, tuple): + raise ParseError.from_node(match.node, context, "placement shape is not a tuple") + if not isinstance(value.layout, runtime.LayoutBase): + raise ParseError.from_node(match.node, context, "placement did not carry a LayoutBase") + return value + + +class PlacedLayoutPattern(ElementPattern): + """The one placement production, including dimensions, strides, and states.""" + + element_name = "placed_layout" + syntax = LazyPattern( + lambda: BindPattern( + ChoicePattern( + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + SequencePattern(_layout_dims(), _layout_strides(), _value_states()), + ), + ), + AstNodePattern( + ast.Tuple, + FieldPattern("elts", SequencePattern(_layout_dims(), _layout_strides())), + ), + AstNodePattern( + ast.Tuple, + FieldPattern("elts", SequencePattern(_layout_dims(), _value_states())), + ), + _layout_dims(), ), PlacedLayoutPattern._bind, ) @@ -556,7 +733,6 @@ class PlacedLayoutPattern(ElementPattern): @staticmethod def _placement_parts(node: ast.AST) -> tuple[ast.AST, tuple[ast.AST, ...]] | None: - """Flatten ``extent @ axis @ axis`` into one extent and its axes.""" if not isinstance(node, ast.BinOp) or not isinstance(node.op, ast.MatMult): return None left = PlacedLayoutPattern._placement_parts(node.left) @@ -570,11 +746,15 @@ def _placement_parts(node: ast.AST) -> tuple[ast.AST, tuple[ast.AST, ...]] | Non @staticmethod def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: - assert isinstance(node, ast.Tuple) + parts = _layout_sugar_parts(node) + if parts is None: + return None + dims_node, strides_node, states_node = parts children: list[AstChild] = [] bindings: list[tuple[str, int]] = [] + states: list[tuple[str, str, str | None]] = [] found_placement = False - for tensor_axis, item in enumerate(node.elts): + for tensor_axis, item in enumerate(dims_node.elts): placement = PlacedLayoutPattern._placement_parts(item) extent_node = item axis_nodes: tuple[ast.AST, ...] = () @@ -584,99 +764,56 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat extent_context = context.child(situation="layout_extent", role="layout_extent") if DimExprPattern().match(extent_node, extent_context) is None: return None - children.append( - AstChild( - f"extent_{tensor_axis}", - DimExprPattern(), - extent_node, - "layout_extent", - "layout_extent", - ) - ) + children.append(AstChild(f"extent_{tensor_axis}", DimExprPattern(), extent_node, "layout_extent", "layout_extent")) for mesh_axis, axis_node in enumerate(axis_nodes): axis_context = context.child(situation="mesh_axis", role="mesh_axis") if MeshAxisPattern().match(axis_node, axis_context) is None: return None child_name = f"binding_{tensor_axis}_{mesh_axis}" bindings.append((child_name, tensor_axis)) - children.append( - AstChild( - child_name, - MeshAxisPattern(), - axis_node, - "mesh_axis", - "mesh_axis", - ) - ) - if not found_placement: + children.append(AstChild(child_name, MeshAxisPattern(), axis_node, "mesh_axis", "mesh_axis")) + if strides_node is not None: + for index, item in enumerate(strides_node.elts): + stride_context = context.child(situation="layout_strides", role="layout_strides") + if DimExprPattern().match(item, stride_context) is None: + return None + children.append(AstChild(f"stride_{index}", DimExprPattern(), item, "layout_strides", "layout_strides")) + if states_node is not None: + for index, item in enumerate(states_node.elts): + state = _value_state_parts(item) + if state is None: + return None + axis_node, kind, reduction = state + axis_context = context.child(situation="mesh_axis", role="mesh_axis") + if MeshAxisPattern().match(axis_node, axis_context) is None: + return None + child_name = f"state_{index}" + states.append((child_name, kind, reduction)) + children.append(AstChild(child_name, MeshAxisPattern(), axis_node, "mesh_axis", "mesh_axis")) + if not found_placement and not states and strides_node is None: return None return dataclasses.replace( - matched, - pattern_id="tensor.layout.placed", - branch_id="placed_layout", + matched, pattern_id="tensor.layout.placed", branch_id="placed_layout", captures={ - **matched.captures, - "rank": len(node.elts), - "bindings": tuple(bindings), - }, - children=tuple(children), + **matched.captures, "rank": len(dims_node.elts), + "stride_rank": None if strides_node is None else len(strides_node.elts), + "bindings": tuple(bindings), "states": tuple(states), + }, children=tuple(children), ) @staticmethod def construct(match, children, context): rank = match.captures["rank"] shape = tuple(children[f"extent_{axis}"] for axis in range(rank)) - bindings = tuple( - (*children[child_name], tensor_axis) - for child_name, tensor_axis in match.captures["bindings"] - ) - referenced_ids = {id(mesh) for mesh, _, _ in bindings} - if context.function is None: - raise ParseError.from_node( - match.node, context, "placed layout requires function context" - ) - meshes = tuple( - mesh for mesh in context.function.state.mesh_stack if id(mesh) in referenced_ids - ) - if len(meshes) != len(referenced_ids): - meshes = tuple(dict.fromkeys(mesh for mesh, _, _ in bindings)) - if len(meshes) != len(referenced_ids): - raise ParseError.from_node( - match.node, context, "placement references an inactive Mesh" - ) - levels = [topology.name for mesh in meshes for topology in mesh.topologies] - if len(levels) != len(set(levels)): - duplicates = sorted(name for name in set(levels) if levels.count(name) > 1) - raise ParseError.from_node( - match.node, - context, - f"a layout can split one level once; two of these meshes name {duplicates}", - ) - mesh = meshes[0] if len(meshes) == 1 else runtime.composed(meshes) - source_offsets: dict[int, int] = {} - offset = 0 - for source in meshes: - source_offsets[id(source)] = offset - offset += len(source.layout.shape) - attrs: list[object] = [runtime.Broadcast() for _ in mesh.layout.shape] - for source, source_axis, tensor_axis in bindings: - target_axis = source_offsets[id(source)] + source_axis - if not isinstance(attrs[target_axis], runtime.Broadcast): - raise ParseError.from_node(match.node, context, "mesh axis is bound more than once") - attrs[target_axis] = runtime.Split(tensor_axis) - try: - canonical = runtime.canonical_shard_layout(shape, mesh, tuple(attrs)) - return runtime.ShardLayout( - layout=runtime.Layout(shape=canonical.layout.shape, strides=None), - attrs=canonical.attrs, - mesh=canonical.mesh, - ) - except (TypeError, ValueError) as error: - raise ParseError.from_node(match.node, context, str(error)) from error + stride_rank = match.captures.get("stride_rank") + strides = None if stride_rank is None else tuple(children[f"stride_{index}"] for index in range(stride_rank)) + splits = tuple((*children[child_name], tensor_axis) for child_name, tensor_axis in match.captures["bindings"]) + states = tuple((*children[child_name], kind, reduction) for child_name, kind, reduction in match.captures.get("states", ())) + return _PlacementCandidate(shape, strides, splits, states) RULES: ClassVar[tuple[AstRule[Any], ...]] = ( - LayoutShapeRule(), - LayoutPositionRule(), + LayoutStrideRankRule(), MeshAxisBoundOnceRule(), PlacementMeshResolutionRule(), + PlacementLevelRule(), PlacementConstructionRule(), PlacementAnswerRule(), ) @@ -711,7 +848,6 @@ class LayoutPattern(ElementPattern): ), pattern_id="tensor.layout.call", ), - ExplicitLayoutPattern(), PlacedLayoutPattern(), PlainLayoutPattern(), ) @@ -831,36 +967,11 @@ def _states_a_layout_slot(elts: object, context: "MatchContext") -> bool: return False -class PlacedShapePattern(ElementPattern): - """Placement sugar read by a slot that states a shape rather than a layout. - - Same syntax as `PlacedLayoutPattern`, different answer. That pattern serves - the layout slot, whose question is only where a value goes, so a layout is - the whole answer. A shape slot has no operand to read a shape from, and the - layout's own shape is the divided one, so it needs the extents as written - kept beside the layout instead of recovered from it. - """ - - element_name = "placed_shape" - syntax = LazyPattern(lambda: PlacedLayoutPattern().syntax) - - @staticmethod - def construct(match, children, context): - layout = PlacedLayoutPattern.construct(match, children, context) - rank = match.captures["rank"] - return PlacedLayout( - shape=tuple(children[f"extent_{axis}"] for axis in range(rank)), - layout=layout, - ) - - RULES: ClassVar[tuple[AstRule[Any], ...]] = (PlacedShapeRule(),) - - class TensorShapeLayoutPattern(ElementPattern): element_name = "tensor_shape_layout" syntax = LazyPattern( lambda: ChoicePattern( - PlacedShapePattern(), + PlacedLayoutPattern(), ShapePattern(), ) ) @@ -1009,6 +1120,8 @@ def construct(match, children, context): shape, layout = placed, None else: shape, layout = children["shape"], children["layout"] + if isinstance(layout, PlacedLayout): + layout = layout.layout storage = children.get("storage") return runtime.TensorType( shape=shape, @@ -2421,7 +2534,9 @@ def construct(match, children, context): value for name, value in children.items() if name.startswith("input_") ) attrs = { - name.removeprefix("attr_"): value + name.removeprefix("attr_"): ( + value.layout if isinstance(value, PlacedLayout) else value + ) for name, value in children.items() if name.startswith("attr_") } @@ -3239,7 +3354,7 @@ def construct(match, children, context): binding = context.values.get("mesh_binding") _enter_mesh_scope(context, mesh, match) if isinstance(binding, str): - context.lexical_scope.define(binding, mesh) + context.lexical_scope.define_mesh(binding, mesh) context.function.state.mesh_stack.append(mesh) return mesh @@ -4846,7 +4961,6 @@ def _body_as_ast_module(body: object, *, strip_docstring: bool = False) -> ast.M "ConstantPattern", "DTypePattern", "DimExprPattern", - "ExplicitLayoutPattern", "ExpressionPattern", "ForPattern", "FunctionDialectRule", diff --git a/tests/analysis/test_render_golden.py b/tests/analysis/test_render_golden.py new file mode 100644 index 00000000..55ffe162 --- /dev/null +++ b/tests/analysis/test_render_golden.py @@ -0,0 +1,40 @@ +"""The annotated analysis surface, locked against one placed fixture.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +from tilefoundry.analysis import analyze +from tilefoundry.inspection import as_script +from tilefoundry.inspection.analysis_report import render_analysis +from tilefoundry.ir.core.module import Module + +FIXTURE = Path(__file__).parents[1] / "fixtures" / "inspection" / "type_printer_sugar.py" + + +def _module_in(path: Path) -> Module: + spec = importlib.util.spec_from_file_location(path.stem, path) + assert spec is not None and spec.loader is not None + loaded = importlib.util.module_from_spec(spec) + spec.loader.exec_module(loaded) + return next(value for value in vars(loaded).values() if isinstance(value, Module)) + + +def _without_comments(source: str) -> str: + return "\n".join(line.split(" # ", 1)[0].rstrip() for line in source.splitlines()) + + +def test_analysis_annotates_the_printed_program_without_changing_it() -> None: + """Annotation adds metadata to canonical source; it does not restate types. + + The golden shares its fixture with the round-trip golden, so a type that + reads one way in emitted code and another in an annotation shows up here as + a diff rather than as two goldens that drifted apart. + """ + module = _module_in(FIXTURE) + result = analyze(module, module.entry_function(), analysis=("compute-cost", "memory")) + annotated = render_analysis(result).annotated + + assert annotated == FIXTURE.with_suffix(".analyzed.txt").read_text() + assert _without_comments(annotated) == _without_comments(as_script(result.function)) diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index f154d9dd..7345d588 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -743,8 +743,8 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> assert annotated.count("reshard(tokens") == 2 assert "v0 = reshard(tokens" in annotated assert "v3 = reshard(tokens" in annotated - assert "offset=0" in annotated - assert "offset=120" in annotated + assert "with cta[:120] as cta_2:" in annotated + assert "with cta_3[120:] as cta_4:" in annotated summary = payload["function_records"]["performance"] cost = payload["function_records"]["compute-cost"] @@ -780,13 +780,19 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> } hoisted = {line.split(" = ", 1)[0] for line in lines if " = Mesh((Topology(" in line} - scoped = {line.lstrip().split()[1] for line in lines if line.lstrip().startswith("with ")} - assert hoisted == {"cta", "cta_2"} | scoped + scoped = { + line.rsplit(" as ", 1)[1].split(":", 1)[0] + for line in lines + if line.lstrip().startswith("with ") + } + assert hoisted == set() + assert scoped == {"cta", "cta_2", "cta_3", "cta_4"} annotated_types = [ line.split(" # ", 1)[1].split("; ", 1)[0] for line in lines if " # Tensor[" in line ] - assert 'Tensor[(120, 64), "f32", (120 @ cta_2.tile, 64)]' in annotated_types - assert 'Tensor[(120, 64), "f32", (12 @ cta_3.tile, 10, 64)]' in annotated_types + assert 'Tensor[(120, 64), "f32", ((120 @ cta_2.tile, 64), (64, 1))]' in annotated_types + assert 'Tensor[(120, 64), "f32", ((12 @ cta_3.tile, 10, 64), (640, 64, 1))]' in annotated_types + assert 'Tensor[(120, 64), "f32", ((120, 64), (64, 1), {cta.tile @ B()})]' in annotated_types rows = payload["calls"] assert len(rows) == 7 @@ -812,9 +818,9 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> else: assert "; performance=" not in statement annotated_meshes = set(re.findall(r"@ (\w+)\.", statement.split(" # ", 1)[1])) - assert annotated_meshes <= hoisted + assert annotated_meshes <= scoped placed_meshes = set(re.findall(r"mesh=(\w+),", statement)) - assert placed_meshes <= hoisted + assert placed_meshes <= scoped assert len([line for line in lines if "; performance=" in line]) == len(timed) assert len({row["value"] for row in rows}) == len(rows) assert "units=" not in annotated diff --git a/tests/fixtures/inspection/__init__.py b/tests/fixtures/inspection/__init__.py new file mode 100644 index 00000000..113592ed --- /dev/null +++ b/tests/fixtures/inspection/__init__.py @@ -0,0 +1 @@ +"""Fixtures whose printed source is the inspection contract under test.""" diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt new file mode 100644 index 00000000..f3f70372 --- /dev/null +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -0,0 +1,39 @@ +from __future__ import annotations + +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 + +@func(mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane'))) +def composed_mesh_pipeline( + x: Tensor[(8, 4, 16), "f32"], + seed: Tensor[(16,), "f32"], + acc: Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem"], + mixed: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((4, 2, 16), None), + attrs=(S(0), B(), P("sum")), + mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), + ), "rmem"] +): + with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: + v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r2048/w0@cta:r512/w0,thread:r64/w0,rmem:r0/w2048@cta:r0/w512,thread:r0/w64 + v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@cta:128,thread:16; traffic traffic=rmem:r2048/w2048@cta:r512/w512,thread:r64/w64 + v2 = reshard(v1, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost; traffic traffic=rmem:r2048/w0@cta:r512/w0,thread:r64/w0,smem:r0/w2048@cta:r0/w512,thread:r0/w64 + v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost flops=bf16:512@cta:128,thread:128; traffic traffic=smem:r2048/w1024@cta:r512/w256,thread:r512/w256 + v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta.tile, 2, 16, 4), (128, 64, 1, 16)), "smem"]; compute-cost; traffic traffic=smem:r1024/w1024@cta:r256/w256,thread:r256/w256 + v5 = reshard(v4, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {cta.tile @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 + v6 = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r64/w0@cta:r64/w0,thread:r8/w0,rmem:r0/w64@cta:r0/w64,thread:r0/w8 + folded = reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic + summed = reshard(mixed, layout=((8, 16), {mesh.warp @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost; traffic + for _ in range(3): # loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete + v10 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost flops=f32:128@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 + v11 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:128@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 + folded = v11 + summed = v10 + v13 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 + v15 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 + v17 = reshard(v6, layout=Layout((16,), (1,)), storage=gmem) # Tensor[(16,), "f32", Layout((16,), (1,))]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 + return (v5, v13, v15, v17) diff --git a/tests/fixtures/inspection/type_printer_sugar.printed.txt b/tests/fixtures/inspection/type_printer_sugar.printed.txt new file mode 100644 index 00000000..fb23fb86 --- /dev/null +++ b/tests/fixtures/inspection/type_printer_sugar.printed.txt @@ -0,0 +1,87 @@ +from __future__ import annotations + +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.module import module +from tilefoundry.target import CudaTarget + +@module(entry="composed_mesh_pipeline", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", 4), Topology("thread", 8),)) +class TypePrinterSugar: + @func(mesh=Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=('lane',))) + def nested_loop_tuple( + x: Tensor[(8, 16), "f32"], + weight: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((8, 16), None), + attrs=(B(), P("max")), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + )] + ): + split = reshard(x, layout=(8 @ mesh.lane, 16), storage=rmem) + whole = reshard(x, layout=((8, 16), {mesh.lane @ B()}), storage=rmem) + for _ in range(3): + whole_2 = add(whole, whole) + split_2 = unary(split, kind="square") + split = split_2 + whole = whole_2 + v1 = reshard(split, layout=((8, 16), (16, 1), {mesh.lane @ B()}), storage=gmem) + v2 = reshard(whole, layout=((8, 16), (16, 1), {mesh.lane @ B()}), storage=gmem) + with Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')) as thread: + per_warp = reshard(weight, layout=(2 @ thread.warp, 4, 16), storage=rmem) + unfolded = reshard(per_warp, layout=((8, 16), {thread.warp @ B()}), storage=gmem) + return (v1, v2, unfolded) + + @func(mesh=Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',))) + def named_and_out_of_scope( + x: Tensor[(8, 16), "f32"], + held: Tensor[(8, 16), "f32", (4 @ mesh.tile, 2, 16), "rmem"], + frag: Tensor[(16,), "f32", + ShardLayout( + layout=Layout((2, 4, 2), (8, 2, 1)), + attrs=(S(0), S(1)), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + ), "rmem"] + ): + mine = reshard(x, layout=(4 @ mesh.tile, 2, 16), storage=rmem) + v0 = reshard(mine, layout=Layout((8, 16), (16, 1)), storage=gmem) + escaped = reshard(x, layout=ShardLayout( + layout=Layout((2, 4, 16), None), + attrs=(S(0), B()), + mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane')), + ), storage=rmem) + return (v0, held, frag, escaped) + + @func(mesh=Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=('warp', 'lane'))) + def composed_mesh_pipeline( + x: Tensor[(8, 4, 16), "f32"], + seed: Tensor[(16,), "f32"], + acc: Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), "rmem"], + mixed: Tensor[(8, 16), "f32", + ShardLayout( + layout=Layout((4, 2, 16), None), + attrs=(S(0), B(), P("sum")), + mesh=Mesh((Topology("cta", 4), Topology("thread", 8)), Layout((4, 2, 4), (8, 4, 1)), names=('tile', 'warp', 'lane')), + ), "rmem"] + ): + with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: + composed = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) + v0 = unary(composed, kind="square") + staged = reshard(v0, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) + narrowed = cast(staged, dtype="bf16") + swapped = transpose(narrowed, perm=(0, 2, 1)) + gathered = reshard(swapped, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) + seeded = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) + folded = reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage=rmem) + summed = reshard(mixed, layout=((8, 16), {mesh.warp @ B()}), storage=rmem) + for _ in range(3): + summed_2 = add(summed, summed) + folded_2 = unary(folded, kind="square") + folded = folded_2 + summed = summed_2 + v4 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) + v5 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) + v6 = reshard(seeded, layout=Layout((16,), (1,)), storage=gmem) + return (gathered, v4, v5, v6) diff --git a/tests/fixtures/inspection/type_printer_sugar.py b/tests/fixtures/inspection/type_printer_sugar.py new file mode 100644 index 00000000..d7161fba --- /dev/null +++ b/tests/fixtures/inspection/type_printer_sugar.py @@ -0,0 +1,125 @@ +"""Placed values covering every shape the printer's layout sugar can take. + +One entry per combination the sugar has to survive: a mesh level alone and two +levels composed on one value, splits on one and on several tensor axes, every +value state a mesh axis can hold, contiguous and explicitly strided layouts +over the same logical shape, and a loop whose carried fields are placed +differently. `_FRAGMENT` is reshard-ed to from both entries: one enters its +mesh, the other never does, so only the first has a binding sugar could name. +""" + +from __future__ import annotations + +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.target import CudaTarget + +_H200 = CudaTarget("nvidia.h200_sxm") +_TOPOLOGIES = (Topology("cta", 4), Topology("thread", 8)) + +_TILE = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=("tile",)) +_WARP_LANE = Mesh((Topology("thread", 8),), Layout((2, 4), (4, 1)), names=("warp", "lane")) +_LANES = Mesh((Topology("thread", 8),), Layout((8,), (1,)), names=("lane",)) + +_FRAGMENT = ShardLayout( + layout=Layout((2, 4, 2), (8, 2, 1)), + attrs=(S(0), S(1)), + mesh=_WARP_LANE, +) +_WEIGHT = ShardLayout( + layout=Layout((8, 16), None), + attrs=(B(), P("max")), + mesh=_WARP_LANE, +) +_MIXED = ShardLayout( + layout=Layout((4, 2, 16), None), + attrs=(S(0), B(), P("sum")), + mesh=Mesh( + (Topology("cta", 4), Topology("thread", 8)), + Layout((4, 2, 4), (8, 4, 1)), + names=("tile", "warp", "lane"), + ), +) +_ESCAPED = ShardLayout( + layout=Layout((2, 4, 16), None), + attrs=(S(0), B()), + mesh=_WARP_LANE, +) + + +@module(entry="composed_mesh_pipeline", target=_H200, topologies=_TOPOLOGIES) +class TypePrinterSugar: + @func(mesh=_WARP_LANE) + def composed_mesh_pipeline( + x: Tensor[(8, 4, 16), "f32"], + seed: Tensor[(16,), "f32"], + acc: Tensor[ + (8, 16), "f32", + ((2 @ mesh.warp, 4, 16), {mesh.lane @ P("sum")}), + "rmem", + ], + mixed: Tensor[(8, 16), "f32", _MIXED, "rmem"], + ): + with _TILE as cta: + composed = tf.reshard( + x, (8 @ cta.tile, 4 @ mesh.warp, 16 @ mesh.lane), "rmem" + ) + staged = tf.reshard(tf.square(composed), (8 @ cta.tile, 4, 16), "smem") + narrowed = tf.cast(staged, dtype="bf16") + swapped = tf.transpose(narrowed, perm=(0, 2, 1)) + gathered = tf.reshard(swapped, (8, 16, 4), "gmem") + seeded = tf.reshard(seed, _FRAGMENT, "rmem") + folded = tf.reshard(acc, (8 @ mesh.warp, 16 @ mesh.lane), "rmem") + summed = tf.reshard( + mixed, ((8, 16), {mesh.warp @ B(), mesh.lane @ B()}), "rmem" + ) + for _ in range(3): + folded = tf.square(folded) + summed = tf.add(summed, summed) + return ( + gathered, + tf.reshard(folded, (8, 16), "gmem"), + tf.reshard(summed, (8, 16), "gmem"), + tf.reshard(seeded, (16,), "gmem"), + ) + + @func + def nested_loop_tuple( + x: Tensor[(8, 16), "f32"], + weight: Tensor[(8, 16), "f32", _WEIGHT], + ): + with _LANES as lanes: + split = tf.reshard(x, (8 @ lanes.lane, 16), "rmem") + whole = tf.reshard(x, ((8, 16), {lanes.lane @ B()}), "rmem") + for _ in range(3): + split = tf.square(split) + whole = tf.add(whole, whole) + with _WARP_LANE as thr: + per_warp = tf.reshard(weight, (8 @ thr.warp, 16), "rmem") + unfolded = tf.reshard( + per_warp, ((8, 16), {thr.warp @ B(), thr.lane @ B()}), "gmem" + ) + return ( + tf.reshard(split, (8, 16), "gmem"), + tf.reshard(whole, (8, 16), "gmem"), + unfolded, + ) + + @func(mesh=_TILE) + def named_and_out_of_scope( + x: Tensor[(8, 16), "f32"], + held: Tensor[(8, 16), "f32", (8 @ mesh.tile, 16), "rmem"], # noqa: F821 + frag: Tensor[(16,), "f32", _FRAGMENT, "rmem"], + ): + """The two placements sugar cannot state. + + ``frag`` holds the same constant `composed_mesh_pipeline` reshards to, + but this function never enters that constant's mesh, so nothing here + binds a name sugar could use. ``escaped`` is resharded onto a mesh this + scope never enters -- a reshard is the one operation allowed to cross + that boundary, so its target names a mesh that is not its own scope. + """ + mine = tf.reshard(x, (8 @ mesh.tile, 16), "rmem") # noqa: F821 + escaped = tf.reshard(x, _ESCAPED, "rmem") + return tf.reshard(mine, (8, 16), "gmem"), held, frag, escaped diff --git a/tests/fixtures/tir/async_sync.py b/tests/fixtures/tir/async_sync.py index 5fae2689..682deefe 100644 --- a/tests/fixtures/tir/async_sync.py +++ b/tests/fixtures/tir/async_sync.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.shard import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -11,14 +11,9 @@ class AsyncStage: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def async_stage_device(a: Tensor[(128, 4), "f32"], b: Tensor[(128, 4), "f32"]): with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as m: - a_view = T.tensor_view(a, layout=ShardLayout(layout=Layout(shape=(128, 4), strides=(4, 1)), attrs=(S(0),), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(128,), strides=(1,)), names=("t",)))) - shared = T.alloc_tensor(tensor_type=Tensor[(128, 4), "f32", - ShardLayout( - layout=Layout((128, 4), (4, 1)), - attrs=(S(0),), - mesh=Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)), - ), "smem"]) - b_view = T.tensor_view(b, layout=ShardLayout(layout=Layout(shape=(128, 4), strides=(4, 1)), attrs=(S(0),), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(128,), strides=(1,)), names=("t",)))) + a_view = T.tensor_view(a, layout=((128 @ m.t, 4), (4, 1))) + shared = T.alloc_tensor(tensor_type=Tensor[(128, 4), "f32", ((128 @ m.t, 4), (4, 1)), "smem"]) + b_view = T.tensor_view(b, layout=((128 @ m.t, 4), (4, 1))) T.copy_async(a_view, shared) T.cp_async_commit() T.cp_async_wait(n=0) diff --git a/tests/fixtures/tir/mma.py b/tests/fixtures/tir/mma.py index 1f42e44b..4ef4038c 100644 --- a/tests/fixtures/tir/mma.py +++ b/tests/fixtures/tir/mma.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.shard import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -10,32 +10,17 @@ class MmHandwritten: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def mm_device(a: Tensor[(16, 16), "bf16"], b: Tensor[(16, 8), "bf16"], c: Tensor[(16, 8), "f32"]): - with Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))) as _warp: - a_view = T.tensor_view(a, layout=ShardLayout(layout=Layout(shape=(2, 4, 2, 8, 2), strides=(1, 2, 8, 16, 128)), attrs=(S(1), S(3)), mesh=Mesh(topologies=(Topology(name="thread", size=32),), layout=Layout(shape=(4, 8), strides=(1, 4)), names=()))) - b_view = T.tensor_view(b, layout=ShardLayout(layout=Layout(shape=(8, 2, 4, 2), strides=(1, 8, 16, 64)), attrs=(S(2), S(0)), mesh=Mesh(topologies=(Topology(name="thread", size=32),), layout=Layout(shape=(4, 8), strides=(1, 4)), names=()))) - a_frag = T.alloc_tensor(tensor_type=Tensor[(16, 16), "bf16", - ShardLayout( - layout=Layout((2, 4, 2, 8, 2), (1, 2, 8, 16, 128)), - attrs=(S(1), S(3)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - ), "rmem"]) - b_frag = T.alloc_tensor(tensor_type=Tensor[(16, 8), "bf16", - ShardLayout( - layout=Layout((8, 2, 4, 2), (1, 8, 16, 64)), - attrs=(S(2), S(0)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - ), "rmem"]) - acc = T.alloc_tensor(tensor_type=Tensor[(16, 8), "f32", - ShardLayout( - layout=Layout((2, 4, 8, 2), (1, 2, 8, 64)), - attrs=(S(1), S(2)), - mesh=Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4))), - ), "rmem"]) + with Mesh((Topology("thread", 32),), Layout((4, 8), (1, 4)), names=('warp', 'lane')) as _warp: + a_view = T.tensor_view(a, layout=((2, 4 @ _warp.warp, 2, 8 @ _warp.lane, 2), (1, 2, 8, 16, 128))) + b_view = T.tensor_view(b, layout=((8 @ _warp.lane, 2, 4 @ _warp.warp, 2), (1, 8, 16, 64))) + a_frag = T.alloc_tensor(tensor_type=Tensor[(16, 16), "bf16", ((2, 4 @ _warp.warp, 2, 8 @ _warp.lane, 2), (1, 2, 8, 16, 128)), "rmem"]) + b_frag = T.alloc_tensor(tensor_type=Tensor[(16, 8), "bf16", ((8 @ _warp.lane, 2, 4 @ _warp.warp, 2), (1, 8, 16, 64)), "rmem"]) + acc = T.alloc_tensor(tensor_type=Tensor[(16, 8), "f32", ((2, 4 @ _warp.warp, 8 @ _warp.lane, 2), (1, 2, 8, 64)), "rmem"]) T.copy(a_view, a_frag) T.copy(b_view, b_frag) T.fill(acc, 0.0) T.mma(acc, a_frag, b_frag, atom=T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN)) - c_view = T.tensor_view(c, layout=ShardLayout(layout=Layout(shape=(2, 4, 8, 2), strides=(1, 2, 8, 64)), attrs=(S(1), S(2)), mesh=Mesh(topologies=(Topology(name="thread", size=32),), layout=Layout(shape=(4, 8), strides=(1, 4)), names=()))) + c_view = T.tensor_view(c, layout=((2, 4 @ _warp.warp, 8 @ _warp.lane, 2), (1, 2, 8, 64))) T.copy(acc, c_view) @prim_func(target=CpuTarget()) diff --git a/tests/fixtures/tir/rmsnorm.py b/tests/fixtures/tir/rmsnorm.py index 4ca68cc8..6e8890d8 100644 --- a/tests/fixtures/tir/rmsnorm.py +++ b/tests/fixtures/tir/rmsnorm.py @@ -2,7 +2,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor -from tilefoundry.ir.types.shard import B, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.shard import B, Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -10,10 +10,10 @@ class TirRmsnorm: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def rmsnorm_device(x: Tensor[(1, 128), "f32"], weight: Tensor[(128,), "f32"], out: Tensor[(1, 128), "f32"]): - with Mesh((Topology("thread", 1),), Layout((1,), (1,))) as thread: - x_view = T.tensor_view(x, layout=ShardLayout(layout=Layout(shape=(1, 128), strides=(128, 1)), attrs=(B(),), mesh=Mesh(topologies=(Topology(name="thread", size=1),), layout=Layout(shape=(1,), strides=(1,)), names=()))) - weight_view = T.tensor_view(weight, layout=ShardLayout(layout=Layout(shape=(128,), strides=(1,)), attrs=(B(),), mesh=Mesh(topologies=(Topology(name="thread", size=1),), layout=Layout(shape=(1,), strides=(1,)), names=()))) - out_view = T.tensor_view(out, layout=ShardLayout(layout=Layout(shape=(1, 128), strides=(128, 1)), attrs=(B(),), mesh=Mesh(topologies=(Topology(name="thread", size=1),), layout=Layout(shape=(1,), strides=(1,)), names=()))) + with Mesh((Topology("thread", 1),), Layout((1,), (1,)), names=('t',)) as thread: + x_view = T.tensor_view(x, layout=((1, 128), (128, 1), {thread.t @ B()})) + weight_view = T.tensor_view(weight, layout=((128,), (1,), {thread.t @ B()})) + out_view = T.tensor_view(out, layout=((1, 128), (128, 1), {thread.t @ B()})) T.rms_norm(x_view, out_view, weight_view, eps=1e-05) T.sync(thread) diff --git a/tests/fixtures/tir/square.py b/tests/fixtures/tir/square.py index e4ee8a52..038e795e 100644 --- a/tests/fixtures/tir/square.py +++ b/tests/fixtures/tir/square.py @@ -4,7 +4,7 @@ from tilefoundry.dsl import DimVar, T, Tensor from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.core.pattern import DimVarRangePat -from tilefoundry.ir.types.shard import Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.shard import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget _S = DimVar("S", 1, 256) @@ -18,14 +18,9 @@ def square_device(x: Tensor[(_S,), "f32"]): @square_device.specialize(DimVarRangePat("S", 1, 127)) def square_small(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=ShardLayout(layout=Layout(shape=(128,), strides=(1,)), attrs=(S(0),), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(128,), strides=(1,)), names=()))) - reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", - ShardLayout( - layout=Layout((128,), (1,)), - attrs=(S(0),), - mesh=Mesh((Topology("thread", 128),), Layout((128,), (1,))), - ), "rmem"]) + with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as thread: + view = T.tensor_view(x, layout=((128 @ thread.t,), (1,))) + reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", ((128 @ thread.t,), (1,)), "rmem"]) for phase in range(0, 2, 1): if phase < 1: T.copy(view, reg) @@ -36,14 +31,9 @@ def square_small(x: Tensor[(_S,), "f32"]): @square_device.specialize(DimVarRangePat("S", 128, 255)) def square_large(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=ShardLayout(layout=Layout(shape=(128,), strides=(1,)), attrs=(S(0),), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(128,), strides=(1,)), names=()))) - reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", - ShardLayout( - layout=Layout((128,), (1,)), - attrs=(S(0),), - mesh=Mesh((Topology("thread", 128),), Layout((128,), (1,))), - ), "rmem"]) + with Mesh((Topology("thread", 128),), Layout((128,), (1,)), names=('t',)) as thread: + view = T.tensor_view(x, layout=((128 @ thread.t,), (1,))) + reg = T.alloc_tensor(tensor_type=Tensor[(128,), "f32", ((128 @ thread.t,), (1,)), "rmem"]) for phase in range(0, 2, 1): if phase < 1: T.copy(view, reg) diff --git a/tests/fixtures/tir/sync.py b/tests/fixtures/tir/sync.py index 9cc3e4fa..7d461e54 100644 --- a/tests/fixtures/tir/sync.py +++ b/tests/fixtures/tir/sync.py @@ -3,7 +3,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core.kinds import BinaryKind -from tilefoundry.ir.types.shard import ComposedLayout, Layout, Mesh, S, ShardLayout, Topology +from tilefoundry.ir.types.shard import Layout, Mesh, Topology from tilefoundry.target import CpuTarget, CudaTarget @@ -12,18 +12,13 @@ class SyncSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def sync_square_device(a: Tensor[(4, 32), "f32"]): with Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1)), names=('w', 't')) as m: - view = T.tensor_view(a, layout=ShardLayout(layout=Layout(shape=(4, 32), strides=(32, 1)), attrs=(S(0), S(1)), mesh=Mesh(topologies=(Topology(name="thread", size=128),), layout=Layout(shape=(4, 32), strides=(32, 1)), names=("w", "t")))) - reg = T.alloc_tensor(tensor_type=Tensor[(4, 32), "f32", - ShardLayout( - layout=Layout((4, 32), (32, 1)), - attrs=(S(0), S(1)), - mesh=Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1)), names=('w', 't')), - ), "rmem"]) + view = T.tensor_view(a, layout=((4 @ m.w, 32 @ m.t), (32, 1))) + reg = T.alloc_tensor(tensor_type=Tensor[(4, 32), "f32", ((4 @ m.w, 32 @ m.t), (32, 1)), "rmem"]) T.copy(view, reg) T.sync(m) - T.sync(Mesh(topologies=(Topology(name="thread", size=128),), layout=ComposedLayout(inner=None, offset=0, outer=Layout(shape=(1, 32), strides=(32, 1))), names=("w", "t"))) - T.sync(Mesh(topologies=(Topology(name="thread", size=128),), layout=ComposedLayout(inner=None, offset=0, outer=Layout(shape=(2, 32), strides=(32, 1))), names=("w", "t"))) - T.sync(Mesh(topologies=(Topology(name="thread", size=128),), layout=ComposedLayout(inner=None, offset=64, outer=Layout(shape=(2, 32), strides=(32, 1))), names=("w", "t"))) + T.sync(m[:1]) + T.sync(m[:2]) + T.sync(m[2:]) T.binary(reg, reg, reg, kind=BinaryKind.MUL) T.copy(reg, view) diff --git a/tests/inspection/test_python_printer.py b/tests/inspection/test_python_printer.py index 94ccfaad..df80c77d 100644 --- a/tests/inspection/test_python_printer.py +++ b/tests/inspection/test_python_printer.py @@ -6,13 +6,9 @@ *same* Target when the emitted source is executed. """ -import re from dataclasses import dataclass, fields, replace from tests._source import import_dsl -from tests.fixtures.placed.gqa_decode import GqaOnline -from tests.fixtures.placed.moe_mega_kernel import MoEMegaKernel -from tests.fixtures.placed.prefill_decode_attention import PrefillDecodeAttention from tests.fixtures.placed.rmsnorm import RmsnormModule from tilefoundry.inspection import PythonPrintOptions, as_script from tilefoundry.inspection.analysis_report import _type_text @@ -51,64 +47,6 @@ def test_inspection_types_are_opt_in_same_line_comments(): ) -def _hoisted_meshes(source: str) -> set[str]: - return { - line.split(" = ", 1)[0] for line in source.splitlines() if " = Mesh((Topology(" in line - } - - -def test_annotated_tuple_fields_each_name_the_hoisted_mesh(): - """Every field of a tuple-typed annotation reaches the prelude's mesh name. - - A loop carrying several tensors annotates a ``Tuple[...]``, so the field - walk has to carry the same mesh name map the signature is rendered from. A - field that lost it would restate the whole mesh as ``ShardLayout(...)``. - """ - annotated = as_script( - PrefillDecodeAttention, options=PythonPrintOptions(show_types=True) - ) - hoisted = _hoisted_meshes(annotated) - tuples = [ - line.split(" # ", 1)[1] for line in annotated.splitlines() if " # Tuple[" in line - ] - - assert len(tuples) == 2 - for annotation in tuples: - assert "ShardLayout(" not in annotation - named = set(re.findall(r"@ (\w+)\.", annotation)) - assert named - assert named <= hoisted - - -def test_an_annotated_layout_sugar_cannot_say_stays_verbose(): - """The fallback keeps the whole layout, and it coexists with sugar. - - Sugar names a mesh axis, so a mesh with no named axes has nothing to name - and stays verbose without dropping what the verbose form carries. - """ - unnamed_axes = as_script(GqaOnline, options=PythonPrintOptions(show_types=True)) - verbose = [ - line for line in unnamed_axes.splitlines() if " # " in line and "ShardLayout(" in line - ] - - assert verbose - for line in verbose: - assert 'mesh=Mesh((Topology("cta", ' in line - assert "names=" not in line - assert "@ " not in line.split(" # ", 1)[1] - - several_meshes = as_script( - MoEMegaKernel, options=PythonPrintOptions(show_types=True) - ) - annotations = [ - line.split(" # ", 1)[1].split("; ", 1)[0] - for line in several_meshes.splitlines() - if " # Tensor[" in line - ] - - assert any("@ cta_2.tile" in annotation for annotation in annotations) - - def test_binding_metadata_names_the_emitted_binding(): tensor_type = TensorType.scalar(DType.f32) source = Var(name="source", type=tensor_type) diff --git a/tests/inspection/test_roundtrip.py b/tests/inspection/test_roundtrip.py index 02dd6428..92e75fe6 100644 --- a/tests/inspection/test_roundtrip.py +++ b/tests/inspection/test_roundtrip.py @@ -217,28 +217,6 @@ def test_nested_composed_shard_layout_roundtrips_without_flattening() -> None: assert as_script(import_dsl(printed)) == printed -def test_compound_split_without_explicit_strides_stays_compact() -> None: - fn = import_dsl( - "from __future__ import annotations\n" - "from tilefoundry import func\n" - "from tilefoundry.dsl import DimVar, Tensor\n" - "from tilefoundry.ir.types.shard import Layout, Mesh, Topology\n" - "\n" - 'n = DimVar("n", 1, 65)\n' - "tiles = ((n - 1) // 8) + 1\n" - 'cta = Mesh((Topology("cta", tiles),), Layout((tiles,), (1,)), names=("tile",))\n' - "\n" - "@func\n" - 'def f(x: Tensor[(tiles, 8), "f32", (tiles @ cta.tile, 8)]):\n' - " return x\n" - ) - - printed = as_script(fn) - - assert "(((n - 1) // 8) + 1) @ cta.tile" in printed - assert as_script(import_dsl(printed)) == printed - - def test_carry_updates_print_last_without_shadowing_the_old_value() -> None: fn = import_dsl( _HEADER + "\n@func\n" diff --git a/tests/inspection/test_tir_roundtrip.py b/tests/inspection/test_tir_roundtrip.py index 9666571a..9104df18 100644 --- a/tests/inspection/test_tir_roundtrip.py +++ b/tests/inspection/test_tir_roundtrip.py @@ -18,6 +18,7 @@ FIXTURES = Path(__file__).parents[1] / "fixtures" CANONICAL = tuple(path for path in (FIXTURES / "tir").glob("*.py") if path.name != "layouts.py") +SUGAR = FIXTURES / "inspection" / "type_printer_sugar.py" def _module_in(path: Path): @@ -37,6 +38,19 @@ def test_fixture_prints_back_to_its_own_source(path: Path) -> None: assert as_script(_module_in(path)) == path.read_text() +def test_placed_types_print_and_reparse_as_layout_sugar() -> None: + """Placement sugar is the whole type surface, and it parses back to itself. + + The golden is the review surface: a value the sugar cannot state would + appear here as the verbose ``ShardLayout(...)`` form instead, and a mesh + axis the sugar named wrongly would not survive the reparse. + """ + printed = as_script(_module_in(SUGAR)) + + assert printed == SUGAR.with_suffix(".printed.txt").read_text() + assert as_script(import_dsl(printed, name="TypePrinterSugar")) == printed + + def test_mixed_hir_tir_module_prints_both_function_families() -> None: mixed = import_dsl( "from tilefoundry import func, module, prim_func\n" @@ -89,7 +103,7 @@ def test_tir_for_if_and_sync_mesh_forms_roundtrip() -> None: ) printed = as_script(function) assert "T.sync(thread)" in printed - assert "T.sync(Mesh(" in printed + assert "T.sync(thread[:])" in printed assert as_script(import_dsl(printed, name="device")) == printed diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 09fa693a..08293a9b 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -46,7 +46,11 @@ def tile_device( c: Tensor[(128,), "f32"], ): atom = T.cuda.mma.atom(op=_OP) - with Mesh((Topology("thread", 32),), _MESH_LAYOUT) as m: + with Mesh( + (Topology("thread", 32),), + _MESH_LAYOUT, + names=("warp", "lane"), + ) as m: a_view = T.tensor_view( a, layout=ShardLayout( diff --git a/tests/parser/test_calls.py b/tests/parser/test_calls.py index 946bf07e..707ac3ad 100644 --- a/tests/parser/test_calls.py +++ b/tests/parser/test_calls.py @@ -25,6 +25,8 @@ from tilefoundry.parser import ParseError from tilefoundry.target import CpuTarget, CudaTarget +_EXTERNAL_PLACEMENT_MESH = Mesh(("cta",), (2,), names=("tile",)) + def test_matmul_layout_literals_are_parser_checked() -> None: assert get_args(MatMul.a_layout.annotation) == ("MK", "KM") @@ -346,7 +348,7 @@ def valueful_escape(x: Tensor[(2,), "f32"]): def test_mesh_binding_does_not_escape_its_with_scope() -> None: """A mesh alias is removed with its lexical frame after the with body.""" - with pytest.raises(ParseError, match="'mesh' is not an active Mesh"): + with pytest.raises(ParseError, match="'mesh' is not a lexical Mesh binding"): @module( entry="escaped_mesh", @@ -499,7 +501,7 @@ def root(x: Tensor[(8, 8), "f32"]): def test_a_boundary_rejects_a_bare_undeclared_mesh_name() -> None: - with pytest.raises(ParseError, match="'nope' is not an active Mesh"): + with pytest.raises(ParseError, match="'nope' is not a lexical Mesh binding"): @module( entry="root", @@ -520,3 +522,17 @@ class MissingTopology: @func(mesh=Mesh(("cta",), (8,), names=("b",))) def run(x: Tensor[(8 @ mesh.b, 8), "f32"]): # noqa: F821 return tf.add(x, x) + + +def test_placement_rejects_an_external_mesh_axis_binding() -> None: + with pytest.raises(ParseError, match="'_EXTERNAL_PLACEMENT_MESH' is not a lexical Mesh binding"): + + @module( + entry="run", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 2),), + ) + class ExternalMeshPlacement: + @func + def run(x: Tensor[(2 @ _EXTERNAL_PLACEMENT_MESH.tile,), "f32"]): # noqa: F821 + return x