lablup / lablup/mlxcel

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

Open
#1,141 0 comments 0 reactions 0 assignees View on GitHub
priority:low status:ready type:chore
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## 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:

```rust
/// 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 #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-#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-#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`.

Contributor guide

Open the contributing guide

Research direction

Start with scripts/ci/check_kernel_dtype_keys.py, Makefile:582-585, .github/workflows/ci.yml:157, and docs/code-guidelines.md to understand the existing checker and documented conventions. Compare the proposed staleness strategies, then validate call-site matching using utils.rs and layers.rs, including the create_causal_mask history and softcap examples. Done means a green checker on main, failure on the pre-#1121 annotation, make verify and CI wiring, and documented strategy and format.

Written by the indexing model from the issue text.

Assessment

Tech stack
github-actions, python, rust
Domain
build-system, ci-cd, tooling
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.