Google Colab で Vulkan で compute kernel を動かす

背景

• Vulkan で機械学習したい
• Colab で開発や動作検証したい

環境

GPU タイプに選択しておきます.

GPU はガチャらしいようで, Tesla K80 だったり Tesla P100 だったり Tesla T4 だったりするようです.

K80 だと動かないかもしれません.
今回は運良く P100 が当たりました. P100 だとメモリ 16GB あっていいですね.

Vulkan 設定　

!apt install vulkan-utils libvulkan-dev

だけでいけます. VK_ICD_FILENAMES 環境変数などの設定は不要でした.
vulkaninfo してみます.

WARNING: [Loader Message] Code 0 : loader_icd_scan: Can not find 'ICD' object in ICD JSON file /usr/share/vulkan/icd.d/nvidia_layers.json.  Skipping ICD JSON
'DISPLAY' environment variable not set... skipping surface info
error: XDG_RUNTIME_DIR not set in the environment.
==========
VULKANINFO
==========

Vulkan Instance Version: 1.1.126

...

GPU0:
VkPhysicalDeviceProperties:
---------------------------
    apiVersion     = 4198495 (1.1.95)
    driverVersion  = 1754316800 (0x6890c000)
    vendorID       = 0x10de
    deviceID       = 0x15f8
    deviceType     = PHYSICAL_DEVICE_TYPE_DISCRETE_GPU
    deviceName     = Tesla P100-PCIE-16GB

...

    VK_KHR_shader_atomic_int64           : extension revision 1
    VK_KHR_shader_draw_parameters        : extension revision 1
    VK_KHR_shader_float16_int8           : extension revision 1
    VK_KHR_shader_float_controls         : extension revision 1

...

Super cool!

機械学習に有益な VK_KHR_shader_float16_int8 も使えました.

動かす

vulkan_minimal_compute をコンパイルして動かしてみます.

その他

リアルタイム描画は難しいですが, Vulkan でグラフィックス描画については, offscreen rendering ならいけるかと思います.

参考: Google Colaboratoryが、3D CGレンダリング用途にも有能すぎた件
https://qiita.com/stakemura/items/71aeefaa927ab277fcf1

TODO

• Colab だと shell 操作がめんどいので, gitpod self-hosted で手元 PC で Vulkan 動かして機械学習開発できるようにする https://www.gitpod.io/