1 /* Copyright 2018 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_XLA_SERVICE_WHILE_UTIL_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_UTIL_H_ 18 19 #include "tensorflow/compiler/xla/service/call_inliner.h" 20 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 21 22 namespace xla { 23 class WhileUtil { 24 public: 25 // Holds a return value from MakeInstructionsLiveIn. 26 struct MakeInstructionsLiveInResult { 27 // The new while operation that has the requested values live in. 28 HloInstruction* new_while_instr; 29 30 // The i'th element of `while_body_live_in_values` is an instruction in the 31 // while body that holds the i'th *newly added* live in value at runtime. 32 std::vector<HloInstruction*> while_body_live_in_values; 33 34 // `while_body_instruction_map` maps instructions in the original while body 35 // to the corresponding instructions in the body for the newly created while 36 // operation. 37 CallInliner::InlinedInstructionMap while_body_instruction_map; 38 }; 39 40 // Replaces `while_instr` with a new while instruction that is equivalent to 41 // `while_instr`, except that it has all of the HLO instructions in 42 // `instructions` as live-in, loop invariant values. These new live in values 43 // are represented as new elements appended to the parameter of the while 44 // loop, which must be of tuple shape. GetTupleElement instructions computing 45 // each new live in value is returned in the `while_body_live_in_values` 46 // vector. 47 // 48 // Precondition: `while_instr` must have a tuple shaped state. 49 // 50 // Every instruction in `instructions` must be contained in the computation 51 // that contains `while_instr`. 52 static StatusOr<MakeInstructionsLiveInResult> MakeInstructionsLiveIn( 53 HloInstruction* while_instr, 54 tensorflow::gtl::ArraySlice<HloInstruction*> instructions); 55 }; 56 } // namespace xla 57 58 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_UTIL_H_ 59