Metal: a run-time index into a function-local fixed-size vector reads 0.0 in a large kernel

[duburcqa]

Summary

On the Metal backend, reading a function-local fixed-size vector at a run-time index inside a dynamic loop returns 0.0 once the enclosing kernel grows past a size threshold. In the same loop iteration, comparisons that read a different local vector at that same run-time index read it correctly, a compile-time index into the affected vector returns the right value, and the loop still selects the right entry. Below the size threshold the identical code is correct, and arm64 is correct at every size.

The MSL that quadrants emits for the affected site is correct, so this is an Apple Metal shader compiler defect rather than a quadrants codegen bug. Opening it here because quadrants is what feeds that compiler, the trigger is a construct quadrants emits routinely (a dynamically indexed spvUnsafeArray in Function storage), and there is a plausible mitigation on the quadrants side.

Reproduction

Self-contained, quadrants only. Each of the NCALL inlined copies of affine_coords runs on its own data so none folds away; the kernel checks itself, since a correct backend must make the run-time-indexed read m_max equal one of the three compile-time-indexed entries of ms.

import os
import sys

import numpy as np
import quadrants as qd

ARCH = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("REPRO_ARCH", "metal")
NCALL = int(os.environ.get("NCALL", "128"))

qd.init(
    arch=getattr(qd, ARCH),
    default_fp=qd.f32,
    default_ip=qd.i32,
    offline_cache=False,
    enable_fallback=False,
    force_scalarize_matrix=True,
    advanced_optimization=True,
    cfg_optimization=False,
    fast_math=True,
)

TRI = np.array(
    [
        [-1.3000001e-01, 8.2301348e-04, -1.0208413e-04],
        [7.0000000e-02, 8.2302000e-04, -1.0208000e-04],
        [-1.3000001e-01, 8.2269000e-04, 3.9897920e-02],
    ],
    dtype=np.float32,
)
POINT = np.array([-1.2e-01, 8.23e-04, 1.0e-03], dtype=np.float32)

B = 8
verts = qd.Vector.field(3, dtype=qd.f32, shape=(NCALL, 4))
out = qd.field(dtype=qd.f32, shape=(B, 8))


@qd.func
def affine_coords(point, tri_v1, tri_v2, tri_v3):
    # Signed areas of the triangle's three axis-aligned projections.
    ms = qd.Vector([0.0, 0.0, 0.0], dt=qd.f32)
    for i in qd.static(range(3)):
        i1, i2 = (i + 1) % 3, (i + 2) % 3
        if i == 1:
            i1, i2 = i2, i1

        ms[i] = (
            tri_v2[i1] * tri_v3[i2]
            - tri_v2[i2] * tri_v3[i1]
            - tri_v1[i1] * tri_v3[i2]
            + tri_v1[i2] * tri_v3[i1]
            + tri_v1[i1] * tri_v2[i2]
            - tri_v1[i2] * tri_v2[i1]
        )

    # Pick the entry that is >= both others. This is the scan that miscompiles.
    m_max = qd.cast(0.0, qd.f32)
    i_x, i_y = qd.cast(0, qd.i32), qd.cast(0, qd.i32)
    absms = qd.abs(ms)
    picked = -1
    for i in range(3):
        if absms[i] >= absms[(i + 1) % 3] and absms[i] >= absms[(i + 2) % 3]:
            m_max = ms[i]
            picked = i
            i_x, i_y = (i + 1) % 3, (i + 2) % 3
            if i == 1:
                i_x, i_y = i_y, i_x
            break

    cs = qd.Vector([0.0, 0.0, 0.0], dt=qd.f32)
    for i in qd.static(range(3)):
        tv1, tv2 = tri_v2, tri_v3
        if i == 1:
            tv1, tv2 = tri_v3, tri_v1
        elif i == 2:
            tv1, tv2 = tri_v1, tri_v2

        cs[i] = (
            point[i_x] * tv1[i_y]
            + point[i_y] * tv2[i_x]
            + tv1[i_x] * tv2[i_y]
            - point[i_x] * tv2[i_y]
            - point[i_y] * tv1[i_x]
            - tv2[i_x] * tv1[i_y]
        )

    # The scan picks one of the three entries, so the run-time-indexed read must equal one of the
    # compile-time-indexed ones.
    is_bad = picked < 0 or (m_max != ms[0] and m_max != ms[1] and m_max != ms[2])

    return cs / m_max, ms, m_max, picked, is_bad


