//===-- cuda/cmath - one division, made fast, and nothing else -----------===//
//
// `cuda::fast_mod_div<unsigned int>`, which is the whole of what our
// FlashInfer closure reaches through `#include <cuda/cmath>`. This is what
// that directive resolves to when the compiler is NVRTC and the include path
// is a header set carried in the binary rather than a directory on a disk.
//
// # Why a shim rather than libcu++
//
// `examples/header_probe.rs` measured NVRTC 13.0 on this box against an empty
// header set: `<cuda/cmath>` fails with *"could not open source file ... (no
// directories in search list)"*, one of 31 external directives out of the
// closure of which **0 were answered**. NVRTC bundles no libcu++ and there is
// no version of this where it finds one by itself.
//
// The alternative is to put CCCL in the header set. That was measured too:
// **13,691,725 bytes across 1,691 files**, against which `cooperative_groups.h`
// and `cuda/std/limits` together came to 14,637 -- a factor of 935. This file
// is the third door in that count, and it is a door the closure opens exactly
// once.
//
// # The name is the trick, again
//
// No extension, a directory in the path, because the directive has both:
// `fastdiv.cuh:19` says `#include <cuda/cmath>`, NVRTC matches
// `includeNames[]` against that literal string, and so the entry -- and, so
// that one file is one spelling, the path on disk -- is `cuda/cmath`. The
// nesting is part of a name; NVRTC walks no directories. See
// `cooperative_groups.h` for the rest of that argument.
//
// # The surface is one class, because one class is what is called
//
// Counted over the 28-file, 17,981-line attention closure and over the whole
// of `kernels-cuda/csrc/src` -- the archive crate's tree, whose `.cuh` half is
// this crate's `csrc/src` now:
//
// | name | closure | our csrc |
// |---|---|---|
// | `cuda::fast_mod_div` | 2 | -- |
// | `cuda::ceil_div`, `cuda::round_up`, `cuda::round_down`, `cuda::div` | 0 | 0 |
// | any other `cuda::` name outside `cuda::std::` | 0 | 0 |
//
// The two mentions are one declaration and one use, both in `fastdiv.cuh`:
//
//     cuda::fast_mod_div<uint32_t> impl_;          // :39
//     __host__ __device__ uint_fastdiv() : impl_(1), d_(0) {}   // :27
//     __host__ uint_fastdiv(uint32_t d) : impl_(d ? d : 1), d_(d) {}  // :29
//     q = n / impl_;  r = n - q * d_;              // :34-35, in divmod()
//
// That is the entire contract: an `explicit` constructor from a 32-bit
// unsigned, callable from `__host__` and from `__device__` because line 27 is
// both, and a division operator, likewise both because `divmod` is. Nothing
// else is touched, and `ceil_div` -- 74 uses in the closure -- is not this
// header's: it is `flashinfer::ceil_div`, defined in `utils.cuh:339`, and it
// would have resolved with an empty file here.
//
// **Deliberately absent, each with the call site that would want it:**
//
//   * **`operator%`** -- upstream has it; `fastdiv.cuh:35` does not use it.
//     It computes the remainder itself as `n - q * d_` off the quotient it
//     already has, which is one multiply cheaper than a second magic
//     multiply. A `%` here would be dead code carrying a second chance to be
//     wrong.
//   * **`operator T()`** -- upstream converts back to the divisor;
//     `fastdiv.cuh:31` returns its OWN `d_` copy instead, so the conversion
//     is never reached.
//   * **`cuda::div`** -- returns a `cuda::std::pair`, which would drag
//     `<cuda/std/utility>` into the set for a function with zero call sites.
//   * **`ceil_div`, `round_up`, `round_down`, `ilog2`, `isqrt`** -- zero uses
//     of the `cuda::` spellings anywhere in either tree.
//
// Absent means a NAME ERROR at the call site, which is the pattern
// `cooperative_groups.h` set with `this_grid()`: the compiler finds the gap
// and names the line. A shim that answers a call it cannot honestly implement
// is the failure this set is built to avoid.
//
// # The arithmetic, and the one input that breaks it
//
// `q = floor(n/d)` for `n < 2^32` and `1 <= d < 2^32` via one 64-bit
// multiply-high, with `M = floor(2^64 / d) + 1`:
//
//     floor(n*M / 2^64) == floor(n/d)
//
// Write `2^64 = q0*d + r0` with `0 <= r0 < d`, so `M = q0 + 1` and
// `M*d = 2^64 + (d - r0)` with `1 <= d - r0 <= d`. Write `n = k*d + s` with
// `0 <= s < d`. Then `n*M / 2^64 = k + s/d + n(d-r0)/(d * 2^64)`, and the
// result floors to `k` exactly when `s * 2^64 + n(d - r0) < d * 2^64`. Since
// `s <= d - 1`, it suffices that `n(d - r0) < 2^64`, and both factors are at
// most `2^32 - 1`, so the product is at most `(2^32-1)^2 < 2^64`. **Strictly**
// -- the bound has room, which is why this holds for every one of the 2^32
// dividends rather than almost all of them.
//
// **`d == 1` is the exception and it is not a rounding matter.** `2^64/d` is
// then `2^64`, which is not a 64-bit number; the magic wraps to 1 and the
// multiply-high returns 0 for every input. So `d == 1` is a comparison in the
// division, not a special magic. Upstream reached the same conclusion from
// the other end: its second template parameter is called `DivisorIsNeverOne`
// and its documented effect is *"skips one comparison in the second stage"*.
// That parameter is honoured here for exactly that reason, and defaulted to
// `false`, because `fastdiv.cuh:39` writes `fast_mod_div<uint32_t>` with one
// argument and `:29` maps a zero divisor onto `1`.
//
// # Host and device compute the quotient differently, on purpose
//
// On the device the multiply-high is `__umul64hi`, an NVRTC builtin needing
// no header. On the host it is `n / d`, the real division -- because the
// class exists to make the DEVICE fast, the host has a divider in silicon,
// and a second implementation of the magic path is a second thing to get
// wrong. The two agree by the proof above; `examples/fp8_pipeline_probe.rs`
// then measured the device path against the DEVICE's own `u32 / u32` over
// three 4,096-wide divisor bands -- `1..4096`, `2^31 ..+4k` and the top 4k,
// the small-magic, the boundary and the every-quotient-is-0-or-1 cases --
// crossed with 65,548 dividends: **805,453,824 pairs, zero mismatches**,
// because a proof of a bound is worth less than a measurement of the
// instruction.
//
// The reference is the hardware divide rather than libcu++ for a reason worth
// recording: the toolkit's own CCCL, 3.0.1 at
// `/usr/local/cuda-13.0/targets/x86_64-linux/include/cccl`, **has no
// `cuda::fast_mod_div`** -- measured, `nvcc` answers *"namespace cuda has no
// member fast_mod_div"*. The class only exists in the newer CCCL that
// FlashInfer vendors under `3rdparty/`, so there is no vendor implementation
// on this machine to compare against, and the correct reference for a
// division is a division.
//
// The `__host__ __device__` on both members is not decoration. NVRTC compiles
// device code only, and an entity marked `__host__` ALONE is a hard error in
// it -- measured here, verbatim: *"A function explicitly marked as a __host__
// function is not allowed in JIT mode"*. An unannotated function is fine
// (NVRTC treats it as `__device__`); a `__host__ __device__` one is fine; a
// `__host__` one stops the compile. That is why this class carries the double
// annotation rather than nothing at all. It also cost a day, because
// **`fastdiv.cuh:29` -- `__host__ uint_fastdiv(uint32_t d)` -- trips the same
// error**, and the failure reads as if `<cuda/cmath>` were at fault when the
// wall is in FlashInfer's execution-space annotations. It is now guarded in
// `csrc/src/attn/flashinfer/fastdiv.cuh` with `#ifndef __CUDACC_RTC__` -- a
// JIT caller computes the divisor host-side and passes it in a params struct,
// so device code never constructs one. Recorded here anyway, because the next
// person to meet it will open this file first.
//
// # Includes nothing
//
// Not `<cstdint>` -- `unsigned int` and `unsigned long long` are the language's
// and `uint32_t` is the includer's problem, which `fastdiv.cuh:18` solves for
// itself. Not `<cuda/std/type_traits>` -- the one constraint worth stating is
// stated with a `static_assert` on `sizeof`.
//
// # It answers the real includer
//
// Not only the probe: the vendored `flashinfer/fastdiv.cuh` compiles under
// NVRTC 13.0 against this file, and so does the whole 15-root closure that
// reaches it -- measured, 796 carried headers, `-default-device`, no include
// path, no toolkit. `fp8_pipeline_probe` reported that compile on every run;
// with it gone, the crate's own NVRTC fixtures still do.
//
//===----------------------------------------------------------------------===//
#pragma once

