Skip to content

fix(parser): give a loop the source span its diagnostics need #149

Description

@zhen8838

Why

A loop carries no source identity, so no diagnostic about a loop can say where it
is. attach_authored_metadata attaches a span to a Call and to nothing else
(src/tilefoundry/parser/ast_pattern.py:201-202):

def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContext") -> object:
    """Attach *node*'s source identity to each unmarked Call it constructs."""
    if not isinstance(value, Call):
        return value

GridRegionExpr is not a Call, and neither is Var. Measured on a 21-line
program with one loop: 4 of 6 Call/Var nodes carry SourceSpanMetadata, and
the three that do not are exactly the loop, its induction_var, and its extent.
diagnostic_location returns None for all three.

This is what stopped #127 from pointing anywhere. That refusal now names the
variables instead:

tilefoundry: error: loop '_step' has a trip count the program computes at run
time from 'reps', so no per-occurrence total can be scaled by it; bind the
extent to a literal, or state it as an open dimension

Readable, but the reader still has to find _step themselves. Both refusals in
_domain_for (src/tilefoundry/analysis/scope.py) are in this position, and so
is any future one about a loop.

What

Give a loop the same source identity a call has. The parser knows the ast.For
node's position at the point it builds the region.

Contract

describe_expr and diagnostic_location (src/tilefoundry/ir/core/metadata.py:65
and :77) start answering for loops instead of returning None.

Risk

Scope needs settling before anyone starts, because Var and GridRegionExpr
are not the same decision.

Attaching to the loop region alone is contained: it buys the diagnostics above
and changes no report.

Attaching to Var also moves report labels. value_label
(src/tilefoundry/ir/core/metadata.py) reads a span when there is one, so a
loop's carried argument would go from acc to acc:19, and a parameter from
k_cache to k_cache:12. Two consequences:

  • tests/installed/models/smoke_qwen3_1_7b.py:73 asserts
    {"k_cache", "v_cache"} <= zero.keys(), which holds today precisely because a
    parameter has no span. It would need to change.
  • value_labels' repeated-name suffix exists for the case where two carried
    arguments share a name and neither has a span
    (('acc', 'rmem') twice, measured on
    tests/fixtures/placed/qwen3_1_7b_pd.py's layer_prefill). Spans would make
    that collision impossible, and the suffix would become unreachable.

A parameter arguably should stay unsuffixed and unlocated: it is a declaration,
not a statement, and its name is already unique. That points at attaching to the
loop region and its induction variable but not to parameters -- but that is the
call this issue is asking someone to make, not one it settles.

Reproducer

from __future__ import annotations
from tilefoundry import func, module
from tilefoundry.dsl import Mesh, Tensor, Topology, tf
from tilefoundry.target import CudaTarget

CTAS = 132
N = CTAS * 128


@module(entry="kernel", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", CTAS),))
class DynTrip:
    @func
    def kernel(x: Tensor[(N,), "f32"], reps: Tensor[(), "i64"]) -> Tensor[(N,), "f32"]:
        with Mesh(("cta",), layout=(CTAS,), names=("block",)) as m:
            placed = tf.reshard(x, (N @ m.block,), "gmem")
            acc = placed
            bound = reps + 0
            for _step in range(bound):
                acc = tf.square(acc)
            return tf.reshard(acc, (N @ m.block,), "gmem")
from tilefoundry.ir.core.metadata import diagnostic_location
from tilefoundry.ir.hir.grid_region import GridRegionExpr
from tilefoundry.ir.visitor import collect_exprs

fn = DynTrip.functions[0]
for expr in collect_exprs(fn.body):
    if isinstance(expr, GridRegionExpr):
        print(diagnostic_location(expr))            # None
        print(diagnostic_location(expr.extent))     # None
        print(expr.metadata)                        # ()

Found while fixing #127.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions