Vengineerの妄想(準備期間)

人生は短いけど、長いです。人生を楽しみましょう！

TensorFlow XLAのTPUのOpsがオープンソース化

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

TensorFlow の TPU 関連コードに、こんなのが追加された。

github.com

HostComputeOp

This TensorFlow pseudo-op is used to record host-side computation.

GetItemXlaOp

The Xla kernel to build up the computation for get_item(data, index).

SendToHostOp

RecvFromHostOp

InfeedDequeueOp

This TensorFlow op supports the XLA Infeed primitve.

InfeedDequeueTupleOp

This TensorFlow op supports the XLA Infeed primitive for tuple types.

OutfeedEnqueueOp

This TensorFlow op implements the XLA Outfeed primitive.

OutfeedEnqueueTupleOp

This TensorFlow op implements the XLA Outfeed primitive for tuple types.

InplaceUpdateOp

InplaceAddOp

UnsortedSegmentSum

WhereOp

が XlaOpKernel クラスを継承して定義されていますね。

この中で、

  • HostComputeOp => XlaHostCompute
  • SendToHostOp => XlaSendToHost
  • RecvFromHostOp => XlaRecvFromHost

は、tensorflow/core/tpu/ops/host_compute_ops.cc の中で、TensorFlowの Op として定義されています。今までは、これらOpの中身は非公開だったのが、この commit にて公開されたのでしょうね。