Files
Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm
T
leonarski_fandClaude Opus 5.5 ecd67ec44f Metal integration engine: back to its panel-s5 kernels; Mac stills keep the old clip, no saturation count
The Metal kernel and host code are restored to their panel-s5 content, so neither the background clip
that follows a handed-over radial curve nor the count of saturated pixels in pass A runs on the Metal
path. This is an accepted, temporary platform difference, to be closed once the kernels can be checked
on a Mac: on macOS a still's background clip stays the plain one (the radial correction itself still
uses the handed-over curve, as on the CPU), and an overloaded observation kept as a lower bound counts
its saturated pixels as 0 instead of at the saturation limit.

One host-side change is kept, with no effect on any result: the radial sums are accumulated on the GPU
only where no curve was handed over. StartPassA now sizes those sums for that case alone, so the old
host code would write past them on every still with a curve.

The shared structs need nothing on the Metal side: BraggPassA.h is compiled into the kernel source, the
box-sum kernel zero-initialises BraggBoxSums (n_saturated = 0) and the host zero-initialises
BraggPassAFrame (n_curve = 0).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
2026-10-22 12:53:37 +02:00

420 lines
21 KiB
Plaintext

// 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;
// The curve is added up on the GPU only where none was handed over (RadialBackgroundFromRings):
// StartPassA sizes its sums for that case alone. The clip does not follow a handed-over curve here.
const bool accumulate_radial = bkg_radial && radial_curve.empty();
f.n_rad = accumulate_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 (accumulate_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 (accumulate_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:(accumulate_radial ? dev->rad_sum : dev->clips) offset:0 atIndex:8]; // not written without the curve
[enc setBuffer:(accumulate_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 (accumulate_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;
}