學習筆記TF062:TensorFlow線性代數編譯框架XLA

XLA(Accelerated Linear Algebra)，線性代數領域專用編譯器(demain-specific compiler)，優化TensorFlow計算。即時(just-in-time,JIT)編譯或提早(ahead-of-time,AOT)編譯實現XLA，有助於硬件加速。XLA還在試驗階段。https://www.tensorflow.org/versions/master/experimental/xla/ 。

XLA優點。線性代數領域專用編譯器，優化TensorFlow計算的執行速度(編譯子圖減小生命週期較短操做執行時間，融合管道化操做減小內存佔用)、內存使用(分析、規劃內存使用需求，消除許多中間結果緩存)、自定義操做依賴(提升自動化融合底層操做low-level op性能，達到手動融合自定義操做custom op效果)、移動端內存佔用(提早AOT編譯子圖減小TensorFlow執行時間，共享頭文件對被其餘程序直接連接)、可移植性方面(爲新硬件開發新後端，TensorFlow不須要更改不少代碼用在新硬件設備上)。

XLA工做原理。LLVM編譯器框架系統，C++編寫，優化任意編程語言縮寫程序編譯時間(compile time)、連接時間(link time)、運行時間(run time)、空閒時間(idle time)。前端解析、驗證、論斷輸入代碼錯誤，解析代碼轉換LLVM中間表示(intermdediate representation,IR)。IR分析、優化改進代碼，發送到代碼生成器，產生本地機器代碼。三相設計LLVM實現。最重要，LLVM IR。編譯器IR表示代碼。C->Clang C/C++/ObjC前端、Fortran->llvm-gcc前端、Haskell->GHC前端 LLVM IR-> LLVM 優化器 ->LLVM IR LLVM X86後端->X8六、LLVM PowerPC後端->PowerPC、LLVM ARM後端->ARM。http://www.aosabook.org/en/llvm.html 。
XLA輸入語言HLO IR，XLA HLO定義圖形，編譯成各類體系結構機器指令。編譯過程。XLA HLO->目標無關優化分析->XLA HLO->XLA後端->目標相關優化分析->目標特定代碼生成。XLA首先進行目標無關優化分析(公共子表達式消除common subexpression elimination CSE，目標無關操做融合，運行時內存緩衝區分析)。XLA將HLO計算髮送到後端。後端執行進一步HLO級目標不相關優化分析。XLA GPU後端執行對GPU編程模型有益操做融合，肯定計算劃分紅流。生成目標特定代碼。XLA CPU、GPU後端用LLVM中間表示、優化、代碼生成。後端用LLVM IR表示XLA HLO計算。XLA 支持x86-64､NVIDIA GPU JIT編譯，x86-64､ARM AOT編譯。AOT更適合移動、嵌入式深度學習應用。

JIT編譯方式。XLA編譯、運行TensorFlow計算圖一部分。XLA 將多個操做(內核)融合到少許編譯內核，融合操做符減小存儲器帶寬提升性能。XLA 運行TensorFlow計算方法。一，打開CPU、GPU設備JIT編譯。二，操做符放在XLA_CPU、XLA_GPU設備。
打開JIT編譯。在會話打開。把全部可能操做符編程成XLA計算。

config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
sess = tf.Session(config=config)
爲一個或多個操做符手動打開JIT編譯。屬性_XlaCompile = true標記編譯操做符。

jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
x = tf.placeholder(np.float32)
with jit_scope():
y = tf.add(x, x)
操做符放在XLA設備。有效設備XLA_CPU、XLA_GPU：

with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"):
output = tf.add(input1, input2)

JIT編譯MNIST實現。https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/tutorials/mnist/mnist_softmax_xla.py 。
不使用XLA運行。

python mnist_softmax_xla.py --xla=false
運行完成生成時間線文件timeline.ctf.json，用Chrome跟蹤事件分析器 chrome://tracing，打開時間線文件，呈現時間線。左側列出GPU，能夠看操做符時間消耗狀況。
用XLA訓練模型。

TF_XLA_FLAGS=--xla_generate_hlo_graph=.* python mnist_softmax_xla.py
XLA框架處於試驗階段，AOT主要應用場景內存較小嵌入式設備、手機、樹莓派。

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import argparse
import sys
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.python.client import timeline
FLAGS = None
def main(_):
  # Import data
  mnist = input_data.read_data_sets(FLAGS.data_dir, one_hot=True)
  # Create the model
  x = tf.placeholder(tf.float32, [None, 784])
  w = tf.Variable(tf.zeros([784, 10]))
  b = tf.Variable(tf.zeros([10]))
  y = tf.matmul(x, w) + b
  # Define loss and optimizer
  y_ = tf.placeholder(tf.float32, [None, 10])
  # The raw formulation of cross-entropy,
  #
  #   tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(tf.nn.softmax(y)),
  #                                 reduction_indices=[1]))
  #
  # can be numerically unstable.
  #
  # So here we use tf.nn.softmax_cross_entropy_with_logits on the raw
  # outputs of 'y', and then average across the batch.
  cross_entropy = tf.reduce_mean(
      tf.nn.softmax_cross_entropy_with_logits(labels=y_, logits=y))
  train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)
  config = tf.ConfigProto()
  jit_level = 0
  if FLAGS.xla:
    # Turns on XLA JIT compilation.
    # 開啓XLA JIT編譯
    jit_level = tf.OptimizerOptions.ON_1
  config.graph_options.optimizer_options.global_jit_level = jit_level
  run_metadata = tf.RunMetadata()
  sess = tf.Session(config=config)
  tf.global_variables_initializer().run(session=sess)
  # Train
  # 訓練
  train_loops = 1000
  for i in range(train_loops):
    batch_xs, batch_ys = mnist.train.next_batch(100)
    # Create a timeline for the last loop and export to json to view with
    # chrome://tracing/.
    # 在最後一次循環建立時間線文件，用chrome://tracing/打開分析
    if i == train_loops - 1:
      sess.run(train_step,
               feed_dict={x: batch_xs,
                          y_: batch_ys},
               options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE),
               run_metadata=run_metadata)
      trace = timeline.Timeline(step_stats=run_metadata.step_stats)
      with open('timeline.ctf.json', 'w') as trace_file:
        trace_file.write(trace.generate_chrome_trace_format())
    else:
      sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})
  # Test trained model
  correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
  accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
  print(sess.run(accuracy,
                 feed_dict={x: mnist.test.images,
                            y_: mnist.test.labels}))
  sess.close()
if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--data_dir',
      type=str,
      default='/tmp/tensorflow/mnist/input_data',
      help='Directory for storing input data')
  parser.add_argument(
      '--xla', type=bool, default=True, help='Turn xla via JIT on')
  FLAGS, unparsed = parser.parse_known_args()
  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)

 

參考資料：
《TensorFlow技術解析與實戰》

歡迎推薦上海機器學習工做機會，個人微信：qingxingfengzi