semq-sys 1.1.0

/*
 * Copyright (c) 2026 The SEMQ Group Inc.
 * Licensed under the PolyForm Noncommercial License 1.0.0. See LICENSE.md for terms.
 */

/*
 * semq_neon.c — NEON backend (ARM64).
 *
 * Implements the SEMQ orbit operator using NEON intrinsics. The
 * inner loop processes 16 input floats per iteration. The `x · scale`
 * step is performed in float64 to match the scalar reference exactly
 * (vcvt_f64_f32 + vmulq_f64). The remaining integer math is done in
 * 32-bit lanes after a saturating narrow from int64.
 *
 * Equivalence band:
 *
 *   This backend produces bit-identical output to the scalar backend
 *   for every input satisfying  |round(x · scale)| ≤ INT32_MAX.
 *   For canonical use (cosine-normalised embeddings, scale = 100)
 *   this band has more than seven orders of magnitude of headroom.
 *   Inputs beyond the band saturate at INT32 boundaries here, while
 *   the scalar backend saturates at INT64 boundaries — outputs may
 *   differ for those extreme values.
 *
 * Tail:
 *
 *   For dim that is not a multiple of 16, the trailing 0..15 lanes
 *   fall through to the scalar fallback at the bottom of this file.
 *   The fallback duplicates the scalar reference (no cross-file
 *   coupling) so each backend is self-contained.
 */

#if defined(__aarch64__) || defined(_M_ARM64)

#include "semq_dispatch.h"

#include <arm_neon.h>
#include <math.h>
#include <stdint.h>

/* -------------------------------------------------------------------------- */
/*  Scalar fallback (per-element)                                             */
/* -------------------------------------------------------------------------- */

static int64_t scaled_round_one(float x, uint32_t scale) {
    if (!isfinite(x)) {
        return 0;
    }
    /* Ties-to-even — matches the vector path's FCVTNS exactly. */
    double y = nearbyint((double)x * (double)scale);
    if (y >= 9.2233720368547748e18) {
        return INT64_MAX;
    }
    if (y <= -9.2233720368547758e18) {
        return INT64_MIN;
    }
    return (int64_t)y;
}

static uint8_t reduce_one(uint64_t a) {
    if (a == 0u) return 0u;
    return (uint8_t)(1u + (uint32_t)((a - 1u) % 9u));
}

static semq_code_t scalar_op(float x, uint32_t scale) {
    int64_t      v = scaled_round_one(x, scale);
    unsigned int s = (v > 0) ? 2u : ((v < 0) ? 0u : 1u);
    uint64_t a;
    if (v == INT64_MIN) {
        a = (uint64_t)INT64_MAX + 1u;
    } else {
        a = (v < 0) ? (uint64_t)(-v) : (uint64_t)v;
    }
    return (semq_code_t)(s * 10u + (unsigned int)reduce_one(a));
}

static double scalar_inverse(semq_code_t code, uint32_t scale) {
    int          sign = (int)(code / 10u) - 1;
    unsigned int g    = (unsigned int)(code % 10u);
    return ((double)g * (double)sign) / (double)scale;
}

/* -------------------------------------------------------------------------- */
/*  Internal helpers                                                          */
/* -------------------------------------------------------------------------- */

/* mod_9_u32(a) = a mod 9, computed branchlessly via magic multiplication. */
static inline uint32x4_t mod_9_u32(uint32x4_t a) {
    /* q = (a * 0x38E38E39) >> 33 ; q = floor(a / 9) for all a in [0, 2^32) */
    uint32x4_t magic = vdupq_n_u32(0x38E38E39u);
    uint64x2_t lo = vmull_u32(vget_low_u32(a),  vget_low_u32(magic));
    uint64x2_t hi = vmull_high_u32(a,           magic);
    uint32x2_t q_lo = vmovn_u64(vshrq_n_u64(lo, 33));
    uint32x2_t q_hi = vmovn_u64(vshrq_n_u64(hi, 33));
    uint32x4_t q    = vcombine_u32(q_lo, q_hi);
    /* r = a - 9*q */
    return vmlsq_u32(a, q, vdupq_n_u32(9u));
}

/* Map sixteen int32 lanes to the SEMQ orbit operator output, returning
 * sixteen uint8 codes packed into a single 128-bit vector.                 */
