experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Dump the computation's SessionModule as part of the tf_compile rule.

Browse Source 
PiperOrigin-RevId: 172946149
This commit is contained in:
Justin Lebar 2017-10-20 15:55:57 -07:00 committed by TensorFlower Gardener
parent ebcae4a5e3
commit 8ff33271ea
4 changed files with 9 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/aot/compile.cc
View File

@ -97,11 +97,11 @@ Status CompileGraph(const GraphDef& graph_def, const tf2xla::Config& config,
TF_RETURN_IF_ERROR(ConvertGraphDefToXla(graph_def, config, client,
&computation,
&compile_result->has_context_arg));
if (!flags.debug_dir.empty()) {
if (!flags.out_session_module.empty()) {
TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::SessionModule> module,
computation.Snapshot());
string file = io::JoinPath(flags.debug_dir, "tfcompile_xla_module.pb");
TF_RETURN_IF_ERROR(WriteBinaryProto(Env::Default(), file, *module));
TF_RETURN_IF_ERROR(
WriteBinaryProto(Env::Default(), flags.out_session_module, *module));
}
xla::cpu::CpuAotCompilationOptions aot_opts(
flags.target_triple, flags.target_cpu, flags.target_features,

5
tensorflow/compiler/aot/flags.cc
View File

@ -33,9 +33,6 @@ void AppendMainFlags(std::vector<Flag>* flag_list, MainFlags* flags) {
"fetch nodes will be dumped to stdout in a comma-separated list. "
"Typically used to format arguments for other tools, e.g. "
"freeze_graph."},
{"debug_dir", &flags->debug_dir,
"Specifies a directory to dump debugging information, including "
"rewritten graphs and the XLA HLO module."},
// Flags controlling the XLA ahead-of-time compilation, that correspond to
// the fields of xla::cpu::CpuAotCompilationOptions.
//
@ -64,6 +61,8 @@ void AppendMainFlags(std::vector<Flag>* flag_list, MainFlags* flags) {
"namespaces are given, within the global namespace."},
{"out_object", &flags->out_object, "Output object file name."},
{"out_header", &flags->out_header, "Output header file name."},
{"out_session_module", &flags->out_session_module,
"Output session module proto."},
{"gen_name_to_index", &flags->gen_name_to_index,
"Generate name-to-index data for Lookup{Arg,Result}Index methods."},
{"gen_program_shape", &flags->gen_program_shape,

2
tensorflow/compiler/aot/flags.h
View File

@ -29,7 +29,6 @@ struct MainFlags {
string graph;
string config;
bool dump_fetch_nodes = false;
string debug_dir;
string target_triple;
string target_cpu;
string target_features;
@ -37,6 +36,7 @@ struct MainFlags {
string cpp_class;
string out_object;
string out_header;
string out_session_module;
// C++ codegen options
bool gen_name_to_index = false;

3
tensorflow/compiler/aot/tfcompile.bzl
View File

@ -129,6 +129,7 @@ def tf_library(name, graph, config,
# Rule that runs tfcompile to produce the header and object file.
header_file = name + ".h"
object_file = name + ".o"
session_module_pb = name + "_session_module.pb"
ep = ("__" + PACKAGE_NAME + "__" + name).replace("/", "_")
native.genrule(
name=("gen_" + name),
@ -139,6 +140,7 @@ def tf_library(name, graph, config,
outs=[
header_file,
object_file,
session_module_pb,
],
cmd=("$(location " + tfcompile_tool + ")" +
" --graph=$(location " + tfcompile_graph + ")" +
@ -148,6 +150,7 @@ def tf_library(name, graph, config,
" --target_triple=" + target_llvm_triple() +
" --out_header=$(@D)/" + header_file +
" --out_object=$(@D)/" + object_file +
" --out_session_module=$(@D)/" + session_module_pb +
" " + (tfcompile_flags or "")),
tools=[tfcompile_tool],
visibility=visibility,