Plan for BF16 datatype ?

[pauldintel]

any optimization on SIMD plan for BFloat16 datatype ?
thanks

[ashvardanian]

Hi @pauldintel! That shouldn’t be too hard to add and can help a lot on older x86 and newer mobile CPUs. Would you like to contribute? Any specific distance functions you are looking for?

[MarkReedZ]

@ashvardanian

I'm adding support for this. Would it make sense for f16 and bf16 to use check_c_source_compiles in cmake to detect compiler support?

check_c_source_compiles(
  [=[
int
main(int argc, char *argv)
{
  __bf16 foo = 1.0;
  return 0;
}
]=]
  HAS_BFLOAT16)

We can retain the ability to disable with #define SIMSIMD_NATIVE_F16 0.

Note the bench disables native F16 - I think we can leave it on by default.

[MarkReedZ]

Current benchmark results vs native f16

dot_bf16_serial_1536d/min_time:10.000/threads:12        1372 ns
cos_bf16_serial_1536d/min_time:10.000/threads:12        1485 ns
l2sq_bf16_serial_1536d/min_time:10.000/threads:12       1393 ns
kl_bf16_serial_1536d/min_time:10.000/threads:12         3352 ns
js_bf16_serial_1536d/min_time:10.000/threads:12         5069 ns

dot_f16_serial_1536d/min_time:10.000/threads:12          264 ns
cos_f16_serial_1536d/min_time:10.000/threads:12          264 ns
l2sq_f16_serial_1536d/min_time:10.000/threads:12         264 ns
kl_f16_serial_1536d/min_time:10.000/threads:12          2983 ns
js_f16_serial_1536d/min_time:10.000/threads:12          7858 ns

[ashvardanian]

Yes, @MarkReedZ, the check_c_source_compiles makes a lot of sense! Can you please clarify the benchmarking results? I'd assume bf16 should be faster than f16`, so the duration/latency should be lower 🤔

[MarkReedZ]

I put bf16, f16, and f32 dot_serial() in godbolt. You can add and remove flags (avx2, avx512fp16, etc) to see whats going on. Without flags the bf16 is longer. Is the compiler using avx2/avx512 on the f16 serial? That would explain the difference.

https://godbolt.org/z/EKE66h9GM

The bf16 and unsigned short f16 have the same performance in dot/cos/l2sg, but bf16 is faster in kl/js.

I'll play around with different compilers.

[MarkReedZ]

Note avx512_bf16 only has support for conversion between bf16 and f32, and a dot product. So I believe our simd accelerated functions will be converting bf16 to f32, and running the f32 algorithms.

Or perhaps it is possible to do a bf16 -> f16 conversion if we can find a way to just shift the exponent.

[ashvardanian]

In most cases it would be better to perform dot products in bf16, upscaling and accumulating in f32 down the road.

[MarkReedZ]

I added the conversion function for compilers that don't support __bf16

SIMSIMD_PUBLIC simsimd_f32_t simsimd_uncompress_bf16(unsigned short x) {
    unsigned int tmp = x << 16; // Zero extends the mantissa
    return *((float*)&tmp);
}

And using the conversion to f32 instead of the native bf16 we get almost the same timings as with plain f32.

unsigned short bf16 -> f32 conversion

dot_bf16_serial_1536d/min_time:10.000/threads:12          183 ns
cos_bf16_serial_1536d/min_time:10.000/threads:12          202 ns
l2sq_bf16_serial_1536d/min_time:10.000/threads:12         166 ns
kl_bf16_serial_1536d/min_time:10.000/threads:12          1505 ns
js_bf16_serial_1536d/min_time:10.000/threads:12          3795 ns

A PR will be up when I have a minute.

[pauldintel]

@MarkReedZ which machine this benchmark is running for ? Intel Bf16 should show better results on 4th Gen Sapphire Rapids (SPR) with AMX accelerators enabled because BF16 supposed to show better results with matrix multiplication comparing FP16 . I am not from above which distance calculation require matrix mul operations.

[ashvardanian]

Alternatively, you can also test on AMD Genoa chips. Like Intel Sapphire Rapids they support AVX-512 BF16, unlike Intel - they don't support F16... so the relative win will be much larger.

[pauldintel]

So far I know Genoa has no BF16 support as at this moment it works on Intel SPR with AMX acceleration only

[ashvardanian]

@pauldintel it should be.