/* qxmx_fa.cpp: M1.7b.7 causal prefill attention, 3-kernel split per K-tile.
 * See qxmx_fa.h for the design.
 *
 * Per K-tile (TILE_T keys):
 *   1. dequant K-tile -> Kt [NQ][HD][TILE_T]  (transposed for Q@K^T; each
 *      q-head's slot holds its kv-head's K-tile, replicated across GQA=6)
 *   2. QK GEMM (gemm_tf32_batched, DPAS): S[hh][n_tok][TILE_T]
 *      = Qsep[hh][n_tok][HD] @ Kt[hh][HD][TILE_T]
 *   3. online-softmax merge (plain SYCL): causal-mask S, exp, merge into
 *      P + (m, l) running state
 *   4. dequant V-tile -> Vt [NQ][TILE_T][HD]  (TurboQuant rotated domain, T13)
 *   5. PV GEMM (gemm_tf32_batched): O[hh][n_tok][HD] += P[hh][n_tok][TILE_T]
 *      @ Vt[hh][TILE_T][HD]   (accumulates across tiles via Oacc)
 *   6. final: butterfly each 32-elem block of O (T13 inverse WHT) x NRM -> out
 *
 * Why the split (vs M1.7b.6 fused scalar): the scalar HD=256 dot per
 * (token,head,key) capped prefill at 230 tok/s.  The QK and PV GEMMs use
 * DPAS at high occupancy (gemm_tf32_batched, M1.7a's proven 4.2 TFLOPs
 * primitive).  Materializes S/P per tile (~3 MB at CHUNK=256, TILE_T=128),
 * bounded regardless of context length.  The K-tile loop walks
 * 0..pos_start+n_tok-1 in TILE_T-aligned blocks; the final partial tile and
 * the per-token causal cap are masked in the softmax-merge step.
 *
 * M1.7b.1-6 history: M1.7b.4 ESIMD (176 tok/s) and M1.7b.6 plain-SYCL scalar
 * (230 tok/s) both did the QK dot as a scalar loop over HD.  The inline-asm
 * DPAS path (M1.7b.5) was blocked by a vISA JIT asm-dialect wall -- see
 * qxmx_m1_7b_5_dpas_attempt.md.  This split reuses the WORKING ESIMD
 * gemm_tf32_batched (which the M1.7b.5 attempt confirmed REQUIRES
 * [[intel::sycl_explicit_simd]] -- the GEMM kernels have it, the softmax
 * merge kernel does not need it). */
#include "qxmx_fa.h"

#include <cmath>
#include <cstdlib>

#include <sycl/sycl.hpp>

#include "qxmx_fp8.h"
#include "qxmx_gemm_tf32.h"

namespace qx {

static constexpr int HD       = QX_HEAD_DIM;          /* 256 */
static constexpr int NQ       = QX_N_HEAD;            /* 24 */
static constexpr int NKV       = QX_N_HEAD_KV;        /* 4 */
static constexpr int GQA       = NQ / NKV;            /* 6 */
static constexpr int KV_DIM    = QX_KV_DIM;           /* 1024 */
static constexpr int VBLK_L    = VBLK;                /* 32 */
static constexpr int BLK_PER_HEAD = HD / VBLK_L;     /* 8 */
static constexpr int TILE_T    = 128;                 /* K-tokens per block */
static constexpr float NRM     = 0.17677669529663687f; /* 1/sqrt(32) */

/* ---- Step 1: dequant Q [n_tok][NQ*HD] -> Qsep [NQ][n_tok][HD] -------------
 * One WG per (q-head, query-token), LWS=HD lanes cooperatively copy + scale.
 * Also initializes m_state=-inf, l_state=0, Oacc=0 for the chunk. */
static void fa_prep_q(sycl::queue& q, const float* Q, int n_tok,
                      fa_split_workspace* ws) {
    /* Copy struct members to locals -- the device kernel can't dereference
     * the host struct pointer. */
    float* Qsep = ws->Qsep;
    float* Oacc = ws->Oacc;
    float* m_state = ws->m_state;
    float* l_state = ws->l_state;
    q.submit([&](sycl::handler& h) {
        h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)NQ * n_tok * HD),
                                         sycl::range<1>(HD)),
            [=](sycl::nd_item<1> it) {
                int gid = it.get_group(0);
                int hh = gid / n_tok;
                int qi = gid % n_tok;
                int d = it.get_local_id(0);
                float v = Q[(size_t)qi * (NQ * HD) + (size_t)hh * HD + d];
                Qsep[(size_t)hh * n_tok * HD + (size_t)qi * HD + d] = v;
                Oacc[(size_t)hh * n_tok * HD + (size_t)qi * HD + d] = 0.0f;
                if (d == 0) {
                    m_state[(size_t)hh * n_tok + qi] = -1e30f;
                    l_state[(size_t)hh * n_tok + qi] = 0.0f;
                }
            });
    });
}

