perf: bf16 kernels

[rok]

See #1651

[wjones127]

I don't quite understand the instructions involved, but it looks like you are frequently casting / reinterpreting bfloat16 as normal float16 and using float16 add / subtract instructions, which I'm worried isn't valid.

[eddyxu]

For bf16, we have https://doc.rust-lang.org/core/arch/x86_64/fn._mm512_dpbf16_ps.html , should we just use rust code instead?

L2 = dot(x, x) + dot(y, y) - 2 * dot(x, y)
Cosine = 1 - dot(x, y)

[rok]

For bf16, we have https://doc.rust-lang.org/core/arch/x86_64/fn._mm512_dpbf16_ps.html , should we just use rust code instead?

L2 = dot(x, x) + dot(y, y) - 2 * dot(x, y)

Cosine = 1 - dot(x, y)

This would indeed simplify things a lot, however I would compare performance to see if loading from memory (e.g. with _mm512_maskz_loadu_epi16) incurs a significant cost (rust option would do it 3 times vs C where we'd do it one time).

[eddyxu]

For bf16, we have https://doc.rust-lang.org/core/arch/x86_64/fn._mm512_dpbf16_ps.html , should we just use rust code instead?
L2 = dot(x, x) + dot(y, y) - 2 * dot(x, y)

Cosine = 1 - dot(x, y)

This would indeed simplify things a lot, however I would compare performance to see if loading from memory (e.g. with _mm512_maskz_loadu_epi16) incurs a significant cost (rust option would do it 3 times vs C where we'd do it one time).

Why we need to load 3 times in rust? Not sure i follow.

[rok]

@eddyxu Oh sorry I misunderstood your question. I thought you were suggesting calling the dot kernel to calculate L2.
If we can use nightly rust functions (e.g. _mm512_dpbf16_ps) that would be preferable and I'll refactor.

[rok]

@eddyxu I've looked at using _mm512_dpbf16_ps via rust but couldn't quite get it right today (e.g. _mm512_castsi512_ph is missing from). Please feel free to take over. I can also do some more work first week of January.