mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 87 Wiki Activity
Files
bb7c9c8fa3b1535979bc2edb14cf6ea40ca57f1d
Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm
T
leonarski_fandClaude Opus 5.5 d6166f6fa5 Panel layout stage 2: prediction and Bragg integration through the panel table 
Everything sparse that took a radial quantity from an image coordinate now takes it from the
reflection's (or spot's) position on the detector plane, through its panel (DetectorModules).
With an empty module table every changed expression reduces to the old one in the same
floating-point order, so flat detectors - every rotation set, EIGER, CCDs, assembled JUNGFRAU -
integrate bit for bit as before.

- Prediction (CPU, GPU, rotation CPU and GPU): bounds test against PlaneBounds instead of the
  image rectangle.
- IntegrateOnImage: a prediction in a gap or off every panel is dropped (PlaneToPanel), no longer
  passed through onto an unrelated stored pixel. The panel id stays local: the engine and the way
  back find it from the image position (the panel boxes tile the image), and a prediction that
  rounding to float put on another panel's edge is dropped so the two agree. Nothing is added to
  Reflection.
- Bragg engines (CPU, CUDA, Metal) take a BraggPanel per reflection (BraggStencil.h): its plane
  position, the panel's pixel axes on the plane and its box. The stencil's radial direction is
  taken on the plane and expressed in the panel's pixel axes; r0, the radial background curve,
  the Gaussian's elongation, the learned moments, tan(2theta) and Finalize's polarisation use the
  plane position; r_max / n_rad come from PlaneBounds. The neighbour mask and the owner map are
  clamped to the reflection's own panel; inside the walked box a pixel off the reflection's panel
  is read as a gap (unreadable, statically masked - not an overload), which is what the assembled
  image has there, so a disk cut by a panel edge is partial rather than silently "full".
- Risk R4: the ProfileEmpirical grid is learned and applied in the plane's axes (the panel axes
  taken to the nearest quarter turn, BraggProfileCell), so turned panels pool one shape. Without
  it the stacked layout read up to 0.3 % off the assembled one in the harness.
- SpotWidth / SpotFootprint: aperture bounds against the spot's panel box, beam distance,
  resolution and the shape's 2theta Jacobian on the plane, moments read in the panel's axes.
- GetDetectorMaxQ_recipA, the far-corner radii (first-pass tilt report, metric-symmetry arm,
  footprint table range) from PlaneBounds; pre-scan spot resolution on the plane;
  LogHeaderBeamCenter: geometric centre of the panels' plane box, and the masked-pixel check on
  the panel the beam lands on (none when it is in a gap).

Harness: PanelLayout_BraggIntegration is now a real [PanelLayout] check, over six settings
(default, bandwidth-elongated ring, radial background, overlap exclude, box sum, empirical
profile, all with polarisation), stacked vs assembled to 1e-5; PanelLayout_BraggIntegrationGPU
runs the same on the CUDA engine. Its spots avoid exact half pixels: the grid centres on the
rounded position, and a half rounds to the other neighbour on a turned panel.

The Metal engine (.metal/.mm) mirrors the CPU change exactly but is UNCOMPILED here (no macOS):
it needs a build and BraggIntegrationEngineMetal_BitIdenticalToCPU on a Mac.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
2026-10-21 10:40:20 +02:00

