numkong 7.0.0 → 7.4.1

package/include/numkong/cast/skylake.h

@@ -8,13 +8,13 @@
  *
  *  @section skylake_cast_instructions AVX-512 Conversion Instructions
  *
- *      Intrinsic                   Instruction                     SKL         ICL         Genoa
- *      _mm512_cvtph_ps             VCVTPH2PS (ZMM, YMM)            5cy @ p05   5cy @ p05   4cy @ p01
- *      _mm512_cvtps_ph             VCVTPS2PH (YMM, ZMM, imm)       5cy @ p05   5cy @ p05   4cy @ p01
- *      _mm512_cvtps_epi32          VCVTPS2DQ (ZMM, ZMM)            4cy @ p0    4cy @ p0    3cy @ p01
- *      _mm512_cvtepi32_ps          VCVTDQ2PS (ZMM, ZMM)            4cy @ p0    4cy @ p0    3cy @ p01
- *      _mm512_cvtepi32_epi16       VPMOVDW (YMM, ZMM)              3cy @ p5    3cy @ p5    2cy @ p12
- *      _mm512_cvtsepi32_epi8       VPMOVSDB (XMM, ZMM)             3cy @ p5    3cy @ p5    2cy @ p12
+ *      Intrinsic              Instruction                SKL        ICL        Genoa
+ *      _mm512_cvtph_ps        VCVTPH2PS (ZMM, YMM)       5cy @ p05  5cy @ p05  4cy @ p01
+ *      _mm512_cvtps_ph        VCVTPS2PH (YMM, ZMM, imm)  5cy @ p05  5cy @ p05  4cy @ p01
+ *      _mm512_cvtps_epi32     VCVTPS2DQ (ZMM, ZMM)       4cy @ p0   4cy @ p0   3cy @ p01
+ *      _mm512_cvtepi32_ps     VCVTDQ2PS (ZMM, ZMM)       4cy @ p0   4cy @ p0   3cy @ p01
+ *      _mm512_cvtepi32_epi16  VPMOVDW (YMM, ZMM)         3cy @ p5   3cy @ p5   2cy @ p12
+ *      _mm512_cvtsepi32_epi8  VPMOVSDB (XMM, ZMM)        3cy @ p5   3cy @ p5   2cy @ p12
  *
  *  F16 conversions use hardware F16C via VCVTPH2PS/VCVTPS2PH. BF16 lacks hardware support on Skylake,
  *  requiring emulation via VPMOVZXWD + VPSLLD for bf16-to-f32, achieving ~4cy total. FP8 (E4M3/E5M2)
