Compiler Explorer

Source code

use std::convert::TryInto;

use std::arch::x86_64::*;

#[target_feature(enable = "sse4.1")]
pub unsafe fn transpose_8by8_sse4(
    in_matrix: &[u8],
    out: &mut [u8],
    in_stride: usize,
    out_stride: usize,
) {
    // Godbolt : https://godbolt.org/z/jan4WP1xr
    // Stack overflow: https://stackoverflow.com/a/42316675

assert!((8 * out_stride) <= out.len());

assert!((8 * in_stride) <= in_matrix.len());

let sv = _mm_set_epi8(15, 7, 14, 6, 13, 5, 12, 4, 11, 3, 10, 2, 9, 1, 8, 0);

let mut pos = 0;

// Load data from memory
    // Load 64 bites to ensure we only take 8 values
    let mn_0 = _mm_loadl_epi64(in_matrix.get_unchecked(pos..).as_ptr().cast());
    pos += in_stride;
    let mn_1 = _mm_loadl_epi64(in_matrix.get_unchecked(pos..).as_ptr().cast());
    pos += in_stride;
    let mv_0 = _mm_unpacklo_epi64(mn_0, mn_1);

let mn_2 = _mm_loadl_epi64(in_matrix.get_unchecked(pos..).as_ptr().cast());
    pos += in_stride;
    let mn_3 = _mm_loadl_epi64(in_matrix.get_unchecked(pos..).as_ptr().cast());
    pos += in_stride;
    let mv_1 = _mm_unpacklo_epi64(mn_2, mn_3);

let mn_4 = _mm_loadl_epi64(in_matrix.get_unchecked(pos..).as_ptr().cast());
    pos += in_stride;
    let mn_5 = _mm_loadl_epi64(in_matrix.get_unchecked(pos..).as_ptr().cast());
    pos += in_stride;
    let mv_2 = _mm_unpacklo_epi64(mn_4, mn_5);

let mn_6 = _mm_loadl_epi64(in_matrix.get_unchecked(pos..).as_ptr().cast());
    pos += in_stride;
    let mn_7 = _mm_loadl_epi64(in_matrix.get_unchecked(pos..).as_ptr().cast());
    let mv_3 = _mm_unpacklo_epi64(mn_6, mn_7);

let ov_0 = _mm_shuffle_epi8(mv_0, sv);
    let ov_1 = _mm_shuffle_epi8(mv_1, sv);
    let ov_2 = _mm_shuffle_epi8(mv_2, sv);
    let ov_3 = _mm_shuffle_epi8(mv_3, sv);

let iv_0 = _mm_unpacklo_epi16(ov_0, ov_1);
    let iv_1 = _mm_unpackhi_epi16(ov_0, ov_1);
    let iv_2 = _mm_unpacklo_epi16(ov_2, ov_3);
    let iv_3 = _mm_unpackhi_epi16(ov_2, ov_3);

let av_0 = _mm_unpacklo_epi32(iv_0, iv_2);
    let av_1 = _mm_unpackhi_epi32(iv_0, iv_2);
    let av_2 = _mm_unpacklo_epi32(iv_1, iv_3);
    let av_3 = _mm_unpackhi_epi32(iv_1, iv_3);

// Now we have av1 having 0-16, av2 - 16-32 etc etc
    // So we want to extract and write only 8 bytes, as that is essentially a matrix
    // transpose of a 8 by 8 matrix writing to different strides.

let sv_0 = _mm_unpackhi_epi64(av_0, _mm_setzero_si128());
    let sv_1 = _mm_unpackhi_epi64(av_1, _mm_setzero_si128());
    let sv_2 = _mm_unpackhi_epi64(av_2, _mm_setzero_si128());
    let sv_3 = _mm_unpackhi_epi64(av_3, _mm_setzero_si128());

pos = 0;
    // Ensure writes are always in bounds
    // Needed to make the below writes unsafe

_mm_storel_epi64(out.get_unchecked_mut(pos..).as_mut_ptr().cast(), av_0);
    pos += out_stride;

_mm_storel_epi64(out.get_unchecked_mut(pos..).as_mut_ptr().cast(), sv_0);
    pos += out_stride;

_mm_storel_epi64(out.get_unchecked_mut(pos..).as_mut_ptr().cast(), av_1);
    pos += out_stride;

_mm_storel_epi64(out.get_unchecked_mut(pos..).as_mut_ptr().cast(), sv_1);
    pos += out_stride;

_mm_storel_epi64(out.get_unchecked_mut(pos..).as_mut_ptr().cast(), av_2);
    pos += out_stride;

_mm_storel_epi64(out.get_unchecked_mut(pos..).as_mut_ptr().cast(), sv_2);
    pos += out_stride;

_mm_storel_epi64(out.get_unchecked_mut(pos..).as_mut_ptr().cast(), av_3);
    pos += out_stride;

_mm_storel_epi64(out.get_unchecked_mut(pos..).as_mut_ptr().cast(), sv_3);
}

fn main() {
    let t = (0..192).collect::<Vec<u8>>();
    let mut out = [0_u8; 64];
    unsafe {
        transpose_8by8_sse4(&t, &mut out, 15, 8);
    }
    print!("[");
    for ch in out.chunks_exact(8) {
        println!("{:?}", ch);
    }
}