static inline uint8x16_t encode_int32_lanes(
        int32x4_t v0, int32x4_t v1, int32x4_t v2, int32x4_t v3) {

    int32x4_t Z = vdupq_n_s32(0);

    /* sign codes ∈ {0,1,2}: s = 1 + (v>0) - (v<0). */
    uint32x4_t pos0 = vcgtq_s32(v0, Z), neg0 = vcltq_s32(v0, Z);
    uint32x4_t pos1 = vcgtq_s32(v1, Z), neg1 = vcltq_s32(v1, Z);
    uint32x4_t pos2 = vcgtq_s32(v2, Z), neg2 = vcltq_s32(v2, Z);
    uint32x4_t pos3 = vcgtq_s32(v3, Z), neg3 = vcltq_s32(v3, Z);
    int32x4_t sgn0 = vsubq_s32(vreinterpretq_s32_u32(neg0),
                                vreinterpretq_s32_u32(pos0));
    int32x4_t sgn1 = vsubq_s32(vreinterpretq_s32_u32(neg1),
                                vreinterpretq_s32_u32(pos1));
    int32x4_t sgn2 = vsubq_s32(vreinterpretq_s32_u32(neg2),
                                vreinterpretq_s32_u32(pos2));
    int32x4_t sgn3 = vsubq_s32(vreinterpretq_s32_u32(neg3),
                                vreinterpretq_s32_u32(pos3));
    uint32x4_t s0 = vreinterpretq_u32_s32(vaddq_s32(vdupq_n_s32(1), sgn0));
    uint32x4_t s1 = vreinterpretq_u32_s32(vaddq_s32(vdupq_n_s32(1), sgn1));
    uint32x4_t s2 = vreinterpretq_u32_s32(vaddq_s32(vdupq_n_s32(1), sgn2));
    uint32x4_t s3 = vreinterpretq_u32_s32(vaddq_s32(vdupq_n_s32(1), sgn3));

    /* a = |v|, then  g  = (a == 0) ? 0 : 1 + ((a-1) mod 9). */
    uint32x4_t a0 = vreinterpretq_u32_s32(vabsq_s32(v0));
    uint32x4_t a1 = vreinterpretq_u32_s32(vabsq_s32(v1));
    uint32x4_t a2 = vreinterpretq_u32_s32(vabsq_s32(v2));
    uint32x4_t a3 = vreinterpretq_u32_s32(vabsq_s32(v3));

    uint32x4_t am1_0 = vsubq_u32(a0, vdupq_n_u32(1u));
    uint32x4_t am1_1 = vsubq_u32(a1, vdupq_n_u32(1u));
    uint32x4_t am1_2 = vsubq_u32(a2, vdupq_n_u32(1u));
    uint32x4_t am1_3 = vsubq_u32(a3, vdupq_n_u32(1u));

    uint32x4_t r0 = mod_9_u32(am1_0);
    uint32x4_t r1 = mod_9_u32(am1_1);
    uint32x4_t r2 = mod_9_u32(am1_2);
    uint32x4_t r3 = mod_9_u32(am1_3);

    uint32x4_t gp0 = vaddq_u32(r0, vdupq_n_u32(1u));
    uint32x4_t gp1 = vaddq_u32(r1, vdupq_n_u32(1u));
    uint32x4_t gp2 = vaddq_u32(r2, vdupq_n_u32(1u));
    uint32x4_t gp3 = vaddq_u32(r3, vdupq_n_u32(1u));

    uint32x4_t Zu = vdupq_n_u32(0u);
    uint32x4_t nz0 = vmvnq_u32(vceqq_u32(a0, Zu));
    uint32x4_t nz1 = vmvnq_u32(vceqq_u32(a1, Zu));
    uint32x4_t nz2 = vmvnq_u32(vceqq_u32(a2, Zu));
    uint32x4_t nz3 = vmvnq_u32(vceqq_u32(a3, Zu));
    uint32x4_t g0 = vandq_u32(gp0, nz0);
    uint32x4_t g1 = vandq_u32(gp1, nz1);
    uint32x4_t g2 = vandq_u32(gp2, nz2);
    uint32x4_t g3 = vandq_u32(gp3, nz3);

    /* code = s · 10 + g ; pack 4 × uint32x4 → uint8x16. */
    uint32x4_t TEN = vdupq_n_u32(10u);
    uint32x4_t c0 = vmlaq_u32(g0, s0, TEN);
    uint32x4_t c1 = vmlaq_u32(g1, s1, TEN);
    uint32x4_t c2 = vmlaq_u32(g2, s2, TEN);
    uint32x4_t c3 = vmlaq_u32(g3, s3, TEN);

    uint16x4_t c0_16 = vmovn_u32(c0);
    uint16x4_t c1_16 = vmovn_u32(c1);
    uint16x4_t c2_16 = vmovn_u32(c2);
    uint16x4_t c3_16 = vmovn_u32(c3);
    uint16x8_t c01   = vcombine_u16(c0_16, c1_16);
    uint16x8_t c23   = vcombine_u16(c2_16, c3_16);
    uint8x8_t  p01   = vmovn_u16(c01);
    uint8x8_t  p23   = vmovn_u16(c23);
    return vcombine_u8(p01, p23);
}

/* Convert four float32 → four int32 with float64 multiplication and
 * saturating narrow. Bit-identical to the scalar path within the band. */
