reg_field: add field_matvec_rls/diag_rls/relax_rls device dispatch + bump pins

reg_field.cpp

@@ -105,4 +105,81 @@ void field_relax(
  throw std::invalid_argument("unsupported device");
 }
 
+void field_matvec_rls(
+ DLTensor & out ,
+ const DLTensor & inp ,
+ const DLTensor & wgt ,
+ const double * voxel_size,
+ const double * absolute ,
+ const double * membrane ,
+ const double * bending ,
+ int8_t bound ,
+ int ndim ,
+ int stream )
+{
+ require_same_device(out, inp);
+ require_same_device(out, wgt);
+#ifdef FF_WITH_CUDA
+ if (IS_CUDA(out))
+ return FF_CUDA::field_matvec_rls(out, inp, wgt, voxel_size, absolute, membrane, bending, bound, ndim, stream);
+#endif
+ if (IS_CPU(out))
+ return FF_CPU::field_matvec_rls(out, inp, wgt, voxel_size, absolute, membrane, bending, bound, ndim, stream);
+
+ if (IS_CUDA(out))
+ throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors");
+ throw std::invalid_argument("unsupported device");
+}
+
+void field_diag_rls(
+ DLTensor & out ,
+ const DLTensor & wgt ,
+ const double * voxel_size,
+ const double * absolute ,
+ const double * membrane ,
+ const double * bending ,
+ int8_t bound ,
+ int ndim ,
+ int stream )
+{
+ require_same_device(out, wgt);
+#ifdef FF_WITH_CUDA
+ if (IS_CUDA(out))
+ return FF_CUDA::field_diag_rls(out, wgt, voxel_size, absolute, membrane, bending, bound, ndim, stream);
+#endif
+ if (IS_CPU(out))
+ return FF_CPU::field_diag_rls(out, wgt, voxel_size, absolute, membrane, bending, bound, ndim, stream);
+
+ if (IS_CUDA(out))
+ throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors");
+ throw std::invalid_argument("unsupported device");
+}
+
+void field_relax_rls(
+ DLTensor & sol ,
+ const DLTensor & hes ,
+ const DLTensor & grd ,
+ const DLTensor & wgt ,
+ const double * voxel_size,
+ const double * absolute ,
+ const double * membrane ,
+ const double * bending ,
+ int8_t bound ,
+ int ndim ,
+ int nb_iter ,
+ int stream )
+{
+ require_same_device(sol, wgt);
+#ifdef FF_WITH_CUDA
+ if (IS_CUDA(sol))
+ return FF_CUDA::field_relax_rls(sol, hes, grd, wgt, voxel_size, absolute, membrane, bending, bound, ndim, nb_iter, stream);
+#endif
+ if (IS_CPU(sol))
+ return FF_CPU::field_relax_rls(sol, hes, grd, wgt, voxel_size, absolute, membrane, bending, bound, ndim, nb_iter, stream);
+
+ if (IS_CUDA(sol))
+ throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors");
+ throw std::invalid_argument("unsupported device");
+}
+
 FF_NAMESPACE_END(FF)

reg_field.h

@@ -134,6 +134,79 @@ void field_relax(
  int stream = 0
 );
 
+/**
+ * @brief Reweighted-least-squares (RLS/JRLS) variant of `field_matvec`.
+ *
+ * Same conventions as `field_matvec`, with an additional per-voxel weight
+ * map `wgt` that spatially modulates the penalty strength (e.g. for
+ * edge-preserving / robust regularisation). `wgt` has shape
+ * `(*batch, *spatial, 1)` for a single weight shared across all channels
+ * (RLS), or `(*batch, *spatial, C)` for a per-channel weight (JRLS, `C`
+ * matching `out`'s channel count) -- the trailing dimension of `wgt`
+ * selects which mode is used.
+ *
+ * @warning The `bending` order (both RLS and JRLS) has a known
+ * self-adjointness bug in its varying-weight coefficient math;
+ * `absolute` and `membrane` are verified self-adjoint.
+ *
+ * @param out Output tensor (*batch, *spatial, C)
+ * @param inp Input tensor (*batch, *spatial, C)
+ * @param wgt Weight tensor (*batch, *spatial, 1 or C)
+ */
+void field_matvec_rls(
+ DLTensor & out ,
+ const DLTensor & inp ,
+ const DLTensor & wgt ,
+ const double * voxel_size = nullptr,
+ const double * absolute = nullptr,
+ const double * membrane = nullptr,
+ const double * bending = nullptr,
+ int8_t bound = bound_t::DCT2,
+ int ndim = 1,
+ int stream = 0
+);
+
+/**
+ * @brief RLS/JRLS variant of `field_diag`, same weight-map conventions as
+ * `field_matvec_rls`. Writes into `out` (*batch, *spatial, C).
+ */
+void field_diag_rls(
+ DLTensor & out ,
+ const DLTensor & wgt ,
+ const double * voxel_size = nullptr,
+ const double * absolute = nullptr,
+ const double * membrane = nullptr,
+ const double * bending = nullptr,
+ int8_t bound = bound_t::DCT2,
+ int ndim = 1,
+ int stream = 0
+);
+
+/**
+ * @brief RLS/JRLS variant of `field_relax`, same weight-map conventions as
+ * `field_matvec_rls`.
+ *
+ * @param sol Field to refine, in/out (*batch, *spatial, C)
+ * @param hes Compact-symmetric Hessian (*batch, *spatial, C*(C+1)/2)
+ * @param grd Gradient (*batch, *spatial, C)
+ * @param wgt Weight tensor (*batch, *spatial, 1 or C)
+ * @param nb_iter Number of relaxation iterations
+ */
+void field_relax_rls(
+ DLTensor & sol ,
+ const DLTensor & hes ,
+ const DLTensor & grd ,
+ const DLTensor & wgt ,
+ const double * voxel_size = nullptr,
+ const double * absolute = nullptr,
+ const double * membrane = nullptr,
+ const double * bending = nullptr,
+ int8_t bound = bound_t::DCT2,
+ int ndim = 1,
+ int nb_iter = 1,
+ int stream = 0
+);
+
 FF_NAMESPACE_END(FF)
 
 #endif // FF_LIB_REG_FIELD