1 /* Copyright 2017 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_RESOURCE_H_ 17 #define TENSORFLOW_COMPILER_TF2XLA_XLA_RESOURCE_H_ 18 19 #include <memory> 20 21 #include "tensorflow/compiler/xla/client/computation_builder.h" 22 #include "tensorflow/compiler/xla/xla_data.pb.h" 23 #include "tensorflow/core/framework/tensor_shape.h" 24 #include "tensorflow/core/framework/types.pb.h" 25 #include "tensorflow/core/lib/core/status.h" 26 27 namespace tensorflow { 28 29 // Represents a resource, such as a Variable or TensorArray. 30 class XlaResource { 31 public: 32 enum Kind { 33 kInvalid, 34 kVariable, 35 kTensorArray, 36 kStack, 37 }; 38 39 XlaResource(Kind kind, int arg_num, string name, DataType type, 40 TensorShape shape, 41 const xla::ComputationDataHandle& initial_value, 42 int64 tensor_array_size, 43 const std::set<string>& tensor_array_gradients); 44 45 XlaResource(const XlaResource&) = delete; 46 XlaResource(XlaResource&&) = delete; 47 XlaResource& operator=(const XlaResource&) = delete; 48 XlaResource& operator=(XlaResource&&) = delete; 49 50 Kind kind() const { return kind_; } 51 52 // If this resource is visible externally to the computation, what was its 53 // argument number? 54 // < 0 means "not visible externally". 55 int arg_num() const { return arg_num_; } 56 57 // A descriptive name for the resource, used in error messages. 58 const string& name() const { return name_; } 59 60 // Current type and value of the resource. Uninitialized resources are 61 // represented by a default (zero) handle and type DT_INVALID. 62 // While the type of a resource is notionally fixed during execution, when 63 // a resource is first initialized we do not yet know its type, so we keep 64 // track of its type dynamically. 65 DataType type() const { return type_; } 66 67 // Shape of the resource. For an uninitialized resource, this is ignored. 68 // For a Variable, this is the shape of the value. For a TensorArray or Stack 69 // this is the shape of each entry in the TensorArray/Stack. 70 const TensorShape& shape() const { return shape_; } 71 72 const xla::ComputationDataHandle& value() const { return value_; } 73 74 // Value of the resource at computation entry. Used to detect which 75 // variables have new values that need to be written back. 76 const xla::ComputationDataHandle& initial_value() const { 77 return initial_value_; 78 } 79 80 // A variable is initialized if it has a value. 81 bool initialized() const { return value_.handle() > 0; } 82 83 // Sets the type and shape of the resource. The type and shape of a resource 84 // must not change once the variable has been initialized. 85 Status SetTypeAndShape(DataType type, const TensorShape& shape); 86 87 // Sets the current value of the resource. Returns an error if the type is not 88 // set to a valid value. 89 Status SetValue(const xla::ComputationDataHandle& value); 90 91 // Sets the current value of the resource to an all-zero value. 92 Status SetZeroValue(xla::ComputationBuilder* builder); 93 94 // Looks up the gradient for `source`, or creates it if it does not already 95 // exist. The call target must be an initialized TensorArray resource. A 96 // TensorArray can have multiple named gradients; see the operator 97 // documentation for TensorArrayGradV3 for details. 98 Status GetOrCreateTensorArrayGradient(const string& source, 99 xla::ComputationBuilder* builder, 100 XlaResource** gradient_out); 101 102 // Packs a resource into a single XLA value `pack`, suitable for use as 103 // an XlaCompiler::Argument. For non-TensorArrays or TensorArrays without 104 // gradients, sets `*pack` to `value`. 105 // For TensorArrays with gradients, packs the value and its gradient values in 106 // a tuple; the gradients values are packed in order by source name. 107 Status Pack(xla::ComputationDataHandle* pack, 108 xla::ComputationBuilder* builder) const; 109 110 // Updates the resource with values from `pack`. If `gradient_sources` is 111 // non-empty, treats `pack` as a tuple that represents a TensorArray and 112 // its gradients, and unpacks and updates the gradient resources. 113 // If `reset_initial_values` is true, sets the initial_values as well as the 114 // values. 115 // Opposite of Pack(). 116 Status SetFromPack(const std::set<string>& gradient_sources, 117 const xla::ComputationDataHandle& pack, 118 xla::ComputationBuilder* builder); 119 120 // TensorArray and Stack specific fields 121 122 // 'tensor_array_size' stores the expected size of the TensorArray or Stack. 123 // We need to store this since sometimes TensorArrays must be initialized 124 // lazily since we do not know the element shape at construction time. 125 // Used by both TensorArrays and Stacks. 126 int64 tensor_array_size() const { return tensor_array_size_; } 127 void set_tensor_array_size(int64 size) { tensor_array_size_ = size; } 128 129 // 'tensor_array_gradient' is a map from TensorArrayGradV3 'source' attributes 130 // to an XlaResource containing the gradient TensorArrays. We store a pointer 131 // here since there should only be one gradient TensorArray per 'source' 132 // string, irrespective of the number of calls to TensorArrayGrad. The map 133 // is ordered since values are packed into tuples by Pack() sorted by name 134 // order. 135 const std::map<string, std::unique_ptr<XlaResource>>& tensor_array_gradients() 136 const { 137 return tensor_array_gradients_; 138 } 139 140 private: 141 const Kind kind_; 142 const int arg_num_; 143 const string name_; 144 145 DataType type_; 146 TensorShape shape_; 147 xla::ComputationDataHandle value_; 148 xla::ComputationDataHandle initial_value_; 149 150 int64 tensor_array_size_ = -1; 151 152 std::map<string, std::unique_ptr<XlaResource>> tensor_array_gradients_; 153 }; 154 155 } // namespace tensorflow 156 157 #endif // TENSORFLOW_COMPILER_TF2XLA_XLA_RESOURCE_H_ 158