[BUG] matmul with size greater than 2^31 gives undermined results - appears to be memory leak

[vansickledavid-sketch]

Describe the bug
A clear and concise description of what the bug is.

To Reproduce

Include code snippet

Expected behavior
A clear and concise description of what you expected to happen.

Desktop (please complete the following information):

• OS Version: [e.g. MacOS 14.1.2]
• Version [e.g. 0.7.0]

Additional context
Add any other context about the problem here.

[vansickledavid-sketch]

matmul at [12_347, 174_000] produces intermittently incorrect results. Some runs are ok, but others are wildly inaccurate. This appears to be a downcast memory address from UInt64 to Int32 somewhere in the kernel. MacBook Pro M5 Max 128 Gb.

 for n in [12000, 24000, 48000, 96000, 144000, 168000, 174000, 192000, 216000, 240000, 264000, 288000] {
        let A = MLXRandom.normal([12347, n], dtype: .float32)
        let v = MLXRandom.normal([n, 1], dtype: .float32)
        let direct = matmul(A, v)
        direct.eval()
        
        var chunked = MLXArray.zeros([12347, 1], dtype: .float32)
        let chunks = 4
        let chunkSize = n / chunks
        for i in 0..<chunks {
            let lo = i * chunkSize
            let hi = (i == chunks-1) ? n : (i+1) * chunkSize
            chunked = chunked + matmul(A[0..., lo..<hi], v[lo..<hi])
        }
        chunked.eval()
        let rel = abs(direct - chunked).max().item(Float.self) / abs(direct).max().item(Float.self)
        print("n=\(n), rel_diff=\(rel)")
    }

Typical output:

n=12000,  rel_diff=3.6e-07
n=24000,  rel_diff=4.7e-07
n=48000,  rel_diff=7.1e-07
n=96000,  rel_diff=1.0e-06
n=144000, rel_diff=1.5e-06
n=168000, rel_diff=1.5e-06
n=174000, rel_diff=0.065    ← intermittent; varies from ~1e-6 to ~0.25 across runs
n=192000, rel_diff=1.6e-06
n=216000, rel_diff=1.1e-06
n=240000, rel_diff=9.5e-07
n=264000, rel_diff=8.4e-07
n=288000, rel_diff=1.1e-06

[vansickledavid-sketch]

Built a workaround using manual block multiplications. Key finding beyond the original report: sliced inputs to matmul also produce incorrect results unless explicitly materialized via slice + 0 (or equivalent). matmul(left[0..., lo..<hi], right[lo..<hi, 0...]) produces wrong results; matmul((left[0..., lo..<hi] + 0), (right[lo..<hi, 0...] + 0)) is correct.
Both conditions (large operand size and strided-view-into-large-parent) appear to trigger the same offset-arithmetic issue in the matmul kernel. The workaround below handles both by chunking and materializing each chunk.

infix operator ×: MultiplicationPrecedence
func ×(left: MLXArray, right: MLXArray) -> MLXArray {
    let safeLimit = 2_000_000_000
    if left.size <= safeLimit && right.size <= safeLimit {
        return matmul(left, right)
    }
    
    assert(left.ndim == 2 && right.ndim == 2, "Only for matrix multiplication")
    
    let m = left.shape[0]
    let k = left.shape[1]
    let n = right.shape[1]
    
    assert(k == right.shape[0], "Dimension Mismatch")
    assert(n <= safeLimit && m <= safeLimit, "Single matrix dimension too large")
 
    let kChunk = safeLimit/max(m, n)
    let numChunks = (k + kChunk - 1)/kChunk
    
    var result = MLXArray.zeros([m, n], dtype: left.dtype)
    for i in 0..<numChunks {
        let lo = i * kChunk
        let hi = min(k, lo + kChunk)
        result = result + matmul(left[0..., lo..<hi] + 0, right[lo..<hi, 0...] + 0)
        result.eval()
    }
    return result
}

[vansickledavid-sketch]

Related
#3248: same overflow pattern in conv_general
#3327: root cause identified as ShapeElem = int32_t in mlx/array.h, affecting multiple operations

[zcbenz]

What is the version of MLX you are using and can you also try some old versions of MLX?