static_embeddings 0.1.1 → 0.1.3

data/ext/static_embeddings/se_tokenizer.c

@@ -4,7 +4,6 @@
 #include <stdlib.h>
 #include <string.h>
 
-#define SE_SIZE_MAX          ((size_t)-1)
 #define SE_CANCEL_CHECK_MASK 0x3ffu
 
 void se_scratch_init(se_scratch_t *s) {
@@ -12,70 +11,58 @@ void se_scratch_init(se_scratch_t *s) {
 }
 
 void se_scratch_free(se_scratch_t *s) {
-    free(s->cps);
-    free(s->cps2);
-    free(s->bytes);
-    free(s->ids);
-    free(s->acc);
+    se_free(s->cps);
+    se_free(s->cps2);
+    se_free(s->bytes);
+    se_free(s->ids);
+    se_free(s->acc);
     memset(s, 0, sizeof(*s));
 }
 
-static int grow_u32(uint32_t **buf, size_t *cap, size_t need) {
+static void *grow_buffer(void *buf, size_t *cap, size_t need, size_t elem_size,
+                         size_t min_capacity) {
     if (*cap >= need)
-        return 1;
+        return buf;
 
-    size_t next = *cap ? *cap : 256;
-    while (next < need) {
-        if (next > SE_SIZE_MAX / 2)
-            return 0;
-        next *= 2;
+    size_t next = 0;
+    size_t bytes = 0;
+    if (!se_next_capacity(*cap, need, min_capacity, &next) ||
+        !se_array_bytes(next, elem_size, &bytes)) {
+        return NULL;
     }
 
-    if (next > SE_SIZE_MAX / sizeof(uint32_t))
-        return 0;
+    void *p = se_realloc(SE_ALLOC_SCRATCH, buf, bytes);
+    if (!p)
+        return NULL;
 
-    uint32_t *p = (uint32_t *)realloc(*buf, next * sizeof(uint32_t));
+    *cap = next;
+    return p;
+}
+
+static int grow_u32(uint32_t **buf, size_t *cap, size_t need) {
+    uint32_t *p = (uint32_t *)grow_buffer(*buf, cap, need, sizeof(uint32_t), 256);
     if (!p)
         return 0;
 
     *buf = p;
-    *cap = next;
     return 1;
 }
 
 static int grow_bytes(uint8_t **buf, size_t *cap, size_t need) {
-    if (*cap >= need)
-        return 1;
-
-    size_t next = *cap ? *cap : 256;
-    while (next < need) {
-        if (next > SE_SIZE_MAX / 2)
-            return 0;
-        next *= 2;
-    }
-
-    uint8_t *p = (uint8_t *)realloc(*buf, next);
+    uint8_t *p = (uint8_t *)grow_buffer(*buf, cap, need, sizeof(uint8_t), 256);
     if (!p)
         return 0;
 
     *buf = p;
-    *cap = next;
     return 1;
 }
 
 static int grow_float(float **buf, size_t *cap, size_t need) {
-    if (*cap >= need)
-        return 1;
-
-    if (need > SE_SIZE_MAX / sizeof(float))
-        return 0;
-
-    float *p = (float *)realloc(*buf, need * sizeof(float));
+    float *p = (float *)grow_buffer(*buf, cap, need, sizeof(float), need);
     if (!p)
         return 0;
 
     *buf = p;
-    *cap = need;
     return 1;
 }
 
@@ -111,7 +98,7 @@ static int is_whitespace(const se_model_t *m, uint32_t cp) {
 
 static int is_control(const se_model_t *m, uint32_t cp) {
     if (cp < 0x80)
-        return cp < 0x20 && cp != '\t' && cp != '\n' && cp != '\r';
+        return (cp < 0x20 || cp == 0x7f) && cp != '\t' && cp != '\n' && cp != '\r';
     return se_range_contains(m->norm.control, m->norm.control_count, cp);
 }
 
@@ -223,6 +210,22 @@ size_t se_prefix_boundary_len(const se_model_t *model, const uint8_t *input, siz
     return 0;
 }
 
+typedef enum {
+    SE_ASCII_DROP = 0,
+    SE_ASCII_SPACE = 1,
+    SE_ASCII_PUNCT = 2,
+    SE_ASCII_WORD = 3
+} se_ascii_class_t;
+
+static const uint8_t SE_ASCII_CLASS[128] = {
+    0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2,
+    2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2,
+    2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 0,
+};
+
+#define SE_ASCII_CANCEL_STRIDE (SE_CANCEL_CHECK_MASK + 1u)
+
 typedef struct {
     const se_model_t *model;
     se_scratch_t *scratch;
@@ -282,9 +285,11 @@ static wordpiece_status_t wordpiece(const se_model_t *m, se_scratch_t *sc, const
         return 1;
     }
 
-    if (word_len > (SE_SIZE_MAX / 4) - 1)
+    size_t bytes_need = 0;
+    if (!se_checked_mul_size(word_len, 4, &bytes_need) ||
+        !se_checked_add_size(bytes_need, 4, &bytes_need))
         return 0;
-    if (!grow_bytes(&sc->bytes, &sc->bytes_cap, word_len * 4 + 4))
+    if (!grow_bytes(&sc->bytes, &sc->bytes_cap, bytes_need))
         return 0;
     if (!grow_u32(&sc->cps2, &sc->cps2_cap, word_len + 1))
         return 0;
@@ -452,6 +457,93 @@ static se_status_t emit_cleaned(token_state_t *st, uint32_t cp, int *stop) {
     return emit_stripped(st, cp, stop);
 }
 
+static inline se_ascii_class_t ascii_class(uint8_t byte) {
+    return (se_ascii_class_t)SE_ASCII_CLASS[byte];
+}
+
+static size_t ascii_word_reserve_len(const uint8_t *input, size_t input_len, size_t i) {
+    size_t run = 1;
+    size_t limit = input_len - i;
+
+    if (limit > SE_ASCII_CANCEL_STRIDE)
+        limit = SE_ASCII_CANCEL_STRIDE;
+
+    while (run < limit && input[i + run] < 0x80 && ascii_class(input[i + run]) == SE_ASCII_WORD)
+        run++;
+
+    return run;
+}
+
+static se_status_t tokenize_ascii_run(token_state_t *st, const uint8_t *input, size_t input_len,
+                                      size_t *ip, int *stop) {
+    size_t i = *ip;
+    size_t next_cancel_check = i + SE_ASCII_CANCEL_STRIDE;
+    const int lc = (int)st->model->meta.do_lower_case;
+    se_scratch_t *sc = st->scratch;
+
+    while (i < input_len) {
+        uint8_t b = input[i];
+        if (b >= 0x80)
+            break;
+
+        if (i >= next_cancel_check) {
+            if (token_cancelled(st)) {
+                *ip = i;
+                se_error_set(st->err, SE_ERR_INTERNAL, "operation cancelled");
+                return SE_ERR_INTERNAL;
+            }
+            next_cancel_check = i + SE_ASCII_CANCEL_STRIDE;
+        }
+
+        se_ascii_class_t cls = ascii_class(b);
+
+        if (cls == SE_ASCII_WORD) {
+            if (st->segment_len >= sc->cps_cap) {
+                if (!grow_u32(&sc->cps, &sc->cps_cap,
+                              st->segment_len + ascii_word_reserve_len(input, input_len, i))) {
+                    *ip = i;
+                    return oom(st, "building token segment");
+                }
+            }
+            sc->cps[st->segment_len++] =
+                (lc && b >= 'A' && b <= 'Z') ? (uint32_t)(b + 32u) : (uint32_t)b;
+            i++;
+            continue;
+        }
+
+        if (cls == SE_ASCII_DROP) {
+            i++;
+            continue;
+        }
+
+        se_status_t rc = flush_segment(st, stop);
+        if (rc != SE_OK) {
+            *ip = i;
+            return rc;
+        }
+        if (*stop) {
+            *ip = i;
+            return SE_OK;
+        }
+
+        if (cls == SE_ASCII_PUNCT) {
+            uint32_t cp = b;
+            rc = append_wordpiece(st, &cp, 1, stop);
+            if (rc != SE_OK) {
+                *ip = i;
+                return rc;
+            }
+        }
+
+        i++;
+        if (*stop)
+            break;
+    }
+
+    *ip = i;
+    return SE_OK;
+}
+
 se_status_t se_tokenize(const se_model_t *model, se_scratch_t *sc, const uint8_t *input,
                         size_t input_len, uint32_t max_tokens, se_token_stats_t *stats,
                         se_error_t *err, volatile sig_atomic_t *cancelled) {
@@ -474,13 +566,24 @@ se_status_t se_tokenize(const se_model_t *model, se_scratch_t *sc, const uint8_t
             return SE_ERR_INTERNAL;
         }
 
+        int stop = 0;
+
+        if (input[i] < 0x80) {
+            se_status_t frc = tokenize_ascii_run(&st, input, input_len, &i, &stop);
+            if (frc != SE_OK)
+                return frc;
+            if (stop)
+                break;
+            if (i >= input_len)
+                break;
+        }
+
         uint32_t cp;
         if (!decode_one(input, input_len, &i, &cp)) {
             se_error_set(err, SE_ERR_INVALID_UTF8, "input is not valid UTF-8");
             return SE_ERR_INVALID_UTF8;
         }
 
-        int stop = 0;
         se_status_t rc = emit_cleaned(&st, cp, &stop);
         if (rc != SE_OK)
             return rc;

data/ext/static_embeddings/se_topk.c

@@ -0,0 +1,236 @@
+#include "se_internal.h"
+
+#include <math.h>
+#include <stdint.h>
+
+#if defined(__ARM_NEON) || defined(__ARM_NEON__)
+#include <arm_neon.h>
+#define SE_HAVE_NEON 1
+#elif defined(__SSE__)
+#include <xmmintrin.h>
+#define SE_HAVE_SSE 1
+#endif
+
+float se_dot_product_f32(const float *q, const float *row, size_t dim) {
+    size_t j = 0;
+#if defined(SE_HAVE_NEON)
+    float32x4_t a0 = vdupq_n_f32(0.0f);
+    float32x4_t a1 = vdupq_n_f32(0.0f);
+    float32x4_t a2 = vdupq_n_f32(0.0f);
+    float32x4_t a3 = vdupq_n_f32(0.0f);
+    for (; j + 15 < dim; j += 16) {
+        a0 = vmlaq_f32(a0, vld1q_f32(q + j), vld1q_f32(row + j));
+        a1 = vmlaq_f32(a1, vld1q_f32(q + j + 4), vld1q_f32(row + j + 4));
+        a2 = vmlaq_f32(a2, vld1q_f32(q + j + 8), vld1q_f32(row + j + 8));
+        a3 = vmlaq_f32(a3, vld1q_f32(q + j + 12), vld1q_f32(row + j + 12));
+    }
+    float32x4_t sumv = vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3));
+#if defined(__aarch64__)
+    float dot = vaddvq_f32(sumv);
+#else
+    float32x2_t pair = vadd_f32(vget_low_f32(sumv), vget_high_f32(sumv));
+    pair = vpadd_f32(pair, pair);
+    float dot = vget_lane_f32(pair, 0);
+#endif
+#elif defined(SE_HAVE_SSE)
+    __m128 a0 = _mm_setzero_ps();
+    __m128 a1 = _mm_setzero_ps();
+    __m128 a2 = _mm_setzero_ps();
+    __m128 a3 = _mm_setzero_ps();
+    for (; j + 15 < dim; j += 16) {
+        a0 = _mm_add_ps(a0, _mm_mul_ps(_mm_loadu_ps(q + j), _mm_loadu_ps(row + j)));
+        a1 = _mm_add_ps(a1, _mm_mul_ps(_mm_loadu_ps(q + j + 4), _mm_loadu_ps(row + j + 4)));
+        a2 = _mm_add_ps(a2, _mm_mul_ps(_mm_loadu_ps(q + j + 8), _mm_loadu_ps(row + j + 8)));
+        a3 = _mm_add_ps(a3, _mm_mul_ps(_mm_loadu_ps(q + j + 12), _mm_loadu_ps(row + j + 12)));
+    }
+    __m128 sumv = _mm_add_ps(_mm_add_ps(a0, a1), _mm_add_ps(a2, a3));
+    float tmp[4];
+    _mm_storeu_ps(tmp, sumv);
+    float dot = (tmp[0] + tmp[1]) + (tmp[2] + tmp[3]);
+#else
+    float s0 = 0.0f;
+    float s1 = 0.0f;
+    float s2 = 0.0f;
+    float s3 = 0.0f;
+    float s4 = 0.0f;
+    float s5 = 0.0f;
+    float s6 = 0.0f;
+    float s7 = 0.0f;
+    for (; j + 7 < dim; j += 8) {
+        s0 += q[j] * row[j];
+        s1 += q[j + 1] * row[j + 1];
+        s2 += q[j + 2] * row[j + 2];
+        s3 += q[j + 3] * row[j + 3];
+        s4 += q[j + 4] * row[j + 4];
+        s5 += q[j + 5] * row[j + 5];
+        s6 += q[j + 6] * row[j + 6];
+        s7 += q[j + 7] * row[j + 7];
+    }
+    float dot = (s0 + s1) + (s2 + s3) + (s4 + s5) + (s6 + s7);
+#endif
+    for (; j < dim; j++)
+        dot += q[j] * row[j];
+    return dot;
+}
+
+float se_dot_and_row_sq_f32(const float *q, const float *row, size_t dim,
+                            float *row_sq_out) {
+    size_t j = 0;
+#if defined(SE_HAVE_NEON)
+    float32x4_t d0 = vdupq_n_f32(0.0f);
+    float32x4_t d1 = vdupq_n_f32(0.0f);
+    float32x4_t d2 = vdupq_n_f32(0.0f);
+    float32x4_t d3 = vdupq_n_f32(0.0f);
+    float32x4_t s0 = vdupq_n_f32(0.0f);
+    float32x4_t s1 = vdupq_n_f32(0.0f);
+    float32x4_t s2 = vdupq_n_f32(0.0f);
+    float32x4_t s3 = vdupq_n_f32(0.0f);
+    for (; j + 15 < dim; j += 16) {
+        float32x4_t q0 = vld1q_f32(q + j);
+        float32x4_t r0 = vld1q_f32(row + j);
+        float32x4_t q1 = vld1q_f32(q + j + 4);
+        float32x4_t r1 = vld1q_f32(row + j + 4);
+        float32x4_t q2 = vld1q_f32(q + j + 8);
+        float32x4_t r2 = vld1q_f32(row + j + 8);
+        float32x4_t q3 = vld1q_f32(q + j + 12);
+        float32x4_t r3 = vld1q_f32(row + j + 12);
+        d0 = vmlaq_f32(d0, q0, r0);
+        d1 = vmlaq_f32(d1, q1, r1);
+        d2 = vmlaq_f32(d2, q2, r2);
+        d3 = vmlaq_f32(d3, q3, r3);
+        s0 = vmlaq_f32(s0, r0, r0);
+        s1 = vmlaq_f32(s1, r1, r1);
+        s2 = vmlaq_f32(s2, r2, r2);
+        s3 = vmlaq_f32(s3, r3, r3);
+    }
+    float32x4_t dotv = vaddq_f32(vaddq_f32(d0, d1), vaddq_f32(d2, d3));
+    float32x4_t sqv = vaddq_f32(vaddq_f32(s0, s1), vaddq_f32(s2, s3));
+#if defined(__aarch64__)
+    float dot = vaddvq_f32(dotv);
+    float row_sq = vaddvq_f32(sqv);
+#else
+    float32x2_t pair = vadd_f32(vget_low_f32(dotv), vget_high_f32(dotv));
+    pair = vpadd_f32(pair, pair);
+    float dot = vget_lane_f32(pair, 0);
+    pair = vadd_f32(vget_low_f32(sqv), vget_high_f32(sqv));
+    pair = vpadd_f32(pair, pair);
+    float row_sq = vget_lane_f32(pair, 0);
+#endif
+#elif defined(SE_HAVE_SSE)
+    __m128 d0 = _mm_setzero_ps();
+    __m128 d1 = _mm_setzero_ps();
+    __m128 d2 = _mm_setzero_ps();
+    __m128 d3 = _mm_setzero_ps();
+    __m128 s0 = _mm_setzero_ps();
+    __m128 s1 = _mm_setzero_ps();
+    __m128 s2 = _mm_setzero_ps();
+    __m128 s3 = _mm_setzero_ps();
+    for (; j + 15 < dim; j += 16) {
+        __m128 q0 = _mm_loadu_ps(q + j);
+        __m128 r0 = _mm_loadu_ps(row + j);
+        __m128 q1 = _mm_loadu_ps(q + j + 4);
+        __m128 r1 = _mm_loadu_ps(row + j + 4);
+        __m128 q2 = _mm_loadu_ps(q + j + 8);
+        __m128 r2 = _mm_loadu_ps(row + j + 8);
+        __m128 q3 = _mm_loadu_ps(q + j + 12);
+        __m128 r3 = _mm_loadu_ps(row + j + 12);
+        d0 = _mm_add_ps(d0, _mm_mul_ps(q0, r0));
+        d1 = _mm_add_ps(d1, _mm_mul_ps(q1, r1));
+        d2 = _mm_add_ps(d2, _mm_mul_ps(q2, r2));
+        d3 = _mm_add_ps(d3, _mm_mul_ps(q3, r3));
+        s0 = _mm_add_ps(s0, _mm_mul_ps(r0, r0));
+        s1 = _mm_add_ps(s1, _mm_mul_ps(r1, r1));
+        s2 = _mm_add_ps(s2, _mm_mul_ps(r2, r2));
+        s3 = _mm_add_ps(s3, _mm_mul_ps(r3, r3));
+    }
+    __m128 dotv = _mm_add_ps(_mm_add_ps(d0, d1), _mm_add_ps(d2, d3));
+    __m128 sqv = _mm_add_ps(_mm_add_ps(s0, s1), _mm_add_ps(s2, s3));
+    float tmp[4];
+    _mm_storeu_ps(tmp, dotv);
+    float dot = (tmp[0] + tmp[1]) + (tmp[2] + tmp[3]);
+    _mm_storeu_ps(tmp, sqv);
+    float row_sq = (tmp[0] + tmp[1]) + (tmp[2] + tmp[3]);
+#else
+    float d0 = 0.0f;
+    float d1 = 0.0f;
+    float d2 = 0.0f;
+    float d3 = 0.0f;
+    float s0 = 0.0f;
+    float s1 = 0.0f;
+    float s2 = 0.0f;
+    float s3 = 0.0f;
+    for (; j + 3 < dim; j += 4) {
+        float r0 = row[j];
+        float r1 = row[j + 1];
+        float r2 = row[j + 2];
+        float r3 = row[j + 3];
+        d0 += q[j] * r0;
+        d1 += q[j + 1] * r1;
+        d2 += q[j + 2] * r2;
+        d3 += q[j + 3] * r3;
+        s0 += r0 * r0;
+        s1 += r1 * r1;
+        s2 += r2 * r2;
+        s3 += r3 * r3;
+    }
+    float dot = (d0 + d1) + (d2 + d3);
+    float row_sq = (s0 + s1) + (s2 + s3);
+#endif
+    for (; j < dim; j++) {
+        float r = row[j];
+        dot += q[j] * r;
+        row_sq += r * r;
+    }
+    *row_sq_out = row_sq;
+    return dot;
+}
+
+static float cosine_score(float dot, float row_sq, float inv_query_norm) {
+    if (row_sq > 0.0f)
+        return dot * inv_query_norm / sqrtf(row_sq);
+    if (row_sq == 0.0f)
+        return 0.0f;
+    return NAN;
+}
+
+void *se_topk_execute(void *arg) {
+    se_topk_job_t *job = (se_topk_job_t *)arg;
+    for (size_t r = 0; r < job->rows; r++) {
+        if ((r & 1023u) == 0 && job->cancelled)
+            return NULL;
+        float score;
+
+        if (job->format == SE_VECTOR_FORMAT_F16) {
+            const uint8_t *row = (const uint8_t *)job->m + r * job->dim * 2u;
+            if (job->cosine) {
+                float row_sq = 0.0f;
+                score = se_dot_and_row_sq_f16(job->q, row, job->dim, &row_sq);
+                score = cosine_score(score, row_sq, job->inv_query_norm);
+            } else {
+                score = se_dot_product_f16(job->q, row, job->dim);
+            }
+        } else {
+            const float *row = (const float *)job->m + r * job->dim;
+            if (job->cosine) {
+                float row_sq = 0.0f;
+                score = se_dot_and_row_sq_f32(job->q, row, job->dim, &row_sq);
+                score = cosine_score(score, row_sq, job->inv_query_norm);
+            } else {
+                score = se_dot_product_f32(job->q, row, job->dim);
+            }
+        }
+
+        if (!(score > job->best_score[job->k - 1]))
+            continue;
+
+        long pos = job->k - 1;
+        while (pos > 0 && job->best_score[pos - 1] < score) {
+            job->best_score[pos] = job->best_score[pos - 1];
+            job->best_idx[pos] = job->best_idx[pos - 1];
+            pos--;
+        }
+        job->best_score[pos] = score;
+        job->best_idx[pos] = r;
+    }
+    return NULL;
+}