@qd.kernel
def probe():
    for i_b in range(B):
        n_bad = 0
        n_no_pick = 0
        checksum = qd.cast(0.0, qd.f32)
        for k in qd.static(range(NCALL)):
            lam, ms, m_max, picked, is_bad = affine_coords(verts[k, 0], verts[k, 1], verts[k, 2], verts[k, 3])
            n_bad += 1 if is_bad else 0
            n_no_pick += 1 if picked < 0 else 0
            checksum += lam[0] + lam[1] + lam[2]
            if is_bad:
                for j in qd.static(range(3)):
                    out[i_b, 3 + j] = ms[j]
                out[i_b, 6] = m_max
                out[i_b, 7] = picked
        out[i_b, 0] = n_bad
        out[i_b, 1] = n_no_pick
        out[i_b, 2] = checksum


data = np.empty((NCALL, 4, 3), dtype=np.float32)
for k in range(NCALL):
    scale = np.float32(1.0 + 0.01 * k)
    data[k, 0] = POINT * scale
    data[k, 1:] = TRI * scale
verts.from_numpy(data)
probe()
rows = out.to_numpy()
print(f"arch={ARCH} NCALL={NCALL}")
print(f"  bad calls per env        = {rows[:, 0].astype(int).tolist()}")
print(f"  calls picking nothing    = {rows[:, 1].astype(int).tolist()}")
print(f"  sum of affine coords     = {rows[0, 2]:.7e}")
print(f"  ms of the last bad call  = ({rows[0, 3]:.7e}, {rows[0, 4]:.7e}, {rows[0, 5]:.7e})")
print(f"  its ms[i], picked        = {rows[0, 6]:.7e}, {int(rows[0, 7])}")
assert (rows[:, 1] == 0).all(), "the scan selected nothing"
assert (rows[:, 0] == 0).all(), f"BUG: ms[i] at a run-time index matched no entry in {int(rows[:, 0].max())} calls"
assert np.isfinite(rows[:, 2]).all(), "BUG: non-finite affine coordinates"
print("OK - not reproduced on this backend")

python repro.py (metal) reports:

arch=metal NCALL=128
  bad calls per env        = [39, 39, 39, 39, 39, 39, 39, 39]
  calls picking nothing    = [0, 0, 0, 0, 0, 0, 0, 0]
  sum of affine coords     = inf
  ms of the last bad call  = (5.1659299e-10, 1.5904801e-02, -1.2863893e-07)
  its ms[i], picked        = 0.0000000e+00, 1
AssertionError: BUG: ms[i] at a run-time index matched no entry in 39 calls

python repro.py arm64 reports bad calls per env = [0, ...] and sum of affine coords = 1.2800000e+02, the exact expected value.

Note the last two lines of the metal output: ms is a perfectly well-formed vector whose entry 1 is the argmax by five orders of magnitude, the scan correctly selected picked = 1, and yet ms[i] in that same iteration came back 0.0.

The size threshold

Sweeping NCALL on metal, everything else fixed. All eight batch entries always agree, so this is deterministic rather than a race:

copies in the kernel	miscompiled copies
64	0
80	0
88	0
96	1
128	39
256	186

So the threshold sits between 89 and 96 copies, and the proportion affected grows with kernel size past it.

The generated MSL is correct

Dumped with QD_DUMP_MSL=1. The excerpt below is from the 29k-line kernel where this was originally found, which shows the pattern at the size that fails in production; the repro above emits the same shape. In the affected region quadrants emits two ordinary function-scope arrays, writes ms at constant indices, and reads it at the very loop index the comparisons use:

