Skip to content

chore: Enforce the // Used by: convention in CI, staleness included #1141

Open

Description

@inureyes

Problem / Background

docs/code-guidelines.md opens with the // Used by: rule: when you touch a shared function that several models depend on, record which models call it. The convention is widely adopted (721 Used by: lines across src/**/*.rs on main at d5abf1fa), but nothing checks it. The two tracked hot spots are annotated unevenly:

File top-level pub fn Used by: lines
src/lib/mlxcel-core/src/utils.rs 25 13
src/lib/mlxcel-core/src/layers.rs 25 36

(layers.rs has 101 pub fn once inherent and trait methods are counted, so its annotations are spread across both free functions and methods.)

Absence is the mild failure. The severe one is a stale annotation, because it does not read as missing information. It reads as an answer, and a contributor sizing up the blast radius of a change trusts it.

The motivating case

Issue #1110 was filed on the belief that create_causal_mask carried no annotation. It carried one, and the one it carried was inverted. Before PR #1121 (merged as f0bf3a2c), utils.rs said:

/// Used by: Llama, Qwen, Mixtral, Gemma, Cohere, Phi, OLMo, Exaone, GLM4,
/// MiniCPM, DeepSeek, Hunyuan, StarCoder2 and other causal attention callers

Grepping that same commit, each of mixtral.rs, phi.rs, phi3small.rs, starcoder2.rs, llama3.rs, gemma.rs, gemma2.rs, cohere.rs, glm4.rs, olmoe.rs and qwen3_moe.rs calls create_causal_mask exactly zero times. Those families had moved to the implicit-causal fused SDPA path, passing mask: None when seq_len > 1. The real caller set at that commit was 44 non-test files under src/models, and 51 non-test files across all of src (55 counting the four test files). Measure call sites with a trailing paren, grep -rln '\bcreate_causal_mask(' src --include='*.rs': without it the pattern also matches the sibling helpers create_causal_mask_with_window and create_causal_mask_with_left_padding and reports 57. It was close to disjoint from the roster.

A checker that only flags absence would have passed this file. That is the point: the check worth building is the one that would have caught this.

PR #1121 also added the "When the caller list is too long to enumerate" section to docs/code-guidelines.md, which is the policy a checker would enforce. Past roughly a dozen callers, do not write a roster. Write the rule that puts a caller on the list, name representatives per group, name the families deliberately absent, and close with the grep one-liner that regenerates the exact set.

Proposed Solution

Add a script under scripts/ci/, wired as a make verify-* target and a CI job, following scripts/ci/check_kernel_dtype_keys.py (131 lines, Makefile:582-585, .github/workflows/ci.yml:157). The same guidelines document states the JIT kernel dtype cache-key rule (enforced, issues #1053 and #1054) and the // Used by: rule (unenforced). Closing that asymmetry is the task.

The design has an easy half and a hard half, and this issue is not asking for a predetermined answer to the hard one. Evaluate the options and record the choice.

Easy half: absence. Flag a top-level pub fn in utils.rs or layers.rs that has more than N cross-module callers and no Used by: line. Pick N against the current tree so the check starts green.

Hard half: staleness. Annotations are prose naming model families, not file paths, so there is no mechanical mapping from Gemma2 to a path. Options worth weighing:

  1. Machine-checkable trailer. Require annotations above the enumeration threshold to end with the regenerating grep one-liner in a fixed form. The checker re-runs it and diffs the result against a count or list recorded in the annotation. Precise, and the only variant that catches the docs(core): the three shared functions code-guidelines.md names as its own example carry no // Used by: annotation #1110 case by construction, at the cost of a format contributors have to keep valid.
  2. Negative claims only. Check just the "not used by" half. A named family that does call the function is unambiguously a bug, needs only a family-name-to-path map plus a grep, and is cheap to run. It catches less, but every hit is real.
  3. Family-name heuristic. Map each name in a roster to candidate model files and warn when a listed family has zero call sites. This would have caught the pre-docs(core): annotate the three tracked shared functions in utils.rs #1121 text; the name-to-path map is the maintenance burden.

Whatever is chosen, counting files that merely mention a bare identifier overcounts. softcap is the example: 49 files under src contain the string, but utils.rs::softcap is one of several distinct definitions, alongside a private fn softcap in src/lib/mlxcel-xla/src/emitter/model.rs:1050, softcap_logits in src/audio/gemma3n/attention.rs:290, softcap_logits in src/lib/mlxcel-xla/src/emitter/gemma3n_math.rs:332, and pub fn softcap_logits in src/models/muse_glimmer.rs:75. Match call sites, not substrings, and scope by import or module path.

Acceptance Criteria

  • A checker script exists under scripts/ci/, is wired into a make verify-* target and into make verify, and runs as a CI job.
  • The check fails on the pre-docs(core): annotate the three tracked shared functions in utils.rs #1121 create_causal_mask annotation. Demonstrate this against f0bf3a2c^ (or a fixture reproducing that text) and paste the output in the PR. A check that only detects absence does not satisfy this criterion, per the evidence above.
  • The check passes on main as it stands, with no annotation edits needed to make it green. If it does flag something, fix the annotation in the same PR and state what was stale.
  • softcap is not miscounted: the check distinguishes utils.rs::softcap from the unrelated same-named functions listed above.
  • docs/code-guidelines.md states which rule is enforced, by which script, and what the enforced annotation format is, matching how the kernel dtype section documents its own checker.
  • The chosen staleness strategy and the rejected options are recorded, in the PR body or in the script's module docstring.
  • Runtime is small enough for the CI job, with the dtype checker as the reference point.

Technical Considerations

  • Scope the rule by a property of the file rather than a hand-maintained list of functions, the way check_kernel_dtype_keys.py scopes by the presence of cuda_kernel(. A list that must be edited by hand goes stale the same way the annotations did.
  • Prefer /// doc comments over // on public items, per the guidelines, so annotations survive into rustdoc. The checker can enforce that cheaply.
  • The counts above were measured at d5abf1fa. Re-measure before setting any threshold.
  • Start with utils.rs and layers.rs, the two files the guidelines name, rather than all of src.

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

    priority:lowLow prioritystatus:readyReady to be worked ontype:choreMaintenance tasks (build, CI, etc.)

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions