MSL: Add cooperative matrix element-wise emulation

[zainsharief]

This PR is designed to fix some of the errors in KhronosGroup/MoltenVK#2753.

Adds support for element-wise operations on CooperativeMatrixKHR types in the MSL backend. Since Metal's simdgroup_matrix API does not natively support element-wise arithmetic, this implementation emulates these operations by lowering them to threadgroup scratch memory, applying the requested operation per-thread, and reloading the result into the matrix.

Some unavoidable limitations:

• simdgroup matrix only allows 8x8 and floating point types (but a well made implementation will account for this by querying device properties)
• metal only supports subgroup scope
• metal only supports compile-time sizes and scope, so we throw an error when these are spec constants

[CLAassistant]

All committers have signed the CLA.

[oscarbg]

just one comment I see this code has in some way bfloat16 support..

SPIRType::BFloat16
			SPIRV_CROSS_THROW("MSL cooperative matrices only support float16, float32, and bfloat16 component types. "

maybe are you interested in adding VK_KHR_shader_bfloat16 support for MoltenVK.. was part of my original cooperative matrix request:
KhronosGroup/MoltenVK#2507

most work should be in SPIRV-Cross side (i.e. support SPV_KHR_bfloat16), right?

there are tests, testing bfloat16 throughput in standard way and coop matrix way.. for example clpeak2 (vulkan mode) shows currently:

 BF16 compute bf16xbf16+fp32 (GFLOPS)
     bf16     : [unsupported] VK_KHR_shader_bfloat16 / shaderBFloat16Type not supported! Skipped
     bf16_2   : [unsupported] VK_KHR_shader_bfloat16 / shaderBFloat16Type not supported! Skipped
     bf16_4   : [unsupported] VK_KHR_shader_bfloat16 / shaderBFloat16Type not supported! Skipped

   Cooperative-matrix bf16xbf16+fp32 (TFLOPS)
     coopmat_bf16 : [unsupported] No bf16xbf16+fp32 coopmat property! Skipped

[oscarbg]

comment #2:

unavoidable limitations:

    simdgroup matrix only allows 8x8 and floating point types (but a well made implementation will account for this by querying device properties)

sadly testing 4 tests I know.. at least 3 don't check that 8x8 is the only supported dim and fail:
clpeak2, llama.cpp and vk_cooperative_matrix_perf fail ..

vkpeak don't know because fails with spec-constant use..

so maybe it's time to start opening issues on these 3 projects on their github sites?

[zainsharief]

Hi @oscarbg, it is definitely possible for me to add some bfloat16 support - but should we leave that for a future PR? It is pretty likely that they do check for dims, but maybe because of the error bug you mentioned in the MoltenVK PR (now fixed), the error appeared because the spec constant default was not 8 - maybe worth me checking quickly but I will also post issues on all 4 now. Thanks for your help! I also fixed the windows issue so hopefully that workflow should pass now.

[oscarbg]

Hi @oscarbg, it is definitely possible for me to add some bfloat16 support - but should we leave that for a future PR? It is pretty likely that they do check for dims, but maybe because of the error bug you mentioned in the MoltenVK PR (now fixed), the error appeared because the spec constant default was not 8 - maybe worth me checking quickly but I will also post issues on all 4 now. Thanks for your help! I also fixed the windows issue so hopefully that workflow should pass now.

@zainsharief

1)yep, bfloat16 for a future PR is OK.. was only trying to gather interest on implementing support for it..
2)thanks for fixing the error bug so fast!!..
3)glad to help.. interested on it :-)
5) also thanks for opening 4 issues!!

posting here for visibilty in the future:

jeffbolznv/vk_cooperative_matrix_perf#13
krrishnarraj/clpeak#190
nihui/vkpeak#49
ggml-org/llama.cpp#25872

[zainsharief]

@oscarbg I was thinking about what 0cc4m was saying under the llama.cpp issue - and maybe it makes more sense that we do hardcode 8x8 + subgroup size and give a warning if the user tries to use spec constants. This way we are not actively punishing developers who check properly, and explicitly stating the issue. I'm curious to hear your thoughts on this?

Only issue I can think of is if Apple ever decides to support more than 8x8, but still doesn't allow spec constants.

[zainsharief]

i have updated the code to reflect this, if you wanted to test it out, but it should be easy to revert if this is unwanted behaviour.

[HansKristian-Work]

This is far too naive to be shippable. 16k LDS block will dramatically lower occupancy and performance for any shader that happens to use element-wise operations. Also, there is real risk of generating a broken shader that exceeds limits for LDS if the shader otherwise uses scratch memory.

[HansKristian-Work]

Don't use fprintf here. I think the other code just faults the shader if this happens.

[HansKristian-Work]

At a glance this looks extremely duplicated. Is there no way to collapse at least some of this code duplication?

[HansKristian-Work]

Did anyone raise this issue with Apple? Not having basic operations like this in the language is kinda ridiculous, and adding yet more massive emulation workarounds that slow things down for basic operations is just not sustainable. ...