spvUnsafeArray<float, 3> tmp26661_unknown;                  // ms
spvUnsafeArray<float, 3> tmp26840_unknown;                  // abs(ms)
tmp26661_unknown[0] = ...; tmp26661_unknown[1] = ...; tmp26661_unknown[2] = ...;
tmp26840_unknown[0] = abs(tmp26661_unknown[0]); ...
for (;;) {
    tmp26853_i32 = _39831 + 1;
    if ((tmp26853_i32 < 3) == false) { ...; break; }
    ...
    bool tmp26873_u1 = tmp26840_unknown[tmp26853_i32] >= tmp26840_unknown[tmp26870_i32];   // reads right
    ...
    if (_39842 != false) {
        ...
        _39851 = tmp26661_unknown[tmp26853_i32];                                           // reads 0.0
        break;
    }
    _39831 = tmp26853_i32;
    continue;
}

Nothing in that MSL is wrong: same index, same storage class, same array size, one array read correctly and the other not. The SPIR-V behind it is the expected OpAccessChain on an OpVariable in Function storage (TaskCodegen::visit(MatrixPtrStmt), the stmt->origin->is<AllocaStmt>() arm). The alloca legitimately stays a tensor alloca here, since GatherScalarizableLocalPointers (quadrants/transforms/scalarize.cpp) declines to scalarize any alloca that carries a non-const MatrixPtrStmt offset.

Ruled out

Measured, each one on its own against a passing baseline:

• buffer-binding pressure: the repro above binds three buffers; the real kernel this came from binds 28 of Metal's 31, and both miscompile;
• control-flow nesting depth: twelve levels of nested run-time loops with a break around the scan, at small kernel size, reads back correctly; the repro above miscompiles at two levels;
• loop unrollability: taking the scan's bound from a field so the trip count is unknown at compile time reads back correctly at small kernel size;
• a dynamically indexed write to a local vector, and run-time indices into vectors selected by run-time branches, both present in the original code, are not needed;
• extra live local vectors around the scan, extra qd.func nesting levels, and a second call site of the same func, all read back correctly on their own.

The only variable that flips the outcome is the size of the enclosing kernel.

Where this bites

Found in Genesis' rigid collider. The affected read is the divisor of a triangle's affine coordinates in EPA's contact construction, so a well-conditioned polytope face produced an infinite contact position, an infinite constraint row, and a halt reported as a force gone to nan. The enclosing kernel is 29k lines of MSL, far past the threshold above, and 21 of its 43 function-scope arrays are float[3]. Genesis works around it by unrolling that scan with qd.static(range(3)) so every read carries a compile-time index.

Suggested mitigation

Since the emitted MSL is correct, the only lever on the quadrants side is to stop emitting the construct Apple miscompiles. For a function-scope alloca of a small fixed-size tensor type, a run-time-indexed load could be lowered to a select chain over per-element loads at constant indices instead of an OpAccessChain with a dynamic index, so no dynamic addressing of thread-local storage survives. Element counts in practice are tiny (2 to 15 in the kernel above), so the select chain is short. Dynamic stores need the mirrored treatment, a per-element predicated store, and are the more invasive half; loads alone would already cover this failure mode. Gating that on the Metal backend keeps CUDA, AMDGPU and CPU on the current path.

If that is judged too costly, the fallback is documenting that a run-time index into a function-local fixed-size vector is unsafe on Metal, since the failure is silent, data-independent, and does not reproduce at small kernel size, which makes it very expensive to find downstream.

Environment

• quadrants 1.1.4, llvm 22.1.0, builds 47e9dfb and 9578f16 both reproduce
• macOS 26.5.2 (25F84), Apple M4 Max, 40 GPU cores
• Python 3.12.8, default_fp=f32, force_scalarize_matrix=True
• arm64 is the passing control at every size tested