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 #ifndef TENSORFLOW_COMPILER_TF2XLA_TF2XLA_DEFS_H_ 17 #define TENSORFLOW_COMPILER_TF2XLA_TF2XLA_DEFS_H_ 18 19 #include <array> 20 21 #include "absl/strings/string_view.h" 22 23 namespace tensorflow { 24 25 // Marks a node for XLA compilation. The attribute value indicates the 26 // compilation device type. 27 inline constexpr absl::string_view kCompileDeviceTypeAttr = 28 "_xla_compile_device_type"; 29 // Marks a node for replication. The attribute value indicates the replication 30 // metadata op. 31 inline constexpr absl::string_view kReplicationInfoAttr = "_replication_info"; 32 // Marks a node for XLA-TPU compilation. The attribute value indicates the 33 // associated compilation cluster and replication metadata op. 34 inline constexpr absl::string_view kTpuReplicateAttr = "_tpu_replicate"; 35 // Marks a node inside of an XLA compilation cluster to be placed outside of the 36 // cluster. 37 inline constexpr absl::string_view kXlaOutsideCompilationAttr = 38 "_xla_outside_compilation"; 39 // Frontend attributes ID. 40 inline constexpr absl::string_view kXlaFrontendAttributesAttrName = 41 "_XlaFrontendAttributes"; 42 // Device types. 43 inline constexpr absl::string_view kCpuDevice = "CPU"; 44 inline constexpr absl::string_view kGpuDevice = "GPU"; 45 inline constexpr absl::string_view kTpuDevice = "TPU"; 46 inline constexpr std::array<absl::string_view, 3> kValidDeviceTypes = { 47 kCpuDevice, kGpuDevice, kTpuDevice}; 48 // Attributes that need to be propagated during rewrites (e.g., in 49 // functionalization). 50 inline constexpr std::array<absl::string_view, 5> kAttrsToPropagate = { 51 kCompileDeviceTypeAttr, 52 kReplicationInfoAttr, 53 kXlaFrontendAttributesAttrName, 54 kXlaOutsideCompilationAttr, 55 kTpuReplicateAttr, 56 }; 57 58 } // namespace tensorflow 59 60 #endif // TENSORFLOW_COMPILER_TF2XLA_TF2XLA_DEFS_H_ 61