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 "absl/container/flat_hash_map.h"
20 #include "absl/container/inlined_vector.h"
21 #include "tensorflow/compiler/xla/service/call_inliner.h"
22 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
23 
24 namespace xla {
25 class WhileUtil {
26  public:
27   // Holds a return value from MakeInstructionsLiveIn.
28   struct MakeInstructionsLiveInResult {
29     // The new while operation that has the requested values live in.
30     HloInstruction* new_while_instr;
31 
32     // The i'th element of `while_body_live_in_values` is an instruction in the
33     // while body that holds the i'th *newly added* live in value at runtime.
34     std::vector<HloInstruction*> while_body_live_in_values;
35 
36     // `while_body_instruction_map` maps instructions in the original while body
37     // to the corresponding instructions in the body for the newly created while
38     // operation.
39     CallInliner::InlinedInstructionMap while_body_instruction_map;
40   };
41 
42   // Replaces `while_instr` with a new while instruction that is equivalent to
43   // `while_instr` except that it has all of the HLO instructions in
44   // `instructions` as live-in, loop invariant values.  These new live in values
45   // are represented as new elements appended to the parameter of the while
46   // loop, which must be of tuple shape.  GetTupleElement instructions computing
47   // each new live in value is returned in the `while_body_live_in_values`
48   // vector.
49   //
50   // Deletes `while_instr` after replacing it.
51   //
52   // Preconditions:
53   //
54   //  `while_instr` must have a tuple shaped state.
55   //
56   //   Every instruction in `instructions` must be contained in the computation
57   //   that contains `while_instr`.
58   static StatusOr<MakeInstructionsLiveInResult> MakeInstructionsLiveIn(
59       HloInstruction* while_instr,
60       absl::Span<HloInstruction* const> instructions);
61 
62   using LoopStateTy = std::vector<HloInstruction*>;
63   using LoopBodyGeneratorTy = std::function<StatusOr<LoopStateTy>(
64       HloInstruction* /*induction_var*/,
65       const LoopStateTy& /*current_values*/)>;
66 
67   // Creates a while loop in `computation` that runs for `trip_count`
68   // iterations.  The structure of the while loop is as follows, in pseudocode:
69   //
70   //  loop_state while_loop() {
71   //    indvar = 0;
72   //    loop_state = init_values
73   //    while (indvar < trip_count) {
74   //      loop_state = loop_body_generator(loop_state)
75   //      indvar++;
76   //    }
77   //    return loop_state;
78   //  }
79   static StatusOr<LoopStateTy> MakeCountedLoop(
80       HloComputation* computation, int32 trip_count,
81       const LoopStateTy& init_values,
82       const LoopBodyGeneratorTy& loop_body_generator,
83       const OpMetadata& metadata);
84 
85   // Returns the GetTupleElement instructions in `while_body` that access
86   // elements in the parameter tuple that don't change across iterations.
87   // Assumes `while_body` is the body computation of the while loop in question.
88   static std::vector<HloInstruction*> GetInvariantGTEsForWhileBody(
89       const HloComputation& while_body);
90 
91   // Returns a map of index to GetTupleElement instructions in
92   // `while_conditional` that access elements in the parameter tuple. Assumes
93   // `while_conditional` is the conditional computation of the while loop in
94   // question.
95   static absl::flat_hash_map<int64, absl::InlinedVector<HloInstruction*, 1>>
96   GetGTEsMapForWhileConditional(const HloComputation& while_conditional);
97 };
98 }  // namespace xla
99 
100 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_UTIL_H_
101