[REFACTOR] Key loop hooks by LoopSite

[mark14wu]

Summary

Refactor loop hook identity from a bare function-relative lineno to a LoopSite that combines the rewritten loop line number with a stable source-file token.

Triton's interpreter parses each jit function with def at line 1, so loops in different files can share a function-relative line number. With bare-lineno identity, two such loops collide in two places inside SymbolicClient:

• the z3 iterator variable loop_i_{lineno} — both loops would share one symbolic variable;
• the iterator-constraint cache keyed by (lineno, start, stop, step) — distinct loop sites could hit each other's cached constraint.

Scope note (verified experimentally): these collisions are an identity-correctness hazard, not a currently reproducible end-to-end mischeck. Device jit functions are executed un-rewritten (_jit_function_call calls the raw fn), so nested cross-file loops never reach the loop hooks today, and sequential launches that share a lineno also share the variable name and the cache key consistently, which yields correct constraints. The fix removes the latent collision and becomes load-bearing if device-function loops are ever wired into the loop hooks.

Changes

• Add LoopSite and loop_file_token in the loop patching layer.
• Pass LoopSite through loop hook dispatch from the Triton frontend and client manager.
• Require LoopSite as the only loop identity: LoopIter raises NotImplementedError on bare line numbers, and ClientManager hook signatures are annotated with LoopSite.
• Rename LoopContext.lineno to loop_site and key symbolic loop context/constraint caches strictly by LoopSite.
• Keep profiler loop reports grouped by integer line number (via loop_site.lineno), dropping its int-compat helper.
• Update sanitizer string expectations and add tests:
  • a unit test asserting two same-lineno sites in different files get distinct z3 variables and distinct cached constraints bound to their own variables (guards against a cache key regressing to lineno-only);
  • an end-to-end cross-file nested-loop OOB test (kernel loop + device-fn helper loop at the same relative line) documenting today's concrete-unroll behavior and guarding the scenario if device-fn loops become symbolic later.

Known limitation

The profiler still aggregates loop_info by bare loop_site.lineno, so loops at the same relative line in different files still merge in profiler reports (pre-existing behavior, kept intentionally; possible follow-up).

Validation

• uv run pytest tests/unit/test_symbolic_client.py
• uv run pytest tests/end_to_end/test_sanitizer.py::test_loop_deferred_checks_after_context tests/end_to_end/test_sanitizer.py::test_cross_file_same_relative_lineno_loops_detect_oob tests/unit/test_profiler.py
• uv run pytest tests/unit/test_sanitizer.py -k loop (loop-affine LoopContext construction)
• uv run pytest tests/unit/test_utils.py -k sanitizer_trace (traceback tests constructing LoopContext)

[github-actions]

Performance Benchmark

Benchmark	main (min)	PR (min)	Change	Samples
gemm	0.097s	0.095s	-1.5%	20 / 20
gemm_oob	0.107s	0.107s	-0.0%	20 / 20
indirect_load	0.019s	0.019s	-0.3%	20 / 20
nested_loop	0.207s	0.206s	-0.2%	20 / 20
block_pointer_loop_advance	0.176s	0.179s	+1.5%	20 / 20
liger_jsd	0.129s	0.128s	-0.5%	20 / 20
flaggems_layernorm	0.355s	0.354s	-0.2%	20 / 20
swiglu	0.160s	0.159s	-0.6%	20 / 20
cross_entropy	0.902s	0.895s	-0.7%	20 / 20
fused_linear_jsd	0.197s	0.195s	-0.6%	20 / 20
Total	2.348s	2.338s	-0.4%	N/A

Iterations: 1 warmup + 20 measured
Samples are shown as main / PR; long pytest benchmarks may use fewer samples.

[Jokeren]

Closing as stale during PR cleanup.