Add option to free up visit state for the DFS HLO visitor
The visit state can be big for large modules and if we keep the
visitor around after use then we often want to deallocate it to
save memory.
PiperOrigin-RevId: 263027955
diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
index 86bed87..94a99c7 100644
--- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
+++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
@@ -307,6 +307,12 @@
visit_state_.erase(visit_state_.begin(), visit_state_.end());
}
+ // Useful when we want to free up the memory used by the visit state without
+ // destroying the actual visitor subclass.
+ void DestroyVisitState() {
+ visit_state_ = absl::flat_hash_map<int, VisitState>{};
+ }
+
void SetVisitState(int id, VisitState state) { visit_state_[id] = state; }
// Sets the visitation state of the given instruction as kVisiting.
