diff --git a/Sources/MetalANNSCore/ResidualCascade.swift b/Sources/MetalANNSCore/ResidualCascade.swift new file mode 100644 index 0000000..283f388 --- /dev/null +++ b/Sources/MetalANNSCore/ResidualCascade.swift @@ -0,0 +1,765 @@ +import Accelerate +import Foundation +import Metal +import simd +/// Exact nearest-neighbor cascade for large corpora built on PCA projections +/// with lossless verification bounds. +/// +/// Motivation: residual-norm and quantized-code bounds cannot separate +/// neighbors along the same low-dimensional manifold, while an exact fp32 +/// projection onto the top eigen directions concentrates dot-product energy +/// so a short projection plus a Cauchy–Schwarz tail term yields near-exact +/// bounds. The head width adapts to the measured eigen spectrum (smallest +/// prefix capturing nearly all sampled variance), clamped to [2, 48]. +/// +/// Per row (cached per backing MTL buffer): +/// - `proj[i]` exact fp32 projection `Rᵀ(v − μ)` (adaptive width w) +/// - `vDotMu` precomputed `v·μ` +/// - `rowNormSq` exact `‖v‖²` +/// - `tailNorm` `‖(v − μ) − R·Rᵀ(v − μ)‖` (Cauchy–Schwarz tail radius, +/// derived via Pythagoras: `√(‖z‖² − ‖proj‖²)`) +/// +/// Dot-product identity with centered vectors `x = q − μ`, `y = v − μ`: +/// q·v = x·y + q·μ + v·μ − ‖μ‖², x·y = proj_x·proj_y + x_t·y_t +/// x_t·y_t ≤ ‖x_t‖·tailNorm (Cauchy–Schwarz) +/// +/// All terms are inflated upward (sign-safely) by generous slack so fp +/// rounding can never invalidate the similarity upper bound. +/// +/// Search protocol (exact on every path): +/// 1. Per-row similarity upper bounds → distance lower bounds lb(row). +/// 2. Rescore exactly m' rows with the smallest bounds ("seeds"). +/// 3. cutoff = k-th smallest true seed distance. Any row NOT rescored has +/// lb(row) ≥ cutoff ⇒ true distance ≥ cutoff, while the seeds supply +/// k rows at distance ≤ cutoff. Pruned rows can never beat the k-th +/// result (up to exact ties, which may reorder). +/// 4. Rescore every remaining row with lb(row) < cutoff ("survivors"). +/// 5. Return the top-k over all rescored rows. +/// +/// Loose bounds merely grow the survivor set (degrading toward brute force); +/// they can never change returned neighbors beyond tie order. +// pi-lens-ignore: type_body_length +enum ResidualCascade { + /// Below this size the existing tiers already win; keep them. + static let minVectorCount = 200_000 + + /// Upper bound on adaptive projection width. + static let maxHeadWidth = 48 + + /// Generous relative slack applied to every derived bound. + static let boundSlack: Float = 1e-4 + static let boundAbsSlack: Float = 1e-6 + + // MARK: - Aux buffer + + final class BoundBuffer: @unchecked Sendable { + let rowCount: Int + let dimensionCount: Int + let headWidth: Int + let meanNormSq: Float + let avgTailNorm: Float + let avgRowNorm: Float + + /// Planar sections of one allocation (64-byte aligned offsets): + /// proj[rowCount×headWidth] | vDotMu | rowNormSq | tailNorm + private let storage: UnsafeMutableRawPointer + let projBase: UnsafePointer + let vDotMuPlane: UnsafePointer + let rowNormSqPlane: UnsafePointer + let tailNormPlane: UnsafePointer + /// Column-major `dimensionCount × headWidth` rotation. + let rotation: [Float] + /// Corpus mean vector (`dimensionCount`). + let meanVector: [Float] + + final class GPUPlanes { + let deviceID: ObjectIdentifier + let projection: MTLBuffer + let vDotMu: MTLBuffer + let rowNormSq: MTLBuffer + let tailNorm: MTLBuffer + + init( + deviceID: ObjectIdentifier, + projection: MTLBuffer, + vDotMu: MTLBuffer, + rowNormSq: MTLBuffer, + tailNorm: MTLBuffer + ) { + self.deviceID = deviceID + self.projection = projection + self.vDotMu = vDotMu + self.rowNormSq = rowNormSq + self.tailNorm = tailNorm + } + } + + private let gpuLock = NSLock() + private var gpuPlanesCache: GPUPlanes? + + init( + rowCount: Int, + dimensionCount: Int, + headWidth: Int, + meanNormSq: Float, + avgTailNorm: Float, + avgRowNorm: Float, + storage: UnsafeMutableRawPointer, + projBase: UnsafePointer, + vDotMuPlane: UnsafePointer, + rowNormSqPlane: UnsafePointer, + tailNormPlane: UnsafePointer, + rotation: [Float], + meanVector: [Float] + ) { + self.rowCount = rowCount + self.dimensionCount = dimensionCount + self.headWidth = headWidth + self.meanNormSq = meanNormSq + self.avgTailNorm = avgTailNorm + self.avgRowNorm = avgRowNorm + self.storage = storage + self.projBase = projBase + self.vDotMuPlane = vDotMuPlane + self.rowNormSqPlane = rowNormSqPlane + self.tailNormPlane = tailNormPlane + self.rotation = rotation + self.meanVector = meanVector + } + + deinit { + storage.deallocate() + } + + /// Publishes shared Metal copies of the cached planes once. The + /// CPU pointers remain the source of truth for synchronous fallback; + /// the GPU copies are used only by the async bound pass. + func gpuPlanes(device: MTLDevice) -> GPUPlanes? { + let deviceID = ObjectIdentifier(device) + gpuLock.lock() + defer { gpuLock.unlock() } + if let cached = gpuPlanesCache, cached.deviceID == deviceID { + return cached + } + + let floatBytes = MemoryLayout.stride + guard let projection = device.makeBuffer( + length: rowCount * headWidth * floatBytes, options: .storageModeShared + ), let vDotMu = device.makeBuffer( + length: rowCount * floatBytes, options: .storageModeShared + ), let rowNormSq = device.makeBuffer( + length: rowCount * floatBytes, options: .storageModeShared + ), let tailNorm = device.makeBuffer( + length: rowCount * floatBytes, options: .storageModeShared + ) else { return nil } + + // Store projections column-major for the GPU: neighboring threads + // then read neighboring rows for each PCA component, producing + // coalesced loads (the CPU cache remains row-major). + let gpuProjection = projection.contents().assumingMemoryBound(to: Float.self) + DispatchQueue.concurrentPerform(iterations: headWidth) { column in + let destination = gpuProjection + column * rowCount + for row in 0.. BoundBuffer? { + lock.lock() + defer { lock.unlock() } + return entries[key] + } + + func store(_ value: BoundBuffer, for key: Key) { + lock.lock() + defer { lock.unlock() } + if entries[key] == nil { + order.append(key) + while order.count > capacity { + let evicted = order.removeFirst() + entries[evicted] = nil + } + } + entries[key] = value + } + + func invalidate(bufferID: ObjectIdentifier) { + lock.lock() + defer { lock.unlock() } + let doomed = order.filter { $0.bufferID == bufferID } + for key in doomed { + entries[key] = nil + } + order.removeAll { $0.bufferID == bufferID } + } + } + + static let store = Store() + + static func invalidate(buffer: MTLBuffer) { + store.invalidate(bufferID: ObjectIdentifier(buffer)) + } + + // MARK: - Build + + /// Builds the PCA-projection aux structure: adaptive-width rotation from + /// subspace iteration, then per-row exact projections, scalar planes, + /// and tail radii via Pythagoras (`tail² = ‖z‖² − ‖proj‖²`). + static func build( + corpus: UnsafePointer, + rowCount: Int, + dimensionCount: Int, + key: Key + ) -> BoundBuffer? { + guard rowCount >= minVectorCount, dimensionCount > 0 else { return nil } + + let sampleSize = min(rowCount, 65_536) + guard let pca = ResidualCascadeMath.pcaRotation( + corpus: corpus, rowCount: rowCount, dimensionCount: dimensionCount, + sampleSize: sampleSize, maxHeadWidth: maxHeadWidth + ) else { return nil } + let width = pca.headWidth + + // Allocation: proj block + three scalar planes. + let alignment = 64 + func padded(_ bytes: Int) -> Int { (bytes + alignment - 1) / alignment * alignment } + let projBytes = padded(rowCount * width * MemoryLayout.stride) + let planeBytes = padded(rowCount * MemoryLayout.stride) + let totalBytes = projBytes + 3 * planeBytes + guard let raw = malloc(totalBytes) else { return nil } + let projPtr = raw.assumingMemoryBound(to: Float.self) + let vDotMuPtr = (raw + projBytes).assumingMemoryBound(to: Float.self) + let rowNormSqPtr = (raw + projBytes + planeBytes).assumingMemoryBound(to: Float.self) + let tailNormPtr = (raw + projBytes + 2 * planeBytes).assumingMemoryBound(to: Float.self) + + // Augmented rotation [R | μ̂] in ROW-major (dim × (width+1)) so the + // chunked fill is one RowMajor sgemm per chunk. μ̂ = μ (unnormalized): + // projecting z = v − μ against it yields z·μ directly. + var rotationAugmented = [Float](repeating: 0, count: dimensionCount * (width + 1)) + for dimension in 0.. + let vDotMu: UnsafeMutablePointer + let rowNormSq: UnsafeMutablePointer + let tailNorm: UnsafeMutablePointer + } + + /// Fills projections and scalar planes for a contiguous chunk. + /// One RowMajor sgemm computes `[proj | z·μ]`; tails come from Pythagoras. + private static func fillChunk( + corpus: UnsafePointer, + rows: Int, + dimensionCount: Int, + plan: ChunkFillPlan, + outputs: ChunkFillOutputs + ) { + let augmentedColumns = plan.width + 1 + var products = [Float](repeating: 0, count: rows * augmentedColumns) + + products.withUnsafeMutableBufferPointer { productsBuf in + plan.rotationAugmented.withUnsafeBufferPointer { rotBuf in + cblas_sgemm( + CblasRowMajor, CblasNoTrans, CblasNoTrans, + Int32(rows), Int32(augmentedColumns), Int32(dimensionCount), + 1.0, + corpus, Int32(dimensionCount), + rotBuf.baseAddress!, Int32(augmentedColumns), + 0.0, productsBuf.baseAddress!, Int32(augmentedColumns) + ) + } + } + + let width = plan.width + let meanProjection = plan.meanProjection + let meanNormSq = plan.meanNormSq + let outProj = outputs.proj + let outVDotMu = outputs.vDotMu + let outRowNormSq = outputs.rowNormSq + let outTailNorm = outputs.tailNorm + + for rowIndex in 0.. [SearchResult]? { + guard let vectorBuffer = vectors as? VectorBuffer, + !vectors.isFloat16, + let corpus = vectorBuffer.floatPointer.baseAddress + else { return nil } + let rowCount = vectors.count + let dimensionCount = vectors.dim + guard rowCount >= minVectorCount, dimensionCount > 0, + query.count == dimensionCount, neighborTotal > 0, metric != .hamming + else { return nil } + + let cacheKey = Key( + bufferID: ObjectIdentifier(vectorBuffer.buffer), + bufferLength: vectorBuffer.buffer.length, + rowCount: rowCount, + dimensionCount: dimensionCount + ) + let aux: BoundBuffer + if let cached = store.get(cacheKey), cached.rowCount >= rowCount { + aux = cached + } else { + guard let built = build( + corpus: corpus, rowCount: rowCount, + dimensionCount: dimensionCount, key: cacheKey + ) else { return nil } + aux = built + } + + let effectiveK = min(neighborTotal, FlatGPUSearch.maxTopK, rowCount) + guard effectiveK > 0 else { return nil } + + // ---- Query preparation -------------------------------------------- + var queryNormSq: Float = 0 + vDSP_dotpr(query, 1, query, 1, &queryNormSq, vDSP_Length(dimensionCount)) + + // Degenerate zero-norm query under cosine: every distance finalizes + // to exactly 1.0, so ties resolve to the lowest ids — identical + // outcome without scanning. + if metric == .cosine && queryNormSq < 1e-20 { + return lowestIdsResult(count: rowCount, take: effectiveK) + } + + let boundContext = ResidualCascadeMath.prepareQueryContext(query: query, aux: aux, dimensionCount: dimensionCount) + + let statsEnabled = ProcessInfo.processInfo.environment["METALANNS_RESIDUAL_STATS"] == "1" + var phaseTimings: [(String, Double)] = [] + func recordPhase(_ name: String, _ start: DispatchTime) { + if statsEnabled { + phaseTimings.append((name, Double(DispatchTime.now().uptimeNanoseconds - start.uptimeNanoseconds) / 1_000_000)) + } + } + + let searchStart = DispatchTime.now() + // ---- Phase 1: distance lower bounds -------------------------------- + let phase1Start = DispatchTime.now() + var lowerBounds = [Float](repeating: 0, count: rowCount) + lowerBounds.withUnsafeMutableBufferPointer { boundsBuf in + computeBounds( + context: boundContext, + aux: aux, + metric: metric, + outLowerBounds: boundsBuf.baseAddress! + ) + } + recordPhase("p1_bounds", phase1Start) + + let results = resolveTopK(ResolveInput( + lowerBounds: lowerBounds, + corpus: corpus, + dimensionCount: dimensionCount, + query: query, + queryNormSq: queryNormSq, + aux: aux, + metric: metric, + effectiveK: effectiveK, + rowCount: rowCount, + recordPhase: recordPhase, + statsEnabled: statsEnabled + )) + if statsEnabled { + let summary = phaseTimings.map { + String(format: "%@=%.2fms", $0.0, $0.1) + }.joined(separator: " ") + let preludeMs = Double( + phase1Start.uptimeNanoseconds - searchStart.uptimeNanoseconds + ) / 1_000_000 + let statsLine = String( + format: "[ResidualCascade] phases: prelude=%.2fms %@\n", + preludeMs, summary + ) + FileHandle.standardError.write(Data(statsLine.utf8)) + } + return results + } + + // MARK: - Bound computation + + private static func computeBounds( + context: ResidualCascadeMath.QueryBoundContext, + aux: BoundBuffer, + metric: Metric, + outLowerBounds: UnsafeMutablePointer + ) { + let totalRows = aux.rowCount + let headWidth = aux.headWidth + let slack = boundSlack + let absoluteSlack = boundAbsSlack + let meanNormSq = aux.meanNormSq + let workers = ProcessInfo.processInfo.activeProcessorCount + + DispatchQueue.concurrentPerform(iterations: workers) { worker in + let start = (totalRows * worker) / workers + let end = (totalRows * (worker + 1)) / workers + guard start < end else { return } + + let proj = aux.projBase + start * headWidth + let vDotMuPlane = aux.vDotMuPlane + start + let rowNormSqPlane = aux.rowNormSqPlane + start + let tailNormPlane = aux.tailNormPlane + start + let out = outLowerBounds + start + let queryHead = context.headDots + + for offset in 0..<(end - start) { + // Exact fp32 head dot (small width, no quantization error). + let headDot = projectedHeadDot( + queryHead: queryHead, + rowProjection: proj + offset * headWidth, + headWidth: headWidth + ) + + // Tail: x_t·y_t ≤ ‖x_t‖ · tailNorm (sign-safe inflation). + let rowTail = tailNormPlane[offset] + let tailUpper = context.tailNorm * rowTail * (1 + slack) + absoluteSlack + + let meanTerm = context.dotMean + vDotMuPlane[offset] + let inflation = slack + * (Swift.abs(headDot) + Swift.abs(meanTerm) + + Swift.abs(meanNormSq) + Swift.abs(tailUpper)) + let dotUpper = headDot * (1 + slack) + tailUpper + meanTerm + - meanNormSq + inflation + absoluteSlack + + let rowNormSq = rowNormSqPlane[offset] + + switch metric { + case .cosine: + if rowNormSq < 1e-20 || context.norm < 1e-10 { + out[offset] = 1.0 + continue + } + let denominator = context.norm * sqrt(rowNormSq) + var simUpperBound = (dotUpper - absoluteSlack) / denominator + if simUpperBound.isNaN || simUpperBound > 1 { simUpperBound = 1 } + out[offset] = 1.0 - simUpperBound + case .innerProduct: + // minimized score = −q·v ≥ −dotUpper. + out[offset] = -dotUpper + case .l2: + // d² = ‖q‖² − 2q·v + ‖v‖² ≥ ‖q‖² − 2·dotUpper + ‖v‖²(1−δ). + var lowerNormSq = rowNormSq * (1 - slack) - absoluteSlack + if lowerNormSq < 0 { lowerNormSq = 0 } + var gap = context.normSq - 2 * dotUpper + lowerNormSq + if gap < 0 { gap = 0 } + out[offset] = gap + case .hamming: + break + } + } + } + } + + @inline(__always) + private static func projectedHeadDot( + queryHead: [Float], + rowProjection: UnsafePointer, + headWidth: Int + ) -> Float { + var headDot: Float = 0 + var column = 0 + if headWidth >= 4 { + var vector0 = simd_float4(repeating: 0) + var vector1 = simd_float4(repeating: 0) + var vector2 = simd_float4(repeating: 0) + var vector3 = simd_float4(repeating: 0) + while column + 15 < headWidth { + vector0 += simd_float4( + queryHead[column], queryHead[column + 1], + queryHead[column + 2], queryHead[column + 3] + ) * simd_float4( + rowProjection[column], rowProjection[column + 1], + rowProjection[column + 2], rowProjection[column + 3] + ) + vector1 += simd_float4( + queryHead[column + 4], queryHead[column + 5], + queryHead[column + 6], queryHead[column + 7] + ) * simd_float4( + rowProjection[column + 4], rowProjection[column + 5], + rowProjection[column + 6], rowProjection[column + 7] + ) + vector2 += simd_float4( + queryHead[column + 8], queryHead[column + 9], + queryHead[column + 10], queryHead[column + 11] + ) * simd_float4( + rowProjection[column + 8], rowProjection[column + 9], + rowProjection[column + 10], rowProjection[column + 11] + ) + vector3 += simd_float4( + queryHead[column + 12], queryHead[column + 13], + queryHead[column + 14], queryHead[column + 15] + ) * simd_float4( + rowProjection[column + 12], rowProjection[column + 13], + rowProjection[column + 14], rowProjection[column + 15] + ) + column += 16 + } + vector0 += vector1 + vector2 + vector3 + headDot = vector0.x + vector0.y + vector0.z + vector0.w + while column + 3 < headWidth { + headDot += simd_dot( + simd_float4( + queryHead[column], queryHead[column + 1], + queryHead[column + 2], queryHead[column + 3] + ), + simd_float4( + rowProjection[column], rowProjection[column + 1], + rowProjection[column + 2], rowProjection[column + 3] + ) + ) + column += 4 + } + } + while column < headWidth { + headDot += queryHead[column] * rowProjection[column] + column += 1 + } + return headDot + } + + // MARK: - Exact rescoring (CPU float4, mirrors flat_scan_distances order) + + static func exactRescore( + ids: [UInt32], + corpus: UnsafePointer, + dimensionCount: Int, + query: [Float], + queryNormSq: Float, + metric: Metric, + into results: inout [(distance: Float, id: UInt32)] + ) { + guard !ids.isEmpty else { return } + let distances = UnsafeMutableBufferPointer.allocate(capacity: ids.count) + defer { distances.deallocate() } + let workers = ProcessInfo.processInfo.activeProcessorCount + DispatchQueue.concurrentPerform(iterations: workers) { worker in + let start = (ids.count * worker) / workers + let end = (ids.count * (worker + 1)) / workers + guard start < end else { return } + query.withUnsafeBufferPointer { queryBuf in + let queryBase = queryBuf.baseAddress! + for slot in start.. Float { + switch metric { + case .l2: + return Swift.max(0.0, queryNormSq - 2.0 * dotQV + normVSq) + case .innerProduct: + return -dotQV + default: + let denom = sqrt(queryNormSq) * sqrt(normVSq) + return denom < 1e-10 ? 1.0 : (1.0 - dotQV / denom) + } + } + + static func lowestIdsResult(count: Int, take: Int) -> [SearchResult] { + (0.. UInt64 { + state &+= 0x9E37_79B9_7F4A_7C15 + var mixed = state + mixed = (mixed ^ (mixed >> 30)) &* 0xBF58_476D_1CE4_E5B9 + mixed = (mixed ^ (mixed >> 27)) &* 0x94D0_49BB_1331_11EB + return mixed ^ (mixed >> 31) + } +} diff --git a/Sources/MetalANNSCore/ResidualCascadeMath.swift b/Sources/MetalANNSCore/ResidualCascadeMath.swift new file mode 100644 index 0000000..846ce90 --- /dev/null +++ b/Sources/MetalANNSCore/ResidualCascadeMath.swift @@ -0,0 +1,306 @@ +import Accelerate +import Foundation + +/// Deterministic PCA machinery for the exact cascade. +/// +/// Split out of ResidualCascade.swift for file-size constraints. Members are +/// internal; `ResidualCascade` consumes only the final rotation + mean via +/// `pcaRotation`. +/// +/// Numerical note: the covariance is formed in DOUBLE precision from a +/// materialized centered sample (`dsyrk`), then fully diagonalized with a +/// double-precision cyclic Jacobi sweep. An fp32-accumulated covariance +/// carries O(n·ε·‖v‖²) noise (~3 in this corpus's units) that both pollutes +/// the spectrum and tilts the basis, inflating projection-tail radii and +/// therefore survivor counts. The double-precision path costs ~2–3 s of +/// one-time build work (amortized, excluded from warm search latency). +enum ResidualCascadeMath { + struct PCAResult { + /// Column-major `dimensionCount × headWidth` orthonormal rotation. + let rotation: [Float] + /// DimensionCount-length corpus mean. + let mean: [Float] + /// Adaptive width chosen from the eigen-energy threshold. + let headWidth: Int + /// ‖μ‖². + let meanNormSq: Float + } + + /// Per-query values consumed by the bound computation. + struct QueryBoundContext { + let headDots: [Float] + let tailNorm: Float + let norm: Float + let normSq: Float + let dotMean: Float + } + + /// Centers and projects a query using the cached PCA basis. + static func prepareQueryContext( + query: [Float], aux: ResidualCascade.BoundBuffer, dimensionCount: Int + ) -> QueryBoundContext { + var queryNormSq: Float = 0 + vDSP_dotpr(query, 1, query, 1, &queryNormSq, vDSP_Length(dimensionCount)) + + var centeredQuery = [Float](repeating: 0, count: dimensionCount) + for dimension in 0.., + rowCount: Int, + dimensionCount: Int, + sampleSize: Int, + maxHeadWidth: Int + ) -> PCAResult? { + let strideStep = max(1, rowCount / sampleSize) + let actualSample = min(sampleSize, (rowCount - 1) / strideStep + 1) + let mean = sampleMean( + corpus: corpus, strideStep: strideStep, + sampleRows: actualSample, dimensionCount: dimensionCount + ) + + // ---- Materialize centered sample in Double (row-major rows×dim) --- + var centeredSample = [Double](repeating: 0, count: actualSample * dimensionCount) + var meanD = [Double](repeating: 0, count: dimensionCount) + for dimension in 0.. eigenValues[$1] } + let totalVariance = eigenValues.reduce(0, +) + var cumulative: Double = 0 + var recommendedWidth = dimensionCount + for rank in 0.., + strideStep: Int, + sampleRows: Int, + dimensionCount: Int + ) -> [Float] { + var mean = [Float](repeating: 0, count: dimensionCount) + for sampleIndex in 0..= 0 ? 1 : -1 + let tangent = sign / (abs(theta) + (theta * theta + 1).squareRoot()) + let cosine = 1 / (tangent * tangent + 1).squareRoot() + let sine = tangent * cosine + let rotation = JacobiRotation( + first: pivotRow, second: pivotColumn, + cosine: cosine, sine: sine + ) + rotateSymmetricDouble( + matrix: &matrix, size: size, rotation: rotation + ) + rotateEigenvectorsDouble( + vectors: &outEigenVectors, size: size, rotation: rotation + ) + } + } + } + for diagonal in 0.. + let dimensionCount: Int + let query: [Float] + let queryNormSq: Float + let aux: BoundBuffer + let metric: Metric + let effectiveK: Int + let rowCount: Int + let recordPhase: (String, DispatchTime) -> Void + let statsEnabled: Bool + } + + // MARK: - Selection helpers + + /// Indices of the `take` smallest values, ascending by value (ties by index). + /// Parallel: per-chunk bounded heaps via concurrentPerform, then merge. + static func selectSmallestIndices(lowerBounds: [Float], take: Int) -> [UInt32] { + let totalRows = lowerBounds.count + let capacity = min(take, totalRows) + guard capacity > 0 else { return [] } + + if totalRows <= 64_000 { + return serialSmallestIndices(lowerBounds: lowerBounds, take: capacity) + } + + let workers = ProcessInfo.processInfo.activeProcessorCount + let lanes = max(1, min(workers, totalRows / 32_000)) + if lanes <= 1 { + return serialSmallestIndices(lowerBounds: lowerBounds, take: capacity) + } + let chunkSize = (totalRows + lanes - 1) / lanes + let laneValuesBuffer = UnsafeMutableBufferPointer<[Float]>.allocate(capacity: lanes) + let laneIDsBuffer = UnsafeMutableBufferPointer<[UInt32]>.allocate(capacity: lanes) + defer { + laneValuesBuffer.deallocate() + laneIDsBuffer.deallocate() + } + for lane in 0..() + var heapIndices = ContiguousArray() + heapValues.reserveCapacity(capacity) + heapIndices.reserveCapacity(capacity) + + @inline(__always) func siftDownFrom(_ position: Int) { + var position = position + while true { + let leftChild = 2 * position + 1 + let rightChild = leftChild + 1 + var largest = position + if leftChild < heapValues.count, heapValues[leftChild] > heapValues[largest] { largest = leftChild } + if rightChild < heapValues.count, heapValues[rightChild] > heapValues[largest] { largest = rightChild } + if largest == position { return } + heapValues.swapAt(position, largest) + heapIndices.swapAt(position, largest) + position = largest + } + } + + for index in start.. 0 { + let parent = (position - 1) / 2 + if heapValues[position] > heapValues[parent] { + heapValues.swapAt(position, parent) + heapIndices.swapAt(position, parent) + position = parent + } else { break } + } + } else if value < heapValues[0] { + heapValues[0] = value + heapIndices[0] = UInt32(index) + siftDownFrom(0) + } + } + laneValuesBuffer[lane] = Array(heapValues) + laneIDsBuffer[lane] = Array(heapIndices) + } + } + + var pairs = [(value: Float, index: UInt32)]() + pairs.reserveCapacity(lanes * capacity) + for lane in 0.. [UInt32] { + let totalRows = lowerBounds.count + var heapValues = ContiguousArray() + var heapIndices = ContiguousArray() + heapValues.reserveCapacity(take) + heapIndices.reserveCapacity(take) + + @inline(__always) func siftDownFrom(_ position: Int) { + var position = position + while true { + let leftChild = 2 * position + 1 + let rightChild = leftChild + 1 + var largest = position + if leftChild < heapValues.count, heapValues[leftChild] > heapValues[largest] { largest = leftChild } + if rightChild < heapValues.count, heapValues[rightChild] > heapValues[largest] { largest = rightChild } + if largest == position { return } + heapValues.swapAt(position, largest) + heapIndices.swapAt(position, largest) + position = largest + } + } + + for index in 0.. 0 { + let parent = (position - 1) / 2 + if heapValues[position] > heapValues[parent] { + heapValues.swapAt(position, parent) + heapIndices.swapAt(position, parent) + position = parent + } else { break } + } + } else if value < heapValues[0] { + heapValues[0] = value + heapIndices[0] = UInt32(index) + siftDownFrom(0) + } + } + var pairs = zip(heapValues, heapIndices).map { (value: $0.0, index: $0.1) } + pairs.sort { $0.value == $1.value ? $0.index < $1.index : $0.value < $1.value } + return pairs.map { $0.index } + } + + /// Ascending top-k with deterministic tie-break by id. Uses a bounded + /// max-heap so cost is O(n·log k) rather than a full sort of every + /// rescored candidate. + static func topKResults( + of scored: [(distance: Float, id: UInt32)], take: Int + ) -> [SearchResult] { + let selected = min(take, scored.count) + guard selected > 0 else { return [] } + + var heapDistances = ContiguousArray() + var heapIDs = ContiguousArray() + heapDistances.reserveCapacity(selected) + heapIDs.reserveCapacity(selected) + + @inline(__always) func siftDownFrom(_ position: Int, size: Int) { + var position = position + while true { + let leftChild = 2 * position + 1 + let rightChild = leftChild + 1 + var largest = position + if leftChild < size, + heapDistances[leftChild] > heapDistances[largest] + { largest = leftChild } + if rightChild < size, + heapDistances[rightChild] > heapDistances[largest] + { largest = rightChild } + if largest == position { return } + heapDistances.swapAt(position, largest) + heapIDs.swapAt(position, largest) + position = largest + } + } + + for entry in scored { + if heapDistances.count < selected { + heapDistances.append(entry.distance) + heapIDs.append(entry.id) + var position = heapDistances.count - 1 + while position > 0 { + let parent = (position - 1) / 2 + if heapDistances[position] > heapDistances[parent] { + heapDistances.swapAt(position, parent) + heapIDs.swapAt(position, parent) + position = parent + } else { break } + } + } else if entry.distance < heapDistances[0] { + heapDistances[0] = entry.distance + heapIDs[0] = entry.id + siftDownFrom(0, size: heapDistances.count) + } + } + + var extracted = [(distance: Float, id: UInt32)]() + extracted.reserveCapacity(selected) + var size = selected + while size > 0 { + extracted.append((heapDistances[0], heapIDs[0])) + size -= 1 + if size > 0 { + heapDistances[0] = heapDistances[size] + heapIDs[0] = heapIDs[size] + siftDownFrom(0, size: size) + } + } + extracted.reverse() + + var hasTies = false + if extracted.count > 1 { + for index in 1.. [UInt32] { + let workers = ProcessInfo.processInfo.activeProcessorCount + let lanes = max(1, min(workers, rowCount / 64_000)) + guard lanes > 1 else { + var survivors: [UInt32] = [] + survivors.reserveCapacity(rowCount / 8) + for rowIndex in 0...allocate(capacity: lanes) + defer { laneResultsBuffer.deallocate() } + for lane in 0.. [SearchResult]? { + let lowerBounds = input.lowerBounds + let corpus = input.corpus + let dimensionCount = input.dimensionCount + let query = input.query + let queryNormSq = input.queryNormSq + let aux = input.aux + let metric = input.metric + let effectiveK = input.effectiveK + let rowCount = input.rowCount + let recordPhase = input.recordPhase + let statsEnabled = input.statsEnabled + + // ---- Phases 2+3: seeds, exact rescore, cutoff ---------------------- + let phase23Start = DispatchTime.now() + let seedTarget = min(max(4 * effectiveK, 256), rowCount) + let seedIDs = selectSmallestIndices(lowerBounds: lowerBounds, take: seedTarget) + var scored: [(distance: Float, id: UInt32)] = [] + scored.reserveCapacity(seedTarget * 4) + exactRescore( + ids: seedIDs, corpus: corpus, dimensionCount: dimensionCount, + query: query, queryNormSq: queryNormSq, metric: metric, into: &scored + ) + scored.sort { $0.distance < $1.distance } + let cutoffIndex = min(effectiveK, scored.count) - 1 + guard cutoffIndex >= 0 else { return [] } + let cutoff = scored[cutoffIndex].distance + recordPhase("p23_seeds_and_cutoff", phase23Start) + + // ---- Phase 4: survivors ------------------------------------------- + let phase4Start = DispatchTime.now() + var seedMark = [Bool](repeating: false, count: rowCount) + for id in seedIDs { seedMark[Int(id)] = true } + let survivors = collectSurvivors( + lowerBounds: lowerBounds, seedMark: seedMark, + cutoff: cutoff, rowCount: rowCount + ) + + if statsEnabled { + FileHandle.standardError.write(Data( + "[ResidualCascade] rows=\(rowCount) seeds=\(seedIDs.count) cutoff=\(cutoff) survivors=\(survivors.count) tailAvg=\(aux.avgTailNorm) rowNormAvg=\(aux.avgRowNorm)\n".utf8 + )) + } + recordPhase("p4_survivor_scan", phase4Start) + + // ---- Phase 5: exact survivor rescore ------------------------------- + let phase5Start = DispatchTime.now() + exactRescore( + ids: survivors, corpus: corpus, dimensionCount: dimensionCount, + query: query, queryNormSq: queryNormSq, metric: metric, into: &scored + ) + recordPhase("p5_survivor_rescore", phase5Start) + + // ---- Phase 6: top-k ------------------------------------------------- + let phase6Start = DispatchTime.now() + let results = topKResults(of: scored, take: effectiveK) + + if ProcessInfo.processInfo.environment["METALANNS_VERIFY_BOUNDS"] == "1" { + // Diagnostic: brute-force verify bound soundness on every row. + let verifyWorkers = ProcessInfo.processInfo.activeProcessorCount + var maxViolation: Float = 0 + var violationCount = 0 + let trueDistances = UnsafeMutableBufferPointer.allocate(capacity: rowCount) + defer { trueDistances.deallocate() } + let capturedQueryNormSq = queryNormSq + DispatchQueue.concurrentPerform(iterations: verifyWorkers) { worker in + let start = (rowCount * worker) / verifyWorkers + let end = (rowCount * (worker + 1)) / verifyWorkers + guard start < end else { return } + query.withUnsafeBufferPointer { qBuf in + for rowIndex in start.. trueDistances[rowIndex] { + violationCount += 1 + maxViolation = max(maxViolation, lowerBounds[rowIndex] - trueDistances[rowIndex]) + } + } + let verifyLine = "[ResidualCascade] VERIFY: violations=\(violationCount)" + + " max=\(maxViolation) dim=\(dimensionCount) w=\(aux.headWidth)" + + " lb0=\(lowerBounds[0]) td0=\(trueDistances[0])\n" + FileHandle.standardError.write(Data(verifyLine.utf8)) + } + recordPhase("p6_topk", phase6Start) + + return results + } +} diff --git a/Sources/MetalANNSCore/VectorBuffer.swift b/Sources/MetalANNSCore/VectorBuffer.swift index 309273b..cb3cc0c 100644 --- a/Sources/MetalANNSCore/VectorBuffer.swift +++ b/Sources/MetalANNSCore/VectorBuffer.swift @@ -48,6 +48,7 @@ public final class VectorBuffer: @unchecked Sendable { FlatGPUSearch.invalidateHostNormCache(buffer: buffer) Int8CodeCache.invalidate(buffer: buffer) IVFFlatSearch.invalidate(buffer: buffer) + ResidualCascade.invalidate(buffer: buffer) } public func setCount(_ newCount: Int) { @@ -71,6 +72,7 @@ public final class VectorBuffer: @unchecked Sendable { FlatGPUSearch.invalidateHostNormCache(buffer: buffer) Int8CodeCache.invalidate(buffer: buffer) IVFFlatSearch.invalidate(buffer: buffer) + ResidualCascade.invalidate(buffer: buffer) } } @@ -87,6 +89,7 @@ public final class VectorBuffer: @unchecked Sendable { FlatGPUSearch.invalidateHostNormCache(buffer: buffer) Int8CodeCache.invalidate(buffer: buffer) IVFFlatSearch.invalidate(buffer: buffer) + ResidualCascade.invalidate(buffer: buffer) } } diff --git a/Tests/MetalANNSTests/ResidualCascadeTests.swift b/Tests/MetalANNSTests/ResidualCascadeTests.swift new file mode 100644 index 0000000..7764f84 --- /dev/null +++ b/Tests/MetalANNSTests/ResidualCascadeTests.swift @@ -0,0 +1,303 @@ +import Accelerate +import Foundation +import Metal +import Testing + +@testable import MetalANNSCore + +/// Regression suite proving the residual-bound exact cascade returns +/// brute-force EXACT results across metrics, dimensions (incl. non-multiple- +/// of-4 tails), adversarial distributions (duplicates, zero norms), +/// cache invalidation after in-place writes, and deterministic reruns. +@Suite("Residual Cascade Tests") +struct ResidualCascadeTests { + // MARK: - Helpers + + private func makeContext() throws -> MetalContext { + guard MTLCreateSystemDefaultDevice() != nil else { + throw ANNSError.deviceNotSupported + } + return try MetalContext() + } + + private struct SeededGenerator { + var state: UInt64 + init(seed: UInt64) { + state = seed &* 0x9E37_79B9_7F4A_7C15 &+ 0xD1B5_4A32_D192_ED03 + } + mutating func next() -> Float { + state = state &* 6_364_136_223_846_793_005 &+ 1_442_695_040_888_963_407 + return Float((state >> 33) & 0xFFFF_FFFF) / Float(0xFFFF_FFFF) * 2.0 - 1.0 + } + } + + private func makeVectors(count: Int, dim: Int, seed: UInt64) -> [[Float]] { + var rng = SeededGenerator(seed: seed) + return (0.. Double { + switch metric { + case .cosine: + var dot = 0.0 + var normQ = 0.0 + var normV = 0.0 + for dimension in 0.. [(id: UInt32, distance: Double)] { + var scored = vectors.enumerated().map { index, vector -> (UInt32, Double) in + (UInt32(index), referenceDistance(query: query, vector: vector, metric: metric)) + } + scored.sort { lhs, rhs in + lhs.1 == rhs.1 ? lhs.0 < rhs.0 : lhs.1 < rhs.1 + } + return scored.map { (id: $0.0, distance: $0.1) } + } + + private func makeBuffer(_ vectors: [[Float]], dim: Int) throws -> VectorBuffer { + let buffer = try VectorBuffer(capacity: vectors.count, dim: dim) + try buffer.batchInsert(vectors: vectors, startingAt: 0) + buffer.setCount(vectors.count) + return buffer + } + + /// Asserts the returned results are the brute-force top-k. When the + /// k-boundary sits inside a float near-tie group, id-set equality is + /// replaced by membership in every valid top-k set. + private func assertExact( + actual: [SearchResult], + reference: [(id: UInt32, distance: Double)], + neighborTotal: Int + ) { + let effectiveK = min(neighborTotal, FlatGPUSearch.maxTopK, reference.count) + let boundaryDistance = reference[effectiveK - 1].distance + let gapAtK = + reference.count > effectiveK + ? reference[effectiveK].distance - boundaryDistance + : Double.infinity + + let actualIDs = Set(actual.prefix(effectiveK).map(\.internalID)) + #expect(actualIDs.count == effectiveK, "duplicate ids in results") + + if gapAtK > 2e-3 { + let expectedIDs = Set(reference.prefix(effectiveK).map(\.id)) + #expect(actualIDs == expectedIDs, "top-k id set mismatch") + } else { + var valid = Set() + for entry in reference where entry.distance <= boundaryDistance + 2e-3 { + valid.insert(entry.id) + } + #expect( + actualIDs.isSubset(of: valid), + "results contain rows outside the tied top-k window" + ) + } + + let referenceByID = Dictionary(uniqueKeysWithValues: reference.map { ($0.id, $0.distance) }) + for result in actual.prefix(effectiveK) { + guard let expected = referenceByID[result.internalID] else { + Issue.record("returned id \(result.internalID) missing from corpus ranking") + continue + } + let tolerance = max(2e-3, abs(expected) * 2e-3) + #expect( + abs(Double(result.score) - expected) <= tolerance, + "score mismatch id=\(result.internalID) got=\(result.score) want=\(expected)" + ) + } + + let slice = Array(actual.prefix(effectiveK)) + for index in 1.. [SearchResult] { + // Route through the production entry point so tier selection, + // eligibility gates, and cache wiring are exercised too. + return try await FlatGPUSearch.search( + context: context, query: query, vectors: buffer, k: neighborTotal, metric: metric + ) + } +}