1 /* Copyright 2021 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 #include "tensorflow/core/tpu/kernels/tpu_fingerprint_lookup.h"
16
17 namespace tensorflow {
18 namespace tpu {
19
Create()20 TpuFingerprintLookup* TpuFingerprintLookup::Create() {
21 return new TpuFingerprintLookup();
22 }
23
RegisterKeyAndIntermediatePair(uint64 key,uint64 intermediate)24 void TpuFingerprintLookup::RegisterKeyAndIntermediatePair(uint64 key,
25 uint64 intermediate) {
26 absl::MutexLock lock(&mu_);
27 auto [it, emplaced] = intermediate_to_key_.try_emplace(intermediate, key);
28 if (it->second != key) {
29 VLOG(2) << "The key (" << it->second
30 << ") is associated with an existing intermediate ( " << it->first
31 << "), which does not match the requesting key (" << key << ").";
32 }
33 }
34
RegisterIntermediateAndValuePair(uint64 intermediate,std::string value)35 bool TpuFingerprintLookup::RegisterIntermediateAndValuePair(uint64 intermediate,
36 std::string value) {
37 absl::MutexLock lock(&mu_);
38 auto it = intermediate_to_key_.find(intermediate);
39 if (it == intermediate_to_key_.end()) {
40 VLOG(2) << "Cannot find the intermediate ( " << intermediate
41 << "). A RegisterKeyAndIntermediatePair must precedes.";
42 return false;
43 } else {
44 uint64 key = it->second;
45 bool is_successful = false;
46 VLOG(2) << "registering key (" << key << ") with value: " << value;
47 auto it = key_to_value_.find(key);
48 if (it == key_to_value_.end()) {
49 // A new key. If the value is not seen before, register key-value and
50 // value-key pairs. Otherwise, skip registration.
51 auto maybe_existing_key = value_to_key_.find(value);
52 if (maybe_existing_key == value_to_key_.end()) {
53 key_to_value_.emplace(key, value);
54 value_to_key_.emplace(value, key);
55 is_successful = true;
56 } else {
57 // The value is registered before with a different key. Skip
58 // registration.
59 if (maybe_existing_key->second != key) {
60 VLOG(2) << "The value (" << value
61 << ") is associated with an existing key ( "
62 << maybe_existing_key->second
63 << "), which does not match the requesting key (" << key
64 << ").";
65 }
66 }
67 } else {
68 // The key is registered before, no actions needed. For debugging purpose,
69 // check if existing value agrees with the value.
70 if (it->second != value) {
71 VLOG(2) << "The key (" << key
72 << ") has been registered and the requesting value ( " << value
73 << " and the existing" << it->second << ") doesn't match.";
74 }
75 }
76 DCHECK(key_to_value_.size() == value_to_key_.size());
77
78 return is_successful;
79 }
80 }
81
Lookup(uint64 key)82 absl::optional<::tensorflow::StringPiece> TpuFingerprintLookup::Lookup(
83 uint64 key) {
84 absl::MutexLock lock(&mu_);
85 auto it = key_to_value_.find(key);
86 if (it == key_to_value_.end()) {
87 return absl::optional<::tensorflow::StringPiece>{};
88 } else {
89 return it->second;
90 }
91 }
92
93 } // namespace tpu
94 } // namespace tensorflow
95