How to dump tensorflow XLA LLVM IR?

937 views Asked by At

I used to use the following command in Tensorflow 1.2:

export TF_XLA_FLAGS='--dump_ir_before_passes=true --dump_temp_products_to=./tmp'

for dumping LLVM IR in Tensorflow. However, the definition file of this flag link_to_the_flag_definition is removed in Tensorflow 1.3 and I wonder now how can I get LLVM IR dump?

Here's a test file for your convenience:

import tensorflow as tf
import numpy as np
import os
from tensorflow.python.client import timeline
import json

run_metadata = tf.RunMetadata()
sess = tf.Session()
jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
x = tf.placeholder(np.float32, shape=[1000000])
y = tf.placeholder(np.float32, shape=[1000000])
c = tf.constant(0.1)

with jit_scope():
  z = tf.add(tf.scalar_mul(0.1,x), y)
  ix = np.ones((1000000), dtype=np.float32)
  iy = np.ones((1000000), dtype=np.float32)
  sess.run(z,
           feed_dict={x: ix, y: iy},
           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())
1

There are 1 answers

0
LukGreg On

I have found the flag --xla_dump_ir_to here: https://github.com/tensorflow/tensorflow/issues/11462. It was added in Tensorflow 1.3.