static inline int32x4_t scale_round_lane(float32x4_t f, float64x2_t dsc) {
    float64x2_t d_lo = vmulq_f64(vcvt_f64_f32(vget_low_f32(f)),  dsc);
    float64x2_t d_hi = vmulq_f64(vcvt_high_f64_f32(f),           dsc);
    int64x2_t   v_lo = vcvtnq_s64_f64(d_lo);
    int64x2_t   v_hi = vcvtnq_s64_f64(d_hi);
    return vcombine_s32(vqmovn_s64(v_lo), vqmovn_s64(v_hi));
}

/* -------------------------------------------------------------------------- */
/*  Backend entry points                                                      */
/* -------------------------------------------------------------------------- */

static semq_status_t neon_encode(
    const float* input,
    uint32_t     dim,
    uint32_t     scale,
    semq_code_t* output) {

    float64x2_t dsc = vdupq_n_f64((double)scale);
    uint32_t i = 0;

    for (; i + 16u <= dim; i += 16u) {
        float32x4_t f0 = vld1q_f32(input + i +  0u);
        float32x4_t f1 = vld1q_f32(input + i +  4u);
        float32x4_t f2 = vld1q_f32(input + i +  8u);
        float32x4_t f3 = vld1q_f32(input + i + 12u);

        int32x4_t v0 = scale_round_lane(f0, dsc);
        int32x4_t v1 = scale_round_lane(f1, dsc);
        int32x4_t v2 = scale_round_lane(f2, dsc);
        int32x4_t v3 = scale_round_lane(f3, dsc);

        vst1q_u8(output + i, encode_int32_lanes(v0, v1, v2, v3));
    }
    for (; i < dim; ++i) {
        output[i] = scalar_op(input[i], scale);
    }
    return SEMQ_OK;
}

static semq_status_t neon_reconstruct(
    const semq_code_t* codes,
    uint32_t           dim,
    uint32_t           scale,
    double*            output) {

    /* Per-lane: x_hat = (g · sign) / scale, where
     *   sign = (code ÷ 10) − 1 ∈ {−1, 0, 1}
     *   g    =  code mod 10   ∈ {0..9}
     * For codes in [0, 29], (c · 26) >> 8 = c ÷ 10 (Hacker's Delight). */

    float64x2_t dsc = vdupq_n_f64((double)scale);
    uint32_t i = 0u;

    for (; i + 8u <= dim; i += 8u) {
        uint8x8_t   c8  = vld1_u8(codes + i);
        uint16x8_t  c16 = vmovl_u8(c8);
        uint16x8_t  q   = vshrq_n_u16(vmulq_u16(c16, vdupq_n_u16(26u)), 8);
        uint16x8_t  r   = vsubq_u16(c16, vmulq_u16(q, vdupq_n_u16(10u)));
        int16x8_t   s16 = vsubq_s16(vreinterpretq_s16_u16(q), vdupq_n_s16(1));
        int16x8_t   g16 = vreinterpretq_s16_u16(r);

        int32x4_t s32_lo = vmovl_s16(vget_low_s16(s16));
        int32x4_t s32_hi = vmovl_s16(vget_high_s16(s16));
        int32x4_t g32_lo = vmovl_s16(vget_low_s16(g16));
        int32x4_t g32_hi = vmovl_s16(vget_high_s16(g16));

        int32x4_t gs_lo = vmulq_s32(g32_lo, s32_lo);
        int32x4_t gs_hi = vmulq_s32(g32_hi, s32_hi);

        int64x2_t l0 = vmovl_s32(vget_low_s32(gs_lo));
        int64x2_t l1 = vmovl_high_s32(gs_lo);
        int64x2_t l2 = vmovl_s32(vget_low_s32(gs_hi));
        int64x2_t l3 = vmovl_high_s32(gs_hi);

        float64x2_t f0 = vdivq_f64(vcvtq_f64_s64(l0), dsc);
        float64x2_t f1 = vdivq_f64(vcvtq_f64_s64(l1), dsc);
        float64x2_t f2 = vdivq_f64(vcvtq_f64_s64(l2), dsc);
        float64x2_t f3 = vdivq_f64(vcvtq_f64_s64(l3), dsc);

        vst1q_f64(output + i + 0, f0);
        vst1q_f64(output + i + 2, f1);
        vst1q_f64(output + i + 4, f2);
        vst1q_f64(output + i + 6, f3);
    }
    for (; i < dim; ++i) {
        output[i] = scalar_inverse(codes[i], scale);
    }
    return SEMQ_OK;
}

/* -------------------------------------------------------------------------- */
/*  Backend descriptor                                                        */
/* -------------------------------------------------------------------------- */

const semq_backend_t SEMQ_NEON_BACKEND = {
    .encode       = neon_encode,
    .reconstruct  = neon_reconstruct
};

#endif /* __aarch64__ */