@@ -41,7 +41,7 @@ extern "C" {
 #pragma GCC target("avx2", "avx512f", "avx512vl", "avx512bw", "avx512dq", "f16c", "fma", "bmi", "bmi2")
 #endif
 
-#pragma region - Type Punned Loads and Stores
+#pragma region Type Punned Loads and Stores
 
 /** @brief Type-agnostic 512-bit full load (Skylake AVX-512). */
 NK_INTERNAL void nk_load_b512_skylake_(void const *src, nk_b512_vec_t *dst) { dst->zmm = _mm512_loadu_si512(src); }
@@ -132,9 +132,32 @@ NK_INTERNAL void nk_partial_store_b64x4_skylake_(nk_b256_vec_t const *src, void
     _mm256_mask_storeu_epi64(dst, mask, src->ymm);
 }
 
-#pragma endregion - Type Punned Loads and Stores
+/** @brief Type-agnostic full store for 512-bit vector (Skylake AVX-512). */
+NK_INTERNAL void nk_store_b512_skylake_(nk_b512_vec_t const *src, void *dst) {
+    _mm512_storeu_si512((__m512i *)dst, src->zmm);
+}
+
+/** @brief Type-agnostic partial store for 16-bit elements (32 elements max) from 512-bit vector (Skylake AVX-512). */
+NK_INTERNAL void nk_partial_store_b16x32_skylake_(nk_b512_vec_t const *src, void *dst, nk_size_t n) {
+    __mmask32 mask = (__mmask32)_bzhi_u32(0xFFFFFFFF, (unsigned int)n);
+    _mm512_mask_storeu_epi16(dst, mask, src->zmm);
+}
+
+/** @brief Type-agnostic partial store for 8-bit elements (64 elements max) from 512-bit vector (Skylake AVX-512). */
+NK_INTERNAL void nk_partial_store_b8x64_skylake_(nk_b512_vec_t const *src, void *dst, nk_size_t n) {
+    __mmask64 mask = _bzhi_u64(0xFFFFFFFFFFFFFFFFULL, (unsigned int)n);
+    _mm512_mask_storeu_epi8(dst, mask, src->zmm);
+}
+
+/** @brief Type-agnostic partial store for 64-bit elements (8 elements max) from 512-bit vector (Skylake AVX-512). */
+NK_INTERNAL void nk_partial_store_b64x8_skylake_(nk_b512_vec_t const *src, void *dst, nk_size_t n) {
+    __mmask8 mask = (__mmask8)_bzhi_u32(0xFF, (unsigned int)n);
+    _mm512_mask_storeu_epi64(dst, mask, src->zmm);
+}
 
-#pragma region - Vectorized Conversions
+#pragma endregion Type Punned Loads and Stores
+
+#pragma region Vectorized Conversions
 
 /** @brief Convert 16x bf16 → 16x f32 (Skylake AVX-512). */
 NK_INTERNAL __m512 nk_bf16x16_to_f32x16_skylake_(__m256i a) {
@@ -169,18 +192,20 @@ NK_INTERNAL __m512 nk_e4m3x16_to_f32x16_skylake_(__m128i e4m3_i8x16) {
     __m512 result_f32x16 = _mm512_castsi512_ps(
         _mm512_ternarylogic_epi32(sign_i32x16, f32_exp_i32x16, f32_mantissa_i32x16, 0xFE));
 
-    // Subnormal fix: for exp==0 lanes, replace with (mantissa / 512) | sign using masked OR
+    // Subnormal fix: vpermps from 8-entry LUT (repeated to fill 16 lanes)
     __mmask16 is_subnormal = _mm512_testn_epi32_mask(e4m3_i32x16, _mm512_set1_epi32(0x78));
-    __m512 subnorm_abs_f32x16 = _mm512_mul_ps(_mm512_cvtepi32_ps(mantissa_i32x16), _mm512_set1_ps(1.0f / 512.0f));
+    __m512 subnorm_lut_f32x16 = _mm512_setr_ps(                                                //
+        0, 1.0f / 512, 2.0f / 512, 3.0f / 512, 4.0f / 512, 5.0f / 512, 6.0f / 512, 7.0f / 512, //
+        0, 1.0f / 512, 2.0f / 512, 3.0f / 512, 4.0f / 512, 5.0f / 512, 6.0f / 512, 7.0f / 512);
+    __m512 subnorm_abs_f32x16 = _mm512_permutexvar_ps(mantissa_i32x16, subnorm_lut_f32x16);
     result_f32x16 = _mm512_mask_or_ps(result_f32x16, is_subnormal, subnorm_abs_f32x16,
                                       _mm512_castsi512_ps(sign_i32x16));
 
-    // NaN path: E4M3FN has NaN only when exp=15 AND mant=7 (0x7F or 0xFF)
-    __mmask16 is_nan = _mm512_mask_cmpeq_epi32_mask(                                //
-        _mm512_cmpeq_epi32_mask(exp_i32x16, _mm512_set1_epi32(15)),                 //
-        mantissa_i32x16, _mm512_set1_epi32(7));                                     //
-    __m512i nan_bits = _mm512_or_si512(sign_i32x16, _mm512_set1_epi32(0x7FC00000)); // F32 quiet NaN
-    return _mm512_mask_blend_ps(is_nan, result_f32x16, _mm512_castsi512_ps(nan_bits));
+    // NaN: E4M3FN has NaN only at magnitude 0x7F (single mask comparison)
+    __m512i lower7_i32x16 = _mm512_and_si512(e4m3_i32x16, _mm512_set1_epi32(0x7F));
+    __mmask16 is_nan = _mm512_cmpeq_epi32_mask(lower7_i32x16, _mm512_set1_epi32(0x7F));
+    __m512i nan_i32x16 = _mm512_or_si512(sign_i32x16, _mm512_set1_epi32(0x7FC00000));
+    return _mm512_mask_blend_ps(is_nan, result_f32x16, _mm512_castsi512_ps(nan_i32x16));
 }
 
 /** @brief Convert 16x e5m2 → 16x f32 via bit manipulation (AVX-512).
@@ -561,9 +586,43 @@ NK_INTERNAL __m256i nk_f64x8_to_u32x8_skylake_(__m512d f64x8) {
     return _mm512_cvtpd_epu32(clamped);
 }
 
-#pragma endregion - Vectorized Conversions
+/**
+ *  @brief Convert 64x E2M3 → 64x I8 using VPSHUFB LUT (Skylake AVX-512).
+ *
+ *  E2M3 format: [sign:1][magnitude:5] where magnitude indexes a 32-entry LUT
+ *  that produces the scaled integer value. Sign bit negates the result.
+ *  The 32-entry LUT is split into two 16-entry halves for VPSHUFB (which
+ *  indexes within 16-byte lanes). Bit 4 of the magnitude selects the half.
+ */
+NK_INTERNAL __m512i nk_e2m3x64_to_i8x64_skylake_(__m512i raw_i8x64) {
+    // lut_magnitude[0..15]  = {0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30}
+    // lut_magnitude[16..31] = {32, 36, 40, 44, 48, 52, 56, 60, 64, 72, 80, 88, 96, 104, 112, 120}
+    // _mm512_set4_epi32(d3, d2, d1, d0) fills bytes [0..3]=d0, [4..7]=d1, [8..11]=d2, [12..15]=d3
+    // per 128-bit lane, matching VPSHUFB's per-lane indexing.
+    __m512i lut_low_i8x64 = _mm512_set4_epi32( //
+        0x1E1C1A18, 0x16141210, 0x0E0C0A08, 0x06040200);
+    __m512i lut_high_i8x64 = _mm512_set4_epi32( //
+        0x78706860, 0x58504840, 0x3C383430, 0x2C282420);
+
+    __m512i magnitude_i8x64 = _mm512_and_si512(raw_i8x64, _mm512_set1_epi8(0x1F));
+    __m512i index_i8x64 = _mm512_and_si512(magnitude_i8x64, _mm512_set1_epi8(0x0F));
+
+    __m512i val_low_i8x64 = _mm512_shuffle_epi8(lut_low_i8x64, index_i8x64);
+    __m512i val_high_i8x64 = _mm512_shuffle_epi8(lut_high_i8x64, index_i8x64);
+
+    // Select high half when bit 4 of magnitude is set (magnitude >= 16)
+    __mmask64 use_high_mask = _mm512_test_epi8_mask(magnitude_i8x64, _mm512_set1_epi8(0x10));
+    __m512i val_i8x64 = _mm512_mask_blend_epi8(use_high_mask, val_low_i8x64, val_high_i8x64);
+
+    // Negate if sign bit (bit 5) is set
+    __mmask64 sign_mask = _mm512_test_epi8_mask(raw_i8x64, _mm512_set1_epi8(0x20));
+    __m512i negated_i8x64 = _mm512_sub_epi8(_mm512_setzero_si512(), val_i8x64);
+    return _mm512_mask_blend_epi8(sign_mask, val_i8x64, negated_i8x64);
+}
+
+#pragma endregion Vectorized Conversions
 
-#pragma region - Converting Loads and Stores
+#pragma region Converting Loads and Stores
 
 /** @brief Load 16 f16 values and convert to 16 f32 (Skylake AVX-512). */
 NK_INTERNAL void nk_load_f16x16_to_f32x16_skylake_(void const *src, nk_b512_vec_t *dst) {
@@ -637,9 +696,9 @@ NK_INTERNAL void nk_partial_load_e3m2x16_to_f32x16_skylake_(void const *src, nk_
     dst->zmm_ps = nk_e3m2x16_to_f32x16_skylake_(e3m2_partial.xmm);
 }
 
-#pragma endregion - Converting Loads and Stores
+#pragma endregion Converting Loads and Stores
 
-#pragma region - Public API
+#pragma region Public API
 
 NK_PUBLIC void nk_cast_skylake(void const *from, nk_dtype_t from_type, nk_size_t n, void *to, nk_dtype_t to_type) {
     // Same-type fast path
@@ -839,7 +898,7 @@ NK_PUBLIC void nk_cast_skylake(void const *from, nk_dtype_t from_type, nk_size_t
     nk_cast_serial(from, from_type, n, to, to_type);
 }
 
-#pragma endregion - Public API
+#pragma endregion Public API
 
 #if defined(__clang__)
 #pragma clang attribute pop