diff --git a/linalg/benches/activations_avx512.rs b/linalg/benches/activations_avx512.rs index f29ca73152..0e88f945f1 100644 --- a/linalg/benches/activations_avx512.rs +++ b/linalg/benches/activations_avx512.rs @@ -2,9 +2,10 @@ // their x86 predecessor. // // sigmoid, tanh : predecessor = FMA (256-bit, 8-wide) kernel -// hardswish, leaky_relu, -// silu, gelu : predecessor = generic scalar kernel +// hardswish, leaky_relu, gelu: predecessor = generic scalar kernel // (no FMA kernel exists on x86) +// silu : predecessor = generic scalar kernel +// (the FMA kernel is benched in silu.rs) // // All buffers are 64-byte aligned (AVX-512 alignment_bytes) and a multiple of // 64 elements so every kernel's nr() divides the length. Criterion reports the diff --git a/linalg/benches/silu.rs b/linalg/benches/silu.rs index acc44ebe20..77a3971e51 100644 --- a/linalg/benches/silu.rs +++ b/linalg/benches/silu.rs @@ -6,7 +6,10 @@ use tract_linalg::element_wise::ElementWiseKer; fn silu_f32(c: &mut Criterion) { let mut group = c.benchmark_group("silu_f32"); group.throughput(Throughput::Elements(1024)); - let mut input = unsafe { Tensor::uninitialized_aligned::(&[1024], 16).unwrap() }; + // The per-arch entries call the kernels through ElementWiseKer::run, which skips + // map_slice_with_alignment: the buffer must meet every kernel's alignment_bytes + // itself (32 for the FMA ymm kernel, 64 for the AVX-512 zmm one) or they fault. + let mut input = unsafe { Tensor::uninitialized_aligned::(&[1024], 64).unwrap() }; let input = unsafe { input.as_slice_mut_unchecked::() }; for (i, x) in input.iter_mut().enumerate() { *x = (i as f32 / 10.0).sin() * 5.0; @@ -21,6 +24,12 @@ fn silu_f32(c: &mut Criterion) { group.bench_function("linalg-asm-fused", |b| { b.iter(|| tract_linalg::arm64::arm64simd_silu_f32_4n_fused::run(input, ())) }); + #[cfg(target_arch = "x86_64")] + if is_x86_feature_detected!("fma") { + group.bench_function("linalg-asm-fused", |b| { + b.iter(|| tract_linalg::x86_64_fma::fma_silu_f32::run(input, ())) + }); + } } #[inline(never)] diff --git a/linalg/build.rs b/linalg/build.rs index d76f5d510d..e769f46731 100644 --- a/linalg/build.rs +++ b/linalg/build.rs @@ -358,6 +358,7 @@ fn main() { let _ = fs::remove_file("fma_mmm_f32_16x6.asm"); let _ = fs::remove_file("fma_mmm_i32_8x8.asm"); let _ = fs::remove_file("fma_sigmoid_f32.asm"); + let _ = fs::remove_file("fma_silu_f32.asm"); let _ = fs::remove_file("fma_tanh_f32.asm"); } } else { diff --git a/linalg/src/x86_64_fma.rs b/linalg/src/x86_64_fma.rs index 21acfe8853..b8c941e449 100644 --- a/linalg/src/x86_64_fma.rs +++ b/linalg/src/x86_64_fma.rs @@ -69,6 +69,7 @@ const AVX512VNNI: fn() -> bool = || is_x86_feature_detected!("avx512vnni"); tanh_impl!(f32, fma_tanh_f32, 8, 8, is_x86_feature_detected!("fma")); sigmoid_impl!(f32, fma_sigmoid_f32, 8, 8, is_x86_feature_detected!("fma")); +silu_impl!(f32, fma_silu_f32, 8, 8, is_x86_feature_detected!("fma")); // AVX-512 (zmm, 16-wide) variants. The assembly lives in x86_64/avx512/; the // main loop handles 64 lanes (4 zmm) per iteration with a 16-lane tail, so @@ -83,13 +84,14 @@ fn plug_fma(ops: &mut Ops) { ops.sigmoid_f32 = Box::new(|| fma_sigmoid_f32::ew()); ops.tanh_f32 = Box::new(|| fma_tanh_f32::ew()); + ops.silu_f32 = Box::new(|| fma_silu_f32::ew()); ops.mul_by_scalar_f32 = Box::new(|| by_scalar::x86_64_avx_f32_mul_by_scalar_32n::ew()); ops.max_f32 = Box::new(|| max::x86_64_fma_max_f32_32n::red()); ops.min_f32 = Box::new(|| min::x86_64_fma_min_f32_32n::red()); ops.softmax2_fastcompact_f32 = Box::new(|| x86_64_fma_softmax2_fastcompact_f32_32n::red()); - log::info!("sigmoid_f32, tanh_f32: x86_64/fma activated"); + log::info!("sigmoid_f32, tanh_f32, silu_f32: x86_64/fma activated"); } /// On hosts that also support AVX-512_FP16 (Sapphire Rapids / Granite Rapids / diff --git a/linalg/x86_64/fma/fma_silu_f32.S.j2 b/linalg/x86_64/fma/fma_silu_f32.S.j2 new file mode 100644 index 0000000000..f79982447f --- /dev/null +++ b/linalg/x86_64/fma/fma_silu_f32.S.j2 @@ -0,0 +1,305 @@ +{# +// vim: set syntax=asm : + +Fused SiLU kernel. + +Per element (z = clamp(x, -18.6, 18.6), w = z^2): + SiLU(x) = f * (0.5 + z * P(w) / Q(w)) +where P is the degree-6 Horner polynomial over the numerator coeffs, Q the +degree-3 Horner polynomial over the denominator coeffs, and f = max(x, -18.6) +the factor the sigmoid multiplies. Reuses the coefficients and the rational +approximation of fma_sigmoid_f32.S.j2. + +f is clamped below at -18.6 but left unclamped above: the upper clamp only +keeps the sigmoid polynomial argument z in range, whereas SiLU(x) ~ x as +x -> +inf so the factor must stay unbounded above. The lower clamp keeps the +negative tail bounded: x < -18.6 saturates at -18.6 * sigmoid(-18.6), which +the true SiLU approaches from below as x -> -inf. + +x86-64 has only 16 ymm registers and the sigmoid kernel already fills them with +a 4-group loop, so holding f costs half the unrolling: the main loop here runs +2 groups (16 elements), with ymm0-3 the rotating coefficients, ymm4-5 z, ymm6-7 +w, ymm8-9 the Horner accumulators and ymm10-11 f. + +System V ABI: + args: rdi, rsi, rdx, rcx, r8, r9 + preserve: rbx, rsp, rbp, r12, r13, r14, r15 + scratch: rax, rdi, rsi, rdx, rcx, r8, r9, r10, r11 + return: rax (+rdx) + +Windows ABI: + args: RCX, RDX, R8, R9 + preserve: RBX, RBP, RDI, RSI, RSP, R12, R13, R14, R15, and XMM6-15 + scratch: RAX, RCX, RDX, R8, R9, R10, R11, XMM0-5, and the upper portions of YMM0-15 and ZMM0-15 + return: rax (+rdx) + +#} + +{% if msvc %} + +_text segment +fma_silu_f32_{{suffix}} proc + +{% else %} + +.intel_syntax noprefix +.text +.p2align 5 +.globl {{G}}fma_silu_f32_{{suffix}} +{{G}}fma_silu_f32_{{suffix}}: +.cfi_startproc +{% endif %} + + push rbp + mov rbp, rsp + + +{% if family == "windows" %} +// https://www.agner.org/optimize/calling_conventions.pdf xmm6-15 are not scratch +// https://stackoverflow.com/questions/43358429/save-value-of-xmm-registers + and rsp,-16 + lea rsp,[rsp-160] + vmovaps [rsp], xmm6 + vmovaps [rsp+16*1],xmm7 + vmovaps [rsp+16*2],xmm8 + vmovaps [rsp+16*3],xmm9 + vmovaps [rsp+16*4],xmm10 + vmovaps [rsp+16*5],xmm11 + vmovaps [rsp+16*6],xmm12 + vmovaps [rsp+16*7],xmm13 + vmovaps [rsp+16*8],xmm14 + vmovaps [rsp+16*9],xmm15 + + // move around arguments to mimick SysV rdi,rsi passing + push rdi + push rsi + mov rdi, rcx + mov rsi, rdx + +{% endif %} + + push rbx + push r12 + push r13 + push r14 + push r15 + + sub rsp, 8 + +{% if family == "unix" %} +// FIXME +// .cfi_def_cfa_offset 64 +{% endif %} + + stmxcsr [rsp + 4] +{% if msvc %} + mov rax, 1FC0h +{% else %} + mov rax, 0x1FC0 +{% endif %} + mov [rsp], eax + ldmxcsr [rsp] +// ---------------------------------------------------------------------- + + cmp rsi, 0 + je {{L}}done + + cmp rsi, 16 + jl {{L}}loop_1 + +{{L}}loop_2: + + vmovaps ymm4, [rdi] + vmovaps ymm5, [rdi + 32] + + vbroadcastss ymm0, dword ptr [{{offset}} {{L}}coeffs_num_low] + vbroadcastss ymm1, dword ptr [{{offset}} {{L}}coeffs_num_high] + vbroadcastss ymm2, dword ptr [{{offset}} {{L}}coeffs_num_alpha_13] + vbroadcastss ymm3, dword ptr [{{offset}} {{L}}coeffs_num_alpha_11] + + vmaxps ymm10, ymm4, ymm0 + vmaxps ymm11, ymm5, ymm0 // ymm10..11 <- f + vbroadcastss ymm0, dword ptr [{{offset}} {{L}}coeffs_num_alpha_9] + + vminps ymm4, ymm10, ymm1 + vminps ymm5, ymm11, ymm1 // ymm4..5 <- z + vbroadcastss ymm1, dword ptr [{{offset}} {{L}}coeffs_num_alpha_7] + + vmulps ymm6, ymm4, ymm4 + vmulps ymm7, ymm5, ymm5 // ymm6..7 <- w + + vmovaps ymm8, ymm2 + vmovaps ymm9, ymm2 + vbroadcastss ymm2, dword ptr [{{offset}} {{L}}coeffs_num_alpha_5] + vfmadd132ps ymm8, ymm3, ymm6 + vfmadd132ps ymm9, ymm3, ymm7 + vbroadcastss ymm3, dword ptr [{{offset}} {{L}}coeffs_num_alpha_3] + vfmadd132ps ymm8, ymm0, ymm6 + vfmadd132ps ymm9, ymm0, ymm7 + vbroadcastss ymm0, dword ptr [{{offset}} {{L}}coeffs_num_alpha_1] + vfmadd132ps ymm8, ymm1, ymm6 + vfmadd132ps ymm9, ymm1, ymm7 + vbroadcastss ymm1, dword ptr [{{offset}} {{L}}coeffs_num_beta_6] + vfmadd132ps ymm8, ymm2, ymm6 + vfmadd132ps ymm9, ymm2, ymm7 + vbroadcastss ymm2, dword ptr [{{offset}} {{L}}coeffs_num_beta_4] + vfmadd132ps ymm8, ymm3, ymm6 + vfmadd132ps ymm9, ymm3, ymm7 + vbroadcastss ymm3, dword ptr [{{offset}} {{L}}coeffs_num_beta_2] + vfmadd132ps ymm8, ymm0, ymm6 + vfmadd132ps ymm9, ymm0, ymm7 + vbroadcastss ymm0, dword ptr [{{offset}} {{L}}coeffs_num_beta_0] + vmulps ymm4, ymm4, ymm8 + vmulps ymm5, ymm5, ymm9 // ymm4..5 <- num + + vmovaps ymm8, ymm1 + vmovaps ymm9, ymm1 + vbroadcastss ymm1, dword ptr [{{offset}} {{L}}coeffs_num_half] + vfmadd132ps ymm8, ymm2, ymm6 + vfmadd132ps ymm9, ymm2, ymm7 + vfmadd132ps ymm8, ymm3, ymm6 + vfmadd132ps ymm9, ymm3, ymm7 + vfmadd132ps ymm8, ymm0, ymm6 + vfmadd132ps ymm9, ymm0, ymm7 // ymm8..9 <- denum + + vdivps ymm4, ymm4, ymm8 + vdivps ymm5, ymm5, ymm9 + vaddps ymm4, ymm4, ymm1 + vaddps ymm5, ymm5, ymm1 // ymm4..5 <- sigmoid + + vmulps ymm4, ymm4, ymm10 + vmulps ymm5, ymm5, ymm11 // ymm4..5 <- SiLU + + vmovaps [rdi], ymm4 + vmovaps [rdi + 32], ymm5 + + add rdi, 64 + sub rsi, 16 + cmp rsi, 16 + jge {{L}}loop_2 + + cmp rsi, 0 + je {{L}}done + +{{L}}loop_1: + vmovaps ymm4, [rdi] + + vbroadcastss ymm0, dword ptr [{{offset}} {{L}}coeffs_num_low] + vbroadcastss ymm1, dword ptr [{{offset}} {{L}}coeffs_num_high] + vbroadcastss ymm2, dword ptr [{{offset}} {{L}}coeffs_num_alpha_13] + vbroadcastss ymm3, dword ptr [{{offset}} {{L}}coeffs_num_alpha_11] + + vmaxps ymm10, ymm4, ymm0 // ymm10 <- f + vbroadcastss ymm0, dword ptr [{{offset}} {{L}}coeffs_num_alpha_9] + + vminps ymm4, ymm10, ymm1 // ymm4 <- z + vbroadcastss ymm1, dword ptr [{{offset}} {{L}}coeffs_num_alpha_7] + + vmulps ymm6, ymm4, ymm4 // ymm6 <- w + + vmovaps ymm8, ymm2 + vbroadcastss ymm2, dword ptr [{{offset}} {{L}}coeffs_num_alpha_5] + vfmadd132ps ymm8, ymm3, ymm6 + vbroadcastss ymm3, dword ptr [{{offset}} {{L}}coeffs_num_alpha_3] + vfmadd132ps ymm8, ymm0, ymm6 + vbroadcastss ymm0, dword ptr [{{offset}} {{L}}coeffs_num_alpha_1] + vfmadd132ps ymm8, ymm1, ymm6 + vbroadcastss ymm1, dword ptr [{{offset}} {{L}}coeffs_num_beta_6] + vfmadd132ps ymm8, ymm2, ymm6 + vbroadcastss ymm2, dword ptr [{{offset}} {{L}}coeffs_num_beta_4] + vfmadd132ps ymm8, ymm3, ymm6 + vbroadcastss ymm3, dword ptr [{{offset}} {{L}}coeffs_num_beta_2] + vfmadd132ps ymm8, ymm0, ymm6 + vbroadcastss ymm0, dword ptr [{{offset}} {{L}}coeffs_num_beta_0] + vmulps ymm4, ymm4, ymm8 // ymm4 <- num + + vmovaps ymm8, ymm1 + vbroadcastss ymm1, dword ptr [{{offset}} {{L}}coeffs_num_half] + vfmadd132ps ymm8, ymm2, ymm6 + vfmadd132ps ymm8, ymm3, ymm6 + vfmadd132ps ymm8, ymm0, ymm6 // ymm8 <- denum + + vdivps ymm4, ymm4, ymm8 + vaddps ymm4, ymm4, ymm1 // ymm4 <- sigmoid + + vmulps ymm4, ymm4, ymm10 // ymm4 <- SiLU + + vmovaps [rdi], ymm4 + add rdi, 32 + sub rsi, 8 + jnz {{L}}loop_1 +{{L}}done: + +// ---------------------------------------------------------------------- + + ldmxcsr [rsp + 4] + + add rsp, 8 + + pop r15 + pop r14 + pop r13 + pop r12 + pop rbx + +{% if family == "windows" %} + pop rsi + pop rdi + + vmovaps xmm15, [rsp+16*9] + vmovaps xmm14, [rsp+16*8] + vmovaps xmm13, [rsp+16*7] + vmovaps xmm12, [rsp+16*6] + vmovaps xmm11, [rsp+16*5] + vmovaps xmm10, [rsp+16*4] + vmovaps xmm9, [rsp+16*3] + vmovaps xmm8, [rsp+16*2] + vmovaps xmm7, [rsp+16*1] + vmovaps xmm6, [rsp] +{% endif %} + + mov rsp, rbp + pop rbp + ret + +{% set float %}{% if msvc %} real4 {%else%} .float {%endif%}{% endset %} + +{{L}}coeffs_num_low: + {{float}} -18.6 // low +{{L}}coeffs_num_high: + {{float}} 18.6 // high + +{{L}}coeffs_num_alpha_13: + {{float}} -4.433153405e-18 +{{L}}coeffs_num_alpha_11: + {{float}} 1.169974371e-14 +{{L}}coeffs_num_alpha_9: + {{float}} -1.875289645e-11 +{{L}}coeffs_num_alpha_7: + {{float}} 4.257889523e-8 +{{L}}coeffs_num_alpha_5: + {{float}} 0.00004811817576 +{{L}}coeffs_num_alpha_3: + {{float}} 0.008163842030 +{{L}}coeffs_num_alpha_1: + {{float}} 0.2499999971 + +{{L}}coeffs_num_beta_6: + {{float}} 3.922935744e-6 +{{L}}coeffs_num_beta_4: + {{float}} 0.001524872358 +{{L}}coeffs_num_beta_2: + {{float}} 0.1159886749 +{{L}}coeffs_num_beta_0: + {{float}} 1.0; + +{{L}}coeffs_num_half: + {{float}} 0.5 + +{% if msvc %} +fma_silu_f32_{{suffix}} endp +_text ends +end +{% else %} +.cfi_endproc +{% endif %}