From c5062327e888ffe083abfd16cb93b7d8e9dcb60b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 28 Jul 2026 05:49:35 +0000 Subject: [PATCH] reg_field: add field_matvec_rls/diag_rls/relax_rls device dispatch + bump pins Wires the reweighted-least-squares (RLS/JRLS) field regulariser API through the device-dispatch boundary: inspects the tensor's device and forwards to FF_CUDA:: (guarded by FF_WITH_CUDA) or FF_CPU::, mirroring the existing field_matvec/field_diag/field_relax dispatch. Declared in reg_field.h with the same signature as the CPU/CUDA libraries (an additional `wgt` tensor threaded alongside out/inp or sol/hes/grd). Bump submodule pins to the commits that add the RLS/JRLS dispatch on each backend: cpu -> fastfields-cpu-lib#40, cuda -> fastfields-cuda-lib#25. Same bending-order correctness caveat as the backend dispatch layers applies (tracked in fastfields-kernels#34) -- absolute/membrane are the verified-symmetric paths; documented in the field_matvec_rls docstring. Verified: `make all CXX=clang++` builds libfastfields-cpu.so and libfastfields.so clean. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_016AjQcY78NgbagPSbPJRr6Z --- cpu | 2 +- cuda | 2 +- reg_field.cpp | 77 +++++++++++++++++++++++++++++++++++++++++++++++++++ reg_field.h | 73 ++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 152 insertions(+), 2 deletions(-) diff --git a/cpu b/cpu index b2aa85f..042a841 160000 --- a/cpu +++ b/cpu @@ -1 +1 @@ -Subproject commit b2aa85f37a980696f040cd8602d4e8188b43acb2 +Subproject commit 042a8416a37fac8f765397c603efa65e31d4bd1c diff --git a/cuda b/cuda index 690d4dd..32a16f4 160000 --- a/cuda +++ b/cuda @@ -1 +1 @@ -Subproject commit 690d4dd8d9612e29d87f6271d9e74fa0e5399373 +Subproject commit 32a16f4f19eb2a23865b1123598d3aa03b13a2aa diff --git a/reg_field.cpp b/reg_field.cpp index e2e4e98..4601e4e 100644 --- a/reg_field.cpp +++ b/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) diff --git a/reg_field.h b/reg_field.h index 67f5940..62a4e2a 100644 --- a/reg_field.h +++ b/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