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_DTENSOR_CC_CONSTANTS_H_ 17 #define TENSORFLOW_DTENSOR_CC_CONSTANTS_H_ 18 19 namespace tensorflow { 20 namespace dtensor { 21 // Constants used within dtensor scope. 22 23 // Qualified attribute without `_` prefix. 24 // Used in Ops attribute registration. 25 static constexpr char kQualifiedLayoutAttr[] = "layout"; 26 27 // Internal attribute to DTensor MLIR passes and Graph nodes. 28 // Prefixed with `_` so that it doesn't require op attribute registration. 29 static constexpr char kLayoutAttr[] = "_layout"; 30 31 // Indicates a non-binding layout hint provided by the user. 32 // `tf` prefix attached in MLIR importer for dialect requirements. 33 static constexpr char kCustomDefaultLayoutAttr[] = "tf._default_layout"; 34 35 // Indicates a non-binding layout hint provided by the user. 36 static constexpr char kDefaultLayoutAttr[] = "_default_layout"; 37 38 // Attribute carries layout information from Custom Device Arguments. 39 // `tf` prefix attached in MLIR importer for dialect requirements. 40 static constexpr char kCustomDeviceAttr[] = "tf._layout"; 41 42 // Attribute attached on _Arg node for the mesh config. 43 static constexpr char kMeshAttr[] = "_mesh"; 44 45 // Attribute carries mesh information from Custom Device Arguments. 46 // `tf` prefix attached in MLIR importer for dialect requirements. 47 static constexpr char kCustomDeviceMeshAttr[] = "tf._mesh"; 48 49 // Attribute carries argument indices for newly inferred layout of resource 50 // handle. 51 static constexpr char kNewResourceLayoutIndices[] = 52 "_inferred_resource_indices"; 53 54 // Attribute carries layout for newly inferred layout of resource handle. 55 static constexpr char kNewResourceArgLayouts[] = "_inferred_resource_layouts"; 56 57 // Attribute carries input layout information for shape op. 58 static constexpr char kShapeOpInputLayout[] = "_shape_input_layout"; 59 60 // Attribute carries input layout index for shape op. This forms a 1 -> 1 61 // mapping for kShapeOpInputLayout above. 62 static constexpr char kShapeOpInputLayoutIndices[] = "_shape_input_indices"; 63 64 // Attribute that carries global shape of operation. Used to preserve global 65 // shape to be used during SPMD expansion. 66 static constexpr char kGlobalShape[] = "_global_shape"; 67 68 // Global shape attribute with `tf.` dialect to be used for annotating func op 69 // arguments/return values. 70 static constexpr char kGlobalShapeDialectAttr[] = "tf._global_shape"; 71 72 // Attribute attached to resource-type function arguments containing the local 73 // shape of the tensor that is being assigned to it. 74 static constexpr char kAssignedResourceLocalShape[] = 75 "tf._assigned_resource_local_shape"; 76 77 // Tensor handles smaller than this is considered as small tensor. We perform 78 // some optimizations around it. For example, will be transformed into constant 79 // values during graph building, instead of being passed as inputs. In addition, 80 // we allow automatical broadcasting small non-DTensor to DTensor device, which 81 // is very useful for shape/axis info tensor in eager mode (eliminating the need 82 // forcing users to do explicit copy-to-mesh). 83 static constexpr int kSmallTensorThreshold = 20; 84 85 // Contains a serialized mesh. Will be attached to a FloorMod op to denote which 86 // mesh the output of the FloorMod op is giving coordinates for. 87 static constexpr char kMeshCoordinatesAttr[] = "_mesh_coordinates"; 88 89 // Attribute used to determine if a module pass should log long form information 90 // such as IR dumps etc. 91 static constexpr char kDoNotLog[] = "dtensor.do_not_log"; 92 93 // The number of TPU cores in a donut. 94 static constexpr int kTpuDonutSize = 8; 95 96 // An attribute used to cache the computation of device seeds, so that we don't 97 // constantly recompute device seeds in a cluster for a given layout. 98 static constexpr char kDeviceSeedForMeshDims[] = 99 "dtensor.device_seed_for_mesh_dims"; 100 101 // Attribute that determines whether to skip XlA compilation. There are some ops 102 // that run on a TPU mesh but are not expected to be compiled by XLA, e.g. 103 // VarHandleOp, DestroyResourceOp, etc. For such an case, set this attribute 104 // to true on the StatefulPartitionedCallOp generated by MLIR lowering. 105 static constexpr char kSkipXlaCompilation[] = "_skip_xla_compilation"; 106 107 // Prefix of pipelining mesh name (kPipelineMeshNamePrefix + composite device 108 // name). 109 static constexpr char kPipelineMeshNamePrefix[] = "pipe_cluster:"; 110 111 // An attribute which stores the cache_key for the graph in the module. Used 112 // to uniquely name functions. 113 static constexpr char kCacheKey[] = "dtensor.cache_key"; 114 115 // An attribute that determines whether a tensor is a sparse tensor. If this 116 // attribute exists in a tensor, then this tensor is a sparse tensor. 117 static constexpr char kSparseValue[] = "tf._sparse"; 118 119 // TPUEmbedding configuration attribute with `tf.` dialect to be used for 120 // annotating func op that contains tpu embedding configuration ops. 121 static constexpr char kTPUEmbeddingConfiguration[] = 122 "tf._tpu_embedding_configuration"; 123 124 // Attribute mapping table_id to func op arguments using as TPUEmbedding tables 125 // `tf` prefix attached in MLIR importer for dialect requirements. 126 static constexpr char kTPUEmbeddingTableID[] = "tf._tpu_embedding_table_id"; 127 128 // Attribute mapping slot_id to func op arguments using as TPUEmbedding slot 129 // variables.`tf` prefix attached in MLIR importer for dialect requirements. 130 static constexpr char kTPUEmbeddingSlotID[] = "tf._tpu_embedding_slot_id"; 131 132 // Name of dtensor load embedding function. 133 static constexpr char kLoadEmbeddingFn[] = "load_embedding_fn"; 134 135 // Name of dtensor retrieve embedding function. 136 static constexpr char kRetrieveEmbeddingFn[] = "retrieve_embedding_fn"; 137 } // namespace dtensor 138 } // namespace tensorflow 139 140 #endif // TENSORFLOW_DTENSOR_CC_CONSTANTS_H_ 141