diff options
author | Benjamin Kramer <kramerb@google.com> | 2018-10-01 13:43:49 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-10-01 13:53:07 -0700 |
commit | 3039a4694e22674b502257ae34b0a5b614a631f3 (patch) | |
tree | 423fdfa7a2e7dd2740af97accfe848bc97b335d0 | |
parent | 44acd839c57494860666c799afd24360f1df3bed (diff) |
[XLA] Migrate from gtl::FlatMap to absl::flat_hash_map
PiperOrigin-RevId: 215272497
80 files changed, 319 insertions, 259 deletions
diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index 5bf4af1014..29b60d1dbe 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -258,6 +258,7 @@ cc_library( "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", "//tensorflow/core/kernels:variable_ops", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -323,6 +324,7 @@ cc_library( "//tensorflow/core:graph", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", @@ -400,6 +402,7 @@ cc_library( "//tensorflow/core:protos_all_cc", "//tensorflow/core/kernels:bounds_check", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], @@ -471,6 +474,7 @@ tf_cc_test( "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -509,6 +513,7 @@ tf_cc_test( "//tensorflow/core:test_main", "//tensorflow/core:testlib", "//tensorflow/core/grappler/optimizers/data:graph_utils", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], diff --git a/tensorflow/compiler/jit/deadness_analysis.cc b/tensorflow/compiler/jit/deadness_analysis.cc index 25e2e9a7af..e63d4b7792 100644 --- a/tensorflow/compiler/jit/deadness_analysis.cc +++ b/tensorflow/compiler/jit/deadness_analysis.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/jit/deadness_analysis.h" #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/jit/deadness_analysis_internal.h" #include "tensorflow/core/graph/algorithm.h" @@ -420,15 +421,15 @@ class PredicateFactory { } }; - gtl::FlatMap<SignatureForAndOr, std::unique_ptr<Predicate>, - HashSignatureForAndOr> + absl::flat_hash_map<SignatureForAndOr, std::unique_ptr<Predicate>, + HashSignatureForAndOr> interned_and_or_instances_; - gtl::FlatMap<SignatureForNot, std::unique_ptr<Predicate>> + absl::flat_hash_map<SignatureForNot, std::unique_ptr<Predicate>> interned_not_instances_; - gtl::FlatMap<SignatureForAndRec, std::unique_ptr<Predicate>> + absl::flat_hash_map<SignatureForAndRec, std::unique_ptr<Predicate>> interned_and_rec_instances_; - gtl::FlatMap<SignatureForSymbol, std::unique_ptr<Predicate>, - HashSignatureForSymbol> + absl::flat_hash_map<SignatureForSymbol, std::unique_ptr<Predicate>, + HashSignatureForSymbol> interned_symbol_instances_; }; @@ -572,7 +573,8 @@ class DeadnessAnalysisImpl : public DeadnessAnalysis { Status PopulateWithReversePostOrder(absl::Span<Node* const> rpo); bool HasInputsWithMismatchingDeadness(const Node& node) override; void Print() const override; - gtl::FlatMap<TensorId, string, TensorId::Hasher> PredicateMapAsString() const; + absl::flat_hash_map<TensorId, string, TensorId::Hasher> PredicateMapAsString() + const; private: enum class EdgeKind { kDataAndControl, kDataOnly, kControlOnly }; @@ -614,7 +616,7 @@ class DeadnessAnalysisImpl : public DeadnessAnalysis { Status HandleNode(Node* n, std::vector<bool>* should_revisit); const Graph& graph_; - gtl::FlatMap<TensorId, Predicate*, TensorId::Hasher> predicate_map_; + absl::flat_hash_map<TensorId, Predicate*, TensorId::Hasher> predicate_map_; PredicateFactory predicate_factory_; bool vlog_; }; @@ -977,9 +979,9 @@ DeadnessAnalysis::~DeadnessAnalysis() {} return Status::OK(); } -gtl::FlatMap<TensorId, string, TensorId::Hasher> +absl::flat_hash_map<TensorId, string, TensorId::Hasher> DeadnessAnalysisImpl::PredicateMapAsString() const { - gtl::FlatMap<TensorId, string, TensorId::Hasher> result; + absl::flat_hash_map<TensorId, string, TensorId::Hasher> result; std::vector<TensorId> tensor_ids; for (const auto& kv_pair : predicate_map_) { CHECK(result.insert({kv_pair.first, kv_pair.second->ToString()}).second); diff --git a/tensorflow/compiler/jit/deadness_analysis_internal.h b/tensorflow/compiler/jit/deadness_analysis_internal.h index 3df2679c62..354782374a 100644 --- a/tensorflow/compiler/jit/deadness_analysis_internal.h +++ b/tensorflow/compiler/jit/deadness_analysis_internal.h @@ -16,15 +16,15 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_ #define TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_ +#include "absl/container/flat_hash_map.h" #include "tensorflow/core/graph/tensor_id.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace tensorflow { namespace deadness_analysis_internal { // Returns a map describing the predicate each Tensor was mapped to. For // testing purposes only. -using PredicateMapTy = gtl::FlatMap<TensorId, string, TensorId::Hasher>; +using PredicateMapTy = absl::flat_hash_map<TensorId, string, TensorId::Hasher>; Status ComputePredicates(const Graph& graph, PredicateMapTy* out_predicate_map); // Returns a map describing the predicate each Tensor was mapped to. For diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index 0839f1cb3d..26cb3af9d6 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -26,6 +26,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core/kernels:variable_ops", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", ], alwayslink = 1, diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index a85006eb03..cfd27a6510 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/jit/kernels/xla_ops.h" +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/jit/defs.h" #include "tensorflow/compiler/tf2xla/shape_util.h" @@ -163,7 +164,7 @@ class XlaExecutableClosureStore { private: mutex mutex_; int64 key_counter_ GUARDED_BY(mutex_); - gtl::FlatMap<KeyT, XlaExecutableClosure> closures_ GUARDED_BY(mutex_); + absl::flat_hash_map<KeyT, XlaExecutableClosure> closures_ GUARDED_BY(mutex_); TF_DISALLOW_COPY_AND_ASSIGN(XlaExecutableClosureStore); }; diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc index 4f9145b479..2a80c745e3 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/jit/mark_for_compilation_pass_test_helper.h" +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "absl/strings/match.h" #include "tensorflow/cc/framework/ops.h" @@ -61,10 +62,10 @@ std::unordered_map<string, string> GetClusters(const Graph& graph) { return ids; } -gtl::FlatMap<string, std::vector<string>> GetClusterSets( +absl::flat_hash_map<string, std::vector<string>> GetClusterSets( const Graph& g, std::vector<string>* cluster_names = nullptr) { CHECK(cluster_names == nullptr || cluster_names->empty()); - gtl::FlatMap<string, std::vector<string>> cluster_sets; + absl::flat_hash_map<string, std::vector<string>> cluster_sets; for (const auto& p : GetClusters(g)) { cluster_sets[p.second].push_back(p.first); } @@ -566,7 +567,7 @@ TEST(XlaCompilationTest, ResourcesClusteringAllowed) { std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global())); TF_EXPECT_OK(root.ToGraph(graph.get())); TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); - gtl::FlatMap<string, std::vector<string>> cluster_sets = + absl::flat_hash_map<string, std::vector<string>> cluster_sets = GetClusterSets(*graph); ASSERT_EQ(cluster_sets.size(), 1); std::vector<string> expected_clustered_nodes = {"AssignmentW", "ReadR", @@ -586,7 +587,7 @@ TEST(XlaCompilationTest, ResourcesClusteringDisallowed) { std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global())); TF_EXPECT_OK(root.ToGraph(graph.get())); TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); - gtl::FlatMap<string, std::vector<string>> cluster_sets = + absl::flat_hash_map<string, std::vector<string>> cluster_sets = GetClusterSets(*graph); ASSERT_EQ(cluster_sets.size(), 1); std::vector<string> expected_clustered_nodes = {"AssignmentW", @@ -616,7 +617,7 @@ TEST(XlaCompilationTest, ChainOfOps) { TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); std::vector<string> cluster_names; - gtl::FlatMap<string, std::vector<string>> cluster_sets = + absl::flat_hash_map<string, std::vector<string>> cluster_sets = GetClusterSets(*graph, &cluster_names); ASSERT_EQ(cluster_sets.size(), 2); diff --git a/tensorflow/compiler/jit/resource_operation_safety_analysis.cc b/tensorflow/compiler/jit/resource_operation_safety_analysis.cc index 56e35c0059..657bb409db 100644 --- a/tensorflow/compiler/jit/resource_operation_safety_analysis.cc +++ b/tensorflow/compiler/jit/resource_operation_safety_analysis.cc @@ -89,7 +89,6 @@ limitations under the License. #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/tensor_id.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/util/ptr_util.h" diff --git a/tensorflow/compiler/jit/xla_compilation_cache.h b/tensorflow/compiler/jit/xla_compilation_cache.h index 10ad87e38c..17c0321c1e 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.h +++ b/tensorflow/compiler/jit/xla_compilation_cache.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_XLA_COMPILATION_CACHE_H_ #define TENSORFLOW_COMPILER_JIT_XLA_COMPILATION_CACHE_H_ +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_context.h" #include "tensorflow/compiler/xla/client/local_client.h" @@ -24,7 +25,6 @@ limitations under the License. #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/lib/core/threadpool.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/thread_annotations.h" @@ -152,7 +152,7 @@ class XlaCompilationCache : public ResourceBase { }; mutex compile_cache_mu_; - gtl::FlatMap<Signature, std::unique_ptr<Entry>, Signature::Hash> cache_ + absl::flat_hash_map<Signature, std::unique_ptr<Entry>, Signature::Hash> cache_ GUARDED_BY(compile_cache_mu_); struct CompileStats { @@ -165,7 +165,7 @@ class XlaCompilationCache : public ResourceBase { mutex compile_stats_mu_; // Maps cluster names to compilation statistics for said cluster. - gtl::FlatMap<string, CompileStats> compile_stats_ + absl::flat_hash_map<string, CompileStats> compile_stats_ GUARDED_BY(compile_stats_mu_); TF_DISALLOW_COPY_AND_ASSIGN(XlaCompilationCache); diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index ba1e3b2b4f..3f631f91ec 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -635,6 +635,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:ops", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", ], ) @@ -649,6 +650,7 @@ tf_cc_test( "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", ], ) diff --git a/tensorflow/compiler/tf2xla/resource_operation_table.cc b/tensorflow/compiler/tf2xla/resource_operation_table.cc index 20f2ce2919..72b240996f 100644 --- a/tensorflow/compiler/tf2xla/resource_operation_table.cc +++ b/tensorflow/compiler/tf2xla/resource_operation_table.cc @@ -15,7 +15,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/resource_operation_table.h" #include "absl/algorithm/container.h" -#include "tensorflow/core/lib/gtl/flatmap.h" +#include "absl/container/flat_hash_map.h" namespace tensorflow { /*static*/ absl::string_view XlaResourceOpInfo::XlaResourceOpKindToString( @@ -30,9 +30,9 @@ namespace tensorflow { } } -static gtl::FlatMap<absl::string_view, XlaResourceOpInfo>* +static absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>* CreateResourceOpInfoMap() { - auto* result = new gtl::FlatMap<absl::string_view, XlaResourceOpInfo>; + auto* result = new absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>; auto add = [&](absl::string_view op, XlaResourceOpKind op_kind, XlaResourceKind resource_kind) { @@ -103,15 +103,15 @@ CreateResourceOpInfoMap() { return result; } -static const gtl::FlatMap<absl::string_view, XlaResourceOpInfo>& +static const absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>& GetStaticResourceOpInfoMap() { - static gtl::FlatMap<absl::string_view, XlaResourceOpInfo>* op_info_map = - CreateResourceOpInfoMap(); + static absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>* + op_info_map = CreateResourceOpInfoMap(); return *op_info_map; } const XlaResourceOpInfo* GetResourceOpInfoForOp(absl::string_view op) { - const gtl::FlatMap<absl::string_view, XlaResourceOpInfo>& op_infos = + const absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>& op_infos = GetStaticResourceOpInfoMap(); auto it = op_infos.find(op); return it == op_infos.end() ? nullptr : &it->second; diff --git a/tensorflow/compiler/tf2xla/resource_operation_table_test.cc b/tensorflow/compiler/tf2xla/resource_operation_table_test.cc index a85ef040a7..956f597301 100644 --- a/tensorflow/compiler/tf2xla/resource_operation_table_test.cc +++ b/tensorflow/compiler/tf2xla/resource_operation_table_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/resource_operation_table.h" #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -33,7 +34,7 @@ bool HasResourceInputOrOutput(const OpDef& op_def) { } TEST(ResourceOperationTableTest, HaveAllResourceOps) { - gtl::FlatMap<string, bool> known_resource_ops; + absl::flat_hash_map<string, bool> known_resource_ops; for (absl::string_view known_resource_op : resource_op_table_internal::GetKnownResourceOps()) { ASSERT_TRUE( diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD index f825f67b44..1191cff109 100644 --- a/tensorflow/compiler/xla/client/BUILD +++ b/tensorflow/compiler/xla/client/BUILD @@ -220,6 +220,7 @@ cc_library( "//tensorflow/compiler/xla/service:shape_inference", "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h index 1da6ddd318..b7295e8a53 100644 --- a/tensorflow/compiler/xla/client/xla_builder.h +++ b/tensorflow/compiler/xla/client/xla_builder.h @@ -21,6 +21,7 @@ limitations under the License. #include <type_traits> #include <utility> +#include "absl/container/flat_hash_map.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/client/padding.h" @@ -34,7 +35,6 @@ limitations under the License. #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/stacktrace.h" @@ -1027,7 +1027,7 @@ class XlaBuilder { // A map from XlaOp::Handle to the index in the instructions_ vector where the // instruction is held. - tensorflow::gtl::FlatMap<int64, int64> handle_to_index_; + absl::flat_hash_map<int64, int64> handle_to_index_; // The embedded computations used by this computation. Each computation was // the entry computation of some XlaComputation, the key is the unique id of diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index e800cf470c..8da6364786 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -146,6 +146,7 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -250,6 +251,7 @@ cc_library( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", @@ -333,6 +335,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", @@ -395,6 +398,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/types:span", ], ) @@ -485,6 +489,7 @@ cc_library( "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -903,6 +908,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], @@ -952,6 +958,7 @@ cc_library( deps = [ "//tensorflow/compiler/xla:types", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", ], ) @@ -987,6 +994,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", ], @@ -1034,6 +1042,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -1087,6 +1096,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", ], @@ -1125,6 +1135,7 @@ cc_library( "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", ], ) @@ -1146,6 +1157,7 @@ tf_cc_test( "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", ], ) @@ -1196,6 +1208,7 @@ cc_library( "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:optional", ], @@ -1216,6 +1229,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", @@ -1260,6 +1274,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -1280,6 +1295,7 @@ tf_cc_test( "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -1304,6 +1320,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", ], ) @@ -1330,6 +1347,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", ], ) @@ -1385,6 +1403,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", ], @@ -1640,6 +1659,7 @@ cc_library( ":while_loop_analysis", "//tensorflow/compiler/xla:statusor", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", ], @@ -1671,6 +1691,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -2203,6 +2224,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", ], @@ -2263,6 +2285,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", @@ -2319,6 +2342,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -2345,6 +2369,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", ], ) @@ -2416,6 +2441,7 @@ cc_library( "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], @@ -2460,6 +2486,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -2588,6 +2615,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -2701,6 +2729,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", ], ) @@ -3147,6 +3176,7 @@ cc_library( ":hlo_pass_pipeline", "//tensorflow/compiler/xla:shape_util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -3269,6 +3299,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", ], ) @@ -3298,6 +3329,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", ], ) @@ -3354,6 +3386,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:ptr_util", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", diff --git a/tensorflow/compiler/xla/service/allocation_tracker.h b/tensorflow/compiler/xla/service/allocation_tracker.h index a7d8927cf7..af227fe4da 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.h +++ b/tensorflow/compiler/xla/service/allocation_tracker.h @@ -22,6 +22,7 @@ limitations under the License. #include <string> #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" @@ -110,7 +111,7 @@ class AllocationTracker { // A map from device memory opaque value to allocation. One such map is // maintained per device ordinal. - using AllocationMap = tensorflow::gtl::FlatMap<const void*, Allocation>; + using AllocationMap = absl::flat_hash_map<const void*, Allocation>; tensorflow::mutex mutex_; @@ -146,7 +147,7 @@ class AllocationTracker { // non-owning "view" into a tuple's sub-buffers. The sub-buffers are then // free'd when both the view *and* the original tuple are Unregistered. This // refcounting is managed in opaque_to_allocation_map_. - tensorflow::gtl::FlatMap<int64, std::vector<std::unique_ptr<ShapedBuffer>>> + absl::flat_hash_map<int64, std::vector<std::unique_ptr<ShapedBuffer>>> handle_to_shaped_buffers_ GUARDED_BY(mutex_); TF_DISALLOW_COPY_AND_ASSIGN(AllocationTracker); diff --git a/tensorflow/compiler/xla/service/batchnorm_expander.cc b/tensorflow/compiler/xla/service/batchnorm_expander.cc index 30d33e0d35..f70f6ddfec 100644 --- a/tensorflow/compiler/xla/service/batchnorm_expander.cc +++ b/tensorflow/compiler/xla/service/batchnorm_expander.cc @@ -35,7 +35,6 @@ limitations under the License. #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation.h b/tensorflow/compiler/xla/service/bfloat16_propagation.h index 6a62439f88..c74326f631 100644 --- a/tensorflow/compiler/xla/service/bfloat16_propagation.h +++ b/tensorflow/compiler/xla/service/bfloat16_propagation.h @@ -21,6 +21,7 @@ limitations under the License. #include <unordered_set> #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/bfloat16_support.h" #include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -186,7 +187,7 @@ class BFloat16Propagation : public HloModulePass { // Mapping from each HloComputation to the number of callers to it in the // module. Populated at the beginning of this pass. - tensorflow::gtl::FlatMap<const HloComputation*, int64> caller_counts_; + absl::flat_hash_map<const HloComputation*, int64> caller_counts_; // We first store the potential F32-to-BF16 changes to changes_to_bf16_, which // are subject to further adjustment, then finally applied to the HLOs. This @@ -195,8 +196,7 @@ class BFloat16Propagation : public HloModulePass { // // For each HloInstruction, changes_to_bf16_ stores the affected buffers in // the output as a map from in-place pointers to subshapes to shape indices. - tensorflow::gtl::FlatMap<HloInstruction*, - tensorflow::gtl::FlatMap<Shape*, ShapeIndex>> + absl::flat_hash_map<HloInstruction*, absl::flat_hash_map<Shape*, ShapeIndex>> changes_to_bf16_; // Whether the last processed HLO module has been changed by this pass. diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc index 34a7be0e9c..3efa0b1dad 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment.cc @@ -22,6 +22,7 @@ limitations under the License. #include <ostream> #include <utility> +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" @@ -41,9 +42,9 @@ limitations under the License. namespace xla { namespace { +using absl::flat_hash_map; using absl::StrAppend; using absl::StrAppendFormat; -using ::tensorflow::gtl::FlatMap; using ::tensorflow::gtl::FlatSet; using ::tensorflow::strings::HumanReadableNumBytes; @@ -519,7 +520,8 @@ void BufferAssignment::AddAssignment(BufferAllocation* allocation, // BufferAllocation. void BufferAssignment::CombineTempAllocations() { VLOG(1) << "CombineTempAllocations()"; - FlatMap<LogicalBuffer::Color, BufferAllocation, LogicalBuffer::Color::Hasher> + flat_hash_map<LogicalBuffer::Color, BufferAllocation, + LogicalBuffer::Color::Hasher> combined_allocation_map; // Move all temp allocations into a single run at the end of the allocations @@ -582,7 +584,8 @@ void BufferAssignment::CombineTempAllocations() { } // Update allocation indices to their new positions. - allocation_index_for_buffer_.clear_no_resize(); + allocation_index_for_buffer_.erase(allocation_index_for_buffer_.begin(), + allocation_index_for_buffer_.end()); for (size_t index = 0; index < allocations_.size(); ++index) { BufferAllocation* allocation = &allocations_[index]; allocation->set_index(index); @@ -814,7 +817,7 @@ Status BufferAssigner::AssignBuffersForComputation( const HloComputation* computation, bool is_thread_local, const FlatSet<const LogicalBuffer*>& colocated_buffers, const FlatSet<BufferAllocation::Index>& colocated_allocations, - FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>>* + flat_hash_map<const HloComputation*, FlatSet<const LogicalBuffer*>>* buffers_to_assign_sequentially, BufferAssignment* assignment) { // Buffers are sorted and assigned to BufferAllocations in decreasing order of @@ -833,7 +836,7 @@ Status BufferAssigner::AssignBuffersForComputation( // Generate a post order sort of instructions for sorting of the // LogicalBuffers. - FlatMap<const HloInstruction*, int> post_order_position; + flat_hash_map<const HloInstruction*, int> post_order_position; int position = 0; for (auto* instruction : computation->MakeInstructionPostOrder()) { post_order_position.emplace(instruction, position); @@ -1043,12 +1046,12 @@ Status BufferAssigner::AssignBuffersForComputation( return Status::OK(); } -FlatMap<LogicalBuffer::Color, FlatSet<const LogicalBuffer*>, - LogicalBuffer::Color::Hasher> +flat_hash_map<LogicalBuffer::Color, FlatSet<const LogicalBuffer*>, + LogicalBuffer::Color::Hasher> BufferAssigner::SplitBuffersByColor( const FlatSet<const LogicalBuffer*>& buffers) { - FlatMap<LogicalBuffer::Color, FlatSet<const LogicalBuffer*>, - LogicalBuffer::Color::Hasher> + flat_hash_map<LogicalBuffer::Color, FlatSet<const LogicalBuffer*>, + LogicalBuffer::Color::Hasher> color_map; for (auto buffer : buffers) { color_map[buffer->color()].insert(buffer); @@ -1057,7 +1060,7 @@ BufferAssigner::SplitBuffersByColor( } Status BufferAssigner::AssignBuffersWithSequentialOrdering( - const FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>>& + const flat_hash_map<const HloComputation*, FlatSet<const LogicalBuffer*>>& buffers_to_assign_sequentially, bool run_whole_module_heap_simulation, BufferAssignment* assignment) { // Run the sequence of instructions through the heap simulator. The heuristic @@ -1155,9 +1158,8 @@ std::vector<const LogicalBuffer*> ComputePeakMemoryLogicalBuffers( const BufferAllocation& allocation, const HeapSimulatorTrace& heap_trace) { // Create a map from LogicalBuffer::Id to LogicalBuffer* for the logical // buffers in this allocation. - tensorflow::gtl::FlatMap<LogicalBuffer::Id, const LogicalBuffer*> - id_to_buffer; - tensorflow::gtl::FlatMap<const LogicalBuffer*, int64> buffer_sizes; + absl::flat_hash_map<LogicalBuffer::Id, const LogicalBuffer*> id_to_buffer; + absl::flat_hash_map<const LogicalBuffer*, int64> buffer_sizes; for (const auto& pair : allocation.assigned_buffers()) { const LogicalBuffer* buffer = pair.first; const BufferAllocation::OffsetSize& offset_size = pair.second; @@ -1679,7 +1681,7 @@ StatusOr<std::unique_ptr<BufferAssignment>> BufferAssigner::CreateAssignment( // First assign buffers for global computatations. Temporary buffers for // sequential computations are collected in 'buffers_to_assign_sequentially'. - FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>> + flat_hash_map<const HloComputation*, FlatSet<const LogicalBuffer*>> buffers_to_assign_sequentially; for (auto* computation : global_computations) { TF_RETURN_IF_ERROR(AssignBuffersForComputation( diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h index 24ba7c16f5..9ba40617a3 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.h +++ b/tensorflow/compiler/xla/service/buffer_assignment.h @@ -22,6 +22,7 @@ limitations under the License. #include <string> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/heap_simulator.h" @@ -33,7 +34,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" @@ -148,7 +148,7 @@ class BufferAllocation { // Access to the logical buffers assigned to this allocation, and their // associated logical offsets and sizes. - const tensorflow::gtl::FlatMap<const LogicalBuffer*, OffsetSize>& + const absl::flat_hash_map<const LogicalBuffer*, OffsetSize>& assigned_buffers() const { return assigned_buffers_; } @@ -323,7 +323,7 @@ class BufferAllocation { // Mapping from the set of buffers assigned to this allocation to their // logical offsets and sizes. - tensorflow::gtl::FlatMap<const LogicalBuffer*, OffsetSize> assigned_buffers_; + absl::flat_hash_map<const LogicalBuffer*, OffsetSize> assigned_buffers_; int64 fragmentation_bytes_ = 0; std::vector<HeapSimulatorTrace> heap_traces_; @@ -500,7 +500,7 @@ class BufferAssignment { int64 temp_allocation_total_size_ = 0; // Maps Buffers to the index of the BufferAllocation which holds the buffer. - tensorflow::gtl::FlatMap<const LogicalBuffer*, BufferAllocation::Index> + absl::flat_hash_map<const LogicalBuffer*, BufferAllocation::Index> allocation_index_for_buffer_; const HloModule* module_; @@ -557,8 +557,8 @@ class BufferAssigner { const tensorflow::gtl::FlatSet<const LogicalBuffer*>& colocated_buffers, const tensorflow::gtl::FlatSet<BufferAllocation::Index>& colocated_allocations, - tensorflow::gtl::FlatMap<const HloComputation*, - tensorflow::gtl::FlatSet<const LogicalBuffer*>>* + absl::flat_hash_map<const HloComputation*, + tensorflow::gtl::FlatSet<const LogicalBuffer*>>* buffers_to_assign_sequentially, BufferAssignment* assignment); @@ -568,9 +568,8 @@ class BufferAssigner { // 'run_whole_module_heap_simulation' is true, the heap simulation will be run // assuming all global computations are sequentially ordered. Status AssignBuffersWithSequentialOrdering( - const tensorflow::gtl::FlatMap< - const HloComputation*, - tensorflow::gtl::FlatSet<const LogicalBuffer*>>& + const absl::flat_hash_map<const HloComputation*, + tensorflow::gtl::FlatSet<const LogicalBuffer*>>& buffers_to_assign_sequentially, bool run_whole_module_heap_simulation, BufferAssignment* assignment); @@ -624,9 +623,9 @@ class BufferAssigner { // Split a set of buffers into several sets, each of which contains buffers // colored with the same color. - tensorflow::gtl::FlatMap<LogicalBuffer::Color, - tensorflow::gtl::FlatSet<const LogicalBuffer*>, - LogicalBuffer::Color::Hasher> + absl::flat_hash_map<LogicalBuffer::Color, + tensorflow::gtl::FlatSet<const LogicalBuffer*>, + LogicalBuffer::Color::Hasher> SplitBuffersByColor( const tensorflow::gtl::FlatSet<const LogicalBuffer*>& buffers); diff --git a/tensorflow/compiler/xla/service/buffer_liveness.h b/tensorflow/compiler/xla/service/buffer_liveness.h index cdd3cf4032..2911bbcfbf 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness.h +++ b/tensorflow/compiler/xla/service/buffer_liveness.h @@ -27,7 +27,6 @@ limitations under the License. #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/call_graph.h b/tensorflow/compiler/xla/service/call_graph.h index 3af2ab5edf..0c2e9b99db 100644 --- a/tensorflow/compiler/xla/service/call_graph.h +++ b/tensorflow/compiler/xla/service/call_graph.h @@ -20,10 +20,10 @@ limitations under the License. #include <ostream> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" namespace xla { @@ -157,7 +157,7 @@ class CallGraphNode { // The map from instruction to index in callsites_ for looking up the callsite // (if any) associated with a particular instruction in this computation. - tensorflow::gtl::FlatMap<const HloInstruction*, int64> callsite_instructions_; + absl::flat_hash_map<const HloInstruction*, int64> callsite_instructions_; // The call sites in other computations which call this computation. std::vector<CallSite> caller_callsites_; @@ -267,7 +267,7 @@ class CallGraph { // Map from HLO computation to the index of the corresponding call graph node // in nodes_. - tensorflow::gtl::FlatMap<const HloComputation*, int64> node_indices_; + absl::flat_hash_map<const HloComputation*, int64> node_indices_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc index b65dfef9c9..7f78412924 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.cc +++ b/tensorflow/compiler/xla/service/copy_insertion.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/copy_insertion.h" +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/xla/service/hlo_alias_analysis.h" @@ -31,7 +32,6 @@ limitations under the License. #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/logging.h" @@ -432,7 +432,7 @@ class CopyRemover { // Construct a list for each HLO buffer in the alias analysis. Maintain a // map from HloValue to the respective list element representing that // value. The map is used to construct the copy info map below. - tensorflow::gtl::FlatMap<const HloValue*, ValueNode*> value_to_node; + absl::flat_hash_map<const HloValue*, ValueNode*> value_to_node; for (const HloBuffer& buffer : alias_analysis.buffers()) { // Verify values contained in the buffer are strictly ordered. This // should always be the case after adding copies to eliminate @@ -480,7 +480,7 @@ class CopyRemover { // respective ValueNode representing that value. void AddValueList( absl::Span<const HloValue* const> values, - tensorflow::gtl::FlatMap<const HloValue*, ValueNode*>* value_to_node) { + absl::flat_hash_map<const HloValue*, ValueNode*>* value_to_node) { ValueNode* tail = nullptr; ValueNode* head = nullptr; for (const HloValue* value : values) { @@ -516,8 +516,7 @@ class CopyRemover { // respective ValueNode. void CreateCopyMap( const HloModule& module, - const tensorflow::gtl::FlatMap<const HloValue*, ValueNode*>& - value_to_node) { + const absl::flat_hash_map<const HloValue*, ValueNode*>& value_to_node) { for (HloComputation* computation : module.computations()) { for (HloInstruction* instruction : computation->instructions()) { // Add copies with unambiguous source values to the map. Copies with @@ -916,7 +915,7 @@ class CopyRemover { ValueNode* src = nullptr; ValueNode* dest = nullptr; }; - tensorflow::gtl::FlatMap<const HloInstruction*, CopyNodes> copy_map_; + absl::flat_hash_map<const HloInstruction*, CopyNodes> copy_map_; }; HloModule* module_; diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD index b7103118ac..6a83909a3b 100644 --- a/tensorflow/compiler/xla/service/cpu/BUILD +++ b/tensorflow/compiler/xla/service/cpu/BUILD @@ -290,6 +290,7 @@ cc_library( "//tensorflow/compiler/xla/service/llvm_ir:loop_emitter", "//tensorflow/compiler/xla/service/llvm_ir:tuple_ops", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/types:span", @@ -309,6 +310,7 @@ cc_library( deps = [ "//tensorflow/compiler/xla:shape_util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@llvm//:analysis", "@llvm//:target", ], @@ -471,6 +473,7 @@ cc_library( "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/core:lib", "//tensorflow/stream_executor", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", ], @@ -762,6 +765,7 @@ cc_library( "//tensorflow/compiler/xla/service:computation_layout", "//tensorflow/compiler/xla/service:layout_assignment", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", ], ) diff --git a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc index bfecbd6e01..c291bf2d1b 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc @@ -17,6 +17,7 @@ limitations under the License. #include <numeric> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/cpu/dot_op_emitter.h" #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" @@ -38,7 +39,7 @@ using absl::nullopt; using absl::optional; using ShouldMakeOperandColMajorCache = - tensorflow::gtl::FlatMap<const HloInstruction*, bool>; + absl::flat_hash_map<const HloInstruction*, bool>; } // namespace static bool ShouldMakeAllUsersColMajor(const HloInstruction* instruction) { diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc index 20cf855735..a9febe891b 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc @@ -17,6 +17,7 @@ limitations under the License. #include <functional> +#include "absl/container/flat_hash_map.h" #include "absl/synchronization/mutex.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/core/platform/dynamic_annotations.h" @@ -30,8 +31,7 @@ namespace cpu { namespace runtime { XfeedManager* GetXfeedManager(int device_ordinal) { - static tensorflow::gtl::FlatMap<int, XfeedManager*>* managers = - new tensorflow::gtl::FlatMap<int, XfeedManager*>(); + static auto* managers = new absl::flat_hash_map<int, XfeedManager*>(); static absl::Mutex* mutex = new absl::Mutex(); absl::MutexLock lock(mutex); diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index c3e8020783..953a75c35f 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc @@ -24,6 +24,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/core/lib/math/math_util.h" #include "tensorflow/core/platform/logging.h" // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" @@ -67,7 +68,6 @@ limitations under the License. #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/core/lib/core/bits.h" #include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" namespace xla { @@ -1398,7 +1398,7 @@ static bool ReductionPreservesLayout(const HloInstruction& reduce) { // // So if we reduce f32[A,B,C,D] on dimensions 1 and 2, this map contains // [0->0, 3->1]. - gtl::FlatMap<int64, int64> unreduced_dim_map; + absl::flat_hash_map<int64, int64> unreduced_dim_map; gtl::FlatSet<int64> reduced_dims(reduce.dimensions().begin(), reduce.dimensions().end()); diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h index daafef4eb3..586f27b104 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h @@ -23,6 +23,7 @@ limitations under the License. #include <unordered_map> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "llvm/ADT/Triple.h" @@ -47,7 +48,6 @@ limitations under the License. #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" @@ -427,7 +427,7 @@ class IrEmitter : public DfsHloVisitorWithDefault, // Maps the buffer allocation slices for the parameters to the computation // being compiled to their parameter numbers. Only relevant for thread local // computations. - tensorflow::gtl::FlatMap<BufferAllocation::Index, int64> + absl::flat_hash_map<BufferAllocation::Index, int64> computation_parameter_allocations_; // Maps HLO instructions to their index into the profile counter array. @@ -567,11 +567,11 @@ class IrEmitter : public DfsHloVisitorWithDefault, } }; - tensorflow::gtl::FlatMap<const Literal*, llvm::Constant*, - LiteralPtrHashFunctor, LiteralPtrEqualityFunctor> + absl::flat_hash_map<const Literal*, llvm::Constant*, LiteralPtrHashFunctor, + LiteralPtrEqualityFunctor> emitted_literals_; - tensorflow::gtl::FlatMap<BufferAllocation::Index, llvm::Constant*> + absl::flat_hash_map<BufferAllocation::Index, llvm::Constant*> constant_buffer_to_global_; std::vector<const HloComputation*> thread_local_computations_; diff --git a/tensorflow/compiler/xla/service/cpu/target_machine_features.cc b/tensorflow/compiler/xla/service/cpu/target_machine_features.cc index a0cd8ee2d2..5cdac203af 100644 --- a/tensorflow/compiler/xla/service/cpu/target_machine_features.cc +++ b/tensorflow/compiler/xla/service/cpu/target_machine_features.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" +#include "tensorflow/core/platform/logging.h" namespace xla { namespace cpu { diff --git a/tensorflow/compiler/xla/service/cpu/target_machine_features.h b/tensorflow/compiler/xla/service/cpu/target_machine_features.h index 8b00ae9e47..a383b4a4a0 100644 --- a/tensorflow/compiler/xla/service/cpu/target_machine_features.h +++ b/tensorflow/compiler/xla/service/cpu/target_machine_features.h @@ -16,10 +16,10 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TARGET_MACHINE_FEATURES_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TARGET_MACHINE_FEATURES_H_ +#include "absl/container/flat_hash_map.h" #include "llvm/Analysis/TargetTransformInfo.h" #include "llvm/Target/TargetMachine.h" #include "tensorflow/compiler/xla/primitive_util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { namespace cpu { @@ -97,8 +97,7 @@ class LLVMTargetMachineFeatures : public TargetMachineFeatures { // This is mutated from within `GetTargetTransformInfoFor` which is // semantically a getter (and thus `const`); and is therefore declared // mutable. Making this mutable is okay because it has cache semantics. - mutable tensorflow::gtl::FlatMap<const llvm::Function*, - llvm::TargetTransformInfo> + mutable absl::flat_hash_map<const llvm::Function*, llvm::TargetTransformInfo> target_transform_info_cache_; llvm::TargetMachine* target_machine_; }; diff --git a/tensorflow/compiler/xla/service/defuser.cc b/tensorflow/compiler/xla/service/defuser.cc index d124f74d19..661539cccb 100644 --- a/tensorflow/compiler/xla/service/defuser.cc +++ b/tensorflow/compiler/xla/service/defuser.cc @@ -22,6 +22,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/call_graph.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -48,7 +49,7 @@ Status Defuse(HloInstruction* fusion_instruction) { fusion_instruction->fused_instructions_computation(); // A map from fused instruction to its defused clone. - tensorflow::gtl::FlatMap<const HloInstruction*, HloInstruction*> + absl::flat_hash_map<const HloInstruction*, HloInstruction*> defused_instructions; // Initialize map to contain the fusion instruction parameters mapping // to the operands of the fusion instruction. diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h index 5761573791..68d01d75a2 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h @@ -27,7 +27,6 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD index 51968d13d4..e65d3fa332 100644 --- a/tensorflow/compiler/xla/service/gpu/BUILD +++ b/tensorflow/compiler/xla/service/gpu/BUILD @@ -91,6 +91,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_reachability", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", ], ) @@ -357,6 +358,7 @@ cc_library( "//tensorflow/core/platform/default/build_config:cufft_plugin", "//tensorflow/core/platform/default/build_config:stream_executor_cuda", # build_cleaner: keep "//tensorflow/stream_executor", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc index 31a9f9b1be..5742632782 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc @@ -19,6 +19,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" @@ -197,7 +198,7 @@ GpuExecutable::ResolveConstantGlobals(se::StreamExecutor* executor) { } module_spec.AddCudaPtxInMemory(ptx().c_str()); - tensorflow::gtl::FlatMap<int64, se::DeviceMemoryBase> globals; + absl::flat_hash_map<int64, se::DeviceMemoryBase> globals; se::ModuleHandle module_handle; executor->LoadModule(module_spec, &module_handle); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.h b/tensorflow/compiler/xla/service/gpu/gpu_executable.h index 38b0f8f15b..0e276282e4 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.h @@ -19,6 +19,7 @@ limitations under the License. #include <memory> #include <string> +#include "absl/container/flat_hash_map.h" #include "absl/strings/string_view.h" #include "absl/types/optional.h" #include "absl/types/span.h" @@ -35,7 +36,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" @@ -101,7 +101,7 @@ class GpuExecutable : public Executable { const PointsToSet& GetRootPointsToSet() const; using BufferAllocToDeviceMemoryMap = - tensorflow::gtl::FlatMap<BufferAllocation::Index, se::DeviceMemoryBase>; + absl::flat_hash_map<BufferAllocation::Index, se::DeviceMemoryBase>; // Loads the PTX or CUBIN for this executable into `executor` and resolves the // globals corresponding to constant buffers. Returns a map mapping buffer diff --git a/tensorflow/compiler/xla/service/gpu/stream_assignment.h b/tensorflow/compiler/xla/service/gpu/stream_assignment.h index c2df83aaa4..52d38b6f20 100644 --- a/tensorflow/compiler/xla/service/gpu/stream_assignment.h +++ b/tensorflow/compiler/xla/service/gpu/stream_assignment.h @@ -16,9 +16,9 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_STREAM_ASSIGNMENT_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_STREAM_ASSIGNMENT_H_ +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { namespace gpu { @@ -34,7 +34,7 @@ class StreamAssignment { private: int stream_count_ = 1; // At least the main stream. - tensorflow::gtl::FlatMap<const HloInstruction*, int> hlo_to_stream_number_; + absl::flat_hash_map<const HloInstruction*, int> hlo_to_stream_number_; }; // Assigns GPU streams to instructions in `module`. diff --git a/tensorflow/compiler/xla/service/heap_simulator.cc b/tensorflow/compiler/xla/service/heap_simulator.cc index 2bd04259c0..147776c8c4 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.cc +++ b/tensorflow/compiler/xla/service/heap_simulator.cc @@ -18,13 +18,14 @@ limitations under the License. #include <algorithm> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/util.h" namespace xla { -using tensorflow::gtl::FlatMap; +using absl::flat_hash_map; using tensorflow::gtl::FlatSet; /*static*/ @@ -56,7 +57,7 @@ StatusOr<int64> HeapSimulator::MinimumMemoryForComputation( const HloComputation& computation, const HloInstructionSequence& sequence, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>* + const absl::flat_hash_map<const HloComputation*, int64>* memory_by_computation) { TF_ASSIGN_OR_RETURN( HeapSimulator::Result result, @@ -88,7 +89,7 @@ StatusOr<HeapSimulator::Result> HeapSimulator::Run( const HloInstructionSequence& instruction_sequence, const TuplePointsToAnalysis& points_to_analysis, const BufferValue::SizeFunction& size_fn, const Options& options, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>* + const absl::flat_hash_map<const HloComputation*, int64>* memory_by_computation) { HeapSimulator heap(std::move(algorithm), size_fn, options, /*schedule=*/nullptr, memory_by_computation); @@ -115,8 +116,10 @@ Status HeapSimulator::RunComputation( // 'used_buffers' is the reverse map - it tracks which buffers were used by an // instruction, so that we can remove the instructions from a buffer's live // set after they are visited. - FlatMap<const BufferValue*, FlatSet<const HloInstruction*>> live_buffers; - FlatMap<const HloInstruction*, FlatSet<const BufferValue*>> used_buffers; + flat_hash_map<const BufferValue*, FlatSet<const HloInstruction*>> + live_buffers; + flat_hash_map<const HloInstruction*, FlatSet<const BufferValue*>> + used_buffers; auto add_user_to_buffer = [this, &live_buffers, &used_buffers]( const HloInstruction* user, const BufferValue* buffer) { @@ -345,7 +348,7 @@ HeapSimulator::HeapSimulator( std::unique_ptr<HeapAlgorithm> algorithm, const BufferValue::SizeFunction& size_fn, const Options& options, const HloSchedule* schedule, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>* + const absl::flat_hash_map<const HloComputation*, int64>* memory_by_computation) : no_fragmentation_stats_(absl::make_unique<NoFragmentationStatsHeap>()), algorithm_(std::move(algorithm)), @@ -536,7 +539,7 @@ void NoFragmentationStatsHeap::Alloc(const BufferValue* buffer, int64 size, void NoFragmentationStatsHeap::AccountForSubcomputationMemory( const HloInstruction* instruction, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) { // We only count the memory usage of the largest subcomputation, instead of // adding them all, because subcomputations won't execute in parallel. diff --git a/tensorflow/compiler/xla/service/heap_simulator.h b/tensorflow/compiler/xla/service/heap_simulator.h index 7d6dcc0dc9..a5bb3f81f7 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.h +++ b/tensorflow/compiler/xla/service/heap_simulator.h @@ -21,6 +21,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/buffer_value.h" #include "tensorflow/compiler/xla/service/buffer_value_containers.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" @@ -30,7 +31,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_schedule.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" namespace xla { @@ -58,7 +58,7 @@ class HeapSimulator { // Result represents the result of the heap simulation. struct Result { // The assignment of buffers to chunks. - tensorflow::gtl::FlatMap<const BufferValue*, Chunk> chunk_map; + absl::flat_hash_map<const BufferValue*, Chunk> chunk_map; // The total size in bytes of the heap, containing all assigned chunks. int64 heap_size = 0; @@ -100,7 +100,7 @@ class HeapSimulator { const HloComputation& computation, const HloInstructionSequence& sequence, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>* + const absl::flat_hash_map<const HloComputation*, int64>* memory_by_computation = nullptr); // Run the heap simulation with the given algorithm, assuming the given @@ -130,7 +130,7 @@ class HeapSimulator { const TuplePointsToAnalysis& points_to_analysis, const BufferValue::SizeFunction& size_fn, const Options& options = Options(), - const tensorflow::gtl::FlatMap<const HloComputation*, int64>* + const absl::flat_hash_map<const HloComputation*, int64>* memory_by_computation = nullptr); private: @@ -140,7 +140,7 @@ class HeapSimulator { HeapSimulator(std::unique_ptr<HeapAlgorithm> algorithm, const BufferValue::SizeFunction& size_fn, const Options& options, const HloSchedule* schedule = nullptr, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>* + const absl::flat_hash_map<const HloComputation*, int64>* memory_by_computation = nullptr); ~HeapSimulator(); @@ -172,7 +172,7 @@ class HeapSimulator { // handle subcomputations. It would be good to unify the handling of // subcomputations, but it's not clear how. const HloSchedule* schedule_; - const tensorflow::gtl::FlatMap<const HloComputation*, int64>* + const absl::flat_hash_map<const HloComputation*, int64>* memory_by_computation_; // In addition to Alloc and Free, the heap simulator exposes a concept of @@ -193,7 +193,7 @@ class HeapSimulator { const BufferValue* canonical = nullptr; int64 refcount = 0; }; - tensorflow::gtl::FlatMap<const BufferValue*, std::shared_ptr<SharedGroup>> + absl::flat_hash_map<const BufferValue*, std::shared_ptr<SharedGroup>> shared_buffers_; // Hold some sets for error-checking the sequence of Alloc and Free calls. @@ -235,7 +235,7 @@ class HeapAlgorithm { // analysis, it's not worth making major changes to HeapSimulator now. virtual void AccountForSubcomputationMemory( const HloInstruction* instruction, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) {} // Free de-allocates a previously allocated buffer. @@ -262,7 +262,7 @@ class NoFragmentationStatsHeap : public HeapAlgorithm { void AccountForSubcomputationMemory( const HloInstruction* instruction, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) override; void Free(const BufferValue* buffer, int64 size) override; @@ -382,8 +382,7 @@ class GlobalDecreasingSizeBestFitHeap : public HeapAlgorithm { // Free time of the buffer. int64 end; }; - tensorflow::gtl::FlatMap<const BufferValue*, BufferInterval> - buffer_intervals_; + absl::flat_hash_map<const BufferValue*, BufferInterval> buffer_intervals_; }; // A heap algorithm that chooses the best results from other algorithms added to diff --git a/tensorflow/compiler/xla/service/heap_simulator_test.cc b/tensorflow/compiler/xla/service/heap_simulator_test.cc index 191fbf8194..ea0bced923 100644 --- a/tensorflow/compiler/xla/service/heap_simulator_test.cc +++ b/tensorflow/compiler/xla/service/heap_simulator_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/buffer_value.h" @@ -31,7 +32,6 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" #include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { namespace { @@ -174,7 +174,7 @@ class HeapSimulatorTracker { // Construct the module sequence grouped by computation. HloSchedule schedule(module_.get()); - tensorflow::gtl::FlatMap<const HloInstruction*, int> reverse_position; + absl::flat_hash_map<const HloInstruction*, int> reverse_position; for (int i = 0; i < full_module_sequence.size(); ++i) { const HloInstruction* instruction = full_module_sequence[i]; schedule.GetOrCreateSequence(instruction->parent()) diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc index 0986da65cb..b6e1f52cf5 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc @@ -20,6 +20,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/xla/map_util.h" @@ -290,13 +291,11 @@ class BufferValueMap { const HloDataflowAnalysis& dataflow_; // A map containing the set of values contained in each buffer. - tensorflow::gtl::FlatMap<BufferNumber, - tensorflow::gtl::FlatSet<const HloValue*>> + absl::flat_hash_map<BufferNumber, tensorflow::gtl::FlatSet<const HloValue*>> buffers_; // A map indicating which buffer each value is contained in. - tensorflow::gtl::FlatMap<const HloValue*, BufferNumber> - value_to_buffer_number_; + absl::flat_hash_map<const HloValue*, BufferNumber> value_to_buffer_number_; // The buffer number of the next buffer to be created. BufferNumber next_buffer_number_ = 0; diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.h b/tensorflow/compiler/xla/service/hlo_alias_analysis.h index e345804537..372f99ff01 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.h @@ -20,6 +20,7 @@ limitations under the License. #include <string> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/service/hlo_buffer.h" #include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h" @@ -110,7 +111,7 @@ class HloAliasAnalysis { std::unique_ptr<HloDataflowAnalysis> dataflow_analysis_; // A map indicating which buffer a value is contained in. - tensorflow::gtl::FlatMap<const HloValue*, HloBuffer*> value_to_buffer_; + absl::flat_hash_map<const HloValue*, HloBuffer*> value_to_buffer_; // A lazily constructed vector containing all HloBuffers sorted by // HloBuffer::Id. diff --git a/tensorflow/compiler/xla/service/hlo_clone_context.h b/tensorflow/compiler/xla/service/hlo_clone_context.h index 658643b427..24910ca07b 100644 --- a/tensorflow/compiler/xla/service/hlo_clone_context.h +++ b/tensorflow/compiler/xla/service/hlo_clone_context.h @@ -18,8 +18,8 @@ limitations under the License. #include <string> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/map_util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { @@ -73,12 +73,12 @@ class HloCloneContext { return FindOrDie(computations_, old_computation); } - const tensorflow::gtl::FlatMap<const HloInstruction*, HloInstruction*>& + const absl::flat_hash_map<const HloInstruction*, HloInstruction*>& cloned_instructions() const { return instructions_; } - const tensorflow::gtl::FlatMap<const HloComputation*, HloComputation*>& + const absl::flat_hash_map<const HloComputation*, HloComputation*>& cloned_computations() const { return computations_; } @@ -86,10 +86,8 @@ class HloCloneContext { private: HloModule* module_; string suffix_; - tensorflow::gtl::FlatMap<const HloInstruction*, HloInstruction*> - instructions_; - tensorflow::gtl::FlatMap<const HloComputation*, HloComputation*> - computations_; + absl::flat_hash_map<const HloInstruction*, HloInstruction*> instructions_; + absl::flat_hash_map<const HloComputation*, HloComputation*> computations_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index 4613d6762e..257dd5876f 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -24,6 +24,7 @@ limitations under the License. #include <sstream> #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" @@ -297,7 +298,7 @@ void ComputeComputationPostOrder( void HloComputation::ComputeInstructionPostOrder( const HloComputation::ChannelDependencyMap& channel_dependency_map, std::vector<HloInstruction*>* post_order, HloInstruction* root, - tensorflow::gtl::FlatMap<HloInstruction*, VisitState>* visited) const { + absl::flat_hash_map<HloInstruction*, VisitState>* visited) const { std::vector<HloInstruction*> dfs_stack; dfs_stack.push_back(root); while (!dfs_stack.empty()) { @@ -394,7 +395,7 @@ std::vector<HloInstruction*> HloComputation::MakeInstructionPostOrder() const { std::vector<HloInstruction*> post_order; post_order.reserve(instruction_count()); std::vector<HloInstruction*> trace_instructions; - tensorflow::gtl::FlatMap<HloInstruction*, VisitState> visited; + absl::flat_hash_map<HloInstruction*, VisitState> visited; for (auto& instruction : instructions_) { if (instruction->opcode() == HloOpcode::kTrace) { // Trace instructions aren't handled by the DFS visitor. Add trace @@ -505,9 +506,9 @@ HloComputationProto HloComputation::ToProto() const { /* static */ StatusOr<std::unique_ptr<HloComputation>> HloComputation::CreateFromProto( const HloComputationProto& proto, - const tensorflow::gtl::FlatMap<int64, HloComputation*>& computation_map) { - tensorflow::gtl::FlatMap<int64, HloInstruction*> instruction_map; - tensorflow::gtl::FlatMap<HloInstruction*, int64> to_proto_id; + const absl::flat_hash_map<int64, HloComputation*>& computation_map) { + absl::flat_hash_map<int64, HloInstruction*> instruction_map; + absl::flat_hash_map<HloInstruction*, int64> to_proto_id; std::vector<std::unique_ptr<HloInstruction>> instructions; int64 parameter_count = 0; for (const HloInstructionProto& instruction_proto : proto.instructions()) { diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h index 936a53bd7e..af929ac009 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.h +++ b/tensorflow/compiler/xla/service/hlo_computation.h @@ -25,6 +25,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/iterator_util.h" #include "tensorflow/compiler/xla/map_util.h" @@ -40,7 +41,6 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" @@ -188,7 +188,7 @@ class HloComputation { // calls. static StatusOr<std::unique_ptr<HloComputation>> CreateFromProto( const HloComputationProto& proto, - const tensorflow::gtl::FlatMap<int64, HloComputation*>& computation_map); + const absl::flat_hash_map<int64, HloComputation*>& computation_map); // Gets the instructions in this computation. // @@ -414,14 +414,14 @@ class HloComputation { // cross-replica-sum the union of the dependencies for all participating // instructions. using ChannelDependencyMap = - tensorflow::gtl::FlatMap<int64, absl::InlinedVector<HloInstruction*, 1>>; + absl::flat_hash_map<int64, absl::InlinedVector<HloInstruction*, 1>>; ChannelDependencyMap ComputeChannelDependencies() const; enum VisitState { kVisiting, kVisited }; void ComputeInstructionPostOrder( const HloComputation::ChannelDependencyMap& channel_dependency_map, std::vector<HloInstruction*>* post_order, HloInstruction* root, - tensorflow::gtl::FlatMap<HloInstruction*, VisitState>* visited) const; + absl::flat_hash_map<HloInstruction*, VisitState>* visited) const; string name_; int64 unique_id_; @@ -439,7 +439,7 @@ class HloComputation { // instruction pointer to location in the list for fast lookup. using InstructionList = std::list<std::unique_ptr<HloInstruction>>; InstructionList instructions_; - tensorflow::gtl::FlatMap<const HloInstruction*, InstructionList::iterator> + absl::flat_hash_map<const HloInstruction*, InstructionList::iterator> instruction_iterators_; std::vector<HloInstruction*> param_instructions_; diff --git a/tensorflow/compiler/xla/service/hlo_domain_map.cc b/tensorflow/compiler/xla/service/hlo_domain_map.cc index 113fd18eae..159c39d557 100644 --- a/tensorflow/compiler/xla/service/hlo_domain_map.cc +++ b/tensorflow/compiler/xla/service/hlo_domain_map.cc @@ -17,6 +17,7 @@ limitations under the License. #include <algorithm> +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" @@ -106,8 +107,8 @@ Status HloDomainMap::PopulateDomainMetadataMap() { auto equal = [](const DomainMetadata* a, const DomainMetadata* b) { return a->Matches(*b); }; - tensorflow::gtl::FlatMap<const DomainMetadata*, int64, decltype(hash), - decltype(equal)> + absl::flat_hash_map<const DomainMetadata*, int64, decltype(hash), + decltype(equal)> domain_metadata(1024, hash, equal); for (auto& domain : instruction_domains_) { diff --git a/tensorflow/compiler/xla/service/hlo_domain_map.h b/tensorflow/compiler/xla/service/hlo_domain_map.h index 56b557d7ce..8584bc021d 100644 --- a/tensorflow/compiler/xla/service/hlo_domain_map.h +++ b/tensorflow/compiler/xla/service/hlo_domain_map.h @@ -19,13 +19,13 @@ limitations under the License. #include <memory> #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_domain_metadata.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" namespace xla { @@ -77,8 +77,7 @@ class HloDomainMap { private: // Map used for representing instruction ordering, i.e. // order_map[a] < order_map[b] means a must be ordered before b. - using InstructionOrderMap = - tensorflow::gtl::FlatMap<const HloInstruction*, int64>; + using InstructionOrderMap = absl::flat_hash_map<const HloInstruction*, int64>; HloDomainMap(string domain_kind) : domain_kind_(std::move(domain_kind)) {} @@ -120,8 +119,8 @@ class HloDomainMap { string domain_kind_; std::vector<std::unique_ptr<DomainMetadata::Domain>> instruction_domains_; - tensorflow::gtl::FlatMap<HloInstruction*, int64> instruction_to_domain_; - tensorflow::gtl::FlatMap<HloInstruction*, int64> domain_metadata_id_; + absl::flat_hash_map<HloInstruction*, int64> instruction_to_domain_; + absl::flat_hash_map<HloInstruction*, int64> domain_metadata_id_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc index 23787dbc8a..5d5c9c7e58 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction.cc @@ -22,6 +22,7 @@ limitations under the License. #include <utility> #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/memory/memory.h" #include "absl/strings/ascii.h" @@ -43,7 +44,6 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/lib/gtl/map_util.h" #include "tensorflow/core/platform/human_readable_json.h" @@ -59,8 +59,8 @@ using absl::StrJoin; /* static */ StatusOr<std::unique_ptr<HloInstruction>> HloInstruction::CreateFromProto( const HloInstructionProto& proto, - const tensorflow::gtl::FlatMap<int64, HloInstruction*>& instruction_map, - const tensorflow::gtl::FlatMap<int64, HloComputation*>& computation_map) { + const absl::flat_hash_map<int64, HloInstruction*>& instruction_map, + const absl::flat_hash_map<int64, HloComputation*>& computation_map) { TF_RET_CHECK(!proto.opcode().empty()); TF_ASSIGN_OR_RETURN(HloOpcode opcode, StringToHloOpcode(proto.opcode())); TF_RET_CHECK(proto.has_shape()); @@ -266,7 +266,8 @@ StatusOr<std::unique_ptr<HloInstruction>> HloInstruction::CreateFromProto( << "Expect 1 called computation for fusion instruction but sees " << proto.called_computation_ids_size(); const int64 fusion_id = proto.called_computation_ids(0); - auto* fused_computation = FindPtrOrNull(computation_map, fusion_id); + auto* fused_computation = + tensorflow::gtl::FindPtrOrNull(computation_map, fusion_id); TF_RET_CHECK(fused_computation != nullptr) << "No fusion computation with id " << fusion_id; instruction = CreateFusion(proto.shape(), fusion_kind, all_operands(), @@ -2661,14 +2662,14 @@ class HloInstruction::FusionReusesParamElements { // the value of this parameter, which would save stack space but not allow us // to finish early if we find a reuse. static UseKind Compute(int64 i, const HloInstruction& hlo) { - tensorflow::gtl::FlatMap<const HloInstruction*, UseKind> memoization_cache; + absl::flat_hash_map<const HloInstruction*, UseKind> memoization_cache; return ComputeInternal(i, hlo, &memoization_cache); } private: static UseKind ComputeInternal( int64 i, const HloInstruction& hlo, - tensorflow::gtl::FlatMap<const HloInstruction*, UseKind>* cache) { + absl::flat_hash_map<const HloInstruction*, UseKind>* cache) { if (auto hlo_param = DynCast<HloParameterInstruction>(&hlo)) { if (hlo_param->parameter_number() == i) { return UseKind::kUse; diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 009bd3bab3..1bfdc88abc 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -32,6 +32,7 @@ limitations under the License. #include <unordered_set> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" @@ -50,7 +51,6 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/iterator_range.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" @@ -247,7 +247,7 @@ class CanonicalNameMap { private: int64 index; - tensorflow::gtl::FlatMap<string, string> canonical_name_map; + absl::flat_hash_map<string, string> canonical_name_map; }; // HLO instructions are the atomic unit of the high-level compiler's IR. @@ -350,8 +350,8 @@ class HloInstruction { // calls. static StatusOr<std::unique_ptr<HloInstruction>> CreateFromProto( const HloInstructionProto& proto, - const tensorflow::gtl::FlatMap<int64, HloInstruction*>& instruction_map, - const tensorflow::gtl::FlatMap<int64, HloComputation*>& computation_map); + const absl::flat_hash_map<int64, HloInstruction*>& instruction_map, + const absl::flat_hash_map<int64, HloComputation*>& computation_map); // Creates a parameter-retrieving instruction. static std::unique_ptr<HloInstruction> CreateParameter(int64 parameter_number, diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc index ad45a82941..1bc168c8b7 100644 --- a/tensorflow/compiler/xla/service/hlo_instructions.cc +++ b/tensorflow/compiler/xla/service/hlo_instructions.cc @@ -18,6 +18,7 @@ limitations under the License. #include <deque> #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "absl/strings/escaping.h" #include "absl/strings/str_cat.h" @@ -28,7 +29,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/window_util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { namespace { @@ -1099,7 +1099,7 @@ void HloFusionInstruction::MergeFusionInstructionIntoMultiOutput( // Note that we add the unfused instructions to this->parent_ computation. // This is necessary because the unique_id needs for an instruction and // it's only added when inserting to the computation. - tensorflow::gtl::FlatMap<HloInstruction*, HloInstruction*> old_to_new; + absl::flat_hash_map<HloInstruction*, HloInstruction*> old_to_new; std::vector<HloInstruction*> unfused_instructions; auto computation_to_merge = instruction_to_merge->fused_instructions_computation(); @@ -1392,7 +1392,7 @@ std::unique_ptr<HloInstruction> HloFusionInstruction::CloneWithNewOperandsImpl( } Status HloFusionInstruction::DeduplicateFusionOperands() { - tensorflow::gtl::FlatMap<const HloInstruction*, int> operand_indices; + absl::flat_hash_map<const HloInstruction*, int> operand_indices; std::vector<int> operands_to_remove; for (int i = 0; i < operand_count(); ++i) { auto emplace_result = operand_indices.emplace(operand(i), i); diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc b/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc index 6a4e766788..1c2b2868fd 100644 --- a/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc +++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc @@ -20,6 +20,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/heap_simulator.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" @@ -74,7 +75,7 @@ class ListScheduler { const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) { ListScheduler scheduler(computation, points_to_analysis, size_function, memory_by_computation); @@ -99,7 +100,7 @@ class ListScheduler { ListScheduler(const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) : computation_(computation), points_to_analysis_(points_to_analysis), @@ -234,8 +235,7 @@ class ListScheduler { // Populate the ready list with instructions which have no operands or // control predecessors. - tensorflow::gtl::FlatMap<const HloInstruction*, int64> - unscheduled_pred_count; + absl::flat_hash_map<const HloInstruction*, int64> unscheduled_pred_count; for (auto* instruction : computation_.instructions()) { // TODO(b/34466113): Replace this and above with successors() or // predecessors() when these methods are added to HloInstruction. @@ -251,8 +251,8 @@ class ListScheduler { std::multimap<Priority, ReadyListEntry> ready_queue; // Map of ready instructions to their iterators in ready_queue. - tensorflow::gtl::FlatMap<const HloInstruction*, - std::multimap<Priority, ReadyListEntry>::iterator> + absl::flat_hash_map<const HloInstruction*, + std::multimap<Priority, ReadyListEntry>::iterator> ready_instructions; auto add_to_ready_queue = [&](HloInstruction* inst) { @@ -347,12 +347,11 @@ class ListScheduler { // Computations are analyzed in post-order. When scheduling an instruction // that includes subcomputations, such as a while loop, we use this map to // look up the memory needed by subcomputations. - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation_; // A map containing the LogicalBuffers that each instruction uses. - tensorflow::gtl::FlatMap<const HloInstruction*, - std::vector<const LogicalBuffer*>> + absl::flat_hash_map<const HloInstruction*, std::vector<const LogicalBuffer*>> buffer_uses_; // A map containing the count of unscheduled HLOs which using a particular @@ -379,7 +378,7 @@ StatusOr<HloInstructionSequence> ScheduleComputationHelper( const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const MemorySchedulerAlgorithm& algorithm, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) { VLOG(2) << "Computation: " << computation.name(); if (algorithm) { @@ -396,13 +395,13 @@ StatusOr<HloInstructionSequence> DFSMemoryScheduler( const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) { // These variables are a hack to prevent overflows. int64 cumulative_total_size = 0; int64 total_hlos = computation.parent()->instruction_count(); - tensorflow::gtl::FlatMap<const HloInstruction*, int64> extra_users; - tensorflow::gtl::FlatMap<const HloInstruction*, int64> total_sizes; + absl::flat_hash_map<const HloInstruction*, int64> extra_users; + absl::flat_hash_map<const HloInstruction*, int64> total_sizes; for (const HloInstruction* hlo : computation.MakeInstructionPostOrder()) { if (ListScheduler::IgnoreInstruction(*hlo)) { extra_users[hlo] = 0; @@ -467,7 +466,7 @@ StatusOr<HloInstructionSequence> ListMemoryScheduler( const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) { return ListScheduler::Run(computation, points_to_analysis, size_function, memory_by_computation); @@ -477,7 +476,7 @@ StatusOr<HloInstructionSequence> PostOrderMemoryScheduler( const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) { return HloInstructionSequence(computation.MakeInstructionPostOrder()); } @@ -486,7 +485,7 @@ StatusOr<HloInstructionSequence> DefaultMemoryScheduler( const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation) { // We try a few schedulers and choose whichever returns a lower min-memory, // not accounting for fragmentation. @@ -549,7 +548,7 @@ StatusOr<HloSchedule> ScheduleModule( HloSchedule schedule(&module); TF_ASSIGN_OR_RETURN(std::unique_ptr<TuplePointsToAnalysis> points_to_analysis, TuplePointsToAnalysis::Run(&module)); - tensorflow::gtl::FlatMap<const HloComputation*, int64> memory_by_computation; + absl::flat_hash_map<const HloComputation*, int64> memory_by_computation; for (const auto* computation : module.MakeComputationPostOrder()) { if (!computation->IsFusionComputation()) { TF_ASSIGN_OR_RETURN(HloInstructionSequence computation_sequence, @@ -577,7 +576,7 @@ StatusOr<HloInstructionSequence> ScheduleComputation( CHECK(!computation.IsFusionComputation()); TF_ASSIGN_OR_RETURN(std::unique_ptr<TuplePointsToAnalysis> points_to_analysis, TuplePointsToAnalysis::Run(computation.parent())); - tensorflow::gtl::FlatMap<const HloComputation*, int64> empty_map; + absl::flat_hash_map<const HloComputation*, int64> empty_map; return ScheduleComputationHelper(computation, *points_to_analysis, size_function, nullptr, empty_map); } diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler.h b/tensorflow/compiler/xla/service/hlo_memory_scheduler.h index 9964c6fdd7..a4c1d3db81 100644 --- a/tensorflow/compiler/xla/service/hlo_memory_scheduler.h +++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler.h @@ -18,6 +18,7 @@ limitations under the License. #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_ordering.h" @@ -37,7 +38,7 @@ namespace xla { typedef std::function<StatusOr<HloInstructionSequence>( const HloComputation&, const TuplePointsToAnalysis&, const LogicalBuffer::SizeFunction&, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>&)> + const absl::flat_hash_map<const HloComputation*, int64>&)> MemorySchedulerAlgorithm; // List scheduler @@ -45,7 +46,7 @@ StatusOr<HloInstructionSequence> ListMemoryScheduler( const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation); // DFS-order scheduler @@ -53,7 +54,7 @@ StatusOr<HloInstructionSequence> DFSMemoryScheduler( const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation); // Naive Post Order scheduler @@ -61,7 +62,7 @@ StatusOr<HloInstructionSequence> PostOrderMemoryScheduler( const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation); // The default scheduling algorithm. Runs both the list scheduler @@ -71,7 +72,7 @@ StatusOr<HloInstructionSequence> DefaultMemoryScheduler( const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, - const tensorflow::gtl::FlatMap<const HloComputation*, int64>& + const absl::flat_hash_map<const HloComputation*, int64>& memory_by_computation); // Returns an HloSchedule which seeks to minimize the memory required for diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc b/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc index 1b9e9bfc77..5a9fccc7dd 100644 --- a/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc +++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include <string> #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/heap_simulator.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_dce.h" @@ -247,7 +248,7 @@ TEST_F(HloSchedulingTest, ListAccountsForSubcomputations) { EXPECT_TRUE(ordering.ExecutesBefore(bcast, add)); EXPECT_TRUE(ordering.ExecutesBefore(transpose, add)); - tensorflow::gtl::FlatMap<const HloComputation*, int64> memory_by_computation; + absl::flat_hash_map<const HloComputation*, int64> memory_by_computation; memory_by_computation[cond_computation] = 17; memory_by_computation[body_computation] = 16; std::unique_ptr<TuplePointsToAnalysis> points_to_analysis = @@ -409,7 +410,7 @@ TEST_F(HloSchedulingTest, HeapSimulatorAccountsForSubcomputations) { EXPECT_EQ(module->entry_computation()->instruction_count(), schedule.sequence(module->entry_computation()).size()); - tensorflow::gtl::FlatMap<const HloComputation*, int64> memory_by_computation; + absl::flat_hash_map<const HloComputation*, int64> memory_by_computation; memory_by_computation[cond_computation] = 17; memory_by_computation[body_computation] = 16; std::unique_ptr<TuplePointsToAnalysis> points_to_analysis = diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc index b3949f3a6d..9359e9a8be 100644 --- a/tensorflow/compiler/xla/service/hlo_module.cc +++ b/tensorflow/compiler/xla/service/hlo_module.cc @@ -23,6 +23,7 @@ limitations under the License. #include <utility> #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/map_util.h" @@ -285,8 +286,8 @@ StatusOr<std::unique_ptr<HloModule>> HloModule::CreateFromProto( << ShapeUtil::HumanStringWithLayout(expected_program_shape.result()) << ", actual: " << ShapeUtil::HumanStringWithLayout(result_shape); - tensorflow::gtl::FlatMap<int64, HloComputation*> computation_map; - tensorflow::gtl::FlatMap<HloComputation*, int64> to_proto_id; + absl::flat_hash_map<int64, HloComputation*> computation_map; + absl::flat_hash_map<HloComputation*, int64> to_proto_id; std::vector<std::unique_ptr<HloComputation>> computations; HloComputation* entry = nullptr; for (const HloComputationProto& computation_proto : proto.computations()) { diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h index 278d94cdd3..0311b73207 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h @@ -22,6 +22,7 @@ limitations under the License. #include <unordered_set> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/types/optional.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -30,7 +31,6 @@ limitations under the License. #include "tensorflow/compiler/xla/status.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -250,25 +250,25 @@ class HloModuleGroupMetadata { std::vector<std::unique_ptr<std::vector<HloInstruction*>>> companion_sets_; // Map from each companion while instruction to the index into companion_set_. - tensorflow::gtl::FlatMap<const HloInstruction*, int64> companion_set_index_; + absl::flat_hash_map<const HloInstruction*, int64> companion_set_index_; // Map from computation to the instruction using it (a kWhile, kConditional). - tensorflow::gtl::FlatMap<const HloComputation*, TrackedInstruction> + absl::flat_hash_map<const HloComputation*, TrackedInstruction> tracked_instructions_; // Maps tracked instructions (kWhile, kConditional, kCall, ...) to the set of // communicating instructions within the proper called computation(s). - tensorflow::gtl::FlatMap<HloInstruction*, std::vector<HloInstruction*>> + absl::flat_hash_map<HloInstruction*, std::vector<HloInstruction*>> tracked_instructions_comms_; // All channels in the module. std::vector<Channel> channels_; // Map from channel ids to the index in channels_. - tensorflow::gtl::FlatMap<int64, int64> channel_id_map_; + absl::flat_hash_map<int64, int64> channel_id_map_; // Map from all-reduce ids to the all reduce instructions. - tensorflow::gtl::FlatMap<int64, std::vector<HloInstruction*>> all_reduce_map_; + absl::flat_hash_map<int64, std::vector<HloInstruction*>> all_reduce_map_; // The maximum channel id used in the module group. int64 max_channel_id_ = -1; @@ -276,7 +276,7 @@ class HloModuleGroupMetadata { // The modules that this metadata was built from. const std::vector<HloModule*>& modules_; - tensorflow::gtl::FlatMap<HloModule*, std::unique_ptr<TuplePointsToAnalysis>> + absl::flat_hash_map<HloModule*, std::unique_ptr<TuplePointsToAnalysis>> points_to_analyses_; }; diff --git a/tensorflow/compiler/xla/service/hlo_module_group_util.h b/tensorflow/compiler/xla/service/hlo_module_group_util.h index 309c23045d..f21b44bcd9 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_util.h +++ b/tensorflow/compiler/xla/service/hlo_module_group_util.h @@ -20,6 +20,7 @@ limitations under the License. #include <memory> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -28,7 +29,6 @@ limitations under the License. #include "tensorflow/compiler/xla/status.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { @@ -87,7 +87,7 @@ class HloModuleGroupUtil { // * visit_state: map from each instruction to its visit state. // * visit_function: function called when each instruction group. // * root: the root instruction of the traversal. - using VisitStates = tensorflow::gtl::FlatMap<HloInstruction*, VisitState>; + using VisitStates = absl::flat_hash_map<HloInstruction*, VisitState>; Status VisitTopologicalOrder(VisitStates* visit_state, const VisitFunction& visit_function, HloInstruction* root); diff --git a/tensorflow/compiler/xla/service/hlo_opcode.cc b/tensorflow/compiler/xla/service/hlo_opcode.cc index 2d4e38589f..4551a1c2e2 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.cc +++ b/tensorflow/compiler/xla/service/hlo_opcode.cc @@ -14,9 +14,9 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { @@ -31,7 +31,7 @@ string HloOpcodeString(HloOpcode opcode) { } StatusOr<HloOpcode> StringToHloOpcode(const string& opcode_name) { - static auto* opcode_map = new tensorflow::gtl::FlatMap<string, HloOpcode>({ + static auto* opcode_map = new absl::flat_hash_map<string, HloOpcode>({ #define STRING_TO_OPCODE_ENTRY(enum_name, opcode_name, ...) \ {opcode_name, HloOpcode::enum_name}, HLO_OPCODE_LIST(STRING_TO_OPCODE_ENTRY) diff --git a/tensorflow/compiler/xla/service/hlo_ordering.h b/tensorflow/compiler/xla/service/hlo_ordering.h index b0361c3f02..66313492eb 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.h +++ b/tensorflow/compiler/xla/service/hlo_ordering.h @@ -20,6 +20,7 @@ limitations under the License. #include <string> #include <utility> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/call_graph.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h" @@ -28,7 +29,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_schedule.h" #include "tensorflow/compiler/xla/service/hlo_value.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { @@ -120,8 +120,8 @@ class PredecessorHloOrdering : public HloOrdering { // predecessors. An instruction is an element of its own predecessor set. // // Subclasses should fill this in to define the desired ordering. - tensorflow::gtl::FlatMap<const HloComputation*, - std::unique_ptr<HloReachabilityMap>> + absl::flat_hash_map<const HloComputation*, + std::unique_ptr<HloReachabilityMap>> predecessors_; }; @@ -204,7 +204,7 @@ class SequentialHloOrdering : public HloOrdering { // this map so more than one instruction may have the same position // value. This is not a problem because ExecutesBefore also verifies // instructions are in the same computation. - tensorflow::gtl::FlatMap<const HloInstruction*, int> order_position_; + absl::flat_hash_map<const HloInstruction*, int> order_position_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc index 8c2f928ca1..59fd01cb58 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc +++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc @@ -17,6 +17,7 @@ limitations under the License. #include <functional> +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" @@ -98,7 +99,7 @@ void HloPassPipeline::MaybeDumpHlo(const HloModule& module, if (!proto_dump_path.empty()) { static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED); static auto* const module_id_to_pass_number = - new tensorflow::gtl::FlatMap<int64, int64>(); + new absl::flat_hash_map<int64, int64>(); tensorflow::mutex_lock lock(mu); const int64 pass_number = (*module_id_to_pass_number)[module.unique_id()]++; diff --git a/tensorflow/compiler/xla/service/hlo_reachability.h b/tensorflow/compiler/xla/service/hlo_reachability.h index b66a2aa4bd..5a5f01f8fd 100644 --- a/tensorflow/compiler/xla/service/hlo_reachability.h +++ b/tensorflow/compiler/xla/service/hlo_reachability.h @@ -19,11 +19,11 @@ limitations under the License. #include <list> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -154,7 +154,7 @@ class HloReachabilityMap { // Dense assignment from HloInstruction* to number. These numbers index // into the bit_vectors_ vector and into the bits within a BitVector. - tensorflow::gtl::FlatMap<const HloInstruction*, int> indices_; + absl::flat_hash_map<const HloInstruction*, int> indices_; // Bitvectors holding the reachability to each instruction. The bit vector for // instruction X includes ones for each instruction which X is reachable from. diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc index a438671936..abdd9a9212 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc @@ -20,6 +20,7 @@ limitations under the License. #include <set> #include <string> +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" @@ -75,7 +76,7 @@ bool IsRematerializable(const HloInstruction* instruction) { // cache before, and eventually calling the IsRematerializable() API. bool CanBeRematerialized( const HloInstruction* instruction, - tensorflow::gtl::FlatMap<const HloInstruction*, bool>* remat_able) { + absl::flat_hash_map<const HloInstruction*, bool>* remat_able) { auto it = remat_able->find(instruction); if (it != remat_able->end()) { return it->second; @@ -268,7 +269,7 @@ class InstructionList { Item* first_; // Item for each instruction. - tensorflow::gtl::FlatMap<const HloInstruction*, Item*> item_map_; + absl::flat_hash_map<const HloInstruction*, Item*> item_map_; }; // Return the items which use the given LogicalBuffer. Sets @@ -503,7 +504,7 @@ MemoryUsageTracker::MemoryUsageTracker( PointsToSet::BufferSet live_out_set = points_to_analysis.GetPointsToSet(computation_->root_instruction()) .CreateFlattenedSet(); - tensorflow::gtl::FlatMap<const LogicalBuffer*, BufferId> + absl::flat_hash_map<const LogicalBuffer*, BufferId> logical_buffer_to_buffer_id; for (auto* item = instruction_list_.first(); item != nullptr; @@ -854,7 +855,7 @@ int64 RematerializationCost(const HloInstruction* instruction, Item* PickRematerializationCandidate( const MemoryUsageTracker& memory_tracker, const InstructionList& instruction_list, int64 memory_limit_bytes, - tensorflow::gtl::FlatMap<const HloInstruction*, bool>* remat_able) { + absl::flat_hash_map<const HloInstruction*, bool>* remat_able) { Item* best_item = nullptr; int64 best_cost = 0; @@ -983,7 +984,7 @@ StatusOr<bool> HloRematerialization::RematerializeComputation( tensorflow::gtl::FlatSet<const HloInstruction*> remat_move_instructions; // The map from instructions to their rematerializable status. - tensorflow::gtl::FlatMap<const HloInstruction*, bool> remat_able; + absl::flat_hash_map<const HloInstruction*, bool> remat_able; // The peak memory of the computation at any point in the instruction // sequence. diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.h b/tensorflow/compiler/xla/service/hlo_rematerialization.h index 7330d73c09..5a02e3a8bb 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.h +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.h @@ -15,6 +15,7 @@ #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_REMATERIALIZATION_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_REMATERIALIZATION_H_ +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/call_graph.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -115,8 +116,7 @@ class HloRematerialization : public HloModulePass { // computations called from sequential context // (CallContext::kSequential). These values are updated as rematerialization // occurs. - tensorflow::gtl::FlatMap<const HloComputation*, int64> - computation_peak_memory_; + absl::flat_hash_map<const HloComputation*, int64> computation_peak_memory_; std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_; diff --git a/tensorflow/compiler/xla/service/hlo_schedule.cc b/tensorflow/compiler/xla/service/hlo_schedule.cc index 3fc5dbeb02..7c5c98f04e 100644 --- a/tensorflow/compiler/xla/service/hlo_schedule.cc +++ b/tensorflow/compiler/xla/service/hlo_schedule.cc @@ -18,6 +18,7 @@ limitations under the License. #include <queue> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/xla/map_util.h" @@ -30,7 +31,7 @@ namespace xla { /* static */ StatusOr<HloSchedule> HloSchedule::CreateFromProto( const HloModule* module, const HloScheduleProto& proto) { - tensorflow::gtl::FlatMap<int64, const HloComputation*> id_to_computation; + absl::flat_hash_map<int64, const HloComputation*> id_to_computation; for (const HloComputation* computation : module->computations()) { id_to_computation[computation->unique_id()] = computation; } @@ -44,7 +45,7 @@ namespace xla { << "No computation exists in HLO module with id " << computation_id; const HloComputation* computation = comp_it->second; - tensorflow::gtl::FlatMap<int64, const HloInstruction*> id_to_instruction; + absl::flat_hash_map<int64, const HloInstruction*> id_to_instruction; for (const HloInstruction* instruction : computation->instructions()) { id_to_instruction[instruction->unique_id()] = instruction; } @@ -112,7 +113,7 @@ Status HloSchedule::UpdateComputationSchedule( const HloComputation* computation) { // Map from unique ID to HloInstruction pointer for instructions in the // computation. - tensorflow::gtl::FlatMap<int, const HloInstruction*> id_to_instruction; + absl::flat_hash_map<int, const HloInstruction*> id_to_instruction; for (const HloInstruction* instruction : computation->instructions()) { InsertOrDie(&id_to_instruction, instruction->unique_id(), instruction); } @@ -126,15 +127,13 @@ Status HloSchedule::UpdateComputationSchedule( // Map from HloInstruction X to newly added instructions (instruction is in // computation, but not in schedule) which use X. If an instruction is not in // the map, then it has no users which are newly added instructions. - tensorflow::gtl::FlatMap<const HloInstruction*, - std::vector<const HloInstruction*>> + absl::flat_hash_map<const HloInstruction*, std::vector<const HloInstruction*>> new_instruction_uses; // For each newly added instruction, this is the count of the instruction's // operands that have not yet been scheduled. When this value reaches zero, // then the instruction may be placed in the schedule. - tensorflow::gtl::FlatMap<const HloInstruction*, int> - unscheduled_operand_count; + absl::flat_hash_map<const HloInstruction*, int> unscheduled_operand_count; // Create a worklist of newly added instructions which are ready to be added // to the schedule. Initialize worklist with those that have zero operands. @@ -217,9 +216,9 @@ Status HloSchedule::Update() { } for (auto it = sequences_.begin(); it != sequences_.end();) { if (nonfusion_computations_ids.count(it->first) == 0) { - it = sequences_.erase(it); + sequences_.erase(it++); } else { - it++; + ++it; } } } @@ -254,7 +253,7 @@ Status HloSchedule::Verify() const { // For each computation verify the set of instructions is the same and that // each dependency and control edge is honored. for (const HloComputation* computation : nonfusion_computations) { - tensorflow::gtl::FlatMap<const HloInstruction*, int> instruction_position; + absl::flat_hash_map<const HloInstruction*, int> instruction_position; int pos = 0; for (const HloInstruction* instruction : sequence(computation).instructions()) { diff --git a/tensorflow/compiler/xla/service/hlo_schedule.h b/tensorflow/compiler/xla/service/hlo_schedule.h index 270fe6039f..0a714101ee 100644 --- a/tensorflow/compiler/xla/service/hlo_schedule.h +++ b/tensorflow/compiler/xla/service/hlo_schedule.h @@ -18,6 +18,7 @@ limitations under the License. #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -103,8 +104,7 @@ class HloSchedule { // Returns a map from HloComputation unique ID to instruction sequence. The // map contains all sequences in the schedule. - const tensorflow::gtl::FlatMap<int64, HloInstructionSequence>& sequences() - const { + const absl::flat_hash_map<int64, HloInstructionSequence>& sequences() const { return sequences_; } @@ -148,7 +148,7 @@ class HloSchedule { // A map from computation unique ID to instruction sequence. Unique IDs are // used rather than HloComputation pointers because HLO pointers are not // unique across HLO transformations because pointers may be recycled. - tensorflow::gtl::FlatMap<int64, HloInstructionSequence> sequences_; + absl::flat_hash_map<int64, HloInstructionSequence> sequences_; }; std::ostream& operator<<(std::ostream& out, const HloSchedule& schedule); diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index 6eb6658904..a7727824fe 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -15,6 +15,7 @@ limitations under the License. #include <set> +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" #include "tensorflow/compiler/xla/service/hlo_instructions.h" @@ -23,7 +24,6 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { @@ -993,7 +993,7 @@ Status CheckSameIsHostTransfer(const HloInstruction* instr1, // Checks various invariants of send and recv instructions. Status VerifySendsAndRecvs(const HloModule& module) { - tensorflow::gtl::FlatMap<int64, const HloInstruction*> host_channels; + absl::flat_hash_map<int64, const HloInstruction*> host_channels; // Host send/recv instructions must have their own unique channel. auto check_unique_host_channel = [&](const HloInstruction* instruction) { const HloSendRecvInstruction* sendrecv = @@ -1061,7 +1061,7 @@ StatusOr<bool> HloVerifier::Run(HloModule* module) { TF_RETURN_IF_ERROR(VerifyHloStructure(module)); TF_RETURN_IF_ERROR(VerifySendsAndRecvs(*module)); - tensorflow::gtl::FlatMap<string, const HloInstruction*> instructions; + absl::flat_hash_map<string, const HloInstruction*> instructions; for (auto* computation : module->computations()) { for (const auto& instruction : computation->instructions()) { diff --git a/tensorflow/compiler/xla/service/indexed_array_analysis.cc b/tensorflow/compiler/xla/service/indexed_array_analysis.cc index 06f0e1ed25..7ee789276d 100644 --- a/tensorflow/compiler/xla/service/indexed_array_analysis.cc +++ b/tensorflow/compiler/xla/service/indexed_array_analysis.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/indexed_array_analysis.h" #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" @@ -95,7 +96,7 @@ Status IndexedArrayAnalysis::TraverseAndPopulateCache( absl::InlinedVector<const HloInstruction*, 4> stack; enum DfsState { kDiscovered, kVisited }; - gtl::FlatMap<const HloInstruction*, DfsState> dfs_state_map; + absl::flat_hash_map<const HloInstruction*, DfsState> dfs_state_map; stack.push_back(root); InsertOrDie(&dfs_state_map, root, kDiscovered); diff --git a/tensorflow/compiler/xla/service/indexed_array_analysis.h b/tensorflow/compiler/xla/service/indexed_array_analysis.h index 3e238f97a0..e5aa67fd85 100644 --- a/tensorflow/compiler/xla/service/indexed_array_analysis.h +++ b/tensorflow/compiler/xla/service/indexed_array_analysis.h @@ -18,10 +18,10 @@ limitations under the License. #include <type_traits> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/util/ptr_util.h" namespace xla { @@ -360,7 +360,7 @@ class IndexedArrayAnalysis { std::vector<std::unique_ptr<Array>> owned_tensors_; std::vector<Literal> owned_literals_; - tensorflow::gtl::FlatMap<const HloInstruction*, Array*> cache_; + absl::flat_hash_map<const HloInstruction*, Array*> cache_; }; // A pass that prints all non-trivial results returned by IndexedArrayAnalysis. diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc index e884122fcb..5a99c40df4 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion.cc @@ -22,11 +22,11 @@ limitations under the License. #include <vector> #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/logging.h" namespace xla { @@ -189,7 +189,7 @@ bool InstructionFusion::EffectivelyAtMostUnary(HloInstruction* hlo) { bool InstructionFusion::CanFuseOnAllPaths( HloInstruction* producer, HloInstruction* consumer, const HloInstructionSet& do_not_fuse, - tensorflow::gtl::FlatMap<std::pair<HloInstruction*, HloInstruction*>, bool>* + absl::flat_hash_map<std::pair<HloInstruction*, HloInstruction*>, bool>* result_cache) { if (consumer == producer) { return true; @@ -241,7 +241,7 @@ InstructionFusion::ComputeGloballyUnfusible( // fusing operations that require duplication later depending on // is_expensive_(). HloInstructionSet do_not_duplicate; - tensorflow::gtl::FlatMap<std::pair<HloInstruction*, HloInstruction*>, bool> + absl::flat_hash_map<std::pair<HloInstruction*, HloInstruction*>, bool> can_fuse_on_all_paths_result_cache; for (HloInstruction* consumer : post_order) { for (HloInstruction* producer : consumer->operands()) { @@ -430,7 +430,7 @@ class ReversePostOrderFusionQueue : public FusionQueue { private: std::vector<HloInstruction*> post_order_; - tensorflow::gtl::FlatMap<HloInstruction*, int> post_order_index_; + absl::flat_hash_map<HloInstruction*, int> post_order_index_; }; } // namespace diff --git a/tensorflow/compiler/xla/service/instruction_fusion.h b/tensorflow/compiler/xla/service/instruction_fusion.h index c1ec3b18a1..da2032f6c7 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.h +++ b/tensorflow/compiler/xla/service/instruction_fusion.h @@ -1,3 +1,4 @@ +#include "absl/container/flat_hash_map.h" /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); @@ -158,8 +159,8 @@ class InstructionFusion : public HloModulePass { bool CanFuseOnAllPaths( HloInstruction* producer, HloInstruction* consumer, const HloInstructionSet& do_not_fuse, - tensorflow::gtl::FlatMap<std::pair<HloInstruction*, HloInstruction*>, - bool>* result_cache); + absl::flat_hash_map<std::pair<HloInstruction*, HloInstruction*>, bool>* + result_cache); // Computes the set of nodes that we do not want to fuse into any of their // consumers based on a global analysis of the HLO graph. diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h index e29c199c42..1591256fad 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.h +++ b/tensorflow/compiler/xla/service/layout_assignment.h @@ -25,6 +25,7 @@ limitations under the License. #include <utility> #include <vector> +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/service/computation_layout.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -38,7 +39,6 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/types.h" @@ -228,8 +228,8 @@ class LayoutConstraints { // Array-shaped buffers which have not yet been constrained. std::set<LogicalBuffer::Id> unconstrained_buffer_ids_; - mutable tensorflow::gtl::FlatMap<const HloInstruction*, - std::unique_ptr<PointsToSet::BufferSet>> + mutable absl::flat_hash_map<const HloInstruction*, + std::unique_ptr<PointsToSet::BufferSet>> buffer_sets_cache_; HloComputation* computation_; diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD index 540bbb7c7a..3934d2e493 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/BUILD +++ b/tensorflow/compiler/xla/service/llvm_ir/BUILD @@ -38,6 +38,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:logical_buffer", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@llvm//:core", ], diff --git a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h index 8d9fa99d82..88cde2d3d9 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h +++ b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h @@ -16,13 +16,13 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_ALIAS_ANALYSIS_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_ALIAS_ANALYSIS_H_ +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_cat.h" #include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" namespace xla { @@ -77,14 +77,14 @@ class AliasAnalysis { // A map from a buffer slice to metadata corresponding to its alias.scope // metadata. The index kParameterAliasSet is used to hold aliasing // information for parameters. - tensorflow::gtl::FlatMap<BufferAllocation::Slice, llvm::MDNode*, - BufferAllocation::Slice::Hasher> + absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*, + BufferAllocation::Slice::Hasher> alias_scope_metadata_; // A map from a buffer slice to metadata corresponding to its noalias // metadata. - tensorflow::gtl::FlatMap<BufferAllocation::Slice, llvm::MDNode*, - BufferAllocation::Slice::Hasher> + absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*, + BufferAllocation::Slice::Hasher> noalias_metadata_; }; diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.cc b/tensorflow/compiler/xla/service/multi_output_fusion.cc index b9ec31c497..95b1c20663 100644 --- a/tensorflow/compiler/xla/service/multi_output_fusion.cc +++ b/tensorflow/compiler/xla/service/multi_output_fusion.cc @@ -18,7 +18,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/types.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.h b/tensorflow/compiler/xla/service/multi_output_fusion.h index 0344626b26..9508ab2ed1 100644 --- a/tensorflow/compiler/xla/service/multi_output_fusion.h +++ b/tensorflow/compiler/xla/service/multi_output_fusion.h @@ -19,6 +19,7 @@ limitations under the License. #include <queue> #include <vector> +#include "absl/container/flat_hash_map.h" #include "absl/strings/string_view.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" @@ -126,7 +127,7 @@ class MultiOutputFusion : public HloModulePass { std::vector<FusionCandidate> candidates_; // A map that maps an instruction to the index_. - tensorflow::gtl::FlatMap<HloInstruction*, int> candidates_index_; + absl::flat_hash_map<HloInstruction*, int> candidates_index_; // The reachability map of current computation. std::unique_ptr<HloReachabilityMap> reachability_; diff --git a/tensorflow/compiler/xla/service/name_uniquer.h b/tensorflow/compiler/xla/service/name_uniquer.h index 6dd89c240f..1ac60f1cf4 100644 --- a/tensorflow/compiler/xla/service/name_uniquer.h +++ b/tensorflow/compiler/xla/service/name_uniquer.h @@ -18,9 +18,9 @@ limitations under the License. #include <string> +#include "absl/container/flat_hash_map.h" #include "absl/strings/string_view.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/macros.h" @@ -78,7 +78,7 @@ class NameUniquer { // Map from name prefix to the generator data structure which tracks used // identifiers and generates new ones. - tensorflow::gtl::FlatMap<string, SequentialIdGenerator> generated_names_; + absl::flat_hash_map<string, SequentialIdGenerator> generated_names_; TF_DISALLOW_COPY_AND_ASSIGN(NameUniquer); }; diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion.h b/tensorflow/compiler/xla/service/reduce_precision_insertion.h index 4bb22428f3..0b4e82e8d6 100644 --- a/tensorflow/compiler/xla/service/reduce_precision_insertion.h +++ b/tensorflow/compiler/xla/service/reduce_precision_insertion.h @@ -22,7 +22,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" #include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h index a9e8a51e09..78392d3bb2 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h @@ -36,7 +36,6 @@ limitations under the License. #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/compactptrset.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc b/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc index 56145822be..067cfcc17d 100644 --- a/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc +++ b/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc @@ -18,7 +18,6 @@ limitations under the License. #include "absl/container/inlined_vector.h" #include "tensorflow/compiler/xla/service/while_util.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc index e8fe33e626..2590473c77 100644 --- a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc @@ -15,17 +15,17 @@ limitations under the License. #include "tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h" #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "tensorflow/compiler/xla/service/tuple_util.h" #include "tensorflow/compiler/xla/service/while_util.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" namespace xla { +using absl::flat_hash_map; using absl::InlinedVector; -using tensorflow::gtl::FlatMap; using tensorflow::gtl::FlatSet; // Copies `to_hoist` to the computation containing `while_instr`, hoisting its @@ -34,7 +34,7 @@ using tensorflow::gtl::FlatSet; // function hoists the operands in `unhoisted_invariant_instructions` and moves // them into `hoisted_instructions`. static void CreateLoopInvariantCopy( - FlatMap<HloInstruction*, HloInstruction*>* hoisted_instructions, + flat_hash_map<HloInstruction*, HloInstruction*>* hoisted_instructions, FlatSet<HloInstruction*>* unhoisted_invariant_instructions, HloInstruction* while_instr, HloInstruction* to_hoist) { HloComputation* parent_of_while = while_instr->parent(); @@ -147,7 +147,7 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( // Maps instructions in the while body to instructions hoisted outside the // while that compute the same value. - FlatMap<HloInstruction*, HloInstruction*> hoisted_instructions; + flat_hash_map<HloInstruction*, HloInstruction*> hoisted_instructions; // Contains instructions that can be legally hoisted, but were deemed to be // unprofitable to be hoisted alone by NotWorthHoistingIndividually. When we diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.cc b/tensorflow/compiler/xla/service/while_loop_simplifier.cc index 9a74f22395..07de8492ba 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier.cc @@ -14,12 +14,12 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/while_loop_simplifier.h" +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/types/optional.h" #include "tensorflow/compiler/xla/service/call_inliner.h" #include "tensorflow/compiler/xla/service/while_loop_analysis.h" -#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { @@ -181,7 +181,7 @@ static StatusOr<bool> TryRemoveDeadWhileParams(HloInstruction* while_op) { used_tuple_indices.end()); std::sort(new_to_old_tuple_idx.begin(), new_to_old_tuple_idx.end()); - tensorflow::gtl::FlatMap<int64, int64> old_to_new_tuple_idx; + absl::flat_hash_map<int64, int64> old_to_new_tuple_idx; for (int64 new_idx = 0; new_idx < new_to_old_tuple_idx.size(); ++new_idx) { int64 old_idx = new_to_old_tuple_idx[new_idx]; old_to_new_tuple_idx[old_idx] = new_idx; @@ -405,7 +405,7 @@ static StatusOr<bool> TryPropagateConstant(HloInstruction* while_op) { // build a map from the tuple element index to the constant value. Limit this // to scalar constant values because propagating array constants can regress // performance by forcing us to copy constants. - tensorflow::gtl::FlatMap<int, const HloInstruction*> index_to_constant; + absl::flat_hash_map<int, const HloInstruction*> index_to_constant; for (int i = 0; i < root_operands.size(); i++) { HloInstruction* instr = root_operands[i]; if (instr->opcode() == HloOpcode::kGetTupleElement && diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD index f474ecb18c..06b6330321 100644 --- a/tensorflow/compiler/xla/tests/BUILD +++ b/tensorflow/compiler/xla/tests/BUILD @@ -422,6 +422,7 @@ xla_test( "//tensorflow/core:regexp_internal", "//tensorflow/core:test", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", ], ) diff --git a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc index db5a824de0..a6e70eb6ca 100644 --- a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc +++ b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include <vector> #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_split.h" @@ -32,7 +33,6 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/tests/test_utils.h" #include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/regexp.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" @@ -83,7 +83,7 @@ struct ParsedProfileOutputLine { Status ParseOneProfileOutputLine( const string& line, bool expect_hlo, - gtl::FlatMap<string, ParsedProfileOutputLine>* parsed_results, + absl::flat_hash_map<string, ParsedProfileOutputLine>* parsed_results, absl::Span<const absl::string_view> opcodes_to_ignore = {}) { string separator = "[^:]*:: +"; string match_percentage = R"(\d+\.\d*% +\d+Σ)"; @@ -208,7 +208,7 @@ XLA_TEST_F(HloProfileTest, ProfileSingleComputation) { std::vector<string> profile_output_lines = absl::StrSplit(profile_output, '\n'); - gtl::FlatMap<string, ParsedProfileOutputLine> parsed_profile_lines; + absl::flat_hash_map<string, ParsedProfileOutputLine> parsed_profile_lines; TF_ASSERT_OK(ParseOneProfileOutputLine( profile_output_lines[1], /*expect_hlo=*/false, &parsed_profile_lines)); @@ -314,7 +314,7 @@ XLA_TEST_F(HloProfileTest, ProfileWhileComputation) { ASSERT_NE(while_body_profile_end, profile_output_lines.end()); - gtl::FlatMap<string, ParsedProfileOutputLine> parsed_profile_lines; + absl::flat_hash_map<string, ParsedProfileOutputLine> parsed_profile_lines; for (auto while_body_profile_i = while_body_profile_start + 1; while_body_profile_i != while_body_profile_end; while_body_profile_i++) { |