diff --git a/docs/develop.md b/docs/develop.md index 19e51b03..750cd1fb 100644 --- a/docs/develop.md +++ b/docs/develop.md @@ -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.` 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. diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 11e9c95c..cff0da4b 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -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` diff --git a/docs/spec/hir.md b/docs/spec/hir.md index aca2e59c..98821842 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -660,7 +660,7 @@ 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 @@ -668,6 +668,11 @@ 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 @@ -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 @@ -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 diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 1b044d70..5fedfe00 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -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') @@ -171,21 +173,10 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | 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 | @@ -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 | ## 3. Implementation Overview diff --git a/docs/tutorial/migrate.md b/docs/tutorial/migrate.md index 21588610..0c94eb94 100644 --- a/docs/tutorial/migrate.md +++ b/docs/tutorial/migrate.md @@ -25,6 +25,9 @@ forest, and `--source` names the directory and lists its files. `@.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 diff --git a/examples/granite_4_0_h_small-cuda/model.py b/examples/granite_4_0_h_small-cuda/model.py index b7cebd26..ace61956 100644 --- a/examples/granite_4_0_h_small-cuda/model.py +++ b/examples/granite_4_0_h_small-cuda/model.py @@ -161,7 +161,7 @@ def conv_step( # closes on this token, so the whole convolution is one multiply # against the kernel and one reduction over it. Channels do not mix -- # that is what depthwise means here, and why no matmul appears. - window = tf.concat(conv_state, entry, axis=2) + window = tf.concat([conv_state, entry], axis=2) weighted = window * tf.reshape(conv_w, new_shape=(1, _CONVD, _KRN)) summed = tf.reduce(weighted, axes=(-1,), keepdim=False, kind="sum") return tf.silu(summed + tf.reshape(conv_b, new_shape=(1, _CONVD))) diff --git a/examples/minicpm3_4b-cutedsl/model.py b/examples/minicpm3_4b-cutedsl/model.py index de67c2e9..714525a0 100644 --- a/examples/minicpm3_4b-cutedsl/model.py +++ b/examples/minicpm3_4b-cutedsl/model.py @@ -258,8 +258,8 @@ def mla_attention( k_rope_b = tf.repeat_interleave(k_rope_e, repeats=_H, axis=2) # Step 6: reassemble nope + rope, each back in its original slot. - query = tf.concat(q_nope, q_rope_e, axis=-1) - k_new = tf.concat(k_nope, k_rope_b, axis=-1) + query = tf.concat([q_nope, q_rope_e], axis=-1) + k_new = tf.concat([k_nope, k_rope_b], axis=-1) # Step 7: attend the cache and the token itself, then project out. q_s = query * scale diff --git a/examples/qwen3_5_35b_a3b-tilelang/model.py b/examples/qwen3_5_35b_a3b-tilelang/model.py index 7aaceb08..7c5e8fe6 100644 --- a/examples/qwen3_5_35b_a3b-tilelang/model.py +++ b/examples/qwen3_5_35b_a3b-tilelang/model.py @@ -126,7 +126,7 @@ def conv_step( # against the kernel and one reduction over it. Channels do not mix # -- that is what depthwise means here, and it is why no matmul # appears. - window = tf.concat(conv_state, entry, axis=2) + window = tf.concat([conv_state, entry], axis=2) weighted = window * tf.reshape(conv_w, new_shape=(1, _CONV, _KERNEL)) summed = tf.reduce(weighted, axes=(-1,), keepdim=False, kind="sum") return tf.silu(summed) @@ -307,7 +307,7 @@ def partial_rope( rot = x[:, :, :, :_ROT] tail = x[:, :, :, _ROT:_D] turned, _ = tf.rope(rot, rot, cos_cache, sin_cache, pos_ids) - return tf.concat(turned, tail, axis=-1) + return tf.concat([turned, tail], axis=-1) @func def partial_rope_kv( @@ -322,7 +322,7 @@ def partial_rope_kv( rot = x[:, :, :, :_ROT] tail = x[:, :, :, _ROT:_D] turned, _ = tf.rope(rot, rot, cos_cache, sin_cache, pos_ids) - return tf.concat(turned, tail, axis=-1) + return tf.concat([turned, tail], axis=-1) @func def full_attention( diff --git a/src/tilefoundry/dsl/_stub_gen.py b/src/tilefoundry/dsl/_stub_gen.py index 31cbd050..28f1e9a2 100644 --- a/src/tilefoundry/dsl/_stub_gen.py +++ b/src/tilefoundry/dsl/_stub_gen.py @@ -12,7 +12,7 @@ import argparse import enum from pathlib import Path -from typing import Iterable +from typing import Iterable, Literal, get_args, get_origin import tilefoundry.ir # noqa: F401 (populates schema registry as a side effect) from tilefoundry.ir.core.expr import Expr @@ -25,6 +25,8 @@ { "Any", "Expr", + "Literal", + "Sequence", "object", "int", "float", @@ -42,7 +44,8 @@ def _expr_type_for_input(pd: ParamDef) -> str: """Input ParamDefs always carry Expr operands at the DSL surface.""" - base = "Expr" + args = get_args(pd.annotation) if get_origin(pd.annotation) is tuple else () + base = "Sequence[Expr]" if len(args) == 1 or args[-1:] == (Ellipsis,) else "Expr" if pd.optional: base = f"{base} | None" return base @@ -73,6 +76,11 @@ def _annotation_type(pd: ParamDef) -> tuple[str, str]: if pd.optional: base = f"{base} | None" return base, name + if get_origin(ann) is Literal: + base = f"Literal[{', '.join(repr(value) for value in get_args(ann))}]" + if pd.optional: + base = f"{base} | None" + return base, "" name = getattr(ann, "__name__", None) if name is None: return "Any", "" @@ -120,6 +128,8 @@ def _function_stub( types_seen.add(type_name) if "Literal[" in rendered: types_seen.add("Literal") + if "Sequence[" in rendered: + types_seen.add("Sequence") parts.append(rendered) sig = ", ".join(parts) head = f"def {schema.name}({sig}) -> Expr: ..." @@ -166,7 +176,12 @@ def _resolve_type_modules(dialect: str, type_names: Iterable[str]) -> list[tuple def _module_header(dialect: str, types_seen: set[str]) -> str: """Build the auto-import header for a generated stub file.""" imports = _resolve_type_modules(dialect, types_seen) - typing_names = "Any, Literal, overload" if "Literal" in types_seen else "Any, overload" + typing_imports = {"Any", "overload"} + if "Literal" in types_seen: + typing_imports.add("Literal") + if "Sequence" in types_seen: + typing_imports.add("Sequence") + typing_names = ", ".join(sorted(typing_imports)) lines = [ "# AUTO-GENERATED by tilefoundry.dsl._stub_gen — do not edit.", "# Regenerate with `python -m tilefoundry.dsl regen`.", diff --git a/src/tilefoundry/ir/hir/nn/matmul.py b/src/tilefoundry/ir/hir/nn/matmul.py index e6cad46a..cc6c33a3 100644 --- a/src/tilefoundry/ir/hir/nn/matmul.py +++ b/src/tilefoundry/ir/hir/nn/matmul.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Literal + import isl import torch @@ -20,7 +22,6 @@ BoundaryRelation, coordinates_of, iterating, - logical_axes_of, register_access_relation, shape_from_relation, ) @@ -29,8 +30,33 @@ @register_op class MatMul(Op): + """Batched matrix multiplication with explicit physical matrix-axis order.""" + lhs = ParamDef(kind="input", pattern=Tensor) rhs = ParamDef(kind="input", pattern=Tensor) + a_layout = ParamDef( + kind="attribute", annotation=Literal["MK", "KM"], default="MK" + ) + b_layout = ParamDef( + kind="attribute", annotation=Literal["NK", "KN"], default="KN" + ) + + +def matmul_axes(op: MatMul) -> tuple[int, int, int, int]: + """Return physical ``(A.M, A.K, B.N, B.K)`` axes for the layout literals.""" + if op.a_layout == "MK": + a_m, a_k = -2, -1 + elif op.a_layout == "KM": + a_m, a_k = -1, -2 + else: + raise ValueError(f"MatMul: a_layout must be 'MK' or 'KM', got {op.a_layout!r}") + if op.b_layout == "NK": + b_n, b_k = -2, -1 + elif op.b_layout == "KN": + b_n, b_k = -1, -2 + else: + raise ValueError(f"MatMul: b_layout must be 'NK' or 'KN', got {op.b_layout!r}") + return a_m, a_k, b_n, b_k def _k_split_axes(t, k_tensor_axis: int) -> "frozenset[int]": @@ -52,17 +78,14 @@ def _broadcast_batch(lhs_batch, rhs_batch): return broadcast_shapes(tuple(lhs_batch), tuple(rhs_batch), raising=False) -def _held_axis(local, logical, axis: int) -> int: - """How much of one logical axis this participant holds.""" - held = 1 - for position, owner in enumerate(logical_axes_of(local, logical)): - if owner == axis: - held *= local.shape[position] - return held - - def _operand_reads( - shape: tuple, out_shape: tuple, inner: str, *, contracts_last: bool + shape: tuple, + out_shape: tuple, + inner: str, + *, + kept_axis: int, + output_axis: int, + contraction_axis: int, ) -> list[str]: """One coordinate per axis of an operand of the contraction. @@ -72,13 +95,15 @@ def _operand_reads( only coordinate rather than the result's. """ rank = len(shape) + kept_axis %= rank + contraction_axis %= rank shift = len(out_shape) - rank reads: list[str] = [] for axis in range(rank): - if axis == rank - 1: - reads.append(inner if contracts_last else f"d{len(out_shape) - 1}") - elif axis == rank - 2: - reads.append(f"d{len(out_shape) - 2}" if contracts_last else inner) + if axis == contraction_axis: + reads.append(inner) + elif axis == kept_axis: + reads.append(f"d{len(out_shape) + output_axis}") elif is_one(shape[axis]) and not is_one(out_shape[axis + shift]): reads.append("0") else: @@ -99,27 +124,34 @@ def _matmul_access_relation(call: "Call", ctx) -> AccessRelations: """ lhs = ctx.type_of(call.args[0]) rhs = ctx.type_of(call.args[1]) + a_m, a_k, b_n, b_k = matmul_axes(call.target) batch = _broadcast_batch(lhs.shape[:-2], rhs.shape[:-2]) if batch is None: raise ValueError( f"MatMul batches {tuple(lhs.shape[:-2])} against " f"{tuple(rhs.shape[:-2])}, which do not broadcast" ) - out_shape = (*batch, lhs.shape[-2], rhs.shape[-1]) - summed = lhs.shape[-1] + out_shape = (*batch, lhs.shape[a_m], rhs.shape[b_n]) + summed = lhs.shape[a_k] rank = len(out_shape) dims = ", ".join((*(f"d{index}" for index in range(rank)), "k")) inner = "0" if is_one(summed) else "k" inputs = [] - for shape, contracts_last, held in ( - (tuple(lhs.shape), True, lhs.shape[-1]), - (tuple(rhs.shape), False, rhs.shape[-2]), + for shape, kept_axis, output_axis, contraction_axis in ( + (tuple(lhs.shape), a_m, -2, a_k), + (tuple(rhs.shape), b_n, -1, b_k), ): - reads = _operand_reads(shape, out_shape, inner, contracts_last=contracts_last) + reads = _operand_reads( + shape, + out_shape, + inner, + kept_axis=kept_axis, + output_axis=output_axis, + contraction_axis=contraction_axis, + ) inputs.append( BoundaryRelation(AffineAccess(isl.map(f"{{ [{dims}] -> [{', '.join(reads)}] }}"))) ) - del held accumulates = ", ".join(f"d{index}" for index in range(rank)) return iterating( (*out_shape, summed), @@ -144,19 +176,26 @@ def _elements(shape: tuple) -> int: def _(call: "Call", ctx: "TypeInferContext") -> TensorType: lhs = ctx.type_of(call.args[0]) rhs = ctx.type_of(call.args[1]) + try: + a_m, a_k, b_n, b_k = matmul_axes(call.target) + except ValueError as error: + ctx.error(call, str(error).removeprefix("MatMul: ")) if lhs.dtype != rhs.dtype: ctx.error(call, f"MatMul dtype mismatch: {lhs.dtype} vs {rhs.dtype}") if len(lhs.shape) < 2 or len(rhs.shape) < 2: ctx.error(call, "MatMul requires rank >= 2 on both operands") if _broadcast_batch(lhs.shape[:-2], rhs.shape[:-2]) is None: ctx.error(call, f"MatMul batch-dim mismatch {lhs.shape[:-2]} vs {rhs.shape[:-2]}") - if lhs.shape[-1] != rhs.shape[-2]: + if lhs.shape[a_k] != rhs.shape[b_k]: ctx.error( call, - f"MatMul contraction-dim mismatch: lhs K={lhs.shape[-1]} vs rhs K={rhs.shape[-2]}", + f"MatMul contraction-dim mismatch: lhs K={lhs.shape[a_k]} vs " + f"rhs K={rhs.shape[b_k]}", ) - if _k_split_axes(lhs, len(lhs.shape) - 1) != _k_split_axes(rhs, len(rhs.shape) - 2): + if _k_split_axes(lhs, a_k % len(lhs.shape)) != _k_split_axes( + rhs, b_k % len(rhs.shape) + ): ctx.error( call, "MatMul contraction dim K must be split on the same mesh axes for both operands", @@ -168,7 +207,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: out_batch = _broadcast_batch(lhs.shape[:-2], rhs.shape[:-2]) out_shape = shape_from_relation( - relation, (*out_batch, lhs.shape[-2], rhs.shape[-1], lhs.shape[-1]) + relation, (*out_batch, lhs.shape[a_m], rhs.shape[b_n], lhs.shape[a_k]) ) k_domain_dim = len(out_shape) try: @@ -187,6 +226,11 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: @register_eval(MatMul) def _eval_matmul(ctx): - - out = torch.matmul(ctx.args[0].data, ctx.args[1].data) + lhs = ctx.args[0].data + rhs = ctx.args[1].data + if ctx.op.a_layout == "KM": + lhs = lhs.transpose(-1, -2) + if ctx.op.b_layout == "NK": + rhs = rhs.transpose(-1, -2) + out = torch.matmul(lhs, rhs) return TensorValue(data=out, type=ctx.result_type) diff --git a/src/tilefoundry/ir/hir/tensor/concat.py b/src/tilefoundry/ir/hir/tensor/concat.py index cd0b75a7..0166f50f 100644 --- a/src/tilefoundry/ir/hir/tensor/concat.py +++ b/src/tilefoundry/ir/hir/tensor/concat.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Tuple + import isl import torch @@ -47,9 +49,7 @@ class Concat(Op): Expr). The lone Param entry documents element type. """ - is_variadic: ClassVar[bool] = True - - inputs = ParamDef(kind="input", pattern=Tensor) + inputs = ParamDef(kind="input", annotation=Tuple[Tensor], pattern=Tensor) axis = ParamDef(kind="attribute", annotation=int) diff --git a/src/tilefoundry/ir/hir/tensor/stack.py b/src/tilefoundry/ir/hir/tensor/stack.py index d60f19b2..806d29b3 100644 --- a/src/tilefoundry/ir/hir/tensor/stack.py +++ b/src/tilefoundry/ir/hir/tensor/stack.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Tuple + import isl import torch @@ -38,9 +40,7 @@ class Stack(Op): """Variadic input op. See Concat for encoding rationale.""" - is_variadic: ClassVar[bool] = True - - inputs = ParamDef(kind="input", pattern=Tensor) + inputs = ParamDef(kind="input", annotation=Tuple[Tensor], pattern=Tensor) axis = ParamDef(kind="attribute", annotation=int) diff --git a/src/tilefoundry/ir/types/dim_isl.py b/src/tilefoundry/ir/types/dim_isl.py index a1d7cdd1..87735d4e 100644 --- a/src/tilefoundry/ir/types/dim_isl.py +++ b/src/tilefoundry/ir/types/dim_isl.py @@ -5,6 +5,7 @@ import isl from tilefoundry.ir.core.expr import Call, Constant, Var +from tilefoundry.ir.core.kinds import BinaryKind from .dim import ( _DIM_OP_TYPES, @@ -17,8 +18,19 @@ DimSub, DimVar, ) +from .dtype import IntegerDType from .tensor_type import TensorType +_INTEGER_BINARY_DIM_OP = { + BinaryKind.ADD: DimAdd, + BinaryKind.SUB: DimSub, + BinaryKind.MUL: DimMul, + BinaryKind.FLOOR_DIV: DimFloorDiv, + BinaryKind.MOD: DimMod, + BinaryKind.MIN: DimMin, + BinaryKind.MAX: DimMax, +} + def _is_const(node) -> bool: if isinstance(node, bool): @@ -92,7 +104,16 @@ def visit_Var(self, dim: Var) -> str: def visit_Call(self, dim: Call) -> str: op = type(dim.target) if op not in _DIM_OP_TYPES: - return _bind_param(dim, self.params, self.param_map, self.identities) + kind = getattr(dim.target, "kind", None) + if not ( + isinstance(dim.type, TensorType) + and dim.type.shape == () + and isinstance(dim.type.dtype, IntegerDType) + and isinstance(kind, BinaryKind) + and kind in _INTEGER_BINARY_DIM_OP + ): + return _bind_param(dim, self.params, self.param_map, self.identities) + op = _INTEGER_BINARY_DIM_OP[kind] a, b = dim.args if op is DimMul and not (_is_const(a) or _is_const(b)): name = _bind_param(dim, self.params, self.param_map, self.identities) diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 610deda3..fe9f4bce 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -1237,6 +1237,7 @@ def child( binding_name: str | None = None, expected_type: object | None = None, values: Mapping[str, object] | None = None, + lexical_bindings: Mapping[str, object] | None = None, isolated_scope: bool = False, function: FuncParserContext | None = None, module: ModuleBuildContext | None = None, @@ -1258,6 +1259,9 @@ def child( ) else: scope = self.lexical_scope.fork() if isolated_scope else self.lexical_scope + if lexical_bindings: + for name, value in lexical_bindings.items(): + scope.define(name, value) if expected_type is None and role == "return_value": expected_type = scope.lookup(_RETURN_TYPE) return MatchContext( @@ -1312,6 +1316,7 @@ class AstChild: isolated_scope: bool = False function_context: FuncParserContext | None = None module_context: ModuleBuildContext | None = None + lexical_bindings: Mapping[str, object] = field(default_factory=dict) @dataclass(frozen=True) @@ -1396,6 +1401,7 @@ def parse_node(pattern: AstPattern[T], node: ast.AST, context: MatchContext) -> binding_name=binding_name, expected_type=child.expected_type, values=child.values, + lexical_bindings=child.lexical_bindings, isolated_scope=child.isolated_scope, function=child.function_context, module=child.module_context, @@ -1590,6 +1596,7 @@ class ShapeDimRule: ) def apply(self, value, *, match, context): + value = runtime.normalize_dim(value) if isinstance(value, bool) or not isinstance( value, (int, runtime.DimVar, runtime.Expr) ): diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index a1512a8a..47b38f64 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -12,7 +12,7 @@ import operator from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, ClassVar +from typing import Any, ClassVar, Literal, get_args, get_origin from tilefoundry.ir.constraints import ( ConstraintProvenance, @@ -78,6 +78,7 @@ _resolve_reference, _slice_size, attach_authored_metadata, + parse_node, runtime, ) @@ -1843,6 +1844,217 @@ def apply(self, value, *, match, context): return value +@dataclass(frozen=True) +class VariadicInputs: + items: tuple[runtime.Expr, ...] + + def __init__(self, values): + items = tuple(values) + if not all(isinstance(value, runtime.Expr) for value in items): + raise TypeError("variadic inputs must contain Expr values") + object.__setattr__(self, "items", items) + + +class VariadicInputsPattern(ElementPattern): + element_name = "variadic_inputs" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "list", + BindPattern( + AstNodePattern( + ast.List, + FieldPattern( + "elts", + RepeatPattern( + ChildPattern( + "input_{index}", + ExpressionPattern(), + "variadic_input", + "input", + ) + ), + ), + ), + VariadicInputsPattern._bind_sequence, + ), + pattern_id="call.variadic.list", + ), + BranchPattern( + "tuple", + BindPattern( + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern( + ChildPattern( + "input_{index}", + ExpressionPattern(), + "variadic_input", + "input", + ) + ), + ), + ), + VariadicInputsPattern._bind_sequence, + ), + pattern_id="call.variadic.tuple", + ), + BindPattern( + AstNodePattern(ast.ListComp), + VariadicInputsPattern._bind_comprehension, + ), + BranchPattern( + "generator_expression", + AstNodePattern(ast.GeneratorExp), + pattern_id="call.variadic.generator_expression", + ), + BranchPattern( + "unsupported", + AstNodePattern(ast.expr), + pattern_id="call.variadic.unsupported", + ), + ) + ) + + @staticmethod + def _bind_sequence( + node: object, context: MatchContext, matched: AstMatch[Any] + ) -> AstMatch[Any]: + assert isinstance(node, (ast.List, ast.Tuple)) + if any(isinstance(element, ast.Starred) for element in node.elts): + raise ParseError.from_node( + node, + context, + "variadic input sequences do not support starred expansion", + ) + return matched + + @staticmethod + def _bind_comprehension( + node: object, context: MatchContext, matched: AstMatch[Any] + ) -> AstMatch[Any]: + assert isinstance(node, ast.ListComp) + if len(node.generators) != 1: + raise ParseError.from_node( + node, + context, + "variadic list comprehension supports exactly one generator", + ) + generator = node.generators[0] + if not isinstance(generator.target, ast.Name): + raise ParseError.from_node( + generator.target, + context, + "variadic list comprehension requires a simple Name target", + ) + if generator.ifs or generator.is_async: + raise ParseError.from_node( + generator, + context, + "variadic list comprehension does not support filters or async generators", + ) + iterator = generator.iter + if not ( + isinstance(iterator, ast.Call) + and isinstance(iterator.func, ast.Name) + and iterator.func.id == "range" + and not iterator.keywords + and 1 <= len(iterator.args) <= 3 + ): + raise ParseError.from_node( + iterator, + context, + "variadic list comprehension requires range with 1 to 3 static integer arguments", + ) + values = [] + for argument in iterator.args: + try: + value = parse_node(StaticValuePattern(), argument, context) + except ParseError as error: + raise ParseError.from_node( + argument, + context, + "variadic list comprehension range arguments must be static integers", + ) from error + if not isinstance(value, int) or isinstance(value, bool): + raise ParseError.from_node( + argument, + context, + "variadic list comprehension range arguments must be static integers", + ) + values.append(value) + children = tuple( + AstChild( + f"input_{index}", + ExpressionPattern(), + node.elt, + "variadic_input", + "input", + lexical_bindings={generator.target.id: value}, + isolated_scope=True, + ) + for index, value in enumerate(range(*values)) + ) + return dataclasses.replace( + matched, + pattern_id="call.variadic.list_comp", + branch_id="list_comp", + children=children, + ) + + @staticmethod + def construct(match, children, context): + return VariadicInputs(tuple(children.values())) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +def _variadic_item_annotation(param: object) -> object | None: + if getattr(param, "kind", None) != "input": + return None + annotation = getattr(param, "annotation", None) + if get_origin(annotation) is not tuple: + return None + args = get_args(annotation) + if len(args) == 1: + return args[0] + if len(args) == 2 and args[1] is Ellipsis: + return args[0] + return None + + +@dataclass(frozen=True) +class CallVariadicInputFormRule: + STATEMENT: ClassVar[str] = ( + "A variadic call must use one explicit list, tuple, or supported static " + "list comprehension." + ) + + def apply(self, value, *, match, context): + schema = match.captures.get("schema") + if not isinstance(schema, runtime.OpSchema): + return value + inputs = tuple(param for param in schema.signature if param.kind == "input") + if len(inputs) != 1 or _variadic_item_annotation(inputs[0]) is None: + return value + argument = match.node.args[0] + if isinstance(argument, ast.GeneratorExp): + raise ParseError.from_node( + argument, + context, + "variadic inputs do not support generator expressions; use a list comprehension", + ) + if not isinstance(argument, (ast.List, ast.Tuple, ast.ListComp)): + raise ParseError.from_node( + argument, + context, + "variadic inputs require an explicit list, tuple, or supported list comprehension", + ) + return value + + class CallPattern(ElementPattern): element_name = "op_call" syntax = LazyPattern( @@ -1875,6 +2087,8 @@ class CallPattern(ElementPattern): @staticmethod def _pattern_for_param(param: object, node: ast.AST) -> AstPattern[Any]: annotation = param.annotation + if _variadic_item_annotation(param) is not None: + return VariadicInputsPattern() if annotation is runtime.TensorType and isinstance(node, ast.Subscript): return TensorPattern() if annotation is runtime.DType: @@ -1886,16 +2100,46 @@ def _pattern_for_param(param: object, node: ast.AST) -> AstPattern[Any]: return StaticValuePattern() @staticmethod - def _schema_children(node: ast.Call, schema: object) -> tuple[AstChild, ...] | None: + def _schema_children( + node: ast.Call, schema: object, context: MatchContext + ) -> tuple[AstChild, ...] | None: + """Bind a call's arguments to one op schema's inputs and attributes. + + A ``Tuple[T]`` input consumes exactly one explicit sequence and flattens + its elements into ``Call.args``. Attributes remain keyword-only, so the + sequence boundary cannot be confused with an attribute position. + """ params = tuple(schema.signature) inputs = [param for param in params if param.kind == "input"] attrs = [param for param in params if param.kind == "attribute"] - variadic = bool(getattr(schema.op_class, "is_variadic", False)) + variadic = ( + len(inputs) == 1 + and _variadic_item_annotation(inputs[0]) is not None + ) + positional = list(node.args) children: list[AstChild] = [] bound_attrs: set[str] = set() - for index, argument in enumerate(node.args): - if variadic or index < len(inputs): - name = inputs[0].name if variadic else inputs[index].name + if variadic: + if len(positional) != 1: + raise ParseError.from_node( + node, + context, + "variadic operations require exactly one list, tuple, or " + "supported list-comprehension input sequence", + ) + children.append( + AstChild( + "variadic_inputs", + CallPattern._pattern_for_param(inputs[0], positional[0]), + positional[0], + "variadic_inputs", + "inputs", + ) + ) + positional = [] + for index, argument in enumerate(positional): + if index < len(inputs): + name = inputs[index].name children.append( AstChild( f"input_{index}", @@ -1936,7 +2180,7 @@ def _schema_children(node: ast.Call, schema: object) -> tuple[AstChild, ...] | N "allocation" if param.annotation is runtime.TensorType else param.name, ) ) - if not variadic and len(node.args) < len(inputs): + if not variadic and len(positional) < len(inputs): return None return tuple(children) @@ -1985,7 +2229,7 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat ) if not isinstance(schema, runtime.OpSchema): return None - children = CallPattern._schema_children(node, schema) + children = CallPattern._schema_children(node, schema, context) if children is None: return None return dataclasses.replace( @@ -2000,7 +2244,11 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat def construct(match, children, context): if match.branch_id == "operation_call": schema = match.captures["schema"] - inputs = tuple(value for name, value in children.items() if name.startswith("input_")) + variadic_inputs = children.get("variadic_inputs") + if isinstance(variadic_inputs, VariadicInputs): + inputs = variadic_inputs.items + else: + inputs = tuple(value for name, value in children.items() if name.startswith("input_")) attrs = { name.removeprefix("attr_"): value for name, value in children.items() @@ -2009,6 +2257,13 @@ def construct(match, children, context): for name, value in tuple(attrs.items()): parameter = next((item for item in schema.signature if item.name == name), None) annotation = None if parameter is None else parameter.annotation + if get_origin(annotation) is Literal and value not in get_args(annotation): + choices = ", ".join(repr(choice) for choice in get_args(annotation)) + raise ParseError.from_node( + match.node, + context, + f"{name} must be one of {choices}, got {value!r}", + ) if ( isinstance(annotation, type) and issubclass(annotation, enum.Enum) @@ -2051,6 +2306,7 @@ def construct(match, children, context): raise RuntimeError(f"no constructor branch for {match.branch_id!r}") RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + CallVariadicInputFormRule(), CallBindingRule(), CallTypeInferenceRule(), CallExpectedTypeRule(), @@ -2315,6 +2571,29 @@ class IndexEndpointPattern(ElementPattern): RULES: ClassVar[tuple[AstRule[Any], ...]] = () +@dataclass(frozen=True) +class TileWindowSliceBoundRule: + STATEMENT: ClassVar[str] = ( + "A tile window cannot be used as a slice bound." + ) + + def apply(self, value, *, match, context): + for authored_name, value_name in ( + ("lower", "start"), + ("upper", "stop"), + ("step", "step"), + ): + if isinstance(getattr(value, value_name), slice): + raise ParseError.from_node( + getattr(match.node, authored_name), + context, + "a tile loop variable is already a window and cannot be " + "used as a slice bound; use x[:, t, :] or bind " + "base = t + 0 before slicing", + ) + return value + + class IndexSlicePattern(ElementPattern): element_name = "index_slice" syntax = LazyPattern( @@ -2345,7 +2624,7 @@ class IndexSlicePattern(ElementPattern): def construct(match, children, context): return slice(children.get("lower"), children.get("upper"), children.get("step")) - RULES: ClassVar[tuple[AstRule[Any], ...]] = () + RULES: ClassVar[tuple[AstRule[Any], ...]] = (TileWindowSliceBoundRule(),) class SubscriptIndexPattern(ElementPattern): @@ -2810,7 +3089,7 @@ class WithPattern(ElementPattern): BlockPattern(), "block", "with_body", - transform=_module_from_body, + transform=_body_as_ast_module, ), ), ), @@ -2841,7 +3120,7 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat AstChild( "body", BlockPattern(), - _module_from_body(node.body), + _body_as_ast_module(node.body), "block", "with_body", ), @@ -2954,7 +3233,7 @@ class LoopCarryStatementPattern(ElementPattern): "nested", LoopCarryPattern(), "loop_carry", - transform=_module_from_body, + transform=_body_as_ast_module, ), ), ), @@ -3019,6 +3298,36 @@ def construct(match, children, context): RULES: ClassVar[tuple[AstRule[Any], ...]] = () +class LoopIteratorPattern(ElementPattern): + element_name = "loop_iterator" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "tile", + AstNodePattern( + ast.Name, + FieldPattern("id", LiteralPattern("tile")), + ), + pattern_id="loop.iterator.tile", + ), + BranchPattern( + "range", + AstNodePattern( + ast.Name, + FieldPattern("id", LiteralPattern("range")), + ), + pattern_id="loop.iterator.range", + ), + ) + ) + + @staticmethod + def construct(match, children, context): + return match.branch_id + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + class LoopHeaderPattern(ElementPattern): element_name = "loop_header" syntax = LazyPattern( @@ -3032,7 +3341,7 @@ class LoopHeaderPattern(ElementPattern): ast.Call, FieldPattern( "func", - AstNodePattern(ast.Name), + LoopIteratorPattern(), ), FieldPattern("keywords", RepeatPattern(AstNodePattern(ast.keyword))), FieldPattern("args", RepeatPattern(AstNodePattern(ast.expr), minimum=1)), @@ -3044,7 +3353,7 @@ class LoopHeaderPattern(ElementPattern): "carry", LoopCarryPattern(), "loop_carry", - transform=_module_from_body, + transform=_body_as_ast_module, ), ), ), @@ -3062,12 +3371,6 @@ def _bind( assert isinstance(node.iter.func, ast.Name) kind = node.iter.func.id count = len(node.iter.args) - if kind not in {"tile", "range"}: - return PatternFailure( - "loop_header", - node.iter.func, - "loop iterator must be tile(...) or range(...)", - ) if node.iter.keywords: return PatternFailure( "loop_header", @@ -3102,7 +3405,7 @@ def _bind( AstChild( "carry", LoopCarryPattern(), - _module_from_body(node.body), + _body_as_ast_module(node.body), "loop_carry", ) ] @@ -3237,7 +3540,7 @@ class ForPattern(ElementPattern): "body", LoopBodyPattern(), "loop_body", - transform=_module_from_body, + transform=_body_as_ast_module, ), ), ), @@ -3742,7 +4045,7 @@ class FunctionPattern(ElementPattern): BlockPattern(), "block", "body", - transform=_module_from_body, + transform=lambda body: _body_as_ast_module(body, strip_docstring=True), ), ), ), @@ -3856,8 +4159,17 @@ def construct(match, children, context): ) -def _module_from_body(body: object) -> ast.Module: +def _body_as_ast_module(body: object, *, strip_docstring: bool = False) -> ast.Module: + """Wrap a statement list as an AST Module with explicit docstring policy.""" assert isinstance(body, list) + if ( + strip_docstring + and body + and isinstance(body[0], ast.Expr) + and isinstance(body[0].value, ast.Constant) + and isinstance(body[0].value.value, str) + ): + body = body[1:] return ast.Module(body=body, type_ignores=[]) @@ -3918,5 +4230,8 @@ def _module_from_body(body: object) -> ast.Module: "TupleExpressionPattern", "TypeAnnotationPattern", "UnaryExpressionPattern", + "CallVariadicInputFormRule", + "VariadicInputs", + "VariadicInputsPattern", "WithPattern", ] diff --git a/src/tilefoundry/parser/spec.py b/src/tilefoundry/parser/spec.py index 7283558d..a11e4c27 100644 --- a/src/tilefoundry/parser/spec.py +++ b/src/tilefoundry/parser/spec.py @@ -140,13 +140,33 @@ def _collect_module_rule_rows() -> tuple[RuleRow, ...]: return tuple(rows) +def _merge_rule_rows(rows: tuple[RuleRow, ...]) -> tuple[RuleRow, ...]: + """Render one row per owning element and rule, retaining every situation.""" + situations: dict[tuple[str, str, str, str], set[str]] = {} + for row in rows: + key = (row.owner, row.rule, row.statement, row.source) + situations.setdefault(key, set()).add(row.situation) + return tuple( + sorted( + RuleRow( + owner=owner, + situation=", ".join(sorted(rule_situations)), + rule=rule, + statement=statement, + source=source, + ) + for (owner, rule, statement, source), rule_situations in situations.items() + ) + ) + + def _escape_cell(value: str) -> str: return value.replace("|", "\\|").replace("\n", " ") def render_spec_content() -> str: root = FunctionPattern() - rows = (*_collect_rule_rows(root), *_collect_module_rule_rows()) + rows = _merge_rule_rows((*_collect_rule_rows(root), *_collect_module_rule_rows())) lines = [ "# Parser Grammar and Constraints", "", @@ -175,85 +195,36 @@ def render_spec_content() -> str: return "\n".join(lines) + "\n" -def render_parser_document() -> str: - """Render the checked three-section Parser Spec document.""" +_GRAMMAR_START = "" +_GRAMMAR_END = "" +_CONSTRAINTS_START = "" +_CONSTRAINTS_END = "" + + +def _replace_generated_section(document: str, start: str, end: str, content: str) -> str: + """Replace exactly one marked generated section in *document*.""" + if document.count(start) != 1 or document.count(end) != 1: + raise ValueError(f"parser spec must contain exactly one {start!r} and {end!r}") + if document.index(start) > document.index(end): + raise ValueError(f"parser spec marker {start!r} must precede {end!r}") + prefix, remainder = document.split(start, 1) + _old, suffix = remainder.split(end, 1) + return f"{prefix}{start}\n{content.rstrip()}\n{end}{suffix}" + + +def render_parser_document(document: str) -> str: + """Update only the marked generated sections of a Parser Spec document.""" generated = render_spec_content().removeprefix("# Parser Grammar and Constraints\n\n") - return f'''# TileFoundry Spec - Parser - -The Parser accepts authored Python functions and produces HIR or TIR through one typed API. - -## 1. Public API - -`@module` executes its Python class body and finalizes the collected Functions, child Modules, -and ordinary methods. `@func` produces an HIR Function; `@prim_func` produces a TIR PrimFunction. -`specialize` and `converter` register variants and weight converters on an existing HIR Function. - -```python -def parse_function( - fn: FunctionType, context: FuncParserContext -) -> hir.Function | tir.PrimFunction: ... -``` - -`FuncParserContext` carries the dialect, Function role, closure, topology scope, target, and -optional base/key for one parse. `FunctionRole` is `ROOT`, `VARIANT`, or `CONVERTER`. -`ParseError` is the single authored-source diagnostic type and includes source location and -recursive parse situation. These are the only public parser symbols. - -## 2. Syntax and Rules - -### 2.1 Syntax - - -{generated.split('| Owner | Situation | Rule | Statement | Source |', 1)[0].rstrip()} - - -### 2.2 Rules - - -| Owner | Situation | Rule | Statement | Source | -| --- | --- | --- | --- | --- | -{generated.split('| Owner | Situation | Rule | Statement | Source |', 1)[1].split('| --- | --- | --- | --- | --- |', 1)[1].lstrip()} - - -## 3. Implementation Overview - -| Component | Responsibility | -| --- | --- | -| Parser API and Context | Receives authored Functions and carries dialect, role, scope, and recursion inputs. | -| Executable Pattern Graph | Composes concrete AST elements into the Function root pattern. | -| Match and Construction | Matches recursively into `AstMatch`, then constructs owner values on return. | -| Ordered Rules | Validates and normalizes each owner value after construction. | -| Module Build | Lets Python execute the class body, records Functions, and finalizes the Module. | -| Pattern Visitor | Traverses the same graph to render this section's generated grammar and constraints. | - -```mermaid -classDiagram - ParserAPI --> FuncParserContext - ParserAPI --> FunctionPattern - AstPattern <|.. Element - Element o-- AstPattern - Element o-- AstRule - AstPattern --> AstMatch - PatternVisitor ..> AstPattern - ParserAPI ..> ModuleBuild -``` - -```mermaid -flowchart TD - API["parse_function(fn, context)"] --> AST["Extract FunctionDef AST"] - AST --> ROOT["FunctionPattern.match"] - ROOT --> TREE["AstMatch tree"] - TREE --> BACKWARD["construct children, then apply Rules"] - BACKWARD --> FUNCTION["HIR Function / TIR PrimFunction"] - FUNCTION --> MODULE{{"Module authoring context?"}} - MODULE -->|yes| FINALIZE["registration / finalization"] - MODULE -->|no| RETURN["return standalone result"] -``` - -Pattern combinators serve both runtime matching and Spec traversal. `AstMatch` separates syntax -matching from object construction, while each Rule reads the recursive context after its owner -value exists. Module class control flow remains Python execution; no Module AST grammar exists. -''' + grammar, rules = generated.split("| Owner | Situation | Rule | Statement | Source |", 1) + rule_rows = rules.split("| --- | --- | --- | --- | --- |", 1)[1].lstrip() + updated = _replace_generated_section(document, _GRAMMAR_START, _GRAMMAR_END, grammar) + return _replace_generated_section( + updated, + _CONSTRAINTS_START, + _CONSTRAINTS_END, + "| Owner | Situation | Rule | Statement | Source |\n" + "| --- | --- | --- | --- | --- |\n" + rule_rows, + ) def _parse_args(argv: list[str] | None) -> argparse.Namespace: @@ -266,13 +237,13 @@ def _parse_args(argv: list[str] | None) -> argparse.Namespace: def _main(argv: list[str] | None = None) -> int: args = _parse_args(argv) - generated = render_spec_content() if args.write is not None: - args.write.write_text(render_parser_document()) + document = args.write.read_text() + args.write.write_text(render_parser_document(document)) return 0 if args.check is not None: - expected = render_parser_document() actual = args.check.read_text() if args.check.exists() else "" + expected = render_parser_document(actual) if actual == expected: return 0 sys.stderr.writelines( @@ -284,7 +255,7 @@ def _main(argv: list[str] | None = None) -> int: ) ) return 1 - sys.stdout.write(generated) + sys.stdout.write(render_spec_content()) return 0 diff --git a/src/tilefoundry/target/amx/atoms.py b/src/tilefoundry/target/amx/atoms.py index 191e475f..092373f0 100644 --- a/src/tilefoundry/target/amx/atoms.py +++ b/src/tilefoundry/target/amx/atoms.py @@ -12,7 +12,7 @@ from dataclasses import dataclass from tilefoundry.ir.core import Call -from tilefoundry.ir.hir.nn.matmul import MatMul +from tilefoundry.ir.hir.nn.matmul import MatMul, matmul_axes from tilefoundry.ir.types import DType, TensorType from tilefoundry.schedule.facts import AtomFact from tilefoundry.target import Target @@ -202,8 +202,9 @@ def candidate_atoms(op: Call, target: Target | None = None) -> list[AtomFact]: ) lhs_type, rhs_type = op.args[0].type, op.args[1].type - m, k = lhs_type.shape[-2], lhs_type.shape[-1] - n = rhs_type.shape[-1] + a_m, a_k, b_n, _b_k = matmul_axes(op.target) + m, k = lhs_type.shape[a_m], lhs_type.shape[a_k] + n = rhs_type.shape[b_n] if not _static_positive(m, n, k) or not _operands_layout_ok(lhs_type, rhs_type): return [] diff --git a/src/tilefoundry/target/cuda/atoms.py b/src/tilefoundry/target/cuda/atoms.py index a94f29ed..d73da84c 100644 --- a/src/tilefoundry/target/cuda/atoms.py +++ b/src/tilefoundry/target/cuda/atoms.py @@ -9,7 +9,7 @@ from __future__ import annotations from tilefoundry.ir.core import Call -from tilefoundry.ir.hir.nn.matmul import MatMul +from tilefoundry.ir.hir.nn.matmul import MatMul, matmul_axes from tilefoundry.ir.tir.cuda.nn.mma import SM80_16x8x16_F32BF16BF16F32_TN, make_atom from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaOpSpec from tilefoundry.ir.types import DType, TensorType, tensor_bytes @@ -129,8 +129,9 @@ def candidate_atoms(op: Call, target: Target | None = None) -> list[AtomFact]: ) lhs_type, rhs_type = op.args[0].type, op.args[1].type - m, k = lhs_type.shape[-2], lhs_type.shape[-1] - n = rhs_type.shape[-1] + a_m, a_k, b_n, _b_k = matmul_axes(op.target) + m, k = lhs_type.shape[a_m], lhs_type.shape[a_k] + n = rhs_type.shape[b_n] if not _static_positive(m, n, k) or not _operands_layout_ok(lhs_type, rhs_type): return [] diff --git a/src/tilefoundry/visitor_registry/op_cost.py b/src/tilefoundry/visitor_registry/op_cost.py index b9ea2418..80d6c339 100644 --- a/src/tilefoundry/visitor_registry/op_cost.py +++ b/src/tilefoundry/visitor_registry/op_cost.py @@ -21,7 +21,7 @@ from tilefoundry.ir.hir.nn.conv2d import Conv2D from tilefoundry.ir.hir.nn.gelu import Gelu from tilefoundry.ir.hir.nn.layer_norm import LayerNorm -from tilefoundry.ir.hir.nn.matmul import MatMul +from tilefoundry.ir.hir.nn.matmul import MatMul, matmul_axes from tilefoundry.ir.hir.nn.relu import ReLU from tilefoundry.ir.hir.nn.rms_norm import RMSNorm from tilefoundry.ir.hir.nn.rope import RoPE @@ -58,6 +58,7 @@ from tilefoundry.ir.types import DType, IntegerDType, TensorType, Type, numel, tensor_bytes from tilefoundry.ir.types.shard import ShardLayout from tilefoundry.ir.types.shard.shard_layout import layout_axis_to_tensor_axis +from tilefoundry.visitor_registry.access_relation import logical_axes_of from .contexts import Cost, CostContext, TrafficBytes from .registries import register_cost_evaluator @@ -115,21 +116,24 @@ def _serviced(call: Call, ctx: CostContext, kind: str) -> Cost: @register_cost_evaluator(MatMul) def _matmul(call: Call, ctx: CostContext) -> Cost: - """One multiply and one add per multiply-accumulate, over every batch. - - The batch comes from the output rather than from the left operand. Either side - may be the one that is broadcast: a block of a weight matrix multiplied by one - token has its batch on the right, and reading the left gave a batch of one -- - the whole block loop's arithmetic charged as a single tile's. The output's batch - is what the call produced, and every batch of it was computed. - """ + """One multiply and one add per multiply-accumulate: 2 * batch * m * k * n.""" lhs, rhs = _input_types(call, ctx) output = _output_type(call, ctx) if not all(isinstance(type, TensorType) for type in (lhs, rhs, output)): raise ValueError("MatMul cost requires tensor inputs and output") - m, k, n = lhs.shape[-2], lhs.shape[-1], rhs.shape[-1] - batch = math.prod(output.shape[:-2]) - flops = 2 * batch * m * k * n + logical_lhs = ctx.type_of(call.args[0]) + if not isinstance(logical_lhs, TensorType): + raise ValueError("MatMul cost requires a tensor lhs") + _a_m, a_k, _b_n, _b_k = matmul_axes(call.target) + k_axis = a_k % len(logical_lhs.shape) + k = math.prod( + extent + for extent, logical_axis in zip( + lhs.shape, logical_axes_of(lhs, logical_lhs) + ) + if logical_axis == k_axis + ) + flops = 2 * numel(output) * k return Cost({lhs.dtype: flops}, _traffic((lhs, rhs), output)) diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index c5ce6e09..9220edb2 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -158,6 +158,83 @@ def test_roofline_uses_exact_integer_ceiling_above_float_precision() -> None: assert bound.ideal_ns == 1_492_537_313_434 +_SPLIT_GRID = 128 +_SPLIT_HIDDEN = 2048 +_SPLIT_OUT = 12288 +_SPLIT_PER = _SPLIT_OUT // _SPLIT_GRID +_SPLIT_BLOCK = 128 + + +@module(entry="last_axis", target=_H200, topologies=(Topology("cta", _SPLIT_GRID),)) +class _SplitLastAxis: + """The weight's N on the mesh, so the result's last axis is the split one.""" + + @func + def last_axis( + x: Tensor[(1, _SPLIT_BLOCK, _SPLIT_HIDDEN), "bf16"], + w: ConstTensor[(_SPLIT_HIDDEN, _SPLIT_OUT), "bf16"], + ) -> Tensor[(1, _SPLIT_BLOCK, _SPLIT_OUT), "bf16"]: + with Mesh(("cta",), layout=(_SPLIT_GRID,), names=("unit",)) as mesh: + rows = tf.reshard( + x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem" + ) + strip = tf.reshard( + w[0:_SPLIT_BLOCK, :], (_SPLIT_BLOCK, _SPLIT_OUT @ mesh.unit), "smem" + ) + return tf.matmul(rows, strip) + + +@module(entry="strip_major", target=_H200, topologies=(Topology("cta", _SPLIT_GRID),)) +class _SplitStripMajor: + """The same gemm with the split axis leading, inside the matmul's batch.""" + + @func + def strip_major( + x: Tensor[(1, _SPLIT_BLOCK, _SPLIT_HIDDEN), "bf16"], + w: ConstTensor[(_SPLIT_GRID, _SPLIT_HIDDEN, _SPLIT_PER), "bf16"], + ) -> Tensor[(_SPLIT_GRID, _SPLIT_BLOCK, _SPLIT_PER), "bf16"]: + with Mesh(("cta",), layout=(_SPLIT_GRID,), names=("unit",)) as mesh: + rows = tf.reshard( + x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem" + ) + strip = tf.reshard( + w[:, 0:_SPLIT_BLOCK, :], + (_SPLIT_GRID @ mesh.unit, _SPLIT_BLOCK, _SPLIT_PER), + "smem", + ) + return tf.matmul(rows, strip) + + +def test_a_matmul_counts_its_rows_once_whichever_axis_the_mesh_split() -> None: + """One gemm, two layouts of the weight, one per-unit answer. + + Split the result's last axis and the axis sharding adds lands where a batch + is read from, so counting ``batch * m * n`` charges the rows twice and never + divides by the grid. Counting the result's elements does not care where the + axis went, and the two spellings agree on the work and the predicted time. + """ + per_layout = {} + for owner, name in ((_SplitLastAxis, "last_axis"), (_SplitStripMajor, "strip_major")): + function = next(item for item in owner.functions if item.name == name) + report = analyze(owner, function, analysis="performance") + product = next( + expr + for expr in postorder(report.function.body) + if isinstance(expr, Call) and type(expr.target).__name__ == "MatMul" + ) + cost = get_metadata(product, ComputeCostMetadata) + summary = get_metadata(report.function, PerformanceSummaryMetadata) + per_layout[name] = ( + dict(cost.flops)["bf16"], + dict(cost.flops_per_unit)["bf16"], + summary.timeline.end_ns, + ) + + for name, (whole, unit, _predicted) in per_layout.items(): + assert whole == unit * _SPLIT_GRID, f"{name} did not divide by the grid" + assert per_layout["last_axis"] == per_layout["strip_major"] + + def test_a_program_whose_buffers_have_nowhere_to_sit_is_refused() -> None: """Placing the buffers is what makes the rest of the answer worth having. diff --git a/tests/models/deepseek_v4_flash/model.py b/tests/models/deepseek_v4_flash/model.py index 37d27d5a..25153104 100644 --- a/tests/models/deepseek_v4_flash/model.py +++ b/tests/models/deepseek_v4_flash/model.py @@ -279,9 +279,9 @@ def mla_kv_update( kv_o1 = tf.cast(kv_o1_f32, dtype="bf16") kv_o0 = tf.reshape(kv_o0, new_shape=(1, 1, 1, config.rope_half, 1)) kv_o1 = tf.reshape(kv_o1, new_shape=(1, 1, 1, config.rope_half, 1)) - kv_interleaved = tf.concat(kv_o0, kv_o1, axis=-1) + kv_interleaved = tf.concat([kv_o0, kv_o1], axis=-1) kv_rope_out = tf.reshape(kv_interleaved, new_shape=(1, 1, 1, config.rope_dim)) - return tf.concat(kv_nope_q, kv_rope_out, axis=-1) + return tf.concat([kv_nope_q, kv_rope_out], axis=-1) @mla_kv_update.converter("w_kv") def _( @@ -336,9 +336,9 @@ def mla_attend( q_o1 = tf.cast(q_o1_f32, dtype="bf16") q_o0 = tf.reshape(q_o0, new_shape=(1, 1, config.n_heads, config.rope_half, 1)) q_o1 = tf.reshape(q_o1, new_shape=(1, 1, config.n_heads, config.rope_half, 1)) - q_interleaved = tf.concat(q_o0, q_o1, axis=-1) + q_interleaved = tf.concat([q_o0, q_o1], axis=-1) q_rope_out = tf.reshape(q_interleaved, new_shape=(1, 1, config.n_heads, config.rope_dim)) - q_final = tf.concat(q_nope, q_rope_out, axis=-1) + q_final = tf.concat([q_nope, q_rope_out], axis=-1) # MQA repeat_interleave to n_heads, for the cache and the new token # alike; the KV latent serves as both K and V (no separate V projection). @@ -395,11 +395,11 @@ def mla_attend( ctx_o1 = tf.cast(ctx_o1_f32, dtype="bf16") ctx_o0 = tf.reshape(ctx_o0, new_shape=(1, 1, config.n_heads, config.rope_half, 1)) ctx_o1 = tf.reshape(ctx_o1, new_shape=(1, 1, config.n_heads, config.rope_half, 1)) - ctx_interleaved = tf.concat(ctx_o0, ctx_o1, axis=-1) + ctx_interleaved = tf.concat([ctx_o0, ctx_o1], axis=-1) ctx_rope_out = tf.reshape( ctx_interleaved, new_shape=(1, 1, config.n_heads, config.rope_dim), ) - ctx_final = tf.concat(ctx_nope, ctx_rope_out, axis=-1) + ctx_final = tf.concat([ctx_nope, ctx_rope_out], axis=-1) o_flat = tf.reshape(ctx_final, new_shape=(1, 1, config.q_proj)) # Grouped low-rank O projection: o_flat's last axis is a contiguous diff --git a/tests/models/kimi_linear_48b_a3b/model.py b/tests/models/kimi_linear_48b_a3b/model.py index b157dd99..065e49f5 100644 --- a/tests/models/kimi_linear_48b_a3b/model.py +++ b/tests/models/kimi_linear_48b_a3b/model.py @@ -342,8 +342,8 @@ def mla_attention( k_rot_h = tf.repeat_interleave(k_rot, repeats=_H, axis=2) q_rot_r, _kr = tf.rope(q_rot, q_rot, cos_cache, sin_cache, pos_ids) - q_full = tf.concat(q_pass, q_rot_r, axis=-1) - k_new = tf.concat(k_nope, k_rot_h, axis=-1) + q_full = tf.concat([q_pass, q_rot_r], axis=-1) + k_new = tf.concat([k_nope, k_rot_h], axis=-1) # Online softmax over two differently shaped score groups: the cache and # the token itself. No mask -- one query at the end of the context may @@ -404,7 +404,7 @@ def short_conv( # the convolution is a weighted sum over the window's time axis rather # than a sliding op. Returns the activation and the window to store next, # which is this window with its oldest position dropped. - window = tf.concat(conv_state, x, axis=1) + window = tf.concat([conv_state, x], axis=1) acc = tf.reduce( window * tf.reshape( diff --git a/tests/models/minicpm3_4b/model.py b/tests/models/minicpm3_4b/model.py index 2ee4f290..6eb649e0 100644 --- a/tests/models/minicpm3_4b/model.py +++ b/tests/models/minicpm3_4b/model.py @@ -257,8 +257,8 @@ def mla_attention( k_rope_b = tf.repeat_interleave(k_rope_e, repeats=_H, axis=2) # Step 6: reassemble nope + rope, each back in its original slot. - query = tf.concat(q_nope, q_rope_e, axis=-1) - k_new = tf.concat(k_nope, k_rope_b, axis=-1) + query = tf.concat([q_nope, q_rope_e], axis=-1) + k_new = tf.concat([k_nope, k_rope_b], axis=-1) # Step 7: attend the cache and the token itself, then project out. q_s = query * scale diff --git a/tests/models/qwen3_5_35b_a3b/model.py b/tests/models/qwen3_5_35b_a3b/model.py index a0da2bf4..e917fa29 100644 --- a/tests/models/qwen3_5_35b_a3b/model.py +++ b/tests/models/qwen3_5_35b_a3b/model.py @@ -130,7 +130,7 @@ def conv_step( # closes on this token, so the whole convolution is one multiply against # the kernel and one reduction over it. Channels do not mix -- that is # what depthwise means here, and it is why no matmul appears. - window = tf.concat(conv_state, entry, axis=2) + window = tf.concat([conv_state, entry], axis=2) weighted = window * tf.reshape(conv_w, new_shape=(1, _CONV, _KERNEL)) summed = tf.reduce(weighted, axes=(-1,), keepdim=False, kind="sum") return tf.silu(summed) @@ -245,7 +245,7 @@ def partial_rope( rot = x[:, :, :, :_ROT] tail = x[:, :, :, _ROT:_D] turned, _ = tf.rope(rot, rot, cos_cache, sin_cache, pos_ids) - return tf.concat(turned, tail, axis=-1) + return tf.concat([turned, tail], axis=-1) @func def partial_rope_kv( @@ -259,7 +259,7 @@ def partial_rope_kv( rot = x[:, :, :, :_ROT] tail = x[:, :, :, _ROT:_D] turned, _ = tf.rope(rot, rot, cos_cache, sin_cache, pos_ids) - return tf.concat(turned, tail, axis=-1) + return tf.concat([turned, tail], axis=-1) @func def full_attention( diff --git a/tests/ops/test_matmul.py b/tests/ops/test_matmul.py index b1d54095..2c0330f2 100644 --- a/tests/ops/test_matmul.py +++ b/tests/ops/test_matmul.py @@ -12,7 +12,9 @@ from __future__ import annotations import pytest +import torch +from tests.evaluator.eval_utils import EvalCase, run_eval_case from tests.ops.cost_utils import CostCase, run_cost_case from tests.ops.typeinfer_utils import ( ExpectedError, @@ -75,6 +77,50 @@ level="cta", topologies=(_CTA,), ), + CostCase( + name="km_nk_projection_reads_the_mapped_k_axis", + op=MatMul(a_layout="KM", b_layout="NK"), + inputs=( + make_shard_tensor_type( + (10, 4), + mesh=_CTA_MESH, + attrs=(Split(axis=0),), + dtype=DType.f32, + ), + make_shard_tensor_type( + (3, 10), + mesh=_CTA_MESH, + attrs=(Split(axis=1),), + dtype=DType.f32, + ), + ), + flops={DType.f32: 2 * 4 * 2 * 3}, + traffic=( + TrafficBytes(read=4 * 2 * 4), + TrafficBytes(read=3 * 2 * 4), + TrafficBytes(write=4 * 3 * 4), + ), + level="cta", + topologies=(_CTA,), + ), +] + + +LAYOUT_CASES = [ + pytest.param(MatMul(), (2, 3), (3, 4), id="mk_kn_default"), + pytest.param(MatMul(b_layout="NK"), (2, 3), (4, 3), id="mk_nk"), + pytest.param( + MatMul(a_layout="KM"), + (3, 2), + (3, 4), + id="km_kn", + ), + pytest.param( + MatMul(a_layout="KM", b_layout="NK"), + (3, 2), + (4, 3), + id="km_nk", + ), ] @@ -83,6 +129,37 @@ def test_matmul_cost(case): run_cost_case(case) +@pytest.mark.parametrize(("op", "lhs_shape", "rhs_shape"), LAYOUT_CASES) +def test_matmul_layouts_share_shape_and_cost(op, lhs_shape, rhs_shape): + lhs = make_tensor_type(lhs_shape, DType.f32) + rhs = make_tensor_type(rhs_shape, DType.f32) + + assert infer_call(op, lhs, rhs) == make_tensor_type((2, 4), DType.f32) + run_cost_case( + CostCase( + name=f"{op.a_layout}_{op.b_layout}", + op=op, + inputs=(lhs, rhs), + flops={DType.f32: 2 * 2 * 3 * 4}, + traffic=( + TrafficBytes(read=2 * 3 * 4), + TrafficBytes(read=4 * 3 * 4), + TrafficBytes(write=2 * 4 * 4), + ), + ) + ) + + +@pytest.mark.parametrize(("op", "lhs_shape", "rhs_shape"), LAYOUT_CASES) +def test_matmul_layouts_evaluate(op, lhs_shape, rhs_shape): + lhs = torch.arange(6, dtype=torch.float32).reshape(lhs_shape) + rhs = torch.arange(12, dtype=torch.float32).reshape(rhs_shape) + logical_lhs = lhs.transpose(-1, -2) if op.a_layout == "KM" else lhs + logical_rhs = rhs.transpose(-1, -2) if op.b_layout == "NK" else rhs + + run_eval_case(EvalCase("matmul_layout", op, (lhs, rhs), logical_lhs @ logical_rhs)) + + def _sharded(shape, attrs): return make_shard_tensor_type(shape, mesh=_M, attrs=attrs, dtype=DType.bf16) @@ -108,6 +185,25 @@ def test_matmul_typeinfer(case): run_typeinfer_case(case) +@pytest.mark.parametrize( + "op", + (MatMul(a_layout="bad"), MatMul(b_layout="bad")), + ids=("invalid_a_layout", "invalid_b_layout"), +) +def test_matmul_rejects_invalid_layouts(op): + run_typeinfer_case( + TypeInferCase( + name="invalid_layout", + op=op, + inputs=( + make_tensor_type((2, 3), DType.f32), + make_tensor_type((4, 3), DType.f32), + ), + expected=ExpectedError(match="layout must be"), + ) + ) + + def test_lhs_splits_k_rhs_unsplit_is_invalid(): lhs = _sharded((16, 8), (Split(axis=1),)) @@ -174,6 +270,39 @@ def test_a_sharded_operand_carries_to_the_output(lhs, rhs, shape, attrs): assert out.layout.attrs == attrs +@pytest.mark.parametrize( + ("op", "lhs", "rhs", "attrs"), + ( + pytest.param( + MatMul(b_layout="NK"), + make_tensor_type((16, 8), DType.bf16), + _sharded((32, 8), (Split(axis=0),)), + (Split(axis=1),), + id="nk_n_split", + ), + pytest.param( + MatMul(a_layout="KM", b_layout="NK"), + _sharded((8, 16), (Split(axis=1),)), + make_tensor_type((32, 8), DType.bf16), + (Split(axis=0),), + id="km_m_split", + ), + pytest.param( + MatMul(a_layout="KM", b_layout="NK"), + _sharded((8, 16), (Split(axis=0),)), + _sharded((32, 8), (Split(axis=1),)), + (Partial(reduction="sum"),), + id="km_nk_k_split", + ), + ), +) +def test_layout_literals_map_sharding_through_access_relations(op, lhs, rhs, attrs): + out = infer_call(op, lhs, rhs) + + assert out.shape == (16, 32) + assert out.layout.attrs == attrs + + def test_double_partial_same_mesh_axis_errors(): lhs = _sharded((16, 8), (Partial("sum"),)) rhs = _sharded((8, 32), (Partial("sum"),)) diff --git a/tests/parser/test_calls.py b/tests/parser/test_calls.py new file mode 100644 index 00000000..13224971 --- /dev/null +++ b/tests/parser/test_calls.py @@ -0,0 +1,207 @@ +"""Parser contracts for explicit variadic input sequences.""" + +from __future__ import annotations + +from typing import get_args, get_origin + +import pytest + +from tilefoundry import func +from tilefoundry.dsl import Tensor, tf +from tilefoundry.ir.core import Call, Constant, Tuple, VerifyError +from tilefoundry.ir.core.pattern import Tensor as TensorPattern +from tilefoundry.ir.hir.nn.matmul import MatMul +from tilefoundry.ir.hir.tensor.concat import Concat +from tilefoundry.ir.hir.tensor.reshape import Reshape +from tilefoundry.ir.hir.tensor.slice import Slice +from tilefoundry.ir.hir.tensor.stack import Stack +from tilefoundry.parser import ParseError + + +def test_matmul_layout_literals_are_parser_checked() -> None: + assert get_args(MatMul.a_layout.annotation) == ("MK", "KM") + assert get_args(MatMul.b_layout.annotation) == ("NK", "KN") + + @func + def default_layout( + a: Tensor[(2, 3), "f32"], b: Tensor[(3, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + return tf.matmul(a, b) + + assert isinstance(default_layout.body, Call) + assert isinstance(default_layout.body.target, MatMul) + assert (default_layout.body.target.a_layout, default_layout.body.target.b_layout) == ( + "MK", + "KN", + ) + + with pytest.raises(ParseError, match="b_layout must be one of 'NK', 'KN'"): + + @func + def invalid_layout( + a: Tensor[(2, 3), "f32"], b: Tensor[(3, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + return tf.matmul(a, b, b_layout="bad") + + +def test_variadic_list_and_tuple_literals_flatten_to_call_args() -> None: + for annotation in (Concat.inputs.annotation, Stack.inputs.annotation): + assert get_origin(annotation) is tuple + assert get_args(annotation) == (TensorPattern,) + + @func + def from_list( + a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + return tf.concat([a, b], axis=0) + + @func + def from_tuple( + a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + return tf.concat((a, b), axis=0) + + for function in (from_list, from_tuple): + assert isinstance(function.body, Call) + assert isinstance(function.body.target, Concat) + assert len(function.body.args) == 2 + + +def test_stack_uses_the_same_variadic_list_contract() -> None: + @func + def stack_rows( + a: Tensor[(4,), "f32"], b: Tensor[(4,), "f32"] + ) -> Tensor[(2, 4), "f32"]: + return tf.stack([a, b], axis=0) + + assert isinstance(stack_rows.body, Call) + assert isinstance(stack_rows.body.target, Stack) + assert len(stack_rows.body.args) == 2 + + +def test_an_empty_variadic_list_reaches_the_operation_verifier() -> None: + with pytest.raises(VerifyError, match="Concat requires at least one input"): + + @func + def empty(x: Tensor[(1, 4), "f32"]) -> Tensor[(1, 4), "f32"]: + return tf.concat([], axis=0) + + +def test_a_static_range_list_comprehension_expands_in_source_order() -> None: + @func + def rows(x: Tensor[(2, 4), "f32"]) -> Tensor[(2, 4), "f32"]: + return tf.stack([x[index, :] for index in range(2)], axis=0) + + assert isinstance(rows.body, Call) + assert isinstance(rows.body.target, Stack) + assert len(rows.body.args) == 2 + indices = [] + for row in rows.body.args: + assert isinstance(row, Call) and isinstance(row.target, Reshape) + sliced = row.args[0] + assert isinstance(sliced, Call) and isinstance(sliced.target, Slice) + starts = sliced.args[1] + assert isinstance(starts, Tuple) + start = starts.elements[0] + assert isinstance(start, Constant) + indices.append(start.value) + assert indices == [0, 1] + + +def test_variadic_direct_positional_inputs_are_rejected() -> None: + with pytest.raises(ParseError, match="require exactly one list, tuple"): + + @func + def direct( + a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + return tf.concat(a, b, axis=0) + + with pytest.raises(ParseError, match="require exactly one list, tuple"): + + @func + def positional_attribute( + a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + return tf.concat([a, b], 0) + + +def test_variadic_generator_and_starred_inputs_name_the_unsupported_form() -> None: + with pytest.raises(ParseError, match="do not support generator expressions"): + + @func + def generator(x: Tensor[(2, 4), "f32"]) -> Tensor[(2, 4), "f32"]: + return tf.stack((x[index] for index in range(2)), axis=0) + + with pytest.raises(ParseError, match="do not support starred expansion"): + + @func + def starred( + a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + return tf.concat([a, *[b]], axis=0) + + with pytest.raises(ParseError, match="require an explicit list, tuple"): + + @func + def expanded( + a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + parts = (a, b) + return tf.concat(*parts, axis=0) + + +def test_variadic_sequence_names_are_not_implicitly_expanded() -> None: + with pytest.raises(ParseError, match="require an explicit list, tuple"): + + @func + def named( + a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"] + ) -> Tensor[(2, 4), "f32"]: + parts = (a, b) + return tf.concat(parts, axis=0) + + +@pytest.mark.parametrize( + "program, message", + [ + ("multiple", "exactly one generator"), + ("target", "simple Name target"), + ("filter", "does not support filters"), + ("iterator", "range with 1 to 3 static integer arguments"), + ("bound", "range arguments must be static integers"), + ], +) +def test_unsupported_list_comprehension_shapes_are_named(program: str, message: str) -> None: + with pytest.raises(ParseError, match=message): + if program == "multiple": + + @func + def rejected(x: Tensor[(1, 4), "f32"]) -> Tensor[(2, 4), "f32"]: + return tf.concat( + [x for first in range(1) for second in range(2)], axis=0 + ) + + elif program == "target": + + @func + def rejected(x: Tensor[(1, 4), "f32"]) -> Tensor[(1, 4), "f32"]: + return tf.concat([x for first, second in range(1)], axis=0) + + elif program == "filter": + + @func + def rejected(x: Tensor[(1, 4), "f32"]) -> Tensor[(1, 4), "f32"]: + return tf.concat([x for index in range(1) if index], axis=0) + + elif program == "iterator": + + @func + def rejected(x: Tensor[(1, 4), "f32"]) -> Tensor[(1, 4), "f32"]: + return tf.concat([x for index in (0,)], axis=0) + + else: + + @func + def rejected(x: Tensor[(1, 4), "f32"]) -> Tensor[(1, 4), "f32"]: + return tf.concat([x for index in range(x)], axis=0) diff --git a/tests/parser/test_dimensions.py b/tests/parser/test_dimensions.py new file mode 100644 index 00000000..36f371fc --- /dev/null +++ b/tests/parser/test_dimensions.py @@ -0,0 +1,33 @@ +"""Parser normalization of authored scalar expressions used as dimensions.""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, Topology, tf +from tilefoundry.ir.core import Call +from tilefoundry.ir.hir.sharding.reshard import Reshard +from tilefoundry.target import CudaTarget + + +def test_body_local_integer_arithmetic_is_normalized_before_layout_construction() -> None: + @module( + entry="f", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 128),), + ) + class Model: + @func + def f( + x: Tensor[(1, 16, 8192), "f32"], + ) -> Tensor[(1, 16, 8192), "f32"]: + with Mesh(("cta",), layout=(128,), names=("unit",)) as mesh: + width = 4096 + 4096 + return tf.reshard( + x, + (1, 16, width @ mesh.unit), + "smem", + ) + + call = Model.functions[0].body + assert isinstance(call, Call) and isinstance(call.target, Reshard) + assert call.target.layout.layout.shape == (1, 16, 128, 64) diff --git a/tests/parser/test_functions.py b/tests/parser/test_functions.py new file mode 100644 index 00000000..a1d2b246 --- /dev/null +++ b/tests/parser/test_functions.py @@ -0,0 +1,30 @@ +"""Parser ownership of function-only docstring semantics.""" + +from __future__ import annotations + +import pytest + +from tilefoundry import func +from tilefoundry.dsl import Mesh, Tensor, Topology +from tilefoundry.inspection import as_script +from tilefoundry.parser import ParseError +from tilefoundry.target import CudaTarget + + +def test_a_function_may_have_a_leading_docstring() -> None: + @func + def documented(x: Tensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: + """This is documentation, not an HIR statement.""" + return x + + assert "This is documentation" not in as_script(documented) + + +def test_a_nested_block_does_not_gain_function_docstring_semantics() -> None: + with pytest.raises(ParseError): + + @func(target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", 1),)) + def nested(x: Tensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: + with Mesh(("cta",), layout=(1,), names=("unit",)) as _mesh: + """A string in a with body remains an ordinary statement.""" + return x diff --git a/tests/parser/test_slices.py b/tests/parser/test_slices.py new file mode 100644 index 00000000..d22df681 --- /dev/null +++ b/tests/parser/test_slices.py @@ -0,0 +1,36 @@ +"""Parser diagnostics for authored tile-window slices.""" + +from __future__ import annotations + +import pytest + +from tilefoundry import func +from tilefoundry.dsl import Mesh, Tensor, Topology, tf +from tilefoundry.parser import ParseError +from tilefoundry.target import CudaTarget + + +def test_a_tile_window_bound_names_the_authored_fix() -> None: + with pytest.raises( + ParseError, + match=( + r"tile loop variable is already a window.*" + r"x\[:, t, :\].*base = t \+ 0" + ), + ): + + @func(target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", 1),)) + def stage( + x: Tensor[(1, 128, 64), "f32"], + out: Tensor[(1, 128, 64), "f32"], + ) -> Tensor[(1, 128, 64), "f32"]: + with Mesh(("cta",), layout=(1,), names=("unit",)) as mesh: + acc = out + for t in tile(128, 128): + window = tf.reshard( + x[:, t:t + 128, :], + (1, 128 @ mesh.unit, 64), + "smem", + ) + acc = tf.insert_slice(acc, window, (0, t, 0)) + return acc