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_LLVM_IR_LLVM_LOOP_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LLVM_LOOP_H_ 18 19 #include <memory> 20 #include <string> 21 22 #include "absl/strings/str_cat.h" 23 #include "absl/strings/string_view.h" 24 #include "absl/types/span.h" 25 #include "llvm/IR/BasicBlock.h" 26 #include "llvm/IR/IRBuilder.h" 27 #include "llvm/IR/Value.h" 28 #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" 29 #include "tensorflow/compiler/xla/types.h" 30 #include "tensorflow/compiler/xla/xla_data.pb.h" 31 32 namespace xla { 33 namespace llvm_ir { 34 35 enum class UnrollMode { 36 kDefaultUnroll, 37 kFullyUnroll, 38 kNoUnroll, 39 }; 40 41 // A class for constructing a for-loop in LLVM IR. 42 class ForLoop { 43 public: 44 ForLoop(const ForLoop&) = delete; 45 ForLoop& operator=(const ForLoop&) = delete; 46 47 // Emit a for-loop at the current insert point of the given IRBuilder. 48 // 49 // start_index and end_index are the loop bounds (end_index is not inclusive). 50 // `step` is the increment of the loop index after each iteration. 51 // 52 // The current insert basic block of the builder is the preheader to the loop 53 // (see below for definition of basic block names). All instructions (if any) 54 // at or after the insert point in the insert basic block are moved to a newly 55 // created exit basic block. Instructions before the insert point remain in 56 // the insert BB: 57 // 58 // +--------------+ +----------------+ 59 // | insert BB | | insert BB | 60 // | ... | | (preheader BB) | 61 // | %foo = ... | | ... | 62 // insert point ->| %bar = ... | ===> | %foo = ... | 63 // | ... | +----------------+ 64 // +--------------+ | 65 // V 66 // [[ LOOP BBs ]] 67 // | 68 // V 69 // +--------------+ 70 // | exit BB | 71 // | %bar = ... | 72 // | ... | 73 // +--------------+ 74 // 75 // `prefix` is used to disambiguate variable and basic block names emitted in 76 // LLVM IR. If non-empty, it is prepended to the name of the induction 77 // variable value and each basic block created for the loop. 78 // 79 // `unroll_mode` specifies the desired LLVM unrolling behavior for generated 80 // loop. 81 static std::unique_ptr<ForLoop> EmitForLoop( 82 absl::string_view prefix, llvm::Value* start_index, 83 llvm::Value* end_index, llvm::Value* step, llvm::IRBuilder<>* b, 84 UnrollMode unroll_mode = llvm_ir::UnrollMode::kDefaultUnroll, 85 bool prevent_vectorization = false); 86 87 // The names of the blocks follow LLVM's conventions. Control flow amongst the 88 // blocks for the example C code looks like: 89 // 90 // for (int i = 0; i < n; ++i) { 91 // do_stuff(i); 92 // } 93 // 94 // +--------------+ 95 // | preheader BB | 96 // | i = 0 | 97 // +--------------+ 98 // | 99 // V 100 // +-------------+ 101 // | header BB |<-+ 102 // | if i < n: | | 103 // | goto body | | 104 // | else: | | 105 // | goto exit | | 106 // +-------------+ | 107 // | | | 108 // +--------+ | | 109 // | V | 110 // | +-------------+ | 111 // | | body BB | | 112 // | | dostuff(i) |--+ 113 // | | ++i | 114 // | +-------------+ 115 // | 116 // | +-------------+ 117 // +->| exit BB | 118 // +-------------+ 119 // 120 // Caller-emitted code to execute within the loop should be placed within the 121 // "body" basic block. 122 // 123 // Return pointers to various blocks in the loop. GetPreheaderBasicBlock()124 llvm::BasicBlock* GetPreheaderBasicBlock() const { return preheader_bb_; } GetHeaderBasicBlock()125 llvm::BasicBlock* GetHeaderBasicBlock() const { return header_bb_; } GetBodyBasicBlock()126 llvm::BasicBlock* GetBodyBasicBlock() const { return body_bb_; } GetExitBasicBlock()127 llvm::BasicBlock* GetExitBasicBlock() const { return exit_bb_; } 128 129 // Return the Value representing the induction variable in the body basic 130 // block of the loop. GetIndVarValue()131 llvm::Value* GetIndVarValue() const { return indvar_; } 132 133 private: 134 // Allow ForLoopNest to call this private constructor. 135 friend class ForLoopNest; 136 137 ForLoop(absl::string_view prefix, absl::string_view suffix, 138 llvm::Value* start_index, llvm::Value* end_index, llvm::Value* step, 139 UnrollMode unroll_mode, bool prevent_vectorization); 140 141 // Emit the loop at the insert point of the builder. 142 void Emit(llvm::IRBuilder<>* b); 143 144 llvm::BasicBlock* CreateLoopBB(absl::string_view name, llvm::IRBuilder<>* b); 145 146 // Creates a name for an LLVM construct, appending prefix_ and suffix_, if 147 // they are set. 148 std::string GetQualifiedName(absl::string_view name); 149 150 // Return a list of metadata nodes that should be associated with the 151 // llvm::Loop for this `ForLoop`. 152 std::vector<llvm::Metadata*> GetLoopMetadata(llvm::IRBuilder<>* b); 153 154 std::string prefix_; 155 std::string suffix_; 156 llvm::Value* start_index_; 157 llvm::Value* end_index_; 158 llvm::Value* step_; 159 160 // To improve readability of the IR, we want the basic blocks to appear 161 // consecutively in the following order: preheader, header, body, loop, 162 // exit. The member insert_before_bb_ points to where the next basic block 163 // should be created to ensure this ordering. 164 llvm::BasicBlock* insert_before_bb_; 165 166 llvm::BasicBlock* preheader_bb_; 167 llvm::BasicBlock* header_bb_; 168 llvm::BasicBlock* body_bb_; 169 llvm::BasicBlock* exit_bb_; 170 llvm::Value* indvar_; 171 UnrollMode unroll_mode_; 172 bool prevent_vectorization_; 173 }; 174 175 // A simple class for constructing nested for-loops. 176 class ForLoopNest { 177 public: 178 ForLoopNest(absl::string_view name, llvm::IRBuilder<>* b, 179 llvm::Type* index_ty = nullptr) name_(name)180 : name_(name), 181 outer_loop_preheader_bb_(nullptr), 182 outer_loop_exit_bb_(nullptr), 183 inner_loop_body_bb_(nullptr), 184 b_(b) { 185 SetIndexType(index_ty); 186 } 187 ForLoopNest(const ForLoopNest&) = delete; 188 ForLoopNest& operator=(const ForLoopNest&) = delete; 189 190 // Adds a loop to the nest. If no loop has been added yet then emit a loop at 191 // the current insert point of the given builder. If one or more loops have 192 // been added then emit loop inside the body of the last added loop. 193 // unroll_mode is used to emit metadata that controls LLVM unrolling. 194 std::unique_ptr<ForLoop> AddLoop( 195 absl::string_view suffix, llvm::Value* start_index, 196 llvm::Value* end_index, llvm::Value* stride, 197 UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll, 198 bool prevent_vectorization = false); 199 200 // Like the above, except that it defaults to a stride of one. 201 std::unique_ptr<ForLoop> AddLoop( 202 absl::string_view suffix, llvm::Value* start_index, 203 llvm::Value* end_index, 204 UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll, 205 bool prevent_vectorization = false); 206 207 // A convenient wrapper of the other flavor of AddLoop. The given start and 208 // end index are constant. 209 std::unique_ptr<ForLoop> AddLoop( 210 int64_t start_index, int64_t end_index, int64_t stride, 211 absl::string_view suffix, 212 UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll, 213 bool prevent_vectorization = false); 214 215 // Like the above, except that it defaults to a stride of one. 216 std::unique_ptr<ForLoop> AddLoop( 217 int64_t start_index, int64_t end_index, absl::string_view suffix, 218 UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll, 219 bool prevent_vectorization = false); 220 221 // Add loops to iterate through the indices within the specified 222 // shape. The returned index collects the induction variables of the 223 // loops so that it will iterate through all coordinates within the 224 // specified shape. 225 // 226 // E.g. if you pass in a 2x3 shape, you will get back an index with 227 // two entries that are induction variables of the two loops that 228 // will be added. That index will iterate through the 6 coordinates 229 // within the shape. One possible order for that sequence would be: 230 // 231 // (0,0), (0,1), (0,2), (1,0), (1,1), (1,2) 232 IrArray::Index AddLoopsForShape(const Shape& shape, absl::string_view suffix); 233 234 // Add a loop for each dimension in "dimensions". "suffix" is the 235 // name suffix of the indvar and basic blocks in this new loop nest. 236 // 237 // The return value is an index with the induction variables. The 238 // size equals the rank of shape and there is a null for each 239 // dimension that is not in "dimensions". 240 std::vector<llvm::Value*> AddLoopsForShapeOnDimensions( 241 const Shape& shape, absl::Span<const int64_t> dimensions, 242 absl::string_view suffix); 243 244 // Emits a series of nested loops for iterating over an operand array. Loops 245 // are constructed in major to minor dimension layout order. No loop is 246 // emitted for the given 'dimension_to_skip'. The function returns an IrArray 247 // index for the given operand_array containing the indvars of the loops. All 248 // dimensions of the index are filled except for 'dimension_to_skip'. 249 // name_suffix is the string to append to the names of LLVM constructs (eg, 250 // basic blocks) constructed by this method. 251 std::vector<llvm::Value*> EmitOperandArrayLoopNest( 252 const llvm_ir::IrArray& operand_array, int64_t dimension_to_skip, 253 absl::string_view name_suffix); 254 255 // Convenience methods which return particular basic blocks of the outermost 256 // or innermost loops. These methods return nullptr if no loops have been 257 // added yet. GetOuterLoopPreheaderBasicBlock()258 llvm::BasicBlock* GetOuterLoopPreheaderBasicBlock() { 259 return outer_loop_preheader_bb_; 260 } GetOuterLoopExitBasicBlock()261 llvm::BasicBlock* GetOuterLoopExitBasicBlock() { return outer_loop_exit_bb_; } GetInnerLoopBodyBasicBlock()262 llvm::BasicBlock* GetInnerLoopBodyBasicBlock() { return inner_loop_body_bb_; } 263 264 private: SetIndexType(llvm::Type * index_ty)265 void SetIndexType(llvm::Type* index_ty) { 266 index_type_ = index_ty == nullptr ? b_->getInt64Ty() : index_ty; 267 } 268 GetConstantWithIndexType(int64_t c)269 llvm::Constant* GetConstantWithIndexType(int64_t c) const { 270 return llvm::ConstantInt::get(index_type_, c); 271 } 272 273 // Human-friendly name of the loop nest. 274 std::string name_; 275 276 // The preheader and exit basic block of the outermost loop, or nullptr if no 277 // loop has been added yet. 278 llvm::BasicBlock* outer_loop_preheader_bb_; 279 llvm::BasicBlock* outer_loop_exit_bb_; 280 281 // The body basic block of the most-recently added loop, or nullptr if no loop 282 // has been added yet. 283 llvm::BasicBlock* inner_loop_body_bb_; 284 285 llvm::IRBuilder<>* b_; 286 287 llvm::Type* index_type_; 288 }; 289 290 } // namespace llvm_ir 291 } // namespace xla 292 293 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LLVM_LOOP_H_ 294