Repository: microsoft/edit
File: src/simd/lines_fwd.rs (Rust)

Code before:
unsafe fn lines_fwd_avx2(
    mut beg: *const u8,
    end: *const u8,
    mut line: CoordType,
    line_stop: CoordType,
) -> (*const u8, CoordType) {
    unsafe {
        use std::arch::x86_64::*;

        #[inline(always)]
        unsafe fn horizontal_sum_i64(v: __m256i) -> i64 {
            unsafe {
                let hi = _mm256_extracti128_si256::<1>(v);
                let lo = _mm256_castsi256_si128(v);
                let sum = _mm_add_epi64(lo, hi);
                let shuf = _mm_shuffle_epi32::<0b11_10_11_10>(sum);
                let sum = _mm_add_epi64(sum, shuf);
                _mm_cvtsi128_si64(sum)
            }
        }

        let lf = _mm256_set1_epi8(b'\n' as i8);
        let mut remaining = end.offset_from_unsigned(beg);

        if line < line_stop {
            // Unrolling the loop by 4x speeds things up by >3x.
            // It allows us to accumulate matches before doing a single `vpsadbw`.
            while remaining >= 128 {
                let v1 = _mm256_loadu_si256(beg.add(0) as *const _);
                let v2 = _mm256_loadu_si256(beg.add(32) as *const _);
                let v3 = _mm256_loadu_si256(beg.add(64) as *const _);
                let v4 = _mm256_loadu_si256(beg.add(96) as *const _);

                // `vpcmpeqb` leaves each comparison result byte as 0 or -1 (0xff).
                // This allows us to accumulate the comparisons by subtracting them.
                let mut sum = _mm256_setzero_si256();
                sum = _mm256_sub_epi8(sum, _mm256_cmpeq_epi8(v1, lf));
                sum = _mm256_sub_epi8(sum, _mm256_cmpeq_epi8(v2, lf));
                sum = _mm256_sub_epi8(sum, _mm256_cmpeq_epi8(v3, lf));
                sum = _mm256_sub_epi8(sum, _mm256_cmpeq_epi8(v4, lf));

                // Calculate the total number of matches in this chunk.
                let sum = _mm256_sad_epu8(sum, _mm256_setzero_si256());
                let sum = horizontal_sum_i64(sum);

                let line_next = line + sum as CoordType;
                if line_next >= line_stop {
                    break;
                }

                beg = beg.add(128);

Review comment (performance):
I've benchmarked this locally and I have to say that I don't understand why and how llvm optimizes the code here. Even something as benign as inlining this call and doing nothing otherwise breaks the performance uplift of this PR in realistic scenarios (tested via `cargo bench TextBuffer`). I suspect that this code is somewhat brittle and may break at any point in the future.

Code after:
unsafe fn lines_fwd_avx2(
    mut beg: *const u8,
    end: *const u8,
    mut line: CoordType,
    line_stop: CoordType,
) -> (*const u8, CoordType) {
    unsafe {
        use std::arch::x86_64::*;

        #[inline(always)]
        unsafe fn horizontal_sum_i64(v: __m256i) -> i64 {
            unsafe {
                let hi = _mm256_extracti128_si256::<1>(v);
                let lo = _mm256_castsi256_si128(v);
                let sum = _mm_add_epi64(lo, hi);
                let shuf = _mm_shuffle_epi32::<0b11_10_11_10>(sum);
                let sum = _mm_add_epi64(sum, shuf);
                _mm_cvtsi128_si64(sum)
            }
        }

        let lf = _mm256_set1_epi8(b'\n' as i8);
        let off = beg.align_offset(32);
        if off != 0 && off < end.offset_from_unsigned(beg) {
            (beg, line) = lines_fwd_fallback(beg, beg.add(off), line, line_stop);
        }

        if line < line_stop {
            // Unrolling the loop by 4x speeds things up by >3x.
            // It allows us to accumulate matches before doing a single `vpsadbw`.
            while end.offset_from_unsigned(beg) >= 128 {
                let v1 = _mm256_loadu_si256(beg.add(0) as *const _);
                let v2 = _mm256_loadu_si256(beg.add(32) as *const _);
                let v3 = _mm256_loadu_si256(beg.add(64) as *const _);
                let v4 = _mm256_loadu_si256(beg.add(96) as *const _);

                // `vpcmpeqb` leaves each comparison result byte as 0 or -1 (0xff).
                // This allows us to accumulate the comparisons by subtracting them.
                let mut sum = _mm256_setzero_si256();
                sum = _mm256_sub_epi8(sum, _mm256_cmpeq_epi8(v1, lf));
                sum = _mm256_sub_epi8(sum, _mm256_cmpeq_epi8(v2, lf));
                sum = _mm256_sub_epi8(sum, _mm256_cmpeq_epi8(v3, lf));
                sum = _mm256_sub_epi8(sum, _mm256_cmpeq_epi8(v4, lf));

                // Calculate the total number of matches in this chunk.
                let sum = _mm256_sad_epu8(sum, _mm256_setzero_si256());
                let sum = horizontal_sum_i64(sum);

                let line_next = line + sum as CoordType;
                if line_next >= line_stop {
                    break;