1 /* Copyright 2020 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 // Common definitions related to TPUs. 17 18 #ifndef TENSORFLOW_CORE_TPU_TPU_DEFS_H_ 19 #define TENSORFLOW_CORE_TPU_TPU_DEFS_H_ 20 21 #include <array> 22 23 #include "tensorflow/core/framework/types.pb.h" 24 25 namespace tensorflow { 26 27 // Name of the TPU device, which corresponds to a single core. 28 extern const char* const DEVICE_TPU_NODE; // "TPU"; 29 30 // The TPU_REPLICATED_CORE device is a virtual device corresponding to one core 31 // of a replicated TPU computation. Only valid within the body of a 32 // TPUReplicate computation. 33 extern const char* const DEVICE_TPU_REPLICATED_CORE; 34 35 // DEVICE_TPU_SYSTEM is now defined in tensorflow/core/framework/types.h/.cc 36 37 // Name of the XLA_TPU_JIT compilation device, which is an internal device to 38 // compile graphs for TPU. Not registered as a device; no operators can be 39 // assigned to this device by a user. 40 extern const char* const DEVICE_TPU_XLA_JIT; // "XLA_TPU_JIT"; 41 42 // Attribute used internally to pass "is_mirrored_variable" attribute on 43 // TPUReplicatedInput nodes to _TPUReplicate. 44 extern const char* const TPUREPLICATE_MIRRORED_VAR_INDICES_ATTR; 45 46 // Attribute used internally to annoate ops which might consume TPU FastMem 47 // variable. 48 extern const char* const TPU_FAST_MEM_ATTR; // "_TPU_FAST_MEM" 49 50 extern const char* const kTPUReplicateAttr; 51 extern const char* const kOutsideCompilationAttr; 52 53 // Supported types for TPUs. 54 static constexpr std::array<DataType, 16> kTpuAllTypes = { 55 {DT_INT32, DT_UINT32, DT_HALF, DT_BFLOAT16, DT_FLOAT, DT_DOUBLE, DT_BOOL, 56 DT_COMPLEX64, DT_INT64, DT_UINT64, DT_QINT8, DT_QUINT8, DT_INT8, DT_UINT8, 57 DT_INT16, DT_UINT16}}; 58 59 } // namespace tensorflow 60 61 #endif // TENSORFLOW_CORE_TPU_TPU_DEFS_H_ 62