Emits a loop for every element in the given shape.
| 37 | |
| 38 | // Emits a loop for every element in the given shape. |
| 39 | class LoopEmitter { |
| 40 | public: |
| 41 | using BodyEmitter = std::function<Status(const IrArray::Index& index)>; |
| 42 | |
| 43 | LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape, |
| 44 | llvm::IRBuilder<>* b); |
| 45 | // Constructs a LoopEmitter from an element generator that generates each |
| 46 | // element of the given target array. |
| 47 | LoopEmitter(const ElementGenerator& target_element_generator, |
| 48 | const IrArray& target_array, llvm::IRBuilder<>* b); |
| 49 | |
| 50 | // Constructs a LoopEmitter that emits one element into each of N separate |
| 51 | // arrays on each iteration of the loop. |
| 52 | // |
| 53 | // This is used for multi-output fusion. target_element_generator must |
| 54 | // produce an LLVM struct with N elements. |
| 55 | LoopEmitter(const ElementGenerator& target_element_generator, |
| 56 | absl::Span<const IrArray> target_arrays, llvm::IRBuilder<>* b); |
| 57 | |
| 58 | LoopEmitter(const LoopEmitter&) = delete; |
| 59 | LoopEmitter& operator=(const LoopEmitter&) = delete; |
| 60 | virtual ~LoopEmitter() = default; |
| 61 | |
| 62 | // Emits a loop nest (with a yet-to-be-filled loop body) that iterates through |
| 63 | // every element in the given shape. Returns the multi-dimensional index that |
| 64 | // specifies the element, will return multiple indices if the loop is |
| 65 | // unrolled. |
| 66 | std::vector<IrArray::Index> EmitIndexAndSetExitBasicBlock() { |
| 67 | return EmitIndexAndSetExitBasicBlock(/*loop_name=*/"", b_->getInt64Ty(), |
| 68 | /*base_index*/nullptr); |
| 69 | } |
| 70 | |
| 71 | virtual std::vector<IrArray::Index> EmitIndexAndSetExitBasicBlock( |
| 72 | absl::string_view loop_name, llvm::Type* index_type, |
| 73 | llvm::Value* base_index); |
| 74 | |
| 75 | // Emits a complete loop nest for every element in the given shape. |
| 76 | Status EmitLoop(absl::string_view loop_name = "", |
| 77 | llvm::Type* index_type = nullptr); |
| 78 | |
| 79 | protected: |
| 80 | // An IR emitter that generates the loop body. |
| 81 | BodyEmitter body_emitter_; |
| 82 | |
| 83 | // The shape that the emitted loop iterates through. |
| 84 | Shape shape_; |
| 85 | |
| 86 | // Points to the exit block of the emitted loop. If the given shape is |
| 87 | // scalar, no loops are emitted and exit_bb_ is nullptr in that case. |
| 88 | llvm::BasicBlock* exit_bb_; |
| 89 | |
| 90 | llvm::IRBuilder<>* b_; |
| 91 | }; |
| 92 | |
| 93 | } // namespace llvm_ir |
| 94 | } // namespace xla |
no outgoing calls
no test coverage detected