MSL: Support for SPV_KHR_cooperative_matrix (and SPV_KHR_shader_bfloat16)

[oscarbg]

Hi,
I opened a few days ago, an equivalent issue at MoltenVK requesting support for these two extensions:
KhronosGroup/MoltenVK#2507
it was mainly motivated by:

1)llama.cpp Vulkan backend using it..
and coop matrix extension being used also in apps like vkpeak it's a good indicator of peak perf(max top/s) achievable on GPUs supporting it..

2)also seems MSL these days should support these two extensions:

2.a)has bfloat intrinsic
2.b)MSL spec has support for simd group matrix functions (section 6.7) which sound like coop mat equivalents..

3)seeing also support added in this project for GLSL backend recently:
c4f82d1
de656c6
which should help someway.. I assume :-)..

as seems, main work of implementing these extension support in MoltenVK is SPIRV-Cross MSL support for this 2 exts.. opening here for visibilty also..

thanks..

[HansKristian-Work]

Not interested in implementing this myself, but if someone implements it natively, it can probably be merged.

[dboyan]

While I don't have plan to do them in the near future, I'm sharing some pointers that might be useful for people that are interested.

In theory, bfloat should be more straightforward because MSL 3.1 starts to have this type. There can be some caveats when dealing with math functions though. Maybe this piece of code in MLX library from Apple can shed some lights on the things to pay attention to.

Cooperative matrix seems more nasty. I previously took a slight look and found two main challenges here (if I understand the msl spec correctly):

• MSL only allows 8x8 simdgroup_matrix
• MSL does not support specifying behavior related to SaturatingAccumulationKHR

I'm not saying that they cannot be done by adding routines to carefully handle them. But it might be quite painful.

[oscarbg]

a little update..

well a "little" changed a month after I opened the request..

Metal4 MSL adds new tensors type (in header <metal_tensor>) (including cooperative_tensor ): https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf see section 2.20

and A19 (and M5?) GPUs adds neural accelerators (new Apple10 Metal family):
https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf
which presumably are used using cooperative_tensor (and simdgroup_matrix?) ops..

so more interested now in this..

[oscarbg]

@kabu1204 thanks for SPIRV-Cross PR!!
now that is merged may I ask if are you interested on implementing MoltenVK side of coop matrix support?
also still not closing issue as bfloat16 was also requested..