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 absl::optional<Tensor> value_bound; 81 82 // The name of this argument, used for debugging. 83 string name; 84 85 // The name of TensorFlow _Arg node, used for debugging. 86 string node_name; 87 88 // For a kResource, what kind of resource is it? 89 XlaResource::Kind resource_kind = XlaResource::kInvalid; 90 91 // For a kResource, has this resource been initialized? 92 bool initialized = false; 93 94 // For a kResource, is this resource on Fast Memory. 95 bool fast_mem = false; 96 97 // For a TensorArray or Stack resource, what is the array's declared size? 98 // (Used for lazy initialization.) 99 int64 max_array_size = -1; 100 101 // TensorArray resource parameters are passed as (array, gradient array 0, 102 // ..., gradient array k), where the gradient arrays are in the same order 103 // as `tensor_array_gradients`. 104 std::set<string> tensor_array_gradients; 105 106 // dynamic dims to arg number map. Empty if no dynamic shapes. 107 std::map<int32, int32> dynamic_dim_to_arg_num_map; 108 bool is_pad_arg = false; 109 110 // Whether this argument will receive the same data across all replicas. 111 bool is_same_data_across_replicas = false; 112 113 bool operator==(const XlaArgument& other) const; 114 115 // Returns a human-readable summary of the argument. 116 string HumanString() const; 117 118 // Returns the dimension sizes for either TensorShape or xla::Shape. 119 std::vector<int64> DimensionSizes() const; 120 absl::InlinedVector<int64, 4> DimensionSizesAsInlinedVector() const; 121 122 // Returns the human-readable string for either TensorShape or xla::Shape. 123 string ShapeHumanString() const; 124 }; 125 126 // Returns true if any of `args` is an uninitialized resource variable. 127 bool AnyUninitializedResourceArg(absl::Span<const XlaArgument> args); 128 129 } // end namespace tensorflow 130 131 #endif // TENSORFLOW_COMPILER_TF2XLA_XLA_ARGUMENT_H_ 132