[FIX] Refuse out-of-bounds accesses in the tracer instead of corrupting memory

[mark14wu]

Summary

Tracing a kernel with an out-of-bounds access, for example a load or store that forgot its mask, made the process abort or segfault at exit (rc 134 corrupted size vs. prev_size / rc 139). The trigger is the out-of-bounds access, not the CPU:

@triton.jit
def oob(x_ptr, out_ptr, n, BLOCK: tl.constexpr):
    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    tl.store(out_ptr + offs, tl.load(x_ptr + offs) + 1)   # 4 * 32 lanes over 100 elements

x = torch.arange(100, dtype=torch.float32); out = torch.zeros(100)
tilelens.trace(Tracer())(oob)[(4,)](x, out, 100, BLOCK=32)   # "done", then crash at exit

Triton's interpreter runs loads, stores and atomics on raw host memory. For CPU tensors it uses the caller's own storage, so the overrun writes past the end of the torch allocation and corrupts glibc heap metadata. The crash only surfaces when the allocator frees at exit. It was first noticed on CPU, but it also happens with CUDA tensors, because the interpreter overruns its host copies the same way. Plain TRITON_INTERPRET=1 without TileLens crashes the same way: this is UB in the kernel, and the interpreter is not patched here.

What the tracer does now (tilelens/clients/tracer/tracer.py):

• Known memory. arg_callback records the storage byte range behind every tensor argument. It recurses into tuples and uses .base for triton.reinterpret wrappers and host TensorDescriptors.
• The check. Before each Triton load, store and atomic, and each Gluon gl.load/gl.store, the tracer's own before-callbacks run _check_in_bounds. If any lane of the pointer block, masked-off lanes included, touches a tensor argument's storage, every active lane must lie fully inside some known storage. Otherwise it raises IndexError before the access runs. The message names the program, the kernel line, the argument, and the offending byte range, and points to the Sanitizer. Masked-off lanes only attribute the access to a tensor. That matters because Triton issues a float atomic_max/atomic_min as two calls split by sign, and one of them can hold only the out-of-bounds lanes.
• What is not judged. Accesses that never touch a known storage, for example a whole program past the end, are not checked. Neither is any access in a program after it casts an integer to a pointer, as pointer tables do. The guard is therefore best-effort, and the README says so; the Sanitizer remains the complete check.
• Shared code. TRITON_ADAPTERS gains argument-normalising adapters for AtomicRMW (ptr, mask) and AtomicCas (ptr), mirroring the Load/Store ones. All policy stays in the tracer. Symbolic clients use op overriders, which do not go through these adapters.

Known gaps, addressed in the stacked follow-up PR: under Gluon, the integer-to-pointer exemption never fires, Gluon atomics are not checked, and Gluon async copies are not checked.

Test Plan

New tests:

• tests/end_to_end/test_tracer.py: refusal of out-of-bounds loads (with and without grid_idx sampling), stores (num_sms 1 and 4), atomic_add, atomic_cas, a sign-split float atomic_max, an int32 word overlapping the end of a byte tensor, a triton.reinterpret argument, and a tuple argument. There are also allow-tests for a view reading its base storage and for a pointer-table gather that mixes an argument row with a non-argument row.
• tests/end_to_end/test_gluon.py: an unmasked Gluon gl.store is refused (as an InterpreterError caused by IndexError).
• tests/unit/test_tracer.py: a Python bool mask; tests/unit/test_adapters.py: the two new adapters.
• Every refusal test uses a tensor from torch.from_numpy on a slice of a larger numpy buffer, so out-of-bounds writes land in sentinels rather than the heap. The tests then assert the sentinels are untouched, which is deterministic.

Runs (CPU, Triton 3.8):

• Repro, 20 runs each. Before: 15/20 crashed on the host and 12/20 in a GPU-less docker image (CPU-only torch, no libcuda). After: 0/20 crashed, and every run raised a clean IndexError.
• Negative controls. On main all new refusal tests fail. Each of the reinterpret, tuple, sign-split and word-overlap tests fails when its specific code path is removed. The pointer-table test fails without the integer-to-pointer exemption.
• pytest tests/ -n 8: 378 passed, 3 skipped.

[github-actions]

Performance Benchmark

Benchmark	main (min)	PR (min)	Change	Samples
gemm	0.105s	0.105s	+0.3%	20 / 20
gemm_oob	0.117s	0.117s	-0.1%	20 / 20
indirect_load	0.022s	0.022s	+1.2%	20 / 20
nested_loop	0.235s	0.236s	+0.4%	20 / 20
block_pointer_loop_advance	0.126s	0.125s	-0.8%	20 / 20
liger_jsd	0.139s	0.139s	-0.4%	20 / 20
flaggems_layernorm	0.395s	0.396s	+0.2%	20 / 20
swiglu	0.171s	0.170s	-0.3%	20 / 20
cross_entropy	0.975s	0.974s	-0.1%	20 / 20
fused_linear_jsd	0.211s	0.210s	-0.3%	20 / 20
Total	2.497s	2.495s	-0.1%	N/A

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