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 #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_ARGUMENT_H_ 17 #define TENSORFLOW_COMPILER_TF2XLA_XLA_ARGUMENT_H_ 18 19 #include "absl/types/optional.h" 20 #include "absl/types/span.h" 21 #include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h" 22 #include "tensorflow/compiler/tf2xla/xla_resource.h" 23 #include "tensorflow/compiler/xla/client/xla_builder.h" 24 #include "tensorflow/compiler/xla/service/hlo_sharding.h" 25 #include "tensorflow/core/framework/tensor.h" 26 27 namespace tensorflow { 28 29 // Describes how to derive the value of each _Arg node in the graph/function 30 // being compiled. There must be one Argument for each _Arg index. 31 struct XlaArgument { 32 enum Kind { 33 // Default value; not a valid kind. 34 kInvalid, 35 36 // Argument is a compile-time constant. No associated runtime parameter. 37 kConstant, 38 39 // Argument is a Variable, TensorArray, or Stack resource. Has an 40 // associated runtime parameter iff `initialized` is true. 41 kResource, 42 43 // A resource variable with a constant value known at compile time. 44 kConstantResource, 45 46 // Argument is a run-time parameter. 47 kParameter, 48 49 // Argument is an XLA token. 50 kToken, 51 52 // Argument is a TensorList. 53 kTensorList, 54 }; 55 56 Kind kind = kInvalid; 57 58 // The type of the argument. If the argument is a resource, this 59 // is the type of the variable's value, not DT_RESOURCE. 60 DataType type = DT_INVALID; 61 62 // The shape of the argument. For: 63 // * a parameter: the shape of the parameter. We allow setting the xla shape 64 // if known. This helps avoid conversions to and from TensorShape. 65 // * a constant: ignored; the shape given by constant_value is used 66 // instead. 67 // * an uninitialized resource: ignored. We don't yet know the shape of an 68 // uninitialized resource (otherwise we would have initialized it!) 69 // * an initialized variable: the shape of the variable's value. 70 // * an initialized TensorArray or Stack resource: the shape of an entry in 71 // the TensorArray/Stack. Note this is the size of a single entry, not the 72 // XLA data structure that represents the complete stack/array. 73 absl::variant<TensorShape, xla::Shape> shape; 74 75 // The value of the argument, if it is a compile-time constant. Must be a 76 // host-memory tensor. 77 Tensor constant_value; 78 79 // The upper bounds of the value. 80 std::optional<Tensor> value_bound; 81 82 // Indicates whether each value is dynamic or constant. 83 std::optional<Tensor> value_dynamism; 84 85 // The name of this argument, used for debugging. 86 string name; 87 88 // The name of TensorFlow _Arg node, used for debugging. 89 string node_name; 90 91 // For a kResource, what kind of resource is it? 92 XlaResource::Kind resource_kind = XlaResource::kInvalid; 93 94 // For a kResource, has this resource been initialized? 95 bool initialized = false; 96 97 // For a kResource, is this resource on Fast Memory. 98 bool fast_mem = false; 99 100 // For a TensorArray or Stack resource, what is the array's declared size? 101 // (Used for lazy initialization.) 102 int64_t max_array_size = -1; 103 104 // TensorArray resource parameters are passed as (array, gradient array 0, 105 // ..., gradient array k), where the gradient arrays are in the same order 106 // as `tensor_array_gradients`. 107 std::set<string> tensor_array_gradients; 108 109 // Whether this argument will receive the same data across all replicas. 110 bool is_same_data_across_replicas = false; 111 112 bool operator==(const XlaArgument& other) const; 113 114 // Returns a human-readable summary of the argument. 115 string HumanString() const; 116 117 // Returns the dimension sizes for either TensorShape or xla::Shape. 118 std::vector<int64_t> DimensionSizes() const; 119 absl::InlinedVector<int64_t, 4> DimensionSizesAsInlinedVector() const; 120 121 // Returns the human-readable string for either TensorShape or xla::Shape. 122 string ShapeHumanString() const; 123 124 // Whether to broadcast this parameter to all replicas before use. 125 // When true, xla_compiler should input/output alias this arg to prevent 126 // unnecessary HBM usage. 127 bool requires_broadcast = false; 128 std::optional<ManagedStackTrace> definition_stack_trace; 129 }; 130 131 // Returns true if any of `args` is an uninitialized resource variable. 132 bool AnyUninitializedResourceArg(absl::Span<const XlaArgument> args); 133 134 } // end namespace tensorflow 135 136 #endif // TENSORFLOW_COMPILER_TF2XLA_XLA_ARGUMENT_H_ 137