SpecForElem for i16/u16 and other digits

[pickfire]

We have specialization for SpecForElem for i8 and u8 (which provides a small performance win over plain extend for vec) but we could look into using specialization for that involves duplicating more digits with rep stosw (or d or q). It may need benchmarking to ensure that it's actually faster as mentioned by @joshtriplett

I am not sure how to reproduce the assembly rep stosw, someone that know could take up this issue.

Previous discussion: https://rust-lang.zulipchat.com/#narrow/stream/219381-t-libs/topic/SpecForElem.20for.20other.20integers

[joshtriplett]

For the purposes of a quick test, at least, we could use inline asm and a target-specific specialization. If that shows no substantial performance increase then there's no point pursuing this further.

[pickfire]

Looking at the generated assembly, I don't know where to insert stosw.

Rust code

#[inline(never)]
fn from_elem_u16_std(n: u16) -> Vec<u16> {
    std::vec::from_elem(n, 100)
}

 fn from_elem_u16_std(n: u16) -> Vec<u16> {
 push    rbx
 mov     rbx, rdi
     unsafe { __rust_alloc(layout.size(), layout.align()) } (liballoc/alloc.rs:80)
     mov     edi, 200
     mov     esi, 2
     call    qword, ptr, [rip, +, __rust_alloc@GOTPCREL]
     Ok(memory) => memory, (liballoc/raw_vec.rs:184)
     test    rax, rax
     je      .LBB7_1
     movabs  rcx, 723401728380766730
     unsafe { intrinsics::move_val_init(&mut *dst, src) } (libcore/ptr/mod.rs:901)
     mov     qword, ptr, [rax, +, 192], rcx
     movaps  xmm0, xmmword, ptr, [rip, +, .LCPI7_0]
     movups  xmmword, ptr, [rax, +, 176], xmm0
     movups  xmmword, ptr, [rax, +, 160], xmm0
     movups  xmmword, ptr, [rax, +, 144], xmm0
     movups  xmmword, ptr, [rax, +, 128], xmm0
     movups  xmmword, ptr, [rax, +, 112], xmm0
     movups  xmmword, ptr, [rax, +, 96], xmm0
     movups  xmmword, ptr, [rax, +, 80], xmm0
     movups  xmmword, ptr, [rax, +, 64], xmm0
     movups  xmmword, ptr, [rax, +, 48], xmm0
     movups  xmmword, ptr, [rax, +, 32], xmm0
     movups  xmmword, ptr, [rax, +, 16], xmm0
     movups  xmmword, ptr, [rax], xmm0
     v (liballoc/vec.rs:1842)
     mov     qword, ptr, [rbx], rax
     movaps  xmm0, xmmword, ptr, [rip, +, .LCPI7_1]
     movups  xmmword, ptr, [rbx, +, 8], xmm0
 }
 pop     rbx
 ret
.LBB7_1:
     Err(_) => handle_alloc_error(layout), (liballoc/raw_vec.rs:185)
     mov     edi, 200
     mov     esi, 2
     call    qword, ptr, [rip, +, _ZN5alloc5alloc18handle_alloc_error17h7f6a7a00afe6ad92E@GOTPCREL]
     ud2

Same version but for u8

 fn from_elem_u8_std(n: u8) -> Vec<u8> {
 push    rbx
 mov     rbx, rdi
     unsafe { __rust_alloc(layout.size(), layout.align()) } (liballoc/alloc.rs:80)
     mov     edi, 100
     mov     esi, 1
     call    qword, ptr, [rip, +, __rust_alloc@GOTPCREL]
     Ok(memory) => memory, (liballoc/raw_vec.rs:184)
     test    rax, rax
     je      .LBB5_1
     unsafe { write_bytes(dst, val, count) } (libcore/intrinsics.rs:2291)
     movaps  xmm0, xmmword, ptr, [rip, +, .LCPI5_0]
     movups  xmmword, ptr, [rax, +, 80], xmm0
     movups  xmmword, ptr, [rax, +, 64], xmm0
     movups  xmmword, ptr, [rax, +, 48], xmm0
     movups  xmmword, ptr, [rax, +, 32], xmm0
     movups  xmmword, ptr, [rax, +, 16], xmm0
     movups  xmmword, ptr, [rax], xmm0
     mov     dword, ptr, [rax, +, 96], 168430090
     v (liballoc/vec.rs:1829)
     mov     qword, ptr, [rbx], rax
     movaps  xmm0, xmmword, ptr, [rip, +, .LCPI5_1]
     movups  xmmword, ptr, [rbx, +, 8], xmm0
 }
 pop     rbx
 ret
