1/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
25namespace tensorflow {
26
27// Name of the TPU device, which corresponds to a single core.
28extern 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.
33extern 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.
40extern 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.
44extern const char* const TPUREPLICATE_MIRRORED_VAR_INDICES_ATTR;
45
46// Attribute used internally to annoate ops which might consume TPU FastMem
47// variable.
48extern const char* const TPU_FAST_MEM_ATTR; // "_TPU_FAST_MEM"
49
50extern const char* const kTPUReplicateAttr;
51extern const char* const kOutsideCompilationAttr;
52
53// Supported types for TPUs.
54static 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