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_KERNELS_WHILE_OP_H_ 17 #define TENSORFLOW_COMPILER_TF2XLA_KERNELS_WHILE_OP_H_ 18 19 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" 20 #include "tensorflow/core/framework/attr_value.pb.h" 21 22 namespace tensorflow { 23 24 // This TensorFlow op provides a functional iteration primitive. 25 // 26 // The inputs and outputs of the loop body must agree on the number, types, and 27 // shapes of the Tensors carried around the loop body. 28 // 29 // Computations in while loops may read from and write to resource variables. 30 // Resource variables may be passed as arguments to a function's body and 31 // condition functions. The XlaCompiler converts resource variable arguments 32 // into parameters to the XLA computation and moves them to the end of the 33 // parameter list, and by using the `return_updated_values_for_all_variables` 34 // we ensure that all variables that appear in the input also appear at the 35 // end of the body's output. This ensures the loop body's input and output 36 // signatures match. 37 // 38 // It is the user's responsibility to ensure that each non-variable _Arg matches 39 // the corresponding _Retval. 40 // 41 // For example, suppose we have a loop body with arguments: 42 // DT_INT32, DT_RESOURCE (pointing to a DT_BOOL var), DT_FLOAT 43 // and return values 44 // DT_INT32, DT_FLOAT 45 // It is an error for the body to return DT_RESOURCE values. 46 // 47 // The body will be lowered into an XLA computation that takes and returns a 48 // tuple with XLA type (I32, F32, PRED). Note the resource variable appears at 49 // the end of both the loop body's input and output argument lists. 50 class XlaWhileOp : public XlaOpKernel { 51 public: 52 explicit XlaWhileOp(OpKernelConstruction* ctx); 53 54 void Compile(XlaOpKernelContext* ctx) override; 55 56 private: 57 NameAttrList cond_name_attr_; 58 NameAttrList body_name_attr_; 59 bool has_token_input_output_; 60 std::vector<string> token_input_nodes_; 61 // Whether to propagate compile time consts into the loop body. 62 // This is not supported by default now since it may cause HBM memory 63 // overheads. 64 bool propagate_compile_time_consts_ = false; 65 66 TF_DISALLOW_COPY_AND_ASSIGN(XlaWhileOp); 67 }; 68 69 } // namespace tensorflow 70 71 #endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_WHILE_OP_H_ 72