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_XLA_SERVICE_CONDITIONAL_CODE_MOTION_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_CONDITIONAL_CODE_MOTION_H_ 18 19 #include "absl/strings/string_view.h" 20 #include "tensorflow/compiler/xla/service/hlo_module.h" 21 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" 22 #include "tensorflow/compiler/xla/statusor.h" 23 24 namespace xla { 25 26 namespace conditional_opt { 27 // At the conceptual level, a boundary can be thought of as representing a 28 // single virtual operation, except this virtual operation is conditionally 29 // instantiated into different concrete operations at each conditional branch. 30 // So a boundary is mapped to a single concrete operation if it is outside of 31 // conditional branches, and is mapped to a list of instructions if inside the 32 // branches. This data structure therefore allows a common data structure 33 // representation of the instructions to be moved, whether they are inside or 34 // outside of the branches. Subsequently, it allows a common implementation 35 // basis to be used for both moving instructions out of and for moving them 36 // inside branches. 37 class Boundary { 38 public: 39 enum class Position { kInsideBranch, kOutsideBranch, kUndefined }; Boundary()40 Boundary() : position_(Position::kUndefined) {} Boundary(Position p)41 explicit Boundary(Position p) : position_(p) {} mutable_operands()42 std::vector<HloInstruction*>& mutable_operands() { return operands_; } operands()43 const std::vector<HloInstruction*>& operands() const { return operands_; } IsInsideBranch()44 bool IsInsideBranch() const { return position_ == Position::kInsideBranch; } IsOutsideBranch()45 bool IsOutsideBranch() const { return position_ == Position::kOutsideBranch; } GetPosition()46 Position GetPosition() const { return position_; } IsEmpty()47 bool IsEmpty() const { return operands_.empty(); } ToString()48 std::string ToString() const { 49 std::string res; 50 for (HloInstruction* op : operands_) { 51 res += op->ToString() + ";"; 52 } 53 return res; 54 } 55 bool operator==(const Boundary& that) { 56 return ContainersEqual(operands_, that.operands_); 57 } 58 59 private: 60 // Boundary instructions in the conditional branches, one from each branch 61 // of the conditional; or a single operand from outside the conditional. 62 std::vector<HloInstruction*> operands_; 63 Position position_; 64 }; 65 66 // HLO pass that moves identical ops in/out of conditional. 67 // - The definition of identical are the shape of the operands are identical 68 // and their properties are identical. 69 // - Only the identical ops that won't share operands with other ops will 70 // be moved out of conditional. 71 // The cost model of the code motion optimization includes two components: 72 // represented by the move_config_ and reuse_config_ arrays of the optimization. 73 // The move_config_ array uses 1 vs 0 to dictate whether each Hlo Opcode, when 74 // used with its first operand being another given Hlo Opcode, is allowed to 75 // move across any conditional boundary; the reuse_config_ array uses an integer 76 // to represent the force between each pair of HloOpcode regarding how 77 // attractive it is to place these instructions together (both inside or outside 78 // of a conditional). Both arrays use Hlo Opcode only to drive the 79 // configuration, regardless of where the operations are located in the 80 // module. 81 class ConditionalCodeMotion : public HloModulePass { 82 public: 83 // If is_layout_sensitive is true, then the hoist process preserves layout 84 // during identical comparison. Otherwise, layout is ignored. 85 // The search configuration is a single integer but is split into four parts: 86 // (sign, n, m, p), where n,m,p each occupy 8 bits and together make the 24 87 // bits at the end of the int32. For the sign part, if search_config is <0, 88 // the reuse_config_ cost model is modified (tuned); if search_config is >0, 89 // the move_config_ cost model is modified (tuned); if search_config == 0, 90 // the default cost model is used with no tuning. When tuning, the entries in 91 // the designated configuration array (move_config_ or reuse_config_) are 92 // flipped between 0 and another default integer, starting from the pth entry 93 // being queried by the optimization and repeated every nth time a new entry 94 // is visited, until a maximal of m entries have been changed. The tuning 95 // start over when optimizing a new model. 96 explicit ConditionalCodeMotion(bool is_layout_sensitive, 97 bool pursue_full_conditional_code_motion, 98 int32 search_config = 0) is_layout_sensitive_(is_layout_sensitive)99 : is_layout_sensitive_(is_layout_sensitive), 100 pursue_full_conditional_code_motion_( 101 /*turn off special case if tuning*/ 102 pursue_full_conditional_code_motion && search_config == 0), 103 search_config_(search_config) {} 104 name()105 absl::string_view name() const override { return "conditional-code-motion"; } 106 StatusOr<bool> Run(HloModule* module) override; 107 108 // Optimization decision for each boundary of the conditional instruction. 109 class Decision { 110 public: 111 enum class Direction : uint8 { 112 kMoveOutOfBranch, 113 kMoveIntoBranch, 114 kNoChange 115 }; 116 117 public: Decision(Direction direction,int benefit)118 Decision(Direction direction, int benefit) 119 : direction_(direction), benefit_(benefit) {} GetDirection()120 Direction GetDirection() const { return direction_; } GetBenefit()121 int GetBenefit() const { return benefit_; } 122 123 private: 124 Direction direction_; 125 int benefit_; 126 }; 127 // If the optimization decision is NO_CHANGE, new_boundary is set to nullptr; 128 // otherwise, it is set to the new boundary after proposed optimization. 129 virtual Decision ConsiderCodeMotion( 130 HloInstruction* conditional, const Boundary& cur_boundary, 131 std::vector<Boundary>& to_move, std::vector<Boundary>& new_boundaries, 132 absl::flat_hash_map<HloInstruction*, int>& visited_count); 133 134 private: 135 const bool is_layout_sensitive_; 136 const bool pursue_full_conditional_code_motion_; 137 int32 search_config_; 138 std::vector<std::vector<int64>> move_config_, reuse_config_; 139 140 StatusOr<bool> MoveInstructionOut(HloInstruction* conditional, 141 std::vector<Boundary>& to_move_out, 142 std::vector<Boundary>& new_boundaries); 143 StatusOr<bool> MoveInstructionIn(HloInstruction* conditional, 144 std::vector<Boundary>& to_move_in, 145 std::vector<Boundary>& new_boundaries); 146 void SetDefaultMoveConfig(); 147 }; 148 } // namespace conditional_opt 149 150 } // namespace xla 151 152 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CONDITIONAL_CODE_MOTION_H_ 153