/* ---- Step 2: dequant K-tile [t0..t0+TILE_T-1] -> Kt [NQ][HD][TILE_T] ------
 * Transposed for Q@K^T: Kt[hh][d][t] = K[hk][t][d].  Each q-head's slot holds
 * its kv-head's K-tile (replicated across GQA).  One WG per (q-head, key-tile
 * element), LWS=HD lanes. */
template <typename Kct>
static void fa_dequant_kt(sycl::queue& q, const uint8_t* kv_k,
                         int64_t t0, int tile_len,
                         fa_split_workspace* ws) {
    int n_wgs = NQ * TILE_T;
    float* Kt = ws->Kt;
    q.submit([&](sycl::handler& h) {
        h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * HD),
                                         sycl::range<1>(HD)),
            [=](sycl::nd_item<1> it) {
                int gid = it.get_group(0);
                int hh = gid / TILE_T;
                int tloc = gid % TILE_T;
                int hk = hh / GQA;
                int d = it.get_local_id(0);
                int64_t kglob = t0 + tloc;
                float kval = 0.0f;
                if (tloc < tile_len) {
                    if constexpr (std::is_same_v<Kct, KFp16>) {
                        const sycl::half* kp =
                            reinterpret_cast<const sycl::half*>(kv_k);
                        kval = float(kp[(size_t)kglob * KV_DIM
                                        + (size_t)hk * HD + d]);
                    } else {
                        const uint8_t* kb = kv_k
                            + (size_t)kglob * Kct::row_bytes
                            + (size_t)(hk * BLK_PER_HEAD + (d >> 5)) * KBLK_BYTES;
                        uint16_t sbits = (uint16_t)kb[0]
                                       | ((uint16_t)kb[1] << 8);
                        float dd = float(sycl::bit_cast<sycl::half>(sbits));
                        uint8_t code = kb[2 + (d & 31)];
                        if constexpr (std::is_same_v<Kct, KQ8_0>)
                            kval = (float)(int8_t)code * dd;
                        else
                            kval = e4m3_to_fp32(code) * dd;
                    }
                }
                /* Kt[hh][d][tloc] = Kt[hh * HD * TILE_T + d * TILE_T + tloc] */
                Kt[(size_t)hh * HD * TILE_T + (size_t)d * TILE_T + tloc] = kval;
            });
    });
}

/* ---- Step 3: QK GEMM via gemm_tf32_batched -------------------------------
 * S[hh][n_tok][TILE_T] = Qsep[hh][n_tok][HD] @ Kt[hh][HD][TILE_T]
 * batch=NQ, M=n_tok, K=HD, N=TILE_T.  Requires n_tok % 8, HD % 8, TILE_T % 16. */
static void fa_qk_gemm(sycl::queue& q, int n_tok, fa_split_workspace* ws) {
    gemm_tf32_batched(q, ws->Qsep, ws->Kt, ws->S, NQ, n_tok, HD, TILE_T);
}

/* ---- Step 4: online-softmax merge (causal-masked) -----------------------
 * One WG per (q-head, query-token), LWS=TILE_T.  For each query token qi, the
 * valid keys in this tile are tloc in [0..tile_len-1] with t0+tloc <= pos_start+qi.
 * S is scaled (1/sqrt(HD)); apply causal mask, merge exp(S - m_new) into P,
 * update m_state/l_state. */
static void fa_softmax_merge(sycl::queue& q, int64_t pos_start, int64_t t0,
                             int tile_len, int n_tok, float scale,
                             fa_split_workspace* ws) {
    int n_wgs = NQ * n_tok;
    const float* S = ws->S;
    float* P = ws->P;
    float* m_state = ws->m_state;
    float* l_state = ws->l_state;
    float* Oalpha = ws->O;  /* temp alpha store */
    q.submit([&](sycl::handler& h) {
        h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * TILE_T),
                                         sycl::range<1>(TILE_T)),
            [=](sycl::nd_item<1> it) {
                int gid = it.get_group(0);
                int hh = gid / n_tok;
                int qi = gid % n_tok;
                int tloc = it.get_local_id(0);
                auto grp = it.get_group();

                int64_t qpos = pos_start + qi;
                int64_t kglob = t0 + tloc;
                bool valid = (tloc < tile_len) && (kglob <= qpos);
                float s = valid ? S[(size_t)hh * n_tok * TILE_T
                                   + (size_t)qi * TILE_T + tloc] * scale
                                : -1e30f;
                float tile_max = sycl::reduce_over_group(grp, s, sycl::maximum<float>());
                float prev_m = m_state[(size_t)hh * n_tok + qi];
                float new_m = sycl::fmax(prev_m, tile_max);
                float alpha = sycl::native::exp(prev_m - new_m);
                float e = (s > -1e29f) ? sycl::native::exp(s - new_m) : 0.0f;
                float tile_sum = sycl::reduce_over_group(grp, e, sycl::plus<float>());
                float prev_l = l_state[(size_t)hh * n_tok + qi];
                float new_l = alpha * prev_l + tile_sum;
                if (tloc == 0) {
                    m_state[(size_t)hh * n_tok + qi] = new_m;
                    l_state[(size_t)hh * n_tok + qi] = new_l;
                    Oalpha[(size_t)hh * n_tok + qi] = alpha;
                }
                P[(size_t)hh * n_tok * TILE_T + (size_t)qi * TILE_T + tloc] = e;
            });
    });
}

