I am too stupid to use AVX-512
Recently I have been working on writing a simple 4x4 matrix implementation in rust using SIMD intrinsics. I defined the following struct to help with this.
#[repr(C, align(16))]
#[derive(Copy, Clone)]
union f32x16 {
#[cfg(target_feature = "sse")]
pub sse: [__m128; 4],
#[cfg(target_feature = "avx")]
pub avx: [__m256; 2],
#[cfg(target_feature = "avx512f")]
pub avx512: __m512,
pub data: [f32; 16],
}
#[repr(transparent)]
#[derive(Copy, Clone)]
pub struct Mat4(f32x16);
My initial implementations of matrix multiplication were slower than a standard avx2 implementation. I brushed it off as a skill issue from my part. But when I tried benchmarking 2 different matrix transpose implementations I realized that something seriously wrong is going on.
The first AVX2 implementation
#[cfg(all(target_feature = "avx2", not(target_feature = "avx512f")))]
#[inline(always)]
pub fn transpose(&self) -> Self {
// Mat4 layout: avx[0] = [c0₀ c0₁ c0₂ c0₃ | c1₀ c1₁ c1₂ c1₃]
// avx[1] = [c2₀ c2₁ c2₂ c2₃ | c3₀ c3₁ c3₂ c3₃]
// Goal rows: res[0] = [c0₀ c1₀ c2₀ c3₀ | c0₁ c1₁ c2₁ c3₁] (Rows 0, 1)
// res[1] = [c0₂ c1₂ c2₂ c3₂ | c0₃ c1₃ c2₃ c3₃] (Rows 2, 3)
use std::arch::x86_64::{_mm256_permutevar8x32_ps, _mm256_blend_ps, _mm256_set_epi32};
unsafe {
// Indices for Row 0/1: [idx5, idx1, idx5, idx1, idx4, idx0, idx4, idx0]
let idx01 = _mm256_set_epi32(5, 1, 5, 1, 4, 0, 4, 0);
let t0 = _mm256_permutevar8x32_ps(self.0.avx[0], idx01);
let t1 = _mm256_permutevar8x32_ps(self.0.avx[1], idx01);
let res01 = _mm256_blend_ps(t0, t1, 0b11001100);
// Indices for Row 2/3: [idx7, idx3, idx7, idx3, idx6, idx2, idx6, idx2]
let idx23 = _mm256_set_epi32(7, 3, 7, 3, 6, 2, 6, 2);
let t2 = _mm256_permutevar8x32_ps(self.0.avx[0], idx23);
let t3 = _mm256_permutevar8x32_ps(self.0.avx[1], idx23);
let res23 = _mm256_blend_ps(t2, t3, 0b11001100);
Mat4(f32x16 { avx: [res01, res23] })
}
}
And the AVX-512 implementation
#[cfg(target_feature = "avx512f")]
#[inline(always)]
pub fn transpose(&self) -> Self {
use std::arch::x86_64::{_mm512_permutexvar_ps, _mm512_set_epi32};
unsafe {
let indices = _mm512_set_epi32(
15, 11, 7, 3, // Row 3 maps to Col 3
14, 10, 6, 2, // Row 2 maps to Col 2
13, 9, 5, 1, // Row 1 maps to Col 1
12, 8, 4, 0 // Row 0 maps to Col 0
);
let transposed = _mm512_permutexvar_ps(indices, self.0.avx512);
Mat4(f32x16 { avx512: transposed })
}
}
Now the AVX-512 implementation is obviously much simpler, its only 1 instruction compared to 6, and I assumed that it would outperform the AVX2 implementation, but I was wrong.
Matrix Transpose/Mat4 Transpose
time: [4.0170 ns 4.0173 ns 4.0176 ns]
change: [+600.21% +600.31% +600.45%] (p = 0.00 < 0.05)
Performance has regressed.
Found 8 outliers among 100 measurements (8.00%)
5 (5.00%) low mild
1 (1.00%) high mild
2 (2.00%) high severe
600% slower. So I decided to paste both implementations in godbolt and see the asm output
AVX2 Implementation
.LCPI0_0:
.zero 4
.zero 4
.long 0
.long 4
.zero 4
.zero 4
.long 1
.long 5
.LCPI0_1:
.long 0
.long 4
.zero 4
.zero 4
.long 1
.long 5
.zero 4
.zero 4
.LCPI0_2:
.zero 4
.zero 4
.long 2
.long 6
.zero 4
.zero 4
.long 3
.long 7
.LCPI0_3:
.long 2
.long 6
.zero 4
.zero 4
.long 3
.long 7
.zero 4
.zero 4
example::main::hb98fc185a6f3c541:
push rbp
mov rbp, rsp
and rsp, -32
sub rsp, 160
vxorps xmm0, xmm0, xmm0
vmovaps ymmword ptr [rsp + 32], ymm0
vmovaps ymmword ptr [rsp], ymm0
mov rax, rsp
vmovaps ymm0, ymmword ptr [rsp]
vmovaps ymm1, ymmword ptr [rsp + 32]
vmovaps ymm2, ymmword ptr [rip + .LCPI0_0]
vpermps ymm2, ymm2, ymm1
vmovaps ymm3, ymmword ptr [rip + .LCPI0_1]
vpermps ymm3, ymm3, ymm0
vblendps ymm2, ymm3, ymm2, 204
vmovaps ymm3, ymmword ptr [rip + .LCPI0_2]
vpermps ymm1, ymm3, ymm1
vmovaps ymm3, ymmword ptr [rip + .LCPI0_3]
vpermps ymm0, ymm3, ymm0
vblendps ymm0, ymm0, ymm1, 204
vmovaps ymmword ptr [rsp + 64], ymm2
vmovaps ymmword ptr [rsp + 96], ymm0
lea rax, [rsp + 64]
mov rsp, rbp
pop rbp
vzeroupper
ret
AVX-512 Implementation
.LCPI0_0:
.long 0
.long 4
.long 8
.long 12
.long 1
.long 5
.long 9
.long 13
.long 2
.long 6
.long 10
.long 14
.long 3
.long 7
.long 11
.long 15
example::main::hb98fc185a6f3c541:
push rbp
mov rbp, rsp
and rsp, -64
sub rsp, 192
vxorps xmm0, xmm0, xmm0
vmovaps ymmword ptr [rsp + 32], ymm0
vmovaps ymmword ptr [rsp], ymm0
mov rax, rsp
vmovaps zmm0, zmmword ptr [rip + .LCPI0_0]
vpermps zmm0, zmm0, zmmword ptr [rsp]
vmovaps zmmword ptr [rsp + 64], zmm0
lea rax, [rsp + 64]
mov rsp, rbp
pop rbp
vzeroupper
ret
both look pretty normal to me. So then I tried letting LLVM autovectorize targeting x86-64-v3 for avx2 and x86-64-v4 for AVX-512.
x86-64-v3
.LCPI0_0:
.zero 4
.zero 4
.long 0
.long 4
.zero 4
.zero 4
.long 1
.long 5
.LCPI0_1:
.long 0
.long 4
.zero 4
.zero 4
.long 1
.long 5
.zero 4
.zero 4
.LCPI0_2:
.zero 4
.zero 4
.zero 4
.long 0
.zero 4
.zero 4
.zero 4
.long 1
.LCPI0_3:
.zero 4
.zero 4
.long 0
.zero 4
.zero 4
.zero 4
.long 1
.zero 4
example::main::hb98fc185a6f3c541:
push rbp
mov rbp, rsp
and rsp, -32
sub rsp, 160
vxorps xmm0, xmm0, xmm0
vmovaps ymmword ptr [rsp + 32], ymm0
vmovaps ymmword ptr [rsp], ymm0
mov rax, rsp
vmovaps ymm0, ymmword ptr [rip + .LCPI0_0]
vpermps ymm0, ymm0, ymmword ptr [rsp + 32]
vmovaps ymm1, ymmword ptr [rip + .LCPI0_1]
vpermps ymm2, ymm1, ymmword ptr [rsp]
vblendps ymm0, ymm2, ymm0, 204
vmovaps ymmword ptr [rsp + 64], ymm0
vmovups xmm0, xmmword ptr [rsp + 40]
vmovsd xmm2, qword ptr [rsp + 56]
vmovaps ymm3, ymmword ptr [rip + .LCPI0_2]
vpermps ymm2, ymm3, ymm2
vmovaps ymm3, ymmword ptr [rip + .LCPI0_3]
vpermps ymm0, ymm3, ymm0
vpermps ymm1, ymm1, ymmword ptr [rsp + 8]
vblendps ymm0, ymm1, ymm0, 204
vblendps ymm0, ymm0, ymm2, 136
vmovaps ymmword ptr [rsp + 96], ymm0
lea rax, [rsp + 64]
mov rsp, rbp
pop rbp
vzeroupper
ret
x86-64-v4
.LCPI0_0:
.long 0
.long 4
.long 8
.long 12
.long 1
.long 5
.long 9
.long 13
.LCPI0_1:
.long 8
.long 12
.long 0
.zero 4
.long 9
.long 13
.long 1
.zero 4
.LCPI0_2:
.long 0
.long 1
.long 2
.long 8
.long 4
.long 5
.long 6
.long 9
example::main::hb98fc185a6f3c541:
push rbp
mov rbp, rsp
and rsp, -64
sub rsp, 192
vxorps xmm0, xmm0, xmm0
vmovaps ymmword ptr [rsp + 32], ymm0
vmovaps ymmword ptr [rsp], ymm0
mov rax, rsp
vmovaps ymm0, ymmword ptr [rip + .LCPI0_0]
vpermps zmm0, zmm0, zmmword ptr [rsp]
vmovaps ymmword ptr [rsp + 64], ymm0
vmovups xmm0, xmmword ptr [rsp + 40]
vmovsd xmm1, qword ptr [rsp + 56]
vmovaps ymm2, ymmword ptr [rip + .LCPI0_1]
vpermi2ps ymm2, ymm0, ymmword ptr [rsp + 8]
vmovaps ymm0, ymmword ptr [rip + .LCPI0_2]
vpermi2ps ymm0, ymm2, ymm1
vmovaps ymmword ptr [rsp + 96], ymm0
lea rax, [rsp + 64]
mov rsp, rbp
pop rbp
vzeroupper
ret
Looks about the same as my implementation, but more importantly, I observed the same performance regression.
For some reason LLVM literally generates SLOWER code if you target x86-64-v4 than if you target x86-64-v3. For context my cpu is a ryzen 7 9800x3d and I used the crate "criterion" for benchmarking.
Can someone more qualified try to explain everything to me? Why is _mm512_permutexvar_ps SO SLOW? Is it only a Zen 5 thing? I assume its an architecture thing considering that LLVM also generated basically the same code as me?