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]) + } +}