417 lines
21 KiB
Plaintext
Raw Blame History

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#import <Metal/Metal.h>
#include "BraggIntegrationEngineMetal.h"
#include <algorithm>
#include <string>
#include <unistd.h>
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
using namespace bragg_engine;
// The kernels' source: BraggStencil.h, BraggPassA.h and BraggIntegrationEngineMetal.metal, copied in
// at configure time (CMakeLists.txt) so the binary needs no file beside it.
extern const char *const bragg_metal_source_stencil;
extern const char *const bragg_metal_source_pass_a;
extern const char *const bragg_metal_source_exp;
extern const char *const bragg_metal_source_pass_b;
extern const char *const bragg_metal_source_kernels;
namespace {
static_assert(MASK_TAIL == 1 && MASK_FLUX == 2, "BraggIntegrationEngineMetal.metal spells the two marks out");
struct Kernels {
id<MTLDevice> device = nil; // nil: no Metal device
id<MTLComputePipelineState> mark = nil, claim = nil, box = nil, clip = nil, fit = nil;
};
// Compiled once per process: a worker makes an engine per pass, and compiling takes a fraction of a
// second. The Metal compiler ships with macOS, so nothing has to be installed to build or run this.
const Kernels &GetKernels() {
static const Kernels kernels = [] {
Kernels k;
k.device = MTLCreateSystemDefaultDevice();
if (k.device == nil)
return k;
// No contraction of a*b+c into a fused multiply-add, as everywhere on the host (CMakeLists.txt).
// Metal would otherwise contract wherever its backend saw fit. BraggStencil.h calls C's math
// functions; the defines are Metal's spellings of them.
const std::string source = std::string("#include <metal_stdlib>\n"
"using namespace metal;\n"
"#pragma clang fp contract(off)\n"
"#define sqrtf(x) sqrt(x)\n"
"#define fmaxf(a, b) fmax(a, b)\n"
"#define floorf(x) floor(x)\n"
"#define lroundf(x) ((long) round(x))\n")
+ bragg_metal_source_stencil + bragg_metal_source_pass_a
+ bragg_metal_source_exp + bragg_metal_source_pass_b
+ bragg_metal_source_kernels;
MTLCompileOptions *options = [MTLCompileOptions new];
// Fast math off: sqrt and division then round correctly and nothing is reassociated, as on the
// CPU. (MTLCompileOptions.mathMode says the same from macOS 15; the build targets 13.) And the
// language of macOS 13 on every macOS, so the kernels mean the same wherever they compile.
options.fastMathEnabled = NO;
options.languageVersion = MTLLanguageVersion3_0;
NSError *error = nil;
id<MTLLibrary> library = [k.device newLibraryWithSource:@(source.c_str()) options:options error:&error];
if (library == nil)
throw JFJochException(JFJochExceptionCategory::GPUMetalError,
std::string("Bragg integration kernels: ") + error.localizedDescription.UTF8String);
auto pipeline = [&](const char *name) {
id<MTLFunction> function = [library newFunctionWithName:@(name)];
id<MTLComputePipelineState> p = [k.device newComputePipelineStateWithFunction:function error:&error];
if (p == nil)
throw JFJochException(JFJochExceptionCategory::GPUMetalError,
std::string(name) + ": " + error.localizedDescription.UTF8String);
return p;
};
k.mark = pipeline("mark_neighbours");
k.claim = pipeline("claim_owner");
k.box = pipeline("box_sums");
k.clip = pipeline("clip_sums");
k.fit = pipeline("profile_fit");
return k;
}();
return kernels;
}
void WaitCommitted(id<MTLCommandBuffer> cb) {
[cb waitUntilCompleted];
if (cb.status != MTLCommandBufferStatusCompleted)
throw JFJochException(JFJochExceptionCategory::GPUMetalError,
cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed");
}
void Wait(id<MTLCommandBuffer> cb) {
[cb commit];
WaitCommitted(cb);
}
} // namespace
struct BraggIntegrationEngineMetal::Device {
id<MTLCommandQueue> queue = nil;
id<MTLBuffer> static_mask = nil;
// Frame-sized, made on the first image: a worker's analysis is built for every pass, and most of
// them never integrate. A frame unmarks what it marked once the CPU has read them (`unmark`, which
// the next frame waits for): clearing all of both took 1.5 ms of a 16M-pixel frame, longer than
// integrating a sparse one. `clear` is true once an unmark completed; the first frame, and one
// after a frame that threw part-way through, clear the whole of both instead.
id<MTLBuffer> mask = nil, owner = nil;
id<MTLCommandBuffer> unmark = nil;
bool clear = false;
id<MTLBuffer> rad_sum = nil, rad_cnt = nil; // the radial background curve, as the clip kernel adds it up
// Per reflection, grown with the frames.
size_t capacity = 0;
id<MTLBuffer> xy = nil, panels = nil, level = nil, sums = nil, clips = nil, thr = nil, clip_out = nil;
id<MTLBuffer> fit_setups = nil, fit_out = nil;
// Pass B's Gaussians, grown with the frames, and the learned profiles.
size_t cells_capacity = 0, profiles_capacity = 0;
id<MTLBuffer> cells = nil, profiles = nil;
// The frame's image while Run hands pass B to FitOnDevice; nil otherwise (the CPU engine's own Run).
id<MTLBuffer> img = nil;
};
bool BraggIntegrationEngineMetal::Available(const DiffractionExperiment &experiment) {
const auto settings = experiment.GetBraggIntegrationSettings();
const bool trimmed_mean = !(settings.GetBackgroundClipNSigma() > 0.0f)
&& settings.GetBackgroundTrimFraction() > 0.0f;
return !gpus_are_hidden() && !trimmed_mean && GetKernels().device != nil;
}
std::string BraggIntegrationEngineMetal::DeviceName() {
if (gpus_are_hidden() || GetKernels().device == nil)
return {};
return GetKernels().device.name.UTF8String;
}
BraggIntegrationEngineMetal::BraggIntegrationEngineMetal(const DiffractionExperiment &experiment,
const PixelMask &mask)
: BraggIntegrationEngineCPU(experiment, mask), dev(std::make_unique<Device>()) {
const Kernels &k = GetKernels();
dev->queue = [k.device newCommandQueue];
dev->static_mask = [k.device newBufferWithBytes:static_mask.data()
length:std::max<size_t>(static_mask.size() * sizeof(uint32_t), 4)
options:MTLResourceStorageModeShared];
}
BraggIntegrationEngineMetal::~BraggIntegrationEngineMetal() = default;
std::vector<Reflection> BraggIntegrationEngineMetal::Run(const ImagePreprocessorBuffer &image,
const std::vector<Reflection> &predicted,
size_t npredicted, int64_t image_number) {
// An image of another size, and a frame with nothing predicted, the CPU engine answers as it would.
if (image.size() != npixel || npredicted == 0)
return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number);
dev->img = nil;
@autoreleasepool {
const Kernels &k = GetKernels();
// The GPU reads the host's own memory, so the image is read where it lies - which Metal allows
// for memory starting on a page, as macOS's allocator hands out any buffer of a frame's size.
// Should one not be, the CPU engine integrates it: the same result.
const size_t page = getpagesize(), bytes = npixel * sizeof(int32_t);
id<MTLBuffer> img = nil;
if (reinterpret_cast<uintptr_t>(image.data()) % page == 0)
img = [k.device newBufferWithBytesNoCopy:const_cast<int32_t *>(image.data())
length:(bytes + page - 1) / page * page
options:MTLResourceStorageModeShared
deallocator:nil];
if (img == nil)
return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number);
counts.predicted += npredicted;
const bool use_owner = overlap != OverlapMode::Off;
if (dev->unmark != nil) { // it reads xy and the panels, which are overwritten below
id<MTLCommandBuffer> unmark = dev->unmark;
dev->unmark = nil;
dev->clear = false; // unless it completed
WaitCommitted(unmark);
dev->clear = true;
}
auto buffer = [&](size_t length) {
return [k.device newBufferWithLength:std::max<size_t>(length, 4) options:MTLResourceStorageModeShared];
};
if (dev->mask == nil) {
dev->mask = buffer(npixel);
if (use_owner)
dev->owner = buffer(npixel * sizeof(uint32_t));
}
const bool was_clear = dev->clear;
dev->clear = false;
if (dev->capacity < npredicted) {
const size_t n = npredicted + npredicted / 2;
dev->xy = buffer(n * 2 * sizeof(float));
dev->panels = buffer(n * sizeof(BraggPanel));
dev->level = buffer(n);
dev->sums = buffer(n * sizeof(BraggBoxSums));
dev->clips = buffer(n);
dev->thr = buffer(n * sizeof(int32_t));
dev->clip_out = buffer(n * sizeof(BraggClipSums));
dev->fit_setups = buffer(n * sizeof(BraggFitSetup));
dev->fit_out = buffer(n * sizeof(BraggFitOut));
dev->capacity = n;
}
auto *xy = static_cast<float *>(dev->xy.contents);
auto *level = static_cast<uint8_t *>(dev->level.contents);
const std::vector<BraggPanel> panels = Panels(predicted, npredicted);
std::copy(panels.begin(), panels.end(), static_cast<BraggPanel *>(dev->panels.contents));
for (size_t i = 0; i < npredicted; ++i) {
xy[2 * i] = predicted[i].predicted_x;
xy[2 * i + 1] = predicted[i].predicted_y;
level[i] = predicted[i].partiality >= NEIGHBOUR_MASK_MIN_PARTIALITY ? MASK_FLUX : MASK_TAIL;
}
PassA a = StartPassA(npredicted);
BraggPassAFrame f{};
f.width = static_cast<int32_t>(xpixel);
f.height = static_cast<int32_t>(ypixel);
f.r1_sq = r1_sq;
f.r2_sq = r2_sq;
f.r3_sq = r3_sq;
f.claim = claim;
f.inv_claim = inv_claim;
f.use_owner = use_owner;
f.exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum;
f.n_rad = bkg_radial ? a.n_rad : 0;
f.static_mask_words = static_cast<uint32_t>(static_mask.size());
f.n = static_cast<uint32_t>(npredicted);
const MTLSize grid = MTLSizeMake(npredicted, 1, 1), group = MTLSizeMake(64, 1, 1);
// --- The neighbour mask, the owner map and pass A's box walk ---
id<MTLCommandBuffer> cb = [dev->queue commandBuffer];
id<MTLBlitCommandEncoder> blit;
if (!was_clear) {
blit = [cb blitCommandEncoder];
[blit fillBuffer:dev->mask range:NSMakeRange(0, npixel) value:0];
if (use_owner) // every byte 0xff: BRAGG_OWNER_NONE
[blit fillBuffer:dev->owner range:NSMakeRange(0, npixel * sizeof(uint32_t)) value:0xff];
[blit endEncoding];
}
const uint32_t claim_not_release = 0;
id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
for (const uint8_t want : {MASK_TAIL, MASK_FLUX}) {
[enc setComputePipelineState:k.mark];
[enc setBuffer:dev->xy offset:0 atIndex:0];
[enc setBuffer:dev->level offset:0 atIndex:1];
[enc setBuffer:dev->mask offset:0 atIndex:2];
[enc setBytes:&stencil length:sizeof(stencil) atIndex:3];
[enc setBytes:&f length:sizeof(f) atIndex:4];
[enc setBytes:&want length:sizeof(want) atIndex:5];
[enc setBuffer:dev->panels offset:0 atIndex:6];
[enc dispatchThreads:grid threadsPerThreadgroup:group];
}
if (use_owner) {
[enc setComputePipelineState:k.claim];
[enc setBuffer:dev->xy offset:0 atIndex:0];
[enc setBuffer:dev->owner offset:0 atIndex:1];
[enc setBytes:&f length:sizeof(f) atIndex:2];
[enc setBytes:&claim_not_release length:sizeof(claim_not_release) atIndex:3];
[enc setBuffer:dev->panels offset:0 atIndex:4];
[enc dispatchThreads:grid threadsPerThreadgroup:group];
}
[enc setComputePipelineState:k.box];
[enc setBuffer:dev->xy offset:0 atIndex:0];
[enc setBuffer:img offset:0 atIndex:1];
[enc setBuffer:dev->mask offset:0 atIndex:2];
[enc setBuffer:(use_owner ? dev->owner : dev->mask) offset:0 atIndex:3]; // not read without an owner map
[enc setBuffer:dev->static_mask offset:0 atIndex:4];
[enc setBytes:&stencil length:sizeof(stencil) atIndex:5];
[enc setBytes:&f length:sizeof(f) atIndex:6];
[enc setBuffer:dev->sums offset:0 atIndex:7];
[enc setBuffer:dev->panels offset:0 atIndex:8];
[enc dispatchThreads:grid threadsPerThreadgroup:group];
[enc endEncoding];
Wait(cb);
// --- The clip thresholds, which are double arithmetic, on the CPU; the clip on the GPU ---
const auto *sums = static_cast<const BraggBoxSums *>(dev->sums.contents);
auto *clips = static_cast<uint8_t *>(dev->clips.contents);
auto *thr = static_cast<int32_t *>(dev->thr.contents);
bool any_clip = false;
for (size_t i = 0; i < npredicted; ++i) {
clips[i] = ClipsRing(sums[i]);
if (clips[i]) {
thr[i] = ClipThreshold(sums[i]);
any_clip = true;
}
}
if (any_clip) {
if (bkg_radial && dev->rad_sum == nil) {
dev->rad_sum = buffer(a.n_rad * 2 * sizeof(uint32_t));
dev->rad_cnt = buffer(a.n_rad * sizeof(int32_t));
}
cb = [dev->queue commandBuffer];
if (bkg_radial) {
blit = [cb blitCommandEncoder];
[blit fillBuffer:dev->rad_sum range:NSMakeRange(0, a.n_rad * 2 * sizeof(uint32_t)) value:0];
[blit fillBuffer:dev->rad_cnt range:NSMakeRange(0, a.n_rad * sizeof(int32_t)) value:0];
[blit endEncoding];
}
enc = [cb computeCommandEncoder];
[enc setComputePipelineState:k.clip];
[enc setBuffer:dev->xy offset:0 atIndex:0];
[enc setBuffer:img offset:0 atIndex:1];
[enc setBuffer:dev->mask offset:0 atIndex:2];
[enc setBytes:&stencil length:sizeof(stencil) atIndex:3];
[enc setBytes:&f length:sizeof(f) atIndex:4];
[enc setBuffer:dev->clips offset:0 atIndex:5];
[enc setBuffer:dev->thr offset:0 atIndex:6];
[enc setBuffer:dev->clip_out offset:0 atIndex:7];
[enc setBuffer:(bkg_radial ? dev->rad_sum : dev->clips) offset:0 atIndex:8]; // not written without the curve
[enc setBuffer:(bkg_radial ? dev->rad_cnt : dev->clips) offset:0 atIndex:9];
[enc setBuffer:dev->panels offset:0 atIndex:10];
[enc dispatchThreads:grid threadsPerThreadgroup:group];
[enc endEncoding];
Wait(cb);
if (bkg_radial) {
const auto *w = static_cast<const uint32_t *>(dev->rad_sum.contents);
const auto *c = static_cast<const int32_t *>(dev->rad_cnt.contents);
for (int b = 0; b < a.n_rad; ++b) {
a.rad_sum[b] = static_cast<int64_t>(w[2 * b] | static_cast<uint64_t>(w[2 * b + 1]) << 32);
a.rad_cnt[b] = c[b];
}
}
}
++frames_on_gpu;
// --- The rest is the CPU engine's ---
const auto *clip_out = static_cast<const BraggClipSums *>(dev->clip_out.contents);
std::vector<int32_t> no_trim; // the trimmed mean is the CPU engine's alone (Available)
for (size_t i = 0; i < npredicted; ++i) {
const Reflection &r = predicted[i];
FinishPassA(a, i, r, MakeBraggStencil(panels[i], stencil), sums[i],
clips[i] ? clip_out[i] : BraggClipSums{}, no_trim);
}
const FlatOwnerMap owner_map{use_owner ? static_cast<const uint32_t *>(dev->owner.contents) : nullptr, xpixel};
dev->img = img;
auto ret = ProfilePasses(BufferSampler{image.data()}, owner_map, predicted, npredicted, panels, a,
image_number);
dev->img = nil;
// --- Unmark this frame's regions and give back its claims, while the CPU goes on ---
dev->unmark = [dev->queue commandBuffer];
enc = [dev->unmark computeCommandEncoder];
const uint8_t unmark = 0;
[enc setComputePipelineState:k.mark];
[enc setBuffer:dev->xy offset:0 atIndex:0];
[enc setBuffer:dev->level offset:0 atIndex:1];
[enc setBuffer:dev->mask offset:0 atIndex:2];
[enc setBytes:&stencil length:sizeof(stencil) atIndex:3];
[enc setBytes:&f length:sizeof(f) atIndex:4];
[enc setBytes:&unmark length:sizeof(unmark) atIndex:5];
[enc setBuffer:dev->panels offset:0 atIndex:6];
[enc dispatchThreads:grid threadsPerThreadgroup:group];
if (use_owner) {
const uint32_t release = 1;
[enc setComputePipelineState:k.claim];
[enc setBuffer:dev->xy offset:0 atIndex:0];
[enc setBuffer:dev->owner offset:0 atIndex:1];
[enc setBytes:&f length:sizeof(f) atIndex:2];
[enc setBytes:&release length:sizeof(release) atIndex:3];
[enc setBuffer:dev->panels offset:0 atIndex:4];
[enc dispatchThreads:grid threadsPerThreadgroup:group];
}
[enc endEncoding];
[dev->unmark commit];
return ret;
}
}
bool BraggIntegrationEngineMetal::FitOnDevice(const std::vector<BraggFitSetup> &setups,
const std::vector<float> &shell_profiles,
std::vector<BraggFitOut> &fits) {
if (dev->img == nil)
return false;
@autoreleasepool {
const Kernels &k = GetKernels();
const size_t n = setups.size();
size_t n_cells = 0;
for (const BraggFitSetup &s : setups)
if (s.fit && s.shell < 0)
n_cells = std::max<size_t>(n_cells, s.cells + static_cast<size_t>(2 * s.Rf + 1) * (2 * s.Rf + 1));
auto grow = [&](__strong id<MTLBuffer> &b, size_t &capacity, size_t bytes) {
if (capacity < bytes) {
capacity = bytes + bytes / 2;
b = [k.device newBufferWithLength:std::max<size_t>(capacity, 4) options:MTLResourceStorageModeShared];
}
};
grow(dev->cells, dev->cells_capacity, n_cells * sizeof(float));
grow(dev->profiles, dev->profiles_capacity, shell_profiles.size() * sizeof(float));
std::copy(setups.begin(), setups.end(), static_cast<BraggFitSetup *>(dev->fit_setups.contents));
std::copy(shell_profiles.begin(), shell_profiles.end(), static_cast<float *>(dev->profiles.contents));
const bool use_owner = overlap != OverlapMode::Off;
const BraggPassBFrame f{static_cast<int32_t>(xpixel), static_cast<int32_t>(ypixel), r1_sq, R, use_owner,
overlap == OverlapMode::Exclude, static_cast<uint32_t>(n)};
id<MTLCommandBuffer> cb = [dev->queue commandBuffer];
id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
[enc setComputePipelineState:k.fit];
[enc setBuffer:dev->fit_setups offset:0 atIndex:0];
[enc setBuffer:dev->img offset:0 atIndex:1];
[enc setBuffer:(use_owner ? dev->owner : dev->mask) offset:0 atIndex:2]; // not read without an owner map
[enc setBuffer:dev->profiles offset:0 atIndex:3];
[enc setBuffer:dev->cells offset:0 atIndex:4];
[enc setBytes:&stencil length:sizeof(stencil) atIndex:5];
[enc setBytes:&f length:sizeof(f) atIndex:6];
[enc setBuffer:dev->fit_out offset:0 atIndex:7];
[enc setBuffer:dev->panels offset:0 atIndex:8];
[enc dispatchThreads:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
[enc endEncoding];
Wait(cb);
const auto *out = static_cast<const BraggFitOut *>(dev->fit_out.contents);
std::copy(out, out + n, fits.begin());
}
return true;
}
Reference in new issue View Git Blame Copy Permalink