normal-grain-merge 0.0.2__cp313-cp313-win_amd64.whl → 0.1.1__cp313-cp313-win_amd64.whl

normal_grain_merge/normal_grain_merge.c

@@ -2,11 +2,17 @@
 #include <stdio.h>
 #include <math.h>
 #include <float.h>
+#include <stdint.h>
+#include <string.h>
 #include <Python.h>
 #include <numpy/arrayobject.h>
 #include <smmintrin.h>
 #include <immintrin.h>  /* AVX2 + SSE4.2 */
 
+#if defined(__FMA__) || (defined(_MSC_VER) && defined(__AVX2__))
+#define NGM_HAS_FMA 1
+#endif
+
 /* ----- Runtime CPU feature detection (GCC/Clang + MSVC) ----- */
 #if defined(_MSC_VER)
   #include <intrin.h>
@@ -126,8 +132,8 @@ static inline int check_shape_requirements(PyArrayObject *base,
     }
     *texture_has_alpha = (tc == 4);
 
-    if (PyArray_NDIM(skin) != 3 || PyArray_DIMS(skin)[2] != 4) {
-        PyErr_SetString(PyExc_ValueError, "skin must have shape (H, W, 4)");
+    if (PyArray_NDIM(skin) != 3 || PyArray_DIMS(skin)[2] != 3) {
+        PyErr_SetString(PyExc_ValueError, "skin must have shape (H, W, 3)");
         return 0;
     }
     if (PyArray_NDIM(im_alpha) != 2) {
@@ -184,10 +190,9 @@ static void kernel_scalar_rgb(const uint8_t *base, const uint8_t *texture,
         const uint8_t t_g = texture[3*i+1];
         const uint8_t t_b = texture[3*i+2];
 
-        const uint8_t s_r = skin[4*i+0];
-        const uint8_t s_g = skin[4*i+1];
-        const uint8_t s_b = skin[4*i+2];
-        const uint8_t s_a = skin[4*i+3];
+        const uint8_t s_r = skin[3*i+0];
+        const uint8_t s_g = skin[3*i+1];
+        const uint8_t s_b = skin[3*i+2];
 
         const uint8_t a_im = im_alpha[i];
 
@@ -203,8 +208,6 @@ static void kernel_scalar_rgb(const uint8_t *base, const uint8_t *texture,
         const float fs_r = s_r * (1.0f/255.0f);
         const float fs_g = s_g * (1.0f/255.0f);
         const float fs_b = s_b * (1.0f/255.0f);
-        const float fs_a = s_a * (1.0f/255.0f);
-
         const float fa_im = a_im * (1.0f/255.0f);
 
         /*
@@ -240,8 +243,8 @@ static void kernel_scalar_rgb(const uint8_t *base, const uint8_t *texture,
 
         /* Normal merge
          * n_out = gm_out * texture_alpha + base * inverse_tpa
-         * 
-         * In this case, texture_alpha is the skin alpha since texture doesn't have an alpha channel here.
+         *
+         * In this case, texture_alpha is supplied by im_alpha since texture doesn't have an alpha channel here.
          */
         fr = fr * fa_im + fb_r * fit_a;
         fg = fg * fa_im + fb_g * fit_a;
@@ -267,10 +270,9 @@ static void kernel_scalar_rgba(const uint8_t *base, const uint8_t *texture,
         const uint8_t t_b = texture[4*i+2];
         const uint8_t t_a = texture[4*i+3];  /* present in RGBA branch */
 
-        const uint8_t s_r = skin[4*i+0];
-        const uint8_t s_g = skin[4*i+1];
-        const uint8_t s_b = skin[4*i+2];
-        const uint8_t s_a = skin[4*i+3];
+        const uint8_t s_r = skin[3*i+0];
+        const uint8_t s_g = skin[3*i+1];
+        const uint8_t s_b = skin[3*i+2];
 
         const uint8_t a_im = im_alpha[i];
 
@@ -286,8 +288,6 @@ static void kernel_scalar_rgba(const uint8_t *base, const uint8_t *texture,
         const float fs_r = s_r * (1.0f/255.0f);
         const float fs_g = s_g * (1.0f/255.0f);
         const float fs_b = s_b * (1.0f/255.0f);
-        const float fs_a = s_a * (1.0f/255.0f);
-
         const float fa_im = a_im * (1.0f/255.0f);
 
         /*
@@ -295,7 +295,7 @@ static void kernel_scalar_rgba(const uint8_t *base, const uint8_t *texture,
          * normal grain merge *
          **********************
          */
-        /* Merge texture and skin alphas */
+        /* Merge texture alpha with the external mask */
 
         /* texture_alpha = texture[..., 3] * im_alpha*/
         ft_a = ft_a * fa_im;
@@ -341,34 +341,48 @@ static void kernel_scalar_rgba(const uint8_t *base, const uint8_t *texture,
 }
 
 /* ---------- AVX2 helpers ----------
-   Interleaved RGB(A) is awkward for SIMD. For a skeleton that still uses AVX2, we use gathers
-   over a stride (3 or 4) to pull 8 pixels for a given channel into a vector.
-   You can later replace gathers with better deinterleaving if needed.
+   Interleaved RGB(A) is awkward for SIMD. We build 8-lane vectors per channel by
+   reusing the scalar u8x4 -> f32 helpers instead of relying on gathers.
 */
 
-/* Convert 8 u8 interleaved channel samples (stride 3 or 4) to float32 in [0,1] via gather. */
-static inline __m256 gather_u8_to_unit_f32_avx2(const uint8_t *base_ptr, int stride,
-                                                npy_intp start_idx) {
-    const int i0 = (int)((start_idx + 0) * stride);
-    const int i1 = (int)((start_idx + 1) * stride);
-    const int i2 = (int)((start_idx + 2) * stride);
-    const int i3 = (int)((start_idx + 3) * stride);
-    const int i4 = (int)((start_idx + 4) * stride);
-    const int i5 = (int)((start_idx + 5) * stride);
-    const int i6 = (int)((start_idx + 6) * stride);
-    const int i7 = (int)((start_idx + 7) * stride);
-
-    __m256i offs = _mm256_setr_epi32(i0, i1, i2, i3, i4, i5, i6, i7);
-    __m256i v32  = _mm256_i32gather_epi32((const int*)base_ptr, offs, 1); /* read 8 x u8 as u32 */
-    v32 = _mm256_and_si256(v32, _mm256_set1_epi32(0xFF));
-    return _mm256_mul_ps(_mm256_cvtepi32_ps(v32), _mm256_set1_ps(1.0f/255.0f));
+/* Forward declaration; definition shared with the SSE helpers later in the file. */
+static inline __m128 u8x4_to_unit_f32(uint8_t a, uint8_t b, uint8_t c, uint8_t d);
+
+/* Build 8-lane [0,1] floats for one channel from interleaved RGB. No gathers. */
+static inline __m256 load8_rgb_channel_to_unit_f32(const uint8_t *p, int ch /*0,1,2*/) {
+    /* pixel i, channel ch is at p[3*i + ch] */
+    __m128 lo = u8x4_to_unit_f32(p[3*0 + ch], p[3*1 + ch], p[3*2 + ch], p[3*3 + ch]);
+    __m128 hi = u8x4_to_unit_f32(p[3*4 + ch], p[3*5 + ch], p[3*6 + ch], p[3*7 + ch]);
+    return _mm256_set_m128(hi, lo);
+}
+
+static inline __m256 load8_rgba_channel_to_unit_f32(const uint8_t *p, int ch /*0..3*/) {
+    __m128 lo = u8x4_to_unit_f32(p[4*0 + ch], p[4*1 + ch], p[4*2 + ch], p[4*3 + ch]);
+    __m128 hi = u8x4_to_unit_f32(p[4*4 + ch], p[4*5 + ch], p[4*6 + ch], p[4*7 + ch]);
+    return _mm256_set_m128(hi, lo);
+}
+
+static inline __m256 mul_add_ps256(__m256 a, __m256 b, __m256 c) {
+#ifdef __FMA__
+    return _mm256_fmadd_ps(a, b, c);
+#else
+    return _mm256_add_ps(_mm256_mul_ps(a, b), c);
+#endif
+}
+
+static inline __m256 fnmadd_ps256(__m256 a, __m256 b, __m256 c) {
+#ifdef __FMA__
+    return _mm256_fnmadd_ps(a, b, c);
+#else
+    return _mm256_sub_ps(c, _mm256_mul_ps(a, b));
+#endif
 }
 
 /* Convert 8 consecutive u8 to float32 in [0,1] (for grayscale im_alpha). */
-static inline __m256 load8_u8_to_unit_f32_avx2(const uint8_t *p) {
+static inline __m256 load8_u8_to_unit_f32_avx2(const uint8_t *p, __m256 inv255) {
     __m128i v8  = _mm_loadl_epi64((const __m128i*)p);        /* 8 bytes -> XMM */
     __m256i v32 = _mm256_cvtepu8_epi32(v8);                  /* widen to 8 x u32 */
-    return _mm256_mul_ps(_mm256_cvtepi32_ps(v32), _mm256_set1_ps(1.0f/255.0f));
+    return _mm256_mul_ps(_mm256_cvtepi32_ps(v32), inv255);
 }
 
 static inline __m256 clamp01_ps(__m256 x) {
@@ -381,32 +395,74 @@ static inline __m256 nan_to_num_ps(__m256 x) {
     return _mm256_blendv_ps(_mm256_set1_ps(0.0f), x, cmp);
 }
 
-/* Truncate [0,1] floats to uint8 and scatter to interleaved RGB output. */
-static inline void store_unit_f32_to_u8_rgb8_avx2(__m256 fr, __m256 fg, __m256 fb,
-                                                  uint8_t *out_ptr, npy_intp start_idx) {
-    __m256 scale = _mm256_set1_ps(255.0f);
-    __m256i ir = _mm256_cvttps_epi32(_mm256_mul_ps(fr, scale));
-    __m256i ig = _mm256_cvttps_epi32(_mm256_mul_ps(fg, scale));
-    __m256i ib = _mm256_cvttps_epi32(_mm256_mul_ps(fb, scale));
-
-    int r[8], g[8], b[8];
-    _mm256_storeu_si256((__m256i*)r, ir);
-    _mm256_storeu_si256((__m256i*)g, ig);
-    _mm256_storeu_si256((__m256i*)b, ib);
-
-    for (int k = 0; k < 8; ++k) {
-        const npy_intp p = start_idx + k;
-        out_ptr[3*p+0] = (uint8_t)(r[k] < 0 ? 0 : r[k] > 255 ? 255 : r[k]);
-        out_ptr[3*p+1] = (uint8_t)(g[k] < 0 ? 0 : g[k] > 255 ? 255 : g[k]);
-        out_ptr[3*p+2] = (uint8_t)(b[k] < 0 ? 0 : b[k] > 255 ? 255 : b[k]);
-    }
+/* Convert 4 float32 RGB vectors in [0,1] to uint8_t RGBRGBRGBRGB without branches. */
+static inline void store_unit_f32_to_u8_rgb4(__m128 fr, __m128 fg, __m128 fb,
+                                             uint8_t *out_ptr) {
+    const __m128 scale = _mm_set1_ps(255.0f);
+    const __m128i zero = _mm_setzero_si128();
+    const __m128i max255 = _mm_set1_epi32(255);
+
+    __m128i ir = _mm_cvttps_epi32(_mm_mul_ps(fr, scale));
+    __m128i ig = _mm_cvttps_epi32(_mm_mul_ps(fg, scale));
+    __m128i ib = _mm_cvttps_epi32(_mm_mul_ps(fb, scale));
+
+    ir = _mm_min_epi32(_mm_max_epi32(ir, zero), max255);
+    ig = _mm_min_epi32(_mm_max_epi32(ig, zero), max255);
+    ib = _mm_min_epi32(_mm_max_epi32(ib, zero), max255);
+
+    __m128i ir16 = _mm_packus_epi32(ir, zero);
+    __m128i ig16 = _mm_packus_epi32(ig, zero);
+    __m128i ib16 = _mm_packus_epi32(ib, zero);
+
+    __m128i ir8 = _mm_packus_epi16(ir16, zero);
+    __m128i ig8 = _mm_packus_epi16(ig16, zero);
+    __m128i ib8 = _mm_packus_epi16(ib16, zero);
+
+    const __m128i mask_r = _mm_setr_epi8(
+        0, (char)0x80, (char)0x80, 1,
+        (char)0x80, (char)0x80, 2, (char)0x80,
+        (char)0x80, 3, (char)0x80, (char)0x80,
+        (char)0x80, (char)0x80, (char)0x80, (char)0x80);
+    const __m128i mask_g = _mm_setr_epi8(
+        (char)0x80, 0, (char)0x80, (char)0x80,
+        1, (char)0x80, (char)0x80, 2,
+        (char)0x80, (char)0x80, 3, (char)0x80,
+        (char)0x80, (char)0x80, (char)0x80, (char)0x80);
+    const __m128i mask_b = _mm_setr_epi8(
+        (char)0x80, (char)0x80, 0, (char)0x80,
+        (char)0x80, 1, (char)0x80, (char)0x80,
+        2, (char)0x80, (char)0x80, 3,
+        (char)0x80, (char)0x80, (char)0x80, (char)0x80);
+
+    __m128i packed = _mm_or_si128(
+        _mm_or_si128(_mm_shuffle_epi8(ir8, mask_r),
+                     _mm_shuffle_epi8(ig8, mask_g)),
+        _mm_shuffle_epi8(ib8, mask_b));
+
+    _mm_storel_epi64((__m128i*)out_ptr, packed);
+    __m128i tail_vec = _mm_srli_si128(packed, 8);
+    uint32_t tail = (uint32_t)_mm_cvtsi128_si32(tail_vec);
+    memcpy(out_ptr + 8, &tail, sizeof(tail));
+}
+
+/* Pack 8 lanes of fr/fg/fb in [0,1] to 24 bytes RGBRGB... and store. */
+static inline void store8_unit_f32_to_u8_rgb(__m256 fr, __m256 fg, __m256 fb, uint8_t *out) {
+    __m128 fr_lo = _mm256_castps256_ps128(fr);
+    __m128 fg_lo = _mm256_castps256_ps128(fg);
+    __m128 fb_lo = _mm256_castps256_ps128(fb);
+    store_unit_f32_to_u8_rgb4(fr_lo, fg_lo, fb_lo, out);
+
+    __m128 fr_hi = _mm256_extractf128_ps(fr, 1);
+    __m128 fg_hi = _mm256_extractf128_ps(fg, 1);
+    __m128 fb_hi = _mm256_extractf128_ps(fb, 1);
+    store_unit_f32_to_u8_rgb4(fr_hi, fg_hi, fb_hi, out + 12);
 }
 
 /* texture is RGB: texture_alpha = im_alpha broadcast, inverse_tpa = 1 - texture_alpha */
 static void kernel_avx2_rgb(const uint8_t *base, const uint8_t *texture,
                             const uint8_t *skin, const uint8_t *im_alpha,
                             uint8_t *out, npy_intp pixels) {
-    const int stride3 = 3, stride4 = 4;
+    const __m256 inv255 = _mm256_set1_ps(1.0f/255.0f);
     const __m256 half = _mm256_set1_ps(0.5f);
     const __m256 one  = _mm256_set1_ps(1.0f);
     const __m256 w    = _mm256_set1_ps((float)SKIN_WEIGHT);
@@ -414,24 +470,28 @@ static void kernel_avx2_rgb(const uint8_t *base, const uint8_t *texture,
 
     npy_intp i = 0;
     for (; i + 8 <= pixels; i += 8) {
+        const uint8_t *base_blk = base + 3*i;
+        const uint8_t *tex_blk  = texture + 3*i;
+        const uint8_t *skin_blk = skin + 3*i;
+
         /* base RGB in [0,1] */
-        __m256 fb_r = gather_u8_to_unit_f32_avx2(base+0, stride3, i);
-        __m256 fb_g = gather_u8_to_unit_f32_avx2(base+1, stride3, i);
-        __m256 fb_b = gather_u8_to_unit_f32_avx2(base+2, stride3, i);
+        __m256 fb_r = load8_rgb_channel_to_unit_f32(base_blk, 0);
+        __m256 fb_g = load8_rgb_channel_to_unit_f32(base_blk, 1);
+        __m256 fb_b = load8_rgb_channel_to_unit_f32(base_blk, 2);
 
         /* texture RGB in [0,1] */
-        __m256 ft_r = gather_u8_to_unit_f32_avx2(texture+0, stride3, i);
-        __m256 ft_g = gather_u8_to_unit_f32_avx2(texture+1, stride3, i);
-        __m256 ft_b = gather_u8_to_unit_f32_avx2(texture+2, stride3, i);
+        __m256 ft_r = load8_rgb_channel_to_unit_f32(tex_blk, 0);
+        __m256 ft_g = load8_rgb_channel_to_unit_f32(tex_blk, 1);
+        __m256 ft_b = load8_rgb_channel_to_unit_f32(tex_blk, 2);
 
         /* skin RGB in [0,1] */
-        __m256 fs_r = gather_u8_to_unit_f32_avx2(skin+0, stride4, i);
-        __m256 fs_g = gather_u8_to_unit_f32_avx2(skin+1, stride4, i);
-        __m256 fs_b = gather_u8_to_unit_f32_avx2(skin+2, stride4, i);
+        __m256 fs_r = load8_rgb_channel_to_unit_f32(skin_blk, 0);
+        __m256 fs_g = load8_rgb_channel_to_unit_f32(skin_blk, 1);
+        __m256 fs_b = load8_rgb_channel_to_unit_f32(skin_blk, 2);
 
         /* texture_alpha = im_alpha */
-        __m256 fa_im = load8_u8_to_unit_f32_avx2(im_alpha + i);
-        __m256 fit_a = _mm256_sub_ps(one, fa_im);
+        __m256 fa_im = load8_u8_to_unit_f32_avx2(im_alpha + i, inv255);
+        __m256 fit_a = fnmadd_ps256(fa_im, one, one);
 
         /* gm_out = clip(texture + skin - 0.5) */
         __m256 gm_r = clamp01_ps(_mm256_sub_ps(_mm256_add_ps(ft_r, fs_r), half));
@@ -439,14 +499,14 @@ static void kernel_avx2_rgb(const uint8_t *base, const uint8_t *texture,
         __m256 gm_b = clamp01_ps(_mm256_sub_ps(_mm256_add_ps(ft_b, fs_b), half));
 
         /* gm_out = gm_out * texture_alpha + texture * inverse_tpa */
-        gm_r = _mm256_add_ps(_mm256_mul_ps(gm_r, fa_im), _mm256_mul_ps(ft_r, fit_a));
-        gm_g = _mm256_add_ps(_mm256_mul_ps(gm_g, fa_im), _mm256_mul_ps(ft_g, fit_a));
-        gm_b = _mm256_add_ps(_mm256_mul_ps(gm_b, fa_im), _mm256_mul_ps(ft_b, fit_a));
+        gm_r = mul_add_ps256(gm_r, fa_im, _mm256_mul_ps(ft_r, fit_a));
+        gm_g = mul_add_ps256(gm_g, fa_im, _mm256_mul_ps(ft_g, fit_a));
+        gm_b = mul_add_ps256(gm_b, fa_im, _mm256_mul_ps(ft_b, fit_a));
 
         /* gm_out = gm_out * (1 - w) + skin * w */
-        gm_r = _mm256_add_ps(_mm256_mul_ps(gm_r, invw), _mm256_mul_ps(fs_r, w));
-        gm_g = _mm256_add_ps(_mm256_mul_ps(gm_g, invw), _mm256_mul_ps(fs_g, w));
-        gm_b = _mm256_add_ps(_mm256_mul_ps(gm_b, invw), _mm256_mul_ps(fs_b, w));
+        gm_r = mul_add_ps256(gm_r, invw, _mm256_mul_ps(fs_r, w));
+        gm_g = mul_add_ps256(gm_g, invw, _mm256_mul_ps(fs_g, w));
+        gm_b = mul_add_ps256(gm_b, invw, _mm256_mul_ps(fs_b, w));
 
         /* nan_to_num */
         gm_r = nan_to_num_ps(gm_r);
@@ -454,15 +514,15 @@ static void kernel_avx2_rgb(const uint8_t *base, const uint8_t *texture,
         gm_b = nan_to_num_ps(gm_b);
 
         /* n_out = gm_out * texture_alpha + base * inverse_tpa */
-        __m256 fr = _mm256_add_ps(_mm256_mul_ps(gm_r, fa_im), _mm256_mul_ps(fb_r, fit_a));
-        __m256 fg = _mm256_add_ps(_mm256_mul_ps(gm_g, fa_im), _mm256_mul_ps(fb_g, fit_a));
-        __m256 fb = _mm256_add_ps(_mm256_mul_ps(gm_b, fa_im), _mm256_mul_ps(fb_b, fit_a));
+        __m256 fr = mul_add_ps256(gm_r, fa_im, _mm256_mul_ps(fb_r, fit_a));
+        __m256 fg = mul_add_ps256(gm_g, fa_im, _mm256_mul_ps(fb_g, fit_a));
+        __m256 fb = mul_add_ps256(gm_b, fa_im, _mm256_mul_ps(fb_b, fit_a));
 
-        store_unit_f32_to_u8_rgb8_avx2(fr, fg, fb, out, i);
+        store8_unit_f32_to_u8_rgb(fr, fg, fb, out + 3*i);
     }
 
     if (i < pixels) {
-        kernel_scalar_rgb(base + 3*i, texture + 3*i, skin + 4*i, im_alpha + i,
+        kernel_scalar_rgb(base + 3*i, texture + 3*i, skin + 3*i, im_alpha + i,
                           out + 3*i, pixels - i);
     }
 }
@@ -471,7 +531,7 @@ static void kernel_avx2_rgb(const uint8_t *base, const uint8_t *texture,
 static void kernel_avx2_rgba(const uint8_t *base, const uint8_t *texture,
                              const uint8_t *skin, const uint8_t *im_alpha,
                              uint8_t *out, npy_intp pixels) {
-    const int stride3 = 3, stride4 = 4;
+    const __m256 inv255 = _mm256_set1_ps(1.0f/255.0f);
     const __m256 half = _mm256_set1_ps(0.5f);
     const __m256 one  = _mm256_set1_ps(1.0f);
     const __m256 w    = _mm256_set1_ps((float)SKIN_WEIGHT);
@@ -479,48 +539,52 @@ static void kernel_avx2_rgba(const uint8_t *base, const uint8_t *texture,
 
     npy_intp i = 0;
     for (; i + 8 <= pixels; i += 8) {
-        __m256 fb_r = gather_u8_to_unit_f32_avx2(base+0, stride3, i);
-        __m256 fb_g = gather_u8_to_unit_f32_avx2(base+1, stride3, i);
-        __m256 fb_b = gather_u8_to_unit_f32_avx2(base+2, stride3, i);
+        const uint8_t *base_blk = base + 3*i;
+        const uint8_t *tex_blk  = texture + 4*i;
+        const uint8_t *skin_blk = skin + 3*i;
+
+        __m256 fb_r = load8_rgb_channel_to_unit_f32(base_blk, 0);
+        __m256 fb_g = load8_rgb_channel_to_unit_f32(base_blk, 1);
+        __m256 fb_b = load8_rgb_channel_to_unit_f32(base_blk, 2);
 
-        __m256 ft_r = gather_u8_to_unit_f32_avx2(texture+0, stride4, i);
-        __m256 ft_g = gather_u8_to_unit_f32_avx2(texture+1, stride4, i);
-        __m256 ft_b = gather_u8_to_unit_f32_avx2(texture+2, stride4, i);
-        __m256 ft_a = gather_u8_to_unit_f32_avx2(texture+3, stride4, i);  /* texture alpha */
+        __m256 ft_r = load8_rgba_channel_to_unit_f32(tex_blk, 0);
+        __m256 ft_g = load8_rgba_channel_to_unit_f32(tex_blk, 1);
+        __m256 ft_b = load8_rgba_channel_to_unit_f32(tex_blk, 2);
+        __m256 ft_a = load8_rgba_channel_to_unit_f32(tex_blk, 3);  /* texture alpha */
 
-        __m256 fs_r = gather_u8_to_unit_f32_avx2(skin+0, stride4, i);
-        __m256 fs_g = gather_u8_to_unit_f32_avx2(skin+1, stride4, i);
-        __m256 fs_b = gather_u8_to_unit_f32_avx2(skin+2, stride4, i);
+        __m256 fs_r = load8_rgb_channel_to_unit_f32(skin_blk, 0);
+        __m256 fs_g = load8_rgb_channel_to_unit_f32(skin_blk, 1);
+        __m256 fs_b = load8_rgb_channel_to_unit_f32(skin_blk, 2);
 
-        __m256 fa_im = load8_u8_to_unit_f32_avx2(im_alpha + i);
+        __m256 fa_im = load8_u8_to_unit_f32_avx2(im_alpha + i, inv255);
         __m256 fta   = _mm256_mul_ps(ft_a, fa_im);           /* texture_alpha */
-        __m256 fit_a = _mm256_sub_ps(one, fta);               /* inverse_tpa  */
+        __m256 fit_a = fnmadd_ps256(fta, one, one);           /* inverse_tpa  */
 
         __m256 gm_r = clamp01_ps(_mm256_sub_ps(_mm256_add_ps(ft_r, fs_r), half));
         __m256 gm_g = clamp01_ps(_mm256_sub_ps(_mm256_add_ps(ft_g, fs_g), half));
         __m256 gm_b = clamp01_ps(_mm256_sub_ps(_mm256_add_ps(ft_b, fs_b), half));
 
-        gm_r = _mm256_add_ps(_mm256_mul_ps(gm_r, fta), _mm256_mul_ps(ft_r, fit_a));
-        gm_g = _mm256_add_ps(_mm256_mul_ps(gm_g, fta), _mm256_mul_ps(ft_g, fit_a));
-        gm_b = _mm256_add_ps(_mm256_mul_ps(gm_b, fta), _mm256_mul_ps(ft_b, fit_a));
+        gm_r = mul_add_ps256(gm_r, fta, _mm256_mul_ps(ft_r, fit_a));
+        gm_g = mul_add_ps256(gm_g, fta, _mm256_mul_ps(ft_g, fit_a));
+        gm_b = mul_add_ps256(gm_b, fta, _mm256_mul_ps(ft_b, fit_a));
 
-        gm_r = _mm256_add_ps(_mm256_mul_ps(gm_r, invw), _mm256_mul_ps(fs_r, w));
-        gm_g = _mm256_add_ps(_mm256_mul_ps(gm_g, invw), _mm256_mul_ps(fs_g, w));
-        gm_b = _mm256_add_ps(_mm256_mul_ps(gm_b, invw), _mm256_mul_ps(fs_b, w));
+        gm_r = mul_add_ps256(gm_r, invw, _mm256_mul_ps(fs_r, w));
+        gm_g = mul_add_ps256(gm_g, invw, _mm256_mul_ps(fs_g, w));
+        gm_b = mul_add_ps256(gm_b, invw, _mm256_mul_ps(fs_b, w));
 
         gm_r = nan_to_num_ps(gm_r);
         gm_g = nan_to_num_ps(gm_g);
         gm_b = nan_to_num_ps(gm_b);
 
-        __m256 fr = _mm256_add_ps(_mm256_mul_ps(gm_r, fta), _mm256_mul_ps(fb_r, fit_a));
-        __m256 fg = _mm256_add_ps(_mm256_mul_ps(gm_g, fta), _mm256_mul_ps(fb_g, fit_a));
-        __m256 fb = _mm256_add_ps(_mm256_mul_ps(gm_b, fta), _mm256_mul_ps(fb_b, fit_a));
+        __m256 fr = mul_add_ps256(gm_r, fta, _mm256_mul_ps(fb_r, fit_a));
+        __m256 fg = mul_add_ps256(gm_g, fta, _mm256_mul_ps(fb_g, fit_a));
+        __m256 fb = mul_add_ps256(gm_b, fta, _mm256_mul_ps(fb_b, fit_a));
 
-        store_unit_f32_to_u8_rgb8_avx2(fr, fg, fb, out, i);
+        store8_unit_f32_to_u8_rgb(fr, fg, fb, out + 3*i);
     }
 
     if (i < pixels) {
-        kernel_scalar_rgba(base + 3*i, texture + 4*i, skin + 4*i, im_alpha + i,
+        kernel_scalar_rgba(base + 3*i, texture + 4*i, skin + 3*i, im_alpha + i,
                            out + 3*i, pixels - i);
     }
 }
@@ -550,6 +614,13 @@ static inline __m128 nan_to_num_ps128(__m128 x) {
     return _mm_blendv_ps(_mm_set1_ps(0.0f), x, cmp);
 }
 
+static inline __m128 mul_add_ps128(__m128 a, __m128 b, __m128 c) {
+#ifdef __FMA__
+    return _mm_fmadd_ps(a, b, c);
+#else
+    return _mm_add_ps(_mm_mul_ps(a, b), c);
+#endif
+}
 
 static void kernel_sse42_rgb(const uint8_t *base, const uint8_t *texture,
                              const uint8_t *skin, const uint8_t *im_alpha,
@@ -575,12 +646,12 @@ static void kernel_sse42_rgb(const uint8_t *base, const uint8_t *texture,
         __m128 ft_b = u8x4_to_unit_f32(texture[3*(i+0)+2], texture[3*(i+1)+2],
                                        texture[3*(i+2)+2], texture[3*(i+3)+2]);
 
-        __m128 fs_r = u8x4_to_unit_f32(skin[4*(i+0)+0], skin[4*(i+1)+0],
-                                       skin[4*(i+2)+0], skin[4*(i+3)+0]);
-        __m128 fs_g = u8x4_to_unit_f32(skin[4*(i+0)+1], skin[4*(i+1)+1],
-                                       skin[4*(i+2)+1], skin[4*(i+3)+1]);
-        __m128 fs_b = u8x4_to_unit_f32(skin[4*(i+0)+2], skin[4*(i+1)+2],
-                                       skin[4*(i+2)+2], skin[4*(i+3)+2]);
+        __m128 fs_r = u8x4_to_unit_f32(skin[3*(i+0)+0], skin[3*(i+1)+0],
+                                       skin[3*(i+2)+0], skin[3*(i+3)+0]);
+        __m128 fs_g = u8x4_to_unit_f32(skin[3*(i+0)+1], skin[3*(i+1)+1],
+                                       skin[3*(i+2)+1], skin[3*(i+3)+1]);
+        __m128 fs_b = u8x4_to_unit_f32(skin[3*(i+0)+2], skin[3*(i+1)+2],
+                                       skin[3*(i+2)+2], skin[3*(i+3)+2]);
 
         __m128 fa_im = load4_u8_to_unit_f32(im_alpha + i);
         __m128 fit_a = _mm_sub_ps(one, fa_im);
@@ -589,39 +660,27 @@ static void kernel_sse42_rgb(const uint8_t *base, const uint8_t *texture,
         __m128 gm_g = clamp01_ps128(_mm_sub_ps(_mm_add_ps(ft_g, fs_g), half));
         __m128 gm_b = clamp01_ps128(_mm_sub_ps(_mm_add_ps(ft_b, fs_b), half));
 
-        gm_r = _mm_add_ps(_mm_mul_ps(gm_r, fa_im), _mm_mul_ps(ft_r, fit_a));
-        gm_g = _mm_add_ps(_mm_mul_ps(gm_g, fa_im), _mm_mul_ps(ft_g, fit_a));
-        gm_b = _mm_add_ps(_mm_mul_ps(gm_b, fa_im), _mm_mul_ps(ft_b, fit_a));
+        gm_r = mul_add_ps128(gm_r, fa_im, _mm_mul_ps(ft_r, fit_a));
+        gm_g = mul_add_ps128(gm_g, fa_im, _mm_mul_ps(ft_g, fit_a));
+        gm_b = mul_add_ps128(gm_b, fa_im, _mm_mul_ps(ft_b, fit_a));
 
-        gm_r = _mm_add_ps(_mm_mul_ps(gm_r, invw), _mm_mul_ps(fs_r, w));
-        gm_g = _mm_add_ps(_mm_mul_ps(gm_g, invw), _mm_mul_ps(fs_g, w));
-        gm_b = _mm_add_ps(_mm_mul_ps(gm_b, invw), _mm_mul_ps(fs_b, w));
+        gm_r = mul_add_ps128(gm_r, invw, _mm_mul_ps(fs_r, w));
+        gm_g = mul_add_ps128(gm_g, invw, _mm_mul_ps(fs_g, w));
+        gm_b = mul_add_ps128(gm_b, invw, _mm_mul_ps(fs_b, w));
 
         gm_r = nan_to_num_ps128(gm_r);
         gm_g = nan_to_num_ps128(gm_g);
         gm_b = nan_to_num_ps128(gm_b);
 
-        __m128 fr = _mm_add_ps(_mm_mul_ps(gm_r, fa_im), _mm_mul_ps(fb_r, fit_a));
-        __m128 fg = _mm_add_ps(_mm_mul_ps(gm_g, fa_im), _mm_mul_ps(fb_g, fit_a));
-        __m128 fb = _mm_add_ps(_mm_mul_ps(gm_b, fa_im), _mm_mul_ps(fb_b, fit_a));
-
-        float rr[4], gg[4], bb[4];
-        _mm_storeu_ps(rr, fr);
-        _mm_storeu_ps(gg, fg);
-        _mm_storeu_ps(bb, fb);
-
-        for (int k = 0; k < 4; ++k) {
-            int r = (int)(rr[k] * 255.0f);
-            int g = (int)(gg[k] * 255.0f);
-            int b = (int)(bb[k] * 255.0f);
-            out[3*(i+k)+0] = (uint8_t)(r < 0 ? 0 : r > 255 ? 255 : r);
-            out[3*(i+k)+1] = (uint8_t)(g < 0 ? 0 : g > 255 ? 255 : g);
-            out[3*(i+k)+2] = (uint8_t)(b < 0 ? 0 : b > 255 ? 255 : b);
-        }
+        __m128 fr = mul_add_ps128(gm_r, fa_im, _mm_mul_ps(fb_r, fit_a));
+        __m128 fg = mul_add_ps128(gm_g, fa_im, _mm_mul_ps(fb_g, fit_a));
+        __m128 fb = mul_add_ps128(gm_b, fa_im, _mm_mul_ps(fb_b, fit_a));
+
+        store_unit_f32_to_u8_rgb4(fr, fg, fb, out + 3*i);
     }
 
     if (i < pixels) {
-        kernel_scalar_rgb(base + 3*i, texture + 3*i, skin + 4*i, im_alpha + i,
+        kernel_scalar_rgb(base + 3*i, texture + 3*i, skin + 3*i, im_alpha + i,
                           out + 3*i, pixels - i);
     }
 }
@@ -652,12 +711,12 @@ static void kernel_sse42_rgba(const uint8_t *base, const uint8_t *texture,
         __m128 ft_a = u8x4_to_unit_f32(texture[4*(i+0)+3], texture[4*(i+1)+3],
                                        texture[4*(i+2)+3], texture[4*(i+3)+3]);
 
-        __m128 fs_r = u8x4_to_unit_f32(skin[4*(i+0)+0], skin[4*(i+1)+0],
-                                       skin[4*(i+2)+0], skin[4*(i+3)+0]);
-        __m128 fs_g = u8x4_to_unit_f32(skin[4*(i+0)+1], skin[4*(i+1)+1],
-                                       skin[4*(i+2)+1], skin[4*(i+3)+1]);
-        __m128 fs_b = u8x4_to_unit_f32(skin[4*(i+0)+2], skin[4*(i+1)+2],
-                                       skin[4*(i+2)+2], skin[4*(i+3)+2]);
+        __m128 fs_r = u8x4_to_unit_f32(skin[3*(i+0)+0], skin[3*(i+1)+0],
+                                       skin[3*(i+2)+0], skin[3*(i+3)+0]);
+        __m128 fs_g = u8x4_to_unit_f32(skin[3*(i+0)+1], skin[3*(i+1)+1],
+                                       skin[3*(i+2)+1], skin[3*(i+3)+1]);
+        __m128 fs_b = u8x4_to_unit_f32(skin[3*(i+0)+2], skin[3*(i+1)+2],
+                                       skin[3*(i+2)+2], skin[3*(i+3)+2]);
 
         __m128 fa_im = load4_u8_to_unit_f32(im_alpha + i);
         __m128 fta   = _mm_mul_ps(ft_a, fa_im);   /* texture_alpha */
@@ -667,39 +726,27 @@ static void kernel_sse42_rgba(const uint8_t *base, const uint8_t *texture,
         __m128 gm_g = clamp01_ps128(_mm_sub_ps(_mm_add_ps(ft_g, fs_g), half));
         __m128 gm_b = clamp01_ps128(_mm_sub_ps(_mm_add_ps(ft_b, fs_b), half));
 
-        gm_r = _mm_add_ps(_mm_mul_ps(gm_r, fta), _mm_mul_ps(ft_r, fit_a));
-        gm_g = _mm_add_ps(_mm_mul_ps(gm_g, fta), _mm_mul_ps(ft_g, fit_a));
-        gm_b = _mm_add_ps(_mm_mul_ps(gm_b, fta), _mm_mul_ps(ft_b, fit_a));
+        gm_r = mul_add_ps128(gm_r, fta, _mm_mul_ps(ft_r, fit_a));
+        gm_g = mul_add_ps128(gm_g, fta, _mm_mul_ps(ft_g, fit_a));
+        gm_b = mul_add_ps128(gm_b, fta, _mm_mul_ps(ft_b, fit_a));
 
-        gm_r = _mm_add_ps(_mm_mul_ps(gm_r, invw), _mm_mul_ps(fs_r, w));
-        gm_g = _mm_add_ps(_mm_mul_ps(gm_g, invw), _mm_mul_ps(fs_g, w));
-        gm_b = _mm_add_ps(_mm_mul_ps(gm_b, invw), _mm_mul_ps(fs_b, w));
+        gm_r = mul_add_ps128(gm_r, invw, _mm_mul_ps(fs_r, w));
+        gm_g = mul_add_ps128(gm_g, invw, _mm_mul_ps(fs_g, w));
+        gm_b = mul_add_ps128(gm_b, invw, _mm_mul_ps(fs_b, w));
 
         gm_r = nan_to_num_ps128(gm_r);
         gm_g = nan_to_num_ps128(gm_g);
         gm_b = nan_to_num_ps128(gm_b);
 
-        __m128 fr = _mm_add_ps(_mm_mul_ps(gm_r, fta), _mm_mul_ps(fb_r, fit_a));
-        __m128 fg = _mm_add_ps(_mm_mul_ps(gm_g, fta), _mm_mul_ps(fb_g, fit_a));
-        __m128 fb = _mm_add_ps(_mm_mul_ps(gm_b, fta), _mm_mul_ps(fb_b, fit_a));
-
-        float rr[4], gg[4], bb[4];
-        _mm_storeu_ps(rr, fr);
-        _mm_storeu_ps(gg, fg);
-        _mm_storeu_ps(bb, fb);
-
-        for (int k = 0; k < 4; ++k) {
-            int r = (int)(rr[k] * 255.0f);
-            int g = (int)(gg[k] * 255.0f);
-            int b = (int)(bb[k] * 255.0f);
-            out[3*(i+k)+0] = (uint8_t)(r < 0 ? 0 : r > 255 ? 255 : r);
-            out[3*(i+k)+1] = (uint8_t)(g < 0 ? 0 : g > 255 ? 255 : g);
-            out[3*(i+k)+2] = (uint8_t)(b < 0 ? 0 : b > 255 ? 255 : b);
-        }
+        __m128 fr = mul_add_ps128(gm_r, fta, _mm_mul_ps(fb_r, fit_a));
+        __m128 fg = mul_add_ps128(gm_g, fta, _mm_mul_ps(fb_g, fit_a));
+        __m128 fb = mul_add_ps128(gm_b, fta, _mm_mul_ps(fb_b, fit_a));
+
+        store_unit_f32_to_u8_rgb4(fr, fg, fb, out + 3*i);
     }
 
     if (i < pixels) {
-        kernel_scalar_rgba(base + 3*i, texture + 4*i, skin + 4*i, im_alpha + i,
+        kernel_scalar_rgba(base + 3*i, texture + 4*i, skin + 3*i, im_alpha + i,
                            out + 3*i, pixels - i);
     }
 }

normal_grain_merge/normal_grain_merge.pyi

@@ -16,7 +16,7 @@ def normal_grain_merge(
     Channel ordering doesn't matter as long as it is consistent.
     :param base: The base RGB image.
     :param texture: The texture, either RGB or RGBA.
-    :param skin: The RGBA skin cutout.
+    :param skin: The RGB skin cutout.
     :param im_alpha: The alpha from the cutout.
     :param kernel: Which kernel to use.
     The `auto` kernel chooses between avx2 and sse4.2 when compiled with gcc and uses `scaler` on Windows.

{normal_grain_merge-0.0.2.dist-info → normal_grain_merge-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: normal_grain_merge
-Version: 0.0.2
+Version: 0.1.1
 Summary: Fused normal and grain merge C extension
 Author: Samuel Howard
 License: MIT
@@ -88,21 +88,22 @@ One of `KernelKind`.
 The entire reason for me writing this was NumPy being slow when this operation is in the hot path.
 So, I decided to write a SIMD version that does the type casting outside NumPy with only the intermediate values being in FP32.
 
-How much of a speedup is this? All numbers are from a Ryzen 7 4800H running Windows 11 and Python 3.12.4.
+How much of a speedup is this? All numbers are from a Ryzen 7 4800H running Ubuntu 24.04 and Python 3.12.3.
 
 | Method/Kernel     | Average Iteration Time |
 |-------------------|------------------------|
-| C scalar kernel   | 0.019565s              |
-| C SSE4.2 kernel   | 0.013705s              |
-| C AVX2 kernel     | 0.016842s              |
-| NumPy version     | 0.228098s              |
-| Old NumPy version | 0.350554s              |
+| C scalar kernel   | 0.016076s              |
+| C SSE4.2 kernel   | 0.007300s              |
+| C AVX2 kernel     | 0.007113s              |
+| NumPy version     | 0.169621s              |
+| Old NumPy version | 0.254648s              |
 
 | Method Comparison  | Speedup  |
 |--------------------|----------|
-| NumPy -> scalar    | 91.4227% |
-| NumPy -> SSE4.2    | 93.9915% |
-| NumPy -> AVX2      | 92.6165% |
-| Old np -> SSE4.2   | 96.0904% |
-| C scalar -> SSE4.2 | 29.9487% |
-| C scalar -> AVX2   | 13.9183% |
+| NumPy -> scalar    | 90.5223% |
+| NumPy -> SSE4.2    | 95.6965% |
+| NumPy -> AVX2      | 95.8063% |
+| Old np -> SSE4.2   | 97.1334% |
+| Old np -> AVX2     | 97.2066% |
+| C scalar -> SSE4.2 | 54.5933% |
+| C scalar -> AVX2   | 55.7525% |

normal_grain_merge-0.1.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+normal_grain_merge/__init__.py,sha256=Roc1wQ7_13LG_Z3Bd82zhk8wn7R1BrcO63fCdsvnnJU,89
+normal_grain_merge/kernel_kind.py,sha256=3cP4WRQSG9ZZeHsrXpXJ5Kcc8wABsmRSgex0rwRT8K4,162
+normal_grain_merge/normal_grain_merge.c,sha256=6bwB2LgwW6jzENlRuZVpzXCdQhsnjYSqIcWb0ua0948,37118
+normal_grain_merge/normal_grain_merge.cp313-win_amd64.pyd,sha256=Vucn7WsXnJBZ0oFpUq8wGPrfuNw9SL2aD9lh_zoKru0,24576
+normal_grain_merge/normal_grain_merge.pyi,sha256=HXa55A0wdcmzPpJzi7qgJws5y2q_uGjdJZQXzTkw9vc,1089
+normal_grain_merge-0.1.1.dist-info/licenses/LICENSE,sha256=qbUDFP46iOpV1ouBhpqjX-kS_cCVMHgrLBNcdTlq7Qc,1089
+normal_grain_merge-0.1.1.dist-info/METADATA,sha256=iyQQ6xdYZSd_NxoT9s1yQ4ZvtlsJm8f2cUw8mG7a_30,3597
+normal_grain_merge-0.1.1.dist-info/WHEEL,sha256=qV0EIPljj1XC_vuSatRWjn02nZIz3N1t8jsZz7HBr2U,101
+normal_grain_merge-0.1.1.dist-info/top_level.txt,sha256=jfUAUKWrxBshHvZ0xTu3uF5VJsUpbWp5NkxUj8OXqu8,19
+normal_grain_merge-0.1.1.dist-info/RECORD,,

normal_grain_merge-0.0.2.dist-info/RECORD

@@ -1,10 +0,0 @@
-normal_grain_merge/__init__.py,sha256=Roc1wQ7_13LG_Z3Bd82zhk8wn7R1BrcO63fCdsvnnJU,89
-normal_grain_merge/kernel_kind.py,sha256=3cP4WRQSG9ZZeHsrXpXJ5Kcc8wABsmRSgex0rwRT8K4,162
-normal_grain_merge/normal_grain_merge.c,sha256=n2dJ-E_DlpKtDoLW2oQ6XoLDrDPKWx_DknA_lbzuB-g,36136
-normal_grain_merge/normal_grain_merge.cp313-win_amd64.pyd,sha256=UAa_pr5gm3Wkgs49Hq7bI3ZMd0CCy16DGnNTiz7ICfc,25088
-normal_grain_merge/normal_grain_merge.pyi,sha256=Tz5RVlNbBqn_MsQ46WikaohEPctHdWsFxK3bloRZl1M,1090
-normal_grain_merge-0.0.2.dist-info/licenses/LICENSE,sha256=qbUDFP46iOpV1ouBhpqjX-kS_cCVMHgrLBNcdTlq7Qc,1089
-normal_grain_merge-0.0.2.dist-info/METADATA,sha256=jNq6oUROWU6TdMcxyrsPzaQIHKUIHRXwNa8riynTC_s,3560
-normal_grain_merge-0.0.2.dist-info/WHEEL,sha256=qV0EIPljj1XC_vuSatRWjn02nZIz3N1t8jsZz7HBr2U,101
-normal_grain_merge-0.0.2.dist-info/top_level.txt,sha256=jfUAUKWrxBshHvZ0xTu3uF5VJsUpbWp5NkxUj8OXqu8,19
-normal_grain_merge-0.0.2.dist-info/RECORD,,