-
Notifications
You must be signed in to change notification settings - Fork 13.5k
vulkan: Add bfloat16 support #12554
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
vulkan: Add bfloat16 support #12554
Conversation
|
The tooling for Vulkan bfloat16 is not all merged yet - see KhronosGroup/SPIRV-Tools#6057 and KhronosGroup/glslang#3905. So if anybody wants to try this locally, you'd need to build a custom glslc. I'll update when it's all merged NVIDIA will release a Vulkan developer driver with support for this extension, hopefully tomorrow. |
9ffb02b to
9cd10f8
Compare
|
I've rebased this, and added a fallback (promote to fp32) when bf16 coopmat support isn't available. Still waiting on the tooling to be merged, so still Draft for now. |
02160db to
21e8793
Compare
f4d8c68 to
7c47fd0
Compare
This adds bfloat16 matrix multiply support based on VK_KHR_shader_bfloat16. The extension is required for coopmat multiply support, but matrix-vector multiply trivially promotes bf16 to fp32 and doesn't require the extension. The copy/get_rows shaders also don't require the extension. It's probably possible to fall back to non-coopmat and promote to fp32 when the extension isn't supported, but this change doesn't do that. The coopmat support also requires a glslc that supports the extension, which currently requires a custom build.
…pport Compile a variant of the scalar mul_mm shader that will promote the bf16 values to float, and use that when either the bf16 extension or the coopmat extensions aren't available.
7c47fd0 to
83771ec
Compare
|
glslc support has landed. I've rebased this change again and it's ready for review. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
This adds bfloat16 matrix multiply support based on VK_KHR_shader_bfloat16. The extension is required for coopmat multiply support, but matrix-vector multiply trivially promotes bf16 to fp32 and doesn't require the extension. The copy/get_rows shaders also don't require the extension.
It's probably possible to fall back to non-coopmat and promote to fp32 when the extension isn't supported, but this change doesn't do that.
The coopmat support also requires a glslc that supports the extension, which currently requires a custom build.