Vengineerの妄想

人生を妄想しています。

Operator Vectorization Library (OVL) と NNVM



TensorFlow XLA と Intel Nervana Graph だけでなく、MXNetのNNVM や HPEの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をPythonDSL(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 があったよ。