diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 2b2720fe..c1935c7b 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -147,11 +147,12 @@ compilation target nests as `ir/{dialect}/{target}/{category}/.py`; target-neutral abstractions stay at `ir/{dialect}/{category}/`. For example the whole MMA surface is target-owned — the `Mma` op, the `MmaOpSpec` / `MmaAtom` descriptors, the CUDA SM80 instruction spec, and its -fragment layouts all live under `ir/tir/cuda/nn/` (`mma.py` + `mma_atom.py`), -and the HIR per-shape `Mma_SM80_*` / `Wgmma_SM90_*` ops under -`ir/hir/cuda/nn/mma.py` — because an MMA instruction fixes a concrete hardware -op. (`codegen/` and `runtime/` are **target-first** instead — their primary -axis is the target — so each tree is organized by its own primary axis.) +fragment layouts all live under `ir/tir/cuda/nn/` (`mma.py` + `mma_atom.py`). +The backend-bound construction stays in TIR: HIR is the checking reference +side, and carrying the instruction name in that reference would make two GPU +targets require different HIR references. (`codegen/` and `runtime/` are +**target-first** instead — their primary axis is the target — so each tree is +organized by its own primary axis.) **Rule 2 — one (node, target) codegen = one file.** Each handler lives at `codegen//tir//.py`. Stmt emitters, diff --git a/docs/spec/hir.md b/docs/spec/hir.md index 98821842..8e223ad9 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -507,9 +507,10 @@ class Binary(Op): - constraints: - Values follow torch pointwise semantics; dtypes do not promote. Both operands MUST already carry the same `dtype`, and typeinfer MUST reject a mismatch. A - Python float scalar is given the other operand's float dtype by the authoring - surface, before it is an operand at all ([parser §2.1](./parser.md#21-syntax)); a Python - integer is not. + Python literal is an ordinary operand of the dtype it is written with — `f32` + for a float, `i64` for an integer — and the authoring surface MUST NOT adapt + it to its peer. The rejection MUST name `Cast` as the remedy, so the dtype a + value carries is the one the author wrote. - The elementwise `min` / `max` kinds are also surfaced as `minimum` / `maximum`. - Equal plain layouts, or one plain layout paired with `layout=None`, pass through only when that layout describes the broadcast result. Otherwise two @@ -1405,91 +1406,6 @@ class RoPE(Op): - A non-`sum` Partial, multiple value-carrying Partials, or a Partial on a secondary cache/index input MUST be rejected with a `Reshard` remedy. -##### CUDA matrix multiply-accumulate family - -```python -class Mma(Op): - """Provide the marker base for HIR CUDA matrix operations.""" - - -class Mma_SM80_16x8x16(Mma): - """Produce an SM80 16-by-8 accumulator fragment. - - Attributes: - a: input; Left 16-by-16 fragment. - b: input; Right 16-by-8 fragment. - dtype_a: attribute; Left operand dtype. - dtype_b: attribute; Right operand dtype. - dtype_acc: attribute; Accumulator dtype. - a_layout: attribute; Left matrix orientation. - b_layout: attribute; Right matrix orientation. - """ - - a: Tensor - b: Tensor - dtype_a: DType - dtype_b: DType - dtype_acc: DType - a_layout: str = "T" - b_layout: str = "N" - - -class Wgmma_SM90_64x128x16(Mma): - """Produce an SM90 64-by-128 accumulator fragment. - - Attributes: - a: input; Left fragment. - b: input; Right fragment. - dtype_a: attribute; Left operand dtype. - dtype_b: attribute; Right operand dtype. - dtype_acc: attribute; Accumulator dtype. - a_layout: attribute; Left matrix orientation. - b_layout: attribute; Right matrix orientation. - """ - - a: Tensor - b: Tensor - dtype_a: DType - dtype_b: DType - dtype_acc: DType - a_layout: str = "T" - b_layout: str = "N" -``` - -- constraints: - - `Mma` is a marker base used for family dispatch and has no independent - callable parameter surface. - - HIR MMA is the two-input value `a @ b`; it has no accumulator input. The - evaluator converts the rounded operands to `dtype_acc` before multiplying. - The independent handwritten TIR MMA surface has an in-place accumulator, - but there is no HIR-to-TIR compile route between these contracts. - - `Mma_SM80_16x8x16` requires `a.shape == (16, 16)` and - `b.shape == (16, 8)`, and returns `(16, 8)`. `Wgmma_SM90_64x128x16` - requires `(64, 16)` and `(16, 128)`, and returns `(64, 128)`. - - Each operand dtype MUST equal its declared `dtype_a` / `dtype_b`. Supported - `(dtype_a, dtype_b, dtype_acc)` combinations are `(f16, f16, f16)`, - `(bf16, bf16, bf16)`, `(f16, f16, f32)`, `(bf16, bf16, f32)`, and - `(f32, f32, f32)`. Each orientation MUST be `N` or `T`; orientation names - the hardware encoding and does not transpose either logical HIR input. - - Plain logical input Types remain valid for evaluation and Analyze. Their - result has a fresh row-major logical layout when either operand has a plain - layout. Fully `Broadcast` `ShardLayout` inputs carry no real ownership and - therefore pin no result mesh. - - Only the known SM80 BF16/BF16/F32 TN A/B fragment pair in RMEM derives the - known C fragment. Its A/B layouts, `Split` bindings, thread topology, and - mesh shape MUST match. Mesh coordinate names do not participate in - identity; the derived C layout uses `a`'s physically compatible mesh. Any - other genuine SM80 shard claim or fragment in non-RMEM storage MUST fail - with an explicit `Reshard` / materialize-to-RMEM remedy. - - WGMMA has no representable fragment contract yet. A genuine WGMMA shard - claim MUST fail with an explicit `Reshard` remedy, while plain logical - WGMMA remains evaluable and analyzable. - - Cost is `2*M*N*K` in the operand dtype with exactly three traffic slots: - A read, B read, and result write. Byte residency and topology projection - come only from the selected operand/result Types. - - `HirToTirPass` MUST reject both concrete HIR MMA Ops by name before - emitting TIR; see [passes §7.1](./passes.md#71-hirtotirpass). - #### `ir/hir/sharding/` `ShardLayout` and `Mesh` are type-system constructs, not Expr inputs diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 5fedfe00..49784319 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -100,14 +100,14 @@ type-annotation ::= tensor | scalar-type signature ::= (name ':' type-annotation (',' name ':' type-annotation)*)? return-type ::= type-annotation -loop-iterator ::= 'tile' - | 'range' +loop-iterator ::= 'tile' '(' expression ',' expression ')' + | 'range' '(' (expression | expression ',' expression | expression ',' + expression ',' expression) ')' 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' loop-iterator '(' (expression | name '=' expression) - (',' (expression | name '=' expression))* ')' ':' loop-carry +loop-header ::= 'for' identifier 'in' loop-iterator ':' loop-carry loop-body ::= (statement (newline statement)*)? for ::= 'for' name 'in' expression ':' loop-body mesh-context ::= ('Mesh' | primary '.' identifier) '(' (expression | ('layout' | 'names') @@ -132,9 +132,11 @@ subscript-index ::= '(' ((index-slice | index-endpoint) (',' (index-slice | index-slice | index-endpoint subscript-expression ::= runtime-expression '[' subscript-index ']' -binary-expression ::= runtime-expression binary-op runtime-expression - | runtime-expression comparison-op runtime-expression - | runtime-expression boolean-op runtime-expression +matmul-expression ::= runtime-expression '@' runtime-expression +binary-expression ::= runtime-expression ('+' | '-' | '*' | '/' | '//' | '%') runtime-expression + | runtime-expression ('==' | '!=' | '<' | '<=' | '>' | '>=') + runtime-expression + | runtime-expression ('and' | 'or') runtime-expression unary-expression ::= unary-op runtime-expression name ::= identifier constant ::= boolean-literal @@ -144,6 +146,7 @@ tuple-expression ::= '(' (runtime-expression (',' runtime-expression)*)? ') runtime-expression ::= op-call | launch | subscript-expression + | matmul-expression | binary-expression | unary-expression | mesh-coordinate @@ -173,39 +176,27 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | Owner | Situation | Rule | Statement | Source | | --- | --- | --- | --- | --- | -| 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 | +| binary_expression, matmul_expression, op_call, unary_expression | expression, slice_endpoint, subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression, matmul_expression, op_call, unary_expression | expression, slice_endpoint, subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression, matmul_expression, op_call, unary_expression | expression, slice_endpoint, subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | | dim_expr | dim_expr, layout_extent, layout_shape, tensor_dim_expr, tensor_optional_slot, tensor_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | | 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 | +| explicit_layout, layout, placed_layout, plain_layout | layout_shape, tensor_optional_slot, tensor_shape | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| explicit_layout, layout, placed_layout, plain_layout | layout_shape, tensor_optional_slot, tensor_shape | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | | function | function | FunctionDialectRule | A function kind and constructed value must agree with the active dialect. | src/tilefoundry/parser/pattern_nodes.py | | function | function | FunctionRegistrationRule | A validated function must be registered exactly once in its owning scope. | src/tilefoundry/parser/pattern_nodes.py | | function | function | 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 | | 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, 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, 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 @@ -218,6 +209,7 @@ function ::= 'def' name '(' signature ')' ('->' return-type)? ':' b | Ordered Rules | Validates and normalizes each owner value after construction. | | Module Build | Lets Python execute the class body, collects declarations, resolves child Modules first, then parses Functions in source order and finalizes the Module. | | Pattern Visitor | Traverses the same graph to render this section's generated grammar and constraints. | +| Refusal | Carries the reason from the pattern that claimed a node and then refused it, so a report names a cause rather than the absence of a match. | ```mermaid classDiagram @@ -227,6 +219,11 @@ classDiagram Element o-- AstPattern Element o-- AstRule AstPattern --> AstMatch + AstPattern --> MatchFailure + MatchFailure <|.. PatternFailure + MatchFailure <|.. ChoiceFailure + ChoiceFailure o-- MatchFailure : causes + ParseError <.. MatchFailure PatternVisitor ..> AstPattern ParserAPI ..> ModuleBuild ``` @@ -247,6 +244,45 @@ flowchart TD MODULE -->|no| RETURN["return standalone result"] ``` +```mermaid +flowchart TD + TRY["alternative.match(node)"] --> OUT{"outcome"} + OUT -->|"AstMatch"| WIN["choice accepts it; pending refusals are discarded"] + OUT -->|"MatchFailure"| CLAIM["claimed the node and refused: reason recorded"] + OUT -->|"None"| PASS["did not recognize the node: nothing recorded"] + PASS --> NEXT["try the next alternative"] + CLAIM --> NEXT + NEXT --> DONE{"any refusal recorded?"} + DONE -->|no| SILENT["return None: no alternative recognized this node"] + DONE -->|yes| COLLECT["ChoiceFailure over the claimants"] + COLLECT --> UP["travels up unchanged; combinators add no wrapping"] + UP --> RENDER["render(): a sole claimant is the whole report"] + RENDER --> RAISE["ParseError with source location"] + SILENT --> RAISE +``` + +A pattern MUST establish that a node is its own before it refuses with a reason. That claim is +what makes the reason trustworthy: it says no remaining alternative can accept this node, so +the refusal is the author's mistake and not another pattern's turn. A callee resolving to an op +schema is such a claim, and a wrong argument count after it is an error. A callee that does not +resolve is not a claim: it may be a bare name or a foreign namespace that another alternative +owns, so the pattern returns `None` and says nothing. Reasons are never reconstructed from the +AST after the fact; an inspection outside the refusing pattern cannot see which step it failed +at, and becomes a second, divergent copy of that knowledge. + +Being last in one choice is not a claim either, because that choice may itself be an +alternative in another. `parse_node` is the single place where no alternative remains, so it +is the only place that MAY describe a node from its shape rather than from a pattern's +statement, and it does so only when the shape says something worth reading. + +`None` and a `MatchFailure` differ only for a choice; every other combinator returns either one +unchanged, so a refusal keeps the identity and the wording of the pattern that produced it all +the way to `parse_node`. Nothing is wrapped, filtered, or re-described on the way up. A +`ChoiceFailure` records only the alternatives that claimed the node, which is normally one, and +it renders as that sole claimant. Two claimants mean two patterns claim overlapping shapes; the +report states both rather than choosing between them, because the ambiguity is in the grammar +and not in the report. + 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. diff --git a/docs/spec/passes.md b/docs/spec/passes.md index 1261413e..a2bdca1b 100644 --- a/docs/spec/passes.md +++ b/docs/spec/passes.md @@ -256,9 +256,9 @@ no return-tensor form. After this pass, `PassManager` reruns HIR Per-op lowering is **registry-dispatched**, not a hand-written `isinstance` chain ([§4](#4-transform-pass-idiom)): each HIR op registers its lowering handler keyed by op class (`register_hir_lowering(OpClass)`), and the pass looks the handler up by -`type(call.target)`. A target-owned op (e.g. the CUDA `Mma`) registers its own -lowering, so the pass core depends on the registry contract, not on importing -target-specific op classes. +`type(call.target)`. An op with a lowering-specific contract (e.g. HIR +`Reshard`) registers its own lowering, so the pass core depends on the registry +contract, not on importing target-specific op classes. A handler is a free function `handler(ctx, target, expr) -> Var`, where `ctx` is the lowering context and `target` is the dispatched `Op` @@ -309,19 +309,6 @@ TIR shape based on whether `storage` is provided: a plain tensor (no `ShardLayout`) and copy from the shard view into the plain storage. -#### CUDA HIR MMA refusal - -`HirToTirPass` MUST reject both `Mma_SM80_16x8x16` and -`Wgmma_SM90_64x128x16` by their concrete HIR Op name before lowering either -operand or emitting allocations, copies, fills, or `TirMma`. Their logical HIR -value/type/cost models are available to evaluation and Analyze only; there is -no HIR compile route. - -The independent handwritten `T.cuda.mma` atom and CUDA runtime surface remains -available and is specified by [tir §2.3](./tir.md#23-tir-ops). This pass MUST -NOT translate HIR fragments into that surface, modify its atom layouts, or rely -on CUDA codegen's legacy `atom=None` fallback. - #### Dispatch lowering The pass lowers each `Module.functions` entry by its shape diff --git a/docs/spec/semantic-analysis.md b/docs/spec/semantic-analysis.md index d7817070..e9fd6519 100644 --- a/docs/spec/semantic-analysis.md +++ b/docs/spec/semantic-analysis.md @@ -167,7 +167,7 @@ def derive_output_shard_layout( ### 3.3 Output storage and mesh/layout compatibility - constraints: - - A symmetric multi-input op (`Binary`, `MatMul`, `Concat`, `Stack`, `Mma`) + - A symmetric multi-input op (`Binary`, `MatMul`, `Concat`, `Stack`) resolves output storage by anchoring on the concrete residency among its operands ([types §2](./types.md#2-tensortype)); the rule is independent of operand order. diff --git a/src/tilefoundry/ir/hir/cuda/__init__.py b/src/tilefoundry/ir/hir/cuda/__init__.py deleted file mode 100644 index 19635214..00000000 --- a/src/tilefoundry/ir/hir/cuda/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""CUDA (NVIDIA) target-specific HIR IR nodes and descriptors.""" diff --git a/src/tilefoundry/ir/hir/cuda/nn/__init__.py b/src/tilefoundry/ir/hir/cuda/nn/__init__.py deleted file mode 100644 index b300cb64..00000000 --- a/src/tilefoundry/ir/hir/cuda/nn/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""CUDA NN instructions (MMA ops).""" diff --git a/src/tilefoundry/ir/hir/cuda/nn/mma.py b/src/tilefoundry/ir/hir/cuda/nn/mma.py deleted file mode 100644 index 377c33f2..00000000 --- a/src/tilefoundry/ir/hir/cuda/nn/mma.py +++ /dev/null @@ -1,296 +0,0 @@ -"""Define shape-specific HIR matrix-multiply-accumulate value Ops. - -``Mma`` is a dispatch marker. Concrete class names encode architecture and -shape while attributes encode dtype and orientation. HIR returns ``A @ B`` as -a value; lowering introduces and zero-initialises the in-place accumulator. -""" - -from __future__ import annotations - -from dataclasses import replace - -import isl -import torch - -from tilefoundry.evaluator import TensorValue, register_eval, to_torch_dtype -from tilefoundry.ir.core import Op -from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor -from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.hir._helpers import resolve_anchor_storage -from tilefoundry.ir.tir.cuda.nn.mma import ( - SM80_16x8x16_F32BF16BF16F32_TN, - make_atom, -) -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import Broadcast, Layout, ShardLayout, try_c_order_strides -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.visitor_registry import register_typeinfer -from tilefoundry.visitor_registry.access_relation import ( - AccessRelations, - AffineAccess, - BoundaryRelation, - identity_access, - iterating, - register_access_relation, -) - -_FLOAT_ACCUMULATOR_COMBINATIONS = frozenset( - { - (DType.f16, DType.f16, DType.f16), - (DType.bf16, DType.bf16, DType.bf16), - (DType.f16, DType.f16, DType.f32), - (DType.bf16, DType.bf16, DType.f32), - (DType.f32, DType.f32, DType.f32), - } -) -_SM80_ATOM = make_atom(SM80_16x8x16_F32BF16BF16F32_TN) - - -class Mma(Op): - """Abstract marker for the family of matrix multiply value Ops.""" - - -@register_op(category="nn") -class Mma_SM80_16x8x16(Mma): - """PTX ``mma.sync.aligned.m16n8k16`` logical value ``A @ B``.""" - - a = ParamDef(kind="input", pattern=Tensor) - b = ParamDef(kind="input", pattern=Tensor) - dtype_a = ParamDef(kind="attribute", annotation=DType) - dtype_b = ParamDef(kind="attribute", annotation=DType) - dtype_acc = ParamDef(kind="attribute", annotation=DType) - a_layout = ParamDef(kind="attribute", annotation=str, default="T") - b_layout = ParamDef(kind="attribute", annotation=str, default="N") - - -@register_op(category="nn") -class Wgmma_SM90_64x128x16(Mma): - """SM90 ``wgmma.mma_async.sync.aligned.m64n128k16`` logical ``A @ B``.""" - - a = ParamDef(kind="input", pattern=Tensor) - b = ParamDef(kind="input", pattern=Tensor) - dtype_a = ParamDef(kind="attribute", annotation=DType) - dtype_b = ParamDef(kind="attribute", annotation=DType) - dtype_acc = ParamDef(kind="attribute", annotation=DType) - a_layout = ParamDef(kind="attribute", annotation=str, default="T") - b_layout = ParamDef(kind="attribute", annotation=str, default="N") - - -def _is_lowerable_sm80_target(op: Mma) -> bool: - return ( - isinstance(op, Mma_SM80_16x8x16) - and op.dtype_a == DType.bf16 - and op.dtype_b == DType.bf16 - and op.dtype_acc == DType.f32 - and op.a_layout == "T" - and op.b_layout == "N" - ) - - -def _is_genuinely_sharded(layout) -> bool: - return isinstance(layout, ShardLayout) and any( - not isinstance(attr, Broadcast) for attr in layout.attrs - ) - - -def _matches_fragment_layout(actual: ShardLayout, expected: ShardLayout) -> bool: - return ( - actual.layout == expected.layout - and actual.attrs == expected.attrs - and actual.mesh.topologies == expected.mesh.topologies - and actual.mesh.layout == expected.mesh.layout - ) - - -def _derive_sm80_fragment_layout(a_ty: TensorType, b_ty: TensorType) -> ShardLayout: - """Return the known C fragment or reject a non-instruction A/B claim.""" - for index, ty, expected, role in ( - (0, a_ty, _SM80_ATOM.A, "A"), - (1, b_ty, _SM80_ATOM.B, "B"), - ): - if not isinstance(ty.layout, ShardLayout): - raise ValueError( - f"input {index} does not carry the known SM80 {role} fragment layout; " - "use an explicit Reshard to that layout and materialize-to-RMEM" - ) - if ty.storage is not StorageKind.RMEM: - raise ValueError( - f"input {index} SM80 {role} fragment is in {ty.storage}, not RMEM; " - "use an explicit Reshard and materialize-to-RMEM" - ) - if not _matches_fragment_layout(ty.layout, expected): - raise ValueError( - f"input {index} does not match the known SM80 {role} fragment layout; " - "use an explicit Reshard to that layout and materialize-to-RMEM" - ) - if ( - a_ty.layout.mesh.topologies != b_ty.layout.mesh.topologies - or a_ty.layout.mesh.layout != b_ty.layout.mesh.layout - ): - raise ValueError( - "input 1 SM80 B fragment references a different physical mesh from input 0; " - "use an explicit Reshard to the common fragment mesh and materialize-to-RMEM" - ) - return replace(_SM80_ATOM.C, mesh=a_ty.layout.mesh) - - -def _validate_contract(call: "Call", ctx: "TypeInferContext", a_shape, b_shape) -> tuple: - op = call.target - a_ty = ctx.type_of(call.args[0]) - b_ty = ctx.type_of(call.args[1]) - if a_ty.shape != a_shape: - ctx.error(call, f"a shape must be {a_shape}, got {a_ty.shape}") - if b_ty.shape != b_shape: - ctx.error(call, f"b shape must be {b_shape}, got {b_ty.shape}") - if a_ty.dtype != op.dtype_a: - ctx.error( - call, - f"dtype_a={op.dtype_a.name} disagrees with input a dtype {a_ty.dtype.name}", - ) - if b_ty.dtype != op.dtype_b: - ctx.error( - call, - f"dtype_b={op.dtype_b.name} disagrees with input b dtype {b_ty.dtype.name}", - ) - combo = (op.dtype_a, op.dtype_b, op.dtype_acc) - if combo not in _FLOAT_ACCUMULATOR_COMBINATIONS: - ctx.error( - call, - "dtype_acc combination " - f"(dtype_a={op.dtype_a.name}, dtype_b={op.dtype_b.name}, " - f"dtype_acc={op.dtype_acc.name}) is unsupported", - ) - for field in ("a_layout", "b_layout"): - value = getattr(op, field) - if value not in ("N", "T"): - ctx.error(call, f"{field} must be 'N' or 'T', got {value!r}") - return a_ty, b_ty - - -def _infer_mma(call, ctx, *, a_shape, b_shape, out_shape) -> TensorType: - a_ty, b_ty = _validate_contract(call, ctx, a_shape, b_shape) - genuine = [ - index - for index, ty in enumerate((a_ty, b_ty)) - if _is_genuinely_sharded(ty.layout) - ] - if genuine: - if isinstance(call.target, Wgmma_SM90_64x128x16): - ctx.error( - call, - f"input {genuine[0]} carries an unrepresentable WGMMA ShardLayout; " - "use an explicit Reshard to a plain logical layout before WGMMA", - ) - if not _is_lowerable_sm80_target(call.target): - ctx.error( - call, - f"input {genuine[0]} claims an SM80 fragment, but only the " - "BF16/BF16/F32 TN fragment contract is representable; use an " - "explicit Reshard to a plain logical layout or materialize-to-RMEM", - ) - try: - out_layout = _derive_sm80_fragment_layout(a_ty, b_ty) - except ValueError as error: - ctx.error(call, str(error)) - elif any(isinstance(ty.layout, ShardLayout) for ty in (a_ty, b_ty)): - out_layout = None - elif a_ty.layout is None and b_ty.layout is None: - out_layout = None - else: - out_layout = Layout(shape=out_shape, strides=try_c_order_strides(out_shape)) - return TensorType( - shape=out_shape, - dtype=call.target.dtype_acc, - layout=out_layout, - storage=resolve_anchor_storage(ctx, call, a_ty.storage, b_ty.storage), - ) - - -@register_typeinfer(Mma_SM80_16x8x16) -def _(call: "Call", ctx: "TypeInferContext") -> TensorType: - return _infer_mma( - call, - ctx, - a_shape=(16, 16), - b_shape=(16, 8), - out_shape=(16, 8), - ) - - -@register_typeinfer(Wgmma_SM90_64x128x16) -def _(call: "Call", ctx: "TypeInferContext") -> TensorType: - return _infer_mma( - call, - ctx, - a_shape=(64, 16), - b_shape=(16, 128), - out_shape=(64, 128), - ) - - -@register_eval(Mma_SM80_16x8x16) -@register_eval(Wgmma_SM90_64x128x16) -def _eval_mma(ctx): - dtype_acc = to_torch_dtype(ctx.op.dtype_acc) - a = ctx.args[0].data.to(dtype_acc) - b = ctx.args[1].data.to(dtype_acc) - return TensorValue(data=torch.matmul(a, b), type=ctx.result_type) - - -__all__ = ["Mma", "Mma_SM80_16x8x16", "Wgmma_SM90_64x128x16"] - - -_TILES = {"Mma_SM80_16x8x16": (16, 8, 16), "Wgmma_SM90_64x128x16": (64, 128, 16)} - - -def _whole_read(held: "Type", rank: int) -> "AffineAccess": - """Every coordinate of one operand, whatever the result coordinate is. - - A tile instruction reads its operands entire, so no result coordinate picks - out part of one. The coordinates are the operand's own axes; which positions - a layout made of them is the reader's question, answered by composing this - with that layout rather than by naming them here. - """ - terms, guards = [], [] - for position, extent in enumerate(held.shape): - if extent == 1: - terms.append("0") - continue - terms.append(f"p{position}") - guards.append(f"0 <= p{position} < {extent}") - dims = ", ".join(f"d{index}" for index in range(rank)) - image = ", ".join(terms) if terms else "0" - where = " and ".join(guards) - return AffineAccess( - isl.map(f"{{ [{dims}] -> [{image}]" + (f" : {where} }}" if where else " }")) - ) - - -def _tile_access(call: "Call", ctx) -> AccessRelations: - """A fixed tile: both operands read whole, the accumulator written whole. - - ``_TILES`` holds each instruction's own count, which is what separates these - from a MatMul: one instruction moves one instruction's elements however many - participants issue it. That count is the accumulator's own extents, so the - space is stated from the instruction rather than from the Type being - derived. Where those elements sit is the reader's answer, so the patterns - are stated in the axes the operands were written in. - """ - m, n, _ = _TILES[type(call.target).__name__] - tile = (m, n) - rank = len(tile) - return iterating( - tile, - AccessRelations( - inputs=( - BoundaryRelation(_whole_read(ctx.type_of(call.args[0]), rank)), - BoundaryRelation(_whole_read(ctx.type_of(call.args[1]), rank)), - ), - outputs=(BoundaryRelation(identity_access(rank)),), - ), - ) - - -register_access_relation(Mma_SM80_16x8x16)(_tile_access) -register_access_relation(Wgmma_SM90_64x128x16)(_tile_access) diff --git a/src/tilefoundry/ir/hir/math/binary.py b/src/tilefoundry/ir/hir/math/binary.py index 28951c70..cb9e14e9 100644 --- a/src/tilefoundry/ir/hir/math/binary.py +++ b/src/tilefoundry/ir/hir/math/binary.py @@ -131,10 +131,10 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: ctx.error( call, f"Binary {op.kind.name}: dtype mismatch " - f"({lhs_ty.dtype.name} vs {rhs_ty.dtype.name}); tensor " - f"operands are never promoted for you. A number you wrote is " - f"f32 -- tf.full_like(x, value=...) gives it x's dtype; " - f"otherwise tf.cast one side explicitly. " + f"({lhs_ty.dtype.name} vs {rhs_ty.dtype.name}); operands are never " + f"promoted, and a Python literal is f32 or i64 like any other. " + f"Write the dtype you want: " + f"tf.cast(, dtype={lhs_ty.dtype.name!r}). " f"See `tilefoundry spec dsl binary`", ) if op.kind in _LOGICAL_KINDS and lhs_ty.dtype != DType.bool: diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index fe9f4bce..cfb8e04f 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -79,42 +79,56 @@ _TYPE_INFER_CONTEXT = "" +class MatchFailure: + """A refusal travelling back up the pattern graph. + + ``match`` has three outcomes: an ``AstMatch``, ``None`` when the pattern did + not recognize the shape it was given, and a ``MatchFailure`` when a pattern + claimed the node and then refused it. Only the third carries a reason, and + only a pattern that has established the node is its own may produce one. + """ + + node: object + + def render(self, indent: int = 0) -> str: + raise NotImplementedError + + @dataclass(frozen=True) -class PatternFailure: - """A recognized pattern whose nested validation failed.""" +class PatternFailure(MatchFailure): + """A pattern that claimed this node, refused it, and said why.""" pattern_id: str node: object detail: str - causes: tuple["PatternFailure", ...] = () def render(self, indent: int = 0) -> str: - lines = [f"{' ' * indent}{self.pattern_id}: {self.detail}"] - for cause in self.causes: - lines.append(cause.render(indent + 1)) - return "\n".join(lines) - - -def _pattern_label(pattern: object) -> str: - label = getattr(pattern, "element_name", None) - if isinstance(label, str) and label: - return label - pattern_id = getattr(pattern, "pattern_id", None) - if isinstance(pattern_id, str) and pattern_id: - return pattern_id - node_type = getattr(pattern, "node_type", None) - if isinstance(node_type, type): - return node_type.__name__ - return type(pattern).__name__ - - -def _wrap_failure(pattern: object, node: object, failure: PatternFailure) -> PatternFailure: - return PatternFailure( - pattern_id=_pattern_label(pattern), - node=node, - detail="nested pattern failed", - causes=(failure,), - ) + return f"{' ' * indent}{self.pattern_id}: {self.detail}" + + +@dataclass(frozen=True) +class ChoiceFailure(MatchFailure): + """Every refusal a choice was handed by the alternatives that claimed the node. + + An alternative that did not recognize the node contributes nothing, so a sole + claimant is the whole report and this wrapper renders as that claimant. + Several claimants mean their claims overlap, which the report states rather + than resolving on its own. + """ + + node: object + causes: tuple[MatchFailure, ...] + + def render(self, indent: int = 0) -> str: + if len(self.causes) == 1: + return self.causes[0].render(indent) + head = f"{' ' * indent}choice: no alternative matched" + return "\n".join([head, *(cause.render(indent + 1) for cause in self.causes)]) + + +def is_matched(result: object) -> bool: + """Whether matching succeeded. Anything else travels back up unchanged.""" + return isinstance(result, AstMatch) def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContext") -> object: @@ -142,63 +156,63 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex runtime = SimpleNamespace( - Call=Call, - Broadcast=Broadcast, - Binary=Binary, - BinaryKind=BinaryKind, - Constant=Constant, - DType=DType, - DimAdd=DimAdd, - DimFloorDiv=DimFloorDiv, - DimMod=DimMod, - DimMul=DimMul, - DimSub=DimSub, - DimVar=DimVar, - Evaluate=Evaluate, - Expr=Expr, - ExecutionDomainMetadata=ExecutionDomainMetadata, - Function=Function, - GridRegionExpr=GridRegionExpr, - Arange=Arange, - IrTuple=IrTuple, - Layout=Layout, - LayoutBase=LayoutBase, - Local=Local, - LetStmt=LetStmt, - MeshScope=MeshScope, - Mesh=Mesh, - Module=Module, - OpSchema=OpSchema, - PrimFunction=PrimFunction, - Reshard=Reshard, - Reshape=Reshape, - Return=Return, - Sequential=Sequential, - ShardLayout=ShardLayout, - Slice=Slice, - Split=Split, - StorageKind=StorageKind, - TensorType=TensorType, - TupleType=TupleType, - TypeInferContext=TypeInferContext, - FunctionScope=FunctionScope, - TypeInferVisitor=TypeInferVisitor, - TupleGetItem=TupleGetItem, - Unary=Unary, - UnaryKind=UnaryKind, - UnitType=UnitType, - Var=Var, - DISPLAY_NAME=DISPLAY_NAME, - c_order_strides=c_order_strides, - canonical_shard_layout=canonical_shard_layout, - composed=composed, - dim_expr=dim_expr, - elaborate=elaborate, - normalize_dim=normalize_dim, - slice_size=slice_size, - simplify_dim=simplify_dim, - resolve_storage=resolve_storage, - ) + Call=Call, + Broadcast=Broadcast, + Binary=Binary, + BinaryKind=BinaryKind, + Constant=Constant, + DType=DType, + DimAdd=DimAdd, + DimFloorDiv=DimFloorDiv, + DimMod=DimMod, + DimMul=DimMul, + DimSub=DimSub, + DimVar=DimVar, + Evaluate=Evaluate, + Expr=Expr, + ExecutionDomainMetadata=ExecutionDomainMetadata, + Function=Function, + GridRegionExpr=GridRegionExpr, + Arange=Arange, + IrTuple=IrTuple, + Layout=Layout, + LayoutBase=LayoutBase, + Local=Local, + LetStmt=LetStmt, + MeshScope=MeshScope, + Mesh=Mesh, + Module=Module, + OpSchema=OpSchema, + PrimFunction=PrimFunction, + Reshard=Reshard, + Reshape=Reshape, + Return=Return, + Sequential=Sequential, + ShardLayout=ShardLayout, + Slice=Slice, + Split=Split, + StorageKind=StorageKind, + TensorType=TensorType, + TupleType=TupleType, + TypeInferContext=TypeInferContext, + FunctionScope=FunctionScope, + TypeInferVisitor=TypeInferVisitor, + TupleGetItem=TupleGetItem, + Unary=Unary, + UnaryKind=UnaryKind, + UnitType=UnitType, + Var=Var, + DISPLAY_NAME=DISPLAY_NAME, + c_order_strides=c_order_strides, + canonical_shard_layout=canonical_shard_layout, + composed=composed, + dim_expr=dim_expr, + elaborate=elaborate, + normalize_dim=normalize_dim, + slice_size=slice_size, + simplify_dim=simplify_dim, + resolve_storage=resolve_storage, +) class AstRule(Protocol[T]): @@ -218,9 +232,7 @@ class AstPattern(Protocol[T]): def accept(self, visitor: PatternVisitor[Any]) -> Any: ... - def match( - self, node: object, context: MatchContext - ) -> AstMatch[T] | PatternFailure | None: ... + def match(self, node: object, context: MatchContext) -> AstMatch[T] | MatchFailure | None: ... class PatternVisitor(Protocol[T]): @@ -296,10 +308,8 @@ def match(self, node: object, context: MatchContext) -> AstMatch[T] | None: if syntax is None: raise TypeError(f"{type(self).__name__} has no executable syntax") matched = syntax.match(node, context) - if isinstance(matched, PatternFailure): - return _wrap_failure(self, node, matched) - if matched is None: - return None + if not is_matched(matched): + return matched if isinstance(matched.pattern, ElementPattern): return matched return AstMatch( @@ -326,9 +336,7 @@ def pattern(self) -> AstPattern[Any]: self._resolved = self.factory() return self._resolved - def match( - self, node: object, context: MatchContext - ) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: matched = self.pattern.match(node, context) return matched @@ -338,16 +346,14 @@ def __init__(self, node_type: type, *parts: AstPattern[Any]): self.node_type = node_type self.parts = tuple(parts) - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: if not isinstance(node, self.node_type): return None matches: list[AstMatch[Any]] = [] for part in self.parts: matched = part.match(node, context) - if isinstance(matched, PatternFailure): - return _wrap_failure(self, node, matched) - if matched is None: - return None + if not is_matched(matched): + return matched matches.append(matched) return self._merge( self, @@ -363,15 +369,13 @@ def __init__(self, name: str, pattern: AstPattern[Any]): self.name = name self.pattern = pattern - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: if not hasattr(node, self.name): return None value = getattr(node, self.name) matched = self.pattern.match(value, context) - if isinstance(matched, PatternFailure): - return _wrap_failure(self, node, matched) - if matched is None: - return None + if not is_matched(matched): + return matched return AstMatch( self, matched.pattern_id, @@ -393,7 +397,7 @@ def __init__( self.value = value self.value_type = value_type - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: raw = node.value if isinstance(node, ast.Constant) else node if self.value is not dataclasses.MISSING and raw != self.value: return None @@ -403,13 +407,11 @@ def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternF class ReferencePattern(CombinatorPattern): - def __init__( - self, *, resolve: bool = False, expected: type | tuple[type, ...] | None = None - ): + def __init__(self, *, resolve: bool = False, expected: type | tuple[type, ...] | None = None): self.resolve = resolve self.expected = expected - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: if not isinstance(node, (ast.Name, ast.Attribute)): return None captures: dict[str, object] = {} @@ -428,16 +430,14 @@ class SequencePattern(CombinatorPattern): def __init__(self, *patterns: AstPattern[Any]): self.patterns = tuple(patterns) - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: if not isinstance(node, (tuple, list)) or len(node) != len(self.patterns): return None matches: list[AstMatch[Any]] = [] for value, pattern in zip(node, self.patterns): matched = pattern.match(value, context) - if isinstance(matched, PatternFailure): - return _wrap_failure(self, node, matched) - if matched is None: - return None + if not is_matched(matched): + return matched matches.append(matched) return self._merge( self, @@ -452,30 +452,21 @@ class ChoicePattern(CombinatorPattern): def __init__(self, *patterns: AstPattern[Any]): self.patterns = tuple(patterns) - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: - failures: list[PatternFailure] = [] + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: + """Take the first alternative that matches, keeping what the others refused. + + An alternative returning ``None`` did not recognize the node and has no + opinion, so nothing is recorded for it; only an alternative that claimed + the node and then refused contributes to the report. + """ + failures: list[MatchFailure] = [] for pattern in self.patterns: matched = pattern.match(node, context) - if isinstance(matched, PatternFailure): - failures.append(matched) - continue - if matched is not None: + if is_matched(matched): return matched - failures.append( - PatternFailure( - pattern_id=_pattern_label(pattern), - node=node, - detail="pattern did not match", - ) - ) - if failures: - return PatternFailure( - pattern_id="choice", - node=node, - detail="no choice matched", - causes=tuple(failures), - ) - return None + if isinstance(matched, MatchFailure): + failures.append(matched) + return ChoiceFailure(node, tuple(failures)) if failures else None class ConditionPattern(CombinatorPattern): @@ -491,7 +482,7 @@ def __init__( self.test = test self.pattern = pattern - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: if not self.test(node, context): return None return self.pattern.match(node, context) @@ -501,14 +492,12 @@ class OptionalPattern(CombinatorPattern): def __init__(self, pattern: AstPattern[Any]): self.pattern = pattern - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: if node is None: return AstMatch(self, "optional", node, {}, "optional") matched = self.pattern.match(node, context) - if isinstance(matched, PatternFailure): - return _wrap_failure(self, node, matched) - if matched is None: - return None + if not is_matched(matched): + return matched return AstMatch( self, matched.pattern_id, @@ -529,22 +518,18 @@ def __init__(self, pattern: AstPattern[Any], *, minimum: int = 0): def _index_child(child: AstChild, index: int) -> AstChild: return dataclasses.replace(child, name=child.name.format(index=index)) - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: if not isinstance(node, (tuple, list)) or len(node) < self.minimum: return None matches: list[AstMatch[Any]] = [] for index, value in enumerate(node): matched = self.pattern.match(value, context) - if isinstance(matched, PatternFailure): - return _wrap_failure(self, node, matched) - if matched is None: - return None + if not is_matched(matched): + return matched matches.append( dataclasses.replace( matched, - children=tuple( - self._index_child(child, index) for child in matched.children - ), + children=tuple(self._index_child(child, index) for child in matched.children), ) ) return self._merge( @@ -561,7 +546,7 @@ def __init__(self, label: str, predicate: Callable[[object, MatchContext], bool] self.label = label self.predicate = predicate - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: if not self.predicate(node, context): return None return AstMatch(self, "predicate", node, {}, "predicate") @@ -572,7 +557,7 @@ def __init__(self, name: str, extractor: Callable[[object, MatchContext], object self.name = name self.extractor = extractor - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: try: value = self.extractor(node, context) except (AttributeError, KeyError, TypeError, ValueError): @@ -612,13 +597,11 @@ def pattern(self) -> AstPattern[Any]: self._pattern = pattern return pattern - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: value = self.transform(node) if self.transform is not None else node if not isinstance(value, ast.AST): return None - values = ( - self.values(value, context) if callable(self.values) else self.values or {} - ) + values = self.values(value, context) if callable(self.values) else self.values or {} expected_type = ( self.expected_type(value, context) if callable(self.expected_type) @@ -638,21 +621,15 @@ def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternF class BranchPattern(CombinatorPattern): - def __init__( - self, branch_id: str, pattern: AstPattern[Any], *, pattern_id: str | None = None - ): + def __init__(self, branch_id: str, pattern: AstPattern[Any], *, pattern_id: str | None = None): self.branch_id = branch_id self.pattern = pattern self.pattern_id = pattern_id or branch_id - def match( - self, node: object, context: MatchContext - ) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: matched = self.pattern.match(node, context) - if isinstance(matched, PatternFailure): - return _wrap_failure(self, node, matched) - if matched is None: - return None + if not is_matched(matched): + return matched return AstMatch( self, self.pattern_id, @@ -670,17 +647,17 @@ class BindPattern(CombinatorPattern): def __init__( self, pattern: AstPattern[Any], - binder: Callable[[object, MatchContext, AstMatch[Any]], AstMatch[Any] | None], + binder: Callable[ + [object, MatchContext, AstMatch[Any]], AstMatch[Any] | MatchFailure | None + ], ): self.pattern = pattern self.binder = binder - def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: matched = self.pattern.match(node, context) - if isinstance(matched, PatternFailure): - return _wrap_failure(self, node, matched) - if matched is None: - return None + if not is_matched(matched): + return matched bound = self.binder(node, context, matched) return bound @@ -773,9 +750,7 @@ def __post_init__(self) -> None: role = self.function_kind if isinstance(role, str) and role not in {item.value for item in FunctionRole}: role = ( - FunctionRole.ROOT - if role in {"func", "prim_func", "kernel"} - else FunctionRole(role) + FunctionRole.ROOT if role in {"func", "prim_func", "kernel"} else FunctionRole(role) ) if role is not self.role: object.__setattr__(self, "role", role) @@ -792,11 +767,7 @@ def __post_init__(self) -> None: @property def specializations(self) -> tuple[object, ...]: - return ( - () - if self.role is not FunctionRole.VARIANT or self.key is None - else (self.key,) - ) + return () if self.role is not FunctionRole.VARIANT or self.key is None else (self.key,) @property def converter(self) -> object | None: @@ -813,9 +784,7 @@ def _child_for(self, callee: object): if self.module_scope is None: return None for _name, child in self.module_scope.items(): - if getattr(child, "owns", lambda *_args, **_kwargs: False)( - callee, derived=True - ): + if getattr(child, "owns", lambda *_args, **_kwargs: False)(callee, derived=True): return child return None @@ -860,9 +829,7 @@ def apply( @dataclass(frozen=True) class ModuleFunctionRegistrationRule: - STATEMENT: ClassVar[str] = ( - "A validated module function must be recorded in declaration order." - ) + STATEMENT: ClassVar[str] = "A validated module function must be recorded in declaration order." def apply( self, @@ -967,13 +934,9 @@ def _validate_function(self, function: object, context: FuncParserContext) -> No for root in self.roots ): raise ValueError( - self._binding_error( - role, getattr(function, "name", binding), self.owner_name - ) + self._binding_error(role, getattr(function, "name", binding), self.owner_name) ) - expected = ( - runtime.PrimFunction if context.dialect == "tir" else runtime.Function - ) + expected = runtime.PrimFunction if context.dialect == "tir" else runtime.Function if not isinstance(function, expected): raise TypeError( f"root {binding!r} constructed {type(function).__name__}, expected {expected.__name__}" @@ -983,13 +946,9 @@ def _validate_function(self, function: object, context: FuncParserContext) -> No if base is None or not isinstance(base, runtime.Function): raise ValueError(f"{role.value} {binding!r}: base is not a HIR Function") if getattr(base, "_sealed", False): - raise RuntimeError( - f"base {base.name!r}: cannot register {role.value} after seal" - ) + raise RuntimeError(f"base {base.name!r}: cannot register {role.value} after seal") if binding == "_" and role is FunctionRole.VARIANT: - raise ValueError( - f"base {base.name!r}: a variant binding may not be named '_'" - ) + raise ValueError(f"base {base.name!r}: a variant binding may not be named '_'") if binding in self.bindings and role is not FunctionRole.CONVERTER: raise ValueError(self._binding_error(role, binding, self.owner_name)) if role is FunctionRole.VARIANT: @@ -998,9 +957,7 @@ def _validate_function(self, function: object, context: FuncParserContext) -> No keys = self.variant_keys.setdefault(id(base), set()) key = context.key if key in keys: - raise ValueError( - f"base {base.name!r}: duplicate specialization key {key!r}" - ) + raise ValueError(f"base {base.name!r}: duplicate specialization key {key!r}") return if getattr(function, "body", None) is None: raise ValueError(f"base {base.name!r}: a converter must have a real body") @@ -1286,13 +1243,9 @@ def resolve_lexical(self, name: str) -> object: ) return value - def resolve_static( - self, node: ast.AST, expected: type[T] | tuple[type[Any], ...] - ) -> T: + def resolve_static(self, node: ast.AST, expected: type[T] | tuple[type[Any], ...]) -> T: if not isinstance(node, (ast.Name, ast.Attribute)): - raise ParseError.from_node( - node, self, "static references use Name or Attribute" - ) + raise ParseError.from_node(node, self, "static references use Name or Attribute") value = _resolve_reference(node, self) if not isinstance(value, expected): raise ParseError.from_node( @@ -1332,9 +1285,7 @@ class AstMatch(Generic[T]): def construct(self, children: Mapping[str, object], context: MatchContext) -> T: pattern_constructor = getattr(self.pattern, "construct", None) if not callable(pattern_constructor): - raise TypeError( - f"Pattern {type(self.pattern).__name__} has no construct method" - ) + raise TypeError(f"Pattern {type(self.pattern).__name__} has no construct method") value = pattern_constructor(self, children, context) for rule in self.pattern.RULES: value = rule.apply(value, match=self, context=context) @@ -1375,14 +1326,32 @@ def from_node( return cls(node=node, context=context, detail=detail) +def _unclaimed_detail(node: ast.AST) -> str | None: + """Describe a node no pattern would take, when its shape says enough. + + This is the only place where "no alternative is left" is a fact rather than + a position in one choice, so it is the only place a description may be + derived from the node instead of stated by a pattern. + """ + if not isinstance(node, ast.Call): + return None + try: + callee = ast.unparse(node.func) + except (TypeError, ValueError): + callee = type(node.func).__name__ + keywords = [keyword.arg or "**" for keyword in node.keywords] + stated = f"keywords {keywords!r}" if keywords else "no keywords" + return f"unsupported call {callee!r} ({len(node.args)} positional, {stated})" + + def parse_node(pattern: AstPattern[T], node: ast.AST, context: MatchContext) -> T: """Select, recursively construct, and apply rules for one local pattern.""" matched = pattern.match(node, context) - if isinstance(matched, PatternFailure): + if isinstance(matched, MatchFailure): raise ParseError.from_node(node, context, matched.render()) if matched is None: - raise ParseError.from_node(node, context) + raise ParseError.from_node(node, context, _unclaimed_detail(node)) active_context = matched.construct_context or context children: dict[str, object] = {} for child in matched.children: @@ -1492,14 +1461,10 @@ class CanonicalDTypeRule: def apply(self, value, *, match, context): if not isinstance(value, runtime.DType): - raise ParseError.from_node( - match.node, context, "dtype did not construct DType" - ) + raise ParseError.from_node(match.node, context, "dtype did not construct DType") canonical = runtime.DType._members().get(value.name) if canonical is not value: - raise ParseError.from_node( - match.node, context, f"non-canonical dtype {value.name!r}" - ) + raise ParseError.from_node(match.node, context, f"non-canonical dtype {value.name!r}") return value @@ -1509,13 +1474,9 @@ class LayoutShapeRule: def apply(self, value, *, match, context): if value is not None and not isinstance(value, runtime.LayoutBase): - raise ParseError.from_node( - match.node, context, "layout did not construct LayoutBase" - ) + raise ParseError.from_node(match.node, context, "layout did not construct LayoutBase") if value is not None and not isinstance(value.shape, tuple): - raise ParseError.from_node( - match.node, context, "layout shape is not a tuple" - ) + raise ParseError.from_node(match.node, context, "layout shape is not a tuple") return value @@ -1525,9 +1486,7 @@ class LayoutPositionRule: def apply(self, value, *, match, context): if context.role in {"storage", "dtype", "shape"}: - raise ParseError.from_node( - match.node, context, "layout used in a non-layout role" - ) + raise ParseError.from_node(match.node, context, "layout used in a non-layout role") return value @@ -1537,47 +1496,31 @@ class StorageValueRule: def apply(self, value, *, match, context): if not isinstance(value, runtime.StorageKind): - raise ParseError.from_node( - match.node, context, "storage did not construct StorageKind" - ) + raise ParseError.from_node(match.node, context, "storage did not construct StorageKind") return value @dataclass(frozen=True) class TensorLayoutStorageRule: - STATEMENT: ClassVar[str] = ( - "A tensor type must contain compatible layout and storage values." - ) + STATEMENT: ClassVar[str] = "A tensor type must contain compatible layout and storage values." def apply(self, value, *, match, context): if not isinstance(value, runtime.TensorType): - raise ParseError.from_node( - match.node, context, "tensor did not construct TensorType" - ) + raise ParseError.from_node(match.node, context, "tensor did not construct TensorType") if not isinstance(value.storage, runtime.StorageKind): - raise ParseError.from_node( - match.node, context, "tensor storage is not StorageKind" - ) - if value.layout is not None and not isinstance( - value.layout, runtime.LayoutBase - ): - raise ParseError.from_node( - match.node, context, "tensor layout is not LayoutBase" - ) + raise ParseError.from_node(match.node, context, "tensor storage is not StorageKind") + if value.layout is not None and not isinstance(value.layout, runtime.LayoutBase): + raise ParseError.from_node(match.node, context, "tensor layout is not LayoutBase") return value @dataclass(frozen=True) class TensorPositionRule: - STATEMENT: ClassVar[str] = ( - "A tensor type's storage must be legal for its dialect and position." - ) + STATEMENT: ClassVar[str] = "A tensor type's storage must be legal for its dialect and position." def apply(self, value, *, match, context): if context.role == "storage": - raise ParseError.from_node( - match.node, context, "TensorType used in a storage role" - ) + raise ParseError.from_node(match.node, context, "TensorType used in a storage role") allowed = context.values.get("allowed_storage") if allowed is not None and value.storage not in allowed: rendered = tuple(str(item) for item in allowed) @@ -1591,15 +1534,11 @@ def apply(self, value, *, match, context): @dataclass(frozen=True) class ShapeDimRule: - STATEMENT: ClassVar[str] = ( - "A shape dimension must be an integer, DimVar, or expression." - ) + STATEMENT: ClassVar[str] = "A shape dimension must be an integer, DimVar, or expression." def apply(self, value, *, match, context): value = runtime.normalize_dim(value) - if isinstance(value, bool) or not isinstance( - value, (int, runtime.DimVar, runtime.Expr) - ): + if isinstance(value, bool) or not isinstance(value, (int, runtime.DimVar, runtime.Expr)): raise ParseError.from_node( match.node, context, @@ -1632,7 +1571,9 @@ def apply(self, value, *, match, context): "FunctionRole", "LayoutPattern", "LexicalScope", + "ChoiceFailure", "MatchContext", + "MatchFailure", "ModuleBuildContext", "NamePattern", "ParseError", @@ -1678,9 +1619,7 @@ def _infer_call(operation, args, context): placeholder_type = runtime.TensorType.scalar(runtime.DType.f32) metadata = () if context.function is not None and context.function.state.mesh_stack: - metadata = ( - runtime.ExecutionDomainMetadata(tuple(context.function.state.mesh_stack)), - ) + metadata = (runtime.ExecutionDomainMetadata(tuple(context.function.state.mesh_stack)),) placeholder = runtime.Call( type=placeholder_type, target=operation, args=tuple(args), metadata=metadata ) diff --git a/src/tilefoundry/parser/grammar_render.py b/src/tilefoundry/parser/grammar_render.py index 53b499d9..844ecad1 100644 --- a/src/tilefoundry/parser/grammar_render.py +++ b/src/tilefoundry/parser/grammar_render.py @@ -117,28 +117,20 @@ def _list_pattern(self, pattern: object) -> _Expr: if isinstance(pattern, RepeatPattern): return _repeated(self.visit(pattern.pattern), pattern.minimum, comma) if isinstance(pattern, SequencePattern): - return _separated( - tuple(self.visit(item) for item in pattern.patterns), comma - ) + return _separated(tuple(self.visit(item) for item in pattern.patterns), comma) if isinstance(pattern, ChoicePattern): return _choice(*(self._list_pattern(item) for item in pattern.patterns)) return self.visit(pattern) @staticmethod def _fields(pattern: AstNodePattern) -> dict[str, object]: - return { - part.name: part.pattern - for part in pattern.parts - if isinstance(part, FieldPattern) - } + return {part.name: part.pattern for part in pattern.parts if isinstance(part, FieldPattern)} def _field(self, fields: dict[str, object], name: str, fallback: str) -> _Expr: pattern = fields.get(name) return _text(fallback) if pattern is None else self.visit(pattern) - def _optional_field( - self, fields: dict[str, object], name: str, fallback: str - ) -> _Expr: + def _optional_field(self, fields: dict[str, object], name: str, fallback: str) -> _Expr: pattern = fields.get(name) if isinstance(pattern, OptionalPattern): pattern = pattern.pattern @@ -151,9 +143,7 @@ def _ast_node(self, pattern: AstNodePattern) -> _Expr: fields = self._fields(pattern) if node_type is ast.Constant: - predicates = [ - part for part in pattern.parts if isinstance(part, PredicatePattern) - ] + predicates = [part for part in pattern.parts if isinstance(part, PredicatePattern)] if predicates: return self.visit(predicates[-1]) value = fields.get("value") @@ -199,9 +189,7 @@ def _ast_node(self, pattern: AstNodePattern) -> _Expr: structural = [ part for part in pattern.parts - if not isinstance( - part, (CapturePattern, FieldPattern, PredicatePattern) - ) + if not isinstance(part, (CapturePattern, FieldPattern, PredicatePattern)) ] if structural: return self.visit(structural[-1]) @@ -251,6 +239,17 @@ def _ast_node(self, pattern: AstNodePattern) -> _Expr: ) return self._list_pattern(values) if node_type is ast.Call: + args = fields.get("args") + keywords = fields.get("keywords") + if not isinstance(args, RepeatPattern) or not isinstance(keywords, RepeatPattern): + arguments = _concat( + self._list_pattern(args) if args is not None else None, + self._list_pattern(keywords) if keywords is not None else None, + ) + return _concat( + self._field(fields, "func", "callee"), + _delimited("(", arguments, ")"), + ) argument_patterns: list[_Expr] = [] minimum = 0 for name in ("args", "keywords"): @@ -344,9 +343,7 @@ def _ast_node(self, pattern: AstNodePattern) -> _Expr: if node_type is ast.Module: body = fields.get("body", SequencePattern()) if isinstance(body, RepeatPattern): - return _repeated( - self.visit(body.pattern), body.minimum, _text("newline") - ) + return _repeated(self.visit(body.pattern), body.minimum, _text("newline")) return self._list_pattern(body) if node_type is ast.FunctionDef: returns = self._optional_field(fields, "returns", "return-type") @@ -368,6 +365,14 @@ def _ast_node(self, pattern: AstNodePattern) -> _Expr: ast.FloorDiv: "//", ast.Mod: "%", ast.Pow: "**", + ast.Eq: "==", + ast.NotEq: "!=", + ast.Lt: "<", + ast.LtE: "<=", + ast.Gt: ">", + ast.GtE: ">=", + ast.And: "and", + ast.Or: "or", ast.UAdd: "+", ast.USub: "-", ast.Not: "not", @@ -384,9 +389,7 @@ def _ast_node(self, pattern: AstNodePattern) -> _Expr: ] if structural: return self.visit(structural[-1]) - predicates = [ - part for part in pattern.parts if isinstance(part, PredicatePattern) - ] + predicates = [part for part in pattern.parts if isinstance(part, PredicatePattern)] if predicates: return self.visit(predicates[-1]) return _text("expression" if node_type is ast.expr else "statement") @@ -415,10 +418,7 @@ def visit(self, pattern: Any) -> _Expr: str: "string-literal", } return _choice( - *( - _text(names.get(item, f"{item.__name__}-literal")) - for item in types - ) + *(_text(names.get(item, f"{item.__name__}-literal")) for item in types) ) if pattern.value is dataclasses.MISSING: return _text("literal") @@ -426,9 +426,7 @@ def visit(self, pattern: Any) -> _Expr: if isinstance(pattern, ReferencePattern): return _text("primary") if isinstance(pattern, SequencePattern): - return _separated( - tuple(self.visit(item) for item in pattern.patterns), _terminal(",") - ) + return _separated(tuple(self.visit(item) for item in pattern.patterns), _terminal(",")) if isinstance(pattern, ChoicePattern): return _choice(*(self.visit(item) for item in pattern.patterns)) if isinstance(pattern, ConditionPattern): @@ -436,9 +434,7 @@ def visit(self, pattern: Any) -> _Expr: if isinstance(pattern, OptionalPattern): return _optional(self.visit(pattern.pattern)) if isinstance(pattern, RepeatPattern): - return _repeated( - self.visit(pattern.pattern), pattern.minimum, _terminal(",") - ) + return _repeated(self.visit(pattern.pattern), pattern.minimum, _terminal(",")) if isinstance(pattern, (ChildPattern, BranchPattern, BindPattern, LazyPattern)): return self.visit(pattern.pattern) if isinstance(pattern, PredicatePattern): diff --git a/src/tilefoundry/parser/parser_visitor.py b/src/tilefoundry/parser/parser_visitor.py index 6d85c4d7..b5ad89ce 100644 --- a/src/tilefoundry/parser/parser_visitor.py +++ b/src/tilefoundry/parser/parser_visitor.py @@ -24,9 +24,7 @@ def __init__(self, context: FuncParserContext): self.root_pattern = FunctionPattern() def visit(self, node: ast.AST) -> Any: - return parse_node( - self.root_pattern, node, MatchContext.from_function(self.context) - ) + return parse_node(self.root_pattern, node, MatchContext.from_function(self.context)) def visit_function(self, node: ast.FunctionDef) -> Any: return self.visit(node) diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index 47b38f64..da6ece1a 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -29,6 +29,7 @@ get_metadata, replace_metadata, ) +from tilefoundry.ir.hir.nn.matmul import MatMul from tilefoundry.ir.tir.launch import launch_call from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import DimVar @@ -61,6 +62,7 @@ LiteralPattern, LoopFrame, MatchContext, + MatchFailure, OptionalPattern, ParseError, PatternFailure, @@ -1401,10 +1403,14 @@ class StaticDictPattern(ElementPattern): ) @staticmethod - def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + def _bind( + node: object, context: MatchContext, matched: AstMatch[Any] + ) -> AstMatch[Any] | MatchFailure | None: assert isinstance(node, ast.Dict) if any(key is None for key in node.keys): - return None + return PatternFailure( + "static_dict", node, "`**` unpacking is not a static dictionary form" + ) children: list[AstChild] = [] for index, (key, value) in enumerate(zip(node.keys, node.values)): assert key is not None @@ -1552,11 +1558,18 @@ class StaticCallPattern(ElementPattern): ) @staticmethod - def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + def _bind( + node: object, context: MatchContext, matched: AstMatch[Any] + ) -> AstMatch[Any] | MatchFailure | None: assert isinstance(node, ast.Call) keyword_names = [keyword.arg for keyword in node.keywords] if len(set(keyword_names)) != len(keyword_names): - return None + repeated = sorted( + {name for name in keyword_names if name and keyword_names.count(name) > 1} + ) + return PatternFailure( + "static_call", node, f"keyword given more than once: {', '.join(repeated)}" + ) children = [AstChild("callee", StaticValuePattern(), node.func, "static_callee")] children.extend( AstChild( @@ -2028,8 +2041,7 @@ def _variadic_item_annotation(param: object) -> object | None: @dataclass(frozen=True) class CallVariadicInputFormRule: STATEMENT: ClassVar[str] = ( - "A variadic call must use one explicit list, tuple, or supported static " - "list comprehension." + "A variadic call must use one explicit list, tuple, or supported static list comprehension." ) def apply(self, value, *, match, context): @@ -2102,7 +2114,7 @@ def _pattern_for_param(param: object, node: ast.AST) -> AstPattern[Any]: @staticmethod def _schema_children( node: ast.Call, schema: object, context: MatchContext - ) -> tuple[AstChild, ...] | None: + ) -> tuple[AstChild, ...] | MatchFailure | 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 @@ -2112,10 +2124,7 @@ def _schema_children( 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 = ( - len(inputs) == 1 - and _variadic_item_annotation(inputs[0]) is not None - ) + 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() @@ -2152,7 +2161,12 @@ def _schema_children( continue attr_index = index - len(inputs) if attr_index >= len(attrs): - return None + return PatternFailure( + "op_call", + node, + f"{schema.name} takes at most {len(inputs) + len(attrs)} positional " + f"arguments, got {len(node.args)}", + ) param = attrs[attr_index] bound_attrs.add(param.name) children.append( @@ -2165,11 +2179,24 @@ def _schema_children( ) ) for keyword in node.keywords: - if keyword.arg is None or keyword.arg in bound_attrs: - return None + if keyword.arg is None: + return PatternFailure( + "op_call", node, "`**` unpacking is not an authored call form" + ) + if keyword.arg in bound_attrs: + return PatternFailure( + "op_call", + node, + f"attribute {keyword.arg!r} is already bound by a positional argument", + ) param = next((item for item in attrs if item.name == keyword.arg), None) if param is None: - return None + known = ", ".join(item.name for item in attrs) or "none" + return PatternFailure( + "op_call", + node, + f"{schema.name} has no attribute {keyword.arg!r}; its attributes are: {known}", + ) bound_attrs.add(param.name) children.append( AstChild( @@ -2181,11 +2208,17 @@ def _schema_children( ) ) if not variadic and len(positional) < len(inputs): - return None + return PatternFailure( + "op_call", + node, + f"{schema.name} takes {len(inputs)} inputs, got {len(positional)}", + ) return tuple(children) @staticmethod - def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + def _bind( + node: object, context: MatchContext, matched: AstMatch[Any] + ) -> AstMatch[Any] | MatchFailure | None: assert isinstance(node, ast.Call) module_owner = None try: @@ -2195,14 +2228,18 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat if isinstance(callee, runtime.Module): module_owner = callee if node.keywords: - return None + return PatternFailure( + "op_call", node, "a module call takes positional arguments only" + ) try: callee = callee.entry_function() - except ValueError: - return None + except ValueError as error: + return PatternFailure("op_call", node.func, str(error)) if isinstance(callee, runtime.Function): if node.keywords: - return None + return PatternFailure( + "op_call", node, "a function call takes positional arguments only" + ) return dataclasses.replace( matched, pattern_id="call.function", @@ -2230,8 +2267,8 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat if not isinstance(schema, runtime.OpSchema): return None children = CallPattern._schema_children(node, schema, context) - if children is None: - return None + if children is None or isinstance(children, MatchFailure): + return children return dataclasses.replace( matched, pattern_id="call.operation", @@ -2248,7 +2285,9 @@ def construct(match, children, context): if isinstance(variadic_inputs, VariadicInputs): inputs = variadic_inputs.items else: - inputs = tuple(value for name, value in children.items() if name.startswith("input_")) + inputs = tuple( + value for name, value in children.items() if name.startswith("input_") + ) attrs = { name.removeprefix("attr_"): value for name, value in children.items() @@ -2331,6 +2370,60 @@ def construct(match, children, context): } +def _expression_operator_pattern(operator_type: type[ast.AST]) -> ChoicePattern: + return ChoicePattern( + *( + AstNodePattern(node_type) + for node_type in _EXPR_BINARY_KINDS + if issubclass(node_type, operator_type) + ) + ) + + +_CALL_RESULT_RULES: tuple[AstRule[Any], ...] = ( + CallBindingRule(), + CallTypeInferenceRule(), + CallExpectedTypeRule(), +) + + +class MatMulExpressionPattern(ElementPattern): + element_name = "matmul_expression" + syntax = LazyPattern( + lambda: BranchPattern( + "matmul_expression", + AstNodePattern( + ast.BinOp, + FieldPattern("op", AstNodePattern(ast.MatMult)), + FieldPattern( + "left", + ChildPattern("left", ExpressionPattern(), "expression"), + ), + FieldPattern( + "right", + ChildPattern("right", ExpressionPattern(), "expression"), + ), + ), + pattern_id="expression.matmul", + ) + ) + + @staticmethod + def construct(match, children, context): + args = (children["left"], children["right"]) + placeholder_type = next( + (arg.type for arg in args if getattr(arg, "type", None) is not None), + runtime.TensorType.scalar(runtime.DType.f32), + ) + return runtime.Call( + type=placeholder_type, + target=MatMul(), + args=args, + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = _CALL_RESULT_RULES + + class BinaryExpressionPattern(ElementPattern): element_name = "binary_expression" syntax = LazyPattern( @@ -2341,10 +2434,7 @@ class BinaryExpressionPattern(ElementPattern): ast.BinOp, FieldPattern( "op", - PredicatePattern( - "binary-op", - lambda op, context: type(op) in _EXPR_BINARY_KINDS, - ), + _expression_operator_pattern(ast.operator), ), CapturePattern( "kind", @@ -2367,12 +2457,7 @@ class BinaryExpressionPattern(ElementPattern): ast.Compare, FieldPattern( "ops", - SequencePattern( - PredicatePattern( - "comparison-op", - lambda op, context: type(op) in _EXPR_BINARY_KINDS, - ) - ), + SequencePattern(_expression_operator_pattern(ast.cmpop)), ), CapturePattern( "kind", @@ -2395,10 +2480,7 @@ class BinaryExpressionPattern(ElementPattern): ast.BoolOp, FieldPattern( "op", - PredicatePattern( - "boolean-op", - lambda op, context: type(op) in _EXPR_BINARY_KINDS, - ), + _expression_operator_pattern(ast.boolop), ), CapturePattern( "kind", @@ -2419,10 +2501,15 @@ class BinaryExpressionPattern(ElementPattern): @staticmethod def construct(match, children, context): + args = (children["left"], children["right"]) + placeholder_type = next( + (arg.type for arg in args if getattr(arg, "type", None) is not None), + runtime.TensorType.scalar(runtime.DType.f32), + ) return runtime.Call( - type=children["left"].type, + type=placeholder_type, target=runtime.Binary(kind=runtime.BinaryKind[match.captures["kind"]]), - args=(children["left"], children["right"]), + args=args, ) RULES: ClassVar[tuple[AstRule[Any], ...]] = ( @@ -2462,9 +2549,10 @@ class UnaryExpressionPattern(ElementPattern): @staticmethod def construct(match, children, context): operand = children["operand"] + target = runtime.Unary(kind=runtime.UnaryKind[match.captures["kind"]]) return runtime.Call( type=operand.type, - target=runtime.Unary(kind=runtime.UnaryKind[match.captures["kind"]]), + target=target, args=(operand,), ) @@ -2573,9 +2661,7 @@ class IndexEndpointPattern(ElementPattern): @dataclass(frozen=True) class TileWindowSliceBoundRule: - STATEMENT: ClassVar[str] = ( - "A tile window cannot be used as a slice bound." - ) + 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 ( @@ -2789,7 +2875,9 @@ class MeshCoordinatePattern(ElementPattern): ) @staticmethod - def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + def _bind( + node: object, context: MatchContext, matched: AstMatch[Any] + ) -> AstMatch[Any] | MatchFailure | None: assert isinstance(node, ast.Attribute) assert isinstance(node.value, ast.Name) if context.function is None: @@ -2806,7 +2894,12 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat if candidate < len(mesh.layout.shape): axis = candidate if axis is None: - return None + named = ", ".join(mesh.names) + return PatternFailure( + "mesh_coordinate", + node, + f"mesh {node.value.id!r} has no axis {node.attr!r}; its axes are: {named}", + ) return dataclasses.replace( matched, pattern_id="expression.mesh_coordinate", @@ -2865,6 +2958,7 @@ class ExpressionPattern(ElementPattern): CallPattern(), LaunchPattern(), SubscriptExpressionPattern(), + MatMulExpressionPattern(), BinaryExpressionPattern(), UnaryExpressionPattern(), MeshCoordinatePattern(), @@ -2954,7 +3048,9 @@ class MeshContextPattern(ElementPattern): ) @staticmethod - def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + def _bind( + node: object, context: MatchContext, matched: AstMatch[Any] + ) -> AstMatch[Any] | MatchFailure | None: assert isinstance(node, ast.Call) if not node.args or not isinstance(node.args[0], ast.Tuple): return None @@ -2964,14 +3060,25 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMat names_node = keywords.get("names") if positional: if layout_node is not None: - return None + return PatternFailure( + "mesh", node, "`layout` is already bound by a positional argument" + ) layout_node = positional.pop(0) if positional: if names_node is not None: - return None + return PatternFailure( + "mesh", node, "`names` is already bound by a positional argument" + ) names_node = positional.pop(0) - if positional or layout_node is None: - return None + if positional: + return PatternFailure( + "mesh", + node, + f"a mesh takes topologies, layout, and names; got " + f"{len(node.args)} positional arguments", + ) + if layout_node is None: + return PatternFailure("mesh", node, "a mesh requires a `layout`") children = [ AstChild( "topology_names", @@ -3155,14 +3262,28 @@ def _bind(node: object, context: MatchContext, matched: AstMatch[Any]): if not isinstance(node.func, ast.Name) or node.func.id != "launch": return None if not node.args: - return None + return PatternFailure("launch", node, "launch(...) requires a callee") keywords = {keyword.arg: keyword.value for keyword in node.keywords} if any(name is None for name in keywords): - return None - if "grid" not in keywords or "block" not in keywords: - return None - if set(keywords) - {"grid", "block", "cluster", "dynamic_smem", "stream", "attrs"}: - return None + return PatternFailure("launch", node, "`**` unpacking is not an authored launch form") + missing = [name for name in ("grid", "block") if name not in keywords] + if missing: + return PatternFailure("launch", node, f"launch(...) requires {' and '.join(missing)}") + unknown = set(keywords) - { + "grid", + "block", + "cluster", + "dynamic_smem", + "stream", + "attrs", + } + if unknown: + return PatternFailure( + "launch", + node, + f"launch(...) has no option {sorted(unknown)!r}; it takes grid, block, " + f"cluster, dynamic_smem, stream, attrs", + ) children = [ AstChild("callee", StaticValuePattern(), node.args[0], "launch_callee"), *( @@ -3305,16 +3426,40 @@ class LoopIteratorPattern(ElementPattern): BranchPattern( "tile", AstNodePattern( - ast.Name, - FieldPattern("id", LiteralPattern("tile")), + ast.Call, + FieldPattern( + "func", + AstNodePattern(ast.Name, FieldPattern("id", LiteralPattern("tile"))), + ), + FieldPattern("keywords", SequencePattern()), + FieldPattern( + "args", + SequencePattern(AstNodePattern(ast.expr), AstNodePattern(ast.expr)), + ), ), pattern_id="loop.iterator.tile", ), BranchPattern( "range", AstNodePattern( - ast.Name, - FieldPattern("id", LiteralPattern("range")), + ast.Call, + FieldPattern( + "func", + AstNodePattern(ast.Name, FieldPattern("id", LiteralPattern("range"))), + ), + FieldPattern("keywords", SequencePattern()), + FieldPattern( + "args", + ChoicePattern( + SequencePattern(AstNodePattern(ast.expr)), + SequencePattern(AstNodePattern(ast.expr), AstNodePattern(ast.expr)), + SequencePattern( + AstNodePattern(ast.expr), + AstNodePattern(ast.expr), + AstNodePattern(ast.expr), + ), + ), + ), ), pattern_id="loop.iterator.range", ), @@ -3335,18 +3480,7 @@ class LoopHeaderPattern(ElementPattern): AstNodePattern( ast.For, FieldPattern("target", AstNodePattern(ast.Name)), - FieldPattern( - "iter", - AstNodePattern( - ast.Call, - FieldPattern( - "func", - LoopIteratorPattern(), - ), - FieldPattern("keywords", RepeatPattern(AstNodePattern(ast.keyword))), - FieldPattern("args", RepeatPattern(AstNodePattern(ast.expr), minimum=1)), - ), - ), + FieldPattern("iter", LoopIteratorPattern()), FieldPattern( "body", ChildPattern( @@ -3361,10 +3495,46 @@ class LoopHeaderPattern(ElementPattern): ) ) + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: + """Name the invalid iterator before the shape-exact syntax rejects it. + + The syntax encodes each iterator's arity so the generated grammar shows + what the parser accepts. A shape mismatch alone would report only that + the pattern did not match, so the specific reason is stated here first. + """ + if ( + isinstance(node, ast.For) + and isinstance(node.iter, ast.Call) + and 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", + node.iter, + "tile()/range() does not accept keyword args (positional-only at the IR level)", + ) + if (kind == "tile" and count != 2) or (kind == "range" and count not in {1, 2, 3}): + if kind == "tile" and count == 1: + detail = "tile(extent) is not supported; use range(extent)" + elif kind == "tile": + detail = f"tile() takes 2 arguments (extent, step), got {count}" + else: + detail = f"range() takes 1 to 3 arguments, got {count}" + return PatternFailure("loop_header", node.iter, detail) + return super().match(node, context) + @staticmethod def _bind( node: object, context: MatchContext, matched: AstMatch[Any] - ) -> AstMatch[Any] | PatternFailure | None: + ) -> AstMatch[Any] | MatchFailure | None: assert isinstance(node, ast.For) assert isinstance(node.target, ast.Name) assert isinstance(node.iter, ast.Call) @@ -3633,7 +3803,11 @@ def construct(match, children, context): raise ParseError.from_node( match.node, context, "tuple assignment arity does not match TupleType" ) - schema = getattr(type(value.target), "_op_schema", None) if isinstance(value, runtime.Call) else None + schema = ( + getattr(type(value.target), "_op_schema", None) + if isinstance(value, runtime.Call) + else None + ) parent_name = getattr(schema, "name", None) or ", ".join(names) value = replace_metadata(value, BindingMetadata(parent_name)) for index, name in enumerate(names): @@ -3875,7 +4049,6 @@ def construct(match, children, context): return None raise ParseError.from_node(match.node, context, "HIR body must end with return") - def fold(index: int) -> list[object]: output: list[object] = [] while index < len(values): diff --git a/src/tilefoundry/parser/spec.py b/src/tilefoundry/parser/spec.py index a11e4c27..039ab7a7 100644 --- a/src/tilefoundry/parser/spec.py +++ b/src/tilefoundry/parser/spec.py @@ -129,10 +129,7 @@ def _collect_rule_rows(root: ElementPattern[Any]) -> tuple[RuleRow, ...]: def _collect_module_rule_rows() -> tuple[RuleRow, ...]: - rows = [ - _row("module", "module_function", rule) - for rule in ModuleBuildContext.FUNCTION_RULES - ] + rows = [_row("module", "module_function", rule) for rule in ModuleBuildContext.FUNCTION_RULES] rows.extend( _row("module", "module_finalization", rule) for rule in ModuleBuildContext.FINALIZATION_RULES @@ -141,21 +138,28 @@ def _collect_module_rule_rows() -> tuple[RuleRow, ...]: 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]] = {} + """Render one row per rule, listing every element and situation it governs. + + A rule shared by several elements states one thing, so it reads as one row. + Repeating it once per owner says nothing new and buries the rules that are + actually distinct. + """ + owners: dict[tuple[str, str, str], set[str]] = {} + situations: dict[tuple[str, str, str], set[str]] = {} for row in rows: - key = (row.owner, row.rule, row.statement, row.source) + key = (row.rule, row.statement, row.source) + owners.setdefault(key, set()).add(row.owner) situations.setdefault(key, set()).add(row.situation) return tuple( sorted( RuleRow( - owner=owner, - situation=", ".join(sorted(rule_situations)), - rule=rule, - statement=statement, - source=source, + owner=", ".join(sorted(owners[key])), + situation=", ".join(sorted(situations[key])), + rule=key[0], + statement=key[1], + source=key[2], ) - for (owner, rule, statement, source), rule_situations in situations.items() + for key in owners ) ) diff --git a/src/tilefoundry/passes/transforms/hir_to_tir.py b/src/tilefoundry/passes/transforms/hir_to_tir.py index b19e1fb7..0a7d684a 100644 --- a/src/tilefoundry/passes/transforms/hir_to_tir.py +++ b/src/tilefoundry/passes/transforms/hir_to_tir.py @@ -14,12 +14,6 @@ from tilefoundry.ir.core import Call, Constant, Expr, Tuple, Var from tilefoundry.ir.core.module import Module from tilefoundry.ir.core.pattern import DimVarRangePat, locate_dim_var -from tilefoundry.ir.hir.cuda.nn.mma import ( - Mma_SM80_16x8x16 as HirMmaSM80_16x8x16, -) -from tilefoundry.ir.hir.cuda.nn.mma import ( - Wgmma_SM90_64x128x16 as HirWgmma_SM90, -) from tilefoundry.ir.hir.function import Function as HirFunction from tilefoundry.ir.hir.grid_region import GridRegionExpr from tilefoundry.ir.hir.math.binary import Binary as HirBinary @@ -1016,16 +1010,6 @@ def _lower_reshard(ctx: "_Lowerer", target, expr) -> Var: -@register_hir_lowering(HirMmaSM80_16x8x16) -@register_hir_lowering(HirWgmma_SM90) -def _lower_mma(ctx: "_Lowerer", target, expr) -> Var: - name = type(target).__name__ - raise ValueError( - f"HirToTirPass: {name} HIR compile route is unsupported; use the " - "independent handwritten TIR MMA atom/runtime surface" - ) - - @register_hir_lowering(HirReLU) def _lower_relu(ctx: "_Lowerer", target, expr) -> Var: x = ctx.lower(expr.args[0]) diff --git a/src/tilefoundry/visitor_registry/op_cost.py b/src/tilefoundry/visitor_registry/op_cost.py index 80d6c339..3f460955 100644 --- a/src/tilefoundry/visitor_registry/op_cost.py +++ b/src/tilefoundry/visitor_registry/op_cost.py @@ -13,7 +13,6 @@ from tilefoundry.ir.core import Call from tilefoundry.ir.core.kinds import BinaryKind, UnaryKind -from tilefoundry.ir.hir.cuda.nn.mma import Mma_SM80_16x8x16, Wgmma_SM90_64x128x16 from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.math.clamp import Clamp from tilefoundry.ir.hir.math.softplus import Softplus @@ -137,20 +136,6 @@ def _matmul(call: Call, ctx: CostContext) -> Cost: return Cost({lhs.dtype: flops}, _traffic((lhs, rhs), output)) -@register_cost_evaluator(Mma_SM80_16x8x16) -@register_cost_evaluator(Wgmma_SM90_64x128x16) -def _mma(call: Call, ctx: CostContext) -> Cost: - a, b = _input_types(call, ctx) - output = _output_type(call, ctx) - if not all(isinstance(type_, TensorType) for type_ in (a, b, output)): - raise ValueError("MMA cost requires tensor inputs and output") - if isinstance(call.target, Mma_SM80_16x8x16): - m, n, k = 16, 8, 16 - else: - m, n, k = 64, 128, 16 - return Cost({a.dtype: 2 * m * n * k}, _traffic((a, b), output)) - - @register_cost_evaluator(Conv2D) def _conv2d(call: Call, ctx: CostContext) -> Cost: input_, weight, bias = _input_types(call, ctx) diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 9220edb2..ff339aa9 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -14,7 +14,6 @@ import pytest -from tests.models.corpus import placed_cases, placed_fixture_roots from tilefoundry import func, module from tilefoundry.analysis import ( ComputeCostMetadata, @@ -278,46 +277,6 @@ def test_a_program_whose_buffers_have_nowhere_to_sit_is_refused() -> None: ] -def test_the_placed_inventory_takes_the_whole_directory() -> None: - """What is asked about is what is there, and what is left out is named. - - The inventory is walked out of the fixture package rather than listed, so a - fixture added to it joins instead of quietly escaping. This pins the walk -- - every root found, a Module reached as somebody's child not counted as one, - and the only roots left out being those whose machine runs no CTAs. Excluding - one for any other reason is the allowlist this prevents, and would show up - here as a number that moved. What those cases answer is asked where the - analyses run; a count is not an answer, so this runs nothing. - """ - roots = placed_fixture_roots() - assert len(roots) == 29 - - outside = [] - for file, name, published in roots: - root = published - try: - root.resolve_target() - except Exception: # noqa: BLE001 -- a root that names no machine - root = replace(root, target=CudaTarget("nvidia.h200_sxm")) - if "cta" not in root.resolve_target().topology_levels: - outside.append(f"{file}.{name}") - assert outside == [ - "fused_twin.Fused", - "nested_twin.Nested", - "square_cpu.Mine", - "square_cpu_runtime.Mine", - "square_twin.Model", - "weighted_twin.WeightedRoot", - ] - - cases = placed_cases() - assert len(roots) - len(outside) == 23, "the roots that answer for CTAs" - assert len({case.id.rsplit("[", 1)[0] for case in cases}) == 29, ( - "one row per selector those roots expose, not one per root" - ) - assert len(cases) == 32, "and one case per stated set of sizes" - - def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> None: """A hole inside a number reads as a program that does less than it does. diff --git a/tests/fixtures/placed/mma_tile.py b/tests/fixtures/placed/mma_tile.py deleted file mode 100644 index 0aaddb74..00000000 --- a/tests/fixtures/placed/mma_tile.py +++ /dev/null @@ -1,54 +0,0 @@ -"""One SM80 16x8x16 BF16 MMA tile with an FP32 result, on one CTA.""" - -from tilefoundry import func, module -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology -from tilefoundry.target import CudaTarget - -_THREAD = Topology("thread", 32) -_THREAD_MESH = Mesh( - topologies=(_THREAD,), - layout=Layout(shape=(4, 8), strides=(1, 4)), - names=("x", "y"), -) - -_A_FRAG = ShardLayout( - layout=Layout(shape=(2, 4, 2, 8, 2), strides=(1, 2, 8, 16, 128)), - attrs=(Split(1), Split(3)), - mesh=_THREAD_MESH, -) -_B_FRAG = ShardLayout( - layout=Layout(shape=(8, 2, 4, 2), strides=(1, 8, 16, 64)), - attrs=(Split(2), Split(0)), - mesh=_THREAD_MESH, -) -_C_FRAG = ShardLayout( - layout=Layout(shape=(2, 4, 8, 2), strides=(1, 2, 8, 64)), - attrs=(Split(1), Split(2)), - mesh=_THREAD_MESH, -) - - -@module( - entry="matmul_16x8x16", - target=CudaTarget("nvidia.h200_sxm"), - topologies=(Topology("cta", 1), _THREAD), -) -class MatmulModule: - @func - def matmul_16x8x16( - a: Tensor[(16, 16), "bf16"], - b: Tensor[(16, 8), "bf16"], - ) -> Tensor[(16, 8), "f32"]: - with Mesh(("cta",), (1,), ("tile",)) as _cta: - a_frag = reshard(a, layout=_A_FRAG, storage="rmem") - b_frag = reshard(b, layout=_B_FRAG, storage="rmem") - c_frag = mma_sm80_16x8x16( - a_frag, - b_frag, - dtype_a="bf16", - dtype_b="bf16", - dtype_acc="f32", - ) - return reshard(c_frag, layout=_C_FRAG, storage="gmem") diff --git a/tests/ops/test_mma.py b/tests/ops/test_mma.py deleted file mode 100644 index e07e4e52..00000000 --- a/tests/ops/test_mma.py +++ /dev/null @@ -1,303 +0,0 @@ -"""HIR CUDA MMA logical value, cost, fragment, and rejection contracts.""" - -from __future__ import annotations - -from dataclasses import replace - -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, TypeInferCase, infer_call, run_typeinfer_case -from tilefoundry.ir.hir.cuda.nn.mma import Mma_SM80_16x8x16, Wgmma_SM90_64x128x16 -from tilefoundry.ir.tir.cuda.nn.mma import ( - SM80_16x8x16_F32BF16BF16F32_TN, - make_atom, -) -from tilefoundry.ir.types import ( - DType, - TensorType, - make_shard_tensor_type, - make_tensor_type, -) -from tilefoundry.ir.types.shard import Broadcast, Layout, Split, make_mesh -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.visitor_registry.contexts import TrafficBytes - -_BF = DType.bf16 -_F32 = DType.f32 -_RMEM = StorageKind.RMEM -_ATOM = make_atom(SM80_16x8x16_F32BF16BF16F32_TN) -_SM80 = Mma_SM80_16x8x16(dtype_a=_BF, dtype_b=_BF, dtype_acc=_F32) -_WGMMA = Wgmma_SM90_64x128x16(dtype_a=_BF, dtype_b=_BF, dtype_acc=_F32) - - -CASES = [ - TypeInferCase( - "sm80_logical", - _SM80, - ( - make_tensor_type((16, 16), _BF, storage=_RMEM), - make_tensor_type((16, 8), _BF, storage=_RMEM), - ), - make_tensor_type((16, 8), _F32, storage=_RMEM), - ), - TypeInferCase( - "wgmma_logical", - _WGMMA, - ( - make_tensor_type((64, 16), _BF, storage=_RMEM), - make_tensor_type((16, 128), _BF, storage=_RMEM), - ), - make_tensor_type((64, 128), _F32, storage=_RMEM), - ), - TypeInferCase( - "wrong_a_shape", - _SM80, - (make_tensor_type((15, 16), _BF), make_tensor_type((16, 8), _BF)), - ExpectedError(match=r"Mma_SM80_16x8x16: a shape"), - ), - TypeInferCase( - "wrong_b_shape", - _SM80, - (make_tensor_type((16, 16), _BF), make_tensor_type((8, 16), _BF)), - ExpectedError(match=r"Mma_SM80_16x8x16: b shape"), - ), - TypeInferCase( - "dtype_a_disagrees", - _SM80, - (make_tensor_type((16, 16), DType.f16), make_tensor_type((16, 8), _BF)), - ExpectedError(match=r"Mma_SM80_16x8x16: dtype_a"), - ), - TypeInferCase( - "unsupported_accumulator", - Mma_SM80_16x8x16(dtype_a=_BF, dtype_b=_BF, dtype_acc=DType.i32), - (make_tensor_type((16, 16), _BF), make_tensor_type((16, 8), _BF)), - ExpectedError(match=r"Mma_SM80_16x8x16: dtype_acc combination"), - ), - TypeInferCase( - "bad_a_orientation", - Mma_SM80_16x8x16( - dtype_a=_BF, dtype_b=_BF, dtype_acc=_F32, a_layout="K" - ), - (make_tensor_type((16, 16), _BF), make_tensor_type((16, 8), _BF)), - ExpectedError(match=r"Mma_SM80_16x8x16: a_layout"), - ), - TypeInferCase( - "bad_b_orientation", - Mma_SM80_16x8x16( - dtype_a=_BF, dtype_b=_BF, dtype_acc=_F32, b_layout="K" - ), - (make_tensor_type((16, 16), _BF), make_tensor_type((16, 8), _BF)), - ExpectedError(match=r"Mma_SM80_16x8x16: b_layout"), - ), -] - - -@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name) -def test_mma_typeinfer(case): - run_typeinfer_case(case) - - -def test_plain_layout_describes_the_logical_accumulator_result(): - a = make_tensor_type( - (16, 16), _BF, storage=_RMEM, layout=Layout((16, 16), (16, 1)) - ) - b = make_tensor_type( - (16, 8), _BF, storage=_RMEM, layout=Layout((16, 8), (8, 1)) - ) - - result = infer_call(_SM80, a, b) - - assert result.layout == Layout(shape=(16, 8), strides=(8, 1)) - - -def test_fragment_mesh_coordinate_names_do_not_affect_compatibility(): - a_mesh = replace(_ATOM.required_scope, names=("a_m", "a_k")) - b_mesh = replace(_ATOM.required_scope, names=("b_k", "b_n")) - a_layout = replace(_ATOM.A, mesh=a_mesh) - b_layout = replace(_ATOM.B, mesh=b_mesh) - expected_c = replace(_ATOM.C, mesh=a_mesh) - a = TensorType((16, 16), _BF, a_layout, _RMEM) - b = TensorType((16, 8), _BF, b_layout, _RMEM) - - result = infer_call(_SM80, a, b) - - assert result.layout == expected_c - assert result.layout is not _ATOM.C - assert result.storage is _RMEM - - -@pytest.mark.parametrize( - ("name", "op", "a", "b", "match"), - [ - pytest.param( - "mismatched_b_fragment", - _SM80, - TensorType((16, 16), _BF, _ATOM.A, _RMEM), - TensorType( - (16, 8), - _BF, - replace( - _ATOM.B, - layout=Layout(_ATOM.B.layout.shape, (2, 8, 16, 64)), - ), - _RMEM, - ), - r"input 1.*known SM80 B fragment layout.*Reshard.*materialize-to-RMEM", - id="mismatched_b_fragment", - ), - pytest.param( - "different_physical_mesh_layout", - _SM80, - TensorType((16, 16), _BF, _ATOM.A, _RMEM), - TensorType( - (16, 8), - _BF, - replace( - _ATOM.B, - mesh=replace( - _ATOM.B.mesh, - layout=Layout(shape=(8, 4), strides=(1, 8)), - ), - ), - _RMEM, - ), - r"input 1.*known SM80 B fragment layout.*Reshard.*materialize-to-RMEM", - id="different_physical_mesh_layout", - ), - pytest.param( - "non_rmem_fragment", - _SM80, - TensorType((16, 16), _BF, _ATOM.A, StorageKind.GMEM), - TensorType((16, 8), _BF, _ATOM.B, StorageKind.GMEM), - r"input 0.*not RMEM.*Reshard.*materialize-to-RMEM", - id="non_rmem_fragment", - ), - pytest.param( - "wgmma_shard_claim", - _WGMMA, - make_shard_tensor_type( - (64, 16), _BF, storage=_RMEM, mesh=make_mesh((4,)), attrs=(Split(0),) - ), - make_tensor_type((16, 128), _BF, storage=_RMEM), - r"input 0.*unrepresentable WGMMA ShardLayout.*Reshard", - id="wgmma_shard_claim", - ), - ], -) -def test_unrepresentable_fragment_claims_fail(name, op, a, b, match): - run_typeinfer_case(TypeInferCase(name, op, (a, b), ExpectedError(match=match))) - - -def test_fully_broadcast_inputs_do_not_pin_an_output_mesh(): - a = make_shard_tensor_type( - (16, 16), - _BF, - storage=_RMEM, - mesh=make_mesh((2,)), - attrs=(Broadcast(),), - ) - b = make_shard_tensor_type( - (16, 8), - _BF, - storage=_RMEM, - mesh=make_mesh((4,)), - attrs=(Broadcast(),), - ) - - result = infer_call(_SM80, a, b) - - assert result.layout is None - - -_SM80_A = torch.arange(16 * 16, dtype=torch.bfloat16).reshape(16, 16) / 32 -_SM80_B = torch.arange(16 * 8, dtype=torch.bfloat16).reshape(16, 8) / 16 -_WGMMA_A = torch.arange(64 * 16, dtype=torch.bfloat16).reshape(64, 16) / 64 -_WGMMA_B = torch.arange(16 * 128, dtype=torch.bfloat16).reshape(16, 128) / 128 - - -@pytest.mark.parametrize( - "case", - [ - EvalCase( - "sm80_a_matmul_b", - _SM80, - (_SM80_A, _SM80_B), - _SM80_A.float() @ _SM80_B.float(), - ), - EvalCase( - "wgmma_a_matmul_b", - _WGMMA, - (_WGMMA_A, _WGMMA_B), - _WGMMA_A.float() @ _WGMMA_B.float(), - ), - EvalCase( - "orientation_is_encoding_not_logical_transpose", - Mma_SM80_16x8x16( - dtype_a=_BF, - dtype_b=_BF, - dtype_acc=_F32, - a_layout="N", - b_layout="T", - ), - (_SM80_A, _SM80_B), - _SM80_A.float() @ _SM80_B.float(), - ), - ], - ids=lambda case: case.name, -) -def test_mma_evaluator_is_the_two_input_product(case): - run_eval_case(case) - - -@pytest.mark.parametrize( - "case", - [ - CostCase( - "sm80_logical", - _SM80, - (make_tensor_type((16, 16), _BF), make_tensor_type((16, 8), _BF)), - flops={_BF: 2 * 16 * 8 * 16}, - traffic=( - TrafficBytes(read=512), - TrafficBytes(read=256), - TrafficBytes(write=512), - ), - ), - CostCase( - "wgmma_logical", - _WGMMA, - ( - make_tensor_type((64, 16), _BF), - make_tensor_type((16, 128), _BF), - ), - flops={_BF: 2 * 64 * 128 * 16}, - traffic=( - TrafficBytes(read=2_048), - TrafficBytes(read=4_096), - TrafficBytes(write=32_768), - ), - ), - CostCase( - "sm80_thread_fragment", - _SM80, - ( - TensorType((16, 16), _BF, _ATOM.A, _RMEM), - TensorType((16, 8), _BF, _ATOM.B, _RMEM), - ), - flops={_BF: 2 * 16 * 8 * 16}, - traffic=( - TrafficBytes(read=16), - TrafficBytes(read=8), - TrafficBytes(write=16), - ), - level="thread", - topologies=_ATOM.required_scope.topologies, - ), - ], - ids=lambda case: case.name, -) -def test_mma_cost(case): - run_cost_case(case) diff --git a/tests/parser/test_calls.py b/tests/parser/test_calls.py index 13224971..043070f2 100644 --- a/tests/parser/test_calls.py +++ b/tests/parser/test_calls.py @@ -2,20 +2,25 @@ from __future__ import annotations +import re from typing import get_args, get_origin import pytest -from tilefoundry import func -from tilefoundry.dsl import Tensor, tf +from tilefoundry import func, module, prim_func +from tilefoundry.dsl import Mesh, Tensor, tf from tilefoundry.ir.core import Call, Constant, Tuple, VerifyError from tilefoundry.ir.core.pattern import Tensor as TensorPattern from tilefoundry.ir.hir.nn.matmul import MatMul +from tilefoundry.ir.hir.sharding.reshard import Reshard 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.ir.types import DType +from tilefoundry.ir.types.shard import Topology from tilefoundry.parser import ParseError +from tilefoundry.target import CpuTarget, CudaTarget def test_matmul_layout_literals_are_parser_checked() -> None: @@ -23,9 +28,7 @@ def test_matmul_layout_literals_are_parser_checked() -> None: 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"]: + 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) @@ -50,15 +53,11 @@ def test_variadic_list_and_tuple_literals_flatten_to_call_args() -> None: assert get_args(annotation) == (TensorPattern,) @func - def from_list( - a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"] - ) -> Tensor[(2, 4), "f32"]: + 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"]: + 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): @@ -69,9 +68,7 @@ def from_tuple( 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"]: + 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) @@ -112,9 +109,7 @@ 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"]: + 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"): @@ -136,17 +131,13 @@ def generator(x: Tensor[(2, 4), "f32"]) -> Tensor[(2, 4), "f32"]: 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"]: + 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"]: + 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) @@ -155,9 +146,7 @@ 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"]: + 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) @@ -178,9 +167,7 @@ def test_unsupported_list_comprehension_shapes_are_named(program: str, message: @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 - ) + return tf.concat([x for first in range(1) for second in range(2)], axis=0) elif program == "target": @@ -205,3 +192,171 @@ def rejected(x: Tensor[(1, 4), "f32"]) -> Tensor[(1, 4), "f32"]: @func def rejected(x: Tensor[(1, 4), "f32"]) -> Tensor[(1, 4), "f32"]: return tf.concat([x for index in range(x)], axis=0) + + +def test_a_python_float_literal_is_f32_like_any_other() -> None: + """A literal carries the dtype it is written with; nothing adapts it.""" + + @func + def add_literal(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: + return x + 1.0 + + assert isinstance(add_literal.body, Call) + constant = next(arg for arg in add_literal.body.args if isinstance(arg, Constant)) + assert constant.type.dtype is DType.f32 + + +def test_a_mismatched_literal_is_rejected_and_names_the_cast() -> None: + """The remedy is the one the author writes, not one the parser guesses.""" + with pytest.raises(VerifyError, match=re.escape("tf.cast(, dtype='bf16')")): + + @func + def add_literal(x: Tensor[(8,), "bf16"]) -> Tensor[(8,), "bf16"]: + return x + 1.0 + + +def test_an_explicit_cast_gives_a_literal_the_operand_dtype() -> None: + @func + def scale(x: Tensor[(8,), "bf16"]) -> Tensor[(8,), "bf16"]: + return x * tf.cast(-0.5, dtype="bf16") + + assert isinstance(scale.body, Call) + assert scale.body.type.dtype is DType.bf16 + + +def test_a_python_integer_operand_is_not_adopted() -> None: + with pytest.raises(VerifyError, match="dtype mismatch \\(bf16 vs i64\\)"): + + @func + def add_integer(x: Tensor[(8,), "bf16"]) -> Tensor[(8,), "bf16"]: + return x + 1 + + +def test_two_tensor_operands_are_never_promoted() -> None: + with pytest.raises(VerifyError, match="dtype mismatch \\(bf16 vs f32\\)"): + + @func + def mixed(x: Tensor[(8,), "bf16"]) -> Tensor[(8,), "f32"]: + return x + tf.cast(x, "f32") + + +def test_the_at_operator_builds_a_matmul() -> None: + @func + def infix(a: Tensor[(2, 4), "bf16"], b: Tensor[(4, 3), "bf16"]) -> Tensor[(2, 3), "bf16"]: + return a @ b + + @func + def spelled(a: Tensor[(2, 4), "bf16"], b: Tensor[(4, 3), "bf16"]) -> Tensor[(2, 3), "bf16"]: + return tf.matmul(a, b) + + assert isinstance(infix.body, Call) + assert isinstance(infix.body.target, MatMul) + assert infix.body.type == spelled.body.type + + +def test_an_unresolved_callee_is_named() -> None: + with pytest.raises(VerifyError, match="unsupported call 'tf.no_such_op'"): + + @func + def missing(x: Tensor[(8,), "bf16"]) -> Tensor[(8,), "bf16"]: + return tf.no_such_op(x, x) + + +def test_placement_at_and_matmul_at_coexist_in_one_function() -> None: + """The two spellings of ``@`` are told apart by position, not by shape.""" + + @module( + entry="placed_matmul", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 2),), + ) + class PlacedMatMul: + @func + def placed_matmul(x: Tensor[(2, 4), "bf16"], b: Tensor[(4, 3), "bf16"]): + with Mesh(("cta",), (2,), ("tile",)) as mesh: + x_local = tf.reshard(x, (2 @ mesh.tile, 4), "rmem") + b_local = tf.reshard(b, (4, 3), "rmem") + return x_local @ b_local + + body = PlacedMatMul.entry_function().body + assert isinstance(body, Call) + assert isinstance(body.target, MatMul) + assert any(isinstance(argument.target, Reshard) for argument in body.args) + + +@pytest.mark.parametrize( + ("program", "message"), + [ + ( + "misspelled", + "matmul has no attribute 'a_layoutt'; its attributes are: a_layout, b_layout", + ), + ("few", "matmul takes 2 inputs, got 1"), + ("many", "matmul takes at most 4 positional arguments, got 5"), + ("twice", "attribute 'a_layout' is already bound by a positional argument"), + ], +) +def test_a_refused_call_states_what_was_wrong_with_it(program: str, message: str) -> None: + """A call the parser recognized and refused names its own reason, not a shape mismatch.""" + with pytest.raises(ParseError, match=re.escape(message)): + if program == "misspelled": + + @func + def refused(a: Tensor[(2, 4), "bf16"], b: Tensor[(4, 3), "bf16"]): + return tf.matmul(a, b, a_layoutt="MK") + + elif program == "few": + + @func + def refused(a: Tensor[(2, 4), "bf16"], b: Tensor[(4, 3), "bf16"]): + return tf.matmul(a) + + elif program == "many": + + @func + def refused(a: Tensor[(2, 4), "bf16"], b: Tensor[(4, 3), "bf16"]): + return tf.matmul(a, b, "MK", "KN", "KN") + + else: + + @func + def refused(a: Tensor[(2, 4), "bf16"], b: Tensor[(4, 3), "bf16"]): + return tf.matmul(a, b, "MK", a_layout="MK") + + +def test_a_refused_call_reports_one_reason_rather_than_every_alternative() -> None: + """A choice drops the alternatives that stated nothing instead of listing them.""" + with pytest.raises(ParseError) as raised: + + @func + def refused(a: Tensor[(2, 4), "bf16"], b: Tensor[(4, 3), "bf16"]): + return tf.matmul(a) + + message = str(raised.value) + assert "pattern did not match" not in message + assert "no alternative matched" not in message + assert message.count("matmul takes 2 inputs, got 1") == 1 + + +def test_a_call_nobody_claims_is_named_without_taking_a_turn() -> None: + """Position is the whole claim for the report of last resort. + + The report runs only after every alternative has declined, so it names any + callee — inside the op namespace or not — while a call another pattern owns + never reaches it. ``launch(...)`` is such a call and stays parseable. + """ + with pytest.raises(ParseError, match=re.escape("unsupported call 'no_such_helper'")): + + @func + def refused(a: Tensor[(2, 4), "bf16"]) -> Tensor[(2, 4), "bf16"]: + return no_such_helper(a) # noqa: F821 + + @func(topologies=(Topology("cta", 2),)) + def device(a: Tensor[(2, 4), "f32"]) -> Tensor[(2, 4), "f32"]: + return tf.mul(a, a) + + @prim_func(target=CpuTarget()) + def host(a: Tensor[(2, 4), "f32"], out: Tensor[(2, 4), "f32"]): + launch(device, a, out, grid=(2, 1, 1), block=(1, 1, 1)) # noqa: F821 + + assert host.body is not None diff --git a/tests/parser/test_functions.py b/tests/parser/test_functions.py index a1d2b246..e1e4bc43 100644 --- a/tests/parser/test_functions.py +++ b/tests/parser/test_functions.py @@ -2,10 +2,12 @@ from __future__ import annotations +import re + import pytest from tilefoundry import func -from tilefoundry.dsl import Mesh, Tensor, Topology +from tilefoundry.dsl import Mesh, Tensor, Topology, tf from tilefoundry.inspection import as_script from tilefoundry.parser import ParseError from tilefoundry.target import CudaTarget @@ -28,3 +30,56 @@ 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 + + +@pytest.mark.parametrize( + ("iterator", "expected"), + ( + ("tile(10)", "tile(extent) is not supported; use range(extent)"), + ("tile(1, 2, 3)", "tile() takes 2 arguments (extent, step), got 3"), + ("range(1, 2, 3, 4)", "range() takes 1 to 3 arguments, got 4"), + ("steps(1, 2)", "loop iterator must be tile(...) or range(...)"), + ), +) +def test_a_loop_iterator_states_why_its_arity_is_invalid(iterator: str, expected: str) -> None: + """The shape-exact syntax must not reduce these to a bare match failure. + + Encoding each iterator's arity is what lets the generated grammar show the + accepted forms, so the reason is stated before the shape rejects. + """ + with pytest.raises(ParseError, match=re.escape(expected)): + if iterator == "tile(10)": + + @func + def looping(x: Tensor[(10, 4), "f32"], seed: Tensor[(4, 4), "f32"]): + out = tf.add(seed, seed) + for row in tile(10): # noqa: F821 + out = tf.add(x[row, :], seed) + return out + + elif iterator == "tile(1, 2, 3)": + + @func + def looping(x: Tensor[(10, 4), "f32"], seed: Tensor[(4, 4), "f32"]): + out = tf.add(seed, seed) + for row in tile(1, 2, 3): # noqa: F821 + out = tf.add(x[row, :], seed) + return out + + elif iterator == "range(1, 2, 3, 4)": + + @func + def looping(x: Tensor[(10, 4), "f32"], seed: Tensor[(4, 4), "f32"]): + out = tf.add(seed, seed) + for row in range(1, 2, 3, 4): + out = tf.add(x[row, :], seed) + return out + + else: + + @func + def looping(x: Tensor[(10, 4), "f32"], seed: Tensor[(4, 4), "f32"]): + out = tf.add(seed, seed) + for row in steps(1, 2): # noqa: F821 + out = tf.add(x[row, :], seed) + return out diff --git a/tests/passes/test_hir_to_tir.py b/tests/passes/test_hir_to_tir.py index bc61630c..0c31e3a7 100644 --- a/tests/passes/test_hir_to_tir.py +++ b/tests/passes/test_hir_to_tir.py @@ -21,7 +21,6 @@ ) from tilefoundry.ir.core import Call, Constant, Var from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.cuda.nn.mma import Mma_SM80_16x8x16, Wgmma_SM90_64x128x16 from tilefoundry.ir.hir.function import Function as HirFunction from tilefoundry.ir.hir.grid_region import GridRegionExpr from tilefoundry.ir.hir.nn.relu import ReLU as HirReLU @@ -312,48 +311,3 @@ def test_non_divisible_tile_window_lowering_fails_closed() -> None: f = tile_window_add with pytest.raises(NotImplementedError, match="requires handwritten tail lowering"): HirToTirPass().run(Module(name="t", functions=(f,), entry=f.name)) - - -def _mma_module(op, a_type, b_type, result_type) -> Module: - a = Var(type=a_type, name="a") - b = Var(type=b_type, name="b") - body = Call(type=result_type, target=op, args=(a, b)) - fn = HirFunction.build( - name="mma_kernel", params=(a, b), body=body, return_type=result_type - ) - return Module(name="mma_test", functions=(fn,), entry=fn.name) - - -@pytest.mark.parametrize( - ("op", "a_type", "b_type", "result_type", "name"), - [ - pytest.param( - Mma_SM80_16x8x16( - dtype_a=DType.bf16, dtype_b=DType.bf16, dtype_acc=DType.f32 - ), - TensorType((16, 16), DType.bf16, None, StorageKind.RMEM), - TensorType((16, 8), DType.bf16, None, StorageKind.RMEM), - TensorType((16, 8), DType.f32, None, StorageKind.RMEM), - "Mma_SM80_16x8x16", - id="sm80", - ), - pytest.param( - Wgmma_SM90_64x128x16( - dtype_a=DType.bf16, dtype_b=DType.bf16, dtype_acc=DType.f32 - ), - TensorType((64, 16), DType.bf16, None, StorageKind.GMEM), - TensorType((16, 128), DType.bf16, None, StorageKind.GMEM), - TensorType((64, 128), DType.f32, None, StorageKind.GMEM), - "Wgmma_SM90_64x128x16", - id="wgmma", - ), - ], -) -def test_hir_mma_lowering_rejects_each_concrete_op_by_name( - op, a_type, b_type, result_type, name -) -> None: - with pytest.raises( - ValueError, - match=rf"HirToTirPass: {name} HIR compile route is unsupported", - ): - HirToTirPass().run(_mma_module(op, a_type, b_type, result_type))