/* ---- Step 5: dequant V-tile -> Vt [NQ][TILE_T][HD] (rotated domain, T13) -
 * V_LM_LUT[code] * vscale * vnorm_corr  (NO per-token butterfly -- T13 defers
 * the inverse WHT to the final O).  One WG per (q-head, key-tile element),
 * LWS=HD lanes. */
static void fa_dequant_vt(sycl::queue& q, const uint32_t* kv_v,
                          int64_t t0, int tile_len,
                          fa_split_workspace* ws) {
    int n_wgs = NQ * TILE_T;
    float* Vt = ws->Vt;
    q.submit([&](sycl::handler& h) {
        h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * HD),
                                         sycl::range<1>(HD)),
            [=](sycl::nd_item<1> it) {
                int gid = it.get_group(0);
                int hh = gid / TILE_T;
                int tloc = gid % TILE_T;
                int hk = hh / GQA;
                int d = it.get_local_id(0);
                int64_t kglob = t0 + tloc;
                float vval = 0.0f;
                if (tloc < tile_len) {
                    int blk = hk * BLK_PER_HEAD + (d >> 5);
                    int vlane = d & 31;
                    const uint32_t* db = kv_v + (size_t)kglob * VROW_U32 + blk * 5;
                    uint16_t sc_u = (uint16_t)(db[0] & 0xFFFFu);
                    uint16_t nc_u = (uint16_t)((db[0] >> 16) & 0xFFFFu);
                    float vscale = float(sycl::bit_cast<sycl::half>(sc_u));
                    float vnorm_corr = float(sycl::bit_cast<sycl::half>(nc_u));
                    uint32_t c0 = db[1], c1 = db[2], c2 = db[3], c3 = db[4];
                    int w = vlane >> 3;
                    int off = (vlane & 7) * 4;
                    uint32_t word = (w == 0) ? c0 : (w == 1) ? c1
                                            : (w == 2) ? c2 : c3;
                    uint32_t code = (word >> off) & 0xFu;
                    vval = V_LM_LUT[code] * vscale * vnorm_corr;
                }
                Vt[(size_t)hh * TILE_T * HD + (size_t)tloc * HD + d] = vval;
            });
    });
}

/* ---- Step 6: PV GEMM + accumulate ---------------------------------------
 * Oacc[hh][n_tok][HD] = alpha * Oacc + P[hh][n_tok][TILE_T] @ Vt[hh][TILE_T][HD]
 * We fold the alpha rescale by pre-multiplying P by alpha (per query token),
 * then call gemm_tf32_batched with accum=Oacc.  batch=NQ, M=n_tok, K=TILE_T,
 * N=HD.  Requires TILE_T % 8, HD % 16. */
static void fa_pv_gemm(sycl::queue& q, int n_tok, fa_split_workspace* ws) {
    /* Online-softmax rescale: Oacc = alpha * Oacc, then Oacc += P @ Vt.
     * alpha is stashed in ws->O[hh*n_tok+qi] by fa_softmax_merge (one value per
     * (q-head, query-token)).  For the first tile alpha=0 (prev_m=-inf), which
     * correctly zeroes Oacc (already 0 from prep) and the tile's P@Vt sets the
     * initial accumulator.  Pre-scale Oacc (not P -- alpha rescales the OLD
     * state, the new P@Vt is already at the new_m scale). */
    int n_wgs = NQ * n_tok;
    const float* Oalpha = ws->O;
    float* Oacc = ws->Oacc;
    q.submit([&](sycl::handler& h) {
        h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * HD),
                                         sycl::range<1>(HD)),
            [=](sycl::nd_item<1> it) {
                int gid = it.get_group(0);
                int hh = gid / n_tok;
                int qi = gid % n_tok;
                int d = it.get_local_id(0);
                float alpha = Oalpha[(size_t)hh * n_tok + qi];
                Oacc[(size_t)hh * n_tok * HD + (size_t)qi * HD + d] *= alpha;
            });
    });
    /* Oacc += P @ Vt  (accum=Oacc, in-place).  M=n_tok, K=TILE_T, N=HD. */
    gemm_tf32_batched(q, ws->P, ws->Vt, ws->Oacc, NQ, n_tok, TILE_T, HD, ws->Oacc);
}

