1 /* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/compiler/xla/service/compilation_environments.h"
17
18 #include <cstdint>
19 #include <memory>
20 #include <string>
21 #include <string_view>
22 #include <utility>
23
24 #include "absl/base/thread_annotations.h"
25 #include "absl/container/flat_hash_map.h"
26 #include "absl/memory/memory.h"
27 #include "absl/strings/str_cat.h"
28 #include "absl/strings/str_join.h"
29 #include "absl/strings/string_view.h"
30 #include "absl/synchronization/mutex.h"
31 #include "tensorflow/core/platform/logging.h"
32 #include "tensorflow/core/platform/protobuf.h"
33
34 namespace xla {
35 namespace {
36
37 // A global singleton stats object for implementing CompilationEnvironments::{
38 // DefaultEnvCreated(), DefaultEnvCreatedByCompilationEnvironments(),
39 // EnvAdded()}.
40 class GlobalCompEnvStats {
41 public:
GetSingleton()42 static GlobalCompEnvStats& GetSingleton() {
43 static GlobalCompEnvStats* singleton = new GlobalCompEnvStats();
44
45 return *singleton;
46 }
47
DefaultEnvCreated(std::string_view env_type)48 void DefaultEnvCreated(std::string_view env_type) ABSL_LOCKS_EXCLUDED(mu_) {
49 {
50 absl::MutexLock l(&mu_);
51 ++stats_[std::string(env_type)].default_env_created;
52 }
53 VLOG(1) << "New GlobalCompEnvStats value: " << ToString();
54 }
55
DefaultEnvCreatedByCompilationEnvironments(std::string_view env_type)56 void DefaultEnvCreatedByCompilationEnvironments(std::string_view env_type)
57 ABSL_LOCKS_EXCLUDED(mu_) {
58 {
59 absl::MutexLock l(&mu_);
60 ++stats_[std::string(env_type)]
61 .default_env_created_by_compilation_environments;
62 }
63 VLOG(1) << "New GlobalCompEnvStats value: " << ToString();
64 }
65
EnvAdded(std::string_view env_type)66 void EnvAdded(std::string_view env_type) ABSL_LOCKS_EXCLUDED(mu_) {
67 {
68 absl::MutexLock l(&mu_);
69 ++stats_[std::string(env_type)].env_added;
70 }
71 VLOG(1) << "New GlobalCompEnvStats value: " << ToString();
72 }
73
ToString() const74 std::string ToString() const ABSL_LOCKS_EXCLUDED(mu_) {
75 absl::ReaderMutexLock l(&mu_);
76 return absl::StrJoin(
77 stats_, "; ",
78 [](std::string* out, const StatMap::value_type& env_stats_pair) {
79 absl::StrAppend(out, env_stats_pair.first, ": { ",
80 env_stats_pair.second.ToString(), " }");
81 });
82 }
83
84 private:
85 struct PerEnvStats {
ToStringxla::__anon391dd9920111::GlobalCompEnvStats::PerEnvStats86 std::string ToString() const {
87 return absl::StrCat(
88 "# default envs created: ", default_env_created, " ",
89 "# default envs created by CompilationEnvironments: ",
90 default_env_created_by_compilation_environments, " ",
91 "# envs added to CompilationEnvironments: ", env_added);
92 }
93
94 unsigned default_env_created = 0;
95 unsigned default_env_created_by_compilation_environments = 0;
96 unsigned env_added = 0;
97 };
98
99 using StatMap = absl::flat_hash_map<std::string, PerEnvStats>;
100
101 GlobalCompEnvStats() = default;
102 GlobalCompEnvStats(const GlobalCompEnvStats&) = delete;
103 GlobalCompEnvStats& operator=(const GlobalCompEnvStats&) = delete;
104 GlobalCompEnvStats(GlobalCompEnvStats&&) = delete;
105 GlobalCompEnvStats& operator=(GlobalCompEnvStats&&) = delete;
106
107 mutable absl::Mutex mu_;
108 StatMap stats_ ABSL_GUARDED_BY(mu_);
109 };
110
111 } // namespace
112
DefaultEnvCreated(std::string_view env_type)113 void CompilationEnvironments::DefaultEnvCreated(std::string_view env_type) {
114 GlobalCompEnvStats::GetSingleton().DefaultEnvCreated(env_type);
115 }
116
DefaultEnvCreatedByCompilationEnvironments(std::string_view env_type)117 void CompilationEnvironments::DefaultEnvCreatedByCompilationEnvironments(
118 std::string_view env_type) {
119 GlobalCompEnvStats::GetSingleton().DefaultEnvCreatedByCompilationEnvironments(
120 env_type);
121 }
122
EnvAdded(std::string_view env_type)123 void CompilationEnvironments::EnvAdded(std::string_view env_type) {
124 GlobalCompEnvStats::GetSingleton().EnvAdded(env_type);
125 }
126
operator =(const CompilationEnvironments & rhs)127 CompilationEnvironments& CompilationEnvironments::operator=(
128 const CompilationEnvironments& rhs) {
129 Clear();
130 for (const auto& descriptor_message_pair : rhs.environments_) {
131 auto env = absl::WrapUnique(descriptor_message_pair.second->New());
132 env->CopyFrom(*descriptor_message_pair.second);
133 environments_.insert({descriptor_message_pair.first, std::move(env)});
134 }
135 return *this;
136 }
137
AddEnv(std::unique_ptr<tensorflow::protobuf::Message> env)138 void CompilationEnvironments::AddEnv(
139 std::unique_ptr<tensorflow::protobuf::Message> env) {
140 auto descriptor = env->GetDescriptor();
141 if (environments_.contains(descriptor)) {
142 LOG(WARNING) << "Replacing CompilationEnvironment of type "
143 << descriptor->full_name();
144 }
145
146 environments_.insert({descriptor, std::move(env)});
147 EnvAdded(descriptor->full_name());
148 }
149
150 } // namespace xla
151