namespace cuda {

/// A divisor with its reciprocal magic precomputed, so that a division on the
/// device is a multiply-high instead of the ~20-instruction sequence
/// `ptxas` emits for a runtime `u32 / u32`.
///
/// `_DivisorIsNeverOne` promises `divisor > 1` and drops the comparison in
/// `operator/`. Upstream's name and upstream's meaning; the promise is
/// unchecked in release builds by upstream and unchecked here, because the
/// only thing a check could do is trap.
///
/// Narrower than upstream in one way that matters: **32-bit unsigned only**.
/// The proof above is a proof about `n < 2^32`, and a 64-bit `T` needs a
/// 128-bit intermediate and a different argument entirely. `fastdiv.cuh:39`
/// instantiates `uint32_t` and nothing in either tree instantiates anything
/// else, so the other widths are a `static_assert` that names the reason
/// rather than an implementation nobody can test.
template <class _Tp, bool _DivisorIsNeverOne = false>
class fast_mod_div {
    static_assert(sizeof(_Tp) == 4 && static_cast<_Tp>(-1) > static_cast<_Tp>(0),
                  "pie shim: cuda::fast_mod_div is implemented for 32-bit unsigned types only "
                  "-- the multiply-high proof is a proof about dividends below 2^32. "
                  "Widen it deliberately, with a test, or use the built-in operator.");

public:
    /// Deleted, as upstream deletes it. A `fast_mod_div` with no divisor has
    /// no meaning, and `fastdiv.cuh:27` initialises `impl_(1)` rather than
    /// leaving it blank precisely because it cannot.
    fast_mod_div() = delete;