/* ---- Step 7: final -- Oacc / l, T13 inverse WHT, store -----------------
 * One WG per (q-head, query-token), LWS=HD.  Normalize by l_state, butterfly
 * each 32-elem block (inverse WHT) x NRM -> value domain, write to out. */
static void fa_finalize(sycl::queue& q, float* out, int n_tok,
                        fa_split_workspace* ws) {
    int n_wgs = NQ * n_tok;
    const float* Oacc = ws->Oacc;
    const float* l_state = ws->l_state;
    q.submit([&](sycl::handler& h) {
        sycl::local_accessor<float, 1> bfly0(sycl::range<1>(HD), h);
        sycl::local_accessor<float, 1> bfly1(sycl::range<1>(HD), h);
        h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * HD),
                                         sycl::range<1>(HD)),
            [=](sycl::nd_item<1> it) {
                int gid = it.get_group(0);
                int hh = gid / n_tok;
                int qi = gid % n_tok;
                int lid = it.get_local_id(0);
                auto grp = it.get_group();
                float l = l_state[(size_t)hh * n_tok + qi];
                float inv_l = 1.0f / l;
                bfly0[lid] = Oacc[(size_t)hh * n_tok * HD
                                  + (size_t)qi * HD + lid] * inv_l;
                sycl::group_barrier(grp);
                int blk = lid >> 5;
                int lane = lid & 31;
                int base = blk * VBLK_L;
                auto* cur = &bfly0;
                auto* nxt = &bfly1;
                for (int bit = 1; bit < VBLK_L; bit <<= 1) {
                    float a = (*cur)[base + lane];
                    float b = (*cur)[base + (lane ^ bit)];
                    sycl::group_barrier(grp);
                    (*nxt)[base + lane] = (lane & bit) ? (b - a) : (a + b);
                    sycl::group_barrier(grp);
                    std::swap(cur, nxt);
                }
                float oval = (*cur)[base + lane] * NRM;
                out[(size_t)qi * (NQ * HD) + (size_t)hh * HD + lid] = oval;
            });
    });
}

/* ---- Entry: fused kernel by default; QXMX_FA_SPLIT=1 keeps the M1.7b.7
 * 3-kernel split as an A/B oracle.  The fused path fills Oacc + l_state and
 * reuses this file's fa_finalize for /l + T13 inverse WHT. */
template <typename Kct>
void gpu_flash_attn(sycl::queue& q, const float* Q,
                    const uint8_t* kv_k, const uint32_t* kv_v,
                    float* out, int64_t pos_start, int n_tok,
                    fa_split_workspace* ws) {
    static const bool use_split = [] {
        const char* e = std::getenv("QXMX_FA_SPLIT");
        return e && std::atoi(e) != 0;
    }();
    if (!use_split) {
        gpu_flash_attn_fused<Kct>(q, Q, kv_k, kv_v, pos_start, n_tok, ws);
        fa_finalize(q, out, n_tok, ws);
        return;
    }
    const float scale = 1.0f / sycl::sqrt((float)HD);
    /* Prep Q + init state. */
    fa_prep_q(q, Q, n_tok, ws);
    /* Walk K in TILE_T-aligned blocks from 0 to pos_start+n_tok-1. */
    int64_t total = pos_start + n_tok;
    int n_tiles = (int)((total + TILE_T - 1) / TILE_T);
    for (int ti = 0; ti < n_tiles; ti++) {
        int64_t t0 = (int64_t)ti * TILE_T;
        int tile_len = (int)(total - t0);
        if (tile_len > TILE_T) tile_len = TILE_T;
        fa_dequant_kt<Kct>(q, kv_k, t0, tile_len, ws);
        fa_qk_gemm(q, n_tok, ws);
        fa_softmax_merge(q, pos_start, t0, tile_len, n_tok, scale, ws);
        fa_dequant_vt(q, kv_v, t0, tile_len, ws);
        fa_pv_gemm(q, n_tok, ws);
    }
    fa_finalize(q, out, n_tok, ws);
}

template void gpu_flash_attn<KFp16>(sycl::queue&, const float*, const uint8_t*,
                                     const uint32_t*, float*, int64_t, int,
                                     fa_split_workspace*);
template void gpu_flash_attn<KQ8_0>(sycl::queue&, const float*, const uint8_t*,
                                     const uint32_t*, float*, int64_t, int,
                                     fa_split_workspace*);
template void gpu_flash_attn<KFp8 >(sycl::queue&, const float*, const uint8_t*,
                                     const uint32_t*, float*, int64_t, int,
                                     fa_split_workspace*);

}  // namespace qx