VkInline で Python で Vulkan compute して GPU 機械学習などに思いを馳せる

背景

• お手軽に python で GPU 計算したい
• Vulkan であれば, Khronos からは安定したコンパイラ, ドライバについては 各 GPU ベンダーから安定(or 定期的なアップデート)したものが出ており, 安心感がある.
  • Windows, Linux, Android で動く
  • macOS は MoltenVK 経由でいけるかも
• Vulkan compute 自体は, 画面出力(Swap chain)は不要にできているので, オフラインでサーバ GPU で動かすのにも使える

少なくとも Colab + Tesla P100 では Vulkan 使えました.

(V100 や A100 はどうかしらん?)

cupy のような感じで, python からインラインで Vulkan compute kernel を呼べる VkInline を試します.

Vulkan 1.2 ドライバが必要になります.
(現状, 1.2　に対応しているのは RADV(AMD OSS(?) Linux Vulkan ドライバ) か Windows Adrenalin ドライバくらいでしょうか?)

VkInline

ライセンスは Anti 996 ライセンスです
(9 時 9 時(21 時), 週 6 日労働のような労働基準法を無視した会社では anti-996 license なソフトウェアは使ってはいけない)

ビルド

特に問題なく, 手順に従えばすんなりいけます.

実行環境セットアップ

• Vulkan 1.2 ドライバ
  • 少なくとも, VK_KHR_buffer_device_address がサポートされているドライバ. amdgpu や古めの AMD GPU で Windows では動きません
    • 現状の vkinline コードでは強制的に exit しているので, ログが出ずにわかりづらい.
  • NVIDIA では, 1.2 対応のドライバは現状(2020 年 6 月 8 日時点) beta driver になります.
• Python cffi, ,numpy, glm

くらいでいけます.

今回は, RADV ドライバで RX5700(Navi) で動作確認しました.

ROCm で RADV で Vulkan(compute kernel) を使うメモ
https://qiita.com/syoyo/items/ce3943757281acbdba49

動かす

test_compute.py を動かしてみます.

// from VkInline test_compute.py
import VkInline as vki
import numpy as np

# interface with numpy
harr = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype='float32')
darr = vki.device_vector_from_numpy(harr)
print(darr.to_host())

# GLSL data type
print(darr.name_view_type())

harr2 = np.array([6,7,8,9,10], dtype='int32')
darr2 = vki.device_vector_from_numpy(harr2)

# kernel with auto parameters, launched twice with different types
kernel = vki.Computer(['arr_in', 'arr_out', 'k'],
'''
void main()
{
    uint id = gl_GlobalInvocationID.x;
    if (id >= get_size(arr_in)) return;
    set_value(arr_out, id, get_value(arr_in, id)*k);
}
''')

darr_out = vki.SVVector('float', 5)
kernel.launch(1,128, [darr, darr_out, vki.SVFloat(10.0)])
print (darr_out.to_host())

darr_out = vki.SVVector('int', 5)
kernel.launch(1,128, [darr2, darr_out, vki.SVInt32(5)])
print (darr_out.to_host())

# create a vector from python list with GLSL type specified
darr3 = vki.device_vector_from_list([3.0, 5.0, 7.0, 9.0 , 11.0], 'float')
print(darr3.to_host())

[1. 2. 3. 4. 5.]
Comb_bb4c7639fd354507
[10. 20. 30. 40. 50.]
[30 35 40 45 50]
[ 3.  5.  7.  9. 11.]

展望

特に問題なくすんなり動いてびっくりです! 期待ができますが, ここから機械学習などのアプリを作っていくにはいろいろ課題があります.

• Vulkan 向けに cuDNN 系とか移植(convolution 系のネットワークを使いたい場合).
• 複数 GPU 対応(Vulkan の機能で GPU to GPU 直通信したいときとか)
• 倍精度や, fp16 対応とか

TODO

• 優秀な GPU 機械学習若人さまが, VkInline をお極めになされることで, 人類史上最速で優秀な Vulkan 機械学習若人さまへと昇華なされるスキームを確立する旅に出たい.