Operator Vectorization Library (OVL) と NNVM

VengineVengineerの戯言 : Twitter、Slideshare
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

One Deep Learning Virtual Machine to Rule Them All

TensorFlow XLA と Intel Nervana Graph だけでなく、MXNetのNNVM や HPEのOperator Vectorization Library (OVL)というのがあるんだね。

知らなかったので、Operator Vectorization Library (OVL)をちょっと調べてみた。

Operator Vectorization Library (OVL)

ここにあるサンプルコード

引用
import tensorflow as tf
import opveclib as ovl


@ovl.operator()
def shift_cyclic(input_tensor):
    # make sure input is 1D
    assert input_tensor.rank == 1

    # define the output tensor
    output_tensor = ovl.output_like(input_tensor)

    # define the workgroup shape and get workgroup position reference
    wg_position = ovl.position_in(input_tensor.shape)

    # read input element at current workgroup position
    input_element = input_tensor[wg_position]

    # define the output position to be 1 greater than the input/workgroup position
    output_position = (wg_position[0] + 1) % input_tensor.size

    # set the output element
    output_tensor[output_position] = input_element

    return output_tensor


a = tf.constant([0, 1, 2, 3, 4], dtype=tf.float32)
b = ovl.as_tensorflow(shift_cyclic(a))

sess = tf.Session()
print(sess.run([b]))

TensorFlowでのOpをPythonのDSL(Domain Specific Language)でサクッと書けるというのが特徴みたい。
@ovl.operatorを付けると、そこがOVLのカスタムOpになって、CPUやCUDA用のコードを自動的に生成するって。
ただし、cuDNNのライブラリとかは利用できないって。。

TensorFlow XLAと基本的に同じだね。
ユーザーがやりたい処理を一つの関数として定義して、それをOVLがカスタムOpにするんだからね。

ドキュメントは、ここ。ホワイトペーパーもあるよ。
コードは、githubにあるよ。でも、公開されてから更新されていないみたい。。。
まー、その頃から、TensorFlow XLAの開発が始まったので、わざわざやる必要が無いからね。。。

そんでもって、NNVMについても調べたよ。
こちらも、NNVM: Build deep learning system by parts　ってあるので、同じだね。。。

こちらは最近も更新されているみたい。

srcディレクトリを覗いたら、

　・c_api
　・core
　・pass
]}}
がある。ここでも pass があるのね。。。

https://github.com/dmlc/nnvm/tree/master/src/pass passディレクトリを覗いたら、
{{{
　・gradient.cc           => Gradient Pass
　・infer_shape_type.cc   => InferShape Pass
　・order_mutation.cc     => OrderMutation Pass
　・place_device.cc       => PlaceDevice Pass
　・plan_memory.cc　　　　=> PlanMemory Pass
　・saveload_json.cc      => LoadJSON Pass / SaveJSON Pass

という pass があったよ。

Deep Teaching: The Sexiest Job of the Futureも参考までに。