Skip to content

fix(ir): correct model analysis and parser surfaces #123

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhen8838 merged 10 commits into from
Aug 26, 2026
Merged

fix(ir): correct model analysis and parser surfaces #123

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
54af552
fix(analysis): resolve cross-scope predecessors in the performance ti…
zhen8838 Aug 25, 2026
1a26895
fix(analysis): count a matmul's rows once whichever axis the mesh split
zhen8838 Aug 25, 2026
7306f84
fix(types): fold an extent written as arithmetic over constants
zhen8838 Aug 25, 2026
98a9597
fix(parser): accept a func docstring and a concat sequence
zhen8838 Aug 25, 2026
6d3be93
fix(parser): explain tile window slice bounds
zhen8838 Aug 26, 2026
d58ddc8
fix(parser): validate tile window slice bounds as a rule
zhen8838 Aug 26, 2026
52af273
feat(parser): make variadic input sequences explicit
zhen8838 Aug 26, 2026
1385518
fix(parser): align scalar and variadic normalization
zhen8838 Aug 26, 2026
83e31b2
fix(ir): refine reviewed parser and analysis changes
zhen8838 Aug 26, 2026
ada0146
feat(hir): add explicit matmul operand layouts
zhen8838 Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions docs/develop.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,9 +138,10 @@ short bullet list — keep it that way.

### DSL / HIR authoring

- No docstring in an `@func` body (parser rejects bare expressions).
- A leading `@func` docstring is accepted; strings in nested blocks remain bare expressions.
- Use `tf.<op>` attribute path; do not alias individual ops.
- Variadic ops take positional inputs; attributes go by keyword.
- Variadic ops take one explicit list, tuple, or supported static list comprehension;
attributes go by keyword.
- Every `@func` parameter MUST reach the return through real ops;
dead `_ = expr` assignments do not count.

Expand Down
9 changes: 5 additions & 4 deletions docs/spec/core-ir.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -393,10 +393,11 @@ class Call(Expr):
- a value-form `Call` is anchored by `LetStmt` in TIR; a Stmt-position effect
invocation is `Evaluate(op, args)`.
- A `Call` MUST NOT appear as a top-level Stmt directly.
- `len(args)` MUST equal the number of `kind="input"` ParamDefs on
`target`.
- Each `args[i].type` MUST satisfy the i-th input ParamDef's pattern
/ typeinfer rule.
- Normally, `len(args)` MUST equal the number of `kind="input"` ParamDefs on
`target`. A sole input annotated `Tuple[T]` describes a variadic sequence,
and every flattened argument corresponds to that ParamDef.
- Each argument MUST satisfy its corresponding input ParamDef's pattern /
typeinfer rule.

### 2.2 `Var` / `Constant` / `Tuple`

Expand Down
19 changes: 15 additions & 4 deletions docs/spec/hir.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -660,14 +660,19 @@ their input when it states one. An input with `layout=None` produces a view with

##### Concat

`Concat(inputs..., axis=a)` materializes a rank-preserving tensor by joining
`Concat([inputs...], axis=a)` materializes a rank-preserving tensor by joining
each input segment along `a`. All inputs MUST have one common rank and dtype,
and every non-concatenated dimension MUST match. Negative `axis` values resolve
against that rank. The output's concatenated extent is the sum of the input
extents; every input access map is defined only on its segment and subtracts
the preceding segments' extent from that axis, while the output map is the
identity.

The authored `inputs` value MUST be one explicit list, tuple, or supported
static list comprehension. Its Tensor elements flatten into `Call.args` in
source order; direct positional tensors and implicit iterable expansion are not
part of this surface.

