// TinyLLama-v0 on the chipformer RV32 machine: int4 weights and 16-bit scales in the 20-bit WROM,
// embedding and scales read with LWROM, WBB emitted whenever a base crosses 32768.
#include <stdint.h>
#include "rt.h"
#include "weights_tl.h"

#ifndef NTOK
#define NTOK 4
#endif
#define W 32

#define VDOT(acc, a, b) __asm__ volatile(".insn r 0x0B, 5, 0, %0, %1, %2" : "+r"(acc) : "r"(a), "r"(b))
#define SMAQA(acc, a, b) __asm__ volatile(".insn r 0x0B, 7, 0, %0, %1, %2" : "+r"(acc) : "r"(a), "r"(b))
#define WBB(v) __asm__ volatile(".insn r 0x0B, 6, 1, zero, %0, zero" :: "r"(v))
#define WBSET(v) __asm__ volatile(".insn r 0x0B, 6, 0, zero, %0, zero" :: "r"(v))
static inline uint32_t lwrom(uint32_t a) {
    uint32_t r;
    __asm__ volatile(".insn r 0x0B, 6, 2, %0, %1, zero" : "=r"(r) : "r"(a));
    return r;
}

static int32_t x[DIM], xb[DIM], q[DIM], kk[KVD], vv[KVD], hb[HID], hb2[HID];
static int8_t xq[256] __attribute__((aligned(4)));
static int8_t kq[NL][W][KVD] __attribute__((aligned(4)));
static int8_t vt[NL][KVD][W] __attribute__((aligned(4)));
static int8_t pp[W] __attribute__((aligned(4)));
static int32_t vmax;
#define DBG(i, v) (((volatile int32_t *)0xE000)[(i)] = (v))

static inline int32_t mulh(int32_t a, int32_t b) { return (int32_t)(((int64_t)a * b) >> 32); }
#define TRACK(v) do { int32_t a_ = (v) < 0 ? -(v) : (v); if (a_ > vmax) vmax = a_; } while (0)

static int32_t quantize(const int32_t *v, int n, int32_t amax) {
    if (amax < 256) amax = 256;
    int32_t r = 0x7F000000 / amax, h = amax / 254;
#pragma GCC unroll 4
    for (int i = 0; i < n; i++) xq[i] = (int8_t)mulh((v[i] + h) << 8, r);
    return amax;
}

static int32_t norm_f(void) {
    int32_t ss = 0;
#pragma GCC unroll 8
    for (int i = 0; i < DIM; i++) ss += xq[i] * xq[i];
    uint32_t m = (uint32_t)ss << 10, r = 0, b = 1u << 30;
    while (b > m) b >>= 2;
    while (b) {
        if (m >= r + b) { m -= r + b; r = (r >> 1) + b; } else r >>= 1;
        b >>= 2;
    }
    return (int32_t)((0x80000000u / r) << 10);
}

static int32_t plain_f(int32_t amax) { return (int32_t)(((uint32_t)amax << 9) / 127 << 8); }

extern void vk_8_init_scaled(int32_t *, int, int, const int32_t *, int32_t);
extern void vk_8_init_raw(int32_t *, int, int, const int32_t *, int32_t);
extern void vk_8_acc_raw(int32_t *, int, int, const int32_t *, int32_t);
const int8_t *vk_x;

static void matvec(int32_t *out, const int32_t *starts, int nchunk, int32_t sc, int rows, int32_t f) {
    int cu[4], co[4];
    for (int c = 0; c < nchunk; c++) { cu[c] = starts[c] >> 10; co[c] = starts[c] & 0x3FF; }
    int su = sc >> 10, so = sc & 0x3FF;
    for (int r = 0; r < rows; r += 4) {
        uint32_t w0 = lwrom((su << 10) | so); so++; if (so >= 960) { su++; so = 0; }
        uint32_t w1 = lwrom((su << 10) | so); so++; if (so >= 960) { su++; so = 0; }
        int32_t S4[4] = {(int32_t)(w0 & 0xFFFF), (int32_t)((w0 >> 16) & 0xFFFF),
                         (int32_t)(w1 & 0xFFFF), (int32_t)((w1 >> 16) & 0xFFFF)};
        if (nchunk == 1) {
            vk_x = xq;
            vk_8_init_scaled(out + r, (cu[0] << 10) | co[0], 4, S4, f);
        } else {
            for (int c = 0; c < nchunk; c++) {
                vk_x = xq + 64 * c;
                int32_t base = (cu[c] << 10) | co[c];
                if (c == 0) vk_8_init_raw(out + r, base, 4, 0, 0);
                else vk_8_acc_raw(out + r, base, 4, 0, 0);
            }
            for (int k = 0; k < 4; k++) out[r + k] = mulh(out[r + k] * S4[k], f);
        }
        for (int c = 0; c < nchunk; c++) {
            co[c] += 32;
            if (co[c] + 32 > 960) { cu[c]++; co[c] = 0; }
        }
    }
}

