[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 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 :
|
||||
|
@ -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;
|
||||
|
||||
|
@ -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.
|
||||
|
@ -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());
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user