diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal index 46d12341f..8063d5efa 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal @@ -125,6 +125,8 @@ kernel void box_sums(device const float2 *xy [[buffer(0)]], && ((static_mask[(uint) idx / 32] >> ((uint) idx % 32)) & 1u); if (v == PX_SATURATED || (v == PX_MASKED && on && !masked)) s.overloaded = 1; + if (v == PX_SATURATED) + ++s.n_saturated; if (!valid_px(v)) continue; s.I_sum += v; s.I_sum_x += (int64_t) x * v; @@ -154,8 +156,17 @@ inline void atomic_add_i64(device atomic_uint *w, int64_t v) { atomic_fetch_add_explicit(&w[1], hi + (old + lo < old ? 1u : 0u), memory_order_relaxed); } -// The background clip: the ring pixels at or below the threshold the CPU computed from box_sums, and -// the radial background curve from the same pixels. +// BraggIntegrationEngine::RadialClipOffset: how far the handed-over radial curve rises from the +// reflection's radius r0 to a ring pixel at radial offset rad, in whole counts. +inline int32_t radial_clip_offset(device const float *curve, int n, float r0, float rad) { + const int b = clamp((int) round(r0 + rad), 0, n - 1); + const int b0 = clamp((int) round(r0), 0, n - 1); + return (int32_t) round(curve[b] - curve[b0]); +} + +// The background clip: the ring pixels at or below the threshold the CPU computed from box_sums - plus +// the rise of the handed-over radial curve where one is in use (n_curve > 0) - and the radial +// background curve from the same pixels where it is accumulated instead (n_rad > 0). kernel void clip_sums(device const float2 *xy [[buffer(0)]], device const int32_t *img [[buffer(1)]], device const uint8_t *mask [[buffer(2)]], @@ -167,6 +178,7 @@ kernel void clip_sums(device const float2 *xy [[buffer(0)]], device atomic_uint *rad_sum [[buffer(8)]], device atomic_int *rad_cnt [[buffer(9)]], device const BraggPanel *panels [[buffer(10)]], + device const float *curve [[buffer(11)]], uint i [[thread_position_in_grid]]) { if (i >= f.n || !clips[i]) return; @@ -185,7 +197,10 @@ kernel void clip_sums(device const float2 *xy [[buffer(0)]], const BraggStencilDist d = BraggStencilDistances(st, x - px, y - py); const int idx = y * f.width + x; const int32_t v = BraggPanelHolds(pn, x, y) ? img[idx] : PX_MASKED; - if (!(d.inner >= f.r2_sq && d.outer < f.r3_sq && valid_px(v) && mask[idx] != MASK_FLUX && v <= t)) + if (!(d.inner >= f.r2_sq && d.outer < f.r3_sq && valid_px(v) && mask[idx] != MASK_FLUX)) + continue; + const int32_t off = f.n_curve > 0 ? radial_clip_offset(curve, f.n_curve, st.r0, d.rad) : 0; + if (!((int64_t) v <= (int64_t) t + off)) continue; c.sum += v; ++c.n; diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm index 7c902dd7b..d4f9c73ed 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm @@ -107,6 +107,7 @@ struct BraggIntegrationEngineMetal::Device { id unmark = nil; bool clear = false; id rad_sum = nil, rad_cnt = nil; // the radial background curve, as the clip kernel adds it up + id curve = nil; // the curve handed over, which the clip follows // Per reflection, grown with the frames. size_t capacity = 0; id xy = nil, panels = nil, level = nil, sums = nil, clips = nil, thr = nil, clip_out = nil; @@ -147,9 +148,7 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor const std::vector &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. - // So does a frame whose radial curve was handed over and is in use: the ring's clip then follows the - // curve (RadialClipOffset), which the clip kernel here does not, and the CPU gives that result exactly. - if (image.size() != npixel || npredicted == 0 || (bkg_radial && !radial_curve.empty())) + if (image.size() != npixel || npredicted == 0) return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number); dev->img = nil; @@ -223,9 +222,16 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor 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. A handed-over curve in use never gets here (above). + // StartPassA sizes its sums for that case alone. const bool accumulate_radial = bkg_radial && radial_curve.empty(); f.n_rad = accumulate_radial ? a.n_rad : 0; + // Where a handed-over curve is in use the clip follows it (RadialClipOffset). + f.n_curve = bkg_radial && !radial_curve.empty() ? static_cast(radial_curve.size()) : 0; + if (f.n_curve > 0) { + if (dev->curve == nil || dev->curve.length < radial_curve.size() * sizeof(float)) + dev->curve = buffer(radial_curve.size() * sizeof(float)); + std::copy(radial_curve.begin(), radial_curve.end(), static_cast(dev->curve.contents)); + } f.static_mask_words = static_cast(static_mask.size()); f.n = static_cast(npredicted); @@ -314,6 +320,7 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor [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 setBuffer:(f.n_curve > 0 ? dev->curve : dev->thr) offset:0 atIndex:11]; // not read without a curve [enc dispatchThreads:grid threadsPerThreadgroup:group]; [enc endEncoding]; Wait(cb); diff --git a/image_analysis/bragg_integration/BraggPassA.h b/image_analysis/bragg_integration/BraggPassA.h index 8c17b108d..da3394e32 100644 --- a/image_analysis/bragg_integration/BraggPassA.h +++ b/image_analysis/bragg_integration/BraggPassA.h @@ -46,3 +46,8 @@ struct BraggPassAFrame { uint32_t static_mask_words; // length of the run's packed pixel mask uint32_t n; // reflections }; + +// Compiled on both sides, so a layout the host and the Metal compiler disagree on fails to compile. +static_assert(sizeof(BraggBoxSums) == 96, "BraggBoxSums layout shared with the Metal kernels"); +static_assert(sizeof(BraggClipSums) == 16, "BraggClipSums layout shared with the Metal kernels"); +static_assert(sizeof(BraggPassAFrame) == 52, "BraggPassAFrame layout shared with the Metal kernels"); diff --git a/image_analysis/bragg_integration/BraggPassB.h b/image_analysis/bragg_integration/BraggPassB.h index e03981a5e..e1bec88c6 100644 --- a/image_analysis/bragg_integration/BraggPassB.h +++ b/image_analysis/bragg_integration/BraggPassB.h @@ -46,6 +46,11 @@ struct BraggPassBFrame { uint32_t n; // reflections }; +// Compiled on both sides, so a layout the host and the Metal compiler disagree on fails to compile. +static_assert(sizeof(BraggFitSetup) == 76, "BraggFitSetup layout shared with the Metal kernels"); +static_assert(sizeof(BraggFitOut) == 40, "BraggFitOut layout shared with the Metal kernels"); +static_assert(sizeof(BraggPassBFrame) == 28, "BraggPassBFrame layout shared with the Metal kernels"); + // The Gaussian profile's (unnormalised) value at grid offset (dx, dy). inline float BraggGaussian(int dx, int dy, float fx, float fy, float ux, float uy, float inv_2s2r, float inv_2s2t) { diff --git a/image_analysis/bragg_integration/BraggStencil.h b/image_analysis/bragg_integration/BraggStencil.h index 2dc44b8ca..e30477268 100644 --- a/image_analysis/bragg_integration/BraggStencil.h +++ b/image_analysis/bragg_integration/BraggStencil.h @@ -185,6 +185,10 @@ struct BraggPanel { float fx, fy, sx, sy; // the plane step of one image pixel along x, and of one along y int32_t x0, y0, x1, y1; // the panel's pixels in the image, inclusive }; +// The Metal engine hands both to its kernels as they lie in host memory, and this header is compiled on +// both sides: a layout the two compilers disagree on fails here, on whichever side it is. +static_assert(sizeof(BraggStencilParams) == 168, "BraggStencilParams layout shared with the Metal kernels"); +static_assert(sizeof(BraggPanel) == 40, "BraggPanel layout shared with the Metal kernels"); BRAGG_STENCIL_HD bool BraggPanelHolds(const BRAGG_STENCIL_THREAD BraggPanel &pn, int x, int y) { return x >= pn.x0 && x <= pn.x1 && y >= pn.y0 && y <= pn.y1; diff --git a/tests/BraggIntegrationEngineMetalTest.cpp b/tests/BraggIntegrationEngineMetalTest.cpp index c602d3132..c7b786dca 100644 --- a/tests/BraggIntegrationEngineMetalTest.cpp +++ b/tests/BraggIntegrationEngineMetalTest.cpp @@ -3,8 +3,6 @@ #include -#ifdef __APPLE__ - #include #include #include @@ -17,12 +15,16 @@ #include "../common/PixelMask.h" #include "../common/Reflection.h" #include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h" +#ifdef __APPLE__ #include "../image_analysis/bragg_integration/BraggIntegrationEngineMetal.h" +#endif // The Metal engine's whole claim is that it IS the CPU engine, bit for bit (BraggIntegrationEngineMetal.h). // So nothing here has a tolerance: every float of every reflection is compared as a bit pattern, and // every count as an integer. A Mac without a Metal device fails rather than skips - the macOS CI job is -// the only place this can be checked, and a skip there would check nothing. +// the only place this can be checked, and a skip there would check nothing. The frames and the cases +// are built everywhere, and one test case checks on the CPU engine alone that they reach what they are +// there for. namespace { @@ -106,20 +108,34 @@ struct Case { float stencil_k = 0.0f; bool footprint = false; float clip_nsigma = 4.0f; + bool still = false; // a still keeps its overloads as lower bounds (saturated pixels at the limit) + float gain = 1.0f; // ADU per photon + bool curve = false; // a radial curve handed over (RadialBackgroundFromRings), which the clip follows }; +// The radial curve a case hands over: rising steeply enough with the radius that the clip it makes the +// ring follow (RadialClipOffset) moves by tens of counts across a ring. +void HandOverCurve(BraggIntegrationEngine &engine) { + std::vector rings(400); + for (size_t i = 0; i < rings.size(); ++i) + rings[i] = 1.0f + 0.01f * static_cast(i * i); + engine.RadialBackgroundFromRings(rings, static_cast(rings.size()), 0.05f, 0.01f); +} + DiffractionExperiment MakeExperiment(const DetectorSetup &det, const Case &c) { DiffractionExperiment x(det); x.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV) .BeamX_pxl(0.4f * static_cast(x.GetXPixelsNum())) .BeamY_pxl(0.55f * static_cast(x.GetYPixelsNum())); - // Rotation: stills keep overloads as lower bounds on CPU/CUDA, but not yet on Metal (platform difference, to be fixed). - x.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {})); - x.ImportIndexingSettings(IndexingSettings().RotationIndexing(true)); + // A rotation drops an overloaded observation, a still keeps it as a lower bound (Finalize). + if (!c.still) { + x.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {})); + x.ImportIndexingSettings(IndexingSettings().RotationIndexing(true)); + } x.BandwidthFWHM(c.bandwidth); BraggIntegrationSettings s; s.Integrator(c.mode).Overlap(c.overlap).BackgroundRadialCorrection(c.radial).StencilKSigma(c.stencil_k) - .BackgroundClipNSigma(c.clip_nsigma); + .BackgroundClipNSigma(c.clip_nsigma).PixelGain(c.gain); if (c.footprint) { // Wider than the r1 disk can hold from the middle out, so both stencils occur on one frame. SpotFootprint fp; @@ -134,6 +150,38 @@ DiffractionExperiment MakeExperiment(const DetectorSetup &det, const Case &c) { return x; } +} // namespace + +// Runs everywhere, on the CPU engine alone: the cases the Metal comparison below adds for the stills, the +// gain and the handed-over curve do reach what they are there for, so that comparison cannot pass by +// comparing nothing. +TEST_CASE("BraggIntegrationEngineMetal_CasesReachTheirModes", "[portable][Integration]") { + Case c; + c.still = true; + c.gain = 4.0f; + c.radial = true; + c.curve = true; + const DiffractionExperiment x = MakeExperiment(DetJF(2), c); + const PixelMask mask(x); + BraggIntegrationEngineCPU cpu(x, mask); + HandOverCurve(cpu); + cpu.BackgroundRadial(true); + // The clip follows the curve: it moves by whole counts across a ring. + CHECK(cpu.RadialClipOffset(500.0f, 10.0f) > 10); + CHECK(cpu.RadialClipOffset(500.0f, -10.0f) < -10); + ImagePreprocessorBuffer image(x.GetPixelsNum()); + const Frame f = RandomFrame(x, 7, {}); + std::copy(f.image.begin(), f.image.end(), image.data()); + const auto out = cpu.Run(image, f.predicted, f.predicted.size(), 0); + CHECK(out.size() > 1000); + // Saturated pixels counted at the limit: overloads kept as lower bounds. + CHECK(cpu.Counts().overloads_kept > 0); +} + +#ifdef __APPLE__ + +namespace { + uint32_t Bits(float f) { return std::bit_cast(f); } void RequireSame(const std::vector &cpu, const std::vector &metal) { @@ -156,6 +204,7 @@ void RequireSame(const std::vector &cpu, const std::vector 0); // Not the CPU engine compared with itself: every frame went through the kernels. REQUIRE(metal.FramesOnGPU() == 3); return n_out; @@ -244,6 +301,39 @@ TEST_CASE("BraggIntegrationEngineMetal_BitIdenticalToCPU", "[portable][Integrati Case no_clip; no_clip.clip_nsigma = 0.0f; cases.emplace_back("no clip", no_clip); + // A radial curve handed over, which the ring's clip follows (RadialClipOffset), under every overlap + // mode and every integrator. + for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical}) + for (OverlapMode overlap : {OverlapMode::Off, OverlapMode::Exclude}) { + Case curve; + curve.mode = mode; + curve.overlap = overlap; + curve.radial = true; + curve.curve = true; + cases.emplace_back("curve mode " + std::to_string(static_cast(mode)) + " overlap " + + std::to_string(static_cast(overlap)), curve); + } + // Stills: overloads kept as lower bounds, every saturated pixel counted. At a gain of one, and in + // detector units at two gains (the clip, the floors and the output's units all move with it), with + // and without the handed-over curve. + for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical}) + for (float gain : {1.0f, 4.0f, 1.7f}) { + Case still; + still.mode = mode; + still.still = true; + still.gain = gain; + still.radial = gain != 1.0f; + still.curve = gain == 4.0f; + cases.emplace_back("still mode " + std::to_string(static_cast(mode)) + " gain " + + std::to_string(gain) + (still.radial ? " radial" : "") + + (still.curve ? " curve" : ""), still); + } + // Footprint stencils on a still with a curve: the clip follows it on a per-reflection stencil. + Case footprint_still = footprint; + footprint_still.still = true; + footprint_still.curve = true; + footprint_still.gain = 4.0f; + cases.emplace_back("footprint still gain 4 curve", footprint_still); for (const auto &[name, c] : cases) { DYNAMIC_SECTION(name) { @@ -253,30 +343,6 @@ TEST_CASE("BraggIntegrationEngineMetal_BitIdenticalToCPU", "[portable][Integrati } } -// A radial curve handed over from the whole frame (RadialBackgroundFromRings) makes the ring's clip follow -// it, which the clip kernel does not: such a frame is integrated by the CPU engine, so the result is the -// CPU's exactly and no frame goes to the GPU. -TEST_CASE("BraggIntegrationEngineMetal_HandedOverRadialCurveStaysOnCPU", "[portable][Integration][Metal]") { - Case c; - c.radial = true; - const DiffractionExperiment x = MakeExperiment(DetJF(2), c); - const PixelMask mask(x); - BraggIntegrationEngineCPU cpu(x, mask); - BraggIntegrationEngineMetal metal(x, mask); - std::vector rings(400); - for (size_t i = 0; i < rings.size(); ++i) - rings[i] = 1.0f + 0.002f * static_cast(i * i); - cpu.RadialBackgroundFromRings(rings, static_cast(rings.size()), 0.05f, 0.01f); - metal.RadialBackgroundFromRings(rings, static_cast(rings.size()), 0.05f, 0.01f); - ImagePreprocessorBuffer image(x.GetPixelsNum()); - const Frame f = RandomFrame(x, 7, {}); - std::copy(f.image.begin(), f.image.end(), image.data()); - const auto out_cpu = cpu.Run(image, f.predicted, f.predicted.size(), 0); - RequireSame(out_cpu, metal.Run(image, f.predicted, f.predicted.size(), 0)); - CHECK(out_cpu.size() > 1000); - CHECK(metal.FramesOnGPU() == 0); -} - // The trimmed-mean background sorts each ring, which only the CPU engine does: such a run keeps the CPU // engine (MXAnalysisWithoutFPGA asks Available). TEST_CASE("BraggIntegrationEngineMetal_TrimmedMeanStaysOnCPU", "[portable][Integration][Metal]") {