static void rope(int32_t *v, int pos, int n) {
    const int16_t *t = ROPE + pos * HS;
    for (int i = 0; i < n; i += 2) {
        int32_t c = t[i % HS], s = t[i % HS + 1], a = v[i] >> 6, b = v[i + 1] >> 6;
        v[i] = (a * c - b * s) >> 8;
        v[i + 1] = (a * s + b * c) >> 8;
    }
}

static inline int8_t q8_of(int32_t v, int32_t r, int32_t h) {
    int32_t t = mulh((v + (v < 0 ? -h : h)) << 8, r);
    return (int8_t)(t > 127 ? 127 : t < -127 ? -127 : t);
}
static inline uint32_t pack4(const int8_t *b) {
    return (uint8_t)b[0] | (uint8_t)b[1] << 8 | (uint8_t)b[2] << 16 | (uint32_t)(uint8_t)b[3] << 24;
}

_Static_assert(HS == 4, "attend packs one 32-bit word per head");

static void attend(int l, int h, int n) {
    const int32_t *qv = q + h * HS;
    int32_t amax = 256;
    for (int j = 0; j < HS; j++) { int32_t a = qv[j] < 0 ? -qv[j] : qv[j]; if (a > amax) amax = a; }
    int32_t r = 0x7F000000 / amax, hf = amax / 254;
    int8_t b[HS];
    for (int j = 0; j < HS; j++) b[j] = q8_of(qv[j], r, hf);
    uint32_t q0 = pack4(b);
    int64_t ca64 = ((((int64_t)amax * KMAX[l]) >> 10) * ATT_C) >> 16;
    int32_t ca = ca64 > INT32_MAX ? INT32_MAX : (int32_t)ca64;
    int32_t sc[W], m = INT32_MIN;
    for (int t = 0; t < n; t++) {
        const uint32_t *kw = (const uint32_t *)&kq[l][t][h * HS];
        int32_t s = 0;
        SMAQA(s, q0, kw[0]);
        sc[t] = s;
        if (s > m) m = s;
    }
    int32_t sum = 0;
    for (int t = 0; t < n; t++) {
        int32_t i = mulh(m - sc[t], ca);
        int32_t p = (i >= 0 && i < 256) ? P8[i] : 0;
        pp[t] = (int8_t)p;
        sum += p;
    }
    int nw = (n + 3) >> 2;
    for (int t = n; t < 4 * nw; t++) pp[t] = 0;
    int32_t o[HS] = {0};
    for (int w = 0; w < nw; w++) {
        uint32_t p0 = ((const uint32_t *)pp)[w];
        for (int j = 0; j < HS; j++) {
            uint32_t vw = ((const uint32_t *)vt[l][h * HS + j])[w];
            SMAQA(o[j], p0, vw);
        }
    }
    int32_t rr = (VMAX[l] << 12) / (127 * sum);
    for (int j = 0; j < HS; j++) {
        xb[h * HS + j] = (int32_t)(((int64_t)o[j] * rr) >> 12);
        TRACK(xb[h * HS + j]);
    }
}

static void pnum(int32_t v) {
    char b[12]; int n = 0, neg = v < 0;
    uint32_t u = neg ? -(uint32_t)v : (uint32_t)v;
    if (!u) b[n++] = '0';
    while (u) { b[n++] = (char)('0' + u % 10); u /= 10; }
    if (neg) rv_putchar('-');
    while (n) rv_putchar(b[--n]);
    rv_putchar(' ');
}

static void embed(int tok, int32_t S) {
    int g = tok >> 2, rr = tok & 3;
    int total = (G_EMB & 0x3FF) + g * 32;
    int unit = (G_EMB >> 10) + total / 960, off = total % 960;
    int32_t base = ((unit << 10) | off) + rr * 8;
    for (int k = 0; k < 8; k++) {
        uint32_t w = lwrom(base + k);
        for (int j = 0; j < 8; j++) {
            int32_t v = (int32_t)((w >> (4 * j)) & 15) - 8;
            x[8 * k + j] = (int32_t)((int64_t)v * S * 2);
        }
    }
    for (int i = 0; i < DIM; i++) TRACK(x[i]);
}

static int32_t emb_scale(int tok) {
    int g = tok >> 2, rr = tok & 3;
    int total = (SC_EMB & 0x3FF) + g * 2 + (rr >> 1);
    int su = (SC_EMB >> 10) + total / 960, so = total % 960;
    uint32_t sw = lwrom((su << 10) | so);
    return (int32_t)((sw >> (16 * (rr & 1))) & 0xFFFF);
}

