usearch 2.21.4 → 2.25.1

package/numkong/include/numkong/cast/sapphire.h

@@ -0,0 +1,262 @@
+/**
+ *  @brief SIMD-accelerated Type Conversions for Sapphire Rapids.
+ *  @file include/numkong/cast/sapphire.h
+ *  @author Ash Vardanian
+ *  @date January 2, 2026
+ *
+ *  @section sapphire_cast_instructions Relevant Instructions
+ *
+ *      Intrinsic                 Instruction                   Sapphire   Genoa
+ *      _mm_cvtss_sh              VCVTSS2SH (XMM, XMM, XMM)     5cy @ p05  5cy @ p01
+ *      _mm_cvtsh_ss              VCVTSH2SS (XMM, XMM, XMM)     5cy @ p05  5cy @ p01
+ *      _mm256_cvtepu8_epi16      VPMOVZXBW (YMM, XMM)          3cy @ p5   3cy @ p12
+ *      _mm256_mul_ph             VMULPH (YMM, YMM, YMM)        4cy @ p05  3cy @ p01
+ *      _mm256_cvtepi16_ph        VCVTW2PH (YMM, YMM)           4cy @ p05  4cy @ p01
+ *      _mm256_cvtph_epi16        VCVTPH2W (YMM, YMM)           4cy @ p05  4cy @ p01
+ *      _mm256_mask_blend_epi16   VPBLENDMW (YMM, K, YMM, YMM)  1cy @ p05  1cy @ p0123
+ *      _mm256_testn_epi16_mask   VPTESTNMW (K, YMM, YMM)       3cy @ p5   3cy @ p0
+ *      _mm256_cvtepi16_epi8      VPMOVWB (XMM, YMM)            4cy @ p5   4cy @ p12
+ *      _mm_maskz_loadu_epi8      VMOVDQU8 (XMM {K}, M128)      7cy @ p23  7cy @ p23
+ *      _mm256_mask_storeu_epi16  VMOVDQU16 (M256 {K}, YMM)     4cy @ p4   4cy @ p4
+ */
+#ifndef NK_CAST_SAPPHIRE_H
+#define NK_CAST_SAPPHIRE_H
+
+#if NK_TARGET_X8664_
+#if NK_TARGET_SAPPHIRE
+
+#include "numkong/types.h"
+#include "numkong/cast/icelake.h" // `nk_cast_icelake`
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if defined(__clang__)
+#pragma clang attribute push(__attribute__((target("avx2,avx512f,avx512vl,avx512bw,avx512fp16,f16c,fma,bmi,bmi2"))), \
+                             apply_to = function)
+#elif defined(__GNUC__)
+#pragma GCC push_options
+#pragma GCC target("avx2", "avx512f", "avx512vl", "avx512bw", "avx512fp16", "f16c", "fma", "bmi", "bmi2")
+#endif
+
+NK_PUBLIC void nk_f32_to_f16_sapphire(nk_f32_t const *from, nk_f16_t *to) {
+    *(nk_u16_t *)to = (nk_u16_t)_mm_cvtsi128_si32(_mm_castph_si128(_mm_cvtss_sh(_mm_setzero_ph(), _mm_set_ss(*from))));
+}
+
+NK_PUBLIC void nk_f16_to_f32_sapphire(nk_f16_t const *from, nk_f32_t *to) {
+    *to = _mm_cvtss_f32(_mm_cvtsh_ss(_mm_setzero_ps(), _mm_castsi128_ph(_mm_cvtsi32_si128(*(nk_u16_t const *)from))));
+}
+
+#pragma region Vectorized Conversions
+
+/** @brief Convert 16x e4m3 → 16x f16 via bit manipulation (AVX-512 FP16).
+ *  E4M3 format: S EEEE MMM (bias=7). F16: S EEEEE MMMMMMMMMM (bias=15).
+ *  Normal: sign | ((exp+8)<<10) | (mant<<7).
+ *  Subnormals (exp=0): value = mantissa ÷ 512, computed via f16 arithmetic. */
+NK_INTERNAL __m256h nk_e4m3x16_to_f16x16_sapphire_(__m128i e4m3_i8x16) {
+    __m256i e4m3_i16x16 = _mm256_cvtepu8_epi16(e4m3_i8x16);
+
+    // Extract fields
+    __m256i mantissa_i16x16 = _mm256_and_si256(e4m3_i16x16, _mm256_set1_epi16(0x07));
+    __m256i sign_i16x16 = _mm256_and_si256(_mm256_slli_epi16(e4m3_i16x16, 8), _mm256_set1_epi16((short)0x8000));
+
+    // Normal path: sign | ((exp+8)<<10) | (mantissa<<7) via single shift + bias add
+    __m256i exp_mantissa_i16x16 = _mm256_slli_epi16(_mm256_and_si256(e4m3_i16x16, _mm256_set1_epi16(0x7F)), 7);
+    __m256i exp_mantissa_biased_i16x16 = _mm256_add_epi16(exp_mantissa_i16x16, _mm256_set1_epi16(0x2000));
+    __m256i normal_i16x16 = _mm256_or_si256(sign_i16x16, exp_mantissa_biased_i16x16);
+
+    // Subnormal fix: for exp==0 lanes, use (subnorm_abs | sign); else keep normal
+    __mmask16 is_subnormal = _mm256_testn_epi16_mask(e4m3_i16x16, _mm256_set1_epi16(0x78));
+    __m256h subnorm_abs_f16x16 = _mm256_mul_ph(_mm256_cvtepi16_ph(mantissa_i16x16),
+                                               _mm256_castsi256_ph(_mm256_set1_epi16(0x1800))); // 1/512
+    __m256i subnorm_signed_i16x16 = _mm256_or_si256(_mm256_castph_si256(subnorm_abs_f16x16), sign_i16x16);
+    __m256i result_i16x16 = _mm256_mask_blend_epi16(is_subnormal, normal_i16x16, subnorm_signed_i16x16);
+
+    // NaN path: E4M3FN has NaN only when exp=15 AND mant=7 (lower 7 bits == 0x7F)
+    __mmask16 is_nan = _mm256_cmpeq_epi16_mask(                                 //
+        _mm256_and_si256(e4m3_i16x16, _mm256_set1_epi16(0x7F)),                 //
+        _mm256_set1_epi16(0x7F));                                               //
+    __m256i nan_bits = _mm256_or_si256(sign_i16x16, _mm256_set1_epi16(0x7E00)); // F16 quiet NaN
+    return _mm256_castsi256_ph(_mm256_mask_blend_epi16(is_nan, result_i16x16, nan_bits));
+}
+
+/** @brief Convert 16x e5m2 → 16x f16 via bit manipulation (AVX-512 FP16).
+ *  E5M2 format: S EEEEE MM (bias=15). F16: S EEEEE MMMMMMMMMM (bias=15).
+ *  Normal: sign | (exp<<10) | (mant<<8) (same exponent bias).
+ *  Subnormals (exp=0): value = mantissa ÷ 65536, computed via f16 arithmetic. */
+NK_INTERNAL __m256h nk_e5m2x16_to_f16x16_sapphire_(__m128i e5m2_i8x16) {
+    __m256i e5m2_i16x16 = _mm256_cvtepu8_epi16(e5m2_i8x16);
+
+    // Extract fields
+    __m256i mantissa_i16x16 = _mm256_and_si256(e5m2_i16x16, _mm256_set1_epi16(0x03));
+    __m256i sign_i16x16 = _mm256_and_si256(_mm256_slli_epi16(e5m2_i16x16, 8), _mm256_set1_epi16((short)0x8000));
+
+    // Normal path: sign | (exp<<10) | (mant<<8) - same exponent bias so just shift lower7 by 8
+    __m256i exp_mantissa_i16x16 = _mm256_slli_epi16(_mm256_and_si256(e5m2_i16x16, _mm256_set1_epi16(0x7F)), 8);
+    __m256i normal_i16x16 = _mm256_or_si256(sign_i16x16, exp_mantissa_i16x16);
+
+    // Subnormal fix: for exp==0 lanes, use (subnorm_abs | sign); else keep normal
+    __mmask16 is_subnormal = _mm256_testn_epi16_mask(e5m2_i16x16, _mm256_set1_epi16(0x7C));
+    __m256h subnorm_abs_f16x16 = _mm256_mul_ph(_mm256_cvtepi16_ph(mantissa_i16x16),
+                                               _mm256_castsi256_ph(_mm256_set1_epi16(0x0100))); // 1/65536
+    __m256i subnorm_signed_i16x16 = _mm256_or_si256(_mm256_castph_si256(subnorm_abs_f16x16), sign_i16x16);
+    return _mm256_castsi256_ph(_mm256_mask_blend_epi16(is_subnormal, normal_i16x16, subnorm_signed_i16x16));
+}
+
+/** @brief Convert 16x f16 → 16x e4m3 via bit manipulation (AVX-512 FP16).
+ *  F16: S EEEEE MMMMMMMMMM (bias=15). E4M3: S EEEE MMM (bias=7).
+ *  Handles normal, subnormal, and overflow cases with RNE rounding. */
+NK_INTERNAL __m128i nk_f16x16_to_e4m3x16_sapphire_(__m256h f16x16) {
+    __m256i bits_i16x16 = _mm256_castph_si256(f16x16);
+    __m256i sign_i16x16 = _mm256_srli_epi16(bits_i16x16, 15);
+    __m256i f16_exp_i16x16 = _mm256_and_si256(_mm256_srli_epi16(bits_i16x16, 10), _mm256_set1_epi16(0x1F));
+
+    // Round mantissa from 10 to 3 bits using RNE (round to nearest, ties to even)
+    __m256i significand_i16x16 = _mm256_or_si256(_mm256_and_si256(bits_i16x16, _mm256_set1_epi16(0x03FF)),
+                                                 _mm256_set1_epi16(0x0400)); // Add implicit 1 bit
+    __m256i lsb_i16x16 = _mm256_and_si256(_mm256_srli_epi16(significand_i16x16, 7), _mm256_set1_epi16(1));
+    __m256i rounding_bias_i16x16 = _mm256_add_epi16(_mm256_set1_epi16(0x003F), lsb_i16x16);
+    __m256i rounded_sig_i16x16 = _mm256_add_epi16(significand_i16x16, rounding_bias_i16x16);
+    __m256i carry_i16x16 = _mm256_srli_epi16(rounded_sig_i16x16, 11); // Carry into exponent if bit 11 set
+    __m256i f16_mantissa_i16x16 = _mm256_and_si256(_mm256_srli_epi16(rounded_sig_i16x16, 7), _mm256_set1_epi16(0x07));
+    // If carry, mantissa becomes 0 (we rounded up to next power of 2)
+    f16_mantissa_i16x16 = _mm256_andnot_si256(_mm256_slli_epi16(carry_i16x16, 15), f16_mantissa_i16x16);
+    __m256i e4m3_exp_i16x16 = _mm256_sub_epi16(_mm256_add_epi16(f16_exp_i16x16, carry_i16x16), _mm256_set1_epi16(8));
+
+    // Detect underflow (exp <= 0) and overflow (exp > 15)
+    __mmask16 is_subnormal = _mm256_cmpgt_epi16_mask(_mm256_set1_epi16(1), e4m3_exp_i16x16);
+    __mmask16 overflow = _mm256_cmpgt_epi16_mask(e4m3_exp_i16x16, _mm256_set1_epi16(15));
+
+    // Normal path: clamp exp to [1,15]
+    // e4m3FN quirk: exp=15 with mantissa=7 is NaN (0x7F), so clamp mantissa to 6 when exp=15.
+    __m256i clamped_exp_i16x16 = _mm256_max_epi16(e4m3_exp_i16x16, _mm256_set1_epi16(1));
+    clamped_exp_i16x16 = _mm256_min_epi16(clamped_exp_i16x16, _mm256_set1_epi16(15));
+    __mmask16 is_max_exp = _mm256_cmpeq_epi16_mask(clamped_exp_i16x16, _mm256_set1_epi16(15));
+    __m256i max_mantissa_i16x16 = _mm256_mask_blend_epi16(is_max_exp, _mm256_set1_epi16(7), _mm256_set1_epi16(6));
+    __m256i normal_mantissa_i16x16 = _mm256_min_epi16(f16_mantissa_i16x16, max_mantissa_i16x16);
+    normal_mantissa_i16x16 = _mm256_mask_blend_epi16(overflow, normal_mantissa_i16x16, _mm256_set1_epi16(0x06));
+    __m256i normal_e4m3_i16x16 = _mm256_or_si256(
+        _mm256_slli_epi16(sign_i16x16, 7),
+        _mm256_or_si256(_mm256_slli_epi16(clamped_exp_i16x16, 3), normal_mantissa_i16x16));
+
+    // Subnormal path: mantissa = round(abs_f16 * 512)
+    __m256h abs_f16x16 = _mm256_castsi256_ph(_mm256_and_si256(_mm256_castph_si256(f16x16), _mm256_set1_epi16(0x7FFF)));
+    __m256h scaled_f16x16 = _mm256_mul_ph(abs_f16x16, _mm256_castsi256_ph(_mm256_set1_epi16(0x6000))); // 512
+    __m256i subnorm_mantissa_i16x16 = _mm256_cvtph_epi16(scaled_f16x16);
+    __mmask16 promotes_to_normal = _mm256_cmpgt_epi16_mask(subnorm_mantissa_i16x16, _mm256_set1_epi16(7));
+    subnorm_mantissa_i16x16 = _mm256_min_epi16(subnorm_mantissa_i16x16, _mm256_set1_epi16(7));
+    subnorm_mantissa_i16x16 = _mm256_max_epi16(subnorm_mantissa_i16x16, _mm256_setzero_si256());
+    __m256i subnorm_e4m3_i16x16 = _mm256_or_si256(_mm256_slli_epi16(sign_i16x16, 7), subnorm_mantissa_i16x16);
+    __m256i first_normal_e4m3_i16x16 = _mm256_or_si256(_mm256_slli_epi16(sign_i16x16, 7), _mm256_set1_epi16(0x08));
+    subnorm_e4m3_i16x16 = _mm256_mask_blend_epi16(promotes_to_normal, subnorm_e4m3_i16x16, first_normal_e4m3_i16x16);
+
+    // Blend: use subnormal result when exp <= 0
+    __m256i e4m3_i16x16 = _mm256_mask_blend_epi16(is_subnormal, normal_e4m3_i16x16, subnorm_e4m3_i16x16);
+
+    // Pack 16 i16s to 16 unsigned i8s via AVX-512BW
+    return _mm256_cvtepi16_epi8(e4m3_i16x16);
+}
+
+/** @brief Convert 16x f16 → 16x e5m2 via bit manipulation (AVX-512 FP16).
+ *  F16: S EEEEE MMMMMMMMMM (bias=15). E5M2: S EEEEE MM (bias=15).
+ *  Same exponent bias, so just round mantissa from 10 to 2 bits. */
+NK_INTERNAL __m128i nk_f16x16_to_e5m2x16_sapphire_(__m256h f16x16) {
+    __m256i bits_i16x16 = _mm256_castph_si256(f16x16);
+    __m256i sign_i16x16 = _mm256_srli_epi16(bits_i16x16, 15);
+    __m256i f16_exp_i16x16 = _mm256_and_si256(_mm256_srli_epi16(bits_i16x16, 10), _mm256_set1_epi16(0x1F));
+
+    // Round mantissa from 10 to 2 bits using RNE (round to nearest, ties to even)
+    __m256i significand_i16x16 = _mm256_or_si256(_mm256_and_si256(bits_i16x16, _mm256_set1_epi16(0x03FF)),
+                                                 _mm256_set1_epi16(0x0400)); // Add implicit 1 bit
+    __m256i lsb_i16x16 = _mm256_and_si256(_mm256_srli_epi16(significand_i16x16, 8), _mm256_set1_epi16(1));
+    __m256i rounding_bias_i16x16 = _mm256_add_epi16(_mm256_set1_epi16(0x007F), lsb_i16x16);
+    __m256i rounded_sig_i16x16 = _mm256_add_epi16(significand_i16x16, rounding_bias_i16x16);
+    __m256i carry_i16x16 = _mm256_srli_epi16(rounded_sig_i16x16, 11); // Carry into exponent if bit 11 set
+    __m256i f16_mantissa_i16x16 = _mm256_and_si256(_mm256_srli_epi16(rounded_sig_i16x16, 8), _mm256_set1_epi16(0x03));
+    // If carry, mantissa becomes 0 (we rounded up to next power of 2)
+    f16_mantissa_i16x16 = _mm256_andnot_si256(_mm256_slli_epi16(carry_i16x16, 15), f16_mantissa_i16x16);
+    __m256i e5m2_exp_i16x16 = _mm256_add_epi16(f16_exp_i16x16, carry_i16x16);
+
+    // Detect subnormal (exp <= 0) and overflow (exp > 31)
+    __mmask16 is_subnormal = _mm256_cmpeq_epi16_mask(f16_exp_i16x16, _mm256_setzero_si256());
+    __mmask16 overflow = _mm256_cmpgt_epi16_mask(e5m2_exp_i16x16, _mm256_set1_epi16(31));
+
+    // Normal path: clamp exp to [1,31], on overflow return infinity
+    __m256i clamped_exp_i16x16 = _mm256_max_epi16(e5m2_exp_i16x16, _mm256_set1_epi16(1));
+    clamped_exp_i16x16 = _mm256_min_epi16(clamped_exp_i16x16, _mm256_set1_epi16(31));
+    __m256i normal_mantissa_i16x16 = _mm256_mask_blend_epi16(overflow, f16_mantissa_i16x16, _mm256_setzero_si256());
+    __m256i normal_e5m2_i16x16 = _mm256_or_si256(
+        _mm256_slli_epi16(sign_i16x16, 7),
+        _mm256_or_si256(_mm256_slli_epi16(clamped_exp_i16x16, 2), normal_mantissa_i16x16));
+
+    // Subnormal path: mantissa = round(abs_f16 * 65536)
+    __m256h abs_f16x16 = _mm256_castsi256_ph(_mm256_and_si256(_mm256_castph_si256(f16x16), _mm256_set1_epi16(0x7FFF)));
+    __m256h scaled_f16x16 = _mm256_mul_ph(abs_f16x16, _mm256_castsi256_ph(_mm256_set1_epi16(0x7C00))); // 65536 (inf)
+    __m256i subnorm_mantissa_i16x16 = _mm256_cvtph_epi16(scaled_f16x16);
+    __mmask16 promotes_to_normal = _mm256_cmpgt_epi16_mask(subnorm_mantissa_i16x16, _mm256_set1_epi16(3));
+    subnorm_mantissa_i16x16 = _mm256_min_epi16(subnorm_mantissa_i16x16, _mm256_set1_epi16(3));
+    subnorm_mantissa_i16x16 = _mm256_max_epi16(subnorm_mantissa_i16x16, _mm256_setzero_si256());
+    __m256i subnorm_e5m2_i16x16 = _mm256_or_si256(_mm256_slli_epi16(sign_i16x16, 7), subnorm_mantissa_i16x16);
+    __m256i first_normal_e5m2_i16x16 = _mm256_or_si256(_mm256_slli_epi16(sign_i16x16, 7), _mm256_set1_epi16(0x04));
+    subnorm_e5m2_i16x16 = _mm256_mask_blend_epi16(promotes_to_normal, subnorm_e5m2_i16x16, first_normal_e5m2_i16x16);
+
+    // Blend: use subnormal result when exp == 0
+    __m256i e5m2_i16x16 = _mm256_mask_blend_epi16(is_subnormal, normal_e5m2_i16x16, subnorm_e5m2_i16x16);
+
+    // Pack 16 i16s to 16 unsigned i8s via AVX-512BW
+    return _mm256_cvtepi16_epi8(e5m2_i16x16);
+}
+
+#pragma endregion Vectorized Conversions
+
+#pragma region Public API
+
+NK_PUBLIC void nk_cast_sapphire(void const *from, nk_dtype_t from_type, nk_size_t n, void *to, nk_dtype_t to_type) {
+    // Group 1: Conversions to f16 (e4m3 → f16, e5m2 → f16)
+    if (to_type == nk_f16_k && (from_type == nk_e4m3_k || from_type == nk_e5m2_k)) {
+        nk_e4m3_t const *from_ptr = (nk_e4m3_t const *)from;
+        nk_f16_t *to_ptr = (nk_f16_t *)to;
+        for (nk_size_t idx = 0; idx < n; idx += 16) {
+            nk_size_t remaining = n - idx;
+            __mmask16 mask = (remaining >= 16) ? 0xFFFF : (unsigned short)_bzhi_u32(0xFFFF, (unsigned)remaining);
+            __m128i in_f8x16 = _mm_maskz_loadu_epi8(mask, from_ptr + idx);
+            __m256h out_f16x16 = (from_type == nk_e4m3_k) ? nk_e4m3x16_to_f16x16_sapphire_(in_f8x16)
+                                                          : nk_e5m2x16_to_f16x16_sapphire_(in_f8x16);
+            _mm256_mask_storeu_epi16(to_ptr + idx, mask, _mm256_castph_si256(out_f16x16));
+        }
+    }
+
+    // Group 2: Conversions from f16 (f16 → e4m3, f16 → e5m2)
+    else if (from_type == nk_f16_k && (to_type == nk_e4m3_k || to_type == nk_e5m2_k)) {
+        nk_f16_t const *from_ptr = (nk_f16_t const *)from;
+        nk_e4m3_t *to_ptr = (nk_e4m3_t *)to;
+        for (nk_size_t idx = 0; idx < n; idx += 16) {
+            nk_size_t remaining = n - idx;
+            __mmask16 mask = (remaining >= 16) ? 0xFFFF : (unsigned short)_bzhi_u32(0xFFFF, (unsigned)remaining);
+            __m256h in_f16x16 = _mm256_castsi256_ph(_mm256_maskz_loadu_epi16(mask, from_ptr + idx));
+            __m128i out_f8x16 = (to_type == nk_e4m3_k) ? nk_f16x16_to_e4m3x16_sapphire_(in_f16x16)
+                                                       : nk_f16x16_to_e5m2x16_sapphire_(in_f16x16);
+            _mm_mask_storeu_epi8(to_ptr + idx, mask, out_f8x16);
+        }
+    }
+
+    // Default: delegate to Ice for all other conversions
+    else nk_cast_icelake(from, from_type, n, to, to_type);
+}
+
+#pragma endregion Public API
+
+#if defined(__clang__)
+#pragma clang attribute pop
+#elif defined(__GNUC__)
+#pragma GCC pop_options
+#endif
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif // NK_TARGET_SAPPHIRE
+#endif // NK_TARGET_X8664_
+#endif // NK_CAST_SAPPHIRE_H