diff options
| author |  Mark Heffernan <meheff@google.com> | 2018-09-19 08:12:29 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-19 08:16:42 -0700 |
| commit | f8655c08cfe3bd99ec1703211e1c9154a14a6150 (patch) | |
| tree | 90bf5c29d3a1f77764c2f2392c4b0564b490c995 /tensorflow/compiler/xla/service/hlo_rematerialization.cc | |
| parent | e1db78697b05be673562fe2b1c9a995d25a71d4c (diff) | |
Add interface for HLO passes which run on HloModuleGroup.
Derive HloModulePass and HloModuleGroupPass from HloPassInterface which run module-scoped and module-group-scoped respectively. Replace all existing uses of HloPassInterface with HloModulePass because all existing passes are module-scoped. Also rewrite HloPassPipeline to support both module-scoped and module-group-scoped passes.
PiperOrigin-RevId: 213629604
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_rematerialization.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_rematerialization.cc | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc index bd6dd79b67..a438671936 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc @@ -1198,6 +1198,12 @@ StatusOr<bool> HloRematerialization::Run(HloModule* module) { << HumanReadableNumBytes(memory_limit_bytes_); XLA_VLOG_LINES(3, "Before HloRematerialization:\n" + module->ToString()); + // Initialize pass object state. + computation_peak_memory_.clear(); + rematerialized_computations_.clear(); + instructions_rematerialized_ = 0; + net_instructions_added_ = 0; + TF_RET_CHECK(module->has_schedule()); TF_ASSIGN_OR_RETURN(points_to_analysis_, TuplePointsToAnalysis::Run(module)); |