    /// `__host__ __device__` because `fastdiv.cuh:27` is, and it constructs
    /// one. `explicit` because upstream is; the direct-initialisations at
    /// `:27` and `:29` are unaffected, and an implicit version would let a
    /// bare integer slide into a parameter expecting a divisor object.
    ///
    /// Precondition, upstream's and ours: `divisor > 0`. Zero produces a
    /// magic of 1 and quotients of 0 -- garbage, not a trap, exactly as a
    /// hardware divide by zero on the device produces garbage.
    /// `fastdiv.cuh:29` writes `d ? d : 1` and never passes one.
    __host__ __device__ explicit fast_mod_div(_Tp __divisor) noexcept : __divisor_(__divisor) {
        // floor(2^64/d) without a 128-bit intermediate: 2^64-1 = q*d + r, so
        // 2^64 = q*d + (r+1), which carries into the quotient exactly when
        // r+1 == d.
        const unsigned long long __all_ones = ~0ull;
        const unsigned long long __q = __all_ones / __divisor;
        const unsigned long long __r = __all_ones % __divisor;
        __magic_ = __q + ((__r + 1ull == static_cast<unsigned long long>(__divisor)) ? 1ull : 0ull) + 1ull;
    }

    /// `n / divisor`, as one `mul.hi.u64` on the device.
    ///
    /// A hidden friend: found by argument-dependent lookup on the
    /// `fast_mod_div` operand, which is how `fastdiv.cuh:34` spells it
    /// (`q = n / impl_`), and invisible to every other overload set.
    ///
    /// The host branch is the hardware divide rather than a second copy of
    /// the magic path -- see the file banner. The two are proved equal on
    /// `0 <= n < 2^32`, `0 < d < 2^32`, and were measured equal by the
    /// now-deleted `fp8_pipeline_probe`.
    template <class _Up>
    [[nodiscard]] __host__ __device__ friend _Tp operator/(_Up __dividend,
                                                           fast_mod_div __divisor) noexcept {
        static_assert(sizeof(_Up) <= 4,
                      "pie shim: cuda::fast_mod_div divides 32-bit dividends. A wider one would "
                      "silently leave the range the multiply-high is exact on.");
        const unsigned int __n = static_cast<unsigned int>(__dividend);
#ifdef __CUDA_ARCH__
        if (!_DivisorIsNeverOne && __divisor.__divisor_ == 1u) {
            return static_cast<_Tp>(__n);
        }
        return static_cast<_Tp>(__umul64hi(static_cast<unsigned long long>(__n), __divisor.__magic_));
#else
        return static_cast<_Tp>(__n / __divisor.__divisor_);
#endif
    }

private:
    _Tp __divisor_;
    unsigned long long __magic_;
};

// The layout is 16 bytes and it is NOT CCCL's 16 bytes. Stated here, loudly,
// because the sizes agreeing is a coincidence that hid the difference once.
//
//   this shim  { u32 __divisor_ @0,  u64 __magic_ @8 }   size 16, align 8
//   CCCL       { u32 __divisor  @0,  u32 __multiplier @4,
//                u32 __add      @8,  i32 __shift     @12 } size 16, align 4
//
// FlashInfer's `uint_fastdiv` embeds one of these, so a struct containing it
// has DIFFERENT interior offsets under the two: `paged_kv_t::num_heads`
// measured at **+24 under this shim and +20 under CCCL**, with the totals
// reconverging at 96 by luck of the pointer block that follows. That was found
// by `tests/flashinfer_decode.rs`, which probes the offsets from inside the
// NVRTC compile and from an nvcc harness and compares 26 of them; an earlier
// measurement had checked only `sizeof` and reported the two as agreeing.
//
// # Why this is stated rather than fixed
//
// Matching CCCL's layout means adopting CCCL's ARITHMETIC — a 32-bit
// multiplier with a separate add and shift, where this uses a 64-bit magic —
// and that is a numerical change to a divisor used on every paged-KV address
// computation. It needs its own parity evidence (§8), not a struct
// rearrangement done in passing.
//
// # What it costs, and the rule that follows
//
// Nothing, while a translation unit uses ONE of the two consistently — which
// is what makes the decode kernel bit-exact against nvcc today, each side
// internally agreeing with itself. It costs a wrong answer the moment a
// struct built against one layout is read against the other: a Rust mirror
// pinned with `offset_of!` from the C++ build, a params block filled by the
// ahead-of-time path and fired by the JIT one, a cached plan crossing between
// them.
//
// **So: a Rust mirror of anything containing a `uint_fastdiv` must be pinned
// against THIS layout, and must say so.** The assertions below are what a
// mirror is checked against, and they fail if the shim's own fields ever move.
static_assert(sizeof(fast_mod_div<unsigned int>) == 16,
              "pie shim: fast_mod_div is 16 bytes -- a mirror pinned on that size "
              "is why this is asserted rather than assumed");
static_assert(alignof(fast_mod_div<unsigned int>) == 8,
              "pie shim: fast_mod_div aligns to 8 because of its 64-bit magic, "
              "where CCCL's aligns to 4 -- see the note above; a mirror that "
              "assumed CCCL's alignment reads the wrong bytes of every struct "
              "that embeds one");

}  // namespace cuda
