1 // Copyright 2019 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 #ifndef TENSORFLOW_COMPILER_XLA_PYTHON_TPU_DRIVER_EVENT_ID_H_ 16 #define TENSORFLOW_COMPILER_XLA_PYTHON_TPU_DRIVER_EVENT_ID_H_ 17 18 #include <cstdint> 19 #include <ostream> 20 #include <string> 21 #include <utility> 22 23 #include "absl/strings/str_cat.h" 24 25 namespace tpu_driver { 26 27 // For gRPC serialization, events are represented as a pair of 28 // {client, operation} ids. To simplify serialization, these are encoded as a 29 // single integer field. 30 // 31 // This class provides a typed interface for these values as well as support for 32 // hashing and ostreams (for logging). 33 struct EventId { 34 uint64_t client_id; 35 uint64_t operation_id; 36 37 template <typename H> AbslHashValueEventId38 friend H AbslHashValue(H h, const EventId& c) { 39 return H::combine(std::move(h), c.client_id, c.operation_id); 40 } 41 42 bool operator==(const EventId& r) const { 43 return r.client_id == client_id && r.operation_id == operation_id; 44 } 45 46 friend std::ostream& operator<<(std::ostream& os, EventId r) { 47 return os << r.client_id << ":" << r.operation_id; 48 } 49 ToStringEventId50 std::string ToString() const { 51 return absl::StrCat(client_id, ":", operation_id); 52 } 53 AsIntEventId54 uint64_t AsInt() const { return client_id << 44 | operation_id; } 55 FromIntEventId56 static EventId FromInt(uint64_t value) { 57 return EventId{value >> 44, value & 0xfffffffffff}; 58 } 59 }; 60 61 } // namespace tpu_driver 62 63 #endif // TENSORFLOW_COMPILER_XLA_PYTHON_TPU_DRIVER_EVENT_ID_H_ 64