TVMを使って演算する（Python編）では、TVMのPythonインタフェースを使って足し算を行う方法を紹介しました。整備されたインタフェースが提供され、ドキュメントもきちんと揃っているPythonからTVMへアクセスすることは、理にかなっています。しかし一方で、様々なデバイスでの演算環境の提供というTVM本来の開発目的を考えると、Pythonが利用できない環境でTVMを利用して演算したい場合もあるはずです。
本記事では、Pythonが入っていない環境でTVMを動かしたい方やPythonを使うことによるオーバーヘッドを抑えたい方、NNVM/TVMをC++のアプリケーションとしてビルドしてデバッグしたい人のために、C++からTVMを使う方法を説明します。
TVMの説明やインストール方法については、Python編の記事を参照していただくことにして、重複を避けるためにここでは説明しません。
TVMのC++インタフェースについて
TVMのPythonインタフェースについても同様のことが言えますが、TVMは比較的最近始まったばかりのプロジェクトであるため、今後インタフェースが追加/変更される可能性が高いと考えられます。事実、この記事を書いている間にも新たなインタフェースが追加されていることから、TVM利用者は今後インタフェースが変わることを強く意識しておく必要があります。
ドキュメントを見ればよいという意見もあると思いますが、TVMから提供されているドキュメントはPythonからアクセスする方法しか紹介されていません。TVMをC++から利用するためのドキュメントは提供されていないのです。このため、Pythonのようにドキュメントを確認しながら実装することができません。
このような状況ですので、今のところC++からTVMを使いたい場合は、PythonからC++のコードを呼び出している箇所を調査し、必要に応じてC++のソースコードを読んでいくしかありません。しかし、TVMプロジェクトのリーダーである Tianqi Chen さんに話を聞いてみると、PythonだけでなくTVMをC++から利用することも想定しているとのことなので、いずれドキュメントも整備されていくことでしょう。
さてドキュメントがないので、ソースコードを読んでC++のインタフェースを把握していくことになるのですが、手掛かりになるのがソースコード中の TVM_DLL というマクロ定義です。TVMが提供しているC++のインタフェースかどうかを調べるためには、このマクロ定義 TVM_DLL が、関数などの先頭に書かれているかを確認すればよいです。
TVM_DLL の定義を見ると以下のように、共有ライブラリ内の関数のexportを行うためのマクロ定義であることが確認できます。TVM_DLL の定義があるからと言って今後もインタフェースが変わらない保証はありませんが、今のところは「TVMの共有ライブラリの関数がexportされている＝TVMのインタフェース」判断してよいでしょう。
#ifndef TVM_DLL
#ifdef _WIN32
#ifdef TVM_EXPORTS
#define TVM_DLL __declspec(dllexport)
#else
#define TVM_DLL __declspec(dllimport)
#endif
#else
#define TVM_DLL
#endif
#endif
マクロ定義 TVM_DLL が利用されている例を1つ示します。
TVM_DLL Tensor placeholder(Array<Expr> shape,
Type dtype = Float(32),
std::string name = "placeholder");
ユーザは、auto A = tvm::placeholder(shape, tvm::Float(32), "A"); のように上記のインタフェースを利用することができます。
TVMを使って演算する
さて、TVMのC++インタフェースと判断した関数などを使って演算してみたいと思います。Python編と同様に、配列の各要素の足し算を行います。また、Python編のサンプル との違いを理解しやすいように、Pythonのコードとの比較も行います。
サンプルコード
前述したように、C++からTVMを利用するための公式のチュートリアルは存在しません。ここで紹介するサンプルコードは、TVMのPythonコードから呼び出される、C++の関数を確認しながら書いたコードになります。そのため、サンプルコードがTVMの使い方の作法に合っているかはわかりません。サンプルコードの実行結果を見る限り問題ないと思いますが、公式の手順ではないことを改めてここで書いておきます。
#include <random>
#include <iomanip>
#include <tvm/tvm.h>
#include <tvm/build_module.h>
const int NUM_TENSORS = 3;
// TVMArray（DLTensor）の要素へアクセスするためのヘルパ関数
template <typename dtype>
inline dtype& TVMARRAY_ADDR(TVMArrayHandle handle, int idx)
{
::DLTensor* dl_tensor = handle;
uint64_t offset = dl_tensor->byte_offset;
float* addr = reinterpret_cast<float*> (static_cast<char*> (dl_tensor->data) + offset);
return addr[idx];
}
// TVMArray（DLTensor）の要素数を取得するためのヘルパ関数
inline uint64_t TVMARRAY_NELM(TVMArrayHandle handle, int axis)
{
::DLTensor* dl_tensor = handle;
return handle->shape[axis];
}
int main()
{
// 演算の定義
// ここでは単純に、C[] = A[] + B[]のようなN個の要素をもつ配列の演算を定義します
auto n = tvm::var("n");
tvm::Array<tvm::Expr> shape;
shape.push_back(n);
auto A = tvm::placeholder(shape, tvm::Float(32), "A");
auto B = tvm::placeholder(shape, tvm::Float(32), "B");
tvm::Tensor C = tvm::compute(A->shape, [&A, &B](tvm::Expr i) {
return A[i] + B[i];
}, "C");
// スケジュールの作成
tvm::Schedule s = tvm::create_schedule({ C->op });
auto cAxis = C->op.as<tvm::ComputeOpNode>()->axis;
tvm::IterVar bx, tx;
s[C].split(cAxis[0], 1, &bx, &tx);
s[C].bind(bx, tvm::thread_axis(tvm::Range(), "blockIdx.x"));
s[C].bind(tx, tvm::thread_axis(tvm::Range(), "threadIdx.x"));
// ターゲットをOpenCLとして、TVMでの実行コードを生成
tvm::BuildConfig config;
auto target = tvm::Target::create("opencl"); // ターゲットデバイス：OpenCL
auto target_host = tvm::Target::create("stackvm"); // ホストデバイス：StackVM
auto args = tvm::Array<tvm::Tensor>({ A, B, C });
std::unordered_map<tvm::Tensor, tvm::Buffer> binds;
tvm::Array<tvm::LoweredFunc> lowered = tvm::lower(s, args, "fadd_cl", binds, config);
tvm::runtime::Module mod = tvm::build(lowered, target, &target_host, config);
auto fadd_cl = mod.GetFunction("fadd_cl");
// OpenCLのカーネルコードを表示
std::vector<tvm::runtime::Module> imported_modules = mod->imports();
std::cout << "=== OpenCL Kernel Code ===" << std::endl;
std::cout << imported_modules[0]->GetSource() << std::endl;
// ホストデバイス、ターゲットデバイスのメモリ確保（配列の要素数は16）
TVMArrayHandle hHostArr[NUM_TENSORS];
TVMArrayHandle hTgtArr[NUM_TENSORS];
tvm_index_t arr_shape = 16;
for (int i = 0; i < NUM_TENSORS; ++i) {
::TVMArrayAlloc(&arr_shape, 1, 2, 32, 1, kDLCPU, 0, &hHostArr[i]);
::TVMArrayAlloc(&arr_shape, 1, 2, 32, 1, kDLOpenCL, 0, &hTgtArr[i]);
}
// ホストデバイスのメモリに値を設定（入力値はランダムに設定）
std::random_device rd;
std::mt19937 mt(rd());
std::uniform_real_distribution<float> urd(0.0, 1.0);
for (int i = 0; i < 2; ++i) {
for (int idx = 0; idx < TVMARRAY_NELM(hHostArr[i], 0); ++idx) {
TVMARRAY_ADDR<float>(hHostArr[i], idx) = urd(mt);
}
}
// ホストデバイスからターゲットデバイスへメモリコピー
for (int i = 0; i < NUM_TENSORS; ++i) {
::TVMArrayCopyFromTo(hHostArr[i], hTgtArr[i], nullptr);
}
// OpenCLで演算実行
fadd_cl(hTgtArr[0], hTgtArr[1], hTgtArr[2]);
// ターゲットデバイスからホストデバイスへメモリコピー
for (int i = 0; i < NUM_TENSORS; ++i) {
::TVMArrayCopyFromTo(hTgtArr[i], hHostArr[i], nullptr);
}
// 演算結果を表示
std::cout << "=== Result ===" << std::endl;
std::cout << std::fixed << std::setprecision(4);
for (int idx = 0; idx < TVMARRAY_NELM(hHostArr[0], 0); ++idx) {
std::cout << "[" << idx << "] " << TVMARRAY_ADDR<float>(hHostArr[0], idx) << " + ";
std::cout << TVMARRAY_ADDR<float>(hHostArr[1], idx) << " = ";
std::cout << TVMARRAY_ADDR<float>(hHostArr[2], idx) << std::endl;
}
// メモリを解放
for (int i = 0; i < NUM_TENSORS; ++i) {
::TVMArrayFree(hHostArr[i]);
::TVMArrayFree(hTgtArr[i]);
}
return 0;
}
実行結果
サンプルコードの実行結果を示します。
Python編 と比較しやすいように、出力する内容を Python編の実行結果 と合わせました。
=== OpenCL Kernel Code ===
__kernel void fadd_cl__kernel0(__global float* restrict C, __global float* restrict A, __global float* restrict B) {
C[((int)get_group_id(0))] = (A[((int)get_group_id(0))] + B[((int)get_group_id(0))]);
}
=== Result ===
[0] 0.8707 + 0.2369 = 1.1076
[1] 0.2981 + 0.1895 = 0.4876
[2] 0.5818 + 0.4378 = 1.0196
[3] 0.4819 + 0.1517 = 0.6336
[4] 0.4473 + 0.9479 = 1.3952
[5] 0.9894 + 0.0608 = 1.0502
[6] 0.8529 + 0.2456 = 1.0985
[7] 0.5195 + 0.9502 = 1.4697
[8] 0.7799 + 0.5020 = 1.2819
[9] 0.1949 + 0.5207 = 0.7157
[10] 0.9008 + 0.7907 = 1.6914
[11] 0.4625 + 0.3332 = 0.7958
[12] 0.5884 + 0.5283 = 1.1167
[13] 0.3682 + 0.8473 = 1.2155
[14] 0.5688 + 0.3947 = 0.9634
[15] 0.3285 + 0.6897 = 1.0182
ランダムに決定される足し算の値を除き、Python編 と同じ結果が得られていることが分かります。
Pythonコードとの比較
ここで Python編のサンプル と、本記事のサンプル を比較してみます。
演算の定義
TVMで演算する内容を定義する部分です。変数定義などの点で、Pythonと比較してC++では多少記述量が増えていますが、C++で特別に書き足す必要のあるコードは特にありません。
// 演算の定義
// ここでは単純に、C[] = A[] + B[]のようなN個の要素をもつ配列の演算を定義します
auto n = tvm::var("n");
tvm::Array<tvm::Expr> shape;
shape.push_back(n);
auto A = tvm::placeholder(shape, tvm::Float(32), "A");
auto B = tvm::placeholder(shape, tvm::Float(32), "B");
tvm::Tensor C = tvm::compute(A->shape, [&A, &B](tvm::Expr i) {
return A[i] + B[i];
}, "C");
# 演算の定義
# ここでは単純に、C[] = A[] + B[]のようなN個の要素をもつ配列の演算を定義します
n = tvm.var("n")
A = tvm.placeholder((n,), name="A")
B = tvm.placeholder((n,), name="B")
C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C")
スケジュールの作成
演算の定義と同じように記述量は多少増えていますが、こちらもC++で特別に書き足す必要はありません。
// スケジュールの作成
tvm::Schedule s = tvm::create_schedule({ C->op });
auto cAxis = C->op.as<tvm::ComputeOpNode>()->axis;
tvm::IterVar bx, tx;
s[C].split(cAxis[0], 1, &bx, &tx);
s[C].bind(bx, tvm::thread_axis(tvm::Range(), "blockIdx.x"));
s[C].bind(tx, tvm::thread_axis(tvm::Range(), "threadIdx.x"));
# スケジュールの作成
s = tvm.create_schedule(C.op)
bx, tx = s[C].split(C.op.axis[0], factor=1)
s[C].bind(bx, tvm.thread_axis("blockIdx.x"))
s[C].bind(tx, tvm.thread_axis("threadIdx.x"))
TVMでの実行コード生成
実行コードを生成する部分は、C++からTVMを利用する場合に、注意しなくてはならない部分が2つあります。
1つ目は、tvm::build() にホストデバイスとして stackvm を指定する必要がある点です。
Python編のサンプルコード のコメントでもStack VMという用語が少しだけ出てきましたが、これはCPU上で動作するTVMの仮想マシンです。OpenCLやCUDAのカーネルコードを実行するために様々な準備が必要であるのと同様、tvm::build() で生成されたOpenCLやCUDAのカーネルコードは、そのままでは実行することができません。そこで、生成されたOpenCLやCUDAのカーネルコードを実行できるように環境を整えたり、カーネルコードを実行したりするのがStack VMの役割になります。
2つ目は、PythonからTVMを利用する際には意識する必要のなかった、tvm::LowerFunc を作る必要がある点です。
tvm::LowerFunc は、OpenCLやCUDAのカーネルコードを生成する前段階の中間言語のようなものです。PythonからTVMを利用する場合は、Pythonのインタフェースの処理の中でうまいこと隠蔽されていたため、このLowerFuncを意識する必要はありませんでしたが、C++では隠蔽されていません。
// ターゲットをOpenCLとして、TVMでの実行コードを生成
tvm::BuildConfig config;
auto target = tvm::Target::create("opencl"); // ターゲットデバイス：OpenCL
auto target_host = tvm::Target::create("stackvm"); // ホストデバイス：StackVM
auto args = tvm::Array<tvm::Tensor>({ A, B, C });
std::unordered_map<tvm::Tensor, tvm::Buffer> binds;
tvm::Array<tvm::LoweredFunc> lowered = tvm::lower(s, args, "fadd_cl", binds, config);
tvm::runtime::Module mod = tvm::build(lowered, target, &target_host, config);
auto fadd_cl = mod.GetFunction("fadd_cl");
# ターゲットをOpenCLとして、TVMでの実行コードを生成
fadd_cl = tvm.build(s, [A, B, C], "opencl", name="fadd_cl")
OpenCLのカーネルコードを表示
ターゲットデバイス向けに生成したカーネルコードを表示する部分についても、PythonとC++の間には特に大きな違いはありません。
std::vector<tvm::runtime::Module> imported_modules = mod->imports();
std::cout << "=== OpenCL Kernel Code ===" << std::endl;
std::cout << imported_modules[0]->GetSource() << std::endl;
# OpenCLのカーネルコードを表示
print("=== OpenCL Kernel Code ===")
print(fadd_cl.imported_modules[0].get_source())
計算に使用するメモリを割り当て＆値を設定
次に、計算に使用するメモリを割り当てたり、値を設定したりする処理です。おそらくこの部分が、C++からTVMを利用する場合とPythonからTVMを利用する場合とで、大きく異なる部分であると思います。
Pythonでは tvm.nd.array が、ホストとターゲットデバイスで利用するメモリの確保や値の設定をうまいこと隠蔽してくれていましたが、C++では tvm.nd.array で隠蔽されていた処理をすべて自分で行う必要があります。また、ホストとターゲットデバイス間のメモリコピーも自分で呼び出さなくてはなりません。
ホストとターゲットデバイスで利用するメモリの確保と解放は、TVMArrayAlloc() と TVMArrayFree() で行います。そして、TVMArrayAlloc 関数の第6引数には、kDLCPU や kDLOpenCL といった、メモリを確保するデバイスを指定します。また、確保したメモリの値を、ホストとターゲットデバイス間でデータ転送するために、TVMArrayCopyFromTo() を使います。このあたりの処理は、CUDAやOpenCLのプログラムを書いたことがある人であれば、ああ、ここでホストとターゲットデバイスのメモリ確保しているのだな、ホストとターゲットデバイス間のメモリコピーを行っているのだな、と何となく予想できるかと思います。
そして一番面倒なところが、確保したメモリに値を設定する処理です。確保したメモリ（DLTensor のメンバ変数 data に確保したメモリの先頭アドレスが保存されている）は、単純な void ポインタ型として提供されるため、プログラム側でアドレスを意識して値を設定する必要があり、苦労します。このため、本サンプルでは、DLTensorの各要素にアクセスするためのヘルパ関数と、DLTensorの要素数を取得するためのヘルパ関数を用意しました。汎用性を考えると（多次元配列が考慮されていないなど）、実装に不備があるのは明らかです。しかし、今回はあくまでサンプル用途として作成しただけに過ぎないので、もしTVMを今後利用される場合は参考程度で見ていただければと思います。
const int NUM_TENSORS = 3;
// TVMArray（DLTensor）の要素へアクセスするためのヘルパ関数
template <typename dtype>
inline dtype& TVMARRAY_ADDR(TVMArrayHandle handle, int idx)
{
::DLTensor* dl_tensor = handle;
uint64_t offset = dl_tensor->byte_offset;
float* addr = reinterpret_cast<float*> (static_cast<char*> (dl_tensor->data) + offset);
return addr[idx];
}
// TVMArray（DLTensor）の要素数を取得するためのヘルパ関数
inline uint64_t TVMARRAY_NELM(TVMArrayHandle handle, int axis)
{
::DLTensor* dl_tensor = handle;
return handle->shape[axis];
}
int main()
{
// <snip>
// ホストデバイス、ターゲットデバイスのメモリ確保（配列の要素数は16）
TVMArrayHandle hHostArr[NUM_TENSORS];
TVMArrayHandle hTgtArr[NUM_TENSORS];
tvm_index_t arr_shape = 16;
for (int i = 0; i < NUM_TENSORS; ++i) {
::TVMArrayAlloc(&arr_shape, 1, 2, 32, 1, kDLCPU, 0, &hHostArr[i]);
::TVMArrayAlloc(&arr_shape, 1, 2, 32, 1, kDLOpenCL, 0, &hTgtArr[i]);
}
// ホストデバイスのメモリに値を設定（入力値はランダムに設定）
std::random_device rd;
std::mt19937 mt(rd());
std::uniform_real_distribution<float> urd(0.0, 1.0);
for (int i = 0; i < 2; ++i) {
for (int idx = 0; idx < TVMARRAY_NELM(hHostArr[i], 0); ++idx) {
TVMARRAY_ADDR<float>(hHostArr[i], idx) = urd(mt);
}
}
// ホストデバイスからターゲットデバイスへメモリコピー
for (int i = 0; i < NUM_TENSORS; ++i) {
::TVMArrayCopyFromTo(hHostArr[i], hTgtArr[i], nullptr);
}
// <snip>
// ターゲットデバイスからホストデバイスへメモリコピー
for (int i = 0; i < NUM_TENSORS; ++i) {
::TVMArrayCopyFromTo(hTgtArr[i], hHostArr[i], nullptr);
}
// <snip>
// メモリを解放
for (int i = 0; i < NUM_TENSORS; ++i) {
::TVMArrayFree(hHostArr[i]);
::TVMArrayFree(hTgtArr[i]);
}
// <snip>
}
# 実行コンテキストをOpenCLに設定
ctx = tvm.opencl(0)
# 計算に使用するメモリを割り当て＆値を設定（配列の要素数は16）
# 入力値はランダムに設定
n = 16
ndarr_A = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
ndarr_B = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
ndarr_C = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
OpenCLで演算実行
演算実行部は、C++とPythonとの間に大きな違いはありません。
// OpenCLで演算実行
fadd_cl(hTgtArr[0], hTgtArr[1], hTgtArr[2]);
# OpenCLで演算実行
fadd_cl(ndarr_A, ndarr_B, ndarr_C)
演算結果を表示
演算結果をメモリから取り出す場合も、メモリに値を設定する場合と同じように、アドレスを意識してアクセスします。値を設定する時も取得する時も、C++では一苦労ですね。
// 演算結果を表示
std::cout << "=== Result ===" << std::endl;
std::cout << std::fixed << std::setprecision(4);
for (int idx = 0; idx < TVMARRAY_NELM(hHostArr[0], 0); ++idx) {
std::cout << "[" << idx << "] " << TVMARRAY_ADDR<float>(hHostArr[0], idx) << " + ";
std::cout << TVMARRAY_ADDR<float>(hHostArr[1], idx) << " = ";
std::cout << TVMARRAY_ADDR<float>(hHostArr[2], idx) << std::endl;
}
# 演算結果を表示
print("=== Result ===")
for i, (a, b, c) in enumerate(zip(np_A, np_B, np_C)):
print("[{}] {:.4f} + {:.4f} = {:.4f}".format(i, a, b, c))
おわりに
TVMをC++から利用する方法を調査し、記事としてまとめました。
C++からTVMを利用する場合、本来ユーザが意識する必要のないTVMの内部構造を理解しなければならず、Pythonと同じ実行結果を得るまでに多くの時間を費やしました。
TVMのPythonインタフェースはユーザが意識する必要のない処理をうまく隠蔽しているため、ユーザが本来行いたい演算の定義やスケジュールの作成に意識を集中することができます。冒頭に書いたような特別な用途でTVMを利用する場合を除き、素直にPythonからTVMを利用した方がよいでしょう。