toy_matmul_cdep.py

import tensorflow as tf
import json
import sys
from tensorflow.python.client import timeline
from protobuf_to_dict import protobuf_to_dict
from google.protobuf.json_format import MessageToJson
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.framework import graph_io
import os

dim = 16384

with tf.device('/gpu:0'):
    X9 = tf.random_uniform([dim, dim], 0, 10)
    _X9 = tf.placeholder(dtype=tf.float32, shape=[dim, dim])
    Z19 = tf.matmul(_X9, _X9)

with tf.device('/gpu:1'):
    Y9 = tf.random_uniform([dim, dim], 0, 10)
    _Y9 = tf.placeholder(dtype=tf.float32, shape=[dim, dim])
    Z29 = tf.matmul(_Y9, _Y9)

    W9 = tf.matmul(Z19, Z29)

n = 8
dim = 32
Z1, Z2, W = [], [], []
X, _X, Y, _Y = [], [], [], []

with tf.device('/gpu:0'):
    for i in range(n):
        dim *= 2

        X.append(tf.random_uniform([dim, dim], 0, 10, name='X' + str(i)))
        Y.append(tf.random_uniform([dim, dim], 0, 10, name='Y' + str(i)))

        with tf.control_dependencies([Z19]):
            _X.append(tf.placeholder(dtype=tf.float32, shape=[dim, dim]))
            Z1.append(tf.matmul(_X[i], _X[i]))

            _Y.append(tf.placeholder(dtype=tf.float32, shape=[dim, dim]))
            Z2.append(tf.matmul(_Y[i], _Y[i]))
            
            W.append(tf.matmul(Z1[i], Z2[i]))


config_proto = tf.ConfigProto(graph_options=tf.GraphOptions(build_cost_model=1))
config_proto.intra_op_parallelism_threads = 1
config_proto.inter_op_parallelism_threads = 1
config_proto.graph_options.optimizer_options.opt_level = -1
config_proto.graph_options.rewrite_options.constant_folding = (rewriter_config_pb2.RewriterConfig.OFF)
config_proto.graph_options.rewrite_options.arithmetic_optimization = (rewriter_config_pb2.RewriterConfig.OFF)

sess = tf.Session(config=config_proto)
sess.run(tf.global_variables_initializer())

X_, Y_, X9_, Y9_ = sess.run([X, Y, X9, Y9])
X_Y_ = X_ + [X9_] + Y_ + [Y9_]
_X_Y = _X + [_X9] + _Y + [_Y9]



run_metadata = tf.RunMetadata()
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, output_partition_graphs=True)


W_ = sess.run(W + [W9],
            {_i: i_ for _i, i_ in zip(_X_Y, X_Y_)},
            options=run_options,
            run_metadata=run_metadata)

jsonObj = MessageToJson(run_metadata)
with open('metadata_matmul.json', 'w') as outfile:
    json.dump(jsonObj, outfile)

trace = timeline.Timeline(step_stats=run_metadata.step_stats)
trace_file = open('timeline_matmul.ctf.json', 'w')
trace_file.write(trace.generate_chrome_trace_format())