Repository navigation
fix: ACOM Metal argmax keeps the lowest bin on an exact tie - #63
Merged
Merged
Conversation
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 <noreply@anthropic.com>
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 #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 <noreply@anthropic.com>
Main now holds #61 as a squash. The conflicted ACOMMetal.swift keeps the branch side: main has it exactly as the first branch commit wrote it, and the tie fix builds on that. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Requested by Paul · project thread
Fixes #62
Stacked on #61 (base
issue-8-metal-timing): merge #61 first. This PR then retargets tomain.What was wrong
The Metal inverse-transform argmax (
acomInverseArgmax) reduced scores in a tree, keeping a value only when it was strictly larger. 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. The CPU path (ACOMMatcher.correlate, the first maximum, as NumPy argmax in py4DSTEM) returns 1. So the two backends could report different in-plane angles for the same pattern.Fix
max(16, n * 4), because Metal requires multiples of 16 bytes and n * 4 is not one below 4 bins (one line).ACOMMetalKernels.bestBins(package access, not public API) runs the kernels on planted ring spectra.ACOMMetalContextgained an init from raw spectra; the library init delegates to it.Tests
ACOMMetalArgmaxTieTests(skipped without a Metal device). It uses one ring, an experimental spectrum of ones, and a template spectrum whose inverse transform is exactly [0, 4, 4, 0] in the kernel Float arithmetic.swift test -c release --parallel --skip DiffractoKitParityTests: all green (28 + 156 + 208 + 173 tests). The existing Metal-against-CPU tests are unchanged, with a worst score gap of 1.19e-7.Review
One numerics review, before the commit. It found nothing blocking. It traced the n=4 planted FFT as exact, and checked the tie rule at every tree level, barrier safety, NaN/-inf parity and the init guards. Its one low finding (scale after the scan, not before) is applied.
Not in this PR: the tie rule belongs in the P31 row of
docs/parity.md. It is left out because other open PRs already edit that file.🤖 Generated with Claude Code