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_LOOP_EMITTER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LOOP_EMITTER_H_ 18 19 #include <functional> 20 21 #include "llvm/IR/BasicBlock.h" 22 #include "llvm/IR/IRBuilder.h" 23 #include "llvm/IR/Value.h" 24 #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" 25 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h" 26 #include "tensorflow/compiler/xla/statusor.h" 27 28 namespace xla { 29 namespace llvm_ir { 30 31 // A function type for emitting code that generates an element in the target 32 // array. The function gets a multi-dimensional index as its only input. This 33 // index specifies the target element for which a value needs to be computed. 34 // The function has to emit code to compute this value and return the resulting 35 // llvm::Value*. 36 using ElementGenerator = 37 std::function<StatusOr<llvm::Value*>(const IrArray::Index& index)>; 38 using BodyEmitter = std::function<Status(const IrArray::Index& index)>; 39 40 // Creates the body emitter from target arrays. 41 BodyEmitter MakeBodyEmitter(const ElementGenerator& target_element_generator, 42 absl::Span<IrArray const> target_arrays, 43 llvm::IRBuilder<>* b, bool is_tuple); 44 45 // Emits a loop for every element in the given shape. 46 class LoopEmitter { 47 public: 48 LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape, 49 llvm::IRBuilder<>* b); 50 51 // Constructs a LoopEmitter from an body_emitter that generates 52 // element of the given target array in the dynamic dimension. 53 LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape, 54 std::vector<llvm::Value*> dynamic_dims, llvm::IRBuilder<>* b); 55 56 // Constructs a LoopEmitter from an element generator that generates each 57 // element of the given target array. 58 LoopEmitter(const ElementGenerator& target_element_generator, 59 const IrArray& target_array, llvm::IRBuilder<>* b); 60 61 // Constructs a LoopEmitter that emits one element into each of N separate 62 // arrays on each iteration of the loop. 63 // 64 // This is used for multi-output fusion. target_element_generator must 65 // produce an LLVM struct with N elements. 66 LoopEmitter(const ElementGenerator& target_element_generator, 67 absl::Span<const IrArray> target_arrays, llvm::IRBuilder<>* b); 68 69 LoopEmitter(const LoopEmitter&) = delete; 70 LoopEmitter& operator=(const LoopEmitter&) = delete; 71 virtual ~LoopEmitter() = default; 72 73 // Emits a loop nest (with a yet-to-be-filled loop body) that iterates through 74 // every element in the given shape. Returns the multi-dimensional index that 75 // specifies the element, will return multiple indices if the loop is 76 // unrolled. 77 virtual std::vector<IrArray::Index> EmitIndexAndSetExitBasicBlock( 78 absl::string_view loop_name, llvm::Type* index_type, 79 llvm::Value* base_index); 80 81 // Emits a complete loop nest for every element in the given shape. 82 Status EmitLoop(absl::string_view loop_name = "", 83 llvm::Type* index_type = nullptr); 84 85 protected: 86 // An IR emitter that generates the loop body. 87 BodyEmitter body_emitter_; 88 89 // The shape that the emitted loop iterates through. 90 Shape shape_; 91 92 // Dynamic dimensions that emitted loop iterates through. Generate the 93 // loop based on the dynamic dimensions if this vector is not empty. 94 std::vector<llvm::Value*> dynamic_dims_; 95 96 // Points to the exit block of the emitted loop. If the given shape is 97 // scalar, no loops are emitted and exit_bb_ is nullptr in that case. 98 llvm::BasicBlock* exit_bb_; 99 100 llvm::IRBuilder<>* b_; 101 102 private: 103 IrArray::Index EmitStaticIndex(ForLoopNest* loop_nest, 104 llvm::Type* index_type); 105 IrArray::Index EmitDynamicIndex(ForLoopNest* loop_nest, 106 llvm::Type* index_type); 107 }; 108 109 } // namespace llvm_ir 110 } // namespace xla 111 112 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LOOP_EMITTER_H_ 113