TensorFlow XLA JIT版で、Pluginの起動に成功しました

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

TensorFlow XLAに動きありに書いたように、
TensrorFlow XLA r1.3のJIT版に Plugin が追加されました。

ThinkPad 13を購入しましたので、Ubuntu 16.04LTS on VirtualBox on Windows 10にて、XLA有りでビルドしました。

先週1週間夏休みだったので、いろいろやってみました。
その結果は、SlidesahreのXLA : JIT編 (r1.3版)にアップしてあります。

下記のサンプルプログラムに対して、"device:CPU:0"、"device:XLA_CPU:0"、"device:XLA_EXEC:0"でテストしました。

import tensorflow as tf

def test_xla():
  config = tf.ConfigProto()
  jit_level = tf.OptimizerOptions.ON_1
  config.graph_options.optimizer_options.global_jit_level = jit_level
  with tf.Session(config=config) as sess:
    x = tf.placeholder(tf.float32, [2], name="x")
    with tf.device("device:XXX:0"):
      y = x * 2
    result = sess.run(y, {x: [1.5, 0.5]})
    print('x * 2 = result : ', result)

if __name__ == '__main__':
  test_xla()

"device:CPU:0"の時は、下記のようになりました。

('x * 2 = result : ', array([ 3.,  1.], dtype=float32))

"device:XLA_CPU:0"の時は、4行のメッセージが表示されましたが、結果は同じです。

I tensorflow/compiler/xla/service/platform_util.cc:58] platform Executor present with 1 visible devices
I tensorflow/compiler/xla/service/platform_util.cc:58] platform Host present with 1 visible devices
I tensorflow/compiler/xla/service/service.cc:187] XLA service 0x3e3e130 executing computations on platform Host. Devices:
I tensorflow/compiler/xla/service/service.cc:195]   StreamExecutor device (0): <undefined>, <undefined>
('x * 2 = result : ', array([ 3.,  1.], dtype=float32))

"device:XLA_EXEC:0"の時は、

F tensorflow/compiler/xla/statusor.cc:41] Attempting to fetch value instead of handling error Not found: could not find registered computation placer for platform Executor -- check target linkage
Aborted (core dumped)

のように、Fatal Errorになり、Aborted (core dumped)になってしまいました。

いろいろ調べた結果、tensorflow/compiler/xla/service/computation_placeer.cc の最後に、
下記のように Plugin 用のコードを追加することで動くようになりました。

変更前
compiler/xla/service/computation_placer.cc

static bool InitModule() {
  xla::ComputationPlacer::RegisterComputationPlacer(se::host::kHostPlatformId, &CreateComputationPlacer);
  xla::ComputationPlacer::RegisterComputationPlacer(se::cuda::kCudaPlatformId, &CreateComputationPlacer);
  return true;
}

static bool module_initialized = InitModule();

変更後
namespace perftools { namespace gputools { namespace executorplugin {
    extern const Platform::Id kExecutorPlatformId;
}  }  } 


static bool InitModule() {
  xla::ComputationPlacer::RegisterComputationPlacer(se::host::kHostPlatformId, &CreateComputationPlacer);
  xla::ComputationPlacer::RegisterComputationPlacer(se::cuda::kCudaPlatformId, &CreateComputationPlacer);
  xla::ComputationPlacer::RegisterComputationPlacer(se::executorplugin::kExecutorPlatformId,
                                                                               &CreateComputationPlacer);
  return true;
}

static bool module_initialized = InitModule();

下記のように、StreamExecutorが起動できるようになりました。

platform Executor present with 1 visible devices
platform Host present with 1 visible devices
XLA service 0x4030a50 executing computations on platform Executor. Devices:
StreamExecutor device (0): Executor, 1.0

ただし、この後に、次のようなエラーが発生してしまいます。
これは、Pluginのコードをそれなりに実装すればなくなると思います。

Traceback (most recent call last):
  File "test_jit.py", line 24, in <module>
    test_xla()
  File "test_jit.py", line 20, in test_xla
    result = sess.run(y, {x: [1.5, 0.5]})
  File ".../tensorflow/python/client/session.py", line 895, in run
    run_metadata_ptr)
  File ".../tensorflow/python/client/session.py", line 1124, in _run
    feed_dict_tensor, options, run_metadata)
  File ".../tensorflow/python/client/session.py", line 1321, in _do_run
    options, run_metadata)
  File ".../tensorflow/python/client/session.py", line 1340, in _do_call
    raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.UnimplementedError: Implicit broadcasting is currently unsupported in HLO evaluator Shape Mismatch: f32[2] vs f32[2] vs f32[]: 
	 [[Node: cluster_0/_0/_1 = _XlaLaunch[Nresources=0, Targs=[DT_FLOAT], Tconstants=[], Tresults=[DT_FLOAT], function=cluster_0[_XlaCompiledKernel=true, _XlaNumConstantArgs=0, _XlaNumResourceArgs=0], _device="/job:localhost/replica:0/task:0/device:XLA_EXEC:0"](_arg_x_0_0/_3)]]

追記)、2017.08.13
昨日、最新コードでやってみたら、上記のcomputation_placer.ccの修正なしでも動くようになりました。
[TF:XLA Fixes to the "evaluator" plugin.]