fix(parser): validate HIR return annotations

docs/spec/hir.md

@@ -111,6 +111,11 @@ parser-lexical mesh binding; `ShardLayout.mesh` MUST point at an active
 binding on the lexical path. A `Mesh` MAY map fewer levels than the domain
 declares, but it MUST NOT create a level or change one's extent.
 
+**Return type.** `Function.return_type` is the HIR result `Type`: a
+`TensorType` for one result or a `TupleType` for multiple results. It is part of
+the function signature and is the result component projected into
+`Function.type`.
+
 **Value type.** `Function.type` is the IR-level `CallableType`
 ([types §7](./types.md#7-callabletype)) projected from `params` +
 `return_type`. The projection is fixed at construction and stays

docs/spec/parser.md

@@ -17,15 +17,43 @@ def parse_function(
 ) -> hir.Function | tir.PrimFunction: ...
 ```
 
-- Every parser-authored `Call` reachable from a Function body carries `SourceSpanMetadata` for
-  the AST expression that constructed it. A parent match fills only Calls without a span, so it
-  cannot replace a more precise child span. Traversal follows `Call` operands and IR `Tuple`
-  values, but does not assign source identity to shared lexical `Var` values. Source spans use
-  physical source-file coordinates with a one-based start column.
-- For `a, b = producer(...)`, detached `TupleGetItem(index=0)` and
-  `TupleGetItem(index=1)` lexical values carry the respective target Name spans (`a` and `b`) and
-  matching `BindingMetadata`; later reads do not replace that identity. A multi-carry loop's
-  derived projections carry the `for` statement span and their carry binding name.
+### 1.1 HIR Return Contracts
+
+An ordinary HIR function and a specialization variant MAY omit `-> return-type`;
+an ordinary function then records its inferred `body.type` as `Function.return_type`.
+When either declares an annotation, the parser MUST require
+`types_compatible(annotation, body.type)`. This is directional compatibility,
+not a parser-only equality relation: an annotation with `layout=None` leaves
+layout unconstrained according to the shared type rule.
+
+A `pass` HIR function is a dispatch prototype and MUST declare `-> return-type`.
+That annotation is its `Function.return_type` and callable return type. Every
+variant body MUST be compatible with that base return type; its own IR return
+type remains the exact base type so all variants share one dispatch signature.
+
+`Tensor[...]` without a storage slot constructs `storage=GMEM`, including in a
+return annotation. It is not an unspecified-storage spelling. Consequently an
+SMEM body under `-> Tensor[...]` must explicitly return a GMEM result; the
+parser reports both the annotation and inferred body type, with the authored
+function location, when they are incompatible. `tuple[...]` annotations are
+accepted and apply this same compatibility rule recursively to every field.
+
+### 1.2 Source Span Metadata
+
+Every parser-authored `Call` reachable from a Function body carries `SourceSpanMetadata` for
+the AST expression that constructed it. A parent match fills only Calls without a span, so it
+cannot replace a more precise child span. Traversal follows `Call` operands and IR `Tuple`
+values, but does not assign source identity to shared lexical `Var` values. Source spans use
+physical source-file coordinates with a one-based start column.
+
+### 1.3 Tuple Binding Metadata
+
+For `a, b = producer(...)`, detached `TupleGetItem(index=0)` and
+`TupleGetItem(index=1)` lexical values carry the respective target Name spans (`a` and `b`) and
+matching `BindingMetadata`; later reads do not replace that identity. A multi-carry loop's
+derived projections carry the `for` statement span and their carry binding name.
+
+### 1.4 Context and Diagnostics
 
 `FuncParserContext` carries the dialect, Function role, closure, topology scope, target, and
 optional base/key for one parse. `FunctionRole` is `ROOT`, `VARIANT`, or `CONVERTER`.
@@ -105,8 +133,11 @@ tensor-optional-slot  ::= layout
 tensor                ::= tensor-head '[' '(' (tensor-shape-layout ',' dtype | tensor-shape-layout
                           ',' dtype ',' tensor-optional-slot | tensor-shape-layout ',' dtype ','
                           tensor-optional-slot ',' tensor-optional-slot) ')' ']'
+tuple-type            ::= 'tuple' '[' '(' type-annotation (',' type-annotation)* ')' ']'
+                          | 'tuple' '[' type-annotation ']'
 scalar-type           ::= primary
 type-annotation       ::= tensor
+                          | tuple-type
                           | scalar-type
 signature             ::= (name ':' type-annotation (',' name ':' type-annotation)*)?
 return-type           ::= type-annotation
@@ -194,6 +225,7 @@ function              ::= 'def' name '(' signature ')' ('->' return-type)? ':' b
 | 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 | FunctionReturnCompatibilityRule | A HIR body with a return annotation must satisfy that annotation; a dispatch prototype must declare one, and each variant body must satisfy the prototype return contract. | src/tilefoundry/parser/pattern_nodes.py |
 | 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 |

src/tilefoundry/ir/types/utils.py

@@ -61,6 +61,15 @@ def layout_compatible(declared_layout, actual_layout) -> bool:
             )
             and layout_compatible(declared.layout, actual.layout)
         )
+    if isinstance(declared, TupleType):
+        return (
+            isinstance(actual, TupleType)
+            and len(declared.fields) == len(actual.fields)
+            and all(
+                types_compatible(declared_field, actual_field)
+                for declared_field, actual_field in zip(declared.fields, actual.fields)
+            )
+        )
     return actual == declared
 
 

src/tilefoundry/parser/pattern_nodes.py

@@ -34,6 +34,8 @@
 from tilefoundry.ir.types import TensorType
 from tilefoundry.ir.types.dim import DimVar
 from tilefoundry.ir.types.shard import Broadcast, Layout, Partial, Split
+from tilefoundry.ir.types.substitute import canonicalize_dims
+from tilefoundry.ir.types.utils import types_compatible
 
 from .ast_pattern import (
     _BINARY_OPERATORS,
@@ -963,11 +965,94 @@ def construct(match, children, context):
     RULES: ClassVar[tuple[AstRule[Any], ...]] = ()
 
 
+class TupleTypePattern(ElementPattern):
+    """Parse ``tuple[T, U, ...]`` as a parser-only type annotation value."""
+
+    element_name = "tuple_type"
+    syntax = LazyPattern(
+        lambda: ChoicePattern(
+            BranchPattern(
+                "tuple",
+                AstNodePattern(
+                    ast.Subscript,
+                    FieldPattern(
+                        "value",
+                        AstNodePattern(
+                            ast.Name,
+                            FieldPattern("id", LiteralPattern("tuple")),
+                        ),
+                    ),
+                    FieldPattern(
+                        "slice",
+                        AstNodePattern(
+                            ast.Tuple,
+                            CapturePattern("field_count", lambda node, context: len(node.elts)),
+                            FieldPattern(
+                                "elts",
+                                RepeatPattern(
+                                    ChildPattern(
+                                        "field_{index}",
+                                        lambda: TypeAnnotationPattern(),
+                                        "type_annotation",
+                                        "tuple_field",
+                                    ),
+                                    minimum=1,
+                                ),
+                            ),
+                        ),
+                    ),
+                ),
+                pattern_id="tuple.annotation",
+            ),
+            BranchPattern(
+                "tuple_single",
+                AstNodePattern(
+                    ast.Subscript,
+                    FieldPattern(
+                        "value",
+                        AstNodePattern(
+                            ast.Name,
+                            FieldPattern("id", LiteralPattern("tuple")),
+                        ),
+                    ),
+                    FieldPattern(
+                        "slice",
+                        ChildPattern(
+                            "field_0",
+                            lambda: TypeAnnotationPattern(),
+                            "type_annotation",
+                            "tuple_field",
+                        ),
+                    ),
+                ),
+                pattern_id="tuple.annotation",
+            ),
+        )
+    )
+
+    @staticmethod
+    def construct(match, children, context):
+        fields = tuple(
+            children[f"field_{index}"]
+            for index in range(match.captures.get("field_count", 1))
+        )
+        if not all(isinstance(field, (runtime.TensorType, runtime.TupleType)) for field in fields):
+            raise ParseError.from_node(
+                match.node,
+                context,
+                "tuple annotation fields must resolve to Tensor or tuple types",
+            )
+        return runtime.TupleType(fields=fields)
+
+    RULES: ClassVar[tuple[AstRule[Any], ...]] = ()
+
+
 class TypeAnnotationPattern(ElementPattern):
     element_name = "type_annotation"
     syntax = LazyPattern(
         lambda: ChoicePattern(
             TensorPattern(),
+            TupleTypePattern(),
             ScalarTypePattern(),
         )
     )
@@ -4062,6 +4147,60 @@ def apply(self, value, *, match, context):
         return value
 
 
+_AUTHORED_RETURN_ANNOTATION = "_tilefoundry_authored_return_annotation"
+_AUTHORED_BODY_TYPE = "_tilefoundry_authored_body_type"
+
+
+@dataclass(frozen=True)
+class FunctionReturnCompatibilityRule:
+    STATEMENT: ClassVar[str] = (
+        "A HIR body with a return annotation must satisfy that annotation; a "
+        "dispatch prototype must declare one, and each variant body must "
+        "satisfy the prototype return contract."
+    )
+
+    def apply(self, value, *, match, context):
+        if context.function is None:
+            raise ParseError.from_node(match.node, context, "function lacks parser context")
+        if context.function.dialect != "hir":
+            return value
+        declared = getattr(value, _AUTHORED_RETURN_ANNOTATION, None)
+        body_type = getattr(value, _AUTHORED_BODY_TYPE, None)
+        try:
+            if body_type is None:
+                if context.function.role is FunctionRole.ROOT and declared is None:
+                    raise ParseError.from_node(
+                        match.node,
+                        context,
+                        "HIR pass prototype requires a return annotation",
+                    )
+                return value
+            if declared is not None and not types_compatible(declared, body_type):
+                raise ParseError.from_node(
+                    match.node,
+                    context,
+                    "return annotation is not compatible with the inferred body type: "
+                    f"annotation {declared!r}, body {body_type!r}",
+                )
+            if context.function.role is FunctionRole.VARIANT:
+                base = context.function.base
+                if not isinstance(base, runtime.Function):
+                    raise ParseError.from_node(
+                        match.node, context, "variant lacks a HIR dispatch prototype"
+                    )
+                if not types_compatible(base.return_type, body_type):
+                    raise ParseError.from_node(
+                        match.node,
+                        context,
+                        f"variant body type {body_type!r} is not compatible with "
+                        f"dispatch return contract {base.return_type!r}",
+                    )
+            return value
+        finally:
+            delattr(value, _AUTHORED_RETURN_ANNOTATION)
+            delattr(value, _AUTHORED_BODY_TYPE)
+
+
 @dataclass(frozen=True)
 class FunctionRoleValidationRule:
     STATEMENT: ClassVar[str] = (
@@ -4196,14 +4335,11 @@ def construct(match, children, context):
         specializations = context.function.specializations
         converter = context.function.converter
         if context.function.dialect == "hir":
+            declared_return = (
+                None if declared_return is None else canonicalize_dims(declared_return)
+            )
             if body is None:
-                if declared_return is None:
-                    raise ParseError.from_node(
-                        match.node,
-                        context,
-                        "HIR pass prototype requires a return annotation",
-                    )
-                return_type = declared_return
+                return_type = declared_return or runtime.UnitType()
             elif context.function.role is FunctionRole.VARIANT:
                 base = context.function.base
                 assert isinstance(base, runtime.Function)
@@ -4222,6 +4358,8 @@ def construct(match, children, context):
                 return_type=return_type,
                 specializations=specializations,
             )
+            setattr(function, _AUTHORED_RETURN_ANNOTATION, declared_return)
+            setattr(function, _AUTHORED_BODY_TYPE, None if body is None else body.type)
             if context.function.role is FunctionRole.VARIANT:
                 setattr(function, runtime.DISPLAY_NAME, match.captures["name"])
                 function.name = function_name
@@ -4252,6 +4390,7 @@ def construct(match, children, context):
     RULES: ClassVar[tuple[AstRule[Any], ...]] = (
         FunctionSignatureRule(),
         FunctionDialectRule(),
+        FunctionReturnCompatibilityRule(),
         FunctionRoleValidationRule(),
         FunctionRegistrationRule(),
     )
@@ -4286,6 +4425,7 @@ def _body_as_ast_module(body: object, *, strip_docstring: bool = False) -> ast.M
     "FunctionDialectRule",
     "FunctionPattern",
     "FunctionRegistrationRule",
+    "FunctionReturnCompatibilityRule",
     "FunctionRoleValidationRule",
     "FunctionSignatureRule",
     "IndexEndpointPattern",
@@ -4324,6 +4464,7 @@ def _body_as_ast_module(body: object, *, strip_docstring: bool = False) -> ast.M
     "TensorShapeLayoutPattern",
     "TupleAssignmentPattern",
     "TupleExpressionPattern",
+    "TupleTypePattern",
     "TypeAnnotationPattern",
     "UnaryExpressionPattern",
     "CallVariadicInputFormRule",

tests/analysis/test_analysis_families.py

@@ -268,7 +268,7 @@ class _SplitLastAxis:
     def last_axis(
         x: Tensor[(1, _SPLIT_BLOCK, _SPLIT_HIDDEN), "bf16"],
         w: ConstTensor[(_SPLIT_HIDDEN, _SPLIT_OUT), "bf16"],
-    ) -> Tensor[(1, _SPLIT_BLOCK, _SPLIT_OUT), "bf16"]:
+    ):
         with Mesh(("cta",), layout=(_SPLIT_GRID,), names=("unit",)) as mesh:
             rows = tf.reshard(
                 x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem"
@@ -287,7 +287,7 @@ class _SplitStripMajor:
     def strip_major(
         x: Tensor[(1, _SPLIT_BLOCK, _SPLIT_HIDDEN), "bf16"],
         w: ConstTensor[(_SPLIT_GRID, _SPLIT_HIDDEN, _SPLIT_PER), "bf16"],
-    ) -> Tensor[(_SPLIT_GRID, _SPLIT_BLOCK, _SPLIT_PER), "bf16"]:
+    ):
         with Mesh(("cta",), layout=(_SPLIT_GRID,), names=("unit",)) as mesh:
             rows = tf.reshard(
                 x[:, :, 0:_SPLIT_BLOCK], (1, _SPLIT_BLOCK, _SPLIT_BLOCK), "smem"

tests/fixtures/placed/prefill_decode_attention.py

@@ -106,7 +106,11 @@ def decode(
                 running_out = next_out
 
             normalized = running_out / running_sum
-            return tf.transpose(tf.cast(normalized, dtype="bf16"), perm=(0, 2, 1, 3))
+            return tf.reshard(
+                tf.transpose(tf.cast(normalized, dtype="bf16"), perm=(0, 2, 1, 3)),
+                (1, SEQ, HEADS, HEAD_DIM),
+                "gmem",
+            )
 
     @attend.specialize(DimVarRangePat("seq", 2, 4097))
     def prefill(

tests/parser/test_dimensions.py

@@ -19,7 +19,7 @@ class Model:
         @func
         def f(
             x: Tensor[(1, 16, 8192), "f32"],
-        ) -> Tensor[(1, 16, 8192), "f32"]:
+        ):
             with Mesh(("cta",), layout=(128,), names=("unit",)) as mesh:
                 width = 4096 + 4096
                 return tf.reshard(