Deformable image registration on the GPU

cpp/core/gpu/gpu.cpp

@@ -870,14 +870,43 @@ Image<float> ComputeContext::download_texture_as_image(const Texture &texture,
   if (texture_channels == 4U && alpha_mode == DownloadTextureAlphaMode::IgnoreAlpha) {
     const std::vector<uint32_t> rgb_channels = {0U, 1U, 2U};
     Adapter::Extract1D source_rgb(source_image, texture_dims, rgb_channels);
-    threaded_copy(source_rgb, image);
+    // Use 2 inner axes so tiny channel loops don't dominate with threading overhead.
+    threaded_copy(source_rgb, image, 0, source_rgb.ndim(), 2);
   } else {
     threaded_copy(source_image, image);
   }
 
   return image;
 }
 
+void ComputeContext::wait_for_all_queue_operations() const {
+  bool callback_invoked = false;
+  wgpu::QueueWorkDoneStatus queue_status = wgpu::QueueWorkDoneStatus::Error;
+  std::string queue_message;
+  const wgpu::Future queue_done_future = m_device.GetQueue().OnSubmittedWorkDone(
+      wgpu::CallbackMode::WaitAnyOnly,
+      [&callback_invoked, &queue_status, &queue_message](wgpu::QueueWorkDoneStatus status, const char *message) {
+        callback_invoked = true;
+        queue_status = status;
+        if (message != nullptr) {
+          queue_message = message;
+        }
+      });
+  const wgpu::WaitStatus wait_status = m_instance.WaitAny(queue_done_future, std::numeric_limits<uint64_t>::max());
+  if (wait_status != wgpu::WaitStatus::Success) {
+    throw MR::Exception("Failed to wait for submitted GPU queue work: wgpu::Instance::WaitAny failed");
+  }
+  if (!callback_invoked) {
+    throw MR::Exception("Failed to wait for submitted GPU queue work: callback was not invoked");
+  }
+  if (queue_status != wgpu::QueueWorkDoneStatus::Success) {
+    const std::string suffix = queue_message.empty() ? std::string() : (": " + queue_message);
+    throw MR::Exception("Failed while waiting for submitted GPU queue work" + suffix);
+  }
+}
+
+void ComputeContext::reduce_memory_usage() const { dawn::native::ReduceMemoryUsage(m_device.Get()); }
+
 Kernel ComputeContext::new_kernel(const KernelSpec &kernel_spec) const {
   struct BindingEntries {
     std::vector<wgpu::BindGroupEntry> bind_group_entries;

cpp/core/gpu/gpu.h

@@ -222,6 +222,7 @@ struct ComputeContext {
   new_buffer_from_host_object(const Object &object, BufferType buffer_type = BufferType::StorageBuffer) const {
     static_assert(std::is_trivially_copyable_v<Object>, "Object must be trivially copyable");
     static_assert(std::is_standard_layout_v<Object>, "Object must be standard layout");
+    static_assert(sizeof(Object) % 4 == 0, "Object size must be a multiple of 4 bytes");
     return {buffer_type, inner_new_buffer_from_host_memory(&object, sizeof(object), buffer_type)};
   }
 
@@ -241,6 +242,16 @@ struct ComputeContext {
     return Buffer<T>{bufferType, std::move(buffer)};
   }
 
+  // Writes a POD-like object into a byte buffer (e.g. uniform buffers).
+  template <typename Object>
+  void write_object_to_buffer(const Buffer<std::byte> &buffer, const Object &object, uint64_t offset_bytes = 0) const {
+    static_assert(std::is_trivially_copyable_v<Object>, "Object must be trivially copyable");
+    static_assert(std::is_standard_layout_v<Object>, "Object must be standard layout");
+    static_assert(sizeof(Object) % 4 == 0, "Object size must be a multiple of 4 bytes");
+    const auto bytes = tcb::as_bytes(tcb::span<const Object>(&object, 1));
+    write_to_buffer<std::byte>(buffer, bytes, offset_bytes);
+  }
+
   // This function blocks until the download is complete.
   template <typename T = float> [[nodiscard]] std::vector<T> download_buffer_as_vector(const Buffer<T> &buffer) const {
     std::vector<T> result(buffer.wgpu_handle.GetSize() / sizeof(T));
@@ -329,6 +340,12 @@ struct ComputeContext {
 
   [[nodiscard]] Sampler new_linear_sampler() const;
 
+  // Block until all queued GPU work has completed.
+  void wait_for_all_queue_operations() const;
+
+  // Ask Dawn to free unused GPU memory (e.g. staging/cached resources).
+  void reduce_memory_usage() const;
+
 private:
   wgpu::Buffer inner_new_empty_buffer(size_t byteSize, BufferType bufferType = BufferType::StorageBuffer) const;
   wgpu::Buffer inner_new_buffer_from_host_memory(const void *srcMemory,

cpp/core/gpu/registration/adabelief.cpp

@@ -0,0 +1,83 @@
+/* Copyright (c) 2008-2026 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "gpu/registration/adabelief.h"
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstddef>
+#include <numeric>
+#include <tcb/span.hpp>
+#include <vector>
+
+AdaBelief::AdaBelief(const std::vector<Parameter> &parameters, float beta1, float beta2, float epsilon)
+    : m_parameters(parameters),
+      m_beta1(beta1),
+      m_beta2(beta2),
+      m_epsilon(epsilon),
+      m_timeStep(0),
+      m_firstMoments(parameters.size(), 0.0F),
+      m_secondMoments(parameters.size(), 0.0F),
+      m_mask(parameters.size(), 0),
+      m_updates(parameters.size(), 0.0F) {}
+
+std::vector<float> AdaBelief::parameterValues() const {
+  std::vector<float> values(m_parameters.size());
+  std::transform(
+      m_parameters.begin(), m_parameters.end(), values.begin(), [](const Parameter &param) { return param.value; });
+  return values;
+}
+
+void AdaBelief::setParameterValues(tcb::span<const float> values) {
+  for (size_t i = 0; i < m_parameters.size() && i < values.size(); ++i) {
+    m_parameters[i].value = values[i];
+  }
+}
+
+void AdaBelief::reset() {
+  m_timeStep = 0;
+  m_firstMoments.assign(m_parameters.size(), 0.0F);
+  m_secondMoments.assign(m_parameters.size(), 0.0F);
+  m_mask.assign(m_parameters.size(), 0);
+  m_updates.assign(m_parameters.size(), 0.0F);
+}
+
+void AdaBelief::step(const std::vector<float> &gradients) {
+  assert(gradients.size() == m_parameters.size());
+  ++m_timeStep;
+
+  // First pass: compute moment estimates, the update direction, and binary masks
+  for (size_t i = 0; i < m_parameters.size(); ++i) {
+    m_firstMoments[i] = m_beta1 * m_firstMoments[i] + (1.0F - m_beta1) * gradients[i];
+    const float diff = gradients[i] - m_firstMoments[i];
+    m_secondMoments[i] = m_beta2 * m_secondMoments[i] + (1.0F - m_beta2) * diff * diff;
+    const float firstMomentCorrected = m_firstMoments[i] / (1.0F - std::pow(m_beta1, m_timeStep));
+    const float secondMomentCorrected = m_secondMoments[i] / (1.0F - std::pow(m_beta2, m_timeStep));
+    m_updates[i] = firstMomentCorrected / (std::sqrt(secondMomentCorrected) + m_epsilon);
+    // Create mask: update only if update and gradient are aligned (i.e. product > 0)
+    m_mask[i] = (m_updates[i] * gradients[i] > 0.0F) ? 1 : 0;
+  }
+
+  // Compute the average mask value across all parameters
+  const float maskSum = std::accumulate(m_mask.begin(), m_mask.end(), 0);
+  const float maskMean = maskSum / static_cast<float>(m_mask.size());
+
+  // Second pass: apply cautious update to each parameter
+  for (size_t i = 0; i < m_parameters.size(); ++i) {
+    // If the mask is 0, no update is applied; otherwise, scale the update to compensate for the overall sparsity.
+    m_parameters[i].value -= m_parameters[i].learning_rate * (m_updates[i] * m_mask[i] / (maskMean + m_epsilon));
+  }
+}

cpp/core/gpu/registration/adabelief.h

@@ -0,0 +1,56 @@
+/* Copyright (c) 2008-2026 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <tcb/span.hpp>
+#include <vector>
+
+// AdaBelief is an improved version of Adam that takes into account the curvature of the loss function.
+// See here https://arxiv.org/abs/2010.07468
+// This version is further enhanced by the idea in this paper https://arxiv.org/abs/2411.16085
+// which consists in performing an element‐wise mask to the update such that only the components
+// where the proposed update direction and the current gradient are aligned
+// (i.e., have the same sign) are applied. This ensures that every step reliably reduces the loss
+// and avoids potential overshooting or oscillations.
+class AdaBelief {
+public:
+  struct Parameter {
+    float value;
+    float learning_rate;
+  };
+
+  AdaBelief(const std::vector<Parameter> &parameters, float beta1 = 0.7F, float beta2 = 0.9999F, float epsilon = 1e-6F);
+
+  void step(const std::vector<float> &gradients);
+  std::vector<float> parameterValues() const;
+  void setParameterValues(tcb::span<const float> values);
+
+  // Resets the optimizer internal state (moments and timestep) but keeps the parameters unchanged
+  void reset();
+
+private:
+  std::vector<Parameter> m_parameters;
+  float m_beta1;
+  float m_beta2;
+  float m_epsilon;
+  int m_timeStep;
+  std::vector<float> m_firstMoments;  // Exponential moving average of gradients (m_t)
+  std::vector<float> m_secondMoments; // Exponential moving average of squared deviations ((g_t - m_t)^2)
+  std::vector<uint32_t> m_mask;
+  std::vector<float> m_updates;
+};

cpp/core/gpu/registration/calculatorinterface.h

@@ -0,0 +1,53 @@
+/* Copyright (c) 2008-2026 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include "gpu/registration/ncccalculator.h"
+#include "gpu/registration/nmicalculator.h"
+#include "gpu/registration/registrationtypes.h"
+#include "gpu/registration/ssdcalculator.h"
+#include "match_variant.h"
+
+#include <cassert>
+#include <variant>
+
+namespace MR::GPU {
+
+class Calculator final : public std::variant<NMICalculator, SSDCalculator, NCCCalculator> {
+public:
+  // Expose std::variant constructors publicly
+  using std::variant<NMICalculator, SSDCalculator, NCCCalculator>::variant;
+
+  struct Config {
+    Texture fixed_texture;
+    Texture moving_texture;
+    MR::Transform fixed_transform;
+    MR::Transform moving_transform;
+    float downscale_factor;
+    GlobalMetric metric;
+  };
+
+  void update(const GlobalTransform &transformation) {
+    MR::match_v(*this, [&](auto &&arg) { arg.update(transformation); });
+  }
+
+  IterationResult get_result() const {
+    return MR::match_v(*this, [&](auto &&arg) { return arg.get_result(); });
+  }
+};
+
+} // namespace MR::GPU

cpp/core/gpu/registration/calculatoroutput.h

@@ -0,0 +1,25 @@
+/* Copyright (c) 2008-2026 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace MR::GPU {
+
+enum class CalculatorOutput : uint8_t { Cost, CostAndGradients };
+
+}

cpp/core/gpu/registration/convergencechecker.cpp

@@ -0,0 +1,76 @@
+/* Copyright (c) 2008-2026 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "gpu/registration/convergencechecker.h"
+#include "exception.h"
+#include <cassert>
+#include <cmath>
+#include <cstddef>
+#include <stdexcept>
+#include <string>
+#include <tcb/span.hpp>
+#include <vector>
+
+namespace MR {
+
+ConvergenceChecker::ConvergenceChecker(const Config &checkerConfiguration) : m_configuration(checkerConfiguration) {
+  assert(!m_configuration.param_thresholds.empty());
+  assert(m_configuration.patienceLimit > 0U);
+}
+
+bool ConvergenceChecker::has_converged(tcb::span<const float> current_params, float current_cost) {
+  if (m_configuration.param_thresholds.size() != current_params.size()) {
+    throw std::invalid_argument("ConvergenceChecker::has_converged: parameter threshold configuration mismatch.");
+  }
+
+  if (!m_initialized) {
+    DEBUG("ConvergenceChecker: Initializing with first parameters and cost.");
+    m_minimum_cost = current_cost;
+    m_best_params.assign(current_params.begin(), current_params.end());
+    m_initialized = true;
+    return false;
+  }
+
+  const bool has_better_cost = current_cost < m_minimum_cost;
+
+  const auto significant_param_improvement = [&]() {
+    for (size_t idx = 0; idx < current_params.size(); ++idx) {
+      const float param_diff = std::fabs(m_best_params[idx] - current_params[idx]);
+      if (param_diff >= m_configuration.param_thresholds[idx]) {
+        return true;
+      }
+    }
+    return false;
+  }();
+
+  if (has_better_cost) {
+    m_minimum_cost = current_cost;
+    m_best_params.assign(current_params.begin(), current_params.end());
+    // For the patience counter, only significant parameter improvements count
+    m_patience_counter = significant_param_improvement ? 0U : m_patience_counter + 1U;
+    DEBUG("ConvergenceChecker: Better cost found. Resetting patience counter to " + std::to_string(m_patience_counter) +
+          ".");
+  } else {
+    DEBUG("ConvergenceChecker: No better cost found. Incrementing patience counter.");
+    ++m_patience_counter;
+  }
+
+  return m_patience_counter >= m_configuration.patienceLimit;
+}
+
+void ConvergenceChecker::reset_patience() { m_patience_counter = 0; }
+
+} // namespace MR