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.

[ashvardanian]

Hey @pauldintel! Have you had the chance to check the bf16 functionality? Have you ever tried to use AMX for vector-matrix operations, aka when one of the arguments contains just one non-zero row/column?

[pauldintel]

@ashvardanian we have tested two matrix with inner product and Before use of AMX Inner product we runtime used Intel oneDNN to reorder F32 data to BF16 . Certain dataset and batch operations we have seen 1.51 to 14x improvement (64 dimension to 1024 dimensions) of FAISS IndexFlat Scalar IP and all with Intel AMX. FAISS IndexFlat BLAS IP shows upto 4.8X gain with AMX.

For single query, comparing to native fp32: FP32->BF16 AMX speedup about 4.85x perf BF16 AMX speedup about 33.7x perf

Hope that helps