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);
        let off = beg.align_offset(32);
        if off != 0 && off < remaining {
            (beg, line) = lines_fwd_fallback(beg, beg.add(off), line, line_stop);
            remaining -= off;
        }

        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;

Review comment (performance):
We can use the SIMD code to do this instead of doing a byte-wise search here. This will hurt the throughput of the `lines_fwd/1` case, but that doesn't matter much, because the latency is much more important (which should be within a few ns). Basically, modify the SIMD loop to this C pseudo-code:
```
beg = (beg + 128) & ~31;
remaining = /* some equivalent of this, but I haven't though it through */
```

I'm curious what difference this would make.

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;