spb/metrika Public

1//2//  MLXBootstrap.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation11import MLX1213/// GPU-batched pairs bootstrap (CLAUDE.md §5).14///15/// Resample indices come from Philox4x32-10 evaluated as vectorized MLX16/// ops with the same key schedule as the CPU reference, so GPU and CPU17/// streams are bit-identical for a given seed — the reproducibility18/// release blocker. The heavy O(reps·n·k²) work (gather + cross-products)19/// runs on the GPU in float32; the tiny O(reps·k³) normal-equation solves20/// run on the CPU in float64 (§5 hybrid).21///22/// Documented tolerance vs the CPU float64 path: coefficient draws agree23/// to ~1e-5 relative for well-conditioned designs (float32 gather sums).24public enum ZQGPUBootstrap {2526    /// Whether the Metal-backed MLX device is usable in this process.27    ///28    /// SwiftPM's command-line build cannot compile Metal shaders, so under29    /// `swift build`/`swift test` the mlx metallib does not exist and any30    /// MLX op would abort. Mirror the C++ loader's search (colocated,31    /// main-bundle Cmlx resource bundle, all loaded bundles) and report32    /// GPU availability only when a metallib is actually present —33    /// xcodebuild-built apps and tests have it, CLI builds fall back to34    /// the CPU path via the planner.35    public static let isAvailable: Bool = {36        #if arch(arm64)37        var roots: [URL] = []38        if let url = Bundle.main.resourceURL { roots.append(url) }39        roots.append(Bundle.main.bundleURL)40        if let executable = Bundle.main.executableURL {41            roots.append(executable.deletingLastPathComponent())42        }43        for bundle in Bundle.allBundles {44            if let url = bundle.resourceURL { roots.append(url) }45        }46        let candidates = roots.flatMap { root in47            [48                root.appendingPathComponent("mlx.metallib"),49                root.appendingPathComponent("default.metallib"),50                root.appendingPathComponent(51                    "mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib"52                ),53            ]54        }55        return candidates.contains { FileManager.default.fileExists(atPath: $0.path) }56        #else57        return false58        #endif59    }()6061    /// Runs `replicates` pairs-bootstrap OLS fits and returns the62    /// coefficient draws, one row per replicate (constant last when63    /// `includeConstant`). Replicates whose resampled design is singular64    /// come back as NaN rows for the caller to drop.65    public static func pairsBootstrapOLS(66        y: [Double],67        predictors: [(name: String, values: [Double])],68        includeConstant: Bool = true,69        replicates: Int,70        seed: UInt64,71        memoryBudgetBytes: Int = 512 << 2072    ) -> [[Double]] {73        let n = y.count74        let p = predictors.count75        let k = p + (includeConstant ? 1 : 0)76        precondition(n > k, "insufficient observations")7778        // Conditioning: the GPU stage forms X'X in float32, which squares79        // the condition number — raw data with an intercept loses 2–380        // digits. Standardizing each predictor by its full-sample mean and81        // scale is an exact reparametrization (the shift/scale constants82        // are fixed across replicates), so the original coefficients are83        // recovered exactly after the float64 solve:84        //   β_j = β̃_j / s_j,   b₀ = α + ȳ − Σ c_j β_j.85        var shifts = [Double](repeating: 0, count: p)86        var scales = [Double](repeating: 1, count: p)87        var yShift = 0.088        if includeConstant {89            yShift = y.reduce(0, +) / Double(n)90            for (j, column) in predictors.enumerated() {91                let mean = column.values.reduce(0, +) / Double(n)92                let variance = column.values.reduce(0) {93                    $0 + ($1 - mean) * ($1 - mean)94                } / Double(n)95                shifts[j] = mean96                scales[j] = variance.squareRoot() > 1e-30 ? variance.squareRoot() : 197            }98        }99100        // Standardized design matrix on the GPU, float32, row-major (n, k).101        var design = [Float]()102        design.reserveCapacity(n * k)103        for row in 0..<n {104            for (j, column) in predictors.enumerated() {105                design.append(Float((column.values[row] - shifts[j]) / scales[j]))106            }107            if includeConstant { design.append(1) }108        }109        let x = MLXArray(design, [n, k])110        let yGPU = MLXArray(y.map { Float($0 - yShift) }, [n, 1])111112        // §5 chunking: keep the gathered design under the memory budget.113        let bytesPerReplicate = n * (k + 1) * 4 * 2   // gathered X & y + transposes114        let chunkSize = max(1, min(replicates, memoryBudgetBytes / bytesPerReplicate))115116        var draws = [[Double]](repeating: [], count: replicates)117        var start = 0118        while start < replicates {119            let count = min(chunkSize, replicates - start)120            let (xtx, xty) = crossProducts(121                x: x, y: yGPU, n: n, k: k,122                firstReplicate: start, replicateCount: count, seed: seed123            )124            solveChunk(125                xtx: xtx, xty: xty, k: k,126                into: &draws, offset: start, count: count127            )128            start += count129        }130131        // Undo the standardization.132        if includeConstant {133            for r in 0..<replicates where !draws[r].isEmpty {134                var recovered = draws[r]135                var interceptShift = yShift136                for j in 0..<p {137                    recovered[j] = draws[r][j] / scales[j]138                    interceptShift -= shifts[j] * recovered[j]139                }140                recovered[p] = draws[r][p] + interceptShift141                draws[r] = recovered142            }143        }144        return draws145    }146147    /// One replicate's resample indices computed on the GPU — exposed so148    /// the cross-backend parity test can compare against the CPU149    /// reference (`ZQResampling.pairsBootstrapIndices`) bit-for-bit.150    public static func gpuIndices(151        replicate: Int, sampleSize: Int, seed: UInt64152    ) -> [Int] {153        let base = UInt64(replicate) * UInt64(sampleSize)154        let positions = MLXArray(0..<Int32(sampleSize)).asType(.uint64)155            + MLXArray(base, dtype: .uint64)156        let indices = philoxBoundedIndices(157            positions: positions, bound: sampleSize, seed: seed158        )159        eval(indices)160        return indices.asArray(Int32.self).map(Int.init)161    }162163    /// Diagnostic hook: one replicate's raw GPU cross-products with the164    /// same standardization as the production path.165    public static func debugCrossProducts(166        y: [Double], predictors: [(name: String, values: [Double])],167        replicate: Int, seed: UInt64168    ) -> (xtx: [Float], xty: [Float]) {169        let n = y.count170        let p = predictors.count171        let k = p + 1172        let meanY = y.reduce(0, +) / Double(n)173        var design = [Float]()174        for row in 0..<n {175            for column in predictors {176                let mean = column.values.reduce(0, +) / Double(n)177                let variance = column.values.reduce(0) {178                    $0 + ($1 - mean) * ($1 - mean)179                } / Double(n)180                design.append(Float((column.values[row] - mean) / variance.squareRoot()))181            }182            design.append(1)183        }184        let x = MLXArray(design, [n, k])185        let yGPU = MLXArray(y.map { Float($0 - meanY) }, [n, 1])186        return crossProducts(187            x: x, y: yGPU, n: n, k: k,188            firstReplicate: replicate, replicateCount: 1, seed: seed189        )190    }191192    /// Diagnostic hook: cross-products for replicate 0 computed inside a193    /// batch of `batchSize`, to expose batch-dependent kernel numerics.194    public static func debugBatchedCrossProducts(195        y: [Double], predictors: [(name: String, values: [Double])],196        batchSize: Int, seed: UInt64197    ) -> (xtx: [Float], xty: [Float]) {198        let n = y.count199        let p = predictors.count200        let k = p + 1201        let meanY = y.reduce(0, +) / Double(n)202        var design = [Float]()203        for row in 0..<n {204            for column in predictors {205                let mean = column.values.reduce(0, +) / Double(n)206                let variance = column.values.reduce(0) {207                    $0 + ($1 - mean) * ($1 - mean)208                } / Double(n)209                design.append(Float((column.values[row] - mean) / variance.squareRoot()))210            }211            design.append(1)212        }213        let x = MLXArray(design, [n, k])214        let yGPU = MLXArray(y.map { Float($0 - meanY) }, [n, 1])215        let (xtx, xty) = crossProducts(216            x: x, y: yGPU, n: n, k: k,217            firstReplicate: 0, replicateCount: batchSize, seed: seed218        )219        return (Array(xtx.prefix(k * k)), Array(xty.prefix(k)))220    }221222    // MARK: - GPU stage223224    /// Gathers the resampled design for a chunk and returns float32225    /// cross-products, pulled back to the CPU: X'X (count·k·k) and X'y226    /// (count·k).227    private static func crossProducts(228        x: MLXArray, y: MLXArray, n: Int, k: Int,229        firstReplicate: Int, replicateCount: Int, seed: UInt64230    ) -> (xtx: [Float], xty: [Float]) {231        // Draw positions: replicate r uses counters r·n ..< (r+1)·n,232        // matching ZQResampling.pairsBootstrapIndices.233        let base = UInt64(firstReplicate) * UInt64(n)234        let total = replicateCount * n235        let positions = MLXArray(0..<Int32(total)).asType(.uint64)236            + MLXArray(base, dtype: .uint64)237238        let indices = philoxBoundedIndices(positions: positions, bound: n, seed: seed)239            .reshaped([replicateCount, n])240241        // Gather rows: (count, n, k) and (count, n, 1).242        let xg = x.take(indices.flattened(), axis: 0)243            .reshaped([replicateCount, n, k])244        let yg = y.take(indices.flattened(), axis: 0)245            .reshaped([replicateCount, n, 1])246247        // Batched cross-products. X'X is computed one column at a time248        // through (count,k,n)@(count,n,1) products: MLX's batched GEMM249        // mis-accumulates small k×k outputs at batch ≥ 2 (~6e-4 relative,250        // verified against float64 sums, independent of strides/fusion),251        // while the single-column path is float32-exact. k is single-digit252        // so the extra dispatches are negligible next to the gather.253        let xt = xg.transposed(0, 2, 1)               // (count, k, n)254        let xColumns = split(xg, parts: k, axis: 2)   // k × (count, n, 1)255        let xtxColumns = xColumns.map { matmul(xt, $0) }256        let xtx = concatenated(xtxColumns, axis: 2)   // (count, k, k)257        let xty = matmul(xt, yg)                      // (count, k, 1)258        eval(xtx, xty)259260        return (xtx.asArray(Float.self), xty.asArray(Float.self))261    }262263    /// Philox4x32-10 as vectorized MLX ops. Given absolute draw positions,264    /// returns bounded int32 indices in 0..<bound — bit-identical to265    /// `Philox4x32.integer(at:bound:)` on the CPU.266    static func philoxBoundedIndices(267        positions: MLXArray, bound: Int, seed: UInt64268    ) -> MLXArray {269        let mask32 = MLXArray(UInt64(0xFFFF_FFFF), dtype: .uint64)270        let block = positions >> 2271        let lane = (positions & MLXArray(UInt64(3), dtype: .uint64)).asType(.uint32)272273        var c0 = (block & mask32)274        var c1 = (block >> 32)275        var c2 = MLXArray.zeros(positions.shape, dtype: .uint64)276        var c3 = MLXArray.zeros(positions.shape, dtype: .uint64)277278        // Key schedule precomputed on the CPU — identical to the279        // reference: key_r = (seedLo + r·W0, seedHi + r·W1) mod 2³².280        let seedLo = UInt32(truncatingIfNeeded: seed)281        let seedHi = UInt32(truncatingIfNeeded: seed >> 32)282        let m0 = MLXArray(UInt64(0xD251_1F53), dtype: .uint64)283        let m1 = MLXArray(UInt64(0xCD9E_8D57), dtype: .uint64)284285        for round in 0..<10 {286            let k0 = seedLo &+ UInt32(round) &* 0x9E37_79B9287            let k1 = seedHi &+ UInt32(round) &* 0xBB67_AE85288            let product0 = m0 * c0                    // c0 < 2³² so exact289            let product1 = m1 * c2290            let high0 = product0 >> 32291            let low0 = product0 & mask32292            let high1 = product1 >> 32293            let low1 = product1 & mask32294            let newC0 = (high1 ^ c1 ^ MLXArray(UInt64(k0), dtype: .uint64)) & mask32295            let newC2 = (high0 ^ c3 ^ MLXArray(UInt64(k1), dtype: .uint64)) & mask32296            c0 = newC0297            c1 = low1298            c2 = newC2299            c3 = low0300        }301302        // Lane select, then bounded multiply-shift like the reference.303        let word = which(304            lane .== MLXArray(UInt32(0), dtype: .uint32), c0,305            which(306                lane .== MLXArray(UInt32(1), dtype: .uint32), c1,307                which(lane .== MLXArray(UInt32(2), dtype: .uint32), c2, c3)308            )309        )310        let bounded = (word * MLXArray(UInt64(bound), dtype: .uint64)) >> 32311        return bounded.asType(.int32)312    }313314    // MARK: - CPU stage315316    /// Solves the chunk's normal equations in float64 via Cholesky.317    /// Singular systems yield NaN rows.318    private static func solveChunk(319        xtx: [Float], xty: [Float], k: Int,320        into draws: inout [[Double]], offset: Int, count: Int321    ) {322        for r in 0..<count {323            var a = [Double](repeating: 0, count: k * k)324            var b = [Double](repeating: 0, count: k)325            // MLX buffers are row-major; LAPACK wants column-major. X'X is326            // symmetric so the transpose is free; X'y is a vector.327            for i in 0..<k {328                b[i] = Double(xty[r * k + i])329                for j in 0..<k {330                    a[j * k + i] = Double(xtx[r * k * k + i * k + j])331                }332            }333            if let solution = try? CholeskySolver.solve(a, k: k, rhs: b) {334                draws[offset + r] = solution335            } else {336                draws[offset + r] = [Double](repeating: .nan, count: k)337            }338        }339    }340}341