diff --git a/src/lib-cuda/Makefile b/src/lib-cuda/Makefile index 7d35525..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 @@ -199,6 +210,127 @@ INDEXFLAGS ?= -DFF_INDEX32=1 # 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 (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 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 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. The exported +# ABI is unchanged -- the slice symbols are hidden-visibility and never reach +# .dynsym. +# +# ~~~ 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 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.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.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. +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 \ @@ -208,6 +340,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 1f4bdba..357ee49 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 #include -#include #include -#include #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 // canUse32BitIndexMath +#include // 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..b4497fe --- /dev/null +++ b/src/lib-cuda/reg_flow_slice.h @@ -0,0 +1,165 @@ +#pragma once +/** + * `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 +#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: +// 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 ) + +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 +// 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 + +FF_NAMESPACE_END(flow_slice) +FF_NAMESPACE_END(FF_DEVICE) +FF_NAMESPACE_END(FF_NS) diff --git a/src/lib-cuda/reg_flow_slice.inl b/src/lib-cuda/reg_flow_slice.inl new file mode 100644 index 0000000..7cccb65 --- /dev/null +++ b/src/lib-cuda/reg_flow_slice.inl @@ -0,0 +1,682 @@ +#pragma once +/** + * 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 +#include +#include +#include +#include +#include +#include +#include +#include +#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)