VK_NV_cooperative_matrix : GeForce RTXシリーズのTensor CoreをVulkanで使用できるようにする拡張

NVIDIAのTuring CoreのGPUより搭載されたTensor CoreをCompute Shaderで使用して演算するための拡張の仕様が公開になっています.

VK_NV_cooperative_matrix仕様書
https://www.khronos.org/registry/vulkan/specs/1.1-extensions/html/vkspec.html#VK_NV_cooperative_matrix

この拡張は,Cooperative matrix型というあらたな行列型をCompute Shader使えるようで,行列の各要素はfloat16,float32,float64,sint8,sint16,sint32,sint64,uint8,uint16,uint32,uint64と8bitから64bitの精度が指定できます.この行列の演算に関しては,TuringアーキテクチャのGeForce RTXシリーズに搭載された Turing Core を使用することができるようです.

Compute Shaderで使用できるということでMLやDLのような用途でなくても良さそうなので,グラフィックス処理に使うということもできそうです.

なお,拡張の仕様は上記のドキュメントで公開され,下記のタスクリストを見るとドライバも対応になったようですが,SPIR-Vやシェーダコンパイラの対応などがまだのようです.

Task list for VK_NV_cooperative_matrix release
https://github.com/KhronosGroup/Vulkan-Docs/issues/923

使用方法に関するサンプルや下記のリポジトリに公開されています.ベンチマークになっているので,パフォーマンスについても見ることができます.シェーダを変更するための
glslangValidator.exe についても情報があります.

シェーダをSPIR-Vバイナリにしたものはリポジトリにありますので動かすことはできるんではないかと思います.

vk_cooperative_matrix_perf
https://github.com/jeffbolznv/vk_cooperative_matrix_perf