fix(check): resolve callee dispatch and optional weights

docs/spec/cli.md

@@ -130,6 +130,14 @@ bounds the caller stated.
     parameter in that same order. Output names MUST come from return position: one tensor is
     `output`; a tuple's tensors are `output[0]`, `output[1]`, and so on in return
     order. These are positions, not names authored in the function.
+  - `--weights` states where weights come from: `random` MUST draw each weight
+    the first time it is asked for, and `ckpt:DIR` MUST read them from a
+    safetensors checkpoint. It is OPTIONAL: omitted, the run has no weight
+    source. A weight reached with no source MUST be refused where it is first
+    asked for, naming the Module that declares it and the weight. It MUST NOT be
+    refused ahead of the run from what the selected Module declares: a Module
+    declares only its own functions' weights, so what a run reaches is not that
+    set.
   - One input file MUST bind one parameter. Its value MAY be a bare tensor or
     an arbitrarily nested tuple or list of tensors; every leaf MUST be a tensor.
   - A target whose step is an orchestration method rather than a `@func` MUST
@@ -157,12 +165,11 @@ bounds the caller stated.
   - Each output MUST report the norm of its reference. Near zero, a relative
     measure divides by nothing, so the report MUST state what it measured instead
     rather than a number with no scale to read it against.
-  - Inputs MUST be stated: random, real weights from a checkpoint, or files, and
-    no form MAY be the default. Weights MUST come from the same draw on both
-    sides, and the report MUST say which form was used and what seed drew it.
-    It MUST also say the actual and declared dtype of every activation, plus the
-    tensor count and shape tree each input
-    file supplied.
+  - Activations MUST be stated -- random or files -- and no form MAY be the
+    default. Weights MUST come from the same draw on both sides, and the report
+    MUST say which form was used and what seed drew it. It MUST also say the
+    actual and declared dtype of every activation, plus the tensor count and
+    shape tree each input file supplied.
   - `--device DEVICE` names where inputs and weights are built, and so where the
     run happens. Omitted, it is the device the selection's Target declares. Given,
     it is honoured as stated: a Target declaring CUDA no longer refuses a machine

docs/spec/hir.md

@@ -304,10 +304,12 @@ executable bodies. There is no base body to fall back to.
 
 *Dispatch resolution.* A `Call` whose target is a dispatch prototype
 (`variants != ()`) is a dispatch call: the variant whose `DimVarRangePat`
