Apache TVMは、CPU、GPU、そして様々な機械学習アクセラレーションチップに適した、エンドツーエンドのディープラーニング構築フレームワークです。中国語版のTVMドキュメントは、→ https://tvm.hyper.ai/ をご覧ください。

著者：Tianqi Chen

TVM は透過的なコード生成をサポートしていますが、一部の畳み込みカーネルに cuDNN を使用し、残りのステージを定義するなど、手書きのコードをパイプラインに統合する必要がある場合があります。

ネイティブTVMはブラックボックス関数呼び出しをサポートしています。具体的には、TVMはDLPackと互換性のあるすべてのテンソル関数をサポートしています。つまり、任意の関数をPOD型（ポインタ、整数、浮動小数点数）を使用して呼び出すか、DLTensorへのポインタを引数として渡すことで呼び出すことができます。

 from __future__ import absolute_import, print_function import tvm from tvm import te import numpy as np from tvm.contrib import cblas import tvm.testing if not tvm.get_global_func("tvm.contrib.cblas.matmul", allow_missing=True): raise Exception("Not compiled with cblas support; can't build this tutorial")

外部テンソル関数の使用

次の例では、 te.externを使用して外部配列関数呼び出しを追加しています。外部呼び出しでは出力テンソルの形状を宣言し、2番目の引数で入力リストを指定します。

ユーザーは、結果の計算方法を記述する関数を提供する必要があります。計算関数は、入力と出力のシンボルプレースホルダーのリストを受け取り、実行されたステートメントを返します。

この場合、登録済みのTVM関数を呼​​び出すだけで十分です。これはCBLASを呼び出すためです。TVMは外部配列関数の内部を制御せず、ブラックボックスとして扱います。スケジュールされたTVM関数をさらに組み合わせることで、結果にバイアス項を追加できます。

 n = 1024 l = 128 m = 235 bias = te.var("bias", dtype="float32") A = te.placeholder((n, l), name="A") B = te.placeholder((l, m), name="B") C = te.extern( (n, m), [A, B], lambda ins, outs: tvm.tir.call_packed( "tvm.contrib.cblas.matmul", ins[0], ins[1], outs[0], False, False ), name="C", ) D = te.compute(C.shape, lambda i, j: C[i, j] + bias, name="D") s = te.create_schedule(D.op)

検証結果

結果が期待どおりであるかどうかを確認します。

 dev = tvm.cpu(0) f = tvm.build(s, [A, B, D, bias], "llvm") a = tvm.nd.array(np.random.uniform(size=(n, l)).astype(A.dtype), dev) b = tvm.nd.array(np.random.uniform(size=(l, m)).astype(B.dtype), dev) d = tvm.nd.array(np.zeros((n, m), dtype=D.dtype), dev) bb = 10.0 f(a, b, d, bb) tvm.testing.assert_allclose(d.numpy(), np.dot(a.numpy(), b.numpy()) + 10, rtol=1e-5)

外部コントリビュートラッパー

TVM は外部呼び出し用の外部コントリビュート ラッパーを提供します。次のコードは前の例と同等です。

 from tvm.contrib import cblas C = cblas.matmul(A, B) D = te.compute(C.shape, lambda i, j: C[i, j] + bias, name="D") s = te.create_schedule(D.op)

Python関数をexternとしてフックする

TVM 内の任意の PackedFunc を呼び出すことができるため、外部関数を使用して Python にコールバックできます。

以下の例では、TVMランタイムシステムにPython関数を登録し、それを用いて計算段階を実行することで、TVMの柔軟性を高めています。例えば、フロントエンドコールバックを挿入することで中間結果をチェックしたり、カスタムコードをTVMに混在させたりすることができます。

 @tvm.register_func("tvm.contrib.my_tvm_addone") def my_tvm_addone(x, y): print("my_tvm_addone signatures: %s, %s" % (type(x), type(y))) tvm.nd.array(x.numpy() + 1).copyto(y) A = te.placeholder((n,), name="A") B = te.extern( A.shape, [A], lambda ins, outs: tvm.tir.call_packed("tvm.contrib.my_tvm_addone", ins[0], outs[0]), name="C", ) s = te.create_schedule(B.op) f = tvm.build(s, [A, B], "llvm") a = tvm.nd.array(np.random.uniform(size=(n,)).astype(A.dtype), dev) b = tvm.nd.array(np.random.uniform(size=(n,)).astype(B.dtype), dev) f(a, b) tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1, rtol=1e-5)

出力結果:

 my_tvm_addone signatures: <class 'tvm.runtime.ndarray.NDArray'>, <class 'tvm.runtime.ndarray.NDArray'>

要約

• TVM は te.extern を介して外部テンソル関数を呼び出します。
• 外部テンソル呼び出しには contrib ラッパーを使用します。
• フロントエンド関数を外部テンソルへのコールバックとしてフックします。

Python ソースコードをダウンロード: extern_op.py

Jupyter ノートブックをダウンロード: extern_op.ipynb