Weird codegen with avx512 feature and masks > 8 lanes

[jhorstmann]

I'm investigating performance differences between packed_simd and this crate. My assumption was that the separate bitmask representation used for avx512 could lead to better performance but instead I see some rather weird assembly for mask handling. This seems to be specific to masks with more than 8 lanes. The following is a reduced example. The same behavior can also be seen with other comparison operations like simd_eq instead of is_nan.

#[inline(never)]
fn nan_bitmask_16(data: &[f32; 16]) -> u16 {
    let chunk = f32x16::from_slice(data);

    let is_nan = chunk.is_nan();
    is_nan.to_bitmask()
}

#[inline(never)]
fn nan_bitmask_8(data: &[f32; 8]) -> u8 {
    let chunk = f32x8::from_slice(data);

    let is_nan = chunk.is_nan();
    is_nan.to_bitmask()
}

The generated code for 8 lanes looks very good:

vxorps xmm0,xmm0,xmm0
vcmpunordps k0,ymm0,YMMWORD PTR [rdi]
kmovd  eax,k0
vzeroupper 
ret

But for 16 lanes something strange is happening:

push   rax
vxorps xmm0,xmm0,xmm0
vcmpunordps k0,zmm0,ZMMWORD PTR [rdi]
kmovw  WORD PTR [rsp],k0
kmovd  eax,k0
movzx  ecx,BYTE PTR [rsp+0x1]
shl    ecx,0x8
movzx  eax,al
or eax,ecx
pop rcx
vzeroupper 
ret 

The mask seems to get spilled to the stack, the high byte gets reloaded and then both halves get combined back into a 16bit word.

This was observed using portable_simd as a crate (actually as an example binary inside a repo checkout) as it was easier to try and simplify the example. I observed similar code in a bigger example using the version from the standard library with Zbuild-std.

Meta

$ rustc +nightly --version
rustc 1.66.0-nightly (81f391930 2022-10-09)

portable_simd commit aad8f0a

[calebzulawski]

Curious, have you ever seen this when not using to_bitmask? That particular function is not quite perfect, yet.

[jhorstmann]

Seems I can even simplify the function to

#[inline(never)]
fn nan_mask_16(chunk: f32x16) -> mask32x16 {
    chunk.is_nan()
}

and still get the exact same pattern. The inline(never) also doesn't change the code, just makes it much easier to find the relevant part in the binary. Unfortunately no build-std available in compiler explorer :)

[calebzulawski]

In that case, this seems like a codegen issue with the bitmask representation in the compiler intrinsics. It's possible the generated LLVM IR is too verbose and not optimizing well, or maybe there's a mismatched layout. I'm not surprised it's moving k0 but I wouldn't expect anything beyond moving it to a GP register.

[jhorstmann]

Managed to extract the relevant llvm ir and put it into compiler explorer

define i16 @mask_from_int_unchecked(ptr %value) {
start:
  %0 = alloca i16, align 1
  %_4 = alloca [2 x i8], align 1
  %mask = alloca [2 x i8], align 1
  %1 = alloca [2 x i8], align 1
  %_3 = load <16 x i32>, ptr %value, align 64
  %2 = lshr <16 x i32> %_3, <i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31, i32 31>
  %3 = trunc <16 x i32> %2 to <16 x i1>
  %4 = bitcast <16 x i1> %3 to i16
  store i16 %4, ptr %0, align 1
  %5 = load [2 x i8], ptr %0, align 1
  store [2 x i8] %5, ptr %mask, align 1
  br label %bb1

bb1:                                              ; preds = %start
  call void @llvm.memcpy.p0.p0.i64(ptr align 1 %_4, ptr align 1 %mask, i64 2, i1 false)
  call void @llvm.memcpy.p0.p0.i64(ptr align 1 %1, ptr align 1 %_4, i64 2, i1 false)
  %6 = load i16, ptr %1, align 1
  ret i16 %6
}

This matches what I see in intrinsics.rs and seems reasonable. It gets optimized by llvm to

define i16 @mask_from_int_unchecked(ptr nocapture readonly %value) local_unnamed_addr #0 {
  %_3 = load <16 x i32>, ptr %value, align 64
  %0 = icmp slt <16 x i32> %_3, zeroinitializer
  %1 = bitcast <16 x i1> %0 to i16
  %bc = bitcast <16 x i1> %0 to <2 x i8>
  %.sroa.2.0.extract.trunc = extractelement <2 x i8> %bc, i64 1
  %_4.sroa.2.0.insert.ext = zext i8 %.sroa.2.0.extract.trunc to i16
  %_4.sroa.2.0.insert.shift = shl nuw i16 %_4.sroa.2.0.insert.ext, 8
  %_4.sroa.0.0.insert.ext = and i16 %1, 255
  %_4.sroa.0.0.insert.insert = or i16 %_4.sroa.2.0.insert.shift, %_4.sroa.0.0.insert.ext
  ret i16 %_4.sroa.0.0.insert.insert
}

I'll probably open an LLVM issue tomorrow, unless you spot any issues in the input llvm.

[calebzulawski]

That looks right to me. I suspect the bitcast is throwing it off, but I'm not sure there's a better way to represent that operation.

[jhorstmann]

I forgot to update this ticket after trying a bunch of things. The linked llvm fixes improved the codegen and makes it consistently use mask registers. In this concrete example with several mask operations that was still slower than the non-avx512 version using vector masks instead of bitmasks. I'm guessing because of more execution resources or lower latency for vector mask operations, I don't think there is much that can be done about that.

I also tried the suggested code gen changes for the simd_bitmask intrinsic. Those did not seem to make any difference in this example. Maybe it does for some other usecase. That experiment can be seen at jhorstmann/rust@d8c401f. If you think that change could still be worthwhile I can clean it up and open a PR on the rust repo.