From 1de1870dd0cd40a43eb841f49b116697181cd85c Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 19 Aug 2026 23:08:25 +0000 Subject: [PATCH] refactor: de-duplicate the dispatch helpers into shared headers Every src/lib-cpu/.cpp and src/lib-cuda/.cpp opened with the same private prologue: pointer marshalling macros, a 32-bit-index predicate, and a family of argument checks. There were 19 copies of VOIDPTR, 19 of CANUSE32BITS, 19 of CHECK_NO_LANES, 17 of CHECK_SAME and 17 of CHECK_SAME_DTYPE, plus 9 copies each of the hub's IS_CPU / IS_CUDA, three of as_weights() and four of _reg_stream(). They had already drifted. Three of them are not a formatting difference: CVOIDPTR posdef.cpp alone guards against a null `data`, because an absent optional operand arrives as a descriptor with data == nullptr and offsetting a null pointer is UB. CHECK_SAME_BATCH distance.cpp and posdef.cpp alone also reject ndim < D, because they derive D from an operand's own ndim and would otherwise read past the end of shape[]. CHECK_SAME_SHAPE was TWO DIFFERENT MACROS SHARING ONE NAME -- a 2-argument whole-shape check in distance.cpp, a 3-argument leading-D check in the regularisers and solve_field. That last one is the argument for doing this at all: with 17 private copies, one name silently meant two things, and nothing could have told you. WHAT MOVED WHERE include/fastfields/core/dispatch.h (new) FF_VOIDPTR, FF_CVOIDPTR, FF_CVOIDPTR_OR_NULL, FF_CANUSE32BITS, the FF_CHECK_* family, and as_weights(). In core/ because src/lib-cuda is compiled by nvcc while src/lib-cpu and src/lib use the host compiler: anything all three share has to be backend-agnostic, and core/ is the only directory that is so by contract. api/checks.h is host-only and was the wrong home. include/fastfields/api/cuda/stream.h (new) _reg_stream(). Names cudaStream_t, so it is CUDA-only by construction and belongs under api/cuda/ -- which also keeps CI's path filter honest. include/fastfields/api/checks.h (extended) is_cpu() / is_cuda() inline functions, replacing IS_CPU / IS_CUDA. Used only by the host-compiled hub. A function in ff:: is collision-safe with no prefix at all, so these two are deleted rather than renamed. All three preserved divergences keep their own name -- FF_CVOIDPTR_OR_NULL, FF_CHECK_SAME_BATCH_ND, FF_CHECK_SAME_SHAPE_N -- so the difference is visible at the call site instead of hiding in one file's prologue. NAMING A macro that lives in a .cpp cannot collide with anything; a macro in an installed header is a name taken from every downstream translation unit. So the moment one of these was hoisted it had to be prefixed, which is why de-duplicating and prefixing are one change and not two. (VOIDPTR and friends were in fact already on the public surface via api/{cpu,cuda}/ pushpull_dispatch.h, which the audit of "17 copies in .cpp files" had missed; the true counts are 19 and 17.) EVIDENCE Not asserted, proved. tools/macro-equivalence.py preprocesses every (file, macro) pair on both sides of the change and compares token streams: 121 macro expansions compared across 23 files: 119 exact, 2 balanced-paren-only, 0 MISMATCHED The two are FF_CVOIDPTR_OR_NULL, which composes FF_CVOIDPTR and so carries one extra balanced parenthesis pair; the tool checks the token streams are equal once parentheses are removed. A whole-translation-unit token diff was tried first and is the wrong instrument -- core/dispatch.h pulls in and makes as_weights() visible everywhere, which swamps the signal with additive, behaviour-free noise. The rewrite itself is tools/dedup-dispatch-helpers.py, committed so that the ~1100 renamed call sites can be re-derived and diffed rather than read. --- CLAUDE.md | 17 +- include/fastfields/api/checks.h | 25 ++ .../fastfields/api/cpu/pushpull_dispatch.h | 27 +- .../fastfields/api/cuda/pushpull_dispatch.h | 27 +- include/fastfields/api/cuda/stream.h | 33 ++ include/fastfields/core/dispatch.h | 149 ++++++++ src/lib-cpu/distance.cpp | 265 ++++++-------- src/lib-cpu/posdef.cpp | 171 ++++------ src/lib-cpu/pushpull.cpp | 74 ++-- src/lib-cpu/pushpull_backward.cpp | 144 ++++---- src/lib-cpu/reg_field.cpp | 217 +++++------- src/lib-cpu/reg_flow.cpp | 237 ++++++------- src/lib-cpu/resize.cpp | 80 ++--- src/lib-cpu/restrict.cpp | 80 ++--- src/lib-cpu/solve_field.cpp | 51 +-- src/lib-cpu/splinc.cpp | 50 ++- src/lib-cuda/distance.cpp | 265 ++++++-------- src/lib-cuda/posdef.cpp | 171 ++++------ src/lib-cuda/pushpull.cpp | 74 ++-- src/lib-cuda/pushpull_backward.cpp | 144 ++++---- src/lib-cuda/reg_field.cpp | 158 ++++----- src/lib-cuda/reg_field_rls.cpp | 111 ++---- src/lib-cuda/reg_flow.cpp | 168 ++++----- src/lib-cuda/reg_flow_rls.cpp | 112 +++--- src/lib-cuda/resize.cpp | 72 ++-- src/lib-cuda/restrict.cpp | 72 ++-- src/lib-cuda/splinc.cpp | 50 ++- src/lib/distance.cpp | 40 +-- src/lib/posdef.cpp | 52 ++- src/lib/pushpull.cpp | 52 ++- src/lib/reg_field.cpp | 100 +++--- src/lib/reg_flow.cpp | 98 +++--- src/lib/resize.cpp | 10 +- src/lib/restrict.cpp | 10 +- src/lib/solve_field.cpp | 8 +- src/lib/splinc.cpp | 11 +- tools/dedup-dispatch-helpers.py | 323 ++++++++++++++++++ tools/macro-equivalence.py | 190 +++++++++++ 38 files changed, 2025 insertions(+), 1913 deletions(-) create mode 100644 include/fastfields/api/cuda/stream.h create mode 100644 include/fastfields/core/dispatch.h create mode 100644 tools/dedup-dispatch-helpers.py create mode 100644 tools/macro-equivalence.py diff --git a/CLAUDE.md b/CLAUDE.md index e439760..c49a1d5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -59,7 +59,11 @@ impl/kernels ─ impl/cpu ─ api/cpu + src/lib-cpu ┐ tensor's device and forwards to `FF_CPU::` or `FF_CUDA::`, the latter guarded by `FF_WITH_CUDA`. - **`include/fastfields/core/`** — `dlpack.h` (vendored, do not edit), - `autocast.h` (32-bit index narrowing), `defines.h`, `cuda_switch.h`. + `autocast.h` (32-bit index narrowing), `defines.h`, `cuda_switch.h`, + `dispatch.h` (the `FF_VOIDPTR` / `FF_CANUSE32BITS` / `FF_CHECK_*` helpers + both dtype-dispatch layers share). `core/` is where anything used by + **both** `src/lib-cpu` (host compiler) and `src/lib-cuda` (nvcc) has to + live: it is the only directory that is backend-agnostic by contract. Also: `make/common.mk` (the shared makefile fragment), `tools/consolidate.sh` (the frozen consolidation rules), `ci/legacy/` and `docs/legacy/` (the six @@ -139,6 +143,17 @@ pushpull's fully-static order×bound compile is nightly `src/lib-cuda`'s `MODULES` split (`reg_field`/`reg_field_rls`, `reg_flow`/`reg_flow_rls`) caps peak memory at ~3.8 GB per module against ~6–7 GB combined. Do not recombine or "tidy" these. +- **Macros in installed headers are `FF_`-prefixed.** Anything `#define`d under + `include/` and not `#undef`'d before the end of that header is inherited by + every downstream translation unit, so it must be namespaced by prefix. Macros + private to a single `.cpp` are exempt — but the moment one is hoisted into a + header it stops being private, which is why de-duplicating a helper and + prefixing it are the same change. Two documented exemptions live in + `core/cuda_switch.h` and must keep their spelling to work at all: the + `__CUDACC_RTC__` fixed-width integer definitions (NVRTC ships no standard + library, so those *are* `` there) and the non-nvcc `__device__` / + `__host__` fallbacks. Prefer an `inline` function to a macro where one will + do — a function in `ff::` is collision-safe without any prefix. - `include/fastfields/core/dlpack.h` is vendored upstream code: do not edit it, and it is skipped by `codespell` (see `.codespellrc`). diff --git a/include/fastfields/api/checks.h b/include/fastfields/api/checks.h index 02b443b..9574b1c 100644 --- a/include/fastfields/api/checks.h +++ b/include/fastfields/api/checks.h @@ -22,6 +22,31 @@ FF_NAMESPACE_BEGIN(FF) */ inline void require_same_device(const DLTensor & /*ref*/) {} +/** + * The hub's device predicates. + * + * Every `src/lib/.cpp` carried its own `IS_CPU` / `IS_CUDA` macro pair + * -- nine identical copies, and two more unprefixed macros leaking out of the + * translation units that defined them. They are plain predicates over a POD + * field, so they are stated here as inline functions instead: a function in + * `ff::` cannot collide with a downstream identifier the way a bare `IS_CPU` + * macro can, and so it needs no prefix to be safe. Two of the 31 unprefixed + * macros are thus removed rather than renamed. + * + * `kDLCUDAHost` is pinned (page-locked) *host* memory: addressable by the CPU, + * so it belongs on the CPU side of the dispatch, not the CUDA one. + */ +inline bool is_cuda(const DLTensor & t) +{ + return t.device.device_type == DLDeviceType::kDLCUDA; +} + +inline bool is_cpu(const DLTensor & t) +{ + return t.device.device_type == DLDeviceType::kDLCPU || + t.device.device_type == DLDeviceType::kDLCUDAHost; +} + template inline void require_same_device(const DLTensor & ref, const DLTensor & t, const Rest &... rest) { if (t.device.device_type != ref.device.device_type || diff --git a/include/fastfields/api/cpu/pushpull_dispatch.h b/include/fastfields/api/cpu/pushpull_dispatch.h index d5da1f6..50ec072 100644 --- a/include/fastfields/api/cpu/pushpull_dispatch.h +++ b/include/fastfields/api/cpu/pushpull_dispatch.h @@ -17,6 +17,7 @@ #include #include #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -25,36 +26,10 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - // reduce/accumulation type used by the sampling kernels. Match jitfields // (float64) for CPU accuracy. typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_BATCH(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same batch shape"); - /*********************************************************************** * DISPATCH * ***********************************************************************/ diff --git a/include/fastfields/api/cuda/pushpull_dispatch.h b/include/fastfields/api/cuda/pushpull_dispatch.h index 32bbf1e..626d0e1 100644 --- a/include/fastfields/api/cuda/pushpull_dispatch.h +++ b/include/fastfields/api/cuda/pushpull_dispatch.h @@ -17,6 +17,7 @@ #include #include #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -25,36 +26,10 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - // reduce/accumulation type used by the sampling kernels. Match jitfields // (float64) for CPU accuracy. typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_BATCH(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same batch shape"); - /*********************************************************************** * DISPATCH * ***********************************************************************/ diff --git a/include/fastfields/api/cuda/stream.h b/include/fastfields/api/cuda/stream.h new file mode 100644 index 0000000..5a0f086 --- /dev/null +++ b/include/fastfields/api/cuda/stream.h @@ -0,0 +1,33 @@ +#pragma once +#ifndef FF_CUDA_STREAM +#define FF_CUDA_STREAM + +/** + * The public ABI carries a CUDA stream as an `intptr_t` (no CUDA types leak + * into the exported signatures); the cuda-impl launchers take a real + * `cudaStream_t`. This is the one-line conversion between them, which had been + * copied verbatim into all four regulariser dispatch sources + * (`reg_field.cpp`, `reg_field_rls.cpp`, `reg_flow.cpp`, `reg_flow_rls.cpp`). + * + * CUDA-only by construction -- it names `cudaStream_t` -- so it lives under + * `api/cuda/` rather than in `core/`, and only `src/lib-cuda` includes it. + * `pushpull` has its own `_pp_stream` in the cuda-impl layer; the two are left + * separate deliberately, as that one sits a layer below. + */ + +#include +#include "fastfields/core/cuda_switch.h" + +FF_NAMESPACE_BEGIN(FF) +FF_NAMESPACE_BEGIN(FF_DEVICE) + +// intptr_t -> cudaStream_t (0 == the default stream). +static inline cudaStream_t _reg_stream(intptr_t stream) +{ + return reinterpret_cast(static_cast(stream)); +} + +FF_NAMESPACE_END(FF_DEVICE) +FF_NAMESPACE_END(FF) + +#endif // FF_CUDA_STREAM diff --git a/include/fastfields/core/dispatch.h b/include/fastfields/core/dispatch.h new file mode 100644 index 0000000..c2a1a74 --- /dev/null +++ b/include/fastfields/core/dispatch.h @@ -0,0 +1,149 @@ +#pragma once +#ifndef FF_CORE_DISPATCH +#define FF_CORE_DISPATCH + +/** + * The helpers shared by both dtype-dispatch layers (`src/lib-cpu` and + * `src/lib-cuda`): unpack the public pointer ABI's arguments into what the + * templated impl layer wants, and reject the combinations it cannot express. + * + * Every `src/lib-cpu/.cpp` and `src/lib-cuda/.cpp` did this + * the same way -- add `byte_offset` to `data`, decide whether the shape and + * stride arrays fit in 32 bits, validate the operands -- and each one carried + * its own copy of the macros. There were 19 copies of `VOIDPTR`, 19 of + * `CANUSE32BITS`, 19 of `CHECK_NO_LANES` and 17 of `CHECK_SAME`, and they had + * already drifted apart; see the three preserved divergences below. + * + * Why `core/` and not `api/` + * -------------------------------------------------------------------------- + * `src/lib-cuda` is compiled by **nvcc** while `src/lib-cpu` and `src/lib` use + * the host compiler, so anything all of them share has to be backend-agnostic. + * That is exactly what `include/fastfields/core/` is for. `api/checks.h` is + * the hub's (host-only) validation header and is the wrong home for macros + * nvcc must also digest. Note the CI consequence, which is intended: a change + * under `core/` triggers every job, CUDA included. + * + * Usage notes + * -------------------------------------------------------------------------- + * * `FF_CANUSE32BITS` calls `canUse32BitIndexMath` unqualified, so it must be + * expanded from inside `ff::cpu` / `ff::cuda` (i.e. `ff::FF_DEVICE`, where + * `core/autocast.h` declares it). Every dispatch source already is. It is + * left unqualified on purpose: qualifying it would change name lookup, and + * this header's contract is that the macros expand token-for-token to what + * the 19 local copies expanded to. + * * The `FF_CHECK_*` macros expand to bare `if` / `for` statements, NOT to a + * `do { ... } while (0)`. That is deliberate: it is what the copies did, and + * wrapping them would silently change which statements a brace-less + * `if (cond) FF_CHECK_...(...);` guards. They are therefore not safe as the + * body of an unbraced `if` / `else`. Tightening that is a behaviour change + * and belongs in its own commit, measured against the CPU suite. + */ + +#include // size_t +#include // int32_t (the shape loops), int64_t +#include // std::invalid_argument +#include // as_weights +#include "fastfields/core/dlpack.h" +#include "fastfields/core/defines.h" // FF_NAMESPACE_* +#include "fastfields/core/autocast.h" // canUse32BitIndexMath + +/*********************************************************************** + * POINTER MARSHALLING * + ***********************************************************************/ + +// DLPack keeps the element offset out of `data`, so every call into the impl +// layer has to fold it back in. +#define FF_VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) +#define FF_CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) + +// PRESERVED DIVERGENCE 1/3. `posdef.cpp` -- and only posdef.cpp -- guarded +// its CVOIDPTR against a null `data`: an absent optional operand is passed as +// a descriptor with `data == nullptr`, and offsetting a null pointer is +// undefined behaviour even when the result is never dereferenced. Given its +// own name rather than folded into FF_CVOIDPTR, so the difference is visible +// at the call site instead of hiding in one file's private prologue. +#define FF_CVOIDPTR_OR_NULL(x) (x.data ? FF_CVOIDPTR(x) : nullptr) + +// Can this tensor's shape/stride arithmetic be narrowed to 32-bit offsets? +// See core/autocast.h. +#define FF_CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) + +/*********************************************************************** + * CHECKS * + ***********************************************************************/ + +#define FF_CHECK_NO_LANES(tensor) \ + if (tensor.dtype.lanes > 1) \ + throw std::invalid_argument("Only scalar data types are supported"); + +#define FF_CHECK_SAME(X, Y, msg) \ + if (X != Y) throw std::invalid_argument(msg); + +#define FF_CHECK_SAME_DTYPE(X, Y) \ + if ((X.dtype.code != Y.dtype.code) || \ + (X.dtype.bits != Y.dtype.bits) || \ + (X.dtype.lanes != Y.dtype.lanes)) \ + throw std::invalid_argument("Tensors do not have the same data type"); + +// Agreement on the leading `D` (batch) dimensions. +#define FF_CHECK_SAME_BATCH(X, Y, D) \ + for (int32_t d=0; d < D; ++d) \ + if (X.shape[d] != Y.shape[d]) \ + throw std::invalid_argument("Tensors do not have the same batch shape"); + +// PRESERVED DIVERGENCE 2/3. As above, but first rejects tensors with fewer +// than `D` dimensions rather than reading past the end of `shape`. +// `distance.cpp` and `posdef.cpp` are the two that need it -- they derive `D` +// from an operand's own `ndim` -- and were the only two that had it. +#define FF_CHECK_SAME_BATCH_ND(X, Y, D) \ + if (X.ndim < D || Y.ndim < D) \ + throw std::invalid_argument("Number of dimensions does not match"); \ + FF_CHECK_SAME_BATCH(X, Y, D) + +// PRESERVED DIVERGENCE 3/3. `CHECK_SAME_SHAPE` was two *different macros +// sharing one name*: a 3-argument leading-D check in the regularisers and +// solve_field (this one), and a 2-argument whole-shape check in distance.cpp +// (below). Merging them under one name would have silently changed one set of +// call sites, which is the sharpest illustration of why 17 private copies of a +// macro is a hazard rather than a tidiness complaint. +#define FF_CHECK_SAME_SHAPE_N(X, Y, D) \ + for (int32_t d=0; d < D; ++d) \ + if (X.shape[d] != Y.shape[d]) \ + throw std::invalid_argument("Tensors do not have the same shape"); + +// Agreement on the WHOLE shape, rank included. +#define FF_CHECK_SAME_SHAPE(X, Y) \ + if (X.ndim != Y.ndim) \ + throw std::invalid_argument("Tensors do not have the same number of dimensions"); \ + FF_CHECK_SAME_BATCH_ND(X, Y, X.ndim) + +/*********************************************************************** + * NON-TENSOR ARGUMENT MARSHALLING * + ***********************************************************************/ + +FF_NAMESPACE_BEGIN(FF) + +/** + * Build a length-`nc` penalty-weight vector from the ABI's `const double *`. + * + * The regulariser entry points take each energy term's weight as a raw pointer + * that is either an `nc`-long array or null ("this term is off"); the impl + * layer takes a filled vector. Was copied verbatim into + * `src/lib-cpu/reg_field.cpp` and `src/lib-cuda/reg_field{,_rls}.cpp`. + * + * Returns `std::vector` rather than the dispatch sources' local + * `reduce_t` typedef, which is `double` in every one of them. If a backend + * ever changes its accumulation type, the assignment at the call site stops + * compiling -- which is the failure you want, rather than a silent narrowing + * inside a shared header. + */ +inline std::vector as_weights(const double * w, int64_t nc) +{ + std::vector v(static_cast(nc), 0.0); + if (w) for (int64_t c = 0; c < nc; ++c) v[static_cast(c)] = w[c]; + return v; +} + +FF_NAMESPACE_END(FF) + +#endif // FF_CORE_DISPATCH diff --git a/src/lib-cpu/distance.cpp b/src/lib-cpu/distance.cpp index 6aea232..56c4421 100644 --- a/src/lib-cpu/distance.cpp +++ b/src/lib-cpu/distance.cpp @@ -2,6 +2,7 @@ #include #include "fastfields/api/cpu/distance.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -13,76 +14,32 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_BATCH(X, Y, D) \ - if (X.ndim < D || Y.ndim < D) \ - throw std::invalid_argument( \ - "Number of dimensions does not match" \ - ); \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument( \ - "Tensors do not have the same batch shape" \ - ); \ - -#define CHECK_SAME_SHAPE(X, Y) \ - if (X.ndim != Y.ndim) \ - throw std::invalid_argument( \ - "Tensors do not have the same number of dimensions" \ - ); \ - CHECK_SAME_BATCH(X, Y, X.ndim) - -#define CHECK_SAME_DTYPE(X, Y) \ - if ( \ - (X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes) \ - ) \ - throw std::invalid_argument( \ - "Tensors do not have the same data type" \ - ); - /*********************************************************************** * EUCLIDEAN * ***********************************************************************/ -#define DISPATCH_DT(func, args...) \ -{ \ - const bool use_32bits = CANUSE32BITS(inp_out); \ - const auto code = static_cast(inp_out.dtype.code); \ - switch (code) { \ - case kDLFloat: switch (inp_out.dtype.bits) { \ - case 32: return ( \ - use_32bits \ - ? func(args) \ - : func(args) \ - ); \ - case 64: return ( \ - use_32bits \ - ? func(args) \ - : func(args) \ - ); \ - default: break; \ - }; \ - default: throw std::invalid_argument( \ - "only floating point data types are supported" \ - ); \ - }; \ +#define DISPATCH_DT(func, args...) \ +{ \ + const bool use_32bits = FF_CANUSE32BITS(inp_out); \ + const auto code = static_cast(inp_out.dtype.code); \ + switch (code) { \ + case kDLFloat: switch (inp_out.dtype.bits) { \ + case 32: return ( \ + use_32bits \ + ? func(args) \ + : func(args) \ + ); \ + case 64: return ( \ + use_32bits \ + ? func(args) \ + : func(args) \ + ); \ + default: break; \ + }; \ + default: throw std::invalid_argument( \ + "only floating point data types are supported" \ + ); \ + }; \ } namespace { @@ -115,11 +72,11 @@ void dt_euclidean( ContiguousStrides _io(inp_out_); DLTensor & inp_out = _io.t; - CHECK_NO_LANES(inp_out) + FF_CHECK_NO_LANES(inp_out) DISPATCH_DT( _dt_euclidean, inp_out.ndim, - VOIDPTR(inp_out), + FF_VOIDPTR(inp_out), voxel_spacing, inp_out.shape, inp_out.strides @@ -156,11 +113,11 @@ void dt_l1( ContiguousStrides _io(inp_out_); DLTensor & inp_out = _io.t; - CHECK_NO_LANES(inp_out) + FF_CHECK_NO_LANES(inp_out) DISPATCH_DT( _dt_l1, inp_out.ndim, - VOIDPTR(inp_out), + FF_VOIDPTR(inp_out), voxel_spacing, inp_out.shape, inp_out.strides @@ -284,39 +241,39 @@ void dt_spline_table( const DLTensor & times = _ti.t; const bool use_32bits = ( - CANUSE32BITS(time) && - CANUSE32BITS(dist) && - CANUSE32BITS(loc) && - CANUSE32BITS(coeff) && - CANUSE32BITS(times) + FF_CANUSE32BITS(time) && + FF_CANUSE32BITS(dist) && + FF_CANUSE32BITS(loc) && + FF_CANUSE32BITS(coeff) && + FF_CANUSE32BITS(times) ); const int32_t ndim = loc.shape[loc.ndim-1]; const int32_t nbatch = loc.ndim - 1; - CHECK_NO_LANES (time) - CHECK_SAME_DTYPE(time, dist) - CHECK_SAME_DTYPE(time, loc) - CHECK_SAME_DTYPE(time, coeff) - CHECK_SAME_DTYPE(time, times) - CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") - CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_NO_LANES (time) + FF_CHECK_SAME_DTYPE(time, dist) + FF_CHECK_SAME_DTYPE(time, loc) + FF_CHECK_SAME_DTYPE(time, coeff) + FF_CHECK_SAME_DTYPE(time, times) + FF_CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") // coeff is (*batch, npoints, ndim) -> nbatch+2 dims // times is (*batch, ntimes) -> nbatch+1 dims - CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") - CHECK_SAME (times.ndim, nbatch+1, "Number of times dimensions does not match") - CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") - CHECK_SAME_BATCH(loc, time, nbatch) - CHECK_SAME_BATCH(loc, dist, nbatch) - CHECK_SAME_BATCH(loc, coeff, nbatch) - CHECK_SAME_BATCH(loc, times, nbatch) + FF_CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") + FF_CHECK_SAME (times.ndim, nbatch+1, "Number of times dimensions does not match") + FF_CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") + FF_CHECK_SAME_BATCH_ND(loc, time, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, dist, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, coeff, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, times, nbatch) DISPATCH_SPLINE( _dt_spline_table, nbatch, // nbatch - VOIDPTR(time), // time - VOIDPTR(dist), // dist - VOIDPTR(loc), // loc - VOIDPTR(coeff), // coeff - VOIDPTR(times), // times + FF_VOIDPTR(time), // time + FF_VOIDPTR(dist), // dist + FF_VOIDPTR(loc), // loc + FF_VOIDPTR(coeff), // coeff + FF_VOIDPTR(times), // times times.shape[times.ndim-1], // ntimes coeff.shape, // size (coeff shape: *batch, npoints, ndim) time.strides, // int64_time @@ -406,33 +363,33 @@ void dt_spline_brent( const DLTensor & coeff = _co.t; const bool use_32bits = ( - CANUSE32BITS(time) && - CANUSE32BITS(dist) && - CANUSE32BITS(loc) && - CANUSE32BITS(coeff) + FF_CANUSE32BITS(time) && + FF_CANUSE32BITS(dist) && + FF_CANUSE32BITS(loc) && + FF_CANUSE32BITS(coeff) ); const int32_t ndim = loc.shape[loc.ndim-1]; const int32_t nbatch = loc.ndim - 1; - CHECK_NO_LANES (time) - CHECK_SAME_DTYPE(time, dist) - CHECK_SAME_DTYPE(time, loc) - CHECK_SAME_DTYPE(time, coeff) - CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") - CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_NO_LANES (time) + FF_CHECK_SAME_DTYPE(time, dist) + FF_CHECK_SAME_DTYPE(time, loc) + FF_CHECK_SAME_DTYPE(time, coeff) + FF_CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") // coeff is (*batch, npoints, ndim) -> nbatch+2 dims - CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") - CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") - CHECK_SAME_BATCH(loc, time, nbatch) - CHECK_SAME_BATCH(loc, dist, nbatch) - CHECK_SAME_BATCH(loc, coeff, nbatch) + FF_CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") + FF_CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") + FF_CHECK_SAME_BATCH_ND(loc, time, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, dist, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, coeff, nbatch) DISPATCH_SPLINE( _dt_spline_brent, nbatch, // nbatch - VOIDPTR(time), // time - VOIDPTR(dist), // dist - VOIDPTR(loc), // loc - VOIDPTR(coeff), // coeff + FF_VOIDPTR(time), // time + FF_VOIDPTR(dist), // dist + FF_VOIDPTR(loc), // loc + FF_VOIDPTR(coeff), // coeff coeff.shape, // size (coeff shape: *batch, npoints, ndim) time.strides, // int64_time dist.strides, // stride_dist @@ -520,33 +477,33 @@ void dt_spline_gaussnewton( const DLTensor & coeff = _co.t; const bool use_32bits = ( - CANUSE32BITS(time) && - CANUSE32BITS(dist) && - CANUSE32BITS(loc) && - CANUSE32BITS(coeff) + FF_CANUSE32BITS(time) && + FF_CANUSE32BITS(dist) && + FF_CANUSE32BITS(loc) && + FF_CANUSE32BITS(coeff) ); const int32_t ndim = loc.shape[loc.ndim-1]; const int32_t nbatch = loc.ndim - 1; - CHECK_NO_LANES (time) - CHECK_SAME_DTYPE(time, dist) - CHECK_SAME_DTYPE(time, loc) - CHECK_SAME_DTYPE(time, coeff) - CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") - CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_NO_LANES (time) + FF_CHECK_SAME_DTYPE(time, dist) + FF_CHECK_SAME_DTYPE(time, loc) + FF_CHECK_SAME_DTYPE(time, coeff) + FF_CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") // coeff is (*batch, npoints, ndim) -> nbatch+2 dims - CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") - CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") - CHECK_SAME_BATCH(loc, time, nbatch) - CHECK_SAME_BATCH(loc, dist, nbatch) - CHECK_SAME_BATCH(loc, coeff, nbatch) + FF_CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") + FF_CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") + FF_CHECK_SAME_BATCH_ND(loc, time, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, dist, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, coeff, nbatch) DISPATCH_SPLINE( _dt_spline_gaussnewton, nbatch, // nbatch - VOIDPTR(time), // time - VOIDPTR(dist), // dist - VOIDPTR(loc), // loc - VOIDPTR(coeff), // coeff + FF_VOIDPTR(time), // time + FF_VOIDPTR(dist), // dist + FF_VOIDPTR(loc), // loc + FF_VOIDPTR(coeff), // coeff coeff.shape, // size (coeff shape: *batch, npoints, ndim) time.strides, // int64_time dist.strides, // stride_dist @@ -701,32 +658,32 @@ void dt_mesh( const DLTensor & faces = _fa.t; bool use_32bits = ( - CANUSE32BITS(dist) && - CANUSE32BITS(loc) && - CANUSE32BITS(vertices) && - CANUSE32BITS(faces) + FF_CANUSE32BITS(dist) && + FF_CANUSE32BITS(loc) && + FF_CANUSE32BITS(vertices) && + FF_CANUSE32BITS(faces) ); const int32_t ndim = loc.shape[loc.ndim-1]; const int32_t nbatch = loc.ndim - 1; - CHECK_NO_LANES (dist) - CHECK_SAME_DTYPE(dist, loc) - CHECK_SAME_DTYPE(dist, vertices) - CHECK_SAME ( dist.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_NO_LANES (dist) + FF_CHECK_SAME_DTYPE(dist, loc) + FF_CHECK_SAME_DTYPE(dist, vertices) + FF_CHECK_SAME ( dist.ndim, nbatch, "Number of batch dimensions does not match") // vertices (N, D) and faces (M, D) describe a single shared mesh and are // always 2D; their leading axis is the vertex/face count, independent of // loc's point batch. Only `loc` and the per-point outputs share a batch. - CHECK_SAME ( vertices.ndim, 2, "Vertices must be a (N, D) tensor") - CHECK_SAME ( faces.ndim, 2, "Faces must be a (M, D) tensor") - CHECK_SAME_BATCH(loc, dist, nbatch) - CHECK_SAME (vertices.shape[vertices.ndim-1], ndim, "Dimensionality of the vertices and location does not match") - CHECK_SAME (faces.shape[faces.ndim-1], ndim, "Dimensionality of the vertices and faces does not match") + FF_CHECK_SAME ( vertices.ndim, 2, "Vertices must be a (N, D) tensor") + FF_CHECK_SAME ( faces.ndim, 2, "Faces must be a (M, D) tensor") + FF_CHECK_SAME_BATCH_ND(loc, dist, nbatch) + FF_CHECK_SAME (vertices.shape[vertices.ndim-1], ndim, "Dimensionality of the vertices and location does not match") + FF_CHECK_SAME (faces.shape[faces.ndim-1], ndim, "Dimensionality of the vertices and faces does not match") if (nearest_vertex.data) { - CHECK_SAME_DTYPE(faces, nearest_vertex) - CHECK_SAME (nearest_vertex.ndim, nbatch, "Number of batch dimensions does not match") - CHECK_SAME_BATCH(loc, nearest_vertex, nbatch) - use_32bits &= CANUSE32BITS(nearest_vertex); + FF_CHECK_SAME_DTYPE(faces, nearest_vertex) + FF_CHECK_SAME (nearest_vertex.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_SAME_BATCH_ND(loc, nearest_vertex, nbatch) + use_32bits &= FF_CANUSE32BITS(nearest_vertex); } // `nearest_vertex` is optional and signalled *only* by `data == nullptr` @@ -736,7 +693,7 @@ void dt_mesh( // `_dt_mesh` keys the optional path off a null `stride_nearest`, and // forwarding a placeholder's raw fields would hand it a wild pointer to // dereference instead. - void * const nearest_data = nearest_vertex.data ? VOIDPTR(nearest_vertex) + void * const nearest_data = nearest_vertex.data ? FF_VOIDPTR(nearest_vertex) : nullptr; int64_t * const nearest_strides = nearest_vertex.data ? nearest_vertex.strides : nullptr; @@ -744,11 +701,11 @@ void dt_mesh( DISPATCH_MESH( _dt_mesh, nbatch, // nbatch - VOIDPTR(dist), // data + FF_VOIDPTR(dist), // data nearest_data, // nearest_vertex - VOIDPTR(loc), // coord - VOIDPTR(vertices), // vertices - VOIDPTR(faces), // faces + FF_VOIDPTR(loc), // coord + FF_VOIDPTR(vertices), // vertices + FF_VOIDPTR(faces), // faces loc.shape, // size faces.shape[0], // nb_faces (M = faces.shape[0]) vertices.shape[0], // nb_vertices (N = vertices.shape[0]) diff --git a/src/lib-cpu/posdef.cpp b/src/lib-cpu/posdef.cpp index 2d61bcd..40d9bf6 100644 --- a/src/lib-cpu/posdef.cpp +++ b/src/lib-cpu/posdef.cpp @@ -3,6 +3,7 @@ #include #include "fastfields/api/cpu/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -11,48 +12,10 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (x.data ? static_cast(static_cast(x.data) + x.byte_offset) : nullptr) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - // reduce/accumulation type used by the compact-symmetric kernels. // jitfields defaults to float64; we do the same for CPU accuracy. typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_BATCH(X, Y, D) \ - if (X.ndim < D || Y.ndim < D) \ - throw std::invalid_argument( \ - "Number of dimensions does not match" \ - ); \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument( \ - "Tensors do not have the same batch shape" \ - ); \ - -#define CHECK_SAME_DTYPE(X, Y) \ - if ( \ - (X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes) \ - ) \ - throw std::invalid_argument( \ - "Tensors do not have the same data type" \ - ); - // C such that C*(C+1)/2 == CC (the compact-symmetric length). static inline int64_t channels_from_packed(int64_t CC) { @@ -165,23 +128,23 @@ void sym_matvec( const DLTensor & hessian = _hes.t; const DLTensor & inp = _inp.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian) && CANUSE32BITS(inp); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian) && FF_CANUSE32BITS(inp); const int32_t nbatch = out.ndim - 1; const int64_t nchannel = out.shape[out.ndim-1]; const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(out, inp, nbatch) - CHECK_SAME_BATCH(out, hessian, nbatch) + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(out, inp, nbatch) + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) DISPATCH_SYM_C( _sym_matvec, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(hessian), CVOIDPTR(inp), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(hessian), FF_CVOIDPTR_OR_NULL(inp), out.shape, out.strides, hessian.strides, inp.strides ) } @@ -247,23 +210,23 @@ void sym_addmatvec_( const DLTensor & hessian = _hes.t; const DLTensor & inp = _inp.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian) && CANUSE32BITS(inp); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian) && FF_CANUSE32BITS(inp); const int32_t nbatch = out.ndim - 1; const int64_t nchannel = out.shape[out.ndim-1]; const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(out, inp, nbatch) - CHECK_SAME_BATCH(out, hessian, nbatch) + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(out, inp, nbatch) + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) DISPATCH_SYM_C( _sym_addmatvec_, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(hessian), CVOIDPTR(inp), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(hessian), FF_CVOIDPTR_OR_NULL(inp), out.shape, out.strides, hessian.strides, inp.strides ) } @@ -281,23 +244,23 @@ void sym_submatvec_( const DLTensor & hessian = _hes.t; const DLTensor & inp = _inp.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian) && CANUSE32BITS(inp); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian) && FF_CANUSE32BITS(inp); const int32_t nbatch = out.ndim - 1; const int64_t nchannel = out.shape[out.ndim-1]; const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(out, inp, nbatch) - CHECK_SAME_BATCH(out, hessian, nbatch) + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(out, inp, nbatch) + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) DISPATCH_SYM_C( _sym_submatvec_, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(hessian), CVOIDPTR(inp), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(hessian), FF_CVOIDPTR_OR_NULL(inp), out.shape, out.strides, hessian.strides, inp.strides ) } @@ -344,23 +307,23 @@ void sym_matvec_backward( const DLTensor & grd = _grd.t; const DLTensor & inp = _inp.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(grd) && CANUSE32BITS(inp); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(grd) && FF_CANUSE32BITS(inp); const int32_t nbatch = grd.ndim - 1; const int64_t nchannel = grd.shape[grd.ndim-1]; const auto code = static_cast(grd.dtype.code); const auto bits = grd.dtype.bits; - CHECK_NO_LANES (grd) - CHECK_SAME_DTYPE(grd, out) - CHECK_SAME_DTYPE(grd, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and grad channel counts differ") - CHECK_SAME (out.shape[out.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(grd, out, nbatch) - CHECK_SAME_BATCH(grd, inp, nbatch) + FF_CHECK_NO_LANES (grd) + FF_CHECK_SAME_DTYPE(grd, out) + FF_CHECK_SAME_DTYPE(grd, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and grad channel counts differ") + FF_CHECK_SAME (out.shape[out.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(grd, out, nbatch) + FF_CHECK_SAME_BATCH_ND(grd, inp, nbatch) DISPATCH_SYM_C( _sym_matvec_backward, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(grd), CVOIDPTR(inp), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(grd), FF_CVOIDPTR_OR_NULL(inp), grd.shape, out.strides, grd.strides, inp.strides ) } @@ -416,26 +379,26 @@ void sym_solve( const DLTensor & weight = _wgt.t; const bool has_wgt = (weight.data != nullptr); - bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian) && CANUSE32BITS(inp); - if (has_wgt) use_32bits = use_32bits && CANUSE32BITS(weight); + bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian) && FF_CANUSE32BITS(inp); + if (has_wgt) use_32bits = use_32bits && FF_CANUSE32BITS(weight); const int32_t nbatch = out.ndim - 1; const int64_t nchannel = out.shape[out.ndim-1]; const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(out, inp, nbatch) - CHECK_SAME_BATCH(out, hessian, nbatch) - if (has_wgt) { CHECK_SAME_DTYPE(out, weight) CHECK_SAME_BATCH(out, weight, nbatch) } + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(out, inp, nbatch) + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) + if (has_wgt) { FF_CHECK_SAME_DTYPE(out, weight) FF_CHECK_SAME_BATCH_ND(out, weight, nbatch) } DISPATCH_SYM( _sym_solve, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(hessian), - has_wgt ? VOIDPTR(weight) : nullptr, + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(inp), FF_CVOIDPTR_OR_NULL(hessian), + has_wgt ? FF_VOIDPTR(weight) : nullptr, out.shape, out.strides, inp.strides, hessian.strides, has_wgt ? weight.strides : nullptr ) @@ -482,23 +445,23 @@ void sym_solve_( const DLTensor & weight = _wgt.t; const bool has_wgt = (weight.data != nullptr); - bool use_32bits = CANUSE32BITS(inp_out) && CANUSE32BITS(hessian); - if (has_wgt) use_32bits = use_32bits && CANUSE32BITS(weight); + bool use_32bits = FF_CANUSE32BITS(inp_out) && FF_CANUSE32BITS(hessian); + if (has_wgt) use_32bits = use_32bits && FF_CANUSE32BITS(weight); const int32_t nbatch = inp_out.ndim - 1; const int64_t nchannel = inp_out.shape[inp_out.ndim-1]; const auto code = static_cast(inp_out.dtype.code); const auto bits = inp_out.dtype.bits; - CHECK_NO_LANES (inp_out) - CHECK_SAME_DTYPE(inp_out, hessian) - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(inp_out, hessian, nbatch) - if (has_wgt) { CHECK_SAME_DTYPE(inp_out, weight) CHECK_SAME_BATCH(inp_out, weight, nbatch) } + FF_CHECK_NO_LANES (inp_out) + FF_CHECK_SAME_DTYPE(inp_out, hessian) + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(inp_out, hessian, nbatch) + if (has_wgt) { FF_CHECK_SAME_DTYPE(inp_out, weight) FF_CHECK_SAME_BATCH_ND(inp_out, weight, nbatch) } DISPATCH_SYM( _sym_solve_, nbatch, nchannel, - VOIDPTR(inp_out), CVOIDPTR(hessian), - has_wgt ? VOIDPTR(weight) : nullptr, + FF_VOIDPTR(inp_out), FF_CVOIDPTR_OR_NULL(hessian), + has_wgt ? FF_VOIDPTR(weight) : nullptr, inp_out.shape, inp_out.strides, hessian.strides, has_wgt ? weight.strides : nullptr ) @@ -540,21 +503,21 @@ void sym_invert( DLTensor & out = _out.t; const DLTensor & hessian = _hes.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian); const int32_t nbatch = out.ndim - 1; const int64_t CC = hessian.shape[hessian.ndim-1]; const int64_t nchannel = channels_from_packed(CC); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME (out.shape[out.ndim-1], CC, "Output and matrix must share the compact layout") - CHECK_SAME_BATCH(out, hessian, nbatch) + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME (out.shape[out.ndim-1], CC, "Output and matrix must share the compact layout") + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) DISPATCH_SYM( _sym_invert, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(hessian), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(hessian), out.shape, out.strides, hessian.strides ) } @@ -586,18 +549,18 @@ void sym_invert_( ContiguousStrides _hes(hessian_); DLTensor & hessian = _hes.t; - const bool use_32bits = CANUSE32BITS(hessian); + const bool use_32bits = FF_CANUSE32BITS(hessian); const int32_t nbatch = hessian.ndim - 1; const int64_t CC = hessian.shape[hessian.ndim-1]; const int64_t nchannel = channels_from_packed(CC); const auto code = static_cast(hessian.dtype.code); const auto bits = hessian.dtype.bits; - CHECK_NO_LANES(hessian) + FF_CHECK_NO_LANES(hessian) DISPATCH_SYM( _sym_invert_, nbatch, nchannel, - VOIDPTR(hessian), + FF_VOIDPTR(hessian), hessian.shape, hessian.strides ) } diff --git a/src/lib-cpu/pushpull.cpp b/src/lib-cpu/pushpull.cpp index e72aa12..f3b267c 100644 --- a/src/lib-cpu/pushpull.cpp +++ b/src/lib-cpu/pushpull.cpp @@ -1,6 +1,6 @@ #include "fastfields/api/cpu/pushpull.h" #include -// VOIDPTR / CHECK_* / DISPATCH_PP and the reduce_t typedef, shared with +// FF_VOIDPTR / CHECK_* / DISPATCH_PP and the reduce_t typedef, shared with // pushpull_backward.cpp so the two translation units cannot drift apart on // which (order, bound) pairs are statically instantiated. #include "fastfields/api/cpu/pushpull_dispatch.h" @@ -152,18 +152,18 @@ void pull( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; // nbatch + ndim + 1 - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); - CHECK_SAME(out.shape[out.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(out, grid, nbatch) - CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME(out.shape[out.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(out, grid, nbatch) + FF_CHECK_SAME_BATCH(inp, grid, nbatch) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -177,7 +177,7 @@ void pull( DISPATCH_PP(_pull, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_CVOIDPTR(inp), FF_CVOIDPTR(grid), grid.shape, inp.shape, out.strides, inp.strides, grid.strides) } @@ -205,18 +205,18 @@ void push( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); - CHECK_SAME(out.shape[out.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(out, grid, nbatch) - CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME(out.shape[out.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(out, grid, nbatch) + FF_CHECK_SAME_BATCH(inp, grid, nbatch) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -231,7 +231,7 @@ void push( DISPATCH_PP(_push, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_CVOIDPTR(inp), FF_CVOIDPTR(grid), grid.shape, out.shape, out.strides, inp.strides, grid.strides) } @@ -257,14 +257,14 @@ void count( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); - CHECK_SAME_BATCH(out, grid, nbatch) + FF_CHECK_SAME_BATCH(out, grid, nbatch) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -278,7 +278,7 @@ void count( DISPATCH_PP(_count, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_CVOIDPTR(grid), grid.shape, out.shape, out.strides, grid.strides) } @@ -307,18 +307,18 @@ void grad( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; // grid/inp rank; out rank == n1 + 1 - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") - CHECK_SAME(out.ndim, grid.ndim + 1, "grad output must have an extra trailing axis") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_SAME(out.ndim, grid.ndim + 1, "grad output must have an extra trailing axis") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); - CHECK_SAME(out.shape[out.ndim-1], ndim, "grad output trailing axis must equal ndim") - CHECK_SAME(out.shape[out.ndim-2], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME(out.shape[out.ndim-1], ndim, "grad output trailing axis must equal ndim") + FF_CHECK_SAME(out.shape[out.ndim-2], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(inp, grid, nbatch) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -332,7 +332,7 @@ void grad( DISPATCH_PP(_grad, bvec, svec, static_cast(nbatch), n1, ex, abs, - VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_CVOIDPTR(inp), FF_CVOIDPTR(grid), grid.shape, inp.shape, out.strides, inp.strides, grid.strides) } diff --git a/src/lib-cpu/pushpull_backward.cpp b/src/lib-cpu/pushpull_backward.cpp index 370bb1c..15cf070 100644 --- a/src/lib-cpu/pushpull_backward.cpp +++ b/src/lib-cpu/pushpull_backward.cpp @@ -199,29 +199,29 @@ void pull_backward( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, gout) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, ginp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") - CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, gout) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, ginp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); // `out` mirrors `inp` (the field), `gout` mirrors `grid`. for (int32_t d = 0; d < out.ndim; ++d) - CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") + FF_CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") for (int32_t d = 0; d < gout.ndim; ++d) - CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") - CHECK_SAME(ginp.shape[ginp.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(inp, grid, nbatch) - CHECK_SAME_BATCH(ginp, grid, nbatch) - - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(gout) - && CANUSE32BITS(inp) && CANUSE32BITS(ginp) - && CANUSE32BITS(grid); + FF_CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") + FF_CHECK_SAME(ginp.shape[ginp.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME_BATCH(ginp, grid, nbatch) + + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(gout) + && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(ginp) + && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -233,8 +233,8 @@ void pull_backward( DISPATCH_PP(_pull_backward, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), VOIDPTR(gout), - CVOIDPTR(inp), CVOIDPTR(ginp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_VOIDPTR(gout), + FF_CVOIDPTR(inp), FF_CVOIDPTR(ginp), FF_CVOIDPTR(grid), grid.shape, inp.shape, out.strides, gout.strides, inp.strides, ginp.strides, grid.strides) } @@ -265,29 +265,29 @@ void push_backward( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, gout) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, ginp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") - CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, gout) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, ginp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); // Here both `out` and `inp` are grid-shaped; `ginp` is the field. for (int32_t d = 0; d < out.ndim; ++d) - CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") + FF_CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") for (int32_t d = 0; d < gout.ndim; ++d) - CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") - CHECK_SAME(ginp.shape[ginp.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(inp, grid, nbatch) - CHECK_SAME_BATCH(ginp, grid, nbatch) - - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(gout) - && CANUSE32BITS(inp) && CANUSE32BITS(ginp) - && CANUSE32BITS(grid); + FF_CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") + FF_CHECK_SAME(ginp.shape[ginp.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME_BATCH(ginp, grid, nbatch) + + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(gout) + && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(ginp) + && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -300,8 +300,8 @@ void push_backward( DISPATCH_PP(_push_backward, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), VOIDPTR(gout), - CVOIDPTR(inp), CVOIDPTR(ginp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_VOIDPTR(gout), + FF_CVOIDPTR(inp), FF_CVOIDPTR(ginp), FF_CVOIDPTR(grid), grid.shape, ginp.shape, out.strides, gout.strides, inp.strides, ginp.strides, grid.strides) } @@ -328,20 +328,20 @@ void count_backward( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (gout) - CHECK_SAME_DTYPE(gout, ginp) - CHECK_SAME_DTYPE(gout, grid) - CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") - CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") + FF_CHECK_NO_LANES (gout) + FF_CHECK_SAME_DTYPE(gout, ginp) + FF_CHECK_SAME_DTYPE(gout, grid) + FF_CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") + FF_CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); for (int32_t d = 0; d < gout.ndim; ++d) - CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") - CHECK_SAME(ginp.shape[ginp.ndim-1], 1, "count gradient must have a single channel") - CHECK_SAME_BATCH(ginp, grid, nbatch) + FF_CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") + FF_CHECK_SAME(ginp.shape[ginp.ndim-1], 1, "count gradient must have a single channel") + FF_CHECK_SAME_BATCH(ginp, grid, nbatch) - const bool use_32bits = CANUSE32BITS(gout) && CANUSE32BITS(ginp) - && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(gout) && FF_CANUSE32BITS(ginp) + && FF_CANUSE32BITS(grid); const auto code = static_cast(gout.dtype.code); const auto bits = gout.dtype.bits; const spline_t spl = static_cast(spline); @@ -353,7 +353,7 @@ void count_backward( DISPATCH_PP(_count_backward, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(gout), CVOIDPTR(ginp), CVOIDPTR(grid), + FF_VOIDPTR(gout), FF_CVOIDPTR(ginp), FF_CVOIDPTR(grid), grid.shape, ginp.shape, gout.strides, ginp.strides, grid.strides) } @@ -385,29 +385,29 @@ void grad_backward( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, gout) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, ginp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") - CHECK_SAME(ginp.ndim, grid.ndim + 1, "ginp must have an extra trailing axis") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, gout) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, ginp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_SAME(ginp.ndim, grid.ndim + 1, "ginp must have an extra trailing axis") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); for (int32_t d = 0; d < out.ndim; ++d) - CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") + FF_CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") for (int32_t d = 0; d < gout.ndim; ++d) - CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") - CHECK_SAME(ginp.shape[ginp.ndim-1], ndim, "ginp trailing axis must equal ndim") - CHECK_SAME(ginp.shape[ginp.ndim-2], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(inp, grid, nbatch) - CHECK_SAME_BATCH(ginp, grid, nbatch) - - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(gout) - && CANUSE32BITS(inp) && CANUSE32BITS(ginp) - && CANUSE32BITS(grid); + FF_CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") + FF_CHECK_SAME(ginp.shape[ginp.ndim-1], ndim, "ginp trailing axis must equal ndim") + FF_CHECK_SAME(ginp.shape[ginp.ndim-2], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME_BATCH(ginp, grid, nbatch) + + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(gout) + && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(ginp) + && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -419,8 +419,8 @@ void grad_backward( DISPATCH_PP(_grad_backward, bvec, svec, static_cast(nbatch), n1, ex, abs, - VOIDPTR(out), VOIDPTR(gout), - CVOIDPTR(inp), CVOIDPTR(ginp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_VOIDPTR(gout), + FF_CVOIDPTR(inp), FF_CVOIDPTR(ginp), FF_CVOIDPTR(grid), grid.shape, inp.shape, out.strides, gout.strides, inp.strides, ginp.strides, grid.strides) } diff --git a/src/lib-cpu/reg_field.cpp b/src/lib-cpu/reg_field.cpp index 092a074..38abb23 100644 --- a/src/lib-cpu/reg_field.cpp +++ b/src/lib-cpu/reg_field.cpp @@ -5,6 +5,7 @@ #include "fastfields/api/cpu/reg_field.h" #include "fastfields/api/cpu/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/bounds.h" @@ -14,48 +15,14 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_SHAPE(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same shape"); - /*********************************************************************** * WRAPPERS * ***********************************************************************/ namespace { -// build a length-nc reduce_t vector from a (possibly null) double array -static inline std::vector as_weights(const double * w, int64_t nc) -{ - std::vector v(static_cast(nc), reduce_t(0)); - if (w) for (int64_t c = 0; c < nc; ++c) v[static_cast(c)] = w[c]; - return v; -} - template inline void _field_matvec( const bound::BoundVec & bvec, @@ -754,22 +721,22 @@ void field_matvec( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ +#define MV_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, inp.strides NDIM_SWITCH(MV_DT) #undef MV_ARGS @@ -796,22 +763,22 @@ void field_addmatvec_( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ +#define MV_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, inp.strides NDIM_SWITCH(ADD_MV_DT) #undef MV_ARGS @@ -838,22 +805,22 @@ void field_submatvec_( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ +#define MV_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, inp.strides NDIM_SWITCH(SUB_MV_DT) #undef MV_ARGS @@ -875,19 +842,19 @@ void field_diag( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define DG_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides NDIM_SWITCH(DG_DT) #undef DG_ARGS @@ -916,19 +883,19 @@ void field_adddiag_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define DG_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides NDIM_SWITCH(ADD_DG_DT) #undef DG_ARGS @@ -953,19 +920,19 @@ void field_subdiag_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define DG_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides NDIM_SWITCH(SUB_DG_DT) #undef DG_ARGS @@ -987,19 +954,19 @@ void field_kernel( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define KN_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides NDIM_SWITCH(KN_DT) #undef KN_ARGS @@ -1025,19 +992,19 @@ void field_addkernel_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define KN_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides NDIM_SWITCH(ADD_KN_DT) #undef KN_ARGS @@ -1063,19 +1030,19 @@ void field_subkernel_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define KN_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides NDIM_SWITCH(SUB_KN_DT) #undef KN_ARGS @@ -1096,25 +1063,25 @@ void field_relax( ) { const int32_t nbatch = sol.ndim - ndim - 1; - CHECK_NO_LANES (sol) - CHECK_SAME_DTYPE(sol, hes) - CHECK_SAME_DTYPE(sol, grd) - CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (sol) + FF_CHECK_SAME_DTYPE(sol, hes) + FF_CHECK_SAME_DTYPE(sol, grd) + FF_CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(sol, grd, sol.ndim) + FF_CHECK_SAME_SHAPE_N(sol, grd, sol.ndim) const int64_t nc = sol.shape[sol.ndim - 1]; - const bool use_32bits = CANUSE32BITS(sol) && CANUSE32BITS(hes) && - CANUSE32BITS(grd); + 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); -#define RX_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(sol), CVOIDPTR(hes), \ - CVOIDPTR(grd), voxel_size, absolute, membrane, bending, \ +#define RX_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(sol), FF_CVOIDPTR(hes), \ + FF_CVOIDPTR(grd), voxel_size, absolute, membrane, bending, \ nb_iter, sol.shape, sol.strides, hes.strides, grd.strides NDIM_SWITCH(RX_DT) #undef RX_ARGS @@ -1180,7 +1147,7 @@ void field_precond( int ndim , intptr_t stream ) { - CHECK_NO_LANES(grd) + FF_CHECK_NO_LANES(grd) if (grd.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); @@ -1201,7 +1168,7 @@ void field_precond_( int ndim , intptr_t stream ) { - CHECK_NO_LANES(sol) + FF_CHECK_NO_LANES(sol) if (sol.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); @@ -1248,27 +1215,27 @@ void field_matvec_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, wgt) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, wgt) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, inp, out.ndim) - CHECK_SAME_SHAPE(out, wgt, out.ndim - 1) + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, wgt, out.ndim - 1) 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 bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_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, \ +#define RLS_MV_ARGS bvec, static_cast(nbatch), nc, is_jrls, FF_VOIDPTR(out), \ + FF_CVOIDPTR(inp), FF_CVOIDPTR(wgt), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, inp.strides, wgt.strides NDIM_SWITCH(RLS_MV_DT) #undef RLS_MV_ARGS @@ -1292,24 +1259,24 @@ void field_diag_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, wgt) - CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, wgt) + FF_CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, wgt, out.ndim - 1) + FF_CHECK_SAME_SHAPE_N(out, wgt, out.ndim - 1) 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 bool use_32bits = FF_CANUSE32BITS(out) && FF_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, \ +#define RLS_DG_ARGS bvec, static_cast(nbatch), nc, is_jrls, FF_VOIDPTR(out), \ + FF_CVOIDPTR(wgt), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, wgt.strides NDIM_SWITCH(RLS_DG_DT) #undef RLS_DG_ARGS @@ -1335,31 +1302,31 @@ void field_relax_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = sol.ndim - ndim - 1; - CHECK_NO_LANES (sol) - CHECK_SAME_DTYPE(sol, hes) - CHECK_SAME_DTYPE(sol, grd) - CHECK_SAME_DTYPE(sol, wgt) - CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (sol.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (sol) + FF_CHECK_SAME_DTYPE(sol, hes) + FF_CHECK_SAME_DTYPE(sol, grd) + FF_CHECK_SAME_DTYPE(sol, wgt) + FF_CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (sol.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(sol, grd, sol.ndim) - CHECK_SAME_SHAPE(sol, wgt, sol.ndim - 1) + FF_CHECK_SAME_SHAPE_N(sol, grd, sol.ndim) + FF_CHECK_SAME_SHAPE_N(sol, wgt, sol.ndim - 1) 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 bool use_32bits = FF_CANUSE32BITS(sol) && FF_CANUSE32BITS(hes) && + FF_CANUSE32BITS(grd) && FF_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, \ +#define RLS_RX_ARGS bvec, static_cast(nbatch), nc, is_jrls, FF_VOIDPTR(sol), \ + FF_CVOIDPTR(hes), FF_CVOIDPTR(grd), FF_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 diff --git a/src/lib-cpu/reg_flow.cpp b/src/lib-cpu/reg_flow.cpp index 6c55856..dd1e02f 100644 --- a/src/lib-cpu/reg_flow.cpp +++ b/src/lib-cpu/reg_flow.cpp @@ -5,6 +5,7 @@ #include "fastfields/api/cpu/reg_flow.h" #include "fastfields/api/cpu/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/bounds.h" @@ -14,35 +15,9 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - // reduction / accumulation type (matches jitfields' float64 default) typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_SHAPE(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same shape"); - /*********************************************************************** * WRAPPERS * ***********************************************************************/ @@ -730,22 +705,22 @@ void flow_matvec( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME_SHAPE(out, inp, out.ndim) + 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 = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ +#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 NDIM_SWITCH(MV_DT) #undef MV_ARGS @@ -775,22 +750,22 @@ void flow_addmatvec_( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME_SHAPE(out, inp, out.ndim) + 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 = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ +#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 NDIM_SWITCH(ADD_MV_DT) #undef MV_ARGS @@ -820,22 +795,22 @@ void flow_submatvec_( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME_SHAPE(out, inp, out.ndim) + 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 = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ +#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 NDIM_SWITCH(SUB_MV_DT) #undef MV_ARGS @@ -859,18 +834,18 @@ void flow_diag( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define DG_ARGS bvec, static_cast(nbatch), VOIDPTR(out), \ +#define DG_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ voxel_size, absolute, membrane, bending, shears, div, \ out.shape, out.strides NDIM_SWITCH(DG_DT) @@ -902,18 +877,18 @@ void flow_adddiag_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define DG_ARGS bvec, static_cast(nbatch), VOIDPTR(out), \ +#define DG_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ voxel_size, absolute, membrane, bending, shears, div, \ out.shape, out.strides NDIM_SWITCH(ADD_DG_DT) @@ -941,18 +916,18 @@ void flow_subdiag_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define DG_ARGS bvec, static_cast(nbatch), VOIDPTR(out), \ +#define DG_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), \ voxel_size, absolute, membrane, bending, shears, div, \ out.shape, out.strides NDIM_SWITCH(SUB_DG_DT) @@ -983,21 +958,21 @@ void flow_kernel( const int ntrail = is_matrix ? 2 : 1; const int32_t nbatch = out.ndim - ndim - ntrail; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") if (is_matrix) - CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, + FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define KN_ARGS bvec, static_cast(nbatch), VOIDPTR(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) NDIM_SWITCH(KN_DT) @@ -1032,21 +1007,21 @@ void flow_addkernel_( const int ntrail = is_matrix ? 2 : 1; const int32_t nbatch = out.ndim - ndim - ntrail; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") if (is_matrix) - CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, + FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define KN_ARGS bvec, static_cast(nbatch), VOIDPTR(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) NDIM_SWITCH(ADD_KN_DT) @@ -1081,21 +1056,21 @@ void flow_subkernel_( const int ntrail = is_matrix ? 2 : 1; const int32_t nbatch = out.ndim - ndim - ntrail; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") if (is_matrix) - CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, + FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define KN_ARGS bvec, static_cast(nbatch), VOIDPTR(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) NDIM_SWITCH(SUB_KN_DT) @@ -1119,27 +1094,27 @@ void flow_relax( ) { const int32_t nbatch = sol.ndim - ndim - 1; - CHECK_NO_LANES (sol) - CHECK_SAME_DTYPE(sol, hes) - CHECK_SAME_DTYPE(sol, grd) - CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (sol) + FF_CHECK_SAME_DTYPE(sol, hes) + FF_CHECK_SAME_DTYPE(sol, grd) + FF_CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (sol.shape[sol.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME (grd.shape[grd.ndim-1], (int64_t)ndim, "Gradient channel dimension must equal ndim") - CHECK_SAME_SHAPE(sol, grd, sol.ndim) + FF_CHECK_SAME (sol.shape[sol.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + 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 = CANUSE32BITS(sol) && CANUSE32BITS(hes) && - CANUSE32BITS(grd); + 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); -#define RX_ARGS bvec, static_cast(nbatch), VOIDPTR(sol), CVOIDPTR(hes), \ - CVOIDPTR(grd), voxel_size, absolute, membrane, bending, \ - shears, div, nb_iter, sol.shape, sol.strides, hes.strides, \ +#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 NDIM_SWITCH(RX_DT) #undef RX_ARGS @@ -1211,7 +1186,7 @@ void flow_precond( int ndim , intptr_t stream ) { - CHECK_NO_LANES(grd) + FF_CHECK_NO_LANES(grd) if (grd.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); @@ -1234,7 +1209,7 @@ void flow_precond_( int ndim , intptr_t stream ) { - CHECK_NO_LANES(sol) + FF_CHECK_NO_LANES(sol) if (sol.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); @@ -1279,26 +1254,26 @@ void flow_matvec_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, wgt) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, wgt) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "Weight tensor must have a trailing size-1 channel axis") - CHECK_SAME_SHAPE(out, inp, out.ndim) - CHECK_SAME_SHAPE(out, wgt, out.ndim - 1) + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "Weight tensor must have a trailing size-1 channel axis") + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, wgt, out.ndim - 1) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(wgt); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_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), VOIDPTR(out), CVOIDPTR(inp), \ - CVOIDPTR(wgt), voxel_size, absolute, membrane, shears, div, \ +#define RLS_MV_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + FF_CVOIDPTR(wgt), voxel_size, absolute, membrane, shears, div, \ out.shape, out.strides, inp.strides, wgt.strides NDIM_SWITCH(RLS_MV_DT) #undef RLS_MV_ARGS @@ -1326,23 +1301,23 @@ void flow_diag_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, wgt) - CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, wgt) + FF_CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "Weight tensor must have a trailing size-1 channel axis") - CHECK_SAME_SHAPE(out, wgt, out.ndim - 1) + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "Weight tensor must have a trailing size-1 channel axis") + FF_CHECK_SAME_SHAPE_N(out, wgt, out.ndim - 1) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(wgt); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_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), VOIDPTR(out), CVOIDPTR(wgt), \ - voxel_size, absolute, membrane, shears, div, \ +#define RLS_DG_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), FF_CVOIDPTR(wgt), \ + voxel_size, absolute, membrane, shears, div, \ out.shape, out.strides, wgt.strides NDIM_SWITCH(RLS_DG_DT) #undef RLS_DG_ARGS @@ -1372,31 +1347,31 @@ void flow_relax_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = sol.ndim - ndim - 1; - CHECK_NO_LANES (sol) - CHECK_SAME_DTYPE(sol, hes) - CHECK_SAME_DTYPE(sol, grd) - CHECK_SAME_DTYPE(sol, wgt) - CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (sol.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (sol) + FF_CHECK_SAME_DTYPE(sol, hes) + FF_CHECK_SAME_DTYPE(sol, grd) + FF_CHECK_SAME_DTYPE(sol, wgt) + FF_CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (sol.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (sol.shape[sol.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME (grd.shape[grd.ndim-1], (int64_t)ndim, "Gradient channel dimension must equal ndim") - CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "Weight tensor must have a trailing size-1 channel axis") - CHECK_SAME_SHAPE(sol, grd, sol.ndim) - CHECK_SAME_SHAPE(sol, wgt, sol.ndim - 1) - - const bool use_32bits = CANUSE32BITS(sol) && CANUSE32BITS(hes) && - CANUSE32BITS(grd) && CANUSE32BITS(wgt); + FF_CHECK_SAME (sol.shape[sol.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (grd.shape[grd.ndim-1], (int64_t)ndim, "Gradient channel dimension must equal ndim") + FF_CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "Weight tensor must have a trailing size-1 channel axis") + FF_CHECK_SAME_SHAPE_N(sol, grd, sol.ndim) + FF_CHECK_SAME_SHAPE_N(sol, wgt, sol.ndim - 1) + + const bool use_32bits = FF_CANUSE32BITS(sol) && FF_CANUSE32BITS(hes) && + FF_CANUSE32BITS(grd) && FF_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), VOIDPTR(sol), CVOIDPTR(hes), \ - CVOIDPTR(grd), CVOIDPTR(wgt), voxel_size, absolute, membrane, \ - shears, div, nb_iter, sol.shape, sol.strides, hes.strides, \ +#define RLS_RX_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(sol), FF_CVOIDPTR(hes), \ + FF_CVOIDPTR(grd), FF_CVOIDPTR(wgt), voxel_size, absolute, membrane, \ + shears, div, nb_iter, sol.shape, sol.strides, hes.strides, \ grd.strides, wgt.strides NDIM_SWITCH(RLS_RX_DT) #undef RLS_RX_ARGS diff --git a/src/lib-cpu/resize.cpp b/src/lib-cpu/resize.cpp index 71fd4cb..6df361b 100644 --- a/src/lib-cpu/resize.cpp +++ b/src/lib-cpu/resize.cpp @@ -2,6 +2,7 @@ #include #include "fastfields/api/cpu/resize.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -10,39 +11,6 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ( \ - (X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes) \ - ) \ - throw std::invalid_argument( \ - "Tensors do not have the same data type" \ - ); - -#define CHECK_SAME_BATCH(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument( \ - "Tensors do not have the same batch shape" \ - ); - /*********************************************************************** * DISPATCH * ***********************************************************************/ @@ -155,23 +123,23 @@ inline void _resample( } #endif -#define DISPATCH_RESIZE(args...) \ -{ \ - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); \ - const auto code = static_cast(inp.dtype.code); \ - const spline_t spl = static_cast(spline); \ - const bound_t bnd = static_cast(bound); \ - switch (ndim) { \ - case 1: RS_ORDER(1, args); break; \ - case 2: RS_ORDER(2, args); break; \ - case 3: RS_ORDER(3, args); break; \ - default: throw std::invalid_argument( \ - "Only 1D, 2D and 3D resize are supported"); \ - }; \ - /* Reached only when a valid dim/spline/bound had an unsupported \ - dtype: the RS_DTYPE switch fell through without returning. */ \ - throw std::invalid_argument( \ - "Unsupported data type for resample (only float32/float64)"); \ +#define DISPATCH_RESIZE(args...) \ +{ \ + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); \ + const auto code = static_cast(inp.dtype.code); \ + const spline_t spl = static_cast(spline); \ + const bound_t bnd = static_cast(bound); \ + switch (ndim) { \ + case 1: RS_ORDER(1, args); break; \ + case 2: RS_ORDER(2, args); break; \ + case 3: RS_ORDER(3, args); break; \ + default: throw std::invalid_argument( \ + "Only 1D, 2D and 3D resize are supported"); \ + }; \ + /* Reached only when a valid dim/spline/bound had an unsupported \ + dtype: the RS_DTYPE switch fell through without returning. */ \ + throw std::invalid_argument( \ + "Unsupported data type for resample (only float32/float64)"); \ } void resample( @@ -192,17 +160,17 @@ void resample( DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_BATCH(out, inp, nbatch) + FF_CHECK_SAME_BATCH(out, inp, nbatch) DISPATCH_RESIZE( static_cast(nbatch), - VOIDPTR(out), - VOIDPTR(inp), + FF_VOIDPTR(out), + FF_VOIDPTR(inp), shift, scale, out.shape, diff --git a/src/lib-cpu/restrict.cpp b/src/lib-cpu/restrict.cpp index f49f144..32ee41c 100644 --- a/src/lib-cpu/restrict.cpp +++ b/src/lib-cpu/restrict.cpp @@ -2,6 +2,7 @@ #include #include "fastfields/api/cpu/restrict.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -10,39 +11,6 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ( \ - (X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes) \ - ) \ - throw std::invalid_argument( \ - "Tensors do not have the same data type" \ - ); - -#define CHECK_SAME_BATCH(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument( \ - "Tensors do not have the same batch shape" \ - ); - /*********************************************************************** * DISPATCH * ***********************************************************************/ @@ -155,23 +123,23 @@ inline void _restriction( } #endif -#define DISPATCH_RESTRICT(args...) \ -{ \ - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); \ - const auto code = static_cast(inp.dtype.code); \ - const spline_t spl = static_cast(spline); \ - const bound_t bnd = static_cast(bound); \ - switch (ndim) { \ - case 1: RT_ORDER(1, args); break; \ - case 2: RT_ORDER(2, args); break; \ - case 3: RT_ORDER(3, args); break; \ - default: throw std::invalid_argument( \ - "Only 1D, 2D and 3D restrict are supported"); \ - }; \ - /* Reached only when a valid dim/spline/bound had an unsupported \ - dtype: the RT_DTYPE switch fell through without returning. */ \ - throw std::invalid_argument( \ - "Unsupported data type for restriction (only float32/float64)");\ +#define DISPATCH_RESTRICT(args...) \ +{ \ + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); \ + const auto code = static_cast(inp.dtype.code); \ + const spline_t spl = static_cast(spline); \ + const bound_t bnd = static_cast(bound); \ + switch (ndim) { \ + case 1: RT_ORDER(1, args); break; \ + case 2: RT_ORDER(2, args); break; \ + case 3: RT_ORDER(3, args); break; \ + default: throw std::invalid_argument( \ + "Only 1D, 2D and 3D restrict are supported"); \ + }; \ + /* Reached only when a valid dim/spline/bound had an unsupported \ + dtype: the RT_DTYPE switch fell through without returning. */ \ + throw std::invalid_argument( \ + "Unsupported data type for restriction (only float32/float64)"); \ } void restriction( @@ -192,17 +160,17 @@ void restriction( DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_BATCH(out, inp, nbatch) + FF_CHECK_SAME_BATCH(out, inp, nbatch) DISPATCH_RESTRICT( static_cast(nbatch), - VOIDPTR(out), - VOIDPTR(inp), + FF_VOIDPTR(out), + FF_VOIDPTR(inp), shift, scale, out.shape, diff --git a/src/lib-cpu/solve_field.cpp b/src/lib-cpu/solve_field.cpp index c30ca4b..d1311f5 100644 --- a/src/lib-cpu/solve_field.cpp +++ b/src/lib-cpu/solve_field.cpp @@ -5,6 +5,7 @@ #include "fastfields/api/cpu/reg_field.h" #include "fastfields/api/cpu/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -13,34 +14,8 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_SHAPE(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same shape"); - /*********************************************************************** * VECTOR PRIMITIVES * ***********************************************************************/ @@ -163,9 +138,9 @@ inline void _axpby_( static inline reduce_t dot(const DLTensor & x, const DLTensor & y) { const int64_t nall = static_cast(x.ndim) - 1; - const bool use_32bits = CANUSE32BITS(x) && CANUSE32BITS(y); + const bool use_32bits = FF_CANUSE32BITS(x) && FF_CANUSE32BITS(y); #define DOT_CALL(SCALAR, OFFSET) \ - _dot(nall, CVOIDPTR(x), CVOIDPTR(y), x.shape, x.strides, y.strides) + _dot(nall, FF_CVOIDPTR(x), FF_CVOIDPTR(y), x.shape, x.strides, y.strides) SOLVE_DT_SWITCH(DOT_CALL) #undef DOT_CALL } @@ -175,9 +150,9 @@ static inline void axpby_(DLTensor & y, const DLTensor & x, reduce_t a, reduce_t b) { const int64_t nall = static_cast(x.ndim) - 1; - const bool use_32bits = CANUSE32BITS(x) && CANUSE32BITS(y); -#define AXPBY_CALL(SCALAR, OFFSET) \ - _axpby_(nall, VOIDPTR(y), CVOIDPTR(x), a, b, \ + const bool use_32bits = FF_CANUSE32BITS(x) && FF_CANUSE32BITS(y); +#define AXPBY_CALL(SCALAR, OFFSET) \ + _axpby_(nall, FF_VOIDPTR(y), FF_CVOIDPTR(x), a, b, \ x.shape, y.strides, x.strides) SOLVE_DT_SWITCH(AXPBY_CALL) #undef AXPBY_CALL @@ -210,17 +185,17 @@ void field_cg( DLTensor & sol = _sol.t; const DLTensor & grd = _grd.t; - CHECK_NO_LANES (sol) - CHECK_SAME_DTYPE(sol, grd) - CHECK_SAME_DTYPE(sol, hes) - CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (sol) + FF_CHECK_SAME_DTYPE(sol, grd) + FF_CHECK_SAME_DTYPE(sol, hes) + FF_CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") if (sol.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(sol, grd, sol.ndim) + FF_CHECK_SAME_SHAPE_N(sol, grd, sol.ndim) const int64_t nc = sol.shape[sol.ndim - 1]; - CHECK_SAME(hes.shape[hes.ndim - 1], nc * (nc + 1) / 2, + FF_CHECK_SAME(hes.shape[hes.ndim - 1], nc * (nc + 1) / 2, "The Hessian's trailing dimension must be C*(C+1)/2") if (nb_iter_out) *nb_iter_out = 0; diff --git a/src/lib-cpu/splinc.cpp b/src/lib-cpu/splinc.cpp index f21172d..161f40d 100644 --- a/src/lib-cpu/splinc.cpp +++ b/src/lib-cpu/splinc.cpp @@ -3,6 +3,7 @@ #include #include "fastfields/api/cpu/splinc.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -11,19 +12,6 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - /*********************************************************************** * POLES * ***********************************************************************/ @@ -107,22 +95,22 @@ inline void _splinc( default: throw std::invalid_argument("Unsupported npoles"); \ } -#define DISPATCH_SPLINC(args...) \ -{ \ - const bool use_32bits = CANUSE32BITS(inp_out); \ - const auto code = static_cast(inp_out.dtype.code); \ - switch (code) { \ - case kDLFloat: switch (inp_out.dtype.bits) { \ - case 32: \ - if (use_32bits) DISPATCH_SPLINC_NPOLES(float, int32_t, args) \ - else DISPATCH_SPLINC_NPOLES(float, int64_t, args) \ - case 64: \ - if (use_32bits) DISPATCH_SPLINC_NPOLES(double, int32_t, args) \ - else DISPATCH_SPLINC_NPOLES(double, int64_t, args) \ - default: break; \ - }; \ - default: break; \ - }; \ +#define DISPATCH_SPLINC(args...) \ +{ \ + const bool use_32bits = FF_CANUSE32BITS(inp_out); \ + const auto code = static_cast(inp_out.dtype.code); \ + switch (code) { \ + case kDLFloat: switch (inp_out.dtype.bits) { \ + case 32: \ + if (use_32bits) DISPATCH_SPLINC_NPOLES(float, int32_t, args) \ + else DISPATCH_SPLINC_NPOLES(float, int64_t, args) \ + case 64: \ + if (use_32bits) DISPATCH_SPLINC_NPOLES(double, int32_t, args) \ + else DISPATCH_SPLINC_NPOLES(double, int64_t, args) \ + default: break; \ + }; \ + default: break; \ + }; \ throw std::invalid_argument("only floating point data types are supported"); \ } @@ -137,7 +125,7 @@ void spline_coeff( ContiguousStrides _io(inp_out_); DLTensor & inp_out = _io.t; - CHECK_NO_LANES(inp_out) + FF_CHECK_NO_LANES(inp_out) double poles[3]; const int npoles = get_poles_host(static_cast(spline), poles); @@ -148,7 +136,7 @@ void spline_coeff( DISPATCH_SPLINC( nbatch, - VOIDPTR(inp_out), + FF_VOIDPTR(inp_out), inp_out.shape, inp_out.strides, poles diff --git a/src/lib-cuda/distance.cpp b/src/lib-cuda/distance.cpp index 08e7e0a..d80dc41 100644 --- a/src/lib-cuda/distance.cpp +++ b/src/lib-cuda/distance.cpp @@ -2,6 +2,7 @@ #include #include "fastfields/api/cuda/distance.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -13,76 +14,32 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_BATCH(X, Y, D) \ - if (X.ndim < D || Y.ndim < D) \ - throw std::invalid_argument( \ - "Number of dimensions does not match" \ - ); \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument( \ - "Tensors do not have the same batch shape" \ - ); \ - -#define CHECK_SAME_SHAPE(X, Y) \ - if (X.ndim != Y.ndim) \ - throw std::invalid_argument( \ - "Tensors do not have the same number of dimensions" \ - ); \ - CHECK_SAME_BATCH(X, Y, X.ndim) - -#define CHECK_SAME_DTYPE(X, Y) \ - if ( \ - (X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes) \ - ) \ - throw std::invalid_argument( \ - "Tensors do not have the same data type" \ - ); - /*********************************************************************** * EUCLIDEAN * ***********************************************************************/ -#define DISPATCH_DT(func, args...) \ -{ \ - const bool use_32bits = CANUSE32BITS(inp_out); \ - const auto code = static_cast(inp_out.dtype.code); \ - switch (code) { \ - case kDLFloat: switch (inp_out.dtype.bits) { \ - case 32: return ( \ - use_32bits \ - ? func(args) \ - : func(args) \ - ); \ - case 64: return ( \ - use_32bits \ - ? func(args) \ - : func(args) \ - ); \ - default: break; \ - }; \ - default: throw std::invalid_argument( \ - "only floating point data types are supported" \ - ); \ - }; \ +#define DISPATCH_DT(func, args...) \ +{ \ + const bool use_32bits = FF_CANUSE32BITS(inp_out); \ + const auto code = static_cast(inp_out.dtype.code); \ + switch (code) { \ + case kDLFloat: switch (inp_out.dtype.bits) { \ + case 32: return ( \ + use_32bits \ + ? func(args) \ + : func(args) \ + ); \ + case 64: return ( \ + use_32bits \ + ? func(args) \ + : func(args) \ + ); \ + default: break; \ + }; \ + default: throw std::invalid_argument( \ + "only floating point data types are supported" \ + ); \ + }; \ } namespace { @@ -116,11 +73,11 @@ void dt_euclidean( ContiguousStrides _io(inp_out_); DLTensor & inp_out = _io.t; - CHECK_NO_LANES(inp_out) + FF_CHECK_NO_LANES(inp_out) DISPATCH_DT( _dt_euclidean, inp_out.ndim, - VOIDPTR(inp_out), + FF_VOIDPTR(inp_out), voxel_spacing, inp_out.shape, inp_out.strides, @@ -159,11 +116,11 @@ void dt_l1( ContiguousStrides _io(inp_out_); DLTensor & inp_out = _io.t; - CHECK_NO_LANES(inp_out) + FF_CHECK_NO_LANES(inp_out) DISPATCH_DT( _dt_l1, inp_out.ndim, - VOIDPTR(inp_out), + FF_VOIDPTR(inp_out), voxel_spacing, inp_out.shape, inp_out.strides, @@ -288,39 +245,39 @@ void dt_spline_table( const DLTensor & times = _ti.t; const bool use_32bits = ( - CANUSE32BITS(time) && - CANUSE32BITS(dist) && - CANUSE32BITS(loc) && - CANUSE32BITS(coeff) && - CANUSE32BITS(times) + FF_CANUSE32BITS(time) && + FF_CANUSE32BITS(dist) && + FF_CANUSE32BITS(loc) && + FF_CANUSE32BITS(coeff) && + FF_CANUSE32BITS(times) ); const int32_t ndim = loc.shape[loc.ndim-1]; const int32_t nbatch = loc.ndim - 1; - CHECK_NO_LANES (time) - CHECK_SAME_DTYPE(time, dist) - CHECK_SAME_DTYPE(time, loc) - CHECK_SAME_DTYPE(time, coeff) - CHECK_SAME_DTYPE(time, times) - CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") - CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_NO_LANES (time) + FF_CHECK_SAME_DTYPE(time, dist) + FF_CHECK_SAME_DTYPE(time, loc) + FF_CHECK_SAME_DTYPE(time, coeff) + FF_CHECK_SAME_DTYPE(time, times) + FF_CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") // coeff is (*batch, npoints, ndim) -> nbatch+2 dims // times is (*batch, ntimes) -> nbatch+1 dims - CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") - CHECK_SAME (times.ndim, nbatch+1, "Number of times dimensions does not match") - CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") - CHECK_SAME_BATCH(loc, time, nbatch) - CHECK_SAME_BATCH(loc, dist, nbatch) - CHECK_SAME_BATCH(loc, coeff, nbatch) - CHECK_SAME_BATCH(loc, times, nbatch) + FF_CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") + FF_CHECK_SAME (times.ndim, nbatch+1, "Number of times dimensions does not match") + FF_CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") + FF_CHECK_SAME_BATCH_ND(loc, time, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, dist, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, coeff, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, times, nbatch) DISPATCH_SPLINE( _dt_spline_table, nbatch, // nbatch - VOIDPTR(time), // time - VOIDPTR(dist), // dist - VOIDPTR(loc), // loc - VOIDPTR(coeff), // coeff - VOIDPTR(times), // times + FF_VOIDPTR(time), // time + FF_VOIDPTR(dist), // dist + FF_VOIDPTR(loc), // loc + FF_VOIDPTR(coeff), // coeff + FF_VOIDPTR(times), // times times.shape[times.ndim-1], // ntimes coeff.shape, // size (coeff shape: *batch, npoints, ndim) time.strides, // int64_time @@ -410,33 +367,33 @@ void dt_spline_brent( const DLTensor & coeff = _co.t; const bool use_32bits = ( - CANUSE32BITS(time) && - CANUSE32BITS(dist) && - CANUSE32BITS(loc) && - CANUSE32BITS(coeff) + FF_CANUSE32BITS(time) && + FF_CANUSE32BITS(dist) && + FF_CANUSE32BITS(loc) && + FF_CANUSE32BITS(coeff) ); const int32_t ndim = loc.shape[loc.ndim-1]; const int32_t nbatch = loc.ndim - 1; - CHECK_NO_LANES (time) - CHECK_SAME_DTYPE(time, dist) - CHECK_SAME_DTYPE(time, loc) - CHECK_SAME_DTYPE(time, coeff) - CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") - CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_NO_LANES (time) + FF_CHECK_SAME_DTYPE(time, dist) + FF_CHECK_SAME_DTYPE(time, loc) + FF_CHECK_SAME_DTYPE(time, coeff) + FF_CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") // coeff is (*batch, npoints, ndim) -> nbatch+2 dims - CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") - CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") - CHECK_SAME_BATCH(loc, time, nbatch) - CHECK_SAME_BATCH(loc, dist, nbatch) - CHECK_SAME_BATCH(loc, coeff, nbatch) + FF_CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") + FF_CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") + FF_CHECK_SAME_BATCH_ND(loc, time, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, dist, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, coeff, nbatch) DISPATCH_SPLINE( _dt_spline_brent, nbatch, // nbatch - VOIDPTR(time), // time - VOIDPTR(dist), // dist - VOIDPTR(loc), // loc - VOIDPTR(coeff), // coeff + FF_VOIDPTR(time), // time + FF_VOIDPTR(dist), // dist + FF_VOIDPTR(loc), // loc + FF_VOIDPTR(coeff), // coeff coeff.shape, // size (coeff shape: *batch, npoints, ndim) time.strides, // int64_time dist.strides, // stride_dist @@ -524,33 +481,33 @@ void dt_spline_gaussnewton( const DLTensor & coeff = _co.t; const bool use_32bits = ( - CANUSE32BITS(time) && - CANUSE32BITS(dist) && - CANUSE32BITS(loc) && - CANUSE32BITS(coeff) + FF_CANUSE32BITS(time) && + FF_CANUSE32BITS(dist) && + FF_CANUSE32BITS(loc) && + FF_CANUSE32BITS(coeff) ); const int32_t ndim = loc.shape[loc.ndim-1]; const int32_t nbatch = loc.ndim - 1; - CHECK_NO_LANES (time) - CHECK_SAME_DTYPE(time, dist) - CHECK_SAME_DTYPE(time, loc) - CHECK_SAME_DTYPE(time, coeff) - CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") - CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_NO_LANES (time) + FF_CHECK_SAME_DTYPE(time, dist) + FF_CHECK_SAME_DTYPE(time, loc) + FF_CHECK_SAME_DTYPE(time, coeff) + FF_CHECK_SAME (time.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_SAME (dist.ndim, nbatch, "Number of batch dimensions does not match") // coeff is (*batch, npoints, ndim) -> nbatch+2 dims - CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") - CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") - CHECK_SAME_BATCH(loc, time, nbatch) - CHECK_SAME_BATCH(loc, dist, nbatch) - CHECK_SAME_BATCH(loc, coeff, nbatch) + FF_CHECK_SAME (coeff.ndim, nbatch+2, "Number of coeff dimensions does not match") + FF_CHECK_SAME (coeff.shape[coeff.ndim-1], ndim, "Dimensionality of coeff and location does not match") + FF_CHECK_SAME_BATCH_ND(loc, time, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, dist, nbatch) + FF_CHECK_SAME_BATCH_ND(loc, coeff, nbatch) DISPATCH_SPLINE( _dt_spline_gaussnewton, nbatch, // nbatch - VOIDPTR(time), // time - VOIDPTR(dist), // dist - VOIDPTR(loc), // loc - VOIDPTR(coeff), // coeff + FF_VOIDPTR(time), // time + FF_VOIDPTR(dist), // dist + FF_VOIDPTR(loc), // loc + FF_VOIDPTR(coeff), // coeff coeff.shape, // size (coeff shape: *batch, npoints, ndim) time.strides, // int64_time dist.strides, // stride_dist @@ -706,42 +663,42 @@ void dt_mesh( const DLTensor & faces = _fa.t; bool use_32bits = ( - CANUSE32BITS(dist) && - CANUSE32BITS(loc) && - CANUSE32BITS(vertices) && - CANUSE32BITS(faces) + FF_CANUSE32BITS(dist) && + FF_CANUSE32BITS(loc) && + FF_CANUSE32BITS(vertices) && + FF_CANUSE32BITS(faces) ); const int32_t ndim = loc.shape[loc.ndim-1]; const int32_t nbatch = loc.ndim - 1; - CHECK_NO_LANES (dist) - CHECK_SAME_DTYPE(dist, loc) - CHECK_SAME_DTYPE(dist, vertices) - CHECK_SAME ( dist.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_NO_LANES (dist) + FF_CHECK_SAME_DTYPE(dist, loc) + FF_CHECK_SAME_DTYPE(dist, vertices) + FF_CHECK_SAME ( dist.ndim, nbatch, "Number of batch dimensions does not match") // vertices (N, D) and faces (M, D) describe a single shared mesh and are // always 2D; their leading axis is the vertex/face count, independent of // loc's point batch. Only `loc` and the per-point outputs share a batch. - CHECK_SAME ( vertices.ndim, 2, "Vertices must be a (N, D) tensor") - CHECK_SAME ( faces.ndim, 2, "Faces must be a (M, D) tensor") - CHECK_SAME_BATCH(loc, dist, nbatch) - CHECK_SAME (vertices.shape[vertices.ndim-1], ndim, "Dimensionality of the vertices and location does not match") - CHECK_SAME (faces.shape[faces.ndim-1], ndim, "Dimensionality of the vertices and faces does not match") + FF_CHECK_SAME ( vertices.ndim, 2, "Vertices must be a (N, D) tensor") + FF_CHECK_SAME ( faces.ndim, 2, "Faces must be a (M, D) tensor") + FF_CHECK_SAME_BATCH_ND(loc, dist, nbatch) + FF_CHECK_SAME (vertices.shape[vertices.ndim-1], ndim, "Dimensionality of the vertices and location does not match") + FF_CHECK_SAME (faces.shape[faces.ndim-1], ndim, "Dimensionality of the vertices and faces does not match") if (nearest_vertex.data) { - CHECK_SAME_DTYPE(faces, nearest_vertex) - CHECK_SAME (nearest_vertex.ndim, nbatch, "Number of batch dimensions does not match") - CHECK_SAME_BATCH(loc, nearest_vertex, nbatch) - use_32bits &= CANUSE32BITS(nearest_vertex); + FF_CHECK_SAME_DTYPE(faces, nearest_vertex) + FF_CHECK_SAME (nearest_vertex.ndim, nbatch, "Number of batch dimensions does not match") + FF_CHECK_SAME_BATCH_ND(loc, nearest_vertex, nbatch) + use_32bits &= FF_CANUSE32BITS(nearest_vertex); } DISPATCH_MESH( _dt_mesh, nbatch, // nbatch - VOIDPTR(dist), // data - VOIDPTR(nearest_vertex), // nearest_vertex - VOIDPTR(loc), // coord - VOIDPTR(vertices), // vertices - VOIDPTR(faces), // faces + FF_VOIDPTR(dist), // data + FF_VOIDPTR(nearest_vertex), // nearest_vertex + FF_VOIDPTR(loc), // coord + FF_VOIDPTR(vertices), // vertices + FF_VOIDPTR(faces), // faces loc.shape, // size faces.shape[0], // nb_faces (M = faces.shape[0]) vertices.shape[0], // nb_vertices (N = vertices.shape[0]) diff --git a/src/lib-cuda/posdef.cpp b/src/lib-cuda/posdef.cpp index 080bf28..7ddf107 100644 --- a/src/lib-cuda/posdef.cpp +++ b/src/lib-cuda/posdef.cpp @@ -3,6 +3,7 @@ #include #include "fastfields/api/cuda/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -11,48 +12,10 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (x.data ? static_cast(static_cast(x.data) + x.byte_offset) : nullptr) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - // reduce/accumulation type used by the compact-symmetric kernels. // jitfields defaults to float64; we do the same for CPU accuracy. typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_BATCH(X, Y, D) \ - if (X.ndim < D || Y.ndim < D) \ - throw std::invalid_argument( \ - "Number of dimensions does not match" \ - ); \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument( \ - "Tensors do not have the same batch shape" \ - ); \ - -#define CHECK_SAME_DTYPE(X, Y) \ - if ( \ - (X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes) \ - ) \ - throw std::invalid_argument( \ - "Tensors do not have the same data type" \ - ); - // C such that C*(C+1)/2 == CC (the compact-symmetric length). static inline int64_t channels_from_packed(int64_t CC) { @@ -165,23 +128,23 @@ void sym_matvec( const DLTensor & hessian = _hes.t; const DLTensor & inp = _inp.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian) && CANUSE32BITS(inp); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian) && FF_CANUSE32BITS(inp); const int32_t nbatch = out.ndim - 1; const int64_t nchannel = out.shape[out.ndim-1]; const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(out, inp, nbatch) - CHECK_SAME_BATCH(out, hessian, nbatch) + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(out, inp, nbatch) + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) DISPATCH_SYM_C( _sym_matvec, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(hessian), CVOIDPTR(inp), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(hessian), FF_CVOIDPTR_OR_NULL(inp), out.shape, out.strides, hessian.strides, inp.strides ) } @@ -247,23 +210,23 @@ void sym_addmatvec_( const DLTensor & hessian = _hes.t; const DLTensor & inp = _inp.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian) && CANUSE32BITS(inp); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian) && FF_CANUSE32BITS(inp); const int32_t nbatch = out.ndim - 1; const int64_t nchannel = out.shape[out.ndim-1]; const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(out, inp, nbatch) - CHECK_SAME_BATCH(out, hessian, nbatch) + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(out, inp, nbatch) + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) DISPATCH_SYM_C( _sym_addmatvec_, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(hessian), CVOIDPTR(inp), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(hessian), FF_CVOIDPTR_OR_NULL(inp), out.shape, out.strides, hessian.strides, inp.strides ) } @@ -281,23 +244,23 @@ void sym_submatvec_( const DLTensor & hessian = _hes.t; const DLTensor & inp = _inp.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian) && CANUSE32BITS(inp); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian) && FF_CANUSE32BITS(inp); const int32_t nbatch = out.ndim - 1; const int64_t nchannel = out.shape[out.ndim-1]; const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(out, inp, nbatch) - CHECK_SAME_BATCH(out, hessian, nbatch) + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(out, inp, nbatch) + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) DISPATCH_SYM_C( _sym_submatvec_, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(hessian), CVOIDPTR(inp), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(hessian), FF_CVOIDPTR_OR_NULL(inp), out.shape, out.strides, hessian.strides, inp.strides ) } @@ -344,23 +307,23 @@ void sym_matvec_backward( const DLTensor & grd = _grd.t; const DLTensor & inp = _inp.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(grd) && CANUSE32BITS(inp); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(grd) && FF_CANUSE32BITS(inp); const int32_t nbatch = grd.ndim - 1; const int64_t nchannel = grd.shape[grd.ndim-1]; const auto code = static_cast(grd.dtype.code); const auto bits = grd.dtype.bits; - CHECK_NO_LANES (grd) - CHECK_SAME_DTYPE(grd, out) - CHECK_SAME_DTYPE(grd, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and grad channel counts differ") - CHECK_SAME (out.shape[out.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(grd, out, nbatch) - CHECK_SAME_BATCH(grd, inp, nbatch) + FF_CHECK_NO_LANES (grd) + FF_CHECK_SAME_DTYPE(grd, out) + FF_CHECK_SAME_DTYPE(grd, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and grad channel counts differ") + FF_CHECK_SAME (out.shape[out.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(grd, out, nbatch) + FF_CHECK_SAME_BATCH_ND(grd, inp, nbatch) DISPATCH_SYM_C( _sym_matvec_backward, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(grd), CVOIDPTR(inp), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(grd), FF_CVOIDPTR_OR_NULL(inp), grd.shape, out.strides, grd.strides, inp.strides ) } @@ -416,26 +379,26 @@ void sym_solve( const DLTensor & weight = _wgt.t; const bool has_wgt = (weight.data != nullptr); - bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian) && CANUSE32BITS(inp); - if (has_wgt) use_32bits = use_32bits && CANUSE32BITS(weight); + bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian) && FF_CANUSE32BITS(inp); + if (has_wgt) use_32bits = use_32bits && FF_CANUSE32BITS(weight); const int32_t nbatch = out.ndim - 1; const int64_t nchannel = out.shape[out.ndim-1]; const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(out, inp, nbatch) - CHECK_SAME_BATCH(out, hessian, nbatch) - if (has_wgt) { CHECK_SAME_DTYPE(out, weight) CHECK_SAME_BATCH(out, weight, nbatch) } + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (inp.shape[inp.ndim-1], nchannel, "Input and output channel counts differ") + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(out, inp, nbatch) + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) + if (has_wgt) { FF_CHECK_SAME_DTYPE(out, weight) FF_CHECK_SAME_BATCH_ND(out, weight, nbatch) } DISPATCH_SYM( _sym_solve, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(hessian), - has_wgt ? VOIDPTR(weight) : nullptr, + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(inp), FF_CVOIDPTR_OR_NULL(hessian), + has_wgt ? FF_VOIDPTR(weight) : nullptr, out.shape, out.strides, inp.strides, hessian.strides, has_wgt ? weight.strides : nullptr ) @@ -482,23 +445,23 @@ void sym_solve_( const DLTensor & weight = _wgt.t; const bool has_wgt = (weight.data != nullptr); - bool use_32bits = CANUSE32BITS(inp_out) && CANUSE32BITS(hessian); - if (has_wgt) use_32bits = use_32bits && CANUSE32BITS(weight); + bool use_32bits = FF_CANUSE32BITS(inp_out) && FF_CANUSE32BITS(hessian); + if (has_wgt) use_32bits = use_32bits && FF_CANUSE32BITS(weight); const int32_t nbatch = inp_out.ndim - 1; const int64_t nchannel = inp_out.shape[inp_out.ndim-1]; const auto code = static_cast(inp_out.dtype.code); const auto bits = inp_out.dtype.bits; - CHECK_NO_LANES (inp_out) - CHECK_SAME_DTYPE(inp_out, hessian) - CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") - CHECK_SAME_BATCH(inp_out, hessian, nbatch) - if (has_wgt) { CHECK_SAME_DTYPE(inp_out, weight) CHECK_SAME_BATCH(inp_out, weight, nbatch) } + FF_CHECK_NO_LANES (inp_out) + FF_CHECK_SAME_DTYPE(inp_out, hessian) + FF_CHECK_SAME (hessian.shape[hessian.ndim-1]*2, nchannel*(nchannel+1), "Matrix is not compatible with the channel count") + FF_CHECK_SAME_BATCH_ND(inp_out, hessian, nbatch) + if (has_wgt) { FF_CHECK_SAME_DTYPE(inp_out, weight) FF_CHECK_SAME_BATCH_ND(inp_out, weight, nbatch) } DISPATCH_SYM( _sym_solve_, nbatch, nchannel, - VOIDPTR(inp_out), CVOIDPTR(hessian), - has_wgt ? VOIDPTR(weight) : nullptr, + FF_VOIDPTR(inp_out), FF_CVOIDPTR_OR_NULL(hessian), + has_wgt ? FF_VOIDPTR(weight) : nullptr, inp_out.shape, inp_out.strides, hessian.strides, has_wgt ? weight.strides : nullptr ) @@ -540,21 +503,21 @@ void sym_invert( DLTensor & out = _out.t; const DLTensor & hessian = _hes.t; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(hessian); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(hessian); const int32_t nbatch = out.ndim - 1; const int64_t CC = hessian.shape[hessian.ndim-1]; const int64_t nchannel = channels_from_packed(CC); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, hessian) - CHECK_SAME (out.shape[out.ndim-1], CC, "Output and matrix must share the compact layout") - CHECK_SAME_BATCH(out, hessian, nbatch) + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, hessian) + FF_CHECK_SAME (out.shape[out.ndim-1], CC, "Output and matrix must share the compact layout") + FF_CHECK_SAME_BATCH_ND(out, hessian, nbatch) DISPATCH_SYM( _sym_invert, nbatch, nchannel, - VOIDPTR(out), CVOIDPTR(hessian), + FF_VOIDPTR(out), FF_CVOIDPTR_OR_NULL(hessian), out.shape, out.strides, hessian.strides ) } @@ -586,18 +549,18 @@ void sym_invert_( ContiguousStrides _hes(hessian_); DLTensor & hessian = _hes.t; - const bool use_32bits = CANUSE32BITS(hessian); + const bool use_32bits = FF_CANUSE32BITS(hessian); const int32_t nbatch = hessian.ndim - 1; const int64_t CC = hessian.shape[hessian.ndim-1]; const int64_t nchannel = channels_from_packed(CC); const auto code = static_cast(hessian.dtype.code); const auto bits = hessian.dtype.bits; - CHECK_NO_LANES(hessian) + FF_CHECK_NO_LANES(hessian) DISPATCH_SYM( _sym_invert_, nbatch, nchannel, - VOIDPTR(hessian), + FF_VOIDPTR(hessian), hessian.shape, hessian.strides ) } diff --git a/src/lib-cuda/pushpull.cpp b/src/lib-cuda/pushpull.cpp index d7853a8..b3f4916 100644 --- a/src/lib-cuda/pushpull.cpp +++ b/src/lib-cuda/pushpull.cpp @@ -1,6 +1,6 @@ #include "fastfields/api/cuda/pushpull.h" #include -// VOIDPTR / CHECK_* / DISPATCH_PP and the reduce_t typedef, shared with +// FF_VOIDPTR / CHECK_* / DISPATCH_PP and the reduce_t typedef, shared with // pushpull_backward.cpp so the two translation units cannot drift apart on // which (order, bound) pairs are statically instantiated. #include "fastfields/api/cuda/pushpull_dispatch.h" @@ -156,18 +156,18 @@ void pull( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; // nbatch + ndim + 1 - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); - CHECK_SAME(out.shape[out.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(out, grid, nbatch) - CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME(out.shape[out.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(out, grid, nbatch) + FF_CHECK_SAME_BATCH(inp, grid, nbatch) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -181,7 +181,7 @@ void pull( DISPATCH_PP(_pull, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_CVOIDPTR(inp), FF_CVOIDPTR(grid), grid.shape, inp.shape, out.strides, inp.strides, grid.strides, stream) } @@ -209,18 +209,18 @@ void push( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); - CHECK_SAME(out.shape[out.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(out, grid, nbatch) - CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME(out.shape[out.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(out, grid, nbatch) + FF_CHECK_SAME_BATCH(inp, grid, nbatch) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -235,7 +235,7 @@ void push( DISPATCH_PP(_push, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_CVOIDPTR(inp), FF_CVOIDPTR(grid), grid.shape, out.shape, out.strides, inp.strides, grid.strides, stream) } @@ -261,14 +261,14 @@ void count( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); - CHECK_SAME_BATCH(out, grid, nbatch) + FF_CHECK_SAME_BATCH(out, grid, nbatch) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -282,7 +282,7 @@ void count( DISPATCH_PP(_count, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_CVOIDPTR(grid), grid.shape, out.shape, out.strides, grid.strides, stream) } @@ -311,18 +311,18 @@ void grad( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; // grid/inp rank; out rank == n1 + 1 - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") - CHECK_SAME(out.ndim, grid.ndim + 1, "grad output must have an extra trailing axis") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_SAME(out.ndim, grid.ndim + 1, "grad output must have an extra trailing axis") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); - CHECK_SAME(out.shape[out.ndim-1], ndim, "grad output trailing axis must equal ndim") - CHECK_SAME(out.shape[out.ndim-2], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME(out.shape[out.ndim-1], ndim, "grad output trailing axis must equal ndim") + FF_CHECK_SAME(out.shape[out.ndim-2], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(inp, grid, nbatch) - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -336,7 +336,7 @@ void grad( DISPATCH_PP(_grad, bvec, svec, static_cast(nbatch), n1, ex, abs, - VOIDPTR(out), CVOIDPTR(inp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_CVOIDPTR(inp), FF_CVOIDPTR(grid), grid.shape, inp.shape, out.strides, inp.strides, grid.strides, stream) } diff --git a/src/lib-cuda/pushpull_backward.cpp b/src/lib-cuda/pushpull_backward.cpp index f0cf88a..ea8f17d 100644 --- a/src/lib-cuda/pushpull_backward.cpp +++ b/src/lib-cuda/pushpull_backward.cpp @@ -207,29 +207,29 @@ void pull_backward( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, gout) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, ginp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") - CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, gout) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, ginp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); // `out` mirrors `inp` (the field), `gout` mirrors `grid`. for (int32_t d = 0; d < out.ndim; ++d) - CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") + FF_CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") for (int32_t d = 0; d < gout.ndim; ++d) - CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") - CHECK_SAME(ginp.shape[ginp.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(inp, grid, nbatch) - CHECK_SAME_BATCH(ginp, grid, nbatch) - - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(gout) - && CANUSE32BITS(inp) && CANUSE32BITS(ginp) - && CANUSE32BITS(grid); + FF_CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") + FF_CHECK_SAME(ginp.shape[ginp.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME_BATCH(ginp, grid, nbatch) + + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(gout) + && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(ginp) + && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -241,8 +241,8 @@ void pull_backward( DISPATCH_PP(_pull_backward, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), VOIDPTR(gout), - CVOIDPTR(inp), CVOIDPTR(ginp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_VOIDPTR(gout), + FF_CVOIDPTR(inp), FF_CVOIDPTR(ginp), FF_CVOIDPTR(grid), grid.shape, inp.shape, out.strides, gout.strides, inp.strides, ginp.strides, grid.strides, stream) } @@ -273,29 +273,29 @@ void push_backward( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, gout) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, ginp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") - CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, gout) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, ginp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); // Here both `out` and `inp` are grid-shaped; `ginp` is the field. for (int32_t d = 0; d < out.ndim; ++d) - CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") + FF_CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") for (int32_t d = 0; d < gout.ndim; ++d) - CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") - CHECK_SAME(ginp.shape[ginp.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(inp, grid, nbatch) - CHECK_SAME_BATCH(ginp, grid, nbatch) - - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(gout) - && CANUSE32BITS(inp) && CANUSE32BITS(ginp) - && CANUSE32BITS(grid); + FF_CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") + FF_CHECK_SAME(ginp.shape[ginp.ndim-1], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME_BATCH(ginp, grid, nbatch) + + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(gout) + && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(ginp) + && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -308,8 +308,8 @@ void push_backward( DISPATCH_PP(_push_backward, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(out), VOIDPTR(gout), - CVOIDPTR(inp), CVOIDPTR(ginp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_VOIDPTR(gout), + FF_CVOIDPTR(inp), FF_CVOIDPTR(ginp), FF_CVOIDPTR(grid), grid.shape, ginp.shape, out.strides, gout.strides, inp.strides, ginp.strides, grid.strides, stream) } @@ -336,20 +336,20 @@ void count_backward( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (gout) - CHECK_SAME_DTYPE(gout, ginp) - CHECK_SAME_DTYPE(gout, grid) - CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") - CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") + FF_CHECK_NO_LANES (gout) + FF_CHECK_SAME_DTYPE(gout, ginp) + FF_CHECK_SAME_DTYPE(gout, grid) + FF_CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") + FF_CHECK_SAME(ginp.ndim, grid.ndim, "ginp and grid must have the same rank") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); for (int32_t d = 0; d < gout.ndim; ++d) - CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") - CHECK_SAME(ginp.shape[ginp.ndim-1], 1, "count gradient must have a single channel") - CHECK_SAME_BATCH(ginp, grid, nbatch) + FF_CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") + FF_CHECK_SAME(ginp.shape[ginp.ndim-1], 1, "count gradient must have a single channel") + FF_CHECK_SAME_BATCH(ginp, grid, nbatch) - const bool use_32bits = CANUSE32BITS(gout) && CANUSE32BITS(ginp) - && CANUSE32BITS(grid); + const bool use_32bits = FF_CANUSE32BITS(gout) && FF_CANUSE32BITS(ginp) + && FF_CANUSE32BITS(grid); const auto code = static_cast(gout.dtype.code); const auto bits = gout.dtype.bits; const spline_t spl = static_cast(spline); @@ -361,7 +361,7 @@ void count_backward( DISPATCH_PP(_count_backward, bvec, svec, static_cast(nbatch), n1, ex, - VOIDPTR(gout), CVOIDPTR(ginp), CVOIDPTR(grid), + FF_VOIDPTR(gout), FF_CVOIDPTR(ginp), FF_CVOIDPTR(grid), grid.shape, ginp.shape, gout.strides, ginp.strides, grid.strides, stream) } @@ -393,29 +393,29 @@ void grad_backward( const int ndim = static_cast(grid.shape[grid.ndim - 1]); const int32_t nbatch = grid.ndim - ndim - 1; const int64_t n1 = grid.ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, gout) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, ginp) - CHECK_SAME_DTYPE(out, grid) - CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") - CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") - CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") - CHECK_SAME(ginp.ndim, grid.ndim + 1, "ginp must have an extra trailing axis") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, gout) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, ginp) + FF_CHECK_SAME_DTYPE(out, grid) + FF_CHECK_SAME(out.ndim, grid.ndim, "out and grid must have the same rank") + FF_CHECK_SAME(gout.ndim, grid.ndim, "gout and grid must have the same rank") + FF_CHECK_SAME(inp.ndim, grid.ndim, "inp and grid must have the same rank") + FF_CHECK_SAME(ginp.ndim, grid.ndim + 1, "ginp must have an extra trailing axis") if (nbatch < 0) throw std::invalid_argument("grid rank is too small for the coordinate dim"); for (int32_t d = 0; d < out.ndim; ++d) - CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") + FF_CHECK_SAME(out.shape[d], inp.shape[d], "out and inp must have the same shape") for (int32_t d = 0; d < gout.ndim; ++d) - CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") - CHECK_SAME(ginp.shape[ginp.ndim-1], ndim, "ginp trailing axis must equal ndim") - CHECK_SAME(ginp.shape[ginp.ndim-2], inp.shape[inp.ndim-1], "channel counts differ") - CHECK_SAME_BATCH(inp, grid, nbatch) - CHECK_SAME_BATCH(ginp, grid, nbatch) - - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(gout) - && CANUSE32BITS(inp) && CANUSE32BITS(ginp) - && CANUSE32BITS(grid); + FF_CHECK_SAME(gout.shape[d], grid.shape[d], "gout and grid must have the same shape") + FF_CHECK_SAME(ginp.shape[ginp.ndim-1], ndim, "ginp trailing axis must equal ndim") + FF_CHECK_SAME(ginp.shape[ginp.ndim-2], inp.shape[inp.ndim-1], "channel counts differ") + FF_CHECK_SAME_BATCH(inp, grid, nbatch) + FF_CHECK_SAME_BATCH(ginp, grid, nbatch) + + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(gout) + && FF_CANUSE32BITS(inp) && FF_CANUSE32BITS(ginp) + && FF_CANUSE32BITS(grid); const auto code = static_cast(out.dtype.code); const auto bits = out.dtype.bits; const spline_t spl = static_cast(spline); @@ -427,8 +427,8 @@ void grad_backward( DISPATCH_PP(_grad_backward, bvec, svec, static_cast(nbatch), n1, ex, abs, - VOIDPTR(out), VOIDPTR(gout), - CVOIDPTR(inp), CVOIDPTR(ginp), CVOIDPTR(grid), + FF_VOIDPTR(out), FF_VOIDPTR(gout), + FF_CVOIDPTR(inp), FF_CVOIDPTR(ginp), FF_CVOIDPTR(grid), grid.shape, inp.shape, out.strides, gout.strides, inp.strides, ginp.strides, grid.strides, stream) } diff --git a/src/lib-cuda/reg_field.cpp b/src/lib-cuda/reg_field.cpp index 47ccdd9..3c84e1d 100644 --- a/src/lib-cuda/reg_field.cpp +++ b/src/lib-cuda/reg_field.cpp @@ -5,6 +5,8 @@ #include "fastfields/api/cuda/reg_field.h" #include "fastfields/api/cuda/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" +#include "fastfields/api/cuda/stream.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/bounds.h" @@ -14,56 +16,14 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_SHAPE(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same shape"); - /*********************************************************************** * WRAPPERS * ***********************************************************************/ namespace { -// int -> cudaStream_t (0 == default stream). The public ABI carries the stream -// as an int; the cuda-impl launchers take a cudaStream_t. Mirrors -// pushpull::_pp_stream in the cuda-impl layer. -static inline cudaStream_t _reg_stream(intptr_t stream) -{ - return reinterpret_cast(static_cast(stream)); -} - -// build a length-nc reduce_t vector from a (possibly null) double array -static inline std::vector as_weights(const double * w, int64_t nc) -{ - std::vector v(static_cast(nc), reduce_t(0)); - if (w) for (int64_t c = 0; c < nc; ++c) v[static_cast(c)] = w[c]; - return v; -} - template inline void _field_matvec( const bound::BoundVec & bvec, @@ -511,23 +471,23 @@ void field_matvec( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ +#define MV_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, inp.strides, cstream NDIM_SWITCH(MV_DT) #undef MV_ARGS @@ -555,23 +515,23 @@ void field_addmatvec_( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ +#define MV_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, inp.strides, cstream NDIM_SWITCH(ADD_MV_DT) #undef MV_ARGS @@ -599,23 +559,23 @@ void field_submatvec_( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) const int64_t nc = out.shape[out.ndim - 1]; - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, \ +#define MV_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, inp.strides, cstream NDIM_SWITCH(SUB_MV_DT) #undef MV_ARGS @@ -637,20 +597,20 @@ void field_diag( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define DG_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, cstream NDIM_SWITCH(DG_DT) #undef DG_ARGS @@ -676,20 +636,20 @@ void field_adddiag_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define DG_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, cstream NDIM_SWITCH(ADD_DG_DT) #undef DG_ARGS @@ -715,20 +675,20 @@ void field_subdiag_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define DG_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define DG_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, cstream NDIM_SWITCH(SUB_DG_DT) #undef DG_ARGS @@ -750,20 +710,20 @@ void field_kernel( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define KN_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, cstream NDIM_SWITCH(KN_DT) #undef KN_ARGS @@ -789,20 +749,20 @@ void field_addkernel_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define KN_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, cstream NDIM_SWITCH(ADD_KN_DT) #undef KN_ARGS @@ -828,20 +788,20 @@ void field_subkernel_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) 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 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); -#define KN_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(out), \ - voxel_size, absolute, membrane, bending, \ +#define KN_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(out), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, cstream NDIM_SWITCH(SUB_KN_DT) #undef KN_ARGS @@ -868,26 +828,26 @@ void field_relax( const DLTensor & g = _grd.t; const int32_t nbatch = s.ndim - ndim - 1; - CHECK_NO_LANES (s) - CHECK_SAME_DTYPE(s, h) - CHECK_SAME_DTYPE(s, g) - CHECK_SAME (s.ndim, g.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (s.ndim, h.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (s) + FF_CHECK_SAME_DTYPE(s, h) + FF_CHECK_SAME_DTYPE(s, g) + FF_CHECK_SAME (s.ndim, g.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (s.ndim, h.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(s, g, s.ndim) + FF_CHECK_SAME_SHAPE_N(s, g, s.ndim) const int64_t nc = s.shape[s.ndim - 1]; - const bool use_32bits = CANUSE32BITS(s) && CANUSE32BITS(h) && - CANUSE32BITS(g); + const bool use_32bits = FF_CANUSE32BITS(s) && FF_CANUSE32BITS(h) && + FF_CANUSE32BITS(g); const auto code = static_cast(s.dtype.code); const auto bits = s.dtype.bits; const bound::type bnd = static_cast(bound); const bound::BoundVec bvec(bnd); const cudaStream_t cstream = _reg_stream(stream); -#define RX_ARGS bvec, static_cast(nbatch), nc, VOIDPTR(s), CVOIDPTR(h), \ - CVOIDPTR(g), voxel_size, absolute, membrane, bending, \ +#define RX_ARGS bvec, static_cast(nbatch), nc, FF_VOIDPTR(s), FF_CVOIDPTR(h), \ + FF_CVOIDPTR(g), voxel_size, absolute, membrane, bending, \ nb_iter, s.shape, s.strides, h.strides, g.strides, cstream NDIM_SWITCH(RX_DT) #undef RX_ARGS @@ -954,7 +914,7 @@ void field_precond( int ndim , intptr_t stream ) { - CHECK_NO_LANES(grd) + FF_CHECK_NO_LANES(grd) if (grd.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); @@ -981,7 +941,7 @@ void field_precond_( int ndim , intptr_t stream ) { - CHECK_NO_LANES(sol) + FF_CHECK_NO_LANES(sol) if (sol.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); diff --git a/src/lib-cuda/reg_field_rls.cpp b/src/lib-cuda/reg_field_rls.cpp index ab7fc8f..ecef8a0 100644 --- a/src/lib-cuda/reg_field_rls.cpp +++ b/src/lib-cuda/reg_field_rls.cpp @@ -5,6 +5,8 @@ #include "fastfields/api/cuda/reg_field.h" #include "fastfields/api/cuda/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" +#include "fastfields/api/cuda/stream.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/bounds.h" @@ -14,55 +16,14 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_SHAPE(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same shape"); - /*********************************************************************** * WRAPPERS * ***********************************************************************/ namespace { -// int -> cudaStream_t (0 == default stream). The public ABI carries the stream -// as an int; the cuda-impl launchers take a cudaStream_t. Mirrors -// pushpull::_pp_stream in the cuda-impl layer. -static inline cudaStream_t _reg_stream(intptr_t stream) -{ - return reinterpret_cast(static_cast(stream)); -} - -// build a length-nc reduce_t vector from a (possibly null) double array -static inline std::vector as_weights(const double * w, int64_t nc) -{ - std::vector v(static_cast(nc), reduce_t(0)); - if (w) for (int64_t c = 0; c < nc; ++c) v[static_cast(c)] = w[c]; - return v; -} // Reweighted-least-squares (RLS/JRLS) variant of `_field_matvec`: an extra // per-voxel weight map `wgt` modulates the penalty strength. `is_jrls` // selects between the per-channel-weight (RLS) and single-shared-weight @@ -393,28 +354,28 @@ void field_matvec_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, wgt) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, wgt) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, inp, out.ndim) - CHECK_SAME_SHAPE(out, wgt, out.ndim - 1) + FF_CHECK_SAME_SHAPE_N(out, inp, out.ndim) + FF_CHECK_SAME_SHAPE_N(out, wgt, out.ndim - 1) 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 bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_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); const cudaStream_t cstream = _reg_stream(stream); -#define RLS_MV_ARGS bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(out), \ - CVOIDPTR(inp), CVOIDPTR(wgt), \ - voxel_size, absolute, membrane, bending, \ +#define RLS_MV_ARGS bvec, static_cast(nbatch), nc, is_jrls, FF_VOIDPTR(out), \ + FF_CVOIDPTR(inp), FF_CVOIDPTR(wgt), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, inp.strides, wgt.strides, cstream NDIM_SWITCH(RLS_MV_DT) #undef RLS_MV_ARGS @@ -438,25 +399,25 @@ void field_diag_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, wgt) - CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, wgt) + FF_CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(out, wgt, out.ndim - 1) + FF_CHECK_SAME_SHAPE_N(out, wgt, out.ndim - 1) 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 bool use_32bits = FF_CANUSE32BITS(out) && FF_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); const cudaStream_t cstream = _reg_stream(stream); -#define RLS_DG_ARGS bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(out), \ - CVOIDPTR(wgt), \ - voxel_size, absolute, membrane, bending, \ +#define RLS_DG_ARGS bvec, static_cast(nbatch), nc, is_jrls, FF_VOIDPTR(out), \ + FF_CVOIDPTR(wgt), \ + voxel_size, absolute, membrane, bending, \ out.shape, out.strides, wgt.strides, cstream NDIM_SWITCH(RLS_DG_DT) #undef RLS_DG_ARGS @@ -485,32 +446,32 @@ void field_relax_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = s.ndim - ndim - 1; - CHECK_NO_LANES (s) - CHECK_SAME_DTYPE(s, h) - CHECK_SAME_DTYPE(s, g) - CHECK_SAME_DTYPE(s, wgt) - CHECK_SAME (s.ndim, g.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (s.ndim, h.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (s.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (s) + FF_CHECK_SAME_DTYPE(s, h) + FF_CHECK_SAME_DTYPE(s, g) + FF_CHECK_SAME_DTYPE(s, wgt) + FF_CHECK_SAME (s.ndim, g.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (s.ndim, h.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (s.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_SHAPE(s, g, s.ndim) - CHECK_SAME_SHAPE(s, wgt, s.ndim - 1) + FF_CHECK_SAME_SHAPE_N(s, g, s.ndim) + FF_CHECK_SAME_SHAPE_N(s, wgt, s.ndim - 1) const int64_t nc = s.shape[s.ndim - 1]; const bool is_jrls = field_rls_is_jrls(wgt, nc, "field_relax_rls"); - const bool use_32bits = CANUSE32BITS(s) && CANUSE32BITS(h) && - CANUSE32BITS(g) && CANUSE32BITS(wgt); + const bool use_32bits = FF_CANUSE32BITS(s) && FF_CANUSE32BITS(h) && + FF_CANUSE32BITS(g) && FF_CANUSE32BITS(wgt); const auto code = static_cast(s.dtype.code); const auto bits = s.dtype.bits; const bound::type bnd = static_cast(bound); const bound::BoundVec bvec(bnd); const cudaStream_t cstream = _reg_stream(stream); -#define RLS_RX_ARGS bvec, static_cast(nbatch), nc, is_jrls, VOIDPTR(s), \ - CVOIDPTR(h), CVOIDPTR(g), CVOIDPTR(wgt), \ - voxel_size, absolute, membrane, bending, \ - nb_iter, s.shape, s.strides, h.strides, g.strides, \ +#define RLS_RX_ARGS bvec, static_cast(nbatch), nc, is_jrls, FF_VOIDPTR(s), \ + FF_CVOIDPTR(h), FF_CVOIDPTR(g), FF_CVOIDPTR(wgt), \ + voxel_size, absolute, membrane, bending, \ + nb_iter, s.shape, s.strides, h.strides, g.strides, \ wgt.strides, cstream NDIM_SWITCH(RLS_RX_DT) #undef RLS_RX_ARGS diff --git a/src/lib-cuda/reg_flow.cpp b/src/lib-cuda/reg_flow.cpp index 57dcf66..e5d8436 100644 --- a/src/lib-cuda/reg_flow.cpp +++ b/src/lib-cuda/reg_flow.cpp @@ -4,6 +4,8 @@ #include "fastfields/api/cuda/reg_flow.h" #include "fastfields/api/cuda/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" +#include "fastfields/api/cuda/stream.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/bounds.h" @@ -13,49 +15,15 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - // reduction / accumulation type (matches jitfields' float64 default) typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_SHAPE(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same shape"); - /*********************************************************************** * WRAPPERS * ***********************************************************************/ namespace { -// int -> cudaStream_t (0 == default stream). The public ABI carries the stream -// as an int; the cuda-impl launchers take a cudaStream_t. Mirrors -// pushpull::_pp_stream in the cuda-impl layer. -static inline cudaStream_t _reg_stream(intptr_t stream) -{ - return reinterpret_cast(static_cast(stream)); -} - // length of the shape/stride arrays: (*batch, *spatial, C) == out.ndim template inline void _flow_matvec( @@ -527,23 +495,23 @@ void flow_matvec( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME_SHAPE(out, inp, out.ndim) + 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 = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ +#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 @@ -573,23 +541,23 @@ void flow_addmatvec_( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME_SHAPE(out, inp, out.ndim) + 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 = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ +#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 @@ -619,23 +587,23 @@ void flow_submatvec_( const DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME_SHAPE(out, inp, out.ndim) + 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 = CANUSE32BITS(out) && CANUSE32BITS(inp); + 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); -#define MV_ARGS bvec, static_cast(nbatch), VOIDPTR(out), CVOIDPTR(inp), \ - voxel_size, absolute, membrane, bending, shears, div, \ +#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 @@ -659,19 +627,19 @@ void flow_diag( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define DG_ARGS bvec, static_cast(nbatch), VOIDPTR(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) @@ -700,19 +668,19 @@ void flow_adddiag_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define DG_ARGS bvec, static_cast(nbatch), VOIDPTR(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) @@ -741,19 +709,19 @@ void flow_subdiag_( DLTensor & out = _out.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define DG_ARGS bvec, static_cast(nbatch), VOIDPTR(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) @@ -784,22 +752,22 @@ void flow_kernel( const int ntrail = is_matrix ? 2 : 1; const int32_t nbatch = out.ndim - ndim - ntrail; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") if (is_matrix) - CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, + FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define KN_ARGS bvec, static_cast(nbatch), VOIDPTR(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) @@ -834,22 +802,22 @@ void flow_addkernel_( const int ntrail = is_matrix ? 2 : 1; const int32_t nbatch = out.ndim - ndim - ntrail; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") if (is_matrix) - CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, + FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define KN_ARGS bvec, static_cast(nbatch), VOIDPTR(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) @@ -884,22 +852,22 @@ void flow_subkernel_( const int ntrail = is_matrix ? 2 : 1; const int32_t nbatch = out.ndim - ndim - ntrail; - CHECK_NO_LANES(out) + FF_CHECK_NO_LANES(out) if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME(out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") if (is_matrix) - CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, + FF_CHECK_SAME(out.shape[out.ndim-2], (int64_t)ndim, "Lamé kernel needs a trailing (ndim, ndim) matrix axis") - const bool use_32bits = CANUSE32BITS(out); + 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); -#define KN_ARGS bvec, static_cast(nbatch), VOIDPTR(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) @@ -923,19 +891,19 @@ void flow_relax( ) { const int32_t nbatch = sol.ndim - ndim - 1; - CHECK_NO_LANES (sol) - CHECK_SAME_DTYPE(sol, hes) - CHECK_SAME_DTYPE(sol, grd) - CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (sol) + FF_CHECK_SAME_DTYPE(sol, hes) + FF_CHECK_SAME_DTYPE(sol, grd) + FF_CHECK_SAME (sol.ndim, grd.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (sol.ndim, hes.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (sol.shape[sol.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME (grd.shape[grd.ndim-1], (int64_t)ndim, "Gradient channel dimension must equal ndim") - CHECK_SAME_SHAPE(sol, grd, sol.ndim) + FF_CHECK_SAME (sol.shape[sol.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + 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 = CANUSE32BITS(sol) && CANUSE32BITS(hes) && - CANUSE32BITS(grd); + 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); @@ -943,9 +911,9 @@ void flow_relax( const cudaStream_t cstream = reinterpret_cast(static_cast(stream)); -#define RX_ARGS bvec, static_cast(nbatch), VOIDPTR(sol), CVOIDPTR(hes), \ - CVOIDPTR(grd), voxel_size, absolute, membrane, bending, \ - shears, div, nb_iter, sol.shape, sol.strides, hes.strides, \ +#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 @@ -1018,7 +986,7 @@ void flow_precond( int ndim , intptr_t stream ) { - CHECK_NO_LANES(grd) + FF_CHECK_NO_LANES(grd) if (grd.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); @@ -1047,7 +1015,7 @@ void flow_precond_( int ndim , intptr_t stream ) { - CHECK_NO_LANES(sol) + FF_CHECK_NO_LANES(sol) if (sol.ndim - ndim - 1 < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); diff --git a/src/lib-cuda/reg_flow_rls.cpp b/src/lib-cuda/reg_flow_rls.cpp index 4dcd61f..5aa6b69 100644 --- a/src/lib-cuda/reg_flow_rls.cpp +++ b/src/lib-cuda/reg_flow_rls.cpp @@ -4,6 +4,8 @@ #include "fastfields/api/cuda/reg_flow.h" #include "fastfields/api/cuda/posdef.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" +#include "fastfields/api/cuda/stream.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/bounds.h" @@ -13,49 +15,15 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CVOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - // reduction / accumulation type (matches jitfields' float64 default) typedef double reduce_t; -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument("Only scalar data types are supported"); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ((X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes)) \ - throw std::invalid_argument("Tensors do not have the same data type"); - -#define CHECK_SAME_SHAPE(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument("Tensors do not have the same shape"); - /*********************************************************************** * WRAPPERS * ***********************************************************************/ namespace { -// int -> cudaStream_t (0 == default stream). The public ABI carries the stream -// as an int; the cuda-impl launchers take a cudaStream_t. Mirrors -// pushpull::_pp_stream in the cuda-impl layer. -static inline cudaStream_t _reg_stream(intptr_t stream) -{ - return reinterpret_cast(static_cast(stream)); -} - // Reweighted-least-squares (JRLS) variant of `_flow_matvec`: an extra // per-voxel weight map `wgt` modulates the penalty strength. A non-zero // shears/div selects the weighted Lamé stencil (which also folds in @@ -315,28 +283,28 @@ void flow_matvec_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME_DTYPE(out, wgt) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME_DTYPE(out, wgt) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME_SHAPE(out, inp, out.ndim) - CHECK_SAME_SHAPE(out, wgt, out.ndim - 1) - CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, + 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) + FF_CHECK_SAME_SHAPE_N(out, wgt, out.ndim - 1) + FF_CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "flow_matvec_rls: weight tensor's trailing dimension must be 1") - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp) && CANUSE32BITS(wgt); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp) && FF_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); const cudaStream_t cstream = _reg_stream(stream); -#define RLS_MV_ARGS bvec, static_cast(nbatch), VOIDPTR(out), CVOIDPTR(inp), \ - CVOIDPTR(wgt), voxel_size, absolute, membrane, shears, div, \ +#define RLS_MV_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), FF_CVOIDPTR(inp), \ + FF_CVOIDPTR(wgt), voxel_size, absolute, membrane, shears, div, \ out.shape, out.strides, inp.strides, wgt.strides, cstream NDIM_SWITCH(RLS_MV_DT) #undef RLS_MV_ARGS @@ -364,25 +332,25 @@ void flow_diag_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = out.ndim - ndim - 1; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, wgt) - CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, wgt) + FF_CHECK_SAME (out.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME_SHAPE(out, wgt, out.ndim - 1) - CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, + FF_CHECK_SAME (out.shape[out.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME_SHAPE_N(out, wgt, out.ndim - 1) + FF_CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "flow_diag_rls: weight tensor's trailing dimension must be 1") - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(wgt); + const bool use_32bits = FF_CANUSE32BITS(out) && FF_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); const cudaStream_t cstream = _reg_stream(stream); -#define RLS_DG_ARGS bvec, static_cast(nbatch), VOIDPTR(out), CVOIDPTR(wgt), \ - voxel_size, absolute, membrane, shears, div, \ +#define RLS_DG_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(out), FF_CVOIDPTR(wgt), \ + voxel_size, absolute, membrane, shears, div, \ out.shape, out.strides, wgt.strides, cstream NDIM_SWITCH(RLS_DG_DT) #undef RLS_DG_ARGS @@ -415,33 +383,33 @@ void flow_relax_rls( const DLTensor & wgt = _wgt.t; const int32_t nbatch = s.ndim - ndim - 1; - CHECK_NO_LANES (s) - CHECK_SAME_DTYPE(s, h) - CHECK_SAME_DTYPE(s, g) - CHECK_SAME_DTYPE(s, wgt) - CHECK_SAME (s.ndim, g.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (s.ndim, h.ndim, "Tensors do not have the same number of dimensions") - CHECK_SAME (s.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (s) + FF_CHECK_SAME_DTYPE(s, h) + FF_CHECK_SAME_DTYPE(s, g) + FF_CHECK_SAME_DTYPE(s, wgt) + FF_CHECK_SAME (s.ndim, g.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (s.ndim, h.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_SAME (s.ndim, wgt.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME (s.shape[s.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") - CHECK_SAME (g.shape[g.ndim-1], (int64_t)ndim, "Gradient channel dimension must equal ndim") - CHECK_SAME_SHAPE(s, g, s.ndim) - CHECK_SAME_SHAPE(s, wgt, s.ndim - 1) - CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, + FF_CHECK_SAME (s.shape[s.ndim-1], (int64_t)ndim, "Channel dimension must equal ndim") + FF_CHECK_SAME (g.shape[g.ndim-1], (int64_t)ndim, "Gradient channel dimension must equal ndim") + FF_CHECK_SAME_SHAPE_N(s, g, s.ndim) + FF_CHECK_SAME_SHAPE_N(s, wgt, s.ndim - 1) + FF_CHECK_SAME (wgt.shape[wgt.ndim-1], (int64_t)1, "flow_relax_rls: weight tensor's trailing dimension must be 1") - const bool use_32bits = CANUSE32BITS(s) && CANUSE32BITS(h) && - CANUSE32BITS(g) && CANUSE32BITS(wgt); + const bool use_32bits = FF_CANUSE32BITS(s) && FF_CANUSE32BITS(h) && + FF_CANUSE32BITS(g) && FF_CANUSE32BITS(wgt); const auto code = static_cast(s.dtype.code); const auto bits = s.dtype.bits; const bound::type bnd = static_cast(bound); const bound::BoundVec bvec(bnd); const cudaStream_t cstream = _reg_stream(stream); -#define RLS_RX_ARGS bvec, static_cast(nbatch), VOIDPTR(s), CVOIDPTR(h), \ - CVOIDPTR(g), CVOIDPTR(wgt), voxel_size, absolute, membrane, \ - shears, div, nb_iter, s.shape, s.strides, h.strides, \ +#define RLS_RX_ARGS bvec, static_cast(nbatch), FF_VOIDPTR(s), FF_CVOIDPTR(h), \ + FF_CVOIDPTR(g), FF_CVOIDPTR(wgt), voxel_size, absolute, membrane, \ + shears, div, nb_iter, s.shape, s.strides, h.strides, \ g.strides, wgt.strides, cstream NDIM_SWITCH(RLS_RX_DT) #undef RLS_RX_ARGS diff --git a/src/lib-cuda/resize.cpp b/src/lib-cuda/resize.cpp index 3bc6c2e..af931c1 100644 --- a/src/lib-cuda/resize.cpp +++ b/src/lib-cuda/resize.cpp @@ -2,6 +2,7 @@ #include #include "fastfields/api/cuda/resize.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -10,39 +11,6 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ( \ - (X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes) \ - ) \ - throw std::invalid_argument( \ - "Tensors do not have the same data type" \ - ); - -#define CHECK_SAME_BATCH(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument( \ - "Tensors do not have the same batch shape" \ - ); - /*********************************************************************** * DISPATCH * ***********************************************************************/ @@ -130,19 +98,19 @@ inline void _resample( default: throw std::invalid_argument("Unsupported spline order"); \ } -#define DISPATCH_RESIZE(args...) \ -{ \ - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); \ - const auto code = static_cast(inp.dtype.code); \ - const spline_t spl = static_cast(spline); \ - const bound_t bnd = static_cast(bound); \ - switch (ndim) { \ - case 1: RS_ORDER(1, args); break; \ - case 2: RS_ORDER(2, args); break; \ - case 3: RS_ORDER(3, args); break; \ - default: throw std::invalid_argument( \ - "Only 1D, 2D and 3D resize are supported"); \ - }; \ +#define DISPATCH_RESIZE(args...) \ +{ \ + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); \ + const auto code = static_cast(inp.dtype.code); \ + const spline_t spl = static_cast(spline); \ + const bound_t bnd = static_cast(bound); \ + switch (ndim) { \ + case 1: RS_ORDER(1, args); break; \ + case 2: RS_ORDER(2, args); break; \ + case 3: RS_ORDER(3, args); break; \ + default: throw std::invalid_argument( \ + "Only 1D, 2D and 3D resize are supported"); \ + }; \ } void resample( @@ -163,17 +131,17 @@ void resample( DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_BATCH(out, inp, nbatch) + FF_CHECK_SAME_BATCH(out, inp, nbatch) DISPATCH_RESIZE( static_cast(nbatch), - VOIDPTR(out), - VOIDPTR(inp), + FF_VOIDPTR(out), + FF_VOIDPTR(inp), shift, scale, out.shape, diff --git a/src/lib-cuda/restrict.cpp b/src/lib-cuda/restrict.cpp index d4f812f..fd0b373 100644 --- a/src/lib-cuda/restrict.cpp +++ b/src/lib-cuda/restrict.cpp @@ -2,6 +2,7 @@ #include #include "fastfields/api/cuda/restrict.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -10,39 +11,6 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - -#define CHECK_SAME(X, Y, msg) \ - if (X != Y) throw std::invalid_argument(msg); - -#define CHECK_SAME_DTYPE(X, Y) \ - if ( \ - (X.dtype.code != Y.dtype.code) || \ - (X.dtype.bits != Y.dtype.bits) || \ - (X.dtype.lanes != Y.dtype.lanes) \ - ) \ - throw std::invalid_argument( \ - "Tensors do not have the same data type" \ - ); - -#define CHECK_SAME_BATCH(X, Y, D) \ - for (int32_t d=0; d < D; ++d) \ - if (X.shape[d] != Y.shape[d]) \ - throw std::invalid_argument( \ - "Tensors do not have the same batch shape" \ - ); - /*********************************************************************** * DISPATCH * ***********************************************************************/ @@ -130,19 +98,19 @@ inline void _restriction( default: throw std::invalid_argument("Unsupported spline order"); \ } -#define DISPATCH_RESTRICT(args...) \ -{ \ - const bool use_32bits = CANUSE32BITS(out) && CANUSE32BITS(inp); \ - const auto code = static_cast(inp.dtype.code); \ - const spline_t spl = static_cast(spline); \ - const bound_t bnd = static_cast(bound); \ - switch (ndim) { \ - case 1: RT_ORDER(1, args); break; \ - case 2: RT_ORDER(2, args); break; \ - case 3: RT_ORDER(3, args); break; \ - default: throw std::invalid_argument( \ - "Only 1D, 2D and 3D restrict are supported"); \ - }; \ +#define DISPATCH_RESTRICT(args...) \ +{ \ + const bool use_32bits = FF_CANUSE32BITS(out) && FF_CANUSE32BITS(inp); \ + const auto code = static_cast(inp.dtype.code); \ + const spline_t spl = static_cast(spline); \ + const bound_t bnd = static_cast(bound); \ + switch (ndim) { \ + case 1: RT_ORDER(1, args); break; \ + case 2: RT_ORDER(2, args); break; \ + case 3: RT_ORDER(3, args); break; \ + default: throw std::invalid_argument( \ + "Only 1D, 2D and 3D restrict are supported"); \ + }; \ } void restriction( @@ -163,17 +131,17 @@ void restriction( DLTensor & inp = _inp.t; const int32_t nbatch = out.ndim - ndim; - CHECK_NO_LANES (out) - CHECK_SAME_DTYPE(out, inp) - CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") + FF_CHECK_NO_LANES (out) + FF_CHECK_SAME_DTYPE(out, inp) + FF_CHECK_SAME (out.ndim, inp.ndim, "Tensors do not have the same number of dimensions") if (nbatch < 0) throw std::invalid_argument("ndim is larger than the tensor rank"); - CHECK_SAME_BATCH(out, inp, nbatch) + FF_CHECK_SAME_BATCH(out, inp, nbatch) DISPATCH_RESTRICT( static_cast(nbatch), - VOIDPTR(out), - VOIDPTR(inp), + FF_VOIDPTR(out), + FF_VOIDPTR(inp), shift, scale, out.shape, diff --git a/src/lib-cuda/splinc.cpp b/src/lib-cuda/splinc.cpp index 39e7382..1b3ac60 100644 --- a/src/lib-cuda/splinc.cpp +++ b/src/lib-cuda/splinc.cpp @@ -3,6 +3,7 @@ #include #include "fastfields/api/cuda/splinc.h" #include "fastfields/core/autocast.h" +#include "fastfields/core/dispatch.h" #include "fastfields/core/dlpack.h" #include "fastfields/core/cuda_switch.h" #include "fastfields/impl/kernels/utils.h" @@ -11,19 +12,6 @@ FF_NAMESPACE_BEGIN(FF) FF_NAMESPACE_BEGIN(FF_DEVICE) -#define VOIDPTR(x) (static_cast(static_cast(x.data) + x.byte_offset)) -#define CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides)) - -/*********************************************************************** - * CHECKS * - ***********************************************************************/ - -#define CHECK_NO_LANES(tensor) \ - if (tensor.dtype.lanes > 1) \ - throw std::invalid_argument( \ - "Only scalar data types are supported" \ - ); - /*********************************************************************** * POLES * ***********************************************************************/ @@ -107,22 +95,22 @@ inline void _splinc( default: throw std::invalid_argument("Unsupported npoles"); \ } -#define DISPATCH_SPLINC(args...) \ -{ \ - const bool use_32bits = CANUSE32BITS(inp_out); \ - const auto code = static_cast(inp_out.dtype.code); \ - switch (code) { \ - case kDLFloat: switch (inp_out.dtype.bits) { \ - case 32: \ - if (use_32bits) DISPATCH_SPLINC_NPOLES(float, int32_t, args) \ - else DISPATCH_SPLINC_NPOLES(float, int64_t, args) \ - case 64: \ - if (use_32bits) DISPATCH_SPLINC_NPOLES(double, int32_t, args) \ - else DISPATCH_SPLINC_NPOLES(double, int64_t, args) \ - default: break; \ - }; \ - default: break; \ - }; \ +#define DISPATCH_SPLINC(args...) \ +{ \ + const bool use_32bits = FF_CANUSE32BITS(inp_out); \ + const auto code = static_cast(inp_out.dtype.code); \ + switch (code) { \ + case kDLFloat: switch (inp_out.dtype.bits) { \ + case 32: \ + if (use_32bits) DISPATCH_SPLINC_NPOLES(float, int32_t, args) \ + else DISPATCH_SPLINC_NPOLES(float, int64_t, args) \ + case 64: \ + if (use_32bits) DISPATCH_SPLINC_NPOLES(double, int32_t, args) \ + else DISPATCH_SPLINC_NPOLES(double, int64_t, args) \ + default: break; \ + }; \ + default: break; \ + }; \ throw std::invalid_argument("only floating point data types are supported"); \ } @@ -137,7 +125,7 @@ void spline_coeff( ContiguousStrides _io(inp_out_); DLTensor & inp_out = _io.t; - CHECK_NO_LANES(inp_out) + FF_CHECK_NO_LANES(inp_out) double poles[3]; const int npoles = get_poles_host(static_cast(spline), poles); @@ -148,7 +136,7 @@ void spline_coeff( DISPATCH_SPLINC( nbatch, - VOIDPTR(inp_out), + FF_VOIDPTR(inp_out), inp_out.shape, inp_out.strides, poles diff --git a/src/lib/distance.cpp b/src/lib/distance.cpp index b38b27e..e053376 100644 --- a/src/lib/distance.cpp +++ b/src/lib/distance.cpp @@ -9,23 +9,19 @@ FF_NAMESPACE_BEGIN(FF) -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - void dt_euclidean( DLTensor & inp_out , double voxel_spacing , intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(inp_out)) + if (is_cuda(inp_out)) return FF_CUDA::dt_euclidean(inp_out, voxel_spacing, stream); #endif - if (IS_CPU(inp_out)) + if (is_cpu(inp_out)) return FF_CPU::dt_euclidean(inp_out, voxel_spacing, stream); - if (IS_CUDA(inp_out)) + if (is_cuda(inp_out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -36,13 +32,13 @@ void dt_l1( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(inp_out)) + if (is_cuda(inp_out)) return FF_CUDA::dt_l1(inp_out, voxel_spacing, stream); #endif - if (IS_CPU(inp_out)) + if (is_cpu(inp_out)) return FF_CPU::dt_l1(inp_out, voxel_spacing, stream); - if (IS_CUDA(inp_out)) + if (is_cuda(inp_out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -59,13 +55,13 @@ void dt_spline_table( { require_same_device(loc, time, dist, coeff, times); #ifdef FF_WITH_CUDA - if (IS_CUDA(loc)) + if (is_cuda(loc)) return FF_CUDA::dt_spline_table(time, dist, loc, coeff, times, spline, bound, stream); #endif - if (IS_CPU(loc)) + if (is_cpu(loc)) return FF_CPU::dt_spline_table(time, dist, loc, coeff, times, spline, bound, stream); - if (IS_CUDA(loc)) + if (is_cuda(loc)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -84,13 +80,13 @@ void dt_spline_brent( { require_same_device(loc, time, dist, coeff); #ifdef FF_WITH_CUDA - if (IS_CUDA(loc)) + if (is_cuda(loc)) return FF_CUDA::dt_spline_brent(time, dist, loc, coeff, max_iter, tol, step, spline, bound, stream); #endif - if (IS_CPU(loc)) + if (is_cpu(loc)) return FF_CPU::dt_spline_brent(time, dist, loc, coeff, max_iter, tol, step, spline, bound, stream); - if (IS_CUDA(loc)) + if (is_cuda(loc)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -108,13 +104,13 @@ void dt_spline_gaussnewton( { require_same_device(loc, time, dist, coeff); #ifdef FF_WITH_CUDA - if (IS_CUDA(loc)) + if (is_cuda(loc)) return FF_CUDA::dt_spline_gaussnewton(time, dist, loc, coeff, max_iter, tol, spline, bound, stream); #endif - if (IS_CPU(loc)) + if (is_cpu(loc)) return FF_CPU::dt_spline_gaussnewton(time, dist, loc, coeff, max_iter, tol, spline, bound, stream); - if (IS_CUDA(loc)) + if (is_cuda(loc)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -136,13 +132,13 @@ void dt_mesh( // it carries data (same convention as posdef's optional `weight`). if (nearest_vertex.data) require_same_device(loc, nearest_vertex); #ifdef FF_WITH_CUDA - if (IS_CUDA(loc)) + if (is_cuda(loc)) return FF_CUDA::dt_mesh(dist, nearest_vertex, loc, vertices, faces, _signed, naive, stream); #endif - if (IS_CPU(loc)) + if (is_cpu(loc)) return FF_CPU::dt_mesh(dist, nearest_vertex, loc, vertices, faces, _signed, naive, stream); - if (IS_CUDA(loc)) + if (is_cuda(loc)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } diff --git a/src/lib/posdef.cpp b/src/lib/posdef.cpp index d0fbbf0..8bea319 100644 --- a/src/lib/posdef.cpp +++ b/src/lib/posdef.cpp @@ -9,10 +9,6 @@ FF_NAMESPACE_BEGIN(FF) -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - void sym_matvec( DLTensor & out , const DLTensor & hessian , @@ -21,13 +17,13 @@ void sym_matvec( { require_same_device(out, hessian, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::sym_matvec(out, hessian, inp, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::sym_matvec(out, hessian, inp, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -40,13 +36,13 @@ void sym_matvec_backward( { require_same_device(out, grd, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::sym_matvec_backward(out, grd, inp, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::sym_matvec_backward(out, grd, inp, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -59,13 +55,13 @@ void sym_addmatvec_( { require_same_device(out, hessian, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::sym_addmatvec_(out, hessian, inp, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::sym_addmatvec_(out, hessian, inp, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -78,13 +74,13 @@ void sym_submatvec_( { require_same_device(out, hessian, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::sym_submatvec_(out, hessian, inp, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::sym_submatvec_(out, hessian, inp, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -99,13 +95,13 @@ void sym_solve( require_same_device(out, hessian, inp); if (weight.data) require_same_device(out, weight); // weight is optional (null-data placeholder) #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::sym_solve(out, hessian, inp, weight, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::sym_solve(out, hessian, inp, weight, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -119,13 +115,13 @@ void sym_solve_( require_same_device(inp_out, hessian); if (weight.data) require_same_device(inp_out, weight); // weight is optional (null-data placeholder) #ifdef FF_WITH_CUDA - if (IS_CUDA(inp_out)) + if (is_cuda(inp_out)) return FF_CUDA::sym_solve_(inp_out, hessian, weight, stream); #endif - if (IS_CPU(inp_out)) + if (is_cpu(inp_out)) return FF_CPU::sym_solve_(inp_out, hessian, weight, stream); - if (IS_CUDA(inp_out)) + if (is_cuda(inp_out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -137,13 +133,13 @@ void sym_invert( { require_same_device(out, hessian); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::sym_invert(out, hessian, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::sym_invert(out, hessian, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -153,13 +149,13 @@ void sym_invert_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(hessian)) + if (is_cuda(hessian)) return FF_CUDA::sym_invert_(hessian, stream); #endif - if (IS_CPU(hessian)) + if (is_cpu(hessian)) return FF_CPU::sym_invert_(hessian, stream); - if (IS_CUDA(hessian)) + if (is_cuda(hessian)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } diff --git a/src/lib/pushpull.cpp b/src/lib/pushpull.cpp index 14a1dbd..5d8e0bd 100644 --- a/src/lib/pushpull.cpp +++ b/src/lib/pushpull.cpp @@ -7,10 +7,6 @@ #include "fastfields/api/cuda/pushpull.h" #endif -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - FF_NAMESPACE_BEGIN(FF) void pull( @@ -24,13 +20,13 @@ void pull( { require_same_device(out, inp, grid); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::pull(out, inp, grid, spline, bound, extrapolate, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::pull(out, inp, grid, spline, bound, extrapolate, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -46,13 +42,13 @@ void push( { require_same_device(out, inp, grid); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::push(out, inp, grid, spline, bound, extrapolate, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::push(out, inp, grid, spline, bound, extrapolate, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -67,13 +63,13 @@ void count( { require_same_device(out, grid); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::count(out, grid, spline, bound, extrapolate, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::count(out, grid, spline, bound, extrapolate, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -90,13 +86,13 @@ void grad( { require_same_device(out, inp, grid); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::grad(out, inp, grid, spline, bound, extrapolate, abs, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::grad(out, inp, grid, spline, bound, extrapolate, abs, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -118,13 +114,13 @@ void pull_backward( { require_same_device(out, gout, inp, ginp, grid); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::pull_backward(out, gout, inp, ginp, grid, spline, bound, extrapolate, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::pull_backward(out, gout, inp, ginp, grid, spline, bound, extrapolate, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -142,13 +138,13 @@ void push_backward( { require_same_device(out, gout, inp, ginp, grid); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::push_backward(out, gout, inp, ginp, grid, spline, bound, extrapolate, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::push_backward(out, gout, inp, ginp, grid, spline, bound, extrapolate, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -164,13 +160,13 @@ void count_backward( { require_same_device(gout, ginp, grid); #ifdef FF_WITH_CUDA - if (IS_CUDA(gout)) + if (is_cuda(gout)) return FF_CUDA::count_backward(gout, ginp, grid, spline, bound, extrapolate, stream); #endif - if (IS_CPU(gout)) + if (is_cpu(gout)) return FF_CPU::count_backward(gout, ginp, grid, spline, bound, extrapolate, stream); - if (IS_CUDA(gout)) + if (is_cuda(gout)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -189,13 +185,13 @@ void grad_backward( { require_same_device(out, gout, inp, ginp, grid); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::grad_backward(out, gout, inp, ginp, grid, spline, bound, extrapolate, abs, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::grad_backward(out, gout, inp, ginp, grid, spline, bound, extrapolate, abs, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } diff --git a/src/lib/reg_field.cpp b/src/lib/reg_field.cpp index a977035..52d23ce 100644 --- a/src/lib/reg_field.cpp +++ b/src/lib/reg_field.cpp @@ -7,10 +7,6 @@ #include "fastfields/api/cuda/reg_field.h" #endif -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - FF_NAMESPACE_BEGIN(FF) void field_matvec( @@ -26,13 +22,13 @@ void field_matvec( { require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_matvec(out, inp, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_matvec(out, inp, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -50,13 +46,13 @@ void field_diag( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_diag(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_diag(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -72,13 +68,13 @@ void field_kernel( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_kernel(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_kernel(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -99,13 +95,13 @@ void field_addmatvec_( { require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_addmatvec_(out, inp, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_addmatvec_(out, inp, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -126,13 +122,13 @@ void field_submatvec_( { require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_submatvec_(out, inp, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_submatvec_(out, inp, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -151,13 +147,13 @@ void field_adddiag_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_adddiag_(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_adddiag_(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -176,13 +172,13 @@ void field_subdiag_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_subdiag_(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_subdiag_(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -201,13 +197,13 @@ void field_addkernel_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_addkernel_(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_addkernel_(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -226,13 +222,13 @@ void field_subkernel_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_subkernel_(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_subkernel_(out, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -251,13 +247,13 @@ void field_relax( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(sol)) + if (is_cuda(sol)) return FF_CUDA::field_relax(sol, hes, grd, voxel_size, absolute, membrane, bending, bound, ndim, nb_iter, stream); #endif - if (IS_CPU(sol)) + if (is_cpu(sol)) return FF_CPU::field_relax(sol, hes, grd, voxel_size, absolute, membrane, bending, bound, ndim, nb_iter, stream); - if (IS_CUDA(sol)) + if (is_cuda(sol)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -277,13 +273,13 @@ void field_forward( require_same_device(out, hes); require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_forward(out, hes, inp, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_forward(out, hes, inp, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -303,13 +299,13 @@ void field_precond( require_same_device(out, hes); require_same_device(out, grd); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_precond(out, hes, grd, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_precond(out, hes, grd, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -327,13 +323,13 @@ void field_precond_( { require_same_device(sol, hes); #ifdef FF_WITH_CUDA - if (IS_CUDA(sol)) + if (is_cuda(sol)) return FF_CUDA::field_precond_(sol, hes, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(sol)) + if (is_cpu(sol)) return FF_CPU::field_precond_(sol, hes, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(sol)) + if (is_cuda(sol)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -353,13 +349,13 @@ void field_matvec_rls( require_same_device(out, inp); require_same_device(out, wgt); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_matvec_rls(out, inp, wgt, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_matvec_rls(out, inp, wgt, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -377,13 +373,13 @@ void field_diag_rls( { require_same_device(out, wgt); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::field_diag_rls(out, wgt, voxel_size, absolute, membrane, bending, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::field_diag_rls(out, wgt, voxel_size, absolute, membrane, bending, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -404,13 +400,13 @@ void field_relax_rls( { require_same_device(sol, wgt); #ifdef FF_WITH_CUDA - if (IS_CUDA(sol)) + if (is_cuda(sol)) return FF_CUDA::field_relax_rls(sol, hes, grd, wgt, voxel_size, absolute, membrane, bending, bound, ndim, nb_iter, stream); #endif - if (IS_CPU(sol)) + if (is_cpu(sol)) return FF_CPU::field_relax_rls(sol, hes, grd, wgt, voxel_size, absolute, membrane, bending, bound, ndim, nb_iter, stream); - if (IS_CUDA(sol)) + if (is_cuda(sol)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } diff --git a/src/lib/reg_flow.cpp b/src/lib/reg_flow.cpp index 072edae..0e89b46 100644 --- a/src/lib/reg_flow.cpp +++ b/src/lib/reg_flow.cpp @@ -7,10 +7,6 @@ #include "fastfields/api/cuda/reg_flow.h" #endif -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - FF_NAMESPACE_BEGIN(FF) void flow_matvec( @@ -28,13 +24,13 @@ void flow_matvec( { require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_matvec(out, inp, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_matvec(out, inp, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -54,13 +50,13 @@ void flow_diag( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_diag(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_diag(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -78,13 +74,13 @@ void flow_kernel( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_kernel(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_kernel(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -107,13 +103,13 @@ void flow_addmatvec_( { require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_addmatvec_(out, inp, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_addmatvec_(out, inp, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -136,13 +132,13 @@ void flow_submatvec_( { require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_submatvec_(out, inp, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_submatvec_(out, inp, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -163,13 +159,13 @@ void flow_adddiag_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_adddiag_(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_adddiag_(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -190,13 +186,13 @@ void flow_subdiag_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_subdiag_(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_subdiag_(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -217,13 +213,13 @@ void flow_addkernel_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_addkernel_(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_addkernel_(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -244,13 +240,13 @@ void flow_subkernel_( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_subkernel_(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_subkernel_(out, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -271,10 +267,10 @@ void flow_relax( intptr_t stream ) { #ifdef FF_WITH_CUDA - if (IS_CUDA(sol)) + if (is_cuda(sol)) return FF_CUDA::flow_relax(sol, hes, grd, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, nb_iter, stream); #endif - if (IS_CPU(sol)) + if (is_cpu(sol)) return FF_CPU::flow_relax(sol, hes, grd, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, nb_iter, stream); throw std::invalid_argument("unsupported device"); @@ -297,13 +293,13 @@ void flow_forward( require_same_device(out, hes); require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_forward(out, hes, inp, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_forward(out, hes, inp, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -325,13 +321,13 @@ void flow_precond( require_same_device(out, hes); require_same_device(out, grd); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_precond(out, hes, grd, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_precond(out, hes, grd, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -351,13 +347,13 @@ void flow_precond_( { require_same_device(sol, hes); #ifdef FF_WITH_CUDA - if (IS_CUDA(sol)) + if (is_cuda(sol)) return FF_CUDA::flow_precond_(sol, hes, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(sol)) + if (is_cpu(sol)) return FF_CPU::flow_precond_(sol, hes, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(sol)) + if (is_cuda(sol)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -379,13 +375,13 @@ void flow_matvec_rls( require_same_device(out, inp); require_same_device(out, wgt); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_matvec_rls(out, inp, wgt, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_matvec_rls(out, inp, wgt, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -405,13 +401,13 @@ void flow_diag_rls( { require_same_device(out, wgt); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::flow_diag_rls(out, wgt, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::flow_diag_rls(out, wgt, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } @@ -434,13 +430,13 @@ void flow_relax_rls( { require_same_device(sol, wgt); #ifdef FF_WITH_CUDA - if (IS_CUDA(sol)) + if (is_cuda(sol)) return FF_CUDA::flow_relax_rls(sol, hes, grd, wgt, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, nb_iter, stream); #endif - if (IS_CPU(sol)) + if (is_cpu(sol)) return FF_CPU::flow_relax_rls(sol, hes, grd, wgt, voxel_size, absolute, membrane, bending, shears, div, bound, ndim, nb_iter, stream); - if (IS_CUDA(sol)) + if (is_cuda(sol)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } diff --git a/src/lib/resize.cpp b/src/lib/resize.cpp index d15ba79..f4f3a26 100644 --- a/src/lib/resize.cpp +++ b/src/lib/resize.cpp @@ -7,10 +7,6 @@ #include "fastfields/api/cuda/resize.h" #endif -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - FF_NAMESPACE_BEGIN(FF) void resample( @@ -25,13 +21,13 @@ void resample( { require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::resample(out, inp, spline, bound, shift, scale, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::resample(out, inp, spline, bound, shift, scale, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } diff --git a/src/lib/restrict.cpp b/src/lib/restrict.cpp index a885d0b..3273fb5 100644 --- a/src/lib/restrict.cpp +++ b/src/lib/restrict.cpp @@ -7,10 +7,6 @@ #include "fastfields/api/cuda/restrict.h" #endif -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - FF_NAMESPACE_BEGIN(FF) void restriction( @@ -25,13 +21,13 @@ void restriction( { require_same_device(out, inp); #ifdef FF_WITH_CUDA - if (IS_CUDA(out)) + if (is_cuda(out)) return FF_CUDA::restriction(out, inp, spline, bound, shift, scale, ndim, stream); #endif - if (IS_CPU(out)) + if (is_cpu(out)) return FF_CPU::restriction(out, inp, spline, bound, shift, scale, ndim, stream); - if (IS_CUDA(out)) + if (is_cuda(out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } diff --git a/src/lib/solve_field.cpp b/src/lib/solve_field.cpp index 7a32684..7aa3c50 100644 --- a/src/lib/solve_field.cpp +++ b/src/lib/solve_field.cpp @@ -3,10 +3,6 @@ #include "fastfields/api/checks.h" #include "fastfields/api/cpu/solve_field.h" -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - FF_NAMESPACE_BEGIN(FF) // The CG driver is CPU-only for now: unlike the other modules there is no @@ -32,12 +28,12 @@ void field_cg( { require_same_device(sol, hes, grd); - if (IS_CPU(sol)) + if (is_cpu(sol)) return FF_CPU::field_cg(sol, hes, grd, voxel_size, absolute, membrane, bending, bound, ndim, nb_iter, tol, nb_iter_out, residual_out, stream); - if (IS_CUDA(sol)) + if (is_cuda(sol)) throw std::invalid_argument( "fastfields: field_cg is not implemented on CUDA yet"); throw std::invalid_argument("unsupported device"); diff --git a/src/lib/splinc.cpp b/src/lib/splinc.cpp index 17b1d5f..ebeb769 100644 --- a/src/lib/splinc.cpp +++ b/src/lib/splinc.cpp @@ -1,15 +1,12 @@ #include #include #include "fastfields/api/splinc.h" +#include "fastfields/api/checks.h" #include "fastfields/api/cpu/splinc.h" #ifdef FF_WITH_CUDA #include "fastfields/api/cuda/splinc.h" #endif -#define IS_CUDA(tensor) (tensor.device.device_type == DLDeviceType::kDLCUDA) -#define IS_CPU(tensor) (tensor.device.device_type == DLDeviceType::kDLCPU || \ - tensor.device.device_type == DLDeviceType::kDLCUDAHost) - FF_NAMESPACE_BEGIN(FF) void spline_coeff( @@ -24,13 +21,13 @@ void spline_coeff( require_splinc_bound(spline, bound); #ifdef FF_WITH_CUDA - if (IS_CUDA(inp_out)) + if (is_cuda(inp_out)) return FF_CUDA::spline_coeff(inp_out, spline, bound, stream); #endif - if (IS_CPU(inp_out)) + if (is_cpu(inp_out)) return FF_CPU::spline_coeff(inp_out, spline, bound, stream); - if (IS_CUDA(inp_out)) + if (is_cuda(inp_out)) throw std::invalid_argument("fastfields: built without CUDA support, cannot operate on CUDA tensors"); throw std::invalid_argument("unsupported device"); } diff --git a/tools/dedup-dispatch-helpers.py b/tools/dedup-dispatch-helpers.py new file mode 100644 index 0000000..27606f9 --- /dev/null +++ b/tools/dedup-dispatch-helpers.py @@ -0,0 +1,323 @@ +#!/usr/bin/env python3 +""" +dedup-dispatch-helpers.py -- hoist the copy-pasted dispatch helpers into the +shared headers, and rewrite every call site onto their FF_-prefixed names. + +This is the script that produced the de-duplication commit. It is committed +rather than merely described because the change touches 32 files and ~1100 call +sites: a reviewer who wants to confirm the rewrite was mechanical can re-run it +against the parent commit and diff, instead of reading 1100 renamed lines. + + git checkout -- include src + python3 tools/dedup-dispatch-helpers.py + git diff # must be empty + +It is idempotent: on an already-converted tree the local #defines are gone and +the call sites already carry the prefix, so a second run is a no-op. + +WHAT MOVES WHERE, AND WHY +-------------------------------------------------------------------------- +Placement is decided by *audience*, and the audience question that actually +bites here is which compiler sees the code: + +include/fastfields/core/dispatch.h (new) + FF_VOIDPTR / FF_CVOIDPTR / FF_CVOIDPTR_OR_NULL / FF_CANUSE32BITS, the + FF_CHECK_* family, and as_weights(). Shared by src/lib-cpu (host compiler) + AND src/lib-cuda (nvcc) AND the two api/*/pushpull_dispatch.h headers, so + it must be backend-agnostic -> core/, which is the only directory that is + backend-agnostic by contract. + +include/fastfields/api/cuda/stream.h (new) + _reg_stream(). Names cudaStream_t, so it is CUDA-only by construction and + belongs under api/cuda/, not core/. Putting it in core/ would also have + mis-tagged it for CI's path filter. + +include/fastfields/api/checks.h (extended) + is_cpu() / is_cuda(), replacing the IS_CPU / IS_CUDA macros. Used only by + src/lib (the hub), which is host-compiled -- the hub's own validation + header is the right home, and functions are collision-safe without a + prefix, so these two macros are deleted rather than renamed. + +THE DIVERGENCES, PRESERVED +-------------------------------------------------------------------------- +Three of the duplicated macros were NOT identical across their copies. Each +keeps its own name so the difference is visible at the call site: + + CVOIDPTR posdef.cpp alone guarded against a null `data` (optional + tensors) -> FF_CVOIDPTR_OR_NULL there, FF_CVOIDPTR + everywhere else. + CHECK_SAME_BATCH distance.cpp / posdef.cpp alone also rejected ndim < D + -> FF_CHECK_SAME_BATCH_ND there, FF_CHECK_SAME_BATCH + everywhere else. + CHECK_SAME_SHAPE TWO DIFFERENT MACROS SHARING ONE NAME: a 2-argument + whole-shape check in distance.cpp, a 3-argument leading-D + check in the regularisers and solve_field + -> FF_CHECK_SAME_SHAPE and FF_CHECK_SAME_SHAPE_N. + +CHECK_NO_LANES and CHECK_SAME_DTYPE also had two spellings each, but those +differ only in line wrapping and expand identically. + +Equivalence is not asserted, it is proved: tools/macro-equivalence.py +preprocesses every (file, macro) pair on both sides and compares token streams. +""" + +import os +import re +import sys + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +# -------------------------------------------------------------------------- +# Which files carry which variant +# -------------------------------------------------------------------------- + +BACKENDS = ("lib-cpu", "lib-cuda") + +DISPATCH_SOURCES = sorted( + os.path.join("src", b, f) + for b in BACKENDS + for f in os.listdir(os.path.join(ROOT, "src", b)) + if f.endswith(".cpp") +) + [ + "include/fastfields/api/cpu/pushpull_dispatch.h", + "include/fastfields/api/cuda/pushpull_dispatch.h", +] + +HUB_SOURCES = sorted( + os.path.join("src", "lib", f) + for f in os.listdir(os.path.join(ROOT, "src", "lib")) + if f.endswith(".cpp") +) + +# The null-tolerant CVOIDPTR lived only in posdef. +CVOIDPTR_OR_NULL = {"src/lib-cpu/posdef.cpp", "src/lib-cuda/posdef.cpp"} + +# The ndim-guarded CHECK_SAME_BATCH lived only in distance and posdef. +BATCH_ND = { + "src/lib-cpu/distance.cpp", + "src/lib-cpu/posdef.cpp", + "src/lib-cuda/distance.cpp", + "src/lib-cuda/posdef.cpp", +} + +# CHECK_SAME_SHAPE's 2-argument whole-shape form lived only in distance; every +# other user had the 3-argument leading-D form. +SHAPE_WHOLE = {"src/lib-cpu/distance.cpp", "src/lib-cuda/distance.cpp"} + +# Macros hoisted into core/dispatch.h. Order is irrelevant to correctness (the +# renames are whole-identifier) but longest-first reads more obviously safe. +HOISTED = [ + "CHECK_SAME_DTYPE", + "CHECK_SAME_BATCH", + "CHECK_SAME_SHAPE", + "CHECK_NO_LANES", + "CHECK_SAME", + "CANUSE32BITS", + "CVOIDPTR", + "VOIDPTR", +] + +HUB_MACROS = ["IS_CUDA", "IS_CPU"] + +# The macro names this script introduces. A backslash-continuation block that +# mentions one of them is a block whose hand-alignment the rename just broke. +NEW_NAMES = re.compile( + r"(? ...`, its backslash continuations, and the single + blank line that separated it from whatever came next.""" + return re.sub( + r"^[ \t]*#[ \t]*define[ \t]+" + re.escape(name) + r"\b" + r"(?:[^\n]*\\\n)*[^\n]*\n" + r"(?:[ \t]*\n)?", + "", + text, + flags=re.MULTILINE, + ) + + +def rename(text, old, new): + """Whole-identifier rename.""" + return re.sub( + r"(?"` right after the `after` include, once.""" + line = '#include "%s"\n' % header + if line in text: + return text + anchor = '#include "%s"\n' % after + if anchor not in text: + raise SystemExit("no anchor %r to insert %r after" % (after, header)) + return text.replace(anchor, anchor + line, 1) + + +def realign_continuations(text): + """Re-align the trailing backslash of every multi-line macro this rewrite + made longer. The tree is hand-column-aligned and CI lints the lines a PR + touches, so leaving hundreds of renamed lines with their backslashes shoved + out of column would be the most visible thing in the diff.""" + lines = text.split("\n") + out, i = [], 0 + while i < len(lines): + if not lines[i].endswith("\\"): + out.append(lines[i]) + i += 1 + continue + j = i + while j < len(lines) and lines[j].endswith("\\"): + j += 1 + block = lines[i : j + 1] # include the block's final, unbackslashed line + if not any(NEW_NAMES.search(ln) for ln in block): + out.extend(block) + else: + bodies = [ln[:-1].rstrip() if ln.endswith("\\") else ln for ln in block] + width = max(len(b) for b in bodies[:-1]) + 1 + for k, body in enumerate(bodies): + out.append(body if k == len(bodies) - 1 else body.ljust(width) + "\\") + i = j + 1 + return "\n".join(out) + + +# -------------------------------------------------------------------------- +# Pass 1 -- the two dtype-dispatch layers +# -------------------------------------------------------------------------- + + +def convert_dispatch(rel): + text = original = read(rel) + + for name in HOISTED: + text = drop_define(text, name) + # pushpull.cpp / pushpull_backward.cpp only *use* these macros -- their + # definitions live in the shared pushpull_dispatch.h, which is itself in + # this list. Such a file gets its call sites renamed but no new include: + # the dispatch header already pulls core/dispatch.h in. + defined_here = text != original + + # Per-file variant selection, before the generic renames below. + if rel in CVOIDPTR_OR_NULL: + text = rename(text, "CVOIDPTR", "FF_CVOIDPTR_OR_NULL") + if rel in BATCH_ND: + text = rename(text, "CHECK_SAME_BATCH", "FF_CHECK_SAME_BATCH_ND") + if rel in SHAPE_WHOLE: + text = rename(text, "CHECK_SAME_SHAPE", "FF_CHECK_SAME_SHAPE") + else: + text = rename(text, "CHECK_SAME_SHAPE", "FF_CHECK_SAME_SHAPE_N") + + for name in HOISTED: + text = rename(text, name, "FF_" + name) + + # as_weights: three byte-identical copies -> core/dispatch.h. + text = re.sub( + r"\n// build a length-nc reduce_t vector from a \(possibly null\) double array\n" + r"static inline std::vector as_weights\(const double \* w, int64_t nc\)\n" + r"\{\n(?:[^\n]*\n)*?\}\n(?:[ \t]*\n)?", + "\n", + text, + ) + + # _reg_stream: four byte-identical copies -> api/cuda/stream.h. + text = re.sub( + r"\n// int -> cudaStream_t \(0 == default stream\)\.[^\n]*\n" + r"(?:// [^\n]*\n)*" + r"static inline cudaStream_t _reg_stream\(intptr_t stream\)\n" + r"\{\n(?:[^\n]*\n)*?\}\n(?:[ \t]*\n)?", + "\n", + text, + ) + + if text == original: + return False + + if defined_here: + text = drop_banner(text, "CHECKS") + text = add_include( + text, "fastfields/core/dispatch.h", "fastfields/core/autocast.h" + ) + if rel.startswith("src/lib-cuda/") and "_reg_stream" in text: + text = add_include( + text, "fastfields/api/cuda/stream.h", "fastfields/core/dispatch.h" + ) + text = realign_continuations(text) + write(rel, text) + return True + + +# -------------------------------------------------------------------------- +# Pass 2 -- the hub +# -------------------------------------------------------------------------- + + +def convert_hub(rel): + text = original = read(rel) + + for name in HUB_MACROS: + text = drop_define(text, name) + text = rename(text, "IS_CUDA", "is_cuda") + text = rename(text, "IS_CPU", "is_cpu") + + if text == original: + return False + + # Every hub source but splinc.cpp already includes api/checks.h for + # require_same_device -- splinc's single in/out tensor has no second + # operand to disagree with it, so it never needed the header before. + if '#include "fastfields/api/checks.h"' not in text: + module = os.path.basename(rel)[: -len(".cpp")] + text = add_include( + text, "fastfields/api/checks.h", "fastfields/api/%s.h" % module + ) + write(rel, text) + return True + + +def main(): + changed = 0 + for rel in DISPATCH_SOURCES: + if convert_dispatch(rel): + changed += 1 + print("dispatch %s" % rel) + for rel in HUB_SOURCES: + if convert_hub(rel): + changed += 1 + print("hub %s" % rel) + print("%d file(s) rewritten" % changed) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/macro-equivalence.py b/tools/macro-equivalence.py new file mode 100644 index 0000000..52acbe3 --- /dev/null +++ b/tools/macro-equivalence.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +""" +macro-equivalence.py -- prove that each hoisted FF_* macro expands +token-for-token to the local macro it replaced. + +WHY THIS EXISTS +-------------------------------------------------------------------------- +De-duplicating a macro that exists in 19 copies is only safe if the surviving +copy expands to the same tokens as every copy it replaced. Eyeballing 19 macro +bodies does not establish that -- three of these had genuinely diverged, and +one pair were two *different* macros sharing a name. So do not eyeball: run +the preprocessor on both sides and compare the token streams. + + python3 tools/macro-equivalence.py + +`` is a checkout of the parent commit (a `git worktree add --detach` +is the cheapest way to get one); `` is the converted tree. Exits +non-zero on any mismatch, so it can be dropped into CI or a pre-merge check. + +WHAT IT COMPARES +-------------------------------------------------------------------------- +For every (dispatch source, macro) pair that existed on the old tree, it +preprocesses `MACRO()` twice: + + * once with only that file's own local #defines in scope (the old body), and + * once with only core/dispatch.h's FF_* #defines in scope (the new body), + +then tokenises both expansions and compares. Placeholder arguments (`T`, `A`, +`B`, `D`, `MSG`) are used rather than real expressions, so the comparison is of +the macro bodies alone and cannot be confounded by the call sites. + +WHY NOT JUST DIFF WHOLE PREPROCESSED TRANSLATION UNITS +-------------------------------------------------------------------------- +Tried that first; it is the wrong instrument. Two unavoidable, harmless +differences swamp the signal: + + * core/dispatch.h pulls in , which reorders standard-library + declarations in the -E output of every TU that now includes it; + * as_weights() is now visible in every dispatch TU rather than only in the + three that defined it. + +Both are additive and irrelevant to behaviour, but they make a whole-TU token +diff useless. Comparing the macro expansions in isolation asks the question +that actually matters. + +THE ONE ACCEPTED DIFFERENCE +-------------------------------------------------------------------------- +FF_CVOIDPTR_OR_NULL is defined by composing FF_CVOIDPTR, so its expansion +carries one extra *balanced parenthesis pair* around the non-null branch that +posdef.cpp's hand-written CVOIDPTR did not have: + + old: (x.data ? static_cast(...) : nullptr) + new: (x.data ? (static_cast(...)) : nullptr) + +That is reported as OK(parens) after checking the two token streams are +identical once parentheses are removed -- i.e. only grouping changed, and +grouping that was already unambiguous. Every other pair must match exactly. +""" + +import difflib +import os +import re +import subprocess +import sys + +TOK = re.compile(r'"(?:\\.|[^"\\])*"|[A-Za-z_]\w*|\d[\w.]*|[^\s]') + +# Placeholder invocation per macro (keyed by the *old* name, with the two +# same-named CHECK_SAME_SHAPE variants disambiguated). +CALLS = { + "VOIDPTR": "(T)", + "CVOIDPTR": "(T)", + "CANUSE32BITS": "(T)", + "CHECK_NO_LANES": "(T)", + "CHECK_SAME": "(A, B, MSG)", + "CHECK_SAME_DTYPE": "(A, B)", + "CHECK_SAME_BATCH": "(A, B, D)", + "CHECK_SAME_SHAPE2": "(A, B)", + "CHECK_SAME_SHAPEN": "(A, B, D)", +} + +HOISTED = ( + "VOIDPTR", "CVOIDPTR", "CANUSE32BITS", "CHECK_NO_LANES", + "CHECK_SAME", "CHECK_SAME_DTYPE", "CHECK_SAME_BATCH", "CHECK_SAME_SHAPE", +) + +# The three preserved divergences (kept in sync with +# tools/dedup-dispatch-helpers.py). +CVOIDPTR_OR_NULL = {"src/lib-cpu/posdef.cpp", "src/lib-cuda/posdef.cpp"} +BATCH_ND = { + "src/lib-cpu/distance.cpp", "src/lib-cpu/posdef.cpp", + "src/lib-cuda/distance.cpp", "src/lib-cuda/posdef.cpp", +} +SHAPE_WHOLE = {"src/lib-cpu/distance.cpp", "src/lib-cuda/distance.cpp"} + + +def defines_in(path): + """name -> the full `#define` text (continuations included), for one file.""" + lines = open(path, encoding="utf-8").read().split("\n") + out, i = {}, 0 + while i < len(lines): + m = re.match(r"\s*#\s*define\s+([A-Za-z_]\w*)", lines[i]) + if m: + body = [lines[i]] + while body[-1].rstrip().endswith("\\") and i + 1 < len(lines): + i += 1 + body.append(lines[i]) + out[m.group(1)] = "\n".join(body) + i += 1 + return out + + +def expand(prelude, call): + """Preprocess `call` with `prelude` in scope; return its token list.""" + src = prelude + "\nSTART " + call + " END\n" + r = subprocess.run( + ["clang++", "-E", "-P", "-x", "c++", "-"], + input=src, capture_output=True, text=True, + ) + if r.returncode: + return None + t = TOK.findall(r.stdout) + return t[t.index("START") + 1 : len(t) - 1 - t[::-1].index("END")] + + +def new_name_for(rel, old): + if old == "CHECK_SAME_SHAPE": + return ("FF_CHECK_SAME_SHAPE", "CHECK_SAME_SHAPE2") if rel in SHAPE_WHOLE \ + else ("FF_CHECK_SAME_SHAPE_N", "CHECK_SAME_SHAPEN") + if old == "CVOIDPTR": + return ("FF_CVOIDPTR_OR_NULL" if rel in CVOIDPTR_OR_NULL + else "FF_CVOIDPTR", old) + if old == "CHECK_SAME_BATCH": + return ("FF_CHECK_SAME_BATCH_ND" if rel in BATCH_ND + else "FF_CHECK_SAME_BATCH", old) + return "FF_" + old, old + + +def main(base, new): + newdefs = defines_in( + os.path.join(new, "include/fastfields/core/dispatch.h")) + new_prelude = "\n".join(v for k, v in newdefs.items() if k.startswith("FF_")) + + files = [] + for d in ("src/lib-cpu", "src/lib-cuda"): + files += [d + "/" + f + for f in sorted(os.listdir(os.path.join(base, d))) + if f.endswith(".cpp")] + files += ["include/fastfields/api/cpu/pushpull_dispatch.h", + "include/fastfields/api/cuda/pushpull_dispatch.h"] + + checked = bad = parens = 0 + for rel in files: + old = defines_in(os.path.join(base, rel)) + old_prelude = "\n".join(v for k, v in old.items() if k in HOISTED) + for name in HOISTED: + if name not in old: + continue + new_name, call_key = new_name_for(rel, name) + args = CALLS[call_key] + a = expand(old_prelude, name + args) + b = expand(new_prelude, new_name + args) + checked += 1 + if a == b: + continue + if (new_name == "FF_CVOIDPTR_OR_NULL" and a and b and + [t for t in a if t not in "()"] == + [t for t in b if t not in "()"]): + parens += 1 + print(" OK(parens) %-24s %s -> %s" % (rel, name, new_name)) + continue + bad += 1 + print("MISMATCH %s %s -> %s" % (rel, name, new_name)) + for tag, i1, i2, j1, j2 in difflib.SequenceMatcher( + None, a or [], b or [], autojunk=False).get_opcodes(): + if tag == "equal": + continue + print(" OLD: %s" % " ".join((a or [])[i1:i2])[:200]) + print(" NEW: %s" % " ".join((b or [])[j1:j2])[:200]) + + print("\n%d macro expansions compared across %d files: " + "%d exact, %d balanced-paren-only, %d MISMATCHED" + % (checked, len(files), checked - parens - bad, parens, bad)) + return 1 if bad else 0 + + +if __name__ == "__main__": + if len(sys.argv) != 3: + sys.exit(__doc__) + sys.exit(main(sys.argv[1], sys.argv[2]))