Hugging Face's logo

phanerozoic
/
oasis

Image-Text-to-Text
oasis
vision
ocr
object-detection
image-segmentation
structured-output
cpu
2-bit
Model card Files
xet
 Community
oasis / oasis.cpp
phanerozoic's picture
phanerozoic
Speed up the engine on 64-bit Arm and re-measure the Arm timings on an idle board
3d79675 verified
Raw History Blame Contribute Delete
192 kB
// Oasis engine: loads an .oasis model and describes RGB images as JSON records, in one C++17 translation
// unit. `python oasis.py build` compiles it into the library that oasis.py loads, as
// c++ -std=c++17 -O3 -ffast-math -fno-finite-math-only -march=native -fopenmp -fPIC -shared -fvisibility=hidden oasis.cpp -o liboasis.so
// cl /std:c++17 /O2 /fp:fast /arch:AVX2 /openmp /EHsc /LD oasis.cpp /Fe:oasis.dll
// with -mcpu=native in place of -march=native for Clang on 64-bit Arm. The C API below is all it exports.
#ifndef OASIS_H
#define OASIS_H
#ifdef __cplusplus
extern "C" {
#endif
#if defined(_WIN32)
#define OASIS_API __declspec(dllexport)
#else
#define OASIS_API __attribute__((visibility("default")))
#endif
typedef struct oasis_model oasis_model;
/* Returns NULL on failure; see oasis_last_error(). */
OASIS_API oasis_model* oasis_load(const char* path);
OASIS_API void oasis_free(oasis_model* model);
/* rgb: packed 8-bit RGB, row-major, width * height * 3 bytes.
options_json: NULL or a JSON object with language (en, de, fr, es, it, nl or pl), keywords (a list
of words or phrases the description should favour, or one newline-separated string), conf, text_conf,
max_objects, max_text_lines, max_tokens, masks, mask_format ("largest", "polygons", "merged",
"bitmap", "none"), mask_threshold, nms, reasoning, sections (any of "scene", "objects", "regions", "text", "relations", "events",
"reasoning"; text alone skips the decoder), fields (any of "bbox", "score", "mask", "position",
"description", "count") and form (an ordered object of fields: {"enum": [...]}, {"type": "integer",
"minimum": a, "maximum": b}, {"type": "string", "maxLength": n} or {"x-source": "bbox", "x-of":
"<earlier enum field>"}), which returns {"form": ...} instead of a description.
Writes a NUL-terminated JSON record to out; returns its length, or a negative value on error
(minus the required size when capacity is too small). */
OASIS_API int oasis_describe(oasis_model* model, const unsigned char* rgb, int width, int height,
const char* options_json, char* out, int capacity);
OASIS_API const char* oasis_last_error(void);
#ifdef __cplusplus
}
#endif
#endif
#include <algorithm>
#include <atomic>
#include <cctype>
#include <chrono>
#include <climits>
#include <cmath>
#include <condition_variable>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <memory>
#include <mutex>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>
#ifdef _OPENMP
#include <omp.h>
#endif
#if defined(__F16C__)
#include <immintrin.h>
#endif
// ==== Internals: model container, tensor access, kernel and pipeline declarations
// 64-bit Arm with NEON: the kernels and the image and text stages take their Arm paths.
#if defined(__aarch64__) && defined(__ARM_NEON)
#include <arm_neon.h>
#define OASIS_NEON 1
#endif
namespace oasis {
class SpinPool;
// Tensor types: 0 f32, 1 f16, 2 cq with fp16 group scales, 3 raw bytes, 4 signed integer rows of `bits`
// bits with one f32 scale per row, 5 cq with 8-bit log-coded group scales.
struct Entry {
int dtype = 0;
int bits = 0;
std::vector<int> shape;
const uint8_t* data = nullptr;
size_t nbytes = 0;
};
// Quantised matrix [rows = outputs, cols = inputs]: packed indices, then one scale per row and 128-group,
// stored either as fp16 (norms) or as an 8-bit code into a 256-entry table (ncode, nlut).
struct CQ {
int rows = 0, cols = 0, pc = 0, bits = 0;
const uint8_t* idx = nullptr;
const uint16_t* norms = nullptr;
const uint8_t* ncode = nullptr;
const float* nlut = nullptr;
const float* cb = nullptr;
const uint8_t* zero_rows = nullptr; // 1 for rows whose every group scale is zero; null when there are none
};
bool inflate_exact(const uint8_t* in, size_t n, uint8_t* out, size_t raw);
struct Model {
std::vector<uint8_t> file;
std::vector<std::vector<uint8_t>> inflated; // tensors stored deflated, expanded at load
std::unordered_map<std::string, std::string> header;
std::unordered_map<std::string, Entry> entries;
std::unordered_map<std::string, std::vector<float>> cache; // fp16 and integer-row tensors as f32, filled at load
std::unordered_map<std::string, std::vector<float>> luts; // decoded 8-bit group-scale codes, filled at load
std::unordered_map<std::string, std::vector<uint8_t>> zero_rows; // quantised matrices with all-zero rows
std::unordered_map<std::string, std::vector<float>> dq; // dequantised matrices, filled on demand
std::mutex dq_mutex;
std::vector<float> codebook;
int cb_offset[5] = {0, 0, 0, 4, 12};
int vocab = 0, n_pieces = 0, bos = 2, eos = 1, pad = 0, ref_base = 0;
std::unordered_map<std::string, int> langs; // every language tag in the vocabulary
std::vector<std::string> languages; // the languages the model describes in
std::vector<std::string> pieces;
std::vector<int> piece_types;
std::unordered_map<std::string, int> piece_id; // normal pieces, for encoding
std::vector<std::pair<std::string, int>> markers; // user-defined pieces, longest first
int byte_id[256] = {0}; // byte-fallback pieces
std::vector<uint32_t> charset;
std::vector<int> perm1, perm2;
std::vector<int> seq_order; // sequency row -> natural row for the 512-point Walsh matrix
std::shared_ptr<SpinPool> pool; // decoder workers, created at the first decode
std::mutex decode_mutex;
bool load(const std::string& path, std::string& err);
bool has(const std::string& name) const { return entries.count(name) != 0; }
const Entry& E(const std::string& name) const;
const float* F(const std::string& name); // f32 view, fp16 converted once
CQ Q(const std::string& name) const;
const float* Wdq(const std::string& name); // dequantised [rows, cols], cached
};
// ---- scalar helpers
// The exponent is capped so that no compiler's fast exp can overflow; below -80 the result is ~0 anyway.
inline float sigmoidf(float x) { return 1.0f / (1.0f + std::exp(std::min(-x, 80.0f))); }
inline float siluf(float x) { return x / (1.0f + std::exp(std::min(-x, 80.0f))); }
inline float gelu_tanh(float x) {
const float k = 0.7978845608028654f;
return 0.5f * x * (1.0f + std::tanh(k * (x + 0.044715f * x * x * x)));
}
#if defined(OASIS_NEON)
// expf as glibc 2.27 and later computes it on 64-bit Arm, four lanes at a time: x times 32/ln 2 rounded to
// k with ties away from zero, 2^(k/32) from its table, and a cubic in the remainder, in double precision.
// Equal to that library's expf for every float that is not a NaN.
inline float32x2_t exp_lanes(float64x2_t xd) {
static const uint64_t T[32] = {
0x3ff0000000000000, 0x3fefd9b0d3158574, 0x3fefb5586cf9890f, 0x3fef9301d0125b51, 0x3fef72b83c7d517b, 0x3fef54873168b9aa,
0x3fef387a6e756238, 0x3fef1e9df51fdee1, 0x3fef06fe0a31b715, 0x3feef1a7373aa9cb, 0x3feedea64c123422, 0x3feece086061892d,
0x3feebfdad5362a27, 0x3feeb42b569d4f82, 0x3feeab07dd485429, 0x3feea47eb03a5585, 0x3feea09e667f3bcd, 0x3fee9f75e8ec5f74,
0x3feea11473eb0187, 0x3feea589994cce13, 0x3feeace5422aa0db, 0x3feeb737b0cdc5e5, 0x3feec49182a3f090, 0x3feed503b23e255d,
0x3feee89f995ad3ad, 0x3feeff76f2fb5e47, 0x3fef199bdd85529c, 0x3fef3720dcef9069, 0x3fef5818dcfba487, 0x3fef7c97337b9b5f,
0x3fefa4afa2a490da, 0x3fefd0765b6e4540};
const float64x2_t z = vmulq_f64(vdupq_n_f64(0x1.71547652b82fep+0 * 32), xd);
const float64x2_t r = vsubq_f64(z, vrndaq_f64(z));
const int64x2_t ki = vcvtaq_s64_f64(z);
const uint64_t k0 = (uint64_t)vgetq_lane_s64(ki, 0), k1 = (uint64_t)vgetq_lane_s64(ki, 1);
const float64x2_t s = vreinterpretq_f64_u64(vcombine_u64(vcreate_u64(T[k0 % 32] + (k0 << 47)), vcreate_u64(T[k1 % 32] + (k1 << 47))));
const float64x2_t c = vfmaq_f64(vdupq_n_f64(0x1.ebfce50fac4f3p-3 / 32 / 32), vdupq_n_f64(0x1.c6af84b912394p-5 / 32 / 32 / 32), r);
const float64x2_t y = vfmaq_f64(vfmaq_f64(vdupq_n_f64(1.0), vdupq_n_f64(0x1.62e42ff0c52d6p-1 / 32), r), c, vmulq_f64(r, r));
return vcvt_f32_f64(vmulq_f64(y, s));
}
inline float32x4_t expf4(float32x4_t x) {
const float32x4_t lo = vdupq_n_f32(-0x1.9fe368p6f), hi = vdupq_n_f32(0x1.62e42ep6f);
const float32x4_t xc = vminq_f32(vmaxq_f32(x, lo), hi); // NaNs pass through
const float32x4_t e = vcombine_f32(exp_lanes(vcvt_f64_f32(vget_low_f32(xc))), exp_lanes(vcvt_high_f64_f32(xc)));
const float32x4_t zero = vbslq_f32(vcltq_f32(x, lo), vdupq_n_f32(0.0f), e);
return vbslq_f32(vcgtq_f32(x, hi), vdupq_n_f32(INFINITY), zero);
}
// tanhf as glibc 2.41 computes it (CORE-MATH's correctly rounded tanhf), four lanes at a time: a rational
// function in double precision, with its own forms for tiny and large arguments. Equal to the library's
// tanhf for every float.
inline float32x2_t tanh_lanes(float64x2_t z) {
static const double cn[] = {0x1p+0, 0x1.30877b8b72d33p-3, 0x1.694aa09ae9e5ep-8, 0x1.4101377abb729p-14,
0x1.e0392b1db0018p-22, 0x1.2533756e546f7p-30, 0x1.d62e5abe6ae8ap-41, 0x1.b06be534182dep-54};
static const double cd[] = {0x1p+0, 0x1.ed99131b0ebeap-2, 0x1.0d27ed6c95a69p-5, 0x1.7cbdaca0e9fccp-11,
0x1.b4e60b892578ep-18, 0x1.a6f707c5c71abp-26, 0x1.35a8b6e2cd94cp-35, 0x1.ca8230677aa01p-47};
const float64x2_t z2 = vmulq_f64(z, z), z4 = vmulq_f64(z2, z2), z8 = vmulq_f64(z4, z4);
auto pair = [&](const double* c, int i) { return vfmaq_f64(vdupq_n_f64(c[i]), z2, vdupq_n_f64(c[i + 1])); };
float64x2_t n0 = pair(cn, 0), n2 = pair(cn, 2), n4 = pair(cn, 4), n6 = pair(cn, 6);
n0 = vfmaq_f64(n0, z4, n2);
n4 = vfmaq_f64(n4, z4, n6);
n0 = vfmaq_f64(n0, z8, n4);
float64x2_t d0 = pair(cd, 0), d2 = pair(cd, 2), d4 = pair(cd, 4), d6 = pair(cd, 6);
d0 = vfmaq_f64(d0, z4, d2);
d4 = vfmaq_f64(d4, z4, d6);
d0 = vfmaq_f64(d0, z8, d4);
return vcvt_f32_f64(vdivq_f64(vmulq_f64(z, n0), d0));
}
inline float32x4_t tanhf4(float32x4_t x) {
const uint32x4_t ax = vandq_u32(vreinterpretq_u32_f32(x), vdupq_n_u32(0x7fffffff));
float32x4_t r = vcombine_f32(tanh_lanes(vcvt_f64_f32(vget_low_f32(x))), tanh_lanes(vcvt_high_f64_f32(x)));
const float32x4_t x2 = vmulq_f32(x, x);
r = vbslq_f32(vcltq_u32(ax, vdupq_n_u32(115u << 23)), vfmaq_f32(x, x, vmulq_n_f32(x2, -0x1.555556p-2f)), r);
r = vbslq_f32(vcltq_u32(ax, vdupq_n_u32(102u << 23)), vfmaq_f32(x, vnegq_f32(x), vabsq_f32(x)), r);
r = vbslq_f32(vceqq_u32(ax, vdupq_n_u32(0)), x, r);
const uint32x4_t sign = vandq_u32(vreinterpretq_u32_f32(x), vdupq_n_u32(0x80000000));
const float32x4_t one = vreinterpretq_f32_u32(vorrq_u32(sign, vreinterpretq_u32_f32(vdupq_n_f32(1.0f))));
const float32x4_t tiny = vreinterpretq_f32_u32(vorrq_u32(sign, vreinterpretq_u32_f32(vdupq_n_f32(0x1p-25f))));
r = vbslq_f32(vcgtq_u32(ax, vdupq_n_u32(0x41102cb3u)), vsubq_f32(one, tiny), r);
return vbslq_f32(vcgtq_u32(ax, vdupq_n_u32(0x7f800000u)), vaddq_f32(x, x), r);
}
// o[i] *= sigmoidf(g[i]) and v[i] = siluf(v[i]) for n a multiple of 4, with the same roundings.
inline void mul_sigmoid(float* o, const float* g, int n) {
const float32x4_t cap = vdupq_n_f32(80.0f), one = vdupq_n_f32(1.0f);
for (int i = 0; i < n; i += 4) {
const float32x4_t e = expf4(vminq_f32(vnegq_f32(vld1q_f32(g + i)), cap));
vst1q_f32(o + i, vmulq_f32(vld1q_f32(o + i), vdivq_f32(one, vaddq_f32(one, e))));
}
}
inline void silu_inplace(float* v, int n) {
const float32x4_t cap = vdupq_n_f32(80.0f), one = vdupq_n_f32(1.0f);
for (int i = 0; i < n; i += 4) {
const float32x4_t x = vld1q_f32(v + i);
vst1q_f32(v + i, vdivq_f32(x, vaddq_f32(one, expf4(vminq_f32(vnegq_f32(x), cap)))));
}
}
#endif
float half_to_float(uint16_t h);
extern bool g_exact_gemv; // float activations in quantised matrix-vector products (OASIS_EXACT=1)
// Threads for one description: OMP_NUM_THREADS when it is set, otherwise every processor, less one on
// machines with four or fewer, where the rest of the system would otherwise preempt a worker at every step.
int engine_threads();
// Scale of group i (row-major over rows and 128-groups) of a quantised matrix.
inline float cq_scale(const CQ& q, size_t i) { return q.ncode ? q.nlut[q.ncode[i]] : half_to_float(q.norms[i]); }
// Activations prepared once for several quantised matrix-vector products with the same input:
// Walsh-rotated per 128-group, int8 per group, reordered to the packed-index layout.
struct QX {
std::vector<int8_t> xs, xa;
std::vector<float> sx;
int pc = 0, bits = 0;
};
// Codebook entries as int8 with one scale, shared by every row of a matrix of the same width.
struct CbTable {
alignas(16) int8_t tbl[16];
float cs;
};
// ---- kernels
void prep_x(const float* x, int cols, int bits, QX& out);
CbTable cb_table(const CQ& q);
float cq_row(const CQ& q, const QX& qx, int r, const CbTable& t); // one output of W x
void cq_gemv_q(const CQ& q, const QX& qx, float* y);
void cq_gemm_q(const CQ& q, const float* A, int M, float* C); // int8-activation matrix product
struct Model;
void qmm(Model& m, const std::string& name, const float* A, int M, float* C); // C = A W^T for a quantised W
float dotf(const float* a, const float* b, int n);
void wht_inplace(float* x, int n); // natural order, normalised
void wht_seq512(const Model& m, const float* x, float* out); // sequency-ordered 512-point transform
void cq_dequant_row(const CQ& q, int row, float* out);
void cq_gemv(const CQ& q, const float* x, float* y); // y[rows] = W x
void gemm_nt(const float* A, int M, int K, const float* W, int N, float* C); // C = A W^T, W [N, K]
void gemm_nn(const float* A, int M, int K, const float* B, int N, float* C); // C = A B, B [K, N]
void zc_norm(const float* x, int n, const float* scale, float* out);
void rms_unit(const float* x, int n, float* out);
void band65(const float* u, const float* w, int n, float* out); // per-output 65-tap band, zero edges
// ---- pipeline
struct Box {
float x0, y0, x1, y1; // canvas-normalised [0, 1]
float score;
int ci, cj; // centre cell on the 64 x 64 grid
};
struct Meta {
int W = 0, H = 0, w = 0, h = 0;
float r = 1.0f;
};
struct Encoded {
Meta meta;
std::vector<float> E; // [256, 512]
std::vector<float> D; // [64, 64, 128]
std::vector<float> S4; // [128, 128, 64]
std::vector<uint8_t> valid32, valid16, valid8, valid4;
std::vector<float> det; // [64, 64, 10]
};
// One field of a form request, in request order.
struct FormField {
enum Kind { CHOICE, COUNT, TEXT, BOX } kind = CHOICE;
std::string name;
std::vector<std::string> choices; // CHOICE: candidate texts
std::vector<std::string> values; // CHOICE: each choice as JSON (string or number)
long lo = 0, hi = 0; // COUNT: inclusive range
int max_length = 60; // TEXT
std::string of; // BOX: the earlier choice field it locates
};
struct Options {
std::string language = "en";
float conf = 0.3f, text_conf = 0.35f, mask_threshold = 0.5f;
int max_objects = 20, max_text_lines = 20, max_tokens = 512;
std::string mask_format = "largest";
bool masks = true, nms = true, reasoning = true;
std::vector<std::string> sections, fields, keywords;
std::vector<FormField> form;
};
Encoded encode(Model& m, const uint8_t* rgb, int W, int H);
void detect(const Encoded& enc, float conf, float text_conf, bool nms, int max_obj, int max_txt,
std::vector<Box>& objects, std::vector<Box>& texts);
std::vector<float> mask_prob(Model& m, const Encoded& enc, const Box& b, std::vector<float>* protos);
std::string read_text(Model& m, const uint8_t* rgb, int W, int H, float x0, float y0, float x1, float y1);
void memory(Model& m, const Encoded& enc, const std::vector<Box>& objects, const std::vector<Box>& texts,
std::vector<float>& M, std::vector<uint8_t>& mvalid, std::vector<int>& refs);
// keywords: words and phrases whose tokens are favoured while the description is decoded.
std::vector<int> generate(Model& m, const std::vector<float>& M, const std::vector<uint8_t>& mvalid,
const std::vector<int>& refs, const std::string& language, int max_tokens,
const std::vector<std::string>& keywords = {});
std::string detokenize(const Model& m, const std::vector<int>& ids);
std::vector<int> encode_text(const Model& m, const std::string& text); // SentencePiece BPE, no dummy prefix
// Fills a form against the decoder memory; ref_boxes are the 0-1000 boxes of refs. Returns the "form" object.
std::string fill_form(Model& m, const std::vector<float>& M, const std::vector<uint8_t>& mvalid, const std::vector<int>& refs,
const std::vector<std::vector<int>>& ref_boxes, const std::string& language, const std::vector<FormField>& form);
std::string describe_json(Model& m, const uint8_t* rgb, int W, int H, const Options& o, double* timings);
} // namespace oasis
// ==== Worker threads that stay awake for the length of one decode
//
// Each decoded token runs dozens of parallel steps that last a few microseconds. Workers that an
// OpenMP runtime parks between such steps take longer to wake than the steps take to run, so the
// decoder shares its work through this pool instead: while a decode is in progress the workers spin
// on a generation counter, and outside it they sleep on a condition variable.
#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86)
#include <immintrin.h>
#define OASIS_PAUSE() _mm_pause()
#elif defined(__aarch64__) && (defined(__GNUC__) || defined(__clang__))
#define OASIS_PAUSE() __asm__ __volatile__("yield" ::: "memory")
#else
#define OASIS_PAUSE() std::this_thread::yield()
#endif
namespace oasis {
class SpinPool {
public:
explicit SpinPool(int threads) : n_(std::max(1, threads)) {
for (int i = 1; i < n_; i++) workers_.emplace_back([this] { loop(); });
}
~SpinPool() {
{
std::lock_guard<std::mutex> lock(mutex_);
stop_.store(true);
awake_.store(true);
}
cv_.notify_all();
for (auto& t : workers_) t.join();
}
SpinPool(const SpinPool&) = delete;
SpinPool& operator=(const SpinPool&) = delete;
int size() const { return n_; }
void wake() {
{
std::lock_guard<std::mutex> lock(mutex_);
awake_.store(true);
}
cv_.notify_all();
}
// Workers notice within a few hundred spins and go back to sleep.
void sleep() {
std::lock_guard<std::mutex> lock(mutex_);
awake_.store(false);
}
// Calls f(i) once for every i in [0, n). Indices are claimed in chunks by whichever thread is free,
// and each is computed entirely by one thread, so results do not depend on the thread count. The call
// returns once every index is done: a worker the system has not scheduled meanwhile takes no part in
// the step instead of holding the others up.
template <class F>
void run(int n, int chunk, F&& f) {
if (n <= 0) return;
if (n_ == 1 || n <= chunk) {
for (int i = 0; i < n; i++) f(i);
return;
}
if (!awake_.load(std::memory_order_relaxed)) wake();
using G = typename std::remove_reference<F>::type;
const uint64_t gen = (ticket_.load(std::memory_order_relaxed) >> 32) + 1;
ticket_.store(gen << 32 | CLOSED, std::memory_order_relaxed);
fn_.store([](void* ctx, int b, int e) {
G& g = *static_cast<G*>(ctx);
for (int i = b; i < e; i++) g(i);
}, std::memory_order_relaxed);
ctx_.store((void*)&f, std::memory_order_relaxed);
chunk_.store(std::max(1, chunk), std::memory_order_relaxed);
done_.store(0, std::memory_order_relaxed);
items_.store(n, std::memory_order_release);
ticket_.store(gen << 32, std::memory_order_release);
work(gen);
while (done_.load(std::memory_order_acquire) < n) OASIS_PAUSE();
}
private:
// The ticket holds the step's generation in its high half and the next unclaimed index in its low half.
// run() closes the new generation before it writes the step and publishes the step's size after that, so
// a worker holding a ticket of an earlier step fails its claim: either it reads that step's size, which
// its index has passed, or it reads the new size, and its claim then meets the closed ticket. A step
// cannot finish while a claimed chunk is outstanding, so the step's function and context stay valid for
// every successful claim.
static constexpr uint32_t CLOSED = 0xffffffffu;
void work(uint64_t gen) {
for (;;) {
uint64_t t = ticket_.load(std::memory_order_acquire);
if ((t >> 32) != gen) return;
if ((uint32_t)t == CLOSED) {
OASIS_PAUSE();
continue;
}
const int items = items_.load(std::memory_order_acquire), chunk = chunk_.load(std::memory_order_relaxed);
const int b = (int)(uint32_t)t;
if (b >= items) return;
if (!ticket_.compare_exchange_weak(t, t + (uint64_t)chunk, std::memory_order_acq_rel, std::memory_order_relaxed)) continue;
const int e = std::min(items, b + chunk);
fn_.load(std::memory_order_relaxed)(ctx_.load(std::memory_order_relaxed), b, e);
done_.fetch_add(e - b, std::memory_order_release);
}
}
void loop() {
uint64_t seen = ticket_.load(std::memory_order_acquire) >> 32;
for (;;) {
uint64_t g;
int spins = 0;
while ((g = ticket_.load(std::memory_order_acquire) >> 32) == seen) {
if (stop_.load(std::memory_order_relaxed)) return;
if (++spins >= 256) {
spins = 0;
if (!awake_.load(std::memory_order_relaxed)) {
std::unique_lock<std::mutex> lock(mutex_);
cv_.wait(lock, [&] {
return awake_.load() || stop_.load() || (ticket_.load(std::memory_order_acquire) >> 32) != seen;
});
}
}
OASIS_PAUSE();
}
seen = g;
work(g);
}
}
int n_;
std::vector<std::thread> workers_;
std::mutex mutex_;
std::condition_variable cv_;
std::atomic<bool> awake_{false}, stop_{false};
std::atomic<uint64_t> ticket_{0};
std::atomic<int> items_{0}, chunk_{1}, done_{0};
std::atomic<void (*)(void*, int, int)> fn_{nullptr};
std::atomic<void*> ctx_{nullptr};
};
} // namespace oasis
// ==== zlib inflate (RFC 1950 and 1951), for the tensors stored deflated
namespace oasis {
namespace {
// Canonical Huffman code. Codes of up to 9 bits resolve in one lookup of the next 9 input bits, longer
// ones by walking the code lengths.
struct Huffman {
uint16_t fast[512]; // length << 9 | symbol, or 0 when the next 9 bits start no code of up to 9 bits
uint16_t count[16]; // codes of each length
uint16_t symbol[288]; // symbols in code order
bool build(const uint8_t* lengths, int n) {
std::memset(count, 0, sizeof count);
for (int i = 0; i < n; i++) count[lengths[i]]++;
count[0] = 0;
int left = 1;
for (int l = 1; l < 16; l++) {
left = 2 * left - count[l];
if (left < 0) return false; // over-subscribed
}
uint16_t next[16] = {0};
for (int l = 1; l < 15; l++) next[l + 1] = (uint16_t)(next[l] + count[l]);
for (int i = 0; i < n; i++)
if (lengths[i]) symbol[next[lengths[i]]++] = (uint16_t)i;
std::memset(fast, 0, sizeof fast);
for (int l = 1, code = 0, k = 0; l <= 9; l++, code <<= 1)
for (int j = 0; j < count[l]; j++, code++, k++) {
int rev = 0; // deflate sends a code's bits from its most significant one
for (int b = 0; b < l; b++) rev |= ((code >> b) & 1) << (l - 1 - b);
for (int f = rev; f < 512; f += 1 << l) fast[f] = (uint16_t)(l << 9 | symbol[k]);
}
return true;
}
};
// Bits in stream order, least significant first. Past the end it supplies zero bytes and counts them, so
// that a stream which ends early is caught by overrun().
struct BitReader {
const uint8_t* p;
const uint8_t* end;
uint64_t buf = 0;
int cnt = 0;
size_t past = 0;
void fill() {
while (cnt <= 56) {
uint64_t byte = 0;
if (p < end) byte = *p++;
else past++;
buf |= byte << cnt;
cnt += 8;
}
}
uint32_t bits(int n) {
if (cnt < n) fill();
uint32_t v = (uint32_t)(buf & ((1ull << n) - 1));
buf >>= n;
cnt -= n;
return v;
}
bool overrun() const { return past * 8 > (size_t)cnt; }
};
int decode(BitReader& br, const Huffman& h) {
if (br.cnt < 16) br.fill();
int e = h.fast[br.buf & 511];
if (e) {
br.buf >>= e >> 9;
br.cnt -= e >> 9;
return e & 511;
}
int code = 0, first = 0, index = 0;
for (int l = 1; l < 16; l++) {
code |= (int)(br.buf & 1);
br.buf >>= 1;
br.cnt--;
int c = h.count[l];
if (code - first < c) return h.symbol[index + code - first];
index += c;
first = (first + c) << 1;
code <<= 1;
}
return -1;
}
// The literal/length and distance codes of a block with dynamic Huffman codes.
bool read_tables(BitReader& br, Huffman& lit, Huffman& dist) {
static const uint8_t ORDER[19] = {16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15};
int nlit = (int)br.bits(5) + 257, ndist = (int)br.bits(5) + 1, ncode = (int)br.bits(4) + 4;
if (nlit > 286 || ndist > 30) return false;
uint8_t cl[19] = {0};
for (int i = 0; i < ncode; i++) cl[ORDER[i]] = (uint8_t)br.bits(3);
Huffman code;
if (!code.build(cl, 19)) return false;
uint8_t lengths[286 + 30];
for (int k = 0; k < nlit + ndist;) {
int s = decode(br, code);
if (s < 0) return false;
if (s < 16) {
lengths[k++] = (uint8_t)s;
continue;
}
int rep, value = 0;
if (s == 16) {
if (k == 0) return false;
value = lengths[k - 1];
rep = 3 + (int)br.bits(2);
} else if (s == 17) {
rep = 3 + (int)br.bits(3);
} else {
rep = 11 + (int)br.bits(7);
}
if (k + rep > nlit + ndist) return false;
while (rep--) lengths[k++] = (uint8_t)value;
}
return lengths[256] != 0 && lit.build(lengths, nlit) && dist.build(lengths + nlit, ndist);
}
const uint16_t LEN_BASE[29] = {3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115,
131, 163, 195, 227, 258};
const uint8_t LEN_EXTRA[29] = {0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0};
const uint16_t DIST_BASE[30] = {1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537,
2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577};
const uint8_t DIST_EXTRA[30] = {0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13};
uint32_t adler32(const uint8_t* p, size_t n) {
uint32_t a = 1, b = 0;
while (n) {
size_t k = std::min<size_t>(n, 5552); // the most bytes before b can overflow
n -= k;
while (k--) {
a += *p++;
b += a;
}
a %= 65521;
b %= 65521;
}
return b << 16 | a;
}
} // namespace
// Inflates a zlib stream into exactly raw bytes; false when the stream is damaged, fails its checksum or
// holds another length.
bool inflate_exact(const uint8_t* in, size_t n, uint8_t* out, size_t raw) {
if (n < 6 || (in[0] & 15) != 8 || (in[0] >> 4) > 7 || ((in[0] << 8) | in[1]) % 31 || (in[1] & 32)) return false;
static const std::pair<Huffman, Huffman> fixed = [] {
uint8_t l[288 + 30];
for (int i = 0; i < 288; i++) l[i] = i < 144 ? 8 : i < 256 ? 9 : i < 280 ? 7 : 8;
for (int i = 0; i < 30; i++) l[288 + i] = 5;
std::pair<Huffman, Huffman> t;
t.first.build(l, 288);
t.second.build(l + 288, 30);
return t;
}();
BitReader br{in + 2, in + n - 4};
Huffman lit, dist;
size_t o = 0;
for (bool last = false; !last;) {
last = br.bits(1) != 0;
int type = (int)br.bits(2);
if (type == 0) {
int drop = br.cnt & 7; // stored blocks start on a byte boundary
br.buf >>= drop;
br.cnt -= drop;
uint32_t len = br.bits(16), nlen = br.bits(16);
if ((len ^ 0xffff) != nlen || len > raw - o) return false;
for (uint32_t i = 0; i < len; i++) out[o++] = (uint8_t)br.bits(8);
} else if (type == 1 || type == 2) {
const Huffman* L = &fixed.first;
const Huffman* D = &fixed.second;
if (type == 2) {
if (!read_tables(br, lit, dist)) return false;
L = &lit;
D = &dist;
}
for (;;) {
int s = decode(br, *L);
if (s < 256) {
if (s < 0 || o == raw) return false;
out[o++] = (uint8_t)s;
} else if (s == 256) {
break;
} else {
s -= 257;
if (s >= 29) return false;
size_t len = LEN_BASE[s] + br.bits(LEN_EXTRA[s]);
int d = decode(br, *D);
if (d < 0 || d >= 30) return false;
size_t back = DIST_BASE[d] + br.bits(DIST_EXTRA[d]);
if (back > o || len > raw - o) return false;
for (size_t i = 0; i < len; i++, o++) out[o] = out[o - back];
}
}
} else {
return false;
}
if (br.overrun()) return false;
}
const uint8_t* t = in + n - 4;
uint32_t want = (uint32_t)t[0] << 24 | (uint32_t)t[1] << 16 | (uint32_t)t[2] << 8 | t[3];
return o == raw && adler32(out, raw) == want;
}
} // namespace oasis
// ==== .oasis container loading and tensor access
namespace oasis {
float half_to_float(uint16_t h) {
#if defined(__F16C__)
return _cvtsh_ss(h);
#elif defined(__aarch64__) && (defined(__GNUC__) || defined(__clang__))
__fp16 v;
std::memcpy(&v, &h, 2);
return (float)v;
#else
uint32_t sign = (uint32_t)(h & 0x8000) << 16;
uint32_t exp = (h >> 10) & 0x1f;
uint32_t man = h & 0x3ff;
uint32_t bits;
if (exp == 0) {
if (man == 0) {
bits = sign;
} else {
exp = 127 - 15 + 1;
while (!(man & 0x400)) {
man <<= 1;
exp--;
}
man &= 0x3ff;
bits = sign | (exp << 23) | (man << 13);
}
} else if (exp == 31) {
bits = sign | 0x7f800000 | (man << 13);
} else {
bits = sign | ((exp + 127 - 15) << 23) | (man << 13);
}
float f;
std::memcpy(&f, &bits, 4);
return f;
#endif
}
template <typename T>
static T rd(const uint8_t*& p) {
T v;
std::memcpy(&v, p, sizeof(T));
p += sizeof(T);
return v;
}
bool Model::load(const std::string& path, std::string& err) {
std::ifstream f(path, std::ios::binary);
if (!f) {
err = "cannot open " + path;
return false;
}
file.assign(std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>());
if (file.size() < 16 || std::memcmp(file.data(), "OASIS\0v1", 8) != 0) {
err = "not an .oasis v1 file";
return false;
}
const uint8_t* p = file.data() + 8;
uint32_t hl = rd<uint32_t>(p);
std::string head((const char*)p, hl);
p += hl;
std::istringstream hs(head);
std::string line;
while (std::getline(hs, line)) {
auto sp = line.find(' ');
if (sp != std::string::npos) header[line.substr(0, sp)] = line.substr(sp + 1);
}
uint32_t n = rd<uint32_t>(p);
struct Raw { std::string name; Entry e; uint64_t off; int flags; };
std::vector<Raw> raws;
for (uint32_t i = 0; i < n; i++) {
uint16_t nl = rd<uint16_t>(p);
Raw r;
r.name.assign((const char*)p, nl);
p += nl;
r.e.dtype = rd<uint8_t>(p);
r.e.bits = rd<uint8_t>(p);
int nd = rd<uint8_t>(p);
r.flags = rd<uint8_t>(p);
for (int d = 0; d < nd; d++) r.e.shape.push_back((int)rd<uint32_t>(p));
r.off = rd<uint64_t>(p);
r.e.nbytes = (size_t)rd<uint64_t>(p);
raws.push_back(std::move(r));
}
size_t data_start = ((size_t)(p - file.data()) + 63) / 64 * 64;
for (auto& r : raws) {
if (data_start + r.off + r.e.nbytes > file.size()) {
err = "truncated tensor " + r.name;
return false;
}
r.e.data = file.data() + data_start + r.off;
if (r.flags & 1) {
// Deflated tensor: u64 size of the stored bytes, then a zlib stream of them.
if (r.e.nbytes < 8) { err = "bad compressed tensor " + r.name; return false; }
uint64_t raw;
std::memcpy(&raw, r.e.data, 8);
inflated.emplace_back(raw);
if (!inflate_exact(r.e.data + 8, r.e.nbytes - 8, inflated.back().data(), raw)) {
err = "cannot inflate tensor " + r.name;
return false;
}
r.e.data = inflated.back().data();
r.e.nbytes = raw;
}
entries[r.name] = r.e;
}
auto num = [&](const char* k) { return std::stoi(header.at(k)); };
vocab = num("vocab");
n_pieces = num("n_pieces");
bos = num("bos");
eos = num("eos");
pad = num("pad");
ref_base = num("ref_base");
std::istringstream ls(header.at("langs"));
std::string tokn;
std::vector<std::string> all;
while (ls >> tokn) {
auto c = tokn.find(':');
langs[tokn.substr(0, c)] = std::stoi(tokn.substr(c + 1));
all.push_back(tokn.substr(0, c));
}
auto lang_line = header.find("languages");
if (lang_line != header.end()) {
std::istringstream as(lang_line->second);
while (as >> tokn)
if (langs.count(tokn)) languages.push_back(tokn);
}
if (languages.empty()) languages = all;
const float* cbp = (const float*)E("quant/codebook").data;
codebook.assign(cbp, cbp + E("quant/codebook").shape[0]);
const Entry& pe = E("tokenizer/pieces");
const uint8_t* q = pe.data;
uint32_t np = rd<uint32_t>(q);
for (uint32_t i = 0; i < np; i++) {
int t = rd<uint8_t>(q);
uint16_t l = rd<uint16_t>(q);
pieces.emplace_back((const char*)q, l);
piece_types.push_back(t);
q += l;
const std::string& p = pieces.back();
if (t == 0) piece_id.emplace(p, (int)i);
else if (t == 3 && !p.empty()) markers.push_back({p, (int)i});
else if (t == 4 && p.size() == 6 && p.compare(0, 3, "<0x") == 0) byte_id[std::stoi(p.substr(3, 2), nullptr, 16) & 255] = (int)i;
}
std::stable_sort(markers.begin(), markers.end(), [](const auto& a, const auto& b) { return a.first.size() > b.first.size(); });
const Entry& ce = E("reader/charset");
for (int i = 0; i < ce.shape[0]; i++) charset.push_back((uint32_t)((const float*)ce.data)[i]);
for (int i = 0; i < 512; i++) {
perm1.push_back((int)((const float*)E("hmlp/perm1").data)[i]);
perm2.push_back((int)((const float*)E("hmlp/perm2").data)[i]);
}
seq_order.resize(512);
for (int j = 0; j < 512; j++) {
int g = j ^ (j >> 1), s = 0;
for (int b = 0; b < 9; b++)
if (g & (1 << b)) s |= 1 << (8 - b);
seq_order[j] = s;
}
// Expand every fp16 and integer-row tensor now, and decode every 8-bit scale table, so that lookups
// during inference never write to the model.
for (auto& kv : entries) {
const Entry& e = kv.second;
if (e.dtype == 1) {
size_t n = e.nbytes / 2;
std::vector<float> v(n);
const uint16_t* h = (const uint16_t*)e.data;
for (size_t i = 0; i < n; i++) v[i] = half_to_float(h[i]);
cache[kv.first] = std::move(v);
} else if (e.dtype == 4) {
// Signed integers of e.bits bits in one LSB-first bit stream, then one f32 scale per row.
int rows = e.shape[0], b = e.bits;
size_t cols = 1;
for (size_t d = 1; d < e.shape.size(); d++) cols *= (size_t)e.shape[d];
size_t n = (size_t)rows * cols, packed = (n * b + 7) / 8;
if (b < 2 || b > 16 || e.nbytes != packed + (size_t)rows * 4) { err = "bad integer tensor " + kv.first; return false; }
std::vector<float> v(n);
for (int r = 0; r < rows; r++) {
float s;
std::memcpy(&s, e.data + packed + (size_t)r * 4, 4);
for (size_t c = 0; c < cols; c++) {
size_t bit = ((size_t)r * cols + c) * b;
uint32_t w = 0;
for (int k = 0; k < b; k++, bit++) w |= (uint32_t)(e.data[bit >> 3] >> (bit & 7) & 1) << k;
int32_t x = (int32_t)(w << (32 - b)) >> (32 - b);
v[(size_t)r * cols + c] = s * (float)x;
}
}
cache[kv.first] = std::move(v);
} else if (e.dtype == 5) {
// Code 0 is a zero scale; codes 1..255 step evenly in log scale from exp(lo) to exp(hi).
float lo, hi;
std::memcpy(&lo, e.data + e.nbytes - 8, 4);
std::memcpy(&hi, e.data + e.nbytes - 4, 4);
std::vector<float> lut(256, 0.0f);
for (int c = 1; c < 256; c++) lut[c] = std::exp(lo + (hi - lo) * (float)(c - 1) / 254.0f);
luts[kv.first] = std::move(lut);
}
if (e.dtype == 2 || e.dtype == 5) {
// A row whose every group scale is zero multiplies to zero; the kernels skip it.
int rows = e.shape[0];
for (size_t d = 1; d + 1 < e.shape.size(); d++) rows *= e.shape[d];
const int groups = (e.shape.back() + 127) / 128;
const uint8_t* sc = e.data + (size_t)rows * groups * 128 * e.bits / 8;
std::vector<uint8_t> z(rows, 1);
bool any = false;
for (int r = 0; r < rows; r++) {
for (int g = 0; g < groups && z[r]; g++) {
size_t i = (size_t)r * groups + g;
z[r] = e.dtype == 5 ? sc[i] == 0 : (((const uint16_t*)sc)[i] & 0x7fff) == 0;
}
any = any || z[r];
}
if (any) zero_rows[kv.first] = std::move(z);
}
}
#if defined(OASIS_NEON)
// On Arm the HConvGLU bands are kept tap-major ([65, rows]), the order the band kernel reads them in.
for (auto& kv : cache) {
const std::string& k = kv.first;
auto ends = [&](const char* s) { size_t n = std::strlen(s); return k.size() >= n && k.compare(k.size() - n, n, s) == 0; };
if (!ends("hconv/gate_conv") && !ends("hconv/value_conv") && !ends("hconv/out_conv")) continue;
const Entry& e = entries.at(k);
if (e.shape.size() != 2 || e.shape[0] != 512 || e.shape[1] != 65) continue;
std::vector<float> t((size_t)65 * 512);
for (int r = 0; r < 512; r++)
for (int j = 0; j < 65; j++) t[(size_t)j * 512 + r] = kv.second[(size_t)r * 65 + j];
kv.second.swap(t);
}
#endif
return true;
}
const Entry& Model::E(const std::string& name) const {
auto it = entries.find(name);
if (it == entries.end()) throw std::runtime_error("missing tensor " + name);
return it->second;
}
const float* Model::F(const std::string& name) {
const Entry& e = E(name);
if (e.dtype == 0) return (const float*)e.data;
if (e.dtype != 1 && e.dtype != 4) throw std::runtime_error("tensor is not a float array: " + name);
return cache.at(name).data();
}
CQ Model::Q(const std::string& name) const {
const Entry& e = E(name);
if (e.dtype != 2 && e.dtype != 5) throw std::runtime_error("tensor is not quantised: " + name);
CQ q;
q.rows = e.shape[0];
for (size_t d = 1; d + 1 < e.shape.size(); d++) q.rows *= e.shape[d];
q.cols = e.shape.back();
q.pc = (q.cols + 127) / 128 * 128;
q.bits = e.bits;
q.idx = e.data;
const uint8_t* scales = e.data + (size_t)q.rows * q.pc * q.bits / 8;
if (e.dtype == 5) {
q.ncode = scales;
q.nlut = luts.at(name).data();
} else {
q.norms = (const uint16_t*)scales;
}
q.cb = codebook.data() + cb_offset[q.bits];
auto zr = zero_rows.find(name);
if (zr != zero_rows.end()) q.zero_rows = zr->second.data();
return q;
}
const float* Model::Wdq(const std::string& name) {
std::lock_guard<std::mutex> lock(dq_mutex);
std::string key = "#dq:" + name;
auto c = dq.find(key);
if (c != dq.end()) return c->second.data();
CQ q = Q(name);
std::vector<float> v((size_t)q.rows * q.cols);
for (int r = 0; r < q.rows; r++) cq_dequant_row(q, r, v.data() + (size_t)r * q.cols);
return (dq[key] = std::move(v)).data();
}
} // namespace oasis
// ==== Numeric kernels: Walsh-Hadamard transforms, codebook matrices, GEMM, norms and bands
namespace oasis {
int engine_threads() {
static const int n = [] {
#ifdef _OPENMP
const char* e = std::getenv("OMP_NUM_THREADS");
if (e && *e) return std::max(1, omp_get_max_threads());
const int p = omp_get_num_procs();
#else
const int p = (int)std::thread::hardware_concurrency();
#endif
return p <= 4 ? std::max(1, p - 1) : p;
}();
return n;
}
void wht_inplace(float* x, int n) {
#if defined(OASIS_NEON)
// The same sums and differences on every element as the loops below, four elements at a time.
if (n >= 8) {
for (int i = 0; i < n; i += 8) { // pairs one apart
const float32x4_t a = vld1q_f32(x + i), b = vld1q_f32(x + i + 4);
const float32x4_t e = vuzp1q_f32(a, b), o = vuzp2q_f32(a, b);
const float32x4_t s = vaddq_f32(e, o), d = vsubq_f32(e, o);
vst1q_f32(x + i, vzip1q_f32(s, d));
vst1q_f32(x + i + 4, vzip2q_f32(s, d));
}
for (int i = 0; i < n; i += 8) { // pairs two apart
const float64x2_t a = vreinterpretq_f64_f32(vld1q_f32(x + i)), b = vreinterpretq_f64_f32(vld1q_f32(x + i + 4));
const float32x4_t lo = vreinterpretq_f32_f64(vzip1q_f64(a, b)), hi = vreinterpretq_f32_f64(vzip2q_f64(a, b));
const float64x2_t s = vreinterpretq_f64_f32(vaddq_f32(lo, hi)), d = vreinterpretq_f64_f32(vsubq_f32(lo, hi));
vst1q_f32(x + i, vreinterpretq_f32_f64(vzip1q_f64(s, d)));
vst1q_f32(x + i + 4, vreinterpretq_f32_f64(vzip2q_f64(s, d)));
}
for (int len = 4; len < n; len <<= 1)
for (int i = 0; i < n; i += len << 1)
for (int j = i; j < i + len; j += 4) {
const float32x4_t a = vld1q_f32(x + j), b = vld1q_f32(x + j + len);
vst1q_f32(x + j, vaddq_f32(a, b));
vst1q_f32(x + j + len, vsubq_f32(a, b));
}
const float s = 1.0f / std::sqrt((float)n);
for (int i = 0; i < n; i += 4) vst1q_f32(x + i, vmulq_n_f32(vld1q_f32(x + i), s));
return;
}
#endif
for (int len = 1; len < n; len <<= 1)
for (int i = 0; i < n; i += len << 1)
for (int j = i; j < i + len; j++) {
float a = x[j], b = x[j + len];
x[j] = a + b;
x[j + len] = a - b;
}
float s = 1.0f / std::sqrt((float)n);
for (int i = 0; i < n; i++) x[i] *= s;
}
void wht_seq512(const Model& m, const float* x, float* out) {
float t[512];
std::memcpy(t, x, sizeof t);
wht_inplace(t, 512);
for (int j = 0; j < 512; j++) out[j] = t[m.seq_order[j]];
}
static inline int cq_index(const uint8_t* base, size_t pos, int bits) {
size_t bit = pos * bits;
size_t byte = bit >> 3;
int shift = (int)(bit & 7);
unsigned v = base[byte] >> shift;
if (shift + bits > 8) v |= (unsigned)base[byte + 1] << (8 - shift);
return (int)(v & ((1u << bits) - 1));
}
void cq_dequant_row(const CQ& q, int row, float* out) {
int groups = q.pc / 128;
float g[128];
for (int gi = 0; gi < groups; gi++) {
float s = cq_scale(q, (size_t)row * groups + gi);
size_t base = (size_t)row * q.pc + (size_t)gi * 128;
for (int i = 0; i < 128; i++) g[i] = q.cb[cq_index(q.idx, base + i, q.bits)] * s;
wht_inplace(g, 128);
int n = std::min(128, q.cols - gi * 128);
if (n > 0) std::memcpy(out + gi * 128, g, sizeof(float) * n);
}
}
#if defined(__AVX2__)
#include <immintrin.h>
// Every AVX2 processor also has FMA3; MSVC enables both with /arch:AVX2 but defines no __FMA__.
#if defined(__FMA__) || defined(_MSC_VER)
#define OASIS_FMA 1
#endif
static inline __m256i dot_i8(__m256i acc, __m256i ax, __m256i ws) {
#if defined(__AVXVNNI__)
return _mm256_dpbusd_avx_epi32(acc, ax, ws);
#else
return _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_maddubs_epi16(ax, ws), _mm256_set1_epi16(1)));
#endif
}
static inline int hsum_i32(__m256i v) {
__m128i s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1));
s = _mm_add_epi32(s, _mm_shuffle_epi32(s, 0x4E));
s = _mm_add_epi32(s, _mm_shuffle_epi32(s, 0xB1));
return _mm_cvtsi128_si32(s);
}
#elif defined(OASIS_NEON)
// Signed int8 products of a and b added in fours into the int32 lanes of acc: SDOT where the core has the
// dot-product instructions, otherwise widening multiplies with pairwise accumulation. Both are exact.
static inline int32x4_t dot_i8(int32x4_t acc, int8x16_t a, int8x16_t b) {
#if defined(__ARM_FEATURE_DOTPROD)
return vdotq_s32(acc, a, b);
#else
acc = vpadalq_s16(acc, vmull_s8(vget_low_s8(a), vget_low_s8(b)));
return vpadalq_s16(acc, vmull_high_s8(a, b));
#endif
}
#endif
bool g_exact_gemv = false;
void prep_x(const float* x, int cols, int bits, QX& o) {
o.pc = (cols + 127) / 128 * 128;
o.bits = bits;
int groups = o.pc / 128;
float xt[2048];
std::memset(xt, 0, sizeof(float) * o.pc);
std::memcpy(xt, x, sizeof(float) * cols);
o.xs.resize(o.pc);
o.xa.resize(o.pc);
o.sx.resize(groups);
for (int g = 0; g < groups; g++) {
float* v = xt + g * 128;
wht_inplace(v, 128);
#if defined(OASIS_NEON)
int8_t* dst = o.xs.data() + g * 128;
// The same scale, roundings and reordering as the loops below, sixteen values at a time.
float32x4_t mx = vdupq_n_f32(0.0f);
for (int i = 0; i < 128; i += 4) mx = vmaxnmq_f32(mx, vabsq_f32(vld1q_f32(v + i)));
const float amax = vmaxnmvq_f32(mx);
const float inv = amax > 0 ? 127.0f / amax : 0.0f;
o.sx[g] = amax / 127.0f;
int8_t tmp[128];
for (int i = 0; i < 128; i += 16) {
int32x4_t q[4];
for (int k = 0; k < 4; k++) q[k] = vcvtnq_s32_f32(vmulq_n_f32(vld1q_f32(v + i + 4 * k), inv));
const int16x8_t h0 = vcombine_s16(vmovn_s32(q[0]), vmovn_s32(q[1])), h1 = vcombine_s16(vmovn_s32(q[2]), vmovn_s32(q[3]));
vst1q_s8(tmp + i, vcombine_s8(vmovn_s16(h0), vmovn_s16(h1)));
}
if (bits == 2) {
const int8x16x4_t a = vld4q_s8(tmp), b = vld4q_s8(tmp + 64);
for (int k = 0; k < 4; k++) {
vst1q_s8(dst + k * 32, a.val[k]);
vst1q_s8(dst + k * 32 + 16, b.val[k]);
}
} else {
for (int c = 0; c < 2; c++) {
const int8x16x2_t a = vld2q_s8(tmp + 64 * c), b = vld2q_s8(tmp + 64 * c + 32);
vst1q_s8(dst + c * 64, a.val[0]);
vst1q_s8(dst + c * 64 + 16, b.val[0]);
vst1q_s8(dst + c * 64 + 32, a.val[1]);
vst1q_s8(dst + c * 64 + 48, b.val[1]);
}
}
for (int i = 0; i < 128; i += 16) vst1q_s8(o.xa.data() + g * 128 + i, vabsq_s8(vld1q_s8(dst + i)));
#else
float amax = 0.0f;
for (int i = 0; i < 128; i++) amax = std::max(amax, std::fabs(v[i]));
float inv = amax > 0 ? 127.0f / amax : 0.0f;
o.sx[g] = amax / 127.0f;
int8_t tmp[128];
for (int i = 0; i < 128; i++) tmp[i] = (int8_t)std::nearbyint(v[i] * inv);
int8_t* dst = o.xs.data() + g * 128;
if (bits == 2) {
for (int k = 0; k < 4; k++)
for (int mm = 0; mm < 32; mm++) dst[k * 32 + mm] = tmp[4 * mm + k];
} else {
for (int c = 0; c < 2; c++)
for (int mm = 0; mm < 32; mm++) {
dst[c * 64 + mm] = tmp[2 * (32 * c + mm)];
dst[c * 64 + 32 + mm] = tmp[2 * (32 * c + mm) + 1];
}
}
for (int i = 0; i < 128; i++) o.xa[g * 128 + i] = (int8_t)std::abs((int)dst[i]);
#endif
}
}
CbTable cb_table(const CQ& q) {
CbTable t;
int ncb = 1 << q.bits;
float cmax = 0.0f;
for (int i = 0; i < ncb; i++) cmax = std::max(cmax, std::fabs(q.cb[i]));
t.cs = cmax / 127.0f;
std::memset(t.tbl, 0, sizeof t.tbl);
for (int i = 0; i < ncb; i++) t.tbl[i] = (int8_t)std::nearbyint(q.cb[i] / t.cs);
return t;
}
// One output of W x for a quantised row: int8 dot products per rotated group, scaled per group.
float cq_row(const CQ& q, const QX& qx, int r, const CbTable& t) {
int groups = q.pc / 128;
size_t row_bytes = (size_t)q.pc * q.bits / 8;
const uint8_t* ib = q.idx + (size_t)r * row_bytes;
float acc = 0.0f;
#if defined(__AVX2__)
const __m256i T = _mm256_broadcastsi128_si256(_mm_load_si128((const __m128i*)t.tbl));
const __m256i m3 = _mm256_set1_epi8(3), m15 = _mm256_set1_epi8(15);
for (int g = 0; g < groups; g++) {
const int8_t* xg = qx.xs.data() + g * 128;
const int8_t* ag = qx.xa.data() + g * 128;
__m256i a = _mm256_setzero_si256();
if (q.bits == 2) {
__m256i b = _mm256_loadu_si256((const __m256i*)(ib + g * 32));
for (int k = 0; k < 4; k++) {
__m256i w = _mm256_shuffle_epi8(T, _mm256_and_si256(_mm256_srli_epi16(b, 2 * k), m3));
a = dot_i8(a, _mm256_loadu_si256((const __m256i*)(ag + k * 32)),
_mm256_sign_epi8(w, _mm256_loadu_si256((const __m256i*)(xg + k * 32))));
}
} else {
for (int c = 0; c < 2; c++) {
__m256i b = _mm256_loadu_si256((const __m256i*)(ib + g * 64 + c * 32));
__m256i lo = _mm256_shuffle_epi8(T, _mm256_and_si256(b, m15));
__m256i hi = _mm256_shuffle_epi8(T, _mm256_and_si256(_mm256_srli_epi16(b, 4), m15));
a = dot_i8(a, _mm256_loadu_si256((const __m256i*)(ag + c * 64)),
_mm256_sign_epi8(lo, _mm256_loadu_si256((const __m256i*)(xg + c * 64))));
a = dot_i8(a, _mm256_loadu_si256((const __m256i*)(ag + c * 64 + 32)),
_mm256_sign_epi8(hi, _mm256_loadu_si256((const __m256i*)(xg + c * 64 + 32))));
}
}
acc += (float)hsum_i32(a) * qx.sx[g] * cq_scale(q, (size_t)r * groups + g);
}
#elif defined(OASIS_NEON)
// The same integer sums as the portable path below, sixteen weights per table lookup.
const int8x16_t T = vld1q_s8(t.tbl);
for (int g = 0; g < groups; g++) {
const int8_t* xg = qx.xs.data() + g * 128;
int32x4_t a = vdupq_n_s32(0);
if (q.bits == 2) {
const uint8x16_t m3 = vdupq_n_u8(3);
const uint8x16_t b0 = vld1q_u8(ib + g * 32), b1 = vld1q_u8(ib + g * 32 + 16);
a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(b0, m3)), vld1q_s8(xg));
a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(b1, m3)), vld1q_s8(xg + 16));
a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b0, 2), m3)), vld1q_s8(xg + 32));
a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b1, 2), m3)), vld1q_s8(xg + 48));
a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b0, 4), m3)), vld1q_s8(xg + 64));
a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b1, 4), m3)), vld1q_s8(xg + 80));
a = dot_i8(a, vqtbl1q_s8(T, vshrq_n_u8(b0, 6)), vld1q_s8(xg + 96));
a = dot_i8(a, vqtbl1q_s8(T, vshrq_n_u8(b1, 6)), vld1q_s8(xg + 112));
} else {
const uint8x16_t m15 = vdupq_n_u8(15);
for (int c = 0; c < 2; c++) {
const uint8_t* bp = ib + g * 64 + c * 32;
const int8_t* xp = xg + c * 64;
const uint8x16_t b0 = vld1q_u8(bp), b1 = vld1q_u8(bp + 16);
a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(b0, m15)), vld1q_s8(xp));
a = dot_i8(a, vqtbl1q_s8(T, vandq_u8(b1, m15)), vld1q_s8(xp + 16));
a = dot_i8(a, vqtbl1q_s8(T, vshrq_n_u8(b0, 4)), vld1q_s8(xp + 32));
a = dot_i8(a, vqtbl1q_s8(T, vshrq_n_u8(b1, 4)), vld1q_s8(xp + 48));
}
}
acc += (float)vaddvq_s32(a) * qx.sx[g] * cq_scale(q, (size_t)r * groups + g);
}
#else
// Portable path: the same int8 arithmetic, one weight at a time.
for (int g = 0; g < groups; g++) {
const int8_t* xg = qx.xs.data() + g * 128;
int a = 0;
if (q.bits == 2) {
for (int k = 0; k < 4; k++)
for (int mm = 0; mm < 32; mm++) a += t.tbl[(ib[g * 32 + mm] >> (2 * k)) & 3] * xg[k * 32 + mm];
} else {
for (int c = 0; c < 2; c++)
for (int mm = 0; mm < 32; mm++) {
uint8_t bb = ib[g * 64 + c * 32 + mm];
a += t.tbl[bb & 15] * xg[c * 64 + mm] + t.tbl[bb >> 4] * xg[c * 64 + 32 + mm];
}
}
acc += (float)a * qx.sx[g] * cq_scale(q, (size_t)r * groups + g);
}
#endif
return acc * t.cs;
}
static void cq_rows_q(const CQ& q, const QX& qx, float* y, bool parallel) {
CbTable t = cb_table(q);
#pragma omp parallel for schedule(static) if (parallel)
for (int r = 0; r < q.rows; r++) y[r] = q.zero_rows && q.zero_rows[r] ? 0.0f : cq_row(q, qx, r, t);
}
void cq_gemv_q(const CQ& q, const QX& qx, float* y) { cq_rows_q(q, qx, y, q.rows >= 256); }
// C[M, rows] = A[M, cols] W^T with int8 activations per rotated 128-group, as in the deployed kernels.
#if defined(OASIS_NEON)
// Every row of W against four prepared inputs at once: a group's weights are looked up once for all four,
// and each output gets the integer sums and the float steps of cq_row as GCC compiles it: with 8-bit scales
// the group's scale times the input's, then a multiply-add of that and the sum; with fp16 scales the sum
// times the input's scale, then a multiply-add of that and the group's scale.
static void cq_rows4(const CQ& q, const QX* qx, const CbTable& t, float* const* y) {
const int groups = q.pc / 128;
const size_t row_bytes = (size_t)q.pc * q.bits / 8;
const int8x16_t T = vld1q_s8(t.tbl);
const uint8x16_t m3 = vdupq_n_u8(3), m15 = vdupq_n_u8(15);
for (int r = 0; r < q.rows; r++) {
if (q.zero_rows && q.zero_rows[r]) {
for (int i = 0; i < 4; i++) y[i][r] = 0.0f;
continue;
}
const uint8_t* ib = q.idx + (size_t)r * row_bytes;
float32x4_t acc = vdupq_n_f32(0.0f);
for (int g = 0; g < groups; g++) {
int8x16_t w[8];
if (q.bits == 2) {
const uint8x16_t b0 = vld1q_u8(ib + g * 32), b1 = vld1q_u8(ib + g * 32 + 16);
w[0] = vqtbl1q_s8(T, vandq_u8(b0, m3));
w[1] = vqtbl1q_s8(T, vandq_u8(b1, m3));
w[2] = vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b0, 2), m3));
w[3] = vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b1, 2), m3));
w[4] = vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b0, 4), m3));
w[5] = vqtbl1q_s8(T, vandq_u8(vshrq_n_u8(b1, 4), m3));
w[6] = vqtbl1q_s8(T, vshrq_n_u8(b0, 6));
w[7] = vqtbl1q_s8(T, vshrq_n_u8(b1, 6));
} else {
for (int c = 0; c < 2; c++) {
const uint8x16_t b0 = vld1q_u8(ib + g * 64 + c * 32), b1 = vld1q_u8(ib + g * 64 + c * 32 + 16);
w[4 * c] = vqtbl1q_s8(T, vandq_u8(b0, m15));
w[4 * c + 1] = vqtbl1q_s8(T, vandq_u8(b1, m15));
w[4 * c + 2] = vqtbl1q_s8(T, vshrq_n_u8(b0, 4));
w[4 * c + 3] = vqtbl1q_s8(T, vshrq_n_u8(b1, 4));
}
}
int32x4_t a[4];
float sx[4];
for (int i = 0; i < 4; i++) {
const int8_t* xg = qx[i].xs.data() + g * 128;
int32x4_t s = vdupq_n_s32(0);
for (int k = 0; k < 8; k++) s = dot_i8(s, w[k], vld1q_s8(xg + 16 * k));
a[i] = s;
sx[i] = qx[i].sx[g];
}
const float32x4_t sum = vcvtq_f32_s32(vpaddq_s32(vpaddq_s32(a[0], a[1]), vpaddq_s32(a[2], a[3])));
const float scale = cq_scale(q, (size_t)r * groups + g);
if (q.ncode) acc = vfmaq_f32(acc, vmulq_n_f32(vld1q_f32(sx), scale), sum);
else acc = vfmaq_f32(acc, vmulq_f32(sum, vld1q_f32(sx)), vdupq_n_f32(scale));
}
float out[4];
vst1q_f32(out, vmulq_n_f32(acc, t.cs));
for (int i = 0; i < 4; i++) y[i][r] = out[i];
}
}
#endif
void cq_gemm_q(const CQ& q, const float* A, int M, float* C) {
CbTable t = cb_table(q);
#if defined(OASIS_NEON)
const int blocks = (M + 3) / 4;
#pragma omp parallel for schedule(static) if (M > 1)
for (int b = 0; b < blocks; b++) {
thread_local QX qx[4];
const int i0 = b * 4, n = std::min(4, M - i0);
for (int i = 0; i < n; i++) prep_x(A + (size_t)(i0 + i) * q.cols, q.cols, q.bits, qx[i]);
if (n == 4) {
float* y[4] = {C + (size_t)i0 * q.rows, C + (size_t)(i0 + 1) * q.rows, C + (size_t)(i0 + 2) * q.rows, C + (size_t)(i0 + 3) * q.rows};
cq_rows4(q, qx, t, y);
continue;
}
for (int i = 0; i < n; i++) {
float* y = C + (size_t)(i0 + i) * q.rows;
for (int r = 0; r < q.rows; r++) y[r] = q.zero_rows && q.zero_rows[r] ? 0.0f : cq_row(q, qx[i], r, t);
}
}
#else
#pragma omp parallel for schedule(static) if (M > 1)
for (int i = 0; i < M; i++) {
thread_local QX qx;
prep_x(A + (size_t)i * q.cols, q.cols, q.bits, qx);
float* y = C + (size_t)i * q.rows;
for (int r = 0; r < q.rows; r++) y[r] = q.zero_rows && q.zero_rows[r] ? 0.0f : cq_row(q, qx, r, t);
}
#endif
}
float dotf(const float* a, const float* b, int n) {
#if defined(OASIS_FMA)
__m256 s0 = _mm256_setzero_ps(), s1 = _mm256_setzero_ps();
int i = 0;
for (; i + 16 <= n; i += 16) {
s0 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i), _mm256_loadu_ps(b + i), s0);
s1 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i + 8), _mm256_loadu_ps(b + i + 8), s1);
}
s0 = _mm256_add_ps(s0, s1);
__m128 h = _mm_add_ps(_mm256_castps256_ps128(s0), _mm256_extractf128_ps(s0, 1));
h = _mm_add_ps(h, _mm_movehl_ps(h, h));
h = _mm_add_ss(h, _mm_shuffle_ps(h, h, 1));
float s = _mm_cvtss_f32(h);
for (; i < n; i++) s += a[i] * b[i];
return s;
#elif defined(OASIS_NEON)
float32x4_t s0 = vdupq_n_f32(0.0f), s1 = s0, s2 = s0, s3 = s0;
int i = 0;
for (; i + 16 <= n; i += 16) {
s0 = vfmaq_f32(s0, vld1q_f32(a + i), vld1q_f32(b + i));
s1 = vfmaq_f32(s1, vld1q_f32(a + i + 4), vld1q_f32(b + i + 4));
s2 = vfmaq_f32(s2, vld1q_f32(a + i + 8), vld1q_f32(b + i + 8));
s3 = vfmaq_f32(s3, vld1q_f32(a + i + 12), vld1q_f32(b + i + 12));
}
float s = vaddvq_f32(vaddq_f32(vaddq_f32(s0, s1), vaddq_f32(s2, s3)));
for (; i < n; i++) s += a[i] * b[i];
return s;
#else
float s = 0.0f;
for (int i = 0; i < n; i++) s += a[i] * b[i];
return s;
#endif
}
void cq_gemv(const CQ& q, const float* x, float* y) {
#if defined(__AVX2__) || defined(OASIS_NEON)
if (!g_exact_gemv && (q.bits == 2 || q.bits == 4)) {
thread_local QX qx;
prep_x(x, q.cols, q.bits, qx);
cq_gemv_q(q, qx, y);
return;
}
#endif
int groups = q.pc / 128;
std::vector<float> xt(q.pc, 0.0f);
std::memcpy(xt.data(), x, sizeof(float) * q.cols);
for (int gi = 0; gi < groups; gi++) wht_inplace(xt.data() + gi * 128, 128);
if (q.bits == 2) {
alignas(16) float tab[256][4];
for (int b = 0; b < 256; b++)
for (int k = 0; k < 4; k++) tab[b][k] = q.cb[(b >> (2 * k)) & 3];
size_t row_bytes = (size_t)q.pc / 4;
#pragma omp parallel for schedule(static) if (q.rows >= 256)
for (int r = 0; r < q.rows; r++) {
const uint8_t* ib = q.idx + (size_t)r * row_bytes;
float acc = 0.0f;
for (int gi = 0; gi < groups; gi++) {
const uint8_t* gb = ib + gi * 32;
const float* xg = xt.data() + gi * 128;
float a0 = 0, a1 = 0, a2 = 0, a3 = 0;
for (int m = 0; m < 32; m++) {
const float* t = tab[gb[m]];
const float* xv = xg + 4 * m;
a0 += t[0] * xv[0];
a1 += t[1] * xv[1];
a2 += t[2] * xv[2];
a3 += t[3] * xv[3];
}
acc += (a0 + a1 + a2 + a3) * cq_scale(q, (size_t)r * groups + gi);
}
y[r] = acc;
}
return;
}
if (q.bits == 4) {
alignas(8) float tab[256][2];
for (int b = 0; b < 256; b++) {
tab[b][0] = q.cb[b & 15];
tab[b][1] = q.cb[b >> 4];
}
size_t row_bytes = (size_t)q.pc / 2;
#pragma omp parallel for schedule(static) if (q.rows >= 256)
for (int r = 0; r < q.rows; r++) {
const uint8_t* ib = q.idx + (size_t)r * row_bytes;
float acc = 0.0f;
for (int gi = 0; gi < groups; gi++) {
const uint8_t* gb = ib + gi * 64;
const float* xg = xt.data() + gi * 128;
float a0 = 0, a1 = 0;
for (int m = 0; m < 64; m++) {
a0 += tab[gb[m]][0] * xg[2 * m];
a1 += tab[gb[m]][1] * xg[2 * m + 1];
}
acc += (a0 + a1) * cq_scale(q, (size_t)r * groups + gi);
}
y[r] = acc;
}
return;
}
#pragma omp parallel for schedule(static) if (q.rows >= 256)
for (int r = 0; r < q.rows; r++) {
float acc = 0.0f;
for (int gi = 0; gi < groups; gi++) {
float a = 0.0f;
size_t base = (size_t)r * q.pc + (size_t)gi * 128;
for (int i = 0; i < 128; i++) a += q.cb[cq_index(q.idx, base + i, q.bits)] * xt[gi * 128 + i];
acc += a * cq_scale(q, (size_t)r * groups + gi);
}
y[r] = acc;
}
}
void gemm_nt(const float* A, int M, int K, const float* W, int N, float* C) {
if (M == 1) {
#pragma omp parallel for schedule(static) if ((long)N * K > 200000)
for (int j = 0; j < N; j++) C[j] = dotf(A, W + (size_t)j * K, K);
return;
}
#pragma omp parallel for schedule(static) if ((long)M * N * K > 200000)
for (int i = 0; i < M; i++)
for (int j = 0; j < N; j++) C[(size_t)i * N + j] = dotf(A + (size_t)i * K, W + (size_t)j * K, K);
}
// C = A B. Every output accumulates its K products in order (fused multiply-adds where the target
// has them), so the register-blocked path and the plain loop give the same values.
static void gemm_nn_rows(const float* A, int i0, int i1, int K, const float* B, int N, float* C, int j0) {
for (int i = i0; i < i1; i++) {
float* c = C + (size_t)i * N;
const float* a = A + (size_t)i * K;
for (int j = j0; j < N; j++) c[j] = 0.0f;
for (int k = 0; k < K; k++) {
float av = a[k];
const float* b = B + (size_t)k * N;
for (int j = j0; j < N; j++) c[j] += av * b[j];
}
}
}
void gemm_nn(const float* A, int M, int K, const float* B, int N, float* C) {
#if defined(OASIS_FMA)
const int NB = N / 16 * 16;
const int blocks = (M + 3) / 4;
#pragma omp parallel for schedule(static) if ((long)M * N * K > 200000)
for (int ib = 0; ib < blocks; ib++) {
int i0 = ib * 4, mr = std::min(4, M - i0);
if (mr == 4) {
const float* a0 = A + (size_t)i0 * K;
const float* a1 = a0 + K;
const float* a2 = a1 + K;
const float* a3 = a2 + K;
for (int j0 = 0; j0 < NB; j0 += 16) {
__m256 c00 = _mm256_setzero_ps(), c01 = _mm256_setzero_ps(), c10 = _mm256_setzero_ps(), c11 = _mm256_setzero_ps();
__m256 c20 = _mm256_setzero_ps(), c21 = _mm256_setzero_ps(), c30 = _mm256_setzero_ps(), c31 = _mm256_setzero_ps();
for (int k = 0; k < K; k++) {
const float* b = B + (size_t)k * N + j0;
__m256 b0 = _mm256_loadu_ps(b), b1 = _mm256_loadu_ps(b + 8), x;
x = _mm256_broadcast_ss(a0 + k);
c00 = _mm256_fmadd_ps(x, b0, c00);
c01 = _mm256_fmadd_ps(x, b1, c01);
x = _mm256_broadcast_ss(a1 + k);
c10 = _mm256_fmadd_ps(x, b0, c10);
c11 = _mm256_fmadd_ps(x, b1, c11);
x = _mm256_broadcast_ss(a2 + k);
c20 = _mm256_fmadd_ps(x, b0, c20);
c21 = _mm256_fmadd_ps(x, b1, c21);
x = _mm256_broadcast_ss(a3 + k);
c30 = _mm256_fmadd_ps(x, b0, c30);
c31 = _mm256_fmadd_ps(x, b1, c31);
}
float* c = C + (size_t)i0 * N + j0;
_mm256_storeu_ps(c, c00);
_mm256_storeu_ps(c + 8, c01);
_mm256_storeu_ps(c + N, c10);
_mm256_storeu_ps(c + N + 8, c11);
_mm256_storeu_ps(c + 2 * N, c20);
_mm256_storeu_ps(c + 2 * N + 8, c21);
_mm256_storeu_ps(c + 3 * N, c30);
_mm256_storeu_ps(c + 3 * N + 8, c31);
}
} else {
for (int i = i0; i < i0 + mr; i++) {
const float* a = A + (size_t)i * K;
for (int j0 = 0; j0 < NB; j0 += 16) {
__m256 c0 = _mm256_setzero_ps(), c1 = _mm256_setzero_ps();
for (int k = 0; k < K; k++) {
const float* b = B + (size_t)k * N + j0;
__m256 x = _mm256_broadcast_ss(a + k);
c0 = _mm256_fmadd_ps(x, _mm256_loadu_ps(b), c0);
c1 = _mm256_fmadd_ps(x, _mm256_loadu_ps(b + 8), c1);
}
_mm256_storeu_ps(C + (size_t)i * N + j0, c0);
_mm256_storeu_ps(C + (size_t)i * N + j0 + 8, c1);
}
}
}
if (NB < N) gemm_nn_rows(A, i0, i0 + mr, K, B, N, C, NB);
}
#elif defined(OASIS_NEON)
const int NB = N / 16 * 16;
const int blocks = (M + 3) / 4;
#pragma omp parallel for schedule(static) if ((long)M * N * K > 200000)
for (int ib = 0; ib < blocks; ib++) {
int i0 = ib * 4, mr = std::min(4, M - i0);
if (mr == 4) {
const float* a0 = A + (size_t)i0 * K;
const float* a1 = a0 + K;
const float* a2 = a1 + K;
const float* a3 = a2 + K;
for (int j0 = 0; j0 < NB; j0 += 16) {
float32x4_t c00 = vdupq_n_f32(0.0f), c01 = c00, c02 = c00, c03 = c00, c10 = c00, c11 = c00, c12 = c00, c13 = c00;
float32x4_t c20 = c00, c21 = c00, c22 = c00, c23 = c00, c30 = c00, c31 = c00, c32 = c00, c33 = c00;
for (int k = 0; k < K; k++) {
const float* b = B + (size_t)k * N + j0;
const float32x4_t b0 = vld1q_f32(b), b1 = vld1q_f32(b + 4), b2 = vld1q_f32(b + 8), b3 = vld1q_f32(b + 12);
float x = a0[k];
c00 = vfmaq_n_f32(c00, b0, x); c01 = vfmaq_n_f32(c01, b1, x); c02 = vfmaq_n_f32(c02, b2, x); c03 = vfmaq_n_f32(c03, b3, x);
x = a1[k];
c10 = vfmaq_n_f32(c10, b0, x); c11 = vfmaq_n_f32(c11, b1, x); c12 = vfmaq_n_f32(c12, b2, x); c13 = vfmaq_n_f32(c13, b3, x);
x = a2[k];
c20 = vfmaq_n_f32(c20, b0, x); c21 = vfmaq_n_f32(c21, b1, x); c22 = vfmaq_n_f32(c22, b2, x); c23 = vfmaq_n_f32(c23, b3, x);
x = a3[k];
c30 = vfmaq_n_f32(c30, b0, x); c31 = vfmaq_n_f32(c31, b1, x); c32 = vfmaq_n_f32(c32, b2, x); c33 = vfmaq_n_f32(c33, b3, x);
}
float* c = C + (size_t)i0 * N + j0;
vst1q_f32(c, c00); vst1q_f32(c + 4, c01); vst1q_f32(c + 8, c02); vst1q_f32(c + 12, c03);
c += N;
vst1q_f32(c, c10); vst1q_f32(c + 4, c11); vst1q_f32(c + 8, c12); vst1q_f32(c + 12, c13);
c += N;
vst1q_f32(c, c20); vst1q_f32(c + 4, c21); vst1q_f32(c + 8, c22); vst1q_f32(c + 12, c23);
c += N;
vst1q_f32(c, c30); vst1q_f32(c + 4, c31); vst1q_f32(c + 8, c32); vst1q_f32(c + 12, c33);
}
} else {
for (int i = i0; i < i0 + mr; i++) {
const float* a = A + (size_t)i * K;
for (int j0 = 0; j0 < NB; j0 += 16) {
float32x4_t c0 = vdupq_n_f32(0.0f), c1 = c0, c2 = c0, c3 = c0;
for (int k = 0; k < K; k++) {
const float* b = B + (size_t)k * N + j0;
const float x = a[k];
c0 = vfmaq_n_f32(c0, vld1q_f32(b), x);
c1 = vfmaq_n_f32(c1, vld1q_f32(b + 4), x);
c2 = vfmaq_n_f32(c2, vld1q_f32(b + 8), x);
c3 = vfmaq_n_f32(c3, vld1q_f32(b + 12), x);
}
float* c = C + (size_t)i * N + j0;
vst1q_f32(c, c0); vst1q_f32(c + 4, c1); vst1q_f32(c + 8, c2); vst1q_f32(c + 12, c3);
}
}
}
if (NB < N) gemm_nn_rows(A, i0, i0 + mr, K, B, N, C, NB);
}
#else
#pragma omp parallel for schedule(static) if ((long)M * N * K > 200000)
for (int i = 0; i < M; i++) gemm_nn_rows(A, i, i + 1, K, B, N, C, 0);
#endif
}
void zc_norm(const float* x, int n, const float* scale, float* out) {
double ss = 0.0;
for (int i = 0; i < n; i++) ss += (double)x[i] * x[i];
float r = 1.0f / std::sqrt((float)(ss / n) + 1e-6f);
for (int i = 0; i < n; i++) out[i] = x[i] * r * (1.0f + scale[i]);
}
void rms_unit(const float* x, int n, float* out) {
double ss = 0.0;
for (int i = 0; i < n; i++) ss += (double)x[i] * x[i];
float r = 1.0f / std::sqrt((float)(ss / n) + 1e-6f);
for (int i = 0; i < n; i++) out[i] = x[i] * r;
}
void band65(const float* u, const float* w, int n, float* out) {
for (int j = 0; j < n; j++) {
const float* wj = w + (size_t)j * 65;
int lo = std::max(0, 32 - j), hi = std::min(65, n + 32 - j);
float s = 0.0f;
for (int t = lo; t < hi; t++) s += wj[t] * u[j + t - 32];
out[j] = s;
}
}
void qmm(Model& m, const std::string& name, const float* A, int M, float* C) {
CQ q = m.Q(name);
if (g_exact_gemv) gemm_nt(A, M, q.cols, m.Wdq(name), q.rows, C);
else cq_gemm_q(q, A, M, C);
}
} // namespace oasis
// ==== Image side: Pillow-exact bilinear resize, encoder, neck, detection, masks, box slots and memory
namespace oasis {
// ------------------------------------------------------------------ resize (matches Pillow 8bpc)
// The resize follows the fixed-point bilinear filter of Pillow, so that the engine reads the same
// pixels as the reference implementation. Pillow carries this notice:
//
// The Python Imaging Library (PIL) is
//
// Copyright © 1997-2011 by Secret Labs AB
// Copyright © 1995-2011 by Fredrik Lundh and contributors
//
// Pillow is the friendly PIL fork. It is
//
// Copyright © 2010 by Jeffrey A. Clark and contributors
//
// Like PIL, Pillow is licensed under the open source MIT-CMU License:
//
// By obtaining, using, and/or copying this software and/or its associated
// documentation, you agree that you have read, understood, and will comply
// with the following terms and conditions:
//
// Permission to use, copy, modify and distribute this software and its
// documentation for any purpose and without fee is hereby granted,
// provided that the above copyright notice appears in all copies, and that
// both that copyright notice and this permission notice appear in supporting
// documentation, and that the name of Secret Labs AB or the author not be
// used in advertising or publicity pertaining to distribution of the software
// without specific, written prior permission.
//
// SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS
// SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.
// IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL,
// INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
// OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
static const int PREC = 22;
static void coeffs(int in_size, int out_size, std::vector<int>& bounds, std::vector<int>& kk, int& ksize) {
double scale = (double)in_size / out_size, fs = scale < 1.0 ? 1.0 : scale;
double support = 1.0 * fs;
ksize = (int)std::ceil(support) * 2 + 1;
bounds.assign(out_size * 2, 0);
kk.assign((size_t)out_size * ksize, 0);
std::vector<double> k(ksize);
for (int xx = 0; xx < out_size; xx++) {
double center = (xx + 0.5) * scale, ww = 0.0, ss = 1.0 / fs;
int xmin = (int)(center - support + 0.5);
if (xmin < 0) xmin = 0;
int xmax = (int)(center + support + 0.5);
if (xmax > in_size) xmax = in_size;
xmax -= xmin;
for (int x = 0; x < xmax; x++) {
double t = (x + xmin - center + 0.5) * ss;
double w = t < 0 ? -t : t;
w = w < 1.0 ? 1.0 - w : 0.0;
k[x] = w;
ww += w;
}
for (int x = 0; x < xmax; x++) {
double v = ww != 0.0 ? k[x] / ww : 0.0;
kk[(size_t)xx * ksize + x] = v < 0 ? (int)(-0.5 + v * (1 << PREC)) : (int)(0.5 + v * (1 << PREC));
}
bounds[xx * 2] = xmin;
bounds[xx * 2 + 1] = xmax;
}
}
static inline uint8_t clip8(long long v) {
if (v >= ((long long)256 << PREC)) return 255;
if (v <= 0) return 0;
return (uint8_t)(v >> PREC);
}
std::vector<uint8_t> resize_rgb(const uint8_t* src, int W0, int H0, int W1, int H1) {
std::vector<uint8_t> tmp((size_t)H0 * W1 * 3), out((size_t)H1 * W1 * 3);
std::vector<int> b, k;
int ks;
if (W1 != W0) {
coeffs(W0, W1, b, k, ks);
#pragma omp parallel for schedule(static) if ((long)H0 * W1 > 65536)
for (int y = 0; y < H0; y++)
for (int x = 0; x < W1; x++)
for (int c = 0; c < 3; c++) {
long long ss = 1LL << (PREC - 1);
for (int t = 0; t < b[x * 2 + 1]; t++) ss += (long long)src[((size_t)y * W0 + b[x * 2] + t) * 3 + c] * k[(size_t)x * ks + t];
tmp[((size_t)y * W1 + x) * 3 + c] = clip8(ss);
}
} else {
std::memcpy(tmp.data(), src, tmp.size());
}
if (H1 != H0) {
coeffs(H0, H1, b, k, ks);
#pragma omp parallel for schedule(static) if ((long)H1 * W1 > 65536)
for (int y = 0; y < H1; y++)
for (int x = 0; x < W1; x++)
for (int c = 0; c < 3; c++) {
long long ss = 1LL << (PREC - 1);
for (int t = 0; t < b[y * 2 + 1]; t++) ss += (long long)tmp[((size_t)(b[y * 2] + t) * W1 + x) * 3 + c] * k[(size_t)y * ks + t];
out[((size_t)y * W1 + x) * 3 + c] = clip8(ss);
}
} else {
out = tmp;
}
return out;
}
// ------------------------------------------------------------------ NHWC building blocks
static std::vector<uint8_t> valid_mask(int n, int vh, int vw) {
std::vector<uint8_t> m((size_t)n * n, 0);
for (int i = 0; i < std::min(n, vh); i++)
for (int j = 0; j < std::min(n, vw); j++) m[(size_t)i * n + j] = 1;
return m;
}
// 3x3 convolution with one pixel of zero padding on every side, HWIO weights.
static void conv3x3(const float* in, int H, int W, int Ci, const float* w, int Co, int sh, int sw, std::vector<float>& out, int& Ho, int& Wo) {
Ho = (H - 1) / sh + 1;
Wo = (W - 1) / sw + 1;
out.assign((size_t)Ho * Wo * Co, 0.0f);
#pragma omp parallel for schedule(static)
for (int y = 0; y < Ho; y++)
for (int x = 0; x < Wo; x++) {
float* o = out.data() + ((size_t)y * Wo + x) * Co;
for (int kh = 0; kh < 3; kh++) {
int iy = y * sh - 1 + kh;
if (iy < 0 || iy >= H) continue;
for (int kw = 0; kw < 3; kw++) {
int ix = x * sw - 1 + kw;
if (ix < 0 || ix >= W) continue;
const float* ip = in + ((size_t)iy * W + ix) * Ci;
const float* wp = w + ((size_t)(kh * 3 + kw) * Ci) * Co;
for (int ci = 0; ci < Ci; ci++) {
float v = ip[ci];
const float* wr = wp + (size_t)ci * Co;
for (int co = 0; co < Co; co++) o[co] += v * wr[co];
}
}
}
}
}
// Depthwise kh x kw convolution, weights [kh, kw, 1, C]; padding kh/2, kw/2.
static void dwconv(const float* in, int H, int W, int C, const float* w, int kh_, int kw_, int sh, int sw, std::vector<float>& out, int& Ho, int& Wo) {
int ph = kh_ / 2, pw_ = kw_ / 2;
Ho = (H + 2 * ph - kh_) / sh + 1;
Wo = (W + 2 * pw_ - kw_) / sw + 1;
out.assign((size_t)Ho * Wo * C, 0.0f);
#pragma omp parallel for schedule(static)
for (int y = 0; y < Ho; y++)
for (int x = 0; x < Wo; x++) {
float* o = out.data() + ((size_t)y * Wo + x) * C;
for (int a = 0; a < kh_; a++) {
int iy = y * sh - ph + a;
if (iy < 0 || iy >= H) continue;
for (int bb = 0; bb < kw_; bb++) {
int ix = x * sw - pw_ + bb;
if (ix < 0 || ix >= W) continue;
const float* ip = in + ((size_t)iy * W + ix) * C;
const float* wp = w + (size_t)(a * kw_ + bb) * C;
for (int c = 0; c < C; c++) o[c] += ip[c] * wp[c];
}
}
}
}
// Fixed 4096-element blocks, so every element is computed the same way whatever the thread count.
static void silu_v(std::vector<float>& v) {
const int n = (int)v.size(), B = 4096, nb = (n + B - 1) / B;
float* p = v.data();
#pragma omp parallel for schedule(static) if (n > 65536)
for (int b = 0; b < nb; b++) {
int e = std::min(n, (b + 1) * B);
#if defined(OASIS_NEON)
const int i0 = b * B, n4 = (e - i0) & ~3;
silu_inplace(p + i0, n4);
for (int i = i0 + n4; i < e; i++) p[i] = siluf(p[i]);
#else
for (int i = b * B; i < e; i++) p[i] = siluf(p[i]);
#endif
}
}
static void norm_rows(std::vector<float>& x, int C, const float* scale) {
size_t n = x.size() / C;
#pragma omp parallel for schedule(static) if (n > 64)
for (long p = 0; p < (long)n; p++) zc_norm(x.data() + p * C, C, scale, x.data() + p * C);
}
static void apply_mask(std::vector<float>& x, int C, const std::vector<uint8_t>& m) {
for (size_t p = 0; p < m.size(); p++)
if (!m[p]) std::memset(x.data() + p * C, 0, sizeof(float) * C);
}
static std::vector<float> pointwise(const std::vector<float>& x, int Ci, const float* k, int Co) {
int P = (int)(x.size() / Ci);
std::vector<float> out((size_t)P * Co);
gemm_nn(x.data(), P, Ci, k, Co, out.data());
return out;
}
static std::vector<float> up2(const std::vector<float>& x, int n, int C) {
std::vector<float> o((size_t)4 * n * n * C);
for (int i = 0; i < 2 * n; i++)
for (int j = 0; j < 2 * n; j++)
std::memcpy(o.data() + ((size_t)i * 2 * n + j) * C, x.data() + ((size_t)(i / 2) * n + j / 2) * C, sizeof(float) * C);
return o;
}
// One HConvGLU's weights, looked up once per block rather than once per cell.
struct HConv {
const float *gate, *gate_bias, *value, *out;
};
static HConv hconv_of(Model& m, const std::string& pre) {
return {m.F(pre + "/gate_conv"), m.F(pre + "/gate_bias"), m.F(pre + "/value_conv"), m.F(pre + "/out_conv")};
}
#if defined(OASIS_NEON)
// The bands of NC cells at once, with tap-major weights w [65, 512] (Model::load stores them so on Arm):
// every tap is one multiply-add across the outputs of all NC cells, so each weight vector is loaded once
// for all of them. up[c] holds a cell's input with 32 zeros on either side.
template <int NC>
static void band_cells(float (*up)[576], const float* w, float (*out)[512]) {
for (int j = 0; j < 512; j += 8) {
float32x4_t a[NC][2];
for (int c = 0; c < NC; c++) a[c][0] = a[c][1] = vdupq_n_f32(0.0f);
for (int t = 0; t < 65; t++) {
const float32x4_t w0 = vld1q_f32(w + (size_t)t * 512 + j), w1 = vld1q_f32(w + (size_t)t * 512 + j + 4);
for (int c = 0; c < NC; c++) {
a[c][0] = vfmaq_f32(a[c][0], w0, vld1q_f32(up[c] + j + t));
a[c][1] = vfmaq_f32(a[c][1], w1, vld1q_f32(up[c] + j + t + 4));
}
}
for (int c = 0; c < NC; c++) {
vst1q_f32(out[c] + j, a[c][0]);
vst1q_f32(out[c] + j + 4, a[c][1]);
}
}
}
static void band_cells(int nc, float (*up)[576], const float* w, float (*out)[512]) {
switch (nc) {
case 4: band_cells<4>(up, w, out); break;
case 3: band_cells<3>(up, w, out); break;
case 2: band_cells<2>(up, w, out); break;
default: band_cells<1>(up, w, out); break;
}
}
// HConvGLU of up to four cells at once: x[c] -> y[c].
static void hconvglu_cells(const Model& m, const HConv& w, const float* const* x, float* const* y, int nc) {
float up[4][576], a[4][512], v[4][512], z[4][512];
for (int c = 0; c < nc; c++) {
std::memset(up[c], 0, sizeof up[c]);
wht_seq512(m, x[c], up[c] + 32);
}
band_cells(nc, up, w.gate, a);
band_cells(nc, up, w.value, v);
// gelu_tanh(a + bias) * v with the operations the scalar loop compiles to, k * fma(0.044715 x, x^2, x) and
// ((1 + tanh) v)(0.5 x), and tanh four at a time.
for (int c = 0; c < nc; c++)
for (int i = 0; i < 512; i += 4) {
const float32x4_t x = vaddq_f32(vld1q_f32(a[c] + i), vld1q_f32(w.gate_bias + i));
const float32x4_t arg = vmulq_n_f32(vfmaq_f32(x, vmulq_n_f32(x, 0.044715f), vmulq_f32(x, x)), 0.7978845608028654f);
const float32x4_t t = vaddq_f32(tanhf4(arg), vdupq_n_f32(1.0f));
vst1q_f32(up[c] + 32 + i, vmulq_f32(vmulq_f32(t, vld1q_f32(v[c] + i)), vmulq_n_f32(x, 0.5f)));
}
band_cells(nc, up, w.out, z);
for (int c = 0; c < nc; c++) wht_seq512(m, z[c], y[c]);
}
#else
static void hconvglu(const Model& m, const HConv& w, const float* x, float* y) {
float u[512], a[512], v[512], h[512], z[512];
wht_seq512(m, x, u);
band65(u, w.gate, 512, a);
band65(u, w.value, 512, v);
for (int i = 0; i < 512; i++) h[i] = gelu_tanh(a[i] + w.gate_bias[i]) * v[i];
band65(h, w.out, 512, z);
wht_seq512(m, z, y);
}
#endif
// residual block B(F) = F + PW(DW3x3(SiLU(N(F)))), masked
static void res_block(Model& m, std::vector<float>& F, int n, int C, const std::string& norm, const std::string& dw,
const std::string& pw, const std::vector<uint8_t>& mask) {
std::vector<float> t = F, d;
norm_rows(t, C, m.F(norm));
silu_v(t);
int ho, wo;
dwconv(t.data(), n, n, C, m.F(dw), 3, 3, 1, 1, d, ho, wo);
std::vector<float> p = pointwise(d, C, m.F(pw), C);
for (size_t i = 0; i < F.size(); i++) F[i] += p[i];
apply_mask(F, C, mask);
}
// Stage checksums on stderr when OASIS_DEBUG=1, for comparing builds.
static void stage(const char* name, const std::vector<float>& v) {
static const bool on = [] { const char* d = std::getenv("OASIS_DEBUG"); return d && d[0] == '1'; }();
if (!on) return;
double s = 0.0, a = 0.0;
for (float x : v) { s += x; a += std::fabs(x); }
std::fprintf(stderr, "%-10s n=%zu sum=%.6g abs=%.6g first=%.6g\n", name, v.size(), s, a, v.empty() ? 0.0 : v[0]);
}
Encoded encode(Model& m, const uint8_t* rgb, int W0, int H0) {
Encoded e;
float r = 512.0f / (float)std::max(W0, H0);
int w = std::max(1, (int)std::nearbyint((double)W0 * r)), h = std::max(1, (int)std::nearbyint((double)H0 * r));
e.meta = {W0, H0, w, h, r};
std::vector<uint8_t> rs = (w == W0 && h == H0) ? std::vector<uint8_t>(rgb, rgb + (size_t)W0 * H0 * 3) : resize_rgb(rgb, W0, H0, w, h);
std::vector<float> canvas((size_t)512 * 512 * 3, 0.0f);
#pragma omp parallel for schedule(static)
for (int y = 0; y < h; y++)
for (int x = 0; x < w; x++)
for (int c = 0; c < 3; c++) canvas[((size_t)y * 512 + x) * 3 + c] = rs[((size_t)y * w + x) * 3 + c] / 127.5f - 1.0f;
auto vm = [&](int s) { return valid_mask(512 / s, (h + s - 1) / s, (w + s - 1) / s); };
auto m2 = vm(2), m4 = vm(4);
e.valid8 = vm(8);
e.valid16 = vm(16);
e.valid32 = vm(32);
e.valid4 = m4;
stage("canvas", canvas);
std::vector<float> f, t, d;
int Ho, Wo;
conv3x3(canvas.data(), 512, 512, 3, m.F("stem/w0"), 32, 2, 2, f, Ho, Wo);
norm_rows(f, 32, m.F("stem/norm_0/scale"));
apply_mask(f, 32, m2);
int C = 32, n = 256;
const int chans[5] = {32, 64, 128, 256, 512};
std::vector<float> feats[6];
for (int j = 1; j <= 4; j++) {
t = f;
silu_v(t);
dwconv(t.data(), n, n, C, m.F("stem/dw_" + std::to_string(j)), 3, 3, 2, 2, d, Ho, Wo);
f = pointwise(d, C, m.F("stem/pw_" + std::to_string(j) + "/kernel"), chans[j]);
C = chans[j];
n = Ho;
norm_rows(f, C, m.F("stem/norm_" + std::to_string(j) + "/scale"));
apply_mask(f, C, j == 1 ? m4 : j == 2 ? e.valid8 : j == 3 ? e.valid16 : e.valid32);
feats[j] = f;
}
for (int j = 1; j <= 4; j++) stage(("stem" + std::to_string(j)).c_str(), feats[j]);
e.S4 = feats[1];
silu_v(e.S4);
std::vector<float> S8 = feats[2], S16 = feats[3], X = feats[4];
silu_v(S8);
silu_v(S16);
silu_v(X);
// ---- convolutional body, 16 x 16 x 512
const auto& v32 = e.valid32;
for (int i = 0; i < 16; i++) {
std::string pre = "body_" + std::to_string(i);
std::vector<float> Z = X, va, ga, mix;
norm_rows(Z, 512, m.F(pre + "/mix_norm/scale"));
dwconv(Z.data(), 16, 16, 512, m.F(pre + "/mix/dw_value"), 3, 3, 1, 1, va, Ho, Wo);
dwconv(Z.data(), 16, 16, 512, m.F(pre + "/mix/dw_gate"), 3, 3, 1, 1, ga, Ho, Wo);
const float* bv = m.F(pre + "/mix/b_value");
const float* bg = m.F(pre + "/mix/b_gate");
#pragma omp parallel for schedule(static)
for (int p = 0; p < 256; p++)
for (int c = 0; c < 512; c++) {
size_t k = (size_t)p * 512 + c;
va[k] = (va[k] + bv[c]) * siluf(ga[k] + bg[c]);
}
dwconv(va.data(), 16, 16, 512, m.F(pre + "/mix/dw_down"), 3, 3, 1, 1, mix, Ho, Wo);
for (int p = 0; p < 256; p++)
if (v32[p])
for (int c = 0; c < 512; c++) X[(size_t)p * 512 + c] += mix[(size_t)p * 512 + c];
std::vector<float> Hh((size_t)256 * 512);
const float* hs = m.F(pre + "/hconv_norm/scale");
const HConv hw = hconv_of(m, pre + "/hconv");
#if defined(OASIS_NEON)
#pragma omp parallel for schedule(static)
for (int b = 0; b < 64; b++) {
float xn[4][512];
const float* xs[4];
float* ys[4];
for (int c = 0; c < 4; c++) {
const int p = b * 4 + c;
zc_norm(X.data() + (size_t)p * 512, 512, hs, xn[c]);
xs[c] = xn[c];
ys[c] = Hh.data() + (size_t)p * 512;
}
hconvglu_cells(m, hw, xs, ys, 4);
}
#else
#pragma omp parallel for schedule(static)
for (int p = 0; p < 256; p++) {
float xn[512];
zc_norm(X.data() + (size_t)p * 512, 512, hs, xn);
hconvglu(m, hw, xn, Hh.data() + (size_t)p * 512);
}
#endif
const float* gam = m.F(pre + "/grn/gamma");
const float* bet = m.F(pre + "/grn/beta");
std::vector<float> rc(512);
#pragma omp parallel for schedule(static)
for (int c = 0; c < 512; c++) {
double s = 0.0;
for (int p = 0; p < 256; p++)
if (v32[p]) s += (double)Hh[(size_t)p * 512 + c] * Hh[(size_t)p * 512 + c];
rc[c] = (float)std::sqrt(s + 1e-6);
}
double rbar = 0.0;
for (int c = 0; c < 512; c++) rbar += rc[c];
float rb = (float)(rbar / 512.0);
#pragma omp parallel for schedule(static)
for (int p = 0; p < 256; p++)
if (v32[p])
for (int c = 0; c < 512; c++) {
float hv = Hh[(size_t)p * 512 + c];
X[(size_t)p * 512 + c] += hv + gam[c] * hv * rc[c] / (rb + 1e-6f) + bet[c];
}
}
stage("body", X);
// ---- attention blocks over 256 cells + 4 registers
const int T = 260;
std::vector<float> Z((size_t)T * 512);
std::memcpy(Z.data(), X.data(), sizeof(float) * 256 * 512);
std::memcpy(Z.data() + 256 * 512, m.F("registers"), sizeof(float) * 4 * 512);
std::vector<uint8_t> tv(T, 1);
for (int p = 0; p < 256; p++) tv[p] = v32[p];
#if defined(OASIS_NEON)
std::vector<int> active; // the valid tokens, in order, for the four-at-a-time HConvGLU
for (int p = 0; p < T; p++)
if (tv[p]) active.push_back(p);
#endif
std::vector<float> E((size_t)256 * 512, 0.0f);
auto add_level = [&](const std::vector<float>& Zl, int l) {
const float* sc = m.F("level_norm_" + std::to_string(l) + "/scale");
const float* g = m.F("level_gates") + (size_t)l * 512;
#pragma omp parallel for schedule(static)
for (int p = 0; p < 256; p++) {
float tmp[512];
zc_norm(Zl.data() + (size_t)p * 512, 512, sc, tmp);
for (int c = 0; c < 512; c++) E[(size_t)p * 512 + c] += g[c] * tmp[c];
}
};
float cosr[256][16], sinr[256][16];
for (int p = 0; p < 256; p++)
for (int r2 = 0; r2 < 8; r2++) {
double fr = 1.0 / std::pow(10000.0, (2.0 * r2) / 16.0);
double ai = (p / 16) * fr, aj = (p % 16) * fr;
cosr[p][r2] = (float)std::cos(ai);
sinr[p][r2] = (float)std::sin(ai);
cosr[p][8 + r2] = (float)std::cos(aj);
sinr[p][8 + r2] = (float)std::sin(aj);
}
std::vector<float> z((size_t)T * 512), q((size_t)T * 128), k((size_t)T * 32), v((size_t)T * 64), gt((size_t)T * 256), o((size_t)T * 256), A((size_t)T * 512);
for (int l = 0; l < 12; l++) {
add_level(Z, l);
std::string pre = "top_" + std::to_string(l) + "/";
const float* an = m.F(pre + "attn_norm/scale");
for (int p = 0; p < T; p++) zc_norm(Z.data() + (size_t)p * 512, 512, an, z.data() + (size_t)p * 512);
qmm(m, pre + "self_attn/q_proj/kernel", z.data(), T, q.data());
qmm(m, pre + "self_attn/k_proj/kernel", z.data(), T, k.data());
qmm(m, pre + "self_attn/v_proj/kernel", z.data(), T, v.data());
qmm(m, pre + "self_attn/gate_proj/kernel", z.data(), T, gt.data());
const float* qn = m.F(pre + "self_attn/q_norm/scale");
const float* kn = m.F(pre + "self_attn/k_norm/scale");
auto rope = [&](float* x, int p) {
if (p >= 256) return;
float x1[16], x2[16];
std::memcpy(x1, x, sizeof x1);
std::memcpy(x2, x + 16, sizeof x2);
for (int i = 0; i < 16; i++) {
x[i] = x1[i] * cosr[p][i] - x2[i] * sinr[p][i];
x[16 + i] = x2[i] * cosr[p][i] + x1[i] * sinr[p][i];
}
};
for (int p = 0; p < T; p++) {
for (int hh = 0; hh < 4; hh++) {
float* qp = q.data() + (size_t)p * 128 + hh * 32;
zc_norm(qp, 32, qn, qp);
rope(qp, p);
}
float* kp = k.data() + (size_t)p * 32;
zc_norm(kp, 32, kn, kp);
rope(kp, p);
}
#pragma omp parallel for schedule(static)
for (int p = 0; p < T; p++) {
#if defined(OASIS_NEON)
std::vector<float> s(T), e(T);
#else
std::vector<float> s(T);
#endif
for (int hh = 0; hh < 4; hh++) {
const float* qp = q.data() + (size_t)p * 128 + hh * 32;
float mx = -1e30f;
for (int j = 0; j < T; j++) {
if (!tv[j]) { s[j] = -1e30f; continue; }
float d = 0.0f;
for (int c = 0; c < 32; c++) d += qp[c] * k[(size_t)j * 32 + c];
s[j] = d / std::sqrt(32.0f);
mx = std::max(mx, s[j]);
}
double den = 0.0;
#if defined(OASIS_NEON)
for (int j = 0; j < T; j += 4) vst1q_f32(e.data() + j, expf4(vsubq_f32(vld1q_f32(s.data() + j), vdupq_n_f32(mx))));
for (int j = 0; j < T; j++) {
s[j] = tv[j] ? e[j] : 0.0f;
den += s[j];
}
#else
for (int j = 0; j < T; j++) {
s[j] = tv[j] ? std::exp(s[j] - mx) : 0.0f;
den += s[j];
}
#endif
float* op = o.data() + (size_t)p * 256 + hh * 64;
for (int c = 0; c < 64; c++) op[c] = 0.0f;
for (int j = 0; j < T; j++) {
if (s[j] == 0.0f) continue;
float wv = (float)(s[j] / den);
for (int c = 0; c < 64; c++) op[c] += wv * v[(size_t)j * 64 + c];
}
#if defined(OASIS_NEON)
mul_sigmoid(op, gt.data() + (size_t)p * 256 + hh * 64, 64);
#else
for (int c = 0; c < 64; c++) op[c] *= sigmoidf(gt[(size_t)p * 256 + hh * 64 + c]);
#endif
}
}
qmm(m, pre + "self_attn/out_proj/kernel", o.data(), T, A.data());
float ag = sigmoidf(m.F(pre + "attn_gate")[0]);
const float* pn = m.F(pre + "post_attn_norm/scale");
const float* hn = m.F(pre + "hconv_norm/scale");
const HConv hw = hconv_of(m, pre + "hconv");
#if defined(OASIS_NEON)
const int nb = ((int)active.size() + 3) / 4;
#pragma omp parallel for schedule(static)
for (int b = 0; b < nb; b++) {
const int nc = std::min(4, (int)active.size() - b * 4);
float tmp[4][512], hy[4][512];
const float* xs[4];
float* ys[4];
for (int c = 0; c < nc; c++) {
float* zp = Z.data() + (size_t)active[b * 4 + c] * 512;
float t1[512];
zc_norm(A.data() + (size_t)active[b * 4 + c] * 512, 512, pn, t1);
for (int k = 0; k < 512; k++) zp[k] += ag * t1[k];
zc_norm(zp, 512, hn, tmp[c]);
xs[c] = tmp[c];
ys[c] = hy[c];
}
hconvglu_cells(m, hw, xs, ys, nc);
for (int c = 0; c < nc; c++) {
float* zp = Z.data() + (size_t)active[b * 4 + c] * 512;
for (int k = 0; k < 512; k++) zp[k] += hy[c][k];
}
}
#else
#pragma omp parallel for schedule(static)
for (int p = 0; p < T; p++) {
if (!tv[p]) continue;
float tmp[512], hy[512];
float* zp = Z.data() + (size_t)p * 512;
zc_norm(A.data() + (size_t)p * 512, 512, pn, tmp);
for (int c = 0; c < 512; c++) zp[c] += ag * tmp[c];
zc_norm(zp, 512, hn, tmp);
hconvglu(m, hw, tmp, hy);
for (int c = 0; c < 512; c++) zp[c] += hy[c];
}
#endif
}
{
const float* en = m.F("enc_norm/scale");
float tmp[512];
for (int p = 0; p < 256; p++) {
zc_norm(Z.data() + (size_t)p * 512, 512, en, tmp);
for (int c = 0; c < 512; c++) E[(size_t)p * 512 + c] += tmp[c];
}
apply_mask(E, 512, v32);
}
stage("E", E);
e.E = E;
// ---- neck to D [64, 64, 128]
std::vector<float> co = pointwise(E, 512, m.F("dense/coarse/kernel"), 128);
std::vector<float> tt = up2(co, 16, 128), mid = pointwise(S16, 256, m.F("dense/mid/kernel"), 128);
for (size_t i = 0; i < tt.size(); i++) tt[i] += mid[i];
apply_mask(tt, 128, e.valid16);
res_block(m, tt, 32, 128, "dense/neck_norm/scale", "dense/neck_dw", "dense/neck_pw/kernel", e.valid16);
std::vector<float> dd = up2(tt, 32, 128), fine = pointwise(S8, 128, m.F("dense/fine/kernel"), 128);
for (size_t i = 0; i < dd.size(); i++) dd[i] += fine[i];
apply_mask(dd, 128, e.valid8);
for (int i = 0; i < 3; i++)
res_block(m, dd, 64, 128, "dense/block_norm_" + std::to_string(i) + "/scale", "dense/dw_" + std::to_string(i),
"dense/pw_" + std::to_string(i) + "/kernel", e.valid8);
norm_rows(dd, 128, m.F("dense/final_norm/scale"));
silu_v(dd);
apply_mask(dd, 128, e.valid8);
stage("D", dd);
e.D = dd;
e.det = pointwise(dd, 128, m.F("dense/out/kernel"), 10);
const float* db = m.F("dense/bias");
for (int p = 0; p < 64 * 64; p++)
for (int c = 0; c < 10; c++) e.det[(size_t)p * 10 + c] += db[c];
stage("det", e.det);
return e;
}
// ------------------------------------------------------------------ detection
static float iou_box(const Box& a, const Box& b) {
float ix = std::max(0.0f, std::min(a.x1, b.x1) - std::max(a.x0, b.x0));
float iy = std::max(0.0f, std::min(a.y1, b.y1) - std::max(a.y0, b.y0));
float inter = ix * iy;
float u = (a.x1 - a.x0) * (a.y1 - a.y0) + (b.x1 - b.x0) * (b.y1 - b.y0) - inter;
return u > 0 ? inter / u : 0.0f;
}
void detect(const Encoded& enc, float conf, float text_conf, bool nms, int max_obj, int max_txt,
std::vector<Box>& objects, std::vector<Box>& texts) {
for (int ch = 0; ch < 2; ch++) {
float thr = ch == 0 ? conf : text_conf;
int limit = ch == 0 ? max_obj : max_txt;
std::vector<float> heat(64 * 64);
for (int p = 0; p < 64 * 64; p++) heat[p] = sigmoidf(enc.det[(size_t)p * 10 + ch]);
struct Cand { float s; int i, j; };
std::vector<Cand> c;
for (int i = 0; i < 64; i++)
for (int j = 0; j < 64; j++) {
int p = i * 64 + j;
if (!enc.valid8[p] || heat[p] < thr) continue;
bool peak = true;
for (int di = -1; di <= 1 && peak; di++)
for (int dj = -1; dj <= 1; dj++) {
int a = i + di, b = j + dj;
if (a < 0 || b < 0 || a >= 64 || b >= 64) continue;
if (heat[a * 64 + b] > heat[p]) { peak = false; break; }
}
if (peak) c.push_back({heat[p], i, j});
}
std::stable_sort(c.begin(), c.end(), [](const Cand& a, const Cand& b) { return a.s > b.s; });
std::vector<Box>& out = ch == 0 ? objects : texts;
float vw = enc.meta.w / 512.0f, vh = enc.meta.h / 512.0f;
for (auto& cd : c) {
const float* d = enc.det.data() + ((size_t)cd.i * 64 + cd.j) * 10;
float w = std::fabs(d[2 + 2 * ch]), h = std::fabs(d[3 + 2 * ch]);
float cx = cd.j + d[6 + 2 * ch], cy = cd.i + d[7 + 2 * ch];
Box b{(cx - w / 2) / 64, (cy - h / 2) / 64, (cx + w / 2) / 64, (cy + h / 2) / 64, cd.s, cd.i, cd.j};
b.x0 = std::clamp(b.x0, 0.0f, vw);
b.x1 = std::clamp(b.x1, 0.0f, vw);
b.y0 = std::clamp(b.y0, 0.0f, vh);
b.y1 = std::clamp(b.y1, 0.0f, vh);
bool dup = false;
if (nms)
for (auto& x : out)
if (iou_box(b, x) > 0.5f) { dup = true; break; }
if (dup) continue;
out.push_back(b);
if ((int)out.size() >= limit) break;
}
}
}
// ------------------------------------------------------------------ masks
static std::vector<float> head_branch(Model& m, std::vector<float> F, int n, const std::string& pre,
const std::vector<uint8_t>& mask, const std::string& outk) {
for (int i = 0; i < 2; i++)
res_block(m, F, n, 128, "dense/" + pre + "_norm_" + std::to_string(i) + "/scale", "dense/" + pre + "_dw_" + std::to_string(i),
"dense/" + pre + "_pw_" + std::to_string(i) + "/kernel", mask);
norm_rows(F, 128, m.F("dense/" + pre + "_final_norm/scale"));
silu_v(F);
return pointwise(F, 128, m.F(outk), 32);
}
// Mask coefficients at one centre cell. The coefficient branch's only spatial operations are two 3x3
// depthwise convolutions, so its output at a cell depends on the 5 x 5 neighbourhood of D alone; cells
// outside the grid enter as masked zeros, as the full branch's padding does.
static std::vector<float> coef_at(Model& m, const Encoded& enc, int ci, int cj) {
std::vector<float> F((size_t)25 * 128, 0.0f);
std::vector<uint8_t> pm(25, 0);
for (int di = -2; di <= 2; di++)
for (int dj = -2; dj <= 2; dj++) {
int i = ci + di, j = cj + dj, p = (di + 2) * 5 + (dj + 2);
if (i < 0 || j < 0 || i >= 64 || j >= 64) continue;
pm[p] = enc.valid8[(size_t)i * 64 + j];
std::memcpy(F.data() + (size_t)p * 128, enc.D.data() + ((size_t)i * 64 + j) * 128, sizeof(float) * 128);
}
std::vector<float> C = head_branch(m, F, 5, "coef", pm, "dense/coef/kernel");
return std::vector<float>(C.begin() + 12 * 32, C.begin() + 13 * 32);
}
std::vector<float> mask_prob(Model& m, const Encoded& enc, const Box& b, std::vector<float>* protos) {
if (protos->empty()) {
std::vector<float> U = up2(enc.D, 64, 128), ed = pointwise(enc.S4, 64, m.F("dense/edge/kernel"), 128);
for (size_t i = 0; i < U.size(); i++) U[i] += ed[i];
apply_mask(U, 128, enc.valid4);
*protos = head_branch(m, U, 128, "proto", enc.valid4, "dense/protos/kernel");
}
const float* P = protos->data();
std::vector<float> coef = coef_at(m, enc, b.ci, b.cj);
const float* a = coef.data();
std::vector<float> out(128 * 128, 0.0f);
int ci = std::min(127, (int)((b.y0 + b.y1) / 2 * 128)), cj = std::min(127, (int)((b.x0 + b.x1) / 2 * 128));
for (int i = 0; i < 128; i++) {
float yc = (i + 0.5f) / 128;
for (int j = 0; j < 128; j++) {
float xc = (j + 0.5f) / 128;
bool inside = (yc >= b.y0 && yc <= b.y1 && xc >= b.x0 && xc <= b.x1) || (i == ci && j == cj);
if (!inside) continue;
const float* pp = P + ((size_t)i * 128 + j) * 32;
float s = 0.0f;
for (int r = 0; r < 32; r++) s += a[r] * pp[r];
out[(size_t)i * 128 + j] = sigmoidf(s);
}
}
return out;
}
// ------------------------------------------------------------------ box slots and decoder memory
void memory(Model& m, const Encoded& enc, const std::vector<Box>& objects, const std::vector<Box>& texts,
std::vector<float>& M, std::vector<uint8_t>& mvalid, std::vector<int>& refs) {
float pe = sigmoidf(m.F("pe_gate")[0]);
M.assign((size_t)256 * 512, 0.0f);
for (int p = 0; p < 256; p++) {
int ri = p / 16, cj = p % 16;
for (int half = 0; half < 2; half++) {
int pos = half == 0 ? ri : cj;
for (int i = 0; i < 128; i++) {
double om = 1.0 / std::pow(10000.0, (double)i / 128.0);
M[(size_t)p * 512 + half * 256 + i] = (float)std::sin(pos * om);
M[(size_t)p * 512 + half * 256 + 128 + i] = (float)std::cos(pos * om);
}
}
for (int c = 0; c < 512; c++) M[(size_t)p * 512 + c] = enc.E[(size_t)p * 512 + c] + pe * M[(size_t)p * 512 + c];
}
mvalid.assign(enc.valid32.begin(), enc.valid32.end());
refs.clear();
const float* sp = m.F("slot_proj/kernel");
const float* sn = m.F("slot_norm/scale");
const float* emb_scale_dummy = nullptr;
(void)emb_scale_dummy;
CQ emb = m.Q("embedding/embedding");
std::vector<float> row(512);
for (int kind = 0; kind < 2; kind++) {
const auto& boxes = kind == 0 ? objects : texts;
for (size_t k = 0; k < boxes.size(); k++) {
const Box& b = boxes[k];
std::vector<float> feat(192);
std::memcpy(feat.data(), enc.D.data() + ((size_t)b.ci * 64 + b.cj) * 128, sizeof(float) * 128);
float bc[4] = {b.x0, b.y0, b.x1, b.y1};
for (int c = 0; c < 4; c++)
for (int r2 = 0; r2 < 8; r2++) {
double a = std::pow(2.0, r2) * 3.14159265358979323846 * bc[c];
feat[128 + c * 16 + r2] = (float)std::sin(a);
feat[128 + c * 16 + 8 + r2] = (float)std::cos(a);
}
std::vector<float> pooled(512, 0.0f);
int count = 0, ci = std::min(15, (int)((b.y0 + b.y1) / 2 * 16)), cc = std::min(15, (int)((b.x0 + b.x1) / 2 * 16));
std::vector<uint8_t> in(256, 0);
for (int i = 0; i < 16; i++)
for (int j = 0; j < 16; j++) {
float yc = (i + 0.5f) / 16, xc = (j + 0.5f) / 16;
in[i * 16 + j] = yc >= b.y0 && yc <= b.y1 && xc >= b.x0 && xc <= b.x1;
}
in[ci * 16 + cc] = 1;
for (int p = 0; p < 256; p++) in[p] = in[p] && enc.valid32[p];
bool any = false;
for (int p = 0; p < 256; p++) any = any || in[p];
if (!any) in[ci * 16 + cc] = 1;
for (int p = 0; p < 256; p++)
if (in[p]) {
count++;
for (int c = 0; c < 512; c++) pooled[c] += enc.E[(size_t)p * 512 + c];
}
for (int c = 0; c < 512; c++) pooled[c] /= count;
std::vector<float> z(512, 0.0f), zp(512);
gemm_nn(feat.data(), 1, 192, sp, 512, z.data());
qmm(m, "slot_pool/kernel", pooled.data(), 1, zp.data());
for (int c = 0; c < 512; c++) z[c] += zp[c];
std::vector<float> s(512);
zc_norm(z.data(), 512, sn, s.data());
int rid = m.ref_base + (kind == 0 ? (int)k : 64 + (int)k);
cq_dequant_row(emb, rid, row.data());
for (int c = 0; c < 512; c++) s[c] += std::sqrt(512.0f) * row[c];
M.insert(M.end(), s.begin(), s.end());
mvalid.push_back(1);
refs.push_back(rid);
}
}
}
} // namespace oasis
// ==== Text side: CTC reader, decoder with caches, detokenisation
namespace oasis {
// ------------------------------------------------------------------ reader
static void conv3x3_s(const float* in, int H, int W, int Ci, const float* w, int Co, int sh, int sw, std::vector<float>& out, int& Ho, int& Wo) {
Ho = (H - 1) / sh + 1;
Wo = (W - 1) / sw + 1;
out.assign((size_t)Ho * Wo * Co, 0.0f);
for (int y = 0; y < Ho; y++)
for (int x = 0; x < Wo; x++) {
float* o = out.data() + ((size_t)y * Wo + x) * Co;
for (int kh = 0; kh < 3; kh++) {
int iy = y * sh - 1 + kh;
if (iy < 0 || iy >= H) continue;
for (int kw = 0; kw < 3; kw++) {
int ix = x * sw - 1 + kw;
if (ix < 0 || ix >= W) continue;
const float* ip = in + ((size_t)iy * W + ix) * Ci;
const float* wp = w + ((size_t)(kh * 3 + kw) * Ci) * Co;
for (int ci = 0; ci < Ci; ci++)
for (int co = 0; co < Co; co++) o[co] += ip[ci] * wp[(size_t)ci * Co + co];
}
}
}
}
static void dw_s(const float* in, int H, int W, int C, const float* w, int kh_, int kw_, int sh, int sw, std::vector<float>& out, int& Ho, int& Wo) {
int ph = kh_ / 2, pw_ = kw_ / 2;
Ho = (H + 2 * ph - kh_) / sh + 1;
Wo = (W + 2 * pw_ - kw_) / sw + 1;
out.assign((size_t)Ho * Wo * C, 0.0f);
for (int y = 0; y < Ho; y++)
for (int x = 0; x < Wo; x++) {
float* o = out.data() + ((size_t)y * Wo + x) * C;
for (int a = 0; a < kh_; a++) {
int iy = y * sh - ph + a;
if (iy < 0 || iy >= H) continue;
for (int b = 0; b < kw_; b++) {
int ix = x * sw - pw_ + b;
if (ix < 0 || ix >= W) continue;
const float* ip = in + ((size_t)iy * W + ix) * C;
const float* wp = w + (size_t)(a * kw_ + b) * C;
for (int c = 0; c < C; c++) o[c] += ip[c] * wp[c];
}
}
}
}
static void norm_all(std::vector<float>& x, int C, const float* s) {
for (size_t p = 0; p < x.size() / C; p++) zc_norm(x.data() + p * C, C, s, x.data() + p * C);
}
// One 48 x 320 window -> argmax label per frame (80 frames).
static std::vector<int> reader_window(Model& m, const std::vector<float>& J) {
std::vector<float> G, t, d;
int H = 48, W = 320, Ho, Wo;
conv3x3_s(J.data(), H, W, 3, m.F("reader/w0"), 32, 2, 2, G, Ho, Wo);
H = Ho;
W = Wo;
norm_all(G, 32, m.F("reader/stage_norm_0/scale"));
const int Cs[4] = {32, 64, 128, 384};
const int sh[3] = {2, 2, 2}, sw[3] = {2, 1, 1};
for (int j = 0; j < 3; j++) {
t = G;
for (auto& v : t) v = siluf(v);
dw_s(t.data(), H, W, Cs[j], m.F("reader/dw_" + std::to_string(j)), 3, 3, sh[j], sw[j], d, Ho, Wo);
H = Ho;
W = Wo;
G.assign((size_t)H * W * Cs[j + 1], 0.0f);
qmm(m, "reader/pw_" + std::to_string(j) + "/kernel", d.data(), H * W, G.data());
norm_all(G, Cs[j + 1], m.F("reader/stage_norm_" + std::to_string(j + 1) + "/scale"));
}
// H == 3, W == 80, C == 384
std::vector<float> Z((size_t)W * 384, 0.0f);
for (int y = 0; y < H; y++)
for (int x = 0; x < W; x++)
for (int c = 0; c < 384; c++) Z[(size_t)x * 384 + c] += siluf(G[((size_t)y * W + x) * 384 + c]) / H;
std::vector<float> h, mixed, vg((size_t)W * 768), gated((size_t)W * 384), dn((size_t)W * 384);
for (int b = 0; b < 4; b++) {
h = Z;
norm_all(h, 384, m.F("reader/norm_" + std::to_string(b) + "/scale"));
int ho, wo;
dw_s(h.data(), 1, W, 384, m.F("reader/mix_" + std::to_string(b)), 1, 5, 1, 1, mixed, ho, wo);
qmm(m, "reader/up_" + std::to_string(b) + "/kernel", mixed.data(), W, vg.data());
for (int x = 0; x < W; x++)
for (int c = 0; c < 384; c++) gated[(size_t)x * 384 + c] = vg[(size_t)x * 768 + c] * siluf(vg[(size_t)x * 768 + 384 + c]);
qmm(m, "reader/down_" + std::to_string(b) + "/kernel", gated.data(), W, dn.data());
for (size_t i = 0; i < Z.size(); i++) Z[i] += dn[i];
}
norm_all(Z, 384, m.F("reader/final_norm/scale"));
std::vector<float> lg((size_t)W * 307);
qmm(m, "reader/out/kernel", Z.data(), W, lg.data());
const float* bias = m.F("reader/out/bias");
std::vector<int> lab(W);
for (int x = 0; x < W; x++) {
int best = 0;
float bv = -1e30f;
for (int c = 0; c < 307; c++) {
float v = lg[(size_t)x * 307 + c] + bias[c];
if (v > bv) { bv = v; best = c; }
}
lab[x] = best;
}
return lab;
}
static void utf8_append(std::string& s, uint32_t cp) {
if (cp < 0x80) s += (char)cp;
else if (cp < 0x800) { s += (char)(0xC0 | (cp >> 6)); s += (char)(0x80 | (cp & 0x3F)); }
else if (cp < 0x10000) { s += (char)(0xE0 | (cp >> 12)); s += (char)(0x80 | ((cp >> 6) & 0x3F)); s += (char)(0x80 | (cp & 0x3F)); }
else { s += (char)(0xF0 | (cp >> 18)); s += (char)(0x80 | ((cp >> 12) & 0x3F)); s += (char)(0x80 | ((cp >> 6) & 0x3F)); s += (char)(0x80 | (cp & 0x3F)); }
}
std::string read_text(Model& m, const uint8_t* rgb, int W, int H, float x0, float y0, float x1, float y1) {
int cx0 = (int)std::floor(x0), cy0 = (int)std::floor(y0), cx1 = (int)std::ceil(x1), cy1 = (int)std::ceil(y1);
int cw = cx1 - cx0, ch = cy1 - cy0;
if (cw < 1 || ch < 1) return "";
std::vector<uint8_t> crop((size_t)cw * ch * 3, 0);
for (int y = 0; y < ch; y++)
for (int x = 0; x < cw; x++) {
int sy = cy0 + y, sx = cx0 + x;
if (sy < 0 || sx < 0 || sy >= H || sx >= W) continue;
std::memcpy(&crop[((size_t)y * cw + x) * 3], &rgb[((size_t)sy * W + sx) * 3], 3);
}
int U = std::max(1, (int)std::nearbyint(48.0 * cw / ch));
std::vector<uint8_t> strip = resize_rgb(crop.data(), cw, ch, U, 48);
std::vector<int> starts;
if (U <= 320) starts.push_back(0);
else for (int s = 0; s < U - 320 + 223; s += 224) starts.push_back(s);
int n_frames = (U + 3) / 4;
std::string out;
int prev = 0;
for (int xs : starts) {
std::vector<float> J((size_t)48 * 320 * 3, 127.0f / 127.5f - 1.0f);
for (int y = 0; y < 48; y++)
for (int x = 0; x < 320 && xs + x < U; x++)
for (int c = 0; c < 3; c++) J[((size_t)y * 320 + x) * 3 + c] = strip[((size_t)y * U + xs + x) * 3 + c] / 127.5f - 1.0f;
std::vector<int> lab = reader_window(m, J);
int lo = xs == 0 ? 0 : 12;
int hi = xs + 320 >= U ? (U - xs + 3) / 4 : 68;
for (int t = lo; t < std::min(hi, 80); t++) {
if (xs / 4 + t >= n_frames) break;
int l = lab[t];
if (l != 0 && l != prev) utf8_append(out, m.charset[l - 1]);
prev = l;
}
}
return out;
}
// ------------------------------------------------------------------ decoder
static const int ORDERS[2] = {2, 3};
static const uint32_t SEED = 0x9E3779B9u, PRIME = 0x01000193u;
static uint32_t engram_index(const std::vector<int>& toks, int t, int oi, int order, int h) {
uint32_t acc = SEED * (uint32_t)(oi * 2 + h + 1);
for (int j = 0; j < order; j++) {
uint32_t tok = t - j >= 0 ? (uint32_t)toks[t - j] : 0u;
acc = (acc ^ tok) * PRIME;
}
acc ^= acc >> 15;
return acc % 18432u;
}
static inline float fexp(float x) {
x = x < -87.0f ? -87.0f : x;
float t = x * 1.4426950408889634f, fi = std::floor(t), f = t - fi;
float p = 1.0f + f * (0.69314718f + f * (0.24022651f + f * (0.05550411f + f * (0.00961813f + f * (0.00133336f + f * 0.00015404f)))));
int32_t bits = ((int32_t)fi + 127) << 23;
float sc;
std::memcpy(&sc, &bits, 4);
return p * sc;
}
// 4 x 4 Sinkhorn: twenty rounds of row then column normalisation, multiplicative form.
static void sinkhorn(float* L) {
double K[16];
float mx = L[0];
for (int i = 1; i < 16; i++) mx = std::max(mx, L[i]);
for (int i = 0; i < 16; i++) K[i] = std::exp((double)L[i] - mx);
for (int it = 0; it < 20; it++) {
for (int i = 0; i < 4; i++) {
double s = K[i * 4] + K[i * 4 + 1] + K[i * 4 + 2] + K[i * 4 + 3];
for (int j = 0; j < 4; j++) K[i * 4 + j] /= s;
}
for (int j = 0; j < 4; j++) {
double s = K[j] + K[4 + j] + K[8 + j] + K[12 + j];
for (int i = 0; i < 4; i++) K[i * 4 + j] /= s;
}
}
for (int i = 0; i < 16; i++) L[i] = (float)K[i];
}
#if defined(OASIS_NEON)
static inline uint16_t to_half(float f) {
__fp16 h = (__fp16)f;
uint16_t b;
std::memcpy(&b, &h, 2);
return b;
}
#endif
struct Decoder {
Model& m;
int nM;
std::vector<uint8_t> mvalid;
// Image-memory keys and values, shared by copies. On Arm, whose caches are small, they are kept in fp16,
// which halves what each token reads, and laid out in the order a head reads them: keys per head as
// [nMp / 16][48][16] (memory entries padded to nMp), values per head as [nM][64].
struct Cross {
std::vector<float> k[8], v[8];
std::vector<uint16_t> kt[8], v16[8];
};
std::shared_ptr<Cross> cross = std::make_shared<Cross>();
int nMp = 0;
std::vector<float>* ck = cross->k;
std::vector<float>* cv = cross->v;
std::vector<int> toks;
// Self-attention keys and values of every position so far: [t][96] and [t][128], both groups interleaved;
// on Arm one array per group, [t][48] and [t][64], so that a head reads its group's history in one run.
#if defined(OASIS_NEON)
std::vector<float> K[8][2], V[8][2];
std::vector<float> raw[8][3];
#else
std::vector<float> K[8], V[8], raw[8][3];
#endif
std::vector<float> ev[2];
CQ emb;
struct Layer {
CQ q, k, v, gate, out, cq_, cgate, cout;
const float *taps[3], *qn, *kn, *pre, *post, *cross, *pcross, *cqn, *hada;
float ag, cg;
const float *d1, *d2, *d3, *d4, *b2, *cond_v, *cond_u, *w1a, *w1b, *w2a, *w2b, *w3a, *w3b;
} Ls[8];
struct Site { CQ tab, key, val; const float* taps; } sites[2];
const float *phi_pre, *phi_post, *phi_res, *a_pre, *a_post, *a_res, *b_pre, *b_post, *b_res, *final_norm;
QX qx;
SpinPool& pool;
std::string L(const char* name, int l) { return std::string("stack/layers/block/") + name + "/" + std::to_string(l); }
// y_i = W_i x for several quantised matrices of one width, x already prepared in qx; rows are shared out.
void rows(std::initializer_list<const CQ*> mats, std::initializer_list<float*> outs, int chunk = 32) {
const CQ* q[4];
float* y[4];
int off[5] = {0}, n = 0;
for (const CQ* p : mats) q[n++] = p;
n = 0;
for (float* p : outs) y[n++] = p;
for (int i = 0; i < n; i++) off[i + 1] = off[i] + q[i]->rows;
#if defined(OASIS_NEON)
// An Arm core reads a long run of rows far faster than scattered pieces of a few kilobytes, so each
// thread claims one contiguous share of the rows; the result does not depend on how rows are shared.
chunk = std::max(chunk, (off[n] + pool.size() - 1) / pool.size());
#endif
const CbTable t = cb_table(*q[0]);
pool.run(off[n], chunk, [&](int g) {
int k = 0;
while (g >= off[k + 1]) k++;
y[k][g - off[k]] = cq_row(*q[k], qx, g - off[k], t);
});
}
Decoder(Model& mm, const std::vector<float>& M, const std::vector<uint8_t>& mv, SpinPool& p) : m(mm), mvalid(mv), pool(p) {
nM = (int)(M.size() / 512);
emb = m.Q("embedding/embedding");
for (int l = 0; l < 8; l++) {
Layer& y = Ls[l];
y.q = m.Q(L("self_attn/q_proj/kernel", l));
y.k = m.Q(L("self_attn/k_proj/kernel", l));
y.v = m.Q(L("self_attn/v_proj/kernel", l));
y.gate = m.Q(L("self_attn/gate_proj/kernel", l));
y.out = m.Q(L("self_attn/out_proj/kernel", l));
y.cq_ = m.Q(L("cross_attn/q_proj/kernel", l));
y.cgate = m.Q(L("cross_attn/gate_proj/kernel", l));
y.cout = m.Q(L("cross_attn/out_proj/kernel", l));
y.taps[0] = m.F(L("self_attn/q_taps", l));
y.taps[1] = m.F(L("self_attn/k_taps", l));
y.taps[2] = m.F(L("self_attn/v_taps", l));
y.qn = m.F(L("self_attn/q_norm/scale", l));
y.kn = m.F(L("self_attn/k_norm/scale", l));
y.pre = m.F(L("ZCRMSNorm_0/scale", l));
y.post = m.F(L("post_attn_norm/scale", l));
y.cross = m.F(L("cross_norm/scale", l));
y.pcross = m.F(L("post_cross_norm/scale", l));
y.cqn = m.F(L("cross_attn/q_norm/scale", l));
y.hada = m.F(L("pre_hada_norm/scale", l));
y.ag = sigmoidf(m.F(L("attn_gate", l))[0]);
y.cg = sigmoidf(m.F(L("cross_gate", l))[0]);
auto P = [&](const char* n) { return m.F(L((std::string("hadamard_mlp/") + n).c_str(), l)); };
y.d1 = P("d1"); y.d2 = P("d2"); y.d3 = P("d3"); y.d4 = P("d4"); y.b2 = P("b2");
y.cond_v = P("cond_v"); y.cond_u = P("cond_u");
y.w1a = P("w1a"); y.w1b = P("w1b"); y.w2a = P("w2a"); y.w2b = P("w2b"); y.w3a = P("w3a"); y.w3b = P("w3b");
ck[l].resize((size_t)nM * 384);
cv[l].resize((size_t)nM * 512);
qmm(m, L("cross_attn/k_proj/kernel", l), M.data(), nM, ck[l].data());
qmm(m, L("cross_attn/v_proj/kernel", l), M.data(), nM, cv[l].data());
const float* kn = m.F(L("cross_attn/k_norm/scale", l));
for (int j = 0; j < nM; j++)
for (int h = 0; h < 8; h++) zc_norm(&ck[l][(size_t)j * 384 + h * 48], 48, kn, &ck[l][(size_t)j * 384 + h * 48]);
#if defined(OASIS_NEON)
nMp = (nM + 15) / 16 * 16;
cross->kt[l].assign((size_t)8 * 48 * nMp, 0);
for (int j = 0; j < nM; j++)
for (int h = 0; h < 8; h++)
for (int c = 0; c < 48; c++)
cross->kt[l][(size_t)h * 48 * nMp + (size_t)(j / 16) * 768 + c * 16 + j % 16] = to_half(ck[l][(size_t)j * 384 + h * 48 + c]);
cross->v16[l].resize((size_t)nM * 512);
for (int h = 0; h < 8; h++)
for (int j = 0; j < nM; j++)
for (int c = 0; c < 64; c++)
cross->v16[l][((size_t)h * nM + j) * 64 + c] = to_half(cv[l][(size_t)j * 512 + h * 64 + c]);
std::vector<float>().swap(ck[l]);
std::vector<float>().swap(cv[l]);
#endif
}
for (int s = 0; s < 2; s++) {
std::string pre = "engrams_" + std::to_string(s) + "/";
sites[s] = {m.Q(pre + "embedding"), m.Q(pre + "key_proj/kernel"), m.Q(pre + "value_proj/kernel"), m.F(pre + "taps")};
}
phi_pre = m.Wdq("stack/mhc_phi_pre");
phi_post = m.Wdq("stack/mhc_phi_post");
phi_res = m.Wdq("stack/mhc_phi_res");
a_pre = m.F("stack/mhc_a_pre"); a_post = m.F("stack/mhc_a_post"); a_res = m.F("stack/mhc_a_res");
b_pre = m.F("stack/mhc_b_pre"); b_post = m.F("stack/mhc_b_post"); b_res = m.F("stack/mhc_b_res");
final_norm = m.F("stack/final_norm/scale");
}
static void kron(const float* in, const float* a, const float* b, float* o) {
#if defined(OASIS_NEON)
// tmp = in [16, 32] times b [32, 32], then o = a^T [16, 16] times tmp, a row of 32 at a time in
// registers; every output adds its products in the same order as the loops below.
float tmp[16 * 32];
for (int i = 0; i < 16; i++) {
float32x4_t c[8];
for (int q = 0; q < 8; q++) c[q] = vdupq_n_f32(0.0f);
for (int j = 0; j < 32; j++) {
const float z = in[i * 32 + j];
const float* br = b + j * 32;
for (int q = 0; q < 8; q++) c[q] = vfmaq_n_f32(c[q], vld1q_f32(br + 4 * q), z);
}
for (int q = 0; q < 8; q++) vst1q_f32(tmp + i * 32 + 4 * q, c[q]);
}
for (int k = 0; k < 16; k++) {
float32x4_t c[8];
for (int q = 0; q < 8; q++) c[q] = vdupq_n_f32(0.0f);
for (int i = 0; i < 16; i++) {
const float av = a[i * 16 + k];
const float* tr = tmp + i * 32;
for (int q = 0; q < 8; q++) c[q] = vfmaq_n_f32(c[q], vld1q_f32(tr + 4 * q), av);
}
for (int q = 0; q < 8; q++) vst1q_f32(o + k * 32 + 4 * q, c[q]);
}
#else
float tmp[16 * 32] = {0};
for (int i = 0; i < 16; i++)
for (int j = 0; j < 32; j++) {
float z = in[i * 32 + j];
const float* br = b + j * 32;
for (int l2 = 0; l2 < 32; l2++) tmp[i * 32 + l2] += z * br[l2];
}
for (int k = 0; k < 16 * 32; k++) o[k] = 0.0f;
for (int i = 0; i < 16; i++)
for (int k = 0; k < 16; k++) {
float av = a[i * 16 + k];
for (int l2 = 0; l2 < 32; l2++) o[k * 32 + l2] += av * tmp[i * 32 + l2];
}
#endif
}
void hmlp(const Layer& y, const float* x, float* out) {
float c8[8], z[512], t[512];
float mx = -1e30f;
for (int r = 0; r < 8; r++) c8[r] = 0.0f;
for (int i = 0; i < 512; i++) {
const float* row = y.cond_v + (size_t)i * 8;
for (int r = 0; r < 8; r++) c8[r] += x[i] * row[r];
}
for (int r = 0; r < 8; r++) mx = std::max(mx, c8[r]);
float den = 0;
for (int r = 0; r < 8; r++) { c8[r] = std::exp(c8[r] - mx); den += c8[r]; }
float cond[512];
for (int i = 0; i < 512; i++) cond[i] = 1.0f;
for (int r = 0; r < 8; r++) {
float w = c8[r] / den;
const float* u = y.cond_u + (size_t)r * 512;
for (int i = 0; i < 512; i++) cond[i] += w * u[i];
}
for (int i = 0; i < 512; i++) t[i] = y.d1[i] * x[i];
kron(t, y.w1a, y.w1b, z);
for (int i = 0; i < 512; i++) t[i] = z[m.perm1[i]];
#if defined(OASIS_NEON)
for (int i = 0; i < 512; i++) z[i] = y.d2[i] * cond[i] * t[i] + y.b2[i];
silu_inplace(z, 512);
#else
for (int i = 0; i < 512; i++) z[i] = siluf(y.d2[i] * cond[i] * t[i] + y.b2[i]);
#endif
kron(z, y.w2a, y.w2b, t);
for (int i = 0; i < 512; i++) z[i] = t[m.perm2[i]];
for (int i = 0; i < 512; i++) t[i] = y.d3[i] * z[i];
kron(t, y.w3a, y.w3b, z);
for (int i = 0; i < 512; i++) out[i] = y.d4[i] * z[i];
}
void softmax_rows(float* s, int n) {
float mx = -1e30f;
for (int j = 0; j < n; j++) mx = std::max(mx, s[j]);
float den = 0.0f;
for (int j = 0; j < n; j++) { s[j] = fexp(s[j] - mx); den += s[j]; }
float inv = 1.0f / den;
for (int j = 0; j < n; j++) s[j] *= inv;
}
void block(int l, const float* u, int t, float* xout) {
const Layer& y = Ls[l];
float x[512], xn[512], tmp[512];
std::memcpy(x, u, sizeof x);
int site = l == 3 ? 0 : l == 7 ? 1 : -1;
if (site >= 0) {
const Site& st = sites[site];
float e[512];
for (int oi = 0; oi < 2; oi++)
for (int h = 0; h < 2; h++) {
int ti = oi * 2 + h;
if (t - (ORDERS[oi] - 1) < 0) std::memset(e + ti * 128, 0, 128 * sizeof(float));
else cq_dequant_row(st.tab, ti * 18432 + (int)engram_index(toks, t, oi, ORDERS[oi], h), e + ti * 128);
}
float k[512], vr[512], v[512] = {0};
prep_x(e, 512, 2, qx);
rows({&st.key, &st.val}, {k, vr});
ev[site].insert(ev[site].end(), vr, vr + 512);
for (int j = 0; j < 4; j++) {
int tt = t - 3 * j;
if (tt < 0) continue;
const float* tp = st.taps + j * 512;
const float* src = ev[site].data() + (size_t)tt * 512;
for (int c = 0; c < 512; c++) v[c] += tp[c] * src[c];
}
float ru[512], rk[512];
rms_unit(x, 512, ru);
rms_unit(k, 512, rk);
float alpha = sigmoidf(dotf(ru, rk, 512) / std::sqrt(512.0f));
for (int c = 0; c < 512; c++) x[c] += alpha * v[c];
}
float skip[512];
std::memcpy(skip, x, sizeof skip);
zc_norm(x, 512, y.pre, xn);
prep_x(xn, 512, 2, qx);
const int dims[3] = {384, 96, 128};
float qkv[3][384], proj[3][384];
float gate[512];
rows({&y.gate, &y.q, &y.k, &y.v}, {gate, proj[0], proj[1], proj[2]});
for (int s = 0; s < 3; s++) {
const float* r = proj[s];
raw[l][s].insert(raw[l][s].end(), r, r + dims[s]);
const float* taps = y.taps[s];
for (int c = 0; c < dims[s]; c++) {
float a = 0;
for (int j = 0; j < 3; j++)
if (t - j >= 0) a += taps[j * dims[s] + c] * raw[l][s][(size_t)(t - j) * dims[s] + c];
qkv[s][c] = a;
}
}
auto rope48 = [&](float* v) {
float a[24], b[24];
std::memcpy(a, v, sizeof a);
std::memcpy(b, v + 24, sizeof b);
for (int i = 0; i < 24; i++) {
float c = rope_c[(size_t)t * 24 + i], s = rope_s[(size_t)t * 24 + i];
v[i] = a[i] * c - b[i] * s;
v[24 + i] = b[i] * c + a[i] * s;
}
};
for (int h = 0; h < 8; h++) { zc_norm(qkv[0] + h * 48, 48, y.qn, qkv[0] + h * 48); rope48(qkv[0] + h * 48); }
for (int h = 0; h < 2; h++) { zc_norm(qkv[1] + h * 48, 48, y.kn, qkv[1] + h * 48); rope48(qkv[1] + h * 48); }
#if defined(OASIS_NEON)
for (int g = 0; g < 2; g++) {
K[l][g].insert(K[l][g].end(), qkv[1] + g * 48, qkv[1] + g * 48 + 48);
V[l][g].insert(V[l][g].end(), qkv[2] + g * 64, qkv[2] + g * 64 + 64);
}
const float* kb[2] = {K[l][0].data(), K[l][1].data()};
const float* vb[2] = {V[l][0].data(), V[l][1].data()};
const size_t ks = 48, vs = 64;
#else
K[l].insert(K[l].end(), qkv[1], qkv[1] + 96);
V[l].insert(V[l].end(), qkv[2], qkv[2] + 128);
#endif
int n = t + 1;
float o[512];
const int span = std::max(n, std::max(nM, nMp));
scratch.resize((size_t)8 * span);
const float isq = 1.0f / std::sqrt(48.0f);
// Heads are independent; each keeps its own score row.
pool.run(8, n >= 64 ? 1 : 8, [&](int h) {
float* s = scratch.data() + (size_t)h * span;
int g = h / 4;
#if defined(OASIS_NEON)
for (int j = 0; j < n; j++) s[j] = dotf(qkv[0] + h * 48, kb[g] + (size_t)j * ks, 48) * isq;
#else
for (int j = 0; j < n; j++) s[j] = dotf(qkv[0] + h * 48, &K[l][(size_t)j * 96 + g * 48], 48) * isq;
#endif
softmax_rows(s, n);
float* oh = o + h * 64;
for (int c = 0; c < 64; c++) oh[c] = 0;
for (int j = 0; j < n; j++) {
#if defined(OASIS_NEON)
const float* vv = vb[g] + (size_t)j * vs;
#else
const float* vv = &V[l][(size_t)j * 128 + g * 64];
#endif
float w = s[j];
for (int c = 0; c < 64; c++) oh[c] += w * vv[c];
}
});
#if defined(OASIS_NEON)
mul_sigmoid(o, gate, 512);
#else
for (int c = 0; c < 512; c++) o[c] *= sigmoidf(gate[c]);
#endif
float A[512];
prep_x(o, 512, 2, qx);
rows({&y.out}, {A});
zc_norm(A, 512, y.post, tmp);
for (int c = 0; c < 512; c++) x[c] = skip[c] + y.ag * tmp[c];
std::memcpy(skip, x, sizeof skip);
float xc[512], qc[384];
zc_norm(x, 512, y.cross, xc);
prep_x(xc, 512, 2, qx);
rows({&y.cq_, &y.cgate}, {qc, gate});
for (int h = 0; h < 8; h++) zc_norm(qc + h * 48, 48, y.cqn, qc + h * 48);
#if defined(OASIS_NEON)
// Scores for sixteen memory entries at a time from the transposed fp16 keys, then the fp16 values.
pool.run(8, nM >= 64 ? 1 : 8, [&](int h) {
float* s = scratch.data() + (size_t)h * span;
const uint16_t* kh = cross->kt[l].data() + (size_t)h * 48 * nMp;
const float* qh = qc + h * 48;
for (int j0 = 0; j0 < nMp; j0 += 16) {
float32x4_t a0 = vdupq_n_f32(0.0f), a1 = a0, a2 = a0, a3 = a0;
for (int c = 0; c < 48; c++) {
const uint16_t* kr = kh + (size_t)(j0 / 16) * 768 + c * 16;
const float16x8_t k0 = vreinterpretq_f16_u16(vld1q_u16(kr)), k1 = vreinterpretq_f16_u16(vld1q_u16(kr + 8));
const float qv = qh[c];
a0 = vfmaq_n_f32(a0, vcvt_f32_f16(vget_low_f16(k0)), qv);
a1 = vfmaq_n_f32(a1, vcvt_high_f32_f16(k0), qv);
a2 = vfmaq_n_f32(a2, vcvt_f32_f16(vget_low_f16(k1)), qv);
a3 = vfmaq_n_f32(a3, vcvt_high_f32_f16(k1), qv);
}
vst1q_f32(s + j0, a0);
vst1q_f32(s + j0 + 4, a1);
vst1q_f32(s + j0 + 8, a2);
vst1q_f32(s + j0 + 12, a3);
}
for (int j = 0; j < nM; j++) s[j] = mvalid[j] ? s[j] * isq : -1e30f;
softmax_rows(s, nM);
float32x4_t acc[16];
for (int q = 0; q < 16; q++) acc[q] = vdupq_n_f32(0.0f);
const uint16_t* vh = cross->v16[l].data() + (size_t)h * nM * 64;
for (int j = 0; j < nM; j++) {
if (!mvalid[j]) continue;
const float w = s[j];
const uint16_t* vv = vh + (size_t)j * 64;
for (int q = 0; q < 8; q++) {
const float16x8_t v8 = vreinterpretq_f16_u16(vld1q_u16(vv + q * 8));
acc[2 * q] = vfmaq_n_f32(acc[2 * q], vcvt_f32_f16(vget_low_f16(v8)), w);
acc[2 * q + 1] = vfmaq_n_f32(acc[2 * q + 1], vcvt_high_f32_f16(v8), w);
}
}
for (int q = 0; q < 16; q++) vst1q_f32(o + h * 64 + q * 4, acc[q]);
});
#else
pool.run(8, nM >= 64 ? 1 : 8, [&](int h) {
float* s = scratch.data() + (size_t)h * span;
for (int j = 0; j < nM; j++) {
if (!mvalid[j]) { s[j] = -1e30f; continue; }
s[j] = dotf(qc + h * 48, &ck[l][(size_t)j * 384 + h * 48], 48) * isq;
}
softmax_rows(s, nM);
float* oh = o + h * 64;
for (int c = 0; c < 64; c++) oh[c] = 0;
for (int j = 0; j < nM; j++) {
if (!mvalid[j]) continue;
const float* vv = &cv[l][(size_t)j * 512 + h * 64];
float w = s[j];
for (int c = 0; c < 64; c++) oh[c] += w * vv[c];
}
});
#endif
#if defined(OASIS_NEON)
mul_sigmoid(o, gate, 512);
#else
for (int c = 0; c < 512; c++) o[c] *= sigmoidf(gate[c]);
#endif
prep_x(o, 512, 2, qx);
rows({&y.cout}, {A});
zc_norm(A, 512, y.pcross, tmp);
for (int c = 0; c < 512; c++) x[c] = skip[c] + y.cg * tmp[c];
zc_norm(x, 512, y.hada, xn);
hmlp(y, xn, tmp);
for (int c = 0; c < 512; c++) xout[c] = x[c] + tmp[c];
}
std::vector<float> scratch, rope_c, rope_s;
void ensure_rope(int t) {
int have = (int)(rope_c.size() / 24);
if (t < have) return;
int want = std::max(t + 1, have * 2 + 64);
rope_c.resize((size_t)want * 24);
rope_s.resize((size_t)want * 24);
for (int p = have; p < want; p++)
for (int i = 0; i < 24; i++) {
double fr = 1.0 / std::pow(100000.0, (2.0 * i) / 48.0);
rope_c[(size_t)p * 24 + i] = (float)std::cos(p * fr);
rope_s[(size_t)p * 24 + i] = (float)std::sin(p * fr);
}
}
// Feed one token; returns logits over the vocabulary.
std::vector<float> step(int tok) {
int t = (int)toks.size();
toks.push_back(tok);
ensure_rope(t);
float X[4][512], row[512];
cq_dequant_row(emb, tok, row);
for (int ln = 0; ln < 4; ln++)
for (int c = 0; c < 512; c++) X[ln][c] = std::sqrt(512.0f) * row[c];
for (int l = 0; l < 8; l++) {
float nx[2048];
rms_unit(&X[0][0], 2048, nx);
float hv[24];
#if defined(OASIS_NEON)
const int phi_chunk = 8; // a contiguous share of the projections for each of three threads
pool.run(24, phi_chunk, [&](int i) {
#else
pool.run(24, 4, [&](int i) {
#endif
const float* row = i < 4 ? phi_pre + (size_t)(4 * l + i) * 2048
: i < 8 ? phi_post + (size_t)(4 * l + i - 4) * 2048
: phi_res + (size_t)(16 * l + i - 8) * 2048;
hv[i] = dotf(row, nx, 2048);
});
float hp[4], hq[4], hr[16];
for (int i = 0; i < 4; i++) {
hp[i] = hv[i];
hq[i] = hv[4 + i];
}
for (int i = 0; i < 16; i++) hr[i] = hv[8 + i];
float u[512] = {0}, yv[512];
for (int i = 0; i < 4; i++) {
float off = (l % 4 == i) ? 4.0f : -4.0f;
hp[i] = sigmoidf(a_pre[l] * hp[i] + b_pre[l * 4 + i] + off);
hq[i] = 2.0f * sigmoidf(a_post[l] * hq[i] + b_post[l * 4 + i] + ((l % 4 == i) ? 0.0f : -4.0f));
for (int c = 0; c < 512; c++) u[c] += hp[i] * X[i][c];
}
for (int i = 0; i < 16; i++) hr[i] = a_res[l] * hr[i] + b_res[l * 16 + i];
sinkhorn(hr);
block(l, u, t, yv);
for (int c = 0; c < 512; c++) yv[c] -= u[c];
float Xn[4][512];
for (int i = 0; i < 4; i++)
for (int c = 0; c < 512; c++)
Xn[i][c] = hq[i] * yv[c] + hr[i * 4] * X[0][c] + hr[i * 4 + 1] * X[1][c] + hr[i * 4 + 2] * X[2][c] + hr[i * 4 + 3] * X[3][c];
std::memcpy(X, Xn, sizeof X);
}
float h[512], hn[512];
for (int c = 0; c < 512; c++) h[c] = 0.25f * (X[0][c] + X[1][c] + X[2][c] + X[3][c]);
zc_norm(h, 512, final_norm, hn);
std::vector<float> logits(emb.rows);
prep_x(hn, 512, emb.bits, qx);
rows({&emb}, {logits.data()}, 128);
return logits;
}
};
static int decode_threads() { return std::max(1, std::min(engine_threads(), 8)); }
// Keeps the pool's workers spinning for the length of one decode.
struct Awake {
SpinPool& p;
explicit Awake(SpinPool& pool) : p(pool) { p.wake(); }
~Awake() { p.sleep(); }
};
// When the sequence ends in one short block written several times over, which the decoder only does
// once it has fallen into a loop (a block of 2 to 8 tokens four times, or one token six times), returns
// how many trailing tokens repeat the first copy; otherwise 0.
static int looping(const std::vector<int>& s) {
const int n = (int)s.size();
for (int p = 1; p <= 8; p++) {
int reps = p == 1 ? 6 : 4;
if (n < p * reps) continue;
bool same = true;
for (int k = 1; k < reps && same; k++)
for (int i = 0; i < p && same; i++) same = s[n - 1 - i] == s[n - 1 - i - k * p];
if (same) return (reps - 1) * p;
}
return 0;
}
// Longest reasoning trace kept; a trace that has not closed by then is closed by the engine.
static const int TRACE_MAX = 256;
// Keyword biasing: a small logit bonus for a keyword's opening word piece, and a larger one for each
// following piece once the earlier ones have been written, so a keyword the model starts is completed
// without common word starts being pushed into every line.
static const float KEYWORD_START_BONUS = 1.0f;
static const float KEYWORD_NEXT_BONUS = 3.0f;
std::vector<int> generate(Model& m, const std::vector<float>& M, const std::vector<uint8_t>& mvalid,
const std::vector<int>& refs, const std::string& language, int max_tokens,
const std::vector<std::string>& keywords) {
std::lock_guard<std::mutex> lock(m.decode_mutex);
if (!m.pool) m.pool = std::make_shared<SpinPool>(decode_threads());
Decoder dec(m, M, mvalid, *m.pool);
Awake awake(*m.pool);
std::vector<uint8_t> banned(m.vocab, 0);
for (auto& kv : m.langs) banned[kv.second] = 1;
for (int i = m.ref_base; i < m.vocab; i++) banned[i] = 1;
for (int r : refs) banned[r] = 0;
auto it = m.langs.find(language);
int lang = it != m.langs.end() ? it->second : m.langs.at("en");
// A keyword follows a space in names and the trace, or "=" in attribute values; each spelling is its
// own token sequence. Only the spaced spelling, which opens on a word start, gets the opening bonus.
std::vector<std::pair<std::vector<int>, bool>> seqs;
for (const auto& k : keywords)
for (int spaced = 1; spaced >= 0; spaced--) {
std::vector<int> s = encode_text(m, spaced ? " " + k : k);
bool dup = false;
for (const auto& q : seqs) dup = dup || q.first == s;
if (!s.empty() && !dup) seqs.push_back({s, spaced == 1});
}
int think = -1, think_end = -1;
for (const auto& mk : m.markers) {
if (mk.first == "<think>") think = mk.second;
if (mk.first == "</think>") think_end = mk.second;
}
std::vector<float> logits = dec.step(m.bos);
logits = dec.step(lang);
std::vector<int> out;
std::vector<std::pair<int, float>> bonus;
bool in_trace = false;
int trace_len = 0;
for (int s = 0; s < max_tokens; s++) {
bonus.clear();
for (const auto& [q, opens] : seqs) {
int n = std::min((int)q.size() - 1, (int)out.size());
for (int len = n; len >= 0; len--)
if (std::equal(q.begin(), q.begin() + len, out.end() - len)) {
if (!len && !opens) break;
float add = len ? KEYWORD_NEXT_BONUS : KEYWORD_START_BONUS;
bool have = false;
for (auto& b : bonus)
if (b.first == q[len]) { b.second = std::max(b.second, add); have = true; }
if (!have) bonus.push_back({q[len], add});
break;
}
}
for (const auto& b : bonus) logits[b.first] += b.second;
int best = -1;
float bv = -1e30f;
for (int i = 0; i < (int)logits.size(); i++)
if (!banned[i] && logits[i] > bv) { bv = logits[i]; best = i; }
if (best < 0 || best == m.eos) break;
// A trace that loops or runs on is closed so that the description still gets written; a
// description that loops ends where the loop began to repeat. Repeated copies are dropped.
out.push_back(best);
if (best == think && think >= 0) {
in_trace = true;
} else if (in_trace) {
int rep = looping(out);
if (best == think_end) {
in_trace = false;
} else if (++trace_len > TRACE_MAX || rep) {
out.resize(out.size() - std::max(rep, 1)); // the newest token was never fed to the decoder
out.push_back(think_end);
in_trace = false;
}
} else if (int rep = looping(out)) {
out.resize(out.size() - rep);
break;
}
logits = dec.step(out.back());
}
return out;
}
// BPE over one marker-free segment: adjacent symbols merged lowest piece id first (piece scores fall
// with id, so this is SentencePiece's highest-score-first order), unknown symbols as bytes.
static void bpe(const Model& m, const std::string& seg, std::vector<int>& ids) {
std::vector<std::string> syms;
for (size_t i = 0; i < seg.size();) {
unsigned char c = (unsigned char)seg[i];
size_t n = c < 0x80 ? 1 : (c >> 5) == 6 ? 2 : (c >> 4) == 14 ? 3 : (c >> 3) == 30 ? 4 : 1;
if (i + n > seg.size()) n = 1;
syms.push_back(seg.substr(i, n));
i += n;
}
for (;;) {
int best = INT_MAX;
size_t at = SIZE_MAX;
for (size_t i = 0; i + 1 < syms.size(); i++) {
auto it = m.piece_id.find(syms[i] + syms[i + 1]);
if (it != m.piece_id.end() && it->second < best) { best = it->second; at = i; }
}
if (at == SIZE_MAX) break;
syms[at] += syms[at + 1];
syms.erase(syms.begin() + (long)at + 1);
}
for (const auto& s : syms) {
auto it = m.piece_id.find(s);
if (it != m.piece_id.end()) ids.push_back(it->second);
else for (unsigned char b : s) ids.push_back(m.byte_id[b]);
}
}
// SentencePiece encoding as the vocabulary specifies it: spaces become U+2581, no dummy prefix,
// user-defined pieces such as <think> or language tags taken whole, the text between them by BPE.
std::vector<int> encode_text(const Model& m, const std::string& text) {
std::string esc;
for (char c : text) {
if (c == ' ') esc += "\xE2\x96\x81";
else esc += c;
}
std::vector<int> ids;
size_t i = 0, start = 0;
while (i < esc.size()) {
const std::pair<std::string, int>* hit = nullptr;
for (const auto& mk : m.markers)
if (esc.compare(i, mk.first.size(), mk.first) == 0) { hit = &mk; break; }
if (!hit) { i++; continue; }
bpe(m, esc.substr(start, i - start), ids);
ids.push_back(hit->second);
i += hit->first.size();
start = i;
}
bpe(m, esc.substr(start), ids);
return ids;
}
static std::string json_str(const std::string& s) {
std::string o = "\"";
for (unsigned char c : s) {
if (c == '"') o += "\\\"";
else if (c == '\\') o += "\\\\";
else if (c == '\n') o += "\\n";
else if (c < 0x20) { char b[8]; std::snprintf(b, sizeof b, "\\u%04x", c); o += b; }
else o += (char)c;
}
return o + "\"";
}
static std::string json_num(double v) {
if (!std::isfinite(v)) return "null";
char b[32];
std::snprintf(b, sizeof b, "%.6g", v);
return b;
}
// Log-probability of tok under the softmax over the tokens not in banned.
static double log_prob(const std::vector<float>& L, const std::vector<uint8_t>& banned, int tok) {
if (banned[tok]) return -INFINITY;
float mx = -1e30f;
for (size_t i = 0; i < L.size(); i++) if (!banned[i] && L[i] > mx) mx = L[i];
double s = 0.0;
for (size_t i = 0; i < L.size(); i++) if (!banned[i]) s += std::exp((double)L[i] - mx);
return (double)L[tok] - mx - std::log(s);
}
// A form is one request line naming every field, then one answer line per field:
// ? sign=stop|yield; n#0..3; note$60; where@sign
// sign=stop
// n=1 ...
// Choices and counts are scored as complete answers with their terminator (a newline, or </s> after the
// last field), choices with "-" for abstention; text is decoded greedily; a box is a softmax over the
// occupied reference tokens. Each answer is written into the context before the next field.
std::string fill_form(Model& m, const std::vector<float>& M, const std::vector<uint8_t>& mvalid, const std::vector<int>& refs,
const std::vector<std::vector<int>>& ref_boxes, const std::string& language, const std::vector<FormField>& form) {
std::lock_guard<std::mutex> lock(m.decode_mutex);
if (!m.pool) m.pool = std::make_shared<SpinPool>(decode_threads());
Decoder dec(m, M, mvalid, *m.pool);
Awake awake(*m.pool);
std::vector<uint8_t> banned(m.vocab, 0), no_refs(m.vocab, 0);
for (auto& kv : m.langs) banned[kv.second] = no_refs[kv.second] = 1;
for (int i = m.ref_base; i < m.vocab; i++) banned[i] = no_refs[i] = 1;
for (int r : refs) banned[r] = 0;
auto lit = m.langs.find(language);
int lang = lit != m.langs.end() ? lit->second : m.langs.at("en");
std::string line = "? ";
for (size_t i = 0; i < form.size(); i++) {
const FormField& f = form[i];
if (i) line += "; ";
line += f.name;
if (f.kind == FormField::CHOICE) {
line += "=";
for (size_t c = 0; c < f.choices.size(); c++) line += (c ? "|" : "") + f.choices[c];
} else if (f.kind == FormField::COUNT) {
line += "#" + std::to_string(f.lo) + ".." + std::to_string(f.hi);
} else if (f.kind == FormField::TEXT) {
line += "$" + std::to_string(f.max_length);
} else {
line += "@" + f.of;
}
}
std::vector<int> prompt = {m.bos, lang};
for (int t : encode_text(m, line + "\n")) prompt.push_back(t);
if (prompt.size() > 1024) throw std::runtime_error("form exceeds the decoder context");
std::vector<float> logits;
for (int t : prompt) logits = dec.step(t);
const std::vector<int> newline = encode_text(m, "\n");
std::ostringstream js;
js << "{";
for (size_t i = 0; i < form.size(); i++) {
const FormField& f = form[i];
const bool last = i + 1 == form.size();
const std::vector<int> term = last ? std::vector<int>{m.eos} : newline;
for (int t : encode_text(m, f.name + "=")) logits = dec.step(t);
js << (i ? "," : "") << json_str(f.name) << ":";
std::vector<int> answer; // tokens written into the context after "name="
if (f.kind == FormField::CHOICE || f.kind == FormField::COUNT) {
std::vector<std::string> cands;
if (f.kind == FormField::CHOICE) {
cands = f.choices;
cands.push_back("-");
} else {
for (long v = f.lo; v <= f.hi; v++) cands.push_back(std::to_string(v));
}
std::vector<std::vector<int>> toks(cands.size());
std::vector<double> score(cands.size());
for (size_t c = 0; c < cands.size(); c++) {
toks[c] = encode_text(m, cands[c]);
toks[c].insert(toks[c].end(), term.begin(), term.end());
Decoder d = dec;
std::vector<float> L = logits;
double s = 0.0;
for (size_t k = 0; k < toks[c].size(); k++) {
s += log_prob(L, banned, toks[c][k]);
if (k + 1 < toks[c].size()) L = d.step(toks[c][k]);
}
score[c] = s;
}
double mx = -INFINITY, den = 0.0;
size_t best = 0;
for (size_t c = 0; c < cands.size(); c++) if (score[c] > mx) { mx = score[c]; best = c; }
for (double s : score) den += std::exp(s - mx);
if (f.kind == FormField::CHOICE)
js << "{\"value\":" << (best == f.choices.size() ? "null" : f.values[best]);
else
js << "{\"value\":" << (f.lo + (long)best);
js << ",\"probabilities\":{";
for (size_t c = 0; c < cands.size(); c++)
js << (c ? "," : "") << json_str(cands[c]) << ":" << json_num(std::exp(score[c] - mx) / den);
js << "}}";
answer = toks[best];
} else if (f.kind == FormField::TEXT) {
std::vector<int> gen;
double s = 0.0;
bool ended = false;
for (int step = 0; step < 4 * f.max_length + 8; step++) {
int tok = -1;
float bv = -1e30f;
for (int v = 0; v < (int)logits.size(); v++)
if (!no_refs[v] && logits[v] > bv) { bv = logits[v]; tok = v; }
s += log_prob(logits, no_refs, tok);
if (tok == term[0] || tok == m.eos || tok == newline[0]) { ended = true; break; }
gen.push_back(tok);
logits = dec.step(tok);
std::string sofar = detokenize(m, gen);
size_t chars = 0;
for (unsigned char c : sofar) chars += (c & 0xC0) != 0x80;
if ((int)chars >= f.max_length) break;
}
if (!ended) s += log_prob(logits, no_refs, term[0]);
js << "{\"value\":" << json_str(detokenize(m, gen)) << ",\"probability\":" << json_num(std::exp(s)) << "}";
answer = term; // the generated tokens are already in the context
} else {
if (refs.empty()) {
js << "{\"value\":null,\"probability\":null}";
answer = encode_text(m, "-");
answer.insert(answer.end(), term.begin(), term.end());
} else {
float mx = -1e30f;
for (int r : refs) mx = std::max(mx, logits[r]);
double den = 0.0;
size_t best = 0;
for (size_t k = 0; k < refs.size(); k++) {
den += std::exp((double)logits[refs[k]] - mx);
if (logits[refs[k]] > logits[refs[best]]) best = k;
}
const std::vector<int>& b = ref_boxes[best];
js << "{\"value\":[" << b[0] << "," << b[1] << "," << b[2] << "," << b[3] << "],\"probability\":"
<< json_num(std::exp((double)logits[refs[best]] - mx) / den) << "}";
answer = {refs[best]};
answer.insert(answer.end(), term.begin(), term.end());
}
}
if (!last)
for (int t : answer) logits = dec.step(t);
}
js << "}";
return js.str();
}
std::string detokenize(const Model& m, const std::vector<int>& ids) {
std::string out, bytes;
auto flush = [&]() { out += bytes; bytes.clear(); };
for (int id : ids) {
if (id >= m.n_pieces) {
flush();
int k = id - m.ref_base;
out += k < 64 ? "<obj" + std::to_string(k) + ">" : "<txt" + std::to_string(k - 64) + ">";
continue;
}
const std::string& p = m.pieces[id];
int t = m.piece_types[id];
if (t == 4 && p.size() == 6 && p.rfind("<0x", 0) == 0) {
bytes += (char)std::stoi(p.substr(3, 2), nullptr, 16);
continue;
}
flush();
if (t == 0) {
std::string s = p;
size_t pos;
while ((pos = s.find("\xE2\x96\x81")) != std::string::npos) s.replace(pos, 3, " ");
out += s;
} else {
out += p;
}
}
flush();
return out;
}
} // namespace oasis
// ==== Record assembly, mask outlines and bitmaps, forms, options parsing and the C API
namespace oasis {
// The calling thread's OpenMP thread count set to engine_threads() for one call, and restored after it.
struct ThreadScope {
#ifdef _OPENMP
int prev = omp_get_max_threads();
ThreadScope() { omp_set_num_threads(engine_threads()); }
~ThreadScope() { omp_set_num_threads(prev); }
#endif
};
// ------------------------------------------------------------------ JSON writing
static std::string esc(const std::string& s) {
std::string o = "\"";
for (unsigned char c : s) {
if (c == '"') o += "\\\"";
else if (c == '\\') o += "\\\\";
else if (c == '\n') o += "\\n";
else if (c == '\t') o += "\\t";
else if (c < 0x20) { char b[8]; std::snprintf(b, sizeof b, "\\u%04x", c); o += b; }
else o += (char)c;
}
return o + "\"";
}
static std::string num(double v, int prec = 3) {
if (!std::isfinite(v)) return "null";
char b[32];
std::snprintf(b, sizeof b, "%.*f", prec, v);
std::string s = b;
if (s.find('.') != std::string::npos) {
while (s.back() == '0') s.pop_back();
if (s.back() == '.') s.pop_back();
}
return s;
}
static std::string trim(const std::string& s) {
size_t a = s.find_first_not_of(" \t\r"), b = s.find_last_not_of(" \t\r");
return a == std::string::npos ? "" : s.substr(a, b - a + 1);
}
// ------------------------------------------------------------------ JSON reading
struct JV {
enum Type { NUL, BOOL, NUM, STR, ARR, OBJ } t = NUL;
bool b = false;
double n = 0.0;
std::string s; // string value, or a number's source text
std::vector<JV> a;
std::vector<std::pair<std::string, JV>> o; // members in document order
const JV* get(const char* k) const {
for (const auto& kv : o)
if (kv.first == k) return &kv.second;
return nullptr;
}
};
struct JReader {
const char* p;
const char* e;
[[noreturn]] void fail(const std::string& what) { throw std::runtime_error("invalid options: " + what); }
void ws() {
while (p < e && (*p == ' ' || *p == '\t' || *p == '\n' || *p == '\r')) p++;
}
static void utf8(std::string& out, unsigned cp) {
if (cp < 0x80) out += (char)cp;
else if (cp < 0x800) { out += (char)(0xC0 | (cp >> 6)); out += (char)(0x80 | (cp & 0x3F)); }
else if (cp < 0x10000) { out += (char)(0xE0 | (cp >> 12)); out += (char)(0x80 | ((cp >> 6) & 0x3F)); out += (char)(0x80 | (cp & 0x3F)); }
else { out += (char)(0xF0 | (cp >> 18)); out += (char)(0x80 | ((cp >> 12) & 0x3F)); out += (char)(0x80 | ((cp >> 6) & 0x3F)); out += (char)(0x80 | (cp & 0x3F)); }
}
unsigned hex4() {
if (e - p < 4) fail("bad \\u escape");
unsigned v = 0;
for (int i = 0; i < 4; i++) {
char c = *p++;
v <<= 4;
if (c >= '0' && c <= '9') v |= (unsigned)(c - '0');
else if (c >= 'a' && c <= 'f') v |= (unsigned)(c - 'a' + 10);
else if (c >= 'A' && c <= 'F') v |= (unsigned)(c - 'A' + 10);
else fail("bad \\u escape");
}
return v;
}
std::string str() {
p++;
std::string out;
while (p < e && *p != '"') {
char c = *p++;
if (c != '\\') { out += c; continue; }
if (p >= e) fail("unterminated string");
char x = *p++;
switch (x) {
case '"': out += '"'; break;
case '\\': out += '\\'; break;
case '/': out += '/'; break;
case 'b': out += '\b'; break;
case 'f': out += '\f'; break;
case 'n': out += '\n'; break;
case 'r': out += '\r'; break;
case 't': out += '\t'; break;
case 'u': {
unsigned cp = hex4();
if (cp >= 0xD800 && cp < 0xDC00 && e - p >= 6 && p[0] == '\\' && p[1] == 'u') {
p += 2;
unsigned lo = hex4();
cp = 0x10000 + ((cp - 0xD800) << 10) + (lo - 0xDC00);
}
utf8(out, cp);
break;
}
default: fail("bad escape in string");
}
}
if (p >= e) fail("unterminated string");
p++;
return out;
}
JV value() {
ws();
if (p >= e) fail("unexpected end of JSON");
JV v;
if (*p == '{') {
p++;
v.t = JV::OBJ;
ws();
if (p < e && *p == '}') { p++; return v; }
for (;;) {
ws();
if (p >= e || *p != '"') fail("expected a key");
std::string k = str();
ws();
if (p >= e || *p != ':') fail("expected ':'");
p++;
JV x = value();
v.o.emplace_back(std::move(k), std::move(x));
ws();
if (p < e && *p == ',') { p++; continue; }
if (p < e && *p == '}') { p++; return v; }
fail("expected ',' or '}'");
}
}
if (*p == '[') {
p++;
v.t = JV::ARR;
ws();
if (p < e && *p == ']') { p++; return v; }
for (;;) {
v.a.push_back(value());
ws();
if (p < e && *p == ',') { p++; continue; }
if (p < e && *p == ']') { p++; return v; }
fail("expected ',' or ']'");
}
}
if (*p == '"') {
v.t = JV::STR;
v.s = str();
return v;
}
auto word = [&](const char* w) {
size_t n = std::strlen(w);
if ((size_t)(e - p) >= n && std::strncmp(p, w, n) == 0) { p += n; return true; }
return false;
};
if (word("true")) { v.t = JV::BOOL; v.b = true; return v; }
if (word("false")) { v.t = JV::BOOL; return v; }
if (word("null")) return v;
const char* q = p;
while (p < e && (std::isdigit((unsigned char)*p) || *p == '-' || *p == '+' || *p == '.' || *p == 'e' || *p == 'E')) p++;
if (q == p) fail("unexpected character");
v.t = JV::NUM;
v.s.assign(q, p);
char* end = nullptr;
v.n = std::strtod(v.s.c_str(), &end);
if (!end || *end) fail("bad number " + v.s);
return v;
}
};
// ------------------------------------------------------------------ mask geometry
using Ring = std::vector<std::pair<float, float>>; // vertices in mask-grid units, cell (i, j) centred at (j + 0.5, i + 0.5)
static double ring_area(const Ring& r) {
double a = 0.0;
for (size_t i = 0, n = r.size(); i < n; i++) {
const auto& p = r[i];
const auto& q = r[(i + 1) % n];
a += (double)p.first * q.second - (double)q.first * p.second;
}
return 0.5 * a;
}
// Douglas-Peucker on a closed ring: split at the vertex farthest from the first, simplify both halves.
static Ring simplify_ring(const Ring& pts, float tol) {
const size_t n = pts.size();
if (n < 4) return pts;
size_t far = 0;
float fd = -1.0f;
for (size_t i = 1; i < n; i++) {
float d = std::hypot(pts[i].first - pts[0].first, pts[i].second - pts[0].second);
if (d > fd) { fd = d; far = i; }
}
std::vector<uint8_t> keep(n, 0);
keep[0] = keep[far] = 1;
auto pt = [&](size_t i) { return pts[i % n]; };
std::vector<std::pair<size_t, size_t>> st{{0, far}, {far, n}};
while (!st.empty()) {
auto [a, c] = st.back();
st.pop_back();
auto A = pt(a), C = pt(c);
float L = std::hypot(C.first - A.first, C.second - A.second), best = 0.0f;
size_t bi = 0;
for (size_t i = a + 1; i < c; i++) {
auto P = pt(i);
float d = L > 0 ? std::fabs((C.first - A.first) * (A.second - P.second) - (A.first - P.first) * (C.second - A.second)) / L
: std::hypot(P.first - A.first, P.second - A.second);
if (d > best) { best = d; bi = i; }
}
if (bi && best > tol) { keep[bi % n] = 1; st.push_back({a, bi}); st.push_back({bi, c}); }
}
Ring out;
for (size_t i = 0; i < n; i++)
if (keep[i]) out.push_back(pts[i]);
return out;
}
// Exterior outlines of the thresholded 128 x 128 probability field, largest first. Marching squares over
// the cell centres, with every crossing interpolated linearly between the two neighbouring probabilities;
// the field is zero outside the grid, so every outline closes. Hole boundaries are left out.
static std::vector<Ring> outlines(const std::vector<float>& prob, float threshold, bool largest_only) {
const int N = 128, M = N + 2; // lattice -1 .. N, stored with an offset of one
const float t = std::max(threshold, 1e-6f);
auto f = [&](int i, int j) -> float {
if (i < 0 || j < 0 || i >= N || j >= N) return -t;
return prob[(size_t)i * N + j] - t;
};
auto inside = [&](int i, int j) { return f(i, j) >= 0.0f; };
// Edge ids: horizontal edge from lattice (i, j) to (i, j + 1), vertical edge from (i, j) to (i + 1, j).
auto hid = [&](int i, int j) { return (((i + 1) * M) + (j + 1)) * 2; };
auto vid = [&](int i, int j) { return (((i + 1) * M) + (j + 1)) * 2 + 1; };
auto cross = [&](int i0, int j0, int i1, int j1) {
float a = f(i0, j0), b = f(i1, j1);
float s = a / (a - b);
return std::make_pair(j0 + 0.5f + s * (j1 - j0), i0 + 0.5f + s * (i1 - i0));
};
enum { T, R, B, L };
struct Seg { int from, to; std::pair<float, float> p; };
std::vector<Seg> segs;
std::vector<int> start_at((size_t)M * M * 2, -1);
for (int i = -1; i < N; i++)
for (int j = -1; j < N; j++) {
int code = (inside(i, j) ? 8 : 0) | (inside(i, j + 1) ? 4 : 0) | (inside(i + 1, j + 1) ? 2 : 0) | (inside(i + 1, j) ? 1 : 0);
if (code == 0 || code == 15) continue;
const int eid[4] = {hid(i, j), vid(i, j + 1), hid(i + 1, j), vid(i, j)};
auto point = [&](int e) {
switch (e) {
case T: return cross(i, j, i, j + 1);
case R: return cross(i, j + 1, i + 1, j + 1);
case B: return cross(i + 1, j, i + 1, j + 1);
default: return cross(i, j, i + 1, j);
}
};
// Directed so that the inside lies to the right of travel on the image grid (y down).
int pairs[2][2], np = 0;
auto add = [&](int a, int b) { pairs[np][0] = a; pairs[np][1] = b; np++; };
float centre = 0.25f * (f(i, j) + f(i, j + 1) + f(i + 1, j + 1) + f(i + 1, j));
switch (code) {
case 1: add(L, B); break;
case 2: add(B, R); break;
case 3: add(L, R); break;
case 4: add(R, T); break;
case 6: add(B, T); break;
case 7: add(L, T); break;
case 8: add(T, L); break;
case 9: add(T, B); break;
case 11: add(T, R); break;
case 12: add(R, L); break;
case 13: add(R, B); break;
case 14: add(B, L); break;
case 5:
if (centre >= 0.0f) { add(L, T); add(R, B); } else { add(R, T); add(L, B); }
break;
case 10:
if (centre >= 0.0f) { add(T, R); add(B, L); } else { add(T, L); add(B, R); }
break;
}
for (int k = 0; k < np; k++) {
start_at[eid[pairs[k][0]]] = (int)segs.size();
segs.push_back({eid[pairs[k][0]], eid[pairs[k][1]], point(pairs[k][0])});
}
}
std::vector<uint8_t> used(segs.size(), 0);
std::vector<Ring> rings;
for (size_t s0 = 0; s0 < segs.size(); s0++) {
if (used[s0]) continue;
Ring r;
int s = (int)s0;
while (s >= 0 && !used[s]) {
used[s] = 1;
r.push_back(segs[s].p);
s = start_at[segs[s].to];
}
if (r.size() >= 3) rings.push_back(std::move(r));
}
if (rings.empty()) return rings;
// Outer boundaries share one orientation and holes the other; the largest ring is always an outer one.
size_t big = 0;
for (size_t k = 1; k < rings.size(); k++)
if (std::fabs(ring_area(rings[k])) > std::fabs(ring_area(rings[big]))) big = k;
const bool positive = ring_area(rings[big]) > 0;
std::vector<std::pair<double, Ring>> outer;
for (auto& r : rings) {
double a = ring_area(r);
if ((a > 0) == positive && std::fabs(a) >= 0.25) outer.push_back({std::fabs(a), simplify_ring(r, 0.5f)});
}
std::stable_sort(outer.begin(), outer.end(), [](const auto& a, const auto& b) { return a.first > b.first; });
std::vector<Ring> out;
for (auto& o : outer) {
if (o.second.size() < 3) continue;
out.push_back(std::move(o.second));
if (largest_only) break;
}
return out;
}
// All pieces joined into one path: each piece is entered from the nearest vertex of the path so far,
// walked once round, and left by the same bridge.
static Ring merged_ring(const std::vector<Ring>& rings) {
if (rings.empty()) return {};
Ring path = rings[0];
for (size_t k = 1; k < rings.size(); k++) {
const Ring& r = rings[k];
size_t bi = 0, bj = 0;
float bd = 1e30f;
for (size_t i = 0; i < path.size(); i++)
for (size_t j = 0; j < r.size(); j++) {
float dx = path[i].first - r[j].first, dy = path[i].second - r[j].second, d = dx * dx + dy * dy;
if (d < bd) { bd = d; bi = i; bj = j; }
}
Ring out(path.begin(), path.begin() + (long)bi + 1);
for (size_t t = 0; t <= r.size(); t++) out.push_back(r[(bj + t) % r.size()]);
out.insert(out.end(), path.begin() + (long)bi, path.end());
path.swap(out);
}
return path;
}
// Column-major run-length encoding at the original image size, starting with a background run: the
// probability field is sampled bilinearly at every pixel centre, thresholded and clipped to the box.
static std::string bitmap_rle(const std::vector<float>& prob, const int bb[4], const Meta& mt, float threshold) {
const int W = mt.W, H = mt.H;
const double bx0 = bb[0] / 1000.0 * W, by0 = bb[1] / 1000.0 * H, bx1 = bb[2] / 1000.0 * W, by1 = bb[3] / 1000.0 * H;
const int x0 = std::max(0, (int)std::floor(bx0)), x1 = std::min(W, (int)std::ceil(bx1));
const int y0 = std::max(0, (int)std::floor(by0)), y1 = std::min(H, (int)std::ceil(by1));
std::vector<long long> counts;
long long run = 0;
bool fg = false;
auto push = [&](bool v, long long n) {
if (n <= 0) return;
if (v == fg) { run += n; return; }
counts.push_back(run);
run = n;
fg = v;
};
const double k = mt.r / 4.0; // original pixels -> mask cells
for (int x = 0; x < W; x++) {
if (x < x0 || x >= x1 || y1 <= y0) { push(false, H); continue; }
push(false, y0);
double u = (x + 0.5) * k - 0.5;
int j0 = (int)std::floor(u);
double fu = u - j0;
int ja = std::clamp(j0, 0, 127), jb = std::clamp(j0 + 1, 0, 127);
bool colin = x + 0.5 >= bx0 && x + 0.5 <= bx1;
for (int y = y0; y < y1; y++) {
bool v = false;
if (colin && y + 0.5 >= by0 && y + 0.5 <= by1) {
double w = (y + 0.5) * k - 0.5;
int i0 = (int)std::floor(w);
double fv = w - i0;
int ia = std::clamp(i0, 0, 127), ib = std::clamp(i0 + 1, 0, 127);
double p = (1 - fv) * ((1 - fu) * prob[ia * 128 + ja] + fu * prob[ia * 128 + jb]) +
fv * ((1 - fu) * prob[ib * 128 + ja] + fu * prob[ib * 128 + jb]);
v = p >= threshold;
}
push(v, 1);
}
push(false, H - y1);
}
counts.push_back(run);
std::ostringstream os;
os << "{\"size\":[" << H << "," << W << "],\"counts\":[";
for (size_t i = 0; i < counts.size(); i++) os << (i ? "," : "") << counts[i];
os << "]}";
return os.str();
}
// ------------------------------------------------------------------ describe
struct Entity {
int n = 0;
std::string name, position, count;
std::vector<std::pair<std::string, std::string>> attrs;
std::vector<std::pair<char, int>> refs;
};
static std::vector<std::pair<std::string, std::string>> parse_attrs(const std::string& s) {
std::vector<std::pair<std::string, std::string>> out;
std::stringstream ss(s);
std::string part;
while (std::getline(ss, part, ';')) {
auto eq = part.find('=');
if (eq == std::string::npos) continue;
out.push_back({trim(part.substr(0, eq)), trim(part.substr(eq + 1))});
}
return out;
}
static bool has(const std::vector<std::string>& v, const char* x) { return std::find(v.begin(), v.end(), x) != v.end(); }
std::string describe_json(Model& m, const uint8_t* rgb, int W, int H, const Options& o, double* tm) {
ThreadScope threads;
using clk = std::chrono::steady_clock;
auto ms = [](clk::time_point x, clk::time_point y) { return std::chrono::duration<double, std::milli>(y - x).count(); };
auto t0 = clk::now();
Encoded enc = encode(m, rgb, W, H);
auto t1 = clk::now();
std::vector<Box> objects, texts;
detect(enc, o.conf, o.text_conf, o.nms, o.max_objects, o.max_text_lines, objects, texts);
const Meta& mt = enc.meta;
auto to1000 = [&](const Box& b, int out[4]) {
double sx = 512.0 / mt.r / mt.W * 1000.0, sy = 512.0 / mt.r / mt.H * 1000.0;
out[0] = (int)std::nearbyint(b.x0 * sx);
out[1] = (int)std::nearbyint(b.y0 * sy);
out[2] = (int)std::nearbyint(b.x1 * sx);
out[3] = (int)std::nearbyint(b.y1 * sy);
};
if (!o.form.empty()) {
std::vector<float> M;
std::vector<uint8_t> mv;
std::vector<int> refs;
memory(m, enc, objects, texts, M, mv, refs);
std::vector<std::vector<int>> ref_boxes;
for (const auto* list : {&objects, &texts})
for (const Box& b : *list) {
int bb[4];
to1000(b, bb);
ref_boxes.push_back({bb[0], bb[1], bb[2], bb[3]});
}
auto t2 = clk::now();
std::string form = fill_form(m, M, mv, refs, ref_boxes, o.language, o.form);
auto t3 = clk::now();
if (tm) { tm[0] = ms(t0, t1); tm[1] = ms(t1, t2); tm[2] = ms(t2, t3); tm[3] = ms(t0, t3); tm[4] = 0; }
return "{\"form\":" + form + ",\"timing_ms\":{\"encode\":" + num(ms(t0, t1), 1) + ",\"detect\":" + num(ms(t1, t2), 1) +
",\"form\":" + num(ms(t2, t3), 1) + ",\"total\":" + num(ms(t0, t3), 1) + "}}";
}
auto wants = [&](const char* s) { return o.sections.empty() || has(o.sections, s); };
auto field = [&](const char* f) { return o.fields.empty() || has(o.fields, f); };
std::vector<std::string> ocr;
for (auto& b : texts) {
if (!wants("text")) { ocr.push_back(""); continue; }
int bb[4];
to1000(b, bb);
ocr.push_back(read_text(m, rgb, W, H, bb[0] / 1000.0f * W, bb[1] / 1000.0f * H, bb[2] / 1000.0f * W, bb[3] / 1000.0f * H));
}
auto t2 = clk::now();
bool need_decoder = wants("scene") || wants("objects") || wants("regions") || wants("relations") || wants("events") || wants("reasoning");
std::vector<float> M;
std::vector<uint8_t> mv;
std::vector<int> refs;
std::vector<int> ids;
if (need_decoder) {
memory(m, enc, objects, texts, M, mv, refs);
ids = generate(m, M, mv, refs, o.language, o.max_tokens, o.keywords);
}
auto t3 = clk::now();
std::string text = detokenize(m, ids);
std::string trace, body = text;
size_t a = text.find("<think>"), b = text.find("</think>");
if (a != std::string::npos && b != std::string::npos && b > a) {
trace = trim(text.substr(a + 7, b - a - 7));
body = text.substr(b + 8);
}
std::vector<std::pair<std::string, std::string>> scene;
std::vector<Entity> ents;
std::vector<int> text_refs;
std::vector<std::string> rel_json, events;
std::stringstream ls(body);
std::string line;
while (std::getline(ls, line)) {
line = trim(line);
if (line.empty()) continue;
if (line.rfind("S ", 0) == 0) {
for (auto& kv : parse_attrs(line.substr(2))) scene.push_back(kv);
} else if (line.size() > 1 && line[0] == 'E' && std::isdigit((unsigned char)line[1])) {
Entity e;
size_t sp = line.find(' ');
e.n = std::atoi(line.substr(1, sp).c_str());
std::string rest = sp == std::string::npos ? "" : line.substr(sp + 1);
size_t br = rest.rfind('{');
if (br != std::string::npos && !rest.empty() && rest.back() == '}') {
e.attrs = parse_attrs(rest.substr(br + 1, rest.size() - br - 2));
rest = trim(rest.substr(0, br));
}
size_t r0;
while ((r0 = rest.find('<')) != std::string::npos) {
size_t r1 = rest.find('>', r0);
if (r1 == std::string::npos) break;
std::string tag = rest.substr(r0 + 1, r1 - r0 - 1);
if (tag.rfind("obj", 0) == 0) e.refs.push_back({'o', std::atoi(tag.c_str() + 3)});
else if (tag.rfind("txt", 0) == 0) e.refs.push_back({'t', std::atoi(tag.c_str() + 3)});
rest.erase(r0, r1 - r0 + 1);
}
for (auto& c : rest)
if (c == '*') c = ' ';
size_t at = rest.find('@');
if (at != std::string::npos) { e.position = trim(rest.substr(at + 1)); rest = rest.substr(0, at); }
rest = trim(rest);
size_t hs = rest.rfind(" #");
if (hs != std::string::npos) { e.count = trim(rest.substr(hs + 2)); rest = trim(rest.substr(0, hs)); }
e.name = rest;
ents.push_back(e);
} else if (line[0] == 'T') {
size_t r0 = 0;
while ((r0 = line.find("<txt", r0)) != std::string::npos) { text_refs.push_back(std::atoi(line.c_str() + r0 + 4)); r0 += 4; }
} else if (line.rfind("R ", 0) == 0) {
std::stringstream rs(line.substr(2));
int s1, s2;
std::string pred;
if (rs >> s1 >> pred >> s2)
rel_json.push_back("{\"subject\":" + std::to_string(s1) + ",\"predicate\":" + esc(pred) + ",\"object\":" + std::to_string(s2) + "}");
} else if (line.rfind("V ", 0) == 0) {
events.push_back(trim(line.substr(2)));
}
}
std::vector<float> protos;
const bool draw = o.masks && o.mask_format != "none" && field("mask");
auto mask_json = [&](const Box& box, const int bb[4]) -> std::string {
std::vector<float> prob = mask_prob(m, enc, box, &protos);
if (o.mask_format == "bitmap") return bitmap_rle(prob, bb, mt, o.mask_threshold);
double sx = 512.0 / 128.0 / mt.r / mt.W * 1000.0, sy = 512.0 / 128.0 / mt.r / mt.H * 1000.0;
auto ring_str = [&](const Ring& rg) {
std::ostringstream rs;
rs << "[";
for (size_t i = 0; i < rg.size(); i++)
rs << (i ? "," : "") << (int)std::nearbyint(std::clamp(rg[i].first * sx, 0.0, 1000.0)) << ","
<< (int)std::nearbyint(std::clamp(rg[i].second * sy, 0.0, 1000.0));
rs << "]";
return rs.str();
};
std::vector<Ring> rings = outlines(prob, o.mask_threshold, o.mask_format == "largest");
if (o.mask_format == "largest") return rings.empty() ? "[]" : ring_str(rings[0]);
if (o.mask_format == "merged") return ring_str(merged_ring(rings));
std::string s = "[";
for (size_t i = 0; i < rings.size(); i++) s += (i ? "," : "") + ring_str(rings[i]);
return s + "]";
};
auto common = [&](const Entity& e) {
std::ostringstream it;
it << "{\"id\":" << e.n << ",\"name\":" << esc(e.name);
if (!e.count.empty() && field("count")) {
bool digits = std::all_of(e.count.begin(), e.count.end(), ::isdigit);
it << ",\"count\":" << (digits ? e.count : esc(e.count));
}
if (!e.position.empty() && field("position")) it << ",\"position\":" << esc(e.position);
if (!e.attrs.empty() && field("description")) {
it << ",\"description\":{";
for (size_t i = 0; i < e.attrs.size(); i++) it << (i ? "," : "") << esc(e.attrs[i].first) << ":" << esc(e.attrs[i].second);
it << "}";
}
return it.str();
};
// One entry per referenced instance, sharing the entity's id, name and description.
std::ostringstream objs_js, regs_js;
bool ofirst = true, rfirst = true;
for (auto& e : ents) {
std::vector<int> inst;
for (auto& r : e.refs)
if (r.first == 'o' && r.second >= 0 && r.second < (int)objects.size()) inst.push_back(r.second);
if (inst.empty()) {
regs_js << (rfirst ? "" : ",") << common(e) << "}";
rfirst = false;
continue;
}
for (int k : inst) {
int bb[4];
to1000(objects[k], bb);
objs_js << (ofirst ? "" : ",") << common(e);
if (field("bbox")) objs_js << ",\"bbox\":[" << bb[0] << "," << bb[1] << "," << bb[2] << "," << bb[3] << "]";
if (field("score")) objs_js << ",\"score\":" << num(objects[k].score);
if (draw) objs_js << ",\"mask\":" << mask_json(objects[k], bb);
objs_js << "}";
ofirst = false;
}
}
std::ostringstream js;
std::vector<std::string> parts;
if (wants("scene")) {
std::string s = "\"scene\":{";
for (size_t i = 0; i < scene.size(); i++) s += (i ? "," : "") + esc(scene[i].first) + ":" + esc(scene[i].second);
parts.push_back(s + "}");
}
if (wants("objects")) parts.push_back("\"objects\":[" + objs_js.str() + "]");
if (wants("regions")) parts.push_back("\"regions\":[" + regs_js.str() + "]");
if (wants("text")) {
std::ostringstream ts;
bool first = true;
for (size_t k = 0; k < texts.size(); k++) {
if (ocr[k].empty()) continue;
int bb[4];
to1000(texts[k], bb);
bool mentioned = std::find(text_refs.begin(), text_refs.end(), (int)k) != text_refs.end();
ts << (first ? "" : ",") << "{\"string\":" << esc(ocr[k]);
if (field("bbox")) ts << ",\"bbox\":[" << bb[0] << "," << bb[1] << "," << bb[2] << "," << bb[3] << "]";
if (field("score")) ts << ",\"score\":" << num(texts[k].score);
if (need_decoder) ts << ",\"mentioned\":" << (mentioned ? "true" : "false");
ts << "}";
first = false;
}
parts.push_back("\"text\":[" + ts.str() + "]");
}
if (wants("relations")) {
std::string s = "\"relations\":[";
for (size_t i = 0; i < rel_json.size(); i++) s += (i ? "," : "") + rel_json[i];
parts.push_back(s + "]");
}
if (wants("events")) {
std::string s = "\"events\":[";
for (size_t i = 0; i < events.size(); i++) s += (i ? "," : "") + esc(events[i]);
parts.push_back(s + "]");
}
if (wants("reasoning") && o.reasoning) parts.push_back("\"reasoning\":" + esc(trace));
auto t4 = clk::now();
parts.push_back("\"timing_ms\":{\"encode\":" + num(ms(t0, t1), 1) + ",\"detect_read\":" + num(ms(t1, t2), 1) + ",\"decode\":" +
num(ms(t2, t3), 1) + ",\"assemble\":" + num(ms(t3, t4), 1) + ",\"total\":" + num(ms(t0, t4), 1) +
",\"tokens\":" + std::to_string(ids.size()) + "}");
js << "{";
for (size_t i = 0; i < parts.size(); i++) js << (i ? "," : "") << parts[i];
js << "}";
if (tm) { tm[0] = ms(t0, t1); tm[1] = ms(t1, t2); tm[2] = ms(t2, t3); tm[3] = ms(t0, t4); tm[4] = (double)ids.size(); }
return js.str();
}
// ------------------------------------------------------------------ options
static void parse_form(const JV& v, Options& o) {
auto bad = [](const std::string& what) { throw std::runtime_error("invalid form: " + what); };
if (v.t != JV::OBJ) bad("form must be an object of fields");
if (v.o.empty()) bad("form has no fields");
for (const auto& kv : v.o) {
const std::string& name = kv.first;
const JV& s = kv.second;
if (name.empty() || name.find_first_of("=;|#$@\n ") != std::string::npos) bad("field name " + esc(name) + " cannot be used");
if (s.t != JV::OBJ) bad("field " + name + " must be an object");
FormField f;
f.name = name;
const JV* en = s.get("enum");
const JV* xs = s.get("x-source");
const JV* ty = s.get("type");
bool done = false;
if (en) {
if (en->t != JV::ARR || en->a.empty()) bad("field " + name + " needs an enum list");
for (const JV& c : en->a) {
if (c.t == JV::STR) { f.choices.push_back(c.s); f.values.push_back(esc(c.s)); }
else if (c.t == JV::NUM) { f.choices.push_back(c.s); f.values.push_back(c.s); }
else bad("field " + name + " has a choice that is not a string or a number");
if (f.choices.back().empty() || f.choices.back().find_first_of("|\n") != std::string::npos)
bad("field " + name + " has a choice that cannot be used");
}
f.kind = FormField::CHOICE;
done = true;
} else if (xs) {
if (xs->t != JV::STR) bad("field " + name + " needs a string x-source");
if (xs->s == "bbox") {
const JV* of = s.get("x-of");
bool ok = false;
if (of && of->t == JV::STR)
for (const FormField& g : o.form) ok = ok || (g.name == of->s && g.kind == FormField::CHOICE);
if (!ok) bad("box field " + name + " must point at an earlier choice field");
f.kind = FormField::BOX;
f.of = of->s;
done = true;
}
}
if (!done && ty && ty->t == JV::STR && ty->s == "integer") {
const JV* lo = s.get("minimum");
const JV* hi = s.get("maximum");
if (!lo || !hi || lo->t != JV::NUM || hi->t != JV::NUM || lo->n != std::floor(lo->n) || hi->n != std::floor(hi->n) || lo->n > hi->n)
bad("field " + name + " needs an integer minimum and maximum");
if (hi->n - lo->n > 1000) bad("field " + name + " spans more than 1000 values");
f.kind = FormField::COUNT;
f.lo = (long)lo->n;
f.hi = (long)hi->n;
done = true;
}
if (!done && ty && ty->t == JV::STR && ty->s == "string") {
const JV* ml = s.get("maxLength");
if (ml) {
if (ml->t != JV::NUM || ml->n < 1 || ml->n > 500) bad("field " + name + " needs a maxLength from 1 to 500");
f.max_length = (int)ml->n;
}
f.kind = FormField::TEXT;
done = true;
}
if (!done) bad("field " + name + " is not a choice, count, box or text field");
o.form.push_back(f);
}
}
static void parse_options(const Model& m, const char* s, Options& o) {
if (!s || !*s) return;
JReader jr{s, s + std::strlen(s)};
JV root = jr.value();
jr.ws();
if (jr.p != jr.e) jr.fail("trailing characters after the JSON object");
if (root.t != JV::OBJ) jr.fail("options must be a JSON object");
static const char* SECTIONS[] = {"scene", "objects", "regions", "text", "relations", "events", "reasoning"};
static const char* FIELDS[] = {"bbox", "score", "mask", "position", "description", "count"};
for (const auto& kv : root.o) {
const std::string& k = kv.first;
const JV& v = kv.second;
auto number = [&](double lo, double hi) {
if (v.t != JV::NUM || !(v.n >= lo && v.n <= hi)) jr.fail(k + " must be a number from " + num(lo, 2) + " to " + num(hi, 2));
return v.n;
};
auto boolean = [&]() {
if (v.t != JV::BOOL) jr.fail(k + " must be true or false");
return v.b;
};
auto names = [&](const char* const* allowed, size_t n, std::vector<std::string>& out) {
if (v.t != JV::ARR) jr.fail(k + " must be a list of names");
for (const JV& x : v.a) {
bool ok = x.t == JV::STR && std::find(allowed, allowed + n, x.s) != allowed + n;
if (!ok) jr.fail(k + " has an unknown name " + (x.t == JV::STR ? x.s : std::string("(not a string)")));
out.push_back(x.s);
}
};
if (k == "language") {
if (v.t != JV::STR || std::find(m.languages.begin(), m.languages.end(), v.s) == m.languages.end()) {
std::string list;
for (const auto& l : m.languages) list += (list.empty() ? "" : ", ") + l;
jr.fail("unsupported language " + (v.t == JV::STR ? v.s : std::string("(not a string)")) + "; choose one of " + list);
}
o.language = v.s;
} else if (k == "keywords") {
std::vector<std::string> ks;
if (v.t == JV::STR) {
std::stringstream ss(v.s);
std::string line;
while (std::getline(ss, line)) ks.push_back(line);
} else if (v.t == JV::ARR) {
for (const JV& x : v.a) {
if (x.t != JV::STR) jr.fail("keywords must be strings");
ks.push_back(x.s);
}
} else {
jr.fail("keywords must be a list of strings or one newline-separated string");
}
for (auto& s : ks) {
s = trim(s);
if (s.empty()) continue;
if (s.size() > 100) jr.fail("a keyword is longer than 100 bytes");
o.keywords.push_back(s);
}
if (o.keywords.size() > 64) jr.fail("more than 64 keywords");
} else if (k == "conf") o.conf = (float)number(0, 1);
else if (k == "text_conf") o.text_conf = (float)number(0, 1);
else if (k == "mask_threshold") o.mask_threshold = (float)number(0, 1);
else if (k == "max_objects") o.max_objects = (int)number(0, 64);
else if (k == "max_text_lines") o.max_text_lines = (int)number(0, 32);
else if (k == "max_tokens") o.max_tokens = (int)number(1, 1536);
else if (k == "masks") o.masks = boolean();
else if (k == "nms") o.nms = boolean();
else if (k == "reasoning") o.reasoning = boolean();
else if (k == "mask_format") {
if (v.t != JV::STR || (v.s != "largest" && v.s != "polygons" && v.s != "merged" && v.s != "bitmap" && v.s != "none"))
jr.fail("mask_format must be largest, polygons, merged, bitmap or none");
o.mask_format = v.s;
} else if (k == "sections") names(SECTIONS, 7, o.sections);
else if (k == "fields") names(FIELDS, 6, o.fields);
else if (k == "form") parse_form(v, o);
}
}
} // namespace oasis
// ------------------------------------------------------------------ C API
static thread_local std::string g_error;
struct oasis_model {
oasis::Model m;
};
extern "C" {
oasis_model* oasis_load(const char* path) {
const char* ex = std::getenv("OASIS_EXACT");
oasis::g_exact_gemv = ex && ex[0] == '1';
auto* h = new oasis_model();
std::string err;
try {
if (!h->m.load(path, err)) {
g_error = err;
delete h;
return nullptr;
}
} catch (const std::exception& e) {
g_error = e.what();
delete h;
return nullptr;
}
return h;
}
void oasis_free(oasis_model* h) { delete h; }
int oasis_describe(oasis_model* h, const unsigned char* rgb, int width, int height, const char* options_json, char* out, int capacity) {
if (!h || !rgb || width <= 0 || height <= 0) { g_error = "invalid arguments"; return -1; }
try {
oasis::Options o;
oasis::parse_options(h->m, options_json, o);
std::string js = oasis::describe_json(h->m, rgb, width, height, o, nullptr);
if ((int)js.size() + 1 > capacity) { g_error = "output buffer too small"; return -(int)js.size() - 1; }
std::memcpy(out, js.c_str(), js.size() + 1);
return (int)js.size();
} catch (const std::exception& e) {
g_error = e.what();
return -1;
}
}
const char* oasis_last_error(void) { return g_error.c_str(); }
} // extern "C"