.LBB5_1:
     Err(_) => handle_alloc_error(layout), (liballoc/raw_vec.rs:185)
     mov     edi, 100
     mov     esi, 1
     call    qword, ptr, [rip, +, _ZN5alloc5alloc18handle_alloc_error17h7f6a7a00afe6ad92E@GOTPCREL]
     ud2

[BenJourdan]

I ran into this because I have to allocate a vec of 5 billion 1.0f64s and found it suspiciously slow compared to filling with zeros. let data = vec![1.0f64;5_015_016_766]; took about 20 seconds. To investigate, I tried adding the following benches to alloc/benches/vec.rs:

#[bench]
fn bench_f64_spec_from_elem_zero(b: &mut Bencher) {
    b.iter(|| black_box(vec![0.0f64; 1_000]));
}

#[bench]
fn bench_f64_spec_from_elem_one(b: &mut Bencher) {
    b.iter(|| black_box(vec![1.0f64; 1_000]));
}

Running these on my machine gave

benchmarks:
    vec::bench_f64_spec_from_elem_one  265.03/iter +/- 65.44
    vec::bench_f64_spec_from_elem_zero 270.11/iter  +/- 3.57

which seems fair enough. However, when I bump the size up to 100_000_000 I get the following:

benchmarks:
    vec::bench_f64_spec_from_elem_one  414710323.70/iter +/- 30228852.17
    vec::bench_f64_spec_from_elem_zero     16351.53/iter    +/- 26020.93

I tried using rep stosq with the following but that made things worse:

#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
impl SpecFromElem for f64{
    #[inline]
    #[track_caller]
    fn from_elem<A: Allocator>(elem: f64, n: usize, alloc: A)-> Vec<f64,A>{
        if elem.to_bits() == 0 {
            return Vec { buf: RawVec::with_capacity_zeroed_in(n,alloc), len: n};
        }
        let mut v: Vec<f64,A> = Vec::with_capacity_in(n,alloc);
        unsafe{
            let dest = v.as_mut_ptr();
            core::arch::asm! {
                "rep stosq",
                inout("rcx") n => _,
                inout("rdi") dest => _,
                in("rax") elem,
                options(nostack, preserves_flags)
            };
            v.set_len(n);
        }
        v
    }
}

benchmarks:
    vec::bench_f64_spec_from_elem_one  459822483.80/iter +/- 57009209.09
    vec::bench_f64_spec_from_elem_zero     12515.96/iter     +/- 1248.89

I ran this on an old server but observed similar behaviour on my modern PC. I am a rust/assembly beginner so I imagine the asm is suboptimal. If we allow the use of multiple threads then

        let mut v = black_box(vec![0.0f64;1_000_000_000]);
        // use rayon to set to 1.0
        v.par_iter_mut().for_each(|x| *x = 1.0);

runs in ~500ms compared to ~4s for

let mut v = black_box(vec![1.0f64;1_000_000_000]);

My use case is constructing the data arrays of sparse CSR matrices. Making this change took the time to allocate the data vec from ~19 seoncds to ~2.5 seconds.

Edit:

I don't think we can do much better than allocating all zeros and then modifying them for nonzero values with types larger than a byte. We can't use std::alloc::alloc_zeroed or memset since that only works for bytes.

It should be possible to use memset with larger data types with values that are repeated bytes. For example 40404040 in hex is about 3.0 in f32. However it seems silly to include optimisations for this.