SPB Git
spboucher.ai

Repository Public

spbmetrika

Commit

feat(gpu): MLX-batched pairs bootstrap with bit-identical Philox streams

- mlx-swift dependency in ZQGPU (the only module allowed to import MLX)
- Philox4x32-10 as vectorized MLX uint64 ops with the reference key
  schedule: GPU resample indices are bit-identical to the CPU generator
  (asserted per replicate) — the §5 reproducibility contract
- batched gather + float32 cross-products on GPU, float64 Cholesky solves
  on CPU (§5 hybrid); chunked to the memory budget
- conditioning: predictors standardized by fixed full-sample mean/scale
  (exact reparametrization) before forming X'X — raw data with an
  intercept lost 2-3 digits in float32
- MLX workaround: batched GEMM mis-accumulates small k x k outputs at
  batch >= 2 (~6e-4 rel, stride/fusion-independent); X'X is built one
  column at a time through the exact (count,k,n)@(count,n,1) path, with
  regression tests pinning the behavior
- planner: gpuAvailable detected from actual metallib presence (SwiftPM
  CLI cannot compile Metal shaders - swift test skips GPU suites, the
  xcodebuild-built app and tests get the GPU); reps >= 500 dispatches
  bootstrap to .gpu
- engine: backend-aware bootstrap, drops singular resamples, output
  titled 'GPU batched' with gpu scalar
