[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 HloModuleConfig::compilation_cache_key() const {
string key = string key =
tensorflow::strings::StrCat("profiling=", hlo_profiling_enabled_); tensorflow::strings::StrCat("profiling=", hlo_profiling_enabled());
StrAppend(&key, "::("); StrAppend(&key, "::(");
std::vector<string> params; std::vector<string> params;
for (const ShapeLayout& param_layout : for (const ShapeLayout& param_layout :

View File

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

View File

@ -158,8 +158,8 @@ StatusOr<std::unique_ptr<Literal>> HloRunner::Execute(
TF_ASSIGN_OR_RETURN( TF_ASSIGN_OR_RETURN(
std::unique_ptr<ShapedBuffer> result, std::unique_ptr<ShapedBuffer> result,
executable->ExecuteOnStream(&service_run_options, argument_buffer_ptrs, executable->ExecuteOnStreamWrapper(
/*hlo_execution_profile=*/nullptr)); &service_run_options, /*profile=*/nullptr, argument_buffer_ptrs));
// Create a ScopedShapedBuffer of the result to manage deallocation. This will // Create a ScopedShapedBuffer of the result to manage deallocation. This will
// deallocate all the device memory when it goes out of scope. // 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) { if (execution_options != nullptr) {
config->set_seed(execution_options->seed()); config->set_seed(execution_options->seed());
config->set_debug_options(execution_options->debug_options()); config->set_debug_options(execution_options->debug_options());
config->enable_hlo_profiling(
execution_options->debug_options().xla_hlo_profile());
} else { } else {
config->set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); config->set_debug_options(legacy_flags::GetDebugOptionsFromFlags());
} }