int main(void) {
    int tok = 1;
    for (int pos = 0; pos < NTOK; pos++) {
        vmax = 0;
        embed(tok, emb_scale(tok));
        DBG(0, x[0]);

        for (int l = 0; l < NL; l++) {
            quantize(x, DIM, vmax);
            int32_t f = norm_f();
            if (l == 0) { DBG(1, xq[0]); DBG(2, f); }
            matvec(q, &G_q[l], 1, SC_q[l], DIM, f);
            matvec(kk, &G_k[l], 1, SC_k[l], DIM, f);
            matvec(vv, &G_v[l], 1, SC_v[l], DIM, f);
            if (l == 0) { DBG(3, q[0]); DBG(4, kk[0]); DBG(5, vv[0]); }

            rope(q, pos, DIM);
            rope(kk, pos, KVD);
            int slot = pos == 0 ? 0 : 1 + (pos - 1) % (W - 1), n = pos + 1 < W ? pos + 1 : W;
            int32_t kr = 0x7F000000 / KMAX[l], kh = KMAX[l] / 254;
            int32_t vr = 0x7F000000 / VMAX[l], vh = VMAX[l] / 254;
            for (int i = 0; i < KVD; i++) {
                kq[l][slot][i] = q8_of(kk[i], kr, kh);
                vt[l][i][slot] = q8_of(vv[i], vr, vh);
            }
            vmax = 0;
            for (int h = 0; h < NH; h++) attend(l, h, n);
            if (l == 0) DBG(6, xb[0]);
            int32_t amax = quantize(xb, DIM, vmax);
            matvec(q, &G_o[l], 1, SC_o[l], DIM, plain_f(amax));
            vmax = 0;
            for (int i = 0; i < DIM; i++) { x[i] += q[i]; TRACK(x[i]); }
            if (l == 0) DBG(7, x[0]);
            quantize(x, DIM, vmax);
            f = norm_f();
            matvec(hb, &G_gate[l], 1, SC_gate[l], HID, f);
            matvec(hb2, &G_up[l], 1, SC_up[l], HID, f);
            vmax = 0;
            for (int i = 0; i < HID; i++) {
                int32_t a = hb[i], s;
                if (a >= 8 << 16) s = a >> 4;
                else if (a < -(8 << 16)) s = 0;
                else s = SILU[(a >> 11) + 256];
                hb[i] = (s * (hb2[i] >> 8)) >> 4;
                TRACK(hb[i]);
            }
            amax = quantize(hb, HID, vmax);
            if (l == 0) DBG(8, hb[0]);
            matvec(q, &G_down[4 * l], 4, SC_down[l], DIM, plain_f(amax));
            vmax = 0;
            for (int i = 0; i < DIM; i++) { x[i] += q[i]; TRACK(x[i]); }
            DBG(9 + l, x[0]);
        }
        quantize(x, DIM, vmax);
        DBG(30, xq[0]); DBG(31, xq[1]); DBG(32, vmax);
        int best = 0, bv = INT32_MIN;
        int unit = G_HEAD >> 10, off = G_HEAD & 0x3FF;
        int su = SC_HEAD >> 10, so = SC_HEAD & 0x3FF;
        int32_t tmp[4];
        for (int g = 0; g < VOCAB / 4; g++) {
            uint32_t w0 = lwrom((su << 10) | so); so++; if (so >= 960) { su++; so = 0; }
            uint32_t w1 = lwrom((su << 10) | so); so++; if (so >= 960) { su++; so = 0; }
            int32_t S4[4] = {(int32_t)(w0 & 0xFFFF), (int32_t)((w0 >> 16) & 0xFFFF),
                             (int32_t)(w1 & 0xFFFF), (int32_t)((w1 >> 16) & 0xFFFF)};
            int32_t base = (unit << 10) | off;
            vk_x = xq;
            vk_8_init_scaled(tmp, base, 4, S4, 1 << 30);
            for (int k = 0; k < 4; k++) if (tmp[k] > bv) { bv = tmp[k]; best = g * 4 + k; }
            off += 32; if (off + 32 > 960) { unit++; off = 0; }
        }
        DBG(100 + pos, best); DBG(200 + pos, bv);
        { int b = best, n = 0; char buf[8];
          if (!b) buf[n++] = '0';
          while (b) { buf[n++] = (char)('0' + b % 10); b /= 10; }
          while (n) rv_putchar(buf[--n]); }
        rv_putchar(' ');
        tok = best;
    }
    rv_putchar('\n');
    return 0;
}