From 69fbabf21f5c5b69af6ba428ad86e322c50bfbb8 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 19 Aug 2026 23:14:38 +0000 Subject: [PATCH] prototype: one shared dispatch driver + RAII autocast, and an FF_INDEX32 knob RFC prototype for discussion, not for merge as-is. One module converted (src/lib-cpu/reg_field.cpp) to show what the other five of its shape would look like. Primary motivation is instantiation count, not line count. Every templated kernel below the dispatch layer is templated on offset_t, offset_t has exactly two values, and which one is used is a runtime property of the tensor (canUse32BitIndexMath). So the 32-bit index path costs exactly x2 instantiations of everything -- on the CUDA side, x2 device code, hence x2 ptxas memory and x2 SASS in a library that is already 166 MB and whose reg_flow module already peaks at 13.6 GB of nvcc RSS on a 16 GB runner, before any -arch/-gencode multiplier. include/fastfields/api/dispatch.h (new) The ndim x bound x dtype x index-width pyramid, stated once. dispatch_nbd(key, args...) replaces NDIM_SWITCH / BOUND_SWITCH / BND1..3 / one _DT per entry point / one _ARGS per entry point. Arguments become an ordinary call argument list, so the GNU `args...` named-variadic-macro syntax (which MSVC's traditional preprocessor rejects) goes with them. FF_INDEX32 turns the index axis into one build flag instead of ~40 hard-coded ternaries across 21 translation units. include/fastfields/core/autocast.h IndexArray: RAII replacement for the copy_if_needed / free_if_needed pair, with a small stack buffer instead of an allocation. Fixes the leak the manual form has when anything between the two throws, and takes cudaMallocHost/cudaFreeHost off the CUDA per-call path. The old helpers are untouched; nothing else uses IndexArray yet. src/lib-cpu/reg_field.cpp Converted. 1369 -> 1036 lines. _field_matvec was byte-identical to _field_matvec_acc<'='> and is gone; the seven remaining wrappers became op structs, and the op character is a template parameter of the struct rather than a macro per variant. tests/dispatch/packcheck.cpp Hand-run, not globbed by any make target, so the recorded baseline is untouched. Checks the driver hands the leaf the same ndim, scalar_t, offset_t and boundary-pack length the macros did, over all 48 (ndim, bound, index-width) configurations, and that the dtype / ndim / bound rejection paths still throw. C++11. Clean under clang++ and g++ at -std=c++11 and -std=c++14, -pedantic, with and without FF_INDEX32. --- include/fastfields/api/dispatch.h | 234 ++++++ include/fastfields/core/autocast.h | 88 ++ src/lib-cpu/reg_field.cpp | 1243 ++++++++++------------------ tests/dispatch/packcheck.cpp | 75 ++ 4 files changed, 852 insertions(+), 788 deletions(-) create mode 100644 include/fastfields/api/dispatch.h create mode 100644 tests/dispatch/packcheck.cpp diff --git a/include/fastfields/api/dispatch.h b/include/fastfields/api/dispatch.h new file mode 100644 index 0000000..db351f6 --- /dev/null +++ b/include/fastfields/api/dispatch.h @@ -0,0 +1,234 @@ +#ifndef FF_API_DISPATCH +#define FF_API_DISPATCH +/** + * Private (not installed) header: the ndim x bound x dtype x index-width + * dispatch pyramid, stated once. + * + * PROTOTYPE -- see the accompanying proposal. Every module under src/lib-cpu + * and src/lib-cuda carries its own hand-written copy of that pyramid as nested + * function-like macros: `NDIM_SWITCH`, `BOUND_SWITCH`, `BND1`/`BND2`/`BND3`, + * one `_DT` per entry point and one `_ARGS` per entry point. The + * switches are byte-identical across modules; only the leaf call differs. This + * header factors the switches out and leaves the leaf as a template argument. + * + * Usage + * ----- + * Wrap the module's leaf in a struct with a single static `run`, templated on + * exactly the axes the pyramid resolves: + * + * struct field_matvec_op { + * template + * static void run() { ... } + * }; + * + * and call it: + * + * dispatch_nbd(key, arg0, arg1, ...); + * + * The arguments are an ordinary function-call argument list, so the + * `#define _ARGS ... #undef` pairs disappear -- and with them the GNU + * `args...` named-variadic-macro syntax, which MSVC's traditional preprocessor + * does not accept. + * + * The boundary condition is expanded to a pack of `ndim` copies, exactly as + * BND1/BND2/BND3 did: `bound::getutils` / `` / `` detect + * isotropy from the pack length, so the length is load-bearing. + * tests/dispatch/packcheck.cpp checks that it still is. + * + * C++11: no `if constexpr`, no fold expressions, no variable templates, no + * generic lambdas. Compiles under clang++ and g++ at -std=c++11 and + * -std=c++14 (the CUDA layer's standard). + */ +#include +#include +#include "fastfields/core/dlpack.h" +#include "fastfields/core/cuda_switch.h" +#include "fastfields/impl/kernels/bounds.h" + +FF_NAMESPACE_BEGIN(FF) +FF_NAMESPACE_BEGIN(FF_DEVICE) + +/*********************************************************************** + * INDEX-WIDTH POLICY * + ***********************************************************************/ +// +// Whether the 32-bit index specialisation (`offset_t == int32_t`) is compiled +// at all. +// +// This is the single largest multiplier in the whole build. Every templated +// kernel below the dispatch layer takes `offset_t`, and `offset_t` has exactly +// two values, chosen at run time by `canUse32BitIndexMath`. So the axis costs +// exactly x2 instantiations of everything -- and on the CUDA side "everything" +// is device code, which is ptxas memory and SASS in the shipped library. +// +// It is stated once, here, instead of being hard-coded into ~40 `_DT` +// macros across 21 translation units, so that turning it off is one flag: +// +// make cpu CXXFLAGS='-std=c++11 -O3 -DFF_INDEX32=0' +// +// With FF_INDEX32=0 the key's `use_32bits` field is ignored, `offset_t` is +// always int64_t, and the whole narrowing path in core/autocast.h collapses to +// a no-op passthrough. +// +// Nothing in this repository measures what the 32-bit path buys. It is +// inherited from ATen's `canUse32BitIndexMath`, where it is a GPU +// register-pressure optimisation; here it is applied to the CPU backend too, +// and on the CUDA side it is paid for with a `cudaMallocHost` per narrowed +// array per call (see core/autocast.h). +#ifndef FF_INDEX32 +# define FF_INDEX32 1 +#endif + +/*********************************************************************** + * THE KEY * + ***********************************************************************/ + +// The runtime axes the pyramid switches on. Built once per public entry point, +// immediately after the argument checks. +struct DispatchKey +{ + int ndim; + bound::type bound; + DLDataTypeCode code; + uint8_t bits; + bool use_32bits; + + inline DispatchKey(int ndim_, int8_t bound_, const DLDataType & dtype, + bool use_32bits_) + : ndim(ndim_), + bound(static_cast(bound_)), + code(static_cast(dtype.code)), + bits(dtype.bits), + use_32bits(use_32bits_) + {} +}; + +/*********************************************************************** + * LEVEL 1 -- dtype x index width * + ***********************************************************************/ +// Innermost level. `Op::run` is called with every compile-time axis fixed; the +// boundary pack `B...` already holds `ndim` entries. +template +struct _dispatch_dtype +{ + template + static inline void call(const DispatchKey & key, A &&... args) + { + switch (key.code) { + case kDLFloat: + switch (key.bits) { +#if FF_INDEX32 + case 32: return key.use_32bits + ? Op::template run(args...) + : Op::template run(args...); + case 64: return key.use_32bits + ? Op::template run(args...) + : Op::template run(args...); +#else + case 32: return Op::template run(args...); + case 64: return Op::template run(args...); +#endif + default: break; + } + break; + default: break; + } + throw std::invalid_argument("only floating point data types are supported"); + } +}; + +/*********************************************************************** + * LEVEL 2 -- expand the boundary pack * + ***********************************************************************/ +// BND1/BND2/BND3, as three explicit specialisations rather than a recursion: +// the recursive form instantiates `ndim + 1` class templates (and as many +// member-function templates) per (op, ndim, bound) triple, which is pure +// front-end cost for no benefit. +template struct _expand_bound; + +template +struct _expand_bound +{ + template + static inline void call(const DispatchKey & key, A &&... args) + { _dispatch_dtype::call(key, args...); } +}; + +template +struct _expand_bound +{ + template + static inline void call(const DispatchKey & key, A &&... args) + { _dispatch_dtype::call(key, args...); } +}; + +template +struct _expand_bound +{ + template + static inline void call(const DispatchKey & key, A &&... args) + { _dispatch_dtype::call(key, args...); } +}; + +/*********************************************************************** + * LEVEL 3 -- boundary condition * + ***********************************************************************/ +// The template argument is `FF_BOUND_` (kernels/bounds.h): the condition +// itself when it is statically compiled, `Dynamic` otherwise, per BOUNDFLAGS. +// The switch labels stay exhaustive on the runtime value either way; only the +// instantiated template argument collapses onto the shared Dynamic path. +template +struct _dispatch_bound +{ + template + static inline void call(const DispatchKey & key, A &&... args) + { + switch (key.bound) { + case bound::type::Zero: + return _expand_bound::call(key, args...); + case bound::type::Replicate: + return _expand_bound::call(key, args...); + case bound::type::DCT1: + return _expand_bound::call(key, args...); + case bound::type::DCT2: + return _expand_bound::call(key, args...); + case bound::type::DST1: + return _expand_bound::call(key, args...); + case bound::type::DST2: + return _expand_bound::call(key, args...); + case bound::type::DFT: + return _expand_bound::call(key, args...); + case bound::type::NoCheck: + return _expand_bound::call(key, args...); + default: + throw std::invalid_argument("Unsupported boundary condition"); + } + } +}; + +/*********************************************************************** + * LEVEL 4 -- ndim * + ***********************************************************************/ + +/** + * @brief ndim x bound x dtype x index-width dispatch. + * + * Replaces `NDIM_SWITCH(_DT)` together with the `_DT` and `_ARGS` + * macro pair. + */ +template +inline void dispatch_nbd(const DispatchKey & key, A &&... args) +{ + switch (key.ndim) { + case 1: return _dispatch_bound::call(key, args...); + case 2: return _dispatch_bound::call(key, args...); + case 3: return _dispatch_bound::call(key, args...); + default: + throw std::invalid_argument("Only 1D, 2D and 3D are supported"); + } +} + +FF_NAMESPACE_END(FF_DEVICE) +FF_NAMESPACE_END(FF) + +#endif // FF_API_DISPATCH diff --git a/include/fastfields/core/autocast.h b/include/fastfields/core/autocast.h index de3af9e..15807a4 100644 --- a/include/fastfields/core/autocast.h +++ b/include/fastfields/core/autocast.h @@ -197,6 +197,94 @@ inline void free_if_needed(OutPointer ptr) _copy_if_needed::free(ptr); } +// ------------------------------------------------------------------ RAII +// +// PROTOTYPE -- see the accompanying proposal. `IndexArray` is the +// RAII form of the hand-managed `copy_if_needed` / `free_if_needed` pair. It +// replaces +// +// const offset_t * _size = copy_if_needed(size, n); +// ... 30 lines, any of which may throw ... +// free_if_needed(_size); // not reached if one did +// +// with +// +// IndexArray _size(size, n); +// +// Three things follow. +// +// 1. The leak on the throwing path goes away. It is not hypothetical: every +// reg_* impl wrapper does `new reduce_t[...]` after the copy, and every +// CUDA launcher throws `std::bad_alloc` / `std::range_error` from +// `copyToDevice` / `GET_BLOCKS`. +// +// 2. The allocator leaves the narrowing path. The arrays being narrowed are +// shape/stride vectors of length `nbatch + ndim (+1)`; FF_INDEX_SBO covers +// every rank the library dispatches (ndim <= 3) with a heap fallback for +// anything larger. +// +// 3. On the CUDA build that deletes a `cudaMallocHost` + `cudaFreeHost` pair +// per narrowed array per call. Page-locked allocation is among the most +// expensive host-side CUDA calls there is and it synchronises the device. +// The header's justification for it -- "so the following H2D copy can be +// async" -- does not hold for the launchers that actually exist: 365 of +// the impl/cuda upload sites call the *synchronous* `copyToDevice`, which +// gains nothing from a pinned source, and the handful that call +// `copyToDeviceAsync` document pageable sources as safe. +// +// A stack buffer is safe for both: `copyToDevice` is synchronous, and +// `copyToDeviceAsync` from pageable memory stages through a driver buffer +// before returning (its own comment says so). +#ifndef FF_INDEX_SBO +# define FF_INDEX_SBO 8 +#endif + +template +class IndexArray +{ + offset_t _sbo[FF_INDEX_SBO]; + offset_t * _heap; + const offset_t * _ptr; + +public: + IndexArray(const int64_t * src, size_t numel) : _heap(nullptr), _ptr(nullptr) + { + if (!src) return; + offset_t * dst = _sbo; + if (numel > FF_INDEX_SBO) + dst = _heap = new offset_t[numel]; + for (size_t i = 0; i < numel; ++i) + dst[i] = static_cast(src[i]); + _ptr = dst; + } + + ~IndexArray() { delete[] _heap; } + + inline operator const offset_t * () const { return _ptr; } + inline const offset_t * get() const { return _ptr; } + + IndexArray(const IndexArray &) = delete; + IndexArray & operator=(const IndexArray &) = delete; +}; + +// Nothing to narrow: borrow the caller's array. Same interface, zero cost. +// This is the only specialisation instantiated when FF_INDEX32 == 0, at which +// point the five `_copy_if_needed` specialisations above have no callers left. +template <> +class IndexArray +{ + const int64_t * _ptr; + +public: + IndexArray(const int64_t * src, size_t /* numel */) : _ptr(src) {} + + inline operator const int64_t * () const { return _ptr; } + inline const int64_t * get() const { return _ptr; } + + IndexArray(const IndexArray &) = delete; + IndexArray & operator=(const IndexArray &) = delete; +}; + FF_NAMESPACE_END(FF_DEVICE) FF_NAMESPACE_END(FF) diff --git a/src/lib-cpu/reg_field.cpp b/src/lib-cpu/reg_field.cpp index 092a074..8fc790b 100644 --- a/src/lib-cpu/reg_field.cpp +++ b/src/lib-cpu/reg_field.cpp @@ -4,6 +4,7 @@ #include #include "fastfields/api/cpu/reg_field.h" #include "fastfields/api/cpu/posdef.h" +#include "fastfields/api/dispatch.h" #include "fastfields/core/autocast.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" @@ -56,455 +57,409 @@ static inline std::vector as_weights(const double * w, int64_t nc) return v; } -template -inline void _field_matvec( - const bound::BoundVec & bvec, - int64_t nbatch , - int64_t nc , - void * out , - const void * inp , - const double * voxel_size , - const double * absolute , - const double * membrane , - const double * bending , - const int64_t * size , - const int64_t * stride_out , - const int64_t * stride_inp ) -{ - 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; - - std::vector a = as_weights(absolute, nc); - std::vector m = as_weights(membrane, nc); - std::vector b = as_weights(bending, nc); - - if (bending) - reg_field::matvec_bending( - bvec, static_cast(nbatch), _out, _inp, - _size, _stride_out, _stride_inp, vx, a.data(), m.data()); - else - reg_field::matvec_absolute( - bvec, static_cast(nbatch), _out, _inp, _wgt, - _size, _stride_out, _stride_inp, _stride_wgt, vx, a.data(), m.data(), b.data()); - else - reg_field::matvec_bending_rls( - bvec, static_cast(nbatch), _out, _inp, _wgt, - _size, _stride_out, _stride_inp, _stride_wgt, vx, a.data(), m.data()); - else - reg_field::matvec_membrane_rls( - bvec, static_cast(nbatch), _out, _inp, _wgt, - _size, _stride_out, _stride_inp, _stride_wgt, a.data()); - else - reg_field::matvec_absolute_rls( + bvec, static_cast(nbatch), _out, _inp, _wgt, + _size, _stride_out, _stride_inp, _stride_wgt, vx, a.data(), m.data(), b.data()); + else + reg_field::matvec_bending_rls( + bvec, static_cast(nbatch), _out, _inp, _wgt, + _size, _stride_out, _stride_inp, _stride_wgt, vx, a.data(), m.data()); + else + reg_field::matvec_membrane_rls( + bvec, static_cast(nbatch), _out, _inp, _wgt, + _size, _stride_out, _stride_inp, _stride_wgt, a.data()); + else + reg_field::matvec_absolute_rls( - bvec, static_cast(nbatch), _out, _wgt, - _size, _stride_out, _stride_wgt, vx, a.data(), m.data(), b.data()); - else - reg_field::diag_bending_rls( - bvec, static_cast(nbatch), _out, _wgt, - _size, _stride_out, _stride_wgt, vx, a.data(), m.data()); - else - reg_field::diag_membrane_rls( - bvec, static_cast(nbatch), _out, _wgt, - _size, _stride_out, _stride_wgt, a.data()); - else - reg_field::diag_absolute_rls( + bvec, static_cast(nbatch), _out, _wgt, + _size, _stride_out, _stride_wgt, vx, a.data(), m.data(), b.data()); + else + reg_field::diag_bending_rls( + bvec, static_cast(nbatch), _out, _wgt, + _size, _stride_out, _stride_wgt, vx, a.data(), m.data()); + else + reg_field::diag_membrane_rls( + bvec, static_cast(nbatch), _out, _wgt, + _size, _stride_out, _stride_wgt, a.data()); + else + reg_field::diag_absolute_rls(MV_ARGS) \ - : _field_matvec_acc(MV_ARGS); \ - case 64: return use_32bits \ - ? _field_matvec_acc(MV_ARGS) \ - : _field_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 \ - ? _field_matvec_acc(MV_ARGS) \ - : _field_matvec_acc(MV_ARGS); \ - case 64: return use_32bits \ - ? _field_matvec_acc(MV_ARGS) \ - : _field_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 \ - ? _field_diag(DG_ARGS); \ - case 64: return use_32bits \ - ? _field_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 \ - ? _field_diag(DG_ARGS) \ - : _field_diag(DG_ARGS); \ - case 64: return use_32bits \ - ? _field_diag(DG_ARGS) \ - : _field_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 \ - ? _field_diag(DG_ARGS) \ - : _field_diag(DG_ARGS); \ - case 64: return use_32bits \ - ? _field_diag(DG_ARGS) \ - : _field_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 \ - ? _field_kernel(KN_ARGS); \ - case 64: return use_32bits \ - ? _field_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 \ - ? _field_kernel(KN_ARGS) \ - : _field_kernel(KN_ARGS); \ - case 64: return use_32bits \ - ? _field_kernel(KN_ARGS) \ - : _field_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 \ - ? _field_kernel(KN_ARGS) \ - : _field_kernel(KN_ARGS); \ - case 64: return use_32bits \ - ? _field_kernel(KN_ARGS) \ - : _field_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 \ - ? _field_relax(RX_ARGS) \ - : _field_relax(RX_ARGS); \ - case 64: return use_32bits \ - ? _field_relax(RX_ARGS) \ - : _field_relax(RX_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define RLS_MV_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _field_matvec_rls(RLS_MV_ARGS) \ - : _field_matvec_rls(RLS_MV_ARGS); \ - case 64: return use_32bits \ - ? _field_matvec_rls(RLS_MV_ARGS) \ - : _field_matvec_rls(RLS_MV_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define RLS_DG_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _field_diag_rls(RLS_DG_ARGS) \ - : _field_diag_rls(RLS_DG_ARGS); \ - case 64: return use_32bits \ - ? _field_diag_rls(RLS_DG_ARGS) \ - : _field_diag_rls(RLS_DG_ARGS); \ - default: break; \ - } break; \ - default: break; \ - } \ - throw std::invalid_argument("only floating point data types are supported"); - -#define RLS_RX_DT(NDIM, BNDS...) \ - switch (code) { \ - case kDLFloat: switch (bits) { \ - case 32: return use_32bits \ - ? _field_relax_rls(RLS_RX_ARGS) \ - : _field_relax_rls(RLS_RX_ARGS); \ - case 64: return use_32bits \ - ? _field_relax_rls(RLS_RX_ARGS) \ - : _field_relax_rls(RLS_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) \ - 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; \ - default: throw std::invalid_argument("Only 1D, 2D and 3D field are supported"); \ - } +// +// `dispatch_nbd(key, ...)` (api/dispatch.h) is the entire pyramid: +// ndim (1/2/3) x boundary condition (8, each expanded to a pack of `ndim` +// copies) x dtype (f32/f64) x index width. It replaces the twelve `_DT` +// macros, `BOUND_SWITCH`, `BND1`/`BND2`/`BND3`, `NDIM_SWITCH` and the six +// `#define _ARGS ... #undef` pairs this file used to carry. void field_matvec( DLTensor & out_ , @@ -762,17 +500,12 @@ void field_matvec( CHECK_SAME_SHAPE(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && 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); - -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides, inp.strides - NDIM_SWITCH(MV_DT) -#undef MV_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out) && CANUSE32BITS(inp)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), voxel_size, absolute, membrane, bending, out.shape, out.strides, inp.strides); } /** @@ -804,17 +537,12 @@ void field_addmatvec_( CHECK_SAME_SHAPE(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && 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); - -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides, inp.strides - NDIM_SWITCH(ADD_MV_DT) -#undef MV_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out) && CANUSE32BITS(inp)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), voxel_size, absolute, membrane, bending, out.shape, out.strides, inp.strides); } /** @@ -846,17 +574,12 @@ void field_submatvec_( CHECK_SAME_SHAPE(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && 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); - -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides, inp.strides - NDIM_SWITCH(SUB_MV_DT) -#undef MV_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out) && CANUSE32BITS(inp)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), voxel_size, absolute, membrane, bending, out.shape, out.strides, inp.strides); } void field_diag( @@ -880,17 +603,12 @@ void field_diag( throw std::invalid_argument("ndim is larger than the tensor rank"); const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = 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); - -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides - NDIM_SWITCH(DG_DT) -#undef DG_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), voxel_size, absolute, membrane, bending, out.shape, out.strides); } /** @@ -921,17 +639,12 @@ void field_adddiag_( throw std::invalid_argument("ndim is larger than the tensor rank"); const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = 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); - -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides - NDIM_SWITCH(ADD_DG_DT) -#undef DG_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), voxel_size, absolute, membrane, bending, out.shape, out.strides); } /** @@ -958,17 +671,12 @@ void field_subdiag_( throw std::invalid_argument("ndim is larger than the tensor rank"); const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = 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); - -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides - NDIM_SWITCH(SUB_DG_DT) -#undef DG_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), voxel_size, absolute, membrane, bending, out.shape, out.strides); } void field_kernel( @@ -992,17 +700,12 @@ void field_kernel( throw std::invalid_argument("ndim is larger than the tensor rank"); const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = 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); - -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides - NDIM_SWITCH(KN_DT) -#undef KN_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), voxel_size, absolute, membrane, bending, out.shape, out.strides); } /** @@ -1030,17 +733,12 @@ void field_addkernel_( throw std::invalid_argument("ndim is larger than the tensor rank"); const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = 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); - -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides - NDIM_SWITCH(ADD_KN_DT) -#undef KN_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), voxel_size, absolute, membrane, bending, out.shape, out.strides); } /** @@ -1068,17 +766,12 @@ void field_subkernel_( throw std::invalid_argument("ndim is larger than the tensor rank"); const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = 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); - -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides - NDIM_SWITCH(SUB_KN_DT) -#undef KN_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out)); + + dispatch_nbd >(key, + bvec, static_cast(nbatch), nc, VOIDPTR(out), voxel_size, absolute, membrane, bending, out.shape, out.strides); } void field_relax( @@ -1106,18 +799,12 @@ void field_relax( CHECK_SAME_SHAPE(sol, grd, sol.ndim) const int64_t nc = sol.shape[sol.ndim - 1]; - const bool use_32bits = CANUSE32BITS(sol) && CANUSE32BITS(hes) && - 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); - -#define RX_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(sol), CVOIDPTR(hes), \ - CVOIDPTR(grd), voxel_size, absolute, membrane, bending, \ - nb_iter, sol.shape, sol.strides, hes.strides, grd.strides - NDIM_SWITCH(RX_DT) -#undef RX_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, sol.dtype, + CANUSE32BITS(sol) && CANUSE32BITS(hes) && CANUSE32BITS(grd)); + + dispatch_nbd(key, + bvec, static_cast(nbatch), nc, VOIDPTR(sol), CVOIDPTR(hes), CVOIDPTR(grd), voxel_size, absolute, membrane, bending, nb_iter, sol.shape, sol.strides, hes.strides, grd.strides); } void field_forward( @@ -1260,18 +947,12 @@ void field_matvec_rls( const int64_t nc = out.shape[out.ndim - 1]; const bool is_jrls = field_rls_is_jrls(wgt, nc, "field_matvec_rls"); - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(wgt); - 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); - -#define RLS_MV_ARGS bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(out), \ - CVOIDPTR(inp), CVOIDPTR(wgt), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides, inp.strides, wgt.strides - NDIM_SWITCH(RLS_MV_DT) -#undef RLS_MV_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(wgt)); + + dispatch_nbd(key, + bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(wgt), voxel_size, absolute, membrane, bending, out.shape, out.strides, inp.strides, wgt.strides); } void field_diag_rls( @@ -1301,18 +982,12 @@ void field_diag_rls( const int64_t nc = out.shape[out.ndim - 1]; const bool is_jrls = field_rls_is_jrls(wgt, nc, "field_diag_rls"); - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(wgt); - 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); - -#define RLS_DG_ARGS bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(out), \ - CVOIDPTR(wgt), \ - voxel_size, absolute, membrane, bending, \ - out.shape, out.strides, wgt.strides - NDIM_SWITCH(RLS_DG_DT) -#undef RLS_DG_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, out.dtype, + CANUSE32BITS(out) && CANUSE32BITS(wgt)); + + dispatch_nbd(key, + bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(out), CVOIDPTR(wgt), voxel_size, absolute, membrane, bending, out.shape, out.strides, wgt.strides); } void field_relax_rls( @@ -1349,20 +1024,12 @@ void field_relax_rls( const int64_t nc = sol.shape[sol.ndim - 1]; const bool is_jrls = field_rls_is_jrls(wgt, nc, "field_relax_rls"); - const bool use_32bits = CANUSE32BITS(sol) && CANUSE32BITS(hes) && - CANUSE32BITS(grd) && CANUSE32BITS(wgt); - 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); - -#define RLS_RX_ARGS bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(sol), \ - CVOIDPTR(hes), CVOIDPTR(grd), CVOIDPTR(wgt), \ - voxel_size, absolute, membrane, bending, \ - nb_iter, sol.shape, sol.strides, hes.strides, grd.strides, \ - wgt.strides - NDIM_SWITCH(RLS_RX_DT) -#undef RLS_RX_ARGS + const bound::BoundVec bvec(static_cast(bound)); + const DispatchKey key(ndim, bound, sol.dtype, + CANUSE32BITS(sol) && CANUSE32BITS(hes) && CANUSE32BITS(grd) && CANUSE32BITS(wgt)); + + dispatch_nbd(key, + bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(sol), CVOIDPTR(hes), CVOIDPTR(grd), CVOIDPTR(wgt), voxel_size, absolute, membrane, bending, nb_iter, sol.shape, sol.strides, hes.strides, grd.strides, wgt.strides); } FF_NAMESPACE_END(FF_DEVICE) diff --git a/tests/dispatch/packcheck.cpp b/tests/dispatch/packcheck.cpp new file mode 100644 index 0000000..04cd33a --- /dev/null +++ b/tests/dispatch/packcheck.cpp @@ -0,0 +1,75 @@ +// Checks that api/dispatch.h hands the leaf exactly what the macro pyramid +// used to: the right ndim, the right scalar_t, the right offset_t, and a +// boundary pack of exactly `ndim` copies (BND1/BND2/BND3) -- the pack length is +// load-bearing, because bound::getutils / / detect isotropy +// from it. Also checks the three rejection paths still throw. +// +// NOT wired into any make target, on purpose: `make test` globs +// tests/lib-cpu/test_*.cpp and tests/lib/test_*.cpp, so this file cannot change +// the recorded baseline (13 suites / 59886 checks, plus 2 hub suites). Run it +// by hand, the way tests/kernels/vector/test.cpp is run: +// +// clang++ -std=c++11 -O1 -Iinclude -o /tmp/packcheck tests/dispatch/packcheck.cpp && /tmp/packcheck +// clang++ -std=c++11 -O1 -Iinclude -DFF_INDEX32=0 -o /tmp/packcheck0 tests/dispatch/packcheck.cpp && /tmp/packcheck0 +// g++ -std=c++11 -O1 -Iinclude -o /tmp/packcheckg tests/dispatch/packcheck.cpp && /tmp/packcheckg +#include +#include +#include +#include "fastfields/api/dispatch.h" + +using namespace ff::cpu; + +static std::vector g_ndim, g_npack, g_off, g_sca; + +struct probe_op { + template + static void run(int /* sentinel */) { + g_ndim .push_back(ndim); + g_npack.push_back(static_cast(sizeof...(B))); + g_off .push_back(static_cast(sizeof(offset_t))); + g_sca .push_back(static_cast(sizeof(scalar_t))); + } +}; + +int main() +{ + int bad = 0, checks = 0; + DLDataType f32; f32.code = kDLFloat; f32.bits = 32; f32.lanes = 1; + DLDataType f64; f64.code = kDLFloat; f64.bits = 64; f64.lanes = 1; + + for (int nd = 1; nd <= 3; ++nd) + for (int b = 0; b < 8; ++b) + for (int w = 0; w < 2; ++w) + { + g_ndim.clear(); g_npack.clear(); g_off.clear(); g_sca.clear(); + DispatchKey k32(nd, static_cast(b), f32, w != 0); + dispatch_nbd(k32, 0); + DispatchKey k64(nd, static_cast(b), f64, w != 0); + dispatch_nbd(k64, 0); + ++checks; + if (g_ndim.size() != 2) { ++bad; continue; } + if (g_ndim [0] != nd || g_ndim [1] != nd) ++bad; // ndim threaded through + if (g_npack[0] != nd || g_npack[1] != nd) ++bad; // BND + if (g_sca [0] != 4 || g_sca [1] != 8 ) ++bad; // float / double +#if FF_INDEX32 + const int want = w ? 4 : 8; +#else + const int want = 8; // axis compiled out +#endif + if (g_off[0] != want || g_off[1] != want) ++bad; + } + + int threw = 0; + DLDataType i32; i32.code = kDLInt; i32.bits = 32; i32.lanes = 1; + try { DispatchKey k(2, 3, i32, true); dispatch_nbd(k, 0); } + catch (const std::invalid_argument &) { ++threw; } // bad dtype + try { DispatchKey k(4, 3, f32, true); dispatch_nbd(k, 0); } + catch (const std::invalid_argument &) { ++threw; } // bad ndim + try { DispatchKey k(2, 99, f32, true); dispatch_nbd(k, 0); } + catch (const std::invalid_argument &) { ++threw; } // bad bound + if (threw != 3) ++bad; + + std::printf("packcheck: %s (%d configurations, %d problems)\n", + bad ? "FAIL" : "OK", checks, bad); + return bad != 0; +}