[XLA:Tools] Make hlo_runner understand --xla_hlo_profile.
PiperOrigin-RevId: 188803724
This commit is contained in:
parent
3b47d5f1e2
commit
d9cbe36d1b
@ -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 :
|
||||||
|
@ -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;
|
||||||
|
|
||||||
|
@ -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.
|
||||||
|
@ -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());
|
||||||
}
|
}
|
||||||
|
Loading…
x
Reference in New Issue
Block a user