From 55b1b9b71c3807c78034dae6169d8deb63273b78 Mon Sep 17 00:00:00 2001 From: Paullo9 Date: Sat, 31 Oct 2026 14:58:08 +0100 Subject: [PATCH 1/2] perf: ACOM Metal path 1.9x faster at 600 beam directions, 3.2x at 2400 Measured first with a skipped-by-default benchmark (ACOMMetalBenchmarkTests, DIFFRACTOKIT_ACOM_BENCHMARK=1), 128 x 128 scan, one CPU preparation worker: 600 axes 0.259 -> 0.133 to 0.136 ms per pattern, 2400 axes 1.131 -> 0.351 to 0.353. The GPU time was the bound: the spectrum kernel read the whole template bank twice per position, and the inverse transform was a direct n^2 sum. - One spectrum thread now computes both passes for a block of 8 positions, so each template value is read once per 8 positions; the sums run over the rings in the same order with the same products. - The inverse transform is a radix-2 FFT in threadgroup memory. - Metal against the CPU reference: worst score gap 1.2e-7 to 1.8e-7 (bound 1e-6). New ACOMMetalAzimuthalLengthTests check 4 to 1024 bins and a partial position block: scores, runner-up, template, and the bin within a pass. Done: Metal timing of issue 8 measured and the kernel sped up. Next: an all-cores CPU against Metal row on a quiet machine for docs/benchmarks.md. Co-Authored-By: Claude Opus 5.5 --- Sources/DiffractoOrientation/ACOMMetal.swift | 165 +++++++++++------- .../ACOMMetalScaleTests.swift | 126 +++++++++++++ 2 files changed, 228 insertions(+), 63 deletions(-) create mode 100644 Tests/DiffractoOrientationTests/ACOMMetalScaleTests.swift diff --git a/Sources/DiffractoOrientation/ACOMMetal.swift b/Sources/DiffractoOrientation/ACOMMetal.swift index aa22325..5cc8535 100644 --- a/Sources/DiffractoOrientation/ACOMMetal.swift +++ b/Sources/DiffractoOrientation/ACOMMetal.swift @@ -5,10 +5,14 @@ // template, the correlation spectrum Σ_ring conj(E)·T and its inverse transform over the azimuth, and returns the best // bin and its score (the forward and the mirrored pass); the CPU then applies the reference's selection of the winner and // the distinct runner-up and its symmetry reduction. Ported in structure from the mac4DSTEM Metal matcher -// (`MetalACOMMatcher`, `ACOMMatching.metal`), with two changes: the inverse transform reads a precomputed Float cos/sin -// table indexed by `(f * a) mod n` instead of calling sin/cos on an argument up to 2π·n (which loses 3 to 4 digits in -// Float), and the shader is compiled at run time from the source below with safe (non-fast) math, so the package needs no -// Metal build step and compiles unchanged where Metal is not used. +// (`MetalACOMMatcher`, `ACOMMatching.metal`), with these changes: the inverse transform reads a precomputed Float cos/sin +// table instead of calling sin/cos on an argument up to 2π·n (which loses 3 to 4 digits in Float); the shader is +// compiled at run time from the source below with safe (non-fast) math, so the package needs no Metal build step and +// compiles unchanged where Metal is not used; and (issue 8, measured) one spectrum thread computes both passes for 8 +// positions, so each template value is read once per 8 positions instead of twice per position, and the inverse +// transform is a radix-2 FFT in threadgroup memory instead of a direct n² sum. At 16 384 patterns that took the Metal +// path from 0.259 to 0.136 ms per pattern at 600 beam directions and from 1.131 to 0.353 at 2 400 (one CPU preparation +// worker; `ACOMMetalBenchmarkTests`), with the scores still within 1e-6 of the CPU reference. // import Foundation @@ -104,7 +108,7 @@ private final class ACOMMetalContext { let templates: Int, radial: Int, azimuthal: Int let batchSize: Int - struct Params { var positions: UInt32, templates: UInt32, radial: UInt32, azimuthal: UInt32, sign: Float } + struct Params { var positions: UInt32, templates: UInt32, radial: UInt32, azimuthal: UInt32 } init?(library: ACOMTemplateLibrary) { guard let device = MTLCreateSystemDefaultDevice(), let queue = device.makeCommandQueue() else { return nil } @@ -130,8 +134,8 @@ private final class ACOMMetalContext { self.device = device; self.queue = queue; self.spectrum = p1; self.argmax = p2 templateRe = tr; templateIm = ti; cosTable = ct; sinTable = st templates = library.count; radial = library.geometry.nRadial; azimuthal = n - // Per position: both spectra, the scores and bins of two passes, and the experimental grid (the reference app's budget). - let perPosition = templates * n * 2 * 4 + 2 * templates * 8 + radial * n * 2 * 4 + // Per position: the spectra of both passes, their scores and bins, and the experimental grid (the reference app's budget). + let perPosition = 2 * templates * n * 2 * 4 + 2 * templates * 8 + radial * n * 2 * 4 let budget = min(Int(device.recommendedMaxWorkingSetSize / 8), 128 << 20) batchSize = max(1, min(64, budget / max(1, perPosition))) } @@ -144,73 +148,94 @@ private final class ACOMMetalContext { private func runUnpooled(re: [Float], im: [Float], positions: Int) -> (forwardScores: [Float], forwardBins: [UInt32], mirrorScores: [Float], mirrorBins: [UInt32])? { - let spectrumCount = positions * templates * azimuthal, resultCount = positions * templates + // One spectrum buffer per pass (pass 0 first, then pass 1), and one score and bin per pass, position and template. + let spectrumCount = 2 * positions * templates * azimuthal, resultCount = positions * templates guard let inRe = device.makeBuffer(bytes: re, length: re.count * 4, options: .storageModeShared), let inIm = device.makeBuffer(bytes: im, length: im.count * 4, options: .storageModeShared), let specRe = device.makeBuffer(length: spectrumCount * 4, options: .storageModePrivate), let specIm = device.makeBuffer(length: spectrumCount * 4, options: .storageModePrivate), - let scores0 = device.makeBuffer(length: resultCount * 4, options: .storageModeShared), - let bins0 = device.makeBuffer(length: resultCount * 4, options: .storageModeShared), - let scores1 = device.makeBuffer(length: resultCount * 4, options: .storageModeShared), - let bins1 = device.makeBuffer(length: resultCount * 4, options: .storageModeShared), + let scores = device.makeBuffer(length: 2 * resultCount * 4, options: .storageModeShared), + let bins = device.makeBuffer(length: 2 * resultCount * 4, options: .storageModeShared), let command = queue.makeCommandBuffer() else { return nil } - for (pass, scores, bins) in [(0, scores0, bins0), (1, scores1, bins1)] { - var params = Params(positions: UInt32(positions), templates: UInt32(templates), radial: UInt32(radial), - azimuthal: UInt32(azimuthal), sign: pass == 0 ? 1 : -1) - guard let e1 = command.makeComputeCommandEncoder() else { return nil } - e1.setComputePipelineState(spectrum) - e1.setBuffer(inRe, offset: 0, index: 0); e1.setBuffer(inIm, offset: 0, index: 1) - e1.setBuffer(templateRe, offset: 0, index: 2); e1.setBuffer(templateIm, offset: 0, index: 3) - e1.setBuffer(specRe, offset: 0, index: 4); e1.setBuffer(specIm, offset: 0, index: 5) - e1.setBytes(¶ms, length: MemoryLayout.stride, index: 6) - e1.dispatchThreads(MTLSize(width: spectrumCount, height: 1, depth: 1), - threadsPerThreadgroup: MTLSize(width: min(spectrum.maxTotalThreadsPerThreadgroup, 256), height: 1, depth: 1)) - e1.endEncoding() - guard let e2 = command.makeComputeCommandEncoder() else { return nil } - e2.setComputePipelineState(argmax) - e2.setBuffer(specRe, offset: 0, index: 0); e2.setBuffer(specIm, offset: 0, index: 1) - e2.setBuffer(scores, offset: 0, index: 2); e2.setBuffer(bins, offset: 0, index: 3) - e2.setBuffer(cosTable, offset: 0, index: 4); e2.setBuffer(sinTable, offset: 0, index: 5) - e2.setBytes(¶ms, length: MemoryLayout.stride, index: 6) - e2.setThreadgroupMemoryLength(azimuthal * 4, index: 0) - e2.setThreadgroupMemoryLength(azimuthal * 4, index: 1) - e2.dispatchThreadgroups(MTLSize(width: 1, height: templates, depth: positions), - threadsPerThreadgroup: MTLSize(width: azimuthal, height: 1, depth: 1)) - e2.endEncoding() - } + var params = Params(positions: UInt32(positions), templates: UInt32(templates), radial: UInt32(radial), + azimuthal: UInt32(azimuthal)) + // Both passes in one sweep: a thread owns one template and frequency and reads that template's ring values once for + // a block of `positionsPerThread` positions (the template bank, not the scan, is the large operand). + guard let e1 = command.makeComputeCommandEncoder() else { return nil } + e1.setComputePipelineState(spectrum) + e1.setBuffer(inRe, offset: 0, index: 0); e1.setBuffer(inIm, offset: 0, index: 1) + e1.setBuffer(templateRe, offset: 0, index: 2); e1.setBuffer(templateIm, offset: 0, index: 3) + e1.setBuffer(specRe, offset: 0, index: 4); e1.setBuffer(specIm, offset: 0, index: 5) + e1.setBytes(¶ms, length: MemoryLayout.stride, index: 6) + let blocks = (positions + Self.positionsPerThread - 1) / Self.positionsPerThread + e1.dispatchThreads(MTLSize(width: azimuthal, height: templates, depth: blocks), + threadsPerThreadgroup: MTLSize(width: min(azimuthal, spectrum.maxTotalThreadsPerThreadgroup), height: 1, depth: 1)) + e1.endEncoding() + guard let e2 = command.makeComputeCommandEncoder() else { return nil } + e2.setComputePipelineState(argmax) + e2.setBuffer(specRe, offset: 0, index: 0); e2.setBuffer(specIm, offset: 0, index: 1) + e2.setBuffer(scores, offset: 0, index: 2); e2.setBuffer(bins, offset: 0, index: 3) + e2.setBuffer(cosTable, offset: 0, index: 4); e2.setBuffer(sinTable, offset: 0, index: 5) + e2.setBytes(¶ms, length: MemoryLayout.stride, index: 6) + for index in 0..<6 { e2.setThreadgroupMemoryLength(azimuthal * 4, index: index) } + e2.dispatchThreadgroups(MTLSize(width: 1, height: templates, depth: 2 * positions), + threadsPerThreadgroup: MTLSize(width: azimuthal, height: 1, depth: 1)) + e2.endEncoding() command.commit() command.waitUntilCompleted() guard command.status == .completed else { return nil } - func floats(_ b: MTLBuffer) -> [Float] { Array(UnsafeBufferPointer(start: b.contents().bindMemory(to: Float.self, capacity: resultCount), count: resultCount)) } - func uints(_ b: MTLBuffer) -> [UInt32] { Array(UnsafeBufferPointer(start: b.contents().bindMemory(to: UInt32.self, capacity: resultCount), count: resultCount)) } - return (floats(scores0), uints(bins0), floats(scores1), uints(bins1)) + func floats(_ b: MTLBuffer, _ pass: Int) -> [Float] { + Array(UnsafeBufferPointer(start: b.contents().bindMemory(to: Float.self, capacity: 2 * resultCount) + pass * resultCount, count: resultCount)) + } + func uints(_ b: MTLBuffer, _ pass: Int) -> [UInt32] { + Array(UnsafeBufferPointer(start: b.contents().bindMemory(to: UInt32.self, capacity: 2 * resultCount) + pass * resultCount, count: resultCount)) + } + return (floats(scores, 0), uints(bins, 0), floats(scores, 1), uints(bins, 1)) } - /// The two kernels. Pass 0 (`sign` +1) is conj(E)·T, pass 1 (`sign` -1) is E·T, as ACOMMatcher.correlate. The inverse DFT - /// is exact over the azimuth (`n` power of two); ties between bins go to the lowest bin, as the reference's first-maximum scan. + /// Positions one spectrum thread serves; must equal `PPT` in the shader. + static let positionsPerThread = 8 + + /// The two kernels. Pass 0 is conj(E)·T, pass 1 is E·T, as ACOMMatcher.correlate; both are summed over the rings in the + /// same order and with the same products as one pass at a time. The inverse transform over the azimuth is a radix-2 FFT + /// (`n` power of two); ties between bins go to the lowest bin, as the reference's first-maximum scan. static let source = """ #include using namespace metal; - struct ACOMParams { uint positions; uint templates; uint radial; uint azimuthal; float sign; }; + struct ACOMParams { uint positions; uint templates; uint radial; uint azimuthal; }; + constant uint PPT = 8; kernel void acomCorrelationSpectrum( const device float *er [[buffer(0)]], const device float *ei [[buffer(1)]], const device float *tr [[buffer(2)]], const device float *ti [[buffer(3)]], device float *cr [[buffer(4)]], device float *ci [[buffer(5)]], - constant ACOMParams &p [[buffer(6)]], uint index [[thread_position_in_grid]]) + constant ACOMParams &p [[buffer(6)]], uint3 index [[thread_position_in_grid]]) { - if (index >= p.positions * p.templates * p.azimuthal) return; - const uint f = index % p.azimuthal; - const uint t = (index / p.azimuthal) % p.templates; - const uint pos = index / (p.azimuthal * p.templates); - const uint eBase = pos * p.radial * p.azimuthal, tBase = t * p.radial * p.azimuthal; - float re = 0.0f, im = 0.0f; + const uint f = index.x, t = index.y, first = index.z * PPT; + if (f >= p.azimuthal || t >= p.templates || first >= p.positions) return; + const uint count = min(PPT, p.positions - first); + const uint grid = p.radial * p.azimuthal, tBase = t * grid; + float fr[PPT], fi[PPT], mr[PPT], mi[PPT]; + for (uint k = 0; k < PPT; ++k) { fr[k] = 0.0f; fi[k] = 0.0f; mr[k] = 0.0f; mi[k] = 0.0f; } for (uint r = 0; r < p.radial; ++r) { const uint o = r * p.azimuthal + f; - const float a = er[eBase + o], b = ei[eBase + o], c = tr[tBase + o], d = ti[tBase + o]; - re += a * c + p.sign * b * d; - im += a * d - p.sign * b * c; + const float c = tr[tBase + o], d = ti[tBase + o]; + for (uint k = 0; k < PPT; ++k) { + if (k < count) { + const uint e = (first + k) * grid + o; + const float a = er[e], b = ei[e]; + fr[k] += a * c + b * d; + fi[k] += a * d - b * c; + mr[k] += a * c + (-1.0f) * b * d; + mi[k] += a * d - (-1.0f) * b * c; + } + } + } + const uint passStride = p.positions * p.templates * p.azimuthal; + for (uint k = 0; k < count; ++k) { + const uint out = ((first + k) * p.templates + t) * p.azimuthal + f; + cr[out] = fr[k]; ci[out] = fi[k]; + cr[passStride + out] = mr[k]; ci[passStride + out] = mi[k]; } - cr[index] = re; ci[index] = im; } kernel void acomInverseArgmax( const device float *cr [[buffer(0)]], const device float *ci [[buffer(1)]], @@ -218,25 +243,39 @@ private final class ACOMMetalContext { const device float *cosT [[buffer(4)]], const device float *sinT [[buffer(5)]], constant ACOMParams &p [[buffer(6)]], threadgroup float *scores [[threadgroup(0)]], threadgroup uint *bins [[threadgroup(1)]], + threadgroup float *sr [[threadgroup(2)]], threadgroup float *si [[threadgroup(3)]], + threadgroup float *sc [[threadgroup(4)]], threadgroup float *ss [[threadgroup(5)]], uint a [[thread_index_in_threadgroup]], uint3 group [[threadgroup_position_in_grid]]) { - const uint pos = group.z, t = group.y; - if (pos >= p.positions || t >= p.templates || a >= p.azimuthal) return; - const uint base = (pos * p.templates + t) * p.azimuthal; - const uint mask = p.azimuthal - 1; - float value = 0.0f; - for (uint f = 0; f < p.azimuthal; ++f) { - const uint k = (f * a) & mask; - value += cr[base + f] * cosT[k] - ci[base + f] * sinT[k]; + // group.z runs over pass 0's positions, then pass 1's + const uint slot = group.z, t = group.y; + if (slot >= 2 * p.positions || t >= p.templates || a >= p.azimuthal) return; + const uint n = p.azimuthal, base = (slot * p.templates + t) * n; + // radix-2 inverse FFT in threadgroup memory: load in bit-reversed order, then log2(n) butterfly stages + const uint bits = ctz(n); + const uint rev = bits == 0 ? 0 : (reverse_bits(a) >> (32 - bits)); + sr[rev] = cr[base + a]; si[rev] = ci[base + a]; sc[a] = cosT[a]; ss[a] = sinT[a]; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint span = 1; span < n; span <<= 1) { + if (a < n / 2) { + const uint k = a % span, i0 = (a / span) * 2 * span + k, i1 = i0 + span; + const uint w = k * (n / (2 * span)); + const float wr = sc[w], wi = ss[w]; + const float xr = sr[i1] * wr - si[i1] * wi, xi = sr[i1] * wi + si[i1] * wr; + const float ur = sr[i0], ui = si[i0]; + sr[i0] = ur + xr; si[i0] = ui + xi; + sr[i1] = ur - xr; si[i1] = ui - xi; + } + threadgroup_barrier(mem_flags::mem_threadgroup); } - scores[a] = value / float(p.azimuthal); + scores[a] = sr[a] / float(n); bins[a] = a; threadgroup_barrier(mem_flags::mem_threadgroup); for (uint stride = p.azimuthal / 2; stride > 0; stride >>= 1) { if (a < stride && scores[a + stride] > scores[a]) { scores[a] = scores[a + stride]; bins[a] = bins[a + stride]; } threadgroup_barrier(mem_flags::mem_threadgroup); } - if (a == 0) { bestScores[pos * p.templates + t] = scores[0]; bestBins[pos * p.templates + t] = bins[0]; } + if (a == 0) { bestScores[slot * p.templates + t] = scores[0]; bestBins[slot * p.templates + t] = bins[0]; } } """ } diff --git a/Tests/DiffractoOrientationTests/ACOMMetalScaleTests.swift b/Tests/DiffractoOrientationTests/ACOMMetalScaleTests.swift new file mode 100644 index 0000000..a817ade --- /dev/null +++ b/Tests/DiffractoOrientationTests/ACOMMetalScaleTests.swift @@ -0,0 +1,126 @@ +// +// ACOMMetalScaleTests.swift +// Role: the Metal option of `ACOMMatcher.matchAll` (P31) at other azimuthal lengths and scan sizes than `ACOMMetalTests`, +// and its timing. +// (1) `ACOMMetalAzimuthalLengthTests` (runs by default): the Metal scores equal the CPU reference within 1e-6, and the +// template, in-plane angle and mirror flag are identical, at 4 to 1 024 azimuthal bins (the inverse transform is a +// radix-2 FFT on the GPU) and at a scan whose size is not a multiple of the 8 positions one spectrum thread serves. +// (2) `ACOMMetalBenchmarkTests` (skipped by default): ms per pattern at a 128 x 128 scan (16 384 patterns) and at 600 +// and 2 400 beam directions. Run on a quiet machine with +// DIFFRACTOKIT_ACOM_BENCHMARK=1 swift test -c release --filter ACOMMetalBenchmarkTests +// The CPU path is timed on the first 1 024 patterns (each position is independent, so its per-pattern cost does +// not depend on the scan size). `DIFFRACTOKIT_ACOM_BENCHMARK_WORKERS` sets the CPU workers of both paths (default 1). +// + +import Foundation +import Testing +import DiffractoCore +import DiffractoOrientation + +private let benchmarkEnabled = ProcessInfo.processInfo.environment["DIFFRACTOKIT_ACOM_BENCHMARK"] != nil + +/// The fcc crystal of tools/acom-benchmark: a = 4.05 angstrom, |g| < 1.6, intensity 1 / |g|^2. +private func fccReflections() -> [ACOMReflection] { + var reflections: [ACOMReflection] = [] + for h in -4...4 { + for k in -4...4 { + for l in -4...4 where (h, k, l) != (0, 0, 0) && (h + k) % 2 == 0 && (k + l) % 2 == 0 { + let g = SIMD3(Double(h), Double(k), Double(l)) / 4.05 + let norm = (g * g).sum().squareRoot() + if norm < 1.6 { reflections.append(ACOMReflection(g: g, intensity: 1 / (norm * norm))) } + } + } + } + return reflections +} + +/// `count` patterns from the library's own templates (every seventh), each turned by a seeded random in-plane angle. +private func scanPeaks(_ library: ACOMTemplateLibrary, count: Int, ipp: Double) -> [[ACOMPeak]] { + var seed: UInt64 = 99 + func rnd() -> Double { seed = seed &* 6364136223846793005 &+ 1442695040888963407; return Double(seed >> 11) / Double(1 << 53) } + return (0.. Void) -> Double { + let t0 = DispatchTime.now().uptimeNanoseconds + body() + return Double(DispatchTime.now().uptimeNanoseconds - t0) / 1e9 + } + + @Test(.enabled(if: benchmarkEnabled)) func metalAgainstCPUAt16384PatternsAnd600Or2400Axes() throws { + let workers = Int(ProcessInfo.processInfo.environment["DIFFRACTOKIT_ACOM_BENCHMARK_WORKERS"] ?? "") ?? 1 + let reflections = fccReflections() + let side = 128, ipp = 0.004, cpuCount = 1024 + for axes in [600, 2400] { + let library = try ACOMTemplateLibrary.generate(reflections: reflections, kMax: 1.6, laueClass: .mThreeBarM, + settings: ACOMTemplateSettings(zoneAxisCount: axes)) + let matcher = try #require(ACOMMatcher(library: library)) + guard matcher.isMetalAvailable else { print("ACOM benchmark: no Metal device"); return } + let peaks = scanPeaks(library, count: side * side, ipp: ipp) + let head = Array(peaks.prefix(cpuCount)) + var cpu: ACOMScanResults?, metalHead: ACOMScanResults?, metal: ACOMScanResults? + let tCPU = seconds { + cpu = matcher.matchAll(peaks: head, width: 32, height: 32, originX: 128, originY: 128, invAngstromPerPixel: ipp, + maxWorkers: workers) + } + _ = matcher.matchAll(peaks: Array(peaks.prefix(64)), width: 8, height: 8, originX: 128, originY: 128, + invAngstromPerPixel: ipp, maxWorkers: workers, backend: .metal) // warm-up: shader compile + let tMetalHead = seconds { + metalHead = matcher.matchAll(peaks: head, width: 32, height: 32, originX: 128, originY: 128, + invAngstromPerPixel: ipp, maxWorkers: workers, backend: .metal) + } + let tMetal = seconds { + metal = matcher.matchAll(peaks: peaks, width: side, height: side, originX: 128, originY: 128, + invAngstromPerPixel: ipp, maxWorkers: workers, backend: .metal) + } + let a = try #require(cpu).results, b = try #require(metalHead).results + #expect(try #require(metal).results.count == side * side) + let worst = zip(a, b).map { abs(Double($0.score) - Double($1.score)) }.max() ?? 0 + print(String(format: "ACOM benchmark %d axes, %d CPU workers: CPU %.3f ms/pattern (1 024), Metal %.3f ms/pattern (1 024), Metal %.3f ms/pattern (16 384); Metal-CPU worst score gap %.2e", + axes, workers, tCPU / Double(cpuCount) * 1000, tMetalHead / Double(cpuCount) * 1000, + tMetal / Double(side * side) * 1000, worst)) + #expect(worst <= 1e-6) + } + } +} From a4cb9fbb65ac7490b1eb46887e067b6a6a220e64 Mon Sep 17 00:00:00 2001 From: Paullo9 Date: Sat, 31 Oct 2026 15:08:21 +0100 Subject: [PATCH 2/2] fix: ACOM Metal argmax keeps the lowest bin on an exact tie The Metal inverse-transform argmax compared scores only (strictly larger wins) in a tree reduction, so on an exact tie the bin it returned depended on the tree layout: over 4 bins with a tie between bins 1 and 2 it returned 2, where the CPU path (and NumPy argmax, as py4DSTEM) returns 1. The two backends could then report different in-plane angles for the same pattern. The reduction now compares (score, bin) pairs: a larger score wins, an equal score keeps the smaller bin. NaN and -inf enter as -FLT_MAX, so they never win, as in the CPU scan that starts at -greatestFiniteMagnitude; an all-NaN row gives bin 0 and -FLT_MAX / n on both paths. The scan runs on unscaled values and only the winner is divided by n, the same order as the CPU. Threadgroup memory lengths are at least 16 bytes (Metal requires multiples of 16; n * 4 is not one below 4 bins). A package seam, ACOMMetalKernels.bestBins, runs the kernels on planted ring spectra; ACOMMetalArgmaxTieTests plants an exact tie (red before: bin 2 in both passes) and an all-NaN correlation. Fixes mac4DSTEM/DiffractoKit#62 Done: Metal argmax ties and NaN follow the CPU rule, tested. Next: add the tie rule to the P31 row of docs/parity.md once the other parity.md PRs land. Co-Authored-By: Claude Opus 5.5 --- Sources/DiffractoOrientation/ACOMMetal.swift | 60 +++++++++++++++---- .../ACOMMetalArgmaxTieTests.swift | 55 +++++++++++++++++ 2 files changed, 103 insertions(+), 12 deletions(-) create mode 100644 Tests/DiffractoOrientationTests/ACOMMetalArgmaxTieTests.swift diff --git a/Sources/DiffractoOrientation/ACOMMetal.swift b/Sources/DiffractoOrientation/ACOMMetal.swift index 5cc8535..449991a 100644 --- a/Sources/DiffractoOrientation/ACOMMetal.swift +++ b/Sources/DiffractoOrientation/ACOMMetal.swift @@ -78,6 +78,24 @@ extension ACOMMatcher { } } +/// The GPU kernels on planted ring spectra, for tests that need an exact correlation (a tie between bins) that no real pattern +/// gives bit for bit. Per position and template, ring major: the forward pass (conj(E)·T) and the mirrored pass (E·T), each its +/// best score and bin. nil where Metal is unavailable. +package enum ACOMMetalKernels { + package static func bestBins(experimentalRe: [Float], experimentalIm: [Float], templateRe: [Float], templateIm: [Float], + templates: Int, radial: Int, azimuthal: Int, positions: Int) + -> (forwardScores: [Float], forwardBins: [UInt32], mirrorScores: [Float], mirrorBins: [UInt32])? { + #if canImport(Metal) + guard positions > 0, experimentalRe.count == positions * radial * azimuthal, experimentalIm.count == experimentalRe.count, + let gpu = ACOMMetalContext(spectraRe: templateRe, spectraIm: templateIm, templates: templates, radial: radial, + azimuthal: azimuthal) else { return nil } + return gpu.run(re: experimentalRe, im: experimentalIm, positions: positions) + #else + return nil + #endif + } +} + #if canImport(Metal) /// A batch's experimental ring spectra, written from several threads into disjoint slices of memory it owns. private final class Slices: @unchecked Sendable { @@ -110,7 +128,13 @@ private final class ACOMMetalContext { struct Params { var positions: UInt32, templates: UInt32, radial: UInt32, azimuthal: UInt32 } - init?(library: ACOMTemplateLibrary) { + convenience init?(library: ACOMTemplateLibrary) { + self.init(spectraRe: library.templateFFTRe, spectraIm: library.templateFFTIm, templates: library.count, + radial: library.geometry.nRadial, azimuthal: library.geometry.nAzimuthal) + } + + /// `templates` ring spectra of `radial` x `azimuthal` values each (ring major), as `ACOMTemplateLibrary.templateFFTRe/Im`. + init?(spectraRe: [Float], spectraIm: [Float], templates: Int, radial: Int, azimuthal: Int) { guard let device = MTLCreateSystemDefaultDevice(), let queue = device.makeCommandQueue() else { return nil } let options = MTLCompileOptions() options.mathMode = .safe @@ -118,13 +142,14 @@ private final class ACOMMetalContext { let f1 = lib.makeFunction(name: "acomCorrelationSpectrum"), let f2 = lib.makeFunction(name: "acomInverseArgmax"), let p1 = try? device.makeComputePipelineState(function: f1), let p2 = try? device.makeComputePipelineState(function: f2) else { return nil } - let n = library.geometry.nAzimuthal - guard n <= p2.maxTotalThreadsPerThreadgroup, n & (n - 1) == 0 else { return nil } - let bytes = library.templateFFTRe.count * MemoryLayout.stride + let n = azimuthal + guard n > 0, n <= p2.maxTotalThreadsPerThreadgroup, n & (n - 1) == 0, templates > 0, radial > 0, + spectraRe.count == templates * radial * n, spectraIm.count == spectraRe.count else { return nil } + let bytes = spectraRe.count * MemoryLayout.stride // The shader indexes with 32-bit unsigned integers. - guard bytes > 0, bytes <= device.maxBufferLength, library.templateFFTRe.count < 1 << 32, - let tr = device.makeBuffer(bytes: library.templateFFTRe, length: bytes, options: .storageModeShared), - let ti = device.makeBuffer(bytes: library.templateFFTIm, length: bytes, options: .storageModeShared) + guard bytes <= device.maxBufferLength, spectraRe.count < 1 << 32, + let tr = device.makeBuffer(bytes: spectraRe, length: bytes, options: .storageModeShared), + let ti = device.makeBuffer(bytes: spectraIm, length: bytes, options: .storageModeShared) else { return nil } // Float tables of the Double cos/sin of 2π k / n, k = 0...stride, index: 6) - for index in 0..<6 { e2.setThreadgroupMemoryLength(azimuthal * 4, index: index) } + // Threadgroup memory lengths must be multiples of 16 bytes (below 4 bins, n * 4 is not). + for index in 0..<6 { e2.setThreadgroupMemoryLength(max(16, azimuthal * 4), index: index) } e2.dispatchThreadgroups(MTLSize(width: 1, height: templates, depth: 2 * positions), threadsPerThreadgroup: MTLSize(width: azimuthal, height: 1, depth: 1)) e2.endEncoding() @@ -268,14 +294,24 @@ private final class ACOMMetalContext { } threadgroup_barrier(mem_flags::mem_threadgroup); } - scores[a] = sr[a] / float(n); + // The CPU scan starts at -FLT_MAX and takes a strictly larger value, so NaN and -inf never win: both enter as -FLT_MAX + // (an all-NaN row then gives bin 0 and -FLT_MAX / n, as on the CPU). Like the CPU, scan the unscaled values and + // scale the winner only, so a division near underflow cannot turn two distinct values into a tie. + const float v = sr[a]; + scores[a] = isnan(v) ? -FLT_MAX : max(v, -FLT_MAX); bins[a] = a; threadgroup_barrier(mem_flags::mem_threadgroup); + // Tree reduction over (score, bin) pairs: the larger score wins, an equal score keeps the smaller bin (the first + // maximum, NumPy argmax), whatever pair the tree happens to put in which slot. for (uint stride = p.azimuthal / 2; stride > 0; stride >>= 1) { - if (a < stride && scores[a + stride] > scores[a]) { scores[a] = scores[a + stride]; bins[a] = bins[a + stride]; } + if (a < stride) { + const float so = scores[a + stride], sm = scores[a]; + const uint bo = bins[a + stride]; + if (so > sm || (so == sm && bo < bins[a])) { scores[a] = so; bins[a] = bo; } + } threadgroup_barrier(mem_flags::mem_threadgroup); } - if (a == 0) { bestScores[slot * p.templates + t] = scores[0]; bestBins[slot * p.templates + t] = bins[0]; } + if (a == 0) { bestScores[slot * p.templates + t] = scores[0] / float(n); bestBins[slot * p.templates + t] = bins[0]; } } """ } diff --git a/Tests/DiffractoOrientationTests/ACOMMetalArgmaxTieTests.swift b/Tests/DiffractoOrientationTests/ACOMMetalArgmaxTieTests.swift new file mode 100644 index 0000000..0871145 --- /dev/null +++ b/Tests/DiffractoOrientationTests/ACOMMetalArgmaxTieTests.swift @@ -0,0 +1,55 @@ +// +// ACOMMetalArgmaxTieTests.swift +// Role: the Metal inverse-transform argmax keeps the reference rule of `ACOMMatcher.correlate` (NumPy argmax, as py4DSTEM): +// on an exact tie between in-plane bins the lowest bin wins, and NaN never wins (an all-NaN row gives bin 0 and the +// CPU start value). The correlation is planted: one ring, an experimental spectrum of ones (so both passes equal the +// template spectrum) and a template spectrum whose inverse transform over 4 bins is exactly 16 x [0, 1, 1, 0] in Float +// arithmetic, a tie between bins 1 and 2. Skipped where Metal is unavailable. +// + +import Foundation +import Testing +import DiffractoOrientation + +#if canImport(Metal) +import Metal +private let hasMetal = MTLCreateSystemDefaultDevice() != nil +#else +private let hasMetal = false +#endif + +/// The reference rule on the CPU path (`ACOMMatcher.correlate`): start below every finite value, a strictly larger value wins. +private func firstMaximum(_ values: [Float]) -> (score: Float, bin: Int) { + var best: Float = -.greatestFiniteMagnitude, bin = 0 + for a in 0.. best { best = values[a]; bin = a } + return (best, bin) +} + +struct ACOMMetalArgmaxTieTests { + @Test(.enabled(if: hasMetal, "no Metal device")) + func anExactTieBetweenBinsGoesToTheLowestBinOnMetalAsOnTheCPU() throws { + // x = [0, 4, 4, 0] over 4 bins; its forward DFT X_k = sum x_a exp(-2 pi i a k / 4). + let x: [Float] = [0, 4, 4, 0] + let templateRe: [Float] = [8, -4, 0, -4], templateIm: [Float] = [0, -4, 0, 4] + let cpu = firstMaximum(x) // the inverse transform divided by n: 16 x [0, 1, 1, 0] / 4 + #expect(cpu.bin == 1 && cpu.score == 4) + let gpu = try #require(ACOMMetalKernels.bestBins(experimentalRe: [1, 1, 1, 1], experimentalIm: [0, 0, 0, 0], + templateRe: templateRe, templateIm: templateIm, + templates: 1, radial: 1, azimuthal: 4, positions: 1)) + #expect(gpu.forwardScores == [cpu.score], "the planted inverse transform is exact") + #expect(gpu.forwardBins == [UInt32(cpu.bin)], "forward pass: the lowest of the tied bins") + #expect(gpu.mirrorScores == [cpu.score]) + #expect(gpu.mirrorBins == [UInt32(cpu.bin)], "mirrored pass: the lowest of the tied bins") + } + + @Test(.enabled(if: hasMetal, "no Metal device")) + func anAllNaNCorrelationGivesBinZeroAndTheCPUStartValueOnMetal() throws { + let cpu = firstMaximum([Float](repeating: .nan, count: 4)) + let gpu = try #require(ACOMMetalKernels.bestBins(experimentalRe: [1, 1, 1, 1], experimentalIm: [0, 0, 0, 0], + templateRe: [.nan, 0, 0, 0], templateIm: [0, 0, 0, 0], + templates: 1, radial: 1, azimuthal: 4, positions: 1)) + #expect(gpu.forwardBins == [UInt32(cpu.bin)] && gpu.mirrorBins == [UInt32(cpu.bin)]) + // The CPU path scales its start value by 1/n after the scan (-greatestFiniteMagnitude / 4 is exact). + #expect(gpu.forwardScores == [cpu.score / 4] && gpu.mirrorScores == [cpu.score / 4]) + } +}