Type inference derives fresh output ownership from those access maps. A
`Split` on a non-concatenated axis MAY propagate when shared ownership
propagation proves a zero-offset projection. A `Split` on the concatenated axis
Expand Down Expand Up @@ -859,15 +864,15 @@ class Stack(Op):
Attributes:
inputs: input; variadic tensors to stack.
axis: attribute; inserted result axis.
is_variadic: attribute; Whether the input parameter consumes all args.
"""

inputs: Tensor
inputs: Tuple[Tensor]
axis: int
is_variadic: ClassVar[bool] = True
```

- constraints:
- The authored `inputs` value MUST be one explicit list, tuple, or supported
static list comprehension. Its Tensor elements flatten into `Call.args`.
- At least one input is required; every input MUST have the same shape and
dtype. `axis` MUST resolve in `[-rank-1, rank]`.
- The operation materializes one distinct result. The inserted axis is local
Expand Down Expand Up @@ -1264,6 +1269,12 @@ Consensus torch.nn.functional ops.
is `Broadcast` / replicated. On each mesh axis, one `Partial(sum)` is
therefore allowed; a double-Partial input or a non-`sum` reduction is
rejected.
- `MatMul.a_layout` is `"MK"` (the default) or `"KM"`; `MatMul.b_layout` is
`"KN"` (the default) or `"NK"`. These literals state the physical order of
each operand's final two axes. The access relation maps them to logical
`(M, K)` and `(K, N)` before deriving the output, contraction ownership, and
`Partial(sum)` state. Its cost is `2 * numel(local_output) * local_K`, where
`local_K` is reconstructed from that same logical-axis mapping.
- `Conv2D` requires rank-4 NCHW input and OIHW weight, a rank-1 bias, and one
common operand dtype. `stride` and `dilation` are positive length-2 tuples,
`padding` is a non-negative length-2 tuple, and `groups` is positive. Input
Expand Down
82 changes: 24 additions & 58 deletions docs/spec/parser.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,12 +100,14 @@ type-annotation ::= tensor
| scalar-type
signature ::= (name ':' type-annotation (',' name ':' type-annotation)*)?
return-type ::= type-annotation
loop-iterator ::= 'tile'
| 'range'
loop-carry-statement ::= expression '=' expression
| 'for' name 'in' expression ':' loop-carry
| statement
loop-carry ::= (loop-carry-statement (newline loop-carry-statement)*)?
loop-header ::= 'for' identifier 'in' ('tile' | 'range') '(' expression (',' expression)*
')' ':' loop-carry
loop-header ::= 'for' identifier 'in' loop-iterator '(' (expression | name '=' expression)
(',' (expression | name '=' expression))* ')' ':' loop-carry
loop-body ::= (statement (newline statement)*)?
for ::= 'for' name 'in' expression ':' loop-body
mesh-context ::= ('Mesh' | primary '.' identifier) '(' (expression | ('layout' | 'names')
Expand Down Expand Up @@ -171,21 +173,10 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b
<!-- parser-constraints:start -->
| Owner | Situation | Rule | Statement | Source |
| --- | --- | --- | --- | --- |
| binary_expression | expression | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| binary_expression | expression | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| binary_expression | expression | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| binary_expression | slice_endpoint | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| binary_expression | slice_endpoint | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| binary_expression | slice_endpoint | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| binary_expression | subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| binary_expression | subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| binary_expression | subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| dim_expr | dim_expr | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py |
| dim_expr | layout_extent | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py |
| dim_expr | layout_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py |
| dim_expr | tensor_dim_expr | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py |
| dim_expr | tensor_optional_slot | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py |
| dim_expr | tensor_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py |
| binary_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 | expression, slice_endpoint, subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| binary_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 |
| dtype | tensor_dtype | CanonicalDTypeRule | A dtype must resolve to a canonical DType. | src/tilefoundry/parser/ast_pattern.py |
| explicit_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py |
| explicit_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py |
Expand All @@ -194,52 +185,27 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b
| function | function | FunctionReturnRule | A HIR function body's inferred type must match its return type. | src/tilefoundry/parser/pattern_nodes.py |
| function | function | FunctionRoleValidationRule | A root, variant, or converter must satisfy its role before registration. | src/tilefoundry/parser/pattern_nodes.py |
| function | function | FunctionSignatureRule | A function must construct an ordered parameter tuple. | 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 | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py |
| layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py |
| op_call | expression | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | expression | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | expression | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | slice_endpoint | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | slice_endpoint | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | slice_endpoint | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| placed_layout | layout_shape | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py |
| placed_layout | layout_shape | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py |
| placed_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py |
| placed_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py |
| placed_layout | tensor_shape | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py |
| placed_layout | tensor_shape | 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 | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | expression, slice_endpoint, subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | expression, slice_endpoint, subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| op_call | expression, slice_endpoint, subscript_index | CallVariadicInputFormRule | A variadic call must use one explicit list, tuple, or supported static list comprehension. | src/tilefoundry/parser/pattern_nodes.py |
| placed_layout | layout_shape, tensor_optional_slot, tensor_shape | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py |
| placed_layout | layout_shape, tensor_optional_slot, tensor_shape | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py |
| plain_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py |
| plain_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py |
| shape | layout_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py |
| shape | layout_strides | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py |
| shape | tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | 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 |
| storage | tensor_optional_slot | StorageValueRule | Storage must resolve to a StorageKind. | src/tilefoundry/parser/ast_pattern.py |
| tensor | annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py |
| tensor | annotation | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py |
| tensor | expression | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py |
| tensor | expression | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py |
| tensor | slice_endpoint | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py |
| tensor | slice_endpoint | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py |
| tensor | subscript_index | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py |
| tensor | subscript_index | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py |
| tensor | type_annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py |
| tensor | type_annotation | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py |
| unary_expression | expression | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | expression | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | expression | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | slice_endpoint | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | slice_endpoint | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | slice_endpoint | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py |
| module | module_function | ModuleFunctionValidationRule | A module function must satisfy its root, variant, or converter role before mutation. | 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_finalization | ModuleFinalizationRule | A module declaration must contain valid unique members and a resolvable entry. | 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 |
| unary_expression | expression, slice_endpoint, subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py |
| unary_expression | expression, slice_endpoint, subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py |
| 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 |
<!-- parser-constraints:end -->

## 3. Implementation Overview
Expand Down
3 changes: 3 additions & 0 deletions docs/tutorial/migrate.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@ forest, and `--source` names the directory and lists its files.
`@<func>.converter` — [runtime §1.1.2](../spec/runtime.md#112-weight-converter-and-prepare--forward).
- Dimensions come from the published config. `head_dim` is a published field, not
`hidden ÷ num_heads`; for this model those differ.
- Variadic tensor operations take one explicit sequence. Write
`tf.concat([left, right], axis=-1)` or `tf.stack((left, right), axis=0)`, rather
than passing tensors as separate positional arguments.

## The five access faces

Expand Down
Loading
Loading