Skip to main content

strat9_kernel/framebuffer/x86/
avx2.rs

1use core::arch::x86_64::*;
2
3#[inline]
4unsafe fn load_u32_unaligned(src: *const u8) -> i32 {
5    core::ptr::read_unaligned(src as *const u32) as i32
6}
7
8#[target_feature(enable = "avx2")]
9pub unsafe fn blend_avx2(dst: *mut u32, src: *const u32, alpha: u8, count: usize) {
10    let alpha_u16 = alpha as u16;
11    let inv_alpha = 255 - alpha_u16;
12
13    let alpha_vec = _mm256_set1_epi16(alpha_u16 as i16);
14    let inv_alpha_vec = _mm256_set1_epi16(inv_alpha as i16);
15    let zero = _mm256_setzero_si256();
16    let ones = _mm256_set1_epi16(1);
17
18    let mut i = 0;
19    while i + 8 <= count {
20        let d = _mm256_loadu_si256(dst.add(i) as *const __m256i);
21        let s = _mm256_loadu_si256(src.add(i) as *const __m256i);
22
23        let d_lo = _mm256_unpacklo_epi8(d, zero);
24        let d_hi = _mm256_unpackhi_epi8(d, zero);
25        let s_lo = _mm256_unpacklo_epi8(s, zero);
26        let s_hi = _mm256_unpackhi_epi8(s, zero);
27
28        let res_lo_s = _mm256_mullo_epi16(s_lo, alpha_vec);
29        let res_lo_d = _mm256_mullo_epi16(d_lo, inv_alpha_vec);
30        let res_lo = _mm256_add_epi16(res_lo_s, res_lo_d);
31        // Exact /255: (x + (x>>8) + 1) >> 8  (dav1d formula)
32        let lo_div = _mm256_srli_epi16(res_lo, 8);
33        let lo_corr = _mm256_add_epi16(res_lo, _mm256_add_epi16(lo_div, ones));
34        let res_lo_final = _mm256_srli_epi16(lo_corr, 8);
35
36        let res_hi_s = _mm256_mullo_epi16(s_hi, alpha_vec);
37        let res_hi_d = _mm256_mullo_epi16(d_hi, inv_alpha_vec);
38        let res_hi = _mm256_add_epi16(res_hi_s, res_hi_d);
39        let hi_div = _mm256_srli_epi16(res_hi, 8);
40        let hi_corr = _mm256_add_epi16(res_hi, _mm256_add_epi16(hi_div, ones));
41        let res_hi_final = _mm256_srli_epi16(hi_corr, 8);
42
43        let res = _mm256_packus_epi16(res_lo_final, res_hi_final);
44        _mm256_storeu_si256(dst.add(i) as *mut __m256i, res);
45        i += 8;
46    }
47
48    if i < count {
49        super::sse2::blend_sse2(dst.add(i), src.add(i), alpha, count - i);
50    }
51}
52
53/// BGR24 => ARGB32 conversion using AVX2 vpshufb (8 pixels/iter)
54///
55/// Load exactly 24 bytes (8 BGR pixels) without buffer overflow
56/// by constructing the YMM from two separately loaded 128-bit lanes.
57/// vpshufb operates per 128-bit lane independently, so each lane must have
58/// its 4 BGR pixels in the first 12 bytes, followed by zeros.
59#[target_feature(enable = "avx2")]
60pub unsafe fn convert_bgr_to_argb_avx2(dst: *mut u32, src: *const u8, count: usize) {
61    // Masque par lane 128-bit: chaque lane traite 4 pixels BGR -> B G R 0
62    // (les index 0-11 réfèrent aux 12 premiers octets de la lane)
63    #[rustfmt::skip]
64    let per_lane = _mm_set_epi8(
65        -1, 11, 10,  9,    // pixel 3: 0, R3, G3, B3
66        -1,  8,  7,  6,    // pixel 2: 0, R2, G2, B2
67        -1,  5,  4,  3,    // pixel 1: 0, R1, G1, B1
68        -1,  2,  1,  0,    // pixel 0: 0, R0, G0, B0
69    );
70    let shuffle_mask = _mm256_broadcastsi128_si256(per_lane);
71    // Alpha mask: sets byte 3,7,11,15,19,23,27,31 to 0xFF
72    let alpha_mask = _mm256_set1_epi32(0xFF000000_u32 as i32);
73
74    let mut i = 0;
75    while i + 8 <= count {
76        // Load 12B low lane  (pixels 0-3)
77        let lo8 = _mm_loadl_epi64(src.add(i * 3) as *const __m128i);
78        let hi4 = _mm_cvtsi32_si128(load_u32_unaligned(src.add(i * 3 + 8)));
79        let lo12 = _mm_unpacklo_epi64(lo8, hi4);
80        // Load 12B high lane (pixels 4-7)
81        let hi8 = _mm_loadl_epi64(src.add(i * 3 + 12) as *const __m128i);
82        let hi4_ = _mm_cvtsi32_si128(load_u32_unaligned(src.add(i * 3 + 20)));
83        let hi12 = _mm_unpacklo_epi64(hi8, hi4_);
84
85        let raw = _mm256_set_m128i(hi12, lo12);
86        let shuffled = _mm256_shuffle_epi8(raw, shuffle_mask);
87        let result = _mm256_or_si256(shuffled, alpha_mask);
88        _mm256_storeu_si256(dst.add(i) as *mut __m256i, result);
89        i += 8;
90    }
91
92    // tail : pixel going through SSE2 -> SSSE3 -> generic -
93    if i < count {
94        super::sse2::convert_bgr_to_argb_sse2(dst.add(i), src.add(i * 3), count - i);
95    }
96}