[XLA:Tools] Make hlo_runner understand --xla_hlo_profile.

PiperOrigin-RevId: 188803724
This commit is contained in:
A. Unique TensorFlower 2018-03-12 17:13:09 -07:00 committed by TensorFlower Gardener
parent 3b47d5f1e2
commit d9cbe36d1b
4 changed files with 7 additions and 11 deletions

View File

@ -40,7 +40,7 @@ void HloModuleConfig::SetDefaultComputationLayout(
string HloModuleConfig::compilation_cache_key() const {
string key =
tensorflow::strings::StrCat("profiling=", hlo_profiling_enabled_);
tensorflow::strings::StrCat("profiling=", hlo_profiling_enabled());
StrAppend(&key, "::(");
std::vector<string> params;
for (const ShapeLayout& param_layout :

View File

@ -63,9 +63,10 @@ class HloModuleConfig {
return &(*entry_computation_layout_);
}
// Sets/returns whether to enable HLO-level profiling.
bool hlo_profiling_enabled() const { return hlo_profiling_enabled_; }
void enable_hlo_profiling(bool enabled) { hlo_profiling_enabled_ = enabled; }
// Returns whether to enable HLO-level profiling.
bool hlo_profiling_enabled() const {
return debug_options_.xla_hlo_profile();
}
// Sets/returns whether this is a "host module". Host modules are used to
// record the data- and control-flow dependencies of host side computation
@ -110,9 +111,6 @@ class HloModuleConfig {
tensorflow::gtl::optional<ComputationLayout> entry_computation_layout_;
// Whether to enable HLO-level profiling.
bool hlo_profiling_enabled_ = false;
// Whether this is a 'host module'.
bool is_host_module_ = false;

View File

@ -158,8 +158,8 @@ StatusOr<std::unique_ptr<Literal>> HloRunner::Execute(
TF_ASSIGN_OR_RETURN(
std::unique_ptr<ShapedBuffer> result,
executable->ExecuteOnStream(&service_run_options, argument_buffer_ptrs,
/*hlo_execution_profile=*/nullptr));
executable->ExecuteOnStreamWrapper(
&service_run_options, /*profile=*/nullptr, argument_buffer_ptrs));
// Create a ScopedShapedBuffer of the result to manage deallocation. This will
// deallocate all the device memory when it goes out of scope.

View File

@ -314,8 +314,6 @@ StatusOr<std::unique_ptr<HloModuleConfig>> Service::CreateModuleConfig(
if (execution_options != nullptr) {
config->set_seed(execution_options->seed());
config->set_debug_options(execution_options->debug_options());
config->enable_hlo_profiling(
execution_options->debug_options().xla_hlo_profile());
} else {
config->set_debug_options(legacy_flags::GetDebugOptionsFromFlags());
}