- 68 tests green under xcodebuild (66 + GPU suites under swift test skip)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed · Aug 5, 2026 parent b24d834
7 changed files +699 −49 Browse files
7 changed files
Unified Split
modified MetrikaKit/Package.resolved +19 −1
@@ -1,5 +1,5 @@
1 { 1 {
2− "originHash" : "e4ae512a67bec7ae66f2abd23f1e87240c39d28811734cdfafc7a91e4323f6db", 2+ "originHash" : "c4143a21eaf19f86159743269f92608859217bba0735131f48d3b12086f2b0c2",
3 "pins" : [ 3 "pins" : [
4 { 4 {
5 "identity" : "duckdb-swift", 5 "identity" : "duckdb-swift",
@@ -10,6 +10,15 @@
10 "version" : "1.1.3" 10 "version" : "1.1.3"
11 } 11 }
12 }, 12 },
13+ {
14+ "identity" : "mlx-swift",
15+ "kind" : "remoteSourceControl",
16+ "location" : "https://github.com/ml-explore/mlx-swift",
17+ "state" : {
18+ "revision" : "0bb916c67f4b9e5c682cbe02a42c701c93ab5021",
19+ "version" : "0.31.6"
20+ }
21+ },
13 { 22 {
14 "identity" : "swift-argument-parser", 23 "identity" : "swift-argument-parser",
15 "kind" : "remoteSourceControl", 24 "kind" : "remoteSourceControl",
@@ -27,6 +36,15 @@
27 "revision" : "a0cb0954ecb21e4e31b0070e6ed5674e8556685a", 36 "revision" : "a0cb0954ecb21e4e31b0070e6ed5674e8556685a",
28 "version" : "1.6.0" 37 "version" : "1.6.0"
29 } 38 }
39+ },
40+ {
41+ "identity" : "swift-numerics",
42+ "kind" : "remoteSourceControl",
43+ "location" : "https://github.com/apple/swift-numerics",
44+ "state" : {
45+ "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2",
46+ "version" : "1.1.1"
47+ }
30 } 48 }
31 ], 49 ],
32 "version" : 3 50 "version" : 3
modified MetrikaKit/Package.swift +5 −1
@@ -31,6 +31,7 @@ let package = Package(
31 .package(url: "https://github.com/duckdb/duckdb-swift", from: "1.0.0"), 31 .package(url: "https://github.com/duckdb/duckdb-swift", from: "1.0.0"),
32 .package(url: "https://github.com/apple/swift-collections", from: "1.1.0"), 32 .package(url: "https://github.com/apple/swift-collections", from: "1.1.0"),
33 .package(url: "https://github.com/apple/swift-argument-parser", from: "1.3.0"), 33 .package(url: "https://github.com/apple/swift-argument-parser", from: "1.3.0"),
34+ .package(url: "https://github.com/ml-explore/mlx-swift", from: "0.18.0"),
34 ], 35 ],
35 targets: [ 36 targets: [
36 .target( 37 .target(
@@ -52,7 +53,10 @@ let package = Package(
52 ), 53 ),
53 .target( 54 .target(
54 name: "ZQGPU", 55 name: "ZQGPU",
55− dependencies: ["ZQData"], 56+ dependencies: [
57+ "ZQData",
58+ .product(name: "MLX", package: "mlx-swift"),
59+ ],
56 swiftSettings: strictConcurrency 60 swiftSettings: strictConcurrency
57 ), 61 ),
58 .target( 62 .target(
modified MetrikaKit/Sources/ZQEngine/Session.swift +81 −47
@@ -35,7 +35,7 @@ public actor ZQSession {
35 35
36 private let store: ZQDataStore 36 private let store: ZQDataStore
37 private let parser = ZQCommandParser() 37 private let parser = ZQCommandParser()
38− private let planner = ZQPlanner() 38+ private let planner = ZQPlanner(gpuAvailable: ZQGPUBootstrap.isAvailable)
39 private let scriptCommands: [String: ZQScriptCommand] 39 private let scriptCommands: [String: ZQScriptCommand]
40 private var logFileURL: URL? 40 private var logFileURL: URL?
41 private var scriptDepth = 0 41 private var scriptDepth = 0
@@ -98,10 +98,12 @@ public actor ZQSession {
98 // User script commands shadow nothing built-in; check after parse 98 // User script commands shadow nothing built-in; check after parse
99 // failure would be better UX, but the verb table already resolved. 99 // failure would be better UX, but the verb table already resolved.
100 let plan = planner.plan(command, rowCount: frame.rowCount) 100 let plan = planner.plan(command, rowCount: frame.rowCount)
101− return try await dispatch(plan.command) 101+ return try await dispatch(plan.command, backend: plan.backend)
102 } 102 }
103 103
104− private func dispatch(_ command: ZQCommand) async throws -> ZQResult { 104+ private func dispatch(
105+ _ command: ZQCommand, backend: ZQBackend = .cpu
106+ ) async throws -> ZQResult {
105 if let script = scriptCommands[command.verb] { 107 if let script = scriptCommands[command.verb] {
106 guard scriptDepth < 8 else { 108 guard scriptDepth < 8 else {
107 throw ZQEngineError("user command recursion too deep") 109 throw ZQEngineError("user command recursion too deep")
@@ -132,7 +134,7 @@ public actor ZQSession {
132 case "set": return try handleSet(command) 134 case "set": return try handleSet(command)
133 case "xtset": return try handleXTSet(command) 135 case "xtset": return try handleXTSet(command)
134 case "display": return try handleDisplay(command) 136 case "display": return try handleDisplay(command)
135− case "bootstrap": return try await handleBootstrap(command) 137+ case "bootstrap": return try await handleBootstrap(command, backend: backend)
136 case "graph": return try handleGraph(command) 138 case "graph": return try handleGraph(command)
137 case "log": return try handleLog(command) 139 case "log": return try handleLog(command)
138 default: 140 default:
@@ -1010,7 +1012,9 @@ public actor ZQSession {
1010 1012
1011 // MARK: - Bootstrap prefix 1013 // MARK: - Bootstrap prefix
1012 1014
1013− private func handleBootstrap(_ command: ZQCommand) async throws -> ZQResult { 1015+ private func handleBootstrap(
1016+ _ command: ZQCommand, backend: ZQBackend
1017+ ) async throws -> ZQResult {
1014 guard let body = command.body else { 1018 guard let body = command.body else {
1015 throw ZQEngineError("bootstrap: syntax is 'bootstrap, reps(#): command'") 1019 throw ZQEngineError("bootstrap: syntax is 'bootstrap, reps(#): command'")
1016 } 1020 }
@@ -1045,57 +1049,83 @@ public actor ZQSession {
1045 1049
1046 // Pairs bootstrap: replicate r regenerates its indices from the 1050 // Pairs bootstrap: replicate r regenerates its indices from the
1047 // Philox counter stream — deterministic and order-independent, so 1051 // Philox counter stream — deterministic and order-independent, so
1048− // replicates can be computed in parallel chunks. 1052+ // replicates can run in parallel CPU chunks or as GPU batches with
1049− let chunkCount = min( 1053+ // identical resamples.
1050− max(1, ProcessInfo.processInfo.activeProcessorCount), reps 1054+ let allDraws: [[Double]]
1051− ) 1055+ if backend == .gpu {
1052− let chunkSize = (reps + chunkCount - 1) / chunkCount 1056+ allDraws = ZQGPUBootstrap.pairsBootstrapOLS(
1053− let draws: [[Double]] = try await withThrowingTaskGroup( 1057+ y: y,
1054− of: [(Int, [Double])].self 1058+ predictors: predictors,
1055− ) { group in 1059+ includeConstant: includeConstant,
1056− for chunk in 0..<chunkCount { 1060+ replicates: reps,
1057− let lower = chunk * chunkSize 1061+ seed: seed
1058− let upper = min(reps, lower + chunkSize) 1062+ )
1059− guard lower < upper else { continue } 1063+ } else {
1060− group.addTask { 1064+ allDraws = try await withThrowingTaskGroup(
1061− var results: [(Int, [Double])] = [] 1065+ of: [(Int, [Double])].self
1062− results.reserveCapacity(upper - lower) 1066+ ) { group in
1063− for replicate in lower..<upper { 1067+ let chunkCount = min(
1064− let indices = ZQResampling.pairsBootstrapIndices( 1068+ max(1, ProcessInfo.processInfo.activeProcessorCount), reps
1065− replicate: replicate, sampleSize: n, generator: generator 1069+ )
1066− ) 1070+ let chunkSize = (reps + chunkCount - 1) / chunkCount
1067− let resampledY = indices.map { y[$0] } 1071+ for chunk in 0..<chunkCount {
1068− let resampledPredictors = predictors.map { column in 1072+ let lower = chunk * chunkSize
1069− (name: column.name, values: indices.map { column.values[$0] }) 1073+ let upper = min(reps, lower + chunkSize)
1074+ guard lower < upper else { continue }
1075+ group.addTask {
1076+ var results: [(Int, [Double])] = []
1077+ results.reserveCapacity(upper - lower)
1078+ for replicate in lower..<upper {
1079+ let indices = ZQResampling.pairsBootstrapIndices(
1080+ replicate: replicate, sampleSize: n, generator: generator
1081+ )
1082+ let resampledY = indices.map { y[$0] }
1083+ let resampledPredictors = predictors.map { column in
1084+ (name: column.name, values: indices.map { column.values[$0] })
1085+ }
1086+ let fit = try ZQOLS.fit(
1087+ y: resampledY,
1088+ predictors: resampledPredictors,
1089+ includeConstant: includeConstant,
1090+ variance: .classical
1091+ )
1092+ results.append((replicate, fit.coefficients.map(\.estimate)))
1070 } 1093 }
1071− let fit = try ZQOLS.fit( 1094+ return results
1072− y: resampledY,
1073− predictors: resampledPredictors,
1074− includeConstant: includeConstant,
1075− variance: .classical
1076− )
1077− results.append((replicate, fit.coefficients.map(\.estimate)))
1078 } 1095 }
1079− return results
1080 } 1096 }
1081− } 1097+ var collected = [[Double]](repeating: [], count: reps)
1082− var collected = [[Double]](repeating: [], count: reps) 1098+ for try await chunkResults in group {
1083− for try await chunkResults in group { 1099+ for (replicate, estimates) in chunkResults {
1084− for (replicate, estimates) in chunkResults { 1100+ collected[replicate] = estimates
1085− collected[replicate] = estimates 1101+ }
1086 } 1102 }
1103+ return collected
1087 } 1104 }
1088− return collected 1105+ }
1106+
1107+ // Singular resamples (NaN rows from the GPU path) are dropped,
1108+ // Stata-style.
1109+ let draws = allDraws.filter { row in !row.contains(where: \.isNaN) }
1110+ let completed = draws.count
1111+ guard completed >= 2 else {
1112+ throw ZQEngineError("bootstrap: too many failed replicates")
1089 } 1113 }
1090 1114
1091 // Bootstrap standard errors: sd of replicate estimates. 1115 // Bootstrap standard errors: sd of replicate estimates.
1116+ let title = backend == .gpu
1117+ ? "Bootstrap results (pairs, GPU batched)"
1118+ : "Bootstrap results (pairs)"
1092 var lines = [ 1119 var lines = [
1093− "Bootstrap results (pairs) Replications = " + 1120+ TableFormatter.pad(title, 46, right: false) + "Replications = " +
1094− TableFormatter.pad("\(reps)", 10), 1121+ TableFormatter.pad("\(completed)", 10),
1095 " Number of obs = " + 1122 " Number of obs = " +
1096 TableFormatter.pad("\(n)", 10), 1123 TableFormatter.pad("\(n)", 10),
1097− "",
1098 ] 1124 ]
1125+ if completed < reps {
1126+ lines.append("(\(reps - completed) replicates dropped: singular resample)")
1127+ }
1128+ lines.append("")
1099 let widths = [12, 12, 12, 8, 8, 22] 1129 let widths = [12, 12, 12, 8, 8, 22]
1100 lines.append( 1130 lines.append(
1101 TableFormatter.pad(sample.responseName, widths[0]) + " | " + 1131 TableFormatter.pad(sample.responseName, widths[0]) + " | " +
@@ -1112,12 +1142,16 @@ public actor ZQSession {
1112 ) 1142 )
1113 lines.append(TableFormatter.rule(widths)) 1143 lines.append(TableFormatter.rule(widths))
1114 1144
1115− var scalars: [String: Double] = ["reps": Double(reps), "N": Double(n)] 1145+ var scalars: [String: Double] = [
1146+ "reps": Double(completed),
1147+ "N": Double(n),
1148+ "gpu": backend == .gpu ? 1 : 0,
1149+ ]
1116 for j in 0..<k { 1150 for j in 0..<k {
1117 let estimates = draws.map { $0[j] } 1151 let estimates = draws.map { $0[j] }
1118− let mean = estimates.reduce(0, +) / Double(reps) 1152+ let mean = estimates.reduce(0, +) / Double(completed)
1119 let variance = estimates.reduce(0) { $0 + ($1 - mean) * ($1 - mean) } 1153 let variance = estimates.reduce(0) { $0 + ($1 - mean) * ($1 - mean) }
1120− / Double(reps - 1) 1154+ / Double(completed - 1)
1121 let se = variance.squareRoot() 1155 let se = variance.squareRoot()
1122 let observed = point.coefficients[j].estimate 1156 let observed = point.coefficients[j].estimate
1123 let z = observed / se 1157 let z = observed / se
added MetrikaKit/Sources/ZQGPU/CholeskySolver.swift +62 −0
@@ -0,0 +1,62 @@
1+//
2+// CholeskySolver.swift
3+// Metrika
4+//
5+// Author: Simon-Pierre Boucher
6+// Contact: contact@spboucher.ai
7+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8+//
9+
10+/// Plain-Swift Cholesky solve for the k×k normal equations of the GPU
11+/// bootstrap. k is the regressor count (single digits in practice), so a
12+/// dependency-free O(k³) routine beats reaching for LAPACK — and ZQStats
13+/// is the only module allowed to import Accelerate (CLAUDE.md §3).
14+enum CholeskySolver {
15+ struct SingularMatrix: Error {}
16+
17+ /// Solves A·x = b for symmetric positive-definite A (column-major,
18+ /// k×k). Throws on non-positive-definite input.
19+ static func solve(_ a: [Double], k: Int, rhs: [Double]) throws -> [Double] {
20+ precondition(a.count == k * k && rhs.count == k)
21+
22+ // Lower-triangular factor L with A = L·Lᵀ.
23+ var l = [Double](repeating: 0, count: k * k)
24+ for j in 0..<k {
25+ var diagonal = a[j * k + j]
26+ for p in 0..<j {
27+ diagonal -= l[p * k + j] * l[p * k + j]
28+ }
29+ guard diagonal > 0 else { throw SingularMatrix() }
30+ let root = diagonal.squareRoot()
31+ l[j * k + j] = root
32+ for i in (j + 1)..<k {
33+ var value = a[j * k + i]
34+ for p in 0..<j {
35+ value -= l[p * k + i] * l[p * k + j]
36+ }
37+ l[j * k + i] = value / root
38+ }
39+ }
40+
41+ // Forward substitution L·z = b.
42+ var z = [Double](repeating: 0, count: k)
43+ for i in 0..<k {
44+ var value = rhs[i]
45+ for p in 0..<i {
46+ value -= l[p * k + i] * z[p]
47+ }
48+ z[i] = value / l[i * k + i]
49+ }
50+
51+ // Back substitution Lᵀ·x = z.
52+ var x = [Double](repeating: 0, count: k)
53+ for i in stride(from: k - 1, through: 0, by: -1) {
54+ var value = z[i]
55+ for p in (i + 1)..<k {
56+ value -= l[i * k + p] * x[p]
57+ }
58+ x[i] = value / l[i * k + i]
59+ }
60+ return x
61+ }
62+}
added MetrikaKit/Sources/ZQGPU/MLXBootstrap.swift +340 −0
@@ -0,0 +1,340 @@
1+//
2+// MLXBootstrap.swift
3+// Metrika
4+//
5+// Author: Simon-Pierre Boucher
6+// Contact: contact@spboucher.ai
7+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8+//
9+
10+import Foundation
11+import MLX
12+
13+/// GPU-batched pairs bootstrap (CLAUDE.md §5).
14+///
15+/// Resample indices come from Philox4x32-10 evaluated as vectorized MLX
16+/// ops with the same key schedule as the CPU reference, so GPU and CPU
17+/// streams are bit-identical for a given seed — the reproducibility
18+/// 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 solves
20+/// run on the CPU in float64 (§5 hybrid).
21+///
22+/// Documented tolerance vs the CPU float64 path: coefficient draws agree
23+/// to ~1e-5 relative for well-conditioned designs (float32 gather sums).
24+public enum ZQGPUBootstrap {
25+
26+ /// Whether the Metal-backed MLX device is usable in this process.
27+ ///
28+ /// SwiftPM's command-line build cannot compile Metal shaders, so under
29+ /// `swift build`/`swift test` the mlx metallib does not exist and any
30+ /// MLX op would abort. Mirror the C++ loader's search (colocated,
31+ /// main-bundle Cmlx resource bundle, all loaded bundles) and report
32+ /// GPU availability only when a metallib is actually present —
33+ /// xcodebuild-built apps and tests have it, CLI builds fall back to
34+ /// 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 in
47+ [
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+ #else
57+ return false
58+ #endif
59+ }()
60+
61+ /// Runs `replicates` pairs-bootstrap OLS fits and returns the
62+ /// coefficient draws, one row per replicate (constant last when
63+ /// `includeConstant`). Replicates whose resampled design is singular
64+ /// 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 << 20
72+ ) -> [[Double]] {
73+ let n = y.count
74+ let p = predictors.count
75+ let k = p + (includeConstant ? 1 : 0)
76+ precondition(n > k, "insufficient observations")
77+
78+ // Conditioning: the GPU stage forms X'X in float32, which squares
79+ // the condition number — raw data with an intercept loses 2–3
80+ // digits. Standardizing each predictor by its full-sample mean and
81+ // scale is an exact reparametrization (the shift/scale constants
82+ // are fixed across replicates), so the original coefficients are
83+ // 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.0
88+ 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] = mean
96+ scales[j] = variance.squareRoot() > 1e-30 ? variance.squareRoot() : 1
97+ }
98+ }
99+
100+ // 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])
111+
112+ // §5 chunking: keep the gathered design under the memory budget.
113+ let bytesPerReplicate = n * (k + 1) * 4 * 2 // gathered X & y + transposes
114+ let chunkSize = max(1, min(replicates, memoryBudgetBytes / bytesPerReplicate))
115+
116+ var draws = [[Double]](repeating: [], count: replicates)
117+ var start = 0
118+ 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: seed
123+ )
124+ solveChunk(
125+ xtx: xtx, xty: xty, k: k,
126+ into: &draws, offset: start, count: count
127+ )
128+ start += count
129+ }
130+
131+ // Undo the standardization.
132+ if includeConstant {
133+ for r in 0..<replicates where !draws[r].isEmpty {
134+ var recovered = draws[r]
135+ var interceptShift = yShift
136+ 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] + interceptShift
141+ draws[r] = recovered
142+ }
143+ }
144+ return draws
145+ }
146+
147+ /// One replicate's resample indices computed on the GPU — exposed so
148+ /// the cross-backend parity test can compare against the CPU
149+ /// reference (`ZQResampling.pairsBootstrapIndices`) bit-for-bit.
150+ public static func gpuIndices(
151+ replicate: Int, sampleSize: Int, seed: UInt64
152+ ) -> [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: seed
158+ )
159+ eval(indices)
160+ return indices.asArray(Int32.self).map(Int.init)
161+ }
162+
163+ /// Diagnostic hook: one replicate's raw GPU cross-products with the
164+ /// same standardization as the production path.
165+ public static func debugCrossProducts(
166+ y: [Double], predictors: [(name: String, values: [Double])],
167+ replicate: Int, seed: UInt64
168+ ) -> (xtx: [Float], xty: [Float]) {
169+ let n = y.count
170+ let p = predictors.count
171+ let k = p + 1
172+ 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: seed
189+ )
190+ }
191+
192+ /// Diagnostic hook: cross-products for replicate 0 computed inside a
193+ /// 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: UInt64
197+ ) -> (xtx: [Float], xty: [Float]) {
198+ let n = y.count
199+ let p = predictors.count
200+ let k = p + 1
201+ 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: seed
218+ )
219+ return (Array(xtx.prefix(k * k)), Array(xty.prefix(k)))
220+ }
221+
222+ // MARK: - GPU stage
223+
224+ /// Gathers the resampled design for a chunk and returns float32
225+ /// cross-products, pulled back to the CPU: X'X (count·k·k) and X'y
226+ /// (count·k).
227+ private static func crossProducts(
228+ x: MLXArray, y: MLXArray, n: Int, k: Int,
229+ firstReplicate: Int, replicateCount: Int, seed: UInt64
230+ ) -> (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 * n
235+ let positions = MLXArray(0..<Int32(total)).asType(.uint64)
236+ + MLXArray(base, dtype: .uint64)
237+
238+ let indices = philoxBoundedIndices(positions: positions, bound: n, seed: seed)
239+ .reshaped([replicateCount, n])
240+
241+ // 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])
246+
247+ // Batched cross-products. X'X is computed one column at a time
248+ // through (count,k,n)@(count,n,1) products: MLX's batched GEMM
249+ // 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-digit
252+ // 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)
259+
260+ return (xtx.asArray(Float.self), xty.asArray(Float.self))
261+ }
262+
263+ /// Philox4x32-10 as vectorized MLX ops. Given absolute draw positions,
264+ /// returns bounded int32 indices in 0..<bound — bit-identical to
265+ /// `Philox4x32.integer(at:bound:)` on the CPU.
266+ static func philoxBoundedIndices(
267+ positions: MLXArray, bound: Int, seed: UInt64
268+ ) -> MLXArray {
269+ let mask32 = MLXArray(UInt64(0xFFFF_FFFF), dtype: .uint64)
270+ let block = positions >> 2
271+ let lane = (positions & MLXArray(UInt64(3), dtype: .uint64)).asType(.uint32)
272+
273+ 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)
277+
278+ // Key schedule precomputed on the CPU — identical to the
279+ // 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)
284+
285+ for round in 0..<10 {
286+ let k0 = seedLo &+ UInt32(round) &* 0x9E37_79B9
287+ let k1 = seedHi &+ UInt32(round) &* 0xBB67_AE85
288+ let product0 = m0 * c0 // c0 < 2³² so exact
289+ let product1 = m1 * c2
290+ let high0 = product0 >> 32
291+ let low0 = product0 & mask32
292+ let high1 = product1 >> 32
293+ let low1 = product1 & mask32
294+ let newC0 = (high1 ^ c1 ^ MLXArray(UInt64(k0), dtype: .uint64)) & mask32
295+ let newC2 = (high0 ^ c3 ^ MLXArray(UInt64(k1), dtype: .uint64)) & mask32
296+ c0 = newC0
297+ c1 = low1
298+ c2 = newC2
299+ c3 = low0
300+ }
301+
302+ // 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)) >> 32
311+ return bounded.asType(.int32)
312+ }
313+
314+ // MARK: - CPU stage
315+
316+ /// 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: Int
321+ ) {
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 is
326+ // 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] = solution
335+ } else {
336+ draws[offset + r] = [Double](repeating: .nan, count: k)
337+ }
338+ }
339+ }
340+}
added MetrikaKit/Tests/MetrikaKitTests/GPUBootstrapTests.swift +119 −0
@@ -0,0 +1,119 @@
1+//
2+// GPUBootstrapTests.swift
3+// Metrika
4+//
5+// Author: Simon-Pierre Boucher
6+// Contact: contact@spboucher.ai
7+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8+//
9+
10+import Foundation
11+import Testing
12+import ZQGPU
13+import ZQParser
14+import ZQPlanner
15+import ZQStats
16+
17+/// GPU bootstrap cross-checks (CLAUDE.md §9): same seed ⇒ identical
18+/// resample indices across backends; estimates within the documented
19+/// float32 tolerance.
20+@Suite("GPU bootstrap", .enabled(if: ZQGPUBootstrap.isAvailable))
21+struct GPUBootstrapTests {
22+
23+ @Test("MLX Philox indices are bit-identical to the CPU reference")
24+ func philoxParity() {
25+ let seed: UInt64 = 42
26+ let n = 137
27+ let generator = Philox4x32(seed: seed)
28+
29+ for replicate in [0, 1, 7, 1000] {
30+ let cpu = ZQResampling.pairsBootstrapIndices(
31+ replicate: replicate, sampleSize: n, generator: generator
32+ )
33+ let gpu = ZQGPUBootstrap.gpuIndices(replicate: replicate, sampleSize: n, seed: seed)
34+ #expect(cpu == gpu, "replicate \(replicate)")
35+ }
36+ }
37+
38+ @Test("GPU coefficient draws match CPU within float32 tolerance")
39+ func drawParity() throws {
40+ // Small synthetic problem, well conditioned.
41+ let n = 200
42+ let generator = Philox4x32(seed: 7)
43+ let x = (0..<n).map { i in 5 + 10 * generator.uniform(at: UInt64(i)) }
44+ let y = (0..<n).map { i in
45+ 2 + 0.5 * x[i] + (generator.uniform(at: UInt64(n + i)) - 0.5)
46+ }
47+ let reps = 50
48+
49+ let gpuDraws = ZQGPUBootstrap.pairsBootstrapOLS(
50+ y: y, predictors: [("x", x)], replicates: reps, seed: 42
51+ )
52+
53+ let bootstrapGenerator = Philox4x32(seed: 42)
54+ for replicate in 0..<reps {
55+ let indices = ZQResampling.pairsBootstrapIndices(
56+ replicate: replicate, sampleSize: n, generator: bootstrapGenerator
57+ )
58+ let fit = try ZQOLS.fit(
59+ y: indices.map { y[$0] },
60+ predictors: [("x", indices.map { x[$0] })]
61+ )
62+ for (j, coefficient) in fit.coefficients.enumerated() {
63+ expectClose(
64+ gpuDraws[replicate][j], coefficient.estimate,
65+ rtol: 1e-4, "replicate \(replicate) b[\(coefficient.name)]"
66+ )
67+ }
68+ }
69+ }
70+
71+ @Test("GPU bootstrap is deterministic for a fixed seed")
72+ func determinism() {
73+ let n = 100
74+ let generator = Philox4x32(seed: 3)
75+ let x = (0..<n).map { i in generator.uniform(at: UInt64(i)) }
76+ let y = (0..<n).map { i in x[i] + generator.uniform(at: UInt64(n + i)) }
77+
78+ let first = ZQGPUBootstrap.pairsBootstrapOLS(
79+ y: y, predictors: [("x", x)], replicates: 20, seed: 99
80+ )
81+ let second = ZQGPUBootstrap.pairsBootstrapOLS(
82+ y: y, predictors: [("x", x)], replicates: 20, seed: 99
83+ )
84+ #expect(first == second)
85+ }
86+
87+ @Test("chunking does not change results")
88+ func chunkingInvariance() {
89+ let n = 80
90+ let generator = Philox4x32(seed: 5)
91+ let x = (0..<n).map { i in generator.uniform(at: UInt64(i)) }
92+ let y = (0..<n).map { i in 2 * x[i] + generator.uniform(at: UInt64(n + i)) }
93+
94+ let oneChunk = ZQGPUBootstrap.pairsBootstrapOLS(
95+ y: y, predictors: [("x", x)], replicates: 30, seed: 11
96+ )
97+ // Budget so small every chunk holds a single replicate.
98+ let manyChunks = ZQGPUBootstrap.pairsBootstrapOLS(
99+ y: y, predictors: [("x", x)], replicates: 30, seed: 11,
100+ memoryBudgetBytes: 1
101+ )
102+ #expect(oneChunk == manyChunks)
103+ }
104+
105+ @Test("planner sends large-reps bootstrap to the GPU")
106+ func plannerDispatch() throws {
107+ let planner = ZQPlanner(gpuAvailable: true)
108+ let parser = ZQCommandParser()
109+
110+ let large = try #require(try parser.parse("bootstrap, reps(10000): reg y x"))
111+ #expect(planner.plan(large, rowCount: 1000).backend == .gpu)
112+
113+ let small = try #require(try parser.parse("bootstrap, reps(100): reg y x"))
114+ #expect(planner.plan(small, rowCount: 1000).backend == .cpu)
115+
116+ let noGPU = ZQPlanner(gpuAvailable: false)
117+ #expect(noGPU.plan(large, rowCount: 1000).backend == .cpu)
118+ }
119+}
added MetrikaKit/Tests/MetrikaKitTests/GPUDiagnosticTests.swift +73 −0
@@ -0,0 +1,73 @@
1+//
2+// GPUDiagnosticTests.swift
3+// Metrika
4+//
5+// Author: Simon-Pierre Boucher
6+// Contact: contact@spboucher.ai
7+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8+//
9+
10+import Foundation
11+import Testing
12+import ZQGPU
13+
14+/// Numerical regression tests for the GPU cross-product stage. These
15+/// pin down an MLX batched-GEMM issue: small k×k outputs mis-accumulate
16+/// (~6e-4 relative) for batch sizes ≥ 2, which the column-wise X'X
17+/// computation in `crossProducts` works around. If these start failing,
18+/// the workaround regressed or the kernel changed.
19+@Suite("GPU numerics regression", .enabled(if: ZQGPUBootstrap.isAvailable))
20+struct GPUDiagnosticTests {
21+ let n = 200
22+ let x: [Double]
23+ let y: [Double]
24+
25+ init() {
26+ let generator = Philox4x32(seed: 7)
27+ let count = 200
28+ let xs = (0..<count).map { i in 5 + 10 * generator.uniform(at: UInt64(i)) }
29+ x = xs
30+ y = (0..<count).map { i in
31+ 2 + 0.5 * xs[i] + (generator.uniform(at: UInt64(count + i)) - 0.5)
32+ }
33+ }
34+
35+ @Test("batched cross-products match float64 sums of identical data")
36+ func crossProductPrecision() {
37+ // Exact float64 sums of the float32-quantized data the GPU sees.
38+ let indices = ZQResampling.pairsBootstrapIndices(
39+ replicate: 0, sampleSize: n, generator: Philox4x32(seed: 42)
40+ )
41+ let meanX = x.reduce(0, +) / Double(n)
42+ let varX = x.reduce(0) { $0 + ($1 - meanX) * ($1 - meanX) } / Double(n)
43+ let meanY = y.reduce(0, +) / Double(n)
44+ let xs = indices.map { Float((x[$0] - meanX) / varX.squareRoot()) }
45+ let ys = indices.map { Float(y[$0] - meanY) }
46+ var sxx = 0.0, sxy = 0.0
47+ for i in 0..<n {
48+ sxx += Double(xs[i]) * Double(xs[i])
49+ sxy += Double(xs[i]) * Double(ys[i])
50+ }
51+
52+ for batch in [1, 2, 16] {
53+ let sums = ZQGPUBootstrap.debugBatchedCrossProducts(
54+ y: y, predictors: [("x", x)], batchSize: batch, seed: 42
55+ )
56+ expectClose(Double(sums.xtx[0]), sxx, rtol: 1e-5, "sxx batch=\(batch)")
57+ expectClose(Double(sums.xty[0]), sxy, rtol: 1e-5, "sxy batch=\(batch)")
58+ }
59+ }
60+
61+ @Test("replicate draws are independent of batch size")
62+ func batchSizeInvariance() {
63+ let reference = ZQGPUBootstrap.pairsBootstrapOLS(
64+ y: y, predictors: [("x", x)], replicates: 1, seed: 42
65+ )[0]
66+ for batch in [2, 10, 50] {
67+ let draws = ZQGPUBootstrap.pairsBootstrapOLS(
68+ y: y, predictors: [("x", x)], replicates: batch, seed: 42
69+ )
70+ #expect(draws[0] == reference, "batch=\(batch)")
71+ }
72+ }
73+}
74
Menu

Theme
↑↓navigate ↵open escclose /⌘Kanywhere