From 8f08c7d6bf23dd94cfdac95cf7d2ed12a592125d Mon Sep 17 00:00:00 2001 From: Yael Date: Thu, 20 Aug 2026 11:44:40 +0000 Subject: [PATCH 1/7] perf(cuda): split reg_flow into per-(family, ndim) translation units reg_flow is the binding constraint on this whole project's CI. nvcc peaks at 12.98 GiB compiling it -- 81% of a 16 GB runner, in one indivisible TU -- and that single number is what pins `make cuda` at -j2, which in turn is what makes a multi-architecture -gencode set unaffordable against the 120-minute timeout. The cost is template instantiation, and it factors. Each exported entry point selects exactly one internal wrapper template (_flow_matvec, _flow_matvec_acc, _flow_diag, _flow_kernel, _flow_relax) and dispatches it over ndim x boundary x dtype x offset width. No two arms of that pyramid share an instantiation, so a TU that only ever reaches `_flow_diag<2, ...>` pays for `_flow_diag<2, ...>` and nothing else. Cutting the file along those seams cuts the instantiation set the same way. What is here ------------ * reg_flow.cpp keeps all 15 exported entry points, every argument check and every error message, and instantiates NOTHING. It normalises strides, validates, decides the offset width, and `switch (ndim)`es into a slice. 205 MB / 2.0 s to compile. * reg_flow_slice.h -- the seam: one hidden-visibility declaration per (entry point, ndim), 30 in all. * reg_flow_slice.inl -- the wrapper templates and the dtype x offset x bound dispatch, verbatim from the single-TU form, plus the slice definitions the including TU selects with FF_FLOW_SLICE_*. * reg_flow_{matvec,diag,kernel,relax}_{1,2,3}d.cpp -- twelve four-line TUs. Why a forwarding seam and not `-D` on one source ------------------------------------------------ Compiling reg_flow.cpp N times with different -D flags, the way BOUNDFLAGS already varies a build, gives N objects that each define ff::cuda::flow_matvec and the link fails. The exported entry point has to live in exactly one TU. `extern template` was the other candidate and was rejected: it has to name every leaf, and with FF_STATIC_BOUNDS=0 six of the eight FF_BOUND_* selectors collapse onto bound::type::Dynamic, so the enumeration would contain the same specialization six times -- ill-formed under [temp.explicit]/5, and which leaves collapse is a build-flag decision. A `switch` has no such problem: duplicate *implicit* instantiations across its arms are the same instantiation. ABI --- The slice symbols are __attribute__((visibility("hidden"))), so libfastfields-cuda.so's dynamic symbol table is unchanged. Every slice function is defined by exactly one TU in MODULES -- too few and the link fails undefined, too many and it fails duplicate. The Makefile list and the FF_FLOW_SLICE_* selection cannot silently disagree. Also: reg_flow.cpp now includes impl/kernels/utils.h explicitly. FF_CANUSE32BITS expands to an unqualified canUse32BitIndexMath, which is declared there and NOT in core/autocast.h, whose include comment in core/dispatch.h says otherwise. Every other dispatch source pulls the kernels in wholesale and never noticed; the front TU is the first that does not. Measured numbers land in the Makefile once build-cuda has run on this branch. --- src/lib-cuda/Makefile | 34 ++ src/lib-cuda/reg_flow.cpp | 634 +++----------------------- src/lib-cuda/reg_flow_diag_1d.cpp | 11 + src/lib-cuda/reg_flow_diag_2d.cpp | 11 + src/lib-cuda/reg_flow_diag_3d.cpp | 11 + src/lib-cuda/reg_flow_kernel_1d.cpp | 11 + src/lib-cuda/reg_flow_kernel_2d.cpp | 11 + src/lib-cuda/reg_flow_kernel_3d.cpp | 11 + src/lib-cuda/reg_flow_matvec_1d.cpp | 11 + src/lib-cuda/reg_flow_matvec_2d.cpp | 11 + src/lib-cuda/reg_flow_matvec_3d.cpp | 11 + src/lib-cuda/reg_flow_relax_1d.cpp | 11 + src/lib-cuda/reg_flow_relax_2d.cpp | 11 + src/lib-cuda/reg_flow_relax_3d.cpp | 11 + src/lib-cuda/reg_flow_slice.h | 169 +++++++ src/lib-cuda/reg_flow_slice.inl | 681 ++++++++++++++++++++++++++++ 16 files changed, 1080 insertions(+), 570 deletions(-) create mode 100644 src/lib-cuda/reg_flow_diag_1d.cpp create mode 100644 src/lib-cuda/reg_flow_diag_2d.cpp create mode 100644 src/lib-cuda/reg_flow_diag_3d.cpp create mode 100644 src/lib-cuda/reg_flow_kernel_1d.cpp create mode 100644 src/lib-cuda/reg_flow_kernel_2d.cpp create mode 100644 src/lib-cuda/reg_flow_kernel_3d.cpp create mode 100644 src/lib-cuda/reg_flow_matvec_1d.cpp create mode 100644 src/lib-cuda/reg_flow_matvec_2d.cpp create mode 100644 src/lib-cuda/reg_flow_matvec_3d.cpp create mode 100644 src/lib-cuda/reg_flow_relax_1d.cpp create mode 100644 src/lib-cuda/reg_flow_relax_2d.cpp create mode 100644 src/lib-cuda/reg_flow_relax_3d.cpp create mode 100644 src/lib-cuda/reg_flow_slice.h create mode 100644 src/lib-cuda/reg_flow_slice.inl diff --git a/src/lib-cuda/Makefile b/src/lib-cuda/Makefile index de2608c..9901bf2 100644 --- a/src/lib-cuda/Makefile +++ b/src/lib-cuda/Makefile @@ -139,6 +139,39 @@ SPLINEFLAGS ?= -DFF_STATIC_SPLINES=0 \ # four cheapest modules in the file (resize 2.00, restrict 1.30, splinc 0.42, # posdef 0.37 GB), together well under a single regulariser. They cost ~7 min # of the ~31 min compile. +# ~~~ reg_flow IS SPLIT ACROSS TRANSLATION UNITS ~~~ +# +# `reg_flow` was the 12.98 GiB entry in the table above -- 81% of a 16 GB +# runner, in one indivisible TU, and the single number that pinned this build +# at -j2. It is now twelve instantiation slices plus a thin front TU +# (`reg_flow.cpp`), which keeps every exported symbol and instantiates nothing; +# see reg_flow_slice.h for the seam and reg_flow_slice.inl for the bodies. +# +# The cut is (operation family) x (ndim), which is exactly how the dispatch +# already factors: each entry point selects one internal wrapper template and +# `switch (ndim)`es it, and no two arms share an instantiation. So a slice pays +# for its own arm and nothing else. The exported ABI is unchanged -- the slice +# symbols are hidden-visibility and never reach .dynsym. +# +# TBD-MEASURED-TABLE +# +# ORDER MATTERS HERE. make dispatches in MODULES order, so the slices are +# listed heaviest-first (3D, then 2D, then 1D): with the long poles started +# first, the short ones backfill the tail instead of extending it. +REG_FLOW_SLICES = \ + reg_flow_matvec_3d \ + reg_flow_kernel_3d \ + reg_flow_relax_3d \ + reg_flow_diag_3d \ + reg_flow_matvec_2d \ + reg_flow_kernel_2d \ + reg_flow_relax_2d \ + reg_flow_diag_2d \ + reg_flow_matvec_1d \ + reg_flow_kernel_1d \ + reg_flow_relax_1d \ + reg_flow_diag_1d + MODULES = \ distance \ posdef \ @@ -148,6 +181,7 @@ MODULES = \ reg_field \ reg_field_rls \ reg_flow \ + $(REG_FLOW_SLICES) \ reg_flow_rls \ pushpull \ pushpull_backward diff --git a/src/lib-cuda/reg_flow.cpp b/src/lib-cuda/reg_flow.cpp index fa7a2b2..a00428b 100644 --- a/src/lib-cuda/reg_flow.cpp +++ b/src/lib-cuda/reg_flow.cpp @@ -1,480 +1,66 @@ +/** + * `reg_flow`'s exported entry points -- and nothing else. + * + * This translation unit instantiates no kernel template. It normalises + * strides, validates arguments, decides the offset width, and hands the call + * to one of the slice functions declared in `reg_flow_slice.h`; the slice TUs + * (`reg_flow__d.cpp`) hold the instantiations. See that header for + * why the seam is a forwarding call rather than `extern template` or a + * per-slice `-D` on one source. + * + * Every argument check, every error message and the order they fire in are + * unchanged from the single-TU form. That matters: they are the observable + * behaviour of these functions and the hub's tests pin them. + */ + #include #include #include #include "fastfields/api/cuda/reg_flow.h" #include "fastfields/api/cuda/posdef.h" -#include "fastfields/core/autocast.h" #include "fastfields/core/dispatch.h" -#include "fastfields/api/cuda/stream.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" -#include "fastfields/impl/kernels/bounds.h" -#include "fastfields/impl/kernels/utils.h" -#include "fastfields/impl/cuda/reg_flow.h" +// `FF_CANUSE32BITS` expands to an unqualified `canUse32BitIndexMath`, which is +// declared in impl/kernels/utils.h -- NOT in core/autocast.h, whose `// +// canUse32BitIndexMath` include comment in core/dispatch.h suggests otherwise. +// Every other dispatch source pulls the kernels in wholesale and never noticed; +// this TU is the first that does not, so it has to name the real home. +#include "fastfields/impl/kernels/utils.h" // canUse32BitIndexMath +#include "fastfields/impl/cuda/utils.h" // allocDevice / freeDevice +#include "reg_flow_slice.h" FF_NAMESPACE_BEGIN(FF_NS) FF_NAMESPACE_BEGIN(FF_DEVICE) -// reduction / accumulation type (matches jitfields' float64 default) -typedef double reduce_t; - -/*********************************************************************** - * WRAPPERS * - ***********************************************************************/ - -namespace { - -// length of the shape/stride arrays: (*batch, *spatial, C) == out.ndim -template -inline void _flow_matvec( - const bound::BoundVec & bvec, - int64_t nbatch , - void * out , - const void * inp , - const double * voxel_size , - double absolute , - double membrane , - double bending , - double shears , - double div , - const int64_t * size , - const int64_t * stride_out , - const int64_t * stride_inp , - cudaStream_t stream ) -{ - const int64_t nall1 = nbatch + ndim + 1; - const offset_t * _size = copy_if_needed(size, nall1); - const offset_t * _stride_out = copy_if_needed(stride_out, nall1); - const offset_t * _stride_inp = copy_if_needed(stride_inp, nall1); - scalar_t * _out = static_cast< scalar_t *>(out); - const scalar_t * _inp = static_cast(inp); - - reduce_t vx[ndim]; - for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; - - // The linear-elastic (Lamé) terms `shears`/`div` couple the flow channels, - // so any non-zero one selects the full combined stencil (matvec_all, which - // also folds in absolute/membrane/bending). Otherwise fall back to the - // cheaper single-penalty stencils (highest-order non-zero wins). - if (shears != 0.0 || div != 0.0) - reg_flow::matvec_all( - bvec, static_cast(nbatch), _out, _inp, - _size, _stride_out, _stride_inp, vx, absolute, membrane, bending, stream); - else if (membrane != 0.0) - reg_flow::matvec_membrane( - bvec, static_cast(nbatch), _out, _inp, - _size, _stride_out, _stride_inp, vx, absolute, stream); - - free_if_needed(_size); - free_if_needed(_stride_out); - free_if_needed(_stride_inp); -} - -// Accumulate variant of _flow_matvec: out += L(inp) (op='+') or out -= L(inp) -// (op='-'), instead of overwriting out. Mirrors the CPU `_flow_matvec_acc`. -template -inline void _flow_matvec_acc( - const bound::BoundVec & bvec, - int64_t nbatch , - void * out , - const void * inp , - const double * voxel_size , - double absolute , - double membrane , - double bending , - double shears , - double div , - const int64_t * size , - const int64_t * stride_out , - const int64_t * stride_inp , - cudaStream_t stream ) -{ - const int64_t nall1 = nbatch + ndim + 1; - const offset_t * _size = copy_if_needed(size, nall1); - const offset_t * _stride_out = copy_if_needed(stride_out, nall1); - const offset_t * _stride_inp = copy_if_needed(stride_inp, nall1); - scalar_t * _out = static_cast< scalar_t *>(out); - const scalar_t * _inp = static_cast(inp); - - reduce_t vx[ndim]; - for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; - - if (shears != 0.0 || div != 0.0) - reg_flow::matvec_all( - bvec, static_cast(nbatch), _out, _inp, - _size, _stride_out, _stride_inp, vx, - absolute, membrane, bending, shears, div, stream); - else if (bending != 0.0) - reg_flow::matvec_bending( - bvec, static_cast(nbatch), _out, _inp, - _size, _stride_out, _stride_inp, vx, absolute, membrane, bending, stream); - else if (membrane != 0.0) - reg_flow::matvec_membrane( - bvec, static_cast(nbatch), _out, _inp, - _size, _stride_out, _stride_inp, vx, absolute, membrane, stream); - else - reg_flow::matvec_absolute( - bvec, static_cast(nbatch), _out, _inp, - _size, _stride_out, _stride_inp, vx, absolute, stream); - - free_if_needed(_size); - free_if_needed(_stride_out); - free_if_needed(_stride_inp); -} - -template -inline void _flow_diag( - const bound::BoundVec & bvec, - int64_t nbatch , - void * out , - const double * voxel_size , - double absolute , - double membrane , - double bending , - double shears , - double div , - const int64_t * size , - const int64_t * stride_out , - cudaStream_t stream ) -{ - const int64_t nall1 = nbatch + ndim + 1; - const offset_t * _size = copy_if_needed(size, nall1); - const offset_t * _stride_out = copy_if_needed(stride_out, nall1); - scalar_t * _out = static_cast(out); - - reduce_t vx[ndim]; - for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; - - if (shears != 0.0 || div != 0.0) - reg_flow::diag_all( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, membrane, bending, shears, div, stream); - else if (bending != 0.0) - reg_flow::diag_bending( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, membrane, bending, stream); - else if (membrane != 0.0) - reg_flow::diag_membrane( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, membrane, stream); - else - reg_flow::diag_absolute( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, stream); - - free_if_needed(_size); - free_if_needed(_stride_out); -} - -// Materialise the Toeplitz convolution kernel (stencil) of the operator (see -// cpu-lib). `nfull` is the length of the size/stride arrays (== out.ndim): -// nbatch+ndim+1 for the per-channel vector stencil, nbatch+ndim+2 for the Lamé -// (cross-channel) matrix stencil. -template -inline void _flow_kernel( - const bound::BoundVec & bvec, - int64_t nbatch , - void * out , - const double * voxel_size , - double absolute , - double membrane , - double bending , - double shears , - double div , - const int64_t * size , - const int64_t * stride_out , - int64_t nfull , - cudaStream_t stream ) -{ - const offset_t * _size = copy_if_needed(size, nfull); - const offset_t * _stride_out = copy_if_needed(stride_out, nfull); - scalar_t * _out = static_cast(out); - - reduce_t vx[ndim]; - for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; - - if (shears != 0.0 || div != 0.0) { - if (bending != 0.0) - reg_flow::kernel_all( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, membrane, bending, shears, div, stream); - else - reg_flow::kernel_lame( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, membrane, shears, div, stream); - } else if (bending != 0.0) - reg_flow::kernel_bending( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, membrane, bending, stream); - else if (membrane != 0.0) - reg_flow::kernel_membrane( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, membrane, stream); - else - reg_flow::kernel_absolute( - bvec, static_cast(nbatch), _out, - _size, _stride_out, vx, absolute, stream); - - free_if_needed(_size); - free_if_needed(_stride_out); -} - -// In-place relaxation sweeps solving `(H + L) x = g` (see cpu-lib). -template -inline void _flow_relax( - const bound::BoundVec & bvec, - int64_t nbatch , - void * sol , - const void * hes , - const void * grd , - const double * voxel_size , - double absolute , - double membrane , - double bending , - double shears , - double div , - int niter , - const int64_t * size , - const int64_t * stride_sol , - const int64_t * stride_hes , - const int64_t * stride_grd , - cudaStream_t stream ) -{ - const int64_t nall1 = nbatch + ndim + 1; - const offset_t * _size = copy_if_needed(size, nall1); - const offset_t * _stride_sol = copy_if_needed(stride_sol, nall1); - const offset_t * _stride_hes = copy_if_needed(stride_hes, nall1); - const offset_t * _stride_grd = copy_if_needed(stride_grd, nall1); - scalar_t * _sol = static_cast< scalar_t *>(sol); - const scalar_t * _hes = static_cast(hes); - const scalar_t * _grd = static_cast(grd); - - reduce_t vx[ndim]; - for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; - - if (shears != 0.0 || div != 0.0) { - if (bending != 0.0) - reg_flow::relax_all_( - bvec, static_cast(nbatch), _sol, _hes, _grd, - _size, _stride_sol, _stride_hes, _stride_grd, vx, - absolute, membrane, bending, shears, div, niter, stream); - else - reg_flow::relax_lame_( - bvec, static_cast(nbatch), _sol, _hes, _grd, - _size, _stride_sol, _stride_hes, _stride_grd, vx, - absolute, membrane, shears, div, niter, stream); - } else if (bending != 0.0) - reg_flow::relax_bending_( - bvec, static_cast(nbatch), _sol, _hes, _grd, - _size, _stride_sol, _stride_hes, _stride_grd, vx, - absolute, membrane, bending, niter, stream); - else - reg_flow::relax_membrane_( - bvec, static_cast(nbatch), _sol, _hes, _grd, - _size, _stride_sol, _stride_hes, _stride_grd, vx, - absolute, membrane, niter, stream); - - free_if_needed(_size); - free_if_needed(_stride_sol); - free_if_needed(_stride_hes); - free_if_needed(_stride_grd); -} - -} // anonymous namespace - /*********************************************************************** * DISPATCH * ***********************************************************************/ -#define BND1(B) B -#define BND2(B) B, B -#define BND3(B) B, B, B - -// matvec dtype x offset dispatch, given ndim and the (repeated) bound pack. -#define MV_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_matvec(MV_ARGS) \ - : _flow_matvec(MV_ARGS); \ - case 64: return use_32bits \ - ? _flow_matvec(MV_ARGS) \ - : _flow_matvec(MV_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define ADD_MV_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_matvec_acc(MV_ARGS) \ - : _flow_matvec_acc(MV_ARGS); \ - case 64: return use_32bits \ - ? _flow_matvec_acc(MV_ARGS) \ - : _flow_matvec_acc(MV_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define SUB_MV_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_matvec_acc(MV_ARGS) \ - : _flow_matvec_acc(MV_ARGS); \ - case 64: return use_32bits \ - ? _flow_matvec_acc(MV_ARGS) \ - : _flow_matvec_acc(MV_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define DG_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_diag(DG_ARGS); \ - case 64: return use_32bits \ - ? _flow_diag(DG_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define ADD_DG_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_diag(DG_ARGS) \ - : _flow_diag(DG_ARGS); \ - case 64: return use_32bits \ - ? _flow_diag(DG_ARGS) \ - : _flow_diag(DG_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define SUB_DG_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_diag(DG_ARGS) \ - : _flow_diag(DG_ARGS); \ - case 64: return use_32bits \ - ? _flow_diag(DG_ARGS) \ - : _flow_diag(DG_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define KN_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_kernel(KN_ARGS); \ - case 64: return use_32bits \ - ? _flow_kernel(KN_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define ADD_KN_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_kernel(KN_ARGS) \ - : _flow_kernel(KN_ARGS); \ - case 64: return use_32bits \ - ? _flow_kernel(KN_ARGS) \ - : _flow_kernel(KN_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define SUB_KN_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_kernel(KN_ARGS) \ - : _flow_kernel(KN_ARGS); \ - case 64: return use_32bits \ - ? _flow_kernel(KN_ARGS) \ - : _flow_kernel(KN_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define RX_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _flow_relax(RX_ARGS) \ - : _flow_relax(RX_ARGS); \ - case 64: return use_32bits \ - ? _flow_relax(RX_ARGS) \ - : _flow_relax(RX_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -// Which of these boundary conditions gets a dedicated (static) instantiation -// and which shares the single Dynamic (runtime) one is a build-time choice -- -// see FF_STATIC_BOUND_* in kernels/bounds.h. `bvec` carries the runtime -// condition for whichever ones fall back to Dynamic. -#define BOUND_SWITCH(DT, NDIM, BND) \ - switch (bnd) { \ - case bound::type::Zero: DT(NDIM, BND(FF_BOUND_ZERO)); break; \ - case bound::type::Replicate: DT(NDIM, BND(FF_BOUND_REPLICATE)); break; \ - case bound::type::DCT1: DT(NDIM, BND(FF_BOUND_DCT1)); break; \ - case bound::type::DCT2: DT(NDIM, BND(FF_BOUND_DCT2)); break; \ - case bound::type::DST1: DT(NDIM, BND(FF_BOUND_DST1)); break; \ - case bound::type::DST2: DT(NDIM, BND(FF_BOUND_DST2)); break; \ - case bound::type::DFT: DT(NDIM, BND(FF_BOUND_DFT)); break; \ - case bound::type::NoCheck: DT(NDIM, BND(FF_BOUND_NOCHECK)); break; \ - default: throw std::invalid_argument("Unsupported boundary condition"); \ - } - -#define NDIM_SWITCH(DT) \ +// What is left of the old `NDIM_SWITCH`: it still picks the spatial rank, but +// the arm is now a call into another translation unit instead of a template +// argument. The dtype x offset x boundary pyramid that used to sit under each +// arm moved with the instantiations, into the slice TUs. +#define FF_FLOW_ND_SWITCH(FN, ARGS) \ switch (ndim) { \ - case 1: BOUND_SWITCH(DT, 1, BND1); break; \ - case 2: BOUND_SWITCH(DT, 2, BND2); break; \ - case 3: BOUND_SWITCH(DT, 3, BND3); break; \ + case 1: return flow_slice::FN##_1d ARGS; \ + case 2: return flow_slice::FN##_2d ARGS; \ + case 3: return flow_slice::FN##_3d ARGS; \ default: throw std::invalid_argument("Only 1D, 2D and 3D flow are supported"); \ } +#define FF_FLOW_MV_CALL \ + (out, inp, voxel_size, absolute, membrane, bending, shears, div, \ + bound, nbatch, use_32bits, stream) + +#define FF_FLOW_DG_CALL \ + (out, voxel_size, absolute, membrane, bending, shears, div, \ + bound, nbatch, use_32bits, stream) + +#define FF_FLOW_RX_CALL \ + (sol, hes, grd, voxel_size, absolute, membrane, bending, shears, \ + div, bound, nb_iter, nbatch, use_32bits, stream) + void flow_matvec( DLTensor & out_ , const DLTensor & inp_ , @@ -503,18 +89,9 @@ void flow_matvec( FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) - const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); -#define MV_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, inp.strides, cstream - NDIM_SWITCH(MV_DT) -#undef MV_ARGS + FF_FLOW_ND_SWITCH(matvec, FF_FLOW_MV_CALL) } /** @@ -549,18 +126,9 @@ void flow_addmatvec_( FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) - const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); -#define MV_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, inp.strides, cstream - NDIM_SWITCH(ADD_MV_DT) -#undef MV_ARGS + FF_FLOW_ND_SWITCH(addmatvec, FF_FLOW_MV_CALL) } /** @@ -595,18 +163,9 @@ void flow_submatvec_( FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) - const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); -#define MV_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, inp.strides, cstream - NDIM_SWITCH(SUB_MV_DT) -#undef MV_ARGS + FF_FLOW_ND_SWITCH(submatvec, FF_FLOW_MV_CALL) } void flow_diag( @@ -632,18 +191,9 @@ void flow_diag( throw std::invalid_argument("ndim is larger than the tensor rank"); FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = FF_CANUSE32BITS(out); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out); -#define DG_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, cstream - NDIM_SWITCH(DG_DT) -#undef DG_ARGS + FF_FLOW_ND_SWITCH(diag, FF_FLOW_DG_CALL) } /** @@ -673,18 +223,9 @@ void flow_adddiag_( throw std::invalid_argument("ndim is larger than the tensor rank"); FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = FF_CANUSE32BITS(out); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out); -#define DG_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, cstream - NDIM_SWITCH(ADD_DG_DT) -#undef DG_ARGS + FF_FLOW_ND_SWITCH(adddiag, FF_FLOW_DG_CALL) } /** @@ -714,18 +255,9 @@ void flow_subdiag_( throw std::invalid_argument("ndim is larger than the tensor rank"); FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = FF_CANUSE32BITS(out); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out); -#define DG_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, cstream - NDIM_SWITCH(SUB_DG_DT) -#undef DG_ARGS + FF_FLOW_ND_SWITCH(subdiag, FF_FLOW_DG_CALL) } void flow_kernel( @@ -760,18 +292,9 @@ void flow_kernel( FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = FF_CANUSE32BITS(out); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out); -#define KN_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, static_cast(out.ndim), cstream - NDIM_SWITCH(KN_DT) -#undef KN_ARGS + FF_FLOW_ND_SWITCH(kernel, FF_FLOW_DG_CALL) } /** @@ -810,18 +333,9 @@ void flow_addkernel_( FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = FF_CANUSE32BITS(out); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out); -#define KN_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, static_cast(out.ndim), cstream - NDIM_SWITCH(ADD_KN_DT) -#undef KN_ARGS + FF_FLOW_ND_SWITCH(addkernel, FF_FLOW_DG_CALL) } /** @@ -860,18 +374,9 @@ void flow_subkernel_( FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = FF_CANUSE32BITS(out); - const auto code = static_cast(out.dtype.code); - const auto bits = out.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = _reg_stream(stream); + const bool use_32bits = FF_CANUSE32BITS(out); -#define KN_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, shears, div, \ - out.shape, out.strides, static_cast(out.ndim), cstream - NDIM_SWITCH(SUB_KN_DT) -#undef KN_ARGS + FF_FLOW_ND_SWITCH(subkernel, FF_FLOW_DG_CALL) } void flow_relax( @@ -902,21 +407,10 @@ void flow_relax( FF_CHECK_SAME (grd.shape[grd.ndim-1], (int64_t)ndim, "Gradient channel dimension must equal ndim") FF_CHECK_SAME_SHAPE_N(sol, grd, sol.ndim) - const bool use_32bits = FF_CANUSE32BITS(sol) && FF_CANUSE32BITS(hes) && - FF_CANUSE32BITS(grd); - const auto code = static_cast(sol.dtype.code); - const auto bits = sol.dtype.bits; - const bound::type bnd = static_cast(bound); - const bound::BoundVec bvec(bnd); - const cudaStream_t cstream = - reinterpret_cast(static_cast(stream)); + const bool use_32bits = FF_CANUSE32BITS(sol) && FF_CANUSE32BITS(hes) && + FF_CANUSE32BITS(grd); -#define RX_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(sol), FF_CVOIDPTR(hes), \ - FF_CVOIDPTR(grd), voxel_size, absolute, membrane, bending, \ - shears, div, nb_iter, sol.shape, sol.strides, hes.strides, \ - grd.strides, cstream - NDIM_SWITCH(RX_DT) -#undef RX_ARGS + FF_FLOW_ND_SWITCH(relax, FF_FLOW_RX_CALL) } void flow_forward( diff --git a/src/lib-cuda/reg_flow_diag_1d.cpp b/src/lib-cuda/reg_flow_diag_1d.cpp new file mode 100644 index 0000000..8a1113f --- /dev/null +++ b/src/lib-cuda/reg_flow_diag_1d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_diag / flow_adddiag_ / flow_subdiag_, 1D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_DIAG 1 +#define FF_FLOW_SLICE_ND1 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_diag_2d.cpp b/src/lib-cuda/reg_flow_diag_2d.cpp new file mode 100644 index 0000000..d109bb9 --- /dev/null +++ b/src/lib-cuda/reg_flow_diag_2d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_diag / flow_adddiag_ / flow_subdiag_, 2D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_DIAG 1 +#define FF_FLOW_SLICE_ND2 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_diag_3d.cpp b/src/lib-cuda/reg_flow_diag_3d.cpp new file mode 100644 index 0000000..e4d1f15 --- /dev/null +++ b/src/lib-cuda/reg_flow_diag_3d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_diag / flow_adddiag_ / flow_subdiag_, 3D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_DIAG 1 +#define FF_FLOW_SLICE_ND3 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_kernel_1d.cpp b/src/lib-cuda/reg_flow_kernel_1d.cpp new file mode 100644 index 0000000..ee7fb90 --- /dev/null +++ b/src/lib-cuda/reg_flow_kernel_1d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_kernel / flow_addkernel_ / flow_subkernel_, 1D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_KERNEL 1 +#define FF_FLOW_SLICE_ND1 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_kernel_2d.cpp b/src/lib-cuda/reg_flow_kernel_2d.cpp new file mode 100644 index 0000000..2de5820 --- /dev/null +++ b/src/lib-cuda/reg_flow_kernel_2d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_kernel / flow_addkernel_ / flow_subkernel_, 2D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_KERNEL 1 +#define FF_FLOW_SLICE_ND2 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_kernel_3d.cpp b/src/lib-cuda/reg_flow_kernel_3d.cpp new file mode 100644 index 0000000..7811494 --- /dev/null +++ b/src/lib-cuda/reg_flow_kernel_3d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_kernel / flow_addkernel_ / flow_subkernel_, 3D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_KERNEL 1 +#define FF_FLOW_SLICE_ND3 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_matvec_1d.cpp b/src/lib-cuda/reg_flow_matvec_1d.cpp new file mode 100644 index 0000000..23c909c --- /dev/null +++ b/src/lib-cuda/reg_flow_matvec_1d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_matvec / flow_addmatvec_ / flow_submatvec_, 1D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_MATVEC 1 +#define FF_FLOW_SLICE_ND1 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_matvec_2d.cpp b/src/lib-cuda/reg_flow_matvec_2d.cpp new file mode 100644 index 0000000..aea82b0 --- /dev/null +++ b/src/lib-cuda/reg_flow_matvec_2d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_matvec / flow_addmatvec_ / flow_submatvec_, 2D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_MATVEC 1 +#define FF_FLOW_SLICE_ND2 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_matvec_3d.cpp b/src/lib-cuda/reg_flow_matvec_3d.cpp new file mode 100644 index 0000000..fb58747 --- /dev/null +++ b/src/lib-cuda/reg_flow_matvec_3d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_matvec / flow_addmatvec_ / flow_submatvec_, 3D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_MATVEC 1 +#define FF_FLOW_SLICE_ND3 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_relax_1d.cpp b/src/lib-cuda/reg_flow_relax_1d.cpp new file mode 100644 index 0000000..034661a --- /dev/null +++ b/src/lib-cuda/reg_flow_relax_1d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_relax, 1D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_RELAX 1 +#define FF_FLOW_SLICE_ND1 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_relax_2d.cpp b/src/lib-cuda/reg_flow_relax_2d.cpp new file mode 100644 index 0000000..65d58a3 --- /dev/null +++ b/src/lib-cuda/reg_flow_relax_2d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_relax, 2D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_RELAX 1 +#define FF_FLOW_SLICE_ND2 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_relax_3d.cpp b/src/lib-cuda/reg_flow_relax_3d.cpp new file mode 100644 index 0000000..9a04a05 --- /dev/null +++ b/src/lib-cuda/reg_flow_relax_3d.cpp @@ -0,0 +1,11 @@ +// One reg_flow instantiation slice: flow_relax, 3D. +// +// The exported symbols live in reg_flow.cpp; this translation unit holds +// only the template instantiations that one arm of its `switch (ndim)` +// reaches. See reg_flow_slice.h for why the seam exists and +// reg_flow_slice.inl for the bodies these two defines select. + +#define FF_FLOW_SLICE_RELAX 1 +#define FF_FLOW_SLICE_ND3 1 + +#include "reg_flow_slice.inl" diff --git a/src/lib-cuda/reg_flow_slice.h b/src/lib-cuda/reg_flow_slice.h new file mode 100644 index 0000000..8b2f674 --- /dev/null +++ b/src/lib-cuda/reg_flow_slice.h @@ -0,0 +1,169 @@ +#pragma once +#ifndef FF_SRC_LIB_CUDA_REG_FLOW_SLICE +#define FF_SRC_LIB_CUDA_REG_FLOW_SLICE + +/** + * `reg_flow`'s internal seam: the boundary between the exported entry points + * (`reg_flow.cpp`, which instantiates nothing) and the template instantiations + * (`reg_flow__d.cpp`, one translation unit per slice). + * + * WHY THIS EXISTS + * ---------------------------------------------------------------------- + * `reg_flow` is the most expensive translation unit in the project: nvcc peaks + * at 12.98 GB compiling it, 81% of a 16 GB CI runner, and it is one + * indivisible TU (see the measured table above MODULES in the Makefile). That + * single number is what pins the CUDA build at `-j2`, and `-j2` is what makes + * a multi-architecture `-gencode` set unaffordable against the 120-minute + * timeout. + * + * The cost is template instantiation, and it factors cleanly. Each exported + * entry point selects exactly one internal wrapper template and dispatches it + * over ndim x boundary x dtype x offset width. Nothing is shared between one + * entry point's leaves and another's, so a TU that only ever calls + * `_flow_diag<2, ...>` pays for `_flow_diag<2, ...>` and nothing else. Cutting + * the file along those seams cuts the instantiation set the same way. + * + * WHY A FORWARDING SEAM AND NOT `-D` ON ONE SOURCE + * ---------------------------------------------------------------------- + * The obvious cheap trick -- compile `reg_flow.cpp` N times with different + * `-D` flags, the way BOUNDFLAGS already varies a build -- produces N objects + * that each define `ff::cuda::flow_matvec`, and the link fails on duplicate + * symbols. So the exported entry point has to live in exactly one TU, and the + * per-slice TUs have to export something else. + * + * That something else is declared here: one ordinary function per + * (entry point, ndim), named for the slice it covers. `reg_flow.cpp` keeps + * every exported symbol, every argument check and every error message, does + * its `switch (ndim)`, and calls one of these. Each slice TU defines the ones + * it owns, and instantiates only what those need. + * + * The alternative -- `extern template` declarations here with explicit + * instantiation definitions in the slice TUs -- was rejected. It has to name + * every leaf, and the leaves are not knowable from the source: with + * `FF_STATIC_BOUNDS=0` six of the eight `FF_BOUND_*` selectors collapse onto + * `bound::type::Dynamic`, so the enumeration would contain the same + * specialization six times, and explicitly instantiating one specialization + * more than once is ill-formed ([temp.explicit]/5). Which leaves collapse is a + * build-flag decision, so the enumeration would have to be written differently + * per BOUNDFLAGS setting. A `switch` has no such problem: duplicate + * *implicit* instantiations across its arms are simply the same instantiation. + * + * THE SYMBOLS DECLARED HERE ARE NOT ABI + * ---------------------------------------------------------------------- + * They are hidden-visibility, so `libfastfields-cuda.so`'s dynamic symbol + * table is byte-identical to what it was before the split -- verify with + * `nm -D --defined-only`. This header is private to `src/lib-cuda/`; it is not + * installed and nothing outside this directory may include it. + */ + +#include +#include "fastfields/core/dlpack.h" + +// Internal linkage would defeat the purpose (the definition and the call are +// in different TUs), but these must not reach the .so's dynamic symbol table: +// the exported ABI is `ff::cuda::flow_*` and nothing else. Hidden visibility +// is the difference. Non-GNU toolchains fall back to ordinary external +// linkage, which costs a few exported symbols and nothing else -- nothing +// implements the Windows build (see make/common.mk), so this only keeps the +// door open the same way that block does. +#if defined(__GNUC__) || defined(__clang__) +# define FF_FLOW_SLICE_HIDDEN __attribute__((visibility("hidden"))) +#else +# define FF_FLOW_SLICE_HIDDEN +#endif + +/*********************************************************************** + * SLICE SIGNATURES * + ***********************************************************************/ + +// The arguments a slice needs that it cannot cheaply re-derive. Everything +// the entry point validated -- rank, channel count, dtype agreement -- has +// already been checked by the time a slice is called, and `nbatch` is passed +// rather than recomputed so that the rank arithmetic still happens in exactly +// one place. +// +// Spelled as macros so the declaration here and the definition in +// `reg_flow_slice.inl` cannot drift: there is one text for each shape. + +#define FF_FLOW_SLICE_SIG_MATVEC(NAME) \ + void NAME( \ + DLTensor & out , \ + const DLTensor & inp , \ + const double * voxel_size , \ + double absolute , \ + double membrane , \ + double bending , \ + double shears , \ + double div , \ + int8_t bound , \ + int32_t nbatch , \ + bool use_32bits , \ + intptr_t stream ) + +#define FF_FLOW_SLICE_SIG_DIAG(NAME) \ + void NAME( \ + DLTensor & out , \ + const double * voxel_size , \ + double absolute , \ + double membrane , \ + double bending , \ + double shears , \ + double div , \ + int8_t bound , \ + int32_t nbatch , \ + bool use_32bits , \ + intptr_t stream ) + +// Same shape as DIAG; named separately because the two families' arguments +// mean different things (`out` is the operator diagonal versus its Toeplitz +// stencil) and are free to diverge. +#define FF_FLOW_SLICE_SIG_KERNEL(NAME) FF_FLOW_SLICE_SIG_DIAG(NAME) + +#define FF_FLOW_SLICE_SIG_RELAX(NAME) \ + void NAME( \ + DLTensor & sol , \ + const DLTensor & hes , \ + const DLTensor & grd , \ + const double * voxel_size , \ + double absolute , \ + double membrane , \ + double bending , \ + double shears , \ + double div , \ + int8_t bound , \ + int nb_iter , \ + int32_t nbatch , \ + bool use_32bits , \ + intptr_t stream ) + +namespace ff { +namespace cuda { +namespace flow_slice { + +// One declaration per (entry point, ndim). The `add`/`sub` prefixes are the +// `op` template argument ('+' / '-') the corresponding exported entry point +// passes; it is part of the slice's identity rather than a runtime argument +// precisely so that a build may put each op in its own TU. +#define FF_FLOW_SLICE_DECL_ND(ND) \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_MATVEC(matvec_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_MATVEC(addmatvec_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_MATVEC(submatvec_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_DIAG (diag_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_DIAG (adddiag_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_DIAG (subdiag_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_KERNEL(kernel_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_KERNEL(addkernel_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_KERNEL(subkernel_##ND##d); \ + FF_FLOW_SLICE_HIDDEN FF_FLOW_SLICE_SIG_RELAX (relax_##ND##d); + +FF_FLOW_SLICE_DECL_ND(1) +FF_FLOW_SLICE_DECL_ND(2) +FF_FLOW_SLICE_DECL_ND(3) + +#undef FF_FLOW_SLICE_DECL_ND + +} // namespace flow_slice +} // namespace cuda +} // namespace ff + +#endif // FF_SRC_LIB_CUDA_REG_FLOW_SLICE diff --git a/src/lib-cuda/reg_flow_slice.inl b/src/lib-cuda/reg_flow_slice.inl new file mode 100644 index 0000000..09220c3 --- /dev/null +++ b/src/lib-cuda/reg_flow_slice.inl @@ -0,0 +1,681 @@ +/** + * The instantiating half of `reg_flow`: the internal wrapper templates, the + * dtype x offset x boundary dispatch they sit under, and the slice functions + * `reg_flow.cpp` calls. See `reg_flow_slice.h` for why the seam exists. + * + * This file is included exactly once per slice translation unit, which selects + * what it owns by defining a subset of these before the include: + * + * FF_FLOW_SLICE_ND1 / _ND2 / _ND3 which spatial ranks (default: none) + * FF_FLOW_SLICE_MATVEC the matvec family (default: off) + * FF_FLOW_SLICE_DIAG the diag family (default: off) + * FF_FLOW_SLICE_KERNEL the stencil family (default: off) + * FF_FLOW_SLICE_RELAX the relaxation sweeps (default: off) + * FF_FLOW_SLICE_OP_SET / _OP_ADD / _OP_SUB which `op` arms (default: all) + * + * The cross product of what is selected is what this TU instantiates, and a + * template nobody in the TU calls costs nothing. Every slice function declared + * in `reg_flow_slice.h` must be defined by exactly one TU in MODULES -- too + * few and the library link fails on an undefined symbol, too many and it fails + * on a duplicate. Both failures are at link time, which is the point: the + * Makefile's MODULES list and this selection cannot silently disagree. + * + * `.inl` and not `.h` deliberately -- it defines functions and is included + * once, the same arrangement `impl/kernels/threadpool.inl` uses. nvcc's -MMD + * records it as a prerequisite of every slice object, so editing it rebuilds + * all of them. + */ + +#include +#include +#include +#include "fastfields/api/cuda/posdef.h" +#include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" +#include "fastfields/api/cuda/stream.h" +#include "fastfields/core/dlpack.h" +#include "fastfields/core/cuda_switch.h" +#include "fastfields/impl/kernels/bounds.h" +#include "fastfields/impl/kernels/utils.h" +#include "fastfields/impl/cuda/reg_flow.h" +#include "reg_flow_slice.h" + +#ifndef FF_FLOW_SLICE_ND1 +# define FF_FLOW_SLICE_ND1 0 +#endif +#ifndef FF_FLOW_SLICE_ND2 +# define FF_FLOW_SLICE_ND2 0 +#endif +#ifndef FF_FLOW_SLICE_ND3 +# define FF_FLOW_SLICE_ND3 0 +#endif +#ifndef FF_FLOW_SLICE_MATVEC +# define FF_FLOW_SLICE_MATVEC 0 +#endif +#ifndef FF_FLOW_SLICE_DIAG +# define FF_FLOW_SLICE_DIAG 0 +#endif +#ifndef FF_FLOW_SLICE_KERNEL +# define FF_FLOW_SLICE_KERNEL 0 +#endif +#ifndef FF_FLOW_SLICE_RELAX +# define FF_FLOW_SLICE_RELAX 0 +#endif +#ifndef FF_FLOW_SLICE_OP_SET +# define FF_FLOW_SLICE_OP_SET 1 +#endif +#ifndef FF_FLOW_SLICE_OP_ADD +# define FF_FLOW_SLICE_OP_ADD 1 +#endif +#ifndef FF_FLOW_SLICE_OP_SUB +# define FF_FLOW_SLICE_OP_SUB 1 +#endif + +#if !(FF_FLOW_SLICE_ND1 || FF_FLOW_SLICE_ND2 || FF_FLOW_SLICE_ND3) +# error "a reg_flow slice must select at least one of FF_FLOW_SLICE_ND1/2/3" +#endif +#if !(FF_FLOW_SLICE_MATVEC || FF_FLOW_SLICE_DIAG || \ + FF_FLOW_SLICE_KERNEL || FF_FLOW_SLICE_RELAX) +# error "a reg_flow slice must select at least one operation family" +#endif + +FF_NAMESPACE_BEGIN(FF_NS) +FF_NAMESPACE_BEGIN(FF_DEVICE) + +// reduction / accumulation type (matches jitfields' float64 default) +typedef double reduce_t; + +/*********************************************************************** + * WRAPPERS * + ***********************************************************************/ + +namespace { + +// length of the shape/stride arrays: (*batch, *spatial, C) == out.ndim +template +inline void _flow_matvec( + const bound::BoundVec & bvec, + int64_t nbatch , + void * out , + const void * inp , + const double * voxel_size , + double absolute , + double membrane , + double bending , + double shears , + double div , + const int64_t * size , + const int64_t * stride_out , + const int64_t * stride_inp , + cudaStream_t stream ) +{ + const int64_t nall1 = nbatch + ndim + 1; + const offset_t * _size = copy_if_needed(size, nall1); + const offset_t * _stride_out = copy_if_needed(stride_out, nall1); + const offset_t * _stride_inp = copy_if_needed(stride_inp, nall1); + scalar_t * _out = static_cast< scalar_t *>(out); + const scalar_t * _inp = static_cast(inp); + + reduce_t vx[ndim]; + for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; + + // The linear-elastic (Lamé) terms `shears`/`div` couple the flow channels, + // so any non-zero one selects the full combined stencil (matvec_all, which + // also folds in absolute/membrane/bending). Otherwise fall back to the + // cheaper single-penalty stencils (highest-order non-zero wins). + if (shears != 0.0 || div != 0.0) + reg_flow::matvec_all( + bvec, static_cast(nbatch), _out, _inp, + _size, _stride_out, _stride_inp, vx, absolute, membrane, bending, stream); + else if (membrane != 0.0) + reg_flow::matvec_membrane( + bvec, static_cast(nbatch), _out, _inp, + _size, _stride_out, _stride_inp, vx, absolute, stream); + + free_if_needed(_size); + free_if_needed(_stride_out); + free_if_needed(_stride_inp); +} + +// Accumulate variant of _flow_matvec: out += L(inp) (op='+') or out -= L(inp) +// (op='-'), instead of overwriting out. Mirrors the CPU `_flow_matvec_acc`. +template +inline void _flow_matvec_acc( + const bound::BoundVec & bvec, + int64_t nbatch , + void * out , + const void * inp , + const double * voxel_size , + double absolute , + double membrane , + double bending , + double shears , + double div , + const int64_t * size , + const int64_t * stride_out , + const int64_t * stride_inp , + cudaStream_t stream ) +{ + const int64_t nall1 = nbatch + ndim + 1; + const offset_t * _size = copy_if_needed(size, nall1); + const offset_t * _stride_out = copy_if_needed(stride_out, nall1); + const offset_t * _stride_inp = copy_if_needed(stride_inp, nall1); + scalar_t * _out = static_cast< scalar_t *>(out); + const scalar_t * _inp = static_cast(inp); + + reduce_t vx[ndim]; + for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; + + if (shears != 0.0 || div != 0.0) + reg_flow::matvec_all( + bvec, static_cast(nbatch), _out, _inp, + _size, _stride_out, _stride_inp, vx, + absolute, membrane, bending, shears, div, stream); + else if (bending != 0.0) + reg_flow::matvec_bending( + bvec, static_cast(nbatch), _out, _inp, + _size, _stride_out, _stride_inp, vx, absolute, membrane, bending, stream); + else if (membrane != 0.0) + reg_flow::matvec_membrane( + bvec, static_cast(nbatch), _out, _inp, + _size, _stride_out, _stride_inp, vx, absolute, membrane, stream); + else + reg_flow::matvec_absolute( + bvec, static_cast(nbatch), _out, _inp, + _size, _stride_out, _stride_inp, vx, absolute, stream); + + free_if_needed(_size); + free_if_needed(_stride_out); + free_if_needed(_stride_inp); +} + +template +inline void _flow_diag( + const bound::BoundVec & bvec, + int64_t nbatch , + void * out , + const double * voxel_size , + double absolute , + double membrane , + double bending , + double shears , + double div , + const int64_t * size , + const int64_t * stride_out , + cudaStream_t stream ) +{ + const int64_t nall1 = nbatch + ndim + 1; + const offset_t * _size = copy_if_needed(size, nall1); + const offset_t * _stride_out = copy_if_needed(stride_out, nall1); + scalar_t * _out = static_cast(out); + + reduce_t vx[ndim]; + for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; + + if (shears != 0.0 || div != 0.0) + reg_flow::diag_all( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, membrane, bending, shears, div, stream); + else if (bending != 0.0) + reg_flow::diag_bending( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, membrane, bending, stream); + else if (membrane != 0.0) + reg_flow::diag_membrane( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, membrane, stream); + else + reg_flow::diag_absolute( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, stream); + + free_if_needed(_size); + free_if_needed(_stride_out); +} + +// Materialise the Toeplitz convolution kernel (stencil) of the operator (see +// cpu-lib). `nfull` is the length of the size/stride arrays (== out.ndim): +// nbatch+ndim+1 for the per-channel vector stencil, nbatch+ndim+2 for the Lamé +// (cross-channel) matrix stencil. +template +inline void _flow_kernel( + const bound::BoundVec & bvec, + int64_t nbatch , + void * out , + const double * voxel_size , + double absolute , + double membrane , + double bending , + double shears , + double div , + const int64_t * size , + const int64_t * stride_out , + int64_t nfull , + cudaStream_t stream ) +{ + const offset_t * _size = copy_if_needed(size, nfull); + const offset_t * _stride_out = copy_if_needed(stride_out, nfull); + scalar_t * _out = static_cast(out); + + reduce_t vx[ndim]; + for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; + + if (shears != 0.0 || div != 0.0) { + if (bending != 0.0) + reg_flow::kernel_all( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, membrane, bending, shears, div, stream); + else + reg_flow::kernel_lame( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, membrane, shears, div, stream); + } else if (bending != 0.0) + reg_flow::kernel_bending( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, membrane, bending, stream); + else if (membrane != 0.0) + reg_flow::kernel_membrane( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, membrane, stream); + else + reg_flow::kernel_absolute( + bvec, static_cast(nbatch), _out, + _size, _stride_out, vx, absolute, stream); + + free_if_needed(_size); + free_if_needed(_stride_out); +} + +// In-place relaxation sweeps solving `(H + L) x = g` (see cpu-lib). +template +inline void _flow_relax( + const bound::BoundVec & bvec, + int64_t nbatch , + void * sol , + const void * hes , + const void * grd , + const double * voxel_size , + double absolute , + double membrane , + double bending , + double shears , + double div , + int niter , + const int64_t * size , + const int64_t * stride_sol , + const int64_t * stride_hes , + const int64_t * stride_grd , + cudaStream_t stream ) +{ + const int64_t nall1 = nbatch + ndim + 1; + const offset_t * _size = copy_if_needed(size, nall1); + const offset_t * _stride_sol = copy_if_needed(stride_sol, nall1); + const offset_t * _stride_hes = copy_if_needed(stride_hes, nall1); + const offset_t * _stride_grd = copy_if_needed(stride_grd, nall1); + scalar_t * _sol = static_cast< scalar_t *>(sol); + const scalar_t * _hes = static_cast(hes); + const scalar_t * _grd = static_cast(grd); + + reduce_t vx[ndim]; + for (int d = 0; d < ndim; ++d) vx[d] = voxel_size ? voxel_size[d] : 1.0; + + if (shears != 0.0 || div != 0.0) { + if (bending != 0.0) + reg_flow::relax_all_( + bvec, static_cast(nbatch), _sol, _hes, _grd, + _size, _stride_sol, _stride_hes, _stride_grd, vx, + absolute, membrane, bending, shears, div, niter, stream); + else + reg_flow::relax_lame_( + bvec, static_cast(nbatch), _sol, _hes, _grd, + _size, _stride_sol, _stride_hes, _stride_grd, vx, + absolute, membrane, shears, div, niter, stream); + } else if (bending != 0.0) + reg_flow::relax_bending_( + bvec, static_cast(nbatch), _sol, _hes, _grd, + _size, _stride_sol, _stride_hes, _stride_grd, vx, + absolute, membrane, bending, niter, stream); + else + reg_flow::relax_membrane_( + bvec, static_cast(nbatch), _sol, _hes, _grd, + _size, _stride_sol, _stride_hes, _stride_grd, vx, + absolute, membrane, niter, stream); + + free_if_needed(_size); + free_if_needed(_stride_sol); + free_if_needed(_stride_hes); + free_if_needed(_stride_grd); +} + +} // anonymous namespace + +/*********************************************************************** + * DISPATCH * + ***********************************************************************/ + +// Unchanged from the single-TU form, except that the ndim arm is now the +// slice's own compile-time constant instead of a `switch (ndim)`: the switch +// moved up into reg_flow.cpp, where it picks a slice rather than a template +// argument. Everything below one arm of that switch is what a slice TU is. + +#define FF_FLOW_BND1(B) B +#define FF_FLOW_BND2(B) B, B +#define FF_FLOW_BND3(B) B, B, B + +// matvec dtype x offset dispatch, given ndim and the (repeated) bound pack. +#define FF_FLOW_MV_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_matvec(FF_FLOW_MV_ARGS) \ + : _flow_matvec(FF_FLOW_MV_ARGS); \ + case 64: return use_32bits \ + ? _flow_matvec(FF_FLOW_MV_ARGS) \ + : _flow_matvec(FF_FLOW_MV_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_ADD_MV_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_matvec_acc(FF_FLOW_MV_ARGS) \ + : _flow_matvec_acc(FF_FLOW_MV_ARGS); \ + case 64: return use_32bits \ + ? _flow_matvec_acc(FF_FLOW_MV_ARGS) \ + : _flow_matvec_acc(FF_FLOW_MV_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_SUB_MV_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_matvec_acc(FF_FLOW_MV_ARGS) \ + : _flow_matvec_acc(FF_FLOW_MV_ARGS); \ + case 64: return use_32bits \ + ? _flow_matvec_acc(FF_FLOW_MV_ARGS) \ + : _flow_matvec_acc(FF_FLOW_MV_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_DG_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_diag(FF_FLOW_DG_ARGS); \ + case 64: return use_32bits \ + ? _flow_diag(FF_FLOW_DG_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_ADD_DG_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_diag(FF_FLOW_DG_ARGS) \ + : _flow_diag(FF_FLOW_DG_ARGS); \ + case 64: return use_32bits \ + ? _flow_diag(FF_FLOW_DG_ARGS) \ + : _flow_diag(FF_FLOW_DG_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_SUB_DG_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_diag(FF_FLOW_DG_ARGS) \ + : _flow_diag(FF_FLOW_DG_ARGS); \ + case 64: return use_32bits \ + ? _flow_diag(FF_FLOW_DG_ARGS) \ + : _flow_diag(FF_FLOW_DG_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_KN_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_kernel(FF_FLOW_KN_ARGS); \ + case 64: return use_32bits \ + ? _flow_kernel(FF_FLOW_KN_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_ADD_KN_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_kernel(FF_FLOW_KN_ARGS) \ + : _flow_kernel(FF_FLOW_KN_ARGS); \ + case 64: return use_32bits \ + ? _flow_kernel(FF_FLOW_KN_ARGS) \ + : _flow_kernel(FF_FLOW_KN_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_SUB_KN_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_kernel(FF_FLOW_KN_ARGS) \ + : _flow_kernel(FF_FLOW_KN_ARGS); \ + case 64: return use_32bits \ + ? _flow_kernel(FF_FLOW_KN_ARGS) \ + : _flow_kernel(FF_FLOW_KN_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +#define FF_FLOW_RX_DT(NDIM, BNDS...) \ + switch (code) { \ + case kDLFloat: switch (bits) { \ + case 32: return use_32bits \ + ? _flow_relax(FF_FLOW_RX_ARGS) \ + : _flow_relax(FF_FLOW_RX_ARGS); \ + case 64: return use_32bits \ + ? _flow_relax(FF_FLOW_RX_ARGS) \ + : _flow_relax(FF_FLOW_RX_ARGS); \ + default: break; \ + } break; \ + default: break; \ + } \ + throw std::invalid_argument("only floating point data types are supported"); + +// Which of these boundary conditions gets a dedicated (static) instantiation +// and which shares the single Dynamic (runtime) one is a build-time choice -- +// see FF_STATIC_BOUND_* in kernels/bounds.h. `bvec` carries the runtime +// condition for whichever ones fall back to Dynamic. +#define FF_FLOW_BOUND_SWITCH(DT, NDIM, BND) \ + switch (bnd) { \ + case bound::type::Zero: DT(NDIM, BND(FF_BOUND_ZERO)); break; \ + case bound::type::Replicate: DT(NDIM, BND(FF_BOUND_REPLICATE)); break; \ + case bound::type::DCT1: DT(NDIM, BND(FF_BOUND_DCT1)); break; \ + case bound::type::DCT2: DT(NDIM, BND(FF_BOUND_DCT2)); break; \ + case bound::type::DST1: DT(NDIM, BND(FF_BOUND_DST1)); break; \ + case bound::type::DST2: DT(NDIM, BND(FF_BOUND_DST2)); break; \ + case bound::type::DFT: DT(NDIM, BND(FF_BOUND_DFT)); break; \ + case bound::type::NoCheck: DT(NDIM, BND(FF_BOUND_NOCHECK)); break; \ + default: throw std::invalid_argument("Unsupported boundary condition"); \ + } + +/*********************************************************************** + * SLICE DEFINITIONS * + ***********************************************************************/ + +// The prologues below are what each exported entry point used to compute +// inline, minus the argument checks (those stayed in reg_flow.cpp, where the +// error messages and their order are unchanged) and minus `nbatch` (passed in, +// so the rank arithmetic still happens exactly once). + +#define FF_FLOW_MV_ARGS \ + bvec, static_cast(nbatch), FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + voxel_size, absolute, membrane, bending, shears, div, \ + out.shape, out.strides, inp.strides, cstream + +#define FF_FLOW_DG_ARGS \ + bvec, static_cast(nbatch), FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, shears, div, \ + out.shape, out.strides, cstream + +#define FF_FLOW_KN_ARGS \ + bvec, static_cast(nbatch), FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, shears, div, \ + out.shape, out.strides, static_cast(out.ndim), cstream + +#define FF_FLOW_RX_ARGS \ + bvec, static_cast(nbatch), FF_VOIDPTR(sol), FF_CVOIDPTR(hes), \ + FF_CVOIDPTR(grd), voxel_size, absolute, membrane, bending, \ + shears, div, nb_iter, sol.shape, sol.strides, hes.strides, \ + grd.strides, cstream + +#define FF_FLOW_PROLOGUE(T) \ + const auto code = static_cast(T.dtype.code); \ + const auto bits = T.dtype.bits; \ + const bound::type bnd = static_cast(bound); \ + const bound::BoundVec bvec(bnd); \ + const cudaStream_t cstream = _reg_stream(stream); + +FF_NAMESPACE_BEGIN(flow_slice) + +// Each of these is one arm of the old `NDIM_SWITCH`, given its own name so it +// can be given its own translation unit. +#define FF_FLOW_DEFINE_SLICE(NAME, SIG, PROLOGUE_T, DT, ND) \ + FF_FLOW_SLICE_HIDDEN SIG(NAME) \ + { \ + FF_FLOW_PROLOGUE(PROLOGUE_T) \ + FF_FLOW_BOUND_SWITCH(DT, ND, FF_FLOW_BND##ND) \ + } + +#if FF_FLOW_SLICE_MATVEC && FF_FLOW_SLICE_OP_SET +# define FF_FLOW_EMIT_MATVEC_SET(ND) \ + FF_FLOW_DEFINE_SLICE(matvec_##ND##d, FF_FLOW_SLICE_SIG_MATVEC, out, FF_FLOW_MV_DT, ND) +#else +# define FF_FLOW_EMIT_MATVEC_SET(ND) +#endif +#if FF_FLOW_SLICE_MATVEC && FF_FLOW_SLICE_OP_ADD +# define FF_FLOW_EMIT_MATVEC_ADD(ND) \ + FF_FLOW_DEFINE_SLICE(addmatvec_##ND##d, FF_FLOW_SLICE_SIG_MATVEC, out, FF_FLOW_ADD_MV_DT, ND) +#else +# define FF_FLOW_EMIT_MATVEC_ADD(ND) +#endif +#if FF_FLOW_SLICE_MATVEC && FF_FLOW_SLICE_OP_SUB +# define FF_FLOW_EMIT_MATVEC_SUB(ND) \ + FF_FLOW_DEFINE_SLICE(submatvec_##ND##d, FF_FLOW_SLICE_SIG_MATVEC, out, FF_FLOW_SUB_MV_DT, ND) +#else +# define FF_FLOW_EMIT_MATVEC_SUB(ND) +#endif + +#if FF_FLOW_SLICE_DIAG && FF_FLOW_SLICE_OP_SET +# define FF_FLOW_EMIT_DIAG_SET(ND) \ + FF_FLOW_DEFINE_SLICE(diag_##ND##d, FF_FLOW_SLICE_SIG_DIAG, out, FF_FLOW_DG_DT, ND) +#else +# define FF_FLOW_EMIT_DIAG_SET(ND) +#endif +#if FF_FLOW_SLICE_DIAG && FF_FLOW_SLICE_OP_ADD +# define FF_FLOW_EMIT_DIAG_ADD(ND) \ + FF_FLOW_DEFINE_SLICE(adddiag_##ND##d, FF_FLOW_SLICE_SIG_DIAG, out, FF_FLOW_ADD_DG_DT, ND) +#else +# define FF_FLOW_EMIT_DIAG_ADD(ND) +#endif +#if FF_FLOW_SLICE_DIAG && FF_FLOW_SLICE_OP_SUB +# define FF_FLOW_EMIT_DIAG_SUB(ND) \ + FF_FLOW_DEFINE_SLICE(subdiag_##ND##d, FF_FLOW_SLICE_SIG_DIAG, out, FF_FLOW_SUB_DG_DT, ND) +#else +# define FF_FLOW_EMIT_DIAG_SUB(ND) +#endif + +#if FF_FLOW_SLICE_KERNEL && FF_FLOW_SLICE_OP_SET +# define FF_FLOW_EMIT_KERNEL_SET(ND) \ + FF_FLOW_DEFINE_SLICE(kernel_##ND##d, FF_FLOW_SLICE_SIG_KERNEL, out, FF_FLOW_KN_DT, ND) +#else +# define FF_FLOW_EMIT_KERNEL_SET(ND) +#endif +#if FF_FLOW_SLICE_KERNEL && FF_FLOW_SLICE_OP_ADD +# define FF_FLOW_EMIT_KERNEL_ADD(ND) \ + FF_FLOW_DEFINE_SLICE(addkernel_##ND##d, FF_FLOW_SLICE_SIG_KERNEL, out, FF_FLOW_ADD_KN_DT, ND) +#else +# define FF_FLOW_EMIT_KERNEL_ADD(ND) +#endif +#if FF_FLOW_SLICE_KERNEL && FF_FLOW_SLICE_OP_SUB +# define FF_FLOW_EMIT_KERNEL_SUB(ND) \ + FF_FLOW_DEFINE_SLICE(subkernel_##ND##d, FF_FLOW_SLICE_SIG_KERNEL, out, FF_FLOW_SUB_KN_DT, ND) +#else +# define FF_FLOW_EMIT_KERNEL_SUB(ND) +#endif + +// `relax` has no op axis -- it solves in place, there is nothing to accumulate +// into -- so FF_FLOW_SLICE_OP_* does not apply to it. +#if FF_FLOW_SLICE_RELAX +# define FF_FLOW_EMIT_RELAX(ND) \ + FF_FLOW_DEFINE_SLICE(relax_##ND##d, FF_FLOW_SLICE_SIG_RELAX, sol, FF_FLOW_RX_DT, ND) +#else +# define FF_FLOW_EMIT_RELAX(ND) +#endif + +#define FF_FLOW_EMIT_ND(ND) \ + FF_FLOW_EMIT_MATVEC_SET(ND) \ + FF_FLOW_EMIT_MATVEC_ADD(ND) \ + FF_FLOW_EMIT_MATVEC_SUB(ND) \ + FF_FLOW_EMIT_DIAG_SET(ND) \ + FF_FLOW_EMIT_DIAG_ADD(ND) \ + FF_FLOW_EMIT_DIAG_SUB(ND) \ + FF_FLOW_EMIT_KERNEL_SET(ND) \ + FF_FLOW_EMIT_KERNEL_ADD(ND) \ + FF_FLOW_EMIT_KERNEL_SUB(ND) \ + FF_FLOW_EMIT_RELAX(ND) + +#if FF_FLOW_SLICE_ND1 +FF_FLOW_EMIT_ND(1) +#endif +#if FF_FLOW_SLICE_ND2 +FF_FLOW_EMIT_ND(2) +#endif +#if FF_FLOW_SLICE_ND3 +FF_FLOW_EMIT_ND(3) +#endif + +FF_NAMESPACE_END(flow_slice) + +FF_NAMESPACE_END(FF_DEVICE) +FF_NAMESPACE_END(FF_NS) From 279e21d8e6e4f866d5dab1c2612ce331c7debc93 Mon Sep 17 00:00:00 2001 From: Yael Date: Thu, 20 Aug 2026 11:50:02 +0000 Subject: [PATCH 2/7] refactor(cuda): adopt #145's and #146's conventions in the new reg_flow files Pre-emptive, so that neither sweep has to touch these files and neither conflicts with them: * #145 -- one include-guard convention, no `#pragma once`. reg_flow_slice.h dropped the pragma and its guard is spelled the way that PR derives new ones from the path: FF_SRC_LIB_CUDA_REG_FLOW_SLICE_H. * #146 -- the public interface is spelled . Same-directory siblings stay quoted, which is what that PR does to impl/cuda/reg_flow.h's `#include "utils.h"`. Include spelling only; the preprocessed output is unchanged. Front TU still compiles and every slice still preprocesses. --- src/lib-cuda/reg_flow.cpp | 14 +++++++------- src/lib-cuda/reg_flow_slice.h | 9 ++++----- src/lib-cuda/reg_flow_slice.inl | 18 +++++++++--------- 3 files changed, 20 insertions(+), 21 deletions(-) diff --git a/src/lib-cuda/reg_flow.cpp b/src/lib-cuda/reg_flow.cpp index a00428b..357ee49 100644 --- a/src/lib-cuda/reg_flow.cpp +++ b/src/lib-cuda/reg_flow.cpp @@ -16,18 +16,18 @@ #include #include #include -#include "fastfields/api/cuda/reg_flow.h" -#include "fastfields/api/cuda/posdef.h" -#include "fastfields/core/dispatch.h" -#include "fastfields/core/dlpack.h" -#include "fastfields/core/cuda_switch.h" +#include +#include +#include +#include +#include // `FF_CANUSE32BITS` expands to an unqualified `canUse32BitIndexMath`, which is // declared in impl/kernels/utils.h -- NOT in core/autocast.h, whose `// // canUse32BitIndexMath` include comment in core/dispatch.h suggests otherwise. // Every other dispatch source pulls the kernels in wholesale and never noticed; // this TU is the first that does not, so it has to name the real home. -#include "fastfields/impl/kernels/utils.h" // canUse32BitIndexMath -#include "fastfields/impl/cuda/utils.h" // allocDevice / freeDevice +#include // canUse32BitIndexMath +#include // allocDevice / freeDevice #include "reg_flow_slice.h" FF_NAMESPACE_BEGIN(FF_NS) diff --git a/src/lib-cuda/reg_flow_slice.h b/src/lib-cuda/reg_flow_slice.h index 8b2f674..f7a8c97 100644 --- a/src/lib-cuda/reg_flow_slice.h +++ b/src/lib-cuda/reg_flow_slice.h @@ -1,6 +1,5 @@ -#pragma once -#ifndef FF_SRC_LIB_CUDA_REG_FLOW_SLICE -#define FF_SRC_LIB_CUDA_REG_FLOW_SLICE +#ifndef FF_SRC_LIB_CUDA_REG_FLOW_SLICE_H +#define FF_SRC_LIB_CUDA_REG_FLOW_SLICE_H /** * `reg_flow`'s internal seam: the boundary between the exported entry points @@ -57,7 +56,7 @@ */ #include -#include "fastfields/core/dlpack.h" +#include // Internal linkage would defeat the purpose (the definition and the call are // in different TUs), but these must not reach the .so's dynamic symbol table: @@ -166,4 +165,4 @@ FF_FLOW_SLICE_DECL_ND(3) } // namespace cuda } // namespace ff -#endif // FF_SRC_LIB_CUDA_REG_FLOW_SLICE +#endif // FF_SRC_LIB_CUDA_REG_FLOW_SLICE_H diff --git a/src/lib-cuda/reg_flow_slice.inl b/src/lib-cuda/reg_flow_slice.inl index 09220c3..e8de015 100644 --- a/src/lib-cuda/reg_flow_slice.inl +++ b/src/lib-cuda/reg_flow_slice.inl @@ -29,15 +29,15 @@ #include #include #include -#include "fastfields/api/cuda/posdef.h" -#include "fastfields/core/autocast.h" -#include "fastfields/core/dispatch.h" -#include "fastfields/api/cuda/stream.h" -#include "fastfields/core/dlpack.h" -#include "fastfields/core/cuda_switch.h" -#include "fastfields/impl/kernels/bounds.h" -#include "fastfields/impl/kernels/utils.h" -#include "fastfields/impl/cuda/reg_flow.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include #include "reg_flow_slice.h" #ifndef FF_FLOW_SLICE_ND1 From 93dfada8682a6273945ac241f36d8b9e9ce362db Mon Sep 17 00:00:00 2001 From: Yael Date: Thu, 20 Aug 2026 12:20:57 +0000 Subject: [PATCH 3/7] docs(cuda): record the measured reg_flow split, and what it does not buy Numbers from `build-cuda` on this branch (run 32365351378) against the same job on 85fdac7 (run 32325044770) -- same runner image, same nvcc 12.0.140, same -O1, same bound/spline policy. heaviest TU 12.98 -> 1.95 GiB (-85.0%) module CPU 1149.89 -> 751.95 s (-34.6%) wall -j2 2243 -> 1672 s (-25.5%, whole `make cuda`) .so link 1.38 -> 1.05 s The header-parse counter-force is recorded because it is the reason this could have gone the other way: a TU that parses the whole template header set and instantiates nothing costs 0.22 GiB and 1.35 s, so twelve extra TUs is ~16 s against a 1149.89 s baseline. Measured, not assumed -- and stated as a ratio so the next module can be judged on its own numbers rather than on this one's result. The section is blunt about the limit: this does NOT raise the -j ceiling. max(peak) is now reg_field at 8.23 GiB and the reg_field/reg_field_rls overlap is unchanged at 15.31 GiB, so -j2 stays the ceiling and stays luck. A scheduling model over the measured table reproduces the observed -j2 wall to 0.25%, and says -j4 would still peak at 19.26 GiB on a 16 GB box. --- src/lib-cuda/Makefile | 55 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 53 insertions(+), 2 deletions(-) diff --git a/src/lib-cuda/Makefile b/src/lib-cuda/Makefile index 9901bf2..5b84987 100644 --- a/src/lib-cuda/Makefile +++ b/src/lib-cuda/Makefile @@ -139,7 +139,7 @@ SPLINEFLAGS ?= -DFF_STATIC_SPLINES=0 \ # four cheapest modules in the file (resize 2.00, restrict 1.30, splinc 0.42, # posdef 0.37 GB), together well under a single regulariser. They cost ~7 min # of the ~31 min compile. -# ~~~ reg_flow IS SPLIT ACROSS TRANSLATION UNITS ~~~ +# ~~~ reg_flow IS SPLIT ACROSS TRANSLATION UNITS (fastfields-lib#147) ~~~ # # `reg_flow` was the 12.98 GiB entry in the table above -- 81% of a 16 GB # runner, in one indivisible TU, and the single number that pinned this build @@ -153,7 +153,58 @@ SPLINEFLAGS ?= -DFF_STATIC_SPLINES=0 \ # for its own arm and nothing else. The exported ABI is unchanged -- the slice # symbols are hidden-visibility and never reach .dynsym. # -# TBD-MEASURED-TABLE +# ~~~ MEASURED, 2026-08-20, build-cuda on this branch vs on 85fdac7 ~~~ +# Same job, same runner image, same nvcc 12.0.140, same -O1 and the +# BOUNDFLAGS/SPLINEFLAGS defaults above. Before, reg_flow was one TU: +# +# reg_flow 12.98 GiB 1149.89 s <-- one indivisible TU +# +# After, it is thirteen (peak RSS / nvcc seconds / role): +# +# reg_flow_matvec_3d 1.95 GiB 174.90 s +# reg_flow_matvec_2d 1.22 GiB 91.61 s +# reg_flow_diag_3d 1.03 GiB 59.12 s +# reg_flow_relax_3d 0.96 GiB 87.40 s +# reg_flow_kernel_3d 0.89 GiB 78.43 s +# reg_flow_diag_2d 0.67 GiB 41.22 s +# reg_flow_matvec_1d 0.61 GiB 43.46 s +# reg_flow_relax_2d 0.60 GiB 45.85 s +# reg_flow_kernel_2d 0.59 GiB 44.29 s +# reg_flow_kernel_1d 0.48 GiB 33.84 s +# reg_flow_diag_1d 0.44 GiB 30.52 s +# reg_flow_relax_1d 0.26 GiB 19.80 s +# reg_flow 0.20 GiB 1.51 s <-- the front TU +# ------- -------- +# heaviest slice 1.95 GiB 751.95 s total +# +# peak 12.98 -> 1.95 GiB (-85.0%; the heaviest slice) +# CPU 1149.89 -> 751.95 s (-34.6%, for this module) +# wall -j2 2243 -> 1672 s (-25.5%, whole `make cuda`) +# link 1.38 -> 1.05 s (the .so link does not care) +# +# THE COUNTER-FORCE IS REAL BUT SMALL, and it was measured rather than +# assumed. Every slice re-parses the whole template header set, so total CPU +# could have gone UP. A TU that parses `impl/cuda/reg_flow.h` plus the kernels +# and instantiates nothing costs 0.22 GiB and 1.35 s. Twelve extra TUs is +# therefore ~16 s of re-parsing against a 1149.89 s baseline -- 1.4%, an order +# of magnitude less than the instantiation work that now divides. That ratio +# is what makes this split pay; it is not a general law, and a module whose +# per-TU instantiation work is comparable to 1.35 s should not be split. +# +# WHAT THIS DOES *NOT* BUY: A HIGHER -j. +# The ceiling is max(peak) over EVERY module, and reg_field (8.23 GiB) and +# reg_field_rls (7.07 GiB) are untouched. reg_flow is simply no longer the +# constraint -- reg_field is, and their overlap (15.31 GiB) is the same -j2 +# hazard the section above describes. Modelling make's dispatch over the +# measured table (the model reproduces the observed -j2 wall to 0.25%): +# +# -j2 1668 s worst coincidence 15.31 GiB reg_field+reg_field_rls +# -j4 869 s worst coincidence 19.26 GiB <-- still over a 16 GB box +# +# So do NOT raise -j on the strength of this change. Splitting reg_field and +# reg_field_rls the same way -- they are structurally identical files, see +# reg_flow_slice.h -- is what would take max(peak) to ~2 GiB and make -j4 +# CPU-bound instead of RAM-bound. # # ORDER MATTERS HERE. make dispatches in MODULES order, so the slices are # listed heaviest-first (3D, then 2D, then 1D): with the long poles started From 7b20ab9a5752bfc925e8f2a519caf8d7c430436d Mon Sep 17 00:00:00 2001 From: Yael Date: Thu, 20 Aug 2026 16:27:09 +0000 Subject: [PATCH 4/7] docs(cuda): measure the index axis inside a split slice, not just before #151 recorded what FF_INDEX32 costs on the unsplit modules. This is the same question asked of a slice, because the interesting case is whether splitting makes the axis cheap enough that the knob stops mattering. It does not. Same TU, index axis toggled: reg_flow_matvec_3d FF_INDEX32=1 1.94 GiB 277.1 s 17.0 MB FF_INDEX32=0 1.05 GiB 174.4 s 10.0 MB -45.8% peak, which is the same proportion #94 measured on the unsplit file with the host compiler (-44.6%). The split and the knob multiply rather than overlap: 12.98 -> 1.94 GiB from one, 1.94 -> 1.05 GiB from the other. --- src/lib-cuda/Makefile | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/src/lib-cuda/Makefile b/src/lib-cuda/Makefile index ef13e14..67b2ee9 100644 --- a/src/lib-cuda/Makefile +++ b/src/lib-cuda/Makefile @@ -249,10 +249,20 @@ INDEXFLAGS ?= -DFF_INDEX32=1 # reg_field_rls the same way (they are structurally identical files) is what # would take max(peak) to ~2 GiB and make -j4 CPU-bound instead of RAM-bound. # -# INTERACTION WITH INDEXFLAGS: they multiply, they do not overlap. The index -# axis is x2 on every leaf, so it is x2 on every slice; FF_INDEX32=0 would take -# the heaviest slice from 1.95 GiB to roughly 1 GiB. What the split changes is -# which resource the knob buys back: memory before, wall-clock now. +# INTERACTION WITH INDEXFLAGS: they MULTIPLY, they do not overlap. Measured on +# the heaviest slice, same TU, index axis toggled: +# +# reg_flow_matvec_3d FF_INDEX32=1 1.94 GiB 277.1 s 17.0 MB +# FF_INDEX32=0 1.05 GiB 174.4 s 10.0 MB +# -45.8% -37.1% -41.1% +# +# That -45.8% is the same proportion #94 measured on the UNSPLIT file with the +# host compiler (-44.6%), which is the point: splitting does not make the index +# axis any cheaper in relative terms, it only shrinks the absolute number the +# proportion applies to. The two compose -- 12.98 -> 1.94 GiB from the split, +# 1.94 -> 1.05 GiB from the knob. So this split is NOT a reason to stop caring +# about FF_INDEX32; what it changes is which resource the knob buys back, since +# memory is no longer what binds this module and wall-clock is. # # ORDER MATTERS HERE. make dispatches in MODULES order, so the slices are # listed heaviest-first (3D, then 2D, then 1D): with the long poles started From 33667799015fe9aa801dfb5187a6e56b142499ce Mon Sep 17 00:00:00 2001 From: Yael Date: Thu, 20 Aug 2026 16:47:27 +0000 Subject: [PATCH 5/7] docs(cuda): the split and FF_INDEX32=0 together are what unlock -j4 build-cuda's index64 leg builds this split with the index axis off, and that combination -- not either half alone -- is the first configuration in which no module in this directory is heavy any more. Measured on that leg: reg_field 3.76 GiB 264.90 s <-- the new maximum reg_field_rls 3.42 GiB 279.38 s resize 1.08 GiB 90.20 s reg_flow_matvec_3d 1.05 GiB 81.45 s <-- heaviest reg_flow slice reg_flow (front) 0.20 GiB 0.88 s Total nvcc CPU 1380.3 s; `make cuda -j2` wall 723 s, measured. Feeding that table through the scheduling model (which reproduces the measured -j2 wall to 0.5%) puts -j4 at a 393 s makespan with a worst coincidence of 9.31 GiB -- comfortably inside a 16 GB runner, with reg_field untouched. Recorded with the caveat that matters: this holds on the index64 leg, not on index32, which is the default and the one that ships. There the same model still says 19.26 GiB at -j4. The note says so explicitly so nobody reads the index64 number as licence to raise -j globally. --- src/lib-cuda/Makefile | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/lib-cuda/Makefile b/src/lib-cuda/Makefile index 67b2ee9..8968ce1 100644 --- a/src/lib-cuda/Makefile +++ b/src/lib-cuda/Makefile @@ -264,6 +264,28 @@ INDEXFLAGS ?= -DFF_INDEX32=1 # about FF_INDEX32; what it changes is which resource the knob buys back, since # memory is no longer what binds this module and wall-clock is. # +# AND THE TWO TOGETHER ARE WHAT WOULD ACTUALLY UNLOCK -j4. build-cuda's +# index64 leg builds this same split with the axis off, and that combination -- +# not either half alone -- is the first configuration in which no module is +# heavy any more: +# +# reg_field 3.76 GiB 264.90 s <-- the new maximum +# reg_field_rls 3.42 GiB 279.38 s +# resize 1.08 GiB 90.20 s +# reg_flow_matvec_3d 1.05 GiB 81.45 s <-- heaviest reg_flow slice +# ... every other module below 1.01 GiB ... +# reg_flow (front) 0.20 GiB 0.88 s +# +# total nvcc CPU 1380.3 s; `make cuda -j2` wall 723 s (measured). +# +# Feeding that table through the same scheduling model (which reproduces the +# measured -j2 wall to 0.5%): -j4 gives a 393 s makespan with a worst +# coincidence of 9.31 GiB, comfortably inside a 16 GB runner. So `-j4` is safe +# on the index64 leg TODAY, with reg_field untouched. It is NOT safe on the +# index32 leg, which is the default and the one that ships -- there the same +# model says 19.26 GiB. Do not raise -j globally on the strength of the +# index64 number. +# # ORDER MATTERS HERE. make dispatches in MODULES order, so the slices are # listed heaviest-first (3D, then 2D, then 1D): with the long poles started # first, the short ones backfill the tail instead of extending it. From 2c4842bd335c9a578e18c6c8b3d939f65211de53 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 20 Aug 2026 21:15:08 +0000 Subject: [PATCH 6/7] refactor(cuda): reg_flow_slice.h opens its namespaces with the macros The seam header spelled `namespace ff { namespace cuda { namespace flow_slice {` by hand -- the one file in the split that did. CLAUDE.md is explicit that the namespace is opened with FF_NAMESPACE_BEGIN(FF_NS) / (FF_DEVICE) / () and never hard-coded, so that a header compiled for either backend lands in the right one rather than asserting `cuda`. Pulls in core/cuda_switch.h, which is where FF_DEVICE comes from and which already provides FF_NAMESPACE_BEGIN via core/defines.h. The other new files (reg_flow_slice.inl and the twelve slice TUs) already used the macros; this makes the set consistent. --- src/lib-cuda/reg_flow_slice.h | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/lib-cuda/reg_flow_slice.h b/src/lib-cuda/reg_flow_slice.h index 1af7751..b4497fe 100644 --- a/src/lib-cuda/reg_flow_slice.h +++ b/src/lib-cuda/reg_flow_slice.h @@ -54,6 +54,7 @@ */ #include +#include #include // Internal linkage would defeat the purpose (the definition and the call are @@ -133,9 +134,9 @@ bool use_32bits , \ intptr_t stream ) -namespace ff { -namespace cuda { -namespace flow_slice { +FF_NAMESPACE_BEGIN(FF_NS) +FF_NAMESPACE_BEGIN(FF_DEVICE) +FF_NAMESPACE_BEGIN(flow_slice) // One declaration per (entry point, ndim). The `add`/`sub` prefixes are the // `op` template argument ('+' / '-') the corresponding exported entry point @@ -159,6 +160,6 @@ FF_FLOW_SLICE_DECL_ND(3) #undef FF_FLOW_SLICE_DECL_ND -} // namespace flow_slice -} // namespace cuda -} // namespace ff +FF_NAMESPACE_END(flow_slice) +FF_NAMESPACE_END(FF_DEVICE) +FF_NAMESPACE_END(FF_NS) From ce69f02eb6f2eed7f3e352af74d4b378044ee434 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 20 Aug 2026 22:03:01 +0000 Subject: [PATCH 7/7] docs(cuda): re-measure the reg_flow split against post-#155 builds #155 fixed nvcc writing `::exit(1)` into the host object in place of every impl/kernels/utils.h helper, which is the state the first measurement of this split was taken in. Both sides of the table are now post-#155 runs of the same job, so nothing quoted here comes from a library that could not run. The finding is that #155 did not move the compile at all: unsplit reg_flow peaks at 12.98 GiB before and after the fix, because the code #155 restored is host code and the peak belongs to a device-side process. The split's case is unchanged, but it is now stated from runs that measure a working library. Re-measured, index32 (BEFORE run 32397332659 / AFTER run 32418602148): heaviest reg_flow TU 12.98 -> 1.94 GiB (-85.0%) reg_flow elapsed, sum 1069.30 -> 803.32 s (-24.9%) `make cuda -j2` wall 2145 -> 1861 s (-13.2%) max peak, any module 12.98 (reg_flow) -> 8.35 (reg_field) and the index64 leg besides, where the elapsed saving is much smaller (-5.7%) because collapsing the offset axis had already removed half the instantiations -- while memory still falls by four fifths. Two corrections to how the previous numbers were framed. The elapsed and wall-clock savings are smaller than first reported (-24.9% and -13.2%, not -34.6% and -25.5%); reg_field, which nothing here touches, moved 848.87 -> 861.05 s between the same two runs, so wall-clock on this job carries several percent of noise and is quoted with that caveat. Peak RSS does not: all thirteen per-slice peaks reproduce the earlier run to within 0.01 GiB, and reproduce off-runner on a different machine to 0.1-0.7%. Also annotates the #80 table above, whose reg_flow row no longer describes a translation unit that exists, with post-#155 figures for the modules this change does not touch. --- src/lib-cuda/Makefile | 160 +++++++++++++++++++++++++----------------- 1 file changed, 94 insertions(+), 66 deletions(-) diff --git a/src/lib-cuda/Makefile b/src/lib-cuda/Makefile index 8968ce1..69dbe5a 100644 --- a/src/lib-cuda/Makefile +++ b/src/lib-cuda/Makefile @@ -146,6 +146,17 @@ INDEXFLAGS ?= -DFF_INDEX32=1 # splinc 0.42 GB 67 s # posdef 0.37 GB 34 s # +# SUPERSEDED IN PART by fastfields-lib#147, below: `reg_flow` is no longer one +# translation unit, so its 12.98 GB row no longer describes anything that gets +# compiled. Everything else in this table is untouched by that change and was +# re-measured post-#155 (run 32397332659) -- reg_field 8.23, reg_field_rls +# 7.07, resize 2.00, reg_flow_rls 1.90, pushpull_backward 1.51, pushpull 1.48, +# restrict 1.30, distance 0.99, posdef 0.43, splinc 0.42 GiB. Two things to +# take from the comparison: the ordering is stable, and the individual numbers +# wander by up to a few tenths of a GB run to run (distance 0.77 -> 0.99 is the +# worst), which is the run-to-run spread the per-module budget below has to +# absorb. +# # (Peak RSS of the largest single process in the nvcc tree -- cicc or ptxas -- # not a sum: `time` reports ru_maxrss, which is a maximum. Across two runs the # heavy modules reproduce to better than 0.1%; the small ones wander a few @@ -213,79 +224,96 @@ INDEXFLAGS ?= -DFF_INDEX32=1 # ABI is unchanged -- the slice symbols are hidden-visibility and never reach # .dynsym. # -# ~~~ MEASURED, build-cuda on #147 vs on 85fdac7 (same image, nvcc 12.0.140, -# same -O1, same bound/spline policy, index32 leg) ~~~ -# -# reg_flow_matvec_3d 1.95 GiB 174.90 s reg_flow_matvec_1d 0.61 43.46 -# reg_flow_matvec_2d 1.22 GiB 91.61 s reg_flow_relax_2d 0.60 45.85 -# reg_flow_diag_3d 1.03 GiB 59.12 s reg_flow_kernel_2d 0.59 44.29 -# reg_flow_relax_3d 0.96 GiB 87.40 s reg_flow_kernel_1d 0.48 33.84 -# reg_flow_kernel_3d 0.89 GiB 78.43 s reg_flow_diag_1d 0.44 30.52 -# reg_flow_diag_2d 0.67 GiB 41.22 s reg_flow_relax_1d 0.26 19.80 -# reg_flow (front) 0.20 1.51 -# -# heaviest TU 12.98 -> 1.95 GiB (-85.0%) -# module CPU 1149.89 -> 751.95 s (-34.6%) -# wall -j2 2243 -> 1672 s (-25.5%, whole `make cuda`) -# .so link 1.38 -> 1.05 s +# ~~~ MEASURED, AND RE-MEASURED AFTER fastfields-lib#155 ~~~ +# +# Both sides below are post-#155 runs of this same job, so nothing here is +# inherited from the build that #155 fixed. That matters because #155 was +# nvcc replacing the body of every impl/kernels/utils.h helper IN THE HOST +# OBJECT with ::exit(1) -- so the first measurement of this split was taken +# against a library that could not run. It turns out not to have moved the +# compile at all: +# +# reg_flow, one TU, index32 peak elapsed +# pre-#155 (run 32365351378) 12.98 GiB 1149.89 s +# post-#155 (run 32397332659) 12.98 GiB 1069.30 s +# +# The peak is identical because the code #155 restored is HOST code, and the +# peak belongs to a device-side process (cicc/ptxas). A correctness +# catastrophe and a compile-cost non-event. +# +# BEFORE run 32397332659, commit 0d40731 (unsplit) +# AFTER run 32418602148, commit 2c4842b (this split) +# Same image, nvcc 12.0.140, -O1, same bound/spline policy, -j2. +# +# ~~~ index32 (the default, and what ships) ~~~ +# +# reg_flow_matvec_3d 1.94 GiB 177.13 s reg_flow_matvec_1d 0.61 48.84 +# reg_flow_matvec_2d 1.22 GiB 98.53 s reg_flow_relax_2d 0.60 48.01 +# reg_flow_diag_3d 1.03 GiB 64.72 s reg_flow_kernel_2d 0.59 49.71 +# reg_flow_relax_3d 0.96 GiB 89.84 s reg_flow_kernel_1d 0.48 38.77 +# reg_flow_kernel_3d 0.89 GiB 83.19 s reg_flow_diag_1d 0.44 35.04 +# reg_flow_diag_2d 0.67 GiB 46.24 s reg_flow_relax_1d 0.26 21.66 +# reg_flow (front) 0.20 1.64 +# +# heaviest reg_flow TU 12.98 -> 1.94 GiB (-85.0%) +# reg_flow elapsed, sum 1069.30 -> 803.32 s (-24.9%) +# `make cuda -j2` wall 2145 -> 1861 s (-13.2%) +# max peak, ANY module 12.98 (reg_flow) -> 8.35 (reg_field) +# +# All thirteen per-slice peaks reproduce the pre-#155 measurement of the same +# split to within 0.01 GiB, and reproduce OFF-RUNNER too: measured locally on +# nvcc 12.0.140 at the same flags, the front TU came out 205532 kB against +# CI's 205332 (0.1%), reg_flow_diag_1d 456904 against 456160 (0.2%). Peak RSS +# is the reproducible number here. WALL-CLOCK IS NOT: reg_field, which this +# change does not touch, moved 8.23 -> 8.35 GiB and 848.87 -> 861.05 s between +# the same two runs, so read the -13.2% as "clearly faster", not as 13.2%. +# +# ~~~ index64 ~~~ +# +# heaviest reg_flow TU 5.69 -> 1.05 GiB (-81.5%) +# reg_flow elapsed, sum 482.78 -> 455.23 s (-5.7%) +# `make cuda -j2` wall 966 -> 918 s (-5.0%) +# max peak, ANY module 5.69 (reg_flow) -> 3.76 (reg_field) +# +# The elapsed saving is far smaller on this leg, and that is the expected +# shape: collapsing the offset axis already removed half the instantiations, +# so there is less left for the split to divide. Memory still falls by 4/5. # # THE COUNTER-FORCE IS REAL BUT SMALL, and was measured rather than assumed. -# Every slice re-parses the whole template header set, so total CPU could have -# risen. A TU that parses impl/cuda/reg_flow.h plus the kernels and -# instantiates nothing costs 0.22 GiB and 1.35 s; twelve extra TUs is therefore -# ~16 s against a 1149.89 s baseline -- 1.4%, an order of magnitude under the -# instantiation work that now divides. That ratio is what makes this split pay. -# It is not a general law: a module whose per-TU instantiation work is -# comparable to 1.35 s should not be split. +# Every slice re-parses the whole template header set, so total elapsed could +# have risen. reg_flow.cpp IS that measurement: it parses impl/cuda/reg_flow.h +# plus the kernels and instantiates nothing, and costs 0.20 GiB and 1.64 s. +# Twelve extra TUs of re-parsing is ~20 s against a 1069 s baseline -- under +# 2%, an order of magnitude below the instantiation work that now divides. +# That ratio is what makes this split pay, and it is NOT a general law: a +# module whose per-TU instantiation work is comparable to ~1.6 s should not be +# split. # # WHAT THIS DOES *NOT* BUY: A HIGHER -j. -# The ceiling is max(peak) over EVERY module, and reg_field (8.23 GiB) and +# The ceiling is max(peak) over EVERY module, and reg_field (8.35 GiB) and # reg_field_rls (7.07 GiB) are untouched. reg_flow simply stops being the # constraint; reg_field becomes it, and the reg_field/reg_field_rls overlap -# (15.31 GiB) is the same -j2 hazard described above. Modelling make's dispatch -# over the measured table -- the model reproduces the observed -j2 wall to -# 0.25% -- gives -j4 a worst coincidence of 19.26 GiB, still over a 16 GB box. -# So do NOT raise -j on the strength of this change. Splitting reg_field and -# reg_field_rls the same way (they are structurally identical files) is what -# would take max(peak) to ~2 GiB and make -j4 CPU-bound instead of RAM-bound. -# -# INTERACTION WITH INDEXFLAGS: they MULTIPLY, they do not overlap. Measured on -# the heaviest slice, same TU, index axis toggled: -# -# reg_flow_matvec_3d FF_INDEX32=1 1.94 GiB 277.1 s 17.0 MB -# FF_INDEX32=0 1.05 GiB 174.4 s 10.0 MB -# -45.8% -37.1% -41.1% -# -# That -45.8% is the same proportion #94 measured on the UNSPLIT file with the -# host compiler (-44.6%), which is the point: splitting does not make the index -# axis any cheaper in relative terms, it only shrinks the absolute number the -# proportion applies to. The two compose -- 12.98 -> 1.94 GiB from the split, -# 1.94 -> 1.05 GiB from the knob. So this split is NOT a reason to stop caring -# about FF_INDEX32; what it changes is which resource the knob buys back, since -# memory is no longer what binds this module and wall-clock is. -# -# AND THE TWO TOGETHER ARE WHAT WOULD ACTUALLY UNLOCK -j4. build-cuda's -# index64 leg builds this same split with the axis off, and that combination -- -# not either half alone -- is the first configuration in which no module is -# heavy any more: -# -# reg_field 3.76 GiB 264.90 s <-- the new maximum -# reg_field_rls 3.42 GiB 279.38 s -# resize 1.08 GiB 90.20 s -# reg_flow_matvec_3d 1.05 GiB 81.45 s <-- heaviest reg_flow slice -# ... every other module below 1.01 GiB ... -# reg_flow (front) 0.20 GiB 0.88 s -# -# total nvcc CPU 1380.3 s; `make cuda -j2` wall 723 s (measured). -# -# Feeding that table through the same scheduling model (which reproduces the -# measured -j2 wall to 0.5%): -j4 gives a 393 s makespan with a worst -# coincidence of 9.31 GiB, comfortably inside a 16 GB runner. So `-j4` is safe -# on the index64 leg TODAY, with reg_field untouched. It is NOT safe on the -# index32 leg, which is the default and the one that ships -- there the same -# model says 19.26 GiB. Do not raise -j globally on the strength of the -# index64 number. +# (15.42 GiB against a 16 GB runner) is the same -j2 hazard described above -- +# unchanged by this PR, because neither module is touched by it. So do NOT +# raise -j on the strength of this change. Splitting reg_field and +# reg_field_rls the same way (they are structurally identical files, and the +# port is a prefix rename of this diff) is what would take max(peak) to ~2 GiB +# and make -j4 CPU-bound instead of RAM-bound. +# +# INTERACTION WITH INDEXFLAGS: they compose, they do not overlap. On the +# heaviest slice, same TU, index axis toggled (the two legs above): +# +# reg_flow_matvec_3d FF_INDEX32=1 1.94 GiB 177.13 s +# FF_INDEX32=0 1.05 GiB 100.01 s +# -45.8% -43.5% +# +# Splitting does not make the index axis relatively cheaper; it shrinks the +# absolute number the proportion applies to. The two compose -- 12.98 -> 1.94 +# GiB from the split, 1.94 -> 1.05 GiB from the knob. So this split is NOT a +# reason to stop caring about FF_INDEX32; what it changes is which resource +# the knob buys back, since memory is no longer what binds this module. # + # ORDER MATTERS HERE. make dispatches in MODULES order, so the slices are # listed heaviest-first (3D, then 2D, then 1D): with the long poles started # first, the short ones backfill the tail instead of extending it.