-matches the call's concrete argument shapes is selected and is the call's
-result. A shape outside the envelope matches no variant and is an error;
-there is no base body to fall back to (the prototype body is `None`). A
-`Call` whose target has `variants == ()` is a direct call to that body.
+matches is selected and is the call's result. Evaluation selects from the
+call's concrete argument shapes; specialization selects from the caller's
+stated dimension bindings. Both use the same variant table. A shape outside
+the envelope matches no variant and is an error; there is no base body to fall
+back to (the prototype body is `None`). A `Call` whose target has
+`variants == ()` is a direct call to that body.
 
 *Authoring freeze.* Variants accumulate during authoring, before the
 base `Function` enters a `Module` ([core-ir §1](./core-ir.md#1-module)). A
@@ -1620,8 +1622,11 @@ def is_concrete(fn: Function) -> bool:
   - `specialize_function` MUST reject an empty binding, an unknown dimension,
     or a selected implementation with no body. It MUST record the chosen
     implementation and sorted bindings on a rebuilt function so `origin_of`
-    and `bound_dims_of` can recover them. Function calls do not rebuild their
-    targets and therefore do not create provenance records.
+    and `bound_dims_of` can recover them. Specialization MUST rebuild called
+    functions affected by the caller's bindings and record their provenance.
+    When a called function is a dispatch prototype, specialization MUST select
+    its implementation from the same bindings by the `variant_for` rule; an
+    unstated dispatch dimension MUST raise `SpecializationError` and name it.
   - `specialize_concretely` MUST require a non-empty string-to-integer mapping
     and MUST reject any residual dimension after specialization.
   - Provenance and bound-dimension records MUST NOT participate in structural

examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md

@@ -10,7 +10,7 @@ Repro files are under `repro/` (TileFoundry) and `kbench/` (TileLang).
 | # | in one line | repro | what it blocked |
 |---|---|---|---|
 | **TF-1** | a slice start carrying a mesh index will not evaluate | `repro/mesh_slice_start.py` | **long-context attention cannot be `check`ed** |
-| TF-2 | dispatch on the callee does not pass `--dim`; on the entry it needs a tuple return annotation | `repro/specialize_through_call.py` | the entry never reaches the dispatch |
+| TF-2 | callee dispatch now passes `--dim`; entry dispatch still needs a tuple return annotation | `repro/specialize_through_call.py` | the entry cannot carry the dispatch; the callee route is fixed |
 | TF-3 | TF-1's error carried no file, line or op | same | locating it |
 | TF-4 | checking one leaf materialises the whole Module's weights | `repro/leaf_weights.py` | no leaf of this model can be checked |
 | TF-5 | `--inputs random` builds states the model cannot be in, and reports an out-of-range first | — | had to dump real activations |
@@ -84,13 +84,15 @@ python repro/specialize_through_call.py
 | where it is put | result |
 |---|---|
 | `Direct`: the entry calls one variant's body directly | PASS |
-| `ToCallee`: variants on the callee, entry calls the prototype | `specialising through 'pick': the callee dispatches on its own variants, which this rebuild does not choose` |
+| `ToCallee`: variants on the callee, entry calls the prototype | PASS |
 | `ToEntry`: variants on the entry, but the entry returns several tensors | `HIR pass prototype requires a return annotation` |
 
-**a. It does not pass through.** `_specialize_callee`
-(`ir/hir/specialize.py:344`) refuses outright to rebuild a callee that carries
-variants of its own. `--dim` is ordinary usage for `check` and `analyze`, so as
-soon as the entry calls a dispatch prototype, both commands are unusable.
+**a. It now passes through — fixed.** Specialization selects a callee's variant
+from the caller's `--dim` bindings and rebuilds through that implementation.
+`check` and `analyze` can therefore both reach a dispatch prototype from the
+entry.
+
+**Fixed in `#145`** — `ToCallee` now passes both commands shown in its repro.
 
 **b. Moving it to the entry does not work either.** The shape
 `tilefoundry tutorial authoring` demonstrates is variants hung on the entry — but

examples/nemotron_3_5_lightning_30b_a3b-tilelang/gen_model.py

@@ -203,17 +203,14 @@ def attn_body(i, var):
         # `attend` is the prototype the two `DimVarRangePat` variants hang off,
         # and it is what the runtime keys its own two bodies on.
         #
-        # The step calls the long placement rather than the prototype, and not
-        # because it wants only that one. Specialising a caller at a bound
-        # dimension refuses to rebuild through a callee that has variants of its
-        # own ("the callee dispatches on its own variants, which this rebuild
-        # does not choose"), and the other way round -- putting the variants on
-        # the entry, which is the shape the authoring tutorial shows -- needs a
-        # return annotation on the prototype, which a step that returns 59
-        # tensors has no way to write. So the dispatch stands where it can be
-        # read and checked (`check model.py:attend --dim ctx_full=0,4096`), and
-        # the body the step names is the one that runs at the lengths this is
-        # about. See ISSUES.md; both limits have a repro under repro/.
+        # The step still calls the long placement rather than the prototype.
+        # The specialization limit that originally required this has been
+        # removed in this PR: a caller can now rebuild through a callee's
+        # variants. Switching this generated call back to the prototype is a
+        # separate model change. Putting the variants on the entry, which is the
+        # shape the authoring tutorial shows, still needs a return annotation on
+        # the prototype, which a step that returns 59 tensors has no way to
+        # write. See ISSUES.md; both shapes have a repro under repro/.
         out += [f"{p}_ctx = attend_by_context({p}_qg, {p}_k_cache, {p}_v_cache,"
                 f" {p}_kta, {p}_vta)"]
         out += proj(var, f"{p}_mx", f"{p}_ctx", f"{p}_w_o", "H", kdim="QP")

examples/nemotron_3_5_lightning_30b_a3b-tilelang/repro/specialize_through_call.py

@@ -1,26 +1,23 @@
 #!/usr/bin/env python
-"""Two ways to state a dispatch, and neither is reachable from a real entry.
+"""Two ways to state a dispatch; callee dispatch now works through a real entry.
 
 `tilefoundry tutorial authoring` puts the variants on the module's entry. That
 works when the entry returns one tensor. A decode step that returns its logits
 *and* the state every layer produced returns a tuple, and a dispatch prototype
 needs a return annotation, whose grammar is `tensor | scalar-type` -- so the
 entry of such a model cannot carry variants at all (`ToEntry` below).
 
-The other placement -- variants on a callee, entry calls it -- parses and types,
-and then fails the moment anything binds a dimension:
-
-    specialising through 'pick': the callee dispatches on its own variants,
-    which this rebuild does not choose
-
-which is every `check --dim` and every `analyze --dim`. `ToCallee` below is the
-smallest program that shows it; `Direct` is the same module with the entry
-calling one variant's body instead of the prototype, and it runs.
+The other placement -- variants on a callee, entry calls it -- parses, types,
+and now specializes through that call. `ToCallee` below is the smallest program
+that exercises it; `Direct` is the same module with the entry calling one
+variant's body instead of the prototype.
 
     $ tilefoundry check repro/specialize_through_call.py:Direct   --inputs random \\
           --dim n=64 --out output --fn nan_inf          # PASS
     $ tilefoundry check repro/specialize_through_call.py:ToCallee --inputs random \\
-          --dim n=64 --out output --fn nan_inf          # the rebuild error
+          --dim n=64 --out output --fn nan_inf          # PASS
+    $ tilefoundry analyze repro/specialize_through_call.py:ToCallee out.md \\
+          --dim n=64 --compute-cost                     # PASS
     $ python repro/specialize_through_call.py           # the return-annotation one
 """
 from __future__ import annotations

src/tilefoundry/cli/check.py

@@ -28,6 +28,7 @@
 from tilefoundry.runtime import PREDICATES, RuntimeModule
 from tilefoundry.runtime.measure import Predicate, check, flatten_outputs
 from tilefoundry.runtime.resource import (
+    DictResource,
     DrawnResource,
     RuntimeResource,
     SafetensorsResource,
@@ -283,7 +284,11 @@ def draw_inputs(module: Module, dims: dict[str, int], seed: int, device: str):
     return _random_activations(concrete, generator, device)
 
 
-def build_resource(spec: str, module: Module, device: str, generator=None) -> RuntimeResource:
+def build_resource(
+    spec: str | None, module: Module, device: str, generator=None
+) -> RuntimeResource:
+    if spec is None:
+        return DictResource({})
     if spec == "random":
         generator = generator or torch.Generator(device=device).manual_seed(SEED)
         return DrawnResource(module, generator, device)
@@ -579,8 +584,6 @@ def run_check(arguments: argparse.Namespace) -> int:
     stated = parse_dims(arguments.dim) or {}
     if arguments.inputs is None:
         raise ValueError("no inputs stated")
-    if arguments.weights is None:
-        raise ValueError(f"needs weights {list(selection.module.weights)!r}")
     device = arguments.device or _device(selection.module)
     runs = []
     for dims in _combinations(stated):

src/tilefoundry/ir/hir/specialize.py

@@ -220,7 +220,7 @@ def visit_Call(self, call: Call, ctx: InstantiateContext) -> Expr:
         new_target = call.target
         if isinstance(new_target, Function):
             new_target = _specialize_callee(
-                new_target, ctx.dims, ctx.type_ctx, call
+                new_target, ctx.dims, ctx.type_ctx
             )
         new_target = _substitute_op_dims(new_target, ctx.dims)
         new_metadata = _substitute_authored_dims(call.metadata, ctx.dims)
@@ -335,18 +335,23 @@ def _specialize_callee(
     callee: Function,
     dims: Mapping[str, int],
     ctx: TypeInferContext,
-    call: Call,
 ) -> Function:
-    """Rebuild a nested callee at the dimensions its caller was given."""
-    if callee.variants:
-        raise ValueError(
-            f"specialising through {call and callee.name!r}: the callee "
-            "dispatches on its own variants, which this rebuild does not choose"
-        )
+    """Rebuild a nested callee at the dimensions its caller was given.
+
+    A dispatching callee picks from the same dims through variant_for. A
+    dimension the caller never bound is still refused and named.
+
+    The dispatch guard on the identity shortcut is defensive: verified variants
+    anchor dispatch dimensions in their parameter types. The guard states that
+    cross-file invariant where the shortcut relies on it.
+    """
+    dispatched = bool(callee.variants)
+    if dispatched:
+        callee = variant_for(callee, dims)
     if callee.body is None:
         return callee
     bound = tuple(substitute_dims(param.type, dims) for param in callee.params)
-    if all(new is param.type for new, param in zip(bound, callee.params)):
+    if not dispatched and all(new is param.type for new, param in zip(bound, callee.params)):
         return callee
     return instantiate_dimensions(callee, bound, ctx, dims)
 

tests/cli/test_cli_check.py

@@ -376,25 +376,8 @@ def test_check_refuses_what_it_cannot_answer(routing, capsys, comparison, refuse
     assert refused in capsys.readouterr().err
 
 
-def test_inputs_must_be_stated_and_weights_must_come_from_somewhere(routing, capsys) -> None:
-    """Neither the inputs nor the weights have a default form."""
-    assert (
-        cli.main(
-            [
-                "check",
-                ROUTING,
-                "--inputs",
-                "random",
-                "--out",
-                "output[0]",
-                "--fn",
-                "nan_inf",
-            ]
-        )
-        == 1
-    )
-    assert "needs weights ['w_router']" in capsys.readouterr().err
-
+def test_inputs_must_be_stated(capsys) -> None:
+    """Activations have no default form."""
     assert (
         cli.main(
             [

tests/fixtures/placed/specialize_through_call.py

@@ -1,18 +1,18 @@
-"""A dispatch on a callee: check runs it, while analyze still refuses it."""
+"""A dispatch on a callee: check and analyze both select its implementation."""
 
 from tilefoundry import func, module
 from tilefoundry.dsl import DimVar, DimVarRangePat, Mesh, Tensor, tf
 from tilefoundry.dsl.tf import *  # noqa: F401, F403
 from tilefoundry.ir.types.shard import Topology
-from tilefoundry.target import CpuTarget
+from tilefoundry.target import CudaTarget
 
 D, W, BOUND = 64, 4, 128
 N = DimVar("n", 1, 1024)
-_CPU = CpuTarget()
+_CUDA = CudaTarget("nvidia.h200_sxm")
 _CTA = Topology("cta", W)
 
 
-@module(entry="run", target=_CPU, topologies=(_CTA,))
+@module(entry="run", target=_CUDA, topologies=(_CTA,))
 class ToCallee:
     """The dispatch is on a callee; the entry calls the prototype."""
 
@@ -34,14 +34,14 @@ def pick_big(
     ) -> Tensor[(1, D), "f32"]:
         with Mesh(("cta",), layout=(W,), names=("w",)) as m:
             xs = tf.reshard(x, (1, D @ m.w), "smem")
-            return tf.reshard(xs + xs, (1, D), "gmem")
+            return tf.reshard(xs + xs + xs, (1, D), "gmem")
 
     @func
     def run(x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"]) -> Tensor[(1, D), "f32"]:
         return pick(x, k)
 
 
-@module(entry="run", target=_CPU, topologies=(_CTA,))
+@module(entry="run", target=_CUDA, topologies=(_CTA,))
 class Direct:
     """The entry calls one variant body directly."""
 
@@ -63,7 +63,7 @@ def pick_big(
     ) -> Tensor[(1, D), "f32"]:
         with Mesh(("cta",), layout=(W,), names=("w",)) as m:
             xs = tf.reshard(x, (1, D @ m.w), "smem")
-            return tf.reshard(xs + xs, (1, D), "gmem")
+            return tf.reshard(xs + xs + xs, (1, D), "gmem")
 
     @func
     def run(x: Tensor[(1, D), "f32"], k: Tensor[(1, N), "f32"]) -> Tensor[(1, D), "f32"]:

tests/installed/conftest.py

@@ -228,6 +228,16 @@ def leaf_weights() -> Path:
     return _fixture_path("placed", "leaf_weights.py")
 
 
+@pytest.fixture(scope="module")
+def square_cpu() -> Path:
+    return _fixture_path("placed", "square_cpu.py")
+
+
+@pytest.fixture(scope="module")
+def hir_composition() -> Path:
+    return _fixture_path("logical", "hir_composition.py")
+
+
 @pytest.fixture(scope="module")
 def specialize_through_call() -> Path:
     return _fixture_path("placed", "specialize_through_call.py")