From e57951d10ef9fc6d5e54eb3e97448e35a73e89c7 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 31 Aug 2026 21:53:56 +0800 Subject: [PATCH 1/2] fix(check): resolve callee dispatch and optional weights --- docs/spec/cli.md | 19 ++-- docs/spec/hir.md | 17 ++-- .../ISSUES.md | 14 +-- .../gen_model.py | 19 ++-- .../repro/specialize_through_call.py | 19 ++-- src/tilefoundry/cli/check.py | 9 +- src/tilefoundry/ir/hir/specialize.py | 23 +++-- tests/cli/test_cli_check.py | 21 +--- .../placed/specialize_through_call.py | 14 +-- tests/installed/conftest.py | 10 ++ tests/installed/smoke_check.py | 92 +++++++++++------- tests/ir/core/test_specialize.py | 95 ++++++++++++++++++- 12 files changed, 240 insertions(+), 112 deletions(-) diff --git a/docs/spec/cli.md b/docs/spec/cli.md index f3a8b229..b233d82b 100644 --- a/docs/spec/cli.md +++ b/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 diff --git a/docs/spec/hir.md b/docs/spec/hir.md index 284f5cd7..c4706cb1 100644 --- a/docs/spec/hir.md +++ b/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 diff --git a/examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md b/examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md index 79feab5a..172def3c 100644 --- a/examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md +++ b/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 this PR** — `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 diff --git a/examples/nemotron_3_5_lightning_30b_a3b-tilelang/gen_model.py b/examples/nemotron_3_5_lightning_30b_a3b-tilelang/gen_model.py index 1dc96bcf..07e64dea 100644 --- a/examples/nemotron_3_5_lightning_30b_a3b-tilelang/gen_model.py +++ b/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") diff --git a/examples/nemotron_3_5_lightning_30b_a3b-tilelang/repro/specialize_through_call.py b/examples/nemotron_3_5_lightning_30b_a3b-tilelang/repro/specialize_through_call.py index 50ba3f2c..d25e7c8c 100644 --- a/examples/nemotron_3_5_lightning_30b_a3b-tilelang/repro/specialize_through_call.py +++ b/examples/nemotron_3_5_lightning_30b_a3b-tilelang/repro/specialize_through_call.py @@ -1,5 +1,5 @@ #!/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 @@ -7,20 +7,17 @@ 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 diff --git a/src/tilefoundry/cli/check.py b/src/tilefoundry/cli/check.py index 52dd79ea..bf9dc5b9 100644 --- a/src/tilefoundry/cli/check.py +++ b/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): diff --git a/src/tilefoundry/ir/hir/specialize.py b/src/tilefoundry/ir/hir/specialize.py index 4008e5a4..ec4e4a2e 100644 --- a/src/tilefoundry/ir/hir/specialize.py +++ b/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) diff --git a/tests/cli/test_cli_check.py b/tests/cli/test_cli_check.py index 42adcdff..10526293 100644 --- a/tests/cli/test_cli_check.py +++ b/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( [ diff --git a/tests/fixtures/placed/specialize_through_call.py b/tests/fixtures/placed/specialize_through_call.py index f2fcfa38..7e915df4 100644 --- a/tests/fixtures/placed/specialize_through_call.py +++ b/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"]: diff --git a/tests/installed/conftest.py b/tests/installed/conftest.py index 11e3aa1f..8642f906 100644 --- a/tests/installed/conftest.py +++ b/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") diff --git a/tests/installed/smoke_check.py b/tests/installed/smoke_check.py index c8b75025..c35b4202 100644 --- a/tests/installed/smoke_check.py +++ b/tests/installed/smoke_check.py @@ -26,6 +26,9 @@ "--inputs", "random", "--weights", "random", "--out", "output", "--fn", "nan_inf" ) +_JUDGE = ("--inputs", "random", "--out", "output", "--fn", "nan_inf") + + def test_check_on_a_leaf_does_not_materialise_its_siblings_weights(tf, leaf_weights): """A leaf check must not allocate its siblings' 96 GiB weight union.""" done = tf( @@ -62,20 +65,54 @@ def test_check_specialises_through_a_dispatching_callee(tf, specialize_through_c assert done.returncode == 0, done.stderr -def test_analyze_through_a_dispatching_callee_is_still_blocked( +def test_analyze_specialises_through_a_dispatching_callee( tf, specialize_through_call, tmp_path ): - """Keep the analyze boundary visible until callee specialization is fixed.""" + """Analyze picks the callee's variant from --dim instead of refusing the rebuild.""" + costs = {} + for extent in ("n=64", "n=512"): + report = tmp_path / f"{extent}.md" + done = tf( + "analyze", + f"{specialize_through_call}:ToCallee", + str(report), + "--dim", + extent, + "--compute-cost", + ) + assert done.returncode == 0, done.stderr + assert "out of memory" not in done.stderr + costs[extent] = next( + line + for line in report.read_text().splitlines() + if line.startswith("# compute-cost") + ) + assert costs["n=64"] != costs["n=512"] + + +def test_weights_are_needed_only_where_one_is_reached(tf, square_cpu, hir_composition): + """A run reaching no weight needs no source; a reached weight reads from one.""" + for extra in ((), ("--weights", "random")): + done = tf("check", f"{square_cpu}:Mine", *_JUDGE, *extra) + assert done.returncode == 0, done.stderr + done = tf( - "analyze", - f"{specialize_through_call}:ToCallee", - str(tmp_path / "out.md"), - "--dim", - "n=64", + "check", + f"{hir_composition}:CrossModule", + *_JUDGE, + "--weights", + "random", ) + assert done.returncode == 0, done.stderr + + +def test_a_reached_weight_with_no_source_is_named_where_it_is_first_asked_for( + tf, hir_composition +): + """The first use names a missing child weight and the Module declaring it.""" + done = tf("check", f"{hir_composition}:CrossModule", *_JUDGE) assert done.returncode != 0 - assert "does not choose" in done.stderr - assert "out of memory" not in done.stderr + assert "'expert'" in done.stderr and "'w'" in done.stderr def test_check_reports_grid_loop_parser_errors_from_the_installed_wheel( @@ -262,24 +299,19 @@ def test_an_orchestration_method_names_the_files_its_inputs_need(tf, orchestrate assert "add_pair, affine_pair" in done.stderr -def test_both_input_axes_are_required(tf, square_twin, tmp_path) -> None: - torch.save(torch.arange(168, dtype=torch.float32), tmp_path / "x.pt") - - for omitted, expected in (("inputs", "no inputs stated"), ("weights", "needs weights")): - argv = ["check", f"{square_twin}:Twin.main"] - if omitted != "inputs": - argv += ["--inputs", f"files:{tmp_path / 'x.pt'}"] - if omitted != "weights": - argv += ["--weights", "random"] - done = tf( - *argv, - "--out", - "output", - "--fn", - "nan_inf", - ) - assert done.returncode == 1 - assert expected in done.stderr +def test_inputs_are_required(tf, square_twin) -> None: + done = tf( + "check", + f"{square_twin}:Twin.main", + "--weights", + "random", + "--out", + "output", + "--fn", + "nan_inf", + ) + assert done.returncode == 1 + assert "no inputs stated" in done.stderr def test_two_entirely_zero_sides_are_a_match_not_a_total_mismatch(tf, square_twin) -> None: @@ -336,12 +368,6 @@ def test_real_weights_come_from_the_checkpoint_and_activations_are_drawn( assert "activations actual f32 (declared f32)" in done.stdout assert "max_violation 0" in done.stdout - refused = tf( - "check", f"{weighted_twin}:WeightedRootTwin.scaled", "--inputs", "random", *_ARGS[4:] - ) - assert refused.returncode == 1 - assert "needs weights" in refused.stderr - def test_a_nested_twin_is_reached_through_the_child_it_is_declared_under(tf, nested_twin) -> None: done = tf("check", f"{nested_twin}:NestedTwin.child.scaled", *_ARGS) diff --git a/tests/ir/core/test_specialize.py b/tests/ir/core/test_specialize.py index 613675ee..1beaf0cc 100644 --- a/tests/ir/core/test_specialize.py +++ b/tests/ir/core/test_specialize.py @@ -11,13 +11,15 @@ SMALL_CONTEXT_T, GqaOnline, ) +from tests.fixtures.placed.specialize_through_call import ToCallee from tilefoundry import func, module -from tilefoundry.dsl import Tensor, Topology, tf +from tilefoundry.dsl import DimVarRangePat, Tensor, Topology, tf from tilefoundry.dsl.tf import * # noqa: F401,F403 -- names resolved dynamically from tilefoundry.evaluator import evaluate from tilefoundry.ir.hir.specialize import ( SpecializationError, bound_dims_of, + display_name, is_concrete, origin_of, residual_dims, @@ -30,6 +32,66 @@ ENTRY = GqaOnline.entry_function() STEADY = {"ctx_len": SMALL_CONTEXT_T} _LOOP_CTX = DimVar("loop_ctx", 1, 4097) +_CALL_M = DimVar("call_m", 1, 17) +_CALL_N = DimVar("call_n", 1, 1024) +_NESTED_N = DimVar("nested_n", 1, 1024) +_DISPATCH_BOUND = 128 + + +@module(entry="run") +class _MissingCalleeDimension: + """A caller dimension independent of the dimension its callee dispatches on.""" + + @func + def pick(x: Tensor[(_CALL_N,), "f32"]) -> Tensor[(_CALL_N,), "f32"]: + pass + + @pick.specialize(DimVarRangePat("call_n", 1, _DISPATCH_BOUND)) + def pick_small(x: Tensor[(_CALL_N,), "f32"]) -> Tensor[(_CALL_N,), "f32"]: + return tf.add(x, x) + + @pick.specialize(DimVarRangePat("call_n", _DISPATCH_BOUND, 1024)) + def pick_big(x: Tensor[(_CALL_N,), "f32"]) -> Tensor[(_CALL_N,), "f32"]: + return tf.add(tf.add(x, x), x) + + @func + def run( + x: Tensor[(_CALL_M,), "f32"], k: Tensor[(_CALL_N,), "f32"] + ) -> Tensor[(_CALL_N,), "f32"]: + return pick(k) # noqa: F821 + + +@module(entry="run") +class _NestedDispatch: + """A dispatch variant whose body calls a second dispatch prototype.""" + + @func + def inner(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: + pass + + @inner.specialize(DimVarRangePat("nested_n", 1, _DISPATCH_BOUND)) + def inner_small(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: + return tf.add(x, x) + + @inner.specialize(DimVarRangePat("nested_n", _DISPATCH_BOUND, 1024)) + def inner_big(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: + return tf.add(tf.add(x, x), x) + + @func + def mid(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: + pass + + @mid.specialize(DimVarRangePat("nested_n", 1, _DISPATCH_BOUND)) + def mid_small(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: + return inner(x) # noqa: F821 + + @mid.specialize(DimVarRangePat("nested_n", _DISPATCH_BOUND, 1024)) + def mid_big(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: + return inner(x) # noqa: F821 + + @func + def run(x: Tensor[(_NESTED_N,), "f32"]) -> Tensor[(_NESTED_N,), "f32"]: + return mid(x) # noqa: F821 def test_the_model_is_dynamic_in_its_context_length_alone() -> None: @@ -79,6 +141,37 @@ def test_specialising_nothing_is_refused() -> None: specialize_function(ENTRY, {}) +def test_a_call_through_a_prototype_selects_by_the_callers_dimensions() -> None: + """The rebuilt call records which variant it chose, not the prototype.""" + run = ToCallee.entry_function() + + for extent, expected in ((64, "pick_small"), (512, "pick_big")): + rebuilt = specialize_function(run, {"n": extent}).body.target + assert display_name(origin_of(rebuilt)) == expected + + +def test_specialising_through_a_call_says_which_dimension_is_missing() -> None: + """A callee dispatching on an unbound dimension names that dimension.""" + caller = _MissingCalleeDimension.entry_function() + + with pytest.raises(SpecializationError, match="'call_n'.*was not given a size"): + specialize_function(caller, {"call_m": 4}) + + +def test_specialising_follows_dispatch_variants_into_nested_dispatches() -> None: + """Selecting a variant keeps rebuilding through dispatches in its body.""" + run = _NestedDispatch.entry_function() + + for extent, suffix in ((64, "small"), (512, "big")): + concrete = specialize_function(run, {"nested_n": extent}) + middle = concrete.body.target + inner = middle.body.target + + assert display_name(origin_of(middle)) == f"mid_{suffix}" + assert display_name(origin_of(inner)) == f"inner_{suffix}" + assert residual_dims(concrete) == () + + def test_a_specialised_function_states_extents_everywhere() -> None: """Nothing may still be a range -- including inside the callees. From 7eefc09f8eb9979a4f6005d3a25dea89254e549a Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 31 Aug 2026 22:02:32 +0800 Subject: [PATCH 2/2] docs(examples): link callee dispatch fix --- examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md b/examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md index 172def3c..cf1dbbd9 100644 --- a/examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md +++ b/examples/nemotron_3_5_lightning_30b_a3b-tilelang/ISSUES.md @@ -92,7 +92,7 @@ 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 this PR** — `ToCallee` now passes both commands shown in its repro. +**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